1use crate::pool::{Pool, matvec_rows, matvec_rows2};
15use cortiq_core::quant::{GROUP_SIZE, Q1_TILE, Q4_TILE, f16_to_f32};
16use cortiq_core::{CmfModel, TensorDtype};
17use std::sync::Arc;
18
19pub enum QTensor {
20 F32 {
21 data: Vec<f32>,
22 rows: usize,
23 cols: usize,
24 },
25 Mapped {
26 model: Arc<CmfModel>,
27 idx: usize,
29 dtype: TensorDtype,
30 rows: usize,
31 cols: usize,
32 row_scale: Vec<f32>,
34 col_field: Vec<f32>,
36 vbit_offsets: Vec<usize>,
40 repack: Vec<u8>,
48 },
49}
50
51fn repack_enabled() -> bool {
58 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
59 *ON.get_or_init(|| {
60 std::env::var("CMF_REPACK")
61 .map(|v| v == "1")
62 .unwrap_or(cfg!(target_os = "android"))
63 })
64}
65
66fn q8_repack(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
70 #[cfg(target_arch = "aarch64")]
71 let arch_ok = sdot_enabled();
72 #[cfg(not(target_arch = "aarch64"))]
73 let arch_ok = false;
74 if !arch_ok || !repack_enabled() || rows < 256 || cols % 16 != 0 {
75 return Vec::new();
76 }
77 q8_repack_layout(bytes, rows, cols)
78}
79
80fn q8_repack_layout(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
83 let groups = rows / 4;
84 let mut rep = vec![0u8; groups * 4 * cols];
85 for g in 0..groups {
86 let dst = &mut rep[g * 4 * cols..(g + 1) * 4 * cols];
87 for c in 0..cols / 16 {
88 for lane in 0..4 {
89 let src = (g * 4 + lane) * cols + c * 16;
90 dst[c * 64 + lane * 16..c * 64 + lane * 16 + 16]
91 .copy_from_slice(&bytes[src..src + 16]);
92 }
93 }
94 }
95 rep
96}
97
98fn vbit_row_offsets(bytes: &[u8], rows: usize, cols: usize) -> Vec<usize> {
101 let ng = cols / GROUP_SIZE;
102 let bits = &bytes[..rows];
103 let mut offsets = Vec::with_capacity(rows + 1);
104 let mut off = rows + rows * ng * 2;
105 for r in 0..rows {
106 offsets.push(off);
107 off += (cols * bits[r] as usize).div_ceil(8);
108 }
109 offsets.push(off);
110 offsets
111}
112
113impl QTensor {
114 pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
115 debug_assert_eq!(data.len(), rows * cols);
116 Self::F32 { data, rows, cols }
117 }
118
119 pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
122 let idx = model
125 .tensor_index(name)
126 .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
127 let entry = &model.tensors[idx];
128 if entry.shape.len() != 2 {
129 return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
130 }
131 let (rows, cols) = (entry.shape[0], entry.shape[1]);
132 let bytes = model.entry_bytes(entry);
133
134 match entry.dtype {
135 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
136 let n = rows * cols;
137 let scales_off = n;
138 let row_scale: Vec<f32> = (0..rows)
139 .map(|o| {
140 f16_to_f32(u16::from_le_bytes([
141 bytes[scales_off + o * 2],
142 bytes[scales_off + o * 2 + 1],
143 ]))
144 })
145 .collect();
146 let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
147 let col_off = n + rows * 2;
148 (0..cols)
149 .map(|i| {
150 f16_to_f32(u16::from_le_bytes([
151 bytes[col_off + i * 2],
152 bytes[col_off + i * 2 + 1],
153 ]))
154 })
155 .collect()
156 } else {
157 Vec::new()
158 };
159 Ok(Self::Mapped {
160 model: model.clone(),
161 idx,
162 dtype: entry.dtype,
163 rows,
164 cols,
165 row_scale,
166 col_field,
167 vbit_offsets: Vec::new(),
168 repack: q8_repack(bytes, rows, cols),
169 })
170 }
171 TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
173 model: model.clone(),
174 idx,
175 dtype: entry.dtype,
176 rows,
177 cols,
178 row_scale: Vec::new(),
179 col_field: Vec::new(),
180 vbit_offsets: vbit_row_offsets(bytes, rows, cols),
181 repack: Vec::new(),
182 }),
183 TensorDtype::VbitRo if cols % GROUP_SIZE == 0 => {
187 let (_, off_off, packed_off) = cortiq_core::quant::vbit_ro_sections(rows, cols);
188 let offsets: Vec<usize> = (0..=rows)
189 .map(|r| packed_off + cortiq_core::quant::vbit_ro_offset(bytes, off_off, r))
190 .collect();
191 Ok(Self::Mapped {
192 model: model.clone(),
193 idx,
194 dtype: entry.dtype,
195 rows,
196 cols,
197 row_scale: Vec::new(),
198 col_field: Vec::new(),
199 vbit_offsets: offsets,
200 repack: Vec::new(),
201 })
202 }
203 TensorDtype::Q4Tiled if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
209 model: model.clone(),
210 idx,
211 dtype: entry.dtype,
212 rows,
213 cols,
214 row_scale: Vec::new(),
215 col_field: Vec::new(),
216 vbit_offsets: Vec::new(),
217 repack: Vec::new(),
218 }),
219 TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
220 model: model.clone(),
221 idx,
222 dtype: entry.dtype,
223 rows,
224 cols,
225 row_scale: Vec::new(),
226 col_field: Vec::new(),
227 vbit_offsets: Vec::new(),
228 repack: Vec::new(),
229 }),
230 TensorDtype::Q1 if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
232 model: model.clone(),
233 idx,
234 dtype: entry.dtype,
235 rows,
236 cols,
237 row_scale: Vec::new(),
238 col_field: Vec::new(),
239 vbit_offsets: Vec::new(),
240 repack: Vec::new(),
241 }),
242 TensorDtype::Q1T if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
246 model: model.clone(),
247 idx,
248 dtype: entry.dtype,
249 rows,
250 cols,
251 row_scale: Vec::new(),
252 col_field: Vec::new(),
253 vbit_offsets: Vec::new(),
254 repack: Vec::new(),
255 }),
256 _ => {
258 let mut data = vec![0.0f32; rows * cols];
259 cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
260 Ok(Self::from_f32(data, rows, cols))
261 }
262 }
263 }
264
265 pub(crate) fn is_q1(&self) -> bool {
268 matches!(
269 self,
270 Self::Mapped {
271 dtype: TensorDtype::Q1,
272 ..
273 }
274 )
275 }
276
277 pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
280 match self {
281 Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
282 _ => None,
283 }
284 }
285
286 pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
293 match self {
294 #[cfg(target_os = "macos")]
295 Self::Mapped {
296 dtype: TensorDtype::Q1T,
297 ..
298 } if !crate::gpu::metal_q1t_enabled() => None,
299 Self::Mapped {
300 idx,
301 dtype:
302 TensorDtype::Q1
303 | TensorDtype::Q1T
304 | TensorDtype::Q4Block
305 | TensorDtype::Q4Tiled
306 | TensorDtype::Q8Row
307 | TensorDtype::Q8_2f,
308 rows,
309 cols,
310 ..
311 } => Some((*idx, *rows, *cols)),
312 _ => None,
313 }
314 }
315
316 pub(crate) fn q4t_parts(&self) -> Option<(usize, usize, usize)> {
322 match self {
323 Self::Mapped {
324 idx,
325 dtype: TensorDtype::Q4Tiled,
326 rows,
327 cols,
328 ..
329 } => Some((*idx, *rows, *cols)),
330 _ => None,
331 }
332 }
333
334 pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
339 match self {
340 Self::Mapped {
341 idx,
342 dtype: TensorDtype::Q8Row,
343 rows,
344 cols,
345 row_scale,
346 col_field,
347 ..
348 } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
349 _ => None,
350 }
351 }
352
353 pub fn rows(&self) -> usize {
354 match self {
355 Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
356 }
357 }
358
359 pub(crate) fn mapped_q4t(&self) -> Option<(&Arc<CmfModel>, usize)> {
362 match self {
363 Self::Mapped {
364 model,
365 idx,
366 dtype: TensorDtype::Q4Tiled,
367 ..
368 } => Some((model, *idx)),
369 _ => None,
370 }
371 }
372
373 pub fn cols(&self) -> usize {
374 match self {
375 Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
376 }
377 }
378
379 pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
382 match self {
383 Self::Mapped {
384 model,
385 idx,
386 dtype: TensorDtype::Q1,
387 ..
388 } => Some((model, *idx)),
389 _ => None,
390 }
391 }
392
393 pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
397 match self {
398 Self::Mapped {
399 model,
400 idx,
401 dtype: TensorDtype::Q8Row,
402 row_scale,
403 ..
404 } => Some((model, *idx, 0, row_scale.as_slice())),
405 Self::Mapped {
406 model,
407 idx,
408 dtype: TensorDtype::Q1,
409 ..
410 } => Some((model, *idx, 1, &[])),
411 Self::Mapped {
416 model,
417 idx,
418 dtype: TensorDtype::Q4Tiled,
419 ..
420 } => Some((model, *idx, 5, &[])),
421 Self::Mapped {
422 model,
423 idx,
424 dtype: TensorDtype::Q4Block,
425 ..
426 } => Some((model, *idx, 2, &[])),
427 Self::Mapped {
428 model,
429 idx,
430 dtype: TensorDtype::Q1T,
431 ..
432 } => Some((model, *idx, 3, &[])),
433 _ => None,
434 }
435 }
436
437 pub fn as_f32(&self) -> Option<&[f32]> {
440 match self {
441 Self::F32 { data, .. } => Some(data),
442 Self::Mapped { .. } => None,
443 }
444 }
445
446 fn quant_bytes(&self) -> &[u8] {
447 match self {
448 Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
449 Self::F32 { .. } => unreachable!("quant_bytes on F32"),
450 }
451 }
452
453 pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
455 let cols = self.cols();
456 debug_assert_eq!(dst.len(), cols);
457 match self {
458 Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
459 Self::Mapped {
460 dtype,
461 row_scale,
462 col_field,
463 vbit_offsets,
464 ..
465 } => {
466 if *dtype == TensorDtype::Q4Tiled {
467 let bytes = self.quant_bytes();
468 let gpr = cols / GROUP_SIZE;
469 for gi in 0..gpr {
470 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
471 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
472 for (k, &b) in tile[2..].iter().enumerate() {
473 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
474 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
475 }
476 }
477 return;
478 }
479 if *dtype == TensorDtype::Q4Block {
480 let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
481 let gpr = cols / GROUP_SIZE;
482 for gi in 0..gpr {
483 let g = r * gpr + gi;
484 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
485 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
486 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
487 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
488 }
489 }
490 return;
491 }
492 if *dtype == TensorDtype::Q1 {
493 let bytes = self.quant_bytes();
494 let gpr = cols / GROUP_SIZE;
495 for gi in 0..gpr {
496 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
497 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
498 for (j, &b) in tile[2..].iter().enumerate() {
499 for k in 0..8 {
500 dst[gi * GROUP_SIZE + j * 8 + k] =
501 (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
502 }
503 }
504 }
505 return;
506 }
507 if *dtype == TensorDtype::Q1T {
508 let bytes = self.quant_bytes();
509 let gpr = cols / GROUP_SIZE;
510 let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
511 for gi in 0..gpr {
512 let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
513 let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
514 bytes[off],
515 bytes[off + 1],
516 ]));
517 let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
518 for k in 0..GROUP_SIZE {
519 dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
520 {
521 1 => s,
522 2 => -s,
523 _ => 0.0,
524 };
525 }
526 }
527 let rows = self.rows();
529 let entries = base_len + (rows + 1) * 4;
530 if entries <= bytes.len() {
531 let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
532 let r0 = u32::from_le_bytes([
533 ptrs[r * 4],
534 ptrs[r * 4 + 1],
535 ptrs[r * 4 + 2],
536 ptrs[r * 4 + 3],
537 ]) as usize;
538 let r1 = u32::from_le_bytes([
539 ptrs[(r + 1) * 4],
540 ptrs[(r + 1) * 4 + 1],
541 ptrs[(r + 1) * 4 + 2],
542 ptrs[(r + 1) * 4 + 3],
543 ]) as usize;
544 let off = entries + r0 * 4;
545 for i in 0..r1 - r0 {
546 let item = &bytes[off + i * 4..off + i * 4 + 4];
547 let c = u16::from_le_bytes([item[0], item[1]]) as usize;
548 let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
549 item[2], item[3],
550 ]));
551 if c < cols {
552 dst[c] = v;
553 }
554 }
555 }
556 return;
557 }
558 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
559 let bytes = self.quant_bytes();
560 let rows = self.rows();
561 let ng = cols / GROUP_SIZE;
562 let bits = &bytes[..rows];
563 let sc_off = rows;
564 let off = vbit_offsets[r];
567 let b = bits[r] as usize;
568 let l = ((1usize << (b - 1)) - 1) as f32;
569 let data = &bytes[off..];
570 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
571 for (i, d) in dst.iter_mut().enumerate() {
572 while nbits < b {
573 acc = (acc << 8) | data[idx] as u64;
574 idx += 1;
575 nbits += 8;
576 }
577 let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
578 nbits -= b;
579 let so = (r * ng + i / GROUP_SIZE) * 2;
580 let sv = f16_to_f32(u16::from_le_bytes([
581 bytes[sc_off + so],
582 bytes[sc_off + so + 1],
583 ]));
584 *d = (u - l) * sv;
585 }
586 return;
587 }
588 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
589 let s = row_scale[r];
590 match dtype {
591 TensorDtype::Q8Row => {
592 for (d, &b) in dst.iter_mut().zip(q) {
593 *d = (b as i8) as f32 * s;
594 }
595 }
596 TensorDtype::Q8_2f => {
597 for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
598 *d = (b as i8) as f32 * s * col_field[i];
599 }
600 }
601 _ => unreachable!(),
602 }
603 }
604 }
605 }
606
607 pub fn sparse_col_ok(&self) -> bool {
612 match self {
613 Self::F32 { .. } => true,
614 Self::Mapped { dtype, .. } => {
615 matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
616 }
617 }
618 }
619
620 pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
624 let inter = self.cols();
625 let hidden = self.rows();
626 debug_assert_eq!(out.len(), hidden);
627 match self {
628 Self::F32 { data, .. } => {
629 for (k, o) in out.iter_mut().enumerate() {
630 *o += w * data[k * inter + c];
631 }
632 }
633 Self::Mapped {
634 dtype,
635 row_scale,
636 col_field,
637 ..
638 } => {
639 let q = self.quant_bytes();
640 let colf = if *dtype == TensorDtype::Q8_2f {
641 col_field[c]
642 } else {
643 1.0
644 };
645 let wc = w * colf;
646 for (k, o) in out.iter_mut().enumerate() {
647 let b = q[k * inter + c] as i8 as f32;
648 *o += wc * b * row_scale[k];
649 }
650 }
651 }
652 }
653
654 pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
658 let cols = self.cols();
659 match self {
660 Self::F32 { data, .. } => {
661 let row = &data[r * cols..(r + 1) * cols];
662 row.iter().zip(x).map(|(w, v)| w * v).sum()
663 }
664 Self::Mapped {
665 dtype,
666 row_scale,
667 col_field,
668 ..
669 } => match dtype {
670 TensorDtype::Q8Row => {
671 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
672 dot_i8_f32(q, x) * row_scale[r]
673 }
674 TensorDtype::Q8_2f => {
675 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
676 dot_i8_col_f32(q, x, col_field) * row_scale[r]
677 }
678 _ => {
679 self.row_f32(r, scratch);
680 scratch.iter().zip(x).map(|(w, v)| w * v).sum()
681 }
682 },
683 }
684 }
685
686 pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
689 match self {
690 Self::F32 { data, .. } => matvec_rows(pool, data, x, out),
691 Self::Mapped {
692 model,
693 idx,
694 dtype,
695 rows,
696 cols,
697 row_scale,
698 col_field,
699 vbit_offsets,
700 repack,
701 } => {
702 let _ = (model, idx);
703 if *dtype == TensorDtype::Q4Block {
704 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
708 let t0 = std::time::Instant::now();
709 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
710 crate::gpu::ProbeArm::Gpu => {
711 if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
712 crate::gpu::probe_record(
713 crate::gpu::OpClass::Matvec,
714 true,
715 t0.elapsed(),
716 );
717 return;
718 }
719 }
720 crate::gpu::ProbeArm::CpuTimed => {
721 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
722 crate::gpu::probe_record(
723 crate::gpu::OpClass::Matvec,
724 false,
725 t0.elapsed(),
726 );
727 return;
728 }
729 crate::gpu::ProbeArm::Cpu => {}
730 }
731 }
732 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
733 return;
734 }
735 if *dtype == TensorDtype::Q4Tiled {
736 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
737 return;
738 }
739 if *dtype == TensorDtype::Q1 {
740 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
745 let t0 = std::time::Instant::now();
746 let arm = if crate::gpu::q1_force() {
747 crate::gpu::ProbeArm::Gpu
748 } else {
749 crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
750 };
751 match arm {
752 crate::gpu::ProbeArm::Gpu => {
753 if crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out) {
754 crate::gpu::probe_record(
755 crate::gpu::OpClass::Matvec,
756 true,
757 t0.elapsed(),
758 );
759 return;
760 }
761 }
762 crate::gpu::ProbeArm::CpuTimed => {
763 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
764 crate::gpu::probe_record(
765 crate::gpu::OpClass::Matvec,
766 false,
767 t0.elapsed(),
768 );
769 return;
770 }
771 crate::gpu::ProbeArm::Cpu => {}
772 }
773 }
774 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
775 return;
776 }
777 if *dtype == TensorDtype::Q1T {
778 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
782 let t0 = std::time::Instant::now();
783 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
784 crate::gpu::ProbeArm::Gpu => {
785 if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
786 q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
787 crate::gpu::probe_record(
788 crate::gpu::OpClass::Matvec,
789 true,
790 t0.elapsed(),
791 );
792 return;
793 }
794 }
795 crate::gpu::ProbeArm::CpuTimed => {
796 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
797 crate::gpu::probe_record(
798 crate::gpu::OpClass::Matvec,
799 false,
800 t0.elapsed(),
801 );
802 return;
803 }
804 crate::gpu::ProbeArm::Cpu => {}
805 }
806 }
807 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
808 return;
809 }
810 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
811 vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
812 return;
813 }
814 let xs = prescale(x, col_field, *dtype);
815 if *rows >= crate::gpu::min_rows()
820 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
821 && std::env::var("CMF_GPU_LMHEAD")
822 .map(|v| v != "0")
823 .unwrap_or(true)
824 && crate::gpu::enabled_here()
825 {
826 let t0 = std::time::Instant::now();
829 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
830 crate::gpu::ProbeArm::Gpu => {}
831 crate::gpu::ProbeArm::CpuTimed => {
832 qmatvec(
833 self.quant_bytes(),
834 repack,
835 row_scale,
836 x,
837 col_field,
838 *dtype,
839 *rows,
840 *cols,
841 out,
842 pool,
843 );
844 crate::gpu::probe_record(
845 crate::gpu::OpClass::Matvec,
846 false,
847 t0.elapsed(),
848 );
849 return;
850 }
851 crate::gpu::ProbeArm::Cpu => {
852 qmatvec(
853 self.quant_bytes(),
854 repack,
855 row_scale,
856 x,
857 col_field,
858 *dtype,
859 *rows,
860 *cols,
861 out,
862 pool,
863 );
864 return;
865 }
866 }
867 let frac = std::env::var("CMF_GPU_SPLIT")
868 .ok()
869 .and_then(|v| v.parse::<f32>().ok())
870 .unwrap_or(0.5)
871 .clamp(0.0, 1.0);
872 let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
873 let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
874 let bytes = self.quant_bytes();
875 let ok = std::thread::scope(|sc| {
876 let g = sc.spawn(|| {
877 crate::gpu::q8_matvec_range(
878 model,
879 *idx,
880 cpu_rows,
881 &row_scale[cpu_rows..],
882 &xs,
883 *rows - cpu_rows,
884 *cols,
885 out_gpu,
886 )
887 });
888 if cpu_rows > 0 {
889 let rep_cpu = if repack.is_empty() {
892 &[][..]
893 } else {
894 &repack[..(cpu_rows / 4) * 4 * *cols]
895 };
896 qmatvec(
897 &bytes[..cpu_rows * *cols],
898 rep_cpu,
899 &row_scale[..cpu_rows],
900 x,
901 col_field,
902 *dtype,
903 cpu_rows,
904 *cols,
905 out_cpu,
906 pool,
907 );
908 }
909 g.join().unwrap_or(false)
910 });
911 if ok {
912 crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
913 return;
914 }
915 qmatvec(
918 &bytes[cpu_rows * *cols..(*rows) * *cols],
919 &[],
920 &row_scale[cpu_rows..],
921 x,
922 col_field,
923 *dtype,
924 *rows - cpu_rows,
925 *cols,
926 out_gpu,
927 pool,
928 );
929 return;
930 }
931 qmatvec(
932 self.quant_bytes(),
933 repack,
934 row_scale,
935 x,
936 col_field,
937 *dtype,
938 *rows,
939 *cols,
940 out,
941 pool,
942 );
943 }
944 }
945 }
946
947 pub fn matvec2(
949 &self,
950 x1: &[f32],
951 x2: &[f32],
952 o1: &mut [f32],
953 o2: &mut [f32],
954 pool: Option<&Pool>,
955 ) {
956 match self {
957 Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
958 Self::Mapped {
959 dtype,
960 rows,
961 cols,
962 row_scale,
963 col_field,
964 vbit_offsets,
965 ..
966 } => {
967 if *dtype == TensorDtype::Q4Block {
968 q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
969 return;
970 }
971 if *dtype == TensorDtype::Q4Tiled {
972 q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
973 return;
974 }
975 if *dtype == TensorDtype::Q1 {
976 q1_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
977 return;
978 }
979 if *dtype == TensorDtype::Q1T {
980 q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
986 return;
987 }
988 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
989 vbitmatvec2(
990 self.quant_bytes(),
991 vbit_offsets,
992 x1,
993 x2,
994 *rows,
995 *cols,
996 o1,
997 o2,
998 pool,
999 );
1000 return;
1001 }
1002 qmatvec2(
1003 self.quant_bytes(),
1004 row_scale,
1005 x1,
1006 x2,
1007 col_field,
1008 *dtype,
1009 *rows,
1010 *cols,
1011 o1,
1012 o2,
1013 pool,
1014 );
1015 }
1016 }
1017 }
1018}
1019
1020impl QTensor {
1021 pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
1026 let cols = self.cols();
1027 let rows = self.rows();
1028 debug_assert_eq!(xs_all.len(), b * cols);
1029 debug_assert_eq!(out.len(), b * rows);
1030 if crate::gptq_capture::capturing() {
1034 if let Self::Mapped { model, idx, .. } = self {
1035 crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
1036 }
1037 }
1038 match self {
1039 Self::F32 { data, .. } => {
1040 let out_addr = SendMut(out.as_mut_ptr());
1041 let run = |start: usize, end: usize| {
1042 for o in start..end {
1043 let row = &data[o * cols..(o + 1) * cols];
1044 for bi in 0..b {
1045 let x = &xs_all[bi * cols..(bi + 1) * cols];
1046 let mut acc = 0f32;
1047 for j in 0..cols {
1048 acc += row[j] * x[j];
1049 }
1050 unsafe { *out_addr.at(bi * rows + o) = acc };
1051 }
1052 }
1053 };
1054 dispatch_rows(pool, rows, &run);
1055 }
1056 Self::Mapped {
1057 dtype,
1058 row_scale,
1059 col_field,
1060 vbit_offsets,
1061 ..
1062 } => {
1063 if *dtype == TensorDtype::Q4Block {
1064 q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1065 return;
1066 }
1067 if *dtype == TensorDtype::Q4Tiled {
1068 if b >= 32
1080 && b * rows * cols >= 128_000_000
1081 && cols % 32 == 0
1082 && !crate::gpu::mm_killed()
1083 && crate::gpu::enabled_here()
1084 {
1085 let class = if b >= 128 {
1086 crate::gpu::OpClass::MatmatWide
1087 } else {
1088 crate::gpu::OpClass::Matmat
1089 };
1090 if let Self::Mapped { model, idx, .. } = self {
1091 let t0 = std::time::Instant::now();
1092 match crate::gpu::probe_arm(class) {
1093 crate::gpu::ProbeArm::Gpu => {
1094 if crate::gpu::q4t_matmat(
1095 model, *idx, xs_all, b, rows, cols, out,
1096 ) {
1097 let el = t0.elapsed();
1098 let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
1108 let budget = std::time::Duration::from_secs_f64(
1109 flops / 1.5e12 * 8.0 + 0.020,
1110 );
1111 if el > budget && !crate::gpu::probe_was_cold() {
1112 tracing::warn!(
1113 "gpu q4t matmat took {el:?} (budget {budget:?}) — \
1114 device contended, CPU for the rest of the process"
1115 );
1116 crate::gpu::mm_kill();
1117 }
1118 crate::gpu::probe_record(class, true, el);
1119 return;
1120 }
1121 }
1122 crate::gpu::ProbeArm::CpuTimed => {
1123 q4t_matmat(
1124 self.quant_bytes(),
1125 xs_all,
1126 b,
1127 rows,
1128 cols,
1129 out,
1130 pool,
1131 );
1132 crate::gpu::probe_record(class, false, t0.elapsed());
1133 return;
1134 }
1135 crate::gpu::ProbeArm::Cpu => {}
1136 }
1137 }
1138 }
1139 q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1140 return;
1141 }
1142 if *dtype == TensorDtype::Q1 {
1143 if b >= 32
1146 && b * rows * cols >= 128_000_000
1147 && cols % 64 == 0
1148 && crate::gpu::enabled_here()
1149 {
1150 if let Self::Mapped { model, idx, .. } = self {
1151 let t0 = std::time::Instant::now();
1152 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1153 crate::gpu::ProbeArm::Gpu => {
1154 if crate::gpu::q1_matmat(
1155 model, *idx, xs_all, b, rows, cols, out,
1156 ) {
1157 crate::gpu::probe_record(
1158 crate::gpu::OpClass::Matmat,
1159 true,
1160 t0.elapsed(),
1161 );
1162 return;
1163 }
1164 }
1165 crate::gpu::ProbeArm::CpuTimed => {
1166 q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1167 crate::gpu::probe_record(
1168 crate::gpu::OpClass::Matmat,
1169 false,
1170 t0.elapsed(),
1171 );
1172 return;
1173 }
1174 crate::gpu::ProbeArm::Cpu => {}
1175 }
1176 }
1177 }
1178 q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1179 return;
1180 }
1181 if *dtype == TensorDtype::Q1T {
1182 if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1185 if let Self::Mapped { model, idx, .. } = self {
1186 let t0 = std::time::Instant::now();
1187 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1188 crate::gpu::ProbeArm::Gpu => {
1189 if crate::gpu::q1t_matmat(
1190 model, *idx, xs_all, b, rows, cols, out,
1191 ) {
1192 crate::gpu::probe_record(
1193 crate::gpu::OpClass::Matmat,
1194 true,
1195 t0.elapsed(),
1196 );
1197 return;
1198 }
1199 }
1200 crate::gpu::ProbeArm::CpuTimed => {
1201 q1t_matmat(
1202 self.quant_bytes(),
1203 xs_all,
1204 b,
1205 rows,
1206 cols,
1207 out,
1208 pool,
1209 );
1210 crate::gpu::probe_record(
1211 crate::gpu::OpClass::Matmat,
1212 false,
1213 t0.elapsed(),
1214 );
1215 return;
1216 }
1217 crate::gpu::ProbeArm::Cpu => {}
1218 }
1219 }
1220 }
1221 q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1222 return;
1223 }
1224 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1225 vbitmatmat(
1226 self.quant_bytes(),
1227 vbit_offsets,
1228 xs_all,
1229 b,
1230 rows,
1231 cols,
1232 out,
1233 pool,
1234 );
1235 return;
1236 }
1237 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
1238 .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
1239 .collect();
1240 if b >= 8 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1246 if let Self::Mapped { model, idx, .. } = self {
1247 let t0 = std::time::Instant::now();
1248 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1249 crate::gpu::ProbeArm::Gpu
1250 if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
1251 && !crate::gpu::q8_resident_or_upload(model, *idx) =>
1252 {
1253 let q = self.quant_bytes();
1257 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1258 return;
1259 }
1260 crate::gpu::ProbeArm::Gpu => {
1261 let flat: Vec<f32> =
1262 pre.iter().flat_map(|v| v.iter().copied()).collect();
1263 if crate::gpu::q8_matmat(
1264 model, *idx, row_scale, &flat, b, rows, cols, out,
1265 ) {
1266 crate::gpu::probe_record(
1267 crate::gpu::OpClass::Matmat,
1268 true,
1269 t0.elapsed(),
1270 );
1271 return;
1272 }
1273 }
1274 crate::gpu::ProbeArm::CpuTimed => {
1275 let q = self.quant_bytes();
1276 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1277 crate::gpu::probe_record(
1278 crate::gpu::OpClass::Matmat,
1279 false,
1280 t0.elapsed(),
1281 );
1282 return;
1283 }
1284 crate::gpu::ProbeArm::Cpu => {}
1285 }
1286 }
1287 }
1288 let q = self.quant_bytes();
1289 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1290 }
1291 }
1292 }
1293}
1294
1295impl QTensor {
1296 pub fn matvec_many<const N: usize>(
1303 ts: [&QTensor; N],
1304 x: &[f32],
1305 mut outs: [&mut [f32]; N],
1306 pool: Option<&Pool>,
1307 ) {
1308 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
1309 let uniform_q8 = ts.iter().all(|t| {
1310 matches!(
1311 t,
1312 Self::Mapped {
1313 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
1314 ..
1315 }
1316 )
1317 });
1318 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
1319 let uniform_q4 = ts.iter().all(|t| {
1320 matches!(
1321 t,
1322 Self::Mapped {
1323 dtype: TensorDtype::Q4Block,
1324 ..
1325 }
1326 )
1327 });
1328 let uniform_vbit = ts.iter().all(|t| {
1329 matches!(
1330 t,
1331 Self::Mapped {
1332 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
1333 ..
1334 }
1335 )
1336 });
1337 let uniform_q1 = ts.iter().all(|t| {
1338 matches!(
1339 t,
1340 Self::Mapped {
1341 dtype: TensorDtype::Q1,
1342 ..
1343 }
1344 )
1345 });
1346 let uniform_q1t = ts.iter().all(|t| {
1347 matches!(
1348 t,
1349 Self::Mapped {
1350 dtype: TensorDtype::Q1T,
1351 ..
1352 }
1353 )
1354 });
1355 let Some(pool) = pool else {
1356 for (t, o) in ts.iter().zip(outs.iter_mut()) {
1357 t.matvec(x, o, None);
1358 }
1359 return;
1360 };
1361 if total_rows < 256
1362 || !(uniform_q8
1363 || uniform_f32
1364 || uniform_q4
1365 || uniform_vbit
1366 || uniform_q1
1367 || uniform_q1t)
1368 {
1369 for (t, o) in ts.iter().zip(outs.iter_mut()) {
1370 t.matvec(x, o, Some(pool));
1371 }
1372 return;
1373 }
1374
1375 if uniform_q1 {
1376 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1379 if a8w8_enabled() {
1380 let act = split_act(x);
1381 let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
1382 let (act, gsum) = (&act, &gsum);
1383 let closures: [_; N] = std::array::from_fn(|i| {
1384 let (bytes, gpr, out) =
1385 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1386 move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
1387 });
1388 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1389 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1390 pool.run_many(&parts);
1391 } else {
1392 let closures: [_; N] = std::array::from_fn(|i| {
1393 let (bytes, gpr, out) =
1394 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1395 move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
1396 });
1397 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1398 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1399 pool.run_many(&parts);
1400 }
1401 return;
1402 }
1403
1404 if uniform_q1t {
1405 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1409 const TILE: usize = cortiq_core::quant::Q1T_TILE;
1410 if a8w8_enabled() {
1411 let act = split_act(x);
1412 let act = &act;
1413 let x_ref = x;
1414 let closures: [_; N] = std::array::from_fn(|i| {
1415 let bytes = ts[i].quant_bytes();
1416 let (rows, cols) = (ts[i].rows(), ts[i].cols());
1417 let gpr = cols / GROUP_SIZE;
1418 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
1419 let out = outs_addr[i];
1420 move |s: usize, e: usize| {
1421 q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
1422 }
1423 });
1424 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1425 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1426 pool.run_many(&parts);
1427 } else {
1428 let x_ref = x;
1429 let closures: [_; N] = std::array::from_fn(|i| {
1430 let bytes = ts[i].quant_bytes();
1431 let (rows, cols) = (ts[i].rows(), ts[i].cols());
1432 let gpr = cols / GROUP_SIZE;
1433 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
1434 let out = outs_addr[i];
1435 move |s: usize, e: usize| {
1436 q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
1437 }
1438 });
1439 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1440 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1441 pool.run_many(&parts);
1442 }
1443 return;
1444 }
1445
1446 if uniform_q4 || uniform_vbit {
1447 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1448 if a8w8_enabled() {
1450 let act = split_act(x);
1451 let act = &act;
1452 if uniform_q4 {
1453 let closures: [_; N] = std::array::from_fn(|i| {
1454 let (packed, scales) =
1455 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1456 let (gpr, cols, out) =
1457 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
1458 move |s: usize, e: usize| {
1459 q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
1460 }
1461 });
1462 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1463 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1464 pool.run_many(&parts);
1465 } else {
1466 let closures: [_; N] = std::array::from_fn(|i| {
1467 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
1468 unreachable!()
1469 };
1470 let (bytes, rows, cols, out) = (
1471 ts[i].quant_bytes(),
1472 ts[i].rows(),
1473 ts[i].cols(),
1474 outs_addr[i],
1475 );
1476 move |s: usize, e: usize| {
1477 vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
1478 }
1479 });
1480 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1481 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1482 pool.run_many(&parts);
1483 }
1484 return;
1485 }
1486 if uniform_q4 {
1487 let closures: [_; N] = std::array::from_fn(|i| {
1488 let (packed, scales) =
1489 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1490 let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1491 move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
1492 });
1493 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1494 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1495 pool.run_many(&parts);
1496 } else {
1497 let closures: [_; N] = std::array::from_fn(|i| {
1498 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
1499 unreachable!()
1500 };
1501 let (bytes, rows, cols, out) = (
1502 ts[i].quant_bytes(),
1503 ts[i].rows(),
1504 ts[i].cols(),
1505 outs_addr[i],
1506 );
1507 move |s: usize, e: usize| {
1508 vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
1509 }
1510 });
1511 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1512 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1513 pool.run_many(&parts);
1514 }
1515 return;
1516 }
1517
1518 if uniform_f32 {
1519 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1520 let closures: [_; N] = std::array::from_fn(|i| {
1521 let Self::F32 { data, cols, .. } = ts[i] else {
1522 unreachable!()
1523 };
1524 let out = outs_addr[i];
1525 move |start: usize, end: usize| {
1526 for o in start..end {
1527 let row = &data[o * cols..(o + 1) * cols];
1528 let mut sum = 0.0f32;
1529 for j in 0..*cols {
1530 sum += row[j] * x[j];
1531 }
1532 unsafe { *out.at(o) = sum };
1534 }
1535 }
1536 });
1537 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1538 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1539 pool.run_many(&parts);
1540 return;
1541 }
1542
1543 struct Ctx<'a> {
1546 bytes: &'a [u8],
1547 #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
1548 rep: &'a [u8],
1549 row_scale: &'a [f32],
1550 cols: usize,
1551 xs: std::borrow::Cow<'a, [f32]>,
1552 }
1553 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
1554 let Self::Mapped {
1555 dtype,
1556 cols,
1557 row_scale,
1558 col_field,
1559 repack,
1560 ..
1561 } = ts[i]
1562 else {
1563 unreachable!()
1564 };
1565 Ctx {
1566 bytes: ts[i].quant_bytes(),
1567 rep: repack,
1568 row_scale,
1569 cols: *cols,
1570 xs: prescale(x, col_field, *dtype),
1571 }
1572 });
1573 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1574 #[cfg(target_arch = "aarch64")]
1575 if sdot_enabled() {
1576 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
1577 let closures: [_; N] = std::array::from_fn(|i| {
1578 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
1579 move |start: usize, end: usize| {
1580 q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
1581 }
1582 });
1583 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1584 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1585 pool.run_many(&parts);
1586 return;
1587 }
1588 #[cfg(target_arch = "x86_64")]
1589 if avx2_a8w8_enabled() {
1590 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
1591 let closures: [_; N] = std::array::from_fn(|i| {
1592 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
1593 move |start: usize, end: usize| {
1594 q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
1595 }
1596 });
1597 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1598 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1599 pool.run_many(&parts);
1600 return;
1601 }
1602 let closures: [_; N] = std::array::from_fn(|i| {
1603 let (c, out) = (&ctxs[i], outs_addr[i]);
1604 move |start: usize, end: usize| {
1605 q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
1606 }
1607 });
1608 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1609 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1610 pool.run_many(&parts);
1611 }
1612}
1613
1614impl QTensor {
1615 #[allow(clippy::needless_range_loop)]
1620 pub fn matvec2_many<const N: usize>(
1621 ts: [&QTensor; N],
1622 x1: &[f32],
1623 x2: &[f32],
1624 mut o1s: [&mut [f32]; N],
1625 mut o2s: [&mut [f32]; N],
1626 pool: Option<&Pool>,
1627 ) {
1628 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
1629 let uniform_q8 = ts.iter().all(|t| {
1630 matches!(
1631 t,
1632 Self::Mapped {
1633 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
1634 ..
1635 }
1636 )
1637 });
1638 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
1639 let uniform_q4 = ts.iter().all(|t| {
1640 matches!(
1641 t,
1642 Self::Mapped {
1643 dtype: TensorDtype::Q4Block,
1644 ..
1645 }
1646 )
1647 });
1648 let uniform_vbit = ts.iter().all(|t| {
1649 matches!(
1650 t,
1651 Self::Mapped {
1652 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
1653 ..
1654 }
1655 )
1656 });
1657 let fusable = pool.is_some()
1658 && total_rows >= 256
1659 && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
1660 if !fusable {
1661 for i in 0..N {
1662 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
1663 }
1664 return;
1665 }
1666 let pool = pool.unwrap();
1667
1668 if uniform_q4 || uniform_vbit {
1669 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1670 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1671 if a8w8_enabled() {
1673 let a1 = split_act(x1);
1674 let a2 = split_act(x2);
1675 let (a1, a2) = (&a1, &a2);
1676 if uniform_q4 {
1677 let closures: [_; N] = std::array::from_fn(|i| {
1678 let (packed, scales) =
1679 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1680 let (gpr, cols, o1, o2) =
1681 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
1682 move |s: usize, e: usize| {
1683 q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
1684 }
1685 });
1686 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1687 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1688 pool.run_many(&parts);
1689 } else {
1690 let closures: [_; N] = std::array::from_fn(|i| {
1691 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
1692 unreachable!()
1693 };
1694 let (bytes, rows, cols, o1, o2) = (
1695 ts[i].quant_bytes(),
1696 ts[i].rows(),
1697 ts[i].cols(),
1698 p1[i],
1699 p2[i],
1700 );
1701 move |s: usize, e: usize| {
1702 vbit_range2_a8w8(
1703 bytes,
1704 vbit_offsets,
1705 x1,
1706 x2,
1707 a1,
1708 a2,
1709 rows,
1710 cols,
1711 o1,
1712 o2,
1713 s,
1714 e,
1715 )
1716 }
1717 });
1718 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1719 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1720 pool.run_many(&parts);
1721 }
1722 return;
1723 }
1724 if uniform_q4 {
1725 let closures: [_; N] = std::array::from_fn(|i| {
1726 let (packed, scales) =
1727 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1728 let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
1729 move |s: usize, e: usize| {
1730 q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
1731 }
1732 });
1733 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1734 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1735 pool.run_many(&parts);
1736 } else {
1737 let closures: [_; N] = std::array::from_fn(|i| {
1738 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
1739 unreachable!()
1740 };
1741 let (bytes, rows, cols, o1, o2) = (
1742 ts[i].quant_bytes(),
1743 ts[i].rows(),
1744 ts[i].cols(),
1745 p1[i],
1746 p2[i],
1747 );
1748 move |s: usize, e: usize| {
1749 vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
1750 }
1751 });
1752 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1753 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1754 pool.run_many(&parts);
1755 }
1756 return;
1757 }
1758
1759 if uniform_f32 {
1760 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1761 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1762 let closures: [_; N] = std::array::from_fn(|i| {
1763 let Self::F32 { data, cols, .. } = ts[i] else {
1764 unreachable!()
1765 };
1766 let (o1, o2) = (p1[i], p2[i]);
1767 move |start: usize, end: usize| {
1768 for o in start..end {
1769 let row = &data[o * cols..(o + 1) * cols];
1770 let (mut s1, mut s2) = (0.0f32, 0.0f32);
1771 for j in 0..*cols {
1772 s1 += row[j] * x1[j];
1773 s2 += row[j] * x2[j];
1774 }
1775 unsafe {
1777 *o1.at(o) = s1;
1778 *o2.at(o) = s2;
1779 }
1780 }
1781 }
1782 });
1783 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1784 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1785 pool.run_many(&parts);
1786 return;
1787 }
1788
1789 struct Ctx<'a> {
1790 bytes: &'a [u8],
1791 row_scale: &'a [f32],
1792 cols: usize,
1793 xs1: std::borrow::Cow<'a, [f32]>,
1794 xs2: std::borrow::Cow<'a, [f32]>,
1795 }
1796 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
1797 let Self::Mapped {
1798 dtype,
1799 cols,
1800 row_scale,
1801 col_field,
1802 ..
1803 } = ts[i]
1804 else {
1805 unreachable!()
1806 };
1807 Ctx {
1808 bytes: ts[i].quant_bytes(),
1809 row_scale,
1810 cols: *cols,
1811 xs1: prescale(x1, col_field, *dtype),
1812 xs2: prescale(x2, col_field, *dtype),
1813 }
1814 });
1815 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1816 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1817 #[cfg(target_arch = "aarch64")]
1818 if sdot_enabled() {
1819 let acts: [(SplitAct, SplitAct); N] =
1820 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
1821 let closures: [_; N] = std::array::from_fn(|i| {
1822 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
1823 move |start: usize, end: usize| {
1824 q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
1825 }
1826 });
1827 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1828 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1829 pool.run_many(&parts);
1830 return;
1831 }
1832 #[cfg(target_arch = "x86_64")]
1833 if avx2_a8w8_enabled() {
1834 let acts: [(SplitAct, SplitAct); N] =
1835 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
1836 let closures: [_; N] = std::array::from_fn(|i| {
1837 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
1838 move |start: usize, end: usize| {
1839 q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
1840 }
1841 });
1842 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1843 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1844 pool.run_many(&parts);
1845 return;
1846 }
1847 let closures: [_; N] = std::array::from_fn(|i| {
1848 let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
1849 move |start: usize, end: usize| {
1850 q8_range2_f32(
1851 c.bytes,
1852 c.row_scale,
1853 &c.xs1,
1854 &c.xs2,
1855 c.cols,
1856 o1,
1857 o2,
1858 start,
1859 end,
1860 )
1861 }
1862 });
1863 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1864 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1865 pool.run_many(&parts);
1866 }
1867
1868 pub fn matvec_silu_mul(
1873 gate: &QTensor,
1874 up: &QTensor,
1875 x: &[f32],
1876 out: &mut [f32],
1877 pool: Option<&Pool>,
1878 ) -> bool {
1879 let inter = gate.rows();
1880 debug_assert_eq!(up.rows(), inter);
1881 debug_assert_eq!(out.len(), inter);
1882 debug_assert_eq!(gate.cols(), up.cols());
1883 if !a8w8_enabled() {
1884 return false;
1885 }
1886 let act = split_act(x);
1887 let act = &act;
1888 let x_ref = x;
1889 let out_addr = SendMut(out.as_mut_ptr());
1890
1891 match (gate, up) {
1892 (
1894 Self::Mapped {
1895 dtype: TensorDtype::Q4Block,
1896 ..
1897 },
1898 Self::Mapped {
1899 dtype: TensorDtype::Q4Block,
1900 ..
1901 },
1902 ) => {
1903 let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
1904 let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
1905 let gpr = gate.cols() / GROUP_SIZE;
1906 let cols = gate.cols();
1907 let run = move |start: usize, end: usize| {
1908 for r in start..end {
1909 let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
1910 let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
1911 for &(j, xv) in &act.outliers {
1912 let flat = r * cols + j;
1913 let gb = gp[flat / 2];
1914 let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
1915 let gsc = f16_to_f32(u16::from_le_bytes([
1916 gs[(flat / GROUP_SIZE) * 2],
1917 gs[(flat / GROUP_SIZE) * 2 + 1],
1918 ]));
1919 gv += ((gn as i32 - 8) as f32) * gsc * xv;
1920 let ub = up_p[flat / 2];
1921 let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
1922 let usc = f16_to_f32(u16::from_le_bytes([
1923 up_s[(flat / GROUP_SIZE) * 2],
1924 up_s[(flat / GROUP_SIZE) * 2 + 1],
1925 ]));
1926 uv += ((un as i32 - 8) as f32) * usc * xv;
1927 }
1928 let silu_g = gv / (1.0 + (-gv).exp());
1929 unsafe { *out_addr.at(r) = silu_g * uv };
1931 }
1932 };
1933 dispatch_rows(pool, inter, &run);
1934 true
1935 }
1936 (
1940 Self::Mapped {
1941 dtype: TensorDtype::Q4Tiled,
1942 ..
1943 },
1944 Self::Mapped {
1945 dtype: TensorDtype::Q4Tiled,
1946 ..
1947 },
1948 ) => {
1949 let g_bytes = gate.quant_bytes();
1950 let u_bytes = up.quant_bytes();
1951 let gpr = gate.cols() / GROUP_SIZE;
1952 let run = move |start: usize, end: usize| {
1953 for r in start..end {
1954 let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
1955 let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
1956 for &(j, xv) in &act.outliers {
1957 let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
1958 gv += w * s * xv;
1959 let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
1960 uv += w * s * xv;
1961 }
1962 let silu_g = gv / (1.0 + (-gv).exp());
1963 unsafe { *out_addr.at(r) = silu_g * uv };
1965 }
1966 };
1967 dispatch_rows(pool, inter, &run);
1968 true
1969 }
1970 (
1972 Self::Mapped {
1973 dtype: TensorDtype::Q1T,
1974 ..
1975 },
1976 Self::Mapped {
1977 dtype: TensorDtype::Q1T,
1978 ..
1979 },
1980 ) => {
1981 const TILE: usize = cortiq_core::quant::Q1T_TILE;
1982 let g_bytes = gate.quant_bytes();
1983 let u_bytes = up.quant_bytes();
1984 let gpr = gate.cols() / GROUP_SIZE;
1985 let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
1986 let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
1987 let run = move |start: usize, end: usize| {
1988 for r in start..end {
1989 let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
1990 let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
1991 for &(j, xv) in &act.outliers {
1992 gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
1993 uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
1994 }
1995 gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
1996 uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
1997 let silu_g = gv / (1.0 + (-gv).exp());
1998 unsafe { *out_addr.at(r) = silu_g * uv };
2000 }
2001 };
2002 dispatch_rows(pool, inter, &run);
2003 true
2004 }
2005 _ => false,
2006 }
2007 }
2008}
2009
2010#[cfg(target_os = "macos")]
2015mod accel_blas {
2016 #[link(name = "Accelerate", kind = "framework")]
2017 unsafe extern "C" {
2018 pub fn cblas_sgemm(
2019 order: i32,
2020 trans_a: i32,
2021 trans_b: i32,
2022 m: i32,
2023 n: i32,
2024 k: i32,
2025 alpha: f32,
2026 a: *const f32,
2027 lda: i32,
2028 b: *const f32,
2029 ldb: i32,
2030 beta: f32,
2031 c: *mut f32,
2032 ldc: i32,
2033 );
2034 }
2035}
2036
2037#[cfg(target_os = "macos")]
2038pub(crate) fn accel_gemm_enabled() -> bool {
2039 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2040 *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
2041}
2042
2043#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
2046pub(crate) fn accel_gemm_enabled() -> bool {
2047 true
2048}
2049
2050#[cfg(target_arch = "aarch64")]
2057#[allow(clippy::too_many_arguments)]
2058pub(crate) fn neon_gemm_rm(
2059 m: usize,
2060 n: usize,
2061 k: usize,
2062 alpha: f32,
2063 a: &[f32],
2064 lda: usize,
2065 b_mat: &[f32],
2066 ldb: usize,
2067 b_rows_are_n: bool,
2068 c: &mut [f32],
2069 ldc: usize,
2070) {
2071 debug_assert!(a.len() >= (m - 1) * lda + k);
2072 debug_assert!(c.len() >= (m - 1) * ldc + n);
2073 unsafe {
2075 use core::arch::aarch64::*;
2076 let mut i = 0usize;
2077 while i < m {
2078 let mi = (m - i).min(4);
2079 let mut j = 0usize;
2080 while j < n {
2081 let nj = (n - j).min(8);
2082 if mi == 4 && nj == 8 {
2083 let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2084 let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2085 let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2086 let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2087 for p in 0..k {
2088 let (b0, b1) = if b_rows_are_n {
2089 let base = b_mat.as_ptr().add(j * ldb + p);
2092 let g = |o: usize| *base.add(o * ldb);
2093 ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
2094 } else {
2095 let base = b_mat.as_ptr().add(p * ldb + j);
2096 (
2097 [*base, *base.add(1), *base.add(2), *base.add(3)],
2098 [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
2099 )
2100 };
2101 let bv0 = vld1q_f32(b0.as_ptr());
2102 let bv1 = vld1q_f32(b1.as_ptr());
2103 let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
2104 let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
2105 let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
2106 let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
2107 c0a = vfmaq_f32(c0a, a0, bv0);
2108 c0b = vfmaq_f32(c0b, a0, bv1);
2109 c1a = vfmaq_f32(c1a, a1, bv0);
2110 c1b = vfmaq_f32(c1b, a1, bv1);
2111 c2a = vfmaq_f32(c2a, a2, bv0);
2112 c2b = vfmaq_f32(c2b, a2, bv1);
2113 c3a = vfmaq_f32(c3a, a3, bv0);
2114 c3b = vfmaq_f32(c3b, a3, bv1);
2115 }
2116 let al = vdupq_n_f32(alpha);
2117 for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
2118 .iter()
2119 .enumerate()
2120 {
2121 let dst = c.as_mut_ptr().add((i + r) * ldc + j);
2122 vst1q_f32(dst, vmulq_f32(*ca, al));
2123 vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
2124 }
2125 } else {
2126 for r in 0..mi {
2127 for q in 0..nj {
2128 let mut acc = 0f32;
2129 for p in 0..k {
2130 let bv = if b_rows_are_n {
2131 b_mat[(j + q) * ldb + p]
2132 } else {
2133 b_mat[p * ldb + j + q]
2134 };
2135 acc += a[(i + r) * lda + p] * bv;
2136 }
2137 c[(i + r) * ldc + j + q] = acc * alpha;
2138 }
2139 }
2140 }
2141 j += nj;
2142 }
2143 i += mi;
2144 }
2145 }
2146}
2147
2148#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
2150#[allow(clippy::too_many_arguments)]
2151pub(crate) fn sgemm_rm(
2152 m: usize,
2153 n: usize,
2154 k: usize,
2155 alpha: f32,
2156 a: &[f32],
2157 lda: usize,
2158 b_mat: &[f32],
2159 ldb: usize,
2160 b_rows_are_n: bool,
2161 c: &mut [f32],
2162 ldc: usize,
2163) {
2164 neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
2165}
2166
2167#[cfg(target_os = "macos")]
2170#[allow(clippy::too_many_arguments)]
2171pub(crate) fn sgemm_rm(
2172 m: usize,
2173 n: usize,
2174 k: usize,
2175 alpha: f32,
2176 a: &[f32],
2177 lda: usize,
2178 b_mat: &[f32],
2179 ldb: usize,
2180 b_rows_are_n: bool,
2181 c: &mut [f32],
2182 ldc: usize,
2183) {
2184 debug_assert!(a.len() >= (m - 1) * lda + k);
2185 debug_assert!(c.len() >= (m - 1) * ldc + n);
2186 #[cfg(target_arch = "aarch64")]
2191 if std::env::var("CMF_FORCE_NEON_GEMM")
2192 .map(|v| v == "1")
2193 .unwrap_or(false)
2194 {
2195 return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
2196 }
2197 unsafe {
2198 accel_blas::cblas_sgemm(
2199 101, 111, if b_rows_are_n { 112 } else { 111 },
2202 m as i32,
2203 n as i32,
2204 k as i32,
2205 alpha,
2206 a.as_ptr(),
2207 lda as i32,
2208 b_mat.as_ptr(),
2209 ldb as i32,
2210 0.0,
2211 c.as_mut_ptr(),
2212 ldc as i32,
2213 );
2214 }
2215}
2216
2217#[cfg(target_os = "macos")]
2224fn qmatmat_accel(
2225 q: &[u8],
2226 row_scale: &[f32],
2227 pre: &[std::borrow::Cow<'_, [f32]>],
2228 rows: usize,
2229 cols: usize,
2230 out: &mut [f32],
2231 pool: Option<&Pool>,
2232) {
2233 const TR: usize = 2048;
2238 let b = pre.len();
2239 thread_local! {
2240 static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
2241 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
2242 }
2243 XPANEL.with(|xp| {
2244 WTILE.with(|wt| {
2245 let mut xpanel = xp.borrow_mut();
2246 xpanel.clear();
2247 for x in pre {
2248 xpanel.extend_from_slice(x);
2249 }
2250 let mut wtile = wt.borrow_mut();
2251 wtile.resize(TR * cols, 0.0);
2252 let mut r0 = 0usize;
2253 while r0 < rows {
2254 let tr = TR.min(rows - r0);
2255 let wt_addr = SendMut(wtile.as_mut_ptr());
2257 let run = |start: usize, end: usize| {
2258 for r in start..end {
2259 let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
2260 let s = row_scale[r0 + r];
2261 let dst =
2263 unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
2264 for (d, &v) in dst.iter_mut().zip(row) {
2265 *d = (v as i8) as f32 * s;
2266 }
2267 }
2268 };
2269 dispatch_rows(pool, tr, &run);
2270 unsafe {
2272 accel_blas::cblas_sgemm(
2273 101, 111, 112, b as i32,
2277 tr as i32,
2278 cols as i32,
2279 1.0,
2280 xpanel.as_ptr(),
2281 cols as i32,
2282 wtile.as_ptr(),
2283 cols as i32,
2284 0.0,
2285 out.as_mut_ptr().add(r0),
2286 rows as i32,
2287 );
2288 }
2289 r0 += tr;
2290 }
2291 })
2292 });
2293}
2294
2295fn qmatmat(
2296 q: &[u8],
2297 row_scale: &[f32],
2298 pre: &[std::borrow::Cow<'_, [f32]>],
2299 rows: usize,
2300 cols: usize,
2301 out: &mut [f32],
2302 pool: Option<&Pool>,
2303) {
2304 let b = pre.len();
2305 debug_assert_eq!(out.len(), b * rows);
2306 #[cfg(target_os = "macos")]
2311 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
2312 qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
2313 return;
2314 }
2315 #[cfg(target_arch = "aarch64")]
2316 if sdot_enabled() {
2317 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
2318 let out_addr = SendMut(out.as_mut_ptr());
2319 let blocked_ok = std::env::var("CMF_X86_BLOCKED")
2322 .map(|v| v != "0")
2323 .unwrap_or(true);
2324 let use_i8mm = i8mm_enabled();
2325 if blocked_ok {
2326 let run = |start: usize, end: usize| {
2327 let mut o = start;
2328 while o < end {
2329 if o + 2 <= end {
2330 let r0 = &q[o * cols..(o + 1) * cols];
2331 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
2332 let mut bi = 0usize;
2333 while bi + 4 <= acts.len() {
2334 let xs = [
2335 acts[bi].xq.as_slice(),
2336 acts[bi + 1].xq.as_slice(),
2337 acts[bi + 2].xq.as_slice(),
2338 acts[bi + 3].xq.as_slice(),
2339 ];
2340 let d = if use_i8mm {
2341 unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
2342 } else {
2343 unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
2344 };
2345 for (r, row) in [r0, r1].into_iter().enumerate() {
2346 for k in 0..4 {
2347 let act = &acts[bi + k];
2348 let mut v = d[r][k] as f32 * act.sx;
2349 for &(j, xv) in &act.outliers {
2350 v += (row[j] as i8) as f32 * xv;
2351 }
2352 unsafe {
2353 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
2354 };
2355 }
2356 }
2357 bi += 4;
2358 }
2359 while bi < acts.len() {
2360 for (r, row) in [r0, r1].into_iter().enumerate() {
2361 let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
2362 unsafe { *out_addr.at(bi * rows + o + r) = v };
2363 }
2364 bi += 1;
2365 }
2366 o += 2;
2367 } else {
2368 let row = &q[o * cols..(o + 1) * cols];
2369 for (bi, act) in acts.iter().enumerate() {
2370 let v = row_dot_sdot(row, act) * row_scale[o];
2371 unsafe { *out_addr.at(bi * rows + o) = v };
2372 }
2373 o += 1;
2374 }
2375 }
2376 };
2377 dispatch_rows(pool, rows, &run);
2378 return;
2379 }
2380 let run = |start: usize, end: usize| {
2381 for o in start..end {
2382 let row = &q[o * cols..(o + 1) * cols];
2383 for (bi, act) in acts.iter().enumerate() {
2384 let v = row_dot_sdot(row, act) * row_scale[o];
2385 unsafe { *out_addr.at(bi * rows + o) = v };
2386 }
2387 }
2388 };
2389 dispatch_rows(pool, rows, &run);
2390 return;
2391 }
2392 #[cfg(target_arch = "x86_64")]
2397 if avx2_a8w8_enabled() {
2398 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
2399 let out_addr = SendMut(out.as_mut_ptr());
2400 let blocked_ok = std::env::var("CMF_X86_BLOCKED")
2403 .map(|v| v != "0")
2404 .unwrap_or(true);
2405 if !avx512vnni_enabled() && blocked_ok {
2406 let run = |start: usize, end: usize| {
2407 let mut o = start;
2408 while o < end {
2409 if o + 2 <= end {
2410 let r0 = &q[o * cols..(o + 1) * cols];
2411 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
2412 let mut bi = 0usize;
2413 while bi + 4 <= acts.len() {
2414 let xs = [
2415 acts[bi].xq.as_slice(),
2416 acts[bi + 1].xq.as_slice(),
2417 acts[bi + 2].xq.as_slice(),
2418 acts[bi + 3].xq.as_slice(),
2419 ];
2420 let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
2421 for (r, row) in [r0, r1].into_iter().enumerate() {
2422 for k in 0..4 {
2423 let act = &acts[bi + k];
2424 let mut v = d[r][k] as f32 * act.sx;
2425 for &(j, xv) in &act.outliers {
2426 v += (row[j] as i8) as f32 * xv;
2427 }
2428 unsafe {
2429 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
2430 };
2431 }
2432 }
2433 bi += 4;
2434 }
2435 while bi < acts.len() {
2436 for (r, row) in [r0, r1].into_iter().enumerate() {
2437 let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
2438 unsafe { *out_addr.at(bi * rows + o + r) = v };
2439 }
2440 bi += 1;
2441 }
2442 o += 2;
2443 } else {
2444 let row = &q[o * cols..(o + 1) * cols];
2445 for (bi, act) in acts.iter().enumerate() {
2446 let v = row_dot_avx2(row, act) * row_scale[o];
2447 unsafe { *out_addr.at(bi * rows + o) = v };
2448 }
2449 o += 1;
2450 }
2451 }
2452 };
2453 dispatch_rows(pool, rows, &run);
2454 return;
2455 }
2456 let run = |start: usize, end: usize| {
2457 for o in start..end {
2458 let row = &q[o * cols..(o + 1) * cols];
2459 for (bi, act) in acts.iter().enumerate() {
2460 let v = row_dot_avx2(row, act) * row_scale[o];
2461 unsafe { *out_addr.at(bi * rows + o) = v };
2462 }
2463 }
2464 };
2465 dispatch_rows(pool, rows, &run);
2466 return;
2467 }
2468 let out_addr = SendMut(out.as_mut_ptr());
2469 let run = |start: usize, end: usize| {
2470 for o in start..end {
2471 let row = &q[o * cols..(o + 1) * cols];
2472 for (bi, x) in pre.iter().enumerate() {
2473 let mut acc = 0f32;
2474 for j in 0..cols {
2475 acc += (row[j] as i8) as f32 * x[j];
2476 }
2477 unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
2478 }
2479 }
2480 };
2481 dispatch_rows(pool, rows, &run);
2482}
2483
2484fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
2487 match pool {
2488 Some(pool) if rows >= 256 => pool.run_rows(rows, run),
2489 _ => run(0, rows),
2490 }
2491}
2492
2493fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
2495 let groups = rows * cols / GROUP_SIZE;
2496 bytes.split_at(groups * 16)
2497}
2498
2499#[inline]
2504fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
2505 #[cfg(target_arch = "aarch64")]
2506 unsafe {
2507 return vbit_fill4_neon(data, buf);
2508 }
2509 #[cfg(target_arch = "x86_64")]
2510 if avx2_enabled() {
2511 return unsafe { vbit_fill4_avx2(data, buf) };
2512 }
2513 #[allow(unreachable_code)]
2514 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
2515 let u = unpack8::<4>(&data[blk * 4..]);
2516 for k in 0..8 {
2517 chunk[k] = (u[k] - 7) as i8 as u8;
2518 }
2519 }
2520}
2521
2522#[cfg(target_arch = "aarch64")]
2523#[target_feature(enable = "neon")]
2524unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
2525 unsafe {
2528 use core::arch::aarch64::*;
2529 let n = buf.len();
2530 let mask = vdupq_n_u8(0x0F);
2531 let seven = vdupq_n_s8(7);
2532 let mut g = 0usize;
2533 while g * 32 + 32 <= n {
2534 let b = vld1q_u8(data.as_ptr().add(g * 16));
2535 let hi = vshrq_n_u8::<4>(b);
2536 let lo = vandq_u8(b, mask);
2537 let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
2538 let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
2539 vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
2540 vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
2541 g += 1;
2542 }
2543 }
2544}
2545
2546#[cfg(target_arch = "x86_64")]
2547#[target_feature(enable = "avx2")]
2548unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
2549 unsafe {
2551 use core::arch::x86_64::*;
2552 let n = buf.len();
2553 let mask = _mm_set1_epi8(0x0F);
2554 let seven = _mm256_set1_epi8(7);
2555 let mut g = 0usize;
2556 while g * 32 + 32 <= n {
2557 let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
2558 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
2559 let lo = _mm_and_si128(b, mask);
2560 let z = _mm256_sub_epi8(
2561 _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
2562 seven,
2563 );
2564 _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
2565 g += 1;
2566 }
2567 }
2568}
2569
2570#[inline(always)]
2575fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
2576 let mut acc = 0u64;
2577 for i in 0..B {
2578 acc = (acc << 8) | data[i] as u64;
2579 }
2580 let mask = (1u64 << B) - 1;
2581 let mut out = [0i32; 8];
2582 for (k, o) in out.iter_mut().enumerate() {
2583 *o = ((acc >> ((7 - k) * B)) & mask) as i32;
2584 }
2585 out
2586}
2587
2588#[allow(clippy::too_many_arguments)]
2594fn vbitmatvec(
2595 bytes: &[u8],
2596 offsets: &[usize],
2597 x: &[f32],
2598 rows: usize,
2599 cols: usize,
2600 out: &mut [f32],
2601 pool: Option<&Pool>,
2602) {
2603 debug_assert_eq!(out.len(), rows);
2604 debug_assert_eq!(offsets.len(), rows + 1);
2605
2606 if a8w8_enabled() {
2610 let act = split_act(x);
2611 let out_addr = SendMut(out.as_mut_ptr());
2612 let run = move |start: usize, end: usize| {
2613 vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
2614 };
2615 dispatch_rows(pool, rows, &run);
2616 return;
2617 }
2618
2619 let out_addr = SendMut(out.as_mut_ptr());
2620 let run = move |start: usize, end: usize| {
2621 vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
2622 };
2623 dispatch_rows(pool, rows, &run);
2624}
2625
2626#[allow(clippy::too_many_arguments)]
2630fn vbit_range_a8w8(
2631 bytes: &[u8],
2632 offsets: &[usize],
2633 x: &[f32],
2634 act: &SplitAct,
2635 rows: usize,
2636 cols: usize,
2637 out: SendMut,
2638 start: usize,
2639 end: usize,
2640) {
2641 let ng = cols / GROUP_SIZE;
2642 let bits = &bytes[..rows];
2643 let sc_off = rows;
2644 let row_dot = |r: usize| -> f32 {
2645 let b = bits[r] as usize;
2646 let l = (1i32 << (b - 1)) - 1;
2647 let mask = (1u64 << b) - 1;
2648 let data = &bytes[offsets[r]..offsets[r + 1]];
2649 if b == 8 {
2650 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
2652 let mut dot = 0f32;
2653 for g in 0..ng {
2654 let so = (r * ng + g) * 2;
2655 let sgf = f16_to_f32(u16::from_le_bytes([
2656 bytes[sc_off + so],
2657 bytes[sc_off + so + 1],
2658 ]));
2659 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2660 let mut gd = 0f32;
2661 for &xv in xg.iter() {
2662 if nbits < 8 {
2663 acc = (acc << 8) | data[idx] as u64;
2664 idx += 1;
2665 nbits += 8;
2666 }
2667 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
2668 nbits -= 8;
2669 gd += (u - l) as f32 * xv;
2670 }
2671 dot += gd * sgf;
2672 }
2673 return dot;
2674 }
2675 thread_local! {
2679 static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
2680 const { std::cell::RefCell::new(Vec::new()) };
2681 }
2682 #[inline(always)]
2683 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
2684 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
2685 let u = unpack8::<B>(&data[blk * B..]);
2686 for k in 0..8 {
2687 chunk[k] = (u[k] - l) as i8 as u8;
2688 }
2689 }
2690 }
2691 let _ = mask;
2692 VBIT_SCRATCH.with(|scratch| {
2693 let mut buf = scratch.borrow_mut();
2694 buf.resize(cols, 0);
2695 match b {
2696 3 => fill::<3>(data, l, &mut buf),
2697 4 => vbit_fill4(data, &mut buf),
2698 5 => fill::<5>(data, l, &mut buf),
2699 6 => fill::<6>(data, l, &mut buf),
2700 _ => unreachable!(),
2701 }
2702 let mut dot = 0f32;
2703 for g in 0..ng {
2704 let so = (r * ng + g) * 2;
2705 let s = f16_to_f32(u16::from_le_bytes([
2706 bytes[sc_off + so],
2707 bytes[sc_off + so + 1],
2708 ]));
2709 let d = dot_i8_i8(
2710 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
2711 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
2712 ) as f32
2713 * act.sx;
2714 dot += d * s;
2715 }
2716 for &(j, xv) in &act.outliers {
2717 let so = (r * ng + j / GROUP_SIZE) * 2;
2718 let s = f16_to_f32(u16::from_le_bytes([
2719 bytes[sc_off + so],
2720 bytes[sc_off + so + 1],
2721 ]));
2722 dot += (buf[j] as i8) as f32 * s * xv;
2724 }
2725 dot
2726 })
2727 };
2728 for r in start..end {
2729 unsafe { *out.at(r) = row_dot(r) };
2731 }
2732}
2733
2734#[allow(clippy::too_many_arguments)]
2736fn vbit_range_f32(
2737 bytes: &[u8],
2738 offsets: &[usize],
2739 x: &[f32],
2740 rows: usize,
2741 cols: usize,
2742 out: SendMut,
2743 start: usize,
2744 end: usize,
2745) {
2746 let ng = cols / GROUP_SIZE;
2747 let bits = &bytes[..rows];
2748 let sc_off = rows;
2749 #[inline(always)]
2753 fn dot_row<const B: usize>(
2754 data: &[u8],
2755 bytes: &[u8],
2756 sc_off: usize,
2757 r: usize,
2758 ng: usize,
2759 x: &[f32],
2760 ) -> f32 {
2761 let l = ((1i32 << (B - 1)) - 1) as f32;
2762 let gbytes = GROUP_SIZE * B / 8;
2763 let mut dot = 0f32;
2764 for g in 0..ng {
2765 let so = (r * ng + g) * 2;
2766 let s = f16_to_f32(u16::from_le_bytes([
2767 bytes[sc_off + so],
2768 bytes[sc_off + so + 1],
2769 ]));
2770 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2771 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
2772 let mut gd = 0f32;
2773 for blk in 0..GROUP_SIZE / 8 {
2774 let u = unpack8::<B>(&gd0[blk * B..]);
2775 let xb = &xg[blk * 8..blk * 8 + 8];
2776 for k in 0..8 {
2777 gd += (u[k] as f32 - l) * xb[k];
2778 }
2779 }
2780 dot += gd * s;
2781 }
2782 dot
2783 }
2784 for r in start..end {
2785 let data = &bytes[offsets[r]..offsets[r + 1]];
2786 let v = match bits[r] {
2787 3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
2788 4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
2789 5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
2790 6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
2791 8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
2792 b => unreachable!("vbit bit-width {b} (validated at load)"),
2793 };
2794 unsafe { *out.at(r) = v };
2796 }
2797}
2798
2799#[allow(clippy::too_many_arguments)]
2804fn vbitmatvec2(
2805 bytes: &[u8],
2806 offsets: &[usize],
2807 x1: &[f32],
2808 x2: &[f32],
2809 rows: usize,
2810 cols: usize,
2811 o1: &mut [f32],
2812 o2: &mut [f32],
2813 pool: Option<&Pool>,
2814) {
2815 debug_assert_eq!(o1.len(), rows);
2816 debug_assert_eq!(o2.len(), rows);
2817
2818 if a8w8_enabled() {
2819 let a1 = split_act(x1);
2820 let a2 = split_act(x2);
2821 let p1 = SendMut(o1.as_mut_ptr());
2822 let p2 = SendMut(o2.as_mut_ptr());
2823 let run = move |start: usize, end: usize| {
2824 vbit_range2_a8w8(
2825 bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
2826 )
2827 };
2828 dispatch_rows(pool, rows, &run);
2829 return;
2830 }
2831
2832 let p1 = SendMut(o1.as_mut_ptr());
2833 let p2 = SendMut(o2.as_mut_ptr());
2834 let run = move |start: usize, end: usize| {
2835 vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
2836 };
2837 dispatch_rows(pool, rows, &run);
2838}
2839
2840#[allow(clippy::too_many_arguments)]
2844fn vbit_range2_a8w8(
2845 bytes: &[u8],
2846 offsets: &[usize],
2847 x1: &[f32],
2848 x2: &[f32],
2849 a1: &SplitAct,
2850 a2: &SplitAct,
2851 rows: usize,
2852 cols: usize,
2853 p1: SendMut,
2854 p2: SendMut,
2855 start: usize,
2856 end: usize,
2857) {
2858 let ng = cols / GROUP_SIZE;
2859 let bits = &bytes[..rows];
2860 let sc_off = rows;
2861 let row_dots = |r: usize| -> (f32, f32) {
2862 let b = bits[r] as usize;
2863 let l = (1i32 << (b - 1)) - 1;
2864 let data = &bytes[offsets[r]..offsets[r + 1]];
2865 if b == 8 {
2866 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
2869 let (mut d1, mut d2) = (0f32, 0f32);
2870 for g in 0..ng {
2871 let so = (r * ng + g) * 2;
2872 let sgf = f16_to_f32(u16::from_le_bytes([
2873 bytes[sc_off + so],
2874 bytes[sc_off + so + 1],
2875 ]));
2876 let (mut g1, mut g2) = (0f32, 0f32);
2877 for k in 0..GROUP_SIZE {
2878 if nbits < 8 {
2879 acc = (acc << 8) | data[idx] as u64;
2880 idx += 1;
2881 nbits += 8;
2882 }
2883 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
2884 nbits -= 8;
2885 let w = (u - l) as f32;
2886 g1 += w * x1[g * GROUP_SIZE + k];
2887 g2 += w * x2[g * GROUP_SIZE + k];
2888 }
2889 d1 += g1 * sgf;
2890 d2 += g2 * sgf;
2891 }
2892 return (d1, d2);
2893 }
2894 thread_local! {
2895 static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
2896 const { std::cell::RefCell::new(Vec::new()) };
2897 }
2898 #[inline(always)]
2899 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
2900 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
2901 let u = unpack8::<B>(&data[blk * B..]);
2902 for k in 0..8 {
2903 chunk[k] = (u[k] - l) as i8 as u8;
2904 }
2905 }
2906 }
2907 VBIT_SCRATCH2.with(|scratch| {
2908 let mut buf = scratch.borrow_mut();
2909 buf.resize(cols, 0);
2910 match b {
2911 3 => fill::<3>(data, l, &mut buf),
2912 4 => vbit_fill4(data, &mut buf),
2913 5 => fill::<5>(data, l, &mut buf),
2914 6 => fill::<6>(data, l, &mut buf),
2915 _ => unreachable!(),
2916 }
2917 let (mut d1, mut d2) = (0f32, 0f32);
2918 for g in 0..ng {
2919 let so = (r * ng + g) * 2;
2920 let s = f16_to_f32(u16::from_le_bytes([
2921 bytes[sc_off + so],
2922 bytes[sc_off + so + 1],
2923 ]));
2924 let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2925 let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
2926 let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
2927 d1 += v1 * s;
2928 d2 += v2 * s;
2929 }
2930 for &(j, xv) in &a1.outliers {
2931 let so = (r * ng + j / GROUP_SIZE) * 2;
2932 let s = f16_to_f32(u16::from_le_bytes([
2933 bytes[sc_off + so],
2934 bytes[sc_off + so + 1],
2935 ]));
2936 d1 += (buf[j] as i8) as f32 * s * xv;
2937 }
2938 for &(j, xv) in &a2.outliers {
2939 let so = (r * ng + j / GROUP_SIZE) * 2;
2940 let s = f16_to_f32(u16::from_le_bytes([
2941 bytes[sc_off + so],
2942 bytes[sc_off + so + 1],
2943 ]));
2944 d2 += (buf[j] as i8) as f32 * s * xv;
2945 }
2946 (d1, d2)
2947 })
2948 };
2949 for r in start..end {
2950 let (v1, v2) = row_dots(r);
2951 unsafe {
2953 *p1.at(r) = v1;
2954 *p2.at(r) = v2;
2955 }
2956 }
2957}
2958
2959#[allow(clippy::too_many_arguments)]
2963fn vbit_range2_f32(
2964 bytes: &[u8],
2965 offsets: &[usize],
2966 x1: &[f32],
2967 x2: &[f32],
2968 rows: usize,
2969 cols: usize,
2970 p1: SendMut,
2971 p2: SendMut,
2972 start: usize,
2973 end: usize,
2974) {
2975 let ng = cols / GROUP_SIZE;
2976 let bits = &bytes[..rows];
2977 let sc_off = rows;
2978 #[inline(always)]
2979 #[allow(clippy::too_many_arguments)]
2980 fn dot_row2<const B: usize>(
2981 data: &[u8],
2982 bytes: &[u8],
2983 sc_off: usize,
2984 r: usize,
2985 ng: usize,
2986 x1: &[f32],
2987 x2: &[f32],
2988 ) -> (f32, f32) {
2989 let l = ((1i32 << (B - 1)) - 1) as f32;
2990 let gbytes = GROUP_SIZE * B / 8;
2991 let (mut d1, mut d2) = (0f32, 0f32);
2992 for g in 0..ng {
2993 let so = (r * ng + g) * 2;
2994 let s = f16_to_f32(u16::from_le_bytes([
2995 bytes[sc_off + so],
2996 bytes[sc_off + so + 1],
2997 ]));
2998 let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2999 let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3000 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
3001 let (mut g1, mut g2) = (0f32, 0f32);
3002 for blk in 0..GROUP_SIZE / 8 {
3003 let u = unpack8::<B>(&gd0[blk * B..]);
3004 for k in 0..8 {
3005 let w = u[k] as f32 - l;
3006 g1 += w * x1g[blk * 8 + k];
3007 g2 += w * x2g[blk * 8 + k];
3008 }
3009 }
3010 d1 += g1 * s;
3011 d2 += g2 * s;
3012 }
3013 (d1, d2)
3014 }
3015 for r in start..end {
3016 let data = &bytes[offsets[r]..offsets[r + 1]];
3017 let (v1, v2) = match bits[r] {
3018 3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
3019 4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
3020 5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
3021 6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
3022 8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
3023 b => unreachable!("vbit bit-width {b} (validated at load)"),
3024 };
3025 unsafe {
3027 *p1.at(r) = v1;
3028 *p2.at(r) = v2;
3029 }
3030 }
3031}
3032
3033#[inline]
3040#[allow(unreachable_code)]
3041fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3042 #[cfg(target_arch = "aarch64")]
3043 unsafe {
3044 return dot_q4t_row_sdot(bytes, r, gpr, xq);
3045 }
3046 #[cfg(target_arch = "x86_64")]
3047 unsafe {
3048 if vnni_tiles_enabled() {
3049 return dot_q4t_row_vnni(bytes, r, gpr, xq);
3050 }
3051 return dot_q4t_row_avx2(bytes, r, gpr, xq);
3052 }
3053 let mut acc = 0f32;
3054 for gi in 0..gpr {
3055 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
3056 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3057 let mut d = 0i32;
3058 for (k, &b) in tile[2..].iter().enumerate() {
3059 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
3060 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
3061 }
3062 acc += d as f32 * s;
3063 }
3064 acc
3065}
3066
3067#[cfg(target_arch = "aarch64")]
3068#[target_feature(enable = "neon,dotprod")]
3069unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3070 unsafe {
3073 use core::arch::aarch64::*;
3074 use core::arch::asm;
3075 let lomask = vdupq_n_u8(0x0F);
3076 let eight = vdupq_n_s8(8);
3077 let mut acc = 0f32;
3078 for gi in 0..gpr {
3079 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3080 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3081 let b = vld1q_u8(t.add(2));
3082 let lo = vandq_u8(b, lomask);
3083 let hi = vshrq_n_u8::<4>(b);
3084 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
3085 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
3086 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
3087 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
3088 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3089 asm!(
3090 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
3091 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
3092 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3093 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
3094 options(pure, nomem, nostack),
3095 );
3096 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
3097 }
3098 acc
3099 }
3100}
3101
3102#[cfg(target_arch = "x86_64")]
3103#[target_feature(enable = "avx2")]
3104unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3105 unsafe {
3107 use core::arch::x86_64::*;
3108 let lomask = _mm_set1_epi8(0x0F);
3109 let eight = _mm256_set1_epi8(8);
3110 let ones = _mm256_set1_epi16(1);
3111 let mut acc = 0f32;
3112 for gi in 0..gpr {
3113 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3114 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3115 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
3116 let lo = _mm_and_si128(b, lomask);
3117 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3118 let w = _mm256_sub_epi8(
3119 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3120 eight,
3121 );
3122 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3123 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
3124 let d = _mm256_madd_epi16(p16, ones);
3125 let hi128 = _mm256_extracti128_si256::<1>(d);
3126 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3127 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3128 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3129 acc += _mm_cvtsi128_si32(s32) as f32 * s;
3130 }
3131 acc
3132 }
3133}
3134
3135#[cfg(target_arch = "x86_64")]
3139#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
3140unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3141 unsafe {
3143 use core::arch::x86_64::*;
3144 let lomask = _mm_set1_epi8(0x0F);
3145 let eight = _mm256_set1_epi8(8);
3146 let mut acc = 0f32;
3147 for gi in 0..gpr {
3148 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3149 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3150 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
3151 let lo = _mm_and_si128(b, lomask);
3152 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3153 let w = _mm256_sub_epi8(
3154 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3155 eight,
3156 );
3157 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3158 let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
3159 acc += d as f32 * s;
3160 }
3161 acc
3162 }
3163}
3164
3165#[cfg(target_arch = "x86_64")]
3170#[target_feature(enable = "avx2,fma")]
3175unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
3176 unsafe {
3178 use core::arch::x86_64::*;
3179 let lomask = _mm_set1_epi8(0x0F);
3180 let eight = _mm256_set1_epi8(8);
3181 let ones = _mm256_set1_epi16(1);
3182 let mut f0 = _mm256_setzero_ps();
3195 let mut f1 = _mm256_setzero_ps();
3196 let mut f2 = _mm256_setzero_ps();
3197 let mut f3 = _mm256_setzero_ps();
3198 for gi in 0..gpr {
3199 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3200 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3201 let sv = _mm256_set1_ps(s);
3202 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
3203 let lo = _mm_and_si128(bb, lomask);
3204 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
3205 let w = _mm256_sub_epi8(
3206 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3207 eight,
3208 );
3209 let aw = _mm256_abs_epi8(w);
3210 let off = gi * GROUP_SIZE;
3211 let dot = |xq: &[i8]| {
3212 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
3213 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
3214 _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
3215 };
3216 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
3217 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
3218 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
3219 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
3220 }
3221 [
3222 hsum256_ps(f0),
3223 hsum256_ps(f1),
3224 hsum256_ps(f2),
3225 hsum256_ps(f3),
3226 ]
3227 }
3228}
3229
3230#[cfg(target_arch = "x86_64")]
3233#[target_feature(enable = "avx2")]
3234#[inline]
3235unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
3236 unsafe {
3238 use core::arch::x86_64::*;
3239 let hi = _mm256_extractf128_ps::<1>(v);
3240 let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
3241 let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
3242 let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
3243 _mm_cvtss_f32(s)
3244 }
3245}
3246
3247#[cfg(target_arch = "x86_64")]
3249#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
3250unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
3251 unsafe {
3253 use core::arch::x86_64::*;
3254 let lomask = _mm_set1_epi8(0x0F);
3255 let eight = _mm256_set1_epi8(8);
3256 let mut f0 = _mm256_setzero_ps();
3259 let mut f1 = _mm256_setzero_ps();
3260 let mut f2 = _mm256_setzero_ps();
3261 let mut f3 = _mm256_setzero_ps();
3262 for gi in 0..gpr {
3263 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3264 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3265 let sv = _mm256_set1_ps(s);
3266 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
3267 let lo = _mm_and_si128(bb, lomask);
3268 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
3269 let w = _mm256_sub_epi8(
3270 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3271 eight,
3272 );
3273 let aw = _mm256_abs_epi8(w);
3274 let off = gi * GROUP_SIZE;
3275 let dot = |xq: &[i8]| {
3276 let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
3277 _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
3278 _mm256_setzero_si256(),
3279 aw,
3280 _mm256_sign_epi8(x, w),
3281 ))
3282 };
3283 f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
3284 f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
3285 f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
3286 f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
3287 }
3288 let acc = [
3289 hsum256_ps(f0),
3290 hsum256_ps(f1),
3291 hsum256_ps(f2),
3292 hsum256_ps(f3),
3293 ];
3294 acc
3295 }
3296}
3297
3298#[cfg(target_arch = "aarch64")]
3303#[target_feature(enable = "neon,dotprod")]
3304unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
3305 unsafe {
3307 use core::arch::aarch64::*;
3308 use core::arch::asm;
3309 let lomask = vdupq_n_u8(0x0F);
3310 let eight = vdupq_n_s8(8);
3311 let mut acc = [0f32; 4];
3312 for gi in 0..gpr {
3313 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3314 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3315 let b = vld1q_u8(t.add(2));
3316 let lo = vandq_u8(b, lomask);
3317 let hi = vshrq_n_u8::<4>(b);
3318 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
3319 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
3320 for (k, xq) in xs.iter().enumerate() {
3321 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
3322 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
3323 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3324 asm!(
3325 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
3326 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
3327 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3328 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
3329 options(pure, nomem, nostack),
3330 );
3331 acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
3332 }
3333 }
3334 acc
3335 }
3336}
3337
3338#[inline]
3340fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
3341 let gi = j / GROUP_SIZE;
3342 let k = j % GROUP_SIZE;
3343 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
3344 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3345 let byte = tile[2 + k / 2];
3346 let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
3347 ((nib as i32 - 8) as f32, s)
3348}
3349
3350#[inline]
3353fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
3354 let mut acc = 0f32;
3355 for gi in 0..gpr {
3356 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
3357 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3358 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
3359 let mut ga = 0f32;
3360 for (k, &b) in tile[2..].iter().enumerate() {
3361 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
3362 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
3363 }
3364 acc += ga * s;
3365 }
3366 acc
3367}
3368
3369fn q4t_matvec(
3371 bytes: &[u8],
3372 x: &[f32],
3373 rows: usize,
3374 cols: usize,
3375 out: &mut [f32],
3376 pool: Option<&Pool>,
3377) {
3378 debug_assert_eq!(out.len(), rows);
3379 let gpr = cols / GROUP_SIZE;
3380 let out_addr = SendMut(out.as_mut_ptr());
3381 if a8w8_enabled() {
3382 let act = split_act(x);
3383 let run = move |start: usize, end: usize| {
3384 for r in start..end {
3385 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
3386 for &(j, xv) in &act.outliers {
3387 let (w, s) = q4t_outlier(bytes, r, gpr, j);
3388 acc += w * s * xv;
3389 }
3390 unsafe { *out_addr.at(r) = acc };
3392 }
3393 };
3394 dispatch_rows(pool, rows, &run);
3395 return;
3396 }
3397 let run = move |start: usize, end: usize| {
3398 for r in start..end {
3399 unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
3401 }
3402 };
3403 dispatch_rows(pool, rows, &run);
3404}
3405
3406#[allow(clippy::too_many_arguments)]
3408fn q4t_matvec2(
3409 bytes: &[u8],
3410 x1: &[f32],
3411 x2: &[f32],
3412 rows: usize,
3413 cols: usize,
3414 o1: &mut [f32],
3415 o2: &mut [f32],
3416 pool: Option<&Pool>,
3417) {
3418 let gpr = cols / GROUP_SIZE;
3419 let p1 = SendMut(o1.as_mut_ptr());
3420 let p2 = SendMut(o2.as_mut_ptr());
3421 if a8w8_enabled() {
3422 let a1 = split_act(x1);
3423 let a2 = split_act(x2);
3424 let run = move |start: usize, end: usize| {
3425 for r in start..end {
3426 let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
3427 let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
3428 for &(j, xv) in &a1.outliers {
3429 let (w, s) = q4t_outlier(bytes, r, gpr, j);
3430 v1 += w * s * xv;
3431 }
3432 for &(j, xv) in &a2.outliers {
3433 let (w, s) = q4t_outlier(bytes, r, gpr, j);
3434 v2 += w * s * xv;
3435 }
3436 unsafe {
3438 *p1.at(r) = v1;
3439 *p2.at(r) = v2;
3440 }
3441 }
3442 };
3443 dispatch_rows(pool, rows, &run);
3444 return;
3445 }
3446 let run = move |start: usize, end: usize| {
3447 for r in start..end {
3448 unsafe {
3450 *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
3451 *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
3452 }
3453 }
3454 };
3455 dispatch_rows(pool, rows, &run);
3456}
3457
3458#[allow(clippy::too_many_arguments)]
3460#[cfg(target_os = "macos")]
3466fn dequant_matmat_accel(
3467 dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
3468 xs_all: &[f32],
3469 b: usize,
3470 rows: usize,
3471 cols: usize,
3472 out: &mut [f32],
3473 pool: Option<&Pool>,
3474) {
3475 const TR: usize = 2048;
3476 thread_local! {
3477 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3478 }
3479 WTILE.with(|wt| {
3480 let mut wtile = wt.borrow_mut();
3481 wtile.resize(TR * cols, 0.0);
3482 let mut r0 = 0usize;
3483 while r0 < rows {
3484 let tr = TR.min(rows - r0);
3485 let wt_addr = SendMut(wtile.as_mut_ptr());
3486 let run = |start: usize, end: usize| {
3487 for r in start..end {
3488 let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
3490 dequant_row(r0 + r, dst);
3491 }
3492 };
3493 dispatch_rows(pool, tr, &run);
3494 unsafe {
3495 accel_blas::cblas_sgemm(
3496 101, 111, 112, b as i32,
3500 tr as i32,
3501 cols as i32,
3502 1.0,
3503 xs_all.as_ptr(),
3504 cols as i32,
3505 wtile.as_ptr(),
3506 cols as i32,
3507 0.0,
3508 out.as_mut_ptr().add(r0),
3509 rows as i32,
3510 );
3511 }
3512 r0 += tr;
3513 }
3514 });
3515}
3516
3517fn q4t_matmat(
3518 bytes: &[u8],
3519 xs_all: &[f32],
3520 b: usize,
3521 rows: usize,
3522 cols: usize,
3523 out: &mut [f32],
3524 pool: Option<&Pool>,
3525) {
3526 debug_assert_eq!(out.len(), b * rows);
3527 let gpr = cols / GROUP_SIZE;
3528 #[cfg(target_os = "macos")]
3532 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
3533 dequant_matmat_accel(
3534 &|r, dst| {
3535 for gi in 0..gpr {
3536 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
3537 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3538 for (k, &bb) in tile[2..].iter().enumerate() {
3539 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
3540 dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
3541 }
3542 }
3543 },
3544 xs_all,
3545 b,
3546 rows,
3547 cols,
3548 out,
3549 pool,
3550 );
3551 return;
3552 }
3553 let out_addr = SendMut(out.as_mut_ptr());
3554 if a8w8_enabled() {
3555 let acts: Vec<SplitAct> = (0..b)
3556 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
3557 .collect();
3558 let acts = &acts;
3559 #[cfg(target_arch = "x86_64")]
3560 let blocked_ok = avx2_enabled()
3561 && std::env::var("CMF_X86_BLOCKED")
3562 .map(|v| v != "0")
3563 .unwrap_or(true);
3564 #[cfg(target_arch = "aarch64")]
3565 let blocked_ok = sdot_enabled()
3566 && std::env::var("CMF_X86_BLOCKED")
3567 .map(|v| v != "0")
3568 .unwrap_or(true);
3569 #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
3570 let blocked_ok = false;
3571 let run = move |start: usize, end: usize| {
3572 for r in start..end {
3573 let mut bi = 0usize;
3574 #[cfg(target_arch = "aarch64")]
3575 if blocked_ok {
3576 while bi + 4 <= acts.len() {
3577 let xs = [
3578 acts[bi].xq.as_slice(),
3579 acts[bi + 1].xq.as_slice(),
3580 acts[bi + 2].xq.as_slice(),
3581 acts[bi + 3].xq.as_slice(),
3582 ];
3583 let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
3584 for k in 0..4 {
3585 let act = &acts[bi + k];
3586 let mut acc = d[k] * act.sx;
3587 for &(j, xv) in &act.outliers {
3588 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
3589 acc += w * sc * xv;
3590 }
3591 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
3593 }
3594 bi += 4;
3595 }
3596 }
3597 #[cfg(target_arch = "x86_64")]
3598 if blocked_ok {
3599 while bi + 4 <= acts.len() {
3600 let xs = [
3601 acts[bi].xq.as_slice(),
3602 acts[bi + 1].xq.as_slice(),
3603 acts[bi + 2].xq.as_slice(),
3604 acts[bi + 3].xq.as_slice(),
3605 ];
3606 let d = unsafe {
3607 if vnni_tiles_enabled() {
3608 dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
3609 } else {
3610 dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
3611 }
3612 };
3613 for k in 0..4 {
3614 let act = &acts[bi + k];
3615 let mut acc = d[k] * act.sx;
3616 for &(j, xv) in &act.outliers {
3617 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
3618 acc += w * sc * xv;
3619 }
3620 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
3622 }
3623 bi += 4;
3624 }
3625 }
3626 let _ = blocked_ok;
3627 while bi < acts.len() {
3628 let act = &acts[bi];
3629 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
3630 for &(j, xv) in &act.outliers {
3631 let (w, s) = q4t_outlier(bytes, r, gpr, j);
3632 acc += w * s * xv;
3633 }
3634 unsafe { *out_addr.at(bi * rows + r) = acc };
3636 bi += 1;
3637 }
3638 }
3639 };
3640 dispatch_rows(pool, rows, &run);
3641 return;
3642 }
3643 let run = move |start: usize, end: usize| {
3644 for r in start..end {
3645 for bi in 0..b {
3646 let x = &xs_all[bi * cols..(bi + 1) * cols];
3647 unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
3649 }
3650 }
3651 };
3652 dispatch_rows(pool, rows, &run);
3653}
3654
3655fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
3664 (0..gpr)
3665 .map(|gi| {
3666 xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
3667 .iter()
3668 .map(|&v| v as i32)
3669 .sum()
3670 })
3671 .collect()
3672}
3673
3674#[inline]
3678#[allow(unreachable_code)]
3679#[cfg(target_arch = "x86_64")]
3684#[target_feature(enable = "avx2")]
3685unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
3686 unsafe {
3688 use core::arch::x86_64::*;
3689 let expand = _mm256_setr_epi8(
3691 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,
3692 3, 3, 3,
3693 );
3694 let bitsel = _mm256_setr_epi8(
3695 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
3696 -128, 1, 2, 4, 8, 16, 32, 64, -128,
3697 );
3698 let ones8 = _mm256_set1_epi8(1);
3699 let ones16 = _mm256_set1_epi16(1);
3700 let mut acc = 0f32;
3701 for gi in 0..gpr {
3702 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
3703 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3704 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
3705 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
3706 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
3707 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3708 let sel = _mm256_and_si256(x, mask);
3709 let p16 = _mm256_maddubs_epi16(ones8, sel);
3711 let d32 = _mm256_madd_epi16(p16, ones16);
3712 let hi128 = _mm256_extracti128_si256::<1>(d32);
3713 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
3714 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3715 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3716 let msum = _mm_cvtsi128_si32(s32);
3717 let d = 2 * msum - gsum[gi];
3720 acc += d as f32 * s;
3721 }
3722 acc
3723 }
3724}
3725
3726#[cfg(target_arch = "x86_64")]
3729#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
3730unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
3731 unsafe {
3733 use core::arch::x86_64::*;
3734 let expand = _mm256_setr_epi8(
3735 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,
3736 3, 3, 3,
3737 );
3738 let bitsel = _mm256_setr_epi8(
3739 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
3740 -128, 1, 2, 4, 8, 16, 32, 64, -128,
3741 );
3742 let ones8 = _mm256_set1_epi8(1);
3743 let mut acc = 0f32;
3744 for gi in 0..gpr {
3745 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
3746 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3747 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
3748 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
3749 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
3750 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3751 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
3752 let d = 2 * msum - gsum[gi];
3753 acc += d as f32 * s;
3754 }
3755 acc
3756 }
3757}
3758
3759#[cfg(target_arch = "x86_64")]
3761#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
3762unsafe fn dot_q1_row_1x4_vnni(
3763 bytes: &[u8],
3764 r: usize,
3765 gpr: usize,
3766 xs: [&[i8]; 4],
3767 gsums: [&[i32]; 4],
3768) -> [f32; 4] {
3769 unsafe {
3771 use core::arch::x86_64::*;
3772 let expand = _mm256_setr_epi8(
3773 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,
3774 3, 3, 3,
3775 );
3776 let bitsel = _mm256_setr_epi8(
3777 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
3778 -128, 1, 2, 4, 8, 16, 32, 64, -128,
3779 );
3780 let ones8 = _mm256_set1_epi8(1);
3781 let mut acc = [0f32; 4];
3782 for gi in 0..gpr {
3783 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
3784 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3785 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
3786 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
3787 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
3788 for (k, xq) in xs.iter().enumerate() {
3789 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3790 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
3791 let d = 2 * msum - gsums[k][gi];
3792 acc[k] += d as f32 * s;
3793 }
3794 }
3795 acc
3796 }
3797}
3798
3799#[cfg(target_arch = "x86_64")]
3802#[target_feature(enable = "avx2")]
3803unsafe fn dot_q1_row_1x4_avx2(
3804 bytes: &[u8],
3805 r: usize,
3806 gpr: usize,
3807 xs: [&[i8]; 4],
3808 gsums: [&[i32]; 4],
3809) -> [f32; 4] {
3810 unsafe {
3812 use core::arch::x86_64::*;
3813 let expand = _mm256_setr_epi8(
3814 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,
3815 3, 3, 3,
3816 );
3817 let bitsel = _mm256_setr_epi8(
3818 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
3819 -128, 1, 2, 4, 8, 16, 32, 64, -128,
3820 );
3821 let ones8 = _mm256_set1_epi8(1);
3822 let ones16 = _mm256_set1_epi16(1);
3823 let mut acc = [0f32; 4];
3824 for gi in 0..gpr {
3825 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
3826 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3827 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
3828 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
3829 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
3830 for (k, xq) in xs.iter().enumerate() {
3831 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3832 let sel = _mm256_and_si256(x, mask);
3833 let p16 = _mm256_maddubs_epi16(ones8, sel);
3834 let d32 = _mm256_madd_epi16(p16, ones16);
3835 let hi128 = _mm256_extracti128_si256::<1>(d32);
3836 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
3837 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3838 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3839 let msum = _mm_cvtsi128_si32(s32);
3840 let d = 2 * msum - gsums[k][gi];
3841 acc[k] += d as f32 * s;
3842 }
3843 }
3844 acc
3845 }
3846}
3847
3848#[allow(unreachable_code)]
3849fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
3850 #[cfg(target_arch = "aarch64")]
3851 unsafe {
3852 return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
3853 }
3854 #[cfg(target_arch = "x86_64")]
3855 if avx2_enabled() {
3856 unsafe {
3857 if vnni_tiles_enabled() {
3858 return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
3859 }
3860 return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
3861 }
3862 }
3863 let _ = gsum;
3864 let mut acc = 0f32;
3865 for gi in 0..gpr {
3866 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
3867 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3868 let mut d = 0i32;
3869 for (j, &b) in tile[2..].iter().enumerate() {
3870 for k in 0..8 {
3871 let w = ((b >> k) & 1) as i32 * 2 - 1;
3872 d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
3873 }
3874 }
3875 acc += d as f32 * s;
3876 }
3877 acc
3878}
3879
3880#[cfg(target_arch = "aarch64")]
3889#[target_feature(enable = "neon,dotprod")]
3890unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
3891 unsafe {
3894 use core::arch::aarch64::*;
3895 use core::arch::asm;
3896 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
3897 let m = vld1q_u8(MASKS.as_ptr());
3898 macro_rules! tile_dot {
3900 ($t:expr, $x:expr) => {{
3901 let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
3902 let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
3903 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
3904 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
3905 let x0 = vld1q_s8($x);
3906 let x1 = vld1q_s8($x.add(16));
3907 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3908 asm!(
3909 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
3910 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
3911 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3912 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
3913 options(pure, nomem, nostack),
3914 );
3915 vaddq_s32(a0, a1)
3916 }};
3917 }
3918 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
3927 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
3928 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
3929 const IW11: [u8; 16] = [
3930 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
3931 ];
3932 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
3933 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
3934 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
3935 let isc = vld1_u8(ISC.as_ptr());
3936 macro_rules! tile_dot_tbl {
3938 ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
3939 let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
3940 let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
3941 let x0 = vld1q_s8($x);
3942 let x1 = vld1q_s8($x.add(16));
3943 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3944 asm!(
3945 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
3946 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
3947 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3948 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
3949 options(pure, nomem, nostack),
3950 );
3951 vaddq_s32(a0, a1)
3952 }};
3953 }
3954 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
3955 let row_base = r * gpr * Q1_TILE;
3956 let abs_end = bytes.len();
3957 let xp = xq.as_ptr();
3958 let gp = gsum.as_ptr();
3959 let mut accv = vdupq_n_f32(0.0);
3960 let mut gi = 0;
3961 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
3964 let t0 = base.add(gi * Q1_TILE);
3965 let ld_a = vld1q_u8(t0);
3966 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
3967 let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
3968 let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
3969 let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
3970 let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
3971 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
3973 let g = vld1q_s32(gp.add(gi));
3974 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
3975 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
3976 let scf: float32x4_t;
3977 asm!(
3978 "fcvtl {o:v}.4s, {i:v}.4h",
3979 o = out(vreg) scf, i = in(vreg) sc16,
3980 options(pure, nomem, nostack),
3981 );
3982 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
3983 gi += 4;
3984 }
3985 let mut acc = vaddvq_f32(accv);
3986 while gi < gpr {
3987 let t = base.add(gi * Q1_TILE);
3988 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3989 let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
3990 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
3991 gi += 1;
3992 }
3993 acc
3994 }
3995}
3996
3997#[cfg(target_arch = "aarch64")]
4002#[target_feature(enable = "neon,dotprod")]
4003unsafe fn dot_q1_row_1x4_sdot(
4004 bytes: &[u8],
4005 r: usize,
4006 gpr: usize,
4007 xs: [&[i8]; 4],
4008 gs: [&[i32]; 4],
4009) -> [f32; 4] {
4010 unsafe {
4012 use core::arch::aarch64::*;
4013 use core::arch::asm;
4014 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
4015 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
4016 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
4017 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
4018 const IW11: [u8; 16] = [
4019 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
4020 ];
4021 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
4022 let m = vld1q_u8(MASKS.as_ptr());
4023 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
4024 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
4025 let isc = vld1_u8(ISC.as_ptr());
4026 macro_rules! sdot2 {
4027 ($w0:expr, $w1:expr, $x:expr) => {{
4028 let x0 = vld1q_s8($x);
4029 let x1 = vld1q_s8($x.add(16));
4030 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4031 asm!(
4032 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
4033 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
4034 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4035 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
4036 options(pure, nomem, nostack),
4037 );
4038 vaddq_s32(a0, a1)
4039 }};
4040 }
4041 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
4042 let row_base = r * gpr * Q1_TILE;
4043 let abs_end = bytes.len();
4044 let mut accv = [vdupq_n_f32(0.0); 4];
4045 let mut gi = 0;
4046 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
4047 let t0 = base.add(gi * Q1_TILE);
4048 let ld_a = vld1q_u8(t0);
4049 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
4050 let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
4052 let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
4053 let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
4054 let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
4055 let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
4056 let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
4057 let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
4058 let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
4059 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
4060 let scf: float32x4_t;
4061 asm!(
4062 "fcvtl {o:v}.4s, {i:v}.4h",
4063 o = out(vreg) scf, i = in(vreg) sc16,
4064 options(pure, nomem, nostack),
4065 );
4066 for k in 0..4 {
4067 let xp = xs[k].as_ptr();
4068 let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
4069 let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
4070 let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
4071 let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
4072 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
4073 let g = vld1q_s32(gs[k].as_ptr().add(gi));
4074 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
4075 accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
4076 }
4077 gi += 4;
4078 }
4079 let mut acc = [
4080 vaddvq_f32(accv[0]),
4081 vaddvq_f32(accv[1]),
4082 vaddvq_f32(accv[2]),
4083 vaddvq_f32(accv[3]),
4084 ];
4085 while gi < gpr {
4086 let t = base.add(gi * Q1_TILE);
4087 let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4088 let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
4089 let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
4090 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
4091 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
4092 for k in 0..4 {
4093 let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
4094 acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
4095 }
4096 gi += 1;
4097 }
4098 acc
4099 }
4100}
4101
4102#[inline]
4104fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
4105 let gi = j / GROUP_SIZE;
4106 let k = j % GROUP_SIZE;
4107 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
4108 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4109 let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
4110 ((bit as i32 * 2 - 1) as f32, s)
4111}
4112
4113#[inline]
4115fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
4116 let mut acc = 0f32;
4117 for gi in 0..gpr {
4118 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
4119 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4120 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4121 let mut ga = 0f32;
4122 for (j, &b) in tile[2..].iter().enumerate() {
4123 for k in 0..8 {
4124 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
4125 }
4126 }
4127 acc += ga * s;
4128 }
4129 acc
4130}
4131
4132#[allow(clippy::too_many_arguments)]
4135fn q1_range_a8w8(
4136 bytes: &[u8],
4137 gpr: usize,
4138 act: &SplitAct,
4139 gsum: &[i32],
4140 out: SendMut,
4141 start: usize,
4142 end: usize,
4143) {
4144 for r in start..end {
4145 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
4146 for &(j, xv) in &act.outliers {
4147 let (w, s) = q1_outlier(bytes, r, gpr, j);
4148 acc += w * s * xv;
4149 }
4150 unsafe { *out.at(r) = acc };
4152 }
4153}
4154
4155fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
4157 for r in start..end {
4158 unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
4160 }
4161}
4162
4163fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
4168 let entries = base_len + (rows + 1) * 4;
4169 (base_len, entries, entries <= bytes.len())
4170}
4171
4172#[inline]
4174fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
4175 let o = rp_off + r * 4;
4176 u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
4177}
4178
4179const SIGN5: [[f32; 5]; 256] = {
4183 let mut lut = [[0.0f32; 5]; 256];
4184 let pow3 = [1u16, 3, 9, 27, 81];
4185 let mut byte = 0usize;
4186 while byte < 256 {
4187 let mut i = 0usize;
4188 while i < 5 {
4189 let code = (byte as u16 / pow3[i]) % 3;
4190 lut[byte][i] = if code == 1 {
4191 1.0
4192 } else if code == 2 {
4193 -1.0
4194 } else {
4195 0.0
4196 };
4197 i += 1;
4198 }
4199 byte += 1;
4200 }
4201 lut
4202};
4203
4204const SIGN5_I8: [[i8; 5]; 256] = {
4206 let mut lut = [[0i8; 5]; 256];
4207 let pow3 = [1u16, 3, 9, 27, 81];
4208 let mut byte = 0usize;
4209 while byte < 256 {
4210 let mut i = 0usize;
4211 while i < 5 {
4212 let code = (byte as u16 / pow3[i]) % 3;
4213 lut[byte][i] = if code == 1 {
4214 1
4215 } else if code == 2 {
4216 -1
4217 } else {
4218 0
4219 };
4220 i += 1;
4221 }
4222 byte += 1;
4223 }
4224 lut
4225};
4226
4227const SIGN5_U64: [u64; 256] = {
4233 let mut lut = [0u64; 256];
4234 let pow3 = [1u16, 3, 9, 27, 81];
4235 let mut byte = 0usize;
4236 while byte < 256 {
4237 let mut v = 0u64;
4238 let mut i = 0usize;
4239 while i < 5 {
4240 let code = (byte as u16 / pow3[i]) % 3;
4241 let s: u8 = if code == 1 {
4242 1
4243 } else if code == 2 {
4244 0xFF
4245 } else {
4246 0
4247 };
4248 v |= (s as u64) << (i * 8);
4249 i += 1;
4250 }
4251 lut[byte] = v;
4252 byte += 1;
4253 }
4254 lut
4255};
4256
4257#[inline]
4262fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
4263 const TILE: usize = cortiq_core::quant::Q1T_TILE;
4264 let off = (r * gpr + j / GROUP_SIZE) * TILE;
4265 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
4266 let within = j % GROUP_SIZE;
4267 SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
4268}
4269
4270#[cfg(target_arch = "aarch64")]
4273#[target_feature(enable = "neon,dotprod")]
4274#[inline]
4275unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
4276 unsafe {
4278 use core::arch::aarch64::*;
4279 use core::arch::asm;
4280 let w0 = vld1q_s8(w);
4281 let w1 = vld1q_s8(w.add(16));
4282 let x0 = vld1q_s8(x);
4283 let x1 = vld1q_s8(x.add(16));
4284 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4285 asm!(
4286 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
4287 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
4288 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4289 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
4290 options(pure, nomem, nostack),
4291 );
4292 vaddvq_s32(vaddq_s32(a0, a1))
4293 }
4294}
4295
4296#[cfg(target_arch = "x86_64")]
4299#[target_feature(enable = "avx2")]
4300#[inline]
4301unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
4302 unsafe {
4304 use core::arch::x86_64::*;
4305 let wv = _mm256_loadu_si256(w as *const __m256i);
4306 let xv = _mm256_loadu_si256(x as *const __m256i);
4307 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
4308 let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
4309 let hi128 = _mm256_extracti128_si256::<1>(d);
4310 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4311 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4312 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4313 _mm_cvtsi128_si32(s32)
4314 }
4315}
4316
4317#[inline]
4322fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
4323 debug_assert!(dst.len() >= 40);
4324 unsafe {
4327 let p = dst.as_mut_ptr();
4328 for bi in 0..7 {
4329 core::ptr::write_unaligned(
4330 p.add(bi * 5) as *mut u64,
4331 SIGN5_U64[*codes.add(bi) as usize],
4332 );
4333 }
4334 }
4335}
4336
4337#[inline]
4342fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
4343 #[cfg(target_arch = "aarch64")]
4344 unsafe {
4345 return sdot32_i8(w, x);
4346 }
4347 #[cfg(target_arch = "x86_64")]
4348 unsafe {
4349 return i8dot32_avx2(w, x);
4350 }
4351 #[allow(unreachable_code)]
4352 unsafe {
4353 let mut s = 0i32;
4354 for k in 0..GROUP_SIZE {
4355 s += *w.add(k) as i32 * *x.add(k) as i32;
4356 }
4357 s
4358 }
4359}
4360
4361#[inline]
4362unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
4363 let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
4364 (
4365 SIGN5_U64[*codes as usize],
4366 SIGN5_U64[*codes.add(1) as usize],
4367 SIGN5_U64[*codes.add(2) as usize],
4368 SIGN5_U64[*codes.add(3) as usize],
4369 SIGN5_U64[*codes.add(4) as usize],
4370 SIGN5_U64[*codes.add(5) as usize],
4371 SIGN5_U64[*codes.add(6) as usize],
4372 )
4373 };
4374
4375 let u0 = s0 | (s1 << 40);
4376 let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
4377 let u2 = (s3 >> 8) | (s4 << 32);
4378 let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
4379
4380 (u0, u1, u2, u3)
4381}
4382
4383#[cfg(target_arch = "aarch64")]
4387#[target_feature(enable = "neon,dotprod")]
4388unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4389 use core::arch::aarch64::*;
4390 use core::arch::asm;
4391 unsafe {
4392 const TILE: usize = cortiq_core::quant::Q1T_TILE;
4393 let mut acc = 0f32;
4394 let bytes_ptr = bytes.as_ptr();
4395 let xq_ptr = xq.as_ptr();
4396 let row_off = r * gpr * TILE;
4397
4398 let gpr2 = gpr & !1;
4399 let mut gi = 0;
4400 while gi < gpr2 {
4401 let off0 = row_off + gi * TILE;
4402 let off1 = off0 + TILE;
4403 let s0 = f16_to_f32(u16::from_le_bytes([
4404 *bytes_ptr.add(off0),
4405 *bytes_ptr.add(off0 + 1),
4406 ]));
4407 let s1 = f16_to_f32(u16::from_le_bytes([
4408 *bytes_ptr.add(off1),
4409 *bytes_ptr.add(off1 + 1),
4410 ]));
4411
4412 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
4413 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
4414
4415 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
4416 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
4417 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
4418 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
4419
4420 let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
4421 let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
4422 let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
4423 let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
4424
4425 let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
4426 let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4427 asm!(
4428 "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
4429 "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
4430 "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
4431 "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
4432 a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
4433 a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
4434 w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
4435 w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
4436 options(pure, nomem, nostack),
4437 );
4438 let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
4439 let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
4440 acc += d0 as f32 * s0 + d1 as f32 * s1;
4441 gi += 2;
4442 }
4443
4444 if gi < gpr {
4445 let off = row_off + gi * TILE;
4446 let s = f16_to_f32(u16::from_le_bytes([
4447 *bytes_ptr.add(off),
4448 *bytes_ptr.add(off + 1),
4449 ]));
4450 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
4451 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
4452 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
4453 let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
4454 let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
4455 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4456 asm!(
4457 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
4458 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
4459 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4460 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
4461 options(pure, nomem, nostack),
4462 );
4463 let d = vaddvq_s32(vaddq_s32(a0, a1));
4464 acc += d as f32 * s;
4465 }
4466 acc
4467 }
4468}
4469
4470#[cfg(target_arch = "x86_64")]
4472#[target_feature(enable = "avx2")]
4473unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4474 use core::arch::x86_64::*;
4475 unsafe {
4476 const TILE: usize = cortiq_core::quant::Q1T_TILE;
4477 let mut acc = 0f32;
4478 let bytes_ptr = bytes.as_ptr();
4479 let xq_ptr = xq.as_ptr();
4480 let row_off = r * gpr * TILE;
4481
4482 let ones = _mm256_set1_epi16(1);
4483 for gi in 0..gpr {
4484 let off = row_off + gi * TILE;
4485 let s = f16_to_f32(u16::from_le_bytes([
4486 *bytes_ptr.add(off),
4487 *bytes_ptr.add(off + 1),
4488 ]));
4489 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
4490 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
4491 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
4492 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
4493 let d256 = _mm256_madd_epi16(p16, ones);
4494 let d128 = _mm_add_epi32(
4495 _mm256_castsi256_si128(d256),
4496 _mm256_extracti128_si256(d256, 1),
4497 );
4498 let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
4499 let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
4500 acc += d32 as f32 * s;
4501 }
4502 acc
4503 }
4504}
4505
4506#[cfg(target_arch = "x86_64")]
4508#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4509unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4510 use core::arch::x86_64::*;
4511 unsafe {
4513 const TILE: usize = cortiq_core::quant::Q1T_TILE;
4514 let mut acc = 0f32;
4515 let bytes_ptr = bytes.as_ptr();
4516 let xq_ptr = xq.as_ptr();
4517 let row_off = r * gpr * TILE;
4518 for gi in 0..gpr {
4519 let off = row_off + gi * TILE;
4520 let s = f16_to_f32(u16::from_le_bytes([
4521 *bytes_ptr.add(off),
4522 *bytes_ptr.add(off + 1),
4523 ]));
4524 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
4525 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
4526 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
4527 let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
4528 acc += d as f32 * s;
4529 }
4530 acc
4531 }
4532}
4533
4534#[inline]
4538fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4539 #[cfg(target_arch = "aarch64")]
4540 unsafe {
4541 return q1t_dot_row_sdot(bytes, r, gpr, xq);
4542 }
4543 #[cfg(target_arch = "x86_64")]
4544 unsafe {
4545 if vnni_tiles_enabled() {
4546 return q1t_dot_row_vnni(bytes, r, gpr, xq);
4547 }
4548 return q1t_dot_row_avx2(bytes, r, gpr, xq);
4549 }
4550 #[allow(unreachable_code)]
4551 {
4552 const TILE: usize = cortiq_core::quant::Q1T_TILE;
4553 let mut acc = 0f32;
4554 let mut sg = [0i8; GROUP_SIZE + 8]; for gi in 0..gpr {
4556 let off = (r * gpr + gi) * TILE;
4557 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
4558 q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
4559 let mut d = 0i32;
4560 for k in 0..GROUP_SIZE {
4561 d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
4562 }
4563 acc += d as f32 * s;
4564 }
4565 acc
4566 }
4567}
4568
4569fn q1t_row_outlier_correction(
4576 bytes: &[u8],
4577 r: usize,
4578 rp_off: usize,
4579 entries_off: usize,
4580 has_ov: bool,
4581 x: &[f32],
4582) -> f32 {
4583 if !has_ov {
4584 return 0.0;
4585 }
4586 let (c0, c1) = (
4587 q1t_rowptr(bytes, rp_off, r),
4588 q1t_rowptr(bytes, rp_off, r + 1),
4589 );
4590 let mut corr = 0f32;
4591 for p in c0..c1 {
4592 let e = entries_off + p * 4;
4593 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
4594 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
4595 corr += val * x[col];
4596 }
4597 corr
4598}
4599
4600fn q1t_dequant_row(
4604 bytes: &[u8],
4605 r: usize,
4606 gpr: usize,
4607 rp_off: usize,
4608 entries_off: usize,
4609 has_ov: bool,
4610 buf: &mut [f32],
4611) {
4612 const TILE: usize = cortiq_core::quant::Q1T_TILE;
4613 for g in 0..gpr {
4614 let off = (r * gpr + g) * TILE;
4615 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
4616 let codes = &bytes[off + 2..off + TILE];
4617 let bc = g * GROUP_SIZE;
4618 for bi in 0..6 {
4620 let lut = &SIGN5[codes[bi] as usize];
4621 let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
4622 for i in 0..5 {
4623 d[i] = lut[i] * s;
4624 }
4625 }
4626 let lut = &SIGN5[codes[6] as usize];
4627 buf[bc + 30] = lut[0] * s;
4628 buf[bc + 31] = lut[1] * s;
4629 }
4630 if !has_ov {
4631 return;
4632 }
4633 let (c0, c1) = (
4634 q1t_rowptr(bytes, rp_off, r),
4635 q1t_rowptr(bytes, rp_off, r + 1),
4636 );
4637 for p in c0..c1 {
4638 let e = entries_off + p * 4;
4639 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
4640 buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
4641 }
4642}
4643
4644fn q1t_add_overlay(
4648 bytes: &[u8],
4649 x: &[f32],
4650 rows: usize,
4651 cols: usize,
4652 out: &mut [f32],
4653 pool: Option<&Pool>,
4654) {
4655 const TILE: usize = cortiq_core::quant::Q1T_TILE;
4656 let gpr = cols / GROUP_SIZE;
4657 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
4658 if !has_ov {
4659 return;
4660 }
4661 let out_addr = SendMut(out.as_mut_ptr());
4662 let run = move |start: usize, end: usize| {
4663 for r in start..end {
4664 let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
4665 unsafe { *out_addr.at(r) += corr };
4667 }
4668 };
4669 dispatch_rows(pool, rows, &run);
4670}
4671
4672#[allow(clippy::too_many_arguments)]
4675fn q1t_range_a8w8(
4676 bytes: &[u8],
4677 gpr: usize,
4678 rp_off: usize,
4679 ent_off: usize,
4680 has_ov: bool,
4681 act: &SplitAct,
4682 x: &[f32],
4683 out: SendMut,
4684 start: usize,
4685 end: usize,
4686) {
4687 for r in start..end {
4688 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
4689 for &(j, xv) in &act.outliers {
4690 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
4691 }
4692 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
4693 unsafe { *out.at(r) = acc };
4695 }
4696}
4697
4698#[allow(clippy::too_many_arguments)]
4701fn q1t_range_f32_batch(
4702 bytes: &[u8],
4703 gpr: usize,
4704 rp_off: usize,
4705 ent_off: usize,
4706 has_ov: bool,
4707 x: &[f32],
4708 out: SendMut,
4709 start: usize,
4710 end: usize,
4711) {
4712 const TILE: usize = cortiq_core::quant::Q1T_TILE;
4713 let mut sg = [0f32; GROUP_SIZE];
4714 for r in start..end {
4715 let mut acc = 0f32;
4716 for g in 0..gpr {
4717 let off = (r * gpr + g) * TILE;
4718 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
4719 let codes = &bytes[off + 2..off + TILE];
4720 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
4721 for bi in 0..6 {
4722 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
4723 }
4724 let lut = &SIGN5[codes[6] as usize];
4725 sg[30] = lut[0];
4726 sg[31] = lut[1];
4727 let mut gsum = 0f32;
4728 for k in 0..GROUP_SIZE {
4729 gsum += sg[k] * xg[k];
4730 }
4731 acc += s * gsum;
4732 }
4733 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
4734 unsafe { *out.at(r) = acc };
4736 }
4737}
4738
4739fn q1t_matvec(
4743 bytes: &[u8],
4744 x: &[f32],
4745 rows: usize,
4746 cols: usize,
4747 out: &mut [f32],
4748 pool: Option<&Pool>,
4749) {
4750 debug_assert_eq!(out.len(), rows);
4751 const TILE: usize = cortiq_core::quant::Q1T_TILE;
4752 let gpr = cols / GROUP_SIZE;
4753 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
4754 let out_addr = SendMut(out.as_mut_ptr());
4755 if a8w8_enabled() {
4759 let act = split_act(x);
4760 let act = &act;
4761 let run = move |start: usize, end: usize| {
4762 for r in start..end {
4763 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
4764 for &(j, xv) in &act.outliers {
4765 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
4766 }
4767 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
4768 unsafe { *out_addr.at(r) = acc };
4770 }
4771 };
4772 dispatch_rows(pool, rows, &run);
4773 return;
4774 }
4775 let run = move |start: usize, end: usize| {
4776 let mut sg = [0f32; GROUP_SIZE];
4780 for r in start..end {
4781 let mut acc = 0f32;
4782 for g in 0..gpr {
4783 let off = (r * gpr + g) * TILE;
4784 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
4785 let codes = &bytes[off + 2..off + TILE];
4786 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
4787 for bi in 0..6 {
4788 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
4789 }
4790 let lut = &SIGN5[codes[6] as usize];
4791 sg[30] = lut[0];
4792 sg[31] = lut[1];
4793 let mut gsum = 0f32;
4794 for k in 0..GROUP_SIZE {
4795 gsum += sg[k] * xg[k];
4796 }
4797 acc += s * gsum;
4798 }
4799 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
4800 unsafe { *out_addr.at(r) = acc };
4801 }
4802 };
4803 dispatch_rows(pool, rows, &run);
4804}
4805
4806#[cfg(target_arch = "aarch64")]
4812#[target_feature(enable = "neon,dotprod")]
4813unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
4814 use core::arch::aarch64::*;
4815 use core::arch::asm;
4816 unsafe {
4818 const TILE: usize = cortiq_core::quant::Q1T_TILE;
4819 let bytes_ptr = bytes.as_ptr();
4820 let row_off = r * gpr * TILE;
4821 let xp = [xa.as_ptr(), xb.as_ptr()];
4822 let mut acc = [0f32; 2];
4823 macro_rules! sdot2 {
4824 ($w0:expr, $w1:expr, $x:expr) => {{
4825 let x0 = vld1q_s8($x);
4826 let x1 = vld1q_s8($x.add(16));
4827 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4828 asm!(
4829 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
4830 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
4831 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4832 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
4833 options(pure, nomem, nostack),
4834 );
4835 vaddvq_s32(vaddq_s32(a0, a1))
4836 }};
4837 }
4838 let gpr2 = gpr & !1;
4839 let mut gi = 0;
4840 while gi < gpr2 {
4841 let off0 = row_off + gi * TILE;
4842 let off1 = off0 + TILE;
4843 let s0 = f16_to_f32(u16::from_le_bytes([
4844 *bytes_ptr.add(off0),
4845 *bytes_ptr.add(off0 + 1),
4846 ]));
4847 let s1 = f16_to_f32(u16::from_le_bytes([
4848 *bytes_ptr.add(off1),
4849 *bytes_ptr.add(off1 + 1),
4850 ]));
4851 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
4852 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
4853 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
4854 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
4855 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
4856 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
4857 for k in 0..2 {
4858 let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
4859 let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
4860 acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
4861 }
4862 gi += 2;
4863 }
4864 if gi < gpr {
4865 let off = row_off + gi * TILE;
4866 let s = f16_to_f32(u16::from_le_bytes([
4867 *bytes_ptr.add(off),
4868 *bytes_ptr.add(off + 1),
4869 ]));
4870 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
4871 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
4872 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
4873 for k in 0..2 {
4874 let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
4875 acc[k] += d as f32 * s;
4876 }
4877 }
4878 acc
4879 }
4880}
4881
4882fn q1t_matvec2(
4888 bytes: &[u8],
4889 x1: &[f32],
4890 x2: &[f32],
4891 rows: usize,
4892 cols: usize,
4893 o1: &mut [f32],
4894 o2: &mut [f32],
4895 pool: Option<&Pool>,
4896) {
4897 debug_assert_eq!(o1.len(), rows);
4898 debug_assert_eq!(o2.len(), rows);
4899 const TILE: usize = cortiq_core::quant::Q1T_TILE;
4900 let gpr = cols / GROUP_SIZE;
4901 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
4902 let out1 = SendMut(o1.as_mut_ptr());
4903 let out2 = SendMut(o2.as_mut_ptr());
4904 if a8w8_enabled() {
4905 let a1 = split_act(x1);
4906 let a2 = split_act(x2);
4907 let (a1, a2) = (&a1, &a2);
4908 let run = move |start: usize, end: usize| {
4909 for r in start..end {
4910 #[cfg(target_arch = "aarch64")]
4911 let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
4914 #[cfg(not(target_arch = "aarch64"))]
4915 let ds = [
4916 q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
4917 q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
4918 ];
4919 let mut acc1 = ds[0] * a1.sx;
4920 for &(j, xv) in &a1.outliers {
4921 acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
4922 }
4923 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
4924 let mut acc2 = ds[1] * a2.sx;
4925 for &(j, xv) in &a2.outliers {
4926 acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
4927 }
4928 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
4929 unsafe {
4931 *out1.at(r) = acc1;
4932 *out2.at(r) = acc2;
4933 }
4934 }
4935 };
4936 dispatch_rows(pool, rows, &run);
4937 return;
4938 }
4939 let run = move |start: usize, end: usize| {
4940 let mut sg = [0f32; GROUP_SIZE];
4943 for r in start..end {
4944 let mut acc1 = 0f32;
4945 let mut acc2 = 0f32;
4946 for g in 0..gpr {
4947 let off = (r * gpr + g) * TILE;
4948 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
4949 let codes = &bytes[off + 2..off + TILE];
4950 for bi in 0..6 {
4951 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
4952 }
4953 let lut = &SIGN5[codes[6] as usize];
4954 sg[30] = lut[0];
4955 sg[31] = lut[1];
4956 let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
4957 let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
4958 let mut gsum1 = 0f32;
4959 for k in 0..GROUP_SIZE {
4960 gsum1 += sg[k] * xg1[k];
4961 }
4962 acc1 += s * gsum1;
4963 let mut gsum2 = 0f32;
4964 for k in 0..GROUP_SIZE {
4965 gsum2 += sg[k] * xg2[k];
4966 }
4967 acc2 += s * gsum2;
4968 }
4969 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
4970 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
4971 unsafe {
4973 *out1.at(r) = acc1;
4974 *out2.at(r) = acc2;
4975 }
4976 }
4977 };
4978 dispatch_rows(pool, rows, &run);
4979}
4980
4981fn q1t_matmat(
4984 bytes: &[u8],
4985 xs: &[f32],
4986 b: usize,
4987 rows: usize,
4988 cols: usize,
4989 out: &mut [f32],
4990 pool: Option<&Pool>,
4991) {
4992 debug_assert_eq!(out.len(), b * rows);
4993 const TILE: usize = cortiq_core::quant::Q1T_TILE;
4994 let gpr = cols / GROUP_SIZE;
4995 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
4996 let out_addr = SendMut(out.as_mut_ptr());
4997 if a8w8_enabled() {
5001 let acts: Vec<SplitAct> = (0..b)
5002 .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
5003 .collect();
5004 let acts = &acts;
5005 let run = move |start: usize, end: usize| {
5006 let mut sg = vec![0i8; cols + 8]; let mut sc = vec![0f32; gpr]; let mut accs = vec![0f32; b]; for r in start..end {
5010 for g in 0..gpr {
5011 let off = (r * gpr + g) * TILE;
5012 sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
5013 q1t_unpack_group_i8(
5014 bytes.as_ptr().wrapping_add(off + 2),
5015 &mut sg[g * GROUP_SIZE..],
5016 );
5017 }
5018 for bi in 0..b {
5019 let act = &acts[bi];
5020 let mut isum = 0f32;
5021 for g in 0..gpr {
5022 let d = q1t_i8dot32(
5023 sg.as_ptr().wrapping_add(g * GROUP_SIZE),
5024 act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
5025 );
5026 isum += d as f32 * sc[g];
5027 }
5028 let mut acc = isum * act.sx;
5029 for &(j, xv) in &act.outliers {
5030 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
5031 }
5032 accs[bi] = acc;
5033 }
5034 if has_ov {
5038 let (c0, c1) = (
5039 q1t_rowptr(bytes, rp_off, r),
5040 q1t_rowptr(bytes, rp_off, r + 1),
5041 );
5042 for p in c0..c1 {
5043 let e = ent_off + p * 4;
5044 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
5045 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
5046 for bi in 0..b {
5047 accs[bi] += val * xs[bi * cols + col];
5048 }
5049 }
5050 }
5051 for bi in 0..b {
5052 unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
5053 }
5054 }
5055 };
5056 dispatch_rows(pool, rows, &run);
5057 return;
5058 }
5059 let run = move |start: usize, end: usize| {
5060 let mut buf = vec![0f32; cols];
5061 for r in start..end {
5062 q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
5063 for bi in 0..b {
5064 let xr = &xs[bi * cols..(bi + 1) * cols];
5065 let mut acc = 0f32;
5066 for j in 0..cols {
5067 acc += buf[j] * xr[j];
5068 }
5069 unsafe { *out_addr.at(bi * rows + r) = acc };
5070 }
5071 }
5072 };
5073 dispatch_rows(pool, rows, &run);
5074}
5075
5076fn q1_matvec(
5077 bytes: &[u8],
5078 x: &[f32],
5079 rows: usize,
5080 cols: usize,
5081 out: &mut [f32],
5082 pool: Option<&Pool>,
5083) {
5084 debug_assert_eq!(out.len(), rows);
5085 let gpr = cols / GROUP_SIZE;
5086 let out_addr = SendMut(out.as_mut_ptr());
5087 if a8w8_enabled() {
5088 let act = split_act(x);
5089 let gsum = q1_group_sums(&act.xq, gpr);
5090 let (act, gsum) = (&act, &gsum);
5091 let run = move |start: usize, end: usize| {
5092 q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end)
5093 };
5094 dispatch_rows(pool, rows, &run);
5095 return;
5096 }
5097 let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
5098 dispatch_rows(pool, rows, &run);
5099}
5100
5101#[allow(clippy::too_many_arguments)]
5103fn q1_matvec2(
5104 bytes: &[u8],
5105 x1: &[f32],
5106 x2: &[f32],
5107 rows: usize,
5108 cols: usize,
5109 o1: &mut [f32],
5110 o2: &mut [f32],
5111 pool: Option<&Pool>,
5112) {
5113 let gpr = cols / GROUP_SIZE;
5114 let p1 = SendMut(o1.as_mut_ptr());
5115 let p2 = SendMut(o2.as_mut_ptr());
5116 if a8w8_enabled() {
5117 let a1 = split_act(x1);
5118 let a2 = split_act(x2);
5119 let g1 = q1_group_sums(&a1.xq, gpr);
5120 let g2 = q1_group_sums(&a2.xq, gpr);
5121 let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
5122 let run = move |start: usize, end: usize| {
5123 for r in start..end {
5124 let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
5125 let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
5126 for &(j, xv) in &a1.outliers {
5127 let (w, s) = q1_outlier(bytes, r, gpr, j);
5128 v1 += w * s * xv;
5129 }
5130 for &(j, xv) in &a2.outliers {
5131 let (w, s) = q1_outlier(bytes, r, gpr, j);
5132 v2 += w * s * xv;
5133 }
5134 unsafe {
5136 *p1.at(r) = v1;
5137 *p2.at(r) = v2;
5138 }
5139 }
5140 };
5141 dispatch_rows(pool, rows, &run);
5142 return;
5143 }
5144 let run = move |start: usize, end: usize| {
5145 for r in start..end {
5146 unsafe {
5148 *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
5149 *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
5150 }
5151 }
5152 };
5153 dispatch_rows(pool, rows, &run);
5154}
5155
5156#[allow(clippy::too_many_arguments)]
5158fn q1_matmat(
5159 bytes: &[u8],
5160 xs_all: &[f32],
5161 b: usize,
5162 rows: usize,
5163 cols: usize,
5164 out: &mut [f32],
5165 pool: Option<&Pool>,
5166) {
5167 debug_assert_eq!(out.len(), b * rows);
5168 let gpr = cols / GROUP_SIZE;
5169 let out_addr = SendMut(out.as_mut_ptr());
5170 if a8w8_enabled() {
5171 let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
5172 .map(|bi| {
5173 let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
5174 let gsum = q1_group_sums(&act.xq, gpr);
5175 (act, gsum)
5176 })
5177 .collect();
5178 let acts = &acts;
5179 #[cfg(target_arch = "x86_64")]
5180 let blocked_ok = avx2_enabled()
5181 && std::env::var("CMF_X86_BLOCKED")
5182 .map(|v| v != "0")
5183 .unwrap_or(true);
5184 #[cfg(target_arch = "aarch64")]
5185 let blocked_ok = sdot_enabled()
5186 && std::env::var("CMF_X86_BLOCKED")
5187 .map(|v| v != "0")
5188 .unwrap_or(true);
5189 let run = move |start: usize, end: usize| {
5190 for r in start..end {
5191 let mut bi = 0usize;
5192 #[cfg(target_arch = "aarch64")]
5195 if blocked_ok {
5196 while bi + 4 <= acts.len() {
5197 let xs = [
5198 acts[bi].0.xq.as_slice(),
5199 acts[bi + 1].0.xq.as_slice(),
5200 acts[bi + 2].0.xq.as_slice(),
5201 acts[bi + 3].0.xq.as_slice(),
5202 ];
5203 let gs = [
5204 acts[bi].1.as_slice(),
5205 acts[bi + 1].1.as_slice(),
5206 acts[bi + 2].1.as_slice(),
5207 acts[bi + 3].1.as_slice(),
5208 ];
5209 let d = unsafe { dot_q1_row_1x4_sdot(bytes, r, gpr, xs, gs) };
5210 for k in 0..4 {
5211 let (act, _) = &acts[bi + k];
5212 let mut acc = d[k] * act.sx;
5213 for &(j, xv) in &act.outliers {
5214 let (w, sc) = q1_outlier(bytes, r, gpr, j);
5215 acc += w * sc * xv;
5216 }
5217 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5219 }
5220 bi += 4;
5221 }
5222 }
5223 #[cfg(target_arch = "x86_64")]
5224 if blocked_ok {
5225 while bi + 4 <= acts.len() {
5226 let xs = [
5227 acts[bi].0.xq.as_slice(),
5228 acts[bi + 1].0.xq.as_slice(),
5229 acts[bi + 2].0.xq.as_slice(),
5230 acts[bi + 3].0.xq.as_slice(),
5231 ];
5232 let gs = [
5233 acts[bi].1.as_slice(),
5234 acts[bi + 1].1.as_slice(),
5235 acts[bi + 2].1.as_slice(),
5236 acts[bi + 3].1.as_slice(),
5237 ];
5238 let d = unsafe {
5239 if vnni_tiles_enabled() {
5240 dot_q1_row_1x4_vnni(bytes, r, gpr, xs, gs)
5241 } else {
5242 dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs)
5243 }
5244 };
5245 for k in 0..4 {
5246 let (act, _) = &acts[bi + k];
5247 let mut acc = d[k] * act.sx;
5248 for &(j, xv) in &act.outliers {
5249 let (w, sc) = q1_outlier(bytes, r, gpr, j);
5250 acc += w * sc * xv;
5251 }
5252 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5254 }
5255 bi += 4;
5256 }
5257 }
5258 while bi < acts.len() {
5259 let (act, gsum) = &acts[bi];
5260 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
5261 for &(j, xv) in &act.outliers {
5262 let (w, s) = q1_outlier(bytes, r, gpr, j);
5263 acc += w * s * xv;
5264 }
5265 unsafe { *out_addr.at(bi * rows + r) = acc };
5267 bi += 1;
5268 }
5269 }
5270 };
5271 dispatch_rows(pool, rows, &run);
5272 return;
5273 }
5274 let run = move |start: usize, end: usize| {
5275 for r in start..end {
5276 for bi in 0..b {
5277 let x = &xs_all[bi * cols..(bi + 1) * cols];
5278 unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
5280 }
5281 }
5282 };
5283 dispatch_rows(pool, rows, &run);
5284}
5285
5286fn q4matvec(
5292 bytes: &[u8],
5293 x: &[f32],
5294 rows: usize,
5295 cols: usize,
5296 out: &mut [f32],
5297 pool: Option<&Pool>,
5298) {
5299 debug_assert_eq!(out.len(), rows);
5300 let (packed, scales) = q4_split(bytes, rows, cols);
5301 let gpr = cols / GROUP_SIZE;
5302 let out_addr = SendMut(out.as_mut_ptr());
5303
5304 if a8w8_enabled() {
5305 let act = split_act(x);
5306 let run = move |start: usize, end: usize| {
5307 q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
5308 };
5309 dispatch_rows(pool, rows, &run);
5310 return;
5311 }
5312
5313 let run =
5314 move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
5315 dispatch_rows(pool, rows, &run);
5316}
5317
5318#[inline]
5321#[allow(unreachable_code)]
5322#[cfg(target_arch = "x86_64")]
5327#[target_feature(enable = "avx2")]
5328unsafe fn dot_q4b_row_1x4_avx2(
5329 buf: &[u8],
5330 scales: &[u8],
5331 g0: usize,
5332 gpr: usize,
5333 xs: [&[i8]; 4],
5334) -> [f32; 4] {
5335 unsafe {
5337 use core::arch::x86_64::*;
5338 let ones = _mm256_set1_epi16(1);
5339 let mut acc = [0f32; 4];
5340 for gi in 0..gpr {
5341 let s = f16_to_f32(u16::from_le_bytes([
5342 scales[(g0 + gi) * 2],
5343 scales[(g0 + gi) * 2 + 1],
5344 ]));
5345 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5346 let aw = _mm256_abs_epi8(w);
5347 for (k, xq) in xs.iter().enumerate() {
5348 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5349 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
5350 let d = _mm256_madd_epi16(p16, ones);
5351 let hi128 = _mm256_extracti128_si256::<1>(d);
5352 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5353 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5354 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5355 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
5356 }
5357 }
5358 acc
5359 }
5360}
5361
5362#[cfg(target_arch = "x86_64")]
5364#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5365unsafe fn dot_q4b_row_1x4_vnni(
5366 buf: &[u8],
5367 scales: &[u8],
5368 g0: usize,
5369 gpr: usize,
5370 xs: [&[i8]; 4],
5371) -> [f32; 4] {
5372 unsafe {
5374 use core::arch::x86_64::*;
5375 let mut acc = [0f32; 4];
5376 for gi in 0..gpr {
5377 let s = f16_to_f32(u16::from_le_bytes([
5378 scales[(g0 + gi) * 2],
5379 scales[(g0 + gi) * 2 + 1],
5380 ]));
5381 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5382 let aw = _mm256_abs_epi8(w);
5383 for (k, xq) in xs.iter().enumerate() {
5384 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5385 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
5386 acc[k] += d as f32 * s;
5387 }
5388 }
5389 acc
5390 }
5391}
5392
5393#[cfg(target_arch = "x86_64")]
5399#[target_feature(enable = "avx2")]
5400unsafe fn dot_q4b_row_1x4_sx_avx2(
5401 buf: &[u8],
5402 scales: &[u8],
5403 g0: usize,
5404 gpr: usize,
5405 xs: [&[i8]; 4],
5406 sxs: [f32; 4],
5407) -> [f32; 4] {
5408 unsafe {
5410 use core::arch::x86_64::*;
5411 let ones = _mm256_set1_epi16(1);
5412 let mut acc = [0f32; 4];
5413 for gi in 0..gpr {
5414 let s = f16_to_f32(u16::from_le_bytes([
5415 scales[(g0 + gi) * 2],
5416 scales[(g0 + gi) * 2 + 1],
5417 ]));
5418 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5419 let aw = _mm256_abs_epi8(w);
5420 for (k, xq) in xs.iter().enumerate() {
5421 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5422 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
5423 let d = _mm256_madd_epi16(p16, ones);
5424 let hi128 = _mm256_extracti128_si256::<1>(d);
5425 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5426 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5427 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5428 acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
5429 }
5430 }
5431 acc
5432 }
5433}
5434
5435#[cfg(target_arch = "x86_64")]
5438#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5439unsafe fn dot_q4b_row_1x4_sx_vnni(
5440 buf: &[u8],
5441 scales: &[u8],
5442 g0: usize,
5443 gpr: usize,
5444 xs: [&[i8]; 4],
5445 sxs: [f32; 4],
5446) -> [f32; 4] {
5447 unsafe {
5449 use core::arch::x86_64::*;
5450 let mut acc = [0f32; 4];
5451 for gi in 0..gpr {
5452 let s = f16_to_f32(u16::from_le_bytes([
5453 scales[(g0 + gi) * 2],
5454 scales[(g0 + gi) * 2 + 1],
5455 ]));
5456 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5457 let aw = _mm256_abs_epi8(w);
5458 for (k, xq) in xs.iter().enumerate() {
5459 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5460 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
5461 acc[k] += (d as f32 * sxs[k]) * s;
5462 }
5463 }
5464 acc
5465 }
5466}
5467
5468#[allow(unreachable_code)]
5469fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
5470 #[cfg(target_arch = "aarch64")]
5471 unsafe {
5472 return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
5473 }
5474 #[cfg(target_arch = "x86_64")]
5475 unsafe {
5476 return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
5477 }
5478 let mut acc = 0f32;
5479 for gi in 0..gpr {
5480 let g = g0 + gi;
5481 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
5482 let mut d = 0i32;
5483 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
5484 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
5485 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
5486 }
5487 acc += d as f32 * s;
5488 }
5489 acc
5490}
5491
5492#[inline]
5494#[allow(unreachable_code)]
5495fn dot_q4_row_i8_2(
5496 packed: &[u8],
5497 scales: &[u8],
5498 g0: usize,
5499 gpr: usize,
5500 xq1: &[i8],
5501 xq2: &[i8],
5502) -> (f32, f32) {
5503 #[cfg(target_arch = "aarch64")]
5504 unsafe {
5505 return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
5506 }
5507 #[cfg(target_arch = "x86_64")]
5508 unsafe {
5509 return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
5510 }
5511 (
5512 dot_q4_row_i8(packed, scales, g0, gpr, xq1),
5513 dot_q4_row_i8(packed, scales, g0, gpr, xq2),
5514 )
5515}
5516
5517#[allow(clippy::too_many_arguments)]
5520fn q4_range_a8w8(
5521 packed: &[u8],
5522 scales: &[u8],
5523 gpr: usize,
5524 cols: usize,
5525 act: &SplitAct,
5526 out: SendMut,
5527 start: usize,
5528 end: usize,
5529) {
5530 for r in start..end {
5531 let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
5532 for &(j, xv) in &act.outliers {
5534 let flat = r * cols + j;
5535 let byte = packed[flat / 2];
5536 let nib = if flat & 1 == 0 {
5537 byte & 0x0F
5538 } else {
5539 byte >> 4
5540 };
5541 let s = f16_to_f32(u16::from_le_bytes([
5542 scales[(flat / GROUP_SIZE) * 2],
5543 scales[(flat / GROUP_SIZE) * 2 + 1],
5544 ]));
5545 acc += ((nib as i32 - 8) as f32) * s * xv;
5546 }
5547 unsafe { *out.at(r) = acc };
5549 }
5550}
5551
5552#[allow(clippy::too_many_arguments)]
5555fn q4_range2_a8w8(
5556 packed: &[u8],
5557 scales: &[u8],
5558 gpr: usize,
5559 cols: usize,
5560 a1: &SplitAct,
5561 a2: &SplitAct,
5562 p1: SendMut,
5563 p2: SendMut,
5564 start: usize,
5565 end: usize,
5566) {
5567 for r in start..end {
5568 let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
5569 let mut acc1 = s1 * a1.sx;
5570 let mut acc2 = s2 * a2.sx;
5571 let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
5573 for &(j, xv) in outliers {
5574 let flat = r * cols + j;
5575 let byte = packed[flat / 2];
5576 let nib = if flat & 1 == 0 {
5577 byte & 0x0F
5578 } else {
5579 byte >> 4
5580 };
5581 let s = f16_to_f32(u16::from_le_bytes([
5582 scales[(flat / GROUP_SIZE) * 2],
5583 scales[(flat / GROUP_SIZE) * 2 + 1],
5584 ]));
5585 *acc += ((nib as i32 - 8) as f32) * s * xv;
5586 }
5587 };
5588 fix(&a1.outliers, &mut acc1);
5589 fix(&a2.outliers, &mut acc2);
5590 unsafe {
5592 *p1.at(r) = acc1;
5593 *p2.at(r) = acc2;
5594 }
5595 }
5596}
5597
5598fn q4_range_f32(
5600 packed: &[u8],
5601 scales: &[u8],
5602 gpr: usize,
5603 x: &[f32],
5604 out: SendMut,
5605 start: usize,
5606 end: usize,
5607) {
5608 for r in start..end {
5609 let mut acc = 0f32;
5610 for gi in 0..gpr {
5611 let g = r * gpr + gi;
5612 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
5613 let pk = &packed[g * 16..(g + 1) * 16];
5614 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5615 let mut ga = 0f32;
5616 for (k, &b) in pk.iter().enumerate() {
5617 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
5618 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
5619 }
5620 acc += ga * s;
5621 }
5622 unsafe { *out.at(r) = acc };
5624 }
5625}
5626
5627#[allow(clippy::too_many_arguments)]
5631fn q4matvec2(
5632 bytes: &[u8],
5633 x1: &[f32],
5634 x2: &[f32],
5635 rows: usize,
5636 cols: usize,
5637 o1: &mut [f32],
5638 o2: &mut [f32],
5639 pool: Option<&Pool>,
5640) {
5641 debug_assert_eq!(o1.len(), rows);
5642 debug_assert_eq!(o2.len(), rows);
5643 let (packed, scales) = q4_split(bytes, rows, cols);
5644 let gpr = cols / GROUP_SIZE;
5645
5646 if a8w8_enabled() {
5647 let a1 = split_act(x1);
5648 let a2 = split_act(x2);
5649 let p1 = SendMut(o1.as_mut_ptr());
5650 let p2 = SendMut(o2.as_mut_ptr());
5651 let run = move |start: usize, end: usize| {
5652 q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
5653 };
5654 dispatch_rows(pool, rows, &run);
5655 return;
5656 }
5657
5658 let p1 = SendMut(o1.as_mut_ptr());
5659 let p2 = SendMut(o2.as_mut_ptr());
5660 let run = move |start: usize, end: usize| {
5661 q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
5662 };
5663 dispatch_rows(pool, rows, &run);
5664}
5665
5666#[allow(clippy::too_many_arguments)]
5668fn q4_range2_f32(
5669 packed: &[u8],
5670 scales: &[u8],
5671 gpr: usize,
5672 x1: &[f32],
5673 x2: &[f32],
5674 p1: SendMut,
5675 p2: SendMut,
5676 start: usize,
5677 end: usize,
5678) {
5679 for r in start..end {
5680 let (mut acc1, mut acc2) = (0f32, 0f32);
5681 for gi in 0..gpr {
5682 let g = r * gpr + gi;
5683 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
5684 let pk = &packed[g * 16..(g + 1) * 16];
5685 let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5686 let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5687 let (mut g1, mut g2) = (0f32, 0f32);
5688 for (k, &b) in pk.iter().enumerate() {
5689 let wl = (b & 0x0F) as f32 - 8.0;
5690 let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
5691 g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
5692 g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
5693 }
5694 acc1 += g1 * s;
5695 acc2 += g2 * s;
5696 }
5697 unsafe {
5699 *p1.at(r) = acc1;
5700 *p2.at(r) = acc2;
5701 }
5702 }
5703}
5704
5705thread_local! {
5706 static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
5709 static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
5710}
5711
5712#[allow(clippy::too_many_arguments)]
5718fn q4matmat(
5719 bytes: &[u8],
5720 xs_all: &[f32],
5721 b: usize,
5722 rows: usize,
5723 cols: usize,
5724 out: &mut [f32],
5725 pool: Option<&Pool>,
5726) {
5727 debug_assert_eq!(xs_all.len(), b * cols);
5728 debug_assert_eq!(out.len(), b * rows);
5729 let (packed, scales) = q4_split(bytes, rows, cols);
5730 let gpr = cols / GROUP_SIZE;
5731 let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
5732
5733 if a8w8_enabled() {
5734 let acts: Vec<SplitAct> = (0..b)
5735 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
5736 .collect();
5737 let acts = &acts;
5738 let out_addr = SendMut(out.as_mut_ptr());
5739 let run = move |start: usize, end: usize| {
5740 ROW_I8.with(|rb| {
5741 let mut buf = rb.borrow_mut();
5742 buf.resize(cols, 0);
5743 for r in start..end {
5744 for gi in 0..gpr {
5748 let g = r * gpr + gi;
5749 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
5750 buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
5751 buf[gi * GROUP_SIZE + k * 2 + 1] =
5752 (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
5753 }
5754 }
5755 let mut bi = 0usize;
5756 #[cfg(target_arch = "x86_64")]
5757 if avx2_enabled()
5758 && std::env::var("CMF_X86_BLOCKED")
5759 .map(|v| v != "0")
5760 .unwrap_or(true)
5761 {
5762 while bi + 4 <= acts.len() {
5763 let xs = [
5764 acts[bi].xq.as_slice(),
5765 acts[bi + 1].xq.as_slice(),
5766 acts[bi + 2].xq.as_slice(),
5767 acts[bi + 3].xq.as_slice(),
5768 ];
5769 let d = unsafe {
5770 if vnni_tiles_enabled() {
5771 dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
5772 } else {
5773 dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
5774 }
5775 };
5776 for k in 0..4 {
5777 let act = &acts[bi + k];
5778 let mut acc = d[k] * act.sx;
5779 for &(j, xv) in &act.outliers {
5780 acc += (buf[j] as i8) as f32
5781 * gscale((r * cols + j) / GROUP_SIZE)
5782 * xv;
5783 }
5784 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5786 }
5787 bi += 4;
5788 }
5789 }
5790 while bi < acts.len() {
5791 let act = &acts[bi];
5792 let mut acc = 0f32;
5793 for gi in 0..gpr {
5794 let d = dot_i8_i8(
5795 &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
5796 &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
5797 );
5798 acc += d as f32 * gscale(r * gpr + gi);
5799 }
5800 acc *= act.sx;
5801 for &(j, xv) in &act.outliers {
5803 acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
5804 }
5805 unsafe { *out_addr.at(bi * rows + r) = acc };
5807 bi += 1;
5808 }
5809 }
5810 })
5811 };
5812 dispatch_rows(pool, rows, &run);
5813 return;
5814 }
5815
5816 let out_addr = SendMut(out.as_mut_ptr());
5817 let run = move |start: usize, end: usize| {
5818 ROW_F32.with(|rb| {
5819 let mut buf = rb.borrow_mut();
5820 buf.resize(cols, 0.0);
5821 for r in start..end {
5822 for gi in 0..gpr {
5825 let g = r * gpr + gi;
5826 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
5827 buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
5828 buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
5829 }
5830 }
5831 for bi in 0..b {
5832 let x = &xs_all[bi * cols..(bi + 1) * cols];
5833 let mut acc = 0f32;
5834 for gi in 0..gpr {
5835 let mut ga = 0f32;
5836 for k in 0..GROUP_SIZE / 2 {
5841 let e = gi * GROUP_SIZE + k * 2;
5842 ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
5843 }
5844 acc += ga * gscale(r * gpr + gi);
5845 }
5846 unsafe { *out_addr.at(bi * rows + r) = acc };
5848 }
5849 }
5850 })
5851 };
5852 dispatch_rows(pool, rows, &run);
5853}
5854
5855#[allow(clippy::too_many_arguments)]
5860fn vbitmatmat(
5861 bytes: &[u8],
5862 offsets: &[usize],
5863 xs_all: &[f32],
5864 b: usize,
5865 rows: usize,
5866 cols: usize,
5867 out: &mut [f32],
5868 pool: Option<&Pool>,
5869) {
5870 debug_assert_eq!(xs_all.len(), b * cols);
5871 debug_assert_eq!(out.len(), b * rows);
5872 debug_assert_eq!(offsets.len(), rows + 1);
5873 let ng = cols / GROUP_SIZE;
5874 let bits = &bytes[..rows];
5875 let sc_off = rows;
5876 let gscale = |r: usize, g: usize| {
5877 let so = (r * ng + g) * 2;
5878 f16_to_f32(u16::from_le_bytes([
5879 bytes[sc_off + so],
5880 bytes[sc_off + so + 1],
5881 ]))
5882 };
5883
5884 let decode_f32 = |r: usize, dst: &mut [f32]| {
5886 let bw = bits[r] as usize;
5887 let l = ((1i32 << (bw - 1)) - 1) as f32;
5888 let data = &bytes[offsets[r]..offsets[r + 1]];
5889 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
5890 for d in dst.iter_mut() {
5891 while nbits < bw {
5892 acc = (acc << 8) | data[idx] as u64;
5893 idx += 1;
5894 nbits += 8;
5895 }
5896 let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
5897 nbits -= bw;
5898 *d = u - l;
5899 }
5900 };
5901
5902 if a8w8_enabled() {
5903 let acts: Vec<SplitAct> = (0..b)
5904 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
5905 .collect();
5906 let acts = &acts;
5907 let out_addr = SendMut(out.as_mut_ptr());
5908 let run = move |start: usize, end: usize| {
5909 for r in start..end {
5910 let bw = bits[r] as usize;
5911 if bw == 8 {
5912 ROW_F32.with(|rb| {
5915 let mut buf = rb.borrow_mut();
5916 buf.resize(cols, 0.0);
5917 decode_f32(r, &mut buf);
5918 for bi in 0..b {
5919 let x = &xs_all[bi * cols..(bi + 1) * cols];
5920 let mut dot = 0f32;
5921 for g in 0..ng {
5922 let mut gd = 0f32;
5923 for k in 0..GROUP_SIZE {
5924 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
5925 }
5926 dot += gd * gscale(r, g);
5927 }
5928 unsafe { *out_addr.at(bi * rows + r) = dot };
5930 }
5931 });
5932 continue;
5933 }
5934 let l = (1i32 << (bw - 1)) - 1;
5935 let data = &bytes[offsets[r]..offsets[r + 1]];
5936 ROW_I8.with(|rb| {
5937 let mut buf = rb.borrow_mut();
5938 buf.resize(cols, 0);
5939 #[inline(always)]
5940 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
5941 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
5942 let u = unpack8::<B>(&data[blk * B..]);
5943 for k in 0..8 {
5944 chunk[k] = (u[k] - l) as i8 as u8;
5945 }
5946 }
5947 }
5948 match bw {
5949 3 => fill::<3>(data, l, &mut buf),
5950 4 => vbit_fill4(data, &mut buf),
5951 5 => fill::<5>(data, l, &mut buf),
5952 6 => fill::<6>(data, l, &mut buf),
5953 _ => unreachable!("vbit bit-width {bw} (validated at load)"),
5954 }
5955 let mut bi = 0usize;
5956 #[cfg(target_arch = "x86_64")]
5960 if avx2_enabled()
5961 && std::env::var("CMF_X86_BLOCKED")
5962 .map(|v| v != "0")
5963 .unwrap_or(true)
5964 {
5965 while bi + 4 <= acts.len() {
5966 let xs = [
5967 acts[bi].xq.as_slice(),
5968 acts[bi + 1].xq.as_slice(),
5969 acts[bi + 2].xq.as_slice(),
5970 acts[bi + 3].xq.as_slice(),
5971 ];
5972 let sxs = [
5973 acts[bi].sx,
5974 acts[bi + 1].sx,
5975 acts[bi + 2].sx,
5976 acts[bi + 3].sx,
5977 ];
5978 let d = unsafe {
5979 if vnni_tiles_enabled() {
5980 dot_q4b_row_1x4_sx_vnni(
5981 &buf,
5982 &bytes[sc_off..],
5983 r * ng,
5984 ng,
5985 xs,
5986 sxs,
5987 )
5988 } else {
5989 dot_q4b_row_1x4_sx_avx2(
5990 &buf,
5991 &bytes[sc_off..],
5992 r * ng,
5993 ng,
5994 xs,
5995 sxs,
5996 )
5997 }
5998 };
5999 for k in 0..4 {
6000 let act = &acts[bi + k];
6001 let mut dot = d[k];
6002 for &(j, xv) in &act.outliers {
6003 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
6004 }
6005 unsafe { *out_addr.at((bi + k) * rows + r) = dot };
6007 }
6008 bi += 4;
6009 }
6010 }
6011 while bi < acts.len() {
6012 let act = &acts[bi];
6013 let mut dot = 0f32;
6014 for g in 0..ng {
6015 let d = dot_i8_i8(
6016 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
6017 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
6018 ) as f32
6019 * act.sx;
6020 dot += d * gscale(r, g);
6021 }
6022 for &(j, xv) in &act.outliers {
6023 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
6024 }
6025 unsafe { *out_addr.at(bi * rows + r) = dot };
6027 bi += 1;
6028 }
6029 });
6030 }
6031 };
6032 dispatch_rows(pool, rows, &run);
6033 return;
6034 }
6035
6036 let out_addr = SendMut(out.as_mut_ptr());
6037 let run = move |start: usize, end: usize| {
6038 ROW_F32.with(|rb| {
6039 let mut buf = rb.borrow_mut();
6040 buf.resize(cols, 0.0);
6041 for r in start..end {
6042 decode_f32(r, &mut buf);
6043 for bi in 0..b {
6044 let x = &xs_all[bi * cols..(bi + 1) * cols];
6045 let mut dot = 0f32;
6046 for g in 0..ng {
6047 let mut gd = 0f32;
6048 for k in 0..GROUP_SIZE {
6049 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
6050 }
6051 dot += gd * gscale(r, g);
6052 }
6053 unsafe { *out_addr.at(bi * rows + r) = dot };
6055 }
6056 }
6057 })
6058 };
6059 dispatch_rows(pool, rows, &run);
6060}
6061
6062pub(crate) fn gpu_batch_job<'a>(
6066 t: &'a QTensor,
6067 x: &[f32],
6068) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
6069 match t {
6070 QTensor::Mapped {
6071 model,
6072 idx,
6073 dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
6074 rows,
6075 cols,
6076 row_scale,
6077 col_field,
6078 ..
6079 } => Some((
6080 model.clone(),
6081 crate::gpu::BatchJob {
6082 idx: *idx,
6083 rows: *rows,
6084 cols: *cols,
6085 row_scale,
6086 xs: prescale(x, col_field, *dt).into_owned(),
6087 q1: false,
6088 },
6089 )),
6090 QTensor::Mapped {
6092 model,
6093 idx,
6094 dtype: TensorDtype::Q1,
6095 rows,
6096 cols,
6097 ..
6098 } => Some((
6099 model.clone(),
6100 crate::gpu::BatchJob {
6101 idx: *idx,
6102 rows: *rows,
6103 cols: *cols,
6104 row_scale: &[],
6105 xs: x.to_vec(),
6106 q1: true,
6107 },
6108 )),
6109 _ => None,
6110 }
6111}
6112
6113thread_local! {
6114 static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
6115 static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
6116}
6117
6118pub(crate) fn prescale<'a>(
6119 x: &'a [f32],
6120 col_field: &[f32],
6121 dtype: TensorDtype,
6122) -> std::borrow::Cow<'a, [f32]> {
6123 if dtype == TensorDtype::Q8_2f {
6124 x.iter().zip(col_field).map(|(a, c)| a * c).collect()
6125 } else {
6126 std::borrow::Cow::Borrowed(x)
6127 }
6128}
6129
6130pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
6133 x: &[f32],
6134 col_field: &[f32],
6135 dtype: TensorDtype,
6136 buf_id: u8,
6137 f: F,
6138) -> R {
6139 if dtype == TensorDtype::Q8_2f {
6140 if buf_id == 1 {
6141 PRESCALE_BUF1.with(|b| {
6142 let mut buf = b.borrow_mut();
6143 buf.clear();
6144 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
6145 f(&buf)
6146 })
6147 } else {
6148 PRESCALE_BUF2.with(|b| {
6149 let mut buf = b.borrow_mut();
6150 buf.clear();
6151 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
6152 f(&buf)
6153 })
6154 }
6155 } else {
6156 f(x)
6157 }
6158}
6159
6160#[cfg(target_arch = "x86_64")]
6165pub(crate) fn avx2_enabled() -> bool {
6166 use std::sync::OnceLock;
6167 static ON: OnceLock<bool> = OnceLock::new();
6168 *ON.get_or_init(|| {
6169 std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
6170 && std::arch::is_x86_feature_detected!("avx2")
6171 && std::arch::is_x86_feature_detected!("fma")
6172 })
6173}
6174
6175#[cfg(target_arch = "x86_64")]
6180fn avx2_a8w8_enabled() -> bool {
6181 use std::sync::OnceLock;
6182 static ON: OnceLock<bool> = OnceLock::new();
6183 *ON.get_or_init(|| {
6184 avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
6185 })
6186}
6187
6188#[inline]
6192pub(crate) fn a8w8_enabled() -> bool {
6193 #[cfg(target_arch = "aarch64")]
6194 {
6195 sdot_enabled()
6196 }
6197 #[cfg(target_arch = "x86_64")]
6198 {
6199 avx2_a8w8_enabled()
6200 }
6201 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
6202 {
6203 false
6204 }
6205}
6206
6207#[inline]
6210#[allow(unreachable_code)]
6211fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
6212 #[cfg(target_arch = "aarch64")]
6213 unsafe {
6214 return dot_i8_sdot(w, xq);
6215 }
6216 #[cfg(target_arch = "x86_64")]
6217 unsafe {
6218 if avx512vnni_enabled() {
6219 return dot_i8_i8_vnni(w, xq);
6220 }
6221 return dot_i8_i8_avx2(w, xq);
6222 }
6223 w.iter()
6224 .zip(xq)
6225 .map(|(&a, &b)| (a as i8) as i32 * b as i32)
6226 .sum()
6227}
6228
6229#[cfg(target_arch = "x86_64")]
6233fn avx512vnni_enabled() -> bool {
6234 use std::sync::OnceLock;
6235 static ON: OnceLock<bool> = OnceLock::new();
6236 *ON.get_or_init(|| {
6237 std::env::var("CMF_AVX512")
6238 .map(|v| v != "0")
6239 .unwrap_or(true)
6240 && std::arch::is_x86_feature_detected!("avx512f")
6241 && std::arch::is_x86_feature_detected!("avx512bw")
6242 && std::arch::is_x86_feature_detected!("avx512vl")
6243 && std::arch::is_x86_feature_detected!("avx512vnni")
6244 })
6245}
6246
6247#[cfg(target_arch = "x86_64")]
6255fn vnni_tiles_enabled() -> bool {
6256 use std::sync::OnceLock;
6257 static ON: OnceLock<bool> = OnceLock::new();
6258 *ON.get_or_init(|| {
6259 std::env::var("CMF_VNNI_TILES")
6260 .map(|v| v != "0")
6261 .unwrap_or(true)
6262 && avx512vnni_enabled()
6263 })
6264}
6265
6266#[cfg(target_arch = "x86_64")]
6271#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6272#[inline]
6273unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
6274 unsafe {
6276 use core::arch::x86_64::*;
6277 let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
6278 let hi128 = _mm256_extracti128_si256::<1>(d);
6279 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6280 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6281 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6282 _mm_cvtsi128_si32(s32)
6283 }
6284}
6285
6286#[cfg(target_arch = "x86_64")]
6291#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6292unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
6293 unsafe {
6295 use core::arch::x86_64::*;
6296 let n = w.len();
6297 let mut j = 0usize;
6298 let mut total: i32;
6299 {
6304 #[inline(always)]
6305 unsafe fn step(
6306 w: *const u8,
6307 x: *const i8,
6308 acc: core::arch::x86_64::__m512i,
6309 ) -> core::arch::x86_64::__m512i {
6310 unsafe {
6311 use core::arch::x86_64::*;
6312 let wv = _mm512_loadu_si512(w as *const _);
6313 let xv = _mm512_loadu_si512(x as *const _);
6314 let aw = _mm512_abs_epi8(wv);
6315 let neg = _mm512_movepi8_mask(wv);
6316 let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
6317 _mm512_dpbusd_epi32(acc, aw, sx)
6318 }
6319 }
6320 let (mut a0, mut a1, mut a2, mut a3) = (
6321 _mm512_setzero_si512(),
6322 _mm512_setzero_si512(),
6323 _mm512_setzero_si512(),
6324 _mm512_setzero_si512(),
6325 );
6326 while j + 256 <= n {
6327 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
6328 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
6329 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
6330 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
6331 j += 256;
6332 }
6333 while j + 64 <= n {
6334 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
6335 j += 64;
6336 }
6337 let s01 = _mm512_add_epi32(a0, a1);
6338 let s23 = _mm512_add_epi32(a2, a3);
6339 total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
6340 }
6341 if j + 32 <= n {
6343 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
6344 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
6345 let d = _mm256_dpbusd_epi32(
6346 _mm256_setzero_si256(),
6347 _mm256_abs_epi8(wv),
6348 _mm256_sign_epi8(xv, wv),
6349 );
6350 let hi128 = _mm256_extracti128_si256::<1>(d);
6351 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6352 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6353 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6354 total += _mm_cvtsi128_si32(s32);
6355 j += 32;
6356 }
6357 while j < n {
6358 total += (w[j] as i8) as i32 * xq[j] as i32;
6359 j += 1;
6360 }
6361 total
6362 }
6363}
6364
6365#[cfg(target_arch = "x86_64")]
6367#[target_feature(enable = "avx2,fma")]
6368unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
6369 unsafe {
6371 use core::arch::x86_64::*;
6372 let n = x.len();
6373 let wp = w.as_ptr();
6374 let xp = x.as_ptr();
6375 let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
6376 let mut j = 0usize;
6377 while j + 16 <= n {
6378 let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
6379 let lo = _mm256_cvtepi8_epi32(wb);
6380 let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
6381 a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
6382 a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
6383 j += 16;
6384 }
6385 let acc = _mm256_add_ps(a0, a1);
6386 let hi128 = _mm256_extractf128_ps::<1>(acc);
6387 let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
6388 let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
6389 let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
6390 let mut sum = _mm_cvtss_f32(s32);
6391 while j < n {
6392 sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
6393 j += 1;
6394 }
6395 sum
6396 }
6397}
6398
6399#[cfg(target_arch = "x86_64")]
6404#[target_feature(enable = "avx2")]
6405unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
6406 unsafe {
6408 use core::arch::x86_64::*;
6409 let n = w.len();
6410 let ones = _mm256_set1_epi16(1);
6411 let mut acc = _mm256_setzero_si256();
6412 let mut j = 0usize;
6413 while j + 32 <= n {
6414 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
6415 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
6416 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6417 acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
6418 j += 32;
6419 }
6420 let hi128 = _mm256_extracti128_si256::<1>(acc);
6421 let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
6422 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6423 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6424 let mut s = _mm_cvtsi128_si32(s32);
6425 while j < n {
6426 s += (w[j] as i8) as i32 * xq[j] as i32;
6427 j += 1;
6428 }
6429 s
6430 }
6431}
6432
6433#[cfg(target_arch = "aarch64")]
6437#[target_feature(enable = "neon,i8mm")]
6438unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
6439 unsafe {
6441 use core::arch::aarch64::*;
6442 use core::arch::asm;
6443 let n = w0.len();
6444 let w0p = w0.as_ptr() as *const i8;
6445 let w1p = w1.as_ptr() as *const i8;
6446 let mut acc01 = vdupq_n_s32(0);
6449 let mut acc23 = vdupq_n_s32(0);
6450 let mut i = 0usize;
6451 while i + 8 <= n {
6452 let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
6453 let xb01 = vcombine_s8(
6454 vld1_s8(xs[0].as_ptr().add(i)),
6455 vld1_s8(xs[1].as_ptr().add(i)),
6456 );
6457 let xb23 = vcombine_s8(
6458 vld1_s8(xs[2].as_ptr().add(i)),
6459 vld1_s8(xs[3].as_ptr().add(i)),
6460 );
6461 asm!(
6462 "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
6463 "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
6464 a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
6465 w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
6466 options(pure, nomem, nostack),
6467 );
6468 i += 8;
6469 }
6470 let mut out = [[0i32; 4]; 2];
6471 let a01: [i32; 4] = core::mem::transmute(acc01);
6472 let a23: [i32; 4] = core::mem::transmute(acc23);
6473 out[0][0] = a01[0];
6474 out[0][1] = a01[1];
6475 out[1][0] = a01[2];
6476 out[1][1] = a01[3];
6477 out[0][2] = a23[0];
6478 out[0][3] = a23[1];
6479 out[1][2] = a23[2];
6480 out[1][3] = a23[3];
6481 if i < n {
6482 for (k, x) in xs.iter().enumerate() {
6483 for j in i..n {
6484 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
6485 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
6486 }
6487 }
6488 }
6489 out
6490 }
6491}
6492
6493#[cfg(target_arch = "aarch64")]
6497#[target_feature(enable = "neon,dotprod")]
6498unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
6499 unsafe {
6501 use core::arch::aarch64::*;
6502 use core::arch::asm;
6503 let n = w0.len();
6504 let w0p = w0.as_ptr() as *const i8;
6505 let w1p = w1.as_ptr() as *const i8;
6506 let mut acc = [[vdupq_n_s32(0); 4]; 2];
6507 let mut i = 0usize;
6508 while i + 16 <= n {
6509 let wv0 = vld1q_s8(w0p.add(i));
6510 let wv1 = vld1q_s8(w1p.add(i));
6511 for (k, x) in xs.iter().enumerate() {
6512 let xv = vld1q_s8(x.as_ptr().add(i));
6513 let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
6514 asm!(
6515 "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
6516 "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
6517 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6518 w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
6519 options(pure, nomem, nostack),
6520 );
6521 acc[0][k] = a0;
6522 acc[1][k] = a1;
6523 }
6524 i += 16;
6525 }
6526 let mut out = [[0i32; 4]; 2];
6527 for r in 0..2 {
6528 for k in 0..4 {
6529 out[r][k] = vaddvq_s32(acc[r][k]);
6530 }
6531 }
6532 if i < n {
6533 for (k, x) in xs.iter().enumerate() {
6534 for j in i..n {
6535 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
6536 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
6537 }
6538 }
6539 }
6540 out
6541 }
6542}
6543
6544#[cfg(target_arch = "x86_64")]
6550#[target_feature(enable = "avx2")]
6551unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
6552 unsafe {
6554 use core::arch::x86_64::*;
6555 let n = w0.len();
6556 let ones = _mm256_set1_epi16(1);
6557 let mut acc = [[_mm256_setzero_si256(); 4]; 2];
6558 let mut j = 0usize;
6559 while j + 32 <= n {
6560 let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
6561 let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
6562 let aw0 = _mm256_abs_epi8(wv0);
6563 let aw1 = _mm256_abs_epi8(wv1);
6564 for (k, x) in xs.iter().enumerate() {
6565 let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
6566 let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
6567 acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
6568 let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
6569 acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
6570 }
6571 j += 32;
6572 }
6573 let mut out = [[0i32; 4]; 2];
6574 for r in 0..2 {
6575 for k in 0..4 {
6576 let a = acc[r][k];
6577 let hi128 = _mm256_extracti128_si256::<1>(a);
6578 let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
6579 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6580 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6581 out[r][k] = _mm_cvtsi128_si32(s32);
6582 }
6583 }
6584 if j < n {
6585 for (k, x) in xs.iter().enumerate() {
6586 for i in j..n {
6587 out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
6588 out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
6589 }
6590 }
6591 }
6592 out
6593 }
6594}
6595
6596#[cfg(target_arch = "x86_64")]
6601#[inline]
6602fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
6603 let dot = if avx512vnni_enabled() && row.len() >= 64 {
6604 (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
6605 } else {
6606 unsafe { dot_i8_i8_avx2(row, &act.xq) }
6607 };
6608 let mut acc = dot as f32 * act.sx;
6609 for &(j, xv) in &act.outliers {
6610 acc += (row[j] as i8) as f32 * xv;
6611 }
6612 acc
6613}
6614
6615#[cfg(target_arch = "x86_64")]
6619#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6620unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
6621 unsafe {
6623 use core::arch::x86_64::*;
6624 let n = w.len();
6625 let flip = _mm512_set1_epi8(-128); #[inline(always)]
6627 unsafe fn step(
6628 w: *const u8,
6629 x: *const i8,
6630 flip: core::arch::x86_64::__m512i,
6631 acc: core::arch::x86_64::__m512i,
6632 ) -> core::arch::x86_64::__m512i {
6633 unsafe {
6634 use core::arch::x86_64::*;
6635 let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
6636 _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
6637 }
6638 }
6639 let (mut a0, mut a1, mut a2, mut a3) = (
6640 _mm512_setzero_si512(),
6641 _mm512_setzero_si512(),
6642 _mm512_setzero_si512(),
6643 _mm512_setzero_si512(),
6644 );
6645 let mut j = 0usize;
6646 while j + 256 <= n {
6647 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
6648 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
6649 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
6650 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
6651 j += 256;
6652 }
6653 while j + 64 <= n {
6654 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
6655 j += 64;
6656 }
6657 let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
6658 _mm512_add_epi32(a0, a1),
6659 _mm512_add_epi32(a2, a3),
6660 ));
6661 while j < n {
6663 total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
6664 j += 1;
6665 }
6666 total
6667 }
6668}
6669
6670#[cfg(target_arch = "x86_64")]
6676#[target_feature(enable = "avx2")]
6677unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
6678 unsafe {
6681 use core::arch::x86_64::*;
6682 let lomask = _mm_set1_epi8(0x0F);
6683 let eight = _mm256_set1_epi8(8);
6684 let ones = _mm256_set1_epi16(1);
6685 let mut acc = 0f32;
6686 for gi in 0..gpr {
6687 let g = g0 + gi;
6688 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
6689 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
6690 let lo = _mm_and_si128(b, lomask);
6691 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
6692 let w = _mm256_sub_epi8(
6693 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
6694 eight,
6695 );
6696 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6697 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
6698 let d = _mm256_madd_epi16(p16, ones);
6699 let hi128 = _mm256_extracti128_si256::<1>(d);
6700 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6701 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6702 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6703 acc += _mm_cvtsi128_si32(s32) as f32 * s;
6704 }
6705 acc
6706 }
6707}
6708
6709#[cfg(target_arch = "x86_64")]
6712#[target_feature(enable = "avx2")]
6713unsafe fn dot_q4_row_avx2_2(
6714 packed: &[u8],
6715 scales: &[u8],
6716 g0: usize,
6717 gpr: usize,
6718 xq1: &[i8],
6719 xq2: &[i8],
6720) -> (f32, f32) {
6721 unsafe {
6723 use core::arch::x86_64::*;
6724 let lomask = _mm_set1_epi8(0x0F);
6725 let eight = _mm256_set1_epi8(8);
6726 let ones = _mm256_set1_epi16(1);
6727 let (mut acc1, mut acc2) = (0f32, 0f32);
6728 #[inline(always)]
6729 unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
6730 unsafe {
6731 use core::arch::x86_64::*;
6732 let hi128 = _mm256_extracti128_si256::<1>(d);
6733 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6734 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6735 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6736 _mm_cvtsi128_si32(s32)
6737 }
6738 }
6739 for gi in 0..gpr {
6740 let g = g0 + gi;
6741 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
6742 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
6743 let lo = _mm_and_si128(b, lomask);
6744 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
6745 let w = _mm256_sub_epi8(
6746 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
6747 eight,
6748 );
6749 let aw = _mm256_abs_epi8(w);
6750 let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6751 let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6752 let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
6753 let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
6754 acc1 += hsum(d1) as f32 * s;
6755 acc2 += hsum(d2) as f32 * s;
6756 }
6757 (acc1, acc2)
6758 }
6759}
6760
6761#[cfg(target_arch = "x86_64")]
6763fn q8_range_avx2(
6764 q: &[u8],
6765 row_scale: &[f32],
6766 act: &SplitAct,
6767 cols: usize,
6768 out_addr: SendMut,
6769 start: usize,
6770 end: usize,
6771) {
6772 for o in start..end {
6773 let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
6774 unsafe { *out_addr.at(o) = v };
6776 }
6777}
6778
6779#[cfg(target_arch = "x86_64")]
6781#[allow(clippy::too_many_arguments)]
6782fn q8_range2_avx2(
6783 q: &[u8],
6784 row_scale: &[f32],
6785 a1: &SplitAct,
6786 a2: &SplitAct,
6787 cols: usize,
6788 p1: SendMut,
6789 p2: SendMut,
6790 start: usize,
6791 end: usize,
6792) {
6793 for o in start..end {
6794 let row = &q[o * cols..(o + 1) * cols];
6795 unsafe {
6797 *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
6798 *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
6799 }
6800 }
6801}
6802
6803#[cfg(target_arch = "aarch64")]
6814fn i8mm_enabled() -> bool {
6815 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6816 *ON.get_or_init(|| {
6817 std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
6818 && std::arch::is_aarch64_feature_detected!("i8mm")
6819 })
6820}
6821
6822#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
6826fn sdot_enabled() -> bool {
6827 use std::sync::OnceLock;
6828 static ON: OnceLock<bool> = OnceLock::new();
6829 *ON.get_or_init(|| {
6830 let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
6831 if !want {
6832 return false;
6833 }
6834
6835 #[cfg(target_arch = "aarch64")]
6836 {
6837 if std::arch::is_aarch64_feature_detected!("dotprod") {
6838 return true;
6839 }
6840 #[cfg(target_os = "android")]
6841 {
6842 if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
6843 if cpuinfo.lines().any(|l| {
6844 (l.starts_with("Features") || l.starts_with("features"))
6845 && l.contains("asimddp")
6846 }) {
6847 return true;
6848 }
6849 }
6850 }
6851 false
6852 }
6853 #[cfg(not(target_arch = "aarch64"))]
6854 {
6855 false
6856 }
6857 })
6858}
6859
6860struct SplitAct {
6865 xq: Vec<i8>,
6866 sx: f32,
6867 outliers: Vec<(usize, f32)>,
6868 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
6871 xsum: i32,
6872}
6873
6874thread_local! {
6875 static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
6878 const { std::cell::RefCell::new(Vec::new()) };
6879}
6880
6881impl Drop for SplitAct {
6882 fn drop(&mut self) {
6883 let buf = std::mem::take(&mut self.xq);
6884 if buf.capacity() > 0 {
6885 XQ_FREE.with(|f| {
6886 let mut f = f.borrow_mut();
6887 if f.len() < 16 {
6888 f.push(buf);
6889 }
6890 });
6891 }
6892 }
6893}
6894
6895fn split_act(x: &[f32]) -> SplitAct {
6896 let n = x.len();
6897 let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
6898 let thr = 8.0 * rms;
6899 let mut outliers: Vec<(usize, f32)> = Vec::new();
6903 let mut amax = 0f32;
6904 for (j, &v) in x.iter().enumerate() {
6905 let a = v.abs();
6906 if a > thr {
6907 outliers.push((j, v));
6908 } else if a > amax {
6909 amax = a;
6910 }
6911 }
6912 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
6913 let inv = 1.0 / sx;
6914 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
6915 xq.clear();
6916 xq.reserve(n);
6917 if outliers.is_empty() {
6918 xq.extend(
6919 x.iter()
6920 .map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8),
6921 );
6922 } else {
6923 xq.extend(x.iter().map(|&v| {
6925 if v.abs() > thr {
6926 0
6927 } else {
6928 (v * inv).round().clamp(-127.0, 127.0) as i8
6929 }
6930 }));
6931 }
6932 let xsum = xq.iter().map(|&v| v as i32).sum();
6933 SplitAct {
6934 xq,
6935 sx,
6936 outliers,
6937 xsum,
6938 }
6939}
6940
6941fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
6942 let n = x.len();
6943 let rms = (x
6944 .iter()
6945 .zip(col)
6946 .map(|(&a, &c)| {
6947 let v = a * c;
6948 (v * v) as f64
6949 })
6950 .sum::<f64>()
6951 / n.max(1) as f64)
6952 .sqrt() as f32;
6953 let thr = 8.0 * rms;
6954
6955 let mut outliers = Vec::new();
6956 let mut amax = 0f32;
6957 for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
6958 let v = a * c;
6959 let s = v.abs();
6960 if s > thr {
6961 outliers.push((j, v));
6962 } else if s > amax {
6963 amax = s;
6964 }
6965 }
6966
6967 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
6968 let inv = 1.0 / sx;
6969 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
6970 xq.clear();
6971 xq.reserve(n);
6972 if outliers.is_empty() {
6973 xq.extend(
6974 x.iter()
6975 .zip(col)
6976 .map(|(&a, &c)| ((a * c) * inv).round().clamp(-127.0, 127.0) as i8),
6977 );
6978 } else {
6979 xq.extend(x.iter().zip(col).map(|(&a, &c)| {
6980 let v = a * c;
6981 if v.abs() > thr {
6982 0
6983 } else {
6984 (v * inv).round().clamp(-127.0, 127.0) as i8
6985 }
6986 }));
6987 }
6988 let xsum = xq.iter().map(|&v| v as i32).sum();
6989 SplitAct {
6990 xq,
6991 sx,
6992 outliers,
6993 xsum,
6994 }
6995}
6996
6997#[cfg(target_arch = "aarch64")]
7000#[target_feature(enable = "neon,dotprod")]
7001unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
7002 unsafe {
7004 use core::arch::aarch64::*;
7005 use core::arch::asm;
7006 let wp = w.as_ptr() as *const i8;
7007 let n = w.len();
7008 let (mut a0, mut a1, mut a2, mut a3) = (
7009 vdupq_n_s32(0),
7010 vdupq_n_s32(0),
7011 vdupq_n_s32(0),
7012 vdupq_n_s32(0),
7013 );
7014 let mut i = 0;
7015 while i + 64 <= n {
7016 let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
7017 let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
7018 let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
7019 let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
7020 asm!(
7021 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7022 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7023 "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
7024 "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
7025 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
7026 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7027 w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
7028 options(pure, nomem, nostack),
7029 );
7030 i += 64;
7031 }
7032 while i + 16 <= n {
7033 let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
7034 asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
7035 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
7036 i += 16;
7037 }
7038 let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
7039 while i < n {
7040 s += (*wp.add(i)) as i32 * xq[i] as i32;
7041 i += 1;
7042 }
7043 s
7044 }
7045}
7046
7047#[cfg(target_arch = "aarch64")]
7051#[target_feature(enable = "neon,dotprod")]
7052unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
7053 unsafe {
7055 use core::arch::aarch64::*;
7056 use core::arch::asm;
7057 let n = xq.len();
7058 let px = xq.as_ptr();
7059 let (p0, p1, p2, p3) = (
7060 w0.as_ptr() as *const i8,
7061 w1.as_ptr() as *const i8,
7062 w2.as_ptr() as *const i8,
7063 w3.as_ptr() as *const i8,
7064 );
7065 let (mut a0, mut a1, mut a2, mut a3) = (
7066 vdupq_n_s32(0),
7067 vdupq_n_s32(0),
7068 vdupq_n_s32(0),
7069 vdupq_n_s32(0),
7070 );
7071 let mut i = 0;
7072 while i + 16 <= n {
7073 let x = vld1q_s8(px.add(i));
7074 let v0 = vld1q_s8(p0.add(i));
7075 let v1 = vld1q_s8(p1.add(i));
7076 let v2 = vld1q_s8(p2.add(i));
7077 let v3 = vld1q_s8(p3.add(i));
7078 asm!(
7079 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
7080 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
7081 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
7082 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
7083 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
7084 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
7085 options(pure, nomem, nostack),
7086 );
7087 i += 16;
7088 }
7089 let mut r = [
7090 vaddvq_s32(a0),
7091 vaddvq_s32(a1),
7092 vaddvq_s32(a2),
7093 vaddvq_s32(a3),
7094 ];
7095 while i < n {
7096 let xi = *px.add(i) as i32;
7097 r[0] += (*p0.add(i)) as i32 * xi;
7098 r[1] += (*p1.add(i)) as i32 * xi;
7099 r[2] += (*p2.add(i)) as i32 * xi;
7100 r[3] += (*p3.add(i)) as i32 * xi;
7101 i += 1;
7102 }
7103 r
7104 }
7105}
7106
7107#[cfg(target_arch = "aarch64")]
7114#[target_feature(enable = "neon,dotprod")]
7115unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
7116 unsafe {
7119 use core::arch::aarch64::*;
7120 use core::arch::asm;
7121 let n = xq.len();
7122 let px = xq.as_ptr();
7123 let pg = g.as_ptr() as *const i8;
7124 let (mut a0, mut a1, mut a2, mut a3) = (
7125 vdupq_n_s32(0),
7126 vdupq_n_s32(0),
7127 vdupq_n_s32(0),
7128 vdupq_n_s32(0),
7129 );
7130 let mut i = 0;
7131 while i + 16 <= n {
7132 let x = vld1q_s8(px.add(i));
7133 let base = pg.add(4 * i);
7134 let v0 = vld1q_s8(base);
7135 let v1 = vld1q_s8(base.add(16));
7136 let v2 = vld1q_s8(base.add(32));
7137 let v3 = vld1q_s8(base.add(48));
7138 asm!(
7139 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
7140 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
7141 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
7142 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
7143 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
7144 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
7145 options(pure, nomem, nostack),
7146 );
7147 i += 16;
7148 }
7149 [
7150 vaddvq_s32(a0),
7151 vaddvq_s32(a1),
7152 vaddvq_s32(a2),
7153 vaddvq_s32(a3),
7154 ]
7155 }
7156}
7157
7158#[cfg(target_arch = "aarch64")]
7164fn q8_range_sdot(
7165 q: &[u8],
7166 rep: &[u8],
7167 row_scale: &[f32],
7168 act: &SplitAct,
7169 cols: usize,
7170 out_addr: SendMut,
7171 start: usize,
7172 end: usize,
7173) {
7174 let mut o = start;
7175 if !rep.is_empty() {
7178 while o < end && o % 4 != 0 {
7179 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
7180 unsafe { *out_addr.at(o) = v };
7181 o += 1;
7182 }
7183 }
7184 while o + 4 <= end {
7185 let r = if rep.is_empty() {
7186 unsafe {
7187 dot_i8_sdot_4rows(
7188 &q[o * cols..(o + 1) * cols],
7189 &q[(o + 1) * cols..(o + 2) * cols],
7190 &q[(o + 2) * cols..(o + 3) * cols],
7191 &q[(o + 3) * cols..(o + 4) * cols],
7192 &act.xq,
7193 )
7194 }
7195 } else {
7196 unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
7197 };
7198 for k in 0..4 {
7199 let mut acc = r[k] as f32 * act.sx;
7200 for &(j, xv) in &act.outliers {
7201 acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
7202 }
7203 unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
7205 }
7206 o += 4;
7207 }
7208 while o < end {
7209 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
7210 unsafe { *out_addr.at(o) = v };
7211 o += 1;
7212 }
7213}
7214
7215#[cfg(target_arch = "aarch64")]
7218#[allow(clippy::too_many_arguments)]
7219fn q8_range2_sdot(
7220 q: &[u8],
7221 row_scale: &[f32],
7222 a1: &SplitAct,
7223 a2: &SplitAct,
7224 cols: usize,
7225 p1: SendMut,
7226 p2: SendMut,
7227 start: usize,
7228 end: usize,
7229) {
7230 for o in start..end {
7231 let row = &q[o * cols..(o + 1) * cols];
7232 unsafe {
7234 *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
7235 *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
7236 }
7237 }
7238}
7239
7240#[allow(clippy::too_many_arguments)]
7242fn q8_range2_f32(
7243 q: &[u8],
7244 row_scale: &[f32],
7245 x1: &[f32],
7246 x2: &[f32],
7247 cols: usize,
7248 p1: SendMut,
7249 p2: SendMut,
7250 start: usize,
7251 end: usize,
7252) {
7253 for o in start..end {
7254 let row = &q[o * cols..(o + 1) * cols];
7255 unsafe {
7257 *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
7258 *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
7259 }
7260 }
7261}
7262
7263fn q8_range_f32(
7265 q: &[u8],
7266 row_scale: &[f32],
7267 xs: &[f32],
7268 cols: usize,
7269 out_addr: SendMut,
7270 start: usize,
7271 end: usize,
7272) {
7273 for o in start..end {
7274 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
7275 unsafe { *out_addr.at(o) = v };
7277 }
7278}
7279
7280#[cfg(target_arch = "aarch64")]
7283#[inline]
7284fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
7285 let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
7286 for &(j, xv) in &act.outliers {
7287 acc += (row[j] as i8) as f32 * xv;
7288 }
7289 acc
7290}
7291
7292#[cfg(target_arch = "aarch64")]
7300#[target_feature(enable = "neon,dotprod")]
7301unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
7302 unsafe {
7305 use core::arch::aarch64::*;
7306 use core::arch::asm;
7307 let lomask = vdupq_n_u8(0x0F);
7308 let eight = vdupq_n_s8(8);
7309 let mut acc = 0f32;
7310 for gi in 0..gpr {
7311 let g = g0 + gi;
7312 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7313 let b = vld1q_u8(packed.as_ptr().add(g * 16));
7314 let lo = vandq_u8(b, lomask);
7315 let hi = vshrq_n_u8::<4>(b);
7316 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
7317 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
7318 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
7319 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
7320 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7321 asm!(
7322 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
7323 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
7324 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7325 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
7326 options(pure, nomem, nostack),
7327 );
7328 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
7329 }
7330 acc
7331 }
7332}
7333
7334#[cfg(target_arch = "aarch64")]
7339#[target_feature(enable = "neon,dotprod")]
7340unsafe fn dot_q4_row_sdot2(
7341 packed: &[u8],
7342 scales: &[u8],
7343 g0: usize,
7344 gpr: usize,
7345 xq1: &[i8],
7346 xq2: &[i8],
7347) -> (f32, f32) {
7348 unsafe {
7351 use core::arch::aarch64::*;
7352 use core::arch::asm;
7353 let lomask = vdupq_n_u8(0x0F);
7354 let eight = vdupq_n_s8(8);
7355 let (mut acc1, mut acc2) = (0f32, 0f32);
7356 for gi in 0..gpr {
7357 let g = g0 + gi;
7358 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7359 let b = vld1q_u8(packed.as_ptr().add(g * 16));
7360 let lo = vandq_u8(b, lomask);
7361 let hi = vshrq_n_u8::<4>(b);
7362 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
7363 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
7364 let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
7365 let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
7366 let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
7367 let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
7368 let (mut a0, mut a1, mut b0, mut b1) = (
7369 vdupq_n_s32(0),
7370 vdupq_n_s32(0),
7371 vdupq_n_s32(0),
7372 vdupq_n_s32(0),
7373 );
7374 asm!(
7375 "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
7376 "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
7377 "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
7378 "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
7379 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7380 b0 = inout(vreg) b0, b1 = inout(vreg) b1,
7381 e0 = in(vreg) e0, e1 = in(vreg) e1,
7382 x10 = in(vreg) x10, x11 = in(vreg) x11,
7383 x20 = in(vreg) x20, x21 = in(vreg) x21,
7384 options(pure, nomem, nostack),
7385 );
7386 acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
7387 acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
7388 }
7389 (acc1, acc2)
7390 }
7391}
7392
7393#[inline]
7398pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
7399 #[cfg(target_arch = "aarch64")]
7400 unsafe {
7401 return axpy_i8_f32_neon(acc, row, w);
7402 }
7403 #[cfg(target_arch = "x86_64")]
7404 if avx2_enabled() {
7405 return unsafe { axpy_i8_f32_avx2(acc, row, w) };
7406 }
7407 #[allow(unreachable_code)]
7408 {
7409 for (a, &b) in acc.iter_mut().zip(row) {
7410 *a += w * b as f32;
7411 }
7412 }
7413}
7414
7415#[cfg(target_arch = "x86_64")]
7417#[target_feature(enable = "avx2,fma")]
7418unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
7419 unsafe {
7421 use core::arch::x86_64::*;
7422 let n = acc.len().min(row.len());
7423 let ap = acc.as_mut_ptr();
7424 let rp = row.as_ptr();
7425 let wv = _mm256_set1_ps(w);
7426 let mut j = 0usize;
7427 while j + 16 <= n {
7428 let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
7429 let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
7430 let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
7431 let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
7432 let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
7433 _mm256_storeu_ps(ap.add(j), v0);
7434 _mm256_storeu_ps(ap.add(j + 8), v1);
7435 j += 16;
7436 }
7437 while j < n {
7438 *ap.add(j) += w * (*rp.add(j)) as f32;
7439 j += 1;
7440 }
7441 }
7442}
7443
7444#[cfg(target_arch = "aarch64")]
7445#[target_feature(enable = "neon")]
7446unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
7447 unsafe {
7449 use core::arch::aarch64::*;
7450 let n = acc.len().min(row.len());
7451 let ap = acc.as_mut_ptr();
7452 let rp = row.as_ptr();
7453 let wv = vdupq_n_f32(w);
7454 let mut j = 0usize;
7455 while j + 16 <= n {
7456 let rb = vld1q_s8(rp.add(j));
7457 let lo = vmovl_s8(vget_low_s8(rb));
7458 let hi = vmovl_s8(vget_high_s8(rb));
7459 for (off, half) in [(0, lo), (8, hi)] {
7460 let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
7461 let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
7462 let o = j + off;
7463 vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
7464 vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
7465 }
7466 j += 16;
7467 }
7468 while j < n {
7469 *ap.add(j) += w * (*rp.add(j)) as f32;
7470 j += 1;
7471 }
7472 }
7473}
7474
7475#[inline]
7478pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
7479 #[cfg(target_arch = "aarch64")]
7480 unsafe {
7481 return dot_i8_f32_neon(w, x);
7482 }
7483 #[cfg(target_arch = "x86_64")]
7484 if avx2_enabled() {
7485 return unsafe { dot_i8_f32_avx2(w, x) };
7486 }
7487 #[allow(unreachable_code)]
7488 {
7489 let mut sum = 0.0f32;
7490 for (j, &b) in w.iter().enumerate() {
7491 sum += (b as i8) as f32 * x[j];
7492 }
7493 sum
7494 }
7495}
7496
7497#[inline]
7501fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
7502 #[cfg(target_arch = "aarch64")]
7503 unsafe {
7504 return dot_i8_col_f32_neon(w, x, col);
7505 }
7506 #[allow(unreachable_code)]
7507 {
7508 let mut sum = 0.0f32;
7509 for (j, &b) in w.iter().enumerate() {
7510 sum += (b as i8) as f32 * x[j] * col[j];
7511 }
7512 sum
7513 }
7514}
7515
7516#[cfg(target_arch = "aarch64")]
7517#[target_feature(enable = "neon")]
7518unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
7519 unsafe {
7521 use core::arch::aarch64::*;
7522 let n = x.len();
7523 let wp = w.as_ptr() as *const i8;
7524 let xp = x.as_ptr();
7525 let cp = col.as_ptr();
7526 let (mut a0, mut a1, mut a2, mut a3) = (
7527 vdupq_n_f32(0.0),
7528 vdupq_n_f32(0.0),
7529 vdupq_n_f32(0.0),
7530 vdupq_n_f32(0.0),
7531 );
7532 let mut j = 0usize;
7533 while j + 16 <= n {
7534 let wb = vld1q_s8(wp.add(j));
7535 let lo = vmovl_s8(vget_low_s8(wb));
7536 let hi = vmovl_s8(vget_high_s8(wb));
7537 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
7538 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
7539 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
7540 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
7541 a0 = vfmaq_f32(
7542 a0,
7543 w0,
7544 vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
7545 );
7546 a1 = vfmaq_f32(
7547 a1,
7548 w1,
7549 vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
7550 );
7551 a2 = vfmaq_f32(
7552 a2,
7553 w2,
7554 vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
7555 );
7556 a3 = vfmaq_f32(
7557 a3,
7558 w3,
7559 vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
7560 );
7561 j += 16;
7562 }
7563 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
7564 while j < n {
7565 sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
7566 j += 1;
7567 }
7568 sum
7569 }
7570}
7571
7572#[cfg(target_arch = "aarch64")]
7573#[target_feature(enable = "neon")]
7574unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
7575 unsafe {
7577 use core::arch::aarch64::*;
7578 let n = x.len();
7579 let wp = w.as_ptr() as *const i8;
7580 let xp = x.as_ptr();
7581 let (mut a0, mut a1, mut a2, mut a3) = (
7582 vdupq_n_f32(0.0),
7583 vdupq_n_f32(0.0),
7584 vdupq_n_f32(0.0),
7585 vdupq_n_f32(0.0),
7586 );
7587 let mut j = 0usize;
7588 while j + 16 <= n {
7589 let wb = vld1q_s8(wp.add(j));
7590 let lo = vmovl_s8(vget_low_s8(wb));
7591 let hi = vmovl_s8(vget_high_s8(wb));
7592 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
7593 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
7594 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
7595 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
7596 a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
7597 a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
7598 a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
7599 a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
7600 j += 16;
7601 }
7602 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
7603 while j < n {
7604 sum += (*wp.add(j)) as f32 * *xp.add(j);
7605 j += 1;
7606 }
7607 sum
7608 }
7609}
7610
7611#[allow(clippy::too_many_arguments)]
7612fn qmatvec(
7613 q: &[u8],
7614 rep: &[u8],
7615 row_scale: &[f32],
7616 x: &[f32],
7617 col_field: &[f32],
7618 dtype: TensorDtype,
7619 rows: usize,
7620 cols: usize,
7621 out: &mut [f32],
7622 pool: Option<&Pool>,
7623) {
7624 debug_assert_eq!(out.len(), rows);
7625 #[cfg(not(target_arch = "aarch64"))]
7626 let _ = rep;
7627
7628 #[cfg(target_arch = "aarch64")]
7629 if sdot_enabled() {
7630 let act = if dtype == TensorDtype::Q8_2f {
7631 split_act_q8_2f(x, col_field)
7632 } else {
7633 split_act(x)
7634 };
7635 let out_addr = SendMut(out.as_mut_ptr());
7636 let run_range = |start: usize, end: usize| {
7637 q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
7638 };
7639 match pool {
7640 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
7641 _ => run_range(0, rows),
7642 }
7643 return;
7644 }
7645 #[cfg(target_arch = "x86_64")]
7648 if avx2_a8w8_enabled() {
7649 let act = if dtype == TensorDtype::Q8_2f {
7650 split_act_q8_2f(x, col_field)
7651 } else {
7652 split_act(x)
7653 };
7654 let out_addr = SendMut(out.as_mut_ptr());
7655 let run_range = |start: usize, end: usize| {
7656 q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
7657 };
7658 match pool {
7659 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
7660 _ => run_range(0, rows),
7661 }
7662 return;
7663 }
7664
7665 prescale_with(x, col_field, dtype, 1, |xs| {
7666 let out_addr = SendMut(out.as_mut_ptr());
7667 let run_range = move |start: usize, end: usize| {
7668 for o in start..end {
7669 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
7670 unsafe { *out_addr.at(o) = v };
7672 }
7673 };
7674 match pool {
7675 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
7676 _ => run_range(0, rows),
7677 }
7678 });
7679}
7680
7681#[allow(clippy::too_many_arguments)]
7682fn qmatvec2(
7683 q: &[u8],
7684 row_scale: &[f32],
7685 x1: &[f32],
7686 x2: &[f32],
7687 col_field: &[f32],
7688 dtype: TensorDtype,
7689 rows: usize,
7690 cols: usize,
7691 o1: &mut [f32],
7692 o2: &mut [f32],
7693 pool: Option<&Pool>,
7694) {
7695 #[cfg(target_arch = "aarch64")]
7696 if sdot_enabled() {
7697 let a1s = if dtype == TensorDtype::Q8_2f {
7698 split_act_q8_2f(x1, col_field)
7699 } else {
7700 split_act(x1)
7701 };
7702 let a2s = if dtype == TensorDtype::Q8_2f {
7703 split_act_q8_2f(x2, col_field)
7704 } else {
7705 split_act(x2)
7706 };
7707 let p1 = SendMut(o1.as_mut_ptr());
7708 let p2 = SendMut(o2.as_mut_ptr());
7709 let run_range = |start: usize, end: usize| {
7710 q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
7711 };
7712 match pool {
7713 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
7714 _ => run_range(0, rows),
7715 }
7716 return;
7717 }
7718 #[cfg(target_arch = "x86_64")]
7719 if avx2_a8w8_enabled() {
7720 let a1s = if dtype == TensorDtype::Q8_2f {
7721 split_act_q8_2f(x1, col_field)
7722 } else {
7723 split_act(x1)
7724 };
7725 let a2s = if dtype == TensorDtype::Q8_2f {
7726 split_act_q8_2f(x2, col_field)
7727 } else {
7728 split_act(x2)
7729 };
7730 let p1 = SendMut(o1.as_mut_ptr());
7731 let p2 = SendMut(o2.as_mut_ptr());
7732 let run_range = |start: usize, end: usize| {
7733 q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
7734 };
7735 match pool {
7736 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
7737 _ => run_range(0, rows),
7738 }
7739 return;
7740 }
7741
7742 prescale_with(x1, col_field, dtype, 1, |x1s| {
7743 prescale_with(x2, col_field, dtype, 2, |x2s| {
7744 let p1 = SendMut(o1.as_mut_ptr());
7745 let p2 = SendMut(o2.as_mut_ptr());
7746 let run_range = move |start: usize, end: usize| {
7747 for o in start..end {
7748 let row = &q[o * cols..(o + 1) * cols];
7749 let s1 = dot_i8_f32(row, x1s) * row_scale[o];
7750 let s2 = dot_i8_f32(row, x2s) * row_scale[o];
7751 unsafe {
7753 *p1.at(o) = s1;
7754 *p2.at(o) = s2;
7755 }
7756 }
7757 };
7758 match pool {
7759 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
7760 _ => run_range(0, rows),
7761 }
7762 });
7763 });
7764}
7765
7766#[derive(Clone, Copy)]
7767struct SendMut(*mut f32);
7768unsafe impl Send for SendMut {}
7769unsafe impl Sync for SendMut {}
7770
7771impl SendMut {
7772 #[inline]
7773 fn at(self, i: usize) -> *mut f32 {
7774 unsafe { self.0.add(i) }
7775 }
7776}
7777
7778#[cfg(test)]
7779mod tests {
7780 use super::*;
7781
7782 #[test]
7783 fn f32_matvec_matches_matvec_rows_bitexact() {
7784 let (rows, cols) = (300, 40);
7785 let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
7786 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
7787 let qt = QTensor::from_f32(w.clone(), rows, cols);
7788
7789 let mut a = vec![0.0f32; rows];
7790 matvec_rows(None, &w, &x, &mut a);
7791 let mut b = vec![0.0f32; rows];
7792 qt.matvec(&x, &mut b, None);
7793 assert_eq!(a, b);
7794 }
7795
7796 #[test]
7797 fn sdot_kernel_exact_on_grid() {
7798 eprintln!("sdot_enabled = {}", sdot_enabled());
7803 let (rows, cols) = (9, 80); let w: Vec<u8> = (0..rows * cols)
7805 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
7806 .collect();
7807 let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
7808 let x: Vec<f32> = (0..cols)
7809 .map(|i| match i % 3 {
7810 0 => 1.0,
7811 1 => -1.0,
7812 _ => 0.0,
7813 })
7814 .collect();
7815 let mut a = vec![0.0f32; rows];
7816 qmatvec(
7817 &w,
7818 &[],
7819 &scales,
7820 &x,
7821 &[],
7822 TensorDtype::Q8Row,
7823 rows,
7824 cols,
7825 &mut a,
7826 None,
7827 );
7828 for o in 0..rows {
7829 let mut acc = 0.0f32;
7830 for j in 0..cols {
7831 acc += (w[o * cols + j] as i8) as f32 * x[j];
7832 }
7833 let expect = acc * scales[o];
7834 assert!(
7835 (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
7836 "row {o}: {} vs {expect}",
7837 a[o]
7838 );
7839 }
7840 }
7841
7842 #[test]
7843 fn q1_tbl_fast_path_matches_reference() {
7844 let (rows, cols) = (5, 256);
7849 let gpr = cols / GROUP_SIZE;
7850 let mut bytes = Vec::new();
7851 for t in 0..rows * gpr {
7852 let s = 0.007 + (t % 11) as f32 * 0.004;
7853 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
7854 for j in 0..4 {
7855 bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
7856 }
7857 }
7858 let x: Vec<f32> = (0..cols)
7859 .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
7860 .collect();
7861 let mut w = vec![0.0f32; rows * cols];
7862 cortiq_core::quant::dequant_q1(&bytes, &mut w);
7863 let mut got = vec![0.0f32; rows];
7864 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
7865 for o in 0..rows {
7866 let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
7867 assert!(
7868 (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
7869 "row {o}: {} vs {expect}",
7870 got[o]
7871 );
7872 }
7873 let b = 5usize;
7876 let mut xs_all = Vec::new();
7877 for bi in 0..b {
7878 xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
7879 }
7880 let mut mm = vec![0.0f32; b * rows];
7881 q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
7882 for bi in 0..b {
7883 let mut single = vec![0.0f32; rows];
7884 q1_matvec(
7885 &bytes,
7886 &xs_all[bi * cols..(bi + 1) * cols],
7887 rows,
7888 cols,
7889 &mut single,
7890 None,
7891 );
7892 assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
7893 }
7894 }
7895
7896 #[test]
7897 fn q1_kernels_match_exact_reference() {
7898 let (rows, cols) = (7, 96);
7900 let gpr = cols / GROUP_SIZE;
7901 let mut bytes = Vec::new();
7902 for t in 0..rows * gpr {
7903 let s = 0.01 + (t % 13) as f32 * 0.003;
7904 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
7905 for j in 0..4 {
7906 bytes.push(((t * 31 + j * 97) % 251) as u8);
7907 }
7908 }
7909 let x: Vec<f32> = (0..cols)
7911 .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
7912 .collect();
7913 let mut w = vec![0.0f32; rows * cols];
7915 cortiq_core::quant::dequant_q1(&bytes, &mut w);
7916 let mut expect = vec![0.0f32; rows];
7917 for o in 0..rows {
7918 expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
7919 }
7920 let mut got = vec![0.0f32; rows];
7921 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
7922 for o in 0..rows {
7923 assert!(
7924 (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
7925 "row {o}: {} vs {}",
7926 got[o],
7927 expect[o]
7928 );
7929 }
7930 let x2: Vec<f32> = x.iter().map(|v| -v).collect();
7932 let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
7933 q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
7934 assert_eq!(a1, got);
7935 let mut xs = x.clone();
7936 xs.extend_from_slice(&x2);
7937 let mut mm = vec![0.0f32; 2 * rows];
7938 q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
7939 assert_eq!(&mm[..rows], got.as_slice());
7940 assert_eq!(&mm[rows..], a2.as_slice());
7941 }
7942
7943 #[test]
7944 fn repack_is_bit_identical() {
7945 let (rows, cols) = (267, 96); let w: Vec<u8> = (0..rows * cols)
7951 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
7952 .collect();
7953 let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
7954 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
7955 let rep = q8_repack_layout(&w, rows, cols);
7956 for g in 0..rows / 4 {
7958 for c in 0..cols / 16 {
7959 for lane in 0..4 {
7960 assert_eq!(
7961 &rep[g * 4 * cols + c * 64 + lane * 16
7962 ..g * 4 * cols + c * 64 + lane * 16 + 16],
7963 &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
7964 );
7965 }
7966 }
7967 }
7968 let mut a = vec![0.0f32; rows];
7969 qmatvec(
7970 &w,
7971 &[],
7972 &scales,
7973 &x,
7974 &[],
7975 TensorDtype::Q8Row,
7976 rows,
7977 cols,
7978 &mut a,
7979 None,
7980 );
7981 let mut b = vec![0.0f32; rows];
7982 qmatvec(
7983 &w,
7984 &rep,
7985 &scales,
7986 &x,
7987 &[],
7988 TensorDtype::Q8Row,
7989 rows,
7990 cols,
7991 &mut b,
7992 None,
7993 );
7994 assert_eq!(a, b, "full-range repack output diverged");
7995
7996 #[cfg(target_arch = "aarch64")]
7997 if sdot_enabled() {
7998 let act = split_act(&x);
8000 let mut c1 = vec![0.0f32; rows];
8001 let mut c2 = vec![0.0f32; rows];
8002 q8_range_sdot(
8003 &w,
8004 &[],
8005 &scales,
8006 &act,
8007 cols,
8008 SendMut(c1.as_mut_ptr()),
8009 3,
8010 rows - 2,
8011 );
8012 q8_range_sdot(
8013 &w,
8014 &rep,
8015 &scales,
8016 &act,
8017 cols,
8018 SendMut(c2.as_mut_ptr()),
8019 3,
8020 rows - 2,
8021 );
8022 assert_eq!(c1, c2, "unaligned-range repack output diverged");
8023 }
8024 }
8025
8026 #[test]
8027 fn sdot_a8w8_noise_is_bounded() {
8028 let (rows, cols) = (16, 512);
8032 let w: Vec<u8> = (0..rows * cols)
8033 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
8034 .collect();
8035 let scales = vec![0.01f32; rows];
8036 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
8037 let mut a = vec![0.0f32; rows];
8038 qmatvec(
8039 &w,
8040 &[],
8041 &scales,
8042 &x,
8043 &[],
8044 TensorDtype::Q8Row,
8045 rows,
8046 cols,
8047 &mut a,
8048 None,
8049 );
8050 let (mut num, mut den) = (0f64, 0f64);
8051 for o in 0..rows {
8052 let mut acc = 0.0f32;
8053 for j in 0..cols {
8054 acc += (w[o * cols + j] as i8) as f32 * x[j];
8055 }
8056 let expect = acc * scales[o];
8057 num += ((a[o] - expect) as f64).powi(2);
8058 den += (expect as f64).powi(2);
8059 }
8060 let rel = (num / den.max(1e-12)).sqrt();
8061 assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
8062 }
8063
8064 #[test]
8065 fn i8_dot_neon_matches_scalar() {
8066 let n = 100;
8067 let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
8068 let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
8069 let mut scalar = 0.0f32;
8070 for j in 0..n {
8071 scalar += (w[j] as i8) as f32 * x[j];
8072 }
8073 let fast = dot_i8_f32(&w, &x);
8074 assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
8075 }
8076
8077 #[test]
8079 fn vbitmatvec_matches_full_dequant() {
8080 let (rows, cols) = (6, 64);
8081 let ng = cols / GROUP_SIZE;
8082 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
8084 let mut bytes = bits.clone();
8085 for g in 0..rows * ng {
8086 let s = 0.02 + 0.001 * g as f32;
8087 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
8088 }
8089 for r in 0..rows {
8090 let b = bits[r] as usize;
8091 let (mut acc, mut nb) = (0u64, 0usize);
8092 let mut rowbytes = Vec::new();
8093 for i in 0..cols {
8094 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
8095 acc = (acc << b) | v;
8096 nb += b;
8097 while nb >= 8 {
8098 nb -= 8;
8099 rowbytes.push(((acc >> nb) & 0xFF) as u8);
8100 }
8101 }
8102 if nb > 0 {
8103 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
8104 }
8105 bytes.extend_from_slice(&rowbytes);
8106 }
8107 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
8108
8109 let mut reference = vec![0f32; rows * cols];
8110 cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
8111 let mut expect = vec![0f32; rows];
8112 for r in 0..rows {
8113 expect[r] = reference[r * cols..(r + 1) * cols]
8114 .iter()
8115 .zip(&x)
8116 .map(|(w, xv)| w * xv)
8117 .sum();
8118 }
8119 let mut got = vec![0f32; rows];
8120 let offsets = vbit_row_offsets(&bytes, rows, cols);
8121 vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
8122 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
8126 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
8127 for r in 0..rows {
8128 assert!(
8129 (got[r] - expect[r]).abs() < tol * scale,
8130 "row {r}: {} vs {}",
8131 got[r],
8132 expect[r]
8133 );
8134 }
8135 }
8136
8137 #[test]
8142 #[cfg(target_arch = "x86_64")]
8143 fn vbit_matmat_blocked_matches_per_row() {
8144 let (rows, cols, b) = (64usize, 128usize, 9usize);
8145 let ng = cols / GROUP_SIZE;
8146 let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
8147 let mut bytes = bits.clone();
8148 for g in 0..rows * ng {
8149 let sc = 0.02 + 0.0005 * g as f32;
8150 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
8151 }
8152 for r in 0..rows {
8153 let bw = bits[r] as usize;
8154 let (mut acc, mut nb) = (0u64, 0usize);
8155 let mut rowbytes = Vec::new();
8156 for i in 0..cols {
8157 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
8158 acc = (acc << bw) | v;
8159 nb += bw;
8160 while nb >= 8 {
8161 nb -= 8;
8162 rowbytes.push(((acc >> nb) & 0xFF) as u8);
8163 }
8164 }
8165 if nb > 0 {
8166 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
8167 }
8168 bytes.extend_from_slice(&rowbytes);
8169 }
8170 let x: Vec<f32> = (0..b * cols)
8171 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
8172 .collect();
8173 let offsets = vbit_row_offsets(&bytes, rows, cols);
8174 let mut y_a = vec![0f32; b * rows];
8175 let mut y_b = vec![0f32; b * rows];
8176 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
8177 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
8178 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
8179 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
8180 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
8181 let max_d = y_a
8182 .iter()
8183 .zip(&y_b)
8184 .map(|(p, q)| (p - q).abs())
8185 .fold(0.0f32, f32::max);
8186 assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
8187 }
8188
8189 #[test]
8197 fn q4t_matmat_blocked_matches_per_row() {
8198 let (rows, cols, b) = (16usize, 64usize, 9usize);
8199 let gpr = cols / GROUP_SIZE;
8200 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
8201 for r in 0..rows {
8202 for g in 0..gpr {
8203 let t = (r * gpr + g) * Q4_TILE;
8204 let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
8205 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
8206 for k in 0..16 {
8207 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
8208 }
8209 }
8210 }
8211 let x: Vec<f32> = (0..b * cols)
8212 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
8213 .collect();
8214 let mut y_blk = vec![0f32; b * rows];
8215 let mut y_row = vec![0f32; b * rows];
8216 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
8217 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
8218 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
8219 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
8220 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
8221 assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
8222 }
8223
8224 #[cfg(target_os = "macos")]
8228 #[test]
8229 fn q4t_matmat_accel_matches_dequant_reference() {
8230 if !accel_gemm_enabled() {
8231 return; }
8233 let (rows, cols, b) = (512usize, 1024usize, 8usize); let gpr = cols / GROUP_SIZE;
8235 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
8236 for r in 0..rows {
8237 for g in 0..gpr {
8238 let t = (r * gpr + g) * Q4_TILE;
8239 let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
8240 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
8241 for k in 0..16 {
8242 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
8243 }
8244 }
8245 }
8246 let x: Vec<f32> = (0..b * cols)
8247 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
8248 .collect();
8249 let mut got = vec![0f32; b * rows];
8250 q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
8251 let mut w = vec![0f32; rows * cols];
8253 for r in 0..rows {
8254 for g in 0..gpr {
8255 let t = (r * gpr + g) * Q4_TILE;
8256 let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
8257 for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
8258 w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
8259 w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
8260 (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
8261 }
8262 }
8263 }
8264 for bi in 0..b {
8265 for r in 0..rows {
8266 let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
8267 let d = (got[bi * rows + r] - want).abs();
8268 assert!(
8269 d <= want.abs().max(1.0) * 1e-4,
8270 "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
8271 got[bi * rows + r]
8272 );
8273 }
8274 }
8275 }
8276
8277 #[test]
8278 fn q4matvec_matches_full_dequant() {
8279 let (rows, cols) = (8, 64);
8280 let groups = rows * cols / GROUP_SIZE;
8281 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
8283 for i in 0..groups * 16 {
8284 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
8285 }
8286 for g in 0..groups {
8287 let s = 0.01 + 0.003 * g as f32;
8288 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
8289 }
8290 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
8291
8292 let mut reference = vec![0.0f32; rows * cols];
8293 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
8294 let mut expect = vec![0.0f32; rows];
8295 for r in 0..rows {
8296 expect[r] = reference[r * cols..(r + 1) * cols]
8297 .iter()
8298 .zip(&x)
8299 .map(|(w, xv)| w * xv)
8300 .sum();
8301 }
8302
8303 let mut got = vec![0.0f32; rows];
8304 q4matvec(&bytes, &x, rows, cols, &mut got, None);
8305 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
8309 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
8310 for r in 0..rows {
8311 assert!(
8312 (got[r] - expect[r]).abs() < tol * scale,
8313 "row {r}: {} vs {}",
8314 got[r],
8315 expect[r]
8316 );
8317 }
8318 }
8319
8320 #[test]
8323 fn vbitmatvec2_equals_two_singles() {
8324 let (rows, cols) = (6, 64);
8325 let ng = cols / GROUP_SIZE;
8326 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
8327 let mut bytes = bits.clone();
8328 for g in 0..rows * ng {
8329 let s = 0.02 + 0.001 * g as f32;
8330 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
8331 }
8332 for r in 0..rows {
8333 let b = bits[r] as usize;
8334 let (mut acc, mut nb) = (0u64, 0usize);
8335 let mut rowbytes = Vec::new();
8336 for i in 0..cols {
8337 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
8338 acc = (acc << b) | v;
8339 nb += b;
8340 while nb >= 8 {
8341 nb -= 8;
8342 rowbytes.push(((acc >> nb) & 0xFF) as u8);
8343 }
8344 }
8345 if nb > 0 {
8346 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
8347 }
8348 bytes.extend_from_slice(&rowbytes);
8349 }
8350 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
8351 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
8352 let offsets = vbit_row_offsets(&bytes, rows, cols);
8353
8354 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
8355 vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
8356 vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
8357 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
8358 vbitmatvec2(
8359 &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
8360 );
8361 assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
8362 assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
8363 }
8364
8365 #[test]
8367 fn q4matvec2_equals_two_singles() {
8368 let (rows, cols) = (8, 128);
8369 let groups = rows * cols / GROUP_SIZE;
8370 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
8371 for i in 0..groups * 16 {
8372 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
8373 }
8374 for g in 0..groups {
8375 let s = 0.01 + 0.003 * g as f32;
8376 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
8377 }
8378 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
8381 x1[9] = 250.0;
8382 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
8383
8384 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
8385 q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
8386 q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
8387 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
8388 q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
8389 assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
8390 assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
8391 }
8392
8393 #[test]
8396 fn matvec_many_equals_separate_matvecs() {
8397 use crate::pool::Pool;
8398 let (r1, r2, cols) = (300, 200, 64);
8399 let mk = |salt: usize, rows: usize| {
8400 QTensor::from_f32(
8401 (0..rows * cols)
8402 .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
8403 .collect(),
8404 rows,
8405 cols,
8406 )
8407 };
8408 let (a, b) = (mk(1, r1), mk(5, r2));
8409 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
8410 let pool = Pool::new(3);
8411
8412 let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
8413 a.matvec(&x, &mut ea, Some(&pool));
8414 b.matvec(&x, &mut eb, Some(&pool));
8415 let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
8416 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
8417 assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
8418 assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
8419 }
8420
8421 #[test]
8424 fn batched_matmat_equals_per_position_matvec() {
8425 let (rows, cols, b) = (8, 64, 5);
8426 let groups = rows * cols / GROUP_SIZE;
8428 let mut q4 = Vec::new();
8429 for i in 0..groups * 16 {
8430 q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
8431 }
8432 for g in 0..groups {
8433 q4.extend_from_slice(
8434 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
8435 );
8436 }
8437 let ng = cols / GROUP_SIZE;
8439 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
8440 let mut vb = bits.clone();
8441 for g in 0..rows * ng {
8442 vb.extend_from_slice(
8443 &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
8444 );
8445 }
8446 for r in 0..rows {
8447 let bw = bits[r] as usize;
8448 let (mut acc, mut nb) = (0u64, 0usize);
8449 let mut rowbytes = Vec::new();
8450 for i in 0..cols {
8451 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
8452 acc = (acc << bw) | v;
8453 nb += bw;
8454 while nb >= 8 {
8455 nb -= 8;
8456 rowbytes.push(((acc >> nb) & 0xFF) as u8);
8457 }
8458 }
8459 if nb > 0 {
8460 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
8461 }
8462 vb.extend_from_slice(&rowbytes);
8463 }
8464 let offsets = vbit_row_offsets(&vb, rows, cols);
8465
8466 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
8467
8468 let mut got = vec![0f32; b * rows];
8470 q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
8471 for bi in 0..b {
8472 let mut expect = vec![0f32; rows];
8473 q4matvec(
8474 &q4,
8475 &xs[bi * cols..(bi + 1) * cols],
8476 rows,
8477 cols,
8478 &mut expect,
8479 None,
8480 );
8481 assert_eq!(
8482 &got[bi * rows..(bi + 1) * rows],
8483 &expect[..],
8484 "q4 batch pos {bi}"
8485 );
8486 }
8487
8488 let mut got = vec![0f32; b * rows];
8490 vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
8491 for bi in 0..b {
8492 let mut expect = vec![0f32; rows];
8493 vbitmatvec(
8494 &vb,
8495 &offsets,
8496 &xs[bi * cols..(bi + 1) * cols],
8497 rows,
8498 cols,
8499 &mut expect,
8500 None,
8501 );
8502 assert_eq!(
8503 &got[bi * rows..(bi + 1) * rows],
8504 &expect[..],
8505 "vbit batch pos {bi}"
8506 );
8507 }
8508 }
8509
8510 #[test]
8514 fn q4_tiled_matches_q4_block_bitexact() {
8515 let (rows, cols, b) = (8usize, 128usize, 3usize);
8516 let groups = rows * cols / GROUP_SIZE;
8517 let mut split = Vec::with_capacity(groups * 18);
8518 for i in 0..groups * 16 {
8519 split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
8520 }
8521 for g in 0..groups {
8522 split.extend_from_slice(
8523 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
8524 );
8525 }
8526 let (packed, scales) = split.split_at(groups * 16);
8528 let mut tiled = Vec::with_capacity(groups * Q4_TILE);
8529 for g in 0..groups {
8530 tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
8531 tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
8532 }
8533
8534 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
8535 x1[9] = 250.0; let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
8537
8538 let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
8539 q4matvec(&split, &x1, rows, cols, &mut a, None);
8540 q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
8541 assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
8542
8543 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
8544 let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
8545 q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
8546 q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
8547 assert_eq!(a1, t1);
8548 assert_eq!(a2, t2);
8549
8550 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
8551 let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
8552 q4matmat(&split, &xs, b, rows, cols, &mut am, None);
8553 q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
8554 assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
8555 }
8556
8557 #[test]
8564 fn q4matvec_sdot_outlier_exact() {
8565 let (rows, cols) = (4, 128);
8566 let groups = rows * cols / GROUP_SIZE;
8567 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
8568 for i in 0..groups * 16 {
8569 bytes.push(((i * 11 + 5) % 256) as u8);
8570 }
8571 for g in 0..groups {
8572 let s = 0.02 + 0.002 * g as f32;
8573 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
8574 }
8575 let mut x: Vec<f32> = (0..cols)
8576 .map(|i| match i % 3 {
8577 0 => 1.0,
8578 1 => -1.0,
8579 _ => 0.0,
8580 })
8581 .collect();
8582 x[17] = 300.0; let mut reference = vec![0.0f32; rows * cols];
8585 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
8586 let mut expect = vec![0.0f32; rows];
8587 for r in 0..rows {
8588 expect[r] = reference[r * cols..(r + 1) * cols]
8589 .iter()
8590 .zip(&x)
8591 .map(|(w, xv)| w * xv)
8592 .sum();
8593 }
8594 let mut got = vec![0.0f32; rows];
8595 q4matvec(&bytes, &x, rows, cols, &mut got, None);
8596 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
8597 for r in 0..rows {
8598 assert!(
8599 (got[r] - expect[r]).abs() < 2e-3 * scale,
8600 "row {r}: {} vs {} (outlier term must be exact)",
8601 got[r],
8602 expect[r]
8603 );
8604 }
8605 }
8606
8607 #[test]
8611 fn q1t_matvec_matches_reference() {
8612 use cortiq_core::quant::{dequant_q1t, f32_to_f16};
8613 let (rows, cols) = (3usize, 64usize); let gpr = cols / GROUP_SIZE;
8615 let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
8616 let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
8618 let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
8619 let mut bytes = Vec::new();
8620 for r in 0..rows {
8621 for g in 0..gpr {
8622 bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
8623 let mut c = [0u8; 7];
8624 for k in 0..GROUP_SIZE {
8625 let code = if is_out(r * cols + g * GROUP_SIZE + k) {
8627 0
8628 } else {
8629 ((k + r * 3 + g) % 3) as u8 };
8631 cortiq_core::quant::q1t_pack(&mut c, k, code);
8632 }
8633 bytes.extend_from_slice(&c);
8634 }
8635 }
8636 let mut row_ptr = vec![0u32; rows + 1];
8639 for &(idx, _) in &outliers {
8640 row_ptr[idx as usize / cols + 1] += 1;
8641 }
8642 for r in 0..rows {
8643 row_ptr[r + 1] += row_ptr[r];
8644 }
8645 for &p in &row_ptr {
8646 bytes.extend_from_slice(&p.to_le_bytes());
8647 }
8648 for &(idx, v) in &outliers {
8649 bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
8650 bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
8651 }
8652
8653 let mut refw = vec![0f32; rows * cols];
8654 dequant_q1t(&bytes, rows, cols, &mut refw);
8655 let x: Vec<f32> = (0..cols)
8658 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
8659 .collect();
8660 let mut expect = vec![0f32; rows];
8661 for r in 0..rows {
8662 let mut a = 0.0f32;
8663 for j in 0..cols {
8664 a += refw[r * cols + j] * x[j];
8665 }
8666 expect[r] = a;
8667 }
8668 let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
8669 let mut got = vec![0f32; rows];
8670 q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
8671 for r in 0..rows {
8672 assert!(
8673 (got[r] - expect[r]).abs() < tol(expect[r]),
8674 "row {r}: {} vs {}",
8675 got[r],
8676 expect[r]
8677 );
8678 }
8679 let x2: Vec<f32> = x.iter().chain(x.iter().map(|v| v)).copied().collect();
8681 let mut gm = vec![0f32; 2 * rows];
8682 q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
8683 for r in 0..rows {
8684 assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
8685 assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
8686 }
8687 let xb: Vec<f32> = (0..cols)
8691 .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
8692 .collect();
8693 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
8694 q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
8695 q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
8696 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
8697 q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
8698 assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
8699 assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
8700 }
8701
8702 #[test]
8705 fn q1t_matvec2_odd_gpr_matches_singles() {
8706 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
8707 let (rows, cols) = (5usize, 96usize); let gpr = cols / GROUP_SIZE;
8709 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
8710 for r in 0..rows {
8711 for g in 0..gpr {
8712 bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
8713 let mut c = [0u8; 7];
8714 for k in 0..GROUP_SIZE {
8715 q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
8716 }
8717 bytes.extend_from_slice(&c);
8718 }
8719 }
8720 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
8721 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
8722 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
8723 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
8724 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
8725 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
8726 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
8727 assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
8728 assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
8729 }
8730
8731 #[test]
8735 #[ignore]
8736 fn q1t_matvec2_speed() {
8737 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
8738 use std::time::Instant;
8739 let (rows, cols) = (8192usize, 4096usize);
8740 let gpr = cols / GROUP_SIZE;
8741 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
8742 for r in 0..rows {
8743 for g in 0..gpr {
8744 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
8745 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
8746 let mut c = [0u8; 7];
8747 for k in 0..GROUP_SIZE {
8748 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
8749 }
8750 bytes.extend_from_slice(&c);
8751 }
8752 }
8753 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
8754 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
8755 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
8756 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
8757 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
8759 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
8760 let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
8761 for _ in 0..8 {
8762 let t0 = Instant::now();
8763 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
8764 t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
8765 let t1 = Instant::now();
8766 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
8767 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
8768 t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
8769 }
8770 assert_eq!(p1, s1);
8771 assert_eq!(p2, s2);
8772 println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
8773 }
8774
8775 #[test]
8779 #[ignore]
8780 fn q1t_matvec_speed() {
8781 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
8782 use std::time::Instant;
8783 let (rows, cols) = (8192usize, 4096usize); let gpr = cols / GROUP_SIZE;
8785 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
8786 for r in 0..rows {
8787 for g in 0..gpr {
8788 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
8789 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
8790 let mut c = [0u8; 7];
8791 for k in 0..GROUP_SIZE {
8792 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
8793 }
8794 bytes.extend_from_slice(&c);
8795 }
8796 }
8797 let (n, stride) = (rows * cols, 40usize); let mut row_ptr = vec![0u32; rows + 1];
8799 let mut idx = 0usize;
8800 while idx < n {
8801 row_ptr[idx / cols + 1] += 1;
8802 idx += stride;
8803 }
8804 for r in 0..rows {
8805 row_ptr[r + 1] += row_ptr[r];
8806 }
8807 for &p in &row_ptr {
8808 bytes.extend_from_slice(&p.to_le_bytes());
8809 }
8810 let mut idx = 0usize;
8811 while idx < n {
8812 bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
8813 bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
8814 idx += stride;
8815 }
8816 let x: Vec<f32> = (0..cols)
8819 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
8820 .collect();
8821 let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
8822
8823 let slow = |out: &mut [f32]| {
8825 let mut buf = vec![0f32; cols];
8826 for r in 0..rows {
8827 for g in 0..gpr {
8828 let off = (r * gpr + g) * Q1T_TILE;
8829 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8830 let codes = &bytes[off + 2..off + Q1T_TILE];
8831 for k in 0..GROUP_SIZE {
8832 buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
8833 1 => s,
8834 2 => -s,
8835 _ => 0.0,
8836 };
8837 }
8838 }
8839 out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
8840 + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
8841 }
8842 };
8843 let iters = 5;
8844 let mut a = vec![0f32; rows];
8845 slow(&mut a); let t = Instant::now();
8847 for _ in 0..iters {
8848 slow(&mut a);
8849 }
8850 let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
8851
8852 let mut b = vec![0f32; rows];
8853 q1t_matvec(&bytes, &x, rows, cols, &mut b, None); let t = Instant::now();
8855 for _ in 0..iters {
8856 q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
8857 }
8858 let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
8859
8860 for r in 0..rows {
8861 assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
8862 }
8863 println!(
8864 "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms fused-LUT {fast_ms:.2} ms => {:.2}x",
8865 slow_ms / fast_ms
8866 );
8867 }
8868}