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::Q8Row
306 | TensorDtype::Q8_2f,
307 rows,
308 cols,
309 ..
310 } => Some((*idx, *rows, *cols)),
311 _ => None,
312 }
313 }
314
315 pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
320 match self {
321 Self::Mapped {
322 idx,
323 dtype: TensorDtype::Q8Row,
324 rows,
325 cols,
326 row_scale,
327 col_field,
328 ..
329 } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
330 _ => None,
331 }
332 }
333
334 pub fn rows(&self) -> usize {
335 match self {
336 Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
337 }
338 }
339
340 pub fn cols(&self) -> usize {
341 match self {
342 Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
343 }
344 }
345
346 pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
349 match self {
350 Self::Mapped {
351 model,
352 idx,
353 dtype: TensorDtype::Q1,
354 ..
355 } => Some((model, *idx)),
356 _ => None,
357 }
358 }
359
360 pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
364 match self {
365 Self::Mapped {
366 model,
367 idx,
368 dtype: TensorDtype::Q8Row,
369 row_scale,
370 ..
371 } => Some((model, *idx, 0, row_scale.as_slice())),
372 Self::Mapped {
373 model,
374 idx,
375 dtype: TensorDtype::Q1,
376 ..
377 } => Some((model, *idx, 1, &[])),
378 Self::Mapped {
379 model,
380 idx,
381 dtype: TensorDtype::Q4Tiled,
382 ..
383 } => Some((model, *idx, 2, &[])),
384 Self::Mapped {
385 model,
386 idx,
387 dtype: TensorDtype::Q4Block,
388 ..
389 } => Some((model, *idx, 2, &[])),
390 Self::Mapped {
391 model,
392 idx,
393 dtype: TensorDtype::Q1T,
394 ..
395 } => Some((model, *idx, 3, &[])),
396 _ => None,
397 }
398 }
399
400 pub fn as_f32(&self) -> Option<&[f32]> {
403 match self {
404 Self::F32 { data, .. } => Some(data),
405 Self::Mapped { .. } => None,
406 }
407 }
408
409 fn quant_bytes(&self) -> &[u8] {
410 match self {
411 Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
412 Self::F32 { .. } => unreachable!("quant_bytes on F32"),
413 }
414 }
415
416 pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
418 let cols = self.cols();
419 debug_assert_eq!(dst.len(), cols);
420 match self {
421 Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
422 Self::Mapped {
423 dtype,
424 row_scale,
425 col_field,
426 vbit_offsets,
427 ..
428 } => {
429 if *dtype == TensorDtype::Q4Tiled {
430 let bytes = self.quant_bytes();
431 let gpr = cols / GROUP_SIZE;
432 for gi in 0..gpr {
433 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
434 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
435 for (k, &b) in tile[2..].iter().enumerate() {
436 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
437 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
438 }
439 }
440 return;
441 }
442 if *dtype == TensorDtype::Q4Block {
443 let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
444 let gpr = cols / GROUP_SIZE;
445 for gi in 0..gpr {
446 let g = r * gpr + gi;
447 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
448 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
449 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
450 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
451 }
452 }
453 return;
454 }
455 if *dtype == TensorDtype::Q1 {
456 let bytes = self.quant_bytes();
457 let gpr = cols / GROUP_SIZE;
458 for gi in 0..gpr {
459 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
460 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
461 for (j, &b) in tile[2..].iter().enumerate() {
462 for k in 0..8 {
463 dst[gi * GROUP_SIZE + j * 8 + k] =
464 (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
465 }
466 }
467 }
468 return;
469 }
470 if *dtype == TensorDtype::Q1T {
471 let bytes = self.quant_bytes();
472 let gpr = cols / GROUP_SIZE;
473 let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
474 for gi in 0..gpr {
475 let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
476 let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
477 bytes[off],
478 bytes[off + 1],
479 ]));
480 let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
481 for k in 0..GROUP_SIZE {
482 dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
483 {
484 1 => s,
485 2 => -s,
486 _ => 0.0,
487 };
488 }
489 }
490 let rows = self.rows();
492 let entries = base_len + (rows + 1) * 4;
493 if entries <= bytes.len() {
494 let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
495 let r0 = u32::from_le_bytes([
496 ptrs[r * 4],
497 ptrs[r * 4 + 1],
498 ptrs[r * 4 + 2],
499 ptrs[r * 4 + 3],
500 ]) as usize;
501 let r1 = u32::from_le_bytes([
502 ptrs[(r + 1) * 4],
503 ptrs[(r + 1) * 4 + 1],
504 ptrs[(r + 1) * 4 + 2],
505 ptrs[(r + 1) * 4 + 3],
506 ]) as usize;
507 let off = entries + r0 * 4;
508 for i in 0..r1 - r0 {
509 let item = &bytes[off + i * 4..off + i * 4 + 4];
510 let c = u16::from_le_bytes([item[0], item[1]]) as usize;
511 let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
512 item[2], item[3],
513 ]));
514 if c < cols {
515 dst[c] = v;
516 }
517 }
518 }
519 return;
520 }
521 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
522 let bytes = self.quant_bytes();
523 let rows = self.rows();
524 let ng = cols / GROUP_SIZE;
525 let bits = &bytes[..rows];
526 let sc_off = rows;
527 let off = vbit_offsets[r];
530 let b = bits[r] as usize;
531 let l = ((1usize << (b - 1)) - 1) as f32;
532 let data = &bytes[off..];
533 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
534 for (i, d) in dst.iter_mut().enumerate() {
535 while nbits < b {
536 acc = (acc << 8) | data[idx] as u64;
537 idx += 1;
538 nbits += 8;
539 }
540 let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
541 nbits -= b;
542 let so = (r * ng + i / GROUP_SIZE) * 2;
543 let sv = f16_to_f32(u16::from_le_bytes([
544 bytes[sc_off + so],
545 bytes[sc_off + so + 1],
546 ]));
547 *d = (u - l) * sv;
548 }
549 return;
550 }
551 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
552 let s = row_scale[r];
553 match dtype {
554 TensorDtype::Q8Row => {
555 for (d, &b) in dst.iter_mut().zip(q) {
556 *d = (b as i8) as f32 * s;
557 }
558 }
559 TensorDtype::Q8_2f => {
560 for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
561 *d = (b as i8) as f32 * s * col_field[i];
562 }
563 }
564 _ => unreachable!(),
565 }
566 }
567 }
568 }
569
570 pub fn sparse_col_ok(&self) -> bool {
575 match self {
576 Self::F32 { .. } => true,
577 Self::Mapped { dtype, .. } => {
578 matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
579 }
580 }
581 }
582
583 pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
587 let inter = self.cols();
588 let hidden = self.rows();
589 debug_assert_eq!(out.len(), hidden);
590 match self {
591 Self::F32 { data, .. } => {
592 for (k, o) in out.iter_mut().enumerate() {
593 *o += w * data[k * inter + c];
594 }
595 }
596 Self::Mapped {
597 dtype,
598 row_scale,
599 col_field,
600 ..
601 } => {
602 let q = self.quant_bytes();
603 let colf = if *dtype == TensorDtype::Q8_2f {
604 col_field[c]
605 } else {
606 1.0
607 };
608 let wc = w * colf;
609 for (k, o) in out.iter_mut().enumerate() {
610 let b = q[k * inter + c] as i8 as f32;
611 *o += wc * b * row_scale[k];
612 }
613 }
614 }
615 }
616
617 pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
621 let cols = self.cols();
622 match self {
623 Self::F32 { data, .. } => {
624 let row = &data[r * cols..(r + 1) * cols];
625 row.iter().zip(x).map(|(w, v)| w * v).sum()
626 }
627 Self::Mapped {
628 dtype,
629 row_scale,
630 col_field,
631 ..
632 } => match dtype {
633 TensorDtype::Q8Row => {
634 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
635 dot_i8_f32(q, x) * row_scale[r]
636 }
637 TensorDtype::Q8_2f => {
638 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
639 dot_i8_col_f32(q, x, col_field) * row_scale[r]
640 }
641 _ => {
642 self.row_f32(r, scratch);
643 scratch.iter().zip(x).map(|(w, v)| w * v).sum()
644 }
645 },
646 }
647 }
648
649 pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
652 match self {
653 Self::F32 { data, .. } => matvec_rows(pool, data, x, out),
654 Self::Mapped {
655 model,
656 idx,
657 dtype,
658 rows,
659 cols,
660 row_scale,
661 col_field,
662 vbit_offsets,
663 repack,
664 } => {
665 let _ = (model, idx);
666 if *dtype == TensorDtype::Q4Block {
667 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
671 let t0 = std::time::Instant::now();
672 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
673 crate::gpu::ProbeArm::Gpu => {
674 if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
675 crate::gpu::probe_record(
676 crate::gpu::OpClass::Matvec,
677 true,
678 t0.elapsed(),
679 );
680 return;
681 }
682 }
683 crate::gpu::ProbeArm::CpuTimed => {
684 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
685 crate::gpu::probe_record(
686 crate::gpu::OpClass::Matvec,
687 false,
688 t0.elapsed(),
689 );
690 return;
691 }
692 crate::gpu::ProbeArm::Cpu => {}
693 }
694 }
695 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
696 return;
697 }
698 if *dtype == TensorDtype::Q4Tiled {
699 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
700 return;
701 }
702 if *dtype == TensorDtype::Q1 {
703 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
708 let t0 = std::time::Instant::now();
709 let arm = if crate::gpu::q1_force() {
710 crate::gpu::ProbeArm::Gpu
711 } else {
712 crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
713 };
714 match arm {
715 crate::gpu::ProbeArm::Gpu => {
716 if crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out) {
717 crate::gpu::probe_record(
718 crate::gpu::OpClass::Matvec,
719 true,
720 t0.elapsed(),
721 );
722 return;
723 }
724 }
725 crate::gpu::ProbeArm::CpuTimed => {
726 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
727 crate::gpu::probe_record(
728 crate::gpu::OpClass::Matvec,
729 false,
730 t0.elapsed(),
731 );
732 return;
733 }
734 crate::gpu::ProbeArm::Cpu => {}
735 }
736 }
737 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
738 return;
739 }
740 if *dtype == TensorDtype::Q1T {
741 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
745 let t0 = std::time::Instant::now();
746 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
747 crate::gpu::ProbeArm::Gpu => {
748 if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
749 q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
750 crate::gpu::probe_record(
751 crate::gpu::OpClass::Matvec,
752 true,
753 t0.elapsed(),
754 );
755 return;
756 }
757 }
758 crate::gpu::ProbeArm::CpuTimed => {
759 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
760 crate::gpu::probe_record(
761 crate::gpu::OpClass::Matvec,
762 false,
763 t0.elapsed(),
764 );
765 return;
766 }
767 crate::gpu::ProbeArm::Cpu => {}
768 }
769 }
770 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
771 return;
772 }
773 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
774 vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
775 return;
776 }
777 let xs = prescale(x, col_field, *dtype);
778 if *rows >= crate::gpu::min_rows()
783 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
784 && std::env::var("CMF_GPU_LMHEAD")
785 .map(|v| v != "0")
786 .unwrap_or(true)
787 && crate::gpu::enabled_here()
788 {
789 let t0 = std::time::Instant::now();
792 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
793 crate::gpu::ProbeArm::Gpu => {}
794 crate::gpu::ProbeArm::CpuTimed => {
795 qmatvec(
796 self.quant_bytes(),
797 repack,
798 row_scale,
799 x,
800 col_field,
801 *dtype,
802 *rows,
803 *cols,
804 out,
805 pool,
806 );
807 crate::gpu::probe_record(
808 crate::gpu::OpClass::Matvec,
809 false,
810 t0.elapsed(),
811 );
812 return;
813 }
814 crate::gpu::ProbeArm::Cpu => {
815 qmatvec(
816 self.quant_bytes(),
817 repack,
818 row_scale,
819 x,
820 col_field,
821 *dtype,
822 *rows,
823 *cols,
824 out,
825 pool,
826 );
827 return;
828 }
829 }
830 let frac = std::env::var("CMF_GPU_SPLIT")
831 .ok()
832 .and_then(|v| v.parse::<f32>().ok())
833 .unwrap_or(0.5)
834 .clamp(0.0, 1.0);
835 let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
836 let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
837 let bytes = self.quant_bytes();
838 let ok = std::thread::scope(|sc| {
839 let g = sc.spawn(|| {
840 crate::gpu::q8_matvec_range(
841 model,
842 *idx,
843 cpu_rows,
844 &row_scale[cpu_rows..],
845 &xs,
846 *rows - cpu_rows,
847 *cols,
848 out_gpu,
849 )
850 });
851 if cpu_rows > 0 {
852 let rep_cpu = if repack.is_empty() {
855 &[][..]
856 } else {
857 &repack[..(cpu_rows / 4) * 4 * *cols]
858 };
859 qmatvec(
860 &bytes[..cpu_rows * *cols],
861 rep_cpu,
862 &row_scale[..cpu_rows],
863 x,
864 col_field,
865 *dtype,
866 cpu_rows,
867 *cols,
868 out_cpu,
869 pool,
870 );
871 }
872 g.join().unwrap_or(false)
873 });
874 if ok {
875 crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
876 return;
877 }
878 qmatvec(
881 &bytes[cpu_rows * *cols..(*rows) * *cols],
882 &[],
883 &row_scale[cpu_rows..],
884 x,
885 col_field,
886 *dtype,
887 *rows - cpu_rows,
888 *cols,
889 out_gpu,
890 pool,
891 );
892 return;
893 }
894 qmatvec(
895 self.quant_bytes(),
896 repack,
897 row_scale,
898 x,
899 col_field,
900 *dtype,
901 *rows,
902 *cols,
903 out,
904 pool,
905 );
906 }
907 }
908 }
909
910 pub fn matvec2(
912 &self,
913 x1: &[f32],
914 x2: &[f32],
915 o1: &mut [f32],
916 o2: &mut [f32],
917 pool: Option<&Pool>,
918 ) {
919 match self {
920 Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
921 Self::Mapped {
922 dtype,
923 rows,
924 cols,
925 row_scale,
926 col_field,
927 vbit_offsets,
928 ..
929 } => {
930 if *dtype == TensorDtype::Q4Block {
931 q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
932 return;
933 }
934 if *dtype == TensorDtype::Q4Tiled {
935 q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
936 return;
937 }
938 if *dtype == TensorDtype::Q1 {
939 q1_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
940 return;
941 }
942 if *dtype == TensorDtype::Q1T {
943 q1t_matvec(self.quant_bytes(), x1, *rows, *cols, o1, pool);
948 q1t_matvec(self.quant_bytes(), x2, *rows, *cols, o2, pool);
949 return;
950 }
951 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
952 vbitmatvec2(
953 self.quant_bytes(),
954 vbit_offsets,
955 x1,
956 x2,
957 *rows,
958 *cols,
959 o1,
960 o2,
961 pool,
962 );
963 return;
964 }
965 qmatvec2(
966 self.quant_bytes(),
967 row_scale,
968 x1,
969 x2,
970 col_field,
971 *dtype,
972 *rows,
973 *cols,
974 o1,
975 o2,
976 pool,
977 );
978 }
979 }
980 }
981}
982
983impl QTensor {
984 pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
989 let cols = self.cols();
990 let rows = self.rows();
991 debug_assert_eq!(xs_all.len(), b * cols);
992 debug_assert_eq!(out.len(), b * rows);
993 if crate::gptq_capture::capturing() {
997 if let Self::Mapped { model, idx, .. } = self {
998 crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
999 }
1000 }
1001 match self {
1002 Self::F32 { data, .. } => {
1003 let out_addr = SendMut(out.as_mut_ptr());
1004 let run = |start: usize, end: usize| {
1005 for o in start..end {
1006 let row = &data[o * cols..(o + 1) * cols];
1007 for bi in 0..b {
1008 let x = &xs_all[bi * cols..(bi + 1) * cols];
1009 let mut acc = 0f32;
1010 for j in 0..cols {
1011 acc += row[j] * x[j];
1012 }
1013 unsafe { *out_addr.at(bi * rows + o) = acc };
1014 }
1015 }
1016 };
1017 dispatch_rows(pool, rows, &run);
1018 }
1019 Self::Mapped {
1020 dtype,
1021 row_scale,
1022 col_field,
1023 vbit_offsets,
1024 ..
1025 } => {
1026 if *dtype == TensorDtype::Q4Block {
1027 q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1028 return;
1029 }
1030 if *dtype == TensorDtype::Q4Tiled {
1031 q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1032 return;
1033 }
1034 if *dtype == TensorDtype::Q1 {
1035 if b >= 32
1038 && b * rows * cols >= 128_000_000
1039 && cols % 64 == 0
1040 && crate::gpu::enabled_here()
1041 {
1042 if let Self::Mapped { model, idx, .. } = self {
1043 let t0 = std::time::Instant::now();
1044 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1045 crate::gpu::ProbeArm::Gpu => {
1046 if crate::gpu::q1_matmat(
1047 model, *idx, xs_all, b, rows, cols, out,
1048 ) {
1049 crate::gpu::probe_record(
1050 crate::gpu::OpClass::Matmat,
1051 true,
1052 t0.elapsed(),
1053 );
1054 return;
1055 }
1056 }
1057 crate::gpu::ProbeArm::CpuTimed => {
1058 q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1059 crate::gpu::probe_record(
1060 crate::gpu::OpClass::Matmat,
1061 false,
1062 t0.elapsed(),
1063 );
1064 return;
1065 }
1066 crate::gpu::ProbeArm::Cpu => {}
1067 }
1068 }
1069 }
1070 q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1071 return;
1072 }
1073 if *dtype == TensorDtype::Q1T {
1074 if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1077 if let Self::Mapped { model, idx, .. } = self {
1078 let t0 = std::time::Instant::now();
1079 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1080 crate::gpu::ProbeArm::Gpu => {
1081 if crate::gpu::q1t_matmat(
1082 model, *idx, xs_all, b, rows, cols, out,
1083 ) {
1084 crate::gpu::probe_record(
1085 crate::gpu::OpClass::Matmat,
1086 true,
1087 t0.elapsed(),
1088 );
1089 return;
1090 }
1091 }
1092 crate::gpu::ProbeArm::CpuTimed => {
1093 q1t_matmat(
1094 self.quant_bytes(),
1095 xs_all,
1096 b,
1097 rows,
1098 cols,
1099 out,
1100 pool,
1101 );
1102 crate::gpu::probe_record(
1103 crate::gpu::OpClass::Matmat,
1104 false,
1105 t0.elapsed(),
1106 );
1107 return;
1108 }
1109 crate::gpu::ProbeArm::Cpu => {}
1110 }
1111 }
1112 }
1113 q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1114 return;
1115 }
1116 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1117 vbitmatmat(
1118 self.quant_bytes(),
1119 vbit_offsets,
1120 xs_all,
1121 b,
1122 rows,
1123 cols,
1124 out,
1125 pool,
1126 );
1127 return;
1128 }
1129 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
1130 .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
1131 .collect();
1132 if b >= 8 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1138 if let Self::Mapped { model, idx, .. } = self {
1139 let t0 = std::time::Instant::now();
1140 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1141 crate::gpu::ProbeArm::Gpu
1142 if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
1143 && !crate::gpu::q8_resident_or_upload(model, *idx) =>
1144 {
1145 let q = self.quant_bytes();
1149 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1150 return;
1151 }
1152 crate::gpu::ProbeArm::Gpu => {
1153 let flat: Vec<f32> =
1154 pre.iter().flat_map(|v| v.iter().copied()).collect();
1155 if crate::gpu::q8_matmat(
1156 model, *idx, row_scale, &flat, b, rows, cols, out,
1157 ) {
1158 crate::gpu::probe_record(
1159 crate::gpu::OpClass::Matmat,
1160 true,
1161 t0.elapsed(),
1162 );
1163 return;
1164 }
1165 }
1166 crate::gpu::ProbeArm::CpuTimed => {
1167 let q = self.quant_bytes();
1168 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1169 crate::gpu::probe_record(
1170 crate::gpu::OpClass::Matmat,
1171 false,
1172 t0.elapsed(),
1173 );
1174 return;
1175 }
1176 crate::gpu::ProbeArm::Cpu => {}
1177 }
1178 }
1179 }
1180 let q = self.quant_bytes();
1181 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1182 }
1183 }
1184 }
1185}
1186
1187impl QTensor {
1188 pub fn matvec_many<const N: usize>(
1195 ts: [&QTensor; N],
1196 x: &[f32],
1197 mut outs: [&mut [f32]; N],
1198 pool: Option<&Pool>,
1199 ) {
1200 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
1201 let uniform_q8 = ts.iter().all(|t| {
1202 matches!(
1203 t,
1204 Self::Mapped {
1205 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
1206 ..
1207 }
1208 )
1209 });
1210 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
1211 let uniform_q4 = ts.iter().all(|t| {
1212 matches!(
1213 t,
1214 Self::Mapped {
1215 dtype: TensorDtype::Q4Block,
1216 ..
1217 }
1218 )
1219 });
1220 let uniform_vbit = ts.iter().all(|t| {
1221 matches!(
1222 t,
1223 Self::Mapped {
1224 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
1225 ..
1226 }
1227 )
1228 });
1229 let uniform_q1 = ts.iter().all(|t| {
1230 matches!(
1231 t,
1232 Self::Mapped {
1233 dtype: TensorDtype::Q1,
1234 ..
1235 }
1236 )
1237 });
1238 let uniform_q1t = ts.iter().all(|t| {
1239 matches!(
1240 t,
1241 Self::Mapped {
1242 dtype: TensorDtype::Q1T,
1243 ..
1244 }
1245 )
1246 });
1247 let Some(pool) = pool else {
1248 for (t, o) in ts.iter().zip(outs.iter_mut()) {
1249 t.matvec(x, o, None);
1250 }
1251 return;
1252 };
1253 if total_rows < 256
1254 || !(uniform_q8
1255 || uniform_f32
1256 || uniform_q4
1257 || uniform_vbit
1258 || uniform_q1
1259 || uniform_q1t)
1260 {
1261 for (t, o) in ts.iter().zip(outs.iter_mut()) {
1262 t.matvec(x, o, Some(pool));
1263 }
1264 return;
1265 }
1266
1267 if uniform_q1 {
1268 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1271 if a8w8_enabled() {
1272 let act = split_act(x);
1273 let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
1274 let (act, gsum) = (&act, &gsum);
1275 let closures: [_; N] = std::array::from_fn(|i| {
1276 let (bytes, gpr, out) =
1277 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1278 move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
1279 });
1280 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1281 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1282 pool.run_many(&parts);
1283 } else {
1284 let closures: [_; N] = std::array::from_fn(|i| {
1285 let (bytes, gpr, out) =
1286 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1287 move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
1288 });
1289 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1290 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1291 pool.run_many(&parts);
1292 }
1293 return;
1294 }
1295
1296 if uniform_q1t {
1297 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1301 const TILE: usize = cortiq_core::quant::Q1T_TILE;
1302 if a8w8_enabled() {
1303 let act = split_act(x);
1304 let act = &act;
1305 let x_ref = x;
1306 let closures: [_; N] = std::array::from_fn(|i| {
1307 let bytes = ts[i].quant_bytes();
1308 let (rows, cols) = (ts[i].rows(), ts[i].cols());
1309 let gpr = cols / GROUP_SIZE;
1310 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
1311 let out = outs_addr[i];
1312 move |s: usize, e: usize| {
1313 q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
1314 }
1315 });
1316 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1317 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1318 pool.run_many(&parts);
1319 } else {
1320 let x_ref = x;
1321 let closures: [_; N] = std::array::from_fn(|i| {
1322 let bytes = ts[i].quant_bytes();
1323 let (rows, cols) = (ts[i].rows(), ts[i].cols());
1324 let gpr = cols / GROUP_SIZE;
1325 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
1326 let out = outs_addr[i];
1327 move |s: usize, e: usize| {
1328 q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
1329 }
1330 });
1331 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1332 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1333 pool.run_many(&parts);
1334 }
1335 return;
1336 }
1337
1338 if uniform_q4 || uniform_vbit {
1339 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1340 if a8w8_enabled() {
1342 let act = split_act(x);
1343 let act = &act;
1344 if uniform_q4 {
1345 let closures: [_; N] = std::array::from_fn(|i| {
1346 let (packed, scales) =
1347 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1348 let (gpr, cols, out) =
1349 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
1350 move |s: usize, e: usize| {
1351 q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
1352 }
1353 });
1354 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1355 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1356 pool.run_many(&parts);
1357 } else {
1358 let closures: [_; N] = std::array::from_fn(|i| {
1359 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
1360 unreachable!()
1361 };
1362 let (bytes, rows, cols, out) = (
1363 ts[i].quant_bytes(),
1364 ts[i].rows(),
1365 ts[i].cols(),
1366 outs_addr[i],
1367 );
1368 move |s: usize, e: usize| {
1369 vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
1370 }
1371 });
1372 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1373 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1374 pool.run_many(&parts);
1375 }
1376 return;
1377 }
1378 if uniform_q4 {
1379 let closures: [_; N] = std::array::from_fn(|i| {
1380 let (packed, scales) =
1381 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1382 let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1383 move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
1384 });
1385 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1386 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1387 pool.run_many(&parts);
1388 } else {
1389 let closures: [_; N] = std::array::from_fn(|i| {
1390 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
1391 unreachable!()
1392 };
1393 let (bytes, rows, cols, out) = (
1394 ts[i].quant_bytes(),
1395 ts[i].rows(),
1396 ts[i].cols(),
1397 outs_addr[i],
1398 );
1399 move |s: usize, e: usize| {
1400 vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
1401 }
1402 });
1403 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1404 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1405 pool.run_many(&parts);
1406 }
1407 return;
1408 }
1409
1410 if uniform_f32 {
1411 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1412 let closures: [_; N] = std::array::from_fn(|i| {
1413 let Self::F32 { data, cols, .. } = ts[i] else {
1414 unreachable!()
1415 };
1416 let out = outs_addr[i];
1417 move |start: usize, end: usize| {
1418 for o in start..end {
1419 let row = &data[o * cols..(o + 1) * cols];
1420 let mut sum = 0.0f32;
1421 for j in 0..*cols {
1422 sum += row[j] * x[j];
1423 }
1424 unsafe { *out.at(o) = sum };
1426 }
1427 }
1428 });
1429 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1430 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1431 pool.run_many(&parts);
1432 return;
1433 }
1434
1435 struct Ctx<'a> {
1438 bytes: &'a [u8],
1439 #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
1440 rep: &'a [u8],
1441 row_scale: &'a [f32],
1442 cols: usize,
1443 xs: std::borrow::Cow<'a, [f32]>,
1444 }
1445 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
1446 let Self::Mapped {
1447 dtype,
1448 cols,
1449 row_scale,
1450 col_field,
1451 repack,
1452 ..
1453 } = ts[i]
1454 else {
1455 unreachable!()
1456 };
1457 Ctx {
1458 bytes: ts[i].quant_bytes(),
1459 rep: repack,
1460 row_scale,
1461 cols: *cols,
1462 xs: prescale(x, col_field, *dtype),
1463 }
1464 });
1465 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1466 #[cfg(target_arch = "aarch64")]
1467 if sdot_enabled() {
1468 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
1469 let closures: [_; N] = std::array::from_fn(|i| {
1470 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
1471 move |start: usize, end: usize| {
1472 q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
1473 }
1474 });
1475 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1476 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1477 pool.run_many(&parts);
1478 return;
1479 }
1480 #[cfg(target_arch = "x86_64")]
1481 if avx2_a8w8_enabled() {
1482 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
1483 let closures: [_; N] = std::array::from_fn(|i| {
1484 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
1485 move |start: usize, end: usize| {
1486 q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
1487 }
1488 });
1489 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1490 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1491 pool.run_many(&parts);
1492 return;
1493 }
1494 let closures: [_; N] = std::array::from_fn(|i| {
1495 let (c, out) = (&ctxs[i], outs_addr[i]);
1496 move |start: usize, end: usize| {
1497 q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
1498 }
1499 });
1500 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1501 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1502 pool.run_many(&parts);
1503 }
1504}
1505
1506impl QTensor {
1507 #[allow(clippy::needless_range_loop)]
1512 pub fn matvec2_many<const N: usize>(
1513 ts: [&QTensor; N],
1514 x1: &[f32],
1515 x2: &[f32],
1516 mut o1s: [&mut [f32]; N],
1517 mut o2s: [&mut [f32]; N],
1518 pool: Option<&Pool>,
1519 ) {
1520 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
1521 let uniform_q8 = ts.iter().all(|t| {
1522 matches!(
1523 t,
1524 Self::Mapped {
1525 dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
1526 ..
1527 }
1528 )
1529 });
1530 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
1531 let uniform_q4 = ts.iter().all(|t| {
1532 matches!(
1533 t,
1534 Self::Mapped {
1535 dtype: TensorDtype::Q4Block,
1536 ..
1537 }
1538 )
1539 });
1540 let uniform_vbit = ts.iter().all(|t| {
1541 matches!(
1542 t,
1543 Self::Mapped {
1544 dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
1545 ..
1546 }
1547 )
1548 });
1549 let fusable = pool.is_some()
1550 && total_rows >= 256
1551 && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
1552 if !fusable {
1553 for i in 0..N {
1554 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
1555 }
1556 return;
1557 }
1558 let pool = pool.unwrap();
1559
1560 if uniform_q4 || uniform_vbit {
1561 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1562 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1563 if a8w8_enabled() {
1565 let a1 = split_act(x1);
1566 let a2 = split_act(x2);
1567 let (a1, a2) = (&a1, &a2);
1568 if uniform_q4 {
1569 let closures: [_; N] = std::array::from_fn(|i| {
1570 let (packed, scales) =
1571 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1572 let (gpr, cols, o1, o2) =
1573 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
1574 move |s: usize, e: usize| {
1575 q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
1576 }
1577 });
1578 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1579 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1580 pool.run_many(&parts);
1581 } else {
1582 let closures: [_; N] = std::array::from_fn(|i| {
1583 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
1584 unreachable!()
1585 };
1586 let (bytes, rows, cols, o1, o2) = (
1587 ts[i].quant_bytes(),
1588 ts[i].rows(),
1589 ts[i].cols(),
1590 p1[i],
1591 p2[i],
1592 );
1593 move |s: usize, e: usize| {
1594 vbit_range2_a8w8(
1595 bytes,
1596 vbit_offsets,
1597 x1,
1598 x2,
1599 a1,
1600 a2,
1601 rows,
1602 cols,
1603 o1,
1604 o2,
1605 s,
1606 e,
1607 )
1608 }
1609 });
1610 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1611 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1612 pool.run_many(&parts);
1613 }
1614 return;
1615 }
1616 if uniform_q4 {
1617 let closures: [_; N] = std::array::from_fn(|i| {
1618 let (packed, scales) =
1619 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1620 let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
1621 move |s: usize, e: usize| {
1622 q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
1623 }
1624 });
1625 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1626 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1627 pool.run_many(&parts);
1628 } else {
1629 let closures: [_; N] = std::array::from_fn(|i| {
1630 let Self::Mapped { vbit_offsets, .. } = ts[i] else {
1631 unreachable!()
1632 };
1633 let (bytes, rows, cols, o1, o2) = (
1634 ts[i].quant_bytes(),
1635 ts[i].rows(),
1636 ts[i].cols(),
1637 p1[i],
1638 p2[i],
1639 );
1640 move |s: usize, e: usize| {
1641 vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
1642 }
1643 });
1644 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1645 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1646 pool.run_many(&parts);
1647 }
1648 return;
1649 }
1650
1651 if uniform_f32 {
1652 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1653 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1654 let closures: [_; N] = std::array::from_fn(|i| {
1655 let Self::F32 { data, cols, .. } = ts[i] else {
1656 unreachable!()
1657 };
1658 let (o1, o2) = (p1[i], p2[i]);
1659 move |start: usize, end: usize| {
1660 for o in start..end {
1661 let row = &data[o * cols..(o + 1) * cols];
1662 let (mut s1, mut s2) = (0.0f32, 0.0f32);
1663 for j in 0..*cols {
1664 s1 += row[j] * x1[j];
1665 s2 += row[j] * x2[j];
1666 }
1667 unsafe {
1669 *o1.at(o) = s1;
1670 *o2.at(o) = s2;
1671 }
1672 }
1673 }
1674 });
1675 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1676 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1677 pool.run_many(&parts);
1678 return;
1679 }
1680
1681 struct Ctx<'a> {
1682 bytes: &'a [u8],
1683 row_scale: &'a [f32],
1684 cols: usize,
1685 xs1: std::borrow::Cow<'a, [f32]>,
1686 xs2: std::borrow::Cow<'a, [f32]>,
1687 }
1688 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
1689 let Self::Mapped {
1690 dtype,
1691 cols,
1692 row_scale,
1693 col_field,
1694 ..
1695 } = ts[i]
1696 else {
1697 unreachable!()
1698 };
1699 Ctx {
1700 bytes: ts[i].quant_bytes(),
1701 row_scale,
1702 cols: *cols,
1703 xs1: prescale(x1, col_field, *dtype),
1704 xs2: prescale(x2, col_field, *dtype),
1705 }
1706 });
1707 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1708 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1709 #[cfg(target_arch = "aarch64")]
1710 if sdot_enabled() {
1711 let acts: [(SplitAct, SplitAct); N] =
1712 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
1713 let closures: [_; N] = std::array::from_fn(|i| {
1714 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
1715 move |start: usize, end: usize| {
1716 q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
1717 }
1718 });
1719 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1720 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1721 pool.run_many(&parts);
1722 return;
1723 }
1724 #[cfg(target_arch = "x86_64")]
1725 if avx2_a8w8_enabled() {
1726 let acts: [(SplitAct, SplitAct); N] =
1727 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
1728 let closures: [_; N] = std::array::from_fn(|i| {
1729 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
1730 move |start: usize, end: usize| {
1731 q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
1732 }
1733 });
1734 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1735 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1736 pool.run_many(&parts);
1737 return;
1738 }
1739 let closures: [_; N] = std::array::from_fn(|i| {
1740 let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
1741 move |start: usize, end: usize| {
1742 q8_range2_f32(
1743 c.bytes,
1744 c.row_scale,
1745 &c.xs1,
1746 &c.xs2,
1747 c.cols,
1748 o1,
1749 o2,
1750 start,
1751 end,
1752 )
1753 }
1754 });
1755 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1756 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1757 pool.run_many(&parts);
1758 }
1759
1760 pub fn matvec_silu_mul(
1765 gate: &QTensor,
1766 up: &QTensor,
1767 x: &[f32],
1768 out: &mut [f32],
1769 pool: Option<&Pool>,
1770 ) -> bool {
1771 let inter = gate.rows();
1772 debug_assert_eq!(up.rows(), inter);
1773 debug_assert_eq!(out.len(), inter);
1774 debug_assert_eq!(gate.cols(), up.cols());
1775 if !a8w8_enabled() {
1776 return false;
1777 }
1778 let act = split_act(x);
1779 let act = &act;
1780 let x_ref = x;
1781 let out_addr = SendMut(out.as_mut_ptr());
1782
1783 match (gate, up) {
1784 (
1786 Self::Mapped {
1787 dtype: TensorDtype::Q4Block,
1788 ..
1789 },
1790 Self::Mapped {
1791 dtype: TensorDtype::Q4Block,
1792 ..
1793 },
1794 ) => {
1795 let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
1796 let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
1797 let gpr = gate.cols() / GROUP_SIZE;
1798 let cols = gate.cols();
1799 let run = move |start: usize, end: usize| {
1800 for r in start..end {
1801 let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
1802 let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
1803 for &(j, xv) in &act.outliers {
1804 let flat = r * cols + j;
1805 let gb = gp[flat / 2];
1806 let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
1807 let gsc = f16_to_f32(u16::from_le_bytes([
1808 gs[(flat / GROUP_SIZE) * 2],
1809 gs[(flat / GROUP_SIZE) * 2 + 1],
1810 ]));
1811 gv += ((gn as i32 - 8) as f32) * gsc * xv;
1812 let ub = up_p[flat / 2];
1813 let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
1814 let usc = f16_to_f32(u16::from_le_bytes([
1815 up_s[(flat / GROUP_SIZE) * 2],
1816 up_s[(flat / GROUP_SIZE) * 2 + 1],
1817 ]));
1818 uv += ((un as i32 - 8) as f32) * usc * xv;
1819 }
1820 let silu_g = gv / (1.0 + (-gv).exp());
1821 unsafe { *out_addr.at(r) = silu_g * uv };
1823 }
1824 };
1825 dispatch_rows(pool, inter, &run);
1826 true
1827 }
1828 (
1830 Self::Mapped {
1831 dtype: TensorDtype::Q1T,
1832 ..
1833 },
1834 Self::Mapped {
1835 dtype: TensorDtype::Q1T,
1836 ..
1837 },
1838 ) => {
1839 const TILE: usize = cortiq_core::quant::Q1T_TILE;
1840 let g_bytes = gate.quant_bytes();
1841 let u_bytes = up.quant_bytes();
1842 let gpr = gate.cols() / GROUP_SIZE;
1843 let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
1844 let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
1845 let run = move |start: usize, end: usize| {
1846 for r in start..end {
1847 let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
1848 let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
1849 for &(j, xv) in &act.outliers {
1850 gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
1851 uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
1852 }
1853 gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
1854 uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
1855 let silu_g = gv / (1.0 + (-gv).exp());
1856 unsafe { *out_addr.at(r) = silu_g * uv };
1858 }
1859 };
1860 dispatch_rows(pool, inter, &run);
1861 true
1862 }
1863 _ => false,
1864 }
1865 }
1866}
1867
1868#[cfg(target_os = "macos")]
1873mod accel_blas {
1874 #[link(name = "Accelerate", kind = "framework")]
1875 unsafe extern "C" {
1876 pub fn cblas_sgemm(
1877 order: i32,
1878 trans_a: i32,
1879 trans_b: i32,
1880 m: i32,
1881 n: i32,
1882 k: i32,
1883 alpha: f32,
1884 a: *const f32,
1885 lda: i32,
1886 b: *const f32,
1887 ldb: i32,
1888 beta: f32,
1889 c: *mut f32,
1890 ldc: i32,
1891 );
1892 }
1893}
1894
1895#[cfg(target_os = "macos")]
1896pub(crate) fn accel_gemm_enabled() -> bool {
1897 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1898 *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
1899}
1900
1901#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
1904pub(crate) fn accel_gemm_enabled() -> bool {
1905 true
1906}
1907
1908#[cfg(target_arch = "aarch64")]
1915#[allow(clippy::too_many_arguments)]
1916pub(crate) fn neon_gemm_rm(
1917 m: usize,
1918 n: usize,
1919 k: usize,
1920 alpha: f32,
1921 a: &[f32],
1922 lda: usize,
1923 b_mat: &[f32],
1924 ldb: usize,
1925 b_rows_are_n: bool,
1926 c: &mut [f32],
1927 ldc: usize,
1928) {
1929 debug_assert!(a.len() >= (m - 1) * lda + k);
1930 debug_assert!(c.len() >= (m - 1) * ldc + n);
1931 unsafe {
1933 use core::arch::aarch64::*;
1934 let mut i = 0usize;
1935 while i < m {
1936 let mi = (m - i).min(4);
1937 let mut j = 0usize;
1938 while j < n {
1939 let nj = (n - j).min(8);
1940 if mi == 4 && nj == 8 {
1941 let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
1942 let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
1943 let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
1944 let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
1945 for p in 0..k {
1946 let (b0, b1) = if b_rows_are_n {
1947 let base = b_mat.as_ptr().add(j * ldb + p);
1950 let g = |o: usize| *base.add(o * ldb);
1951 ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
1952 } else {
1953 let base = b_mat.as_ptr().add(p * ldb + j);
1954 (
1955 [*base, *base.add(1), *base.add(2), *base.add(3)],
1956 [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
1957 )
1958 };
1959 let bv0 = vld1q_f32(b0.as_ptr());
1960 let bv1 = vld1q_f32(b1.as_ptr());
1961 let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
1962 let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
1963 let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
1964 let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
1965 c0a = vfmaq_f32(c0a, a0, bv0);
1966 c0b = vfmaq_f32(c0b, a0, bv1);
1967 c1a = vfmaq_f32(c1a, a1, bv0);
1968 c1b = vfmaq_f32(c1b, a1, bv1);
1969 c2a = vfmaq_f32(c2a, a2, bv0);
1970 c2b = vfmaq_f32(c2b, a2, bv1);
1971 c3a = vfmaq_f32(c3a, a3, bv0);
1972 c3b = vfmaq_f32(c3b, a3, bv1);
1973 }
1974 let al = vdupq_n_f32(alpha);
1975 for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
1976 .iter()
1977 .enumerate()
1978 {
1979 let dst = c.as_mut_ptr().add((i + r) * ldc + j);
1980 vst1q_f32(dst, vmulq_f32(*ca, al));
1981 vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
1982 }
1983 } else {
1984 for r in 0..mi {
1985 for q in 0..nj {
1986 let mut acc = 0f32;
1987 for p in 0..k {
1988 let bv = if b_rows_are_n {
1989 b_mat[(j + q) * ldb + p]
1990 } else {
1991 b_mat[p * ldb + j + q]
1992 };
1993 acc += a[(i + r) * lda + p] * bv;
1994 }
1995 c[(i + r) * ldc + j + q] = acc * alpha;
1996 }
1997 }
1998 }
1999 j += nj;
2000 }
2001 i += mi;
2002 }
2003 }
2004}
2005
2006#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
2008#[allow(clippy::too_many_arguments)]
2009pub(crate) fn sgemm_rm(
2010 m: usize,
2011 n: usize,
2012 k: usize,
2013 alpha: f32,
2014 a: &[f32],
2015 lda: usize,
2016 b_mat: &[f32],
2017 ldb: usize,
2018 b_rows_are_n: bool,
2019 c: &mut [f32],
2020 ldc: usize,
2021) {
2022 neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
2023}
2024
2025#[cfg(target_os = "macos")]
2028#[allow(clippy::too_many_arguments)]
2029pub(crate) fn sgemm_rm(
2030 m: usize,
2031 n: usize,
2032 k: usize,
2033 alpha: f32,
2034 a: &[f32],
2035 lda: usize,
2036 b_mat: &[f32],
2037 ldb: usize,
2038 b_rows_are_n: bool,
2039 c: &mut [f32],
2040 ldc: usize,
2041) {
2042 debug_assert!(a.len() >= (m - 1) * lda + k);
2043 debug_assert!(c.len() >= (m - 1) * ldc + n);
2044 #[cfg(target_arch = "aarch64")]
2049 if std::env::var("CMF_FORCE_NEON_GEMM")
2050 .map(|v| v == "1")
2051 .unwrap_or(false)
2052 {
2053 return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
2054 }
2055 unsafe {
2056 accel_blas::cblas_sgemm(
2057 101, 111, if b_rows_are_n { 112 } else { 111 },
2060 m as i32,
2061 n as i32,
2062 k as i32,
2063 alpha,
2064 a.as_ptr(),
2065 lda as i32,
2066 b_mat.as_ptr(),
2067 ldb as i32,
2068 0.0,
2069 c.as_mut_ptr(),
2070 ldc as i32,
2071 );
2072 }
2073}
2074
2075#[cfg(target_os = "macos")]
2082fn qmatmat_accel(
2083 q: &[u8],
2084 row_scale: &[f32],
2085 pre: &[std::borrow::Cow<'_, [f32]>],
2086 rows: usize,
2087 cols: usize,
2088 out: &mut [f32],
2089 pool: Option<&Pool>,
2090) {
2091 const TR: usize = 2048;
2096 let b = pre.len();
2097 thread_local! {
2098 static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
2099 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
2100 }
2101 XPANEL.with(|xp| {
2102 WTILE.with(|wt| {
2103 let mut xpanel = xp.borrow_mut();
2104 xpanel.clear();
2105 for x in pre {
2106 xpanel.extend_from_slice(x);
2107 }
2108 let mut wtile = wt.borrow_mut();
2109 wtile.resize(TR * cols, 0.0);
2110 let mut r0 = 0usize;
2111 while r0 < rows {
2112 let tr = TR.min(rows - r0);
2113 let wt_addr = SendMut(wtile.as_mut_ptr());
2115 let run = |start: usize, end: usize| {
2116 for r in start..end {
2117 let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
2118 let s = row_scale[r0 + r];
2119 let dst =
2121 unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
2122 for (d, &v) in dst.iter_mut().zip(row) {
2123 *d = (v as i8) as f32 * s;
2124 }
2125 }
2126 };
2127 dispatch_rows(pool, tr, &run);
2128 unsafe {
2130 accel_blas::cblas_sgemm(
2131 101, 111, 112, b as i32,
2135 tr as i32,
2136 cols as i32,
2137 1.0,
2138 xpanel.as_ptr(),
2139 cols as i32,
2140 wtile.as_ptr(),
2141 cols as i32,
2142 0.0,
2143 out.as_mut_ptr().add(r0),
2144 rows as i32,
2145 );
2146 }
2147 r0 += tr;
2148 }
2149 })
2150 });
2151}
2152
2153fn qmatmat(
2154 q: &[u8],
2155 row_scale: &[f32],
2156 pre: &[std::borrow::Cow<'_, [f32]>],
2157 rows: usize,
2158 cols: usize,
2159 out: &mut [f32],
2160 pool: Option<&Pool>,
2161) {
2162 let b = pre.len();
2163 debug_assert_eq!(out.len(), b * rows);
2164 #[cfg(target_os = "macos")]
2169 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
2170 qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
2171 return;
2172 }
2173 #[cfg(target_arch = "aarch64")]
2174 if sdot_enabled() {
2175 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
2176 let out_addr = SendMut(out.as_mut_ptr());
2177 let blocked_ok = std::env::var("CMF_X86_BLOCKED")
2180 .map(|v| v != "0")
2181 .unwrap_or(true);
2182 let use_i8mm = i8mm_enabled();
2183 if blocked_ok {
2184 let run = |start: usize, end: usize| {
2185 let mut o = start;
2186 while o < end {
2187 if o + 2 <= end {
2188 let r0 = &q[o * cols..(o + 1) * cols];
2189 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
2190 let mut bi = 0usize;
2191 while bi + 4 <= acts.len() {
2192 let xs = [
2193 acts[bi].xq.as_slice(),
2194 acts[bi + 1].xq.as_slice(),
2195 acts[bi + 2].xq.as_slice(),
2196 acts[bi + 3].xq.as_slice(),
2197 ];
2198 let d = if use_i8mm {
2199 unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
2200 } else {
2201 unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
2202 };
2203 for (r, row) in [r0, r1].into_iter().enumerate() {
2204 for k in 0..4 {
2205 let act = &acts[bi + k];
2206 let mut v = d[r][k] as f32 * act.sx;
2207 for &(j, xv) in &act.outliers {
2208 v += (row[j] as i8) as f32 * xv;
2209 }
2210 unsafe {
2211 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
2212 };
2213 }
2214 }
2215 bi += 4;
2216 }
2217 while bi < acts.len() {
2218 for (r, row) in [r0, r1].into_iter().enumerate() {
2219 let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
2220 unsafe { *out_addr.at(bi * rows + o + r) = v };
2221 }
2222 bi += 1;
2223 }
2224 o += 2;
2225 } else {
2226 let row = &q[o * cols..(o + 1) * cols];
2227 for (bi, act) in acts.iter().enumerate() {
2228 let v = row_dot_sdot(row, act) * row_scale[o];
2229 unsafe { *out_addr.at(bi * rows + o) = v };
2230 }
2231 o += 1;
2232 }
2233 }
2234 };
2235 dispatch_rows(pool, rows, &run);
2236 return;
2237 }
2238 let run = |start: usize, end: usize| {
2239 for o in start..end {
2240 let row = &q[o * cols..(o + 1) * cols];
2241 for (bi, act) in acts.iter().enumerate() {
2242 let v = row_dot_sdot(row, act) * row_scale[o];
2243 unsafe { *out_addr.at(bi * rows + o) = v };
2244 }
2245 }
2246 };
2247 dispatch_rows(pool, rows, &run);
2248 return;
2249 }
2250 #[cfg(target_arch = "x86_64")]
2255 if avx2_a8w8_enabled() {
2256 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
2257 let out_addr = SendMut(out.as_mut_ptr());
2258 let blocked_ok = std::env::var("CMF_X86_BLOCKED")
2261 .map(|v| v != "0")
2262 .unwrap_or(true);
2263 if !avx512vnni_enabled() && blocked_ok {
2264 let run = |start: usize, end: usize| {
2265 let mut o = start;
2266 while o < end {
2267 if o + 2 <= end {
2268 let r0 = &q[o * cols..(o + 1) * cols];
2269 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
2270 let mut bi = 0usize;
2271 while bi + 4 <= acts.len() {
2272 let xs = [
2273 acts[bi].xq.as_slice(),
2274 acts[bi + 1].xq.as_slice(),
2275 acts[bi + 2].xq.as_slice(),
2276 acts[bi + 3].xq.as_slice(),
2277 ];
2278 let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
2279 for (r, row) in [r0, r1].into_iter().enumerate() {
2280 for k in 0..4 {
2281 let act = &acts[bi + k];
2282 let mut v = d[r][k] as f32 * act.sx;
2283 for &(j, xv) in &act.outliers {
2284 v += (row[j] as i8) as f32 * xv;
2285 }
2286 unsafe {
2287 *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
2288 };
2289 }
2290 }
2291 bi += 4;
2292 }
2293 while bi < acts.len() {
2294 for (r, row) in [r0, r1].into_iter().enumerate() {
2295 let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
2296 unsafe { *out_addr.at(bi * rows + o + r) = v };
2297 }
2298 bi += 1;
2299 }
2300 o += 2;
2301 } else {
2302 let row = &q[o * cols..(o + 1) * cols];
2303 for (bi, act) in acts.iter().enumerate() {
2304 let v = row_dot_avx2(row, act) * row_scale[o];
2305 unsafe { *out_addr.at(bi * rows + o) = v };
2306 }
2307 o += 1;
2308 }
2309 }
2310 };
2311 dispatch_rows(pool, rows, &run);
2312 return;
2313 }
2314 let run = |start: usize, end: usize| {
2315 for o in start..end {
2316 let row = &q[o * cols..(o + 1) * cols];
2317 for (bi, act) in acts.iter().enumerate() {
2318 let v = row_dot_avx2(row, act) * row_scale[o];
2319 unsafe { *out_addr.at(bi * rows + o) = v };
2320 }
2321 }
2322 };
2323 dispatch_rows(pool, rows, &run);
2324 return;
2325 }
2326 let out_addr = SendMut(out.as_mut_ptr());
2327 let run = |start: usize, end: usize| {
2328 for o in start..end {
2329 let row = &q[o * cols..(o + 1) * cols];
2330 for (bi, x) in pre.iter().enumerate() {
2331 let mut acc = 0f32;
2332 for j in 0..cols {
2333 acc += (row[j] as i8) as f32 * x[j];
2334 }
2335 unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
2336 }
2337 }
2338 };
2339 dispatch_rows(pool, rows, &run);
2340}
2341
2342fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
2345 match pool {
2346 Some(pool) if rows >= 256 => pool.run_rows(rows, run),
2347 _ => run(0, rows),
2348 }
2349}
2350
2351fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
2353 let groups = rows * cols / GROUP_SIZE;
2354 bytes.split_at(groups * 16)
2355}
2356
2357#[inline]
2362fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
2363 #[cfg(target_arch = "aarch64")]
2364 unsafe {
2365 return vbit_fill4_neon(data, buf);
2366 }
2367 #[cfg(target_arch = "x86_64")]
2368 if avx2_enabled() {
2369 return unsafe { vbit_fill4_avx2(data, buf) };
2370 }
2371 #[allow(unreachable_code)]
2372 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
2373 let u = unpack8::<4>(&data[blk * 4..]);
2374 for k in 0..8 {
2375 chunk[k] = (u[k] - 7) as i8 as u8;
2376 }
2377 }
2378}
2379
2380#[cfg(target_arch = "aarch64")]
2381#[target_feature(enable = "neon")]
2382unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
2383 unsafe {
2386 use core::arch::aarch64::*;
2387 let n = buf.len();
2388 let mask = vdupq_n_u8(0x0F);
2389 let seven = vdupq_n_s8(7);
2390 let mut g = 0usize;
2391 while g * 32 + 32 <= n {
2392 let b = vld1q_u8(data.as_ptr().add(g * 16));
2393 let hi = vshrq_n_u8::<4>(b);
2394 let lo = vandq_u8(b, mask);
2395 let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
2396 let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
2397 vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
2398 vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
2399 g += 1;
2400 }
2401 }
2402}
2403
2404#[cfg(target_arch = "x86_64")]
2405#[target_feature(enable = "avx2")]
2406unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
2407 unsafe {
2409 use core::arch::x86_64::*;
2410 let n = buf.len();
2411 let mask = _mm_set1_epi8(0x0F);
2412 let seven = _mm256_set1_epi8(7);
2413 let mut g = 0usize;
2414 while g * 32 + 32 <= n {
2415 let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
2416 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
2417 let lo = _mm_and_si128(b, mask);
2418 let z = _mm256_sub_epi8(
2419 _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
2420 seven,
2421 );
2422 _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
2423 g += 1;
2424 }
2425 }
2426}
2427
2428#[inline(always)]
2433fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
2434 let mut acc = 0u64;
2435 for i in 0..B {
2436 acc = (acc << 8) | data[i] as u64;
2437 }
2438 let mask = (1u64 << B) - 1;
2439 let mut out = [0i32; 8];
2440 for (k, o) in out.iter_mut().enumerate() {
2441 *o = ((acc >> ((7 - k) * B)) & mask) as i32;
2442 }
2443 out
2444}
2445
2446#[allow(clippy::too_many_arguments)]
2452fn vbitmatvec(
2453 bytes: &[u8],
2454 offsets: &[usize],
2455 x: &[f32],
2456 rows: usize,
2457 cols: usize,
2458 out: &mut [f32],
2459 pool: Option<&Pool>,
2460) {
2461 debug_assert_eq!(out.len(), rows);
2462 debug_assert_eq!(offsets.len(), rows + 1);
2463
2464 if a8w8_enabled() {
2468 let act = split_act(x);
2469 let out_addr = SendMut(out.as_mut_ptr());
2470 let run = move |start: usize, end: usize| {
2471 vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
2472 };
2473 dispatch_rows(pool, rows, &run);
2474 return;
2475 }
2476
2477 let out_addr = SendMut(out.as_mut_ptr());
2478 let run = move |start: usize, end: usize| {
2479 vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
2480 };
2481 dispatch_rows(pool, rows, &run);
2482}
2483
2484#[allow(clippy::too_many_arguments)]
2488fn vbit_range_a8w8(
2489 bytes: &[u8],
2490 offsets: &[usize],
2491 x: &[f32],
2492 act: &SplitAct,
2493 rows: usize,
2494 cols: usize,
2495 out: SendMut,
2496 start: usize,
2497 end: usize,
2498) {
2499 let ng = cols / GROUP_SIZE;
2500 let bits = &bytes[..rows];
2501 let sc_off = rows;
2502 let row_dot = |r: usize| -> f32 {
2503 let b = bits[r] as usize;
2504 let l = (1i32 << (b - 1)) - 1;
2505 let mask = (1u64 << b) - 1;
2506 let data = &bytes[offsets[r]..offsets[r + 1]];
2507 if b == 8 {
2508 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
2510 let mut dot = 0f32;
2511 for g in 0..ng {
2512 let so = (r * ng + g) * 2;
2513 let sgf = f16_to_f32(u16::from_le_bytes([
2514 bytes[sc_off + so],
2515 bytes[sc_off + so + 1],
2516 ]));
2517 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2518 let mut gd = 0f32;
2519 for &xv in xg.iter() {
2520 if nbits < 8 {
2521 acc = (acc << 8) | data[idx] as u64;
2522 idx += 1;
2523 nbits += 8;
2524 }
2525 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
2526 nbits -= 8;
2527 gd += (u - l) as f32 * xv;
2528 }
2529 dot += gd * sgf;
2530 }
2531 return dot;
2532 }
2533 thread_local! {
2537 static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
2538 const { std::cell::RefCell::new(Vec::new()) };
2539 }
2540 #[inline(always)]
2541 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
2542 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
2543 let u = unpack8::<B>(&data[blk * B..]);
2544 for k in 0..8 {
2545 chunk[k] = (u[k] - l) as i8 as u8;
2546 }
2547 }
2548 }
2549 let _ = mask;
2550 VBIT_SCRATCH.with(|scratch| {
2551 let mut buf = scratch.borrow_mut();
2552 buf.resize(cols, 0);
2553 match b {
2554 3 => fill::<3>(data, l, &mut buf),
2555 4 => vbit_fill4(data, &mut buf),
2556 5 => fill::<5>(data, l, &mut buf),
2557 6 => fill::<6>(data, l, &mut buf),
2558 _ => unreachable!(),
2559 }
2560 let mut dot = 0f32;
2561 for g in 0..ng {
2562 let so = (r * ng + g) * 2;
2563 let s = f16_to_f32(u16::from_le_bytes([
2564 bytes[sc_off + so],
2565 bytes[sc_off + so + 1],
2566 ]));
2567 let d = dot_i8_i8(
2568 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
2569 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
2570 ) as f32
2571 * act.sx;
2572 dot += d * s;
2573 }
2574 for &(j, xv) in &act.outliers {
2575 let so = (r * ng + j / GROUP_SIZE) * 2;
2576 let s = f16_to_f32(u16::from_le_bytes([
2577 bytes[sc_off + so],
2578 bytes[sc_off + so + 1],
2579 ]));
2580 dot += (buf[j] as i8) as f32 * s * xv;
2582 }
2583 dot
2584 })
2585 };
2586 for r in start..end {
2587 unsafe { *out.at(r) = row_dot(r) };
2589 }
2590}
2591
2592#[allow(clippy::too_many_arguments)]
2594fn vbit_range_f32(
2595 bytes: &[u8],
2596 offsets: &[usize],
2597 x: &[f32],
2598 rows: usize,
2599 cols: usize,
2600 out: SendMut,
2601 start: usize,
2602 end: usize,
2603) {
2604 let ng = cols / GROUP_SIZE;
2605 let bits = &bytes[..rows];
2606 let sc_off = rows;
2607 #[inline(always)]
2611 fn dot_row<const B: usize>(
2612 data: &[u8],
2613 bytes: &[u8],
2614 sc_off: usize,
2615 r: usize,
2616 ng: usize,
2617 x: &[f32],
2618 ) -> f32 {
2619 let l = ((1i32 << (B - 1)) - 1) as f32;
2620 let gbytes = GROUP_SIZE * B / 8;
2621 let mut dot = 0f32;
2622 for g in 0..ng {
2623 let so = (r * ng + g) * 2;
2624 let s = f16_to_f32(u16::from_le_bytes([
2625 bytes[sc_off + so],
2626 bytes[sc_off + so + 1],
2627 ]));
2628 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2629 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
2630 let mut gd = 0f32;
2631 for blk in 0..GROUP_SIZE / 8 {
2632 let u = unpack8::<B>(&gd0[blk * B..]);
2633 let xb = &xg[blk * 8..blk * 8 + 8];
2634 for k in 0..8 {
2635 gd += (u[k] as f32 - l) * xb[k];
2636 }
2637 }
2638 dot += gd * s;
2639 }
2640 dot
2641 }
2642 for r in start..end {
2643 let data = &bytes[offsets[r]..offsets[r + 1]];
2644 let v = match bits[r] {
2645 3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
2646 4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
2647 5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
2648 6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
2649 8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
2650 b => unreachable!("vbit bit-width {b} (validated at load)"),
2651 };
2652 unsafe { *out.at(r) = v };
2654 }
2655}
2656
2657#[allow(clippy::too_many_arguments)]
2662fn vbitmatvec2(
2663 bytes: &[u8],
2664 offsets: &[usize],
2665 x1: &[f32],
2666 x2: &[f32],
2667 rows: usize,
2668 cols: usize,
2669 o1: &mut [f32],
2670 o2: &mut [f32],
2671 pool: Option<&Pool>,
2672) {
2673 debug_assert_eq!(o1.len(), rows);
2674 debug_assert_eq!(o2.len(), rows);
2675
2676 if a8w8_enabled() {
2677 let a1 = split_act(x1);
2678 let a2 = split_act(x2);
2679 let p1 = SendMut(o1.as_mut_ptr());
2680 let p2 = SendMut(o2.as_mut_ptr());
2681 let run = move |start: usize, end: usize| {
2682 vbit_range2_a8w8(
2683 bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
2684 )
2685 };
2686 dispatch_rows(pool, rows, &run);
2687 return;
2688 }
2689
2690 let p1 = SendMut(o1.as_mut_ptr());
2691 let p2 = SendMut(o2.as_mut_ptr());
2692 let run = move |start: usize, end: usize| {
2693 vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
2694 };
2695 dispatch_rows(pool, rows, &run);
2696}
2697
2698#[allow(clippy::too_many_arguments)]
2702fn vbit_range2_a8w8(
2703 bytes: &[u8],
2704 offsets: &[usize],
2705 x1: &[f32],
2706 x2: &[f32],
2707 a1: &SplitAct,
2708 a2: &SplitAct,
2709 rows: usize,
2710 cols: usize,
2711 p1: SendMut,
2712 p2: SendMut,
2713 start: usize,
2714 end: usize,
2715) {
2716 let ng = cols / GROUP_SIZE;
2717 let bits = &bytes[..rows];
2718 let sc_off = rows;
2719 let row_dots = |r: usize| -> (f32, f32) {
2720 let b = bits[r] as usize;
2721 let l = (1i32 << (b - 1)) - 1;
2722 let data = &bytes[offsets[r]..offsets[r + 1]];
2723 if b == 8 {
2724 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
2727 let (mut d1, mut d2) = (0f32, 0f32);
2728 for g in 0..ng {
2729 let so = (r * ng + g) * 2;
2730 let sgf = f16_to_f32(u16::from_le_bytes([
2731 bytes[sc_off + so],
2732 bytes[sc_off + so + 1],
2733 ]));
2734 let (mut g1, mut g2) = (0f32, 0f32);
2735 for k in 0..GROUP_SIZE {
2736 if nbits < 8 {
2737 acc = (acc << 8) | data[idx] as u64;
2738 idx += 1;
2739 nbits += 8;
2740 }
2741 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
2742 nbits -= 8;
2743 let w = (u - l) as f32;
2744 g1 += w * x1[g * GROUP_SIZE + k];
2745 g2 += w * x2[g * GROUP_SIZE + k];
2746 }
2747 d1 += g1 * sgf;
2748 d2 += g2 * sgf;
2749 }
2750 return (d1, d2);
2751 }
2752 thread_local! {
2753 static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
2754 const { std::cell::RefCell::new(Vec::new()) };
2755 }
2756 #[inline(always)]
2757 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
2758 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
2759 let u = unpack8::<B>(&data[blk * B..]);
2760 for k in 0..8 {
2761 chunk[k] = (u[k] - l) as i8 as u8;
2762 }
2763 }
2764 }
2765 VBIT_SCRATCH2.with(|scratch| {
2766 let mut buf = scratch.borrow_mut();
2767 buf.resize(cols, 0);
2768 match b {
2769 3 => fill::<3>(data, l, &mut buf),
2770 4 => vbit_fill4(data, &mut buf),
2771 5 => fill::<5>(data, l, &mut buf),
2772 6 => fill::<6>(data, l, &mut buf),
2773 _ => unreachable!(),
2774 }
2775 let (mut d1, mut d2) = (0f32, 0f32);
2776 for g in 0..ng {
2777 let so = (r * ng + g) * 2;
2778 let s = f16_to_f32(u16::from_le_bytes([
2779 bytes[sc_off + so],
2780 bytes[sc_off + so + 1],
2781 ]));
2782 let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2783 let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
2784 let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
2785 d1 += v1 * s;
2786 d2 += v2 * s;
2787 }
2788 for &(j, xv) in &a1.outliers {
2789 let so = (r * ng + j / GROUP_SIZE) * 2;
2790 let s = f16_to_f32(u16::from_le_bytes([
2791 bytes[sc_off + so],
2792 bytes[sc_off + so + 1],
2793 ]));
2794 d1 += (buf[j] as i8) as f32 * s * xv;
2795 }
2796 for &(j, xv) in &a2.outliers {
2797 let so = (r * ng + j / GROUP_SIZE) * 2;
2798 let s = f16_to_f32(u16::from_le_bytes([
2799 bytes[sc_off + so],
2800 bytes[sc_off + so + 1],
2801 ]));
2802 d2 += (buf[j] as i8) as f32 * s * xv;
2803 }
2804 (d1, d2)
2805 })
2806 };
2807 for r in start..end {
2808 let (v1, v2) = row_dots(r);
2809 unsafe {
2811 *p1.at(r) = v1;
2812 *p2.at(r) = v2;
2813 }
2814 }
2815}
2816
2817#[allow(clippy::too_many_arguments)]
2821fn vbit_range2_f32(
2822 bytes: &[u8],
2823 offsets: &[usize],
2824 x1: &[f32],
2825 x2: &[f32],
2826 rows: usize,
2827 cols: usize,
2828 p1: SendMut,
2829 p2: SendMut,
2830 start: usize,
2831 end: usize,
2832) {
2833 let ng = cols / GROUP_SIZE;
2834 let bits = &bytes[..rows];
2835 let sc_off = rows;
2836 #[inline(always)]
2837 #[allow(clippy::too_many_arguments)]
2838 fn dot_row2<const B: usize>(
2839 data: &[u8],
2840 bytes: &[u8],
2841 sc_off: usize,
2842 r: usize,
2843 ng: usize,
2844 x1: &[f32],
2845 x2: &[f32],
2846 ) -> (f32, f32) {
2847 let l = ((1i32 << (B - 1)) - 1) as f32;
2848 let gbytes = GROUP_SIZE * B / 8;
2849 let (mut d1, mut d2) = (0f32, 0f32);
2850 for g in 0..ng {
2851 let so = (r * ng + g) * 2;
2852 let s = f16_to_f32(u16::from_le_bytes([
2853 bytes[sc_off + so],
2854 bytes[sc_off + so + 1],
2855 ]));
2856 let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2857 let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2858 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
2859 let (mut g1, mut g2) = (0f32, 0f32);
2860 for blk in 0..GROUP_SIZE / 8 {
2861 let u = unpack8::<B>(&gd0[blk * B..]);
2862 for k in 0..8 {
2863 let w = u[k] as f32 - l;
2864 g1 += w * x1g[blk * 8 + k];
2865 g2 += w * x2g[blk * 8 + k];
2866 }
2867 }
2868 d1 += g1 * s;
2869 d2 += g2 * s;
2870 }
2871 (d1, d2)
2872 }
2873 for r in start..end {
2874 let data = &bytes[offsets[r]..offsets[r + 1]];
2875 let (v1, v2) = match bits[r] {
2876 3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
2877 4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
2878 5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
2879 6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
2880 8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
2881 b => unreachable!("vbit bit-width {b} (validated at load)"),
2882 };
2883 unsafe {
2885 *p1.at(r) = v1;
2886 *p2.at(r) = v2;
2887 }
2888 }
2889}
2890
2891#[inline]
2898#[allow(unreachable_code)]
2899fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
2900 #[cfg(target_arch = "aarch64")]
2901 unsafe {
2902 return dot_q4t_row_sdot(bytes, r, gpr, xq);
2903 }
2904 #[cfg(target_arch = "x86_64")]
2905 unsafe {
2906 return dot_q4t_row_avx2(bytes, r, gpr, xq);
2907 }
2908 let mut acc = 0f32;
2909 for gi in 0..gpr {
2910 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
2911 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2912 let mut d = 0i32;
2913 for (k, &b) in tile[2..].iter().enumerate() {
2914 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
2915 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
2916 }
2917 acc += d as f32 * s;
2918 }
2919 acc
2920}
2921
2922#[cfg(target_arch = "aarch64")]
2923#[target_feature(enable = "neon,dotprod")]
2924unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
2925 unsafe {
2928 use core::arch::aarch64::*;
2929 use core::arch::asm;
2930 let lomask = vdupq_n_u8(0x0F);
2931 let eight = vdupq_n_s8(8);
2932 let mut acc = 0f32;
2933 for gi in 0..gpr {
2934 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
2935 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2936 let b = vld1q_u8(t.add(2));
2937 let lo = vandq_u8(b, lomask);
2938 let hi = vshrq_n_u8::<4>(b);
2939 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
2940 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
2941 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
2942 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
2943 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
2944 asm!(
2945 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
2946 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
2947 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
2948 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
2949 options(pure, nomem, nostack),
2950 );
2951 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
2952 }
2953 acc
2954 }
2955}
2956
2957#[cfg(target_arch = "x86_64")]
2958#[target_feature(enable = "avx2")]
2959unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
2960 unsafe {
2962 use core::arch::x86_64::*;
2963 let lomask = _mm_set1_epi8(0x0F);
2964 let eight = _mm256_set1_epi8(8);
2965 let ones = _mm256_set1_epi16(1);
2966 let mut acc = 0f32;
2967 for gi in 0..gpr {
2968 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
2969 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2970 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
2971 let lo = _mm_and_si128(b, lomask);
2972 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
2973 let w = _mm256_sub_epi8(
2974 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
2975 eight,
2976 );
2977 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2978 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
2979 let d = _mm256_madd_epi16(p16, ones);
2980 let hi128 = _mm256_extracti128_si256::<1>(d);
2981 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
2982 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2983 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2984 acc += _mm_cvtsi128_si32(s32) as f32 * s;
2985 }
2986 acc
2987 }
2988}
2989
2990#[cfg(target_arch = "x86_64")]
2995#[target_feature(enable = "avx2")]
2996unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
2997 unsafe {
2999 use core::arch::x86_64::*;
3000 let lomask = _mm_set1_epi8(0x0F);
3001 let eight = _mm256_set1_epi8(8);
3002 let ones = _mm256_set1_epi16(1);
3003 let mut acc = [0f32; 4];
3004 for gi in 0..gpr {
3005 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3006 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3007 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
3008 let lo = _mm_and_si128(bb, lomask);
3009 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
3010 let w = _mm256_sub_epi8(
3011 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3012 eight,
3013 );
3014 let aw = _mm256_abs_epi8(w);
3015 for (k, xq) in xs.iter().enumerate() {
3016 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3017 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
3018 let d = _mm256_madd_epi16(p16, ones);
3019 let hi128 = _mm256_extracti128_si256::<1>(d);
3020 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3021 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3022 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3023 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
3024 }
3025 }
3026 acc
3027 }
3028}
3029
3030#[inline]
3032fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
3033 let gi = j / GROUP_SIZE;
3034 let k = j % GROUP_SIZE;
3035 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
3036 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3037 let byte = tile[2 + k / 2];
3038 let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
3039 ((nib as i32 - 8) as f32, s)
3040}
3041
3042#[inline]
3045fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
3046 let mut acc = 0f32;
3047 for gi in 0..gpr {
3048 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
3049 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3050 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
3051 let mut ga = 0f32;
3052 for (k, &b) in tile[2..].iter().enumerate() {
3053 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
3054 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
3055 }
3056 acc += ga * s;
3057 }
3058 acc
3059}
3060
3061fn q4t_matvec(
3063 bytes: &[u8],
3064 x: &[f32],
3065 rows: usize,
3066 cols: usize,
3067 out: &mut [f32],
3068 pool: Option<&Pool>,
3069) {
3070 debug_assert_eq!(out.len(), rows);
3071 let gpr = cols / GROUP_SIZE;
3072 let out_addr = SendMut(out.as_mut_ptr());
3073 if a8w8_enabled() {
3074 let act = split_act(x);
3075 let run = move |start: usize, end: usize| {
3076 for r in start..end {
3077 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
3078 for &(j, xv) in &act.outliers {
3079 let (w, s) = q4t_outlier(bytes, r, gpr, j);
3080 acc += w * s * xv;
3081 }
3082 unsafe { *out_addr.at(r) = acc };
3084 }
3085 };
3086 dispatch_rows(pool, rows, &run);
3087 return;
3088 }
3089 let run = move |start: usize, end: usize| {
3090 for r in start..end {
3091 unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
3093 }
3094 };
3095 dispatch_rows(pool, rows, &run);
3096}
3097
3098#[allow(clippy::too_many_arguments)]
3100fn q4t_matvec2(
3101 bytes: &[u8],
3102 x1: &[f32],
3103 x2: &[f32],
3104 rows: usize,
3105 cols: usize,
3106 o1: &mut [f32],
3107 o2: &mut [f32],
3108 pool: Option<&Pool>,
3109) {
3110 let gpr = cols / GROUP_SIZE;
3111 let p1 = SendMut(o1.as_mut_ptr());
3112 let p2 = SendMut(o2.as_mut_ptr());
3113 if a8w8_enabled() {
3114 let a1 = split_act(x1);
3115 let a2 = split_act(x2);
3116 let run = move |start: usize, end: usize| {
3117 for r in start..end {
3118 let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
3119 let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
3120 for &(j, xv) in &a1.outliers {
3121 let (w, s) = q4t_outlier(bytes, r, gpr, j);
3122 v1 += w * s * xv;
3123 }
3124 for &(j, xv) in &a2.outliers {
3125 let (w, s) = q4t_outlier(bytes, r, gpr, j);
3126 v2 += w * s * xv;
3127 }
3128 unsafe {
3130 *p1.at(r) = v1;
3131 *p2.at(r) = v2;
3132 }
3133 }
3134 };
3135 dispatch_rows(pool, rows, &run);
3136 return;
3137 }
3138 let run = move |start: usize, end: usize| {
3139 for r in start..end {
3140 unsafe {
3142 *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
3143 *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
3144 }
3145 }
3146 };
3147 dispatch_rows(pool, rows, &run);
3148}
3149
3150#[allow(clippy::too_many_arguments)]
3152fn q4t_matmat(
3153 bytes: &[u8],
3154 xs_all: &[f32],
3155 b: usize,
3156 rows: usize,
3157 cols: usize,
3158 out: &mut [f32],
3159 pool: Option<&Pool>,
3160) {
3161 debug_assert_eq!(out.len(), b * rows);
3162 let gpr = cols / GROUP_SIZE;
3163 let out_addr = SendMut(out.as_mut_ptr());
3164 if a8w8_enabled() {
3165 let acts: Vec<SplitAct> = (0..b)
3166 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
3167 .collect();
3168 let acts = &acts;
3169 #[cfg(target_arch = "x86_64")]
3170 let blocked_ok = avx2_enabled()
3171 && std::env::var("CMF_X86_BLOCKED")
3172 .map(|v| v != "0")
3173 .unwrap_or(true);
3174 #[cfg(not(target_arch = "x86_64"))]
3175 let blocked_ok = false;
3176 let run = move |start: usize, end: usize| {
3177 for r in start..end {
3178 let mut bi = 0usize;
3179 #[cfg(target_arch = "x86_64")]
3180 if blocked_ok {
3181 while bi + 4 <= acts.len() {
3182 let xs = [
3183 acts[bi].xq.as_slice(),
3184 acts[bi + 1].xq.as_slice(),
3185 acts[bi + 2].xq.as_slice(),
3186 acts[bi + 3].xq.as_slice(),
3187 ];
3188 let d = unsafe { dot_q4t_row_1x4_avx2(bytes, r, gpr, xs) };
3189 for k in 0..4 {
3190 let act = &acts[bi + k];
3191 let mut acc = d[k] * act.sx;
3192 for &(j, xv) in &act.outliers {
3193 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
3194 acc += w * sc * xv;
3195 }
3196 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
3198 }
3199 bi += 4;
3200 }
3201 }
3202 let _ = blocked_ok;
3203 while bi < acts.len() {
3204 let act = &acts[bi];
3205 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
3206 for &(j, xv) in &act.outliers {
3207 let (w, s) = q4t_outlier(bytes, r, gpr, j);
3208 acc += w * s * xv;
3209 }
3210 unsafe { *out_addr.at(bi * rows + r) = acc };
3212 bi += 1;
3213 }
3214 }
3215 };
3216 dispatch_rows(pool, rows, &run);
3217 return;
3218 }
3219 let run = move |start: usize, end: usize| {
3220 for r in start..end {
3221 for bi in 0..b {
3222 let x = &xs_all[bi * cols..(bi + 1) * cols];
3223 unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
3225 }
3226 }
3227 };
3228 dispatch_rows(pool, rows, &run);
3229}
3230
3231fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
3240 (0..gpr)
3241 .map(|gi| {
3242 xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
3243 .iter()
3244 .map(|&v| v as i32)
3245 .sum()
3246 })
3247 .collect()
3248}
3249
3250#[inline]
3254#[allow(unreachable_code)]
3255#[cfg(target_arch = "x86_64")]
3260#[target_feature(enable = "avx2")]
3261unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
3262 unsafe {
3264 use core::arch::x86_64::*;
3265 let expand = _mm256_setr_epi8(
3267 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,
3268 3, 3, 3,
3269 );
3270 let bitsel = _mm256_setr_epi8(
3271 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
3272 -128, 1, 2, 4, 8, 16, 32, 64, -128,
3273 );
3274 let ones8 = _mm256_set1_epi8(1);
3275 let ones16 = _mm256_set1_epi16(1);
3276 let mut acc = 0f32;
3277 for gi in 0..gpr {
3278 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
3279 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3280 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
3281 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
3282 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
3283 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3284 let sel = _mm256_and_si256(x, mask);
3285 let p16 = _mm256_maddubs_epi16(ones8, sel);
3287 let d32 = _mm256_madd_epi16(p16, ones16);
3288 let hi128 = _mm256_extracti128_si256::<1>(d32);
3289 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
3290 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3291 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3292 let msum = _mm_cvtsi128_si32(s32);
3293 let d = 2 * msum - gsum[gi];
3296 acc += d as f32 * s;
3297 }
3298 acc
3299 }
3300}
3301
3302#[cfg(target_arch = "x86_64")]
3305#[target_feature(enable = "avx2")]
3306unsafe fn dot_q1_row_1x4_avx2(
3307 bytes: &[u8],
3308 r: usize,
3309 gpr: usize,
3310 xs: [&[i8]; 4],
3311 gsums: [&[i32]; 4],
3312) -> [f32; 4] {
3313 unsafe {
3315 use core::arch::x86_64::*;
3316 let expand = _mm256_setr_epi8(
3317 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,
3318 3, 3, 3,
3319 );
3320 let bitsel = _mm256_setr_epi8(
3321 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
3322 -128, 1, 2, 4, 8, 16, 32, 64, -128,
3323 );
3324 let ones8 = _mm256_set1_epi8(1);
3325 let ones16 = _mm256_set1_epi16(1);
3326 let mut acc = [0f32; 4];
3327 for gi in 0..gpr {
3328 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
3329 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3330 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
3331 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
3332 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
3333 for (k, xq) in xs.iter().enumerate() {
3334 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3335 let sel = _mm256_and_si256(x, mask);
3336 let p16 = _mm256_maddubs_epi16(ones8, sel);
3337 let d32 = _mm256_madd_epi16(p16, ones16);
3338 let hi128 = _mm256_extracti128_si256::<1>(d32);
3339 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
3340 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3341 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3342 let msum = _mm_cvtsi128_si32(s32);
3343 let d = 2 * msum - gsums[k][gi];
3344 acc[k] += d as f32 * s;
3345 }
3346 }
3347 acc
3348 }
3349}
3350
3351#[allow(unreachable_code)]
3352fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
3353 #[cfg(target_arch = "aarch64")]
3354 unsafe {
3355 return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
3356 }
3357 #[cfg(target_arch = "x86_64")]
3358 if avx2_enabled() {
3359 unsafe {
3360 return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
3361 }
3362 }
3363 let _ = gsum;
3364 let mut acc = 0f32;
3365 for gi in 0..gpr {
3366 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
3367 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3368 let mut d = 0i32;
3369 for (j, &b) in tile[2..].iter().enumerate() {
3370 for k in 0..8 {
3371 let w = ((b >> k) & 1) as i32 * 2 - 1;
3372 d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
3373 }
3374 }
3375 acc += d as f32 * s;
3376 }
3377 acc
3378}
3379
3380#[cfg(target_arch = "aarch64")]
3389#[target_feature(enable = "neon,dotprod")]
3390unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
3391 unsafe {
3394 use core::arch::aarch64::*;
3395 use core::arch::asm;
3396 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
3397 let m = vld1q_u8(MASKS.as_ptr());
3398 macro_rules! tile_dot {
3400 ($t:expr, $x:expr) => {{
3401 let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
3402 let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
3403 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
3404 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
3405 let x0 = vld1q_s8($x);
3406 let x1 = vld1q_s8($x.add(16));
3407 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3408 asm!(
3409 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
3410 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
3411 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3412 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
3413 options(pure, nomem, nostack),
3414 );
3415 vaddq_s32(a0, a1)
3416 }};
3417 }
3418 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
3427 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
3428 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
3429 const IW11: [u8; 16] = [
3430 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
3431 ];
3432 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
3433 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
3434 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
3435 let isc = vld1_u8(ISC.as_ptr());
3436 macro_rules! tile_dot_tbl {
3438 ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
3439 let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
3440 let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
3441 let x0 = vld1q_s8($x);
3442 let x1 = vld1q_s8($x.add(16));
3443 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3444 asm!(
3445 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
3446 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
3447 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3448 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
3449 options(pure, nomem, nostack),
3450 );
3451 vaddq_s32(a0, a1)
3452 }};
3453 }
3454 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
3455 let row_base = r * gpr * Q1_TILE;
3456 let abs_end = bytes.len();
3457 let xp = xq.as_ptr();
3458 let gp = gsum.as_ptr();
3459 let mut accv = vdupq_n_f32(0.0);
3460 let mut gi = 0;
3461 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
3464 let t0 = base.add(gi * Q1_TILE);
3465 let ld_a = vld1q_u8(t0);
3466 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
3467 let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
3468 let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
3469 let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
3470 let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
3471 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
3473 let g = vld1q_s32(gp.add(gi));
3474 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
3475 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
3476 let scf: float32x4_t;
3477 asm!(
3478 "fcvtl {o:v}.4s, {i:v}.4h",
3479 o = out(vreg) scf, i = in(vreg) sc16,
3480 options(pure, nomem, nostack),
3481 );
3482 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
3483 gi += 4;
3484 }
3485 let mut acc = vaddvq_f32(accv);
3486 while gi < gpr {
3487 let t = base.add(gi * Q1_TILE);
3488 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3489 let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
3490 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
3491 gi += 1;
3492 }
3493 acc
3494 }
3495}
3496
3497#[cfg(target_arch = "aarch64")]
3502#[target_feature(enable = "neon,dotprod")]
3503unsafe fn dot_q1_row_1x4_sdot(
3504 bytes: &[u8],
3505 r: usize,
3506 gpr: usize,
3507 xs: [&[i8]; 4],
3508 gs: [&[i32]; 4],
3509) -> [f32; 4] {
3510 unsafe {
3512 use core::arch::aarch64::*;
3513 use core::arch::asm;
3514 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
3515 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
3516 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
3517 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
3518 const IW11: [u8; 16] = [
3519 10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
3520 ];
3521 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
3522 let m = vld1q_u8(MASKS.as_ptr());
3523 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
3524 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
3525 let isc = vld1_u8(ISC.as_ptr());
3526 macro_rules! sdot2 {
3527 ($w0:expr, $w1:expr, $x:expr) => {{
3528 let x0 = vld1q_s8($x);
3529 let x1 = vld1q_s8($x.add(16));
3530 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3531 asm!(
3532 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
3533 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
3534 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3535 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
3536 options(pure, nomem, nostack),
3537 );
3538 vaddq_s32(a0, a1)
3539 }};
3540 }
3541 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
3542 let row_base = r * gpr * Q1_TILE;
3543 let abs_end = bytes.len();
3544 let mut accv = [vdupq_n_f32(0.0); 4];
3545 let mut gi = 0;
3546 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
3547 let t0 = base.add(gi * Q1_TILE);
3548 let ld_a = vld1q_u8(t0);
3549 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
3550 let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
3552 let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
3553 let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
3554 let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
3555 let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
3556 let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
3557 let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
3558 let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
3559 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
3560 let scf: float32x4_t;
3561 asm!(
3562 "fcvtl {o:v}.4s, {i:v}.4h",
3563 o = out(vreg) scf, i = in(vreg) sc16,
3564 options(pure, nomem, nostack),
3565 );
3566 for k in 0..4 {
3567 let xp = xs[k].as_ptr();
3568 let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
3569 let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
3570 let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
3571 let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
3572 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
3573 let g = vld1q_s32(gs[k].as_ptr().add(gi));
3574 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
3575 accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
3576 }
3577 gi += 4;
3578 }
3579 let mut acc = [
3580 vaddvq_f32(accv[0]),
3581 vaddvq_f32(accv[1]),
3582 vaddvq_f32(accv[2]),
3583 vaddvq_f32(accv[3]),
3584 ];
3585 while gi < gpr {
3586 let t = base.add(gi * Q1_TILE);
3587 let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3588 let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
3589 let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
3590 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
3591 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
3592 for k in 0..4 {
3593 let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
3594 acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
3595 }
3596 gi += 1;
3597 }
3598 acc
3599 }
3600}
3601
3602#[inline]
3604fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
3605 let gi = j / GROUP_SIZE;
3606 let k = j % GROUP_SIZE;
3607 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
3608 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3609 let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
3610 ((bit as i32 * 2 - 1) as f32, s)
3611}
3612
3613#[inline]
3615fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
3616 let mut acc = 0f32;
3617 for gi in 0..gpr {
3618 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
3619 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3620 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
3621 let mut ga = 0f32;
3622 for (j, &b) in tile[2..].iter().enumerate() {
3623 for k in 0..8 {
3624 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
3625 }
3626 }
3627 acc += ga * s;
3628 }
3629 acc
3630}
3631
3632#[allow(clippy::too_many_arguments)]
3635fn q1_range_a8w8(
3636 bytes: &[u8],
3637 gpr: usize,
3638 act: &SplitAct,
3639 gsum: &[i32],
3640 out: SendMut,
3641 start: usize,
3642 end: usize,
3643) {
3644 for r in start..end {
3645 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
3646 for &(j, xv) in &act.outliers {
3647 let (w, s) = q1_outlier(bytes, r, gpr, j);
3648 acc += w * s * xv;
3649 }
3650 unsafe { *out.at(r) = acc };
3652 }
3653}
3654
3655fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
3657 for r in start..end {
3658 unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
3660 }
3661}
3662
3663fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
3668 let entries = base_len + (rows + 1) * 4;
3669 (base_len, entries, entries <= bytes.len())
3670}
3671
3672#[inline]
3674fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
3675 let o = rp_off + r * 4;
3676 u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
3677}
3678
3679const SIGN5: [[f32; 5]; 256] = {
3683 let mut lut = [[0.0f32; 5]; 256];
3684 let pow3 = [1u16, 3, 9, 27, 81];
3685 let mut byte = 0usize;
3686 while byte < 256 {
3687 let mut i = 0usize;
3688 while i < 5 {
3689 let code = (byte as u16 / pow3[i]) % 3;
3690 lut[byte][i] = if code == 1 {
3691 1.0
3692 } else if code == 2 {
3693 -1.0
3694 } else {
3695 0.0
3696 };
3697 i += 1;
3698 }
3699 byte += 1;
3700 }
3701 lut
3702};
3703
3704const SIGN5_I8: [[i8; 5]; 256] = {
3706 let mut lut = [[0i8; 5]; 256];
3707 let pow3 = [1u16, 3, 9, 27, 81];
3708 let mut byte = 0usize;
3709 while byte < 256 {
3710 let mut i = 0usize;
3711 while i < 5 {
3712 let code = (byte as u16 / pow3[i]) % 3;
3713 lut[byte][i] = if code == 1 {
3714 1
3715 } else if code == 2 {
3716 -1
3717 } else {
3718 0
3719 };
3720 i += 1;
3721 }
3722 byte += 1;
3723 }
3724 lut
3725};
3726
3727const SIGN5_U64: [u64; 256] = {
3733 let mut lut = [0u64; 256];
3734 let pow3 = [1u16, 3, 9, 27, 81];
3735 let mut byte = 0usize;
3736 while byte < 256 {
3737 let mut v = 0u64;
3738 let mut i = 0usize;
3739 while i < 5 {
3740 let code = (byte as u16 / pow3[i]) % 3;
3741 let s: u8 = if code == 1 {
3742 1
3743 } else if code == 2 {
3744 0xFF
3745 } else {
3746 0
3747 };
3748 v |= (s as u64) << (i * 8);
3749 i += 1;
3750 }
3751 lut[byte] = v;
3752 byte += 1;
3753 }
3754 lut
3755};
3756
3757#[inline]
3762fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
3763 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3764 let off = (r * gpr + j / GROUP_SIZE) * TILE;
3765 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
3766 let within = j % GROUP_SIZE;
3767 SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
3768}
3769
3770#[cfg(target_arch = "aarch64")]
3773#[target_feature(enable = "neon,dotprod")]
3774#[inline]
3775unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
3776 unsafe {
3778 use core::arch::aarch64::*;
3779 use core::arch::asm;
3780 let w0 = vld1q_s8(w);
3781 let w1 = vld1q_s8(w.add(16));
3782 let x0 = vld1q_s8(x);
3783 let x1 = vld1q_s8(x.add(16));
3784 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3785 asm!(
3786 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
3787 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
3788 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3789 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
3790 options(pure, nomem, nostack),
3791 );
3792 vaddvq_s32(vaddq_s32(a0, a1))
3793 }
3794}
3795
3796#[cfg(target_arch = "x86_64")]
3799#[target_feature(enable = "avx2")]
3800#[inline]
3801unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
3802 unsafe {
3804 use core::arch::x86_64::*;
3805 let wv = _mm256_loadu_si256(w as *const __m256i);
3806 let xv = _mm256_loadu_si256(x as *const __m256i);
3807 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
3808 let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
3809 let hi128 = _mm256_extracti128_si256::<1>(d);
3810 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3811 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3812 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3813 _mm_cvtsi128_si32(s32)
3814 }
3815}
3816
3817#[inline]
3822fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
3823 debug_assert!(dst.len() >= 40);
3824 unsafe {
3827 let p = dst.as_mut_ptr();
3828 for bi in 0..7 {
3829 core::ptr::write_unaligned(
3830 p.add(bi * 5) as *mut u64,
3831 SIGN5_U64[*codes.add(bi) as usize],
3832 );
3833 }
3834 }
3835}
3836
3837#[inline]
3842fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
3843 #[cfg(target_arch = "aarch64")]
3844 unsafe {
3845 return sdot32_i8(w, x);
3846 }
3847 #[cfg(target_arch = "x86_64")]
3848 unsafe {
3849 return i8dot32_avx2(w, x);
3850 }
3851 #[allow(unreachable_code)]
3852 unsafe {
3853 let mut s = 0i32;
3854 for k in 0..GROUP_SIZE {
3855 s += *w.add(k) as i32 * *x.add(k) as i32;
3856 }
3857 s
3858 }
3859}
3860
3861#[inline]
3862unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
3863 let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
3864 (
3865 SIGN5_U64[*codes as usize],
3866 SIGN5_U64[*codes.add(1) as usize],
3867 SIGN5_U64[*codes.add(2) as usize],
3868 SIGN5_U64[*codes.add(3) as usize],
3869 SIGN5_U64[*codes.add(4) as usize],
3870 SIGN5_U64[*codes.add(5) as usize],
3871 SIGN5_U64[*codes.add(6) as usize],
3872 )
3873 };
3874
3875 let u0 = s0 | (s1 << 40);
3876 let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
3877 let u2 = (s3 >> 8) | (s4 << 32);
3878 let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
3879
3880 (u0, u1, u2, u3)
3881}
3882
3883#[cfg(target_arch = "aarch64")]
3887#[target_feature(enable = "neon,dotprod")]
3888unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3889 use core::arch::aarch64::*;
3890 use core::arch::asm;
3891 unsafe {
3892 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3893 let mut acc = 0f32;
3894 let bytes_ptr = bytes.as_ptr();
3895 let xq_ptr = xq.as_ptr();
3896 let row_off = r * gpr * TILE;
3897
3898 let gpr2 = gpr & !1;
3899 let mut gi = 0;
3900 while gi < gpr2 {
3901 let off0 = row_off + gi * TILE;
3902 let off1 = off0 + TILE;
3903 let s0 = f16_to_f32(u16::from_le_bytes([
3904 *bytes_ptr.add(off0),
3905 *bytes_ptr.add(off0 + 1),
3906 ]));
3907 let s1 = f16_to_f32(u16::from_le_bytes([
3908 *bytes_ptr.add(off1),
3909 *bytes_ptr.add(off1 + 1),
3910 ]));
3911
3912 let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
3913 let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
3914
3915 let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
3916 let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
3917 let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
3918 let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
3919
3920 let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
3921 let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
3922 let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
3923 let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
3924
3925 let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
3926 let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3927 asm!(
3928 "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
3929 "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
3930 "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
3931 "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
3932 a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
3933 a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
3934 w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
3935 w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
3936 options(pure, nomem, nostack),
3937 );
3938 let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
3939 let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
3940 acc += d0 as f32 * s0 + d1 as f32 * s1;
3941 gi += 2;
3942 }
3943
3944 if gi < gpr {
3945 let off = row_off + gi * TILE;
3946 let s = f16_to_f32(u16::from_le_bytes([
3947 *bytes_ptr.add(off),
3948 *bytes_ptr.add(off + 1),
3949 ]));
3950 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
3951 let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
3952 let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
3953 let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
3954 let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
3955 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3956 asm!(
3957 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
3958 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
3959 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3960 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
3961 options(pure, nomem, nostack),
3962 );
3963 let d = vaddvq_s32(vaddq_s32(a0, a1));
3964 acc += d as f32 * s;
3965 }
3966 acc
3967 }
3968}
3969
3970#[cfg(target_arch = "x86_64")]
3972#[target_feature(enable = "avx2")]
3973unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3974 use core::arch::x86_64::*;
3975 unsafe {
3976 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3977 let mut acc = 0f32;
3978 let bytes_ptr = bytes.as_ptr();
3979 let xq_ptr = xq.as_ptr();
3980 let row_off = r * gpr * TILE;
3981
3982 let ones = _mm256_set1_epi16(1);
3983 for gi in 0..gpr {
3984 let off = row_off + gi * TILE;
3985 let s = f16_to_f32(u16::from_le_bytes([
3986 *bytes_ptr.add(off),
3987 *bytes_ptr.add(off + 1),
3988 ]));
3989 let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
3990 let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
3991 let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
3992 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
3993 let d256 = _mm256_madd_epi16(p16, ones);
3994 let d128 = _mm_add_epi32(
3995 _mm256_castsi256_si128(d256),
3996 _mm256_extracti128_si256(d256, 1),
3997 );
3998 let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
3999 let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
4000 acc += d32 as f32 * s;
4001 }
4002 acc
4003 }
4004}
4005
4006#[inline]
4010fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4011 #[cfg(target_arch = "aarch64")]
4012 unsafe {
4013 return q1t_dot_row_sdot(bytes, r, gpr, xq);
4014 }
4015 #[cfg(target_arch = "x86_64")]
4016 unsafe {
4017 return q1t_dot_row_avx2(bytes, r, gpr, xq);
4018 }
4019 #[allow(unreachable_code)]
4020 {
4021 const TILE: usize = cortiq_core::quant::Q1T_TILE;
4022 let mut acc = 0f32;
4023 let mut sg = [0i8; GROUP_SIZE + 8]; for gi in 0..gpr {
4025 let off = (r * gpr + gi) * TILE;
4026 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
4027 q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
4028 let mut d = 0i32;
4029 for k in 0..GROUP_SIZE {
4030 d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
4031 }
4032 acc += d as f32 * s;
4033 }
4034 acc
4035 }
4036}
4037
4038fn q1t_row_outlier_correction(
4045 bytes: &[u8],
4046 r: usize,
4047 rp_off: usize,
4048 entries_off: usize,
4049 has_ov: bool,
4050 x: &[f32],
4051) -> f32 {
4052 if !has_ov {
4053 return 0.0;
4054 }
4055 let (c0, c1) = (
4056 q1t_rowptr(bytes, rp_off, r),
4057 q1t_rowptr(bytes, rp_off, r + 1),
4058 );
4059 let mut corr = 0f32;
4060 for p in c0..c1 {
4061 let e = entries_off + p * 4;
4062 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
4063 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
4064 corr += val * x[col];
4065 }
4066 corr
4067}
4068
4069fn q1t_dequant_row(
4073 bytes: &[u8],
4074 r: usize,
4075 gpr: usize,
4076 rp_off: usize,
4077 entries_off: usize,
4078 has_ov: bool,
4079 buf: &mut [f32],
4080) {
4081 const TILE: usize = cortiq_core::quant::Q1T_TILE;
4082 for g in 0..gpr {
4083 let off = (r * gpr + g) * TILE;
4084 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
4085 let codes = &bytes[off + 2..off + TILE];
4086 let bc = g * GROUP_SIZE;
4087 for bi in 0..6 {
4089 let lut = &SIGN5[codes[bi] as usize];
4090 let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
4091 for i in 0..5 {
4092 d[i] = lut[i] * s;
4093 }
4094 }
4095 let lut = &SIGN5[codes[6] as usize];
4096 buf[bc + 30] = lut[0] * s;
4097 buf[bc + 31] = lut[1] * s;
4098 }
4099 if !has_ov {
4100 return;
4101 }
4102 let (c0, c1) = (
4103 q1t_rowptr(bytes, rp_off, r),
4104 q1t_rowptr(bytes, rp_off, r + 1),
4105 );
4106 for p in c0..c1 {
4107 let e = entries_off + p * 4;
4108 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
4109 buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
4110 }
4111}
4112
4113fn q1t_add_overlay(
4117 bytes: &[u8],
4118 x: &[f32],
4119 rows: usize,
4120 cols: usize,
4121 out: &mut [f32],
4122 pool: Option<&Pool>,
4123) {
4124 const TILE: usize = cortiq_core::quant::Q1T_TILE;
4125 let gpr = cols / GROUP_SIZE;
4126 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
4127 if !has_ov {
4128 return;
4129 }
4130 let out_addr = SendMut(out.as_mut_ptr());
4131 let run = move |start: usize, end: usize| {
4132 for r in start..end {
4133 let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
4134 unsafe { *out_addr.at(r) += corr };
4136 }
4137 };
4138 dispatch_rows(pool, rows, &run);
4139}
4140
4141#[allow(clippy::too_many_arguments)]
4144fn q1t_range_a8w8(
4145 bytes: &[u8],
4146 gpr: usize,
4147 rp_off: usize,
4148 ent_off: usize,
4149 has_ov: bool,
4150 act: &SplitAct,
4151 x: &[f32],
4152 out: SendMut,
4153 start: usize,
4154 end: usize,
4155) {
4156 for r in start..end {
4157 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
4158 for &(j, xv) in &act.outliers {
4159 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
4160 }
4161 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
4162 unsafe { *out.at(r) = acc };
4164 }
4165}
4166
4167#[allow(clippy::too_many_arguments)]
4170fn q1t_range_f32_batch(
4171 bytes: &[u8],
4172 gpr: usize,
4173 rp_off: usize,
4174 ent_off: usize,
4175 has_ov: bool,
4176 x: &[f32],
4177 out: SendMut,
4178 start: usize,
4179 end: usize,
4180) {
4181 const TILE: usize = cortiq_core::quant::Q1T_TILE;
4182 let mut sg = [0f32; GROUP_SIZE];
4183 for r in start..end {
4184 let mut acc = 0f32;
4185 for g in 0..gpr {
4186 let off = (r * gpr + g) * TILE;
4187 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
4188 let codes = &bytes[off + 2..off + TILE];
4189 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
4190 for bi in 0..6 {
4191 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
4192 }
4193 let lut = &SIGN5[codes[6] as usize];
4194 sg[30] = lut[0];
4195 sg[31] = lut[1];
4196 let mut gsum = 0f32;
4197 for k in 0..GROUP_SIZE {
4198 gsum += sg[k] * xg[k];
4199 }
4200 acc += s * gsum;
4201 }
4202 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
4203 unsafe { *out.at(r) = acc };
4205 }
4206}
4207
4208fn q1t_matvec(
4212 bytes: &[u8],
4213 x: &[f32],
4214 rows: usize,
4215 cols: usize,
4216 out: &mut [f32],
4217 pool: Option<&Pool>,
4218) {
4219 debug_assert_eq!(out.len(), rows);
4220 const TILE: usize = cortiq_core::quant::Q1T_TILE;
4221 let gpr = cols / GROUP_SIZE;
4222 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
4223 let out_addr = SendMut(out.as_mut_ptr());
4224 if a8w8_enabled() {
4228 let act = split_act(x);
4229 let act = &act;
4230 let run = move |start: usize, end: usize| {
4231 for r in start..end {
4232 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
4233 for &(j, xv) in &act.outliers {
4234 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
4235 }
4236 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
4237 unsafe { *out_addr.at(r) = acc };
4239 }
4240 };
4241 dispatch_rows(pool, rows, &run);
4242 return;
4243 }
4244 let run = move |start: usize, end: usize| {
4245 let mut sg = [0f32; GROUP_SIZE];
4249 for r in start..end {
4250 let mut acc = 0f32;
4251 for g in 0..gpr {
4252 let off = (r * gpr + g) * TILE;
4253 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
4254 let codes = &bytes[off + 2..off + TILE];
4255 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
4256 for bi in 0..6 {
4257 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
4258 }
4259 let lut = &SIGN5[codes[6] as usize];
4260 sg[30] = lut[0];
4261 sg[31] = lut[1];
4262 let mut gsum = 0f32;
4263 for k in 0..GROUP_SIZE {
4264 gsum += sg[k] * xg[k];
4265 }
4266 acc += s * gsum;
4267 }
4268 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
4269 unsafe { *out_addr.at(r) = acc };
4270 }
4271 };
4272 dispatch_rows(pool, rows, &run);
4273}
4274
4275fn q1t_matmat(
4278 bytes: &[u8],
4279 xs: &[f32],
4280 b: usize,
4281 rows: usize,
4282 cols: usize,
4283 out: &mut [f32],
4284 pool: Option<&Pool>,
4285) {
4286 debug_assert_eq!(out.len(), b * rows);
4287 const TILE: usize = cortiq_core::quant::Q1T_TILE;
4288 let gpr = cols / GROUP_SIZE;
4289 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
4290 let out_addr = SendMut(out.as_mut_ptr());
4291 if a8w8_enabled() {
4295 let acts: Vec<SplitAct> = (0..b)
4296 .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
4297 .collect();
4298 let acts = &acts;
4299 let run = move |start: usize, end: usize| {
4300 let mut sg = vec![0i8; cols + 8]; let mut sc = vec![0f32; gpr]; let mut accs = vec![0f32; b]; for r in start..end {
4304 for g in 0..gpr {
4305 let off = (r * gpr + g) * TILE;
4306 sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
4307 q1t_unpack_group_i8(
4308 bytes.as_ptr().wrapping_add(off + 2),
4309 &mut sg[g * GROUP_SIZE..],
4310 );
4311 }
4312 for bi in 0..b {
4313 let act = &acts[bi];
4314 let mut isum = 0f32;
4315 for g in 0..gpr {
4316 let d = q1t_i8dot32(
4317 sg.as_ptr().wrapping_add(g * GROUP_SIZE),
4318 act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
4319 );
4320 isum += d as f32 * sc[g];
4321 }
4322 let mut acc = isum * act.sx;
4323 for &(j, xv) in &act.outliers {
4324 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
4325 }
4326 accs[bi] = acc;
4327 }
4328 if has_ov {
4332 let (c0, c1) = (
4333 q1t_rowptr(bytes, rp_off, r),
4334 q1t_rowptr(bytes, rp_off, r + 1),
4335 );
4336 for p in c0..c1 {
4337 let e = ent_off + p * 4;
4338 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
4339 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
4340 for bi in 0..b {
4341 accs[bi] += val * xs[bi * cols + col];
4342 }
4343 }
4344 }
4345 for bi in 0..b {
4346 unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
4347 }
4348 }
4349 };
4350 dispatch_rows(pool, rows, &run);
4351 return;
4352 }
4353 let run = move |start: usize, end: usize| {
4354 let mut buf = vec![0f32; cols];
4355 for r in start..end {
4356 q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
4357 for bi in 0..b {
4358 let xr = &xs[bi * cols..(bi + 1) * cols];
4359 let mut acc = 0f32;
4360 for j in 0..cols {
4361 acc += buf[j] * xr[j];
4362 }
4363 unsafe { *out_addr.at(bi * rows + r) = acc };
4364 }
4365 }
4366 };
4367 dispatch_rows(pool, rows, &run);
4368}
4369
4370fn q1_matvec(
4371 bytes: &[u8],
4372 x: &[f32],
4373 rows: usize,
4374 cols: usize,
4375 out: &mut [f32],
4376 pool: Option<&Pool>,
4377) {
4378 debug_assert_eq!(out.len(), rows);
4379 let gpr = cols / GROUP_SIZE;
4380 let out_addr = SendMut(out.as_mut_ptr());
4381 if a8w8_enabled() {
4382 let act = split_act(x);
4383 let gsum = q1_group_sums(&act.xq, gpr);
4384 let (act, gsum) = (&act, &gsum);
4385 let run = move |start: usize, end: usize| {
4386 q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end)
4387 };
4388 dispatch_rows(pool, rows, &run);
4389 return;
4390 }
4391 let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
4392 dispatch_rows(pool, rows, &run);
4393}
4394
4395#[allow(clippy::too_many_arguments)]
4397fn q1_matvec2(
4398 bytes: &[u8],
4399 x1: &[f32],
4400 x2: &[f32],
4401 rows: usize,
4402 cols: usize,
4403 o1: &mut [f32],
4404 o2: &mut [f32],
4405 pool: Option<&Pool>,
4406) {
4407 let gpr = cols / GROUP_SIZE;
4408 let p1 = SendMut(o1.as_mut_ptr());
4409 let p2 = SendMut(o2.as_mut_ptr());
4410 if a8w8_enabled() {
4411 let a1 = split_act(x1);
4412 let a2 = split_act(x2);
4413 let g1 = q1_group_sums(&a1.xq, gpr);
4414 let g2 = q1_group_sums(&a2.xq, gpr);
4415 let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
4416 let run = move |start: usize, end: usize| {
4417 for r in start..end {
4418 let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
4419 let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
4420 for &(j, xv) in &a1.outliers {
4421 let (w, s) = q1_outlier(bytes, r, gpr, j);
4422 v1 += w * s * xv;
4423 }
4424 for &(j, xv) in &a2.outliers {
4425 let (w, s) = q1_outlier(bytes, r, gpr, j);
4426 v2 += w * s * xv;
4427 }
4428 unsafe {
4430 *p1.at(r) = v1;
4431 *p2.at(r) = v2;
4432 }
4433 }
4434 };
4435 dispatch_rows(pool, rows, &run);
4436 return;
4437 }
4438 let run = move |start: usize, end: usize| {
4439 for r in start..end {
4440 unsafe {
4442 *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
4443 *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
4444 }
4445 }
4446 };
4447 dispatch_rows(pool, rows, &run);
4448}
4449
4450#[allow(clippy::too_many_arguments)]
4452fn q1_matmat(
4453 bytes: &[u8],
4454 xs_all: &[f32],
4455 b: usize,
4456 rows: usize,
4457 cols: usize,
4458 out: &mut [f32],
4459 pool: Option<&Pool>,
4460) {
4461 debug_assert_eq!(out.len(), b * rows);
4462 let gpr = cols / GROUP_SIZE;
4463 let out_addr = SendMut(out.as_mut_ptr());
4464 if a8w8_enabled() {
4465 let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
4466 .map(|bi| {
4467 let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
4468 let gsum = q1_group_sums(&act.xq, gpr);
4469 (act, gsum)
4470 })
4471 .collect();
4472 let acts = &acts;
4473 #[cfg(target_arch = "x86_64")]
4474 let blocked_ok = avx2_enabled()
4475 && std::env::var("CMF_X86_BLOCKED")
4476 .map(|v| v != "0")
4477 .unwrap_or(true);
4478 #[cfg(target_arch = "aarch64")]
4479 let blocked_ok = sdot_enabled()
4480 && std::env::var("CMF_X86_BLOCKED")
4481 .map(|v| v != "0")
4482 .unwrap_or(true);
4483 let run = move |start: usize, end: usize| {
4484 for r in start..end {
4485 let mut bi = 0usize;
4486 #[cfg(target_arch = "aarch64")]
4489 if blocked_ok {
4490 while bi + 4 <= acts.len() {
4491 let xs = [
4492 acts[bi].0.xq.as_slice(),
4493 acts[bi + 1].0.xq.as_slice(),
4494 acts[bi + 2].0.xq.as_slice(),
4495 acts[bi + 3].0.xq.as_slice(),
4496 ];
4497 let gs = [
4498 acts[bi].1.as_slice(),
4499 acts[bi + 1].1.as_slice(),
4500 acts[bi + 2].1.as_slice(),
4501 acts[bi + 3].1.as_slice(),
4502 ];
4503 let d = unsafe { dot_q1_row_1x4_sdot(bytes, r, gpr, xs, gs) };
4504 for k in 0..4 {
4505 let (act, _) = &acts[bi + k];
4506 let mut acc = d[k] * act.sx;
4507 for &(j, xv) in &act.outliers {
4508 let (w, sc) = q1_outlier(bytes, r, gpr, j);
4509 acc += w * sc * xv;
4510 }
4511 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
4513 }
4514 bi += 4;
4515 }
4516 }
4517 #[cfg(target_arch = "x86_64")]
4518 if blocked_ok {
4519 while bi + 4 <= acts.len() {
4520 let xs = [
4521 acts[bi].0.xq.as_slice(),
4522 acts[bi + 1].0.xq.as_slice(),
4523 acts[bi + 2].0.xq.as_slice(),
4524 acts[bi + 3].0.xq.as_slice(),
4525 ];
4526 let gs = [
4527 acts[bi].1.as_slice(),
4528 acts[bi + 1].1.as_slice(),
4529 acts[bi + 2].1.as_slice(),
4530 acts[bi + 3].1.as_slice(),
4531 ];
4532 let d = unsafe { dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs) };
4533 for k in 0..4 {
4534 let (act, _) = &acts[bi + k];
4535 let mut acc = d[k] * act.sx;
4536 for &(j, xv) in &act.outliers {
4537 let (w, sc) = q1_outlier(bytes, r, gpr, j);
4538 acc += w * sc * xv;
4539 }
4540 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
4542 }
4543 bi += 4;
4544 }
4545 }
4546 while bi < acts.len() {
4547 let (act, gsum) = &acts[bi];
4548 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
4549 for &(j, xv) in &act.outliers {
4550 let (w, s) = q1_outlier(bytes, r, gpr, j);
4551 acc += w * s * xv;
4552 }
4553 unsafe { *out_addr.at(bi * rows + r) = acc };
4555 bi += 1;
4556 }
4557 }
4558 };
4559 dispatch_rows(pool, rows, &run);
4560 return;
4561 }
4562 let run = move |start: usize, end: usize| {
4563 for r in start..end {
4564 for bi in 0..b {
4565 let x = &xs_all[bi * cols..(bi + 1) * cols];
4566 unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
4568 }
4569 }
4570 };
4571 dispatch_rows(pool, rows, &run);
4572}
4573
4574fn q4matvec(
4580 bytes: &[u8],
4581 x: &[f32],
4582 rows: usize,
4583 cols: usize,
4584 out: &mut [f32],
4585 pool: Option<&Pool>,
4586) {
4587 debug_assert_eq!(out.len(), rows);
4588 let (packed, scales) = q4_split(bytes, rows, cols);
4589 let gpr = cols / GROUP_SIZE;
4590 let out_addr = SendMut(out.as_mut_ptr());
4591
4592 if a8w8_enabled() {
4593 let act = split_act(x);
4594 let run = move |start: usize, end: usize| {
4595 q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
4596 };
4597 dispatch_rows(pool, rows, &run);
4598 return;
4599 }
4600
4601 let run =
4602 move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
4603 dispatch_rows(pool, rows, &run);
4604}
4605
4606#[inline]
4609#[allow(unreachable_code)]
4610#[cfg(target_arch = "x86_64")]
4615#[target_feature(enable = "avx2")]
4616unsafe fn dot_q4b_row_1x4_avx2(
4617 buf: &[u8],
4618 scales: &[u8],
4619 g0: usize,
4620 gpr: usize,
4621 xs: [&[i8]; 4],
4622) -> [f32; 4] {
4623 unsafe {
4625 use core::arch::x86_64::*;
4626 let ones = _mm256_set1_epi16(1);
4627 let mut acc = [0f32; 4];
4628 for gi in 0..gpr {
4629 let s = f16_to_f32(u16::from_le_bytes([
4630 scales[(g0 + gi) * 2],
4631 scales[(g0 + gi) * 2 + 1],
4632 ]));
4633 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4634 let aw = _mm256_abs_epi8(w);
4635 for (k, xq) in xs.iter().enumerate() {
4636 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4637 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
4638 let d = _mm256_madd_epi16(p16, ones);
4639 let hi128 = _mm256_extracti128_si256::<1>(d);
4640 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4641 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4642 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4643 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
4644 }
4645 }
4646 acc
4647 }
4648}
4649
4650#[cfg(target_arch = "x86_64")]
4656#[target_feature(enable = "avx2")]
4657unsafe fn dot_q4b_row_1x4_sx_avx2(
4658 buf: &[u8],
4659 scales: &[u8],
4660 g0: usize,
4661 gpr: usize,
4662 xs: [&[i8]; 4],
4663 sxs: [f32; 4],
4664) -> [f32; 4] {
4665 unsafe {
4667 use core::arch::x86_64::*;
4668 let ones = _mm256_set1_epi16(1);
4669 let mut acc = [0f32; 4];
4670 for gi in 0..gpr {
4671 let s = f16_to_f32(u16::from_le_bytes([
4672 scales[(g0 + gi) * 2],
4673 scales[(g0 + gi) * 2 + 1],
4674 ]));
4675 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4676 let aw = _mm256_abs_epi8(w);
4677 for (k, xq) in xs.iter().enumerate() {
4678 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4679 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
4680 let d = _mm256_madd_epi16(p16, ones);
4681 let hi128 = _mm256_extracti128_si256::<1>(d);
4682 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4683 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4684 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4685 acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
4686 }
4687 }
4688 acc
4689 }
4690}
4691
4692#[allow(unreachable_code)]
4693fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
4694 #[cfg(target_arch = "aarch64")]
4695 unsafe {
4696 return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
4697 }
4698 #[cfg(target_arch = "x86_64")]
4699 unsafe {
4700 return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
4701 }
4702 let mut acc = 0f32;
4703 for gi in 0..gpr {
4704 let g = g0 + gi;
4705 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4706 let mut d = 0i32;
4707 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
4708 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
4709 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
4710 }
4711 acc += d as f32 * s;
4712 }
4713 acc
4714}
4715
4716#[inline]
4718#[allow(unreachable_code)]
4719fn dot_q4_row_i8_2(
4720 packed: &[u8],
4721 scales: &[u8],
4722 g0: usize,
4723 gpr: usize,
4724 xq1: &[i8],
4725 xq2: &[i8],
4726) -> (f32, f32) {
4727 #[cfg(target_arch = "aarch64")]
4728 unsafe {
4729 return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
4730 }
4731 #[cfg(target_arch = "x86_64")]
4732 unsafe {
4733 return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
4734 }
4735 (
4736 dot_q4_row_i8(packed, scales, g0, gpr, xq1),
4737 dot_q4_row_i8(packed, scales, g0, gpr, xq2),
4738 )
4739}
4740
4741#[allow(clippy::too_many_arguments)]
4744fn q4_range_a8w8(
4745 packed: &[u8],
4746 scales: &[u8],
4747 gpr: usize,
4748 cols: usize,
4749 act: &SplitAct,
4750 out: SendMut,
4751 start: usize,
4752 end: usize,
4753) {
4754 for r in start..end {
4755 let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
4756 for &(j, xv) in &act.outliers {
4758 let flat = r * cols + j;
4759 let byte = packed[flat / 2];
4760 let nib = if flat & 1 == 0 {
4761 byte & 0x0F
4762 } else {
4763 byte >> 4
4764 };
4765 let s = f16_to_f32(u16::from_le_bytes([
4766 scales[(flat / GROUP_SIZE) * 2],
4767 scales[(flat / GROUP_SIZE) * 2 + 1],
4768 ]));
4769 acc += ((nib as i32 - 8) as f32) * s * xv;
4770 }
4771 unsafe { *out.at(r) = acc };
4773 }
4774}
4775
4776#[allow(clippy::too_many_arguments)]
4779fn q4_range2_a8w8(
4780 packed: &[u8],
4781 scales: &[u8],
4782 gpr: usize,
4783 cols: usize,
4784 a1: &SplitAct,
4785 a2: &SplitAct,
4786 p1: SendMut,
4787 p2: SendMut,
4788 start: usize,
4789 end: usize,
4790) {
4791 for r in start..end {
4792 let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
4793 let mut acc1 = s1 * a1.sx;
4794 let mut acc2 = s2 * a2.sx;
4795 let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
4797 for &(j, xv) in outliers {
4798 let flat = r * cols + j;
4799 let byte = packed[flat / 2];
4800 let nib = if flat & 1 == 0 {
4801 byte & 0x0F
4802 } else {
4803 byte >> 4
4804 };
4805 let s = f16_to_f32(u16::from_le_bytes([
4806 scales[(flat / GROUP_SIZE) * 2],
4807 scales[(flat / GROUP_SIZE) * 2 + 1],
4808 ]));
4809 *acc += ((nib as i32 - 8) as f32) * s * xv;
4810 }
4811 };
4812 fix(&a1.outliers, &mut acc1);
4813 fix(&a2.outliers, &mut acc2);
4814 unsafe {
4816 *p1.at(r) = acc1;
4817 *p2.at(r) = acc2;
4818 }
4819 }
4820}
4821
4822fn q4_range_f32(
4824 packed: &[u8],
4825 scales: &[u8],
4826 gpr: usize,
4827 x: &[f32],
4828 out: SendMut,
4829 start: usize,
4830 end: usize,
4831) {
4832 for r in start..end {
4833 let mut acc = 0f32;
4834 for gi in 0..gpr {
4835 let g = r * gpr + gi;
4836 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4837 let pk = &packed[g * 16..(g + 1) * 16];
4838 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4839 let mut ga = 0f32;
4840 for (k, &b) in pk.iter().enumerate() {
4841 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4842 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4843 }
4844 acc += ga * s;
4845 }
4846 unsafe { *out.at(r) = acc };
4848 }
4849}
4850
4851#[allow(clippy::too_many_arguments)]
4855fn q4matvec2(
4856 bytes: &[u8],
4857 x1: &[f32],
4858 x2: &[f32],
4859 rows: usize,
4860 cols: usize,
4861 o1: &mut [f32],
4862 o2: &mut [f32],
4863 pool: Option<&Pool>,
4864) {
4865 debug_assert_eq!(o1.len(), rows);
4866 debug_assert_eq!(o2.len(), rows);
4867 let (packed, scales) = q4_split(bytes, rows, cols);
4868 let gpr = cols / GROUP_SIZE;
4869
4870 if a8w8_enabled() {
4871 let a1 = split_act(x1);
4872 let a2 = split_act(x2);
4873 let p1 = SendMut(o1.as_mut_ptr());
4874 let p2 = SendMut(o2.as_mut_ptr());
4875 let run = move |start: usize, end: usize| {
4876 q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
4877 };
4878 dispatch_rows(pool, rows, &run);
4879 return;
4880 }
4881
4882 let p1 = SendMut(o1.as_mut_ptr());
4883 let p2 = SendMut(o2.as_mut_ptr());
4884 let run = move |start: usize, end: usize| {
4885 q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
4886 };
4887 dispatch_rows(pool, rows, &run);
4888}
4889
4890#[allow(clippy::too_many_arguments)]
4892fn q4_range2_f32(
4893 packed: &[u8],
4894 scales: &[u8],
4895 gpr: usize,
4896 x1: &[f32],
4897 x2: &[f32],
4898 p1: SendMut,
4899 p2: SendMut,
4900 start: usize,
4901 end: usize,
4902) {
4903 for r in start..end {
4904 let (mut acc1, mut acc2) = (0f32, 0f32);
4905 for gi in 0..gpr {
4906 let g = r * gpr + gi;
4907 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4908 let pk = &packed[g * 16..(g + 1) * 16];
4909 let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4910 let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4911 let (mut g1, mut g2) = (0f32, 0f32);
4912 for (k, &b) in pk.iter().enumerate() {
4913 let wl = (b & 0x0F) as f32 - 8.0;
4914 let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
4915 g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
4916 g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
4917 }
4918 acc1 += g1 * s;
4919 acc2 += g2 * s;
4920 }
4921 unsafe {
4923 *p1.at(r) = acc1;
4924 *p2.at(r) = acc2;
4925 }
4926 }
4927}
4928
4929thread_local! {
4930 static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
4933 static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4934}
4935
4936#[allow(clippy::too_many_arguments)]
4942fn q4matmat(
4943 bytes: &[u8],
4944 xs_all: &[f32],
4945 b: usize,
4946 rows: usize,
4947 cols: usize,
4948 out: &mut [f32],
4949 pool: Option<&Pool>,
4950) {
4951 debug_assert_eq!(xs_all.len(), b * cols);
4952 debug_assert_eq!(out.len(), b * rows);
4953 let (packed, scales) = q4_split(bytes, rows, cols);
4954 let gpr = cols / GROUP_SIZE;
4955 let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4956
4957 if a8w8_enabled() {
4958 let acts: Vec<SplitAct> = (0..b)
4959 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
4960 .collect();
4961 let acts = &acts;
4962 let out_addr = SendMut(out.as_mut_ptr());
4963 let run = move |start: usize, end: usize| {
4964 ROW_I8.with(|rb| {
4965 let mut buf = rb.borrow_mut();
4966 buf.resize(cols, 0);
4967 for r in start..end {
4968 for gi in 0..gpr {
4972 let g = r * gpr + gi;
4973 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
4974 buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
4975 buf[gi * GROUP_SIZE + k * 2 + 1] =
4976 (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
4977 }
4978 }
4979 let mut bi = 0usize;
4980 #[cfg(target_arch = "x86_64")]
4981 if avx2_enabled()
4982 && std::env::var("CMF_X86_BLOCKED")
4983 .map(|v| v != "0")
4984 .unwrap_or(true)
4985 {
4986 while bi + 4 <= acts.len() {
4987 let xs = [
4988 acts[bi].xq.as_slice(),
4989 acts[bi + 1].xq.as_slice(),
4990 acts[bi + 2].xq.as_slice(),
4991 acts[bi + 3].xq.as_slice(),
4992 ];
4993 let d = unsafe { dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs) };
4994 for k in 0..4 {
4995 let act = &acts[bi + k];
4996 let mut acc = d[k] * act.sx;
4997 for &(j, xv) in &act.outliers {
4998 acc += (buf[j] as i8) as f32
4999 * gscale((r * cols + j) / GROUP_SIZE)
5000 * xv;
5001 }
5002 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5004 }
5005 bi += 4;
5006 }
5007 }
5008 while bi < acts.len() {
5009 let act = &acts[bi];
5010 let mut acc = 0f32;
5011 for gi in 0..gpr {
5012 let d = dot_i8_i8(
5013 &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
5014 &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
5015 );
5016 acc += d as f32 * gscale(r * gpr + gi);
5017 }
5018 acc *= act.sx;
5019 for &(j, xv) in &act.outliers {
5021 acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
5022 }
5023 unsafe { *out_addr.at(bi * rows + r) = acc };
5025 bi += 1;
5026 }
5027 }
5028 })
5029 };
5030 dispatch_rows(pool, rows, &run);
5031 return;
5032 }
5033
5034 let out_addr = SendMut(out.as_mut_ptr());
5035 let run = move |start: usize, end: usize| {
5036 ROW_F32.with(|rb| {
5037 let mut buf = rb.borrow_mut();
5038 buf.resize(cols, 0.0);
5039 for r in start..end {
5040 for gi in 0..gpr {
5043 let g = r * gpr + gi;
5044 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
5045 buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
5046 buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
5047 }
5048 }
5049 for bi in 0..b {
5050 let x = &xs_all[bi * cols..(bi + 1) * cols];
5051 let mut acc = 0f32;
5052 for gi in 0..gpr {
5053 let mut ga = 0f32;
5054 for k in 0..GROUP_SIZE / 2 {
5059 let e = gi * GROUP_SIZE + k * 2;
5060 ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
5061 }
5062 acc += ga * gscale(r * gpr + gi);
5063 }
5064 unsafe { *out_addr.at(bi * rows + r) = acc };
5066 }
5067 }
5068 })
5069 };
5070 dispatch_rows(pool, rows, &run);
5071}
5072
5073#[allow(clippy::too_many_arguments)]
5078fn vbitmatmat(
5079 bytes: &[u8],
5080 offsets: &[usize],
5081 xs_all: &[f32],
5082 b: usize,
5083 rows: usize,
5084 cols: usize,
5085 out: &mut [f32],
5086 pool: Option<&Pool>,
5087) {
5088 debug_assert_eq!(xs_all.len(), b * cols);
5089 debug_assert_eq!(out.len(), b * rows);
5090 debug_assert_eq!(offsets.len(), rows + 1);
5091 let ng = cols / GROUP_SIZE;
5092 let bits = &bytes[..rows];
5093 let sc_off = rows;
5094 let gscale = |r: usize, g: usize| {
5095 let so = (r * ng + g) * 2;
5096 f16_to_f32(u16::from_le_bytes([
5097 bytes[sc_off + so],
5098 bytes[sc_off + so + 1],
5099 ]))
5100 };
5101
5102 let decode_f32 = |r: usize, dst: &mut [f32]| {
5104 let bw = bits[r] as usize;
5105 let l = ((1i32 << (bw - 1)) - 1) as f32;
5106 let data = &bytes[offsets[r]..offsets[r + 1]];
5107 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
5108 for d in dst.iter_mut() {
5109 while nbits < bw {
5110 acc = (acc << 8) | data[idx] as u64;
5111 idx += 1;
5112 nbits += 8;
5113 }
5114 let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
5115 nbits -= bw;
5116 *d = u - l;
5117 }
5118 };
5119
5120 if a8w8_enabled() {
5121 let acts: Vec<SplitAct> = (0..b)
5122 .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
5123 .collect();
5124 let acts = &acts;
5125 let out_addr = SendMut(out.as_mut_ptr());
5126 let run = move |start: usize, end: usize| {
5127 for r in start..end {
5128 let bw = bits[r] as usize;
5129 if bw == 8 {
5130 ROW_F32.with(|rb| {
5133 let mut buf = rb.borrow_mut();
5134 buf.resize(cols, 0.0);
5135 decode_f32(r, &mut buf);
5136 for bi in 0..b {
5137 let x = &xs_all[bi * cols..(bi + 1) * cols];
5138 let mut dot = 0f32;
5139 for g in 0..ng {
5140 let mut gd = 0f32;
5141 for k in 0..GROUP_SIZE {
5142 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
5143 }
5144 dot += gd * gscale(r, g);
5145 }
5146 unsafe { *out_addr.at(bi * rows + r) = dot };
5148 }
5149 });
5150 continue;
5151 }
5152 let l = (1i32 << (bw - 1)) - 1;
5153 let data = &bytes[offsets[r]..offsets[r + 1]];
5154 ROW_I8.with(|rb| {
5155 let mut buf = rb.borrow_mut();
5156 buf.resize(cols, 0);
5157 #[inline(always)]
5158 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
5159 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
5160 let u = unpack8::<B>(&data[blk * B..]);
5161 for k in 0..8 {
5162 chunk[k] = (u[k] - l) as i8 as u8;
5163 }
5164 }
5165 }
5166 match bw {
5167 3 => fill::<3>(data, l, &mut buf),
5168 4 => vbit_fill4(data, &mut buf),
5169 5 => fill::<5>(data, l, &mut buf),
5170 6 => fill::<6>(data, l, &mut buf),
5171 _ => unreachable!("vbit bit-width {bw} (validated at load)"),
5172 }
5173 let mut bi = 0usize;
5174 #[cfg(target_arch = "x86_64")]
5178 if avx2_enabled()
5179 && std::env::var("CMF_X86_BLOCKED")
5180 .map(|v| v != "0")
5181 .unwrap_or(true)
5182 {
5183 while bi + 4 <= acts.len() {
5184 let xs = [
5185 acts[bi].xq.as_slice(),
5186 acts[bi + 1].xq.as_slice(),
5187 acts[bi + 2].xq.as_slice(),
5188 acts[bi + 3].xq.as_slice(),
5189 ];
5190 let sxs = [
5191 acts[bi].sx,
5192 acts[bi + 1].sx,
5193 acts[bi + 2].sx,
5194 acts[bi + 3].sx,
5195 ];
5196 let d = unsafe {
5197 dot_q4b_row_1x4_sx_avx2(&buf, &bytes[sc_off..], r * ng, ng, xs, sxs)
5198 };
5199 for k in 0..4 {
5200 let act = &acts[bi + k];
5201 let mut dot = d[k];
5202 for &(j, xv) in &act.outliers {
5203 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
5204 }
5205 unsafe { *out_addr.at((bi + k) * rows + r) = dot };
5207 }
5208 bi += 4;
5209 }
5210 }
5211 while bi < acts.len() {
5212 let act = &acts[bi];
5213 let mut dot = 0f32;
5214 for g in 0..ng {
5215 let d = dot_i8_i8(
5216 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
5217 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
5218 ) as f32
5219 * act.sx;
5220 dot += d * gscale(r, g);
5221 }
5222 for &(j, xv) in &act.outliers {
5223 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
5224 }
5225 unsafe { *out_addr.at(bi * rows + r) = dot };
5227 bi += 1;
5228 }
5229 });
5230 }
5231 };
5232 dispatch_rows(pool, rows, &run);
5233 return;
5234 }
5235
5236 let out_addr = SendMut(out.as_mut_ptr());
5237 let run = move |start: usize, end: usize| {
5238 ROW_F32.with(|rb| {
5239 let mut buf = rb.borrow_mut();
5240 buf.resize(cols, 0.0);
5241 for r in start..end {
5242 decode_f32(r, &mut buf);
5243 for bi in 0..b {
5244 let x = &xs_all[bi * cols..(bi + 1) * cols];
5245 let mut dot = 0f32;
5246 for g in 0..ng {
5247 let mut gd = 0f32;
5248 for k in 0..GROUP_SIZE {
5249 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
5250 }
5251 dot += gd * gscale(r, g);
5252 }
5253 unsafe { *out_addr.at(bi * rows + r) = dot };
5255 }
5256 }
5257 })
5258 };
5259 dispatch_rows(pool, rows, &run);
5260}
5261
5262pub(crate) fn gpu_batch_job<'a>(
5266 t: &'a QTensor,
5267 x: &[f32],
5268) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
5269 match t {
5270 QTensor::Mapped {
5271 model,
5272 idx,
5273 dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
5274 rows,
5275 cols,
5276 row_scale,
5277 col_field,
5278 ..
5279 } => Some((
5280 model.clone(),
5281 crate::gpu::BatchJob {
5282 idx: *idx,
5283 rows: *rows,
5284 cols: *cols,
5285 row_scale,
5286 xs: prescale(x, col_field, *dt).into_owned(),
5287 q1: false,
5288 },
5289 )),
5290 QTensor::Mapped {
5292 model,
5293 idx,
5294 dtype: TensorDtype::Q1,
5295 rows,
5296 cols,
5297 ..
5298 } => Some((
5299 model.clone(),
5300 crate::gpu::BatchJob {
5301 idx: *idx,
5302 rows: *rows,
5303 cols: *cols,
5304 row_scale: &[],
5305 xs: x.to_vec(),
5306 q1: true,
5307 },
5308 )),
5309 _ => None,
5310 }
5311}
5312
5313thread_local! {
5314 static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
5315 static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
5316}
5317
5318pub(crate) fn prescale<'a>(
5319 x: &'a [f32],
5320 col_field: &[f32],
5321 dtype: TensorDtype,
5322) -> std::borrow::Cow<'a, [f32]> {
5323 if dtype == TensorDtype::Q8_2f {
5324 x.iter().zip(col_field).map(|(a, c)| a * c).collect()
5325 } else {
5326 std::borrow::Cow::Borrowed(x)
5327 }
5328}
5329
5330pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
5333 x: &[f32],
5334 col_field: &[f32],
5335 dtype: TensorDtype,
5336 buf_id: u8,
5337 f: F,
5338) -> R {
5339 if dtype == TensorDtype::Q8_2f {
5340 if buf_id == 1 {
5341 PRESCALE_BUF1.with(|b| {
5342 let mut buf = b.borrow_mut();
5343 buf.clear();
5344 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
5345 f(&buf)
5346 })
5347 } else {
5348 PRESCALE_BUF2.with(|b| {
5349 let mut buf = b.borrow_mut();
5350 buf.clear();
5351 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
5352 f(&buf)
5353 })
5354 }
5355 } else {
5356 f(x)
5357 }
5358}
5359
5360#[cfg(target_arch = "x86_64")]
5365pub(crate) fn avx2_enabled() -> bool {
5366 use std::sync::OnceLock;
5367 static ON: OnceLock<bool> = OnceLock::new();
5368 *ON.get_or_init(|| {
5369 std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
5370 && std::arch::is_x86_feature_detected!("avx2")
5371 && std::arch::is_x86_feature_detected!("fma")
5372 })
5373}
5374
5375#[cfg(target_arch = "x86_64")]
5380fn avx2_a8w8_enabled() -> bool {
5381 use std::sync::OnceLock;
5382 static ON: OnceLock<bool> = OnceLock::new();
5383 *ON.get_or_init(|| {
5384 avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
5385 })
5386}
5387
5388#[inline]
5392pub(crate) fn a8w8_enabled() -> bool {
5393 #[cfg(target_arch = "aarch64")]
5394 {
5395 sdot_enabled()
5396 }
5397 #[cfg(target_arch = "x86_64")]
5398 {
5399 avx2_a8w8_enabled()
5400 }
5401 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5402 {
5403 false
5404 }
5405}
5406
5407#[inline]
5410#[allow(unreachable_code)]
5411fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
5412 #[cfg(target_arch = "aarch64")]
5413 unsafe {
5414 return dot_i8_sdot(w, xq);
5415 }
5416 #[cfg(target_arch = "x86_64")]
5417 unsafe {
5418 if avx512vnni_enabled() {
5419 return dot_i8_i8_vnni(w, xq);
5420 }
5421 return dot_i8_i8_avx2(w, xq);
5422 }
5423 w.iter()
5424 .zip(xq)
5425 .map(|(&a, &b)| (a as i8) as i32 * b as i32)
5426 .sum()
5427}
5428
5429#[cfg(target_arch = "x86_64")]
5433fn avx512vnni_enabled() -> bool {
5434 use std::sync::OnceLock;
5435 static ON: OnceLock<bool> = OnceLock::new();
5436 *ON.get_or_init(|| {
5437 std::env::var("CMF_AVX512")
5438 .map(|v| v != "0")
5439 .unwrap_or(true)
5440 && std::arch::is_x86_feature_detected!("avx512f")
5441 && std::arch::is_x86_feature_detected!("avx512bw")
5442 && std::arch::is_x86_feature_detected!("avx512vl")
5443 && std::arch::is_x86_feature_detected!("avx512vnni")
5444 })
5445}
5446
5447#[cfg(target_arch = "x86_64")]
5452#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5453unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
5454 unsafe {
5456 use core::arch::x86_64::*;
5457 let n = w.len();
5458 let mut j = 0usize;
5459 let mut total: i32;
5460 {
5465 #[inline(always)]
5466 unsafe fn step(
5467 w: *const u8,
5468 x: *const i8,
5469 acc: core::arch::x86_64::__m512i,
5470 ) -> core::arch::x86_64::__m512i {
5471 unsafe {
5472 use core::arch::x86_64::*;
5473 let wv = _mm512_loadu_si512(w as *const _);
5474 let xv = _mm512_loadu_si512(x as *const _);
5475 let aw = _mm512_abs_epi8(wv);
5476 let neg = _mm512_movepi8_mask(wv);
5477 let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
5478 _mm512_dpbusd_epi32(acc, aw, sx)
5479 }
5480 }
5481 let (mut a0, mut a1, mut a2, mut a3) = (
5482 _mm512_setzero_si512(),
5483 _mm512_setzero_si512(),
5484 _mm512_setzero_si512(),
5485 _mm512_setzero_si512(),
5486 );
5487 while j + 256 <= n {
5488 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
5489 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
5490 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
5491 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
5492 j += 256;
5493 }
5494 while j + 64 <= n {
5495 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
5496 j += 64;
5497 }
5498 let s01 = _mm512_add_epi32(a0, a1);
5499 let s23 = _mm512_add_epi32(a2, a3);
5500 total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
5501 }
5502 if j + 32 <= n {
5504 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
5505 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
5506 let d = _mm256_dpbusd_epi32(
5507 _mm256_setzero_si256(),
5508 _mm256_abs_epi8(wv),
5509 _mm256_sign_epi8(xv, wv),
5510 );
5511 let hi128 = _mm256_extracti128_si256::<1>(d);
5512 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5513 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5514 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5515 total += _mm_cvtsi128_si32(s32);
5516 j += 32;
5517 }
5518 while j < n {
5519 total += (w[j] as i8) as i32 * xq[j] as i32;
5520 j += 1;
5521 }
5522 total
5523 }
5524}
5525
5526#[cfg(target_arch = "x86_64")]
5528#[target_feature(enable = "avx2,fma")]
5529unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
5530 unsafe {
5532 use core::arch::x86_64::*;
5533 let n = x.len();
5534 let wp = w.as_ptr();
5535 let xp = x.as_ptr();
5536 let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
5537 let mut j = 0usize;
5538 while j + 16 <= n {
5539 let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
5540 let lo = _mm256_cvtepi8_epi32(wb);
5541 let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
5542 a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
5543 a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
5544 j += 16;
5545 }
5546 let acc = _mm256_add_ps(a0, a1);
5547 let hi128 = _mm256_extractf128_ps::<1>(acc);
5548 let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
5549 let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
5550 let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
5551 let mut sum = _mm_cvtss_f32(s32);
5552 while j < n {
5553 sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
5554 j += 1;
5555 }
5556 sum
5557 }
5558}
5559
5560#[cfg(target_arch = "x86_64")]
5565#[target_feature(enable = "avx2")]
5566unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
5567 unsafe {
5569 use core::arch::x86_64::*;
5570 let n = w.len();
5571 let ones = _mm256_set1_epi16(1);
5572 let mut acc = _mm256_setzero_si256();
5573 let mut j = 0usize;
5574 while j + 32 <= n {
5575 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
5576 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
5577 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
5578 acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
5579 j += 32;
5580 }
5581 let hi128 = _mm256_extracti128_si256::<1>(acc);
5582 let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
5583 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5584 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5585 let mut s = _mm_cvtsi128_si32(s32);
5586 while j < n {
5587 s += (w[j] as i8) as i32 * xq[j] as i32;
5588 j += 1;
5589 }
5590 s
5591 }
5592}
5593
5594#[cfg(target_arch = "aarch64")]
5598#[target_feature(enable = "neon,i8mm")]
5599unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
5600 unsafe {
5602 use core::arch::aarch64::*;
5603 use core::arch::asm;
5604 let n = w0.len();
5605 let w0p = w0.as_ptr() as *const i8;
5606 let w1p = w1.as_ptr() as *const i8;
5607 let mut acc01 = vdupq_n_s32(0);
5610 let mut acc23 = vdupq_n_s32(0);
5611 let mut i = 0usize;
5612 while i + 8 <= n {
5613 let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
5614 let xb01 = vcombine_s8(
5615 vld1_s8(xs[0].as_ptr().add(i)),
5616 vld1_s8(xs[1].as_ptr().add(i)),
5617 );
5618 let xb23 = vcombine_s8(
5619 vld1_s8(xs[2].as_ptr().add(i)),
5620 vld1_s8(xs[3].as_ptr().add(i)),
5621 );
5622 asm!(
5623 "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
5624 "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
5625 a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
5626 w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
5627 options(pure, nomem, nostack),
5628 );
5629 i += 8;
5630 }
5631 let mut out = [[0i32; 4]; 2];
5632 let a01: [i32; 4] = core::mem::transmute(acc01);
5633 let a23: [i32; 4] = core::mem::transmute(acc23);
5634 out[0][0] = a01[0];
5635 out[0][1] = a01[1];
5636 out[1][0] = a01[2];
5637 out[1][1] = a01[3];
5638 out[0][2] = a23[0];
5639 out[0][3] = a23[1];
5640 out[1][2] = a23[2];
5641 out[1][3] = a23[3];
5642 if i < n {
5643 for (k, x) in xs.iter().enumerate() {
5644 for j in i..n {
5645 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
5646 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
5647 }
5648 }
5649 }
5650 out
5651 }
5652}
5653
5654#[cfg(target_arch = "aarch64")]
5658#[target_feature(enable = "neon,dotprod")]
5659unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
5660 unsafe {
5662 use core::arch::aarch64::*;
5663 use core::arch::asm;
5664 let n = w0.len();
5665 let w0p = w0.as_ptr() as *const i8;
5666 let w1p = w1.as_ptr() as *const i8;
5667 let mut acc = [[vdupq_n_s32(0); 4]; 2];
5668 let mut i = 0usize;
5669 while i + 16 <= n {
5670 let wv0 = vld1q_s8(w0p.add(i));
5671 let wv1 = vld1q_s8(w1p.add(i));
5672 for (k, x) in xs.iter().enumerate() {
5673 let xv = vld1q_s8(x.as_ptr().add(i));
5674 let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
5675 asm!(
5676 "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
5677 "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
5678 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5679 w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
5680 options(pure, nomem, nostack),
5681 );
5682 acc[0][k] = a0;
5683 acc[1][k] = a1;
5684 }
5685 i += 16;
5686 }
5687 let mut out = [[0i32; 4]; 2];
5688 for r in 0..2 {
5689 for k in 0..4 {
5690 out[r][k] = vaddvq_s32(acc[r][k]);
5691 }
5692 }
5693 if i < n {
5694 for (k, x) in xs.iter().enumerate() {
5695 for j in i..n {
5696 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
5697 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
5698 }
5699 }
5700 }
5701 out
5702 }
5703}
5704
5705#[cfg(target_arch = "x86_64")]
5711#[target_feature(enable = "avx2")]
5712unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
5713 unsafe {
5715 use core::arch::x86_64::*;
5716 let n = w0.len();
5717 let ones = _mm256_set1_epi16(1);
5718 let mut acc = [[_mm256_setzero_si256(); 4]; 2];
5719 let mut j = 0usize;
5720 while j + 32 <= n {
5721 let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
5722 let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
5723 let aw0 = _mm256_abs_epi8(wv0);
5724 let aw1 = _mm256_abs_epi8(wv1);
5725 for (k, x) in xs.iter().enumerate() {
5726 let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
5727 let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
5728 acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
5729 let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
5730 acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
5731 }
5732 j += 32;
5733 }
5734 let mut out = [[0i32; 4]; 2];
5735 for r in 0..2 {
5736 for k in 0..4 {
5737 let a = acc[r][k];
5738 let hi128 = _mm256_extracti128_si256::<1>(a);
5739 let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
5740 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5741 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5742 out[r][k] = _mm_cvtsi128_si32(s32);
5743 }
5744 }
5745 if j < n {
5746 for (k, x) in xs.iter().enumerate() {
5747 for i in j..n {
5748 out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
5749 out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
5750 }
5751 }
5752 }
5753 out
5754 }
5755}
5756
5757#[cfg(target_arch = "x86_64")]
5762#[inline]
5763fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
5764 let dot = if avx512vnni_enabled() && row.len() >= 64 {
5765 (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
5766 } else {
5767 unsafe { dot_i8_i8_avx2(row, &act.xq) }
5768 };
5769 let mut acc = dot as f32 * act.sx;
5770 for &(j, xv) in &act.outliers {
5771 acc += (row[j] as i8) as f32 * xv;
5772 }
5773 acc
5774}
5775
5776#[cfg(target_arch = "x86_64")]
5780#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5781unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
5782 unsafe {
5784 use core::arch::x86_64::*;
5785 let n = w.len();
5786 let flip = _mm512_set1_epi8(-128); #[inline(always)]
5788 unsafe fn step(
5789 w: *const u8,
5790 x: *const i8,
5791 flip: core::arch::x86_64::__m512i,
5792 acc: core::arch::x86_64::__m512i,
5793 ) -> core::arch::x86_64::__m512i {
5794 unsafe {
5795 use core::arch::x86_64::*;
5796 let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
5797 _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
5798 }
5799 }
5800 let (mut a0, mut a1, mut a2, mut a3) = (
5801 _mm512_setzero_si512(),
5802 _mm512_setzero_si512(),
5803 _mm512_setzero_si512(),
5804 _mm512_setzero_si512(),
5805 );
5806 let mut j = 0usize;
5807 while j + 256 <= n {
5808 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
5809 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
5810 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
5811 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
5812 j += 256;
5813 }
5814 while j + 64 <= n {
5815 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
5816 j += 64;
5817 }
5818 let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
5819 _mm512_add_epi32(a0, a1),
5820 _mm512_add_epi32(a2, a3),
5821 ));
5822 while j < n {
5824 total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
5825 j += 1;
5826 }
5827 total
5828 }
5829}
5830
5831#[cfg(target_arch = "x86_64")]
5837#[target_feature(enable = "avx2")]
5838unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
5839 unsafe {
5842 use core::arch::x86_64::*;
5843 let lomask = _mm_set1_epi8(0x0F);
5844 let eight = _mm256_set1_epi8(8);
5845 let ones = _mm256_set1_epi16(1);
5846 let mut acc = 0f32;
5847 for gi in 0..gpr {
5848 let g = g0 + gi;
5849 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
5850 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
5851 let lo = _mm_and_si128(b, lomask);
5852 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5853 let w = _mm256_sub_epi8(
5854 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5855 eight,
5856 );
5857 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5858 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5859 let d = _mm256_madd_epi16(p16, ones);
5860 let hi128 = _mm256_extracti128_si256::<1>(d);
5861 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5862 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5863 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5864 acc += _mm_cvtsi128_si32(s32) as f32 * s;
5865 }
5866 acc
5867 }
5868}
5869
5870#[cfg(target_arch = "x86_64")]
5873#[target_feature(enable = "avx2")]
5874unsafe fn dot_q4_row_avx2_2(
5875 packed: &[u8],
5876 scales: &[u8],
5877 g0: usize,
5878 gpr: usize,
5879 xq1: &[i8],
5880 xq2: &[i8],
5881) -> (f32, f32) {
5882 unsafe {
5884 use core::arch::x86_64::*;
5885 let lomask = _mm_set1_epi8(0x0F);
5886 let eight = _mm256_set1_epi8(8);
5887 let ones = _mm256_set1_epi16(1);
5888 let (mut acc1, mut acc2) = (0f32, 0f32);
5889 #[inline(always)]
5890 unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
5891 unsafe {
5892 use core::arch::x86_64::*;
5893 let hi128 = _mm256_extracti128_si256::<1>(d);
5894 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5895 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5896 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5897 _mm_cvtsi128_si32(s32)
5898 }
5899 }
5900 for gi in 0..gpr {
5901 let g = g0 + gi;
5902 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
5903 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
5904 let lo = _mm_and_si128(b, lomask);
5905 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5906 let w = _mm256_sub_epi8(
5907 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5908 eight,
5909 );
5910 let aw = _mm256_abs_epi8(w);
5911 let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5912 let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5913 let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
5914 let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
5915 acc1 += hsum(d1) as f32 * s;
5916 acc2 += hsum(d2) as f32 * s;
5917 }
5918 (acc1, acc2)
5919 }
5920}
5921
5922#[cfg(target_arch = "x86_64")]
5924fn q8_range_avx2(
5925 q: &[u8],
5926 row_scale: &[f32],
5927 act: &SplitAct,
5928 cols: usize,
5929 out_addr: SendMut,
5930 start: usize,
5931 end: usize,
5932) {
5933 for o in start..end {
5934 let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
5935 unsafe { *out_addr.at(o) = v };
5937 }
5938}
5939
5940#[cfg(target_arch = "x86_64")]
5942#[allow(clippy::too_many_arguments)]
5943fn q8_range2_avx2(
5944 q: &[u8],
5945 row_scale: &[f32],
5946 a1: &SplitAct,
5947 a2: &SplitAct,
5948 cols: usize,
5949 p1: SendMut,
5950 p2: SendMut,
5951 start: usize,
5952 end: usize,
5953) {
5954 for o in start..end {
5955 let row = &q[o * cols..(o + 1) * cols];
5956 unsafe {
5958 *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
5959 *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
5960 }
5961 }
5962}
5963
5964#[cfg(target_arch = "aarch64")]
5975fn i8mm_enabled() -> bool {
5976 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5977 *ON.get_or_init(|| {
5978 std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
5979 && std::arch::is_aarch64_feature_detected!("i8mm")
5980 })
5981}
5982
5983#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
5987fn sdot_enabled() -> bool {
5988 use std::sync::OnceLock;
5989 static ON: OnceLock<bool> = OnceLock::new();
5990 *ON.get_or_init(|| {
5991 let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
5992 if !want {
5993 return false;
5994 }
5995
5996 #[cfg(target_arch = "aarch64")]
5997 {
5998 if std::arch::is_aarch64_feature_detected!("dotprod") {
5999 return true;
6000 }
6001 #[cfg(target_os = "android")]
6002 {
6003 if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
6004 if cpuinfo.lines().any(|l| {
6005 (l.starts_with("Features") || l.starts_with("features"))
6006 && l.contains("asimddp")
6007 }) {
6008 return true;
6009 }
6010 }
6011 }
6012 false
6013 }
6014 #[cfg(not(target_arch = "aarch64"))]
6015 {
6016 false
6017 }
6018 })
6019}
6020
6021struct SplitAct {
6026 xq: Vec<i8>,
6027 sx: f32,
6028 outliers: Vec<(usize, f32)>,
6029 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
6032 xsum: i32,
6033}
6034
6035thread_local! {
6036 static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
6039 const { std::cell::RefCell::new(Vec::new()) };
6040}
6041
6042impl Drop for SplitAct {
6043 fn drop(&mut self) {
6044 let buf = std::mem::take(&mut self.xq);
6045 if buf.capacity() > 0 {
6046 XQ_FREE.with(|f| {
6047 let mut f = f.borrow_mut();
6048 if f.len() < 16 {
6049 f.push(buf);
6050 }
6051 });
6052 }
6053 }
6054}
6055
6056fn split_act(x: &[f32]) -> SplitAct {
6057 let n = x.len();
6058 let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
6059 let thr = 8.0 * rms;
6060 let mut outliers: Vec<(usize, f32)> = Vec::new();
6064 let mut amax = 0f32;
6065 for (j, &v) in x.iter().enumerate() {
6066 let a = v.abs();
6067 if a > thr {
6068 outliers.push((j, v));
6069 } else if a > amax {
6070 amax = a;
6071 }
6072 }
6073 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
6074 let inv = 1.0 / sx;
6075 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
6076 xq.clear();
6077 xq.reserve(n);
6078 if outliers.is_empty() {
6079 xq.extend(
6080 x.iter()
6081 .map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8),
6082 );
6083 } else {
6084 xq.extend(x.iter().map(|&v| {
6086 if v.abs() > thr {
6087 0
6088 } else {
6089 (v * inv).round().clamp(-127.0, 127.0) as i8
6090 }
6091 }));
6092 }
6093 let xsum = xq.iter().map(|&v| v as i32).sum();
6094 SplitAct {
6095 xq,
6096 sx,
6097 outliers,
6098 xsum,
6099 }
6100}
6101
6102fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
6103 let n = x.len();
6104 let rms = (x
6105 .iter()
6106 .zip(col)
6107 .map(|(&a, &c)| {
6108 let v = a * c;
6109 (v * v) as f64
6110 })
6111 .sum::<f64>()
6112 / n.max(1) as f64)
6113 .sqrt() as f32;
6114 let thr = 8.0 * rms;
6115
6116 let mut outliers = Vec::new();
6117 let mut amax = 0f32;
6118 for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
6119 let v = a * c;
6120 let s = v.abs();
6121 if s > thr {
6122 outliers.push((j, v));
6123 } else if s > amax {
6124 amax = s;
6125 }
6126 }
6127
6128 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
6129 let inv = 1.0 / sx;
6130 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
6131 xq.clear();
6132 xq.reserve(n);
6133 if outliers.is_empty() {
6134 xq.extend(
6135 x.iter()
6136 .zip(col)
6137 .map(|(&a, &c)| ((a * c) * inv).round().clamp(-127.0, 127.0) as i8),
6138 );
6139 } else {
6140 xq.extend(x.iter().zip(col).map(|(&a, &c)| {
6141 let v = a * c;
6142 if v.abs() > thr {
6143 0
6144 } else {
6145 (v * inv).round().clamp(-127.0, 127.0) as i8
6146 }
6147 }));
6148 }
6149 let xsum = xq.iter().map(|&v| v as i32).sum();
6150 SplitAct {
6151 xq,
6152 sx,
6153 outliers,
6154 xsum,
6155 }
6156}
6157
6158#[cfg(target_arch = "aarch64")]
6161#[target_feature(enable = "neon,dotprod")]
6162unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
6163 unsafe {
6165 use core::arch::aarch64::*;
6166 use core::arch::asm;
6167 let wp = w.as_ptr() as *const i8;
6168 let n = w.len();
6169 let (mut a0, mut a1, mut a2, mut a3) = (
6170 vdupq_n_s32(0),
6171 vdupq_n_s32(0),
6172 vdupq_n_s32(0),
6173 vdupq_n_s32(0),
6174 );
6175 let mut i = 0;
6176 while i + 64 <= n {
6177 let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
6178 let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
6179 let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
6180 let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
6181 asm!(
6182 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6183 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6184 "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
6185 "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
6186 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
6187 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6188 w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
6189 options(pure, nomem, nostack),
6190 );
6191 i += 64;
6192 }
6193 while i + 16 <= n {
6194 let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
6195 asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
6196 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
6197 i += 16;
6198 }
6199 let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
6200 while i < n {
6201 s += (*wp.add(i)) as i32 * xq[i] as i32;
6202 i += 1;
6203 }
6204 s
6205 }
6206}
6207
6208#[cfg(target_arch = "aarch64")]
6212#[target_feature(enable = "neon,dotprod")]
6213unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
6214 unsafe {
6216 use core::arch::aarch64::*;
6217 use core::arch::asm;
6218 let n = xq.len();
6219 let px = xq.as_ptr();
6220 let (p0, p1, p2, p3) = (
6221 w0.as_ptr() as *const i8,
6222 w1.as_ptr() as *const i8,
6223 w2.as_ptr() as *const i8,
6224 w3.as_ptr() as *const i8,
6225 );
6226 let (mut a0, mut a1, mut a2, mut a3) = (
6227 vdupq_n_s32(0),
6228 vdupq_n_s32(0),
6229 vdupq_n_s32(0),
6230 vdupq_n_s32(0),
6231 );
6232 let mut i = 0;
6233 while i + 16 <= n {
6234 let x = vld1q_s8(px.add(i));
6235 let v0 = vld1q_s8(p0.add(i));
6236 let v1 = vld1q_s8(p1.add(i));
6237 let v2 = vld1q_s8(p2.add(i));
6238 let v3 = vld1q_s8(p3.add(i));
6239 asm!(
6240 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
6241 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
6242 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
6243 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
6244 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
6245 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
6246 options(pure, nomem, nostack),
6247 );
6248 i += 16;
6249 }
6250 let mut r = [
6251 vaddvq_s32(a0),
6252 vaddvq_s32(a1),
6253 vaddvq_s32(a2),
6254 vaddvq_s32(a3),
6255 ];
6256 while i < n {
6257 let xi = *px.add(i) as i32;
6258 r[0] += (*p0.add(i)) as i32 * xi;
6259 r[1] += (*p1.add(i)) as i32 * xi;
6260 r[2] += (*p2.add(i)) as i32 * xi;
6261 r[3] += (*p3.add(i)) as i32 * xi;
6262 i += 1;
6263 }
6264 r
6265 }
6266}
6267
6268#[cfg(target_arch = "aarch64")]
6275#[target_feature(enable = "neon,dotprod")]
6276unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
6277 unsafe {
6280 use core::arch::aarch64::*;
6281 use core::arch::asm;
6282 let n = xq.len();
6283 let px = xq.as_ptr();
6284 let pg = g.as_ptr() as *const i8;
6285 let (mut a0, mut a1, mut a2, mut a3) = (
6286 vdupq_n_s32(0),
6287 vdupq_n_s32(0),
6288 vdupq_n_s32(0),
6289 vdupq_n_s32(0),
6290 );
6291 let mut i = 0;
6292 while i + 16 <= n {
6293 let x = vld1q_s8(px.add(i));
6294 let base = pg.add(4 * i);
6295 let v0 = vld1q_s8(base);
6296 let v1 = vld1q_s8(base.add(16));
6297 let v2 = vld1q_s8(base.add(32));
6298 let v3 = vld1q_s8(base.add(48));
6299 asm!(
6300 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
6301 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
6302 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
6303 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
6304 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
6305 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
6306 options(pure, nomem, nostack),
6307 );
6308 i += 16;
6309 }
6310 [
6311 vaddvq_s32(a0),
6312 vaddvq_s32(a1),
6313 vaddvq_s32(a2),
6314 vaddvq_s32(a3),
6315 ]
6316 }
6317}
6318
6319#[cfg(target_arch = "aarch64")]
6325fn q8_range_sdot(
6326 q: &[u8],
6327 rep: &[u8],
6328 row_scale: &[f32],
6329 act: &SplitAct,
6330 cols: usize,
6331 out_addr: SendMut,
6332 start: usize,
6333 end: usize,
6334) {
6335 let mut o = start;
6336 if !rep.is_empty() {
6339 while o < end && o % 4 != 0 {
6340 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
6341 unsafe { *out_addr.at(o) = v };
6342 o += 1;
6343 }
6344 }
6345 while o + 4 <= end {
6346 let r = if rep.is_empty() {
6347 unsafe {
6348 dot_i8_sdot_4rows(
6349 &q[o * cols..(o + 1) * cols],
6350 &q[(o + 1) * cols..(o + 2) * cols],
6351 &q[(o + 2) * cols..(o + 3) * cols],
6352 &q[(o + 3) * cols..(o + 4) * cols],
6353 &act.xq,
6354 )
6355 }
6356 } else {
6357 unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
6358 };
6359 for k in 0..4 {
6360 let mut acc = r[k] as f32 * act.sx;
6361 for &(j, xv) in &act.outliers {
6362 acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
6363 }
6364 unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
6366 }
6367 o += 4;
6368 }
6369 while o < end {
6370 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
6371 unsafe { *out_addr.at(o) = v };
6372 o += 1;
6373 }
6374}
6375
6376#[cfg(target_arch = "aarch64")]
6379#[allow(clippy::too_many_arguments)]
6380fn q8_range2_sdot(
6381 q: &[u8],
6382 row_scale: &[f32],
6383 a1: &SplitAct,
6384 a2: &SplitAct,
6385 cols: usize,
6386 p1: SendMut,
6387 p2: SendMut,
6388 start: usize,
6389 end: usize,
6390) {
6391 for o in start..end {
6392 let row = &q[o * cols..(o + 1) * cols];
6393 unsafe {
6395 *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
6396 *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
6397 }
6398 }
6399}
6400
6401#[allow(clippy::too_many_arguments)]
6403fn q8_range2_f32(
6404 q: &[u8],
6405 row_scale: &[f32],
6406 x1: &[f32],
6407 x2: &[f32],
6408 cols: usize,
6409 p1: SendMut,
6410 p2: SendMut,
6411 start: usize,
6412 end: usize,
6413) {
6414 for o in start..end {
6415 let row = &q[o * cols..(o + 1) * cols];
6416 unsafe {
6418 *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
6419 *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
6420 }
6421 }
6422}
6423
6424fn q8_range_f32(
6426 q: &[u8],
6427 row_scale: &[f32],
6428 xs: &[f32],
6429 cols: usize,
6430 out_addr: SendMut,
6431 start: usize,
6432 end: usize,
6433) {
6434 for o in start..end {
6435 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
6436 unsafe { *out_addr.at(o) = v };
6438 }
6439}
6440
6441#[cfg(target_arch = "aarch64")]
6444#[inline]
6445fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
6446 let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
6447 for &(j, xv) in &act.outliers {
6448 acc += (row[j] as i8) as f32 * xv;
6449 }
6450 acc
6451}
6452
6453#[cfg(target_arch = "aarch64")]
6461#[target_feature(enable = "neon,dotprod")]
6462unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
6463 unsafe {
6466 use core::arch::aarch64::*;
6467 use core::arch::asm;
6468 let lomask = vdupq_n_u8(0x0F);
6469 let eight = vdupq_n_s8(8);
6470 let mut acc = 0f32;
6471 for gi in 0..gpr {
6472 let g = g0 + gi;
6473 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
6474 let b = vld1q_u8(packed.as_ptr().add(g * 16));
6475 let lo = vandq_u8(b, lomask);
6476 let hi = vshrq_n_u8::<4>(b);
6477 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6478 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6479 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
6480 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
6481 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6482 asm!(
6483 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6484 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6485 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6486 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6487 options(pure, nomem, nostack),
6488 );
6489 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
6490 }
6491 acc
6492 }
6493}
6494
6495#[cfg(target_arch = "aarch64")]
6500#[target_feature(enable = "neon,dotprod")]
6501unsafe fn dot_q4_row_sdot2(
6502 packed: &[u8],
6503 scales: &[u8],
6504 g0: usize,
6505 gpr: usize,
6506 xq1: &[i8],
6507 xq2: &[i8],
6508) -> (f32, f32) {
6509 unsafe {
6512 use core::arch::aarch64::*;
6513 use core::arch::asm;
6514 let lomask = vdupq_n_u8(0x0F);
6515 let eight = vdupq_n_s8(8);
6516 let (mut acc1, mut acc2) = (0f32, 0f32);
6517 for gi in 0..gpr {
6518 let g = g0 + gi;
6519 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
6520 let b = vld1q_u8(packed.as_ptr().add(g * 16));
6521 let lo = vandq_u8(b, lomask);
6522 let hi = vshrq_n_u8::<4>(b);
6523 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6524 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6525 let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
6526 let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
6527 let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
6528 let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
6529 let (mut a0, mut a1, mut b0, mut b1) = (
6530 vdupq_n_s32(0),
6531 vdupq_n_s32(0),
6532 vdupq_n_s32(0),
6533 vdupq_n_s32(0),
6534 );
6535 asm!(
6536 "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
6537 "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
6538 "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
6539 "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
6540 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6541 b0 = inout(vreg) b0, b1 = inout(vreg) b1,
6542 e0 = in(vreg) e0, e1 = in(vreg) e1,
6543 x10 = in(vreg) x10, x11 = in(vreg) x11,
6544 x20 = in(vreg) x20, x21 = in(vreg) x21,
6545 options(pure, nomem, nostack),
6546 );
6547 acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
6548 acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
6549 }
6550 (acc1, acc2)
6551 }
6552}
6553
6554#[inline]
6559pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
6560 #[cfg(target_arch = "aarch64")]
6561 unsafe {
6562 return axpy_i8_f32_neon(acc, row, w);
6563 }
6564 #[cfg(target_arch = "x86_64")]
6565 if avx2_enabled() {
6566 return unsafe { axpy_i8_f32_avx2(acc, row, w) };
6567 }
6568 #[allow(unreachable_code)]
6569 {
6570 for (a, &b) in acc.iter_mut().zip(row) {
6571 *a += w * b as f32;
6572 }
6573 }
6574}
6575
6576#[cfg(target_arch = "x86_64")]
6578#[target_feature(enable = "avx2,fma")]
6579unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
6580 unsafe {
6582 use core::arch::x86_64::*;
6583 let n = acc.len().min(row.len());
6584 let ap = acc.as_mut_ptr();
6585 let rp = row.as_ptr();
6586 let wv = _mm256_set1_ps(w);
6587 let mut j = 0usize;
6588 while j + 16 <= n {
6589 let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
6590 let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
6591 let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
6592 let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
6593 let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
6594 _mm256_storeu_ps(ap.add(j), v0);
6595 _mm256_storeu_ps(ap.add(j + 8), v1);
6596 j += 16;
6597 }
6598 while j < n {
6599 *ap.add(j) += w * (*rp.add(j)) as f32;
6600 j += 1;
6601 }
6602 }
6603}
6604
6605#[cfg(target_arch = "aarch64")]
6606#[target_feature(enable = "neon")]
6607unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
6608 unsafe {
6610 use core::arch::aarch64::*;
6611 let n = acc.len().min(row.len());
6612 let ap = acc.as_mut_ptr();
6613 let rp = row.as_ptr();
6614 let wv = vdupq_n_f32(w);
6615 let mut j = 0usize;
6616 while j + 16 <= n {
6617 let rb = vld1q_s8(rp.add(j));
6618 let lo = vmovl_s8(vget_low_s8(rb));
6619 let hi = vmovl_s8(vget_high_s8(rb));
6620 for (off, half) in [(0, lo), (8, hi)] {
6621 let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
6622 let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
6623 let o = j + off;
6624 vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
6625 vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
6626 }
6627 j += 16;
6628 }
6629 while j < n {
6630 *ap.add(j) += w * (*rp.add(j)) as f32;
6631 j += 1;
6632 }
6633 }
6634}
6635
6636#[inline]
6639pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
6640 #[cfg(target_arch = "aarch64")]
6641 unsafe {
6642 return dot_i8_f32_neon(w, x);
6643 }
6644 #[cfg(target_arch = "x86_64")]
6645 if avx2_enabled() {
6646 return unsafe { dot_i8_f32_avx2(w, x) };
6647 }
6648 #[allow(unreachable_code)]
6649 {
6650 let mut sum = 0.0f32;
6651 for (j, &b) in w.iter().enumerate() {
6652 sum += (b as i8) as f32 * x[j];
6653 }
6654 sum
6655 }
6656}
6657
6658#[inline]
6662fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
6663 #[cfg(target_arch = "aarch64")]
6664 unsafe {
6665 return dot_i8_col_f32_neon(w, x, col);
6666 }
6667 #[allow(unreachable_code)]
6668 {
6669 let mut sum = 0.0f32;
6670 for (j, &b) in w.iter().enumerate() {
6671 sum += (b as i8) as f32 * x[j] * col[j];
6672 }
6673 sum
6674 }
6675}
6676
6677#[cfg(target_arch = "aarch64")]
6678#[target_feature(enable = "neon")]
6679unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
6680 unsafe {
6682 use core::arch::aarch64::*;
6683 let n = x.len();
6684 let wp = w.as_ptr() as *const i8;
6685 let xp = x.as_ptr();
6686 let cp = col.as_ptr();
6687 let (mut a0, mut a1, mut a2, mut a3) = (
6688 vdupq_n_f32(0.0),
6689 vdupq_n_f32(0.0),
6690 vdupq_n_f32(0.0),
6691 vdupq_n_f32(0.0),
6692 );
6693 let mut j = 0usize;
6694 while j + 16 <= n {
6695 let wb = vld1q_s8(wp.add(j));
6696 let lo = vmovl_s8(vget_low_s8(wb));
6697 let hi = vmovl_s8(vget_high_s8(wb));
6698 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
6699 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
6700 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
6701 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
6702 a0 = vfmaq_f32(
6703 a0,
6704 w0,
6705 vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
6706 );
6707 a1 = vfmaq_f32(
6708 a1,
6709 w1,
6710 vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
6711 );
6712 a2 = vfmaq_f32(
6713 a2,
6714 w2,
6715 vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
6716 );
6717 a3 = vfmaq_f32(
6718 a3,
6719 w3,
6720 vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
6721 );
6722 j += 16;
6723 }
6724 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
6725 while j < n {
6726 sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
6727 j += 1;
6728 }
6729 sum
6730 }
6731}
6732
6733#[cfg(target_arch = "aarch64")]
6734#[target_feature(enable = "neon")]
6735unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
6736 unsafe {
6738 use core::arch::aarch64::*;
6739 let n = x.len();
6740 let wp = w.as_ptr() as *const i8;
6741 let xp = x.as_ptr();
6742 let (mut a0, mut a1, mut a2, mut a3) = (
6743 vdupq_n_f32(0.0),
6744 vdupq_n_f32(0.0),
6745 vdupq_n_f32(0.0),
6746 vdupq_n_f32(0.0),
6747 );
6748 let mut j = 0usize;
6749 while j + 16 <= n {
6750 let wb = vld1q_s8(wp.add(j));
6751 let lo = vmovl_s8(vget_low_s8(wb));
6752 let hi = vmovl_s8(vget_high_s8(wb));
6753 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
6754 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
6755 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
6756 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
6757 a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
6758 a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
6759 a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
6760 a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
6761 j += 16;
6762 }
6763 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
6764 while j < n {
6765 sum += (*wp.add(j)) as f32 * *xp.add(j);
6766 j += 1;
6767 }
6768 sum
6769 }
6770}
6771
6772#[allow(clippy::too_many_arguments)]
6773fn qmatvec(
6774 q: &[u8],
6775 rep: &[u8],
6776 row_scale: &[f32],
6777 x: &[f32],
6778 col_field: &[f32],
6779 dtype: TensorDtype,
6780 rows: usize,
6781 cols: usize,
6782 out: &mut [f32],
6783 pool: Option<&Pool>,
6784) {
6785 debug_assert_eq!(out.len(), rows);
6786 #[cfg(not(target_arch = "aarch64"))]
6787 let _ = rep;
6788
6789 #[cfg(target_arch = "aarch64")]
6790 if sdot_enabled() {
6791 let act = if dtype == TensorDtype::Q8_2f {
6792 split_act_q8_2f(x, col_field)
6793 } else {
6794 split_act(x)
6795 };
6796 let out_addr = SendMut(out.as_mut_ptr());
6797 let run_range = |start: usize, end: usize| {
6798 q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
6799 };
6800 match pool {
6801 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
6802 _ => run_range(0, rows),
6803 }
6804 return;
6805 }
6806 #[cfg(target_arch = "x86_64")]
6809 if avx2_a8w8_enabled() {
6810 let act = if dtype == TensorDtype::Q8_2f {
6811 split_act_q8_2f(x, col_field)
6812 } else {
6813 split_act(x)
6814 };
6815 let out_addr = SendMut(out.as_mut_ptr());
6816 let run_range = |start: usize, end: usize| {
6817 q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
6818 };
6819 match pool {
6820 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
6821 _ => run_range(0, rows),
6822 }
6823 return;
6824 }
6825
6826 prescale_with(x, col_field, dtype, 1, |xs| {
6827 let out_addr = SendMut(out.as_mut_ptr());
6828 let run_range = move |start: usize, end: usize| {
6829 for o in start..end {
6830 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
6831 unsafe { *out_addr.at(o) = v };
6833 }
6834 };
6835 match pool {
6836 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
6837 _ => run_range(0, rows),
6838 }
6839 });
6840}
6841
6842#[allow(clippy::too_many_arguments)]
6843fn qmatvec2(
6844 q: &[u8],
6845 row_scale: &[f32],
6846 x1: &[f32],
6847 x2: &[f32],
6848 col_field: &[f32],
6849 dtype: TensorDtype,
6850 rows: usize,
6851 cols: usize,
6852 o1: &mut [f32],
6853 o2: &mut [f32],
6854 pool: Option<&Pool>,
6855) {
6856 #[cfg(target_arch = "aarch64")]
6857 if sdot_enabled() {
6858 let a1s = if dtype == TensorDtype::Q8_2f {
6859 split_act_q8_2f(x1, col_field)
6860 } else {
6861 split_act(x1)
6862 };
6863 let a2s = if dtype == TensorDtype::Q8_2f {
6864 split_act_q8_2f(x2, col_field)
6865 } else {
6866 split_act(x2)
6867 };
6868 let p1 = SendMut(o1.as_mut_ptr());
6869 let p2 = SendMut(o2.as_mut_ptr());
6870 let run_range = |start: usize, end: usize| {
6871 q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
6872 };
6873 match pool {
6874 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
6875 _ => run_range(0, rows),
6876 }
6877 return;
6878 }
6879 #[cfg(target_arch = "x86_64")]
6880 if avx2_a8w8_enabled() {
6881 let a1s = if dtype == TensorDtype::Q8_2f {
6882 split_act_q8_2f(x1, col_field)
6883 } else {
6884 split_act(x1)
6885 };
6886 let a2s = if dtype == TensorDtype::Q8_2f {
6887 split_act_q8_2f(x2, col_field)
6888 } else {
6889 split_act(x2)
6890 };
6891 let p1 = SendMut(o1.as_mut_ptr());
6892 let p2 = SendMut(o2.as_mut_ptr());
6893 let run_range = |start: usize, end: usize| {
6894 q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
6895 };
6896 match pool {
6897 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
6898 _ => run_range(0, rows),
6899 }
6900 return;
6901 }
6902
6903 prescale_with(x1, col_field, dtype, 1, |x1s| {
6904 prescale_with(x2, col_field, dtype, 2, |x2s| {
6905 let p1 = SendMut(o1.as_mut_ptr());
6906 let p2 = SendMut(o2.as_mut_ptr());
6907 let run_range = move |start: usize, end: usize| {
6908 for o in start..end {
6909 let row = &q[o * cols..(o + 1) * cols];
6910 let s1 = dot_i8_f32(row, x1s) * row_scale[o];
6911 let s2 = dot_i8_f32(row, x2s) * row_scale[o];
6912 unsafe {
6914 *p1.at(o) = s1;
6915 *p2.at(o) = s2;
6916 }
6917 }
6918 };
6919 match pool {
6920 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
6921 _ => run_range(0, rows),
6922 }
6923 });
6924 });
6925}
6926
6927#[derive(Clone, Copy)]
6928struct SendMut(*mut f32);
6929unsafe impl Send for SendMut {}
6930unsafe impl Sync for SendMut {}
6931
6932impl SendMut {
6933 #[inline]
6934 fn at(self, i: usize) -> *mut f32 {
6935 unsafe { self.0.add(i) }
6936 }
6937}
6938
6939#[cfg(test)]
6940mod tests {
6941 use super::*;
6942
6943 #[test]
6944 fn f32_matvec_matches_matvec_rows_bitexact() {
6945 let (rows, cols) = (300, 40);
6946 let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
6947 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
6948 let qt = QTensor::from_f32(w.clone(), rows, cols);
6949
6950 let mut a = vec![0.0f32; rows];
6951 matvec_rows(None, &w, &x, &mut a);
6952 let mut b = vec![0.0f32; rows];
6953 qt.matvec(&x, &mut b, None);
6954 assert_eq!(a, b);
6955 }
6956
6957 #[test]
6958 fn sdot_kernel_exact_on_grid() {
6959 eprintln!("sdot_enabled = {}", sdot_enabled());
6964 let (rows, cols) = (9, 80); let w: Vec<u8> = (0..rows * cols)
6966 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
6967 .collect();
6968 let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
6969 let x: Vec<f32> = (0..cols)
6970 .map(|i| match i % 3 {
6971 0 => 1.0,
6972 1 => -1.0,
6973 _ => 0.0,
6974 })
6975 .collect();
6976 let mut a = vec![0.0f32; rows];
6977 qmatvec(
6978 &w,
6979 &[],
6980 &scales,
6981 &x,
6982 &[],
6983 TensorDtype::Q8Row,
6984 rows,
6985 cols,
6986 &mut a,
6987 None,
6988 );
6989 for o in 0..rows {
6990 let mut acc = 0.0f32;
6991 for j in 0..cols {
6992 acc += (w[o * cols + j] as i8) as f32 * x[j];
6993 }
6994 let expect = acc * scales[o];
6995 assert!(
6996 (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
6997 "row {o}: {} vs {expect}",
6998 a[o]
6999 );
7000 }
7001 }
7002
7003 #[test]
7004 fn q1_tbl_fast_path_matches_reference() {
7005 let (rows, cols) = (5, 256);
7010 let gpr = cols / GROUP_SIZE;
7011 let mut bytes = Vec::new();
7012 for t in 0..rows * gpr {
7013 let s = 0.007 + (t % 11) as f32 * 0.004;
7014 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
7015 for j in 0..4 {
7016 bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
7017 }
7018 }
7019 let x: Vec<f32> = (0..cols)
7020 .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
7021 .collect();
7022 let mut w = vec![0.0f32; rows * cols];
7023 cortiq_core::quant::dequant_q1(&bytes, &mut w);
7024 let mut got = vec![0.0f32; rows];
7025 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
7026 for o in 0..rows {
7027 let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
7028 assert!(
7029 (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
7030 "row {o}: {} vs {expect}",
7031 got[o]
7032 );
7033 }
7034 let b = 5usize;
7037 let mut xs_all = Vec::new();
7038 for bi in 0..b {
7039 xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
7040 }
7041 let mut mm = vec![0.0f32; b * rows];
7042 q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
7043 for bi in 0..b {
7044 let mut single = vec![0.0f32; rows];
7045 q1_matvec(
7046 &bytes,
7047 &xs_all[bi * cols..(bi + 1) * cols],
7048 rows,
7049 cols,
7050 &mut single,
7051 None,
7052 );
7053 assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
7054 }
7055 }
7056
7057 #[test]
7058 fn q1_kernels_match_exact_reference() {
7059 let (rows, cols) = (7, 96);
7061 let gpr = cols / GROUP_SIZE;
7062 let mut bytes = Vec::new();
7063 for t in 0..rows * gpr {
7064 let s = 0.01 + (t % 13) as f32 * 0.003;
7065 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
7066 for j in 0..4 {
7067 bytes.push(((t * 31 + j * 97) % 251) as u8);
7068 }
7069 }
7070 let x: Vec<f32> = (0..cols)
7072 .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
7073 .collect();
7074 let mut w = vec![0.0f32; rows * cols];
7076 cortiq_core::quant::dequant_q1(&bytes, &mut w);
7077 let mut expect = vec![0.0f32; rows];
7078 for o in 0..rows {
7079 expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
7080 }
7081 let mut got = vec![0.0f32; rows];
7082 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
7083 for o in 0..rows {
7084 assert!(
7085 (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
7086 "row {o}: {} vs {}",
7087 got[o],
7088 expect[o]
7089 );
7090 }
7091 let x2: Vec<f32> = x.iter().map(|v| -v).collect();
7093 let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
7094 q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
7095 assert_eq!(a1, got);
7096 let mut xs = x.clone();
7097 xs.extend_from_slice(&x2);
7098 let mut mm = vec![0.0f32; 2 * rows];
7099 q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
7100 assert_eq!(&mm[..rows], got.as_slice());
7101 assert_eq!(&mm[rows..], a2.as_slice());
7102 }
7103
7104 #[test]
7105 fn repack_is_bit_identical() {
7106 let (rows, cols) = (267, 96); let w: Vec<u8> = (0..rows * cols)
7112 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
7113 .collect();
7114 let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
7115 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
7116 let rep = q8_repack_layout(&w, rows, cols);
7117 for g in 0..rows / 4 {
7119 for c in 0..cols / 16 {
7120 for lane in 0..4 {
7121 assert_eq!(
7122 &rep[g * 4 * cols + c * 64 + lane * 16
7123 ..g * 4 * cols + c * 64 + lane * 16 + 16],
7124 &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
7125 );
7126 }
7127 }
7128 }
7129 let mut a = vec![0.0f32; rows];
7130 qmatvec(
7131 &w,
7132 &[],
7133 &scales,
7134 &x,
7135 &[],
7136 TensorDtype::Q8Row,
7137 rows,
7138 cols,
7139 &mut a,
7140 None,
7141 );
7142 let mut b = vec![0.0f32; rows];
7143 qmatvec(
7144 &w,
7145 &rep,
7146 &scales,
7147 &x,
7148 &[],
7149 TensorDtype::Q8Row,
7150 rows,
7151 cols,
7152 &mut b,
7153 None,
7154 );
7155 assert_eq!(a, b, "full-range repack output diverged");
7156
7157 #[cfg(target_arch = "aarch64")]
7158 if sdot_enabled() {
7159 let act = split_act(&x);
7161 let mut c1 = vec![0.0f32; rows];
7162 let mut c2 = vec![0.0f32; rows];
7163 q8_range_sdot(
7164 &w,
7165 &[],
7166 &scales,
7167 &act,
7168 cols,
7169 SendMut(c1.as_mut_ptr()),
7170 3,
7171 rows - 2,
7172 );
7173 q8_range_sdot(
7174 &w,
7175 &rep,
7176 &scales,
7177 &act,
7178 cols,
7179 SendMut(c2.as_mut_ptr()),
7180 3,
7181 rows - 2,
7182 );
7183 assert_eq!(c1, c2, "unaligned-range repack output diverged");
7184 }
7185 }
7186
7187 #[test]
7188 fn sdot_a8w8_noise_is_bounded() {
7189 let (rows, cols) = (16, 512);
7193 let w: Vec<u8> = (0..rows * cols)
7194 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
7195 .collect();
7196 let scales = vec![0.01f32; rows];
7197 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
7198 let mut a = vec![0.0f32; rows];
7199 qmatvec(
7200 &w,
7201 &[],
7202 &scales,
7203 &x,
7204 &[],
7205 TensorDtype::Q8Row,
7206 rows,
7207 cols,
7208 &mut a,
7209 None,
7210 );
7211 let (mut num, mut den) = (0f64, 0f64);
7212 for o in 0..rows {
7213 let mut acc = 0.0f32;
7214 for j in 0..cols {
7215 acc += (w[o * cols + j] as i8) as f32 * x[j];
7216 }
7217 let expect = acc * scales[o];
7218 num += ((a[o] - expect) as f64).powi(2);
7219 den += (expect as f64).powi(2);
7220 }
7221 let rel = (num / den.max(1e-12)).sqrt();
7222 assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
7223 }
7224
7225 #[test]
7226 fn i8_dot_neon_matches_scalar() {
7227 let n = 100;
7228 let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
7229 let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
7230 let mut scalar = 0.0f32;
7231 for j in 0..n {
7232 scalar += (w[j] as i8) as f32 * x[j];
7233 }
7234 let fast = dot_i8_f32(&w, &x);
7235 assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
7236 }
7237
7238 #[test]
7240 fn vbitmatvec_matches_full_dequant() {
7241 let (rows, cols) = (6, 64);
7242 let ng = cols / GROUP_SIZE;
7243 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
7245 let mut bytes = bits.clone();
7246 for g in 0..rows * ng {
7247 let s = 0.02 + 0.001 * g as f32;
7248 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
7249 }
7250 for r in 0..rows {
7251 let b = bits[r] as usize;
7252 let (mut acc, mut nb) = (0u64, 0usize);
7253 let mut rowbytes = Vec::new();
7254 for i in 0..cols {
7255 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
7256 acc = (acc << b) | v;
7257 nb += b;
7258 while nb >= 8 {
7259 nb -= 8;
7260 rowbytes.push(((acc >> nb) & 0xFF) as u8);
7261 }
7262 }
7263 if nb > 0 {
7264 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
7265 }
7266 bytes.extend_from_slice(&rowbytes);
7267 }
7268 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
7269
7270 let mut reference = vec![0f32; rows * cols];
7271 cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
7272 let mut expect = vec![0f32; rows];
7273 for r in 0..rows {
7274 expect[r] = reference[r * cols..(r + 1) * cols]
7275 .iter()
7276 .zip(&x)
7277 .map(|(w, xv)| w * xv)
7278 .sum();
7279 }
7280 let mut got = vec![0f32; rows];
7281 let offsets = vbit_row_offsets(&bytes, rows, cols);
7282 vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
7283 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
7287 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
7288 for r in 0..rows {
7289 assert!(
7290 (got[r] - expect[r]).abs() < tol * scale,
7291 "row {r}: {} vs {}",
7292 got[r],
7293 expect[r]
7294 );
7295 }
7296 }
7297
7298 #[test]
7303 #[cfg(target_arch = "x86_64")]
7304 fn vbit_matmat_blocked_matches_per_row() {
7305 let (rows, cols, b) = (64usize, 128usize, 9usize);
7306 let ng = cols / GROUP_SIZE;
7307 let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
7308 let mut bytes = bits.clone();
7309 for g in 0..rows * ng {
7310 let sc = 0.02 + 0.0005 * g as f32;
7311 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
7312 }
7313 for r in 0..rows {
7314 let bw = bits[r] as usize;
7315 let (mut acc, mut nb) = (0u64, 0usize);
7316 let mut rowbytes = Vec::new();
7317 for i in 0..cols {
7318 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
7319 acc = (acc << bw) | v;
7320 nb += bw;
7321 while nb >= 8 {
7322 nb -= 8;
7323 rowbytes.push(((acc >> nb) & 0xFF) as u8);
7324 }
7325 }
7326 if nb > 0 {
7327 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
7328 }
7329 bytes.extend_from_slice(&rowbytes);
7330 }
7331 let x: Vec<f32> = (0..b * cols)
7332 .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
7333 .collect();
7334 let offsets = vbit_row_offsets(&bytes, rows, cols);
7335 let mut y_a = vec![0f32; b * rows];
7336 let mut y_b = vec![0f32; b * rows];
7337 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
7338 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
7339 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
7340 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
7341 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
7342 let max_d = y_a
7343 .iter()
7344 .zip(&y_b)
7345 .map(|(p, q)| (p - q).abs())
7346 .fold(0.0f32, f32::max);
7347 assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
7348 }
7349
7350 #[test]
7351 fn q4matvec_matches_full_dequant() {
7352 let (rows, cols) = (8, 64);
7353 let groups = rows * cols / GROUP_SIZE;
7354 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
7356 for i in 0..groups * 16 {
7357 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
7358 }
7359 for g in 0..groups {
7360 let s = 0.01 + 0.003 * g as f32;
7361 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
7362 }
7363 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
7364
7365 let mut reference = vec![0.0f32; rows * cols];
7366 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
7367 let mut expect = vec![0.0f32; rows];
7368 for r in 0..rows {
7369 expect[r] = reference[r * cols..(r + 1) * cols]
7370 .iter()
7371 .zip(&x)
7372 .map(|(w, xv)| w * xv)
7373 .sum();
7374 }
7375
7376 let mut got = vec![0.0f32; rows];
7377 q4matvec(&bytes, &x, rows, cols, &mut got, None);
7378 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
7382 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
7383 for r in 0..rows {
7384 assert!(
7385 (got[r] - expect[r]).abs() < tol * scale,
7386 "row {r}: {} vs {}",
7387 got[r],
7388 expect[r]
7389 );
7390 }
7391 }
7392
7393 #[test]
7396 fn vbitmatvec2_equals_two_singles() {
7397 let (rows, cols) = (6, 64);
7398 let ng = cols / GROUP_SIZE;
7399 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
7400 let mut bytes = bits.clone();
7401 for g in 0..rows * ng {
7402 let s = 0.02 + 0.001 * g as f32;
7403 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
7404 }
7405 for r in 0..rows {
7406 let b = bits[r] as usize;
7407 let (mut acc, mut nb) = (0u64, 0usize);
7408 let mut rowbytes = Vec::new();
7409 for i in 0..cols {
7410 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
7411 acc = (acc << b) | v;
7412 nb += b;
7413 while nb >= 8 {
7414 nb -= 8;
7415 rowbytes.push(((acc >> nb) & 0xFF) as u8);
7416 }
7417 }
7418 if nb > 0 {
7419 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
7420 }
7421 bytes.extend_from_slice(&rowbytes);
7422 }
7423 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
7424 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
7425 let offsets = vbit_row_offsets(&bytes, rows, cols);
7426
7427 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
7428 vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
7429 vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
7430 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
7431 vbitmatvec2(
7432 &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
7433 );
7434 assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
7435 assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
7436 }
7437
7438 #[test]
7440 fn q4matvec2_equals_two_singles() {
7441 let (rows, cols) = (8, 128);
7442 let groups = rows * cols / GROUP_SIZE;
7443 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
7444 for i in 0..groups * 16 {
7445 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
7446 }
7447 for g in 0..groups {
7448 let s = 0.01 + 0.003 * g as f32;
7449 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
7450 }
7451 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
7454 x1[9] = 250.0;
7455 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
7456
7457 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
7458 q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
7459 q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
7460 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
7461 q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
7462 assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
7463 assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
7464 }
7465
7466 #[test]
7469 fn matvec_many_equals_separate_matvecs() {
7470 use crate::pool::Pool;
7471 let (r1, r2, cols) = (300, 200, 64);
7472 let mk = |salt: usize, rows: usize| {
7473 QTensor::from_f32(
7474 (0..rows * cols)
7475 .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
7476 .collect(),
7477 rows,
7478 cols,
7479 )
7480 };
7481 let (a, b) = (mk(1, r1), mk(5, r2));
7482 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
7483 let pool = Pool::new(3);
7484
7485 let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
7486 a.matvec(&x, &mut ea, Some(&pool));
7487 b.matvec(&x, &mut eb, Some(&pool));
7488 let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
7489 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
7490 assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
7491 assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
7492 }
7493
7494 #[test]
7497 fn batched_matmat_equals_per_position_matvec() {
7498 let (rows, cols, b) = (8, 64, 5);
7499 let groups = rows * cols / GROUP_SIZE;
7501 let mut q4 = Vec::new();
7502 for i in 0..groups * 16 {
7503 q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
7504 }
7505 for g in 0..groups {
7506 q4.extend_from_slice(
7507 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
7508 );
7509 }
7510 let ng = cols / GROUP_SIZE;
7512 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
7513 let mut vb = bits.clone();
7514 for g in 0..rows * ng {
7515 vb.extend_from_slice(
7516 &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
7517 );
7518 }
7519 for r in 0..rows {
7520 let bw = bits[r] as usize;
7521 let (mut acc, mut nb) = (0u64, 0usize);
7522 let mut rowbytes = Vec::new();
7523 for i in 0..cols {
7524 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
7525 acc = (acc << bw) | v;
7526 nb += bw;
7527 while nb >= 8 {
7528 nb -= 8;
7529 rowbytes.push(((acc >> nb) & 0xFF) as u8);
7530 }
7531 }
7532 if nb > 0 {
7533 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
7534 }
7535 vb.extend_from_slice(&rowbytes);
7536 }
7537 let offsets = vbit_row_offsets(&vb, rows, cols);
7538
7539 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
7540
7541 let mut got = vec![0f32; b * rows];
7543 q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
7544 for bi in 0..b {
7545 let mut expect = vec![0f32; rows];
7546 q4matvec(
7547 &q4,
7548 &xs[bi * cols..(bi + 1) * cols],
7549 rows,
7550 cols,
7551 &mut expect,
7552 None,
7553 );
7554 assert_eq!(
7555 &got[bi * rows..(bi + 1) * rows],
7556 &expect[..],
7557 "q4 batch pos {bi}"
7558 );
7559 }
7560
7561 let mut got = vec![0f32; b * rows];
7563 vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
7564 for bi in 0..b {
7565 let mut expect = vec![0f32; rows];
7566 vbitmatvec(
7567 &vb,
7568 &offsets,
7569 &xs[bi * cols..(bi + 1) * cols],
7570 rows,
7571 cols,
7572 &mut expect,
7573 None,
7574 );
7575 assert_eq!(
7576 &got[bi * rows..(bi + 1) * rows],
7577 &expect[..],
7578 "vbit batch pos {bi}"
7579 );
7580 }
7581 }
7582
7583 #[test]
7587 fn q4_tiled_matches_q4_block_bitexact() {
7588 let (rows, cols, b) = (8usize, 128usize, 3usize);
7589 let groups = rows * cols / GROUP_SIZE;
7590 let mut split = Vec::with_capacity(groups * 18);
7591 for i in 0..groups * 16 {
7592 split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
7593 }
7594 for g in 0..groups {
7595 split.extend_from_slice(
7596 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
7597 );
7598 }
7599 let (packed, scales) = split.split_at(groups * 16);
7601 let mut tiled = Vec::with_capacity(groups * Q4_TILE);
7602 for g in 0..groups {
7603 tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
7604 tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
7605 }
7606
7607 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
7608 x1[9] = 250.0; let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
7610
7611 let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
7612 q4matvec(&split, &x1, rows, cols, &mut a, None);
7613 q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
7614 assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
7615
7616 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
7617 let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
7618 q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
7619 q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
7620 assert_eq!(a1, t1);
7621 assert_eq!(a2, t2);
7622
7623 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
7624 let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
7625 q4matmat(&split, &xs, b, rows, cols, &mut am, None);
7626 q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
7627 assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
7628 }
7629
7630 #[test]
7637 fn q4matvec_sdot_outlier_exact() {
7638 let (rows, cols) = (4, 128);
7639 let groups = rows * cols / GROUP_SIZE;
7640 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
7641 for i in 0..groups * 16 {
7642 bytes.push(((i * 11 + 5) % 256) as u8);
7643 }
7644 for g in 0..groups {
7645 let s = 0.02 + 0.002 * g as f32;
7646 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
7647 }
7648 let mut x: Vec<f32> = (0..cols)
7649 .map(|i| match i % 3 {
7650 0 => 1.0,
7651 1 => -1.0,
7652 _ => 0.0,
7653 })
7654 .collect();
7655 x[17] = 300.0; let mut reference = vec![0.0f32; rows * cols];
7658 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
7659 let mut expect = vec![0.0f32; rows];
7660 for r in 0..rows {
7661 expect[r] = reference[r * cols..(r + 1) * cols]
7662 .iter()
7663 .zip(&x)
7664 .map(|(w, xv)| w * xv)
7665 .sum();
7666 }
7667 let mut got = vec![0.0f32; rows];
7668 q4matvec(&bytes, &x, rows, cols, &mut got, None);
7669 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
7670 for r in 0..rows {
7671 assert!(
7672 (got[r] - expect[r]).abs() < 2e-3 * scale,
7673 "row {r}: {} vs {} (outlier term must be exact)",
7674 got[r],
7675 expect[r]
7676 );
7677 }
7678 }
7679
7680 #[test]
7684 fn q1t_matvec_matches_reference() {
7685 use cortiq_core::quant::{dequant_q1t, f32_to_f16};
7686 let (rows, cols) = (3usize, 64usize); let gpr = cols / GROUP_SIZE;
7688 let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
7689 let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
7691 let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
7692 let mut bytes = Vec::new();
7693 for r in 0..rows {
7694 for g in 0..gpr {
7695 bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
7696 let mut c = [0u8; 7];
7697 for k in 0..GROUP_SIZE {
7698 let code = if is_out(r * cols + g * GROUP_SIZE + k) {
7700 0
7701 } else {
7702 ((k + r * 3 + g) % 3) as u8 };
7704 cortiq_core::quant::q1t_pack(&mut c, k, code);
7705 }
7706 bytes.extend_from_slice(&c);
7707 }
7708 }
7709 let mut row_ptr = vec![0u32; rows + 1];
7712 for &(idx, _) in &outliers {
7713 row_ptr[idx as usize / cols + 1] += 1;
7714 }
7715 for r in 0..rows {
7716 row_ptr[r + 1] += row_ptr[r];
7717 }
7718 for &p in &row_ptr {
7719 bytes.extend_from_slice(&p.to_le_bytes());
7720 }
7721 for &(idx, v) in &outliers {
7722 bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
7723 bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
7724 }
7725
7726 let mut refw = vec![0f32; rows * cols];
7727 dequant_q1t(&bytes, rows, cols, &mut refw);
7728 let x: Vec<f32> = (0..cols)
7731 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
7732 .collect();
7733 let mut expect = vec![0f32; rows];
7734 for r in 0..rows {
7735 let mut a = 0.0f32;
7736 for j in 0..cols {
7737 a += refw[r * cols + j] * x[j];
7738 }
7739 expect[r] = a;
7740 }
7741 let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
7742 let mut got = vec![0f32; rows];
7743 q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
7744 for r in 0..rows {
7745 assert!(
7746 (got[r] - expect[r]).abs() < tol(expect[r]),
7747 "row {r}: {} vs {}",
7748 got[r],
7749 expect[r]
7750 );
7751 }
7752 let x2: Vec<f32> = x.iter().chain(x.iter().map(|v| v)).copied().collect();
7754 let mut gm = vec![0f32; 2 * rows];
7755 q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
7756 for r in 0..rows {
7757 assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
7758 assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
7759 }
7760 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
7765 q1t_matvec(&bytes, &x, rows, cols, &mut p1, None);
7766 q1t_matvec(&bytes, &x, rows, cols, &mut p2, None);
7767 for r in 0..rows {
7768 assert!(
7769 (p1[r] - expect[r]).abs() < tol(expect[r])
7770 && (p2[r] - expect[r]).abs() < tol(expect[r])
7771 );
7772 }
7773 }
7774
7775 #[test]
7779 #[ignore]
7780 fn q1t_matvec_speed() {
7781 use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
7782 use std::time::Instant;
7783 let (rows, cols) = (8192usize, 4096usize); let gpr = cols / GROUP_SIZE;
7785 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
7786 for r in 0..rows {
7787 for g in 0..gpr {
7788 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
7789 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
7790 let mut c = [0u8; 7];
7791 for k in 0..GROUP_SIZE {
7792 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
7793 }
7794 bytes.extend_from_slice(&c);
7795 }
7796 }
7797 let (n, stride) = (rows * cols, 40usize); let mut row_ptr = vec![0u32; rows + 1];
7799 let mut idx = 0usize;
7800 while idx < n {
7801 row_ptr[idx / cols + 1] += 1;
7802 idx += stride;
7803 }
7804 for r in 0..rows {
7805 row_ptr[r + 1] += row_ptr[r];
7806 }
7807 for &p in &row_ptr {
7808 bytes.extend_from_slice(&p.to_le_bytes());
7809 }
7810 let mut idx = 0usize;
7811 while idx < n {
7812 bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
7813 bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
7814 idx += stride;
7815 }
7816 let x: Vec<f32> = (0..cols)
7819 .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
7820 .collect();
7821 let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
7822
7823 let slow = |out: &mut [f32]| {
7825 let mut buf = vec![0f32; cols];
7826 for r in 0..rows {
7827 for g in 0..gpr {
7828 let off = (r * gpr + g) * Q1T_TILE;
7829 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7830 let codes = &bytes[off + 2..off + Q1T_TILE];
7831 for k in 0..GROUP_SIZE {
7832 buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
7833 1 => s,
7834 2 => -s,
7835 _ => 0.0,
7836 };
7837 }
7838 }
7839 out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
7840 + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
7841 }
7842 };
7843 let iters = 5;
7844 let mut a = vec![0f32; rows];
7845 slow(&mut a); let t = Instant::now();
7847 for _ in 0..iters {
7848 slow(&mut a);
7849 }
7850 let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
7851
7852 let mut b = vec![0f32; rows];
7853 q1t_matvec(&bytes, &x, rows, cols, &mut b, None); let t = Instant::now();
7855 for _ in 0..iters {
7856 q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
7857 }
7858 let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
7859
7860 for r in 0..rows {
7861 assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
7862 }
7863 println!(
7864 "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms fused-LUT {fast_ms:.2} ms => {:.2}x",
7865 slow_ms / fast_ms
7866 );
7867 }
7868}