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 _ => {
245 let mut data = vec![0.0f32; rows * cols];
246 cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
247 Ok(Self::from_f32(data, rows, cols))
248 }
249 }
250 }
251
252 pub(crate) fn is_q1(&self) -> bool {
255 matches!(self, Self::Mapped { dtype: TensorDtype::Q1, .. })
256 }
257
258 pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
261 match self {
262 Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
263 _ => None,
264 }
265 }
266
267 pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
270 match self {
271 Self::Mapped { idx, dtype: TensorDtype::Q1, rows, cols, .. } => {
272 Some((*idx, *rows, *cols))
273 }
274 _ => None,
275 }
276 }
277
278 pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
283 match self {
284 Self::Mapped {
285 idx,
286 dtype: TensorDtype::Q8Row,
287 rows,
288 cols,
289 row_scale,
290 col_field,
291 ..
292 } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
293 _ => None,
294 }
295 }
296
297 pub fn rows(&self) -> usize {
298 match self {
299 Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
300 }
301 }
302
303 pub fn cols(&self) -> usize {
304 match self {
305 Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
306 }
307 }
308
309 pub fn as_f32(&self) -> Option<&[f32]> {
312 match self {
313 Self::F32 { data, .. } => Some(data),
314 Self::Mapped { .. } => None,
315 }
316 }
317
318 fn quant_bytes(&self) -> &[u8] {
319 match self {
320 Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
321 Self::F32 { .. } => unreachable!("quant_bytes on F32"),
322 }
323 }
324
325 pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
327 let cols = self.cols();
328 debug_assert_eq!(dst.len(), cols);
329 match self {
330 Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
331 Self::Mapped {
332 dtype,
333 row_scale,
334 col_field,
335 vbit_offsets,
336 ..
337 } => {
338 if *dtype == TensorDtype::Q4Tiled {
339 let bytes = self.quant_bytes();
340 let gpr = cols / GROUP_SIZE;
341 for gi in 0..gpr {
342 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
343 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
344 for (k, &b) in tile[2..].iter().enumerate() {
345 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
346 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
347 }
348 }
349 return;
350 }
351 if *dtype == TensorDtype::Q4Block {
352 let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
353 let gpr = cols / GROUP_SIZE;
354 for gi in 0..gpr {
355 let g = r * gpr + gi;
356 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
357 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
358 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
359 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
360 }
361 }
362 return;
363 }
364 if *dtype == TensorDtype::Q1 {
365 let bytes = self.quant_bytes();
366 let gpr = cols / GROUP_SIZE;
367 for gi in 0..gpr {
368 let tile =
369 &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
370 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
371 for (j, &b) in tile[2..].iter().enumerate() {
372 for k in 0..8 {
373 dst[gi * GROUP_SIZE + j * 8 + k] =
374 (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
375 }
376 }
377 }
378 return;
379 }
380 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
381 let bytes = self.quant_bytes();
382 let rows = self.rows();
383 let ng = cols / GROUP_SIZE;
384 let bits = &bytes[..rows];
385 let sc_off = rows;
386 let off = vbit_offsets[r];
389 let b = bits[r] as usize;
390 let l = ((1usize << (b - 1)) - 1) as f32;
391 let data = &bytes[off..];
392 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
393 for (i, d) in dst.iter_mut().enumerate() {
394 while nbits < b {
395 acc = (acc << 8) | data[idx] as u64;
396 idx += 1;
397 nbits += 8;
398 }
399 let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
400 nbits -= b;
401 let so = (r * ng + i / GROUP_SIZE) * 2;
402 let sv = f16_to_f32(u16::from_le_bytes([
403 bytes[sc_off + so],
404 bytes[sc_off + so + 1],
405 ]));
406 *d = (u - l) * sv;
407 }
408 return;
409 }
410 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
411 let s = row_scale[r];
412 match dtype {
413 TensorDtype::Q8Row => {
414 for (d, &b) in dst.iter_mut().zip(q) {
415 *d = (b as i8) as f32 * s;
416 }
417 }
418 TensorDtype::Q8_2f => {
419 for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
420 *d = (b as i8) as f32 * s * col_field[i];
421 }
422 }
423 _ => unreachable!(),
424 }
425 }
426 }
427 }
428
429 pub fn sparse_col_ok(&self) -> bool {
434 match self {
435 Self::F32 { .. } => true,
436 Self::Mapped { dtype, .. } => {
437 matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
438 }
439 }
440 }
441
442 pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
446 let inter = self.cols();
447 let hidden = self.rows();
448 debug_assert_eq!(out.len(), hidden);
449 match self {
450 Self::F32 { data, .. } => {
451 for (k, o) in out.iter_mut().enumerate() {
452 *o += w * data[k * inter + c];
453 }
454 }
455 Self::Mapped {
456 dtype, row_scale, col_field, ..
457 } => {
458 let q = self.quant_bytes();
459 let colf = if *dtype == TensorDtype::Q8_2f {
460 col_field[c]
461 } else {
462 1.0
463 };
464 let wc = w * colf;
465 for (k, o) in out.iter_mut().enumerate() {
466 let b = q[k * inter + c] as i8 as f32;
467 *o += wc * b * row_scale[k];
468 }
469 }
470 }
471 }
472
473 pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
477 let cols = self.cols();
478 match self {
479 Self::F32 { data, .. } => {
480 let row = &data[r * cols..(r + 1) * cols];
481 row.iter().zip(x).map(|(w, v)| w * v).sum()
482 }
483 Self::Mapped { dtype, row_scale, col_field, .. } => match dtype {
484 TensorDtype::Q8Row => {
485 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
486 dot_i8_f32(q, x) * row_scale[r]
487 }
488 TensorDtype::Q8_2f => {
489 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
490 dot_i8_col_f32(q, x, col_field) * row_scale[r]
491 }
492 _ => {
493 self.row_f32(r, scratch);
494 scratch.iter().zip(x).map(|(w, v)| w * v).sum()
495 }
496 },
497 }
498 }
499
500 pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
503 match self {
504 Self::F32 { data, .. } => matvec_rows(pool, data, x, out),
505 Self::Mapped {
506 model,
507 idx,
508 dtype,
509 rows,
510 cols,
511 row_scale,
512 col_field,
513 vbit_offsets,
514 repack,
515 } => {
516 let _ = (model, idx);
517 if *dtype == TensorDtype::Q4Block {
518 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
519 return;
520 }
521 if *dtype == TensorDtype::Q4Tiled {
522 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
523 return;
524 }
525 if *dtype == TensorDtype::Q1 {
526 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
531 let t0 = std::time::Instant::now();
532 let arm = if crate::gpu::q1_force() {
533 crate::gpu::ProbeArm::Gpu
534 } else {
535 crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
536 };
537 match arm {
538 crate::gpu::ProbeArm::Gpu => {
539 if crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out) {
540 crate::gpu::probe_record(
541 crate::gpu::OpClass::Matvec,
542 true,
543 t0.elapsed(),
544 );
545 return;
546 }
547 }
548 crate::gpu::ProbeArm::CpuTimed => {
549 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
550 crate::gpu::probe_record(
551 crate::gpu::OpClass::Matvec,
552 false,
553 t0.elapsed(),
554 );
555 return;
556 }
557 crate::gpu::ProbeArm::Cpu => {}
558 }
559 }
560 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
561 return;
562 }
563 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
564 vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
565 return;
566 }
567 let xs = prescale(x, col_field, *dtype);
568 if *rows >= crate::gpu::min_rows()
573 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
574 && std::env::var("CMF_GPU_LMHEAD").map(|v| v != "0").unwrap_or(true)
575 && crate::gpu::enabled_here()
576 {
577 let t0 = std::time::Instant::now();
580 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
581 crate::gpu::ProbeArm::Gpu => {}
582 crate::gpu::ProbeArm::CpuTimed => {
583 qmatvec(self.quant_bytes(), repack, row_scale, &xs, *rows, *cols, out, pool);
584 crate::gpu::probe_record(
585 crate::gpu::OpClass::Matvec,
586 false,
587 t0.elapsed(),
588 );
589 return;
590 }
591 crate::gpu::ProbeArm::Cpu => {
592 qmatvec(self.quant_bytes(), repack, row_scale, &xs, *rows, *cols, out, pool);
593 return;
594 }
595 }
596 let frac = std::env::var("CMF_GPU_SPLIT")
597 .ok()
598 .and_then(|v| v.parse::<f32>().ok())
599 .unwrap_or(0.5)
600 .clamp(0.0, 1.0);
601 let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
602 let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
603 let bytes = self.quant_bytes();
604 let ok = std::thread::scope(|sc| {
605 let g = sc.spawn(|| {
606 crate::gpu::q8_matvec_range(
607 model,
608 *idx,
609 cpu_rows,
610 &row_scale[cpu_rows..],
611 &xs,
612 *rows - cpu_rows,
613 *cols,
614 out_gpu,
615 )
616 });
617 if cpu_rows > 0 {
618 let rep_cpu = if repack.is_empty() {
621 &[][..]
622 } else {
623 &repack[..(cpu_rows / 4) * 4 * *cols]
624 };
625 qmatvec(
626 &bytes[..cpu_rows * *cols],
627 rep_cpu,
628 &row_scale[..cpu_rows],
629 &xs,
630 cpu_rows,
631 *cols,
632 out_cpu,
633 pool,
634 );
635 }
636 g.join().unwrap_or(false)
637 });
638 if ok {
639 crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
640 return;
641 }
642 qmatvec(
645 &bytes[cpu_rows * *cols..(*rows) * *cols],
646 &[],
647 &row_scale[cpu_rows..],
648 &xs,
649 *rows - cpu_rows,
650 *cols,
651 out_gpu,
652 pool,
653 );
654 return;
655 }
656 qmatvec(self.quant_bytes(), repack, row_scale, &xs, *rows, *cols, out, pool);
657 }
658 }
659 }
660
661 pub fn matvec2(&self, x1: &[f32], x2: &[f32], o1: &mut [f32], o2: &mut [f32], pool: Option<&Pool>) {
663 match self {
664 Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
665 Self::Mapped {
666 dtype,
667 rows,
668 cols,
669 row_scale,
670 col_field,
671 vbit_offsets,
672 ..
673 } => {
674 if *dtype == TensorDtype::Q4Block {
675 q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
676 return;
677 }
678 if *dtype == TensorDtype::Q4Tiled {
679 q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
680 return;
681 }
682 if *dtype == TensorDtype::Q1 {
683 q1_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
684 return;
685 }
686 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
687 vbitmatvec2(self.quant_bytes(), vbit_offsets, x1, x2, *rows, *cols, o1, o2, pool);
688 return;
689 }
690 let x1s = prescale(x1, col_field, *dtype);
691 let x2s = prescale(x2, col_field, *dtype);
692 qmatvec2(self.quant_bytes(), row_scale, &x1s, &x2s, *rows, *cols, o1, o2, pool);
693 }
694 }
695 }
696}
697
698impl QTensor {
699 pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
704 let cols = self.cols();
705 let rows = self.rows();
706 debug_assert_eq!(xs_all.len(), b * cols);
707 debug_assert_eq!(out.len(), b * rows);
708 match self {
709 Self::F32 { data, .. } => {
710 let out_addr = SendMut(out.as_mut_ptr());
711 let run = |start: usize, end: usize| {
712 for o in start..end {
713 let row = &data[o * cols..(o + 1) * cols];
714 for bi in 0..b {
715 let x = &xs_all[bi * cols..(bi + 1) * cols];
716 let mut acc = 0f32;
717 for j in 0..cols {
718 acc += row[j] * x[j];
719 }
720 unsafe { *out_addr.at(bi * rows + o) = acc };
721 }
722 }
723 };
724 dispatch_rows(pool, rows, &run);
725 }
726 Self::Mapped {
727 dtype,
728 row_scale,
729 col_field,
730 vbit_offsets,
731 ..
732 } => {
733 if *dtype == TensorDtype::Q4Block {
734 q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
735 return;
736 }
737 if *dtype == TensorDtype::Q4Tiled {
738 q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
739 return;
740 }
741 if *dtype == TensorDtype::Q1 {
742 q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
743 return;
744 }
745 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
746 vbitmatmat(self.quant_bytes(), vbit_offsets, xs_all, b, rows, cols, out, pool);
747 return;
748 }
749 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
750 .map(|bi| {
751 prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype)
752 })
753 .collect();
754 if b >= 8
760 && b * rows * cols >= 128_000_000
761 && crate::gpu::enabled_here()
762 {
763 if let Self::Mapped { model, idx, .. } = self {
764 let t0 = std::time::Instant::now();
765 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
766 crate::gpu::ProbeArm::Gpu
767 if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
768 && !crate::gpu::q8_resident_or_upload(model, *idx) =>
769 {
770 let q = self.quant_bytes();
774 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
775 return;
776 }
777 crate::gpu::ProbeArm::Gpu => {
778 let flat: Vec<f32> =
779 pre.iter().flat_map(|v| v.iter().copied()).collect();
780 if crate::gpu::q8_matmat(
781 model, *idx, row_scale, &flat, b, rows, cols, out)
782 {
783 crate::gpu::probe_record(
784 crate::gpu::OpClass::Matmat,
785 true,
786 t0.elapsed(),
787 );
788 return;
789 }
790 }
791 crate::gpu::ProbeArm::CpuTimed => {
792 let q = self.quant_bytes();
793 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
794 crate::gpu::probe_record(
795 crate::gpu::OpClass::Matmat,
796 false,
797 t0.elapsed(),
798 );
799 return;
800 }
801 crate::gpu::ProbeArm::Cpu => {}
802 }
803 }
804 }
805 let q = self.quant_bytes();
806 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
807 }
808 }
809 }
810}
811
812impl QTensor {
813 pub fn matvec_many<const N: usize>(
820 ts: [&QTensor; N],
821 x: &[f32],
822 mut outs: [&mut [f32]; N],
823 pool: Option<&Pool>,
824 ) {
825 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
826 let uniform_q8 = ts.iter().all(|t| {
827 matches!(
828 t,
829 Self::Mapped { dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f, .. }
830 )
831 });
832 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
833 let uniform_q4 = ts
834 .iter()
835 .all(|t| matches!(t, Self::Mapped { dtype: TensorDtype::Q4Block, .. }));
836 let uniform_vbit = ts
837 .iter()
838 .all(|t| matches!(
839 t,
840 Self::Mapped { dtype: TensorDtype::Vbit | TensorDtype::VbitRo, .. }
841 ));
842 let uniform_q1 = ts
843 .iter()
844 .all(|t| matches!(t, Self::Mapped { dtype: TensorDtype::Q1, .. }));
845 let Some(pool) = pool else {
846 for (t, o) in ts.iter().zip(outs.iter_mut()) {
847 t.matvec(x, o, None);
848 }
849 return;
850 };
851 if total_rows < 256
852 || !(uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit || uniform_q1)
853 {
854 for (t, o) in ts.iter().zip(outs.iter_mut()) {
855 t.matvec(x, o, Some(pool));
856 }
857 return;
858 }
859
860 if uniform_q1 {
861 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
864 if a8w8_enabled() {
865 let act = split_act(x);
866 let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
867 let (act, gsum) = (&act, &gsum);
868 let closures: [_; N] = std::array::from_fn(|i| {
869 let (bytes, gpr, out) =
870 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
871 move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
872 });
873 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
874 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
875 pool.run_many(&parts);
876 } else {
877 let closures: [_; N] = std::array::from_fn(|i| {
878 let (bytes, gpr, out) =
879 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
880 move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
881 });
882 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
883 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
884 pool.run_many(&parts);
885 }
886 return;
887 }
888
889 if uniform_q4 || uniform_vbit {
890 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
891 if a8w8_enabled() {
893 let act = split_act(x);
894 let act = &act;
895 if uniform_q4 {
896 let closures: [_; N] = std::array::from_fn(|i| {
897 let (packed, scales) =
898 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
899 let (gpr, cols, out) =
900 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
901 move |s: usize, e: usize| {
902 q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
903 }
904 });
905 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
906 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
907 pool.run_many(&parts);
908 } else {
909 let closures: [_; N] = std::array::from_fn(|i| {
910 let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
911 let (bytes, rows, cols, out) =
912 (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), outs_addr[i]);
913 move |s: usize, e: usize| {
914 vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
915 }
916 });
917 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
918 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
919 pool.run_many(&parts);
920 }
921 return;
922 }
923 if uniform_q4 {
924 let closures: [_; N] = std::array::from_fn(|i| {
925 let (packed, scales) =
926 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
927 let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
928 move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
929 });
930 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
931 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
932 pool.run_many(&parts);
933 } else {
934 let closures: [_; N] = std::array::from_fn(|i| {
935 let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
936 let (bytes, rows, cols, out) =
937 (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), outs_addr[i]);
938 move |s: usize, e: usize| {
939 vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
940 }
941 });
942 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
943 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
944 pool.run_many(&parts);
945 }
946 return;
947 }
948
949 if uniform_f32 {
950 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
951 let closures: [_; N] = std::array::from_fn(|i| {
952 let Self::F32 { data, cols, .. } = ts[i] else { unreachable!() };
953 let out = outs_addr[i];
954 move |start: usize, end: usize| {
955 for o in start..end {
956 let row = &data[o * cols..(o + 1) * cols];
957 let mut sum = 0.0f32;
958 for j in 0..*cols {
959 sum += row[j] * x[j];
960 }
961 unsafe { *out.at(o) = sum };
963 }
964 }
965 });
966 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
967 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
968 pool.run_many(&parts);
969 return;
970 }
971
972 struct Ctx<'a> {
975 bytes: &'a [u8],
976 #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
977 rep: &'a [u8],
978 row_scale: &'a [f32],
979 cols: usize,
980 xs: std::borrow::Cow<'a, [f32]>,
981 }
982 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
983 let Self::Mapped { dtype, cols, row_scale, col_field, repack, .. } = ts[i] else {
984 unreachable!()
985 };
986 Ctx {
987 bytes: ts[i].quant_bytes(),
988 rep: repack,
989 row_scale,
990 cols: *cols,
991 xs: prescale(x, col_field, *dtype),
992 }
993 });
994 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
995 #[cfg(target_arch = "aarch64")]
996 if sdot_enabled() {
997 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
998 let closures: [_; N] = std::array::from_fn(|i| {
999 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
1000 move |start: usize, end: usize| {
1001 q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
1002 }
1003 });
1004 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1005 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1006 pool.run_many(&parts);
1007 return;
1008 }
1009 #[cfg(target_arch = "x86_64")]
1010 if avx2_a8w8_enabled() {
1011 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
1012 let closures: [_; N] = std::array::from_fn(|i| {
1013 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
1014 move |start: usize, end: usize| {
1015 q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
1016 }
1017 });
1018 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1019 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1020 pool.run_many(&parts);
1021 return;
1022 }
1023 let closures: [_; N] = std::array::from_fn(|i| {
1024 let (c, out) = (&ctxs[i], outs_addr[i]);
1025 move |start: usize, end: usize| {
1026 q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
1027 }
1028 });
1029 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1030 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1031 pool.run_many(&parts);
1032 }
1033}
1034
1035impl QTensor {
1036 #[allow(clippy::needless_range_loop)]
1041 pub fn matvec2_many<const N: usize>(
1042 ts: [&QTensor; N],
1043 x1: &[f32],
1044 x2: &[f32],
1045 mut o1s: [&mut [f32]; N],
1046 mut o2s: [&mut [f32]; N],
1047 pool: Option<&Pool>,
1048 ) {
1049 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
1050 let uniform_q8 = ts.iter().all(|t| {
1051 matches!(
1052 t,
1053 Self::Mapped { dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f, .. }
1054 )
1055 });
1056 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
1057 let uniform_q4 = ts
1058 .iter()
1059 .all(|t| matches!(t, Self::Mapped { dtype: TensorDtype::Q4Block, .. }));
1060 let uniform_vbit = ts
1061 .iter()
1062 .all(|t| matches!(
1063 t,
1064 Self::Mapped { dtype: TensorDtype::Vbit | TensorDtype::VbitRo, .. }
1065 ));
1066 let fusable = pool.is_some()
1067 && total_rows >= 256
1068 && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
1069 if !fusable {
1070 for i in 0..N {
1071 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
1072 }
1073 return;
1074 }
1075 let pool = pool.unwrap();
1076
1077 if uniform_q4 || uniform_vbit {
1078 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1079 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1080 if a8w8_enabled() {
1082 let a1 = split_act(x1);
1083 let a2 = split_act(x2);
1084 let (a1, a2) = (&a1, &a2);
1085 if uniform_q4 {
1086 let closures: [_; N] = std::array::from_fn(|i| {
1087 let (packed, scales) =
1088 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1089 let (gpr, cols, o1, o2) =
1090 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
1091 move |s: usize, e: usize| {
1092 q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
1093 }
1094 });
1095 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1096 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1097 pool.run_many(&parts);
1098 } else {
1099 let closures: [_; N] = std::array::from_fn(|i| {
1100 let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
1101 let (bytes, rows, cols, o1, o2) =
1102 (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), p1[i], p2[i]);
1103 move |s: usize, e: usize| {
1104 vbit_range2_a8w8(
1105 bytes, vbit_offsets, x1, x2, a1, a2, rows, cols, o1, o2, s, e,
1106 )
1107 }
1108 });
1109 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1110 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1111 pool.run_many(&parts);
1112 }
1113 return;
1114 }
1115 if uniform_q4 {
1116 let closures: [_; N] = std::array::from_fn(|i| {
1117 let (packed, scales) =
1118 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1119 let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
1120 move |s: usize, e: usize| {
1121 q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
1122 }
1123 });
1124 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1125 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1126 pool.run_many(&parts);
1127 } else {
1128 let closures: [_; N] = std::array::from_fn(|i| {
1129 let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
1130 let (bytes, rows, cols, o1, o2) =
1131 (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), p1[i], p2[i]);
1132 move |s: usize, e: usize| {
1133 vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
1134 }
1135 });
1136 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1137 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1138 pool.run_many(&parts);
1139 }
1140 return;
1141 }
1142
1143 if uniform_f32 {
1144 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1145 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1146 let closures: [_; N] = std::array::from_fn(|i| {
1147 let Self::F32 { data, cols, .. } = ts[i] else { unreachable!() };
1148 let (o1, o2) = (p1[i], p2[i]);
1149 move |start: usize, end: usize| {
1150 for o in start..end {
1151 let row = &data[o * cols..(o + 1) * cols];
1152 let (mut s1, mut s2) = (0.0f32, 0.0f32);
1153 for j in 0..*cols {
1154 s1 += row[j] * x1[j];
1155 s2 += row[j] * x2[j];
1156 }
1157 unsafe {
1159 *o1.at(o) = s1;
1160 *o2.at(o) = s2;
1161 }
1162 }
1163 }
1164 });
1165 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1166 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1167 pool.run_many(&parts);
1168 return;
1169 }
1170
1171 struct Ctx<'a> {
1172 bytes: &'a [u8],
1173 row_scale: &'a [f32],
1174 cols: usize,
1175 xs1: std::borrow::Cow<'a, [f32]>,
1176 xs2: std::borrow::Cow<'a, [f32]>,
1177 }
1178 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
1179 let Self::Mapped { dtype, cols, row_scale, col_field, .. } = ts[i] else {
1180 unreachable!()
1181 };
1182 Ctx {
1183 bytes: ts[i].quant_bytes(),
1184 row_scale,
1185 cols: *cols,
1186 xs1: prescale(x1, col_field, *dtype),
1187 xs2: prescale(x2, col_field, *dtype),
1188 }
1189 });
1190 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1191 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1192 #[cfg(target_arch = "aarch64")]
1193 if sdot_enabled() {
1194 let acts: [(SplitAct, SplitAct); N] =
1195 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
1196 let closures: [_; N] = std::array::from_fn(|i| {
1197 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
1198 move |start: usize, end: usize| {
1199 q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
1200 }
1201 });
1202 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1203 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1204 pool.run_many(&parts);
1205 return;
1206 }
1207 #[cfg(target_arch = "x86_64")]
1208 if avx2_a8w8_enabled() {
1209 let acts: [(SplitAct, SplitAct); N] =
1210 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
1211 let closures: [_; N] = std::array::from_fn(|i| {
1212 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
1213 move |start: usize, end: usize| {
1214 q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
1215 }
1216 });
1217 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1218 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1219 pool.run_many(&parts);
1220 return;
1221 }
1222 let closures: [_; N] = std::array::from_fn(|i| {
1223 let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
1224 move |start: usize, end: usize| {
1225 q8_range2_f32(c.bytes, c.row_scale, &c.xs1, &c.xs2, c.cols, o1, o2, start, end)
1226 }
1227 });
1228 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1229 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1230 pool.run_many(&parts);
1231 }
1232}
1233
1234#[cfg(target_os = "macos")]
1239mod accel_blas {
1240 #[link(name = "Accelerate", kind = "framework")]
1241 unsafe extern "C" {
1242 pub fn cblas_sgemm(
1243 order: i32,
1244 trans_a: i32,
1245 trans_b: i32,
1246 m: i32,
1247 n: i32,
1248 k: i32,
1249 alpha: f32,
1250 a: *const f32,
1251 lda: i32,
1252 b: *const f32,
1253 ldb: i32,
1254 beta: f32,
1255 c: *mut f32,
1256 ldc: i32,
1257 );
1258 }
1259}
1260
1261#[cfg(target_os = "macos")]
1262pub(crate) fn accel_gemm_enabled() -> bool {
1263 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1264 *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
1265}
1266
1267#[cfg(target_os = "macos")]
1270#[allow(clippy::too_many_arguments)]
1271pub(crate) fn sgemm_rm(
1272 m: usize,
1273 n: usize,
1274 k: usize,
1275 alpha: f32,
1276 a: &[f32],
1277 lda: usize,
1278 b_mat: &[f32],
1279 ldb: usize,
1280 b_rows_are_n: bool,
1281 c: &mut [f32],
1282 ldc: usize,
1283) {
1284 debug_assert!(a.len() >= (m - 1) * lda + k);
1285 debug_assert!(c.len() >= (m - 1) * ldc + n);
1286 unsafe {
1287 accel_blas::cblas_sgemm(
1288 101, 111, if b_rows_are_n { 112 } else { 111 },
1291 m as i32,
1292 n as i32,
1293 k as i32,
1294 alpha,
1295 a.as_ptr(),
1296 lda as i32,
1297 b_mat.as_ptr(),
1298 ldb as i32,
1299 0.0,
1300 c.as_mut_ptr(),
1301 ldc as i32,
1302 );
1303 }
1304}
1305
1306#[cfg(target_os = "macos")]
1313fn qmatmat_accel(
1314 q: &[u8],
1315 row_scale: &[f32],
1316 pre: &[std::borrow::Cow<'_, [f32]>],
1317 rows: usize,
1318 cols: usize,
1319 out: &mut [f32],
1320 pool: Option<&Pool>,
1321) {
1322 const TR: usize = 2048;
1327 let b = pre.len();
1328 thread_local! {
1329 static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
1330 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
1331 }
1332 XPANEL.with(|xp| {
1333 WTILE.with(|wt| {
1334 let mut xpanel = xp.borrow_mut();
1335 xpanel.clear();
1336 for x in pre {
1337 xpanel.extend_from_slice(x);
1338 }
1339 let mut wtile = wt.borrow_mut();
1340 wtile.resize(TR * cols, 0.0);
1341 let mut r0 = 0usize;
1342 while r0 < rows {
1343 let tr = TR.min(rows - r0);
1344 let wt_addr = SendMut(wtile.as_mut_ptr());
1346 let run = |start: usize, end: usize| {
1347 for r in start..end {
1348 let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
1349 let s = row_scale[r0 + r];
1350 let dst = unsafe {
1352 std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols)
1353 };
1354 for (d, &v) in dst.iter_mut().zip(row) {
1355 *d = (v as i8) as f32 * s;
1356 }
1357 }
1358 };
1359 dispatch_rows(pool, tr, &run);
1360 unsafe {
1362 accel_blas::cblas_sgemm(
1363 101, 111, 112, b as i32,
1367 tr as i32,
1368 cols as i32,
1369 1.0,
1370 xpanel.as_ptr(),
1371 cols as i32,
1372 wtile.as_ptr(),
1373 cols as i32,
1374 0.0,
1375 out.as_mut_ptr().add(r0),
1376 rows as i32,
1377 );
1378 }
1379 r0 += tr;
1380 }
1381 })
1382 });
1383}
1384
1385fn qmatmat(
1386 q: &[u8],
1387 row_scale: &[f32],
1388 pre: &[std::borrow::Cow<'_, [f32]>],
1389 rows: usize,
1390 cols: usize,
1391 out: &mut [f32],
1392 pool: Option<&Pool>,
1393) {
1394 let b = pre.len();
1395 debug_assert_eq!(out.len(), b * rows);
1396 #[cfg(target_os = "macos")]
1401 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
1402 qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
1403 return;
1404 }
1405 #[cfg(target_arch = "aarch64")]
1406 if sdot_enabled() {
1407 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
1408 let out_addr = SendMut(out.as_mut_ptr());
1409 let run = |start: usize, end: usize| {
1410 for o in start..end {
1411 let row = &q[o * cols..(o + 1) * cols];
1412 for (bi, act) in acts.iter().enumerate() {
1413 let v = row_dot_sdot(row, act) * row_scale[o];
1414 unsafe { *out_addr.at(bi * rows + o) = v };
1415 }
1416 }
1417 };
1418 dispatch_rows(pool, rows, &run);
1419 return;
1420 }
1421 #[cfg(target_arch = "x86_64")]
1425 if avx2_a8w8_enabled() {
1426 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
1427 let out_addr = SendMut(out.as_mut_ptr());
1428 let run = |start: usize, end: usize| {
1429 for o in start..end {
1430 let row = &q[o * cols..(o + 1) * cols];
1431 for (bi, act) in acts.iter().enumerate() {
1432 let v = row_dot_avx2(row, act) * row_scale[o];
1433 unsafe { *out_addr.at(bi * rows + o) = v };
1434 }
1435 }
1436 };
1437 dispatch_rows(pool, rows, &run);
1438 return;
1439 }
1440 let out_addr = SendMut(out.as_mut_ptr());
1441 let run = |start: usize, end: usize| {
1442 for o in start..end {
1443 let row = &q[o * cols..(o + 1) * cols];
1444 for (bi, x) in pre.iter().enumerate() {
1445 let mut acc = 0f32;
1446 for j in 0..cols {
1447 acc += (row[j] as i8) as f32 * x[j];
1448 }
1449 unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
1450 }
1451 }
1452 };
1453 dispatch_rows(pool, rows, &run);
1454}
1455
1456fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
1459 match pool {
1460 Some(pool) if rows >= 256 => pool.run_rows(rows, run),
1461 _ => run(0, rows),
1462 }
1463}
1464
1465fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
1467 let groups = rows * cols / GROUP_SIZE;
1468 bytes.split_at(groups * 16)
1469}
1470
1471#[inline]
1476fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
1477 #[cfg(target_arch = "aarch64")]
1478 unsafe {
1479 return vbit_fill4_neon(data, buf);
1480 }
1481 #[cfg(target_arch = "x86_64")]
1482 if avx2_enabled() {
1483 return unsafe { vbit_fill4_avx2(data, buf) };
1484 }
1485 #[allow(unreachable_code)]
1486 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
1487 let u = unpack8::<4>(&data[blk * 4..]);
1488 for k in 0..8 {
1489 chunk[k] = (u[k] - 7) as i8 as u8;
1490 }
1491 }
1492}
1493
1494#[cfg(target_arch = "aarch64")]
1495#[target_feature(enable = "neon")]
1496unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
1497 unsafe {
1500 use core::arch::aarch64::*;
1501 let n = buf.len();
1502 let mask = vdupq_n_u8(0x0F);
1503 let seven = vdupq_n_s8(7);
1504 let mut g = 0usize;
1505 while g * 32 + 32 <= n {
1506 let b = vld1q_u8(data.as_ptr().add(g * 16));
1507 let hi = vshrq_n_u8::<4>(b);
1508 let lo = vandq_u8(b, mask);
1509 let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
1510 let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
1511 vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
1512 vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
1513 g += 1;
1514 }
1515 }
1516}
1517
1518#[cfg(target_arch = "x86_64")]
1519#[target_feature(enable = "avx2")]
1520unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
1521 unsafe {
1523 use core::arch::x86_64::*;
1524 let n = buf.len();
1525 let mask = _mm_set1_epi8(0x0F);
1526 let seven = _mm256_set1_epi8(7);
1527 let mut g = 0usize;
1528 while g * 32 + 32 <= n {
1529 let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
1530 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
1531 let lo = _mm_and_si128(b, mask);
1532 let z = _mm256_sub_epi8(
1533 _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
1534 seven,
1535 );
1536 _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
1537 g += 1;
1538 }
1539 }
1540}
1541
1542#[inline(always)]
1547fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
1548 let mut acc = 0u64;
1549 for i in 0..B {
1550 acc = (acc << 8) | data[i] as u64;
1551 }
1552 let mask = (1u64 << B) - 1;
1553 let mut out = [0i32; 8];
1554 for (k, o) in out.iter_mut().enumerate() {
1555 *o = ((acc >> ((7 - k) * B)) & mask) as i32;
1556 }
1557 out
1558}
1559
1560#[allow(clippy::too_many_arguments)]
1566fn vbitmatvec(
1567 bytes: &[u8],
1568 offsets: &[usize],
1569 x: &[f32],
1570 rows: usize,
1571 cols: usize,
1572 out: &mut [f32],
1573 pool: Option<&Pool>,
1574) {
1575 debug_assert_eq!(out.len(), rows);
1576 debug_assert_eq!(offsets.len(), rows + 1);
1577
1578 if a8w8_enabled() {
1582 let act = split_act(x);
1583 let out_addr = SendMut(out.as_mut_ptr());
1584 let run = move |start: usize, end: usize| {
1585 vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
1586 };
1587 dispatch_rows(pool, rows, &run);
1588 return;
1589 }
1590
1591 let out_addr = SendMut(out.as_mut_ptr());
1592 let run = move |start: usize, end: usize| {
1593 vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
1594 };
1595 dispatch_rows(pool, rows, &run);
1596}
1597
1598#[allow(clippy::too_many_arguments)]
1602fn vbit_range_a8w8(
1603 bytes: &[u8],
1604 offsets: &[usize],
1605 x: &[f32],
1606 act: &SplitAct,
1607 rows: usize,
1608 cols: usize,
1609 out: SendMut,
1610 start: usize,
1611 end: usize,
1612) {
1613 let ng = cols / GROUP_SIZE;
1614 let bits = &bytes[..rows];
1615 let sc_off = rows;
1616 let row_dot = |r: usize| -> f32 {
1617 let b = bits[r] as usize;
1618 let l = ((1i32 << (b - 1)) - 1) as i32;
1619 let mask = (1u64 << b) - 1;
1620 let data = &bytes[offsets[r]..offsets[r + 1]];
1621 if b == 8 {
1622 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
1624 let mut dot = 0f32;
1625 for g in 0..ng {
1626 let so = (r * ng + g) * 2;
1627 let sgf = f16_to_f32(u16::from_le_bytes([
1628 bytes[sc_off + so],
1629 bytes[sc_off + so + 1],
1630 ]));
1631 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
1632 let mut gd = 0f32;
1633 for &xv in xg.iter() {
1634 if nbits < 8 {
1635 acc = (acc << 8) | data[idx] as u64;
1636 idx += 1;
1637 nbits += 8;
1638 }
1639 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
1640 nbits -= 8;
1641 gd += (u - l) as f32 * xv;
1642 }
1643 dot += gd * sgf;
1644 }
1645 return dot;
1646 }
1647 thread_local! {
1651 static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
1652 const { std::cell::RefCell::new(Vec::new()) };
1653 }
1654 #[inline(always)]
1655 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
1656 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
1657 let u = unpack8::<B>(&data[blk * B..]);
1658 for k in 0..8 {
1659 chunk[k] = (u[k] - l) as i8 as u8;
1660 }
1661 }
1662 }
1663 let _ = mask;
1664 VBIT_SCRATCH.with(|scratch| {
1665 let mut buf = scratch.borrow_mut();
1666 buf.resize(cols, 0);
1667 match b {
1668 3 => fill::<3>(data, l, &mut buf),
1669 4 => vbit_fill4(data, &mut buf),
1670 5 => fill::<5>(data, l, &mut buf),
1671 6 => fill::<6>(data, l, &mut buf),
1672 _ => unreachable!(),
1673 }
1674 let mut dot = 0f32;
1675 for g in 0..ng {
1676 let so = (r * ng + g) * 2;
1677 let s = f16_to_f32(u16::from_le_bytes([
1678 bytes[sc_off + so],
1679 bytes[sc_off + so + 1],
1680 ]));
1681 let d = dot_i8_i8(
1682 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
1683 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
1684 ) as f32
1685 * act.sx;
1686 dot += d * s;
1687 }
1688 for &(j, xv) in &act.outliers {
1689 let so = (r * ng + j / GROUP_SIZE) * 2;
1690 let s = f16_to_f32(u16::from_le_bytes([
1691 bytes[sc_off + so],
1692 bytes[sc_off + so + 1],
1693 ]));
1694 dot += (buf[j] as i8) as f32 * s * xv;
1696 }
1697 dot
1698 })
1699 };
1700 for r in start..end {
1701 unsafe { *out.at(r) = row_dot(r) };
1703 }
1704}
1705
1706#[allow(clippy::too_many_arguments)]
1708fn vbit_range_f32(
1709 bytes: &[u8],
1710 offsets: &[usize],
1711 x: &[f32],
1712 rows: usize,
1713 cols: usize,
1714 out: SendMut,
1715 start: usize,
1716 end: usize,
1717) {
1718 let ng = cols / GROUP_SIZE;
1719 let bits = &bytes[..rows];
1720 let sc_off = rows;
1721 #[inline(always)]
1725 fn dot_row<const B: usize>(
1726 data: &[u8],
1727 bytes: &[u8],
1728 sc_off: usize,
1729 r: usize,
1730 ng: usize,
1731 x: &[f32],
1732 ) -> f32 {
1733 let l = ((1i32 << (B - 1)) - 1) as f32;
1734 let gbytes = GROUP_SIZE * B / 8;
1735 let mut dot = 0f32;
1736 for g in 0..ng {
1737 let so = (r * ng + g) * 2;
1738 let s = f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]));
1739 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
1740 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
1741 let mut gd = 0f32;
1742 for blk in 0..GROUP_SIZE / 8 {
1743 let u = unpack8::<B>(&gd0[blk * B..]);
1744 let xb = &xg[blk * 8..blk * 8 + 8];
1745 for k in 0..8 {
1746 gd += (u[k] as f32 - l) * xb[k];
1747 }
1748 }
1749 dot += gd * s;
1750 }
1751 dot
1752 }
1753 for r in start..end {
1754 let data = &bytes[offsets[r]..offsets[r + 1]];
1755 let v = match bits[r] {
1756 3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
1757 4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
1758 5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
1759 6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
1760 8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
1761 b => unreachable!("vbit bit-width {b} (validated at load)"),
1762 };
1763 unsafe { *out.at(r) = v };
1765 }
1766}
1767
1768#[allow(clippy::too_many_arguments)]
1773fn vbitmatvec2(
1774 bytes: &[u8],
1775 offsets: &[usize],
1776 x1: &[f32],
1777 x2: &[f32],
1778 rows: usize,
1779 cols: usize,
1780 o1: &mut [f32],
1781 o2: &mut [f32],
1782 pool: Option<&Pool>,
1783) {
1784 debug_assert_eq!(o1.len(), rows);
1785 debug_assert_eq!(o2.len(), rows);
1786
1787 if a8w8_enabled() {
1788 let a1 = split_act(x1);
1789 let a2 = split_act(x2);
1790 let p1 = SendMut(o1.as_mut_ptr());
1791 let p2 = SendMut(o2.as_mut_ptr());
1792 let run = move |start: usize, end: usize| {
1793 vbit_range2_a8w8(bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end)
1794 };
1795 dispatch_rows(pool, rows, &run);
1796 return;
1797 }
1798
1799 let p1 = SendMut(o1.as_mut_ptr());
1800 let p2 = SendMut(o2.as_mut_ptr());
1801 let run = move |start: usize, end: usize| {
1802 vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
1803 };
1804 dispatch_rows(pool, rows, &run);
1805}
1806
1807#[allow(clippy::too_many_arguments)]
1811fn vbit_range2_a8w8(
1812 bytes: &[u8],
1813 offsets: &[usize],
1814 x1: &[f32],
1815 x2: &[f32],
1816 a1: &SplitAct,
1817 a2: &SplitAct,
1818 rows: usize,
1819 cols: usize,
1820 p1: SendMut,
1821 p2: SendMut,
1822 start: usize,
1823 end: usize,
1824) {
1825 let ng = cols / GROUP_SIZE;
1826 let bits = &bytes[..rows];
1827 let sc_off = rows;
1828 let row_dots = |r: usize| -> (f32, f32) {
1829 let b = bits[r] as usize;
1830 let l = (1i32 << (b - 1)) - 1;
1831 let data = &bytes[offsets[r]..offsets[r + 1]];
1832 if b == 8 {
1833 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
1836 let (mut d1, mut d2) = (0f32, 0f32);
1837 for g in 0..ng {
1838 let so = (r * ng + g) * 2;
1839 let sgf = f16_to_f32(u16::from_le_bytes([
1840 bytes[sc_off + so],
1841 bytes[sc_off + so + 1],
1842 ]));
1843 let (mut g1, mut g2) = (0f32, 0f32);
1844 for k in 0..GROUP_SIZE {
1845 if nbits < 8 {
1846 acc = (acc << 8) | data[idx] as u64;
1847 idx += 1;
1848 nbits += 8;
1849 }
1850 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
1851 nbits -= 8;
1852 let w = (u - l) as f32;
1853 g1 += w * x1[g * GROUP_SIZE + k];
1854 g2 += w * x2[g * GROUP_SIZE + k];
1855 }
1856 d1 += g1 * sgf;
1857 d2 += g2 * sgf;
1858 }
1859 return (d1, d2);
1860 }
1861 thread_local! {
1862 static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
1863 const { std::cell::RefCell::new(Vec::new()) };
1864 }
1865 #[inline(always)]
1866 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
1867 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
1868 let u = unpack8::<B>(&data[blk * B..]);
1869 for k in 0..8 {
1870 chunk[k] = (u[k] - l) as i8 as u8;
1871 }
1872 }
1873 }
1874 VBIT_SCRATCH2.with(|scratch| {
1875 let mut buf = scratch.borrow_mut();
1876 buf.resize(cols, 0);
1877 match b {
1878 3 => fill::<3>(data, l, &mut buf),
1879 4 => vbit_fill4(data, &mut buf),
1880 5 => fill::<5>(data, l, &mut buf),
1881 6 => fill::<6>(data, l, &mut buf),
1882 _ => unreachable!(),
1883 }
1884 let (mut d1, mut d2) = (0f32, 0f32);
1885 for g in 0..ng {
1886 let so = (r * ng + g) * 2;
1887 let s = f16_to_f32(u16::from_le_bytes([
1888 bytes[sc_off + so],
1889 bytes[sc_off + so + 1],
1890 ]));
1891 let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
1892 let v1 =
1893 dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
1894 let v2 =
1895 dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
1896 d1 += v1 * s;
1897 d2 += v2 * s;
1898 }
1899 for &(j, xv) in &a1.outliers {
1900 let so = (r * ng + j / GROUP_SIZE) * 2;
1901 let s = f16_to_f32(u16::from_le_bytes([
1902 bytes[sc_off + so],
1903 bytes[sc_off + so + 1],
1904 ]));
1905 d1 += (buf[j] as i8) as f32 * s * xv;
1906 }
1907 for &(j, xv) in &a2.outliers {
1908 let so = (r * ng + j / GROUP_SIZE) * 2;
1909 let s = f16_to_f32(u16::from_le_bytes([
1910 bytes[sc_off + so],
1911 bytes[sc_off + so + 1],
1912 ]));
1913 d2 += (buf[j] as i8) as f32 * s * xv;
1914 }
1915 (d1, d2)
1916 })
1917 };
1918 for r in start..end {
1919 let (v1, v2) = row_dots(r);
1920 unsafe {
1922 *p1.at(r) = v1;
1923 *p2.at(r) = v2;
1924 }
1925 }
1926}
1927
1928#[allow(clippy::too_many_arguments)]
1932fn vbit_range2_f32(
1933 bytes: &[u8],
1934 offsets: &[usize],
1935 x1: &[f32],
1936 x2: &[f32],
1937 rows: usize,
1938 cols: usize,
1939 p1: SendMut,
1940 p2: SendMut,
1941 start: usize,
1942 end: usize,
1943) {
1944 let ng = cols / GROUP_SIZE;
1945 let bits = &bytes[..rows];
1946 let sc_off = rows;
1947 #[inline(always)]
1948 #[allow(clippy::too_many_arguments)]
1949 fn dot_row2<const B: usize>(
1950 data: &[u8],
1951 bytes: &[u8],
1952 sc_off: usize,
1953 r: usize,
1954 ng: usize,
1955 x1: &[f32],
1956 x2: &[f32],
1957 ) -> (f32, f32) {
1958 let l = ((1i32 << (B - 1)) - 1) as f32;
1959 let gbytes = GROUP_SIZE * B / 8;
1960 let (mut d1, mut d2) = (0f32, 0f32);
1961 for g in 0..ng {
1962 let so = (r * ng + g) * 2;
1963 let s = f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]));
1964 let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
1965 let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
1966 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
1967 let (mut g1, mut g2) = (0f32, 0f32);
1968 for blk in 0..GROUP_SIZE / 8 {
1969 let u = unpack8::<B>(&gd0[blk * B..]);
1970 for k in 0..8 {
1971 let w = u[k] as f32 - l;
1972 g1 += w * x1g[blk * 8 + k];
1973 g2 += w * x2g[blk * 8 + k];
1974 }
1975 }
1976 d1 += g1 * s;
1977 d2 += g2 * s;
1978 }
1979 (d1, d2)
1980 }
1981 for r in start..end {
1982 let data = &bytes[offsets[r]..offsets[r + 1]];
1983 let (v1, v2) = match bits[r] {
1984 3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
1985 4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
1986 5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
1987 6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
1988 8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
1989 b => unreachable!("vbit bit-width {b} (validated at load)"),
1990 };
1991 unsafe {
1993 *p1.at(r) = v1;
1994 *p2.at(r) = v2;
1995 }
1996 }
1997}
1998
1999#[inline]
2006#[allow(unreachable_code)]
2007fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
2008 #[cfg(target_arch = "aarch64")]
2009 unsafe {
2010 return dot_q4t_row_sdot(bytes, r, gpr, xq);
2011 }
2012 #[cfg(target_arch = "x86_64")]
2013 unsafe {
2014 return dot_q4t_row_avx2(bytes, r, gpr, xq);
2015 }
2016 let mut acc = 0f32;
2017 for gi in 0..gpr {
2018 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
2019 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2020 let mut d = 0i32;
2021 for (k, &b) in tile[2..].iter().enumerate() {
2022 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
2023 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
2024 }
2025 acc += d as f32 * s;
2026 }
2027 acc
2028}
2029
2030#[cfg(target_arch = "aarch64")]
2031#[target_feature(enable = "neon,dotprod")]
2032unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
2033 unsafe {
2036 use core::arch::aarch64::*;
2037 use core::arch::asm;
2038 let lomask = vdupq_n_u8(0x0F);
2039 let eight = vdupq_n_s8(8);
2040 let mut acc = 0f32;
2041 for gi in 0..gpr {
2042 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
2043 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2044 let b = vld1q_u8(t.add(2));
2045 let lo = vandq_u8(b, lomask);
2046 let hi = vshrq_n_u8::<4>(b);
2047 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
2048 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
2049 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
2050 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
2051 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
2052 asm!(
2053 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
2054 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
2055 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
2056 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
2057 options(pure, nomem, nostack),
2058 );
2059 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
2060 }
2061 acc
2062 }
2063}
2064
2065#[cfg(target_arch = "x86_64")]
2066#[target_feature(enable = "avx2")]
2067unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
2068 unsafe {
2070 use core::arch::x86_64::*;
2071 let lomask = _mm_set1_epi8(0x0F);
2072 let eight = _mm256_set1_epi8(8);
2073 let ones = _mm256_set1_epi16(1);
2074 let mut acc = 0f32;
2075 for gi in 0..gpr {
2076 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
2077 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2078 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
2079 let lo = _mm_and_si128(b, lomask);
2080 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
2081 let w = _mm256_sub_epi8(
2082 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
2083 eight,
2084 );
2085 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2086 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
2087 let d = _mm256_madd_epi16(p16, ones);
2088 let hi128 = _mm256_extracti128_si256::<1>(d);
2089 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
2090 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2091 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2092 acc += _mm_cvtsi128_si32(s32) as f32 * s;
2093 }
2094 acc
2095 }
2096}
2097
2098#[inline]
2100fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
2101 let gi = j / GROUP_SIZE;
2102 let k = j % GROUP_SIZE;
2103 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
2104 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2105 let byte = tile[2 + k / 2];
2106 let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
2107 ((nib as i32 - 8) as f32, s)
2108}
2109
2110#[inline]
2113fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
2114 let mut acc = 0f32;
2115 for gi in 0..gpr {
2116 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
2117 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2118 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
2119 let mut ga = 0f32;
2120 for (k, &b) in tile[2..].iter().enumerate() {
2121 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
2122 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
2123 }
2124 acc += ga * s;
2125 }
2126 acc
2127}
2128
2129fn q4t_matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
2131 debug_assert_eq!(out.len(), rows);
2132 let gpr = cols / GROUP_SIZE;
2133 let out_addr = SendMut(out.as_mut_ptr());
2134 if a8w8_enabled() {
2135 let act = split_act(x);
2136 let run = move |start: usize, end: usize| {
2137 for r in start..end {
2138 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
2139 for &(j, xv) in &act.outliers {
2140 let (w, s) = q4t_outlier(bytes, r, gpr, j);
2141 acc += w * s * xv;
2142 }
2143 unsafe { *out_addr.at(r) = acc };
2145 }
2146 };
2147 dispatch_rows(pool, rows, &run);
2148 return;
2149 }
2150 let run = move |start: usize, end: usize| {
2151 for r in start..end {
2152 unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
2154 }
2155 };
2156 dispatch_rows(pool, rows, &run);
2157}
2158
2159#[allow(clippy::too_many_arguments)]
2161fn q4t_matvec2(
2162 bytes: &[u8],
2163 x1: &[f32],
2164 x2: &[f32],
2165 rows: usize,
2166 cols: usize,
2167 o1: &mut [f32],
2168 o2: &mut [f32],
2169 pool: Option<&Pool>,
2170) {
2171 let gpr = cols / GROUP_SIZE;
2172 let p1 = SendMut(o1.as_mut_ptr());
2173 let p2 = SendMut(o2.as_mut_ptr());
2174 if a8w8_enabled() {
2175 let a1 = split_act(x1);
2176 let a2 = split_act(x2);
2177 let run = move |start: usize, end: usize| {
2178 for r in start..end {
2179 let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
2180 let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
2181 for &(j, xv) in &a1.outliers {
2182 let (w, s) = q4t_outlier(bytes, r, gpr, j);
2183 v1 += w * s * xv;
2184 }
2185 for &(j, xv) in &a2.outliers {
2186 let (w, s) = q4t_outlier(bytes, r, gpr, j);
2187 v2 += w * s * xv;
2188 }
2189 unsafe {
2191 *p1.at(r) = v1;
2192 *p2.at(r) = v2;
2193 }
2194 }
2195 };
2196 dispatch_rows(pool, rows, &run);
2197 return;
2198 }
2199 let run = move |start: usize, end: usize| {
2200 for r in start..end {
2201 unsafe {
2203 *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
2204 *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
2205 }
2206 }
2207 };
2208 dispatch_rows(pool, rows, &run);
2209}
2210
2211#[allow(clippy::too_many_arguments)]
2213fn q4t_matmat(
2214 bytes: &[u8],
2215 xs_all: &[f32],
2216 b: usize,
2217 rows: usize,
2218 cols: usize,
2219 out: &mut [f32],
2220 pool: Option<&Pool>,
2221) {
2222 debug_assert_eq!(out.len(), b * rows);
2223 let gpr = cols / GROUP_SIZE;
2224 let out_addr = SendMut(out.as_mut_ptr());
2225 if a8w8_enabled() {
2226 let acts: Vec<SplitAct> =
2227 (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
2228 let acts = &acts;
2229 let run = move |start: usize, end: usize| {
2230 for r in start..end {
2231 for (bi, act) in acts.iter().enumerate() {
2232 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
2233 for &(j, xv) in &act.outliers {
2234 let (w, s) = q4t_outlier(bytes, r, gpr, j);
2235 acc += w * s * xv;
2236 }
2237 unsafe { *out_addr.at(bi * rows + r) = acc };
2239 }
2240 }
2241 };
2242 dispatch_rows(pool, rows, &run);
2243 return;
2244 }
2245 let run = move |start: usize, end: usize| {
2246 for r in start..end {
2247 for bi in 0..b {
2248 let x = &xs_all[bi * cols..(bi + 1) * cols];
2249 unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
2251 }
2252 }
2253 };
2254 dispatch_rows(pool, rows, &run);
2255}
2256
2257fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
2266 (0..gpr)
2267 .map(|gi| {
2268 xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
2269 .iter()
2270 .map(|&v| v as i32)
2271 .sum()
2272 })
2273 .collect()
2274}
2275
2276#[inline]
2280#[allow(unreachable_code)]
2281fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
2282 #[cfg(target_arch = "aarch64")]
2283 unsafe {
2284 return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
2285 }
2286 let _ = gsum;
2287 let mut acc = 0f32;
2288 for gi in 0..gpr {
2289 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
2290 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2291 let mut d = 0i32;
2292 for (j, &b) in tile[2..].iter().enumerate() {
2293 for k in 0..8 {
2294 let w = ((b >> k) & 1) as i32 * 2 - 1;
2295 d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
2296 }
2297 }
2298 acc += d as f32 * s;
2299 }
2300 acc
2301}
2302
2303#[cfg(target_arch = "aarch64")]
2312#[target_feature(enable = "neon,dotprod")]
2313unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
2314 unsafe {
2317 use core::arch::aarch64::*;
2318 use core::arch::asm;
2319 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
2320 let m = vld1q_u8(MASKS.as_ptr());
2321 macro_rules! tile_dot {
2323 ($t:expr, $x:expr) => {{
2324 let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
2325 let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
2326 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
2327 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
2328 let x0 = vld1q_s8($x);
2329 let x1 = vld1q_s8($x.add(16));
2330 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
2331 asm!(
2332 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
2333 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
2334 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
2335 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
2336 options(pure, nomem, nostack),
2337 );
2338 vaddq_s32(a0, a1)
2339 }};
2340 }
2341 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
2342 let xp = xq.as_ptr();
2343 let gp = gsum.as_ptr();
2344 let mut accv = vdupq_n_f32(0.0);
2345 let mut gi = 0;
2346 while gi + 4 <= gpr {
2347 let t0 = base.add(gi * Q1_TILE);
2348 let t1 = base.add((gi + 1) * Q1_TILE);
2349 let t2 = base.add((gi + 2) * Q1_TILE);
2350 let t3 = base.add((gi + 3) * Q1_TILE);
2351 let d0 = tile_dot!(t0, xp.add(gi * GROUP_SIZE));
2352 let d1 = tile_dot!(t1, xp.add((gi + 1) * GROUP_SIZE));
2353 let d2 = tile_dot!(t2, xp.add((gi + 2) * GROUP_SIZE));
2354 let d3 = tile_dot!(t3, xp.add((gi + 3) * GROUP_SIZE));
2355 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
2357 let g = vld1q_s32(gp.add(gi));
2358 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
2359 let sc = [
2360 f16_to_f32(u16::from_le_bytes([*t0, *t0.add(1)])),
2361 f16_to_f32(u16::from_le_bytes([*t1, *t1.add(1)])),
2362 f16_to_f32(u16::from_le_bytes([*t2, *t2.add(1)])),
2363 f16_to_f32(u16::from_le_bytes([*t3, *t3.add(1)])),
2364 ];
2365 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), vld1q_f32(sc.as_ptr()));
2366 gi += 4;
2367 }
2368 let mut acc = vaddvq_f32(accv);
2369 while gi < gpr {
2370 let t = base.add(gi * Q1_TILE);
2371 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2372 let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
2373 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
2374 gi += 1;
2375 }
2376 acc
2377 }
2378}
2379
2380#[inline]
2382fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
2383 let gi = j / GROUP_SIZE;
2384 let k = j % GROUP_SIZE;
2385 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
2386 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2387 let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
2388 ((bit as i32 * 2 - 1) as f32, s)
2389}
2390
2391#[inline]
2393fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
2394 let mut acc = 0f32;
2395 for gi in 0..gpr {
2396 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
2397 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2398 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
2399 let mut ga = 0f32;
2400 for (j, &b) in tile[2..].iter().enumerate() {
2401 for k in 0..8 {
2402 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
2403 }
2404 }
2405 acc += ga * s;
2406 }
2407 acc
2408}
2409
2410#[allow(clippy::too_many_arguments)]
2413fn q1_range_a8w8(
2414 bytes: &[u8],
2415 gpr: usize,
2416 act: &SplitAct,
2417 gsum: &[i32],
2418 out: SendMut,
2419 start: usize,
2420 end: usize,
2421) {
2422 for r in start..end {
2423 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
2424 for &(j, xv) in &act.outliers {
2425 let (w, s) = q1_outlier(bytes, r, gpr, j);
2426 acc += w * s * xv;
2427 }
2428 unsafe { *out.at(r) = acc };
2430 }
2431}
2432
2433fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
2435 for r in start..end {
2436 unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
2438 }
2439}
2440
2441fn q1_matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
2443 debug_assert_eq!(out.len(), rows);
2444 let gpr = cols / GROUP_SIZE;
2445 let out_addr = SendMut(out.as_mut_ptr());
2446 if a8w8_enabled() {
2447 let act = split_act(x);
2448 let gsum = q1_group_sums(&act.xq, gpr);
2449 let (act, gsum) = (&act, &gsum);
2450 let run =
2451 move |start: usize, end: usize| q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end);
2452 dispatch_rows(pool, rows, &run);
2453 return;
2454 }
2455 let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
2456 dispatch_rows(pool, rows, &run);
2457}
2458
2459#[allow(clippy::too_many_arguments)]
2461fn q1_matvec2(
2462 bytes: &[u8],
2463 x1: &[f32],
2464 x2: &[f32],
2465 rows: usize,
2466 cols: usize,
2467 o1: &mut [f32],
2468 o2: &mut [f32],
2469 pool: Option<&Pool>,
2470) {
2471 let gpr = cols / GROUP_SIZE;
2472 let p1 = SendMut(o1.as_mut_ptr());
2473 let p2 = SendMut(o2.as_mut_ptr());
2474 if a8w8_enabled() {
2475 let a1 = split_act(x1);
2476 let a2 = split_act(x2);
2477 let g1 = q1_group_sums(&a1.xq, gpr);
2478 let g2 = q1_group_sums(&a2.xq, gpr);
2479 let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
2480 let run = move |start: usize, end: usize| {
2481 for r in start..end {
2482 let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
2483 let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
2484 for &(j, xv) in &a1.outliers {
2485 let (w, s) = q1_outlier(bytes, r, gpr, j);
2486 v1 += w * s * xv;
2487 }
2488 for &(j, xv) in &a2.outliers {
2489 let (w, s) = q1_outlier(bytes, r, gpr, j);
2490 v2 += w * s * xv;
2491 }
2492 unsafe {
2494 *p1.at(r) = v1;
2495 *p2.at(r) = v2;
2496 }
2497 }
2498 };
2499 dispatch_rows(pool, rows, &run);
2500 return;
2501 }
2502 let run = move |start: usize, end: usize| {
2503 for r in start..end {
2504 unsafe {
2506 *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
2507 *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
2508 }
2509 }
2510 };
2511 dispatch_rows(pool, rows, &run);
2512}
2513
2514#[allow(clippy::too_many_arguments)]
2516fn q1_matmat(
2517 bytes: &[u8],
2518 xs_all: &[f32],
2519 b: usize,
2520 rows: usize,
2521 cols: usize,
2522 out: &mut [f32],
2523 pool: Option<&Pool>,
2524) {
2525 debug_assert_eq!(out.len(), b * rows);
2526 let gpr = cols / GROUP_SIZE;
2527 let out_addr = SendMut(out.as_mut_ptr());
2528 if a8w8_enabled() {
2529 let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
2530 .map(|bi| {
2531 let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
2532 let gsum = q1_group_sums(&act.xq, gpr);
2533 (act, gsum)
2534 })
2535 .collect();
2536 let acts = &acts;
2537 let run = move |start: usize, end: usize| {
2538 for r in start..end {
2539 for (bi, (act, gsum)) in acts.iter().enumerate() {
2540 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
2541 for &(j, xv) in &act.outliers {
2542 let (w, s) = q1_outlier(bytes, r, gpr, j);
2543 acc += w * s * xv;
2544 }
2545 unsafe { *out_addr.at(bi * rows + r) = acc };
2547 }
2548 }
2549 };
2550 dispatch_rows(pool, rows, &run);
2551 return;
2552 }
2553 let run = move |start: usize, end: usize| {
2554 for r in start..end {
2555 for bi in 0..b {
2556 let x = &xs_all[bi * cols..(bi + 1) * cols];
2557 unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
2559 }
2560 }
2561 };
2562 dispatch_rows(pool, rows, &run);
2563}
2564
2565fn q4matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
2571 debug_assert_eq!(out.len(), rows);
2572 let (packed, scales) = q4_split(bytes, rows, cols);
2573 let gpr = cols / GROUP_SIZE;
2574 let out_addr = SendMut(out.as_mut_ptr());
2575
2576 if a8w8_enabled() {
2577 let act = split_act(x);
2578 let run = move |start: usize, end: usize| {
2579 q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
2580 };
2581 dispatch_rows(pool, rows, &run);
2582 return;
2583 }
2584
2585 let run = move |start: usize, end: usize| {
2586 q4_range_f32(packed, scales, gpr, x, out_addr, start, end)
2587 };
2588 dispatch_rows(pool, rows, &run);
2589}
2590
2591#[inline]
2594#[allow(unreachable_code)]
2595fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
2596 #[cfg(target_arch = "aarch64")]
2597 unsafe {
2598 return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
2599 }
2600 #[cfg(target_arch = "x86_64")]
2601 unsafe {
2602 return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
2603 }
2604 let mut acc = 0f32;
2605 for gi in 0..gpr {
2606 let g = g0 + gi;
2607 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
2608 let mut d = 0i32;
2609 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
2610 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
2611 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
2612 }
2613 acc += d as f32 * s;
2614 }
2615 acc
2616}
2617
2618#[inline]
2620#[allow(unreachable_code)]
2621fn dot_q4_row_i8_2(
2622 packed: &[u8],
2623 scales: &[u8],
2624 g0: usize,
2625 gpr: usize,
2626 xq1: &[i8],
2627 xq2: &[i8],
2628) -> (f32, f32) {
2629 #[cfg(target_arch = "aarch64")]
2630 unsafe {
2631 return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
2632 }
2633 #[cfg(target_arch = "x86_64")]
2634 unsafe {
2635 return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
2636 }
2637 (
2638 dot_q4_row_i8(packed, scales, g0, gpr, xq1),
2639 dot_q4_row_i8(packed, scales, g0, gpr, xq2),
2640 )
2641}
2642
2643#[allow(clippy::too_many_arguments)]
2646fn q4_range_a8w8(
2647 packed: &[u8],
2648 scales: &[u8],
2649 gpr: usize,
2650 cols: usize,
2651 act: &SplitAct,
2652 out: SendMut,
2653 start: usize,
2654 end: usize,
2655) {
2656 for r in start..end {
2657 let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
2658 for &(j, xv) in &act.outliers {
2660 let flat = r * cols + j;
2661 let byte = packed[flat / 2];
2662 let nib = if flat & 1 == 0 { byte & 0x0F } else { byte >> 4 };
2663 let s = f16_to_f32(u16::from_le_bytes([
2664 scales[(flat / GROUP_SIZE) * 2],
2665 scales[(flat / GROUP_SIZE) * 2 + 1],
2666 ]));
2667 acc += ((nib as i32 - 8) as f32) * s * xv;
2668 }
2669 unsafe { *out.at(r) = acc };
2671 }
2672}
2673
2674#[allow(clippy::too_many_arguments)]
2677fn q4_range2_a8w8(
2678 packed: &[u8],
2679 scales: &[u8],
2680 gpr: usize,
2681 cols: usize,
2682 a1: &SplitAct,
2683 a2: &SplitAct,
2684 p1: SendMut,
2685 p2: SendMut,
2686 start: usize,
2687 end: usize,
2688) {
2689 for r in start..end {
2690 let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
2691 let mut acc1 = s1 * a1.sx;
2692 let mut acc2 = s2 * a2.sx;
2693 let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
2695 for &(j, xv) in outliers {
2696 let flat = r * cols + j;
2697 let byte = packed[flat / 2];
2698 let nib = if flat & 1 == 0 { byte & 0x0F } else { byte >> 4 };
2699 let s = f16_to_f32(u16::from_le_bytes([
2700 scales[(flat / GROUP_SIZE) * 2],
2701 scales[(flat / GROUP_SIZE) * 2 + 1],
2702 ]));
2703 *acc += ((nib as i32 - 8) as f32) * s * xv;
2704 }
2705 };
2706 fix(&a1.outliers, &mut acc1);
2707 fix(&a2.outliers, &mut acc2);
2708 unsafe {
2710 *p1.at(r) = acc1;
2711 *p2.at(r) = acc2;
2712 }
2713 }
2714}
2715
2716fn q4_range_f32(
2718 packed: &[u8],
2719 scales: &[u8],
2720 gpr: usize,
2721 x: &[f32],
2722 out: SendMut,
2723 start: usize,
2724 end: usize,
2725) {
2726 for r in start..end {
2727 let mut acc = 0f32;
2728 for gi in 0..gpr {
2729 let g = r * gpr + gi;
2730 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
2731 let pk = &packed[g * 16..(g + 1) * 16];
2732 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
2733 let mut ga = 0f32;
2734 for (k, &b) in pk.iter().enumerate() {
2735 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
2736 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
2737 }
2738 acc += ga * s;
2739 }
2740 unsafe { *out.at(r) = acc };
2742 }
2743}
2744
2745#[allow(clippy::too_many_arguments)]
2749fn q4matvec2(
2750 bytes: &[u8],
2751 x1: &[f32],
2752 x2: &[f32],
2753 rows: usize,
2754 cols: usize,
2755 o1: &mut [f32],
2756 o2: &mut [f32],
2757 pool: Option<&Pool>,
2758) {
2759 debug_assert_eq!(o1.len(), rows);
2760 debug_assert_eq!(o2.len(), rows);
2761 let (packed, scales) = q4_split(bytes, rows, cols);
2762 let gpr = cols / GROUP_SIZE;
2763
2764 if a8w8_enabled() {
2765 let a1 = split_act(x1);
2766 let a2 = split_act(x2);
2767 let p1 = SendMut(o1.as_mut_ptr());
2768 let p2 = SendMut(o2.as_mut_ptr());
2769 let run = move |start: usize, end: usize| {
2770 q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
2771 };
2772 dispatch_rows(pool, rows, &run);
2773 return;
2774 }
2775
2776 let p1 = SendMut(o1.as_mut_ptr());
2777 let p2 = SendMut(o2.as_mut_ptr());
2778 let run = move |start: usize, end: usize| {
2779 q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
2780 };
2781 dispatch_rows(pool, rows, &run);
2782}
2783
2784#[allow(clippy::too_many_arguments)]
2786fn q4_range2_f32(
2787 packed: &[u8],
2788 scales: &[u8],
2789 gpr: usize,
2790 x1: &[f32],
2791 x2: &[f32],
2792 p1: SendMut,
2793 p2: SendMut,
2794 start: usize,
2795 end: usize,
2796) {
2797 for r in start..end {
2798 let (mut acc1, mut acc2) = (0f32, 0f32);
2799 for gi in 0..gpr {
2800 let g = r * gpr + gi;
2801 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
2802 let pk = &packed[g * 16..(g + 1) * 16];
2803 let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
2804 let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
2805 let (mut g1, mut g2) = (0f32, 0f32);
2806 for (k, &b) in pk.iter().enumerate() {
2807 let wl = (b & 0x0F) as f32 - 8.0;
2808 let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
2809 g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
2810 g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
2811 }
2812 acc1 += g1 * s;
2813 acc2 += g2 * s;
2814 }
2815 unsafe {
2817 *p1.at(r) = acc1;
2818 *p2.at(r) = acc2;
2819 }
2820 }
2821}
2822
2823thread_local! {
2824 static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
2827 static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
2828}
2829
2830#[allow(clippy::too_many_arguments)]
2836fn q4matmat(
2837 bytes: &[u8],
2838 xs_all: &[f32],
2839 b: usize,
2840 rows: usize,
2841 cols: usize,
2842 out: &mut [f32],
2843 pool: Option<&Pool>,
2844) {
2845 debug_assert_eq!(xs_all.len(), b * cols);
2846 debug_assert_eq!(out.len(), b * rows);
2847 let (packed, scales) = q4_split(bytes, rows, cols);
2848 let gpr = cols / GROUP_SIZE;
2849 let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
2850
2851 if a8w8_enabled() {
2852 let acts: Vec<SplitAct> =
2853 (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
2854 let acts = &acts;
2855 let out_addr = SendMut(out.as_mut_ptr());
2856 let run = move |start: usize, end: usize| {
2857 ROW_I8.with(|rb| {
2858 let mut buf = rb.borrow_mut();
2859 buf.resize(cols, 0);
2860 for r in start..end {
2861 for gi in 0..gpr {
2865 let g = r * gpr + gi;
2866 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
2867 buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
2868 buf[gi * GROUP_SIZE + k * 2 + 1] =
2869 (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
2870 }
2871 }
2872 for (bi, act) in acts.iter().enumerate() {
2873 let mut acc = 0f32;
2874 for gi in 0..gpr {
2875 let d = dot_i8_i8(
2876 &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
2877 &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
2878 );
2879 acc += d as f32 * gscale(r * gpr + gi);
2880 }
2881 acc *= act.sx;
2882 for &(j, xv) in &act.outliers {
2884 acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
2885 }
2886 unsafe { *out_addr.at(bi * rows + r) = acc };
2888 }
2889 }
2890 })
2891 };
2892 dispatch_rows(pool, rows, &run);
2893 return;
2894 }
2895
2896 let out_addr = SendMut(out.as_mut_ptr());
2897 let run = move |start: usize, end: usize| {
2898 ROW_F32.with(|rb| {
2899 let mut buf = rb.borrow_mut();
2900 buf.resize(cols, 0.0);
2901 for r in start..end {
2902 for gi in 0..gpr {
2905 let g = r * gpr + gi;
2906 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
2907 buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
2908 buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
2909 }
2910 }
2911 for bi in 0..b {
2912 let x = &xs_all[bi * cols..(bi + 1) * cols];
2913 let mut acc = 0f32;
2914 for gi in 0..gpr {
2915 let mut ga = 0f32;
2916 for k in 0..GROUP_SIZE / 2 {
2921 let e = gi * GROUP_SIZE + k * 2;
2922 ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
2923 }
2924 acc += ga * gscale(r * gpr + gi);
2925 }
2926 unsafe { *out_addr.at(bi * rows + r) = acc };
2928 }
2929 }
2930 })
2931 };
2932 dispatch_rows(pool, rows, &run);
2933}
2934
2935#[allow(clippy::too_many_arguments)]
2940fn vbitmatmat(
2941 bytes: &[u8],
2942 offsets: &[usize],
2943 xs_all: &[f32],
2944 b: usize,
2945 rows: usize,
2946 cols: usize,
2947 out: &mut [f32],
2948 pool: Option<&Pool>,
2949) {
2950 debug_assert_eq!(xs_all.len(), b * cols);
2951 debug_assert_eq!(out.len(), b * rows);
2952 debug_assert_eq!(offsets.len(), rows + 1);
2953 let ng = cols / GROUP_SIZE;
2954 let bits = &bytes[..rows];
2955 let sc_off = rows;
2956 let gscale = |r: usize, g: usize| {
2957 let so = (r * ng + g) * 2;
2958 f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]))
2959 };
2960
2961 let decode_f32 = |r: usize, dst: &mut [f32]| {
2963 let bw = bits[r] as usize;
2964 let l = ((1i32 << (bw - 1)) - 1) as f32;
2965 let data = &bytes[offsets[r]..offsets[r + 1]];
2966 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
2967 for d in dst.iter_mut() {
2968 while nbits < bw {
2969 acc = (acc << 8) | data[idx] as u64;
2970 idx += 1;
2971 nbits += 8;
2972 }
2973 let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
2974 nbits -= bw;
2975 *d = u - l;
2976 }
2977 };
2978
2979 if a8w8_enabled() {
2980 let acts: Vec<SplitAct> =
2981 (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
2982 let acts = &acts;
2983 let out_addr = SendMut(out.as_mut_ptr());
2984 let run = move |start: usize, end: usize| {
2985 for r in start..end {
2986 let bw = bits[r] as usize;
2987 if bw == 8 {
2988 ROW_F32.with(|rb| {
2991 let mut buf = rb.borrow_mut();
2992 buf.resize(cols, 0.0);
2993 decode_f32(r, &mut buf);
2994 for bi in 0..b {
2995 let x = &xs_all[bi * cols..(bi + 1) * cols];
2996 let mut dot = 0f32;
2997 for g in 0..ng {
2998 let mut gd = 0f32;
2999 for k in 0..GROUP_SIZE {
3000 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
3001 }
3002 dot += gd * gscale(r, g);
3003 }
3004 unsafe { *out_addr.at(bi * rows + r) = dot };
3006 }
3007 });
3008 continue;
3009 }
3010 let l = (1i32 << (bw - 1)) - 1;
3011 let data = &bytes[offsets[r]..offsets[r + 1]];
3012 ROW_I8.with(|rb| {
3013 let mut buf = rb.borrow_mut();
3014 buf.resize(cols, 0);
3015 #[inline(always)]
3016 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3017 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3018 let u = unpack8::<B>(&data[blk * B..]);
3019 for k in 0..8 {
3020 chunk[k] = (u[k] - l) as i8 as u8;
3021 }
3022 }
3023 }
3024 match bw {
3025 3 => fill::<3>(data, l, &mut buf),
3026 4 => vbit_fill4(data, &mut buf),
3027 5 => fill::<5>(data, l, &mut buf),
3028 6 => fill::<6>(data, l, &mut buf),
3029 _ => unreachable!("vbit bit-width {bw} (validated at load)"),
3030 }
3031 for (bi, act) in acts.iter().enumerate() {
3032 let mut dot = 0f32;
3033 for g in 0..ng {
3034 let d = dot_i8_i8(
3035 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3036 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3037 ) as f32
3038 * act.sx;
3039 dot += d * gscale(r, g);
3040 }
3041 for &(j, xv) in &act.outliers {
3042 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
3043 }
3044 unsafe { *out_addr.at(bi * rows + r) = dot };
3046 }
3047 });
3048 }
3049 };
3050 dispatch_rows(pool, rows, &run);
3051 return;
3052 }
3053
3054 let out_addr = SendMut(out.as_mut_ptr());
3055 let run = move |start: usize, end: usize| {
3056 ROW_F32.with(|rb| {
3057 let mut buf = rb.borrow_mut();
3058 buf.resize(cols, 0.0);
3059 for r in start..end {
3060 decode_f32(r, &mut buf);
3061 for bi in 0..b {
3062 let x = &xs_all[bi * cols..(bi + 1) * cols];
3063 let mut dot = 0f32;
3064 for g in 0..ng {
3065 let mut gd = 0f32;
3066 for k in 0..GROUP_SIZE {
3067 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
3068 }
3069 dot += gd * gscale(r, g);
3070 }
3071 unsafe { *out_addr.at(bi * rows + r) = dot };
3073 }
3074 }
3075 })
3076 };
3077 dispatch_rows(pool, rows, &run);
3078}
3079
3080pub(crate) fn gpu_batch_job<'a>(
3084 t: &'a QTensor,
3085 x: &[f32],
3086) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
3087 match t {
3088 QTensor::Mapped {
3089 model,
3090 idx,
3091 dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
3092 rows,
3093 cols,
3094 row_scale,
3095 col_field,
3096 ..
3097 } => Some((
3098 model.clone(),
3099 crate::gpu::BatchJob {
3100 idx: *idx,
3101 rows: *rows,
3102 cols: *cols,
3103 row_scale,
3104 xs: prescale(x, col_field, *dt).into_owned(),
3105 q1: false,
3106 },
3107 )),
3108 QTensor::Mapped {
3110 model,
3111 idx,
3112 dtype: TensorDtype::Q1,
3113 rows,
3114 cols,
3115 ..
3116 } => Some((
3117 model.clone(),
3118 crate::gpu::BatchJob {
3119 idx: *idx,
3120 rows: *rows,
3121 cols: *cols,
3122 row_scale: &[],
3123 xs: x.to_vec(),
3124 q1: true,
3125 },
3126 )),
3127 _ => None,
3128 }
3129}
3130
3131pub(crate) fn prescale<'a>(
3135 x: &'a [f32],
3136 col_field: &[f32],
3137 dtype: TensorDtype,
3138) -> std::borrow::Cow<'a, [f32]> {
3139 if dtype == TensorDtype::Q8_2f {
3140 x.iter().zip(col_field).map(|(a, c)| a * c).collect()
3141 } else {
3142 std::borrow::Cow::Borrowed(x)
3143 }
3144}
3145
3146#[cfg(target_arch = "x86_64")]
3151pub(crate) fn avx2_enabled() -> bool {
3152 use std::sync::OnceLock;
3153 static ON: OnceLock<bool> = OnceLock::new();
3154 *ON.get_or_init(|| {
3155 std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
3156 && std::arch::is_x86_feature_detected!("avx2")
3157 && std::arch::is_x86_feature_detected!("fma")
3158 })
3159}
3160
3161#[cfg(target_arch = "x86_64")]
3166fn avx2_a8w8_enabled() -> bool {
3167 use std::sync::OnceLock;
3168 static ON: OnceLock<bool> = OnceLock::new();
3169 *ON.get_or_init(|| {
3170 avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
3171 })
3172}
3173
3174#[inline]
3178pub(crate) fn a8w8_enabled() -> bool {
3179 #[cfg(target_arch = "aarch64")]
3180 {
3181 sdot_enabled()
3182 }
3183 #[cfg(target_arch = "x86_64")]
3184 {
3185 avx2_a8w8_enabled()
3186 }
3187 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
3188 {
3189 false
3190 }
3191}
3192
3193#[inline]
3196#[allow(unreachable_code)]
3197fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
3198 #[cfg(target_arch = "aarch64")]
3199 unsafe {
3200 return dot_i8_sdot(w, xq);
3201 }
3202 #[cfg(target_arch = "x86_64")]
3203 unsafe {
3204 if avx512vnni_enabled() {
3205 return dot_i8_i8_vnni(w, xq);
3206 }
3207 return dot_i8_i8_avx2(w, xq);
3208 }
3209 w.iter().zip(xq).map(|(&a, &b)| (a as i8) as i32 * b as i32).sum()
3210}
3211
3212#[cfg(target_arch = "x86_64")]
3216fn avx512vnni_enabled() -> bool {
3217 use std::sync::OnceLock;
3218 static ON: OnceLock<bool> = OnceLock::new();
3219 *ON.get_or_init(|| {
3220 std::env::var("CMF_AVX512").map(|v| v != "0").unwrap_or(true)
3221 && std::arch::is_x86_feature_detected!("avx512f")
3222 && std::arch::is_x86_feature_detected!("avx512bw")
3223 && std::arch::is_x86_feature_detected!("avx512vl")
3224 && std::arch::is_x86_feature_detected!("avx512vnni")
3225 })
3226}
3227
3228#[cfg(target_arch = "x86_64")]
3233#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
3234unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
3235 unsafe {
3237 use core::arch::x86_64::*;
3238 let n = w.len();
3239 let mut j = 0usize;
3240 let mut total: i32;
3241 {
3246 #[inline(always)]
3247 unsafe fn step(
3248 w: *const u8,
3249 x: *const i8,
3250 acc: core::arch::x86_64::__m512i,
3251 ) -> core::arch::x86_64::__m512i {
3252 unsafe {
3253 use core::arch::x86_64::*;
3254 let wv = _mm512_loadu_si512(w as *const _);
3255 let xv = _mm512_loadu_si512(x as *const _);
3256 let aw = _mm512_abs_epi8(wv);
3257 let neg = _mm512_movepi8_mask(wv);
3258 let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
3259 _mm512_dpbusd_epi32(acc, aw, sx)
3260 }
3261 }
3262 let (mut a0, mut a1, mut a2, mut a3) = (
3263 _mm512_setzero_si512(),
3264 _mm512_setzero_si512(),
3265 _mm512_setzero_si512(),
3266 _mm512_setzero_si512(),
3267 );
3268 while j + 256 <= n {
3269 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
3270 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
3271 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
3272 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
3273 j += 256;
3274 }
3275 while j + 64 <= n {
3276 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
3277 j += 64;
3278 }
3279 let s01 = _mm512_add_epi32(a0, a1);
3280 let s23 = _mm512_add_epi32(a2, a3);
3281 total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
3282 }
3283 if j + 32 <= n {
3285 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
3286 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
3287 let d = _mm256_dpbusd_epi32(
3288 _mm256_setzero_si256(),
3289 _mm256_abs_epi8(wv),
3290 _mm256_sign_epi8(xv, wv),
3291 );
3292 let hi128 = _mm256_extracti128_si256::<1>(d);
3293 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3294 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3295 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3296 total += _mm_cvtsi128_si32(s32);
3297 j += 32;
3298 }
3299 while j < n {
3300 total += (w[j] as i8) as i32 * xq[j] as i32;
3301 j += 1;
3302 }
3303 total
3304 }
3305}
3306
3307#[cfg(target_arch = "x86_64")]
3309#[target_feature(enable = "avx2,fma")]
3310unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
3311 unsafe {
3313 use core::arch::x86_64::*;
3314 let n = x.len();
3315 let wp = w.as_ptr();
3316 let xp = x.as_ptr();
3317 let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
3318 let mut j = 0usize;
3319 while j + 16 <= n {
3320 let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
3321 let lo = _mm256_cvtepi8_epi32(wb);
3322 let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
3323 a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
3324 a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
3325 j += 16;
3326 }
3327 let acc = _mm256_add_ps(a0, a1);
3328 let hi128 = _mm256_extractf128_ps::<1>(acc);
3329 let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
3330 let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
3331 let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
3332 let mut sum = _mm_cvtss_f32(s32);
3333 while j < n {
3334 sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
3335 j += 1;
3336 }
3337 sum
3338 }
3339}
3340
3341#[cfg(target_arch = "x86_64")]
3346#[target_feature(enable = "avx2")]
3347unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
3348 unsafe {
3350 use core::arch::x86_64::*;
3351 let n = w.len();
3352 let ones = _mm256_set1_epi16(1);
3353 let mut acc = _mm256_setzero_si256();
3354 let mut j = 0usize;
3355 while j + 32 <= n {
3356 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
3357 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
3358 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
3359 acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
3360 j += 32;
3361 }
3362 let hi128 = _mm256_extracti128_si256::<1>(acc);
3363 let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
3364 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3365 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3366 let mut s = _mm_cvtsi128_si32(s32);
3367 while j < n {
3368 s += (w[j] as i8) as i32 * xq[j] as i32;
3369 j += 1;
3370 }
3371 s
3372 }
3373}
3374
3375#[cfg(target_arch = "x86_64")]
3380#[inline]
3381fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
3382 let dot = if avx512vnni_enabled() && row.len() >= 64 {
3383 (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
3384 } else {
3385 unsafe { dot_i8_i8_avx2(row, &act.xq) }
3386 };
3387 let mut acc = dot as f32 * act.sx;
3388 for &(j, xv) in &act.outliers {
3389 acc += (row[j] as i8) as f32 * xv;
3390 }
3391 acc
3392}
3393
3394#[cfg(target_arch = "x86_64")]
3398#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
3399unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
3400 unsafe {
3402 use core::arch::x86_64::*;
3403 let n = w.len();
3404 let flip = _mm512_set1_epi8(-128); #[inline(always)]
3406 unsafe fn step(
3407 w: *const u8,
3408 x: *const i8,
3409 flip: core::arch::x86_64::__m512i,
3410 acc: core::arch::x86_64::__m512i,
3411 ) -> core::arch::x86_64::__m512i {
3412 unsafe {
3413 use core::arch::x86_64::*;
3414 let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
3415 _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
3416 }
3417 }
3418 let (mut a0, mut a1, mut a2, mut a3) = (
3419 _mm512_setzero_si512(),
3420 _mm512_setzero_si512(),
3421 _mm512_setzero_si512(),
3422 _mm512_setzero_si512(),
3423 );
3424 let mut j = 0usize;
3425 while j + 256 <= n {
3426 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
3427 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
3428 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
3429 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
3430 j += 256;
3431 }
3432 while j + 64 <= n {
3433 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
3434 j += 64;
3435 }
3436 let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
3437 _mm512_add_epi32(a0, a1),
3438 _mm512_add_epi32(a2, a3),
3439 ));
3440 while j < n {
3442 total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
3443 j += 1;
3444 }
3445 total
3446 }
3447}
3448
3449#[cfg(target_arch = "x86_64")]
3455#[target_feature(enable = "avx2")]
3456unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
3457 unsafe {
3460 use core::arch::x86_64::*;
3461 let lomask = _mm_set1_epi8(0x0F);
3462 let eight = _mm256_set1_epi8(8);
3463 let ones = _mm256_set1_epi16(1);
3464 let mut acc = 0f32;
3465 for gi in 0..gpr {
3466 let g = g0 + gi;
3467 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
3468 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
3469 let lo = _mm_and_si128(b, lomask);
3470 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3471 let w = _mm256_sub_epi8(
3472 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3473 eight,
3474 );
3475 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3476 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
3477 let d = _mm256_madd_epi16(p16, ones);
3478 let hi128 = _mm256_extracti128_si256::<1>(d);
3479 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3480 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3481 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3482 acc += _mm_cvtsi128_si32(s32) as f32 * s;
3483 }
3484 acc
3485 }
3486}
3487
3488#[cfg(target_arch = "x86_64")]
3491#[target_feature(enable = "avx2")]
3492unsafe fn dot_q4_row_avx2_2(
3493 packed: &[u8],
3494 scales: &[u8],
3495 g0: usize,
3496 gpr: usize,
3497 xq1: &[i8],
3498 xq2: &[i8],
3499) -> (f32, f32) {
3500 unsafe {
3502 use core::arch::x86_64::*;
3503 let lomask = _mm_set1_epi8(0x0F);
3504 let eight = _mm256_set1_epi8(8);
3505 let ones = _mm256_set1_epi16(1);
3506 let (mut acc1, mut acc2) = (0f32, 0f32);
3507 #[inline(always)]
3508 unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
3509 unsafe {
3510 use core::arch::x86_64::*;
3511 let hi128 = _mm256_extracti128_si256::<1>(d);
3512 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3513 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3514 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3515 _mm_cvtsi128_si32(s32)
3516 }
3517 }
3518 for gi in 0..gpr {
3519 let g = g0 + gi;
3520 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
3521 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
3522 let lo = _mm_and_si128(b, lomask);
3523 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3524 let w = _mm256_sub_epi8(
3525 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3526 eight,
3527 );
3528 let aw = _mm256_abs_epi8(w);
3529 let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3530 let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3531 let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
3532 let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
3533 acc1 += hsum(d1) as f32 * s;
3534 acc2 += hsum(d2) as f32 * s;
3535 }
3536 (acc1, acc2)
3537 }
3538}
3539
3540#[cfg(target_arch = "x86_64")]
3542fn q8_range_avx2(
3543 q: &[u8],
3544 row_scale: &[f32],
3545 act: &SplitAct,
3546 cols: usize,
3547 out_addr: SendMut,
3548 start: usize,
3549 end: usize,
3550) {
3551 for o in start..end {
3552 let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
3553 unsafe { *out_addr.at(o) = v };
3555 }
3556}
3557
3558#[cfg(target_arch = "x86_64")]
3560#[allow(clippy::too_many_arguments)]
3561fn q8_range2_avx2(
3562 q: &[u8],
3563 row_scale: &[f32],
3564 a1: &SplitAct,
3565 a2: &SplitAct,
3566 cols: usize,
3567 p1: SendMut,
3568 p2: SendMut,
3569 start: usize,
3570 end: usize,
3571) {
3572 for o in start..end {
3573 let row = &q[o * cols..(o + 1) * cols];
3574 unsafe {
3576 *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
3577 *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
3578 }
3579 }
3580}
3581
3582#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
3588fn sdot_enabled() -> bool {
3589 use std::sync::OnceLock;
3590 static ON: OnceLock<bool> = OnceLock::new();
3591 *ON.get_or_init(|| {
3592 let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
3593 #[cfg(target_arch = "aarch64")]
3594 {
3595 want && std::arch::is_aarch64_feature_detected!("dotprod")
3596 }
3597 #[cfg(not(target_arch = "aarch64"))]
3598 {
3599 let _ = want;
3600 false
3601 }
3602 })
3603}
3604
3605struct SplitAct {
3610 xq: Vec<i8>,
3611 sx: f32,
3612 outliers: Vec<(usize, f32)>,
3613 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
3616 xsum: i32,
3617}
3618
3619thread_local! {
3620 static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
3623 const { std::cell::RefCell::new(Vec::new()) };
3624}
3625
3626impl Drop for SplitAct {
3627 fn drop(&mut self) {
3628 let buf = std::mem::take(&mut self.xq);
3629 if buf.capacity() > 0 {
3630 XQ_FREE.with(|f| {
3631 let mut f = f.borrow_mut();
3632 if f.len() < 16 {
3633 f.push(buf);
3634 }
3635 });
3636 }
3637 }
3638}
3639
3640fn split_act(x: &[f32]) -> SplitAct {
3641 let n = x.len();
3642 let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
3643 let thr = 8.0 * rms;
3644 let mut outliers: Vec<(usize, f32)> = Vec::new();
3648 let mut amax = 0f32;
3649 for (j, &v) in x.iter().enumerate() {
3650 let a = v.abs();
3651 if a > thr {
3652 outliers.push((j, v));
3653 } else if a > amax {
3654 amax = a;
3655 }
3656 }
3657 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
3658 let inv = 1.0 / sx;
3659 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
3660 xq.clear();
3661 xq.reserve(n);
3662 if outliers.is_empty() {
3663 xq.extend(x.iter().map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8));
3664 } else {
3665 xq.extend(x.iter().map(|&v| {
3667 if v.abs() > thr {
3668 0
3669 } else {
3670 (v * inv).round().clamp(-127.0, 127.0) as i8
3671 }
3672 }));
3673 }
3674 let xsum = xq.iter().map(|&v| v as i32).sum();
3675 SplitAct { xq, sx, outliers, xsum }
3676}
3677
3678#[cfg(target_arch = "aarch64")]
3681#[target_feature(enable = "neon,dotprod")]
3682unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
3683 unsafe {
3685 use core::arch::aarch64::*;
3686 use core::arch::asm;
3687 let wp = w.as_ptr() as *const i8;
3688 let n = w.len();
3689 let (mut a0, mut a1, mut a2, mut a3) =
3690 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
3691 let mut i = 0;
3692 while i + 64 <= n {
3693 let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
3694 let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
3695 let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
3696 let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
3697 asm!(
3698 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
3699 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
3700 "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
3701 "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
3702 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
3703 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
3704 w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
3705 options(pure, nomem, nostack),
3706 );
3707 i += 64;
3708 }
3709 while i + 16 <= n {
3710 let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
3711 asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
3712 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
3713 i += 16;
3714 }
3715 let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
3716 while i < n {
3717 s += (*wp.add(i)) as i32 * xq[i] as i32;
3718 i += 1;
3719 }
3720 s
3721}
3722}
3723
3724#[cfg(target_arch = "aarch64")]
3728#[target_feature(enable = "neon,dotprod")]
3729unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
3730 unsafe {
3732 use core::arch::aarch64::*;
3733 use core::arch::asm;
3734 let n = xq.len();
3735 let px = xq.as_ptr();
3736 let (p0, p1, p2, p3) = (
3737 w0.as_ptr() as *const i8,
3738 w1.as_ptr() as *const i8,
3739 w2.as_ptr() as *const i8,
3740 w3.as_ptr() as *const i8,
3741 );
3742 let (mut a0, mut a1, mut a2, mut a3) =
3743 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
3744 let mut i = 0;
3745 while i + 16 <= n {
3746 let x = vld1q_s8(px.add(i));
3747 let v0 = vld1q_s8(p0.add(i));
3748 let v1 = vld1q_s8(p1.add(i));
3749 let v2 = vld1q_s8(p2.add(i));
3750 let v3 = vld1q_s8(p3.add(i));
3751 asm!(
3752 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
3753 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
3754 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
3755 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
3756 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
3757 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
3758 options(pure, nomem, nostack),
3759 );
3760 i += 16;
3761 }
3762 let mut r = [vaddvq_s32(a0), vaddvq_s32(a1), vaddvq_s32(a2), vaddvq_s32(a3)];
3763 while i < n {
3764 let xi = *px.add(i) as i32;
3765 r[0] += (*p0.add(i)) as i32 * xi;
3766 r[1] += (*p1.add(i)) as i32 * xi;
3767 r[2] += (*p2.add(i)) as i32 * xi;
3768 r[3] += (*p3.add(i)) as i32 * xi;
3769 i += 1;
3770 }
3771 r
3772}
3773}
3774
3775#[cfg(target_arch = "aarch64")]
3782#[target_feature(enable = "neon,dotprod")]
3783unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
3784 unsafe {
3787 use core::arch::aarch64::*;
3788 use core::arch::asm;
3789 let n = xq.len();
3790 let px = xq.as_ptr();
3791 let pg = g.as_ptr() as *const i8;
3792 let (mut a0, mut a1, mut a2, mut a3) =
3793 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
3794 let mut i = 0;
3795 while i + 16 <= n {
3796 let x = vld1q_s8(px.add(i));
3797 let base = pg.add(4 * i);
3798 let v0 = vld1q_s8(base);
3799 let v1 = vld1q_s8(base.add(16));
3800 let v2 = vld1q_s8(base.add(32));
3801 let v3 = vld1q_s8(base.add(48));
3802 asm!(
3803 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
3804 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
3805 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
3806 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
3807 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
3808 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
3809 options(pure, nomem, nostack),
3810 );
3811 i += 16;
3812 }
3813 [vaddvq_s32(a0), vaddvq_s32(a1), vaddvq_s32(a2), vaddvq_s32(a3)]
3814 }
3815}
3816
3817#[cfg(target_arch = "aarch64")]
3823fn q8_range_sdot(
3824 q: &[u8],
3825 rep: &[u8],
3826 row_scale: &[f32],
3827 act: &SplitAct,
3828 cols: usize,
3829 out_addr: SendMut,
3830 start: usize,
3831 end: usize,
3832) {
3833 let mut o = start;
3834 if !rep.is_empty() {
3837 while o < end && o % 4 != 0 {
3838 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
3839 unsafe { *out_addr.at(o) = v };
3840 o += 1;
3841 }
3842 }
3843 while o + 4 <= end {
3844 let r = if rep.is_empty() {
3845 unsafe {
3846 dot_i8_sdot_4rows(
3847 &q[o * cols..(o + 1) * cols],
3848 &q[(o + 1) * cols..(o + 2) * cols],
3849 &q[(o + 2) * cols..(o + 3) * cols],
3850 &q[(o + 3) * cols..(o + 4) * cols],
3851 &act.xq,
3852 )
3853 }
3854 } else {
3855 unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
3856 };
3857 for k in 0..4 {
3858 let mut acc = r[k] as f32 * act.sx;
3859 for &(j, xv) in &act.outliers {
3860 acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
3861 }
3862 unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
3864 }
3865 o += 4;
3866 }
3867 while o < end {
3868 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
3869 unsafe { *out_addr.at(o) = v };
3870 o += 1;
3871 }
3872}
3873
3874#[cfg(target_arch = "aarch64")]
3877#[allow(clippy::too_many_arguments)]
3878fn q8_range2_sdot(
3879 q: &[u8],
3880 row_scale: &[f32],
3881 a1: &SplitAct,
3882 a2: &SplitAct,
3883 cols: usize,
3884 p1: SendMut,
3885 p2: SendMut,
3886 start: usize,
3887 end: usize,
3888) {
3889 for o in start..end {
3890 let row = &q[o * cols..(o + 1) * cols];
3891 unsafe {
3893 *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
3894 *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
3895 }
3896 }
3897}
3898
3899#[allow(clippy::too_many_arguments)]
3901fn q8_range2_f32(
3902 q: &[u8],
3903 row_scale: &[f32],
3904 x1: &[f32],
3905 x2: &[f32],
3906 cols: usize,
3907 p1: SendMut,
3908 p2: SendMut,
3909 start: usize,
3910 end: usize,
3911) {
3912 for o in start..end {
3913 let row = &q[o * cols..(o + 1) * cols];
3914 unsafe {
3916 *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
3917 *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
3918 }
3919 }
3920}
3921
3922fn q8_range_f32(
3924 q: &[u8],
3925 row_scale: &[f32],
3926 xs: &[f32],
3927 cols: usize,
3928 out_addr: SendMut,
3929 start: usize,
3930 end: usize,
3931) {
3932 for o in start..end {
3933 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
3934 unsafe { *out_addr.at(o) = v };
3936 }
3937}
3938
3939#[cfg(target_arch = "aarch64")]
3942#[inline]
3943fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
3944 let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
3945 for &(j, xv) in &act.outliers {
3946 acc += (row[j] as i8) as f32 * xv;
3947 }
3948 acc
3949}
3950
3951#[cfg(target_arch = "aarch64")]
3959#[target_feature(enable = "neon,dotprod")]
3960unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
3961 unsafe {
3964 use core::arch::aarch64::*;
3965 use core::arch::asm;
3966 let lomask = vdupq_n_u8(0x0F);
3967 let eight = vdupq_n_s8(8);
3968 let mut acc = 0f32;
3969 for gi in 0..gpr {
3970 let g = g0 + gi;
3971 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
3972 let b = vld1q_u8(packed.as_ptr().add(g * 16));
3973 let lo = vandq_u8(b, lomask);
3974 let hi = vshrq_n_u8::<4>(b);
3975 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
3976 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
3977 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
3978 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
3979 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3980 asm!(
3981 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
3982 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
3983 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3984 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
3985 options(pure, nomem, nostack),
3986 );
3987 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
3988 }
3989 acc
3990 }
3991}
3992
3993#[cfg(target_arch = "aarch64")]
3998#[target_feature(enable = "neon,dotprod")]
3999unsafe fn dot_q4_row_sdot2(
4000 packed: &[u8],
4001 scales: &[u8],
4002 g0: usize,
4003 gpr: usize,
4004 xq1: &[i8],
4005 xq2: &[i8],
4006) -> (f32, f32) {
4007 unsafe {
4010 use core::arch::aarch64::*;
4011 use core::arch::asm;
4012 let lomask = vdupq_n_u8(0x0F);
4013 let eight = vdupq_n_s8(8);
4014 let (mut acc1, mut acc2) = (0f32, 0f32);
4015 for gi in 0..gpr {
4016 let g = g0 + gi;
4017 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4018 let b = vld1q_u8(packed.as_ptr().add(g * 16));
4019 let lo = vandq_u8(b, lomask);
4020 let hi = vshrq_n_u8::<4>(b);
4021 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4022 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4023 let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
4024 let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
4025 let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
4026 let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
4027 let (mut a0, mut a1, mut b0, mut b1) =
4028 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
4029 asm!(
4030 "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
4031 "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
4032 "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
4033 "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
4034 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4035 b0 = inout(vreg) b0, b1 = inout(vreg) b1,
4036 e0 = in(vreg) e0, e1 = in(vreg) e1,
4037 x10 = in(vreg) x10, x11 = in(vreg) x11,
4038 x20 = in(vreg) x20, x21 = in(vreg) x21,
4039 options(pure, nomem, nostack),
4040 );
4041 acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4042 acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
4043 }
4044 (acc1, acc2)
4045 }
4046}
4047
4048#[inline]
4053pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
4054 #[cfg(target_arch = "aarch64")]
4055 unsafe {
4056 return axpy_i8_f32_neon(acc, row, w);
4057 }
4058 #[cfg(target_arch = "x86_64")]
4059 if avx2_enabled() {
4060 return unsafe { axpy_i8_f32_avx2(acc, row, w) };
4061 }
4062 #[allow(unreachable_code)]
4063 {
4064 for (a, &b) in acc.iter_mut().zip(row) {
4065 *a += w * b as f32;
4066 }
4067 }
4068}
4069
4070#[cfg(target_arch = "x86_64")]
4072#[target_feature(enable = "avx2,fma")]
4073unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
4074 unsafe {
4076 use core::arch::x86_64::*;
4077 let n = acc.len().min(row.len());
4078 let ap = acc.as_mut_ptr();
4079 let rp = row.as_ptr();
4080 let wv = _mm256_set1_ps(w);
4081 let mut j = 0usize;
4082 while j + 16 <= n {
4083 let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
4084 let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
4085 let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
4086 let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
4087 let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
4088 _mm256_storeu_ps(ap.add(j), v0);
4089 _mm256_storeu_ps(ap.add(j + 8), v1);
4090 j += 16;
4091 }
4092 while j < n {
4093 *ap.add(j) += w * (*rp.add(j)) as f32;
4094 j += 1;
4095 }
4096 }
4097}
4098
4099#[cfg(target_arch = "aarch64")]
4100#[target_feature(enable = "neon")]
4101unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
4102 unsafe {
4104 use core::arch::aarch64::*;
4105 let n = acc.len().min(row.len());
4106 let ap = acc.as_mut_ptr();
4107 let rp = row.as_ptr();
4108 let wv = vdupq_n_f32(w);
4109 let mut j = 0usize;
4110 while j + 16 <= n {
4111 let rb = vld1q_s8(rp.add(j));
4112 let lo = vmovl_s8(vget_low_s8(rb));
4113 let hi = vmovl_s8(vget_high_s8(rb));
4114 for (off, half) in [(0, lo), (8, hi)] {
4115 let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
4116 let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
4117 let o = j + off;
4118 vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
4119 vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
4120 }
4121 j += 16;
4122 }
4123 while j < n {
4124 *ap.add(j) += w * (*rp.add(j)) as f32;
4125 j += 1;
4126 }
4127}
4128}
4129
4130#[inline]
4133pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
4134 #[cfg(target_arch = "aarch64")]
4135 unsafe {
4136 return dot_i8_f32_neon(w, x);
4137 }
4138 #[cfg(target_arch = "x86_64")]
4139 if avx2_enabled() {
4140 return unsafe { dot_i8_f32_avx2(w, x) };
4141 }
4142 #[allow(unreachable_code)]
4143 {
4144 let mut sum = 0.0f32;
4145 for (j, &b) in w.iter().enumerate() {
4146 sum += (b as i8) as f32 * x[j];
4147 }
4148 sum
4149 }
4150}
4151
4152#[inline]
4156fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
4157 #[cfg(target_arch = "aarch64")]
4158 unsafe {
4159 return dot_i8_col_f32_neon(w, x, col);
4160 }
4161 #[allow(unreachable_code)]
4162 {
4163 let mut sum = 0.0f32;
4164 for (j, &b) in w.iter().enumerate() {
4165 sum += (b as i8) as f32 * x[j] * col[j];
4166 }
4167 sum
4168 }
4169}
4170
4171#[cfg(target_arch = "aarch64")]
4172#[target_feature(enable = "neon")]
4173unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
4174 unsafe {
4176 use core::arch::aarch64::*;
4177 let n = x.len();
4178 let wp = w.as_ptr() as *const i8;
4179 let xp = x.as_ptr();
4180 let cp = col.as_ptr();
4181 let (mut a0, mut a1, mut a2, mut a3) =
4182 (vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4183 let mut j = 0usize;
4184 while j + 16 <= n {
4185 let wb = vld1q_s8(wp.add(j));
4186 let lo = vmovl_s8(vget_low_s8(wb));
4187 let hi = vmovl_s8(vget_high_s8(wb));
4188 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
4189 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
4190 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
4191 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
4192 a0 = vfmaq_f32(a0, w0, vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))));
4193 a1 = vfmaq_f32(a1, w1, vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))));
4194 a2 = vfmaq_f32(a2, w2, vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))));
4195 a3 = vfmaq_f32(a3, w3, vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))));
4196 j += 16;
4197 }
4198 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
4199 while j < n {
4200 sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
4201 j += 1;
4202 }
4203 sum
4204}
4205}
4206
4207#[cfg(target_arch = "aarch64")]
4208#[target_feature(enable = "neon")]
4209unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
4210 unsafe {
4212 use core::arch::aarch64::*;
4213 let n = x.len();
4214 let wp = w.as_ptr() as *const i8;
4215 let xp = x.as_ptr();
4216 let (mut a0, mut a1, mut a2, mut a3) =
4217 (vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4218 let mut j = 0usize;
4219 while j + 16 <= n {
4220 let wb = vld1q_s8(wp.add(j));
4221 let lo = vmovl_s8(vget_low_s8(wb));
4222 let hi = vmovl_s8(vget_high_s8(wb));
4223 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
4224 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
4225 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
4226 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
4227 a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
4228 a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
4229 a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
4230 a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
4231 j += 16;
4232 }
4233 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
4234 while j < n {
4235 sum += (*wp.add(j)) as f32 * *xp.add(j);
4236 j += 1;
4237 }
4238 sum
4239}
4240}
4241
4242#[allow(clippy::too_many_arguments)]
4243fn qmatvec(
4244 q: &[u8],
4245 rep: &[u8],
4246 row_scale: &[f32],
4247 xs: &[f32],
4248 rows: usize,
4249 cols: usize,
4250 out: &mut [f32],
4251 pool: Option<&Pool>,
4252) {
4253 debug_assert_eq!(out.len(), rows);
4254 #[cfg(not(target_arch = "aarch64"))]
4255 let _ = rep;
4256
4257 #[cfg(target_arch = "aarch64")]
4258 if sdot_enabled() {
4259 let act = split_act(xs);
4260 let out_addr = SendMut(out.as_mut_ptr());
4261 let run_range = |start: usize, end: usize| {
4262 q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
4263 };
4264 match pool {
4265 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
4266 _ => run_range(0, rows),
4267 }
4268 return;
4269 }
4270 #[cfg(target_arch = "x86_64")]
4273 if avx2_a8w8_enabled() {
4274 let act = split_act(xs);
4275 let out_addr = SendMut(out.as_mut_ptr());
4276 let run_range =
4277 |start: usize, end: usize| q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end);
4278 match pool {
4279 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
4280 _ => run_range(0, rows),
4281 }
4282 return;
4283 }
4284
4285 let row_dot = |o: usize| -> f32 { dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o] };
4286 match pool {
4287 Some(pool) if rows >= 256 => {
4288 let out_addr = SendMut(out.as_mut_ptr());
4289 pool.run(&move |widx, n| {
4290 let chunk = rows.div_ceil(n);
4291 let start = widx * chunk;
4292 let end = (start + chunk).min(rows);
4293 for o in start..end {
4294 unsafe { *out_addr.at(o) = row_dot(o) };
4296 }
4297 });
4298 }
4299 _ => {
4300 for (o, dst) in out.iter_mut().enumerate() {
4301 *dst = row_dot(o);
4302 }
4303 }
4304 }
4305}
4306
4307#[allow(clippy::too_many_arguments)]
4308fn qmatvec2(
4309 q: &[u8],
4310 row_scale: &[f32],
4311 x1: &[f32],
4312 x2: &[f32],
4313 rows: usize,
4314 cols: usize,
4315 o1: &mut [f32],
4316 o2: &mut [f32],
4317 pool: Option<&Pool>,
4318) {
4319 #[cfg(target_arch = "aarch64")]
4320 if sdot_enabled() {
4321 let a1s = split_act(x1);
4322 let a2s = split_act(x2);
4323 let p1 = SendMut(o1.as_mut_ptr());
4324 let p2 = SendMut(o2.as_mut_ptr());
4325 let run_range = |start: usize, end: usize| {
4326 q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
4327 };
4328 match pool {
4329 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
4330 _ => run_range(0, rows),
4331 }
4332 return;
4333 }
4334 #[cfg(target_arch = "x86_64")]
4335 if avx2_a8w8_enabled() {
4336 let a1s = split_act(x1);
4337 let a2s = split_act(x2);
4338 let p1 = SendMut(o1.as_mut_ptr());
4339 let p2 = SendMut(o2.as_mut_ptr());
4340 let run_range = |start: usize, end: usize| {
4341 q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
4342 };
4343 match pool {
4344 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
4345 _ => run_range(0, rows),
4346 }
4347 return;
4348 }
4349
4350 let row_dots = |o: usize| -> (f32, f32) {
4351 let row = &q[o * cols..(o + 1) * cols];
4352 (dot_i8_f32(row, x1) * row_scale[o], dot_i8_f32(row, x2) * row_scale[o])
4353 };
4354 match pool {
4355 Some(pool) if rows >= 256 => {
4356 let p1 = SendMut(o1.as_mut_ptr());
4357 let p2 = SendMut(o2.as_mut_ptr());
4358 pool.run(&move |widx, n| {
4359 let chunk = rows.div_ceil(n);
4360 let start = widx * chunk;
4361 let end = (start + chunk).min(rows);
4362 for o in start..end {
4363 let (s1, s2) = row_dots(o);
4364 unsafe {
4366 *p1.at(o) = s1;
4367 *p2.at(o) = s2;
4368 }
4369 }
4370 });
4371 }
4372 _ => {
4373 for o in 0..rows {
4374 let (s1, s2) = row_dots(o);
4375 o1[o] = s1;
4376 o2[o] = s2;
4377 }
4378 }
4379 }
4380}
4381
4382#[derive(Clone, Copy)]
4383struct SendMut(*mut f32);
4384unsafe impl Send for SendMut {}
4385unsafe impl Sync for SendMut {}
4386
4387impl SendMut {
4388 #[inline]
4389 fn at(self, i: usize) -> *mut f32 {
4390 unsafe { self.0.add(i) }
4391 }
4392}
4393
4394#[cfg(test)]
4395mod tests {
4396 use super::*;
4397
4398 #[test]
4399 fn f32_matvec_matches_matvec_rows_bitexact() {
4400 let (rows, cols) = (300, 40);
4401 let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
4402 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
4403 let qt = QTensor::from_f32(w.clone(), rows, cols);
4404
4405 let mut a = vec![0.0f32; rows];
4406 matvec_rows(None, &w, &x, &mut a);
4407 let mut b = vec![0.0f32; rows];
4408 qt.matvec(&x, &mut b, None);
4409 assert_eq!(a, b);
4410 }
4411
4412 #[test]
4413 fn sdot_kernel_exact_on_grid() {
4414 eprintln!("sdot_enabled = {}", sdot_enabled());
4419 let (rows, cols) = (9, 80); let w: Vec<u8> = (0..rows * cols)
4421 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
4422 .collect();
4423 let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
4424 let x: Vec<f32> = (0..cols)
4425 .map(|i| match i % 3 {
4426 0 => 1.0,
4427 1 => -1.0,
4428 _ => 0.0,
4429 })
4430 .collect();
4431 let mut a = vec![0.0f32; rows];
4432 qmatvec(&w, &[], &scales, &x, rows, cols, &mut a, None);
4433 for o in 0..rows {
4434 let mut acc = 0.0f32;
4435 for j in 0..cols {
4436 acc += (w[o * cols + j] as i8) as f32 * x[j];
4437 }
4438 let expect = acc * scales[o];
4439 assert!(
4440 (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
4441 "row {o}: {} vs {expect}",
4442 a[o]
4443 );
4444 }
4445 }
4446
4447 #[test]
4448 fn q1_kernels_match_exact_reference() {
4449 let (rows, cols) = (7, 96);
4451 let gpr = cols / GROUP_SIZE;
4452 let mut bytes = Vec::new();
4453 for t in 0..rows * gpr {
4454 let s = 0.01 + (t % 13) as f32 * 0.003;
4455 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
4456 for j in 0..4 {
4457 bytes.push(((t * 31 + j * 97) % 251) as u8);
4458 }
4459 }
4460 let x: Vec<f32> = (0..cols)
4462 .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
4463 .collect();
4464 let mut w = vec![0.0f32; rows * cols];
4466 cortiq_core::quant::dequant_q1(&bytes, &mut w);
4467 let mut expect = vec![0.0f32; rows];
4468 for o in 0..rows {
4469 expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
4470 }
4471 let mut got = vec![0.0f32; rows];
4472 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
4473 for o in 0..rows {
4474 assert!(
4475 (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
4476 "row {o}: {} vs {}",
4477 got[o],
4478 expect[o]
4479 );
4480 }
4481 let x2: Vec<f32> = x.iter().map(|v| -v).collect();
4483 let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
4484 q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
4485 assert_eq!(a1, got);
4486 let mut xs = x.clone();
4487 xs.extend_from_slice(&x2);
4488 let mut mm = vec![0.0f32; 2 * rows];
4489 q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
4490 assert_eq!(&mm[..rows], got.as_slice());
4491 assert_eq!(&mm[rows..], a2.as_slice());
4492 }
4493
4494 #[test]
4495 fn repack_is_bit_identical() {
4496 let (rows, cols) = (267, 96); let w: Vec<u8> = (0..rows * cols)
4502 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
4503 .collect();
4504 let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
4505 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
4506 let rep = q8_repack_layout(&w, rows, cols);
4507 for g in 0..rows / 4 {
4509 for c in 0..cols / 16 {
4510 for lane in 0..4 {
4511 assert_eq!(
4512 &rep[g * 4 * cols + c * 64 + lane * 16..g * 4 * cols + c * 64 + lane * 16 + 16],
4513 &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
4514 );
4515 }
4516 }
4517 }
4518 let mut a = vec![0.0f32; rows];
4519 qmatvec(&w, &[], &scales, &x, rows, cols, &mut a, None);
4520 let mut b = vec![0.0f32; rows];
4521 qmatvec(&w, &rep, &scales, &x, rows, cols, &mut b, None);
4522 assert_eq!(a, b, "full-range repack output diverged");
4523
4524 #[cfg(target_arch = "aarch64")]
4525 if sdot_enabled() {
4526 let act = split_act(&x);
4528 let mut c1 = vec![0.0f32; rows];
4529 let mut c2 = vec![0.0f32; rows];
4530 q8_range_sdot(&w, &[], &scales, &act, cols, SendMut(c1.as_mut_ptr()), 3, rows - 2);
4531 q8_range_sdot(&w, &rep, &scales, &act, cols, SendMut(c2.as_mut_ptr()), 3, rows - 2);
4532 assert_eq!(c1, c2, "unaligned-range repack output diverged");
4533 }
4534 }
4535
4536 #[test]
4537 fn sdot_a8w8_noise_is_bounded() {
4538 let (rows, cols) = (16, 512);
4542 let w: Vec<u8> = (0..rows * cols)
4543 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
4544 .collect();
4545 let scales = vec![0.01f32; rows];
4546 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
4547 let mut a = vec![0.0f32; rows];
4548 qmatvec(&w, &[], &scales, &x, rows, cols, &mut a, None);
4549 let (mut num, mut den) = (0f64, 0f64);
4550 for o in 0..rows {
4551 let mut acc = 0.0f32;
4552 for j in 0..cols {
4553 acc += (w[o * cols + j] as i8) as f32 * x[j];
4554 }
4555 let expect = acc * scales[o];
4556 num += ((a[o] - expect) as f64).powi(2);
4557 den += (expect as f64).powi(2);
4558 }
4559 let rel = (num / den.max(1e-12)).sqrt();
4560 assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
4561 }
4562
4563 #[test]
4564 fn i8_dot_neon_matches_scalar() {
4565 let n = 100;
4566 let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
4567 let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
4568 let mut scalar = 0.0f32;
4569 for j in 0..n {
4570 scalar += (w[j] as i8) as f32 * x[j];
4571 }
4572 let fast = dot_i8_f32(&w, &x);
4573 assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
4574 }
4575
4576 #[test]
4578 fn vbitmatvec_matches_full_dequant() {
4579 let (rows, cols) = (6, 64);
4580 let ng = cols / GROUP_SIZE;
4581 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
4583 let mut bytes = bits.clone();
4584 for g in 0..rows * ng {
4585 let s = 0.02 + 0.001 * g as f32;
4586 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
4587 }
4588 for r in 0..rows {
4589 let b = bits[r] as usize;
4590 let (mut acc, mut nb) = (0u64, 0usize);
4591 let mut rowbytes = Vec::new();
4592 for i in 0..cols {
4593 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
4594 acc = (acc << b) | v;
4595 nb += b;
4596 while nb >= 8 {
4597 nb -= 8;
4598 rowbytes.push(((acc >> nb) & 0xFF) as u8);
4599 }
4600 }
4601 if nb > 0 {
4602 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
4603 }
4604 bytes.extend_from_slice(&rowbytes);
4605 }
4606 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
4607
4608 let mut reference = vec![0f32; rows * cols];
4609 cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
4610 let mut expect = vec![0f32; rows];
4611 for r in 0..rows {
4612 expect[r] = reference[r * cols..(r + 1) * cols]
4613 .iter()
4614 .zip(&x)
4615 .map(|(w, xv)| w * xv)
4616 .sum();
4617 }
4618 let mut got = vec![0f32; rows];
4619 let offsets = vbit_row_offsets(&bytes, rows, cols);
4620 vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
4621 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
4625 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
4626 for r in 0..rows {
4627 assert!(
4628 (got[r] - expect[r]).abs() < tol * scale,
4629 "row {r}: {} vs {}",
4630 got[r],
4631 expect[r]
4632 );
4633 }
4634 }
4635
4636 #[test]
4639 fn q4matvec_matches_full_dequant() {
4640 let (rows, cols) = (8, 64);
4641 let groups = rows * cols / GROUP_SIZE;
4642 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
4644 for i in 0..groups * 16 {
4645 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
4646 }
4647 for g in 0..groups {
4648 let s = 0.01 + 0.003 * g as f32;
4649 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
4650 }
4651 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
4652
4653 let mut reference = vec![0.0f32; rows * cols];
4654 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
4655 let mut expect = vec![0.0f32; rows];
4656 for r in 0..rows {
4657 expect[r] = reference[r * cols..(r + 1) * cols]
4658 .iter()
4659 .zip(&x)
4660 .map(|(w, xv)| w * xv)
4661 .sum();
4662 }
4663
4664 let mut got = vec![0.0f32; rows];
4665 q4matvec(&bytes, &x, rows, cols, &mut got, None);
4666 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
4670 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
4671 for r in 0..rows {
4672 assert!(
4673 (got[r] - expect[r]).abs() < tol * scale,
4674 "row {r}: {} vs {}",
4675 got[r],
4676 expect[r]
4677 );
4678 }
4679 }
4680
4681 #[test]
4684 fn vbitmatvec2_equals_two_singles() {
4685 let (rows, cols) = (6, 64);
4686 let ng = cols / GROUP_SIZE;
4687 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
4688 let mut bytes = bits.clone();
4689 for g in 0..rows * ng {
4690 let s = 0.02 + 0.001 * g as f32;
4691 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
4692 }
4693 for r in 0..rows {
4694 let b = bits[r] as usize;
4695 let (mut acc, mut nb) = (0u64, 0usize);
4696 let mut rowbytes = Vec::new();
4697 for i in 0..cols {
4698 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
4699 acc = (acc << b) | v;
4700 nb += b;
4701 while nb >= 8 {
4702 nb -= 8;
4703 rowbytes.push(((acc >> nb) & 0xFF) as u8);
4704 }
4705 }
4706 if nb > 0 {
4707 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
4708 }
4709 bytes.extend_from_slice(&rowbytes);
4710 }
4711 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
4712 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
4713 let offsets = vbit_row_offsets(&bytes, rows, cols);
4714
4715 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
4716 vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
4717 vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
4718 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
4719 vbitmatvec2(&bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
4720 assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
4721 assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
4722 }
4723
4724 #[test]
4726 fn q4matvec2_equals_two_singles() {
4727 let (rows, cols) = (8, 128);
4728 let groups = rows * cols / GROUP_SIZE;
4729 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
4730 for i in 0..groups * 16 {
4731 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
4732 }
4733 for g in 0..groups {
4734 let s = 0.01 + 0.003 * g as f32;
4735 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
4736 }
4737 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
4740 x1[9] = 250.0;
4741 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
4742
4743 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
4744 q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
4745 q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
4746 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
4747 q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
4748 assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
4749 assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
4750 }
4751
4752 #[test]
4755 fn matvec_many_equals_separate_matvecs() {
4756 use crate::pool::Pool;
4757 let (r1, r2, cols) = (300, 200, 64);
4758 let mk = |salt: usize, rows: usize| {
4759 QTensor::from_f32(
4760 (0..rows * cols).map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5).collect(),
4761 rows,
4762 cols,
4763 )
4764 };
4765 let (a, b) = (mk(1, r1), mk(5, r2));
4766 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
4767 let pool = Pool::new(3);
4768
4769 let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
4770 a.matvec(&x, &mut ea, Some(&pool));
4771 b.matvec(&x, &mut eb, Some(&pool));
4772 let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
4773 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
4774 assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
4775 assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
4776 }
4777
4778 #[test]
4781 fn batched_matmat_equals_per_position_matvec() {
4782 let (rows, cols, b) = (8, 64, 5);
4783 let groups = rows * cols / GROUP_SIZE;
4785 let mut q4 = Vec::new();
4786 for i in 0..groups * 16 {
4787 q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
4788 }
4789 for g in 0..groups {
4790 q4.extend_from_slice(
4791 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
4792 );
4793 }
4794 let ng = cols / GROUP_SIZE;
4796 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
4797 let mut vb = bits.clone();
4798 for g in 0..rows * ng {
4799 vb.extend_from_slice(
4800 &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
4801 );
4802 }
4803 for r in 0..rows {
4804 let bw = bits[r] as usize;
4805 let (mut acc, mut nb) = (0u64, 0usize);
4806 let mut rowbytes = Vec::new();
4807 for i in 0..cols {
4808 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
4809 acc = (acc << bw) | v;
4810 nb += bw;
4811 while nb >= 8 {
4812 nb -= 8;
4813 rowbytes.push(((acc >> nb) & 0xFF) as u8);
4814 }
4815 }
4816 if nb > 0 {
4817 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
4818 }
4819 vb.extend_from_slice(&rowbytes);
4820 }
4821 let offsets = vbit_row_offsets(&vb, rows, cols);
4822
4823 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
4824
4825 let mut got = vec![0f32; b * rows];
4827 q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
4828 for bi in 0..b {
4829 let mut expect = vec![0f32; rows];
4830 q4matvec(&q4, &xs[bi * cols..(bi + 1) * cols], rows, cols, &mut expect, None);
4831 assert_eq!(&got[bi * rows..(bi + 1) * rows], &expect[..], "q4 batch pos {bi}");
4832 }
4833
4834 let mut got = vec![0f32; b * rows];
4836 vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
4837 for bi in 0..b {
4838 let mut expect = vec![0f32; rows];
4839 vbitmatvec(
4840 &vb, &offsets, &xs[bi * cols..(bi + 1) * cols], rows, cols, &mut expect, None,
4841 );
4842 assert_eq!(&got[bi * rows..(bi + 1) * rows], &expect[..], "vbit batch pos {bi}");
4843 }
4844 }
4845
4846 #[test]
4850 fn q4_tiled_matches_q4_block_bitexact() {
4851 let (rows, cols, b) = (8usize, 128usize, 3usize);
4852 let groups = rows * cols / GROUP_SIZE;
4853 let mut split = Vec::with_capacity(groups * 18);
4854 for i in 0..groups * 16 {
4855 split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
4856 }
4857 for g in 0..groups {
4858 split.extend_from_slice(
4859 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
4860 );
4861 }
4862 let (packed, scales) = split.split_at(groups * 16);
4864 let mut tiled = Vec::with_capacity(groups * Q4_TILE);
4865 for g in 0..groups {
4866 tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
4867 tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
4868 }
4869
4870 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
4871 x1[9] = 250.0; let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
4873
4874 let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
4875 q4matvec(&split, &x1, rows, cols, &mut a, None);
4876 q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
4877 assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
4878
4879 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
4880 let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
4881 q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
4882 q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
4883 assert_eq!(a1, t1);
4884 assert_eq!(a2, t2);
4885
4886 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
4887 let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
4888 q4matmat(&split, &xs, b, rows, cols, &mut am, None);
4889 q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
4890 assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
4891 }
4892
4893 #[test]
4900 fn q4matvec_sdot_outlier_exact() {
4901 let (rows, cols) = (4, 128);
4902 let groups = rows * cols / GROUP_SIZE;
4903 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
4904 for i in 0..groups * 16 {
4905 bytes.push(((i * 11 + 5) % 256) as u8);
4906 }
4907 for g in 0..groups {
4908 let s = 0.02 + 0.002 * g as f32;
4909 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
4910 }
4911 let mut x: Vec<f32> = (0..cols)
4912 .map(|i| match i % 3 {
4913 0 => 1.0,
4914 1 => -1.0,
4915 _ => 0.0,
4916 })
4917 .collect();
4918 x[17] = 300.0; let mut reference = vec![0.0f32; rows * cols];
4921 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
4922 let mut expect = vec![0.0f32; rows];
4923 for r in 0..rows {
4924 expect[r] = reference[r * cols..(r + 1) * cols]
4925 .iter()
4926 .zip(&x)
4927 .map(|(w, xv)| w * xv)
4928 .sum();
4929 }
4930 let mut got = vec![0.0f32; rows];
4931 q4matvec(&bytes, &x, rows, cols, &mut got, None);
4932 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
4933 for r in 0..rows {
4934 assert!(
4935 (got[r] - expect[r]).abs() < 2e-3 * scale,
4936 "row {r}: {} vs {} (outlier term must be exact)",
4937 got[r],
4938 expect[r]
4939 );
4940 }
4941 }
4942}