1use crate::pool::{matvec_rows, matvec_rows2, Pool};
15use cortiq_core::quant::{f16_to_f32, GROUP_SIZE, Q1_TILE, Q4_TILE};
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").map(|v| v == "1").unwrap_or(false)
61 })
62}
63
64fn q8_repack(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
68 #[cfg(target_arch = "aarch64")]
69 let arch_ok = sdot_enabled();
70 #[cfg(not(target_arch = "aarch64"))]
71 let arch_ok = false;
72 if !arch_ok || !repack_enabled() || rows < 256 || cols % 16 != 0 {
73 return Vec::new();
74 }
75 q8_repack_layout(bytes, rows, cols)
76}
77
78fn q8_repack_layout(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
81 let groups = rows / 4;
82 let mut rep = vec![0u8; groups * 4 * cols];
83 for g in 0..groups {
84 let dst = &mut rep[g * 4 * cols..(g + 1) * 4 * cols];
85 for c in 0..cols / 16 {
86 for lane in 0..4 {
87 let src = (g * 4 + lane) * cols + c * 16;
88 dst[c * 64 + lane * 16..c * 64 + lane * 16 + 16]
89 .copy_from_slice(&bytes[src..src + 16]);
90 }
91 }
92 }
93 rep
94}
95
96fn vbit_row_offsets(bytes: &[u8], rows: usize, cols: usize) -> Vec<usize> {
99 let ng = cols / GROUP_SIZE;
100 let bits = &bytes[..rows];
101 let mut offsets = Vec::with_capacity(rows + 1);
102 let mut off = rows + rows * ng * 2;
103 for r in 0..rows {
104 offsets.push(off);
105 off += (cols * bits[r] as usize + 7) / 8;
106 }
107 offsets.push(off);
108 offsets
109}
110
111impl QTensor {
112 pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
113 debug_assert_eq!(data.len(), rows * cols);
114 Self::F32 { data, rows, cols }
115 }
116
117 pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
120 let idx = model
123 .tensor_index(name)
124 .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
125 let entry = &model.tensors[idx];
126 if entry.shape.len() != 2 {
127 return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
128 }
129 let (rows, cols) = (entry.shape[0], entry.shape[1]);
130 let bytes = model.entry_bytes(entry);
131
132 match entry.dtype {
133 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
134 let n = rows * cols;
135 let scales_off = n;
136 let row_scale: Vec<f32> = (0..rows)
137 .map(|o| {
138 f16_to_f32(u16::from_le_bytes([
139 bytes[scales_off + o * 2],
140 bytes[scales_off + o * 2 + 1],
141 ]))
142 })
143 .collect();
144 let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
145 let col_off = n + rows * 2;
146 (0..cols)
147 .map(|i| {
148 f16_to_f32(u16::from_le_bytes([
149 bytes[col_off + i * 2],
150 bytes[col_off + i * 2 + 1],
151 ]))
152 })
153 .collect()
154 } else {
155 Vec::new()
156 };
157 Ok(Self::Mapped {
158 model: model.clone(),
159 idx,
160 dtype: entry.dtype,
161 rows,
162 cols,
163 row_scale,
164 col_field,
165 vbit_offsets: Vec::new(),
166 repack: q8_repack(bytes, rows, cols),
167 })
168 }
169 TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
171 model: model.clone(),
172 idx,
173 dtype: entry.dtype,
174 rows,
175 cols,
176 row_scale: Vec::new(),
177 col_field: Vec::new(),
178 vbit_offsets: vbit_row_offsets(bytes, rows, cols),
179 repack: Vec::new(),
180 }),
181 TensorDtype::VbitRo if cols % GROUP_SIZE == 0 => {
185 let (_, off_off, packed_off) =
186 cortiq_core::quant::vbit_ro_sections(rows, cols);
187 let offsets: Vec<usize> = (0..=rows)
188 .map(|r| {
189 packed_off + cortiq_core::quant::vbit_ro_offset(bytes, off_off, r)
190 })
191 .collect();
192 Ok(Self::Mapped {
193 model: model.clone(),
194 idx,
195 dtype: entry.dtype,
196 rows,
197 cols,
198 row_scale: Vec::new(),
199 col_field: Vec::new(),
200 vbit_offsets: offsets,
201 repack: Vec::new(),
202 })
203 }
204 TensorDtype::Q4Tiled if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
210 model: model.clone(),
211 idx,
212 dtype: entry.dtype,
213 rows,
214 cols,
215 row_scale: Vec::new(),
216 col_field: Vec::new(),
217 vbit_offsets: Vec::new(),
218 repack: Vec::new(),
219 }),
220 TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
221 model: model.clone(),
222 idx,
223 dtype: entry.dtype,
224 rows,
225 cols,
226 row_scale: Vec::new(),
227 col_field: Vec::new(),
228 vbit_offsets: Vec::new(),
229 repack: Vec::new(),
230 }),
231 TensorDtype::Q1 if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
233 model: model.clone(),
234 idx,
235 dtype: entry.dtype,
236 rows,
237 cols,
238 row_scale: Vec::new(),
239 col_field: Vec::new(),
240 vbit_offsets: Vec::new(),
241 repack: Vec::new(),
242 }),
243 TensorDtype::Q1T if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
247 model: model.clone(),
248 idx,
249 dtype: entry.dtype,
250 rows,
251 cols,
252 row_scale: Vec::new(),
253 col_field: Vec::new(),
254 vbit_offsets: Vec::new(),
255 repack: Vec::new(),
256 }),
257 _ => {
259 let mut data = vec![0.0f32; rows * cols];
260 cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
261 Ok(Self::from_f32(data, rows, cols))
262 }
263 }
264 }
265
266 pub(crate) fn is_q1(&self) -> bool {
269 matches!(self, Self::Mapped { dtype: TensorDtype::Q1, .. })
270 }
271
272 pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
275 match self {
276 Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
277 _ => None,
278 }
279 }
280
281 pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
288 match self {
289 Self::Mapped {
290 idx,
291 dtype: TensorDtype::Q1 | TensorDtype::Q1T | TensorDtype::Q4Block,
292 rows,
293 cols,
294 ..
295 } => Some((*idx, *rows, *cols)),
296 _ => None,
297 }
298 }
299
300 pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
305 match self {
306 Self::Mapped {
307 idx,
308 dtype: TensorDtype::Q8Row,
309 rows,
310 cols,
311 row_scale,
312 col_field,
313 ..
314 } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
315 _ => None,
316 }
317 }
318
319 pub fn rows(&self) -> usize {
320 match self {
321 Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
322 }
323 }
324
325 pub fn cols(&self) -> usize {
326 match self {
327 Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
328 }
329 }
330
331 pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
334 match self {
335 Self::Mapped { model, idx, dtype: TensorDtype::Q1, .. } => Some((model, *idx)),
336 _ => None,
337 }
338 }
339
340 pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
344 match self {
345 Self::Mapped { model, idx, dtype: TensorDtype::Q8Row, row_scale, .. } => {
346 Some((model, *idx, 0, row_scale.as_slice()))
347 }
348 Self::Mapped { model, idx, dtype: TensorDtype::Q1, .. } => Some((model, *idx, 1, &[])),
349 Self::Mapped { model, idx, dtype: TensorDtype::Q4Tiled, .. } => Some((model, *idx, 2, &[])),
350 Self::Mapped { model, idx, dtype: TensorDtype::Q1T, .. } => Some((model, *idx, 3, &[])),
351 _ => None,
352 }
353 }
354
355 pub fn as_f32(&self) -> Option<&[f32]> {
358 match self {
359 Self::F32 { data, .. } => Some(data),
360 Self::Mapped { .. } => None,
361 }
362 }
363
364 fn quant_bytes(&self) -> &[u8] {
365 match self {
366 Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
367 Self::F32 { .. } => unreachable!("quant_bytes on F32"),
368 }
369 }
370
371 pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
373 let cols = self.cols();
374 debug_assert_eq!(dst.len(), cols);
375 match self {
376 Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
377 Self::Mapped {
378 dtype,
379 row_scale,
380 col_field,
381 vbit_offsets,
382 ..
383 } => {
384 if *dtype == TensorDtype::Q4Tiled {
385 let bytes = self.quant_bytes();
386 let gpr = cols / GROUP_SIZE;
387 for gi in 0..gpr {
388 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
389 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
390 for (k, &b) in tile[2..].iter().enumerate() {
391 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
392 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
393 }
394 }
395 return;
396 }
397 if *dtype == TensorDtype::Q4Block {
398 let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
399 let gpr = cols / GROUP_SIZE;
400 for gi in 0..gpr {
401 let g = r * gpr + gi;
402 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
403 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
404 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
405 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
406 }
407 }
408 return;
409 }
410 if *dtype == TensorDtype::Q1 {
411 let bytes = self.quant_bytes();
412 let gpr = cols / GROUP_SIZE;
413 for gi in 0..gpr {
414 let tile =
415 &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
416 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
417 for (j, &b) in tile[2..].iter().enumerate() {
418 for k in 0..8 {
419 dst[gi * GROUP_SIZE + j * 8 + k] =
420 (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
421 }
422 }
423 }
424 return;
425 }
426 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
427 let bytes = self.quant_bytes();
428 let rows = self.rows();
429 let ng = cols / GROUP_SIZE;
430 let bits = &bytes[..rows];
431 let sc_off = rows;
432 let off = vbit_offsets[r];
435 let b = bits[r] as usize;
436 let l = ((1usize << (b - 1)) - 1) as f32;
437 let data = &bytes[off..];
438 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
439 for (i, d) in dst.iter_mut().enumerate() {
440 while nbits < b {
441 acc = (acc << 8) | data[idx] as u64;
442 idx += 1;
443 nbits += 8;
444 }
445 let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
446 nbits -= b;
447 let so = (r * ng + i / GROUP_SIZE) * 2;
448 let sv = f16_to_f32(u16::from_le_bytes([
449 bytes[sc_off + so],
450 bytes[sc_off + so + 1],
451 ]));
452 *d = (u - l) * sv;
453 }
454 return;
455 }
456 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
457 let s = row_scale[r];
458 match dtype {
459 TensorDtype::Q8Row => {
460 for (d, &b) in dst.iter_mut().zip(q) {
461 *d = (b as i8) as f32 * s;
462 }
463 }
464 TensorDtype::Q8_2f => {
465 for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
466 *d = (b as i8) as f32 * s * col_field[i];
467 }
468 }
469 _ => unreachable!(),
470 }
471 }
472 }
473 }
474
475 pub fn sparse_col_ok(&self) -> bool {
480 match self {
481 Self::F32 { .. } => true,
482 Self::Mapped { dtype, .. } => {
483 matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
484 }
485 }
486 }
487
488 pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
492 let inter = self.cols();
493 let hidden = self.rows();
494 debug_assert_eq!(out.len(), hidden);
495 match self {
496 Self::F32 { data, .. } => {
497 for (k, o) in out.iter_mut().enumerate() {
498 *o += w * data[k * inter + c];
499 }
500 }
501 Self::Mapped {
502 dtype, row_scale, col_field, ..
503 } => {
504 let q = self.quant_bytes();
505 let colf = if *dtype == TensorDtype::Q8_2f {
506 col_field[c]
507 } else {
508 1.0
509 };
510 let wc = w * colf;
511 for (k, o) in out.iter_mut().enumerate() {
512 let b = q[k * inter + c] as i8 as f32;
513 *o += wc * b * row_scale[k];
514 }
515 }
516 }
517 }
518
519 pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
523 let cols = self.cols();
524 match self {
525 Self::F32 { data, .. } => {
526 let row = &data[r * cols..(r + 1) * cols];
527 row.iter().zip(x).map(|(w, v)| w * v).sum()
528 }
529 Self::Mapped { dtype, row_scale, col_field, .. } => match dtype {
530 TensorDtype::Q8Row => {
531 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
532 dot_i8_f32(q, x) * row_scale[r]
533 }
534 TensorDtype::Q8_2f => {
535 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
536 dot_i8_col_f32(q, x, col_field) * row_scale[r]
537 }
538 _ => {
539 self.row_f32(r, scratch);
540 scratch.iter().zip(x).map(|(w, v)| w * v).sum()
541 }
542 },
543 }
544 }
545
546 pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
549 match self {
550 Self::F32 { data, .. } => matvec_rows(pool, data, x, out),
551 Self::Mapped {
552 model,
553 idx,
554 dtype,
555 rows,
556 cols,
557 row_scale,
558 col_field,
559 vbit_offsets,
560 repack,
561 } => {
562 let _ = (model, idx);
563 if *dtype == TensorDtype::Q4Block {
564 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
568 let t0 = std::time::Instant::now();
569 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
570 crate::gpu::ProbeArm::Gpu => {
571 if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
572 crate::gpu::probe_record(
573 crate::gpu::OpClass::Matvec,
574 true,
575 t0.elapsed(),
576 );
577 return;
578 }
579 }
580 crate::gpu::ProbeArm::CpuTimed => {
581 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
582 crate::gpu::probe_record(
583 crate::gpu::OpClass::Matvec,
584 false,
585 t0.elapsed(),
586 );
587 return;
588 }
589 crate::gpu::ProbeArm::Cpu => {}
590 }
591 }
592 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
593 return;
594 }
595 if *dtype == TensorDtype::Q4Tiled {
596 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
597 return;
598 }
599 if *dtype == TensorDtype::Q1 {
600 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
605 let t0 = std::time::Instant::now();
606 let arm = if crate::gpu::q1_force() {
607 crate::gpu::ProbeArm::Gpu
608 } else {
609 crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
610 };
611 match arm {
612 crate::gpu::ProbeArm::Gpu => {
613 if crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out) {
614 crate::gpu::probe_record(
615 crate::gpu::OpClass::Matvec,
616 true,
617 t0.elapsed(),
618 );
619 return;
620 }
621 }
622 crate::gpu::ProbeArm::CpuTimed => {
623 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
624 crate::gpu::probe_record(
625 crate::gpu::OpClass::Matvec,
626 false,
627 t0.elapsed(),
628 );
629 return;
630 }
631 crate::gpu::ProbeArm::Cpu => {}
632 }
633 }
634 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
635 return;
636 }
637 if *dtype == TensorDtype::Q1T {
638 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
642 let t0 = std::time::Instant::now();
643 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
644 crate::gpu::ProbeArm::Gpu => {
645 if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
646 q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
647 crate::gpu::probe_record(
648 crate::gpu::OpClass::Matvec,
649 true,
650 t0.elapsed(),
651 );
652 return;
653 }
654 }
655 crate::gpu::ProbeArm::CpuTimed => {
656 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
657 crate::gpu::probe_record(
658 crate::gpu::OpClass::Matvec,
659 false,
660 t0.elapsed(),
661 );
662 return;
663 }
664 crate::gpu::ProbeArm::Cpu => {}
665 }
666 }
667 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
668 return;
669 }
670 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
671 vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
672 return;
673 }
674 let xs = prescale(x, col_field, *dtype);
675 if *rows >= crate::gpu::min_rows()
680 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
681 && std::env::var("CMF_GPU_LMHEAD").map(|v| v != "0").unwrap_or(true)
682 && crate::gpu::enabled_here()
683 {
684 let t0 = std::time::Instant::now();
687 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
688 crate::gpu::ProbeArm::Gpu => {}
689 crate::gpu::ProbeArm::CpuTimed => {
690 qmatvec(self.quant_bytes(), repack, row_scale, &xs, *rows, *cols, out, pool);
691 crate::gpu::probe_record(
692 crate::gpu::OpClass::Matvec,
693 false,
694 t0.elapsed(),
695 );
696 return;
697 }
698 crate::gpu::ProbeArm::Cpu => {
699 qmatvec(self.quant_bytes(), repack, row_scale, &xs, *rows, *cols, out, pool);
700 return;
701 }
702 }
703 let frac = std::env::var("CMF_GPU_SPLIT")
704 .ok()
705 .and_then(|v| v.parse::<f32>().ok())
706 .unwrap_or(0.5)
707 .clamp(0.0, 1.0);
708 let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
709 let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
710 let bytes = self.quant_bytes();
711 let ok = std::thread::scope(|sc| {
712 let g = sc.spawn(|| {
713 crate::gpu::q8_matvec_range(
714 model,
715 *idx,
716 cpu_rows,
717 &row_scale[cpu_rows..],
718 &xs,
719 *rows - cpu_rows,
720 *cols,
721 out_gpu,
722 )
723 });
724 if cpu_rows > 0 {
725 let rep_cpu = if repack.is_empty() {
728 &[][..]
729 } else {
730 &repack[..(cpu_rows / 4) * 4 * *cols]
731 };
732 qmatvec(
733 &bytes[..cpu_rows * *cols],
734 rep_cpu,
735 &row_scale[..cpu_rows],
736 &xs,
737 cpu_rows,
738 *cols,
739 out_cpu,
740 pool,
741 );
742 }
743 g.join().unwrap_or(false)
744 });
745 if ok {
746 crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
747 return;
748 }
749 qmatvec(
752 &bytes[cpu_rows * *cols..(*rows) * *cols],
753 &[],
754 &row_scale[cpu_rows..],
755 &xs,
756 *rows - cpu_rows,
757 *cols,
758 out_gpu,
759 pool,
760 );
761 return;
762 }
763 qmatvec(self.quant_bytes(), repack, row_scale, &xs, *rows, *cols, out, pool);
764 }
765 }
766 }
767
768 pub fn matvec2(&self, x1: &[f32], x2: &[f32], o1: &mut [f32], o2: &mut [f32], pool: Option<&Pool>) {
770 match self {
771 Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
772 Self::Mapped {
773 dtype,
774 rows,
775 cols,
776 row_scale,
777 col_field,
778 vbit_offsets,
779 ..
780 } => {
781 if *dtype == TensorDtype::Q4Block {
782 q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
783 return;
784 }
785 if *dtype == TensorDtype::Q4Tiled {
786 q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
787 return;
788 }
789 if *dtype == TensorDtype::Q1 {
790 q1_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
791 return;
792 }
793 if *dtype == TensorDtype::Q1T {
794 q1t_matvec(self.quant_bytes(), x1, *rows, *cols, o1, pool);
799 q1t_matvec(self.quant_bytes(), x2, *rows, *cols, o2, pool);
800 return;
801 }
802 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
803 vbitmatvec2(self.quant_bytes(), vbit_offsets, x1, x2, *rows, *cols, o1, o2, pool);
804 return;
805 }
806 let x1s = prescale(x1, col_field, *dtype);
807 let x2s = prescale(x2, col_field, *dtype);
808 qmatvec2(self.quant_bytes(), row_scale, &x1s, &x2s, *rows, *cols, o1, o2, pool);
809 }
810 }
811 }
812}
813
814impl QTensor {
815 pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
820 let cols = self.cols();
821 let rows = self.rows();
822 debug_assert_eq!(xs_all.len(), b * cols);
823 debug_assert_eq!(out.len(), b * rows);
824 if crate::gptq_capture::capturing() {
828 if let Self::Mapped { model, idx, .. } = self {
829 crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
830 }
831 }
832 match self {
833 Self::F32 { data, .. } => {
834 let out_addr = SendMut(out.as_mut_ptr());
835 let run = |start: usize, end: usize| {
836 for o in start..end {
837 let row = &data[o * cols..(o + 1) * cols];
838 for bi in 0..b {
839 let x = &xs_all[bi * cols..(bi + 1) * cols];
840 let mut acc = 0f32;
841 for j in 0..cols {
842 acc += row[j] * x[j];
843 }
844 unsafe { *out_addr.at(bi * rows + o) = acc };
845 }
846 }
847 };
848 dispatch_rows(pool, rows, &run);
849 }
850 Self::Mapped {
851 dtype,
852 row_scale,
853 col_field,
854 vbit_offsets,
855 ..
856 } => {
857 if *dtype == TensorDtype::Q4Block {
858 q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
859 return;
860 }
861 if *dtype == TensorDtype::Q4Tiled {
862 q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
863 return;
864 }
865 if *dtype == TensorDtype::Q1 {
866 if b >= 32
869 && b * rows * cols >= 128_000_000
870 && cols % 64 == 0
871 && crate::gpu::enabled_here()
872 {
873 if let Self::Mapped { model, idx, .. } = self {
874 let t0 = std::time::Instant::now();
875 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
876 crate::gpu::ProbeArm::Gpu => {
877 if crate::gpu::q1_matmat(model, *idx, xs_all, b, rows, cols, out) {
878 crate::gpu::probe_record(
879 crate::gpu::OpClass::Matmat,
880 true,
881 t0.elapsed(),
882 );
883 return;
884 }
885 }
886 crate::gpu::ProbeArm::CpuTimed => {
887 q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
888 crate::gpu::probe_record(
889 crate::gpu::OpClass::Matmat,
890 false,
891 t0.elapsed(),
892 );
893 return;
894 }
895 crate::gpu::ProbeArm::Cpu => {}
896 }
897 }
898 }
899 q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
900 return;
901 }
902 if *dtype == TensorDtype::Q1T {
903 if b >= 32
906 && b * rows * cols >= 128_000_000
907 && crate::gpu::enabled_here()
908 {
909 if let Self::Mapped { model, idx, .. } = self {
910 let t0 = std::time::Instant::now();
911 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
912 crate::gpu::ProbeArm::Gpu => {
913 if crate::gpu::q1t_matmat(model, *idx, xs_all, b, rows, cols, out) {
914 crate::gpu::probe_record(
915 crate::gpu::OpClass::Matmat,
916 true,
917 t0.elapsed(),
918 );
919 return;
920 }
921 }
922 crate::gpu::ProbeArm::CpuTimed => {
923 q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
924 crate::gpu::probe_record(
925 crate::gpu::OpClass::Matmat,
926 false,
927 t0.elapsed(),
928 );
929 return;
930 }
931 crate::gpu::ProbeArm::Cpu => {}
932 }
933 }
934 }
935 q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
936 return;
937 }
938 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
939 vbitmatmat(self.quant_bytes(), vbit_offsets, xs_all, b, rows, cols, out, pool);
940 return;
941 }
942 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
943 .map(|bi| {
944 prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype)
945 })
946 .collect();
947 if b >= 8
953 && b * rows * cols >= 128_000_000
954 && crate::gpu::enabled_here()
955 {
956 if let Self::Mapped { model, idx, .. } = self {
957 let t0 = std::time::Instant::now();
958 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
959 crate::gpu::ProbeArm::Gpu
960 if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
961 && !crate::gpu::q8_resident_or_upload(model, *idx) =>
962 {
963 let q = self.quant_bytes();
967 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
968 return;
969 }
970 crate::gpu::ProbeArm::Gpu => {
971 let flat: Vec<f32> =
972 pre.iter().flat_map(|v| v.iter().copied()).collect();
973 if crate::gpu::q8_matmat(
974 model, *idx, row_scale, &flat, b, rows, cols, out)
975 {
976 crate::gpu::probe_record(
977 crate::gpu::OpClass::Matmat,
978 true,
979 t0.elapsed(),
980 );
981 return;
982 }
983 }
984 crate::gpu::ProbeArm::CpuTimed => {
985 let q = self.quant_bytes();
986 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
987 crate::gpu::probe_record(
988 crate::gpu::OpClass::Matmat,
989 false,
990 t0.elapsed(),
991 );
992 return;
993 }
994 crate::gpu::ProbeArm::Cpu => {}
995 }
996 }
997 }
998 let q = self.quant_bytes();
999 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1000 }
1001 }
1002 }
1003}
1004
1005impl QTensor {
1006 pub fn matvec_many<const N: usize>(
1013 ts: [&QTensor; N],
1014 x: &[f32],
1015 mut outs: [&mut [f32]; N],
1016 pool: Option<&Pool>,
1017 ) {
1018 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
1019 let uniform_q8 = ts.iter().all(|t| {
1020 matches!(
1021 t,
1022 Self::Mapped { dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f, .. }
1023 )
1024 });
1025 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
1026 let uniform_q4 = ts
1027 .iter()
1028 .all(|t| matches!(t, Self::Mapped { dtype: TensorDtype::Q4Block, .. }));
1029 let uniform_vbit = ts
1030 .iter()
1031 .all(|t| matches!(
1032 t,
1033 Self::Mapped { dtype: TensorDtype::Vbit | TensorDtype::VbitRo, .. }
1034 ));
1035 let uniform_q1 = ts
1036 .iter()
1037 .all(|t| matches!(t, Self::Mapped { dtype: TensorDtype::Q1, .. }));
1038 let Some(pool) = pool else {
1039 for (t, o) in ts.iter().zip(outs.iter_mut()) {
1040 t.matvec(x, o, None);
1041 }
1042 return;
1043 };
1044 if total_rows < 256
1045 || !(uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit || uniform_q1)
1046 {
1047 for (t, o) in ts.iter().zip(outs.iter_mut()) {
1048 t.matvec(x, o, Some(pool));
1049 }
1050 return;
1051 }
1052
1053 if uniform_q1 {
1054 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1057 if a8w8_enabled() {
1058 let act = split_act(x);
1059 let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
1060 let (act, gsum) = (&act, &gsum);
1061 let closures: [_; N] = std::array::from_fn(|i| {
1062 let (bytes, gpr, out) =
1063 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1064 move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
1065 });
1066 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1067 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1068 pool.run_many(&parts);
1069 } else {
1070 let closures: [_; N] = std::array::from_fn(|i| {
1071 let (bytes, gpr, out) =
1072 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1073 move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
1074 });
1075 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1076 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1077 pool.run_many(&parts);
1078 }
1079 return;
1080 }
1081
1082 if uniform_q4 || uniform_vbit {
1083 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1084 if a8w8_enabled() {
1086 let act = split_act(x);
1087 let act = &act;
1088 if uniform_q4 {
1089 let closures: [_; N] = std::array::from_fn(|i| {
1090 let (packed, scales) =
1091 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1092 let (gpr, cols, out) =
1093 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
1094 move |s: usize, e: usize| {
1095 q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
1096 }
1097 });
1098 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1099 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1100 pool.run_many(&parts);
1101 } else {
1102 let closures: [_; N] = std::array::from_fn(|i| {
1103 let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
1104 let (bytes, rows, cols, out) =
1105 (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), outs_addr[i]);
1106 move |s: usize, e: usize| {
1107 vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
1108 }
1109 });
1110 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1111 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1112 pool.run_many(&parts);
1113 }
1114 return;
1115 }
1116 if uniform_q4 {
1117 let closures: [_; N] = std::array::from_fn(|i| {
1118 let (packed, scales) =
1119 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1120 let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1121 move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
1122 });
1123 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1124 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1125 pool.run_many(&parts);
1126 } else {
1127 let closures: [_; N] = std::array::from_fn(|i| {
1128 let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
1129 let (bytes, rows, cols, out) =
1130 (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), outs_addr[i]);
1131 move |s: usize, e: usize| {
1132 vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
1133 }
1134 });
1135 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1136 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1137 pool.run_many(&parts);
1138 }
1139 return;
1140 }
1141
1142 if uniform_f32 {
1143 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1144 let closures: [_; N] = std::array::from_fn(|i| {
1145 let Self::F32 { data, cols, .. } = ts[i] else { unreachable!() };
1146 let out = outs_addr[i];
1147 move |start: usize, end: usize| {
1148 for o in start..end {
1149 let row = &data[o * cols..(o + 1) * cols];
1150 let mut sum = 0.0f32;
1151 for j in 0..*cols {
1152 sum += row[j] * x[j];
1153 }
1154 unsafe { *out.at(o) = sum };
1156 }
1157 }
1158 });
1159 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1160 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1161 pool.run_many(&parts);
1162 return;
1163 }
1164
1165 struct Ctx<'a> {
1168 bytes: &'a [u8],
1169 #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
1170 rep: &'a [u8],
1171 row_scale: &'a [f32],
1172 cols: usize,
1173 xs: std::borrow::Cow<'a, [f32]>,
1174 }
1175 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
1176 let Self::Mapped { dtype, cols, row_scale, col_field, repack, .. } = ts[i] else {
1177 unreachable!()
1178 };
1179 Ctx {
1180 bytes: ts[i].quant_bytes(),
1181 rep: repack,
1182 row_scale,
1183 cols: *cols,
1184 xs: prescale(x, col_field, *dtype),
1185 }
1186 });
1187 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1188 #[cfg(target_arch = "aarch64")]
1189 if sdot_enabled() {
1190 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
1191 let closures: [_; N] = std::array::from_fn(|i| {
1192 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
1193 move |start: usize, end: usize| {
1194 q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
1195 }
1196 });
1197 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1198 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1199 pool.run_many(&parts);
1200 return;
1201 }
1202 #[cfg(target_arch = "x86_64")]
1203 if avx2_a8w8_enabled() {
1204 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
1205 let closures: [_; N] = std::array::from_fn(|i| {
1206 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
1207 move |start: usize, end: usize| {
1208 q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
1209 }
1210 });
1211 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1212 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1213 pool.run_many(&parts);
1214 return;
1215 }
1216 let closures: [_; N] = std::array::from_fn(|i| {
1217 let (c, out) = (&ctxs[i], outs_addr[i]);
1218 move |start: usize, end: usize| {
1219 q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
1220 }
1221 });
1222 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1223 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1224 pool.run_many(&parts);
1225 }
1226}
1227
1228impl QTensor {
1229 #[allow(clippy::needless_range_loop)]
1234 pub fn matvec2_many<const N: usize>(
1235 ts: [&QTensor; N],
1236 x1: &[f32],
1237 x2: &[f32],
1238 mut o1s: [&mut [f32]; N],
1239 mut o2s: [&mut [f32]; N],
1240 pool: Option<&Pool>,
1241 ) {
1242 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
1243 let uniform_q8 = ts.iter().all(|t| {
1244 matches!(
1245 t,
1246 Self::Mapped { dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f, .. }
1247 )
1248 });
1249 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
1250 let uniform_q4 = ts
1251 .iter()
1252 .all(|t| matches!(t, Self::Mapped { dtype: TensorDtype::Q4Block, .. }));
1253 let uniform_vbit = ts
1254 .iter()
1255 .all(|t| matches!(
1256 t,
1257 Self::Mapped { dtype: TensorDtype::Vbit | TensorDtype::VbitRo, .. }
1258 ));
1259 let fusable = pool.is_some()
1260 && total_rows >= 256
1261 && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
1262 if !fusable {
1263 for i in 0..N {
1264 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
1265 }
1266 return;
1267 }
1268 let pool = pool.unwrap();
1269
1270 if uniform_q4 || uniform_vbit {
1271 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1272 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1273 if a8w8_enabled() {
1275 let a1 = split_act(x1);
1276 let a2 = split_act(x2);
1277 let (a1, a2) = (&a1, &a2);
1278 if uniform_q4 {
1279 let closures: [_; N] = std::array::from_fn(|i| {
1280 let (packed, scales) =
1281 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1282 let (gpr, cols, o1, o2) =
1283 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
1284 move |s: usize, e: usize| {
1285 q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
1286 }
1287 });
1288 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1289 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1290 pool.run_many(&parts);
1291 } else {
1292 let closures: [_; N] = std::array::from_fn(|i| {
1293 let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
1294 let (bytes, rows, cols, o1, o2) =
1295 (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), p1[i], p2[i]);
1296 move |s: usize, e: usize| {
1297 vbit_range2_a8w8(
1298 bytes, vbit_offsets, x1, x2, a1, a2, rows, cols, o1, o2, s, e,
1299 )
1300 }
1301 });
1302 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1303 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1304 pool.run_many(&parts);
1305 }
1306 return;
1307 }
1308 if uniform_q4 {
1309 let closures: [_; N] = std::array::from_fn(|i| {
1310 let (packed, scales) =
1311 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1312 let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
1313 move |s: usize, e: usize| {
1314 q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
1315 }
1316 });
1317 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1318 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1319 pool.run_many(&parts);
1320 } else {
1321 let closures: [_; N] = std::array::from_fn(|i| {
1322 let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
1323 let (bytes, rows, cols, o1, o2) =
1324 (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), p1[i], p2[i]);
1325 move |s: usize, e: usize| {
1326 vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
1327 }
1328 });
1329 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1330 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1331 pool.run_many(&parts);
1332 }
1333 return;
1334 }
1335
1336 if uniform_f32 {
1337 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1338 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1339 let closures: [_; N] = std::array::from_fn(|i| {
1340 let Self::F32 { data, cols, .. } = ts[i] else { unreachable!() };
1341 let (o1, o2) = (p1[i], p2[i]);
1342 move |start: usize, end: usize| {
1343 for o in start..end {
1344 let row = &data[o * cols..(o + 1) * cols];
1345 let (mut s1, mut s2) = (0.0f32, 0.0f32);
1346 for j in 0..*cols {
1347 s1 += row[j] * x1[j];
1348 s2 += row[j] * x2[j];
1349 }
1350 unsafe {
1352 *o1.at(o) = s1;
1353 *o2.at(o) = s2;
1354 }
1355 }
1356 }
1357 });
1358 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1359 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1360 pool.run_many(&parts);
1361 return;
1362 }
1363
1364 struct Ctx<'a> {
1365 bytes: &'a [u8],
1366 row_scale: &'a [f32],
1367 cols: usize,
1368 xs1: std::borrow::Cow<'a, [f32]>,
1369 xs2: std::borrow::Cow<'a, [f32]>,
1370 }
1371 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
1372 let Self::Mapped { dtype, cols, row_scale, col_field, .. } = ts[i] else {
1373 unreachable!()
1374 };
1375 Ctx {
1376 bytes: ts[i].quant_bytes(),
1377 row_scale,
1378 cols: *cols,
1379 xs1: prescale(x1, col_field, *dtype),
1380 xs2: prescale(x2, col_field, *dtype),
1381 }
1382 });
1383 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1384 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1385 #[cfg(target_arch = "aarch64")]
1386 if sdot_enabled() {
1387 let acts: [(SplitAct, SplitAct); N] =
1388 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
1389 let closures: [_; N] = std::array::from_fn(|i| {
1390 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
1391 move |start: usize, end: usize| {
1392 q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
1393 }
1394 });
1395 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1396 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1397 pool.run_many(&parts);
1398 return;
1399 }
1400 #[cfg(target_arch = "x86_64")]
1401 if avx2_a8w8_enabled() {
1402 let acts: [(SplitAct, SplitAct); N] =
1403 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
1404 let closures: [_; N] = std::array::from_fn(|i| {
1405 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
1406 move |start: usize, end: usize| {
1407 q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
1408 }
1409 });
1410 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1411 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1412 pool.run_many(&parts);
1413 return;
1414 }
1415 let closures: [_; N] = std::array::from_fn(|i| {
1416 let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
1417 move |start: usize, end: usize| {
1418 q8_range2_f32(c.bytes, c.row_scale, &c.xs1, &c.xs2, c.cols, o1, o2, start, end)
1419 }
1420 });
1421 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1422 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1423 pool.run_many(&parts);
1424 }
1425}
1426
1427#[cfg(target_os = "macos")]
1432mod accel_blas {
1433 #[link(name = "Accelerate", kind = "framework")]
1434 unsafe extern "C" {
1435 pub fn cblas_sgemm(
1436 order: i32,
1437 trans_a: i32,
1438 trans_b: i32,
1439 m: i32,
1440 n: i32,
1441 k: i32,
1442 alpha: f32,
1443 a: *const f32,
1444 lda: i32,
1445 b: *const f32,
1446 ldb: i32,
1447 beta: f32,
1448 c: *mut f32,
1449 ldc: i32,
1450 );
1451 }
1452}
1453
1454#[cfg(target_os = "macos")]
1455pub(crate) fn accel_gemm_enabled() -> bool {
1456 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1457 *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
1458}
1459
1460#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
1463pub(crate) fn accel_gemm_enabled() -> bool {
1464 true
1465}
1466
1467#[cfg(target_arch = "aarch64")]
1474#[allow(clippy::too_many_arguments)]
1475pub(crate) fn neon_gemm_rm(
1476 m: usize,
1477 n: usize,
1478 k: usize,
1479 alpha: f32,
1480 a: &[f32],
1481 lda: usize,
1482 b_mat: &[f32],
1483 ldb: usize,
1484 b_rows_are_n: bool,
1485 c: &mut [f32],
1486 ldc: usize,
1487) {
1488 debug_assert!(a.len() >= (m - 1) * lda + k);
1489 debug_assert!(c.len() >= (m - 1) * ldc + n);
1490 unsafe {
1492 use core::arch::aarch64::*;
1493 let mut i = 0usize;
1494 while i < m {
1495 let mi = (m - i).min(4);
1496 let mut j = 0usize;
1497 while j < n {
1498 let nj = (n - j).min(8);
1499 if mi == 4 && nj == 8 {
1500 let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
1501 let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
1502 let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
1503 let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
1504 for p in 0..k {
1505 let (b0, b1) = if b_rows_are_n {
1506 let base = b_mat.as_ptr().add(j * ldb + p);
1509 let g = |o: usize| *base.add(o * ldb);
1510 (
1511 [g(0), g(1), g(2), g(3)],
1512 [g(4), g(5), g(6), g(7)],
1513 )
1514 } else {
1515 let base = b_mat.as_ptr().add(p * ldb + j);
1516 (
1517 [*base, *base.add(1), *base.add(2), *base.add(3)],
1518 [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
1519 )
1520 };
1521 let bv0 = vld1q_f32(b0.as_ptr());
1522 let bv1 = vld1q_f32(b1.as_ptr());
1523 let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
1524 let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
1525 let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
1526 let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
1527 c0a = vfmaq_f32(c0a, a0, bv0);
1528 c0b = vfmaq_f32(c0b, a0, bv1);
1529 c1a = vfmaq_f32(c1a, a1, bv0);
1530 c1b = vfmaq_f32(c1b, a1, bv1);
1531 c2a = vfmaq_f32(c2a, a2, bv0);
1532 c2b = vfmaq_f32(c2b, a2, bv1);
1533 c3a = vfmaq_f32(c3a, a3, bv0);
1534 c3b = vfmaq_f32(c3b, a3, bv1);
1535 }
1536 let al = vdupq_n_f32(alpha);
1537 for (r, (ca, cb)) in
1538 [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)].iter().enumerate()
1539 {
1540 let dst = c.as_mut_ptr().add((i + r) * ldc + j);
1541 vst1q_f32(dst, vmulq_f32(*ca, al));
1542 vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
1543 }
1544 } else {
1545 for r in 0..mi {
1546 for q in 0..nj {
1547 let mut acc = 0f32;
1548 for p in 0..k {
1549 let bv = if b_rows_are_n {
1550 b_mat[(j + q) * ldb + p]
1551 } else {
1552 b_mat[p * ldb + j + q]
1553 };
1554 acc += a[(i + r) * lda + p] * bv;
1555 }
1556 c[(i + r) * ldc + j + q] = acc * alpha;
1557 }
1558 }
1559 }
1560 j += nj;
1561 }
1562 i += mi;
1563 }
1564 }
1565}
1566
1567#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
1569#[allow(clippy::too_many_arguments)]
1570pub(crate) fn sgemm_rm(
1571 m: usize,
1572 n: usize,
1573 k: usize,
1574 alpha: f32,
1575 a: &[f32],
1576 lda: usize,
1577 b_mat: &[f32],
1578 ldb: usize,
1579 b_rows_are_n: bool,
1580 c: &mut [f32],
1581 ldc: usize,
1582) {
1583 neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
1584}
1585
1586#[cfg(target_os = "macos")]
1589#[allow(clippy::too_many_arguments)]
1590pub(crate) fn sgemm_rm(
1591 m: usize,
1592 n: usize,
1593 k: usize,
1594 alpha: f32,
1595 a: &[f32],
1596 lda: usize,
1597 b_mat: &[f32],
1598 ldb: usize,
1599 b_rows_are_n: bool,
1600 c: &mut [f32],
1601 ldc: usize,
1602) {
1603 debug_assert!(a.len() >= (m - 1) * lda + k);
1604 debug_assert!(c.len() >= (m - 1) * ldc + n);
1605 #[cfg(target_arch = "aarch64")]
1610 if std::env::var("CMF_FORCE_NEON_GEMM").map(|v| v == "1").unwrap_or(false) {
1611 return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
1612 }
1613 unsafe {
1614 accel_blas::cblas_sgemm(
1615 101, 111, if b_rows_are_n { 112 } else { 111 },
1618 m as i32,
1619 n as i32,
1620 k as i32,
1621 alpha,
1622 a.as_ptr(),
1623 lda as i32,
1624 b_mat.as_ptr(),
1625 ldb as i32,
1626 0.0,
1627 c.as_mut_ptr(),
1628 ldc as i32,
1629 );
1630 }
1631}
1632
1633#[cfg(target_os = "macos")]
1640fn qmatmat_accel(
1641 q: &[u8],
1642 row_scale: &[f32],
1643 pre: &[std::borrow::Cow<'_, [f32]>],
1644 rows: usize,
1645 cols: usize,
1646 out: &mut [f32],
1647 pool: Option<&Pool>,
1648) {
1649 const TR: usize = 2048;
1654 let b = pre.len();
1655 thread_local! {
1656 static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
1657 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
1658 }
1659 XPANEL.with(|xp| {
1660 WTILE.with(|wt| {
1661 let mut xpanel = xp.borrow_mut();
1662 xpanel.clear();
1663 for x in pre {
1664 xpanel.extend_from_slice(x);
1665 }
1666 let mut wtile = wt.borrow_mut();
1667 wtile.resize(TR * cols, 0.0);
1668 let mut r0 = 0usize;
1669 while r0 < rows {
1670 let tr = TR.min(rows - r0);
1671 let wt_addr = SendMut(wtile.as_mut_ptr());
1673 let run = |start: usize, end: usize| {
1674 for r in start..end {
1675 let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
1676 let s = row_scale[r0 + r];
1677 let dst = unsafe {
1679 std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols)
1680 };
1681 for (d, &v) in dst.iter_mut().zip(row) {
1682 *d = (v as i8) as f32 * s;
1683 }
1684 }
1685 };
1686 dispatch_rows(pool, tr, &run);
1687 unsafe {
1689 accel_blas::cblas_sgemm(
1690 101, 111, 112, b as i32,
1694 tr as i32,
1695 cols as i32,
1696 1.0,
1697 xpanel.as_ptr(),
1698 cols as i32,
1699 wtile.as_ptr(),
1700 cols as i32,
1701 0.0,
1702 out.as_mut_ptr().add(r0),
1703 rows as i32,
1704 );
1705 }
1706 r0 += tr;
1707 }
1708 })
1709 });
1710}
1711
1712fn qmatmat(
1713 q: &[u8],
1714 row_scale: &[f32],
1715 pre: &[std::borrow::Cow<'_, [f32]>],
1716 rows: usize,
1717 cols: usize,
1718 out: &mut [f32],
1719 pool: Option<&Pool>,
1720) {
1721 let b = pre.len();
1722 debug_assert_eq!(out.len(), b * rows);
1723 #[cfg(target_os = "macos")]
1728 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
1729 qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
1730 return;
1731 }
1732 #[cfg(target_arch = "aarch64")]
1733 if sdot_enabled() {
1734 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
1735 let out_addr = SendMut(out.as_mut_ptr());
1736 let blocked_ok =
1739 std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true);
1740 let use_i8mm = i8mm_enabled();
1741 if blocked_ok {
1742 let run = |start: usize, end: usize| {
1743 let mut o = start;
1744 while o < end {
1745 if o + 2 <= end {
1746 let r0 = &q[o * cols..(o + 1) * cols];
1747 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
1748 let mut bi = 0usize;
1749 while bi + 4 <= acts.len() {
1750 let xs = [
1751 acts[bi].xq.as_slice(),
1752 acts[bi + 1].xq.as_slice(),
1753 acts[bi + 2].xq.as_slice(),
1754 acts[bi + 3].xq.as_slice(),
1755 ];
1756 let d = if use_i8mm {
1757 unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
1758 } else {
1759 unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
1760 };
1761 for (r, row) in [r0, r1].into_iter().enumerate() {
1762 for k in 0..4 {
1763 let act = &acts[bi + k];
1764 let mut v = d[r][k] as f32 * act.sx;
1765 for &(j, xv) in &act.outliers {
1766 v += (row[j] as i8) as f32 * xv;
1767 }
1768 unsafe {
1769 *out_addr.at((bi + k) * rows + o + r) =
1770 v * row_scale[o + r]
1771 };
1772 }
1773 }
1774 bi += 4;
1775 }
1776 while bi < acts.len() {
1777 for (r, row) in [r0, r1].into_iter().enumerate() {
1778 let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
1779 unsafe { *out_addr.at(bi * rows + o + r) = v };
1780 }
1781 bi += 1;
1782 }
1783 o += 2;
1784 } else {
1785 let row = &q[o * cols..(o + 1) * cols];
1786 for (bi, act) in acts.iter().enumerate() {
1787 let v = row_dot_sdot(row, act) * row_scale[o];
1788 unsafe { *out_addr.at(bi * rows + o) = v };
1789 }
1790 o += 1;
1791 }
1792 }
1793 };
1794 dispatch_rows(pool, rows, &run);
1795 return;
1796 }
1797 let run = |start: usize, end: usize| {
1798 for o in start..end {
1799 let row = &q[o * cols..(o + 1) * cols];
1800 for (bi, act) in acts.iter().enumerate() {
1801 let v = row_dot_sdot(row, act) * row_scale[o];
1802 unsafe { *out_addr.at(bi * rows + o) = v };
1803 }
1804 }
1805 };
1806 dispatch_rows(pool, rows, &run);
1807 return;
1808 }
1809 #[cfg(target_arch = "x86_64")]
1814 if avx2_a8w8_enabled() {
1815 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
1816 let out_addr = SendMut(out.as_mut_ptr());
1817 let blocked_ok = std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true);
1820 if !avx512vnni_enabled() && blocked_ok {
1821 let run = |start: usize, end: usize| {
1822 let mut o = start;
1823 while o < end {
1824 if o + 2 <= end {
1825 let r0 = &q[o * cols..(o + 1) * cols];
1826 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
1827 let mut bi = 0usize;
1828 while bi + 4 <= acts.len() {
1829 let xs = [
1830 acts[bi].xq.as_slice(),
1831 acts[bi + 1].xq.as_slice(),
1832 acts[bi + 2].xq.as_slice(),
1833 acts[bi + 3].xq.as_slice(),
1834 ];
1835 let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
1836 for (r, row) in [r0, r1].into_iter().enumerate() {
1837 for k in 0..4 {
1838 let act = &acts[bi + k];
1839 let mut v = d[r][k] as f32 * act.sx;
1840 for &(j, xv) in &act.outliers {
1841 v += (row[j] as i8) as f32 * xv;
1842 }
1843 unsafe {
1844 *out_addr.at((bi + k) * rows + o + r) =
1845 v * row_scale[o + r]
1846 };
1847 }
1848 }
1849 bi += 4;
1850 }
1851 while bi < acts.len() {
1852 for (r, row) in [r0, r1].into_iter().enumerate() {
1853 let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
1854 unsafe { *out_addr.at(bi * rows + o + r) = v };
1855 }
1856 bi += 1;
1857 }
1858 o += 2;
1859 } else {
1860 let row = &q[o * cols..(o + 1) * cols];
1861 for (bi, act) in acts.iter().enumerate() {
1862 let v = row_dot_avx2(row, act) * row_scale[o];
1863 unsafe { *out_addr.at(bi * rows + o) = v };
1864 }
1865 o += 1;
1866 }
1867 }
1868 };
1869 dispatch_rows(pool, rows, &run);
1870 return;
1871 }
1872 let run = |start: usize, end: usize| {
1873 for o in start..end {
1874 let row = &q[o * cols..(o + 1) * cols];
1875 for (bi, act) in acts.iter().enumerate() {
1876 let v = row_dot_avx2(row, act) * row_scale[o];
1877 unsafe { *out_addr.at(bi * rows + o) = v };
1878 }
1879 }
1880 };
1881 dispatch_rows(pool, rows, &run);
1882 return;
1883 }
1884 let out_addr = SendMut(out.as_mut_ptr());
1885 let run = |start: usize, end: usize| {
1886 for o in start..end {
1887 let row = &q[o * cols..(o + 1) * cols];
1888 for (bi, x) in pre.iter().enumerate() {
1889 let mut acc = 0f32;
1890 for j in 0..cols {
1891 acc += (row[j] as i8) as f32 * x[j];
1892 }
1893 unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
1894 }
1895 }
1896 };
1897 dispatch_rows(pool, rows, &run);
1898}
1899
1900fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
1903 match pool {
1904 Some(pool) if rows >= 256 => pool.run_rows(rows, run),
1905 _ => run(0, rows),
1906 }
1907}
1908
1909fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
1911 let groups = rows * cols / GROUP_SIZE;
1912 bytes.split_at(groups * 16)
1913}
1914
1915#[inline]
1920fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
1921 #[cfg(target_arch = "aarch64")]
1922 unsafe {
1923 return vbit_fill4_neon(data, buf);
1924 }
1925 #[cfg(target_arch = "x86_64")]
1926 if avx2_enabled() {
1927 return unsafe { vbit_fill4_avx2(data, buf) };
1928 }
1929 #[allow(unreachable_code)]
1930 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
1931 let u = unpack8::<4>(&data[blk * 4..]);
1932 for k in 0..8 {
1933 chunk[k] = (u[k] - 7) as i8 as u8;
1934 }
1935 }
1936}
1937
1938#[cfg(target_arch = "aarch64")]
1939#[target_feature(enable = "neon")]
1940unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
1941 unsafe {
1944 use core::arch::aarch64::*;
1945 let n = buf.len();
1946 let mask = vdupq_n_u8(0x0F);
1947 let seven = vdupq_n_s8(7);
1948 let mut g = 0usize;
1949 while g * 32 + 32 <= n {
1950 let b = vld1q_u8(data.as_ptr().add(g * 16));
1951 let hi = vshrq_n_u8::<4>(b);
1952 let lo = vandq_u8(b, mask);
1953 let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
1954 let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
1955 vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
1956 vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
1957 g += 1;
1958 }
1959 }
1960}
1961
1962#[cfg(target_arch = "x86_64")]
1963#[target_feature(enable = "avx2")]
1964unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
1965 unsafe {
1967 use core::arch::x86_64::*;
1968 let n = buf.len();
1969 let mask = _mm_set1_epi8(0x0F);
1970 let seven = _mm256_set1_epi8(7);
1971 let mut g = 0usize;
1972 while g * 32 + 32 <= n {
1973 let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
1974 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
1975 let lo = _mm_and_si128(b, mask);
1976 let z = _mm256_sub_epi8(
1977 _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
1978 seven,
1979 );
1980 _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
1981 g += 1;
1982 }
1983 }
1984}
1985
1986#[inline(always)]
1991fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
1992 let mut acc = 0u64;
1993 for i in 0..B {
1994 acc = (acc << 8) | data[i] as u64;
1995 }
1996 let mask = (1u64 << B) - 1;
1997 let mut out = [0i32; 8];
1998 for (k, o) in out.iter_mut().enumerate() {
1999 *o = ((acc >> ((7 - k) * B)) & mask) as i32;
2000 }
2001 out
2002}
2003
2004#[allow(clippy::too_many_arguments)]
2010fn vbitmatvec(
2011 bytes: &[u8],
2012 offsets: &[usize],
2013 x: &[f32],
2014 rows: usize,
2015 cols: usize,
2016 out: &mut [f32],
2017 pool: Option<&Pool>,
2018) {
2019 debug_assert_eq!(out.len(), rows);
2020 debug_assert_eq!(offsets.len(), rows + 1);
2021
2022 if a8w8_enabled() {
2026 let act = split_act(x);
2027 let out_addr = SendMut(out.as_mut_ptr());
2028 let run = move |start: usize, end: usize| {
2029 vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
2030 };
2031 dispatch_rows(pool, rows, &run);
2032 return;
2033 }
2034
2035 let out_addr = SendMut(out.as_mut_ptr());
2036 let run = move |start: usize, end: usize| {
2037 vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
2038 };
2039 dispatch_rows(pool, rows, &run);
2040}
2041
2042#[allow(clippy::too_many_arguments)]
2046fn vbit_range_a8w8(
2047 bytes: &[u8],
2048 offsets: &[usize],
2049 x: &[f32],
2050 act: &SplitAct,
2051 rows: usize,
2052 cols: usize,
2053 out: SendMut,
2054 start: usize,
2055 end: usize,
2056) {
2057 let ng = cols / GROUP_SIZE;
2058 let bits = &bytes[..rows];
2059 let sc_off = rows;
2060 let row_dot = |r: usize| -> f32 {
2061 let b = bits[r] as usize;
2062 let l = ((1i32 << (b - 1)) - 1) as i32;
2063 let mask = (1u64 << b) - 1;
2064 let data = &bytes[offsets[r]..offsets[r + 1]];
2065 if b == 8 {
2066 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
2068 let mut dot = 0f32;
2069 for g in 0..ng {
2070 let so = (r * ng + g) * 2;
2071 let sgf = f16_to_f32(u16::from_le_bytes([
2072 bytes[sc_off + so],
2073 bytes[sc_off + so + 1],
2074 ]));
2075 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2076 let mut gd = 0f32;
2077 for &xv in xg.iter() {
2078 if nbits < 8 {
2079 acc = (acc << 8) | data[idx] as u64;
2080 idx += 1;
2081 nbits += 8;
2082 }
2083 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
2084 nbits -= 8;
2085 gd += (u - l) as f32 * xv;
2086 }
2087 dot += gd * sgf;
2088 }
2089 return dot;
2090 }
2091 thread_local! {
2095 static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
2096 const { std::cell::RefCell::new(Vec::new()) };
2097 }
2098 #[inline(always)]
2099 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
2100 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
2101 let u = unpack8::<B>(&data[blk * B..]);
2102 for k in 0..8 {
2103 chunk[k] = (u[k] - l) as i8 as u8;
2104 }
2105 }
2106 }
2107 let _ = mask;
2108 VBIT_SCRATCH.with(|scratch| {
2109 let mut buf = scratch.borrow_mut();
2110 buf.resize(cols, 0);
2111 match b {
2112 3 => fill::<3>(data, l, &mut buf),
2113 4 => vbit_fill4(data, &mut buf),
2114 5 => fill::<5>(data, l, &mut buf),
2115 6 => fill::<6>(data, l, &mut buf),
2116 _ => unreachable!(),
2117 }
2118 let mut dot = 0f32;
2119 for g in 0..ng {
2120 let so = (r * ng + g) * 2;
2121 let s = f16_to_f32(u16::from_le_bytes([
2122 bytes[sc_off + so],
2123 bytes[sc_off + so + 1],
2124 ]));
2125 let d = dot_i8_i8(
2126 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
2127 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
2128 ) as f32
2129 * act.sx;
2130 dot += d * s;
2131 }
2132 for &(j, xv) in &act.outliers {
2133 let so = (r * ng + j / GROUP_SIZE) * 2;
2134 let s = f16_to_f32(u16::from_le_bytes([
2135 bytes[sc_off + so],
2136 bytes[sc_off + so + 1],
2137 ]));
2138 dot += (buf[j] as i8) as f32 * s * xv;
2140 }
2141 dot
2142 })
2143 };
2144 for r in start..end {
2145 unsafe { *out.at(r) = row_dot(r) };
2147 }
2148}
2149
2150#[allow(clippy::too_many_arguments)]
2152fn vbit_range_f32(
2153 bytes: &[u8],
2154 offsets: &[usize],
2155 x: &[f32],
2156 rows: usize,
2157 cols: usize,
2158 out: SendMut,
2159 start: usize,
2160 end: usize,
2161) {
2162 let ng = cols / GROUP_SIZE;
2163 let bits = &bytes[..rows];
2164 let sc_off = rows;
2165 #[inline(always)]
2169 fn dot_row<const B: usize>(
2170 data: &[u8],
2171 bytes: &[u8],
2172 sc_off: usize,
2173 r: usize,
2174 ng: usize,
2175 x: &[f32],
2176 ) -> f32 {
2177 let l = ((1i32 << (B - 1)) - 1) as f32;
2178 let gbytes = GROUP_SIZE * B / 8;
2179 let mut dot = 0f32;
2180 for g in 0..ng {
2181 let so = (r * ng + g) * 2;
2182 let s = f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]));
2183 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2184 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
2185 let mut gd = 0f32;
2186 for blk in 0..GROUP_SIZE / 8 {
2187 let u = unpack8::<B>(&gd0[blk * B..]);
2188 let xb = &xg[blk * 8..blk * 8 + 8];
2189 for k in 0..8 {
2190 gd += (u[k] as f32 - l) * xb[k];
2191 }
2192 }
2193 dot += gd * s;
2194 }
2195 dot
2196 }
2197 for r in start..end {
2198 let data = &bytes[offsets[r]..offsets[r + 1]];
2199 let v = match bits[r] {
2200 3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
2201 4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
2202 5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
2203 6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
2204 8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
2205 b => unreachable!("vbit bit-width {b} (validated at load)"),
2206 };
2207 unsafe { *out.at(r) = v };
2209 }
2210}
2211
2212#[allow(clippy::too_many_arguments)]
2217fn vbitmatvec2(
2218 bytes: &[u8],
2219 offsets: &[usize],
2220 x1: &[f32],
2221 x2: &[f32],
2222 rows: usize,
2223 cols: usize,
2224 o1: &mut [f32],
2225 o2: &mut [f32],
2226 pool: Option<&Pool>,
2227) {
2228 debug_assert_eq!(o1.len(), rows);
2229 debug_assert_eq!(o2.len(), rows);
2230
2231 if a8w8_enabled() {
2232 let a1 = split_act(x1);
2233 let a2 = split_act(x2);
2234 let p1 = SendMut(o1.as_mut_ptr());
2235 let p2 = SendMut(o2.as_mut_ptr());
2236 let run = move |start: usize, end: usize| {
2237 vbit_range2_a8w8(bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end)
2238 };
2239 dispatch_rows(pool, rows, &run);
2240 return;
2241 }
2242
2243 let p1 = SendMut(o1.as_mut_ptr());
2244 let p2 = SendMut(o2.as_mut_ptr());
2245 let run = move |start: usize, end: usize| {
2246 vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
2247 };
2248 dispatch_rows(pool, rows, &run);
2249}
2250
2251#[allow(clippy::too_many_arguments)]
2255fn vbit_range2_a8w8(
2256 bytes: &[u8],
2257 offsets: &[usize],
2258 x1: &[f32],
2259 x2: &[f32],
2260 a1: &SplitAct,
2261 a2: &SplitAct,
2262 rows: usize,
2263 cols: usize,
2264 p1: SendMut,
2265 p2: SendMut,
2266 start: usize,
2267 end: usize,
2268) {
2269 let ng = cols / GROUP_SIZE;
2270 let bits = &bytes[..rows];
2271 let sc_off = rows;
2272 let row_dots = |r: usize| -> (f32, f32) {
2273 let b = bits[r] as usize;
2274 let l = (1i32 << (b - 1)) - 1;
2275 let data = &bytes[offsets[r]..offsets[r + 1]];
2276 if b == 8 {
2277 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
2280 let (mut d1, mut d2) = (0f32, 0f32);
2281 for g in 0..ng {
2282 let so = (r * ng + g) * 2;
2283 let sgf = f16_to_f32(u16::from_le_bytes([
2284 bytes[sc_off + so],
2285 bytes[sc_off + so + 1],
2286 ]));
2287 let (mut g1, mut g2) = (0f32, 0f32);
2288 for k in 0..GROUP_SIZE {
2289 if nbits < 8 {
2290 acc = (acc << 8) | data[idx] as u64;
2291 idx += 1;
2292 nbits += 8;
2293 }
2294 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
2295 nbits -= 8;
2296 let w = (u - l) as f32;
2297 g1 += w * x1[g * GROUP_SIZE + k];
2298 g2 += w * x2[g * GROUP_SIZE + k];
2299 }
2300 d1 += g1 * sgf;
2301 d2 += g2 * sgf;
2302 }
2303 return (d1, d2);
2304 }
2305 thread_local! {
2306 static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
2307 const { std::cell::RefCell::new(Vec::new()) };
2308 }
2309 #[inline(always)]
2310 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
2311 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
2312 let u = unpack8::<B>(&data[blk * B..]);
2313 for k in 0..8 {
2314 chunk[k] = (u[k] - l) as i8 as u8;
2315 }
2316 }
2317 }
2318 VBIT_SCRATCH2.with(|scratch| {
2319 let mut buf = scratch.borrow_mut();
2320 buf.resize(cols, 0);
2321 match b {
2322 3 => fill::<3>(data, l, &mut buf),
2323 4 => vbit_fill4(data, &mut buf),
2324 5 => fill::<5>(data, l, &mut buf),
2325 6 => fill::<6>(data, l, &mut buf),
2326 _ => unreachable!(),
2327 }
2328 let (mut d1, mut d2) = (0f32, 0f32);
2329 for g in 0..ng {
2330 let so = (r * ng + g) * 2;
2331 let s = f16_to_f32(u16::from_le_bytes([
2332 bytes[sc_off + so],
2333 bytes[sc_off + so + 1],
2334 ]));
2335 let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2336 let v1 =
2337 dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
2338 let v2 =
2339 dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
2340 d1 += v1 * s;
2341 d2 += v2 * s;
2342 }
2343 for &(j, xv) in &a1.outliers {
2344 let so = (r * ng + j / GROUP_SIZE) * 2;
2345 let s = f16_to_f32(u16::from_le_bytes([
2346 bytes[sc_off + so],
2347 bytes[sc_off + so + 1],
2348 ]));
2349 d1 += (buf[j] as i8) as f32 * s * xv;
2350 }
2351 for &(j, xv) in &a2.outliers {
2352 let so = (r * ng + j / GROUP_SIZE) * 2;
2353 let s = f16_to_f32(u16::from_le_bytes([
2354 bytes[sc_off + so],
2355 bytes[sc_off + so + 1],
2356 ]));
2357 d2 += (buf[j] as i8) as f32 * s * xv;
2358 }
2359 (d1, d2)
2360 })
2361 };
2362 for r in start..end {
2363 let (v1, v2) = row_dots(r);
2364 unsafe {
2366 *p1.at(r) = v1;
2367 *p2.at(r) = v2;
2368 }
2369 }
2370}
2371
2372#[allow(clippy::too_many_arguments)]
2376fn vbit_range2_f32(
2377 bytes: &[u8],
2378 offsets: &[usize],
2379 x1: &[f32],
2380 x2: &[f32],
2381 rows: usize,
2382 cols: usize,
2383 p1: SendMut,
2384 p2: SendMut,
2385 start: usize,
2386 end: usize,
2387) {
2388 let ng = cols / GROUP_SIZE;
2389 let bits = &bytes[..rows];
2390 let sc_off = rows;
2391 #[inline(always)]
2392 #[allow(clippy::too_many_arguments)]
2393 fn dot_row2<const B: usize>(
2394 data: &[u8],
2395 bytes: &[u8],
2396 sc_off: usize,
2397 r: usize,
2398 ng: usize,
2399 x1: &[f32],
2400 x2: &[f32],
2401 ) -> (f32, f32) {
2402 let l = ((1i32 << (B - 1)) - 1) as f32;
2403 let gbytes = GROUP_SIZE * B / 8;
2404 let (mut d1, mut d2) = (0f32, 0f32);
2405 for g in 0..ng {
2406 let so = (r * ng + g) * 2;
2407 let s = f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]));
2408 let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2409 let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2410 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
2411 let (mut g1, mut g2) = (0f32, 0f32);
2412 for blk in 0..GROUP_SIZE / 8 {
2413 let u = unpack8::<B>(&gd0[blk * B..]);
2414 for k in 0..8 {
2415 let w = u[k] as f32 - l;
2416 g1 += w * x1g[blk * 8 + k];
2417 g2 += w * x2g[blk * 8 + k];
2418 }
2419 }
2420 d1 += g1 * s;
2421 d2 += g2 * s;
2422 }
2423 (d1, d2)
2424 }
2425 for r in start..end {
2426 let data = &bytes[offsets[r]..offsets[r + 1]];
2427 let (v1, v2) = match bits[r] {
2428 3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
2429 4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
2430 5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
2431 6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
2432 8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
2433 b => unreachable!("vbit bit-width {b} (validated at load)"),
2434 };
2435 unsafe {
2437 *p1.at(r) = v1;
2438 *p2.at(r) = v2;
2439 }
2440 }
2441}
2442
2443#[inline]
2450#[allow(unreachable_code)]
2451fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
2452 #[cfg(target_arch = "aarch64")]
2453 unsafe {
2454 return dot_q4t_row_sdot(bytes, r, gpr, xq);
2455 }
2456 #[cfg(target_arch = "x86_64")]
2457 unsafe {
2458 return dot_q4t_row_avx2(bytes, r, gpr, xq);
2459 }
2460 let mut acc = 0f32;
2461 for gi in 0..gpr {
2462 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
2463 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2464 let mut d = 0i32;
2465 for (k, &b) in tile[2..].iter().enumerate() {
2466 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
2467 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
2468 }
2469 acc += d as f32 * s;
2470 }
2471 acc
2472}
2473
2474#[cfg(target_arch = "aarch64")]
2475#[target_feature(enable = "neon,dotprod")]
2476unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
2477 unsafe {
2480 use core::arch::aarch64::*;
2481 use core::arch::asm;
2482 let lomask = vdupq_n_u8(0x0F);
2483 let eight = vdupq_n_s8(8);
2484 let mut acc = 0f32;
2485 for gi in 0..gpr {
2486 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
2487 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2488 let b = vld1q_u8(t.add(2));
2489 let lo = vandq_u8(b, lomask);
2490 let hi = vshrq_n_u8::<4>(b);
2491 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
2492 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
2493 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
2494 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
2495 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
2496 asm!(
2497 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
2498 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
2499 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
2500 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
2501 options(pure, nomem, nostack),
2502 );
2503 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
2504 }
2505 acc
2506 }
2507}
2508
2509#[cfg(target_arch = "x86_64")]
2510#[target_feature(enable = "avx2")]
2511unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
2512 unsafe {
2514 use core::arch::x86_64::*;
2515 let lomask = _mm_set1_epi8(0x0F);
2516 let eight = _mm256_set1_epi8(8);
2517 let ones = _mm256_set1_epi16(1);
2518 let mut acc = 0f32;
2519 for gi in 0..gpr {
2520 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
2521 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2522 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
2523 let lo = _mm_and_si128(b, lomask);
2524 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
2525 let w = _mm256_sub_epi8(
2526 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
2527 eight,
2528 );
2529 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2530 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
2531 let d = _mm256_madd_epi16(p16, ones);
2532 let hi128 = _mm256_extracti128_si256::<1>(d);
2533 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
2534 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2535 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2536 acc += _mm_cvtsi128_si32(s32) as f32 * s;
2537 }
2538 acc
2539 }
2540}
2541
2542#[cfg(target_arch = "x86_64")]
2547#[target_feature(enable = "avx2")]
2548unsafe fn dot_q4t_row_1x4_avx2(
2549 bytes: &[u8],
2550 r: usize,
2551 gpr: usize,
2552 xs: [&[i8]; 4],
2553) -> [f32; 4] {
2554 unsafe {
2556 use core::arch::x86_64::*;
2557 let lomask = _mm_set1_epi8(0x0F);
2558 let eight = _mm256_set1_epi8(8);
2559 let ones = _mm256_set1_epi16(1);
2560 let mut acc = [0f32; 4];
2561 for gi in 0..gpr {
2562 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
2563 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2564 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
2565 let lo = _mm_and_si128(bb, lomask);
2566 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
2567 let w = _mm256_sub_epi8(
2568 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
2569 eight,
2570 );
2571 let aw = _mm256_abs_epi8(w);
2572 for (k, xq) in xs.iter().enumerate() {
2573 let x =
2574 _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2575 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
2576 let d = _mm256_madd_epi16(p16, ones);
2577 let hi128 = _mm256_extracti128_si256::<1>(d);
2578 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
2579 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2580 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2581 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
2582 }
2583 }
2584 acc
2585 }
2586}
2587
2588#[inline]
2590fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
2591 let gi = j / GROUP_SIZE;
2592 let k = j % GROUP_SIZE;
2593 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
2594 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2595 let byte = tile[2 + k / 2];
2596 let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
2597 ((nib as i32 - 8) as f32, s)
2598}
2599
2600#[inline]
2603fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
2604 let mut acc = 0f32;
2605 for gi in 0..gpr {
2606 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
2607 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2608 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
2609 let mut ga = 0f32;
2610 for (k, &b) in tile[2..].iter().enumerate() {
2611 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
2612 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
2613 }
2614 acc += ga * s;
2615 }
2616 acc
2617}
2618
2619fn q4t_matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
2621 debug_assert_eq!(out.len(), rows);
2622 let gpr = cols / GROUP_SIZE;
2623 let out_addr = SendMut(out.as_mut_ptr());
2624 if a8w8_enabled() {
2625 let act = split_act(x);
2626 let run = move |start: usize, end: usize| {
2627 for r in start..end {
2628 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
2629 for &(j, xv) in &act.outliers {
2630 let (w, s) = q4t_outlier(bytes, r, gpr, j);
2631 acc += w * s * xv;
2632 }
2633 unsafe { *out_addr.at(r) = acc };
2635 }
2636 };
2637 dispatch_rows(pool, rows, &run);
2638 return;
2639 }
2640 let run = move |start: usize, end: usize| {
2641 for r in start..end {
2642 unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
2644 }
2645 };
2646 dispatch_rows(pool, rows, &run);
2647}
2648
2649#[allow(clippy::too_many_arguments)]
2651fn q4t_matvec2(
2652 bytes: &[u8],
2653 x1: &[f32],
2654 x2: &[f32],
2655 rows: usize,
2656 cols: usize,
2657 o1: &mut [f32],
2658 o2: &mut [f32],
2659 pool: Option<&Pool>,
2660) {
2661 let gpr = cols / GROUP_SIZE;
2662 let p1 = SendMut(o1.as_mut_ptr());
2663 let p2 = SendMut(o2.as_mut_ptr());
2664 if a8w8_enabled() {
2665 let a1 = split_act(x1);
2666 let a2 = split_act(x2);
2667 let run = move |start: usize, end: usize| {
2668 for r in start..end {
2669 let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
2670 let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
2671 for &(j, xv) in &a1.outliers {
2672 let (w, s) = q4t_outlier(bytes, r, gpr, j);
2673 v1 += w * s * xv;
2674 }
2675 for &(j, xv) in &a2.outliers {
2676 let (w, s) = q4t_outlier(bytes, r, gpr, j);
2677 v2 += w * s * xv;
2678 }
2679 unsafe {
2681 *p1.at(r) = v1;
2682 *p2.at(r) = v2;
2683 }
2684 }
2685 };
2686 dispatch_rows(pool, rows, &run);
2687 return;
2688 }
2689 let run = move |start: usize, end: usize| {
2690 for r in start..end {
2691 unsafe {
2693 *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
2694 *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
2695 }
2696 }
2697 };
2698 dispatch_rows(pool, rows, &run);
2699}
2700
2701#[allow(clippy::too_many_arguments)]
2703fn q4t_matmat(
2704 bytes: &[u8],
2705 xs_all: &[f32],
2706 b: usize,
2707 rows: usize,
2708 cols: usize,
2709 out: &mut [f32],
2710 pool: Option<&Pool>,
2711) {
2712 debug_assert_eq!(out.len(), b * rows);
2713 let gpr = cols / GROUP_SIZE;
2714 let out_addr = SendMut(out.as_mut_ptr());
2715 if a8w8_enabled() {
2716 let acts: Vec<SplitAct> =
2717 (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
2718 let acts = &acts;
2719 #[cfg(target_arch = "x86_64")]
2720 let blocked_ok = avx2_enabled()
2721 && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true);
2722 #[cfg(not(target_arch = "x86_64"))]
2723 let blocked_ok = false;
2724 let run = move |start: usize, end: usize| {
2725 for r in start..end {
2726 let mut bi = 0usize;
2727 #[cfg(target_arch = "x86_64")]
2728 if blocked_ok {
2729 while bi + 4 <= acts.len() {
2730 let xs = [
2731 acts[bi].xq.as_slice(),
2732 acts[bi + 1].xq.as_slice(),
2733 acts[bi + 2].xq.as_slice(),
2734 acts[bi + 3].xq.as_slice(),
2735 ];
2736 let d = unsafe { dot_q4t_row_1x4_avx2(bytes, r, gpr, xs) };
2737 for k in 0..4 {
2738 let act = &acts[bi + k];
2739 let mut acc = d[k] * act.sx;
2740 for &(j, xv) in &act.outliers {
2741 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
2742 acc += w * sc * xv;
2743 }
2744 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
2746 }
2747 bi += 4;
2748 }
2749 }
2750 let _ = blocked_ok;
2751 while bi < acts.len() {
2752 let act = &acts[bi];
2753 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
2754 for &(j, xv) in &act.outliers {
2755 let (w, s) = q4t_outlier(bytes, r, gpr, j);
2756 acc += w * s * xv;
2757 }
2758 unsafe { *out_addr.at(bi * rows + r) = acc };
2760 bi += 1;
2761 }
2762 }
2763 };
2764 dispatch_rows(pool, rows, &run);
2765 return;
2766 }
2767 let run = move |start: usize, end: usize| {
2768 for r in start..end {
2769 for bi in 0..b {
2770 let x = &xs_all[bi * cols..(bi + 1) * cols];
2771 unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
2773 }
2774 }
2775 };
2776 dispatch_rows(pool, rows, &run);
2777}
2778
2779fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
2788 (0..gpr)
2789 .map(|gi| {
2790 xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
2791 .iter()
2792 .map(|&v| v as i32)
2793 .sum()
2794 })
2795 .collect()
2796}
2797
2798#[inline]
2802#[allow(unreachable_code)]
2803#[cfg(target_arch = "x86_64")]
2808#[target_feature(enable = "avx2")]
2809unsafe fn dot_q1_row_avx2(
2810 bytes: &[u8],
2811 r: usize,
2812 gpr: usize,
2813 xq: &[i8],
2814 gsum: &[i32],
2815) -> f32 {
2816 unsafe {
2818 use core::arch::x86_64::*;
2819 let expand = _mm256_setr_epi8(
2821 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1,
2822 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3,
2823 );
2824 let bitsel = _mm256_setr_epi8(
2825 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128,
2826 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128,
2827 );
2828 let ones8 = _mm256_set1_epi8(1);
2829 let ones16 = _mm256_set1_epi16(1);
2830 let mut acc = 0f32;
2831 for gi in 0..gpr {
2832 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
2833 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2834 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
2835 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
2836 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
2837 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2838 let sel = _mm256_and_si256(x, mask);
2839 let p16 = _mm256_maddubs_epi16(ones8, sel);
2841 let d32 = _mm256_madd_epi16(p16, ones16);
2842 let hi128 = _mm256_extracti128_si256::<1>(d32);
2843 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
2844 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2845 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2846 let msum = _mm_cvtsi128_si32(s32);
2847 let d = 2 * msum - gsum[gi];
2850 acc += d as f32 * s;
2851 }
2852 acc
2853 }
2854}
2855
2856#[cfg(target_arch = "x86_64")]
2859#[target_feature(enable = "avx2")]
2860unsafe fn dot_q1_row_1x4_avx2(
2861 bytes: &[u8],
2862 r: usize,
2863 gpr: usize,
2864 xs: [&[i8]; 4],
2865 gsums: [&[i32]; 4],
2866) -> [f32; 4] {
2867 unsafe {
2869 use core::arch::x86_64::*;
2870 let expand = _mm256_setr_epi8(
2871 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1,
2872 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3,
2873 );
2874 let bitsel = _mm256_setr_epi8(
2875 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128,
2876 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128,
2877 );
2878 let ones8 = _mm256_set1_epi8(1);
2879 let ones16 = _mm256_set1_epi16(1);
2880 let mut acc = [0f32; 4];
2881 for gi in 0..gpr {
2882 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
2883 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2884 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
2885 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
2886 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
2887 for (k, xq) in xs.iter().enumerate() {
2888 let x =
2889 _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2890 let sel = _mm256_and_si256(x, mask);
2891 let p16 = _mm256_maddubs_epi16(ones8, sel);
2892 let d32 = _mm256_madd_epi16(p16, ones16);
2893 let hi128 = _mm256_extracti128_si256::<1>(d32);
2894 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
2895 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2896 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2897 let msum = _mm_cvtsi128_si32(s32);
2898 let d = 2 * msum - gsums[k][gi];
2899 acc[k] += d as f32 * s;
2900 }
2901 }
2902 acc
2903 }
2904}
2905
2906fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
2907 #[cfg(target_arch = "aarch64")]
2908 unsafe {
2909 return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
2910 }
2911 #[cfg(target_arch = "x86_64")]
2912 if avx2_enabled() {
2913 unsafe {
2914 return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
2915 }
2916 }
2917 let _ = gsum;
2918 let mut acc = 0f32;
2919 for gi in 0..gpr {
2920 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
2921 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2922 let mut d = 0i32;
2923 for (j, &b) in tile[2..].iter().enumerate() {
2924 for k in 0..8 {
2925 let w = ((b >> k) & 1) as i32 * 2 - 1;
2926 d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
2927 }
2928 }
2929 acc += d as f32 * s;
2930 }
2931 acc
2932}
2933
2934#[cfg(target_arch = "aarch64")]
2943#[target_feature(enable = "neon,dotprod")]
2944unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
2945 unsafe {
2948 use core::arch::aarch64::*;
2949 use core::arch::asm;
2950 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
2951 let m = vld1q_u8(MASKS.as_ptr());
2952 macro_rules! tile_dot {
2954 ($t:expr, $x:expr) => {{
2955 let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
2956 let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
2957 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
2958 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
2959 let x0 = vld1q_s8($x);
2960 let x1 = vld1q_s8($x.add(16));
2961 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
2962 asm!(
2963 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
2964 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
2965 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
2966 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
2967 options(pure, nomem, nostack),
2968 );
2969 vaddq_s32(a0, a1)
2970 }};
2971 }
2972 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
2981 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
2982 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
2983 const IW11: [u8; 16] = [10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11];
2984 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
2985 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
2986 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
2987 let isc = vld1_u8(ISC.as_ptr());
2988 macro_rules! tile_dot_tbl {
2990 ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
2991 let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
2992 let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
2993 let x0 = vld1q_s8($x);
2994 let x1 = vld1q_s8($x.add(16));
2995 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
2996 asm!(
2997 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
2998 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
2999 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3000 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
3001 options(pure, nomem, nostack),
3002 );
3003 vaddq_s32(a0, a1)
3004 }};
3005 }
3006 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
3007 let row_base = r * gpr * Q1_TILE;
3008 let abs_end = bytes.len();
3009 let xp = xq.as_ptr();
3010 let gp = gsum.as_ptr();
3011 let mut accv = vdupq_n_f32(0.0);
3012 let mut gi = 0;
3013 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
3016 let t0 = base.add(gi * Q1_TILE);
3017 let ld_a = vld1q_u8(t0);
3018 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
3019 let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
3020 let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
3021 let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
3022 let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
3023 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
3025 let g = vld1q_s32(gp.add(gi));
3026 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
3027 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
3028 let scf: float32x4_t;
3029 asm!(
3030 "fcvtl {o:v}.4s, {i:v}.4h",
3031 o = out(vreg) scf, i = in(vreg) sc16,
3032 options(pure, nomem, nostack),
3033 );
3034 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
3035 gi += 4;
3036 }
3037 let mut acc = vaddvq_f32(accv);
3038 while gi < gpr {
3039 let t = base.add(gi * Q1_TILE);
3040 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3041 let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
3042 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
3043 gi += 1;
3044 }
3045 acc
3046 }
3047}
3048
3049#[cfg(target_arch = "aarch64")]
3054#[target_feature(enable = "neon,dotprod")]
3055unsafe fn dot_q1_row_1x4_sdot(
3056 bytes: &[u8],
3057 r: usize,
3058 gpr: usize,
3059 xs: [&[i8]; 4],
3060 gs: [&[i32]; 4],
3061) -> [f32; 4] {
3062 unsafe {
3064 use core::arch::aarch64::*;
3065 use core::arch::asm;
3066 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
3067 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
3068 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
3069 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
3070 const IW11: [u8; 16] = [10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11];
3071 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
3072 let m = vld1q_u8(MASKS.as_ptr());
3073 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
3074 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
3075 let isc = vld1_u8(ISC.as_ptr());
3076 macro_rules! sdot2 {
3077 ($w0:expr, $w1:expr, $x:expr) => {{
3078 let x0 = vld1q_s8($x);
3079 let x1 = vld1q_s8($x.add(16));
3080 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3081 asm!(
3082 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
3083 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
3084 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3085 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
3086 options(pure, nomem, nostack),
3087 );
3088 vaddq_s32(a0, a1)
3089 }};
3090 }
3091 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
3092 let row_base = r * gpr * Q1_TILE;
3093 let abs_end = bytes.len();
3094 let mut accv = [vdupq_n_f32(0.0); 4];
3095 let mut gi = 0;
3096 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
3097 let t0 = base.add(gi * Q1_TILE);
3098 let ld_a = vld1q_u8(t0);
3099 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
3100 let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
3102 let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
3103 let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
3104 let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
3105 let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
3106 let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
3107 let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
3108 let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
3109 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
3110 let scf: float32x4_t;
3111 asm!(
3112 "fcvtl {o:v}.4s, {i:v}.4h",
3113 o = out(vreg) scf, i = in(vreg) sc16,
3114 options(pure, nomem, nostack),
3115 );
3116 for k in 0..4 {
3117 let xp = xs[k].as_ptr();
3118 let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
3119 let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
3120 let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
3121 let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
3122 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
3123 let g = vld1q_s32(gs[k].as_ptr().add(gi));
3124 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
3125 accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
3126 }
3127 gi += 4;
3128 }
3129 let mut acc = [
3130 vaddvq_f32(accv[0]),
3131 vaddvq_f32(accv[1]),
3132 vaddvq_f32(accv[2]),
3133 vaddvq_f32(accv[3]),
3134 ];
3135 while gi < gpr {
3136 let t = base.add(gi * Q1_TILE);
3137 let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3138 let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
3139 let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
3140 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
3141 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
3142 for k in 0..4 {
3143 let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
3144 acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
3145 }
3146 gi += 1;
3147 }
3148 acc
3149 }
3150}
3151
3152#[inline]
3154fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
3155 let gi = j / GROUP_SIZE;
3156 let k = j % GROUP_SIZE;
3157 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
3158 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3159 let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
3160 ((bit as i32 * 2 - 1) as f32, s)
3161}
3162
3163#[inline]
3165fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
3166 let mut acc = 0f32;
3167 for gi in 0..gpr {
3168 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
3169 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3170 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
3171 let mut ga = 0f32;
3172 for (j, &b) in tile[2..].iter().enumerate() {
3173 for k in 0..8 {
3174 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
3175 }
3176 }
3177 acc += ga * s;
3178 }
3179 acc
3180}
3181
3182#[allow(clippy::too_many_arguments)]
3185fn q1_range_a8w8(
3186 bytes: &[u8],
3187 gpr: usize,
3188 act: &SplitAct,
3189 gsum: &[i32],
3190 out: SendMut,
3191 start: usize,
3192 end: usize,
3193) {
3194 for r in start..end {
3195 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
3196 for &(j, xv) in &act.outliers {
3197 let (w, s) = q1_outlier(bytes, r, gpr, j);
3198 acc += w * s * xv;
3199 }
3200 unsafe { *out.at(r) = acc };
3202 }
3203}
3204
3205fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
3207 for r in start..end {
3208 unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
3210 }
3211}
3212
3213fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
3218 let entries = base_len + (rows + 1) * 4;
3219 (base_len, entries, entries <= bytes.len())
3220}
3221
3222#[inline]
3224fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
3225 let o = rp_off + r * 4;
3226 u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
3227}
3228
3229const SIGN5: [[f32; 5]; 256] = {
3233 let mut lut = [[0.0f32; 5]; 256];
3234 let pow3 = [1u16, 3, 9, 27, 81];
3235 let mut byte = 0usize;
3236 while byte < 256 {
3237 let mut i = 0usize;
3238 while i < 5 {
3239 let code = (byte as u16 / pow3[i]) % 3;
3240 lut[byte][i] = if code == 1 {
3241 1.0
3242 } else if code == 2 {
3243 -1.0
3244 } else {
3245 0.0
3246 };
3247 i += 1;
3248 }
3249 byte += 1;
3250 }
3251 lut
3252};
3253
3254const SIGN5_I8: [[i8; 5]; 256] = {
3256 let mut lut = [[0i8; 5]; 256];
3257 let pow3 = [1u16, 3, 9, 27, 81];
3258 let mut byte = 0usize;
3259 while byte < 256 {
3260 let mut i = 0usize;
3261 while i < 5 {
3262 let code = (byte as u16 / pow3[i]) % 3;
3263 lut[byte][i] = if code == 1 {
3264 1
3265 } else if code == 2 {
3266 -1
3267 } else {
3268 0
3269 };
3270 i += 1;
3271 }
3272 byte += 1;
3273 }
3274 lut
3275};
3276
3277const SIGN5_U64: [u64; 256] = {
3283 let mut lut = [0u64; 256];
3284 let pow3 = [1u16, 3, 9, 27, 81];
3285 let mut byte = 0usize;
3286 while byte < 256 {
3287 let mut v = 0u64;
3288 let mut i = 0usize;
3289 while i < 5 {
3290 let code = (byte as u16 / pow3[i]) % 3;
3291 let s: u8 = if code == 1 {
3292 1
3293 } else if code == 2 {
3294 0xFF
3295 } else {
3296 0
3297 };
3298 v |= (s as u64) << (i * 8);
3299 i += 1;
3300 }
3301 lut[byte] = v;
3302 byte += 1;
3303 }
3304 lut
3305};
3306
3307#[inline]
3312fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
3313 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3314 let off = (r * gpr + j / GROUP_SIZE) * TILE;
3315 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
3316 let within = j % GROUP_SIZE;
3317 SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
3318}
3319
3320#[cfg(target_arch = "aarch64")]
3323#[target_feature(enable = "neon,dotprod")]
3324#[inline]
3325unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
3326 unsafe {
3328 use core::arch::aarch64::*;
3329 use core::arch::asm;
3330 let w0 = vld1q_s8(w);
3331 let w1 = vld1q_s8(w.add(16));
3332 let x0 = vld1q_s8(x);
3333 let x1 = vld1q_s8(x.add(16));
3334 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3335 asm!(
3336 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
3337 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
3338 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3339 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
3340 options(pure, nomem, nostack),
3341 );
3342 vaddvq_s32(vaddq_s32(a0, a1))
3343 }
3344}
3345
3346#[cfg(target_arch = "x86_64")]
3349#[target_feature(enable = "avx2")]
3350#[inline]
3351unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
3352 unsafe {
3354 use core::arch::x86_64::*;
3355 let wv = _mm256_loadu_si256(w as *const __m256i);
3356 let xv = _mm256_loadu_si256(x as *const __m256i);
3357 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
3358 let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
3359 let hi128 = _mm256_extracti128_si256::<1>(d);
3360 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3361 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3362 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3363 _mm_cvtsi128_si32(s32)
3364 }
3365}
3366
3367#[inline]
3372fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
3373 debug_assert!(dst.len() >= 40);
3374 unsafe {
3377 let p = dst.as_mut_ptr();
3378 for bi in 0..7 {
3379 core::ptr::write_unaligned(
3380 p.add(bi * 5) as *mut u64,
3381 SIGN5_U64[*codes.add(bi) as usize],
3382 );
3383 }
3384 }
3385}
3386
3387#[inline]
3392fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
3393 #[cfg(target_arch = "aarch64")]
3394 unsafe {
3395 return sdot32_i8(w, x);
3396 }
3397 #[cfg(target_arch = "x86_64")]
3398 unsafe {
3399 return i8dot32_avx2(w, x);
3400 }
3401 #[allow(unreachable_code)]
3402 unsafe {
3403 let mut s = 0i32;
3404 for k in 0..GROUP_SIZE {
3405 s += *w.add(k) as i32 * *x.add(k) as i32;
3406 }
3407 s
3408 }
3409}
3410
3411#[cfg(target_arch = "aarch64")]
3415#[target_feature(enable = "neon,dotprod")]
3416unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3417 unsafe {
3419 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3420 let mut acc = 0f32;
3421 let mut sg = [0i8; GROUP_SIZE + 8]; for gi in 0..gpr {
3423 let off = (r * gpr + gi) * TILE;
3424 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
3425 q1t_unpack_group_i8(bytes.as_ptr().add(off + 2), &mut sg);
3426 acc += sdot32_i8(sg.as_ptr(), xq.as_ptr().add(gi * GROUP_SIZE)) as f32 * s;
3427 }
3428 acc
3429 }
3430}
3431
3432#[cfg(target_arch = "x86_64")]
3434#[target_feature(enable = "avx2")]
3435unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3436 unsafe {
3438 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3439 let mut acc = 0f32;
3440 let mut sg = [0i8; GROUP_SIZE + 8]; for gi in 0..gpr {
3442 let off = (r * gpr + gi) * TILE;
3443 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
3444 q1t_unpack_group_i8(bytes.as_ptr().add(off + 2), &mut sg);
3445 acc += i8dot32_avx2(sg.as_ptr(), xq.as_ptr().add(gi * GROUP_SIZE)) as f32 * s;
3446 }
3447 acc
3448 }
3449}
3450
3451#[inline]
3455fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3456 #[cfg(target_arch = "aarch64")]
3457 unsafe {
3458 return q1t_dot_row_sdot(bytes, r, gpr, xq);
3459 }
3460 #[cfg(target_arch = "x86_64")]
3461 unsafe {
3462 return q1t_dot_row_avx2(bytes, r, gpr, xq);
3463 }
3464 #[allow(unreachable_code)]
3465 {
3466 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3467 let mut acc = 0f32;
3468 let mut sg = [0i8; GROUP_SIZE + 8]; for gi in 0..gpr {
3470 let off = (r * gpr + gi) * TILE;
3471 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
3472 q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
3473 let mut d = 0i32;
3474 for k in 0..GROUP_SIZE {
3475 d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
3476 }
3477 acc += d as f32 * s;
3478 }
3479 acc
3480 }
3481}
3482
3483fn q1t_row_outlier_correction(
3490 bytes: &[u8],
3491 r: usize,
3492 rp_off: usize,
3493 entries_off: usize,
3494 has_ov: bool,
3495 x: &[f32],
3496) -> f32 {
3497 if !has_ov {
3498 return 0.0;
3499 }
3500 let (c0, c1) = (q1t_rowptr(bytes, rp_off, r), q1t_rowptr(bytes, rp_off, r + 1));
3501 let mut corr = 0f32;
3502 for p in c0..c1 {
3503 let e = entries_off + p * 4;
3504 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
3505 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
3506 corr += val * x[col];
3507 }
3508 corr
3509}
3510
3511fn q1t_dequant_row(
3515 bytes: &[u8],
3516 r: usize,
3517 gpr: usize,
3518 rp_off: usize,
3519 entries_off: usize,
3520 has_ov: bool,
3521 buf: &mut [f32],
3522) {
3523 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3524 for g in 0..gpr {
3525 let off = (r * gpr + g) * TILE;
3526 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
3527 let codes = &bytes[off + 2..off + TILE];
3528 let bc = g * GROUP_SIZE;
3529 for bi in 0..6 {
3531 let lut = &SIGN5[codes[bi] as usize];
3532 let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
3533 for i in 0..5 {
3534 d[i] = lut[i] * s;
3535 }
3536 }
3537 let lut = &SIGN5[codes[6] as usize];
3538 buf[bc + 30] = lut[0] * s;
3539 buf[bc + 31] = lut[1] * s;
3540 }
3541 if !has_ov {
3542 return;
3543 }
3544 let (c0, c1) = (q1t_rowptr(bytes, rp_off, r), q1t_rowptr(bytes, rp_off, r + 1));
3545 for p in c0..c1 {
3546 let e = entries_off + p * 4;
3547 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
3548 buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
3549 }
3550}
3551
3552fn q1t_add_overlay(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
3556 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3557 let gpr = cols / GROUP_SIZE;
3558 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
3559 if !has_ov {
3560 return;
3561 }
3562 let out_addr = SendMut(out.as_mut_ptr());
3563 let run = move |start: usize, end: usize| {
3564 for r in start..end {
3565 let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
3566 unsafe { *out_addr.at(r) += corr };
3568 }
3569 };
3570 dispatch_rows(pool, rows, &run);
3571}
3572
3573fn q1t_matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
3577 debug_assert_eq!(out.len(), rows);
3578 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3579 let gpr = cols / GROUP_SIZE;
3580 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
3581 let out_addr = SendMut(out.as_mut_ptr());
3582 if a8w8_enabled() {
3586 let act = split_act(x);
3587 let act = &act;
3588 let run = move |start: usize, end: usize| {
3589 for r in start..end {
3590 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
3591 for &(j, xv) in &act.outliers {
3592 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
3593 }
3594 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
3595 unsafe { *out_addr.at(r) = acc };
3597 }
3598 };
3599 dispatch_rows(pool, rows, &run);
3600 return;
3601 }
3602 let run = move |start: usize, end: usize| {
3603 let mut sg = [0f32; GROUP_SIZE];
3607 for r in start..end {
3608 let mut acc = 0f32;
3609 for g in 0..gpr {
3610 let off = (r * gpr + g) * TILE;
3611 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
3612 let codes = &bytes[off + 2..off + TILE];
3613 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
3614 for bi in 0..6 {
3615 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
3616 }
3617 let lut = &SIGN5[codes[6] as usize];
3618 sg[30] = lut[0];
3619 sg[31] = lut[1];
3620 let mut gsum = 0f32;
3621 for k in 0..GROUP_SIZE {
3622 gsum += sg[k] * xg[k];
3623 }
3624 acc += s * gsum;
3625 }
3626 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
3627 unsafe { *out_addr.at(r) = acc };
3628 }
3629 };
3630 dispatch_rows(pool, rows, &run);
3631}
3632
3633fn q1t_matmat(
3636 bytes: &[u8],
3637 xs: &[f32],
3638 b: usize,
3639 rows: usize,
3640 cols: usize,
3641 out: &mut [f32],
3642 pool: Option<&Pool>,
3643) {
3644 debug_assert_eq!(out.len(), b * rows);
3645 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3646 let gpr = cols / GROUP_SIZE;
3647 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
3648 let out_addr = SendMut(out.as_mut_ptr());
3649 if a8w8_enabled() {
3653 let acts: Vec<SplitAct> = (0..b)
3654 .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
3655 .collect();
3656 let acts = &acts;
3657 let run = move |start: usize, end: usize| {
3658 let mut sg = vec![0i8; cols + 8]; let mut sc = vec![0f32; gpr]; let mut accs = vec![0f32; b]; for r in start..end {
3662 for g in 0..gpr {
3663 let off = (r * gpr + g) * TILE;
3664 sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
3665 q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg[g * GROUP_SIZE..]);
3666 }
3667 for bi in 0..b {
3668 let act = &acts[bi];
3669 let mut isum = 0f32;
3670 for g in 0..gpr {
3671 let d = q1t_i8dot32(
3672 sg.as_ptr().wrapping_add(g * GROUP_SIZE),
3673 act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
3674 );
3675 isum += d as f32 * sc[g];
3676 }
3677 let mut acc = isum * act.sx;
3678 for &(j, xv) in &act.outliers {
3679 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
3680 }
3681 accs[bi] = acc;
3682 }
3683 if has_ov {
3687 let (c0, c1) =
3688 (q1t_rowptr(bytes, rp_off, r), q1t_rowptr(bytes, rp_off, r + 1));
3689 for p in c0..c1 {
3690 let e = ent_off + p * 4;
3691 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
3692 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
3693 for bi in 0..b {
3694 accs[bi] += val * xs[bi * cols + col];
3695 }
3696 }
3697 }
3698 for bi in 0..b {
3699 unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
3700 }
3701 }
3702 };
3703 dispatch_rows(pool, rows, &run);
3704 return;
3705 }
3706 let run = move |start: usize, end: usize| {
3707 let mut buf = vec![0f32; cols];
3708 for r in start..end {
3709 q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
3710 for bi in 0..b {
3711 let xr = &xs[bi * cols..(bi + 1) * cols];
3712 let mut acc = 0f32;
3713 for j in 0..cols {
3714 acc += buf[j] * xr[j];
3715 }
3716 unsafe { *out_addr.at(bi * rows + r) = acc };
3717 }
3718 }
3719 };
3720 dispatch_rows(pool, rows, &run);
3721}
3722
3723fn q1_matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
3724 debug_assert_eq!(out.len(), rows);
3725 let gpr = cols / GROUP_SIZE;
3726 let out_addr = SendMut(out.as_mut_ptr());
3727 if a8w8_enabled() {
3728 let act = split_act(x);
3729 let gsum = q1_group_sums(&act.xq, gpr);
3730 let (act, gsum) = (&act, &gsum);
3731 let run =
3732 move |start: usize, end: usize| q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end);
3733 dispatch_rows(pool, rows, &run);
3734 return;
3735 }
3736 let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
3737 dispatch_rows(pool, rows, &run);
3738}
3739
3740#[allow(clippy::too_many_arguments)]
3742fn q1_matvec2(
3743 bytes: &[u8],
3744 x1: &[f32],
3745 x2: &[f32],
3746 rows: usize,
3747 cols: usize,
3748 o1: &mut [f32],
3749 o2: &mut [f32],
3750 pool: Option<&Pool>,
3751) {
3752 let gpr = cols / GROUP_SIZE;
3753 let p1 = SendMut(o1.as_mut_ptr());
3754 let p2 = SendMut(o2.as_mut_ptr());
3755 if a8w8_enabled() {
3756 let a1 = split_act(x1);
3757 let a2 = split_act(x2);
3758 let g1 = q1_group_sums(&a1.xq, gpr);
3759 let g2 = q1_group_sums(&a2.xq, gpr);
3760 let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
3761 let run = move |start: usize, end: usize| {
3762 for r in start..end {
3763 let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
3764 let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
3765 for &(j, xv) in &a1.outliers {
3766 let (w, s) = q1_outlier(bytes, r, gpr, j);
3767 v1 += w * s * xv;
3768 }
3769 for &(j, xv) in &a2.outliers {
3770 let (w, s) = q1_outlier(bytes, r, gpr, j);
3771 v2 += w * s * xv;
3772 }
3773 unsafe {
3775 *p1.at(r) = v1;
3776 *p2.at(r) = v2;
3777 }
3778 }
3779 };
3780 dispatch_rows(pool, rows, &run);
3781 return;
3782 }
3783 let run = move |start: usize, end: usize| {
3784 for r in start..end {
3785 unsafe {
3787 *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
3788 *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
3789 }
3790 }
3791 };
3792 dispatch_rows(pool, rows, &run);
3793}
3794
3795#[allow(clippy::too_many_arguments)]
3797fn q1_matmat(
3798 bytes: &[u8],
3799 xs_all: &[f32],
3800 b: usize,
3801 rows: usize,
3802 cols: usize,
3803 out: &mut [f32],
3804 pool: Option<&Pool>,
3805) {
3806 debug_assert_eq!(out.len(), b * rows);
3807 let gpr = cols / GROUP_SIZE;
3808 let out_addr = SendMut(out.as_mut_ptr());
3809 if a8w8_enabled() {
3810 let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
3811 .map(|bi| {
3812 let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
3813 let gsum = q1_group_sums(&act.xq, gpr);
3814 (act, gsum)
3815 })
3816 .collect();
3817 let acts = &acts;
3818 #[cfg(target_arch = "x86_64")]
3819 let blocked_ok = avx2_enabled()
3820 && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true);
3821 #[cfg(target_arch = "aarch64")]
3822 let blocked_ok = sdot_enabled()
3823 && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true);
3824 let run = move |start: usize, end: usize| {
3825 for r in start..end {
3826 let mut bi = 0usize;
3827 #[cfg(target_arch = "aarch64")]
3830 if blocked_ok {
3831 while bi + 4 <= acts.len() {
3832 let xs = [
3833 acts[bi].0.xq.as_slice(),
3834 acts[bi + 1].0.xq.as_slice(),
3835 acts[bi + 2].0.xq.as_slice(),
3836 acts[bi + 3].0.xq.as_slice(),
3837 ];
3838 let gs = [
3839 acts[bi].1.as_slice(),
3840 acts[bi + 1].1.as_slice(),
3841 acts[bi + 2].1.as_slice(),
3842 acts[bi + 3].1.as_slice(),
3843 ];
3844 let d = unsafe { dot_q1_row_1x4_sdot(bytes, r, gpr, xs, gs) };
3845 for k in 0..4 {
3846 let (act, _) = &acts[bi + k];
3847 let mut acc = d[k] * act.sx;
3848 for &(j, xv) in &act.outliers {
3849 let (w, sc) = q1_outlier(bytes, r, gpr, j);
3850 acc += w * sc * xv;
3851 }
3852 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
3854 }
3855 bi += 4;
3856 }
3857 }
3858 #[cfg(target_arch = "x86_64")]
3859 if blocked_ok {
3860 while bi + 4 <= acts.len() {
3861 let xs = [
3862 acts[bi].0.xq.as_slice(),
3863 acts[bi + 1].0.xq.as_slice(),
3864 acts[bi + 2].0.xq.as_slice(),
3865 acts[bi + 3].0.xq.as_slice(),
3866 ];
3867 let gs = [
3868 acts[bi].1.as_slice(),
3869 acts[bi + 1].1.as_slice(),
3870 acts[bi + 2].1.as_slice(),
3871 acts[bi + 3].1.as_slice(),
3872 ];
3873 let d = unsafe { dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs) };
3874 for k in 0..4 {
3875 let (act, _) = &acts[bi + k];
3876 let mut acc = d[k] * act.sx;
3877 for &(j, xv) in &act.outliers {
3878 let (w, sc) = q1_outlier(bytes, r, gpr, j);
3879 acc += w * sc * xv;
3880 }
3881 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
3883 }
3884 bi += 4;
3885 }
3886 }
3887 while bi < acts.len() {
3888 let (act, gsum) = &acts[bi];
3889 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
3890 for &(j, xv) in &act.outliers {
3891 let (w, s) = q1_outlier(bytes, r, gpr, j);
3892 acc += w * s * xv;
3893 }
3894 unsafe { *out_addr.at(bi * rows + r) = acc };
3896 bi += 1;
3897 }
3898 }
3899 };
3900 dispatch_rows(pool, rows, &run);
3901 return;
3902 }
3903 let run = move |start: usize, end: usize| {
3904 for r in start..end {
3905 for bi in 0..b {
3906 let x = &xs_all[bi * cols..(bi + 1) * cols];
3907 unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
3909 }
3910 }
3911 };
3912 dispatch_rows(pool, rows, &run);
3913}
3914
3915fn q4matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
3921 debug_assert_eq!(out.len(), rows);
3922 let (packed, scales) = q4_split(bytes, rows, cols);
3923 let gpr = cols / GROUP_SIZE;
3924 let out_addr = SendMut(out.as_mut_ptr());
3925
3926 if a8w8_enabled() {
3927 let act = split_act(x);
3928 let run = move |start: usize, end: usize| {
3929 q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
3930 };
3931 dispatch_rows(pool, rows, &run);
3932 return;
3933 }
3934
3935 let run = move |start: usize, end: usize| {
3936 q4_range_f32(packed, scales, gpr, x, out_addr, start, end)
3937 };
3938 dispatch_rows(pool, rows, &run);
3939}
3940
3941#[inline]
3944#[allow(unreachable_code)]
3945#[cfg(target_arch = "x86_64")]
3950#[target_feature(enable = "avx2")]
3951unsafe fn dot_q4b_row_1x4_avx2(
3952 buf: &[u8],
3953 scales: &[u8],
3954 g0: usize,
3955 gpr: usize,
3956 xs: [&[i8]; 4],
3957) -> [f32; 4] {
3958 unsafe {
3960 use core::arch::x86_64::*;
3961 let ones = _mm256_set1_epi16(1);
3962 let mut acc = [0f32; 4];
3963 for gi in 0..gpr {
3964 let s = f16_to_f32(u16::from_le_bytes([
3965 scales[(g0 + gi) * 2],
3966 scales[(g0 + gi) * 2 + 1],
3967 ]));
3968 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3969 let aw = _mm256_abs_epi8(w);
3970 for (k, xq) in xs.iter().enumerate() {
3971 let x =
3972 _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3973 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
3974 let d = _mm256_madd_epi16(p16, ones);
3975 let hi128 = _mm256_extracti128_si256::<1>(d);
3976 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3977 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3978 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3979 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
3980 }
3981 }
3982 acc
3983 }
3984}
3985
3986#[cfg(target_arch = "x86_64")]
3992#[target_feature(enable = "avx2")]
3993unsafe fn dot_q4b_row_1x4_sx_avx2(
3994 buf: &[u8],
3995 scales: &[u8],
3996 g0: usize,
3997 gpr: usize,
3998 xs: [&[i8]; 4],
3999 sxs: [f32; 4],
4000) -> [f32; 4] {
4001 unsafe {
4003 use core::arch::x86_64::*;
4004 let ones = _mm256_set1_epi16(1);
4005 let mut acc = [0f32; 4];
4006 for gi in 0..gpr {
4007 let s = f16_to_f32(u16::from_le_bytes([
4008 scales[(g0 + gi) * 2],
4009 scales[(g0 + gi) * 2 + 1],
4010 ]));
4011 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4012 let aw = _mm256_abs_epi8(w);
4013 for (k, xq) in xs.iter().enumerate() {
4014 let x =
4015 _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4016 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
4017 let d = _mm256_madd_epi16(p16, ones);
4018 let hi128 = _mm256_extracti128_si256::<1>(d);
4019 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4020 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4021 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4022 acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
4023 }
4024 }
4025 acc
4026 }
4027}
4028
4029fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
4030 #[cfg(target_arch = "aarch64")]
4031 unsafe {
4032 return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
4033 }
4034 #[cfg(target_arch = "x86_64")]
4035 unsafe {
4036 return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
4037 }
4038 let mut acc = 0f32;
4039 for gi in 0..gpr {
4040 let g = g0 + gi;
4041 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4042 let mut d = 0i32;
4043 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
4044 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
4045 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
4046 }
4047 acc += d as f32 * s;
4048 }
4049 acc
4050}
4051
4052#[inline]
4054#[allow(unreachable_code)]
4055fn dot_q4_row_i8_2(
4056 packed: &[u8],
4057 scales: &[u8],
4058 g0: usize,
4059 gpr: usize,
4060 xq1: &[i8],
4061 xq2: &[i8],
4062) -> (f32, f32) {
4063 #[cfg(target_arch = "aarch64")]
4064 unsafe {
4065 return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
4066 }
4067 #[cfg(target_arch = "x86_64")]
4068 unsafe {
4069 return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
4070 }
4071 (
4072 dot_q4_row_i8(packed, scales, g0, gpr, xq1),
4073 dot_q4_row_i8(packed, scales, g0, gpr, xq2),
4074 )
4075}
4076
4077#[allow(clippy::too_many_arguments)]
4080fn q4_range_a8w8(
4081 packed: &[u8],
4082 scales: &[u8],
4083 gpr: usize,
4084 cols: usize,
4085 act: &SplitAct,
4086 out: SendMut,
4087 start: usize,
4088 end: usize,
4089) {
4090 for r in start..end {
4091 let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
4092 for &(j, xv) in &act.outliers {
4094 let flat = r * cols + j;
4095 let byte = packed[flat / 2];
4096 let nib = if flat & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4097 let s = f16_to_f32(u16::from_le_bytes([
4098 scales[(flat / GROUP_SIZE) * 2],
4099 scales[(flat / GROUP_SIZE) * 2 + 1],
4100 ]));
4101 acc += ((nib as i32 - 8) as f32) * s * xv;
4102 }
4103 unsafe { *out.at(r) = acc };
4105 }
4106}
4107
4108#[allow(clippy::too_many_arguments)]
4111fn q4_range2_a8w8(
4112 packed: &[u8],
4113 scales: &[u8],
4114 gpr: usize,
4115 cols: usize,
4116 a1: &SplitAct,
4117 a2: &SplitAct,
4118 p1: SendMut,
4119 p2: SendMut,
4120 start: usize,
4121 end: usize,
4122) {
4123 for r in start..end {
4124 let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
4125 let mut acc1 = s1 * a1.sx;
4126 let mut acc2 = s2 * a2.sx;
4127 let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
4129 for &(j, xv) in outliers {
4130 let flat = r * cols + j;
4131 let byte = packed[flat / 2];
4132 let nib = if flat & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4133 let s = f16_to_f32(u16::from_le_bytes([
4134 scales[(flat / GROUP_SIZE) * 2],
4135 scales[(flat / GROUP_SIZE) * 2 + 1],
4136 ]));
4137 *acc += ((nib as i32 - 8) as f32) * s * xv;
4138 }
4139 };
4140 fix(&a1.outliers, &mut acc1);
4141 fix(&a2.outliers, &mut acc2);
4142 unsafe {
4144 *p1.at(r) = acc1;
4145 *p2.at(r) = acc2;
4146 }
4147 }
4148}
4149
4150fn q4_range_f32(
4152 packed: &[u8],
4153 scales: &[u8],
4154 gpr: usize,
4155 x: &[f32],
4156 out: SendMut,
4157 start: usize,
4158 end: usize,
4159) {
4160 for r in start..end {
4161 let mut acc = 0f32;
4162 for gi in 0..gpr {
4163 let g = r * gpr + gi;
4164 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4165 let pk = &packed[g * 16..(g + 1) * 16];
4166 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4167 let mut ga = 0f32;
4168 for (k, &b) in pk.iter().enumerate() {
4169 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4170 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4171 }
4172 acc += ga * s;
4173 }
4174 unsafe { *out.at(r) = acc };
4176 }
4177}
4178
4179#[allow(clippy::too_many_arguments)]
4183fn q4matvec2(
4184 bytes: &[u8],
4185 x1: &[f32],
4186 x2: &[f32],
4187 rows: usize,
4188 cols: usize,
4189 o1: &mut [f32],
4190 o2: &mut [f32],
4191 pool: Option<&Pool>,
4192) {
4193 debug_assert_eq!(o1.len(), rows);
4194 debug_assert_eq!(o2.len(), rows);
4195 let (packed, scales) = q4_split(bytes, rows, cols);
4196 let gpr = cols / GROUP_SIZE;
4197
4198 if a8w8_enabled() {
4199 let a1 = split_act(x1);
4200 let a2 = split_act(x2);
4201 let p1 = SendMut(o1.as_mut_ptr());
4202 let p2 = SendMut(o2.as_mut_ptr());
4203 let run = move |start: usize, end: usize| {
4204 q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
4205 };
4206 dispatch_rows(pool, rows, &run);
4207 return;
4208 }
4209
4210 let p1 = SendMut(o1.as_mut_ptr());
4211 let p2 = SendMut(o2.as_mut_ptr());
4212 let run = move |start: usize, end: usize| {
4213 q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
4214 };
4215 dispatch_rows(pool, rows, &run);
4216}
4217
4218#[allow(clippy::too_many_arguments)]
4220fn q4_range2_f32(
4221 packed: &[u8],
4222 scales: &[u8],
4223 gpr: usize,
4224 x1: &[f32],
4225 x2: &[f32],
4226 p1: SendMut,
4227 p2: SendMut,
4228 start: usize,
4229 end: usize,
4230) {
4231 for r in start..end {
4232 let (mut acc1, mut acc2) = (0f32, 0f32);
4233 for gi in 0..gpr {
4234 let g = r * gpr + gi;
4235 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4236 let pk = &packed[g * 16..(g + 1) * 16];
4237 let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4238 let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4239 let (mut g1, mut g2) = (0f32, 0f32);
4240 for (k, &b) in pk.iter().enumerate() {
4241 let wl = (b & 0x0F) as f32 - 8.0;
4242 let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
4243 g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
4244 g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
4245 }
4246 acc1 += g1 * s;
4247 acc2 += g2 * s;
4248 }
4249 unsafe {
4251 *p1.at(r) = acc1;
4252 *p2.at(r) = acc2;
4253 }
4254 }
4255}
4256
4257thread_local! {
4258 static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
4261 static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4262}
4263
4264#[allow(clippy::too_many_arguments)]
4270fn q4matmat(
4271 bytes: &[u8],
4272 xs_all: &[f32],
4273 b: usize,
4274 rows: usize,
4275 cols: usize,
4276 out: &mut [f32],
4277 pool: Option<&Pool>,
4278) {
4279 debug_assert_eq!(xs_all.len(), b * cols);
4280 debug_assert_eq!(out.len(), b * rows);
4281 let (packed, scales) = q4_split(bytes, rows, cols);
4282 let gpr = cols / GROUP_SIZE;
4283 let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4284
4285 if a8w8_enabled() {
4286 let acts: Vec<SplitAct> =
4287 (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
4288 let acts = &acts;
4289 let out_addr = SendMut(out.as_mut_ptr());
4290 let run = move |start: usize, end: usize| {
4291 ROW_I8.with(|rb| {
4292 let mut buf = rb.borrow_mut();
4293 buf.resize(cols, 0);
4294 for r in start..end {
4295 for gi in 0..gpr {
4299 let g = r * gpr + gi;
4300 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
4301 buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
4302 buf[gi * GROUP_SIZE + k * 2 + 1] =
4303 (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
4304 }
4305 }
4306 let mut bi = 0usize;
4307 #[cfg(target_arch = "x86_64")]
4308 if avx2_enabled()
4309 && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true)
4310 {
4311 while bi + 4 <= acts.len() {
4312 let xs = [
4313 acts[bi].xq.as_slice(),
4314 acts[bi + 1].xq.as_slice(),
4315 acts[bi + 2].xq.as_slice(),
4316 acts[bi + 3].xq.as_slice(),
4317 ];
4318 let d = unsafe {
4319 dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
4320 };
4321 for k in 0..4 {
4322 let act = &acts[bi + k];
4323 let mut acc = d[k] * act.sx;
4324 for &(j, xv) in &act.outliers {
4325 acc += (buf[j] as i8) as f32
4326 * gscale((r * cols + j) / GROUP_SIZE)
4327 * xv;
4328 }
4329 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
4331 }
4332 bi += 4;
4333 }
4334 }
4335 while bi < acts.len() {
4336 let act = &acts[bi];
4337 let mut acc = 0f32;
4338 for gi in 0..gpr {
4339 let d = dot_i8_i8(
4340 &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
4341 &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
4342 );
4343 acc += d as f32 * gscale(r * gpr + gi);
4344 }
4345 acc *= act.sx;
4346 for &(j, xv) in &act.outliers {
4348 acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
4349 }
4350 unsafe { *out_addr.at(bi * rows + r) = acc };
4352 bi += 1;
4353 }
4354 }
4355 })
4356 };
4357 dispatch_rows(pool, rows, &run);
4358 return;
4359 }
4360
4361 let out_addr = SendMut(out.as_mut_ptr());
4362 let run = move |start: usize, end: usize| {
4363 ROW_F32.with(|rb| {
4364 let mut buf = rb.borrow_mut();
4365 buf.resize(cols, 0.0);
4366 for r in start..end {
4367 for gi in 0..gpr {
4370 let g = r * gpr + gi;
4371 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
4372 buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
4373 buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
4374 }
4375 }
4376 for bi in 0..b {
4377 let x = &xs_all[bi * cols..(bi + 1) * cols];
4378 let mut acc = 0f32;
4379 for gi in 0..gpr {
4380 let mut ga = 0f32;
4381 for k in 0..GROUP_SIZE / 2 {
4386 let e = gi * GROUP_SIZE + k * 2;
4387 ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
4388 }
4389 acc += ga * gscale(r * gpr + gi);
4390 }
4391 unsafe { *out_addr.at(bi * rows + r) = acc };
4393 }
4394 }
4395 })
4396 };
4397 dispatch_rows(pool, rows, &run);
4398}
4399
4400#[allow(clippy::too_many_arguments)]
4405fn vbitmatmat(
4406 bytes: &[u8],
4407 offsets: &[usize],
4408 xs_all: &[f32],
4409 b: usize,
4410 rows: usize,
4411 cols: usize,
4412 out: &mut [f32],
4413 pool: Option<&Pool>,
4414) {
4415 debug_assert_eq!(xs_all.len(), b * cols);
4416 debug_assert_eq!(out.len(), b * rows);
4417 debug_assert_eq!(offsets.len(), rows + 1);
4418 let ng = cols / GROUP_SIZE;
4419 let bits = &bytes[..rows];
4420 let sc_off = rows;
4421 let gscale = |r: usize, g: usize| {
4422 let so = (r * ng + g) * 2;
4423 f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]))
4424 };
4425
4426 let decode_f32 = |r: usize, dst: &mut [f32]| {
4428 let bw = bits[r] as usize;
4429 let l = ((1i32 << (bw - 1)) - 1) as f32;
4430 let data = &bytes[offsets[r]..offsets[r + 1]];
4431 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
4432 for d in dst.iter_mut() {
4433 while nbits < bw {
4434 acc = (acc << 8) | data[idx] as u64;
4435 idx += 1;
4436 nbits += 8;
4437 }
4438 let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
4439 nbits -= bw;
4440 *d = u - l;
4441 }
4442 };
4443
4444 if a8w8_enabled() {
4445 let acts: Vec<SplitAct> =
4446 (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
4447 let acts = &acts;
4448 let out_addr = SendMut(out.as_mut_ptr());
4449 let run = move |start: usize, end: usize| {
4450 for r in start..end {
4451 let bw = bits[r] as usize;
4452 if bw == 8 {
4453 ROW_F32.with(|rb| {
4456 let mut buf = rb.borrow_mut();
4457 buf.resize(cols, 0.0);
4458 decode_f32(r, &mut buf);
4459 for bi in 0..b {
4460 let x = &xs_all[bi * cols..(bi + 1) * cols];
4461 let mut dot = 0f32;
4462 for g in 0..ng {
4463 let mut gd = 0f32;
4464 for k in 0..GROUP_SIZE {
4465 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
4466 }
4467 dot += gd * gscale(r, g);
4468 }
4469 unsafe { *out_addr.at(bi * rows + r) = dot };
4471 }
4472 });
4473 continue;
4474 }
4475 let l = (1i32 << (bw - 1)) - 1;
4476 let data = &bytes[offsets[r]..offsets[r + 1]];
4477 ROW_I8.with(|rb| {
4478 let mut buf = rb.borrow_mut();
4479 buf.resize(cols, 0);
4480 #[inline(always)]
4481 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
4482 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4483 let u = unpack8::<B>(&data[blk * B..]);
4484 for k in 0..8 {
4485 chunk[k] = (u[k] - l) as i8 as u8;
4486 }
4487 }
4488 }
4489 match bw {
4490 3 => fill::<3>(data, l, &mut buf),
4491 4 => vbit_fill4(data, &mut buf),
4492 5 => fill::<5>(data, l, &mut buf),
4493 6 => fill::<6>(data, l, &mut buf),
4494 _ => unreachable!("vbit bit-width {bw} (validated at load)"),
4495 }
4496 let mut bi = 0usize;
4497 #[cfg(target_arch = "x86_64")]
4501 if avx2_enabled()
4502 && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true)
4503 {
4504 while bi + 4 <= acts.len() {
4505 let xs = [
4506 acts[bi].xq.as_slice(),
4507 acts[bi + 1].xq.as_slice(),
4508 acts[bi + 2].xq.as_slice(),
4509 acts[bi + 3].xq.as_slice(),
4510 ];
4511 let sxs = [
4512 acts[bi].sx,
4513 acts[bi + 1].sx,
4514 acts[bi + 2].sx,
4515 acts[bi + 3].sx,
4516 ];
4517 let d = unsafe {
4518 dot_q4b_row_1x4_sx_avx2(&buf, &bytes[sc_off..], r * ng, ng, xs, sxs)
4519 };
4520 for k in 0..4 {
4521 let act = &acts[bi + k];
4522 let mut dot = d[k];
4523 for &(j, xv) in &act.outliers {
4524 dot +=
4525 (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
4526 }
4527 unsafe { *out_addr.at((bi + k) * rows + r) = dot };
4529 }
4530 bi += 4;
4531 }
4532 }
4533 while bi < acts.len() {
4534 let act = &acts[bi];
4535 let mut dot = 0f32;
4536 for g in 0..ng {
4537 let d = dot_i8_i8(
4538 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4539 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4540 ) as f32
4541 * act.sx;
4542 dot += d * gscale(r, g);
4543 }
4544 for &(j, xv) in &act.outliers {
4545 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
4546 }
4547 unsafe { *out_addr.at(bi * rows + r) = dot };
4549 bi += 1;
4550 }
4551 });
4552 }
4553 };
4554 dispatch_rows(pool, rows, &run);
4555 return;
4556 }
4557
4558 let out_addr = SendMut(out.as_mut_ptr());
4559 let run = move |start: usize, end: usize| {
4560 ROW_F32.with(|rb| {
4561 let mut buf = rb.borrow_mut();
4562 buf.resize(cols, 0.0);
4563 for r in start..end {
4564 decode_f32(r, &mut buf);
4565 for bi in 0..b {
4566 let x = &xs_all[bi * cols..(bi + 1) * cols];
4567 let mut dot = 0f32;
4568 for g in 0..ng {
4569 let mut gd = 0f32;
4570 for k in 0..GROUP_SIZE {
4571 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
4572 }
4573 dot += gd * gscale(r, g);
4574 }
4575 unsafe { *out_addr.at(bi * rows + r) = dot };
4577 }
4578 }
4579 })
4580 };
4581 dispatch_rows(pool, rows, &run);
4582}
4583
4584pub(crate) fn gpu_batch_job<'a>(
4588 t: &'a QTensor,
4589 x: &[f32],
4590) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
4591 match t {
4592 QTensor::Mapped {
4593 model,
4594 idx,
4595 dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
4596 rows,
4597 cols,
4598 row_scale,
4599 col_field,
4600 ..
4601 } => Some((
4602 model.clone(),
4603 crate::gpu::BatchJob {
4604 idx: *idx,
4605 rows: *rows,
4606 cols: *cols,
4607 row_scale,
4608 xs: prescale(x, col_field, *dt).into_owned(),
4609 q1: false,
4610 },
4611 )),
4612 QTensor::Mapped {
4614 model,
4615 idx,
4616 dtype: TensorDtype::Q1,
4617 rows,
4618 cols,
4619 ..
4620 } => Some((
4621 model.clone(),
4622 crate::gpu::BatchJob {
4623 idx: *idx,
4624 rows: *rows,
4625 cols: *cols,
4626 row_scale: &[],
4627 xs: x.to_vec(),
4628 q1: true,
4629 },
4630 )),
4631 _ => None,
4632 }
4633}
4634
4635pub(crate) fn prescale<'a>(
4639 x: &'a [f32],
4640 col_field: &[f32],
4641 dtype: TensorDtype,
4642) -> std::borrow::Cow<'a, [f32]> {
4643 if dtype == TensorDtype::Q8_2f {
4644 x.iter().zip(col_field).map(|(a, c)| a * c).collect()
4645 } else {
4646 std::borrow::Cow::Borrowed(x)
4647 }
4648}
4649
4650#[cfg(target_arch = "x86_64")]
4655pub(crate) fn avx2_enabled() -> bool {
4656 use std::sync::OnceLock;
4657 static ON: OnceLock<bool> = OnceLock::new();
4658 *ON.get_or_init(|| {
4659 std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
4660 && std::arch::is_x86_feature_detected!("avx2")
4661 && std::arch::is_x86_feature_detected!("fma")
4662 })
4663}
4664
4665#[cfg(target_arch = "x86_64")]
4670fn avx2_a8w8_enabled() -> bool {
4671 use std::sync::OnceLock;
4672 static ON: OnceLock<bool> = OnceLock::new();
4673 *ON.get_or_init(|| {
4674 avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
4675 })
4676}
4677
4678#[inline]
4682pub(crate) fn a8w8_enabled() -> bool {
4683 #[cfg(target_arch = "aarch64")]
4684 {
4685 sdot_enabled()
4686 }
4687 #[cfg(target_arch = "x86_64")]
4688 {
4689 avx2_a8w8_enabled()
4690 }
4691 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
4692 {
4693 false
4694 }
4695}
4696
4697#[inline]
4700#[allow(unreachable_code)]
4701fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
4702 #[cfg(target_arch = "aarch64")]
4703 unsafe {
4704 return dot_i8_sdot(w, xq);
4705 }
4706 #[cfg(target_arch = "x86_64")]
4707 unsafe {
4708 if avx512vnni_enabled() {
4709 return dot_i8_i8_vnni(w, xq);
4710 }
4711 return dot_i8_i8_avx2(w, xq);
4712 }
4713 w.iter().zip(xq).map(|(&a, &b)| (a as i8) as i32 * b as i32).sum()
4714}
4715
4716#[cfg(target_arch = "x86_64")]
4720fn avx512vnni_enabled() -> bool {
4721 use std::sync::OnceLock;
4722 static ON: OnceLock<bool> = OnceLock::new();
4723 *ON.get_or_init(|| {
4724 std::env::var("CMF_AVX512").map(|v| v != "0").unwrap_or(true)
4725 && std::arch::is_x86_feature_detected!("avx512f")
4726 && std::arch::is_x86_feature_detected!("avx512bw")
4727 && std::arch::is_x86_feature_detected!("avx512vl")
4728 && std::arch::is_x86_feature_detected!("avx512vnni")
4729 })
4730}
4731
4732#[cfg(target_arch = "x86_64")]
4737#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4738unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
4739 unsafe {
4741 use core::arch::x86_64::*;
4742 let n = w.len();
4743 let mut j = 0usize;
4744 let mut total: i32;
4745 {
4750 #[inline(always)]
4751 unsafe fn step(
4752 w: *const u8,
4753 x: *const i8,
4754 acc: core::arch::x86_64::__m512i,
4755 ) -> core::arch::x86_64::__m512i {
4756 unsafe {
4757 use core::arch::x86_64::*;
4758 let wv = _mm512_loadu_si512(w as *const _);
4759 let xv = _mm512_loadu_si512(x as *const _);
4760 let aw = _mm512_abs_epi8(wv);
4761 let neg = _mm512_movepi8_mask(wv);
4762 let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
4763 _mm512_dpbusd_epi32(acc, aw, sx)
4764 }
4765 }
4766 let (mut a0, mut a1, mut a2, mut a3) = (
4767 _mm512_setzero_si512(),
4768 _mm512_setzero_si512(),
4769 _mm512_setzero_si512(),
4770 _mm512_setzero_si512(),
4771 );
4772 while j + 256 <= n {
4773 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
4774 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
4775 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
4776 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
4777 j += 256;
4778 }
4779 while j + 64 <= n {
4780 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
4781 j += 64;
4782 }
4783 let s01 = _mm512_add_epi32(a0, a1);
4784 let s23 = _mm512_add_epi32(a2, a3);
4785 total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
4786 }
4787 if j + 32 <= n {
4789 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
4790 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
4791 let d = _mm256_dpbusd_epi32(
4792 _mm256_setzero_si256(),
4793 _mm256_abs_epi8(wv),
4794 _mm256_sign_epi8(xv, wv),
4795 );
4796 let hi128 = _mm256_extracti128_si256::<1>(d);
4797 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4798 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4799 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4800 total += _mm_cvtsi128_si32(s32);
4801 j += 32;
4802 }
4803 while j < n {
4804 total += (w[j] as i8) as i32 * xq[j] as i32;
4805 j += 1;
4806 }
4807 total
4808 }
4809}
4810
4811#[cfg(target_arch = "x86_64")]
4813#[target_feature(enable = "avx2,fma")]
4814unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
4815 unsafe {
4817 use core::arch::x86_64::*;
4818 let n = x.len();
4819 let wp = w.as_ptr();
4820 let xp = x.as_ptr();
4821 let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
4822 let mut j = 0usize;
4823 while j + 16 <= n {
4824 let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
4825 let lo = _mm256_cvtepi8_epi32(wb);
4826 let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
4827 a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
4828 a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
4829 j += 16;
4830 }
4831 let acc = _mm256_add_ps(a0, a1);
4832 let hi128 = _mm256_extractf128_ps::<1>(acc);
4833 let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
4834 let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
4835 let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
4836 let mut sum = _mm_cvtss_f32(s32);
4837 while j < n {
4838 sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
4839 j += 1;
4840 }
4841 sum
4842 }
4843}
4844
4845#[cfg(target_arch = "x86_64")]
4850#[target_feature(enable = "avx2")]
4851unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
4852 unsafe {
4854 use core::arch::x86_64::*;
4855 let n = w.len();
4856 let ones = _mm256_set1_epi16(1);
4857 let mut acc = _mm256_setzero_si256();
4858 let mut j = 0usize;
4859 while j + 32 <= n {
4860 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
4861 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
4862 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
4863 acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
4864 j += 32;
4865 }
4866 let hi128 = _mm256_extracti128_si256::<1>(acc);
4867 let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
4868 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4869 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4870 let mut s = _mm_cvtsi128_si32(s32);
4871 while j < n {
4872 s += (w[j] as i8) as i32 * xq[j] as i32;
4873 j += 1;
4874 }
4875 s
4876 }
4877}
4878
4879#[cfg(target_arch = "aarch64")]
4883#[target_feature(enable = "neon,i8mm")]
4884unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
4885 unsafe {
4887 use core::arch::aarch64::*;
4888 use core::arch::asm;
4889 let n = w0.len();
4890 let w0p = w0.as_ptr() as *const i8;
4891 let w1p = w1.as_ptr() as *const i8;
4892 let mut acc01 = vdupq_n_s32(0);
4895 let mut acc23 = vdupq_n_s32(0);
4896 let mut i = 0usize;
4897 while i + 8 <= n {
4898 let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
4899 let xb01 =
4900 vcombine_s8(vld1_s8(xs[0].as_ptr().add(i)), vld1_s8(xs[1].as_ptr().add(i)));
4901 let xb23 =
4902 vcombine_s8(vld1_s8(xs[2].as_ptr().add(i)), vld1_s8(xs[3].as_ptr().add(i)));
4903 asm!(
4904 "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
4905 "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
4906 a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
4907 w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
4908 options(pure, nomem, nostack),
4909 );
4910 i += 8;
4911 }
4912 let mut out = [[0i32; 4]; 2];
4913 let a01: [i32; 4] = core::mem::transmute(acc01);
4914 let a23: [i32; 4] = core::mem::transmute(acc23);
4915 out[0][0] = a01[0];
4916 out[0][1] = a01[1];
4917 out[1][0] = a01[2];
4918 out[1][1] = a01[3];
4919 out[0][2] = a23[0];
4920 out[0][3] = a23[1];
4921 out[1][2] = a23[2];
4922 out[1][3] = a23[3];
4923 if i < n {
4924 for (k, x) in xs.iter().enumerate() {
4925 for j in i..n {
4926 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
4927 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
4928 }
4929 }
4930 }
4931 out
4932 }
4933}
4934
4935#[cfg(target_arch = "aarch64")]
4939#[target_feature(enable = "neon,dotprod")]
4940unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
4941 unsafe {
4943 use core::arch::aarch64::*;
4944 use core::arch::asm;
4945 let n = w0.len();
4946 let w0p = w0.as_ptr() as *const i8;
4947 let w1p = w1.as_ptr() as *const i8;
4948 let mut acc = [[vdupq_n_s32(0); 4]; 2];
4949 let mut i = 0usize;
4950 while i + 16 <= n {
4951 let wv0 = vld1q_s8(w0p.add(i));
4952 let wv1 = vld1q_s8(w1p.add(i));
4953 for (k, x) in xs.iter().enumerate() {
4954 let xv = vld1q_s8(x.as_ptr().add(i));
4955 let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
4956 asm!(
4957 "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
4958 "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
4959 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4960 w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
4961 options(pure, nomem, nostack),
4962 );
4963 acc[0][k] = a0;
4964 acc[1][k] = a1;
4965 }
4966 i += 16;
4967 }
4968 let mut out = [[0i32; 4]; 2];
4969 for r in 0..2 {
4970 for k in 0..4 {
4971 out[r][k] = vaddvq_s32(acc[r][k]);
4972 }
4973 }
4974 if i < n {
4975 for (k, x) in xs.iter().enumerate() {
4976 for j in i..n {
4977 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
4978 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
4979 }
4980 }
4981 }
4982 out
4983 }
4984}
4985
4986#[cfg(target_arch = "x86_64")]
4992#[target_feature(enable = "avx2")]
4993unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
4994 unsafe {
4996 use core::arch::x86_64::*;
4997 let n = w0.len();
4998 let ones = _mm256_set1_epi16(1);
4999 let mut acc = [[_mm256_setzero_si256(); 4]; 2];
5000 let mut j = 0usize;
5001 while j + 32 <= n {
5002 let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
5003 let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
5004 let aw0 = _mm256_abs_epi8(wv0);
5005 let aw1 = _mm256_abs_epi8(wv1);
5006 for (k, x) in xs.iter().enumerate() {
5007 let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
5008 let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
5009 acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
5010 let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
5011 acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
5012 }
5013 j += 32;
5014 }
5015 let mut out = [[0i32; 4]; 2];
5016 for r in 0..2 {
5017 for k in 0..4 {
5018 let a = acc[r][k];
5019 let hi128 = _mm256_extracti128_si256::<1>(a);
5020 let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
5021 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5022 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5023 out[r][k] = _mm_cvtsi128_si32(s32);
5024 }
5025 }
5026 if j < n {
5027 for (k, x) in xs.iter().enumerate() {
5028 for i in j..n {
5029 out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
5030 out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
5031 }
5032 }
5033 }
5034 out
5035 }
5036}
5037
5038#[cfg(target_arch = "x86_64")]
5043#[inline]
5044fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
5045 let dot = if avx512vnni_enabled() && row.len() >= 64 {
5046 (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
5047 } else {
5048 unsafe { dot_i8_i8_avx2(row, &act.xq) }
5049 };
5050 let mut acc = dot as f32 * act.sx;
5051 for &(j, xv) in &act.outliers {
5052 acc += (row[j] as i8) as f32 * xv;
5053 }
5054 acc
5055}
5056
5057#[cfg(target_arch = "x86_64")]
5061#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5062unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
5063 unsafe {
5065 use core::arch::x86_64::*;
5066 let n = w.len();
5067 let flip = _mm512_set1_epi8(-128); #[inline(always)]
5069 unsafe fn step(
5070 w: *const u8,
5071 x: *const i8,
5072 flip: core::arch::x86_64::__m512i,
5073 acc: core::arch::x86_64::__m512i,
5074 ) -> core::arch::x86_64::__m512i {
5075 unsafe {
5076 use core::arch::x86_64::*;
5077 let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
5078 _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
5079 }
5080 }
5081 let (mut a0, mut a1, mut a2, mut a3) = (
5082 _mm512_setzero_si512(),
5083 _mm512_setzero_si512(),
5084 _mm512_setzero_si512(),
5085 _mm512_setzero_si512(),
5086 );
5087 let mut j = 0usize;
5088 while j + 256 <= n {
5089 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
5090 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
5091 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
5092 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
5093 j += 256;
5094 }
5095 while j + 64 <= n {
5096 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
5097 j += 64;
5098 }
5099 let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
5100 _mm512_add_epi32(a0, a1),
5101 _mm512_add_epi32(a2, a3),
5102 ));
5103 while j < n {
5105 total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
5106 j += 1;
5107 }
5108 total
5109 }
5110}
5111
5112#[cfg(target_arch = "x86_64")]
5118#[target_feature(enable = "avx2")]
5119unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
5120 unsafe {
5123 use core::arch::x86_64::*;
5124 let lomask = _mm_set1_epi8(0x0F);
5125 let eight = _mm256_set1_epi8(8);
5126 let ones = _mm256_set1_epi16(1);
5127 let mut acc = 0f32;
5128 for gi in 0..gpr {
5129 let g = g0 + gi;
5130 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
5131 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
5132 let lo = _mm_and_si128(b, lomask);
5133 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5134 let w = _mm256_sub_epi8(
5135 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5136 eight,
5137 );
5138 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5139 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5140 let d = _mm256_madd_epi16(p16, ones);
5141 let hi128 = _mm256_extracti128_si256::<1>(d);
5142 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5143 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5144 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5145 acc += _mm_cvtsi128_si32(s32) as f32 * s;
5146 }
5147 acc
5148 }
5149}
5150
5151#[cfg(target_arch = "x86_64")]
5154#[target_feature(enable = "avx2")]
5155unsafe fn dot_q4_row_avx2_2(
5156 packed: &[u8],
5157 scales: &[u8],
5158 g0: usize,
5159 gpr: usize,
5160 xq1: &[i8],
5161 xq2: &[i8],
5162) -> (f32, f32) {
5163 unsafe {
5165 use core::arch::x86_64::*;
5166 let lomask = _mm_set1_epi8(0x0F);
5167 let eight = _mm256_set1_epi8(8);
5168 let ones = _mm256_set1_epi16(1);
5169 let (mut acc1, mut acc2) = (0f32, 0f32);
5170 #[inline(always)]
5171 unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
5172 unsafe {
5173 use core::arch::x86_64::*;
5174 let hi128 = _mm256_extracti128_si256::<1>(d);
5175 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5176 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5177 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5178 _mm_cvtsi128_si32(s32)
5179 }
5180 }
5181 for gi in 0..gpr {
5182 let g = g0 + gi;
5183 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
5184 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
5185 let lo = _mm_and_si128(b, lomask);
5186 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5187 let w = _mm256_sub_epi8(
5188 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5189 eight,
5190 );
5191 let aw = _mm256_abs_epi8(w);
5192 let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5193 let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5194 let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
5195 let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
5196 acc1 += hsum(d1) as f32 * s;
5197 acc2 += hsum(d2) as f32 * s;
5198 }
5199 (acc1, acc2)
5200 }
5201}
5202
5203#[cfg(target_arch = "x86_64")]
5205fn q8_range_avx2(
5206 q: &[u8],
5207 row_scale: &[f32],
5208 act: &SplitAct,
5209 cols: usize,
5210 out_addr: SendMut,
5211 start: usize,
5212 end: usize,
5213) {
5214 for o in start..end {
5215 let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
5216 unsafe { *out_addr.at(o) = v };
5218 }
5219}
5220
5221#[cfg(target_arch = "x86_64")]
5223#[allow(clippy::too_many_arguments)]
5224fn q8_range2_avx2(
5225 q: &[u8],
5226 row_scale: &[f32],
5227 a1: &SplitAct,
5228 a2: &SplitAct,
5229 cols: usize,
5230 p1: SendMut,
5231 p2: SendMut,
5232 start: usize,
5233 end: usize,
5234) {
5235 for o in start..end {
5236 let row = &q[o * cols..(o + 1) * cols];
5237 unsafe {
5239 *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
5240 *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
5241 }
5242 }
5243}
5244
5245#[cfg(target_arch = "aarch64")]
5256fn i8mm_enabled() -> bool {
5257 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5258 *ON.get_or_init(|| {
5259 std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
5260 && std::arch::is_aarch64_feature_detected!("i8mm")
5261 })
5262}
5263
5264#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
5268fn sdot_enabled() -> bool {
5269 use std::sync::OnceLock;
5270 static ON: OnceLock<bool> = OnceLock::new();
5271 *ON.get_or_init(|| {
5272 let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
5273 #[cfg(target_arch = "aarch64")]
5274 {
5275 want && std::arch::is_aarch64_feature_detected!("dotprod")
5276 }
5277 #[cfg(not(target_arch = "aarch64"))]
5278 {
5279 let _ = want;
5280 false
5281 }
5282 })
5283}
5284
5285struct SplitAct {
5290 xq: Vec<i8>,
5291 sx: f32,
5292 outliers: Vec<(usize, f32)>,
5293 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
5296 xsum: i32,
5297}
5298
5299thread_local! {
5300 static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
5303 const { std::cell::RefCell::new(Vec::new()) };
5304}
5305
5306impl Drop for SplitAct {
5307 fn drop(&mut self) {
5308 let buf = std::mem::take(&mut self.xq);
5309 if buf.capacity() > 0 {
5310 XQ_FREE.with(|f| {
5311 let mut f = f.borrow_mut();
5312 if f.len() < 16 {
5313 f.push(buf);
5314 }
5315 });
5316 }
5317 }
5318}
5319
5320fn split_act(x: &[f32]) -> SplitAct {
5321 let n = x.len();
5322 let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
5323 let thr = 8.0 * rms;
5324 let mut outliers: Vec<(usize, f32)> = Vec::new();
5328 let mut amax = 0f32;
5329 for (j, &v) in x.iter().enumerate() {
5330 let a = v.abs();
5331 if a > thr {
5332 outliers.push((j, v));
5333 } else if a > amax {
5334 amax = a;
5335 }
5336 }
5337 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
5338 let inv = 1.0 / sx;
5339 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
5340 xq.clear();
5341 xq.reserve(n);
5342 if outliers.is_empty() {
5343 xq.extend(x.iter().map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8));
5344 } else {
5345 xq.extend(x.iter().map(|&v| {
5347 if v.abs() > thr {
5348 0
5349 } else {
5350 (v * inv).round().clamp(-127.0, 127.0) as i8
5351 }
5352 }));
5353 }
5354 let xsum = xq.iter().map(|&v| v as i32).sum();
5355 SplitAct { xq, sx, outliers, xsum }
5356}
5357
5358#[cfg(target_arch = "aarch64")]
5361#[target_feature(enable = "neon,dotprod")]
5362unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
5363 unsafe {
5365 use core::arch::aarch64::*;
5366 use core::arch::asm;
5367 let wp = w.as_ptr() as *const i8;
5368 let n = w.len();
5369 let (mut a0, mut a1, mut a2, mut a3) =
5370 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
5371 let mut i = 0;
5372 while i + 64 <= n {
5373 let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
5374 let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
5375 let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
5376 let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
5377 asm!(
5378 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
5379 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
5380 "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
5381 "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
5382 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
5383 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
5384 w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
5385 options(pure, nomem, nostack),
5386 );
5387 i += 64;
5388 }
5389 while i + 16 <= n {
5390 let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
5391 asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
5392 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
5393 i += 16;
5394 }
5395 let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
5396 while i < n {
5397 s += (*wp.add(i)) as i32 * xq[i] as i32;
5398 i += 1;
5399 }
5400 s
5401}
5402}
5403
5404#[cfg(target_arch = "aarch64")]
5408#[target_feature(enable = "neon,dotprod")]
5409unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
5410 unsafe {
5412 use core::arch::aarch64::*;
5413 use core::arch::asm;
5414 let n = xq.len();
5415 let px = xq.as_ptr();
5416 let (p0, p1, p2, p3) = (
5417 w0.as_ptr() as *const i8,
5418 w1.as_ptr() as *const i8,
5419 w2.as_ptr() as *const i8,
5420 w3.as_ptr() as *const i8,
5421 );
5422 let (mut a0, mut a1, mut a2, mut a3) =
5423 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
5424 let mut i = 0;
5425 while i + 16 <= n {
5426 let x = vld1q_s8(px.add(i));
5427 let v0 = vld1q_s8(p0.add(i));
5428 let v1 = vld1q_s8(p1.add(i));
5429 let v2 = vld1q_s8(p2.add(i));
5430 let v3 = vld1q_s8(p3.add(i));
5431 asm!(
5432 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
5433 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
5434 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
5435 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
5436 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
5437 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
5438 options(pure, nomem, nostack),
5439 );
5440 i += 16;
5441 }
5442 let mut r = [vaddvq_s32(a0), vaddvq_s32(a1), vaddvq_s32(a2), vaddvq_s32(a3)];
5443 while i < n {
5444 let xi = *px.add(i) as i32;
5445 r[0] += (*p0.add(i)) as i32 * xi;
5446 r[1] += (*p1.add(i)) as i32 * xi;
5447 r[2] += (*p2.add(i)) as i32 * xi;
5448 r[3] += (*p3.add(i)) as i32 * xi;
5449 i += 1;
5450 }
5451 r
5452}
5453}
5454
5455#[cfg(target_arch = "aarch64")]
5462#[target_feature(enable = "neon,dotprod")]
5463unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
5464 unsafe {
5467 use core::arch::aarch64::*;
5468 use core::arch::asm;
5469 let n = xq.len();
5470 let px = xq.as_ptr();
5471 let pg = g.as_ptr() as *const i8;
5472 let (mut a0, mut a1, mut a2, mut a3) =
5473 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
5474 let mut i = 0;
5475 while i + 16 <= n {
5476 let x = vld1q_s8(px.add(i));
5477 let base = pg.add(4 * i);
5478 let v0 = vld1q_s8(base);
5479 let v1 = vld1q_s8(base.add(16));
5480 let v2 = vld1q_s8(base.add(32));
5481 let v3 = vld1q_s8(base.add(48));
5482 asm!(
5483 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
5484 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
5485 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
5486 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
5487 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
5488 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
5489 options(pure, nomem, nostack),
5490 );
5491 i += 16;
5492 }
5493 [vaddvq_s32(a0), vaddvq_s32(a1), vaddvq_s32(a2), vaddvq_s32(a3)]
5494 }
5495}
5496
5497#[cfg(target_arch = "aarch64")]
5503fn q8_range_sdot(
5504 q: &[u8],
5505 rep: &[u8],
5506 row_scale: &[f32],
5507 act: &SplitAct,
5508 cols: usize,
5509 out_addr: SendMut,
5510 start: usize,
5511 end: usize,
5512) {
5513 let mut o = start;
5514 if !rep.is_empty() {
5517 while o < end && o % 4 != 0 {
5518 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
5519 unsafe { *out_addr.at(o) = v };
5520 o += 1;
5521 }
5522 }
5523 while o + 4 <= end {
5524 let r = if rep.is_empty() {
5525 unsafe {
5526 dot_i8_sdot_4rows(
5527 &q[o * cols..(o + 1) * cols],
5528 &q[(o + 1) * cols..(o + 2) * cols],
5529 &q[(o + 2) * cols..(o + 3) * cols],
5530 &q[(o + 3) * cols..(o + 4) * cols],
5531 &act.xq,
5532 )
5533 }
5534 } else {
5535 unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
5536 };
5537 for k in 0..4 {
5538 let mut acc = r[k] as f32 * act.sx;
5539 for &(j, xv) in &act.outliers {
5540 acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
5541 }
5542 unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
5544 }
5545 o += 4;
5546 }
5547 while o < end {
5548 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
5549 unsafe { *out_addr.at(o) = v };
5550 o += 1;
5551 }
5552}
5553
5554#[cfg(target_arch = "aarch64")]
5557#[allow(clippy::too_many_arguments)]
5558fn q8_range2_sdot(
5559 q: &[u8],
5560 row_scale: &[f32],
5561 a1: &SplitAct,
5562 a2: &SplitAct,
5563 cols: usize,
5564 p1: SendMut,
5565 p2: SendMut,
5566 start: usize,
5567 end: usize,
5568) {
5569 for o in start..end {
5570 let row = &q[o * cols..(o + 1) * cols];
5571 unsafe {
5573 *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
5574 *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
5575 }
5576 }
5577}
5578
5579#[allow(clippy::too_many_arguments)]
5581fn q8_range2_f32(
5582 q: &[u8],
5583 row_scale: &[f32],
5584 x1: &[f32],
5585 x2: &[f32],
5586 cols: usize,
5587 p1: SendMut,
5588 p2: SendMut,
5589 start: usize,
5590 end: usize,
5591) {
5592 for o in start..end {
5593 let row = &q[o * cols..(o + 1) * cols];
5594 unsafe {
5596 *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
5597 *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
5598 }
5599 }
5600}
5601
5602fn q8_range_f32(
5604 q: &[u8],
5605 row_scale: &[f32],
5606 xs: &[f32],
5607 cols: usize,
5608 out_addr: SendMut,
5609 start: usize,
5610 end: usize,
5611) {
5612 for o in start..end {
5613 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
5614 unsafe { *out_addr.at(o) = v };
5616 }
5617}
5618
5619#[cfg(target_arch = "aarch64")]
5622#[inline]
5623fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
5624 let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
5625 for &(j, xv) in &act.outliers {
5626 acc += (row[j] as i8) as f32 * xv;
5627 }
5628 acc
5629}
5630
5631#[cfg(target_arch = "aarch64")]
5639#[target_feature(enable = "neon,dotprod")]
5640unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
5641 unsafe {
5644 use core::arch::aarch64::*;
5645 use core::arch::asm;
5646 let lomask = vdupq_n_u8(0x0F);
5647 let eight = vdupq_n_s8(8);
5648 let mut acc = 0f32;
5649 for gi in 0..gpr {
5650 let g = g0 + gi;
5651 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
5652 let b = vld1q_u8(packed.as_ptr().add(g * 16));
5653 let lo = vandq_u8(b, lomask);
5654 let hi = vshrq_n_u8::<4>(b);
5655 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5656 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5657 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5658 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5659 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5660 asm!(
5661 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5662 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5663 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5664 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5665 options(pure, nomem, nostack),
5666 );
5667 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5668 }
5669 acc
5670 }
5671}
5672
5673#[cfg(target_arch = "aarch64")]
5678#[target_feature(enable = "neon,dotprod")]
5679unsafe fn dot_q4_row_sdot2(
5680 packed: &[u8],
5681 scales: &[u8],
5682 g0: usize,
5683 gpr: usize,
5684 xq1: &[i8],
5685 xq2: &[i8],
5686) -> (f32, f32) {
5687 unsafe {
5690 use core::arch::aarch64::*;
5691 use core::arch::asm;
5692 let lomask = vdupq_n_u8(0x0F);
5693 let eight = vdupq_n_s8(8);
5694 let (mut acc1, mut acc2) = (0f32, 0f32);
5695 for gi in 0..gpr {
5696 let g = g0 + gi;
5697 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
5698 let b = vld1q_u8(packed.as_ptr().add(g * 16));
5699 let lo = vandq_u8(b, lomask);
5700 let hi = vshrq_n_u8::<4>(b);
5701 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5702 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5703 let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
5704 let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
5705 let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
5706 let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
5707 let (mut a0, mut a1, mut b0, mut b1) =
5708 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
5709 asm!(
5710 "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
5711 "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
5712 "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
5713 "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
5714 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5715 b0 = inout(vreg) b0, b1 = inout(vreg) b1,
5716 e0 = in(vreg) e0, e1 = in(vreg) e1,
5717 x10 = in(vreg) x10, x11 = in(vreg) x11,
5718 x20 = in(vreg) x20, x21 = in(vreg) x21,
5719 options(pure, nomem, nostack),
5720 );
5721 acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5722 acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
5723 }
5724 (acc1, acc2)
5725 }
5726}
5727
5728#[inline]
5733pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
5734 #[cfg(target_arch = "aarch64")]
5735 unsafe {
5736 return axpy_i8_f32_neon(acc, row, w);
5737 }
5738 #[cfg(target_arch = "x86_64")]
5739 if avx2_enabled() {
5740 return unsafe { axpy_i8_f32_avx2(acc, row, w) };
5741 }
5742 #[allow(unreachable_code)]
5743 {
5744 for (a, &b) in acc.iter_mut().zip(row) {
5745 *a += w * b as f32;
5746 }
5747 }
5748}
5749
5750#[cfg(target_arch = "x86_64")]
5752#[target_feature(enable = "avx2,fma")]
5753unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
5754 unsafe {
5756 use core::arch::x86_64::*;
5757 let n = acc.len().min(row.len());
5758 let ap = acc.as_mut_ptr();
5759 let rp = row.as_ptr();
5760 let wv = _mm256_set1_ps(w);
5761 let mut j = 0usize;
5762 while j + 16 <= n {
5763 let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
5764 let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
5765 let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
5766 let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
5767 let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
5768 _mm256_storeu_ps(ap.add(j), v0);
5769 _mm256_storeu_ps(ap.add(j + 8), v1);
5770 j += 16;
5771 }
5772 while j < n {
5773 *ap.add(j) += w * (*rp.add(j)) as f32;
5774 j += 1;
5775 }
5776 }
5777}
5778
5779#[cfg(target_arch = "aarch64")]
5780#[target_feature(enable = "neon")]
5781unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
5782 unsafe {
5784 use core::arch::aarch64::*;
5785 let n = acc.len().min(row.len());
5786 let ap = acc.as_mut_ptr();
5787 let rp = row.as_ptr();
5788 let wv = vdupq_n_f32(w);
5789 let mut j = 0usize;
5790 while j + 16 <= n {
5791 let rb = vld1q_s8(rp.add(j));
5792 let lo = vmovl_s8(vget_low_s8(rb));
5793 let hi = vmovl_s8(vget_high_s8(rb));
5794 for (off, half) in [(0, lo), (8, hi)] {
5795 let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
5796 let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
5797 let o = j + off;
5798 vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
5799 vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
5800 }
5801 j += 16;
5802 }
5803 while j < n {
5804 *ap.add(j) += w * (*rp.add(j)) as f32;
5805 j += 1;
5806 }
5807}
5808}
5809
5810#[inline]
5813pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
5814 #[cfg(target_arch = "aarch64")]
5815 unsafe {
5816 return dot_i8_f32_neon(w, x);
5817 }
5818 #[cfg(target_arch = "x86_64")]
5819 if avx2_enabled() {
5820 return unsafe { dot_i8_f32_avx2(w, x) };
5821 }
5822 #[allow(unreachable_code)]
5823 {
5824 let mut sum = 0.0f32;
5825 for (j, &b) in w.iter().enumerate() {
5826 sum += (b as i8) as f32 * x[j];
5827 }
5828 sum
5829 }
5830}
5831
5832#[inline]
5836fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
5837 #[cfg(target_arch = "aarch64")]
5838 unsafe {
5839 return dot_i8_col_f32_neon(w, x, col);
5840 }
5841 #[allow(unreachable_code)]
5842 {
5843 let mut sum = 0.0f32;
5844 for (j, &b) in w.iter().enumerate() {
5845 sum += (b as i8) as f32 * x[j] * col[j];
5846 }
5847 sum
5848 }
5849}
5850
5851#[cfg(target_arch = "aarch64")]
5852#[target_feature(enable = "neon")]
5853unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
5854 unsafe {
5856 use core::arch::aarch64::*;
5857 let n = x.len();
5858 let wp = w.as_ptr() as *const i8;
5859 let xp = x.as_ptr();
5860 let cp = col.as_ptr();
5861 let (mut a0, mut a1, mut a2, mut a3) =
5862 (vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0));
5863 let mut j = 0usize;
5864 while j + 16 <= n {
5865 let wb = vld1q_s8(wp.add(j));
5866 let lo = vmovl_s8(vget_low_s8(wb));
5867 let hi = vmovl_s8(vget_high_s8(wb));
5868 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
5869 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
5870 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
5871 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
5872 a0 = vfmaq_f32(a0, w0, vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))));
5873 a1 = vfmaq_f32(a1, w1, vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))));
5874 a2 = vfmaq_f32(a2, w2, vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))));
5875 a3 = vfmaq_f32(a3, w3, vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))));
5876 j += 16;
5877 }
5878 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
5879 while j < n {
5880 sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
5881 j += 1;
5882 }
5883 sum
5884}
5885}
5886
5887#[cfg(target_arch = "aarch64")]
5888#[target_feature(enable = "neon")]
5889unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
5890 unsafe {
5892 use core::arch::aarch64::*;
5893 let n = x.len();
5894 let wp = w.as_ptr() as *const i8;
5895 let xp = x.as_ptr();
5896 let (mut a0, mut a1, mut a2, mut a3) =
5897 (vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0));
5898 let mut j = 0usize;
5899 while j + 16 <= n {
5900 let wb = vld1q_s8(wp.add(j));
5901 let lo = vmovl_s8(vget_low_s8(wb));
5902 let hi = vmovl_s8(vget_high_s8(wb));
5903 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
5904 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
5905 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
5906 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
5907 a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
5908 a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
5909 a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
5910 a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
5911 j += 16;
5912 }
5913 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
5914 while j < n {
5915 sum += (*wp.add(j)) as f32 * *xp.add(j);
5916 j += 1;
5917 }
5918 sum
5919}
5920}
5921
5922#[allow(clippy::too_many_arguments)]
5923fn qmatvec(
5924 q: &[u8],
5925 rep: &[u8],
5926 row_scale: &[f32],
5927 xs: &[f32],
5928 rows: usize,
5929 cols: usize,
5930 out: &mut [f32],
5931 pool: Option<&Pool>,
5932) {
5933 debug_assert_eq!(out.len(), rows);
5934 #[cfg(not(target_arch = "aarch64"))]
5935 let _ = rep;
5936
5937 #[cfg(target_arch = "aarch64")]
5938 if sdot_enabled() {
5939 let act = split_act(xs);
5940 let out_addr = SendMut(out.as_mut_ptr());
5941 let run_range = |start: usize, end: usize| {
5942 q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
5943 };
5944 match pool {
5945 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
5946 _ => run_range(0, rows),
5947 }
5948 return;
5949 }
5950 #[cfg(target_arch = "x86_64")]
5953 if avx2_a8w8_enabled() {
5954 let act = split_act(xs);
5955 let out_addr = SendMut(out.as_mut_ptr());
5956 let run_range =
5957 |start: usize, end: usize| q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end);
5958 match pool {
5959 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
5960 _ => run_range(0, rows),
5961 }
5962 return;
5963 }
5964
5965 let row_dot = |o: usize| -> f32 { dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o] };
5966 match pool {
5967 Some(pool) if rows >= 256 => {
5968 let out_addr = SendMut(out.as_mut_ptr());
5969 pool.run(&move |widx, n| {
5970 let chunk = rows.div_ceil(n);
5971 let start = widx * chunk;
5972 let end = (start + chunk).min(rows);
5973 for o in start..end {
5974 unsafe { *out_addr.at(o) = row_dot(o) };
5976 }
5977 });
5978 }
5979 _ => {
5980 for (o, dst) in out.iter_mut().enumerate() {
5981 *dst = row_dot(o);
5982 }
5983 }
5984 }
5985}
5986
5987#[allow(clippy::too_many_arguments)]
5988fn qmatvec2(
5989 q: &[u8],
5990 row_scale: &[f32],
5991 x1: &[f32],
5992 x2: &[f32],
5993 rows: usize,
5994 cols: usize,
5995 o1: &mut [f32],
5996 o2: &mut [f32],
5997 pool: Option<&Pool>,
5998) {
5999 #[cfg(target_arch = "aarch64")]
6000 if sdot_enabled() {
6001 let a1s = split_act(x1);
6002 let a2s = split_act(x2);
6003 let p1 = SendMut(o1.as_mut_ptr());
6004 let p2 = SendMut(o2.as_mut_ptr());
6005 let run_range = |start: usize, end: usize| {
6006 q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
6007 };
6008 match pool {
6009 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
6010 _ => run_range(0, rows),
6011 }
6012 return;
6013 }
6014 #[cfg(target_arch = "x86_64")]
6015 if avx2_a8w8_enabled() {
6016 let a1s = split_act(x1);
6017 let a2s = split_act(x2);
6018 let p1 = SendMut(o1.as_mut_ptr());
6019 let p2 = SendMut(o2.as_mut_ptr());
6020 let run_range = |start: usize, end: usize| {
6021 q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
6022 };
6023 match pool {
6024 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
6025 _ => run_range(0, rows),
6026 }
6027 return;
6028 }
6029
6030 let row_dots = |o: usize| -> (f32, f32) {
6031 let row = &q[o * cols..(o + 1) * cols];
6032 (dot_i8_f32(row, x1) * row_scale[o], dot_i8_f32(row, x2) * row_scale[o])
6033 };
6034 match pool {
6035 Some(pool) if rows >= 256 => {
6036 let p1 = SendMut(o1.as_mut_ptr());
6037 let p2 = SendMut(o2.as_mut_ptr());
6038 pool.run(&move |widx, n| {
6039 let chunk = rows.div_ceil(n);
6040 let start = widx * chunk;
6041 let end = (start + chunk).min(rows);
6042 for o in start..end {
6043 let (s1, s2) = row_dots(o);
6044 unsafe {
6046 *p1.at(o) = s1;
6047 *p2.at(o) = s2;
6048 }
6049 }
6050 });
6051 }
6052 _ => {
6053 for o in 0..rows {
6054 let (s1, s2) = row_dots(o);
6055 o1[o] = s1;
6056 o2[o] = s2;
6057 }
6058 }
6059 }
6060}
6061
6062#[derive(Clone, Copy)]
6063struct SendMut(*mut f32);
6064unsafe impl Send for SendMut {}
6065unsafe impl Sync for SendMut {}
6066
6067impl SendMut {
6068 #[inline]
6069 fn at(self, i: usize) -> *mut f32 {
6070 unsafe { self.0.add(i) }
6071 }
6072}
6073
6074#[cfg(test)]
6075mod tests {
6076 use super::*;
6077
6078 #[test]
6079 fn f32_matvec_matches_matvec_rows_bitexact() {
6080 let (rows, cols) = (300, 40);
6081 let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
6082 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
6083 let qt = QTensor::from_f32(w.clone(), rows, cols);
6084
6085 let mut a = vec![0.0f32; rows];
6086 matvec_rows(None, &w, &x, &mut a);
6087 let mut b = vec![0.0f32; rows];
6088 qt.matvec(&x, &mut b, None);
6089 assert_eq!(a, b);
6090 }
6091
6092 #[test]
6093 fn sdot_kernel_exact_on_grid() {
6094 eprintln!("sdot_enabled = {}", sdot_enabled());
6099 let (rows, cols) = (9, 80); let w: Vec<u8> = (0..rows * cols)
6101 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
6102 .collect();
6103 let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
6104 let x: Vec<f32> = (0..cols)
6105 .map(|i| match i % 3 {
6106 0 => 1.0,
6107 1 => -1.0,
6108 _ => 0.0,
6109 })
6110 .collect();
6111 let mut a = vec![0.0f32; rows];
6112 qmatvec(&w, &[], &scales, &x, rows, cols, &mut a, None);
6113 for o in 0..rows {
6114 let mut acc = 0.0f32;
6115 for j in 0..cols {
6116 acc += (w[o * cols + j] as i8) as f32 * x[j];
6117 }
6118 let expect = acc * scales[o];
6119 assert!(
6120 (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
6121 "row {o}: {} vs {expect}",
6122 a[o]
6123 );
6124 }
6125 }
6126
6127 #[test]
6128 fn q1_tbl_fast_path_matches_reference() {
6129 let (rows, cols) = (5, 256);
6134 let gpr = cols / GROUP_SIZE;
6135 let mut bytes = Vec::new();
6136 for t in 0..rows * gpr {
6137 let s = 0.007 + (t % 11) as f32 * 0.004;
6138 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
6139 for j in 0..4 {
6140 bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
6141 }
6142 }
6143 let x: Vec<f32> = (0..cols)
6144 .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
6145 .collect();
6146 let mut w = vec![0.0f32; rows * cols];
6147 cortiq_core::quant::dequant_q1(&bytes, &mut w);
6148 let mut got = vec![0.0f32; rows];
6149 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
6150 for o in 0..rows {
6151 let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
6152 assert!(
6153 (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
6154 "row {o}: {} vs {expect}",
6155 got[o]
6156 );
6157 }
6158 let b = 5usize;
6161 let mut xs_all = Vec::new();
6162 for bi in 0..b {
6163 xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
6164 }
6165 let mut mm = vec![0.0f32; b * rows];
6166 q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
6167 for bi in 0..b {
6168 let mut single = vec![0.0f32; rows];
6169 q1_matvec(&bytes, &xs_all[bi * cols..(bi + 1) * cols], rows, cols, &mut single, None);
6170 assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
6171 }
6172 }
6173
6174 #[test]
6175 fn q1_kernels_match_exact_reference() {
6176 let (rows, cols) = (7, 96);
6178 let gpr = cols / GROUP_SIZE;
6179 let mut bytes = Vec::new();
6180 for t in 0..rows * gpr {
6181 let s = 0.01 + (t % 13) as f32 * 0.003;
6182 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
6183 for j in 0..4 {
6184 bytes.push(((t * 31 + j * 97) % 251) as u8);
6185 }
6186 }
6187 let x: Vec<f32> = (0..cols)
6189 .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
6190 .collect();
6191 let mut w = vec![0.0f32; rows * cols];
6193 cortiq_core::quant::dequant_q1(&bytes, &mut w);
6194 let mut expect = vec![0.0f32; rows];
6195 for o in 0..rows {
6196 expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
6197 }
6198 let mut got = vec![0.0f32; rows];
6199 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
6200 for o in 0..rows {
6201 assert!(
6202 (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
6203 "row {o}: {} vs {}",
6204 got[o],
6205 expect[o]
6206 );
6207 }
6208 let x2: Vec<f32> = x.iter().map(|v| -v).collect();
6210 let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
6211 q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
6212 assert_eq!(a1, got);
6213 let mut xs = x.clone();
6214 xs.extend_from_slice(&x2);
6215 let mut mm = vec![0.0f32; 2 * rows];
6216 q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
6217 assert_eq!(&mm[..rows], got.as_slice());
6218 assert_eq!(&mm[rows..], a2.as_slice());
6219 }
6220
6221 #[test]
6222 fn repack_is_bit_identical() {
6223 let (rows, cols) = (267, 96); let w: Vec<u8> = (0..rows * cols)
6229 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
6230 .collect();
6231 let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
6232 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
6233 let rep = q8_repack_layout(&w, rows, cols);
6234 for g in 0..rows / 4 {
6236 for c in 0..cols / 16 {
6237 for lane in 0..4 {
6238 assert_eq!(
6239 &rep[g * 4 * cols + c * 64 + lane * 16..g * 4 * cols + c * 64 + lane * 16 + 16],
6240 &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
6241 );
6242 }
6243 }
6244 }
6245 let mut a = vec![0.0f32; rows];
6246 qmatvec(&w, &[], &scales, &x, rows, cols, &mut a, None);
6247 let mut b = vec![0.0f32; rows];
6248 qmatvec(&w, &rep, &scales, &x, rows, cols, &mut b, None);
6249 assert_eq!(a, b, "full-range repack output diverged");
6250
6251 #[cfg(target_arch = "aarch64")]
6252 if sdot_enabled() {
6253 let act = split_act(&x);
6255 let mut c1 = vec![0.0f32; rows];
6256 let mut c2 = vec![0.0f32; rows];
6257 q8_range_sdot(&w, &[], &scales, &act, cols, SendMut(c1.as_mut_ptr()), 3, rows - 2);
6258 q8_range_sdot(&w, &rep, &scales, &act, cols, SendMut(c2.as_mut_ptr()), 3, rows - 2);
6259 assert_eq!(c1, c2, "unaligned-range repack output diverged");
6260 }
6261 }
6262
6263 #[test]
6264 fn sdot_a8w8_noise_is_bounded() {
6265 let (rows, cols) = (16, 512);
6269 let w: Vec<u8> = (0..rows * cols)
6270 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
6271 .collect();
6272 let scales = vec![0.01f32; rows];
6273 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
6274 let mut a = vec![0.0f32; rows];
6275 qmatvec(&w, &[], &scales, &x, rows, cols, &mut a, None);
6276 let (mut num, mut den) = (0f64, 0f64);
6277 for o in 0..rows {
6278 let mut acc = 0.0f32;
6279 for j in 0..cols {
6280 acc += (w[o * cols + j] as i8) as f32 * x[j];
6281 }
6282 let expect = acc * scales[o];
6283 num += ((a[o] - expect) as f64).powi(2);
6284 den += (expect as f64).powi(2);
6285 }
6286 let rel = (num / den.max(1e-12)).sqrt();
6287 assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
6288 }
6289
6290 #[test]
6291 fn i8_dot_neon_matches_scalar() {
6292 let n = 100;
6293 let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
6294 let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
6295 let mut scalar = 0.0f32;
6296 for j in 0..n {
6297 scalar += (w[j] as i8) as f32 * x[j];
6298 }
6299 let fast = dot_i8_f32(&w, &x);
6300 assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
6301 }
6302
6303 #[test]
6305 fn vbitmatvec_matches_full_dequant() {
6306 let (rows, cols) = (6, 64);
6307 let ng = cols / GROUP_SIZE;
6308 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
6310 let mut bytes = bits.clone();
6311 for g in 0..rows * ng {
6312 let s = 0.02 + 0.001 * g as f32;
6313 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
6314 }
6315 for r in 0..rows {
6316 let b = bits[r] as usize;
6317 let (mut acc, mut nb) = (0u64, 0usize);
6318 let mut rowbytes = Vec::new();
6319 for i in 0..cols {
6320 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
6321 acc = (acc << b) | v;
6322 nb += b;
6323 while nb >= 8 {
6324 nb -= 8;
6325 rowbytes.push(((acc >> nb) & 0xFF) as u8);
6326 }
6327 }
6328 if nb > 0 {
6329 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
6330 }
6331 bytes.extend_from_slice(&rowbytes);
6332 }
6333 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
6334
6335 let mut reference = vec![0f32; rows * cols];
6336 cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
6337 let mut expect = vec![0f32; rows];
6338 for r in 0..rows {
6339 expect[r] = reference[r * cols..(r + 1) * cols]
6340 .iter()
6341 .zip(&x)
6342 .map(|(w, xv)| w * xv)
6343 .sum();
6344 }
6345 let mut got = vec![0f32; rows];
6346 let offsets = vbit_row_offsets(&bytes, rows, cols);
6347 vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
6348 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
6352 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
6353 for r in 0..rows {
6354 assert!(
6355 (got[r] - expect[r]).abs() < tol * scale,
6356 "row {r}: {} vs {}",
6357 got[r],
6358 expect[r]
6359 );
6360 }
6361 }
6362
6363 #[test]
6368 #[cfg(target_arch = "x86_64")]
6369 fn vbit_matmat_blocked_matches_per_row() {
6370 let (rows, cols, b) = (64usize, 128usize, 9usize);
6371 let ng = cols / GROUP_SIZE;
6372 let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
6373 let mut bytes = bits.clone();
6374 for g in 0..rows * ng {
6375 let sc = 0.02 + 0.0005 * g as f32;
6376 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
6377 }
6378 for r in 0..rows {
6379 let bw = bits[r] as usize;
6380 let (mut acc, mut nb) = (0u64, 0usize);
6381 let mut rowbytes = Vec::new();
6382 for i in 0..cols {
6383 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
6384 acc = (acc << bw) | v;
6385 nb += bw;
6386 while nb >= 8 {
6387 nb -= 8;
6388 rowbytes.push(((acc >> nb) & 0xFF) as u8);
6389 }
6390 }
6391 if nb > 0 {
6392 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
6393 }
6394 bytes.extend_from_slice(&rowbytes);
6395 }
6396 let x: Vec<f32> =
6397 (0..b * cols).map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5).collect();
6398 let offsets = vbit_row_offsets(&bytes, rows, cols);
6399 let mut y_a = vec![0f32; b * rows];
6400 let mut y_b = vec![0f32; b * rows];
6401 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
6402 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
6403 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
6404 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
6405 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
6406 let max_d =
6407 y_a.iter().zip(&y_b).map(|(p, q)| (p - q).abs()).fold(0.0f32, f32::max);
6408 assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
6409 }
6410
6411 #[test]
6412 fn q4matvec_matches_full_dequant() {
6413 let (rows, cols) = (8, 64);
6414 let groups = rows * cols / GROUP_SIZE;
6415 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
6417 for i in 0..groups * 16 {
6418 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
6419 }
6420 for g in 0..groups {
6421 let s = 0.01 + 0.003 * g as f32;
6422 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
6423 }
6424 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
6425
6426 let mut reference = vec![0.0f32; rows * cols];
6427 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
6428 let mut expect = vec![0.0f32; rows];
6429 for r in 0..rows {
6430 expect[r] = reference[r * cols..(r + 1) * cols]
6431 .iter()
6432 .zip(&x)
6433 .map(|(w, xv)| w * xv)
6434 .sum();
6435 }
6436
6437 let mut got = vec![0.0f32; rows];
6438 q4matvec(&bytes, &x, rows, cols, &mut got, None);
6439 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
6443 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
6444 for r in 0..rows {
6445 assert!(
6446 (got[r] - expect[r]).abs() < tol * scale,
6447 "row {r}: {} vs {}",
6448 got[r],
6449 expect[r]
6450 );
6451 }
6452 }
6453
6454 #[test]
6457 fn vbitmatvec2_equals_two_singles() {
6458 let (rows, cols) = (6, 64);
6459 let ng = cols / GROUP_SIZE;
6460 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
6461 let mut bytes = bits.clone();
6462 for g in 0..rows * ng {
6463 let s = 0.02 + 0.001 * g as f32;
6464 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
6465 }
6466 for r in 0..rows {
6467 let b = bits[r] as usize;
6468 let (mut acc, mut nb) = (0u64, 0usize);
6469 let mut rowbytes = Vec::new();
6470 for i in 0..cols {
6471 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
6472 acc = (acc << b) | v;
6473 nb += b;
6474 while nb >= 8 {
6475 nb -= 8;
6476 rowbytes.push(((acc >> nb) & 0xFF) as u8);
6477 }
6478 }
6479 if nb > 0 {
6480 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
6481 }
6482 bytes.extend_from_slice(&rowbytes);
6483 }
6484 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
6485 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
6486 let offsets = vbit_row_offsets(&bytes, rows, cols);
6487
6488 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
6489 vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
6490 vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
6491 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
6492 vbitmatvec2(&bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
6493 assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
6494 assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
6495 }
6496
6497 #[test]
6499 fn q4matvec2_equals_two_singles() {
6500 let (rows, cols) = (8, 128);
6501 let groups = rows * cols / GROUP_SIZE;
6502 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
6503 for i in 0..groups * 16 {
6504 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
6505 }
6506 for g in 0..groups {
6507 let s = 0.01 + 0.003 * g as f32;
6508 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
6509 }
6510 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
6513 x1[9] = 250.0;
6514 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
6515
6516 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
6517 q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
6518 q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
6519 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
6520 q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
6521 assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
6522 assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
6523 }
6524
6525 #[test]
6528 fn matvec_many_equals_separate_matvecs() {
6529 use crate::pool::Pool;
6530 let (r1, r2, cols) = (300, 200, 64);
6531 let mk = |salt: usize, rows: usize| {
6532 QTensor::from_f32(
6533 (0..rows * cols).map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5).collect(),
6534 rows,
6535 cols,
6536 )
6537 };
6538 let (a, b) = (mk(1, r1), mk(5, r2));
6539 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
6540 let pool = Pool::new(3);
6541
6542 let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
6543 a.matvec(&x, &mut ea, Some(&pool));
6544 b.matvec(&x, &mut eb, Some(&pool));
6545 let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
6546 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
6547 assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
6548 assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
6549 }
6550
6551 #[test]
6554 fn batched_matmat_equals_per_position_matvec() {
6555 let (rows, cols, b) = (8, 64, 5);
6556 let groups = rows * cols / GROUP_SIZE;
6558 let mut q4 = Vec::new();
6559 for i in 0..groups * 16 {
6560 q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
6561 }
6562 for g in 0..groups {
6563 q4.extend_from_slice(
6564 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
6565 );
6566 }
6567 let ng = cols / GROUP_SIZE;
6569 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
6570 let mut vb = bits.clone();
6571 for g in 0..rows * ng {
6572 vb.extend_from_slice(
6573 &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
6574 );
6575 }
6576 for r in 0..rows {
6577 let bw = bits[r] as usize;
6578 let (mut acc, mut nb) = (0u64, 0usize);
6579 let mut rowbytes = Vec::new();
6580 for i in 0..cols {
6581 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
6582 acc = (acc << bw) | v;
6583 nb += bw;
6584 while nb >= 8 {
6585 nb -= 8;
6586 rowbytes.push(((acc >> nb) & 0xFF) as u8);
6587 }
6588 }
6589 if nb > 0 {
6590 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
6591 }
6592 vb.extend_from_slice(&rowbytes);
6593 }
6594 let offsets = vbit_row_offsets(&vb, rows, cols);
6595
6596 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
6597
6598 let mut got = vec![0f32; b * rows];
6600 q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
6601 for bi in 0..b {
6602 let mut expect = vec![0f32; rows];
6603 q4matvec(&q4, &xs[bi * cols..(bi + 1) * cols], rows, cols, &mut expect, None);
6604 assert_eq!(&got[bi * rows..(bi + 1) * rows], &expect[..], "q4 batch pos {bi}");
6605 }
6606
6607 let mut got = vec![0f32; b * rows];
6609 vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
6610 for bi in 0..b {
6611 let mut expect = vec![0f32; rows];
6612 vbitmatvec(
6613 &vb, &offsets, &xs[bi * cols..(bi + 1) * cols], rows, cols, &mut expect, None,
6614 );
6615 assert_eq!(&got[bi * rows..(bi + 1) * rows], &expect[..], "vbit batch pos {bi}");
6616 }
6617 }
6618
6619 #[test]
6623 fn q4_tiled_matches_q4_block_bitexact() {
6624 let (rows, cols, b) = (8usize, 128usize, 3usize);
6625 let groups = rows * cols / GROUP_SIZE;
6626 let mut split = Vec::with_capacity(groups * 18);
6627 for i in 0..groups * 16 {
6628 split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
6629 }
6630 for g in 0..groups {
6631 split.extend_from_slice(
6632 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
6633 );
6634 }
6635 let (packed, scales) = split.split_at(groups * 16);
6637 let mut tiled = Vec::with_capacity(groups * Q4_TILE);
6638 for g in 0..groups {
6639 tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
6640 tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
6641 }
6642
6643 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
6644 x1[9] = 250.0; let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
6646
6647 let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
6648 q4matvec(&split, &x1, rows, cols, &mut a, None);
6649 q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
6650 assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
6651
6652 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
6653 let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
6654 q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
6655 q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
6656 assert_eq!(a1, t1);
6657 assert_eq!(a2, t2);
6658
6659 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
6660 let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
6661 q4matmat(&split, &xs, b, rows, cols, &mut am, None);
6662 q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
6663 assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
6664 }
6665
6666 #[test]
6673 fn q4matvec_sdot_outlier_exact() {
6674 let (rows, cols) = (4, 128);
6675 let groups = rows * cols / GROUP_SIZE;
6676 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
6677 for i in 0..groups * 16 {
6678 bytes.push(((i * 11 + 5) % 256) as u8);
6679 }
6680 for g in 0..groups {
6681 let s = 0.02 + 0.002 * g as f32;
6682 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
6683 }
6684 let mut x: Vec<f32> = (0..cols)
6685 .map(|i| match i % 3 {
6686 0 => 1.0,
6687 1 => -1.0,
6688 _ => 0.0,
6689 })
6690 .collect();
6691 x[17] = 300.0; let mut reference = vec![0.0f32; rows * cols];
6694 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
6695 let mut expect = vec![0.0f32; rows];
6696 for r in 0..rows {
6697 expect[r] = reference[r * cols..(r + 1) * cols]
6698 .iter()
6699 .zip(&x)
6700 .map(|(w, xv)| w * xv)
6701 .sum();
6702 }
6703 let mut got = vec![0.0f32; rows];
6704 q4matvec(&bytes, &x, rows, cols, &mut got, None);
6705 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
6706 for r in 0..rows {
6707 assert!(
6708 (got[r] - expect[r]).abs() < 2e-3 * scale,
6709 "row {r}: {} vs {} (outlier term must be exact)",
6710 got[r],
6711 expect[r]
6712 );
6713 }
6714 }
6715
6716 #[test]
6720 fn q1t_matvec_matches_reference() {
6721 use cortiq_core::quant::{dequant_q1t, f32_to_f16};
6722 let (rows, cols) = (3usize, 64usize); let gpr = cols / GROUP_SIZE;
6724 let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
6725 let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
6727 let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
6728 let mut bytes = Vec::new();
6729 for r in 0..rows {
6730 for g in 0..gpr {
6731 bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
6732 let mut c = [0u8; 7];
6733 for k in 0..GROUP_SIZE {
6734 let code = if is_out(r * cols + g * GROUP_SIZE + k) {
6736 0
6737 } else {
6738 ((k + r * 3 + g) % 3) as u8 };
6740 cortiq_core::quant::q1t_pack(&mut c, k, code);
6741 }
6742 bytes.extend_from_slice(&c);
6743 }
6744 }
6745 let mut row_ptr = vec![0u32; rows + 1];
6748 for &(idx, _) in &outliers {
6749 row_ptr[idx as usize / cols + 1] += 1;
6750 }
6751 for r in 0..rows {
6752 row_ptr[r + 1] += row_ptr[r];
6753 }
6754 for &p in &row_ptr {
6755 bytes.extend_from_slice(&p.to_le_bytes());
6756 }
6757 for &(idx, v) in &outliers {
6758 bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
6759 bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
6760 }
6761
6762 let mut refw = vec![0f32; rows * cols];
6763 dequant_q1t(&bytes, rows, cols, &mut refw);
6764 let x: Vec<f32> = (0..cols).map(|j| if j % 3 == 0 { 1.0 } else { -1.0 }).collect();
6767 let mut expect = vec![0f32; rows];
6768 for r in 0..rows {
6769 let mut a = 0.0f32;
6770 for j in 0..cols {
6771 a += refw[r * cols + j] * x[j];
6772 }
6773 expect[r] = a;
6774 }
6775 let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
6776 let mut got = vec![0f32; rows];
6777 q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
6778 for r in 0..rows {
6779 assert!((got[r] - expect[r]).abs() < tol(expect[r]), "row {r}: {} vs {}", got[r], expect[r]);
6780 }
6781 let x2: Vec<f32> = x.iter().chain(x.iter().map(|v| v)).copied().collect();
6783 let mut gm = vec![0f32; 2 * rows];
6784 q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
6785 for r in 0..rows {
6786 assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
6787 assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
6788 }
6789 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
6794 q1t_matvec(&bytes, &x, rows, cols, &mut p1, None);
6795 q1t_matvec(&bytes, &x, rows, cols, &mut p2, None);
6796 for r in 0..rows {
6797 assert!((p1[r] - expect[r]).abs() < tol(expect[r]) && (p2[r] - expect[r]).abs() < tol(expect[r]));
6798 }
6799 }
6800
6801 #[test]
6805 #[ignore]
6806 fn q1t_matvec_speed() {
6807 use cortiq_core::quant::{f32_to_f16, q1t_code, q1t_pack, Q1T_TILE};
6808 use std::time::Instant;
6809 let (rows, cols) = (8192usize, 4096usize); let gpr = cols / GROUP_SIZE;
6811 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
6812 for r in 0..rows {
6813 for g in 0..gpr {
6814 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
6815 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
6816 let mut c = [0u8; 7];
6817 for k in 0..GROUP_SIZE {
6818 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
6819 }
6820 bytes.extend_from_slice(&c);
6821 }
6822 }
6823 let (n, stride) = (rows * cols, 40usize); let mut row_ptr = vec![0u32; rows + 1];
6825 let mut idx = 0usize;
6826 while idx < n {
6827 row_ptr[idx / cols + 1] += 1;
6828 idx += stride;
6829 }
6830 for r in 0..rows {
6831 row_ptr[r + 1] += row_ptr[r];
6832 }
6833 for &p in &row_ptr {
6834 bytes.extend_from_slice(&p.to_le_bytes());
6835 }
6836 let mut idx = 0usize;
6837 while idx < n {
6838 bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
6839 bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
6840 idx += stride;
6841 }
6842 let x: Vec<f32> = (0..cols).map(|j| if j % 3 == 0 { 1.0 } else { -1.0 }).collect();
6845 let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
6846
6847 let slow = |out: &mut [f32]| {
6849 let mut buf = vec![0f32; cols];
6850 for r in 0..rows {
6851 for g in 0..gpr {
6852 let off = (r * gpr + g) * Q1T_TILE;
6853 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6854 let codes = &bytes[off + 2..off + Q1T_TILE];
6855 for k in 0..GROUP_SIZE {
6856 buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
6857 1 => s,
6858 2 => -s,
6859 _ => 0.0,
6860 };
6861 }
6862 }
6863 out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
6864 + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
6865 }
6866 };
6867 let iters = 5;
6868 let mut a = vec![0f32; rows];
6869 slow(&mut a); let t = Instant::now();
6871 for _ in 0..iters {
6872 slow(&mut a);
6873 }
6874 let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
6875
6876 let mut b = vec![0f32; rows];
6877 q1t_matvec(&bytes, &x, rows, cols, &mut b, None); let t = Instant::now();
6879 for _ in 0..iters {
6880 q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
6881 }
6882 let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
6883
6884 for r in 0..rows {
6885 assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
6886 }
6887 println!(
6888 "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms fused-LUT {fast_ms:.2} ms => {:.2}x",
6889 slow_ms / fast_ms
6890 );
6891 }
6892}