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")]
3171unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
3172 unsafe {
3174 use core::arch::x86_64::*;
3175 let lomask = _mm_set1_epi8(0x0F);
3176 let eight = _mm256_set1_epi8(8);
3177 let ones = _mm256_set1_epi16(1);
3178 let mut acc = [0f32; 4];
3179 for gi in 0..gpr {
3180 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3181 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3182 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
3183 let lo = _mm_and_si128(bb, lomask);
3184 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
3185 let w = _mm256_sub_epi8(
3186 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3187 eight,
3188 );
3189 let aw = _mm256_abs_epi8(w);
3190 for (k, xq) in xs.iter().enumerate() {
3191 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3192 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
3193 let d = _mm256_madd_epi16(p16, ones);
3194 let hi128 = _mm256_extracti128_si256::<1>(d);
3195 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3196 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3197 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3198 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
3199 }
3200 }
3201 acc
3202 }
3203}
3204
3205#[cfg(target_arch = "x86_64")]
3207#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
3208unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
3209 unsafe {
3211 use core::arch::x86_64::*;
3212 let lomask = _mm_set1_epi8(0x0F);
3213 let eight = _mm256_set1_epi8(8);
3214 let mut acc = [0f32; 4];
3215 for gi in 0..gpr {
3216 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3217 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3218 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
3219 let lo = _mm_and_si128(bb, lomask);
3220 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
3221 let w = _mm256_sub_epi8(
3222 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3223 eight,
3224 );
3225 let aw = _mm256_abs_epi8(w);
3226 for (k, xq) in xs.iter().enumerate() {
3227 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3228 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
3229 acc[k] += d as f32 * s;
3230 }
3231 }
3232 acc
3233 }
3234}
3235
3236#[cfg(target_arch = "aarch64")]
3241#[target_feature(enable = "neon,dotprod")]
3242unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
3243 unsafe {
3245 use core::arch::aarch64::*;
3246 use core::arch::asm;
3247 let lomask = vdupq_n_u8(0x0F);
3248 let eight = vdupq_n_s8(8);
3249 let mut acc = [0f32; 4];
3250 for gi in 0..gpr {
3251 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3252 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3253 let b = vld1q_u8(t.add(2));
3254 let lo = vandq_u8(b, lomask);
3255 let hi = vshrq_n_u8::<4>(b);
3256 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
3257 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
3258 for (k, xq) in xs.iter().enumerate() {
3259 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
3260 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
3261 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3262 asm!(
3263 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
3264 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
3265 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3266 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
3267 options(pure, nomem, nostack),
3268 );
3269 acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
3270 }
3271 }
3272 acc
3273 }
3274}
3275
3276#[inline]
3278fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
3279 let gi = j / GROUP_SIZE;
3280 let k = j % GROUP_SIZE;
3281 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
3282 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3283 let byte = tile[2 + k / 2];
3284 let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
3285 ((nib as i32 - 8) as f32, s)
3286}
3287
3288#[inline]
3291fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
3292 let mut acc = 0f32;
3293 for gi in 0..gpr {
3294 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
3295 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3296 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
3297 let mut ga = 0f32;
3298 for (k, &b) in tile[2..].iter().enumerate() {
3299 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
3300 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
3301 }
3302 acc += ga * s;
3303 }
3304 acc
3305}
3306
3307fn q4t_matvec(
3309 bytes: &[u8],
3310 x: &[f32],
3311 rows: usize,
3312 cols: usize,
3313 out: &mut [f32],
3314 pool: Option<&Pool>,
3315) {
3316 debug_assert_eq!(out.len(), rows);
3317 let gpr = cols / GROUP_SIZE;
3318 let out_addr = SendMut(out.as_mut_ptr());
3319 if a8w8_enabled() {
3320 let act = split_act(x);
3321 let run = move |start: usize, end: usize| {
3322 for r in start..end {
3323 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
3324 for &(j, xv) in &act.outliers {
3325 let (w, s) = q4t_outlier(bytes, r, gpr, j);
3326 acc += w * s * xv;
3327 }
3328 unsafe { *out_addr.at(r) = acc };
3330 }
3331 };
3332 dispatch_rows(pool, rows, &run);
3333 return;
3334 }
3335 let run = move |start: usize, end: usize| {
3336 for r in start..end {
3337 unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
3339 }
3340 };
3341 dispatch_rows(pool, rows, &run);
3342}
3343
3344#[allow(clippy::too_many_arguments)]
3346fn q4t_matvec2(
3347 bytes: &[u8],
3348 x1: &[f32],
3349 x2: &[f32],
3350 rows: usize,
3351 cols: usize,
3352 o1: &mut [f32],
3353 o2: &mut [f32],
3354 pool: Option<&Pool>,
3355) {
3356 let gpr = cols / GROUP_SIZE;
3357 let p1 = SendMut(o1.as_mut_ptr());
3358 let p2 = SendMut(o2.as_mut_ptr());
3359 if a8w8_enabled() {
3360 let a1 = split_act(x1);
3361 let a2 = split_act(x2);
3362 let run = move |start: usize, end: usize| {
3363 for r in start..end {
3364 let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
3365 let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
3366 for &(j, xv) in &a1.outliers {
3367 let (w, s) = q4t_outlier(bytes, r, gpr, j);
3368 v1 += w * s * xv;
3369 }
3370 for &(j, xv) in &a2.outliers {
3371 let (w, s) = q4t_outlier(bytes, r, gpr, j);
3372 v2 += w * s * xv;
3373 }
3374 unsafe {
3376 *p1.at(r) = v1;
3377 *p2.at(r) = v2;
3378 }
3379 }
3380 };
3381 dispatch_rows(pool, rows, &run);
3382 return;
3383 }
3384 let run = move |start: usize, end: usize| {
3385 for r in start..end {
3386 unsafe {
3388 *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
3389 *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
3390 }
3391 }
3392 };
3393 dispatch_rows(pool, rows, &run);
3394}
3395
3396#[allow(clippy::too_many_arguments)]
3398#[cfg(target_os = "macos")]
3404fn dequant_matmat_accel(
3405 dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
3406 xs_all: &[f32],
3407 b: usize,
3408 rows: usize,
3409 cols: usize,
3410 out: &mut [f32],
3411 pool: Option<&Pool>,
3412) {
3413 const TR: usize = 2048;
3414 thread_local! {
3415 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3416 }
3417 WTILE.with(|wt| {
3418 let mut wtile = wt.borrow_mut();
3419 wtile.resize(TR * cols, 0.0);
3420 let mut r0 = 0usize;
3421 while r0 < rows {
3422 let tr = TR.min(rows - r0);
3423 let wt_addr = SendMut(wtile.as_mut_ptr());
3424 let run = |start: usize, end: usize| {
3425 for r in start..end {
3426 let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
3428 dequant_row(r0 + r, dst);
3429 }
3430 };
3431 dispatch_rows(pool, tr, &run);
3432 unsafe {
3433 accel_blas::cblas_sgemm(
3434 101, 111, 112, b as i32,
3438 tr as i32,
3439 cols as i32,
3440 1.0,
3441 xs_all.as_ptr(),
3442 cols as i32,
3443 wtile.as_ptr(),
3444 cols as i32,
3445 0.0,
3446 out.as_mut_ptr().add(r0),
3447 rows as i32,
3448 );
3449 }
3450 r0 += tr;
3451 }
3452 });
3453}
3454
3455fn q4t_matmat(
3456 bytes: &[u8],
3457 xs_all: &[f32],
3458 b: usize,
3459 rows: usize,
3460 cols: usize,
3461 out: &mut [f32],
3462 pool: Option<&Pool>,
3463) {
3464 debug_assert_eq!(out.len(), b * rows);
3465 let gpr = cols / GROUP_SIZE;
3466 #[cfg(target_os = "macos")]
3470 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
3471 dequant_matmat_accel(
3472 &|r, dst| {
3473 for gi in 0..gpr {
3474 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
3475 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3476 for (k, &bb) in tile[2..].iter().enumerate() {
3477 dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
3478 dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
3479 }
3480 }
3481 },
3482 xs_all,
3483 b,
3484 rows,
3485 cols,
3486 out,
3487 pool,
3488 );
3489 return;
3490 }
3491 let out_addr = SendMut(out.as_mut_ptr());
3492 if a8w8_enabled() {
3493 let acts: Vec<SplitAct> = (0..b)
3494 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
3495 .collect();
3496 let acts = &acts;
3497 #[cfg(target_arch = "x86_64")]
3498 let blocked_ok = avx2_enabled()
3499 && std::env::var("CMF_X86_BLOCKED")
3500 .map(|v| v != "0")
3501 .unwrap_or(true);
3502 #[cfg(target_arch = "aarch64")]
3503 let blocked_ok = sdot_enabled()
3504 && std::env::var("CMF_X86_BLOCKED")
3505 .map(|v| v != "0")
3506 .unwrap_or(true);
3507 #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
3508 let blocked_ok = false;
3509 let run = move |start: usize, end: usize| {
3510 for r in start..end {
3511 let mut bi = 0usize;
3512 #[cfg(target_arch = "aarch64")]
3513 if blocked_ok {
3514 while bi + 4 <= acts.len() {
3515 let xs = [
3516 acts[bi].xq.as_slice(),
3517 acts[bi + 1].xq.as_slice(),
3518 acts[bi + 2].xq.as_slice(),
3519 acts[bi + 3].xq.as_slice(),
3520 ];
3521 let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
3522 for k in 0..4 {
3523 let act = &acts[bi + k];
3524 let mut acc = d[k] * act.sx;
3525 for &(j, xv) in &act.outliers {
3526 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
3527 acc += w * sc * xv;
3528 }
3529 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
3531 }
3532 bi += 4;
3533 }
3534 }
3535 #[cfg(target_arch = "x86_64")]
3536 if blocked_ok {
3537 while bi + 4 <= acts.len() {
3538 let xs = [
3539 acts[bi].xq.as_slice(),
3540 acts[bi + 1].xq.as_slice(),
3541 acts[bi + 2].xq.as_slice(),
3542 acts[bi + 3].xq.as_slice(),
3543 ];
3544 let d = unsafe {
3545 if vnni_tiles_enabled() {
3546 dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
3547 } else {
3548 dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
3549 }
3550 };
3551 for k in 0..4 {
3552 let act = &acts[bi + k];
3553 let mut acc = d[k] * act.sx;
3554 for &(j, xv) in &act.outliers {
3555 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
3556 acc += w * sc * xv;
3557 }
3558 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
3560 }
3561 bi += 4;
3562 }
3563 }
3564 let _ = blocked_ok;
3565 while bi < acts.len() {
3566 let act = &acts[bi];
3567 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
3568 for &(j, xv) in &act.outliers {
3569 let (w, s) = q4t_outlier(bytes, r, gpr, j);
3570 acc += w * s * xv;
3571 }
3572 unsafe { *out_addr.at(bi * rows + r) = acc };
3574 bi += 1;
3575 }
3576 }
3577 };
3578 dispatch_rows(pool, rows, &run);
3579 return;
3580 }
3581 let run = move |start: usize, end: usize| {
3582 for r in start..end {
3583 for bi in 0..b {
3584 let x = &xs_all[bi * cols..(bi + 1) * cols];
3585 unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
3587 }
3588 }
3589 };
3590 dispatch_rows(pool, rows, &run);
3591}
3592
3593fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
3602 (0..gpr)
3603 .map(|gi| {
3604 xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
3605 .iter()
3606 .map(|&v| v as i32)
3607 .sum()
3608 })
3609 .collect()
3610}
3611
3612#[inline]
3616#[allow(unreachable_code)]
3617#[cfg(target_arch = "x86_64")]
3622#[target_feature(enable = "avx2")]
3623unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
3624 unsafe {
3626 use core::arch::x86_64::*;
3627 let expand = _mm256_setr_epi8(
3629 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,
3630 3, 3, 3,
3631 );
3632 let bitsel = _mm256_setr_epi8(
3633 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
3634 -128, 1, 2, 4, 8, 16, 32, 64, -128,
3635 );
3636 let ones8 = _mm256_set1_epi8(1);
3637 let ones16 = _mm256_set1_epi16(1);
3638 let mut acc = 0f32;
3639 for gi in 0..gpr {
3640 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
3641 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3642 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
3643 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
3644 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
3645 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3646 let sel = _mm256_and_si256(x, mask);
3647 let p16 = _mm256_maddubs_epi16(ones8, sel);
3649 let d32 = _mm256_madd_epi16(p16, ones16);
3650 let hi128 = _mm256_extracti128_si256::<1>(d32);
3651 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
3652 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3653 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3654 let msum = _mm_cvtsi128_si32(s32);
3655 let d = 2 * msum - gsum[gi];
3658 acc += d as f32 * s;
3659 }
3660 acc
3661 }
3662}
3663
3664#[cfg(target_arch = "x86_64")]
3667#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
3668unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
3669 unsafe {
3671 use core::arch::x86_64::*;
3672 let expand = _mm256_setr_epi8(
3673 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,
3674 3, 3, 3,
3675 );
3676 let bitsel = _mm256_setr_epi8(
3677 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
3678 -128, 1, 2, 4, 8, 16, 32, 64, -128,
3679 );
3680 let ones8 = _mm256_set1_epi8(1);
3681 let mut acc = 0f32;
3682 for gi in 0..gpr {
3683 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
3684 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3685 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
3686 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
3687 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
3688 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3689 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
3690 let d = 2 * msum - gsum[gi];
3691 acc += d as f32 * s;
3692 }
3693 acc
3694 }
3695}
3696
3697#[cfg(target_arch = "x86_64")]
3699#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
3700unsafe fn dot_q1_row_1x4_vnni(
3701 bytes: &[u8],
3702 r: usize,
3703 gpr: usize,
3704 xs: [&[i8]; 4],
3705 gsums: [&[i32]; 4],
3706) -> [f32; 4] {
3707 unsafe {
3709 use core::arch::x86_64::*;
3710 let expand = _mm256_setr_epi8(
3711 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,
3712 3, 3, 3,
3713 );
3714 let bitsel = _mm256_setr_epi8(
3715 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
3716 -128, 1, 2, 4, 8, 16, 32, 64, -128,
3717 );
3718 let ones8 = _mm256_set1_epi8(1);
3719 let mut acc = [0f32; 4];
3720 for gi in 0..gpr {
3721 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
3722 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3723 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
3724 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
3725 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
3726 for (k, xq) in xs.iter().enumerate() {
3727 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3728 let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
3729 let d = 2 * msum - gsums[k][gi];
3730 acc[k] += d as f32 * s;
3731 }
3732 }
3733 acc
3734 }
3735}
3736
3737#[cfg(target_arch = "x86_64")]
3740#[target_feature(enable = "avx2")]
3741unsafe fn dot_q1_row_1x4_avx2(
3742 bytes: &[u8],
3743 r: usize,
3744 gpr: usize,
3745 xs: [&[i8]; 4],
3746 gsums: [&[i32]; 4],
3747) -> [f32; 4] {
3748 unsafe {
3750 use core::arch::x86_64::*;
3751 let expand = _mm256_setr_epi8(
3752 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,
3753 3, 3, 3,
3754 );
3755 let bitsel = _mm256_setr_epi8(
3756 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
3757 -128, 1, 2, 4, 8, 16, 32, 64, -128,
3758 );
3759 let ones8 = _mm256_set1_epi8(1);
3760 let ones16 = _mm256_set1_epi16(1);
3761 let mut acc = [0f32; 4];
3762 for gi in 0..gpr {
3763 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
3764 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3765 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
3766 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
3767 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
3768 for (k, xq) in xs.iter().enumerate() {
3769 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3770 let sel = _mm256_and_si256(x, mask);
3771 let p16 = _mm256_maddubs_epi16(ones8, sel);
3772 let d32 = _mm256_madd_epi16(p16, ones16);
3773 let hi128 = _mm256_extracti128_si256::<1>(d32);
3774 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
3775 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3776 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3777 let msum = _mm_cvtsi128_si32(s32);
3778 let d = 2 * msum - gsums[k][gi];
3779 acc[k] += d as f32 * s;
3780 }
3781 }
3782 acc
3783 }
3784}
3785
3786#[allow(unreachable_code)]
3787fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
3788 #[cfg(target_arch = "aarch64")]
3789 unsafe {
3790 return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
3791 }
3792 #[cfg(target_arch = "x86_64")]
3793 if avx2_enabled() {
3794 unsafe {
3795 if vnni_tiles_enabled() {
3796 return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
3797 }
3798 return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
3799 }
3800 }
3801 let _ = gsum;
3802 let mut acc = 0f32;
3803 for gi in 0..gpr {
3804 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
3805 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3806 let mut d = 0i32;
3807 for (j, &b) in tile[2..].iter().enumerate() {
3808 for k in 0..8 {
3809 let w = ((b >> k) & 1) as i32 * 2 - 1;
3810 d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
3811 }
3812 }
3813 acc += d as f32 * s;
3814 }
3815 acc
3816}
3817
3818#[cfg(target_arch = "aarch64")]
3827#[target_feature(enable = "neon,dotprod")]
3828unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
3829 unsafe {
3832 use core::arch::aarch64::*;
3833 use core::arch::asm;
3834 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
3835 let m = vld1q_u8(MASKS.as_ptr());
3836 macro_rules! tile_dot {
3838 ($t:expr, $x:expr) => {{
3839 let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
3840 let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
3841 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
3842 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
3843 let x0 = vld1q_s8($x);
3844 let x1 = vld1q_s8($x.add(16));
3845 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3846 asm!(
3847 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
3848 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
3849 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3850 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
3851 options(pure, nomem, nostack),
3852 );
3853 vaddq_s32(a0, a1)
3854 }};
3855 }
3856 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
3865 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
3866 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
3867 const IW11: [u8; 16] = [
3868 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
3869 ];
3870 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
3871 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
3872 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
3873 let isc = vld1_u8(ISC.as_ptr());
3874 macro_rules! tile_dot_tbl {
3876 ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
3877 let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
3878 let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
3879 let x0 = vld1q_s8($x);
3880 let x1 = vld1q_s8($x.add(16));
3881 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3882 asm!(
3883 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
3884 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
3885 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3886 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
3887 options(pure, nomem, nostack),
3888 );
3889 vaddq_s32(a0, a1)
3890 }};
3891 }
3892 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
3893 let row_base = r * gpr * Q1_TILE;
3894 let abs_end = bytes.len();
3895 let xp = xq.as_ptr();
3896 let gp = gsum.as_ptr();
3897 let mut accv = vdupq_n_f32(0.0);
3898 let mut gi = 0;
3899 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
3902 let t0 = base.add(gi * Q1_TILE);
3903 let ld_a = vld1q_u8(t0);
3904 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
3905 let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
3906 let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
3907 let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
3908 let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
3909 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
3911 let g = vld1q_s32(gp.add(gi));
3912 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
3913 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
3914 let scf: float32x4_t;
3915 asm!(
3916 "fcvtl {o:v}.4s, {i:v}.4h",
3917 o = out(vreg) scf, i = in(vreg) sc16,
3918 options(pure, nomem, nostack),
3919 );
3920 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
3921 gi += 4;
3922 }
3923 let mut acc = vaddvq_f32(accv);
3924 while gi < gpr {
3925 let t = base.add(gi * Q1_TILE);
3926 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3927 let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
3928 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
3929 gi += 1;
3930 }
3931 acc
3932 }
3933}
3934
3935#[cfg(target_arch = "aarch64")]
3940#[target_feature(enable = "neon,dotprod")]
3941unsafe fn dot_q1_row_1x4_sdot(
3942 bytes: &[u8],
3943 r: usize,
3944 gpr: usize,
3945 xs: [&[i8]; 4],
3946 gs: [&[i32]; 4],
3947) -> [f32; 4] {
3948 unsafe {
3950 use core::arch::aarch64::*;
3951 use core::arch::asm;
3952 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
3953 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
3954 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
3955 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
3956 const IW11: [u8; 16] = [
3957 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
3958 ];
3959 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
3960 let m = vld1q_u8(MASKS.as_ptr());
3961 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
3962 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
3963 let isc = vld1_u8(ISC.as_ptr());
3964 macro_rules! sdot2 {
3965 ($w0:expr, $w1:expr, $x:expr) => {{
3966 let x0 = vld1q_s8($x);
3967 let x1 = vld1q_s8($x.add(16));
3968 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3969 asm!(
3970 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
3971 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
3972 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3973 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
3974 options(pure, nomem, nostack),
3975 );
3976 vaddq_s32(a0, a1)
3977 }};
3978 }
3979 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
3980 let row_base = r * gpr * Q1_TILE;
3981 let abs_end = bytes.len();
3982 let mut accv = [vdupq_n_f32(0.0); 4];
3983 let mut gi = 0;
3984 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
3985 let t0 = base.add(gi * Q1_TILE);
3986 let ld_a = vld1q_u8(t0);
3987 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
3988 let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
3990 let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
3991 let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
3992 let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
3993 let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
3994 let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
3995 let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
3996 let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
3997 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
3998 let scf: float32x4_t;
3999 asm!(
4000 "fcvtl {o:v}.4s, {i:v}.4h",
4001 o = out(vreg) scf, i = in(vreg) sc16,
4002 options(pure, nomem, nostack),
4003 );
4004 for k in 0..4 {
4005 let xp = xs[k].as_ptr();
4006 let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
4007 let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
4008 let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
4009 let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
4010 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
4011 let g = vld1q_s32(gs[k].as_ptr().add(gi));
4012 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
4013 accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
4014 }
4015 gi += 4;
4016 }
4017 let mut acc = [
4018 vaddvq_f32(accv[0]),
4019 vaddvq_f32(accv[1]),
4020 vaddvq_f32(accv[2]),
4021 vaddvq_f32(accv[3]),
4022 ];
4023 while gi < gpr {
4024 let t = base.add(gi * Q1_TILE);
4025 let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4026 let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
4027 let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
4028 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
4029 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
4030 for k in 0..4 {
4031 let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
4032 acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
4033 }
4034 gi += 1;
4035 }
4036 acc
4037 }
4038}
4039
4040#[inline]
4042fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
4043 let gi = j / GROUP_SIZE;
4044 let k = j % GROUP_SIZE;
4045 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
4046 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4047 let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
4048 ((bit as i32 * 2 - 1) as f32, s)
4049}
4050
4051#[inline]
4053fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
4054 let mut acc = 0f32;
4055 for gi in 0..gpr {
4056 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
4057 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4058 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4059 let mut ga = 0f32;
4060 for (j, &b) in tile[2..].iter().enumerate() {
4061 for k in 0..8 {
4062 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
4063 }
4064 }
4065 acc += ga * s;
4066 }
4067 acc
4068}
4069
4070#[allow(clippy::too_many_arguments)]
4073fn q1_range_a8w8(
4074 bytes: &[u8],
4075 gpr: usize,
4076 act: &SplitAct,
4077 gsum: &[i32],
4078 out: SendMut,
4079 start: usize,
4080 end: usize,
4081) {
4082 for r in start..end {
4083 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
4084 for &(j, xv) in &act.outliers {
4085 let (w, s) = q1_outlier(bytes, r, gpr, j);
4086 acc += w * s * xv;
4087 }
4088 unsafe { *out.at(r) = acc };
4090 }
4091}
4092
4093fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
4095 for r in start..end {
4096 unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
4098 }
4099}
4100
4101fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
4106 let entries = base_len + (rows + 1) * 4;
4107 (base_len, entries, entries <= bytes.len())
4108}
4109
4110#[inline]
4112fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
4113 let o = rp_off + r * 4;
4114 u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
4115}
4116
4117const SIGN5: [[f32; 5]; 256] = {
4121 let mut lut = [[0.0f32; 5]; 256];
4122 let pow3 = [1u16, 3, 9, 27, 81];
4123 let mut byte = 0usize;
4124 while byte < 256 {
4125 let mut i = 0usize;
4126 while i < 5 {
4127 let code = (byte as u16 / pow3[i]) % 3;
4128 lut[byte][i] = if code == 1 {
4129 1.0
4130 } else if code == 2 {
4131 -1.0
4132 } else {
4133 0.0
4134 };
4135 i += 1;
4136 }
4137 byte += 1;
4138 }
4139 lut
4140};
4141
4142const SIGN5_I8: [[i8; 5]; 256] = {
4144 let mut lut = [[0i8; 5]; 256];
4145 let pow3 = [1u16, 3, 9, 27, 81];
4146 let mut byte = 0usize;
4147 while byte < 256 {
4148 let mut i = 0usize;
4149 while i < 5 {
4150 let code = (byte as u16 / pow3[i]) % 3;
4151 lut[byte][i] = if code == 1 {
4152 1
4153 } else if code == 2 {
4154 -1
4155 } else {
4156 0
4157 };
4158 i += 1;
4159 }
4160 byte += 1;
4161 }
4162 lut
4163};
4164
4165const SIGN5_U64: [u64; 256] = {
4171 let mut lut = [0u64; 256];
4172 let pow3 = [1u16, 3, 9, 27, 81];
4173 let mut byte = 0usize;
4174 while byte < 256 {
4175 let mut v = 0u64;
4176 let mut i = 0usize;
4177 while i < 5 {
4178 let code = (byte as u16 / pow3[i]) % 3;
4179 let s: u8 = if code == 1 {
4180 1
4181 } else if code == 2 {
4182 0xFF
4183 } else {
4184 0
4185 };
4186 v |= (s as u64) << (i * 8);
4187 i += 1;
4188 }
4189 lut[byte] = v;
4190 byte += 1;
4191 }
4192 lut
4193};
4194
4195#[inline]
4200fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
4201 const TILE: usize = cortiq_core::quant::Q1T_TILE;
4202 let off = (r * gpr + j / GROUP_SIZE) * TILE;
4203 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
4204 let within = j % GROUP_SIZE;
4205 SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
4206}
4207
4208#[cfg(target_arch = "aarch64")]
4211#[target_feature(enable = "neon,dotprod")]
4212#[inline]
4213unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
4214 unsafe {
4216 use core::arch::aarch64::*;
4217 use core::arch::asm;
4218 let w0 = vld1q_s8(w);
4219 let w1 = vld1q_s8(w.add(16));
4220 let x0 = vld1q_s8(x);
4221 let x1 = vld1q_s8(x.add(16));
4222 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4223 asm!(
4224 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
4225 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
4226 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4227 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
4228 options(pure, nomem, nostack),
4229 );
4230 vaddvq_s32(vaddq_s32(a0, a1))
4231 }
4232}
4233
4234#[cfg(target_arch = "x86_64")]
4237#[target_feature(enable = "avx2")]
4238#[inline]
4239unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
4240 unsafe {
4242 use core::arch::x86_64::*;
4243 let wv = _mm256_loadu_si256(w as *const __m256i);
4244 let xv = _mm256_loadu_si256(x as *const __m256i);
4245 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
4246 let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
4247 let hi128 = _mm256_extracti128_si256::<1>(d);
4248 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4249 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4250 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4251 _mm_cvtsi128_si32(s32)
4252 }
4253}
4254
4255#[inline]
4260fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
4261 debug_assert!(dst.len() >= 40);
4262 unsafe {
4265 let p = dst.as_mut_ptr();
4266 for bi in 0..7 {
4267 core::ptr::write_unaligned(
4268 p.add(bi * 5) as *mut u64,
4269 SIGN5_U64[*codes.add(bi) as usize],
4270 );
4271 }
4272 }
4273}
4274
4275#[inline]
4280fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
4281 #[cfg(target_arch = "aarch64")]
4282 unsafe {
4283 return sdot32_i8(w, x);
4284 }
4285 #[cfg(target_arch = "x86_64")]
4286 unsafe {
4287 return i8dot32_avx2(w, x);
4288 }
4289 #[allow(unreachable_code)]
4290 unsafe {
4291 let mut s = 0i32;
4292 for k in 0..GROUP_SIZE {
4293 s += *w.add(k) as i32 * *x.add(k) as i32;
4294 }
4295 s
4296 }
4297}
4298
4299#[inline]
4300unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
4301 let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
4302 (
4303 SIGN5_U64[*codes as usize],
4304 SIGN5_U64[*codes.add(1) as usize],
4305 SIGN5_U64[*codes.add(2) as usize],
4306 SIGN5_U64[*codes.add(3) as usize],
4307 SIGN5_U64[*codes.add(4) as usize],
4308 SIGN5_U64[*codes.add(5) as usize],
4309 SIGN5_U64[*codes.add(6) as usize],
4310 )
4311 };
4312
4313 let u0 = s0 | (s1 << 40);
4314 let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
4315 let u2 = (s3 >> 8) | (s4 << 32);
4316 let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
4317
4318 (u0, u1, u2, u3)
4319}
4320
4321#[cfg(target_arch = "aarch64")]
4325#[target_feature(enable = "neon,dotprod")]
4326unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4327 use core::arch::aarch64::*;
4328 use core::arch::asm;
4329 unsafe {
4330 const TILE: usize = cortiq_core::quant::Q1T_TILE;
4331 let mut acc = 0f32;
4332 let bytes_ptr = bytes.as_ptr();
4333 let xq_ptr = xq.as_ptr();
4334 let row_off = r * gpr * TILE;
4335
4336 let gpr2 = gpr & !1;
4337 let mut gi = 0;
4338 while gi < gpr2 {
4339 let off0 = row_off + gi * TILE;
4340 let off1 = off0 + TILE;
4341 let s0 = f16_to_f32(u16::from_le_bytes([
4342 *bytes_ptr.add(off0),
4343 *bytes_ptr.add(off0 + 1),
4344 ]));
4345 let s1 = f16_to_f32(u16::from_le_bytes([
4346 *bytes_ptr.add(off1),
4347 *bytes_ptr.add(off1 + 1),
4348 ]));
4349
4350 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
4351 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
4352
4353 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
4354 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
4355 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
4356 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
4357
4358 let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
4359 let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
4360 let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
4361 let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
4362
4363 let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
4364 let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4365 asm!(
4366 "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
4367 "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
4368 "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
4369 "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
4370 a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
4371 a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
4372 w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
4373 w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
4374 options(pure, nomem, nostack),
4375 );
4376 let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
4377 let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
4378 acc += d0 as f32 * s0 + d1 as f32 * s1;
4379 gi += 2;
4380 }
4381
4382 if gi < gpr {
4383 let off = row_off + gi * TILE;
4384 let s = f16_to_f32(u16::from_le_bytes([
4385 *bytes_ptr.add(off),
4386 *bytes_ptr.add(off + 1),
4387 ]));
4388 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
4389 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
4390 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
4391 let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
4392 let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
4393 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4394 asm!(
4395 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
4396 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
4397 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4398 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
4399 options(pure, nomem, nostack),
4400 );
4401 let d = vaddvq_s32(vaddq_s32(a0, a1));
4402 acc += d as f32 * s;
4403 }
4404 acc
4405 }
4406}
4407
4408#[cfg(target_arch = "x86_64")]
4410#[target_feature(enable = "avx2")]
4411unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4412 use core::arch::x86_64::*;
4413 unsafe {
4414 const TILE: usize = cortiq_core::quant::Q1T_TILE;
4415 let mut acc = 0f32;
4416 let bytes_ptr = bytes.as_ptr();
4417 let xq_ptr = xq.as_ptr();
4418 let row_off = r * gpr * TILE;
4419
4420 let ones = _mm256_set1_epi16(1);
4421 for gi in 0..gpr {
4422 let off = row_off + gi * TILE;
4423 let s = f16_to_f32(u16::from_le_bytes([
4424 *bytes_ptr.add(off),
4425 *bytes_ptr.add(off + 1),
4426 ]));
4427 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
4428 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
4429 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
4430 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
4431 let d256 = _mm256_madd_epi16(p16, ones);
4432 let d128 = _mm_add_epi32(
4433 _mm256_castsi256_si128(d256),
4434 _mm256_extracti128_si256(d256, 1),
4435 );
4436 let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
4437 let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
4438 acc += d32 as f32 * s;
4439 }
4440 acc
4441 }
4442}
4443
4444#[cfg(target_arch = "x86_64")]
4446#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4447unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4448 use core::arch::x86_64::*;
4449 unsafe {
4451 const TILE: usize = cortiq_core::quant::Q1T_TILE;
4452 let mut acc = 0f32;
4453 let bytes_ptr = bytes.as_ptr();
4454 let xq_ptr = xq.as_ptr();
4455 let row_off = r * gpr * TILE;
4456 for gi in 0..gpr {
4457 let off = row_off + gi * TILE;
4458 let s = f16_to_f32(u16::from_le_bytes([
4459 *bytes_ptr.add(off),
4460 *bytes_ptr.add(off + 1),
4461 ]));
4462 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
4463 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
4464 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
4465 let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
4466 acc += d as f32 * s;
4467 }
4468 acc
4469 }
4470}
4471
4472#[inline]
4476fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4477 #[cfg(target_arch = "aarch64")]
4478 unsafe {
4479 return q1t_dot_row_sdot(bytes, r, gpr, xq);
4480 }
4481 #[cfg(target_arch = "x86_64")]
4482 unsafe {
4483 if vnni_tiles_enabled() {
4484 return q1t_dot_row_vnni(bytes, r, gpr, xq);
4485 }
4486 return q1t_dot_row_avx2(bytes, r, gpr, xq);
4487 }
4488 #[allow(unreachable_code)]
4489 {
4490 const TILE: usize = cortiq_core::quant::Q1T_TILE;
4491 let mut acc = 0f32;
4492 let mut sg = [0i8; GROUP_SIZE + 8]; for gi in 0..gpr {
4494 let off = (r * gpr + gi) * TILE;
4495 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
4496 q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
4497 let mut d = 0i32;
4498 for k in 0..GROUP_SIZE {
4499 d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
4500 }
4501 acc += d as f32 * s;
4502 }
4503 acc
4504 }
4505}
4506
4507fn q1t_row_outlier_correction(
4514 bytes: &[u8],
4515 r: usize,
4516 rp_off: usize,
4517 entries_off: usize,
4518 has_ov: bool,
4519 x: &[f32],
4520) -> f32 {
4521 if !has_ov {
4522 return 0.0;
4523 }
4524 let (c0, c1) = (
4525 q1t_rowptr(bytes, rp_off, r),
4526 q1t_rowptr(bytes, rp_off, r + 1),
4527 );
4528 let mut corr = 0f32;
4529 for p in c0..c1 {
4530 let e = entries_off + p * 4;
4531 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
4532 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
4533 corr += val * x[col];
4534 }
4535 corr
4536}
4537
4538fn q1t_dequant_row(
4542 bytes: &[u8],
4543 r: usize,
4544 gpr: usize,
4545 rp_off: usize,
4546 entries_off: usize,
4547 has_ov: bool,
4548 buf: &mut [f32],
4549) {
4550 const TILE: usize = cortiq_core::quant::Q1T_TILE;
4551 for g in 0..gpr {
4552 let off = (r * gpr + g) * TILE;
4553 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
4554 let codes = &bytes[off + 2..off + TILE];
4555 let bc = g * GROUP_SIZE;
4556 for bi in 0..6 {
4558 let lut = &SIGN5[codes[bi] as usize];
4559 let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
4560 for i in 0..5 {
4561 d[i] = lut[i] * s;
4562 }
4563 }
4564 let lut = &SIGN5[codes[6] as usize];
4565 buf[bc + 30] = lut[0] * s;
4566 buf[bc + 31] = lut[1] * s;
4567 }
4568 if !has_ov {
4569 return;
4570 }
4571 let (c0, c1) = (
4572 q1t_rowptr(bytes, rp_off, r),
4573 q1t_rowptr(bytes, rp_off, r + 1),
4574 );
4575 for p in c0..c1 {
4576 let e = entries_off + p * 4;
4577 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
4578 buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
4579 }
4580}
4581
4582fn q1t_add_overlay(
4586 bytes: &[u8],
4587 x: &[f32],
4588 rows: usize,
4589 cols: usize,
4590 out: &mut [f32],
4591 pool: Option<&Pool>,
4592) {
4593 const TILE: usize = cortiq_core::quant::Q1T_TILE;
4594 let gpr = cols / GROUP_SIZE;
4595 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
4596 if !has_ov {
4597 return;
4598 }
4599 let out_addr = SendMut(out.as_mut_ptr());
4600 let run = move |start: usize, end: usize| {
4601 for r in start..end {
4602 let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
4603 unsafe { *out_addr.at(r) += corr };
4605 }
4606 };
4607 dispatch_rows(pool, rows, &run);
4608}
4609
4610#[allow(clippy::too_many_arguments)]
4613fn q1t_range_a8w8(
4614 bytes: &[u8],
4615 gpr: usize,
4616 rp_off: usize,
4617 ent_off: usize,
4618 has_ov: bool,
4619 act: &SplitAct,
4620 x: &[f32],
4621 out: SendMut,
4622 start: usize,
4623 end: usize,
4624) {
4625 for r in start..end {
4626 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
4627 for &(j, xv) in &act.outliers {
4628 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
4629 }
4630 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
4631 unsafe { *out.at(r) = acc };
4633 }
4634}
4635
4636#[allow(clippy::too_many_arguments)]
4639fn q1t_range_f32_batch(
4640 bytes: &[u8],
4641 gpr: usize,
4642 rp_off: usize,
4643 ent_off: usize,
4644 has_ov: bool,
4645 x: &[f32],
4646 out: SendMut,
4647 start: usize,
4648 end: usize,
4649) {
4650 const TILE: usize = cortiq_core::quant::Q1T_TILE;
4651 let mut sg = [0f32; GROUP_SIZE];
4652 for r in start..end {
4653 let mut acc = 0f32;
4654 for g in 0..gpr {
4655 let off = (r * gpr + g) * TILE;
4656 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
4657 let codes = &bytes[off + 2..off + TILE];
4658 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
4659 for bi in 0..6 {
4660 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
4661 }
4662 let lut = &SIGN5[codes[6] as usize];
4663 sg[30] = lut[0];
4664 sg[31] = lut[1];
4665 let mut gsum = 0f32;
4666 for k in 0..GROUP_SIZE {
4667 gsum += sg[k] * xg[k];
4668 }
4669 acc += s * gsum;
4670 }
4671 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
4672 unsafe { *out.at(r) = acc };
4674 }
4675}
4676
4677fn q1t_matvec(
4681 bytes: &[u8],
4682 x: &[f32],
4683 rows: usize,
4684 cols: usize,
4685 out: &mut [f32],
4686 pool: Option<&Pool>,
4687) {
4688 debug_assert_eq!(out.len(), rows);
4689 const TILE: usize = cortiq_core::quant::Q1T_TILE;
4690 let gpr = cols / GROUP_SIZE;
4691 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
4692 let out_addr = SendMut(out.as_mut_ptr());
4693 if a8w8_enabled() {
4697 let act = split_act(x);
4698 let act = &act;
4699 let run = move |start: usize, end: usize| {
4700 for r in start..end {
4701 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
4702 for &(j, xv) in &act.outliers {
4703 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
4704 }
4705 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
4706 unsafe { *out_addr.at(r) = acc };
4708 }
4709 };
4710 dispatch_rows(pool, rows, &run);
4711 return;
4712 }
4713 let run = move |start: usize, end: usize| {
4714 let mut sg = [0f32; GROUP_SIZE];
4718 for r in start..end {
4719 let mut acc = 0f32;
4720 for g in 0..gpr {
4721 let off = (r * gpr + g) * TILE;
4722 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
4723 let codes = &bytes[off + 2..off + TILE];
4724 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
4725 for bi in 0..6 {
4726 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
4727 }
4728 let lut = &SIGN5[codes[6] as usize];
4729 sg[30] = lut[0];
4730 sg[31] = lut[1];
4731 let mut gsum = 0f32;
4732 for k in 0..GROUP_SIZE {
4733 gsum += sg[k] * xg[k];
4734 }
4735 acc += s * gsum;
4736 }
4737 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
4738 unsafe { *out_addr.at(r) = acc };
4739 }
4740 };
4741 dispatch_rows(pool, rows, &run);
4742}
4743
4744#[cfg(target_arch = "aarch64")]
4750#[target_feature(enable = "neon,dotprod")]
4751unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
4752 use core::arch::aarch64::*;
4753 use core::arch::asm;
4754 unsafe {
4756 const TILE: usize = cortiq_core::quant::Q1T_TILE;
4757 let bytes_ptr = bytes.as_ptr();
4758 let row_off = r * gpr * TILE;
4759 let xp = [xa.as_ptr(), xb.as_ptr()];
4760 let mut acc = [0f32; 2];
4761 macro_rules! sdot2 {
4762 ($w0:expr, $w1:expr, $x:expr) => {{
4763 let x0 = vld1q_s8($x);
4764 let x1 = vld1q_s8($x.add(16));
4765 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4766 asm!(
4767 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
4768 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
4769 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4770 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
4771 options(pure, nomem, nostack),
4772 );
4773 vaddvq_s32(vaddq_s32(a0, a1))
4774 }};
4775 }
4776 let gpr2 = gpr & !1;
4777 let mut gi = 0;
4778 while gi < gpr2 {
4779 let off0 = row_off + gi * TILE;
4780 let off1 = off0 + TILE;
4781 let s0 = f16_to_f32(u16::from_le_bytes([
4782 *bytes_ptr.add(off0),
4783 *bytes_ptr.add(off0 + 1),
4784 ]));
4785 let s1 = f16_to_f32(u16::from_le_bytes([
4786 *bytes_ptr.add(off1),
4787 *bytes_ptr.add(off1 + 1),
4788 ]));
4789 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
4790 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
4791 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
4792 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
4793 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
4794 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
4795 for k in 0..2 {
4796 let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
4797 let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
4798 acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
4799 }
4800 gi += 2;
4801 }
4802 if gi < gpr {
4803 let off = row_off + gi * TILE;
4804 let s = f16_to_f32(u16::from_le_bytes([
4805 *bytes_ptr.add(off),
4806 *bytes_ptr.add(off + 1),
4807 ]));
4808 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
4809 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
4810 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
4811 for k in 0..2 {
4812 let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
4813 acc[k] += d as f32 * s;
4814 }
4815 }
4816 acc
4817 }
4818}
4819
4820fn q1t_matvec2(
4826 bytes: &[u8],
4827 x1: &[f32],
4828 x2: &[f32],
4829 rows: usize,
4830 cols: usize,
4831 o1: &mut [f32],
4832 o2: &mut [f32],
4833 pool: Option<&Pool>,
4834) {
4835 debug_assert_eq!(o1.len(), rows);
4836 debug_assert_eq!(o2.len(), rows);
4837 const TILE: usize = cortiq_core::quant::Q1T_TILE;
4838 let gpr = cols / GROUP_SIZE;
4839 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
4840 let out1 = SendMut(o1.as_mut_ptr());
4841 let out2 = SendMut(o2.as_mut_ptr());
4842 if a8w8_enabled() {
4843 let a1 = split_act(x1);
4844 let a2 = split_act(x2);
4845 let (a1, a2) = (&a1, &a2);
4846 let run = move |start: usize, end: usize| {
4847 for r in start..end {
4848 #[cfg(target_arch = "aarch64")]
4849 let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
4852 #[cfg(not(target_arch = "aarch64"))]
4853 let ds = [
4854 q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
4855 q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
4856 ];
4857 let mut acc1 = ds[0] * a1.sx;
4858 for &(j, xv) in &a1.outliers {
4859 acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
4860 }
4861 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
4862 let mut acc2 = ds[1] * a2.sx;
4863 for &(j, xv) in &a2.outliers {
4864 acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
4865 }
4866 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
4867 unsafe {
4869 *out1.at(r) = acc1;
4870 *out2.at(r) = acc2;
4871 }
4872 }
4873 };
4874 dispatch_rows(pool, rows, &run);
4875 return;
4876 }
4877 let run = move |start: usize, end: usize| {
4878 let mut sg = [0f32; GROUP_SIZE];
4881 for r in start..end {
4882 let mut acc1 = 0f32;
4883 let mut acc2 = 0f32;
4884 for g in 0..gpr {
4885 let off = (r * gpr + g) * TILE;
4886 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
4887 let codes = &bytes[off + 2..off + TILE];
4888 for bi in 0..6 {
4889 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
4890 }
4891 let lut = &SIGN5[codes[6] as usize];
4892 sg[30] = lut[0];
4893 sg[31] = lut[1];
4894 let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
4895 let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
4896 let mut gsum1 = 0f32;
4897 for k in 0..GROUP_SIZE {
4898 gsum1 += sg[k] * xg1[k];
4899 }
4900 acc1 += s * gsum1;
4901 let mut gsum2 = 0f32;
4902 for k in 0..GROUP_SIZE {
4903 gsum2 += sg[k] * xg2[k];
4904 }
4905 acc2 += s * gsum2;
4906 }
4907 acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
4908 acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
4909 unsafe {
4911 *out1.at(r) = acc1;
4912 *out2.at(r) = acc2;
4913 }
4914 }
4915 };
4916 dispatch_rows(pool, rows, &run);
4917}
4918
4919fn q1t_matmat(
4922 bytes: &[u8],
4923 xs: &[f32],
4924 b: usize,
4925 rows: usize,
4926 cols: usize,
4927 out: &mut [f32],
4928 pool: Option<&Pool>,
4929) {
4930 debug_assert_eq!(out.len(), b * rows);
4931 const TILE: usize = cortiq_core::quant::Q1T_TILE;
4932 let gpr = cols / GROUP_SIZE;
4933 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
4934 let out_addr = SendMut(out.as_mut_ptr());
4935 if a8w8_enabled() {
4939 let acts: Vec<SplitAct> = (0..b)
4940 .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
4941 .collect();
4942 let acts = &acts;
4943 let run = move |start: usize, end: usize| {
4944 let mut sg = vec![0i8; cols + 8]; let mut sc = vec![0f32; gpr]; let mut accs = vec![0f32; b]; for r in start..end {
4948 for g in 0..gpr {
4949 let off = (r * gpr + g) * TILE;
4950 sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
4951 q1t_unpack_group_i8(
4952 bytes.as_ptr().wrapping_add(off + 2),
4953 &mut sg[g * GROUP_SIZE..],
4954 );
4955 }
4956 for bi in 0..b {
4957 let act = &acts[bi];
4958 let mut isum = 0f32;
4959 for g in 0..gpr {
4960 let d = q1t_i8dot32(
4961 sg.as_ptr().wrapping_add(g * GROUP_SIZE),
4962 act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
4963 );
4964 isum += d as f32 * sc[g];
4965 }
4966 let mut acc = isum * act.sx;
4967 for &(j, xv) in &act.outliers {
4968 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
4969 }
4970 accs[bi] = acc;
4971 }
4972 if has_ov {
4976 let (c0, c1) = (
4977 q1t_rowptr(bytes, rp_off, r),
4978 q1t_rowptr(bytes, rp_off, r + 1),
4979 );
4980 for p in c0..c1 {
4981 let e = ent_off + p * 4;
4982 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
4983 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
4984 for bi in 0..b {
4985 accs[bi] += val * xs[bi * cols + col];
4986 }
4987 }
4988 }
4989 for bi in 0..b {
4990 unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
4991 }
4992 }
4993 };
4994 dispatch_rows(pool, rows, &run);
4995 return;
4996 }
4997 let run = move |start: usize, end: usize| {
4998 let mut buf = vec![0f32; cols];
4999 for r in start..end {
5000 q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
5001 for bi in 0..b {
5002 let xr = &xs[bi * cols..(bi + 1) * cols];
5003 let mut acc = 0f32;
5004 for j in 0..cols {
5005 acc += buf[j] * xr[j];
5006 }
5007 unsafe { *out_addr.at(bi * rows + r) = acc };
5008 }
5009 }
5010 };
5011 dispatch_rows(pool, rows, &run);
5012}
5013
5014fn q1_matvec(
5015 bytes: &[u8],
5016 x: &[f32],
5017 rows: usize,
5018 cols: usize,
5019 out: &mut [f32],
5020 pool: Option<&Pool>,
5021) {
5022 debug_assert_eq!(out.len(), rows);
5023 let gpr = cols / GROUP_SIZE;
5024 let out_addr = SendMut(out.as_mut_ptr());
5025 if a8w8_enabled() {
5026 let act = split_act(x);
5027 let gsum = q1_group_sums(&act.xq, gpr);
5028 let (act, gsum) = (&act, &gsum);
5029 let run = move |start: usize, end: usize| {
5030 q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end)
5031 };
5032 dispatch_rows(pool, rows, &run);
5033 return;
5034 }
5035 let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
5036 dispatch_rows(pool, rows, &run);
5037}
5038
5039#[allow(clippy::too_many_arguments)]
5041fn q1_matvec2(
5042 bytes: &[u8],
5043 x1: &[f32],
5044 x2: &[f32],
5045 rows: usize,
5046 cols: usize,
5047 o1: &mut [f32],
5048 o2: &mut [f32],
5049 pool: Option<&Pool>,
5050) {
5051 let gpr = cols / GROUP_SIZE;
5052 let p1 = SendMut(o1.as_mut_ptr());
5053 let p2 = SendMut(o2.as_mut_ptr());
5054 if a8w8_enabled() {
5055 let a1 = split_act(x1);
5056 let a2 = split_act(x2);
5057 let g1 = q1_group_sums(&a1.xq, gpr);
5058 let g2 = q1_group_sums(&a2.xq, gpr);
5059 let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
5060 let run = move |start: usize, end: usize| {
5061 for r in start..end {
5062 let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
5063 let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
5064 for &(j, xv) in &a1.outliers {
5065 let (w, s) = q1_outlier(bytes, r, gpr, j);
5066 v1 += w * s * xv;
5067 }
5068 for &(j, xv) in &a2.outliers {
5069 let (w, s) = q1_outlier(bytes, r, gpr, j);
5070 v2 += w * s * xv;
5071 }
5072 unsafe {
5074 *p1.at(r) = v1;
5075 *p2.at(r) = v2;
5076 }
5077 }
5078 };
5079 dispatch_rows(pool, rows, &run);
5080 return;
5081 }
5082 let run = move |start: usize, end: usize| {
5083 for r in start..end {
5084 unsafe {
5086 *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
5087 *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
5088 }
5089 }
5090 };
5091 dispatch_rows(pool, rows, &run);
5092}
5093
5094#[allow(clippy::too_many_arguments)]
5096fn q1_matmat(
5097 bytes: &[u8],
5098 xs_all: &[f32],
5099 b: usize,
5100 rows: usize,
5101 cols: usize,
5102 out: &mut [f32],
5103 pool: Option<&Pool>,
5104) {
5105 debug_assert_eq!(out.len(), b * rows);
5106 let gpr = cols / GROUP_SIZE;
5107 let out_addr = SendMut(out.as_mut_ptr());
5108 if a8w8_enabled() {
5109 let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
5110 .map(|bi| {
5111 let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
5112 let gsum = q1_group_sums(&act.xq, gpr);
5113 (act, gsum)
5114 })
5115 .collect();
5116 let acts = &acts;
5117 #[cfg(target_arch = "x86_64")]
5118 let blocked_ok = avx2_enabled()
5119 && std::env::var("CMF_X86_BLOCKED")
5120 .map(|v| v != "0")
5121 .unwrap_or(true);
5122 #[cfg(target_arch = "aarch64")]
5123 let blocked_ok = sdot_enabled()
5124 && std::env::var("CMF_X86_BLOCKED")
5125 .map(|v| v != "0")
5126 .unwrap_or(true);
5127 let run = move |start: usize, end: usize| {
5128 for r in start..end {
5129 let mut bi = 0usize;
5130 #[cfg(target_arch = "aarch64")]
5133 if blocked_ok {
5134 while bi + 4 <= acts.len() {
5135 let xs = [
5136 acts[bi].0.xq.as_slice(),
5137 acts[bi + 1].0.xq.as_slice(),
5138 acts[bi + 2].0.xq.as_slice(),
5139 acts[bi + 3].0.xq.as_slice(),
5140 ];
5141 let gs = [
5142 acts[bi].1.as_slice(),
5143 acts[bi + 1].1.as_slice(),
5144 acts[bi + 2].1.as_slice(),
5145 acts[bi + 3].1.as_slice(),
5146 ];
5147 let d = unsafe { dot_q1_row_1x4_sdot(bytes, r, gpr, xs, gs) };
5148 for k in 0..4 {
5149 let (act, _) = &acts[bi + k];
5150 let mut acc = d[k] * act.sx;
5151 for &(j, xv) in &act.outliers {
5152 let (w, sc) = q1_outlier(bytes, r, gpr, j);
5153 acc += w * sc * xv;
5154 }
5155 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5157 }
5158 bi += 4;
5159 }
5160 }
5161 #[cfg(target_arch = "x86_64")]
5162 if blocked_ok {
5163 while bi + 4 <= acts.len() {
5164 let xs = [
5165 acts[bi].0.xq.as_slice(),
5166 acts[bi + 1].0.xq.as_slice(),
5167 acts[bi + 2].0.xq.as_slice(),
5168 acts[bi + 3].0.xq.as_slice(),
5169 ];
5170 let gs = [
5171 acts[bi].1.as_slice(),
5172 acts[bi + 1].1.as_slice(),
5173 acts[bi + 2].1.as_slice(),
5174 acts[bi + 3].1.as_slice(),
5175 ];
5176 let d = unsafe {
5177 if vnni_tiles_enabled() {
5178 dot_q1_row_1x4_vnni(bytes, r, gpr, xs, gs)
5179 } else {
5180 dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs)
5181 }
5182 };
5183 for k in 0..4 {
5184 let (act, _) = &acts[bi + k];
5185 let mut acc = d[k] * act.sx;
5186 for &(j, xv) in &act.outliers {
5187 let (w, sc) = q1_outlier(bytes, r, gpr, j);
5188 acc += w * sc * xv;
5189 }
5190 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5192 }
5193 bi += 4;
5194 }
5195 }
5196 while bi < acts.len() {
5197 let (act, gsum) = &acts[bi];
5198 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
5199 for &(j, xv) in &act.outliers {
5200 let (w, s) = q1_outlier(bytes, r, gpr, j);
5201 acc += w * s * xv;
5202 }
5203 unsafe { *out_addr.at(bi * rows + r) = acc };
5205 bi += 1;
5206 }
5207 }
5208 };
5209 dispatch_rows(pool, rows, &run);
5210 return;
5211 }
5212 let run = move |start: usize, end: usize| {
5213 for r in start..end {
5214 for bi in 0..b {
5215 let x = &xs_all[bi * cols..(bi + 1) * cols];
5216 unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
5218 }
5219 }
5220 };
5221 dispatch_rows(pool, rows, &run);
5222}
5223
5224fn q4matvec(
5230 bytes: &[u8],
5231 x: &[f32],
5232 rows: usize,
5233 cols: usize,
5234 out: &mut [f32],
5235 pool: Option<&Pool>,
5236) {
5237 debug_assert_eq!(out.len(), rows);
5238 let (packed, scales) = q4_split(bytes, rows, cols);
5239 let gpr = cols / GROUP_SIZE;
5240 let out_addr = SendMut(out.as_mut_ptr());
5241
5242 if a8w8_enabled() {
5243 let act = split_act(x);
5244 let run = move |start: usize, end: usize| {
5245 q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
5246 };
5247 dispatch_rows(pool, rows, &run);
5248 return;
5249 }
5250
5251 let run =
5252 move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
5253 dispatch_rows(pool, rows, &run);
5254}
5255
5256#[inline]
5259#[allow(unreachable_code)]
5260#[cfg(target_arch = "x86_64")]
5265#[target_feature(enable = "avx2")]
5266unsafe fn dot_q4b_row_1x4_avx2(
5267 buf: &[u8],
5268 scales: &[u8],
5269 g0: usize,
5270 gpr: usize,
5271 xs: [&[i8]; 4],
5272) -> [f32; 4] {
5273 unsafe {
5275 use core::arch::x86_64::*;
5276 let ones = _mm256_set1_epi16(1);
5277 let mut acc = [0f32; 4];
5278 for gi in 0..gpr {
5279 let s = f16_to_f32(u16::from_le_bytes([
5280 scales[(g0 + gi) * 2],
5281 scales[(g0 + gi) * 2 + 1],
5282 ]));
5283 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5284 let aw = _mm256_abs_epi8(w);
5285 for (k, xq) in xs.iter().enumerate() {
5286 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5287 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
5288 let d = _mm256_madd_epi16(p16, ones);
5289 let hi128 = _mm256_extracti128_si256::<1>(d);
5290 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5291 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5292 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5293 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
5294 }
5295 }
5296 acc
5297 }
5298}
5299
5300#[cfg(target_arch = "x86_64")]
5302#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5303unsafe fn dot_q4b_row_1x4_vnni(
5304 buf: &[u8],
5305 scales: &[u8],
5306 g0: usize,
5307 gpr: usize,
5308 xs: [&[i8]; 4],
5309) -> [f32; 4] {
5310 unsafe {
5312 use core::arch::x86_64::*;
5313 let mut acc = [0f32; 4];
5314 for gi in 0..gpr {
5315 let s = f16_to_f32(u16::from_le_bytes([
5316 scales[(g0 + gi) * 2],
5317 scales[(g0 + gi) * 2 + 1],
5318 ]));
5319 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5320 let aw = _mm256_abs_epi8(w);
5321 for (k, xq) in xs.iter().enumerate() {
5322 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5323 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
5324 acc[k] += d as f32 * s;
5325 }
5326 }
5327 acc
5328 }
5329}
5330
5331#[cfg(target_arch = "x86_64")]
5337#[target_feature(enable = "avx2")]
5338unsafe fn dot_q4b_row_1x4_sx_avx2(
5339 buf: &[u8],
5340 scales: &[u8],
5341 g0: usize,
5342 gpr: usize,
5343 xs: [&[i8]; 4],
5344 sxs: [f32; 4],
5345) -> [f32; 4] {
5346 unsafe {
5348 use core::arch::x86_64::*;
5349 let ones = _mm256_set1_epi16(1);
5350 let mut acc = [0f32; 4];
5351 for gi in 0..gpr {
5352 let s = f16_to_f32(u16::from_le_bytes([
5353 scales[(g0 + gi) * 2],
5354 scales[(g0 + gi) * 2 + 1],
5355 ]));
5356 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5357 let aw = _mm256_abs_epi8(w);
5358 for (k, xq) in xs.iter().enumerate() {
5359 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5360 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
5361 let d = _mm256_madd_epi16(p16, ones);
5362 let hi128 = _mm256_extracti128_si256::<1>(d);
5363 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5364 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5365 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5366 acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
5367 }
5368 }
5369 acc
5370 }
5371}
5372
5373#[cfg(target_arch = "x86_64")]
5376#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5377unsafe fn dot_q4b_row_1x4_sx_vnni(
5378 buf: &[u8],
5379 scales: &[u8],
5380 g0: usize,
5381 gpr: usize,
5382 xs: [&[i8]; 4],
5383 sxs: [f32; 4],
5384) -> [f32; 4] {
5385 unsafe {
5387 use core::arch::x86_64::*;
5388 let mut acc = [0f32; 4];
5389 for gi in 0..gpr {
5390 let s = f16_to_f32(u16::from_le_bytes([
5391 scales[(g0 + gi) * 2],
5392 scales[(g0 + gi) * 2 + 1],
5393 ]));
5394 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5395 let aw = _mm256_abs_epi8(w);
5396 for (k, xq) in xs.iter().enumerate() {
5397 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5398 let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
5399 acc[k] += (d as f32 * sxs[k]) * s;
5400 }
5401 }
5402 acc
5403 }
5404}
5405
5406#[allow(unreachable_code)]
5407fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
5408 #[cfg(target_arch = "aarch64")]
5409 unsafe {
5410 return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
5411 }
5412 #[cfg(target_arch = "x86_64")]
5413 unsafe {
5414 return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
5415 }
5416 let mut acc = 0f32;
5417 for gi in 0..gpr {
5418 let g = g0 + gi;
5419 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
5420 let mut d = 0i32;
5421 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
5422 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
5423 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
5424 }
5425 acc += d as f32 * s;
5426 }
5427 acc
5428}
5429
5430#[inline]
5432#[allow(unreachable_code)]
5433fn dot_q4_row_i8_2(
5434 packed: &[u8],
5435 scales: &[u8],
5436 g0: usize,
5437 gpr: usize,
5438 xq1: &[i8],
5439 xq2: &[i8],
5440) -> (f32, f32) {
5441 #[cfg(target_arch = "aarch64")]
5442 unsafe {
5443 return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
5444 }
5445 #[cfg(target_arch = "x86_64")]
5446 unsafe {
5447 return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
5448 }
5449 (
5450 dot_q4_row_i8(packed, scales, g0, gpr, xq1),
5451 dot_q4_row_i8(packed, scales, g0, gpr, xq2),
5452 )
5453}
5454
5455#[allow(clippy::too_many_arguments)]
5458fn q4_range_a8w8(
5459 packed: &[u8],
5460 scales: &[u8],
5461 gpr: usize,
5462 cols: usize,
5463 act: &SplitAct,
5464 out: SendMut,
5465 start: usize,
5466 end: usize,
5467) {
5468 for r in start..end {
5469 let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
5470 for &(j, xv) in &act.outliers {
5472 let flat = r * cols + j;
5473 let byte = packed[flat / 2];
5474 let nib = if flat & 1 == 0 {
5475 byte & 0x0F
5476 } else {
5477 byte >> 4
5478 };
5479 let s = f16_to_f32(u16::from_le_bytes([
5480 scales[(flat / GROUP_SIZE) * 2],
5481 scales[(flat / GROUP_SIZE) * 2 + 1],
5482 ]));
5483 acc += ((nib as i32 - 8) as f32) * s * xv;
5484 }
5485 unsafe { *out.at(r) = acc };
5487 }
5488}
5489
5490#[allow(clippy::too_many_arguments)]
5493fn q4_range2_a8w8(
5494 packed: &[u8],
5495 scales: &[u8],
5496 gpr: usize,
5497 cols: usize,
5498 a1: &SplitAct,
5499 a2: &SplitAct,
5500 p1: SendMut,
5501 p2: SendMut,
5502 start: usize,
5503 end: usize,
5504) {
5505 for r in start..end {
5506 let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
5507 let mut acc1 = s1 * a1.sx;
5508 let mut acc2 = s2 * a2.sx;
5509 let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
5511 for &(j, xv) in outliers {
5512 let flat = r * cols + j;
5513 let byte = packed[flat / 2];
5514 let nib = if flat & 1 == 0 {
5515 byte & 0x0F
5516 } else {
5517 byte >> 4
5518 };
5519 let s = f16_to_f32(u16::from_le_bytes([
5520 scales[(flat / GROUP_SIZE) * 2],
5521 scales[(flat / GROUP_SIZE) * 2 + 1],
5522 ]));
5523 *acc += ((nib as i32 - 8) as f32) * s * xv;
5524 }
5525 };
5526 fix(&a1.outliers, &mut acc1);
5527 fix(&a2.outliers, &mut acc2);
5528 unsafe {
5530 *p1.at(r) = acc1;
5531 *p2.at(r) = acc2;
5532 }
5533 }
5534}
5535
5536fn q4_range_f32(
5538 packed: &[u8],
5539 scales: &[u8],
5540 gpr: usize,
5541 x: &[f32],
5542 out: SendMut,
5543 start: usize,
5544 end: usize,
5545) {
5546 for r in start..end {
5547 let mut acc = 0f32;
5548 for gi in 0..gpr {
5549 let g = r * gpr + gi;
5550 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
5551 let pk = &packed[g * 16..(g + 1) * 16];
5552 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5553 let mut ga = 0f32;
5554 for (k, &b) in pk.iter().enumerate() {
5555 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
5556 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
5557 }
5558 acc += ga * s;
5559 }
5560 unsafe { *out.at(r) = acc };
5562 }
5563}
5564
5565#[allow(clippy::too_many_arguments)]
5569fn q4matvec2(
5570 bytes: &[u8],
5571 x1: &[f32],
5572 x2: &[f32],
5573 rows: usize,
5574 cols: usize,
5575 o1: &mut [f32],
5576 o2: &mut [f32],
5577 pool: Option<&Pool>,
5578) {
5579 debug_assert_eq!(o1.len(), rows);
5580 debug_assert_eq!(o2.len(), rows);
5581 let (packed, scales) = q4_split(bytes, rows, cols);
5582 let gpr = cols / GROUP_SIZE;
5583
5584 if a8w8_enabled() {
5585 let a1 = split_act(x1);
5586 let a2 = split_act(x2);
5587 let p1 = SendMut(o1.as_mut_ptr());
5588 let p2 = SendMut(o2.as_mut_ptr());
5589 let run = move |start: usize, end: usize| {
5590 q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
5591 };
5592 dispatch_rows(pool, rows, &run);
5593 return;
5594 }
5595
5596 let p1 = SendMut(o1.as_mut_ptr());
5597 let p2 = SendMut(o2.as_mut_ptr());
5598 let run = move |start: usize, end: usize| {
5599 q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
5600 };
5601 dispatch_rows(pool, rows, &run);
5602}
5603
5604#[allow(clippy::too_many_arguments)]
5606fn q4_range2_f32(
5607 packed: &[u8],
5608 scales: &[u8],
5609 gpr: usize,
5610 x1: &[f32],
5611 x2: &[f32],
5612 p1: SendMut,
5613 p2: SendMut,
5614 start: usize,
5615 end: usize,
5616) {
5617 for r in start..end {
5618 let (mut acc1, mut acc2) = (0f32, 0f32);
5619 for gi in 0..gpr {
5620 let g = r * gpr + gi;
5621 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
5622 let pk = &packed[g * 16..(g + 1) * 16];
5623 let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5624 let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5625 let (mut g1, mut g2) = (0f32, 0f32);
5626 for (k, &b) in pk.iter().enumerate() {
5627 let wl = (b & 0x0F) as f32 - 8.0;
5628 let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
5629 g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
5630 g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
5631 }
5632 acc1 += g1 * s;
5633 acc2 += g2 * s;
5634 }
5635 unsafe {
5637 *p1.at(r) = acc1;
5638 *p2.at(r) = acc2;
5639 }
5640 }
5641}
5642
5643thread_local! {
5644 static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
5647 static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
5648}
5649
5650#[allow(clippy::too_many_arguments)]
5656fn q4matmat(
5657 bytes: &[u8],
5658 xs_all: &[f32],
5659 b: usize,
5660 rows: usize,
5661 cols: usize,
5662 out: &mut [f32],
5663 pool: Option<&Pool>,
5664) {
5665 debug_assert_eq!(xs_all.len(), b * cols);
5666 debug_assert_eq!(out.len(), b * rows);
5667 let (packed, scales) = q4_split(bytes, rows, cols);
5668 let gpr = cols / GROUP_SIZE;
5669 let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
5670
5671 if a8w8_enabled() {
5672 let acts: Vec<SplitAct> = (0..b)
5673 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
5674 .collect();
5675 let acts = &acts;
5676 let out_addr = SendMut(out.as_mut_ptr());
5677 let run = move |start: usize, end: usize| {
5678 ROW_I8.with(|rb| {
5679 let mut buf = rb.borrow_mut();
5680 buf.resize(cols, 0);
5681 for r in start..end {
5682 for gi in 0..gpr {
5686 let g = r * gpr + gi;
5687 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
5688 buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
5689 buf[gi * GROUP_SIZE + k * 2 + 1] =
5690 (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
5691 }
5692 }
5693 let mut bi = 0usize;
5694 #[cfg(target_arch = "x86_64")]
5695 if avx2_enabled()
5696 && std::env::var("CMF_X86_BLOCKED")
5697 .map(|v| v != "0")
5698 .unwrap_or(true)
5699 {
5700 while bi + 4 <= acts.len() {
5701 let xs = [
5702 acts[bi].xq.as_slice(),
5703 acts[bi + 1].xq.as_slice(),
5704 acts[bi + 2].xq.as_slice(),
5705 acts[bi + 3].xq.as_slice(),
5706 ];
5707 let d = unsafe {
5708 if vnni_tiles_enabled() {
5709 dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
5710 } else {
5711 dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
5712 }
5713 };
5714 for k in 0..4 {
5715 let act = &acts[bi + k];
5716 let mut acc = d[k] * act.sx;
5717 for &(j, xv) in &act.outliers {
5718 acc += (buf[j] as i8) as f32
5719 * gscale((r * cols + j) / GROUP_SIZE)
5720 * xv;
5721 }
5722 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5724 }
5725 bi += 4;
5726 }
5727 }
5728 while bi < acts.len() {
5729 let act = &acts[bi];
5730 let mut acc = 0f32;
5731 for gi in 0..gpr {
5732 let d = dot_i8_i8(
5733 &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
5734 &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
5735 );
5736 acc += d as f32 * gscale(r * gpr + gi);
5737 }
5738 acc *= act.sx;
5739 for &(j, xv) in &act.outliers {
5741 acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
5742 }
5743 unsafe { *out_addr.at(bi * rows + r) = acc };
5745 bi += 1;
5746 }
5747 }
5748 })
5749 };
5750 dispatch_rows(pool, rows, &run);
5751 return;
5752 }
5753
5754 let out_addr = SendMut(out.as_mut_ptr());
5755 let run = move |start: usize, end: usize| {
5756 ROW_F32.with(|rb| {
5757 let mut buf = rb.borrow_mut();
5758 buf.resize(cols, 0.0);
5759 for r in start..end {
5760 for gi in 0..gpr {
5763 let g = r * gpr + gi;
5764 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
5765 buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
5766 buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
5767 }
5768 }
5769 for bi in 0..b {
5770 let x = &xs_all[bi * cols..(bi + 1) * cols];
5771 let mut acc = 0f32;
5772 for gi in 0..gpr {
5773 let mut ga = 0f32;
5774 for k in 0..GROUP_SIZE / 2 {
5779 let e = gi * GROUP_SIZE + k * 2;
5780 ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
5781 }
5782 acc += ga * gscale(r * gpr + gi);
5783 }
5784 unsafe { *out_addr.at(bi * rows + r) = acc };
5786 }
5787 }
5788 })
5789 };
5790 dispatch_rows(pool, rows, &run);
5791}
5792
5793#[allow(clippy::too_many_arguments)]
5798fn vbitmatmat(
5799 bytes: &[u8],
5800 offsets: &[usize],
5801 xs_all: &[f32],
5802 b: usize,
5803 rows: usize,
5804 cols: usize,
5805 out: &mut [f32],
5806 pool: Option<&Pool>,
5807) {
5808 debug_assert_eq!(xs_all.len(), b * cols);
5809 debug_assert_eq!(out.len(), b * rows);
5810 debug_assert_eq!(offsets.len(), rows + 1);
5811 let ng = cols / GROUP_SIZE;
5812 let bits = &bytes[..rows];
5813 let sc_off = rows;
5814 let gscale = |r: usize, g: usize| {
5815 let so = (r * ng + g) * 2;
5816 f16_to_f32(u16::from_le_bytes([
5817 bytes[sc_off + so],
5818 bytes[sc_off + so + 1],
5819 ]))
5820 };
5821
5822 let decode_f32 = |r: usize, dst: &mut [f32]| {
5824 let bw = bits[r] as usize;
5825 let l = ((1i32 << (bw - 1)) - 1) as f32;
5826 let data = &bytes[offsets[r]..offsets[r + 1]];
5827 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
5828 for d in dst.iter_mut() {
5829 while nbits < bw {
5830 acc = (acc << 8) | data[idx] as u64;
5831 idx += 1;
5832 nbits += 8;
5833 }
5834 let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
5835 nbits -= bw;
5836 *d = u - l;
5837 }
5838 };
5839
5840 if a8w8_enabled() {
5841 let acts: Vec<SplitAct> = (0..b)
5842 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
5843 .collect();
5844 let acts = &acts;
5845 let out_addr = SendMut(out.as_mut_ptr());
5846 let run = move |start: usize, end: usize| {
5847 for r in start..end {
5848 let bw = bits[r] as usize;
5849 if bw == 8 {
5850 ROW_F32.with(|rb| {
5853 let mut buf = rb.borrow_mut();
5854 buf.resize(cols, 0.0);
5855 decode_f32(r, &mut buf);
5856 for bi in 0..b {
5857 let x = &xs_all[bi * cols..(bi + 1) * cols];
5858 let mut dot = 0f32;
5859 for g in 0..ng {
5860 let mut gd = 0f32;
5861 for k in 0..GROUP_SIZE {
5862 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
5863 }
5864 dot += gd * gscale(r, g);
5865 }
5866 unsafe { *out_addr.at(bi * rows + r) = dot };
5868 }
5869 });
5870 continue;
5871 }
5872 let l = (1i32 << (bw - 1)) - 1;
5873 let data = &bytes[offsets[r]..offsets[r + 1]];
5874 ROW_I8.with(|rb| {
5875 let mut buf = rb.borrow_mut();
5876 buf.resize(cols, 0);
5877 #[inline(always)]
5878 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
5879 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
5880 let u = unpack8::<B>(&data[blk * B..]);
5881 for k in 0..8 {
5882 chunk[k] = (u[k] - l) as i8 as u8;
5883 }
5884 }
5885 }
5886 match bw {
5887 3 => fill::<3>(data, l, &mut buf),
5888 4 => vbit_fill4(data, &mut buf),
5889 5 => fill::<5>(data, l, &mut buf),
5890 6 => fill::<6>(data, l, &mut buf),
5891 _ => unreachable!("vbit bit-width {bw} (validated at load)"),
5892 }
5893 let mut bi = 0usize;
5894 #[cfg(target_arch = "x86_64")]
5898 if avx2_enabled()
5899 && std::env::var("CMF_X86_BLOCKED")
5900 .map(|v| v != "0")
5901 .unwrap_or(true)
5902 {
5903 while bi + 4 <= acts.len() {
5904 let xs = [
5905 acts[bi].xq.as_slice(),
5906 acts[bi + 1].xq.as_slice(),
5907 acts[bi + 2].xq.as_slice(),
5908 acts[bi + 3].xq.as_slice(),
5909 ];
5910 let sxs = [
5911 acts[bi].sx,
5912 acts[bi + 1].sx,
5913 acts[bi + 2].sx,
5914 acts[bi + 3].sx,
5915 ];
5916 let d = unsafe {
5917 if vnni_tiles_enabled() {
5918 dot_q4b_row_1x4_sx_vnni(
5919 &buf,
5920 &bytes[sc_off..],
5921 r * ng,
5922 ng,
5923 xs,
5924 sxs,
5925 )
5926 } else {
5927 dot_q4b_row_1x4_sx_avx2(
5928 &buf,
5929 &bytes[sc_off..],
5930 r * ng,
5931 ng,
5932 xs,
5933 sxs,
5934 )
5935 }
5936 };
5937 for k in 0..4 {
5938 let act = &acts[bi + k];
5939 let mut dot = d[k];
5940 for &(j, xv) in &act.outliers {
5941 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
5942 }
5943 unsafe { *out_addr.at((bi + k) * rows + r) = dot };
5945 }
5946 bi += 4;
5947 }
5948 }
5949 while bi < acts.len() {
5950 let act = &acts[bi];
5951 let mut dot = 0f32;
5952 for g in 0..ng {
5953 let d = dot_i8_i8(
5954 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
5955 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
5956 ) as f32
5957 * act.sx;
5958 dot += d * gscale(r, g);
5959 }
5960 for &(j, xv) in &act.outliers {
5961 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
5962 }
5963 unsafe { *out_addr.at(bi * rows + r) = dot };
5965 bi += 1;
5966 }
5967 });
5968 }
5969 };
5970 dispatch_rows(pool, rows, &run);
5971 return;
5972 }
5973
5974 let out_addr = SendMut(out.as_mut_ptr());
5975 let run = move |start: usize, end: usize| {
5976 ROW_F32.with(|rb| {
5977 let mut buf = rb.borrow_mut();
5978 buf.resize(cols, 0.0);
5979 for r in start..end {
5980 decode_f32(r, &mut buf);
5981 for bi in 0..b {
5982 let x = &xs_all[bi * cols..(bi + 1) * cols];
5983 let mut dot = 0f32;
5984 for g in 0..ng {
5985 let mut gd = 0f32;
5986 for k in 0..GROUP_SIZE {
5987 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
5988 }
5989 dot += gd * gscale(r, g);
5990 }
5991 unsafe { *out_addr.at(bi * rows + r) = dot };
5993 }
5994 }
5995 })
5996 };
5997 dispatch_rows(pool, rows, &run);
5998}
5999
6000pub(crate) fn gpu_batch_job<'a>(
6004 t: &'a QTensor,
6005 x: &[f32],
6006) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
6007 match t {
6008 QTensor::Mapped {
6009 model,
6010 idx,
6011 dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
6012 rows,
6013 cols,
6014 row_scale,
6015 col_field,
6016 ..
6017 } => Some((
6018 model.clone(),
6019 crate::gpu::BatchJob {
6020 idx: *idx,
6021 rows: *rows,
6022 cols: *cols,
6023 row_scale,
6024 xs: prescale(x, col_field, *dt).into_owned(),
6025 q1: false,
6026 },
6027 )),
6028 QTensor::Mapped {
6030 model,
6031 idx,
6032 dtype: TensorDtype::Q1,
6033 rows,
6034 cols,
6035 ..
6036 } => Some((
6037 model.clone(),
6038 crate::gpu::BatchJob {
6039 idx: *idx,
6040 rows: *rows,
6041 cols: *cols,
6042 row_scale: &[],
6043 xs: x.to_vec(),
6044 q1: true,
6045 },
6046 )),
6047 _ => None,
6048 }
6049}
6050
6051thread_local! {
6052 static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
6053 static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
6054}
6055
6056pub(crate) fn prescale<'a>(
6057 x: &'a [f32],
6058 col_field: &[f32],
6059 dtype: TensorDtype,
6060) -> std::borrow::Cow<'a, [f32]> {
6061 if dtype == TensorDtype::Q8_2f {
6062 x.iter().zip(col_field).map(|(a, c)| a * c).collect()
6063 } else {
6064 std::borrow::Cow::Borrowed(x)
6065 }
6066}
6067
6068pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
6071 x: &[f32],
6072 col_field: &[f32],
6073 dtype: TensorDtype,
6074 buf_id: u8,
6075 f: F,
6076) -> R {
6077 if dtype == TensorDtype::Q8_2f {
6078 if buf_id == 1 {
6079 PRESCALE_BUF1.with(|b| {
6080 let mut buf = b.borrow_mut();
6081 buf.clear();
6082 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
6083 f(&buf)
6084 })
6085 } else {
6086 PRESCALE_BUF2.with(|b| {
6087 let mut buf = b.borrow_mut();
6088 buf.clear();
6089 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
6090 f(&buf)
6091 })
6092 }
6093 } else {
6094 f(x)
6095 }
6096}
6097
6098#[cfg(target_arch = "x86_64")]
6103pub(crate) fn avx2_enabled() -> bool {
6104 use std::sync::OnceLock;
6105 static ON: OnceLock<bool> = OnceLock::new();
6106 *ON.get_or_init(|| {
6107 std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
6108 && std::arch::is_x86_feature_detected!("avx2")
6109 && std::arch::is_x86_feature_detected!("fma")
6110 })
6111}
6112
6113#[cfg(target_arch = "x86_64")]
6118fn avx2_a8w8_enabled() -> bool {
6119 use std::sync::OnceLock;
6120 static ON: OnceLock<bool> = OnceLock::new();
6121 *ON.get_or_init(|| {
6122 avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
6123 })
6124}
6125
6126#[inline]
6130pub(crate) fn a8w8_enabled() -> bool {
6131 #[cfg(target_arch = "aarch64")]
6132 {
6133 sdot_enabled()
6134 }
6135 #[cfg(target_arch = "x86_64")]
6136 {
6137 avx2_a8w8_enabled()
6138 }
6139 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
6140 {
6141 false
6142 }
6143}
6144
6145#[inline]
6148#[allow(unreachable_code)]
6149fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
6150 #[cfg(target_arch = "aarch64")]
6151 unsafe {
6152 return dot_i8_sdot(w, xq);
6153 }
6154 #[cfg(target_arch = "x86_64")]
6155 unsafe {
6156 if avx512vnni_enabled() {
6157 return dot_i8_i8_vnni(w, xq);
6158 }
6159 return dot_i8_i8_avx2(w, xq);
6160 }
6161 w.iter()
6162 .zip(xq)
6163 .map(|(&a, &b)| (a as i8) as i32 * b as i32)
6164 .sum()
6165}
6166
6167#[cfg(target_arch = "x86_64")]
6171fn avx512vnni_enabled() -> bool {
6172 use std::sync::OnceLock;
6173 static ON: OnceLock<bool> = OnceLock::new();
6174 *ON.get_or_init(|| {
6175 std::env::var("CMF_AVX512")
6176 .map(|v| v != "0")
6177 .unwrap_or(true)
6178 && std::arch::is_x86_feature_detected!("avx512f")
6179 && std::arch::is_x86_feature_detected!("avx512bw")
6180 && std::arch::is_x86_feature_detected!("avx512vl")
6181 && std::arch::is_x86_feature_detected!("avx512vnni")
6182 })
6183}
6184
6185#[cfg(target_arch = "x86_64")]
6193fn vnni_tiles_enabled() -> bool {
6194 use std::sync::OnceLock;
6195 static ON: OnceLock<bool> = OnceLock::new();
6196 *ON.get_or_init(|| {
6197 std::env::var("CMF_VNNI_TILES")
6198 .map(|v| v != "0")
6199 .unwrap_or(true)
6200 && avx512vnni_enabled()
6201 })
6202}
6203
6204#[cfg(target_arch = "x86_64")]
6209#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6210#[inline]
6211unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
6212 unsafe {
6214 use core::arch::x86_64::*;
6215 let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
6216 let hi128 = _mm256_extracti128_si256::<1>(d);
6217 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6218 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6219 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6220 _mm_cvtsi128_si32(s32)
6221 }
6222}
6223
6224#[cfg(target_arch = "x86_64")]
6229#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6230unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
6231 unsafe {
6233 use core::arch::x86_64::*;
6234 let n = w.len();
6235 let mut j = 0usize;
6236 let mut total: i32;
6237 {
6242 #[inline(always)]
6243 unsafe fn step(
6244 w: *const u8,
6245 x: *const i8,
6246 acc: core::arch::x86_64::__m512i,
6247 ) -> core::arch::x86_64::__m512i {
6248 unsafe {
6249 use core::arch::x86_64::*;
6250 let wv = _mm512_loadu_si512(w as *const _);
6251 let xv = _mm512_loadu_si512(x as *const _);
6252 let aw = _mm512_abs_epi8(wv);
6253 let neg = _mm512_movepi8_mask(wv);
6254 let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
6255 _mm512_dpbusd_epi32(acc, aw, sx)
6256 }
6257 }
6258 let (mut a0, mut a1, mut a2, mut a3) = (
6259 _mm512_setzero_si512(),
6260 _mm512_setzero_si512(),
6261 _mm512_setzero_si512(),
6262 _mm512_setzero_si512(),
6263 );
6264 while j + 256 <= n {
6265 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
6266 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
6267 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
6268 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
6269 j += 256;
6270 }
6271 while j + 64 <= n {
6272 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
6273 j += 64;
6274 }
6275 let s01 = _mm512_add_epi32(a0, a1);
6276 let s23 = _mm512_add_epi32(a2, a3);
6277 total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
6278 }
6279 if j + 32 <= n {
6281 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
6282 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
6283 let d = _mm256_dpbusd_epi32(
6284 _mm256_setzero_si256(),
6285 _mm256_abs_epi8(wv),
6286 _mm256_sign_epi8(xv, wv),
6287 );
6288 let hi128 = _mm256_extracti128_si256::<1>(d);
6289 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6290 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6291 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6292 total += _mm_cvtsi128_si32(s32);
6293 j += 32;
6294 }
6295 while j < n {
6296 total += (w[j] as i8) as i32 * xq[j] as i32;
6297 j += 1;
6298 }
6299 total
6300 }
6301}
6302
6303#[cfg(target_arch = "x86_64")]
6305#[target_feature(enable = "avx2,fma")]
6306unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
6307 unsafe {
6309 use core::arch::x86_64::*;
6310 let n = x.len();
6311 let wp = w.as_ptr();
6312 let xp = x.as_ptr();
6313 let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
6314 let mut j = 0usize;
6315 while j + 16 <= n {
6316 let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
6317 let lo = _mm256_cvtepi8_epi32(wb);
6318 let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
6319 a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
6320 a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
6321 j += 16;
6322 }
6323 let acc = _mm256_add_ps(a0, a1);
6324 let hi128 = _mm256_extractf128_ps::<1>(acc);
6325 let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
6326 let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
6327 let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
6328 let mut sum = _mm_cvtss_f32(s32);
6329 while j < n {
6330 sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
6331 j += 1;
6332 }
6333 sum
6334 }
6335}
6336
6337#[cfg(target_arch = "x86_64")]
6342#[target_feature(enable = "avx2")]
6343unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
6344 unsafe {
6346 use core::arch::x86_64::*;
6347 let n = w.len();
6348 let ones = _mm256_set1_epi16(1);
6349 let mut acc = _mm256_setzero_si256();
6350 let mut j = 0usize;
6351 while j + 32 <= n {
6352 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
6353 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
6354 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6355 acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
6356 j += 32;
6357 }
6358 let hi128 = _mm256_extracti128_si256::<1>(acc);
6359 let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
6360 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6361 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6362 let mut s = _mm_cvtsi128_si32(s32);
6363 while j < n {
6364 s += (w[j] as i8) as i32 * xq[j] as i32;
6365 j += 1;
6366 }
6367 s
6368 }
6369}
6370
6371#[cfg(target_arch = "aarch64")]
6375#[target_feature(enable = "neon,i8mm")]
6376unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
6377 unsafe {
6379 use core::arch::aarch64::*;
6380 use core::arch::asm;
6381 let n = w0.len();
6382 let w0p = w0.as_ptr() as *const i8;
6383 let w1p = w1.as_ptr() as *const i8;
6384 let mut acc01 = vdupq_n_s32(0);
6387 let mut acc23 = vdupq_n_s32(0);
6388 let mut i = 0usize;
6389 while i + 8 <= n {
6390 let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
6391 let xb01 = vcombine_s8(
6392 vld1_s8(xs[0].as_ptr().add(i)),
6393 vld1_s8(xs[1].as_ptr().add(i)),
6394 );
6395 let xb23 = vcombine_s8(
6396 vld1_s8(xs[2].as_ptr().add(i)),
6397 vld1_s8(xs[3].as_ptr().add(i)),
6398 );
6399 asm!(
6400 "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
6401 "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
6402 a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
6403 w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
6404 options(pure, nomem, nostack),
6405 );
6406 i += 8;
6407 }
6408 let mut out = [[0i32; 4]; 2];
6409 let a01: [i32; 4] = core::mem::transmute(acc01);
6410 let a23: [i32; 4] = core::mem::transmute(acc23);
6411 out[0][0] = a01[0];
6412 out[0][1] = a01[1];
6413 out[1][0] = a01[2];
6414 out[1][1] = a01[3];
6415 out[0][2] = a23[0];
6416 out[0][3] = a23[1];
6417 out[1][2] = a23[2];
6418 out[1][3] = a23[3];
6419 if i < n {
6420 for (k, x) in xs.iter().enumerate() {
6421 for j in i..n {
6422 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
6423 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
6424 }
6425 }
6426 }
6427 out
6428 }
6429}
6430
6431#[cfg(target_arch = "aarch64")]
6435#[target_feature(enable = "neon,dotprod")]
6436unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
6437 unsafe {
6439 use core::arch::aarch64::*;
6440 use core::arch::asm;
6441 let n = w0.len();
6442 let w0p = w0.as_ptr() as *const i8;
6443 let w1p = w1.as_ptr() as *const i8;
6444 let mut acc = [[vdupq_n_s32(0); 4]; 2];
6445 let mut i = 0usize;
6446 while i + 16 <= n {
6447 let wv0 = vld1q_s8(w0p.add(i));
6448 let wv1 = vld1q_s8(w1p.add(i));
6449 for (k, x) in xs.iter().enumerate() {
6450 let xv = vld1q_s8(x.as_ptr().add(i));
6451 let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
6452 asm!(
6453 "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
6454 "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
6455 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6456 w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
6457 options(pure, nomem, nostack),
6458 );
6459 acc[0][k] = a0;
6460 acc[1][k] = a1;
6461 }
6462 i += 16;
6463 }
6464 let mut out = [[0i32; 4]; 2];
6465 for r in 0..2 {
6466 for k in 0..4 {
6467 out[r][k] = vaddvq_s32(acc[r][k]);
6468 }
6469 }
6470 if i < n {
6471 for (k, x) in xs.iter().enumerate() {
6472 for j in i..n {
6473 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
6474 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
6475 }
6476 }
6477 }
6478 out
6479 }
6480}
6481
6482#[cfg(target_arch = "x86_64")]
6488#[target_feature(enable = "avx2")]
6489unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
6490 unsafe {
6492 use core::arch::x86_64::*;
6493 let n = w0.len();
6494 let ones = _mm256_set1_epi16(1);
6495 let mut acc = [[_mm256_setzero_si256(); 4]; 2];
6496 let mut j = 0usize;
6497 while j + 32 <= n {
6498 let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
6499 let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
6500 let aw0 = _mm256_abs_epi8(wv0);
6501 let aw1 = _mm256_abs_epi8(wv1);
6502 for (k, x) in xs.iter().enumerate() {
6503 let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
6504 let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
6505 acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
6506 let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
6507 acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
6508 }
6509 j += 32;
6510 }
6511 let mut out = [[0i32; 4]; 2];
6512 for r in 0..2 {
6513 for k in 0..4 {
6514 let a = acc[r][k];
6515 let hi128 = _mm256_extracti128_si256::<1>(a);
6516 let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
6517 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6518 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6519 out[r][k] = _mm_cvtsi128_si32(s32);
6520 }
6521 }
6522 if j < n {
6523 for (k, x) in xs.iter().enumerate() {
6524 for i in j..n {
6525 out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
6526 out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
6527 }
6528 }
6529 }
6530 out
6531 }
6532}
6533
6534#[cfg(target_arch = "x86_64")]
6539#[inline]
6540fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
6541 let dot = if avx512vnni_enabled() && row.len() >= 64 {
6542 (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
6543 } else {
6544 unsafe { dot_i8_i8_avx2(row, &act.xq) }
6545 };
6546 let mut acc = dot as f32 * act.sx;
6547 for &(j, xv) in &act.outliers {
6548 acc += (row[j] as i8) as f32 * xv;
6549 }
6550 acc
6551}
6552
6553#[cfg(target_arch = "x86_64")]
6557#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6558unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
6559 unsafe {
6561 use core::arch::x86_64::*;
6562 let n = w.len();
6563 let flip = _mm512_set1_epi8(-128); #[inline(always)]
6565 unsafe fn step(
6566 w: *const u8,
6567 x: *const i8,
6568 flip: core::arch::x86_64::__m512i,
6569 acc: core::arch::x86_64::__m512i,
6570 ) -> core::arch::x86_64::__m512i {
6571 unsafe {
6572 use core::arch::x86_64::*;
6573 let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
6574 _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
6575 }
6576 }
6577 let (mut a0, mut a1, mut a2, mut a3) = (
6578 _mm512_setzero_si512(),
6579 _mm512_setzero_si512(),
6580 _mm512_setzero_si512(),
6581 _mm512_setzero_si512(),
6582 );
6583 let mut j = 0usize;
6584 while j + 256 <= n {
6585 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
6586 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
6587 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
6588 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
6589 j += 256;
6590 }
6591 while j + 64 <= n {
6592 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
6593 j += 64;
6594 }
6595 let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
6596 _mm512_add_epi32(a0, a1),
6597 _mm512_add_epi32(a2, a3),
6598 ));
6599 while j < n {
6601 total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
6602 j += 1;
6603 }
6604 total
6605 }
6606}
6607
6608#[cfg(target_arch = "x86_64")]
6614#[target_feature(enable = "avx2")]
6615unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
6616 unsafe {
6619 use core::arch::x86_64::*;
6620 let lomask = _mm_set1_epi8(0x0F);
6621 let eight = _mm256_set1_epi8(8);
6622 let ones = _mm256_set1_epi16(1);
6623 let mut acc = 0f32;
6624 for gi in 0..gpr {
6625 let g = g0 + gi;
6626 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
6627 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
6628 let lo = _mm_and_si128(b, lomask);
6629 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
6630 let w = _mm256_sub_epi8(
6631 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
6632 eight,
6633 );
6634 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6635 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
6636 let d = _mm256_madd_epi16(p16, ones);
6637 let hi128 = _mm256_extracti128_si256::<1>(d);
6638 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6639 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6640 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6641 acc += _mm_cvtsi128_si32(s32) as f32 * s;
6642 }
6643 acc
6644 }
6645}
6646
6647#[cfg(target_arch = "x86_64")]
6650#[target_feature(enable = "avx2")]
6651unsafe fn dot_q4_row_avx2_2(
6652 packed: &[u8],
6653 scales: &[u8],
6654 g0: usize,
6655 gpr: usize,
6656 xq1: &[i8],
6657 xq2: &[i8],
6658) -> (f32, f32) {
6659 unsafe {
6661 use core::arch::x86_64::*;
6662 let lomask = _mm_set1_epi8(0x0F);
6663 let eight = _mm256_set1_epi8(8);
6664 let ones = _mm256_set1_epi16(1);
6665 let (mut acc1, mut acc2) = (0f32, 0f32);
6666 #[inline(always)]
6667 unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
6668 unsafe {
6669 use core::arch::x86_64::*;
6670 let hi128 = _mm256_extracti128_si256::<1>(d);
6671 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6672 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6673 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6674 _mm_cvtsi128_si32(s32)
6675 }
6676 }
6677 for gi in 0..gpr {
6678 let g = g0 + gi;
6679 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
6680 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
6681 let lo = _mm_and_si128(b, lomask);
6682 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
6683 let w = _mm256_sub_epi8(
6684 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
6685 eight,
6686 );
6687 let aw = _mm256_abs_epi8(w);
6688 let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6689 let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6690 let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
6691 let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
6692 acc1 += hsum(d1) as f32 * s;
6693 acc2 += hsum(d2) as f32 * s;
6694 }
6695 (acc1, acc2)
6696 }
6697}
6698
6699#[cfg(target_arch = "x86_64")]
6701fn q8_range_avx2(
6702 q: &[u8],
6703 row_scale: &[f32],
6704 act: &SplitAct,
6705 cols: usize,
6706 out_addr: SendMut,
6707 start: usize,
6708 end: usize,
6709) {
6710 for o in start..end {
6711 let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
6712 unsafe { *out_addr.at(o) = v };
6714 }
6715}
6716
6717#[cfg(target_arch = "x86_64")]
6719#[allow(clippy::too_many_arguments)]
6720fn q8_range2_avx2(
6721 q: &[u8],
6722 row_scale: &[f32],
6723 a1: &SplitAct,
6724 a2: &SplitAct,
6725 cols: usize,
6726 p1: SendMut,
6727 p2: SendMut,
6728 start: usize,
6729 end: usize,
6730) {
6731 for o in start..end {
6732 let row = &q[o * cols..(o + 1) * cols];
6733 unsafe {
6735 *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
6736 *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
6737 }
6738 }
6739}
6740
6741#[cfg(target_arch = "aarch64")]
6752fn i8mm_enabled() -> bool {
6753 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6754 *ON.get_or_init(|| {
6755 std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
6756 && std::arch::is_aarch64_feature_detected!("i8mm")
6757 })
6758}
6759
6760#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
6764fn sdot_enabled() -> bool {
6765 use std::sync::OnceLock;
6766 static ON: OnceLock<bool> = OnceLock::new();
6767 *ON.get_or_init(|| {
6768 let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
6769 if !want {
6770 return false;
6771 }
6772
6773 #[cfg(target_arch = "aarch64")]
6774 {
6775 if std::arch::is_aarch64_feature_detected!("dotprod") {
6776 return true;
6777 }
6778 #[cfg(target_os = "android")]
6779 {
6780 if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
6781 if cpuinfo.lines().any(|l| {
6782 (l.starts_with("Features") || l.starts_with("features"))
6783 && l.contains("asimddp")
6784 }) {
6785 return true;
6786 }
6787 }
6788 }
6789 false
6790 }
6791 #[cfg(not(target_arch = "aarch64"))]
6792 {
6793 false
6794 }
6795 })
6796}
6797
6798struct SplitAct {
6803 xq: Vec<i8>,
6804 sx: f32,
6805 outliers: Vec<(usize, f32)>,
6806 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
6809 xsum: i32,
6810}
6811
6812thread_local! {
6813 static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
6816 const { std::cell::RefCell::new(Vec::new()) };
6817}
6818
6819impl Drop for SplitAct {
6820 fn drop(&mut self) {
6821 let buf = std::mem::take(&mut self.xq);
6822 if buf.capacity() > 0 {
6823 XQ_FREE.with(|f| {
6824 let mut f = f.borrow_mut();
6825 if f.len() < 16 {
6826 f.push(buf);
6827 }
6828 });
6829 }
6830 }
6831}
6832
6833fn split_act(x: &[f32]) -> SplitAct {
6834 let n = x.len();
6835 let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
6836 let thr = 8.0 * rms;
6837 let mut outliers: Vec<(usize, f32)> = Vec::new();
6841 let mut amax = 0f32;
6842 for (j, &v) in x.iter().enumerate() {
6843 let a = v.abs();
6844 if a > thr {
6845 outliers.push((j, v));
6846 } else if a > amax {
6847 amax = a;
6848 }
6849 }
6850 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
6851 let inv = 1.0 / sx;
6852 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
6853 xq.clear();
6854 xq.reserve(n);
6855 if outliers.is_empty() {
6856 xq.extend(
6857 x.iter()
6858 .map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8),
6859 );
6860 } else {
6861 xq.extend(x.iter().map(|&v| {
6863 if v.abs() > thr {
6864 0
6865 } else {
6866 (v * inv).round().clamp(-127.0, 127.0) as i8
6867 }
6868 }));
6869 }
6870 let xsum = xq.iter().map(|&v| v as i32).sum();
6871 SplitAct {
6872 xq,
6873 sx,
6874 outliers,
6875 xsum,
6876 }
6877}
6878
6879fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
6880 let n = x.len();
6881 let rms = (x
6882 .iter()
6883 .zip(col)
6884 .map(|(&a, &c)| {
6885 let v = a * c;
6886 (v * v) as f64
6887 })
6888 .sum::<f64>()
6889 / n.max(1) as f64)
6890 .sqrt() as f32;
6891 let thr = 8.0 * rms;
6892
6893 let mut outliers = Vec::new();
6894 let mut amax = 0f32;
6895 for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
6896 let v = a * c;
6897 let s = v.abs();
6898 if s > thr {
6899 outliers.push((j, v));
6900 } else if s > amax {
6901 amax = s;
6902 }
6903 }
6904
6905 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
6906 let inv = 1.0 / sx;
6907 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
6908 xq.clear();
6909 xq.reserve(n);
6910 if outliers.is_empty() {
6911 xq.extend(
6912 x.iter()
6913 .zip(col)
6914 .map(|(&a, &c)| ((a * c) * inv).round().clamp(-127.0, 127.0) as i8),
6915 );
6916 } else {
6917 xq.extend(x.iter().zip(col).map(|(&a, &c)| {
6918 let v = a * c;
6919 if v.abs() > thr {
6920 0
6921 } else {
6922 (v * inv).round().clamp(-127.0, 127.0) as i8
6923 }
6924 }));
6925 }
6926 let xsum = xq.iter().map(|&v| v as i32).sum();
6927 SplitAct {
6928 xq,
6929 sx,
6930 outliers,
6931 xsum,
6932 }
6933}
6934
6935#[cfg(target_arch = "aarch64")]
6938#[target_feature(enable = "neon,dotprod")]
6939unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
6940 unsafe {
6942 use core::arch::aarch64::*;
6943 use core::arch::asm;
6944 let wp = w.as_ptr() as *const i8;
6945 let n = w.len();
6946 let (mut a0, mut a1, mut a2, mut a3) = (
6947 vdupq_n_s32(0),
6948 vdupq_n_s32(0),
6949 vdupq_n_s32(0),
6950 vdupq_n_s32(0),
6951 );
6952 let mut i = 0;
6953 while i + 64 <= n {
6954 let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
6955 let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
6956 let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
6957 let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
6958 asm!(
6959 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6960 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6961 "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
6962 "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
6963 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
6964 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6965 w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
6966 options(pure, nomem, nostack),
6967 );
6968 i += 64;
6969 }
6970 while i + 16 <= n {
6971 let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
6972 asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
6973 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
6974 i += 16;
6975 }
6976 let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
6977 while i < n {
6978 s += (*wp.add(i)) as i32 * xq[i] as i32;
6979 i += 1;
6980 }
6981 s
6982 }
6983}
6984
6985#[cfg(target_arch = "aarch64")]
6989#[target_feature(enable = "neon,dotprod")]
6990unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
6991 unsafe {
6993 use core::arch::aarch64::*;
6994 use core::arch::asm;
6995 let n = xq.len();
6996 let px = xq.as_ptr();
6997 let (p0, p1, p2, p3) = (
6998 w0.as_ptr() as *const i8,
6999 w1.as_ptr() as *const i8,
7000 w2.as_ptr() as *const i8,
7001 w3.as_ptr() as *const i8,
7002 );
7003 let (mut a0, mut a1, mut a2, mut a3) = (
7004 vdupq_n_s32(0),
7005 vdupq_n_s32(0),
7006 vdupq_n_s32(0),
7007 vdupq_n_s32(0),
7008 );
7009 let mut i = 0;
7010 while i + 16 <= n {
7011 let x = vld1q_s8(px.add(i));
7012 let v0 = vld1q_s8(p0.add(i));
7013 let v1 = vld1q_s8(p1.add(i));
7014 let v2 = vld1q_s8(p2.add(i));
7015 let v3 = vld1q_s8(p3.add(i));
7016 asm!(
7017 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
7018 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
7019 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
7020 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
7021 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
7022 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
7023 options(pure, nomem, nostack),
7024 );
7025 i += 16;
7026 }
7027 let mut r = [
7028 vaddvq_s32(a0),
7029 vaddvq_s32(a1),
7030 vaddvq_s32(a2),
7031 vaddvq_s32(a3),
7032 ];
7033 while i < n {
7034 let xi = *px.add(i) as i32;
7035 r[0] += (*p0.add(i)) as i32 * xi;
7036 r[1] += (*p1.add(i)) as i32 * xi;
7037 r[2] += (*p2.add(i)) as i32 * xi;
7038 r[3] += (*p3.add(i)) as i32 * xi;
7039 i += 1;
7040 }
7041 r
7042 }
7043}
7044
7045#[cfg(target_arch = "aarch64")]
7052#[target_feature(enable = "neon,dotprod")]
7053unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
7054 unsafe {
7057 use core::arch::aarch64::*;
7058 use core::arch::asm;
7059 let n = xq.len();
7060 let px = xq.as_ptr();
7061 let pg = g.as_ptr() as *const i8;
7062 let (mut a0, mut a1, mut a2, mut a3) = (
7063 vdupq_n_s32(0),
7064 vdupq_n_s32(0),
7065 vdupq_n_s32(0),
7066 vdupq_n_s32(0),
7067 );
7068 let mut i = 0;
7069 while i + 16 <= n {
7070 let x = vld1q_s8(px.add(i));
7071 let base = pg.add(4 * i);
7072 let v0 = vld1q_s8(base);
7073 let v1 = vld1q_s8(base.add(16));
7074 let v2 = vld1q_s8(base.add(32));
7075 let v3 = vld1q_s8(base.add(48));
7076 asm!(
7077 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
7078 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
7079 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
7080 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
7081 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
7082 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
7083 options(pure, nomem, nostack),
7084 );
7085 i += 16;
7086 }
7087 [
7088 vaddvq_s32(a0),
7089 vaddvq_s32(a1),
7090 vaddvq_s32(a2),
7091 vaddvq_s32(a3),
7092 ]
7093 }
7094}
7095
7096#[cfg(target_arch = "aarch64")]
7102fn q8_range_sdot(
7103 q: &[u8],
7104 rep: &[u8],
7105 row_scale: &[f32],
7106 act: &SplitAct,
7107 cols: usize,
7108 out_addr: SendMut,
7109 start: usize,
7110 end: usize,
7111) {
7112 let mut o = start;
7113 if !rep.is_empty() {
7116 while o < end && o % 4 != 0 {
7117 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
7118 unsafe { *out_addr.at(o) = v };
7119 o += 1;
7120 }
7121 }
7122 while o + 4 <= end {
7123 let r = if rep.is_empty() {
7124 unsafe {
7125 dot_i8_sdot_4rows(
7126 &q[o * cols..(o + 1) * cols],
7127 &q[(o + 1) * cols..(o + 2) * cols],
7128 &q[(o + 2) * cols..(o + 3) * cols],
7129 &q[(o + 3) * cols..(o + 4) * cols],
7130 &act.xq,
7131 )
7132 }
7133 } else {
7134 unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
7135 };
7136 for k in 0..4 {
7137 let mut acc = r[k] as f32 * act.sx;
7138 for &(j, xv) in &act.outliers {
7139 acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
7140 }
7141 unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
7143 }
7144 o += 4;
7145 }
7146 while o < end {
7147 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
7148 unsafe { *out_addr.at(o) = v };
7149 o += 1;
7150 }
7151}
7152
7153#[cfg(target_arch = "aarch64")]
7156#[allow(clippy::too_many_arguments)]
7157fn q8_range2_sdot(
7158 q: &[u8],
7159 row_scale: &[f32],
7160 a1: &SplitAct,
7161 a2: &SplitAct,
7162 cols: usize,
7163 p1: SendMut,
7164 p2: SendMut,
7165 start: usize,
7166 end: usize,
7167) {
7168 for o in start..end {
7169 let row = &q[o * cols..(o + 1) * cols];
7170 unsafe {
7172 *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
7173 *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
7174 }
7175 }
7176}
7177
7178#[allow(clippy::too_many_arguments)]
7180fn q8_range2_f32(
7181 q: &[u8],
7182 row_scale: &[f32],
7183 x1: &[f32],
7184 x2: &[f32],
7185 cols: usize,
7186 p1: SendMut,
7187 p2: SendMut,
7188 start: usize,
7189 end: usize,
7190) {
7191 for o in start..end {
7192 let row = &q[o * cols..(o + 1) * cols];
7193 unsafe {
7195 *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
7196 *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
7197 }
7198 }
7199}
7200
7201fn q8_range_f32(
7203 q: &[u8],
7204 row_scale: &[f32],
7205 xs: &[f32],
7206 cols: usize,
7207 out_addr: SendMut,
7208 start: usize,
7209 end: usize,
7210) {
7211 for o in start..end {
7212 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
7213 unsafe { *out_addr.at(o) = v };
7215 }
7216}
7217
7218#[cfg(target_arch = "aarch64")]
7221#[inline]
7222fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
7223 let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
7224 for &(j, xv) in &act.outliers {
7225 acc += (row[j] as i8) as f32 * xv;
7226 }
7227 acc
7228}
7229
7230#[cfg(target_arch = "aarch64")]
7238#[target_feature(enable = "neon,dotprod")]
7239unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
7240 unsafe {
7243 use core::arch::aarch64::*;
7244 use core::arch::asm;
7245 let lomask = vdupq_n_u8(0x0F);
7246 let eight = vdupq_n_s8(8);
7247 let mut acc = 0f32;
7248 for gi in 0..gpr {
7249 let g = g0 + gi;
7250 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7251 let b = vld1q_u8(packed.as_ptr().add(g * 16));
7252 let lo = vandq_u8(b, lomask);
7253 let hi = vshrq_n_u8::<4>(b);
7254 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
7255 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
7256 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
7257 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
7258 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7259 asm!(
7260 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
7261 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
7262 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7263 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
7264 options(pure, nomem, nostack),
7265 );
7266 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
7267 }
7268 acc
7269 }
7270}
7271
7272#[cfg(target_arch = "aarch64")]
7277#[target_feature(enable = "neon,dotprod")]
7278unsafe fn dot_q4_row_sdot2(
7279 packed: &[u8],
7280 scales: &[u8],
7281 g0: usize,
7282 gpr: usize,
7283 xq1: &[i8],
7284 xq2: &[i8],
7285) -> (f32, f32) {
7286 unsafe {
7289 use core::arch::aarch64::*;
7290 use core::arch::asm;
7291 let lomask = vdupq_n_u8(0x0F);
7292 let eight = vdupq_n_s8(8);
7293 let (mut acc1, mut acc2) = (0f32, 0f32);
7294 for gi in 0..gpr {
7295 let g = g0 + gi;
7296 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7297 let b = vld1q_u8(packed.as_ptr().add(g * 16));
7298 let lo = vandq_u8(b, lomask);
7299 let hi = vshrq_n_u8::<4>(b);
7300 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
7301 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
7302 let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
7303 let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
7304 let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
7305 let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
7306 let (mut a0, mut a1, mut b0, mut b1) = (
7307 vdupq_n_s32(0),
7308 vdupq_n_s32(0),
7309 vdupq_n_s32(0),
7310 vdupq_n_s32(0),
7311 );
7312 asm!(
7313 "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
7314 "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
7315 "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
7316 "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
7317 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7318 b0 = inout(vreg) b0, b1 = inout(vreg) b1,
7319 e0 = in(vreg) e0, e1 = in(vreg) e1,
7320 x10 = in(vreg) x10, x11 = in(vreg) x11,
7321 x20 = in(vreg) x20, x21 = in(vreg) x21,
7322 options(pure, nomem, nostack),
7323 );
7324 acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
7325 acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
7326 }
7327 (acc1, acc2)
7328 }
7329}
7330
7331#[inline]
7336pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
7337 #[cfg(target_arch = "aarch64")]
7338 unsafe {
7339 return axpy_i8_f32_neon(acc, row, w);
7340 }
7341 #[cfg(target_arch = "x86_64")]
7342 if avx2_enabled() {
7343 return unsafe { axpy_i8_f32_avx2(acc, row, w) };
7344 }
7345 #[allow(unreachable_code)]
7346 {
7347 for (a, &b) in acc.iter_mut().zip(row) {
7348 *a += w * b as f32;
7349 }
7350 }
7351}
7352
7353#[cfg(target_arch = "x86_64")]
7355#[target_feature(enable = "avx2,fma")]
7356unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
7357 unsafe {
7359 use core::arch::x86_64::*;
7360 let n = acc.len().min(row.len());
7361 let ap = acc.as_mut_ptr();
7362 let rp = row.as_ptr();
7363 let wv = _mm256_set1_ps(w);
7364 let mut j = 0usize;
7365 while j + 16 <= n {
7366 let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
7367 let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
7368 let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
7369 let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
7370 let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
7371 _mm256_storeu_ps(ap.add(j), v0);
7372 _mm256_storeu_ps(ap.add(j + 8), v1);
7373 j += 16;
7374 }
7375 while j < n {
7376 *ap.add(j) += w * (*rp.add(j)) as f32;
7377 j += 1;
7378 }
7379 }
7380}
7381
7382#[cfg(target_arch = "aarch64")]
7383#[target_feature(enable = "neon")]
7384unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
7385 unsafe {
7387 use core::arch::aarch64::*;
7388 let n = acc.len().min(row.len());
7389 let ap = acc.as_mut_ptr();
7390 let rp = row.as_ptr();
7391 let wv = vdupq_n_f32(w);
7392 let mut j = 0usize;
7393 while j + 16 <= n {
7394 let rb = vld1q_s8(rp.add(j));
7395 let lo = vmovl_s8(vget_low_s8(rb));
7396 let hi = vmovl_s8(vget_high_s8(rb));
7397 for (off, half) in [(0, lo), (8, hi)] {
7398 let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
7399 let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
7400 let o = j + off;
7401 vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
7402 vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
7403 }
7404 j += 16;
7405 }
7406 while j < n {
7407 *ap.add(j) += w * (*rp.add(j)) as f32;
7408 j += 1;
7409 }
7410 }
7411}
7412
7413#[inline]
7416pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
7417 #[cfg(target_arch = "aarch64")]
7418 unsafe {
7419 return dot_i8_f32_neon(w, x);
7420 }
7421 #[cfg(target_arch = "x86_64")]
7422 if avx2_enabled() {
7423 return unsafe { dot_i8_f32_avx2(w, x) };
7424 }
7425 #[allow(unreachable_code)]
7426 {
7427 let mut sum = 0.0f32;
7428 for (j, &b) in w.iter().enumerate() {
7429 sum += (b as i8) as f32 * x[j];
7430 }
7431 sum
7432 }
7433}
7434
7435#[inline]
7439fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
7440 #[cfg(target_arch = "aarch64")]
7441 unsafe {
7442 return dot_i8_col_f32_neon(w, x, col);
7443 }
7444 #[allow(unreachable_code)]
7445 {
7446 let mut sum = 0.0f32;
7447 for (j, &b) in w.iter().enumerate() {
7448 sum += (b as i8) as f32 * x[j] * col[j];
7449 }
7450 sum
7451 }
7452}
7453
7454#[cfg(target_arch = "aarch64")]
7455#[target_feature(enable = "neon")]
7456unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
7457 unsafe {
7459 use core::arch::aarch64::*;
7460 let n = x.len();
7461 let wp = w.as_ptr() as *const i8;
7462 let xp = x.as_ptr();
7463 let cp = col.as_ptr();
7464 let (mut a0, mut a1, mut a2, mut a3) = (
7465 vdupq_n_f32(0.0),
7466 vdupq_n_f32(0.0),
7467 vdupq_n_f32(0.0),
7468 vdupq_n_f32(0.0),
7469 );
7470 let mut j = 0usize;
7471 while j + 16 <= n {
7472 let wb = vld1q_s8(wp.add(j));
7473 let lo = vmovl_s8(vget_low_s8(wb));
7474 let hi = vmovl_s8(vget_high_s8(wb));
7475 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
7476 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
7477 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
7478 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
7479 a0 = vfmaq_f32(
7480 a0,
7481 w0,
7482 vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
7483 );
7484 a1 = vfmaq_f32(
7485 a1,
7486 w1,
7487 vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
7488 );
7489 a2 = vfmaq_f32(
7490 a2,
7491 w2,
7492 vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
7493 );
7494 a3 = vfmaq_f32(
7495 a3,
7496 w3,
7497 vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
7498 );
7499 j += 16;
7500 }
7501 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
7502 while j < n {
7503 sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
7504 j += 1;
7505 }
7506 sum
7507 }
7508}
7509
7510#[cfg(target_arch = "aarch64")]
7511#[target_feature(enable = "neon")]
7512unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
7513 unsafe {
7515 use core::arch::aarch64::*;
7516 let n = x.len();
7517 let wp = w.as_ptr() as *const i8;
7518 let xp = x.as_ptr();
7519 let (mut a0, mut a1, mut a2, mut a3) = (
7520 vdupq_n_f32(0.0),
7521 vdupq_n_f32(0.0),
7522 vdupq_n_f32(0.0),
7523 vdupq_n_f32(0.0),
7524 );
7525 let mut j = 0usize;
7526 while j + 16 <= n {
7527 let wb = vld1q_s8(wp.add(j));
7528 let lo = vmovl_s8(vget_low_s8(wb));
7529 let hi = vmovl_s8(vget_high_s8(wb));
7530 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
7531 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
7532 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
7533 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
7534 a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
7535 a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
7536 a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
7537 a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
7538 j += 16;
7539 }
7540 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
7541 while j < n {
7542 sum += (*wp.add(j)) as f32 * *xp.add(j);
7543 j += 1;
7544 }
7545 sum
7546 }
7547}
7548
7549#[allow(clippy::too_many_arguments)]
7550fn qmatvec(
7551 q: &[u8],
7552 rep: &[u8],
7553 row_scale: &[f32],
7554 x: &[f32],
7555 col_field: &[f32],
7556 dtype: TensorDtype,
7557 rows: usize,
7558 cols: usize,
7559 out: &mut [f32],
7560 pool: Option<&Pool>,
7561) {
7562 debug_assert_eq!(out.len(), rows);
7563 #[cfg(not(target_arch = "aarch64"))]
7564 let _ = rep;
7565
7566 #[cfg(target_arch = "aarch64")]
7567 if sdot_enabled() {
7568 let act = if dtype == TensorDtype::Q8_2f {
7569 split_act_q8_2f(x, col_field)
7570 } else {
7571 split_act(x)
7572 };
7573 let out_addr = SendMut(out.as_mut_ptr());
7574 let run_range = |start: usize, end: usize| {
7575 q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
7576 };
7577 match pool {
7578 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
7579 _ => run_range(0, rows),
7580 }
7581 return;
7582 }
7583 #[cfg(target_arch = "x86_64")]
7586 if avx2_a8w8_enabled() {
7587 let act = if dtype == TensorDtype::Q8_2f {
7588 split_act_q8_2f(x, col_field)
7589 } else {
7590 split_act(x)
7591 };
7592 let out_addr = SendMut(out.as_mut_ptr());
7593 let run_range = |start: usize, end: usize| {
7594 q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
7595 };
7596 match pool {
7597 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
7598 _ => run_range(0, rows),
7599 }
7600 return;
7601 }
7602
7603 prescale_with(x, col_field, dtype, 1, |xs| {
7604 let out_addr = SendMut(out.as_mut_ptr());
7605 let run_range = move |start: usize, end: usize| {
7606 for o in start..end {
7607 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
7608 unsafe { *out_addr.at(o) = v };
7610 }
7611 };
7612 match pool {
7613 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
7614 _ => run_range(0, rows),
7615 }
7616 });
7617}
7618
7619#[allow(clippy::too_many_arguments)]
7620fn qmatvec2(
7621 q: &[u8],
7622 row_scale: &[f32],
7623 x1: &[f32],
7624 x2: &[f32],
7625 col_field: &[f32],
7626 dtype: TensorDtype,
7627 rows: usize,
7628 cols: usize,
7629 o1: &mut [f32],
7630 o2: &mut [f32],
7631 pool: Option<&Pool>,
7632) {
7633 #[cfg(target_arch = "aarch64")]
7634 if sdot_enabled() {
7635 let a1s = if dtype == TensorDtype::Q8_2f {
7636 split_act_q8_2f(x1, col_field)
7637 } else {
7638 split_act(x1)
7639 };
7640 let a2s = if dtype == TensorDtype::Q8_2f {
7641 split_act_q8_2f(x2, col_field)
7642 } else {
7643 split_act(x2)
7644 };
7645 let p1 = SendMut(o1.as_mut_ptr());
7646 let p2 = SendMut(o2.as_mut_ptr());
7647 let run_range = |start: usize, end: usize| {
7648 q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
7649 };
7650 match pool {
7651 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
7652 _ => run_range(0, rows),
7653 }
7654 return;
7655 }
7656 #[cfg(target_arch = "x86_64")]
7657 if avx2_a8w8_enabled() {
7658 let a1s = if dtype == TensorDtype::Q8_2f {
7659 split_act_q8_2f(x1, col_field)
7660 } else {
7661 split_act(x1)
7662 };
7663 let a2s = if dtype == TensorDtype::Q8_2f {
7664 split_act_q8_2f(x2, col_field)
7665 } else {
7666 split_act(x2)
7667 };
7668 let p1 = SendMut(o1.as_mut_ptr());
7669 let p2 = SendMut(o2.as_mut_ptr());
7670 let run_range = |start: usize, end: usize| {
7671 q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
7672 };
7673 match pool {
7674 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
7675 _ => run_range(0, rows),
7676 }
7677 return;
7678 }
7679
7680 prescale_with(x1, col_field, dtype, 1, |x1s| {
7681 prescale_with(x2, col_field, dtype, 2, |x2s| {
7682 let p1 = SendMut(o1.as_mut_ptr());
7683 let p2 = SendMut(o2.as_mut_ptr());
7684 let run_range = move |start: usize, end: usize| {
7685 for o in start..end {
7686 let row = &q[o * cols..(o + 1) * cols];
7687 let s1 = dot_i8_f32(row, x1s) * row_scale[o];
7688 let s2 = dot_i8_f32(row, x2s) * row_scale[o];
7689 unsafe {
7691 *p1.at(o) = s1;
7692 *p2.at(o) = s2;
7693 }
7694 }
7695 };
7696 match pool {
7697 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
7698 _ => run_range(0, rows),
7699 }
7700 });
7701 });
7702}
7703
7704#[derive(Clone, Copy)]
7705struct SendMut(*mut f32);
7706unsafe impl Send for SendMut {}
7707unsafe impl Sync for SendMut {}
7708
7709impl SendMut {
7710 #[inline]
7711 fn at(self, i: usize) -> *mut f32 {
7712 unsafe { self.0.add(i) }
7713 }
7714}
7715
7716#[cfg(test)]
7717mod tests {
7718 use super::*;
7719
7720 #[test]
7721 fn f32_matvec_matches_matvec_rows_bitexact() {
7722 let (rows, cols) = (300, 40);
7723 let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
7724 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
7725 let qt = QTensor::from_f32(w.clone(), rows, cols);
7726
7727 let mut a = vec![0.0f32; rows];
7728 matvec_rows(None, &w, &x, &mut a);
7729 let mut b = vec![0.0f32; rows];
7730 qt.matvec(&x, &mut b, None);
7731 assert_eq!(a, b);
7732 }
7733
7734 #[test]
7735 fn sdot_kernel_exact_on_grid() {
7736 eprintln!("sdot_enabled = {}", sdot_enabled());
7741 let (rows, cols) = (9, 80); let w: Vec<u8> = (0..rows * cols)
7743 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
7744 .collect();
7745 let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
7746 let x: Vec<f32> = (0..cols)
7747 .map(|i| match i % 3 {
7748 0 => 1.0,
7749 1 => -1.0,
7750 _ => 0.0,
7751 })
7752 .collect();
7753 let mut a = vec![0.0f32; rows];
7754 qmatvec(
7755 &w,
7756 &[],
7757 &scales,
7758 &x,
7759 &[],
7760 TensorDtype::Q8Row,
7761 rows,
7762 cols,
7763 &mut a,
7764 None,
7765 );
7766 for o in 0..rows {
7767 let mut acc = 0.0f32;
7768 for j in 0..cols {
7769 acc += (w[o * cols + j] as i8) as f32 * x[j];
7770 }
7771 let expect = acc * scales[o];
7772 assert!(
7773 (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
7774 "row {o}: {} vs {expect}",
7775 a[o]
7776 );
7777 }
7778 }
7779
7780 #[test]
7781 fn q1_tbl_fast_path_matches_reference() {
7782 let (rows, cols) = (5, 256);
7787 let gpr = cols / GROUP_SIZE;
7788 let mut bytes = Vec::new();
7789 for t in 0..rows * gpr {
7790 let s = 0.007 + (t % 11) as f32 * 0.004;
7791 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
7792 for j in 0..4 {
7793 bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
7794 }
7795 }
7796 let x: Vec<f32> = (0..cols)
7797 .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
7798 .collect();
7799 let mut w = vec![0.0f32; rows * cols];
7800 cortiq_core::quant::dequant_q1(&bytes, &mut w);
7801 let mut got = vec![0.0f32; rows];
7802 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
7803 for o in 0..rows {
7804 let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
7805 assert!(
7806 (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
7807 "row {o}: {} vs {expect}",
7808 got[o]
7809 );
7810 }
7811 let b = 5usize;
7814 let mut xs_all = Vec::new();
7815 for bi in 0..b {
7816 xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
7817 }
7818 let mut mm = vec![0.0f32; b * rows];
7819 q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
7820 for bi in 0..b {
7821 let mut single = vec![0.0f32; rows];
7822 q1_matvec(
7823 &bytes,
7824 &xs_all[bi * cols..(bi + 1) * cols],
7825 rows,
7826 cols,
7827 &mut single,
7828 None,
7829 );
7830 assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
7831 }
7832 }
7833
7834 #[test]
7835 fn q1_kernels_match_exact_reference() {
7836 let (rows, cols) = (7, 96);
7838 let gpr = cols / GROUP_SIZE;
7839 let mut bytes = Vec::new();
7840 for t in 0..rows * gpr {
7841 let s = 0.01 + (t % 13) as f32 * 0.003;
7842 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
7843 for j in 0..4 {
7844 bytes.push(((t * 31 + j * 97) % 251) as u8);
7845 }
7846 }
7847 let x: Vec<f32> = (0..cols)
7849 .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
7850 .collect();
7851 let mut w = vec![0.0f32; rows * cols];
7853 cortiq_core::quant::dequant_q1(&bytes, &mut w);
7854 let mut expect = vec![0.0f32; rows];
7855 for o in 0..rows {
7856 expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
7857 }
7858 let mut got = vec![0.0f32; rows];
7859 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
7860 for o in 0..rows {
7861 assert!(
7862 (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
7863 "row {o}: {} vs {}",
7864 got[o],
7865 expect[o]
7866 );
7867 }
7868 let x2: Vec<f32> = x.iter().map(|v| -v).collect();
7870 let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
7871 q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
7872 assert_eq!(a1, got);
7873 let mut xs = x.clone();
7874 xs.extend_from_slice(&x2);
7875 let mut mm = vec![0.0f32; 2 * rows];
7876 q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
7877 assert_eq!(&mm[..rows], got.as_slice());
7878 assert_eq!(&mm[rows..], a2.as_slice());
7879 }
7880
7881 #[test]
7882 fn repack_is_bit_identical() {
7883 let (rows, cols) = (267, 96); let w: Vec<u8> = (0..rows * cols)
7889 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
7890 .collect();
7891 let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
7892 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
7893 let rep = q8_repack_layout(&w, rows, cols);
7894 for g in 0..rows / 4 {
7896 for c in 0..cols / 16 {
7897 for lane in 0..4 {
7898 assert_eq!(
7899 &rep[g * 4 * cols + c * 64 + lane * 16
7900 ..g * 4 * cols + c * 64 + lane * 16 + 16],
7901 &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
7902 );
7903 }
7904 }
7905 }
7906 let mut a = vec![0.0f32; rows];
7907 qmatvec(
7908 &w,
7909 &[],
7910 &scales,
7911 &x,
7912 &[],
7913 TensorDtype::Q8Row,
7914 rows,
7915 cols,
7916 &mut a,
7917 None,
7918 );
7919 let mut b = vec![0.0f32; rows];
7920 qmatvec(
7921 &w,
7922 &rep,
7923 &scales,
7924 &x,
7925 &[],
7926 TensorDtype::Q8Row,
7927 rows,
7928 cols,
7929 &mut b,
7930 None,
7931 );
7932 assert_eq!(a, b, "full-range repack output diverged");
7933
7934 #[cfg(target_arch = "aarch64")]
7935 if sdot_enabled() {
7936 let act = split_act(&x);
7938 let mut c1 = vec![0.0f32; rows];
7939 let mut c2 = vec![0.0f32; rows];
7940 q8_range_sdot(
7941 &w,
7942 &[],
7943 &scales,
7944 &act,
7945 cols,
7946 SendMut(c1.as_mut_ptr()),
7947 3,
7948 rows - 2,
7949 );
7950 q8_range_sdot(
7951 &w,
7952 &rep,
7953 &scales,
7954 &act,
7955 cols,
7956 SendMut(c2.as_mut_ptr()),
7957 3,
7958 rows - 2,
7959 );
7960 assert_eq!(c1, c2, "unaligned-range repack output diverged");
7961 }
7962 }
7963
7964 #[test]
7965 fn sdot_a8w8_noise_is_bounded() {
7966 let (rows, cols) = (16, 512);
7970 let w: Vec<u8> = (0..rows * cols)
7971 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
7972 .collect();
7973 let scales = vec![0.01f32; rows];
7974 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
7975 let mut a = vec![0.0f32; rows];
7976 qmatvec(
7977 &w,
7978 &[],
7979 &scales,
7980 &x,
7981 &[],
7982 TensorDtype::Q8Row,
7983 rows,
7984 cols,
7985 &mut a,
7986 None,
7987 );
7988 let (mut num, mut den) = (0f64, 0f64);
7989 for o in 0..rows {
7990 let mut acc = 0.0f32;
7991 for j in 0..cols {
7992 acc += (w[o * cols + j] as i8) as f32 * x[j];
7993 }
7994 let expect = acc * scales[o];
7995 num += ((a[o] - expect) as f64).powi(2);
7996 den += (expect as f64).powi(2);
7997 }
7998 let rel = (num / den.max(1e-12)).sqrt();
7999 assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
8000 }
8001
8002 #[test]
8003 fn i8_dot_neon_matches_scalar() {
8004 let n = 100;
8005 let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
8006 let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
8007 let mut scalar = 0.0f32;
8008 for j in 0..n {
8009 scalar += (w[j] as i8) as f32 * x[j];
8010 }
8011 let fast = dot_i8_f32(&w, &x);
8012 assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
8013 }
8014
8015 #[test]
8017 fn vbitmatvec_matches_full_dequant() {
8018 let (rows, cols) = (6, 64);
8019 let ng = cols / GROUP_SIZE;
8020 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
8022 let mut bytes = bits.clone();
8023 for g in 0..rows * ng {
8024 let s = 0.02 + 0.001 * g as f32;
8025 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
8026 }
8027 for r in 0..rows {
8028 let b = bits[r] as usize;
8029 let (mut acc, mut nb) = (0u64, 0usize);
8030 let mut rowbytes = Vec::new();
8031 for i in 0..cols {
8032 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
8033 acc = (acc << b) | v;
8034 nb += b;
8035 while nb >= 8 {
8036 nb -= 8;
8037 rowbytes.push(((acc >> nb) & 0xFF) as u8);
8038 }
8039 }
8040 if nb > 0 {
8041 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
8042 }
8043 bytes.extend_from_slice(&rowbytes);
8044 }
8045 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
8046
8047 let mut reference = vec![0f32; rows * cols];
8048 cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
8049 let mut expect = vec![0f32; rows];
8050 for r in 0..rows {
8051 expect[r] = reference[r * cols..(r + 1) * cols]
8052 .iter()
8053 .zip(&x)
8054 .map(|(w, xv)| w * xv)
8055 .sum();
8056 }
8057 let mut got = vec![0f32; rows];
8058 let offsets = vbit_row_offsets(&bytes, rows, cols);
8059 vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
8060 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
8064 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
8065 for r in 0..rows {
8066 assert!(
8067 (got[r] - expect[r]).abs() < tol * scale,
8068 "row {r}: {} vs {}",
8069 got[r],
8070 expect[r]
8071 );
8072 }
8073 }
8074
8075 #[test]
8080 #[cfg(target_arch = "x86_64")]
8081 fn vbit_matmat_blocked_matches_per_row() {
8082 let (rows, cols, b) = (64usize, 128usize, 9usize);
8083 let ng = cols / GROUP_SIZE;
8084 let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
8085 let mut bytes = bits.clone();
8086 for g in 0..rows * ng {
8087 let sc = 0.02 + 0.0005 * g as f32;
8088 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
8089 }
8090 for r in 0..rows {
8091 let bw = bits[r] as usize;
8092 let (mut acc, mut nb) = (0u64, 0usize);
8093 let mut rowbytes = Vec::new();
8094 for i in 0..cols {
8095 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
8096 acc = (acc << bw) | v;
8097 nb += bw;
8098 while nb >= 8 {
8099 nb -= 8;
8100 rowbytes.push(((acc >> nb) & 0xFF) as u8);
8101 }
8102 }
8103 if nb > 0 {
8104 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
8105 }
8106 bytes.extend_from_slice(&rowbytes);
8107 }
8108 let x: Vec<f32> = (0..b * cols)
8109 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
8110 .collect();
8111 let offsets = vbit_row_offsets(&bytes, rows, cols);
8112 let mut y_a = vec![0f32; b * rows];
8113 let mut y_b = vec![0f32; b * rows];
8114 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
8115 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
8116 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
8117 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
8118 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
8119 let max_d = y_a
8120 .iter()
8121 .zip(&y_b)
8122 .map(|(p, q)| (p - q).abs())
8123 .fold(0.0f32, f32::max);
8124 assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
8125 }
8126
8127 #[test]
8135 fn q4t_matmat_blocked_matches_per_row() {
8136 let (rows, cols, b) = (16usize, 64usize, 9usize);
8137 let gpr = cols / GROUP_SIZE;
8138 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
8139 for r in 0..rows {
8140 for g in 0..gpr {
8141 let t = (r * gpr + g) * Q4_TILE;
8142 let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
8143 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
8144 for k in 0..16 {
8145 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
8146 }
8147 }
8148 }
8149 let x: Vec<f32> = (0..b * cols)
8150 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
8151 .collect();
8152 let mut y_blk = vec![0f32; b * rows];
8153 let mut y_row = vec![0f32; b * rows];
8154 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
8155 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
8156 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
8157 q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
8158 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
8159 assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
8160 }
8161
8162 #[cfg(target_os = "macos")]
8166 #[test]
8167 fn q4t_matmat_accel_matches_dequant_reference() {
8168 if !accel_gemm_enabled() {
8169 return; }
8171 let (rows, cols, b) = (512usize, 1024usize, 8usize); let gpr = cols / GROUP_SIZE;
8173 let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
8174 for r in 0..rows {
8175 for g in 0..gpr {
8176 let t = (r * gpr + g) * Q4_TILE;
8177 let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
8178 bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
8179 for k in 0..16 {
8180 bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
8181 }
8182 }
8183 }
8184 let x: Vec<f32> = (0..b * cols)
8185 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
8186 .collect();
8187 let mut got = vec![0f32; b * rows];
8188 q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
8189 let mut w = vec![0f32; rows * cols];
8191 for r in 0..rows {
8192 for g in 0..gpr {
8193 let t = (r * gpr + g) * Q4_TILE;
8194 let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
8195 for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
8196 w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
8197 w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
8198 (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
8199 }
8200 }
8201 }
8202 for bi in 0..b {
8203 for r in 0..rows {
8204 let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
8205 let d = (got[bi * rows + r] - want).abs();
8206 assert!(
8207 d <= want.abs().max(1.0) * 1e-4,
8208 "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
8209 got[bi * rows + r]
8210 );
8211 }
8212 }
8213 }
8214
8215 #[test]
8216 fn q4matvec_matches_full_dequant() {
8217 let (rows, cols) = (8, 64);
8218 let groups = rows * cols / GROUP_SIZE;
8219 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
8221 for i in 0..groups * 16 {
8222 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
8223 }
8224 for g in 0..groups {
8225 let s = 0.01 + 0.003 * g as f32;
8226 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
8227 }
8228 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
8229
8230 let mut reference = vec![0.0f32; rows * cols];
8231 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
8232 let mut expect = vec![0.0f32; rows];
8233 for r in 0..rows {
8234 expect[r] = reference[r * cols..(r + 1) * cols]
8235 .iter()
8236 .zip(&x)
8237 .map(|(w, xv)| w * xv)
8238 .sum();
8239 }
8240
8241 let mut got = vec![0.0f32; rows];
8242 q4matvec(&bytes, &x, rows, cols, &mut got, None);
8243 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
8247 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
8248 for r in 0..rows {
8249 assert!(
8250 (got[r] - expect[r]).abs() < tol * scale,
8251 "row {r}: {} vs {}",
8252 got[r],
8253 expect[r]
8254 );
8255 }
8256 }
8257
8258 #[test]
8261 fn vbitmatvec2_equals_two_singles() {
8262 let (rows, cols) = (6, 64);
8263 let ng = cols / GROUP_SIZE;
8264 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
8265 let mut bytes = bits.clone();
8266 for g in 0..rows * ng {
8267 let s = 0.02 + 0.001 * g as f32;
8268 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
8269 }
8270 for r in 0..rows {
8271 let b = bits[r] as usize;
8272 let (mut acc, mut nb) = (0u64, 0usize);
8273 let mut rowbytes = Vec::new();
8274 for i in 0..cols {
8275 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
8276 acc = (acc << b) | v;
8277 nb += b;
8278 while nb >= 8 {
8279 nb -= 8;
8280 rowbytes.push(((acc >> nb) & 0xFF) as u8);
8281 }
8282 }
8283 if nb > 0 {
8284 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
8285 }
8286 bytes.extend_from_slice(&rowbytes);
8287 }
8288 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
8289 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
8290 let offsets = vbit_row_offsets(&bytes, rows, cols);
8291
8292 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
8293 vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
8294 vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
8295 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
8296 vbitmatvec2(
8297 &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
8298 );
8299 assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
8300 assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
8301 }
8302
8303 #[test]
8305 fn q4matvec2_equals_two_singles() {
8306 let (rows, cols) = (8, 128);
8307 let groups = rows * cols / GROUP_SIZE;
8308 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
8309 for i in 0..groups * 16 {
8310 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
8311 }
8312 for g in 0..groups {
8313 let s = 0.01 + 0.003 * g as f32;
8314 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
8315 }
8316 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
8319 x1[9] = 250.0;
8320 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
8321
8322 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
8323 q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
8324 q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
8325 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
8326 q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
8327 assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
8328 assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
8329 }
8330
8331 #[test]
8334 fn matvec_many_equals_separate_matvecs() {
8335 use crate::pool::Pool;
8336 let (r1, r2, cols) = (300, 200, 64);
8337 let mk = |salt: usize, rows: usize| {
8338 QTensor::from_f32(
8339 (0..rows * cols)
8340 .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
8341 .collect(),
8342 rows,
8343 cols,
8344 )
8345 };
8346 let (a, b) = (mk(1, r1), mk(5, r2));
8347 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
8348 let pool = Pool::new(3);
8349
8350 let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
8351 a.matvec(&x, &mut ea, Some(&pool));
8352 b.matvec(&x, &mut eb, Some(&pool));
8353 let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
8354 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
8355 assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
8356 assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
8357 }
8358
8359 #[test]
8362 fn batched_matmat_equals_per_position_matvec() {
8363 let (rows, cols, b) = (8, 64, 5);
8364 let groups = rows * cols / GROUP_SIZE;
8366 let mut q4 = Vec::new();
8367 for i in 0..groups * 16 {
8368 q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
8369 }
8370 for g in 0..groups {
8371 q4.extend_from_slice(
8372 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
8373 );
8374 }
8375 let ng = cols / GROUP_SIZE;
8377 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
8378 let mut vb = bits.clone();
8379 for g in 0..rows * ng {
8380 vb.extend_from_slice(
8381 &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
8382 );
8383 }
8384 for r in 0..rows {
8385 let bw = bits[r] as usize;
8386 let (mut acc, mut nb) = (0u64, 0usize);
8387 let mut rowbytes = Vec::new();
8388 for i in 0..cols {
8389 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
8390 acc = (acc << bw) | v;
8391 nb += bw;
8392 while nb >= 8 {
8393 nb -= 8;
8394 rowbytes.push(((acc >> nb) & 0xFF) as u8);
8395 }
8396 }
8397 if nb > 0 {
8398 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
8399 }
8400 vb.extend_from_slice(&rowbytes);
8401 }
8402 let offsets = vbit_row_offsets(&vb, rows, cols);
8403
8404 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
8405
8406 let mut got = vec![0f32; b * rows];
8408 q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
8409 for bi in 0..b {
8410 let mut expect = vec![0f32; rows];
8411 q4matvec(
8412 &q4,
8413 &xs[bi * cols..(bi + 1) * cols],
8414 rows,
8415 cols,
8416 &mut expect,
8417 None,
8418 );
8419 assert_eq!(
8420 &got[bi * rows..(bi + 1) * rows],
8421 &expect[..],
8422 "q4 batch pos {bi}"
8423 );
8424 }
8425
8426 let mut got = vec![0f32; b * rows];
8428 vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
8429 for bi in 0..b {
8430 let mut expect = vec![0f32; rows];
8431 vbitmatvec(
8432 &vb,
8433 &offsets,
8434 &xs[bi * cols..(bi + 1) * cols],
8435 rows,
8436 cols,
8437 &mut expect,
8438 None,
8439 );
8440 assert_eq!(
8441 &got[bi * rows..(bi + 1) * rows],
8442 &expect[..],
8443 "vbit batch pos {bi}"
8444 );
8445 }
8446 }
8447
8448 #[test]
8452 fn q4_tiled_matches_q4_block_bitexact() {
8453 let (rows, cols, b) = (8usize, 128usize, 3usize);
8454 let groups = rows * cols / GROUP_SIZE;
8455 let mut split = Vec::with_capacity(groups * 18);
8456 for i in 0..groups * 16 {
8457 split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
8458 }
8459 for g in 0..groups {
8460 split.extend_from_slice(
8461 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
8462 );
8463 }
8464 let (packed, scales) = split.split_at(groups * 16);
8466 let mut tiled = Vec::with_capacity(groups * Q4_TILE);
8467 for g in 0..groups {
8468 tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
8469 tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
8470 }
8471
8472 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
8473 x1[9] = 250.0; let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
8475
8476 let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
8477 q4matvec(&split, &x1, rows, cols, &mut a, None);
8478 q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
8479 assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
8480
8481 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
8482 let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
8483 q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
8484 q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
8485 assert_eq!(a1, t1);
8486 assert_eq!(a2, t2);
8487
8488 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
8489 let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
8490 q4matmat(&split, &xs, b, rows, cols, &mut am, None);
8491 q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
8492 assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
8493 }
8494
8495 #[test]
8502 fn q4matvec_sdot_outlier_exact() {
8503 let (rows, cols) = (4, 128);
8504 let groups = rows * cols / GROUP_SIZE;
8505 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
8506 for i in 0..groups * 16 {
8507 bytes.push(((i * 11 + 5) % 256) as u8);
8508 }
8509 for g in 0..groups {
8510 let s = 0.02 + 0.002 * g as f32;
8511 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
8512 }
8513 let mut x: Vec<f32> = (0..cols)
8514 .map(|i| match i % 3 {
8515 0 => 1.0,
8516 1 => -1.0,
8517 _ => 0.0,
8518 })
8519 .collect();
8520 x[17] = 300.0; let mut reference = vec![0.0f32; rows * cols];
8523 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
8524 let mut expect = vec![0.0f32; rows];
8525 for r in 0..rows {
8526 expect[r] = reference[r * cols..(r + 1) * cols]
8527 .iter()
8528 .zip(&x)
8529 .map(|(w, xv)| w * xv)
8530 .sum();
8531 }
8532 let mut got = vec![0.0f32; rows];
8533 q4matvec(&bytes, &x, rows, cols, &mut got, None);
8534 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
8535 for r in 0..rows {
8536 assert!(
8537 (got[r] - expect[r]).abs() < 2e-3 * scale,
8538 "row {r}: {} vs {} (outlier term must be exact)",
8539 got[r],
8540 expect[r]
8541 );
8542 }
8543 }
8544
8545 #[test]
8549 fn q1t_matvec_matches_reference() {
8550 use cortiq_core::quant::{dequant_q1t, f32_to_f16};
8551 let (rows, cols) = (3usize, 64usize); let gpr = cols / GROUP_SIZE;
8553 let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
8554 let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
8556 let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
8557 let mut bytes = Vec::new();
8558 for r in 0..rows {
8559 for g in 0..gpr {
8560 bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
8561 let mut c = [0u8; 7];
8562 for k in 0..GROUP_SIZE {
8563 let code = if is_out(r * cols + g * GROUP_SIZE + k) {
8565 0
8566 } else {
8567 ((k + r * 3 + g) % 3) as u8 };
8569 cortiq_core::quant::q1t_pack(&mut c, k, code);
8570 }
8571 bytes.extend_from_slice(&c);
8572 }
8573 }
8574 let mut row_ptr = vec![0u32; rows + 1];
8577 for &(idx, _) in &outliers {
8578 row_ptr[idx as usize / cols + 1] += 1;
8579 }
8580 for r in 0..rows {
8581 row_ptr[r + 1] += row_ptr[r];
8582 }
8583 for &p in &row_ptr {
8584 bytes.extend_from_slice(&p.to_le_bytes());
8585 }
8586 for &(idx, v) in &outliers {
8587 bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
8588 bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
8589 }
8590
8591 let mut refw = vec![0f32; rows * cols];
8592 dequant_q1t(&bytes, rows, cols, &mut refw);
8593 let x: Vec<f32> = (0..cols)
8596 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
8597 .collect();
8598 let mut expect = vec![0f32; rows];
8599 for r in 0..rows {
8600 let mut a = 0.0f32;
8601 for j in 0..cols {
8602 a += refw[r * cols + j] * x[j];
8603 }
8604 expect[r] = a;
8605 }
8606 let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
8607 let mut got = vec![0f32; rows];
8608 q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
8609 for r in 0..rows {
8610 assert!(
8611 (got[r] - expect[r]).abs() < tol(expect[r]),
8612 "row {r}: {} vs {}",
8613 got[r],
8614 expect[r]
8615 );
8616 }
8617 let x2: Vec<f32> = x.iter().chain(x.iter().map(|v| v)).copied().collect();
8619 let mut gm = vec![0f32; 2 * rows];
8620 q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
8621 for r in 0..rows {
8622 assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
8623 assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
8624 }
8625 let xb: Vec<f32> = (0..cols)
8629 .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
8630 .collect();
8631 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
8632 q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
8633 q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
8634 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
8635 q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
8636 assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
8637 assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
8638 }
8639
8640 #[test]
8643 fn q1t_matvec2_odd_gpr_matches_singles() {
8644 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
8645 let (rows, cols) = (5usize, 96usize); let gpr = cols / GROUP_SIZE;
8647 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
8648 for r in 0..rows {
8649 for g in 0..gpr {
8650 bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
8651 let mut c = [0u8; 7];
8652 for k in 0..GROUP_SIZE {
8653 q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
8654 }
8655 bytes.extend_from_slice(&c);
8656 }
8657 }
8658 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
8659 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
8660 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
8661 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
8662 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
8663 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
8664 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
8665 assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
8666 assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
8667 }
8668
8669 #[test]
8673 #[ignore]
8674 fn q1t_matvec2_speed() {
8675 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
8676 use std::time::Instant;
8677 let (rows, cols) = (8192usize, 4096usize);
8678 let gpr = cols / GROUP_SIZE;
8679 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
8680 for r in 0..rows {
8681 for g in 0..gpr {
8682 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
8683 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
8684 let mut c = [0u8; 7];
8685 for k in 0..GROUP_SIZE {
8686 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
8687 }
8688 bytes.extend_from_slice(&c);
8689 }
8690 }
8691 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
8692 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
8693 let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
8694 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
8695 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
8697 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
8698 let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
8699 for _ in 0..8 {
8700 let t0 = Instant::now();
8701 q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
8702 t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
8703 let t1 = Instant::now();
8704 q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
8705 q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
8706 t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
8707 }
8708 assert_eq!(p1, s1);
8709 assert_eq!(p2, s2);
8710 println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
8711 }
8712
8713 #[test]
8717 #[ignore]
8718 fn q1t_matvec_speed() {
8719 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
8720 use std::time::Instant;
8721 let (rows, cols) = (8192usize, 4096usize); let gpr = cols / GROUP_SIZE;
8723 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
8724 for r in 0..rows {
8725 for g in 0..gpr {
8726 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
8727 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
8728 let mut c = [0u8; 7];
8729 for k in 0..GROUP_SIZE {
8730 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
8731 }
8732 bytes.extend_from_slice(&c);
8733 }
8734 }
8735 let (n, stride) = (rows * cols, 40usize); let mut row_ptr = vec![0u32; rows + 1];
8737 let mut idx = 0usize;
8738 while idx < n {
8739 row_ptr[idx / cols + 1] += 1;
8740 idx += stride;
8741 }
8742 for r in 0..rows {
8743 row_ptr[r + 1] += row_ptr[r];
8744 }
8745 for &p in &row_ptr {
8746 bytes.extend_from_slice(&p.to_le_bytes());
8747 }
8748 let mut idx = 0usize;
8749 while idx < n {
8750 bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
8751 bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
8752 idx += stride;
8753 }
8754 let x: Vec<f32> = (0..cols)
8757 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
8758 .collect();
8759 let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
8760
8761 let slow = |out: &mut [f32]| {
8763 let mut buf = vec![0f32; cols];
8764 for r in 0..rows {
8765 for g in 0..gpr {
8766 let off = (r * gpr + g) * Q1T_TILE;
8767 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8768 let codes = &bytes[off + 2..off + Q1T_TILE];
8769 for k in 0..GROUP_SIZE {
8770 buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
8771 1 => s,
8772 2 => -s,
8773 _ => 0.0,
8774 };
8775 }
8776 }
8777 out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
8778 + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
8779 }
8780 };
8781 let iters = 5;
8782 let mut a = vec![0f32; rows];
8783 slow(&mut a); let t = Instant::now();
8785 for _ in 0..iters {
8786 slow(&mut a);
8787 }
8788 let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
8789
8790 let mut b = vec![0f32; rows];
8791 q1t_matvec(&bytes, &x, rows, cols, &mut b, None); let t = Instant::now();
8793 for _ in 0..iters {
8794 q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
8795 }
8796 let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
8797
8798 for r in 0..rows {
8799 assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
8800 }
8801 println!(
8802 "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms fused-LUT {fast_ms:.2} ms => {:.2}x",
8803 slow_ms / fast_ms
8804 );
8805 }
8806}