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")]
1426 if avx2_a8w8_enabled() {
1427 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
1428 let out_addr = SendMut(out.as_mut_ptr());
1429 let blocked_ok = std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true);
1432 if !avx512vnni_enabled() && blocked_ok {
1433 let run = |start: usize, end: usize| {
1434 let mut o = start;
1435 while o < end {
1436 if o + 2 <= end {
1437 let r0 = &q[o * cols..(o + 1) * cols];
1438 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
1439 let mut bi = 0usize;
1440 while bi + 4 <= acts.len() {
1441 let xs = [
1442 acts[bi].xq.as_slice(),
1443 acts[bi + 1].xq.as_slice(),
1444 acts[bi + 2].xq.as_slice(),
1445 acts[bi + 3].xq.as_slice(),
1446 ];
1447 let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
1448 for (r, row) in [r0, r1].into_iter().enumerate() {
1449 for k in 0..4 {
1450 let act = &acts[bi + k];
1451 let mut v = d[r][k] as f32 * act.sx;
1452 for &(j, xv) in &act.outliers {
1453 v += (row[j] as i8) as f32 * xv;
1454 }
1455 unsafe {
1456 *out_addr.at((bi + k) * rows + o + r) =
1457 v * row_scale[o + r]
1458 };
1459 }
1460 }
1461 bi += 4;
1462 }
1463 while bi < acts.len() {
1464 for (r, row) in [r0, r1].into_iter().enumerate() {
1465 let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
1466 unsafe { *out_addr.at(bi * rows + o + r) = v };
1467 }
1468 bi += 1;
1469 }
1470 o += 2;
1471 } else {
1472 let row = &q[o * cols..(o + 1) * cols];
1473 for (bi, act) in acts.iter().enumerate() {
1474 let v = row_dot_avx2(row, act) * row_scale[o];
1475 unsafe { *out_addr.at(bi * rows + o) = v };
1476 }
1477 o += 1;
1478 }
1479 }
1480 };
1481 dispatch_rows(pool, rows, &run);
1482 return;
1483 }
1484 let run = |start: usize, end: usize| {
1485 for o in start..end {
1486 let row = &q[o * cols..(o + 1) * cols];
1487 for (bi, act) in acts.iter().enumerate() {
1488 let v = row_dot_avx2(row, act) * row_scale[o];
1489 unsafe { *out_addr.at(bi * rows + o) = v };
1490 }
1491 }
1492 };
1493 dispatch_rows(pool, rows, &run);
1494 return;
1495 }
1496 let out_addr = SendMut(out.as_mut_ptr());
1497 let run = |start: usize, end: usize| {
1498 for o in start..end {
1499 let row = &q[o * cols..(o + 1) * cols];
1500 for (bi, x) in pre.iter().enumerate() {
1501 let mut acc = 0f32;
1502 for j in 0..cols {
1503 acc += (row[j] as i8) as f32 * x[j];
1504 }
1505 unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
1506 }
1507 }
1508 };
1509 dispatch_rows(pool, rows, &run);
1510}
1511
1512fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
1515 match pool {
1516 Some(pool) if rows >= 256 => pool.run_rows(rows, run),
1517 _ => run(0, rows),
1518 }
1519}
1520
1521fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
1523 let groups = rows * cols / GROUP_SIZE;
1524 bytes.split_at(groups * 16)
1525}
1526
1527#[inline]
1532fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
1533 #[cfg(target_arch = "aarch64")]
1534 unsafe {
1535 return vbit_fill4_neon(data, buf);
1536 }
1537 #[cfg(target_arch = "x86_64")]
1538 if avx2_enabled() {
1539 return unsafe { vbit_fill4_avx2(data, buf) };
1540 }
1541 #[allow(unreachable_code)]
1542 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
1543 let u = unpack8::<4>(&data[blk * 4..]);
1544 for k in 0..8 {
1545 chunk[k] = (u[k] - 7) as i8 as u8;
1546 }
1547 }
1548}
1549
1550#[cfg(target_arch = "aarch64")]
1551#[target_feature(enable = "neon")]
1552unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
1553 unsafe {
1556 use core::arch::aarch64::*;
1557 let n = buf.len();
1558 let mask = vdupq_n_u8(0x0F);
1559 let seven = vdupq_n_s8(7);
1560 let mut g = 0usize;
1561 while g * 32 + 32 <= n {
1562 let b = vld1q_u8(data.as_ptr().add(g * 16));
1563 let hi = vshrq_n_u8::<4>(b);
1564 let lo = vandq_u8(b, mask);
1565 let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
1566 let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
1567 vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
1568 vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
1569 g += 1;
1570 }
1571 }
1572}
1573
1574#[cfg(target_arch = "x86_64")]
1575#[target_feature(enable = "avx2")]
1576unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
1577 unsafe {
1579 use core::arch::x86_64::*;
1580 let n = buf.len();
1581 let mask = _mm_set1_epi8(0x0F);
1582 let seven = _mm256_set1_epi8(7);
1583 let mut g = 0usize;
1584 while g * 32 + 32 <= n {
1585 let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
1586 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
1587 let lo = _mm_and_si128(b, mask);
1588 let z = _mm256_sub_epi8(
1589 _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
1590 seven,
1591 );
1592 _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
1593 g += 1;
1594 }
1595 }
1596}
1597
1598#[inline(always)]
1603fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
1604 let mut acc = 0u64;
1605 for i in 0..B {
1606 acc = (acc << 8) | data[i] as u64;
1607 }
1608 let mask = (1u64 << B) - 1;
1609 let mut out = [0i32; 8];
1610 for (k, o) in out.iter_mut().enumerate() {
1611 *o = ((acc >> ((7 - k) * B)) & mask) as i32;
1612 }
1613 out
1614}
1615
1616#[allow(clippy::too_many_arguments)]
1622fn vbitmatvec(
1623 bytes: &[u8],
1624 offsets: &[usize],
1625 x: &[f32],
1626 rows: usize,
1627 cols: usize,
1628 out: &mut [f32],
1629 pool: Option<&Pool>,
1630) {
1631 debug_assert_eq!(out.len(), rows);
1632 debug_assert_eq!(offsets.len(), rows + 1);
1633
1634 if a8w8_enabled() {
1638 let act = split_act(x);
1639 let out_addr = SendMut(out.as_mut_ptr());
1640 let run = move |start: usize, end: usize| {
1641 vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
1642 };
1643 dispatch_rows(pool, rows, &run);
1644 return;
1645 }
1646
1647 let out_addr = SendMut(out.as_mut_ptr());
1648 let run = move |start: usize, end: usize| {
1649 vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
1650 };
1651 dispatch_rows(pool, rows, &run);
1652}
1653
1654#[allow(clippy::too_many_arguments)]
1658fn vbit_range_a8w8(
1659 bytes: &[u8],
1660 offsets: &[usize],
1661 x: &[f32],
1662 act: &SplitAct,
1663 rows: usize,
1664 cols: usize,
1665 out: SendMut,
1666 start: usize,
1667 end: usize,
1668) {
1669 let ng = cols / GROUP_SIZE;
1670 let bits = &bytes[..rows];
1671 let sc_off = rows;
1672 let row_dot = |r: usize| -> f32 {
1673 let b = bits[r] as usize;
1674 let l = ((1i32 << (b - 1)) - 1) as i32;
1675 let mask = (1u64 << b) - 1;
1676 let data = &bytes[offsets[r]..offsets[r + 1]];
1677 if b == 8 {
1678 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
1680 let mut dot = 0f32;
1681 for g in 0..ng {
1682 let so = (r * ng + g) * 2;
1683 let sgf = f16_to_f32(u16::from_le_bytes([
1684 bytes[sc_off + so],
1685 bytes[sc_off + so + 1],
1686 ]));
1687 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
1688 let mut gd = 0f32;
1689 for &xv in xg.iter() {
1690 if nbits < 8 {
1691 acc = (acc << 8) | data[idx] as u64;
1692 idx += 1;
1693 nbits += 8;
1694 }
1695 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
1696 nbits -= 8;
1697 gd += (u - l) as f32 * xv;
1698 }
1699 dot += gd * sgf;
1700 }
1701 return dot;
1702 }
1703 thread_local! {
1707 static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
1708 const { std::cell::RefCell::new(Vec::new()) };
1709 }
1710 #[inline(always)]
1711 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
1712 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
1713 let u = unpack8::<B>(&data[blk * B..]);
1714 for k in 0..8 {
1715 chunk[k] = (u[k] - l) as i8 as u8;
1716 }
1717 }
1718 }
1719 let _ = mask;
1720 VBIT_SCRATCH.with(|scratch| {
1721 let mut buf = scratch.borrow_mut();
1722 buf.resize(cols, 0);
1723 match b {
1724 3 => fill::<3>(data, l, &mut buf),
1725 4 => vbit_fill4(data, &mut buf),
1726 5 => fill::<5>(data, l, &mut buf),
1727 6 => fill::<6>(data, l, &mut buf),
1728 _ => unreachable!(),
1729 }
1730 let mut dot = 0f32;
1731 for g in 0..ng {
1732 let so = (r * ng + g) * 2;
1733 let s = f16_to_f32(u16::from_le_bytes([
1734 bytes[sc_off + so],
1735 bytes[sc_off + so + 1],
1736 ]));
1737 let d = dot_i8_i8(
1738 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
1739 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
1740 ) as f32
1741 * act.sx;
1742 dot += d * s;
1743 }
1744 for &(j, xv) in &act.outliers {
1745 let so = (r * ng + j / GROUP_SIZE) * 2;
1746 let s = f16_to_f32(u16::from_le_bytes([
1747 bytes[sc_off + so],
1748 bytes[sc_off + so + 1],
1749 ]));
1750 dot += (buf[j] as i8) as f32 * s * xv;
1752 }
1753 dot
1754 })
1755 };
1756 for r in start..end {
1757 unsafe { *out.at(r) = row_dot(r) };
1759 }
1760}
1761
1762#[allow(clippy::too_many_arguments)]
1764fn vbit_range_f32(
1765 bytes: &[u8],
1766 offsets: &[usize],
1767 x: &[f32],
1768 rows: usize,
1769 cols: usize,
1770 out: SendMut,
1771 start: usize,
1772 end: usize,
1773) {
1774 let ng = cols / GROUP_SIZE;
1775 let bits = &bytes[..rows];
1776 let sc_off = rows;
1777 #[inline(always)]
1781 fn dot_row<const B: usize>(
1782 data: &[u8],
1783 bytes: &[u8],
1784 sc_off: usize,
1785 r: usize,
1786 ng: usize,
1787 x: &[f32],
1788 ) -> f32 {
1789 let l = ((1i32 << (B - 1)) - 1) as f32;
1790 let gbytes = GROUP_SIZE * B / 8;
1791 let mut dot = 0f32;
1792 for g in 0..ng {
1793 let so = (r * ng + g) * 2;
1794 let s = f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]));
1795 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
1796 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
1797 let mut gd = 0f32;
1798 for blk in 0..GROUP_SIZE / 8 {
1799 let u = unpack8::<B>(&gd0[blk * B..]);
1800 let xb = &xg[blk * 8..blk * 8 + 8];
1801 for k in 0..8 {
1802 gd += (u[k] as f32 - l) * xb[k];
1803 }
1804 }
1805 dot += gd * s;
1806 }
1807 dot
1808 }
1809 for r in start..end {
1810 let data = &bytes[offsets[r]..offsets[r + 1]];
1811 let v = match bits[r] {
1812 3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
1813 4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
1814 5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
1815 6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
1816 8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
1817 b => unreachable!("vbit bit-width {b} (validated at load)"),
1818 };
1819 unsafe { *out.at(r) = v };
1821 }
1822}
1823
1824#[allow(clippy::too_many_arguments)]
1829fn vbitmatvec2(
1830 bytes: &[u8],
1831 offsets: &[usize],
1832 x1: &[f32],
1833 x2: &[f32],
1834 rows: usize,
1835 cols: usize,
1836 o1: &mut [f32],
1837 o2: &mut [f32],
1838 pool: Option<&Pool>,
1839) {
1840 debug_assert_eq!(o1.len(), rows);
1841 debug_assert_eq!(o2.len(), rows);
1842
1843 if a8w8_enabled() {
1844 let a1 = split_act(x1);
1845 let a2 = split_act(x2);
1846 let p1 = SendMut(o1.as_mut_ptr());
1847 let p2 = SendMut(o2.as_mut_ptr());
1848 let run = move |start: usize, end: usize| {
1849 vbit_range2_a8w8(bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end)
1850 };
1851 dispatch_rows(pool, rows, &run);
1852 return;
1853 }
1854
1855 let p1 = SendMut(o1.as_mut_ptr());
1856 let p2 = SendMut(o2.as_mut_ptr());
1857 let run = move |start: usize, end: usize| {
1858 vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
1859 };
1860 dispatch_rows(pool, rows, &run);
1861}
1862
1863#[allow(clippy::too_many_arguments)]
1867fn vbit_range2_a8w8(
1868 bytes: &[u8],
1869 offsets: &[usize],
1870 x1: &[f32],
1871 x2: &[f32],
1872 a1: &SplitAct,
1873 a2: &SplitAct,
1874 rows: usize,
1875 cols: usize,
1876 p1: SendMut,
1877 p2: SendMut,
1878 start: usize,
1879 end: usize,
1880) {
1881 let ng = cols / GROUP_SIZE;
1882 let bits = &bytes[..rows];
1883 let sc_off = rows;
1884 let row_dots = |r: usize| -> (f32, f32) {
1885 let b = bits[r] as usize;
1886 let l = (1i32 << (b - 1)) - 1;
1887 let data = &bytes[offsets[r]..offsets[r + 1]];
1888 if b == 8 {
1889 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
1892 let (mut d1, mut d2) = (0f32, 0f32);
1893 for g in 0..ng {
1894 let so = (r * ng + g) * 2;
1895 let sgf = f16_to_f32(u16::from_le_bytes([
1896 bytes[sc_off + so],
1897 bytes[sc_off + so + 1],
1898 ]));
1899 let (mut g1, mut g2) = (0f32, 0f32);
1900 for k in 0..GROUP_SIZE {
1901 if nbits < 8 {
1902 acc = (acc << 8) | data[idx] as u64;
1903 idx += 1;
1904 nbits += 8;
1905 }
1906 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
1907 nbits -= 8;
1908 let w = (u - l) as f32;
1909 g1 += w * x1[g * GROUP_SIZE + k];
1910 g2 += w * x2[g * GROUP_SIZE + k];
1911 }
1912 d1 += g1 * sgf;
1913 d2 += g2 * sgf;
1914 }
1915 return (d1, d2);
1916 }
1917 thread_local! {
1918 static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
1919 const { std::cell::RefCell::new(Vec::new()) };
1920 }
1921 #[inline(always)]
1922 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
1923 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
1924 let u = unpack8::<B>(&data[blk * B..]);
1925 for k in 0..8 {
1926 chunk[k] = (u[k] - l) as i8 as u8;
1927 }
1928 }
1929 }
1930 VBIT_SCRATCH2.with(|scratch| {
1931 let mut buf = scratch.borrow_mut();
1932 buf.resize(cols, 0);
1933 match b {
1934 3 => fill::<3>(data, l, &mut buf),
1935 4 => vbit_fill4(data, &mut buf),
1936 5 => fill::<5>(data, l, &mut buf),
1937 6 => fill::<6>(data, l, &mut buf),
1938 _ => unreachable!(),
1939 }
1940 let (mut d1, mut d2) = (0f32, 0f32);
1941 for g in 0..ng {
1942 let so = (r * ng + g) * 2;
1943 let s = f16_to_f32(u16::from_le_bytes([
1944 bytes[sc_off + so],
1945 bytes[sc_off + so + 1],
1946 ]));
1947 let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
1948 let v1 =
1949 dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
1950 let v2 =
1951 dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
1952 d1 += v1 * s;
1953 d2 += v2 * s;
1954 }
1955 for &(j, xv) in &a1.outliers {
1956 let so = (r * ng + j / GROUP_SIZE) * 2;
1957 let s = f16_to_f32(u16::from_le_bytes([
1958 bytes[sc_off + so],
1959 bytes[sc_off + so + 1],
1960 ]));
1961 d1 += (buf[j] as i8) as f32 * s * xv;
1962 }
1963 for &(j, xv) in &a2.outliers {
1964 let so = (r * ng + j / GROUP_SIZE) * 2;
1965 let s = f16_to_f32(u16::from_le_bytes([
1966 bytes[sc_off + so],
1967 bytes[sc_off + so + 1],
1968 ]));
1969 d2 += (buf[j] as i8) as f32 * s * xv;
1970 }
1971 (d1, d2)
1972 })
1973 };
1974 for r in start..end {
1975 let (v1, v2) = row_dots(r);
1976 unsafe {
1978 *p1.at(r) = v1;
1979 *p2.at(r) = v2;
1980 }
1981 }
1982}
1983
1984#[allow(clippy::too_many_arguments)]
1988fn vbit_range2_f32(
1989 bytes: &[u8],
1990 offsets: &[usize],
1991 x1: &[f32],
1992 x2: &[f32],
1993 rows: usize,
1994 cols: usize,
1995 p1: SendMut,
1996 p2: SendMut,
1997 start: usize,
1998 end: usize,
1999) {
2000 let ng = cols / GROUP_SIZE;
2001 let bits = &bytes[..rows];
2002 let sc_off = rows;
2003 #[inline(always)]
2004 #[allow(clippy::too_many_arguments)]
2005 fn dot_row2<const B: usize>(
2006 data: &[u8],
2007 bytes: &[u8],
2008 sc_off: usize,
2009 r: usize,
2010 ng: usize,
2011 x1: &[f32],
2012 x2: &[f32],
2013 ) -> (f32, f32) {
2014 let l = ((1i32 << (B - 1)) - 1) as f32;
2015 let gbytes = GROUP_SIZE * B / 8;
2016 let (mut d1, mut d2) = (0f32, 0f32);
2017 for g in 0..ng {
2018 let so = (r * ng + g) * 2;
2019 let s = f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]));
2020 let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2021 let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2022 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
2023 let (mut g1, mut g2) = (0f32, 0f32);
2024 for blk in 0..GROUP_SIZE / 8 {
2025 let u = unpack8::<B>(&gd0[blk * B..]);
2026 for k in 0..8 {
2027 let w = u[k] as f32 - l;
2028 g1 += w * x1g[blk * 8 + k];
2029 g2 += w * x2g[blk * 8 + k];
2030 }
2031 }
2032 d1 += g1 * s;
2033 d2 += g2 * s;
2034 }
2035 (d1, d2)
2036 }
2037 for r in start..end {
2038 let data = &bytes[offsets[r]..offsets[r + 1]];
2039 let (v1, v2) = match bits[r] {
2040 3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
2041 4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
2042 5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
2043 6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
2044 8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
2045 b => unreachable!("vbit bit-width {b} (validated at load)"),
2046 };
2047 unsafe {
2049 *p1.at(r) = v1;
2050 *p2.at(r) = v2;
2051 }
2052 }
2053}
2054
2055#[inline]
2062#[allow(unreachable_code)]
2063fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
2064 #[cfg(target_arch = "aarch64")]
2065 unsafe {
2066 return dot_q4t_row_sdot(bytes, r, gpr, xq);
2067 }
2068 #[cfg(target_arch = "x86_64")]
2069 unsafe {
2070 return dot_q4t_row_avx2(bytes, r, gpr, xq);
2071 }
2072 let mut acc = 0f32;
2073 for gi in 0..gpr {
2074 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
2075 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2076 let mut d = 0i32;
2077 for (k, &b) in tile[2..].iter().enumerate() {
2078 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
2079 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
2080 }
2081 acc += d as f32 * s;
2082 }
2083 acc
2084}
2085
2086#[cfg(target_arch = "aarch64")]
2087#[target_feature(enable = "neon,dotprod")]
2088unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
2089 unsafe {
2092 use core::arch::aarch64::*;
2093 use core::arch::asm;
2094 let lomask = vdupq_n_u8(0x0F);
2095 let eight = vdupq_n_s8(8);
2096 let mut acc = 0f32;
2097 for gi in 0..gpr {
2098 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
2099 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2100 let b = vld1q_u8(t.add(2));
2101 let lo = vandq_u8(b, lomask);
2102 let hi = vshrq_n_u8::<4>(b);
2103 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
2104 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
2105 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
2106 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
2107 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
2108 asm!(
2109 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
2110 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
2111 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
2112 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
2113 options(pure, nomem, nostack),
2114 );
2115 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
2116 }
2117 acc
2118 }
2119}
2120
2121#[cfg(target_arch = "x86_64")]
2122#[target_feature(enable = "avx2")]
2123unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
2124 unsafe {
2126 use core::arch::x86_64::*;
2127 let lomask = _mm_set1_epi8(0x0F);
2128 let eight = _mm256_set1_epi8(8);
2129 let ones = _mm256_set1_epi16(1);
2130 let mut acc = 0f32;
2131 for gi in 0..gpr {
2132 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
2133 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2134 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
2135 let lo = _mm_and_si128(b, lomask);
2136 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
2137 let w = _mm256_sub_epi8(
2138 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
2139 eight,
2140 );
2141 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2142 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
2143 let d = _mm256_madd_epi16(p16, ones);
2144 let hi128 = _mm256_extracti128_si256::<1>(d);
2145 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
2146 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2147 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2148 acc += _mm_cvtsi128_si32(s32) as f32 * s;
2149 }
2150 acc
2151 }
2152}
2153
2154#[cfg(target_arch = "x86_64")]
2159#[target_feature(enable = "avx2")]
2160unsafe fn dot_q4t_row_1x4_avx2(
2161 bytes: &[u8],
2162 r: usize,
2163 gpr: usize,
2164 xs: [&[i8]; 4],
2165) -> [f32; 4] {
2166 unsafe {
2168 use core::arch::x86_64::*;
2169 let lomask = _mm_set1_epi8(0x0F);
2170 let eight = _mm256_set1_epi8(8);
2171 let ones = _mm256_set1_epi16(1);
2172 let mut acc = [0f32; 4];
2173 for gi in 0..gpr {
2174 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
2175 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2176 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
2177 let lo = _mm_and_si128(bb, lomask);
2178 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
2179 let w = _mm256_sub_epi8(
2180 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
2181 eight,
2182 );
2183 let aw = _mm256_abs_epi8(w);
2184 for (k, xq) in xs.iter().enumerate() {
2185 let x =
2186 _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2187 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
2188 let d = _mm256_madd_epi16(p16, ones);
2189 let hi128 = _mm256_extracti128_si256::<1>(d);
2190 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
2191 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2192 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2193 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
2194 }
2195 }
2196 acc
2197 }
2198}
2199
2200#[inline]
2202fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
2203 let gi = j / GROUP_SIZE;
2204 let k = j % GROUP_SIZE;
2205 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
2206 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2207 let byte = tile[2 + k / 2];
2208 let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
2209 ((nib as i32 - 8) as f32, s)
2210}
2211
2212#[inline]
2215fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
2216 let mut acc = 0f32;
2217 for gi in 0..gpr {
2218 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
2219 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2220 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
2221 let mut ga = 0f32;
2222 for (k, &b) in tile[2..].iter().enumerate() {
2223 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
2224 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
2225 }
2226 acc += ga * s;
2227 }
2228 acc
2229}
2230
2231fn q4t_matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
2233 debug_assert_eq!(out.len(), rows);
2234 let gpr = cols / GROUP_SIZE;
2235 let out_addr = SendMut(out.as_mut_ptr());
2236 if a8w8_enabled() {
2237 let act = split_act(x);
2238 let run = move |start: usize, end: usize| {
2239 for r in start..end {
2240 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
2241 for &(j, xv) in &act.outliers {
2242 let (w, s) = q4t_outlier(bytes, r, gpr, j);
2243 acc += w * s * xv;
2244 }
2245 unsafe { *out_addr.at(r) = acc };
2247 }
2248 };
2249 dispatch_rows(pool, rows, &run);
2250 return;
2251 }
2252 let run = move |start: usize, end: usize| {
2253 for r in start..end {
2254 unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
2256 }
2257 };
2258 dispatch_rows(pool, rows, &run);
2259}
2260
2261#[allow(clippy::too_many_arguments)]
2263fn q4t_matvec2(
2264 bytes: &[u8],
2265 x1: &[f32],
2266 x2: &[f32],
2267 rows: usize,
2268 cols: usize,
2269 o1: &mut [f32],
2270 o2: &mut [f32],
2271 pool: Option<&Pool>,
2272) {
2273 let gpr = cols / GROUP_SIZE;
2274 let p1 = SendMut(o1.as_mut_ptr());
2275 let p2 = SendMut(o2.as_mut_ptr());
2276 if a8w8_enabled() {
2277 let a1 = split_act(x1);
2278 let a2 = split_act(x2);
2279 let run = move |start: usize, end: usize| {
2280 for r in start..end {
2281 let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
2282 let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
2283 for &(j, xv) in &a1.outliers {
2284 let (w, s) = q4t_outlier(bytes, r, gpr, j);
2285 v1 += w * s * xv;
2286 }
2287 for &(j, xv) in &a2.outliers {
2288 let (w, s) = q4t_outlier(bytes, r, gpr, j);
2289 v2 += w * s * xv;
2290 }
2291 unsafe {
2293 *p1.at(r) = v1;
2294 *p2.at(r) = v2;
2295 }
2296 }
2297 };
2298 dispatch_rows(pool, rows, &run);
2299 return;
2300 }
2301 let run = move |start: usize, end: usize| {
2302 for r in start..end {
2303 unsafe {
2305 *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
2306 *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
2307 }
2308 }
2309 };
2310 dispatch_rows(pool, rows, &run);
2311}
2312
2313#[allow(clippy::too_many_arguments)]
2315fn q4t_matmat(
2316 bytes: &[u8],
2317 xs_all: &[f32],
2318 b: usize,
2319 rows: usize,
2320 cols: usize,
2321 out: &mut [f32],
2322 pool: Option<&Pool>,
2323) {
2324 debug_assert_eq!(out.len(), b * rows);
2325 let gpr = cols / GROUP_SIZE;
2326 let out_addr = SendMut(out.as_mut_ptr());
2327 if a8w8_enabled() {
2328 let acts: Vec<SplitAct> =
2329 (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
2330 let acts = &acts;
2331 #[cfg(target_arch = "x86_64")]
2332 let blocked_ok = avx2_enabled()
2333 && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true);
2334 #[cfg(not(target_arch = "x86_64"))]
2335 let blocked_ok = false;
2336 let run = move |start: usize, end: usize| {
2337 for r in start..end {
2338 let mut bi = 0usize;
2339 #[cfg(target_arch = "x86_64")]
2340 if blocked_ok {
2341 while bi + 4 <= acts.len() {
2342 let xs = [
2343 acts[bi].xq.as_slice(),
2344 acts[bi + 1].xq.as_slice(),
2345 acts[bi + 2].xq.as_slice(),
2346 acts[bi + 3].xq.as_slice(),
2347 ];
2348 let d = unsafe { dot_q4t_row_1x4_avx2(bytes, r, gpr, xs) };
2349 for k in 0..4 {
2350 let act = &acts[bi + k];
2351 let mut acc = d[k] * act.sx;
2352 for &(j, xv) in &act.outliers {
2353 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
2354 acc += w * sc * xv;
2355 }
2356 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
2358 }
2359 bi += 4;
2360 }
2361 }
2362 let _ = blocked_ok;
2363 while bi < acts.len() {
2364 let act = &acts[bi];
2365 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
2366 for &(j, xv) in &act.outliers {
2367 let (w, s) = q4t_outlier(bytes, r, gpr, j);
2368 acc += w * s * xv;
2369 }
2370 unsafe { *out_addr.at(bi * rows + r) = acc };
2372 bi += 1;
2373 }
2374 }
2375 };
2376 dispatch_rows(pool, rows, &run);
2377 return;
2378 }
2379 let run = move |start: usize, end: usize| {
2380 for r in start..end {
2381 for bi in 0..b {
2382 let x = &xs_all[bi * cols..(bi + 1) * cols];
2383 unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
2385 }
2386 }
2387 };
2388 dispatch_rows(pool, rows, &run);
2389}
2390
2391fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
2400 (0..gpr)
2401 .map(|gi| {
2402 xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
2403 .iter()
2404 .map(|&v| v as i32)
2405 .sum()
2406 })
2407 .collect()
2408}
2409
2410#[inline]
2414#[allow(unreachable_code)]
2415fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
2416 #[cfg(target_arch = "aarch64")]
2417 unsafe {
2418 return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
2419 }
2420 let _ = gsum;
2421 let mut acc = 0f32;
2422 for gi in 0..gpr {
2423 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
2424 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2425 let mut d = 0i32;
2426 for (j, &b) in tile[2..].iter().enumerate() {
2427 for k in 0..8 {
2428 let w = ((b >> k) & 1) as i32 * 2 - 1;
2429 d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
2430 }
2431 }
2432 acc += d as f32 * s;
2433 }
2434 acc
2435}
2436
2437#[cfg(target_arch = "aarch64")]
2446#[target_feature(enable = "neon,dotprod")]
2447unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
2448 unsafe {
2451 use core::arch::aarch64::*;
2452 use core::arch::asm;
2453 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
2454 let m = vld1q_u8(MASKS.as_ptr());
2455 macro_rules! tile_dot {
2457 ($t:expr, $x:expr) => {{
2458 let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
2459 let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
2460 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
2461 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
2462 let x0 = vld1q_s8($x);
2463 let x1 = vld1q_s8($x.add(16));
2464 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
2465 asm!(
2466 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
2467 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
2468 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
2469 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
2470 options(pure, nomem, nostack),
2471 );
2472 vaddq_s32(a0, a1)
2473 }};
2474 }
2475 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
2476 let xp = xq.as_ptr();
2477 let gp = gsum.as_ptr();
2478 let mut accv = vdupq_n_f32(0.0);
2479 let mut gi = 0;
2480 while gi + 4 <= gpr {
2481 let t0 = base.add(gi * Q1_TILE);
2482 let t1 = base.add((gi + 1) * Q1_TILE);
2483 let t2 = base.add((gi + 2) * Q1_TILE);
2484 let t3 = base.add((gi + 3) * Q1_TILE);
2485 let d0 = tile_dot!(t0, xp.add(gi * GROUP_SIZE));
2486 let d1 = tile_dot!(t1, xp.add((gi + 1) * GROUP_SIZE));
2487 let d2 = tile_dot!(t2, xp.add((gi + 2) * GROUP_SIZE));
2488 let d3 = tile_dot!(t3, xp.add((gi + 3) * GROUP_SIZE));
2489 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
2491 let g = vld1q_s32(gp.add(gi));
2492 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
2493 let sc = [
2494 f16_to_f32(u16::from_le_bytes([*t0, *t0.add(1)])),
2495 f16_to_f32(u16::from_le_bytes([*t1, *t1.add(1)])),
2496 f16_to_f32(u16::from_le_bytes([*t2, *t2.add(1)])),
2497 f16_to_f32(u16::from_le_bytes([*t3, *t3.add(1)])),
2498 ];
2499 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), vld1q_f32(sc.as_ptr()));
2500 gi += 4;
2501 }
2502 let mut acc = vaddvq_f32(accv);
2503 while gi < gpr {
2504 let t = base.add(gi * Q1_TILE);
2505 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2506 let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
2507 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
2508 gi += 1;
2509 }
2510 acc
2511 }
2512}
2513
2514#[inline]
2516fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
2517 let gi = j / GROUP_SIZE;
2518 let k = j % GROUP_SIZE;
2519 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
2520 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2521 let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
2522 ((bit as i32 * 2 - 1) as f32, s)
2523}
2524
2525#[inline]
2527fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
2528 let mut acc = 0f32;
2529 for gi in 0..gpr {
2530 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
2531 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2532 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
2533 let mut ga = 0f32;
2534 for (j, &b) in tile[2..].iter().enumerate() {
2535 for k in 0..8 {
2536 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
2537 }
2538 }
2539 acc += ga * s;
2540 }
2541 acc
2542}
2543
2544#[allow(clippy::too_many_arguments)]
2547fn q1_range_a8w8(
2548 bytes: &[u8],
2549 gpr: usize,
2550 act: &SplitAct,
2551 gsum: &[i32],
2552 out: SendMut,
2553 start: usize,
2554 end: usize,
2555) {
2556 for r in start..end {
2557 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
2558 for &(j, xv) in &act.outliers {
2559 let (w, s) = q1_outlier(bytes, r, gpr, j);
2560 acc += w * s * xv;
2561 }
2562 unsafe { *out.at(r) = acc };
2564 }
2565}
2566
2567fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
2569 for r in start..end {
2570 unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
2572 }
2573}
2574
2575fn q1_matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
2577 debug_assert_eq!(out.len(), rows);
2578 let gpr = cols / GROUP_SIZE;
2579 let out_addr = SendMut(out.as_mut_ptr());
2580 if a8w8_enabled() {
2581 let act = split_act(x);
2582 let gsum = q1_group_sums(&act.xq, gpr);
2583 let (act, gsum) = (&act, &gsum);
2584 let run =
2585 move |start: usize, end: usize| q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end);
2586 dispatch_rows(pool, rows, &run);
2587 return;
2588 }
2589 let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
2590 dispatch_rows(pool, rows, &run);
2591}
2592
2593#[allow(clippy::too_many_arguments)]
2595fn q1_matvec2(
2596 bytes: &[u8],
2597 x1: &[f32],
2598 x2: &[f32],
2599 rows: usize,
2600 cols: usize,
2601 o1: &mut [f32],
2602 o2: &mut [f32],
2603 pool: Option<&Pool>,
2604) {
2605 let gpr = cols / GROUP_SIZE;
2606 let p1 = SendMut(o1.as_mut_ptr());
2607 let p2 = SendMut(o2.as_mut_ptr());
2608 if a8w8_enabled() {
2609 let a1 = split_act(x1);
2610 let a2 = split_act(x2);
2611 let g1 = q1_group_sums(&a1.xq, gpr);
2612 let g2 = q1_group_sums(&a2.xq, gpr);
2613 let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
2614 let run = move |start: usize, end: usize| {
2615 for r in start..end {
2616 let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
2617 let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
2618 for &(j, xv) in &a1.outliers {
2619 let (w, s) = q1_outlier(bytes, r, gpr, j);
2620 v1 += w * s * xv;
2621 }
2622 for &(j, xv) in &a2.outliers {
2623 let (w, s) = q1_outlier(bytes, r, gpr, j);
2624 v2 += w * s * xv;
2625 }
2626 unsafe {
2628 *p1.at(r) = v1;
2629 *p2.at(r) = v2;
2630 }
2631 }
2632 };
2633 dispatch_rows(pool, rows, &run);
2634 return;
2635 }
2636 let run = move |start: usize, end: usize| {
2637 for r in start..end {
2638 unsafe {
2640 *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
2641 *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
2642 }
2643 }
2644 };
2645 dispatch_rows(pool, rows, &run);
2646}
2647
2648#[allow(clippy::too_many_arguments)]
2650fn q1_matmat(
2651 bytes: &[u8],
2652 xs_all: &[f32],
2653 b: usize,
2654 rows: usize,
2655 cols: usize,
2656 out: &mut [f32],
2657 pool: Option<&Pool>,
2658) {
2659 debug_assert_eq!(out.len(), b * rows);
2660 let gpr = cols / GROUP_SIZE;
2661 let out_addr = SendMut(out.as_mut_ptr());
2662 if a8w8_enabled() {
2663 let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
2664 .map(|bi| {
2665 let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
2666 let gsum = q1_group_sums(&act.xq, gpr);
2667 (act, gsum)
2668 })
2669 .collect();
2670 let acts = &acts;
2671 let run = move |start: usize, end: usize| {
2672 for r in start..end {
2673 for (bi, (act, gsum)) in acts.iter().enumerate() {
2674 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
2675 for &(j, xv) in &act.outliers {
2676 let (w, s) = q1_outlier(bytes, r, gpr, j);
2677 acc += w * s * xv;
2678 }
2679 unsafe { *out_addr.at(bi * rows + r) = acc };
2681 }
2682 }
2683 };
2684 dispatch_rows(pool, rows, &run);
2685 return;
2686 }
2687 let run = move |start: usize, end: usize| {
2688 for r in start..end {
2689 for bi in 0..b {
2690 let x = &xs_all[bi * cols..(bi + 1) * cols];
2691 unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
2693 }
2694 }
2695 };
2696 dispatch_rows(pool, rows, &run);
2697}
2698
2699fn q4matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
2705 debug_assert_eq!(out.len(), rows);
2706 let (packed, scales) = q4_split(bytes, rows, cols);
2707 let gpr = cols / GROUP_SIZE;
2708 let out_addr = SendMut(out.as_mut_ptr());
2709
2710 if a8w8_enabled() {
2711 let act = split_act(x);
2712 let run = move |start: usize, end: usize| {
2713 q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
2714 };
2715 dispatch_rows(pool, rows, &run);
2716 return;
2717 }
2718
2719 let run = move |start: usize, end: usize| {
2720 q4_range_f32(packed, scales, gpr, x, out_addr, start, end)
2721 };
2722 dispatch_rows(pool, rows, &run);
2723}
2724
2725#[inline]
2728#[allow(unreachable_code)]
2729#[cfg(target_arch = "x86_64")]
2734#[target_feature(enable = "avx2")]
2735unsafe fn dot_q4b_row_1x4_avx2(
2736 buf: &[u8],
2737 scales: &[u8],
2738 g0: usize,
2739 gpr: usize,
2740 xs: [&[i8]; 4],
2741) -> [f32; 4] {
2742 unsafe {
2744 use core::arch::x86_64::*;
2745 let ones = _mm256_set1_epi16(1);
2746 let mut acc = [0f32; 4];
2747 for gi in 0..gpr {
2748 let s = f16_to_f32(u16::from_le_bytes([
2749 scales[(g0 + gi) * 2],
2750 scales[(g0 + gi) * 2 + 1],
2751 ]));
2752 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2753 let aw = _mm256_abs_epi8(w);
2754 for (k, xq) in xs.iter().enumerate() {
2755 let x =
2756 _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2757 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
2758 let d = _mm256_madd_epi16(p16, ones);
2759 let hi128 = _mm256_extracti128_si256::<1>(d);
2760 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
2761 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2762 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2763 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
2764 }
2765 }
2766 acc
2767 }
2768}
2769
2770#[cfg(target_arch = "x86_64")]
2776#[target_feature(enable = "avx2")]
2777unsafe fn dot_q4b_row_1x4_sx_avx2(
2778 buf: &[u8],
2779 scales: &[u8],
2780 g0: usize,
2781 gpr: usize,
2782 xs: [&[i8]; 4],
2783 sxs: [f32; 4],
2784) -> [f32; 4] {
2785 unsafe {
2787 use core::arch::x86_64::*;
2788 let ones = _mm256_set1_epi16(1);
2789 let mut acc = [0f32; 4];
2790 for gi in 0..gpr {
2791 let s = f16_to_f32(u16::from_le_bytes([
2792 scales[(g0 + gi) * 2],
2793 scales[(g0 + gi) * 2 + 1],
2794 ]));
2795 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2796 let aw = _mm256_abs_epi8(w);
2797 for (k, xq) in xs.iter().enumerate() {
2798 let x =
2799 _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2800 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
2801 let d = _mm256_madd_epi16(p16, ones);
2802 let hi128 = _mm256_extracti128_si256::<1>(d);
2803 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
2804 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2805 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2806 acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
2807 }
2808 }
2809 acc
2810 }
2811}
2812
2813fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
2814 #[cfg(target_arch = "aarch64")]
2815 unsafe {
2816 return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
2817 }
2818 #[cfg(target_arch = "x86_64")]
2819 unsafe {
2820 return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
2821 }
2822 let mut acc = 0f32;
2823 for gi in 0..gpr {
2824 let g = g0 + gi;
2825 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
2826 let mut d = 0i32;
2827 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
2828 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
2829 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
2830 }
2831 acc += d as f32 * s;
2832 }
2833 acc
2834}
2835
2836#[inline]
2838#[allow(unreachable_code)]
2839fn dot_q4_row_i8_2(
2840 packed: &[u8],
2841 scales: &[u8],
2842 g0: usize,
2843 gpr: usize,
2844 xq1: &[i8],
2845 xq2: &[i8],
2846) -> (f32, f32) {
2847 #[cfg(target_arch = "aarch64")]
2848 unsafe {
2849 return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
2850 }
2851 #[cfg(target_arch = "x86_64")]
2852 unsafe {
2853 return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
2854 }
2855 (
2856 dot_q4_row_i8(packed, scales, g0, gpr, xq1),
2857 dot_q4_row_i8(packed, scales, g0, gpr, xq2),
2858 )
2859}
2860
2861#[allow(clippy::too_many_arguments)]
2864fn q4_range_a8w8(
2865 packed: &[u8],
2866 scales: &[u8],
2867 gpr: usize,
2868 cols: usize,
2869 act: &SplitAct,
2870 out: SendMut,
2871 start: usize,
2872 end: usize,
2873) {
2874 for r in start..end {
2875 let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
2876 for &(j, xv) in &act.outliers {
2878 let flat = r * cols + j;
2879 let byte = packed[flat / 2];
2880 let nib = if flat & 1 == 0 { byte & 0x0F } else { byte >> 4 };
2881 let s = f16_to_f32(u16::from_le_bytes([
2882 scales[(flat / GROUP_SIZE) * 2],
2883 scales[(flat / GROUP_SIZE) * 2 + 1],
2884 ]));
2885 acc += ((nib as i32 - 8) as f32) * s * xv;
2886 }
2887 unsafe { *out.at(r) = acc };
2889 }
2890}
2891
2892#[allow(clippy::too_many_arguments)]
2895fn q4_range2_a8w8(
2896 packed: &[u8],
2897 scales: &[u8],
2898 gpr: usize,
2899 cols: usize,
2900 a1: &SplitAct,
2901 a2: &SplitAct,
2902 p1: SendMut,
2903 p2: SendMut,
2904 start: usize,
2905 end: usize,
2906) {
2907 for r in start..end {
2908 let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
2909 let mut acc1 = s1 * a1.sx;
2910 let mut acc2 = s2 * a2.sx;
2911 let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
2913 for &(j, xv) in outliers {
2914 let flat = r * cols + j;
2915 let byte = packed[flat / 2];
2916 let nib = if flat & 1 == 0 { byte & 0x0F } else { byte >> 4 };
2917 let s = f16_to_f32(u16::from_le_bytes([
2918 scales[(flat / GROUP_SIZE) * 2],
2919 scales[(flat / GROUP_SIZE) * 2 + 1],
2920 ]));
2921 *acc += ((nib as i32 - 8) as f32) * s * xv;
2922 }
2923 };
2924 fix(&a1.outliers, &mut acc1);
2925 fix(&a2.outliers, &mut acc2);
2926 unsafe {
2928 *p1.at(r) = acc1;
2929 *p2.at(r) = acc2;
2930 }
2931 }
2932}
2933
2934fn q4_range_f32(
2936 packed: &[u8],
2937 scales: &[u8],
2938 gpr: usize,
2939 x: &[f32],
2940 out: SendMut,
2941 start: usize,
2942 end: usize,
2943) {
2944 for r in start..end {
2945 let mut acc = 0f32;
2946 for gi in 0..gpr {
2947 let g = r * gpr + gi;
2948 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
2949 let pk = &packed[g * 16..(g + 1) * 16];
2950 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
2951 let mut ga = 0f32;
2952 for (k, &b) in pk.iter().enumerate() {
2953 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
2954 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
2955 }
2956 acc += ga * s;
2957 }
2958 unsafe { *out.at(r) = acc };
2960 }
2961}
2962
2963#[allow(clippy::too_many_arguments)]
2967fn q4matvec2(
2968 bytes: &[u8],
2969 x1: &[f32],
2970 x2: &[f32],
2971 rows: usize,
2972 cols: usize,
2973 o1: &mut [f32],
2974 o2: &mut [f32],
2975 pool: Option<&Pool>,
2976) {
2977 debug_assert_eq!(o1.len(), rows);
2978 debug_assert_eq!(o2.len(), rows);
2979 let (packed, scales) = q4_split(bytes, rows, cols);
2980 let gpr = cols / GROUP_SIZE;
2981
2982 if a8w8_enabled() {
2983 let a1 = split_act(x1);
2984 let a2 = split_act(x2);
2985 let p1 = SendMut(o1.as_mut_ptr());
2986 let p2 = SendMut(o2.as_mut_ptr());
2987 let run = move |start: usize, end: usize| {
2988 q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
2989 };
2990 dispatch_rows(pool, rows, &run);
2991 return;
2992 }
2993
2994 let p1 = SendMut(o1.as_mut_ptr());
2995 let p2 = SendMut(o2.as_mut_ptr());
2996 let run = move |start: usize, end: usize| {
2997 q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
2998 };
2999 dispatch_rows(pool, rows, &run);
3000}
3001
3002#[allow(clippy::too_many_arguments)]
3004fn q4_range2_f32(
3005 packed: &[u8],
3006 scales: &[u8],
3007 gpr: usize,
3008 x1: &[f32],
3009 x2: &[f32],
3010 p1: SendMut,
3011 p2: SendMut,
3012 start: usize,
3013 end: usize,
3014) {
3015 for r in start..end {
3016 let (mut acc1, mut acc2) = (0f32, 0f32);
3017 for gi in 0..gpr {
3018 let g = r * gpr + gi;
3019 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
3020 let pk = &packed[g * 16..(g + 1) * 16];
3021 let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
3022 let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
3023 let (mut g1, mut g2) = (0f32, 0f32);
3024 for (k, &b) in pk.iter().enumerate() {
3025 let wl = (b & 0x0F) as f32 - 8.0;
3026 let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
3027 g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
3028 g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
3029 }
3030 acc1 += g1 * s;
3031 acc2 += g2 * s;
3032 }
3033 unsafe {
3035 *p1.at(r) = acc1;
3036 *p2.at(r) = acc2;
3037 }
3038 }
3039}
3040
3041thread_local! {
3042 static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
3045 static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3046}
3047
3048#[allow(clippy::too_many_arguments)]
3054fn q4matmat(
3055 bytes: &[u8],
3056 xs_all: &[f32],
3057 b: usize,
3058 rows: usize,
3059 cols: usize,
3060 out: &mut [f32],
3061 pool: Option<&Pool>,
3062) {
3063 debug_assert_eq!(xs_all.len(), b * cols);
3064 debug_assert_eq!(out.len(), b * rows);
3065 let (packed, scales) = q4_split(bytes, rows, cols);
3066 let gpr = cols / GROUP_SIZE;
3067 let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
3068
3069 if a8w8_enabled() {
3070 let acts: Vec<SplitAct> =
3071 (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
3072 let acts = &acts;
3073 let out_addr = SendMut(out.as_mut_ptr());
3074 let run = move |start: usize, end: usize| {
3075 ROW_I8.with(|rb| {
3076 let mut buf = rb.borrow_mut();
3077 buf.resize(cols, 0);
3078 for r in start..end {
3079 for gi in 0..gpr {
3083 let g = r * gpr + gi;
3084 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
3085 buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
3086 buf[gi * GROUP_SIZE + k * 2 + 1] =
3087 (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
3088 }
3089 }
3090 let mut bi = 0usize;
3091 #[cfg(target_arch = "x86_64")]
3092 if avx2_enabled()
3093 && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true)
3094 {
3095 while bi + 4 <= acts.len() {
3096 let xs = [
3097 acts[bi].xq.as_slice(),
3098 acts[bi + 1].xq.as_slice(),
3099 acts[bi + 2].xq.as_slice(),
3100 acts[bi + 3].xq.as_slice(),
3101 ];
3102 let d = unsafe {
3103 dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
3104 };
3105 for k in 0..4 {
3106 let act = &acts[bi + k];
3107 let mut acc = d[k] * act.sx;
3108 for &(j, xv) in &act.outliers {
3109 acc += (buf[j] as i8) as f32
3110 * gscale((r * cols + j) / GROUP_SIZE)
3111 * xv;
3112 }
3113 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
3115 }
3116 bi += 4;
3117 }
3118 }
3119 while bi < acts.len() {
3120 let act = &acts[bi];
3121 let mut acc = 0f32;
3122 for gi in 0..gpr {
3123 let d = dot_i8_i8(
3124 &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
3125 &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
3126 );
3127 acc += d as f32 * gscale(r * gpr + gi);
3128 }
3129 acc *= act.sx;
3130 for &(j, xv) in &act.outliers {
3132 acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
3133 }
3134 unsafe { *out_addr.at(bi * rows + r) = acc };
3136 bi += 1;
3137 }
3138 }
3139 })
3140 };
3141 dispatch_rows(pool, rows, &run);
3142 return;
3143 }
3144
3145 let out_addr = SendMut(out.as_mut_ptr());
3146 let run = move |start: usize, end: usize| {
3147 ROW_F32.with(|rb| {
3148 let mut buf = rb.borrow_mut();
3149 buf.resize(cols, 0.0);
3150 for r in start..end {
3151 for gi in 0..gpr {
3154 let g = r * gpr + gi;
3155 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
3156 buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
3157 buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
3158 }
3159 }
3160 for bi in 0..b {
3161 let x = &xs_all[bi * cols..(bi + 1) * cols];
3162 let mut acc = 0f32;
3163 for gi in 0..gpr {
3164 let mut ga = 0f32;
3165 for k in 0..GROUP_SIZE / 2 {
3170 let e = gi * GROUP_SIZE + k * 2;
3171 ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
3172 }
3173 acc += ga * gscale(r * gpr + gi);
3174 }
3175 unsafe { *out_addr.at(bi * rows + r) = acc };
3177 }
3178 }
3179 })
3180 };
3181 dispatch_rows(pool, rows, &run);
3182}
3183
3184#[allow(clippy::too_many_arguments)]
3189fn vbitmatmat(
3190 bytes: &[u8],
3191 offsets: &[usize],
3192 xs_all: &[f32],
3193 b: usize,
3194 rows: usize,
3195 cols: usize,
3196 out: &mut [f32],
3197 pool: Option<&Pool>,
3198) {
3199 debug_assert_eq!(xs_all.len(), b * cols);
3200 debug_assert_eq!(out.len(), b * rows);
3201 debug_assert_eq!(offsets.len(), rows + 1);
3202 let ng = cols / GROUP_SIZE;
3203 let bits = &bytes[..rows];
3204 let sc_off = rows;
3205 let gscale = |r: usize, g: usize| {
3206 let so = (r * ng + g) * 2;
3207 f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]))
3208 };
3209
3210 let decode_f32 = |r: usize, dst: &mut [f32]| {
3212 let bw = bits[r] as usize;
3213 let l = ((1i32 << (bw - 1)) - 1) as f32;
3214 let data = &bytes[offsets[r]..offsets[r + 1]];
3215 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3216 for d in dst.iter_mut() {
3217 while nbits < bw {
3218 acc = (acc << 8) | data[idx] as u64;
3219 idx += 1;
3220 nbits += 8;
3221 }
3222 let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
3223 nbits -= bw;
3224 *d = u - l;
3225 }
3226 };
3227
3228 if a8w8_enabled() {
3229 let acts: Vec<SplitAct> =
3230 (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
3231 let acts = &acts;
3232 let out_addr = SendMut(out.as_mut_ptr());
3233 let run = move |start: usize, end: usize| {
3234 for r in start..end {
3235 let bw = bits[r] as usize;
3236 if bw == 8 {
3237 ROW_F32.with(|rb| {
3240 let mut buf = rb.borrow_mut();
3241 buf.resize(cols, 0.0);
3242 decode_f32(r, &mut buf);
3243 for bi in 0..b {
3244 let x = &xs_all[bi * cols..(bi + 1) * cols];
3245 let mut dot = 0f32;
3246 for g in 0..ng {
3247 let mut gd = 0f32;
3248 for k in 0..GROUP_SIZE {
3249 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
3250 }
3251 dot += gd * gscale(r, g);
3252 }
3253 unsafe { *out_addr.at(bi * rows + r) = dot };
3255 }
3256 });
3257 continue;
3258 }
3259 let l = (1i32 << (bw - 1)) - 1;
3260 let data = &bytes[offsets[r]..offsets[r + 1]];
3261 ROW_I8.with(|rb| {
3262 let mut buf = rb.borrow_mut();
3263 buf.resize(cols, 0);
3264 #[inline(always)]
3265 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3266 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3267 let u = unpack8::<B>(&data[blk * B..]);
3268 for k in 0..8 {
3269 chunk[k] = (u[k] - l) as i8 as u8;
3270 }
3271 }
3272 }
3273 match bw {
3274 3 => fill::<3>(data, l, &mut buf),
3275 4 => vbit_fill4(data, &mut buf),
3276 5 => fill::<5>(data, l, &mut buf),
3277 6 => fill::<6>(data, l, &mut buf),
3278 _ => unreachable!("vbit bit-width {bw} (validated at load)"),
3279 }
3280 let mut bi = 0usize;
3281 #[cfg(target_arch = "x86_64")]
3285 if avx2_enabled()
3286 && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true)
3287 {
3288 while bi + 4 <= acts.len() {
3289 let xs = [
3290 acts[bi].xq.as_slice(),
3291 acts[bi + 1].xq.as_slice(),
3292 acts[bi + 2].xq.as_slice(),
3293 acts[bi + 3].xq.as_slice(),
3294 ];
3295 let sxs = [
3296 acts[bi].sx,
3297 acts[bi + 1].sx,
3298 acts[bi + 2].sx,
3299 acts[bi + 3].sx,
3300 ];
3301 let d = unsafe {
3302 dot_q4b_row_1x4_sx_avx2(&buf, &bytes[sc_off..], r * ng, ng, xs, sxs)
3303 };
3304 for k in 0..4 {
3305 let act = &acts[bi + k];
3306 let mut dot = d[k];
3307 for &(j, xv) in &act.outliers {
3308 dot +=
3309 (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
3310 }
3311 unsafe { *out_addr.at((bi + k) * rows + r) = dot };
3313 }
3314 bi += 4;
3315 }
3316 }
3317 while bi < acts.len() {
3318 let act = &acts[bi];
3319 let mut dot = 0f32;
3320 for g in 0..ng {
3321 let d = dot_i8_i8(
3322 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3323 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3324 ) as f32
3325 * act.sx;
3326 dot += d * gscale(r, g);
3327 }
3328 for &(j, xv) in &act.outliers {
3329 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
3330 }
3331 unsafe { *out_addr.at(bi * rows + r) = dot };
3333 bi += 1;
3334 }
3335 });
3336 }
3337 };
3338 dispatch_rows(pool, rows, &run);
3339 return;
3340 }
3341
3342 let out_addr = SendMut(out.as_mut_ptr());
3343 let run = move |start: usize, end: usize| {
3344 ROW_F32.with(|rb| {
3345 let mut buf = rb.borrow_mut();
3346 buf.resize(cols, 0.0);
3347 for r in start..end {
3348 decode_f32(r, &mut buf);
3349 for bi in 0..b {
3350 let x = &xs_all[bi * cols..(bi + 1) * cols];
3351 let mut dot = 0f32;
3352 for g in 0..ng {
3353 let mut gd = 0f32;
3354 for k in 0..GROUP_SIZE {
3355 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
3356 }
3357 dot += gd * gscale(r, g);
3358 }
3359 unsafe { *out_addr.at(bi * rows + r) = dot };
3361 }
3362 }
3363 })
3364 };
3365 dispatch_rows(pool, rows, &run);
3366}
3367
3368pub(crate) fn gpu_batch_job<'a>(
3372 t: &'a QTensor,
3373 x: &[f32],
3374) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
3375 match t {
3376 QTensor::Mapped {
3377 model,
3378 idx,
3379 dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
3380 rows,
3381 cols,
3382 row_scale,
3383 col_field,
3384 ..
3385 } => Some((
3386 model.clone(),
3387 crate::gpu::BatchJob {
3388 idx: *idx,
3389 rows: *rows,
3390 cols: *cols,
3391 row_scale,
3392 xs: prescale(x, col_field, *dt).into_owned(),
3393 q1: false,
3394 },
3395 )),
3396 QTensor::Mapped {
3398 model,
3399 idx,
3400 dtype: TensorDtype::Q1,
3401 rows,
3402 cols,
3403 ..
3404 } => Some((
3405 model.clone(),
3406 crate::gpu::BatchJob {
3407 idx: *idx,
3408 rows: *rows,
3409 cols: *cols,
3410 row_scale: &[],
3411 xs: x.to_vec(),
3412 q1: true,
3413 },
3414 )),
3415 _ => None,
3416 }
3417}
3418
3419pub(crate) fn prescale<'a>(
3423 x: &'a [f32],
3424 col_field: &[f32],
3425 dtype: TensorDtype,
3426) -> std::borrow::Cow<'a, [f32]> {
3427 if dtype == TensorDtype::Q8_2f {
3428 x.iter().zip(col_field).map(|(a, c)| a * c).collect()
3429 } else {
3430 std::borrow::Cow::Borrowed(x)
3431 }
3432}
3433
3434#[cfg(target_arch = "x86_64")]
3439pub(crate) fn avx2_enabled() -> bool {
3440 use std::sync::OnceLock;
3441 static ON: OnceLock<bool> = OnceLock::new();
3442 *ON.get_or_init(|| {
3443 std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
3444 && std::arch::is_x86_feature_detected!("avx2")
3445 && std::arch::is_x86_feature_detected!("fma")
3446 })
3447}
3448
3449#[cfg(target_arch = "x86_64")]
3454fn avx2_a8w8_enabled() -> bool {
3455 use std::sync::OnceLock;
3456 static ON: OnceLock<bool> = OnceLock::new();
3457 *ON.get_or_init(|| {
3458 avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
3459 })
3460}
3461
3462#[inline]
3466pub(crate) fn a8w8_enabled() -> bool {
3467 #[cfg(target_arch = "aarch64")]
3468 {
3469 sdot_enabled()
3470 }
3471 #[cfg(target_arch = "x86_64")]
3472 {
3473 avx2_a8w8_enabled()
3474 }
3475 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
3476 {
3477 false
3478 }
3479}
3480
3481#[inline]
3484#[allow(unreachable_code)]
3485fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
3486 #[cfg(target_arch = "aarch64")]
3487 unsafe {
3488 return dot_i8_sdot(w, xq);
3489 }
3490 #[cfg(target_arch = "x86_64")]
3491 unsafe {
3492 if avx512vnni_enabled() {
3493 return dot_i8_i8_vnni(w, xq);
3494 }
3495 return dot_i8_i8_avx2(w, xq);
3496 }
3497 w.iter().zip(xq).map(|(&a, &b)| (a as i8) as i32 * b as i32).sum()
3498}
3499
3500#[cfg(target_arch = "x86_64")]
3504fn avx512vnni_enabled() -> bool {
3505 use std::sync::OnceLock;
3506 static ON: OnceLock<bool> = OnceLock::new();
3507 *ON.get_or_init(|| {
3508 std::env::var("CMF_AVX512").map(|v| v != "0").unwrap_or(true)
3509 && std::arch::is_x86_feature_detected!("avx512f")
3510 && std::arch::is_x86_feature_detected!("avx512bw")
3511 && std::arch::is_x86_feature_detected!("avx512vl")
3512 && std::arch::is_x86_feature_detected!("avx512vnni")
3513 })
3514}
3515
3516#[cfg(target_arch = "x86_64")]
3521#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
3522unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
3523 unsafe {
3525 use core::arch::x86_64::*;
3526 let n = w.len();
3527 let mut j = 0usize;
3528 let mut total: i32;
3529 {
3534 #[inline(always)]
3535 unsafe fn step(
3536 w: *const u8,
3537 x: *const i8,
3538 acc: core::arch::x86_64::__m512i,
3539 ) -> core::arch::x86_64::__m512i {
3540 unsafe {
3541 use core::arch::x86_64::*;
3542 let wv = _mm512_loadu_si512(w as *const _);
3543 let xv = _mm512_loadu_si512(x as *const _);
3544 let aw = _mm512_abs_epi8(wv);
3545 let neg = _mm512_movepi8_mask(wv);
3546 let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
3547 _mm512_dpbusd_epi32(acc, aw, sx)
3548 }
3549 }
3550 let (mut a0, mut a1, mut a2, mut a3) = (
3551 _mm512_setzero_si512(),
3552 _mm512_setzero_si512(),
3553 _mm512_setzero_si512(),
3554 _mm512_setzero_si512(),
3555 );
3556 while j + 256 <= n {
3557 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
3558 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
3559 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
3560 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
3561 j += 256;
3562 }
3563 while j + 64 <= n {
3564 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
3565 j += 64;
3566 }
3567 let s01 = _mm512_add_epi32(a0, a1);
3568 let s23 = _mm512_add_epi32(a2, a3);
3569 total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
3570 }
3571 if j + 32 <= n {
3573 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
3574 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
3575 let d = _mm256_dpbusd_epi32(
3576 _mm256_setzero_si256(),
3577 _mm256_abs_epi8(wv),
3578 _mm256_sign_epi8(xv, wv),
3579 );
3580 let hi128 = _mm256_extracti128_si256::<1>(d);
3581 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3582 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3583 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3584 total += _mm_cvtsi128_si32(s32);
3585 j += 32;
3586 }
3587 while j < n {
3588 total += (w[j] as i8) as i32 * xq[j] as i32;
3589 j += 1;
3590 }
3591 total
3592 }
3593}
3594
3595#[cfg(target_arch = "x86_64")]
3597#[target_feature(enable = "avx2,fma")]
3598unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
3599 unsafe {
3601 use core::arch::x86_64::*;
3602 let n = x.len();
3603 let wp = w.as_ptr();
3604 let xp = x.as_ptr();
3605 let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
3606 let mut j = 0usize;
3607 while j + 16 <= n {
3608 let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
3609 let lo = _mm256_cvtepi8_epi32(wb);
3610 let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
3611 a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
3612 a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
3613 j += 16;
3614 }
3615 let acc = _mm256_add_ps(a0, a1);
3616 let hi128 = _mm256_extractf128_ps::<1>(acc);
3617 let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
3618 let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
3619 let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
3620 let mut sum = _mm_cvtss_f32(s32);
3621 while j < n {
3622 sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
3623 j += 1;
3624 }
3625 sum
3626 }
3627}
3628
3629#[cfg(target_arch = "x86_64")]
3634#[target_feature(enable = "avx2")]
3635unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
3636 unsafe {
3638 use core::arch::x86_64::*;
3639 let n = w.len();
3640 let ones = _mm256_set1_epi16(1);
3641 let mut acc = _mm256_setzero_si256();
3642 let mut j = 0usize;
3643 while j + 32 <= n {
3644 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
3645 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
3646 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
3647 acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
3648 j += 32;
3649 }
3650 let hi128 = _mm256_extracti128_si256::<1>(acc);
3651 let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
3652 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3653 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3654 let mut s = _mm_cvtsi128_si32(s32);
3655 while j < n {
3656 s += (w[j] as i8) as i32 * xq[j] as i32;
3657 j += 1;
3658 }
3659 s
3660 }
3661}
3662
3663#[cfg(target_arch = "x86_64")]
3669#[target_feature(enable = "avx2")]
3670unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
3671 unsafe {
3673 use core::arch::x86_64::*;
3674 let n = w0.len();
3675 let ones = _mm256_set1_epi16(1);
3676 let mut acc = [[_mm256_setzero_si256(); 4]; 2];
3677 let mut j = 0usize;
3678 while j + 32 <= n {
3679 let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
3680 let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
3681 let aw0 = _mm256_abs_epi8(wv0);
3682 let aw1 = _mm256_abs_epi8(wv1);
3683 for (k, x) in xs.iter().enumerate() {
3684 let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
3685 let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
3686 acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
3687 let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
3688 acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
3689 }
3690 j += 32;
3691 }
3692 let mut out = [[0i32; 4]; 2];
3693 for r in 0..2 {
3694 for k in 0..4 {
3695 let a = acc[r][k];
3696 let hi128 = _mm256_extracti128_si256::<1>(a);
3697 let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
3698 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3699 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3700 out[r][k] = _mm_cvtsi128_si32(s32);
3701 }
3702 }
3703 if j < n {
3704 for (k, x) in xs.iter().enumerate() {
3705 for i in j..n {
3706 out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
3707 out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
3708 }
3709 }
3710 }
3711 out
3712 }
3713}
3714
3715#[cfg(target_arch = "x86_64")]
3720#[inline]
3721fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
3722 let dot = if avx512vnni_enabled() && row.len() >= 64 {
3723 (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
3724 } else {
3725 unsafe { dot_i8_i8_avx2(row, &act.xq) }
3726 };
3727 let mut acc = dot as f32 * act.sx;
3728 for &(j, xv) in &act.outliers {
3729 acc += (row[j] as i8) as f32 * xv;
3730 }
3731 acc
3732}
3733
3734#[cfg(target_arch = "x86_64")]
3738#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
3739unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
3740 unsafe {
3742 use core::arch::x86_64::*;
3743 let n = w.len();
3744 let flip = _mm512_set1_epi8(-128); #[inline(always)]
3746 unsafe fn step(
3747 w: *const u8,
3748 x: *const i8,
3749 flip: core::arch::x86_64::__m512i,
3750 acc: core::arch::x86_64::__m512i,
3751 ) -> core::arch::x86_64::__m512i {
3752 unsafe {
3753 use core::arch::x86_64::*;
3754 let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
3755 _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
3756 }
3757 }
3758 let (mut a0, mut a1, mut a2, mut a3) = (
3759 _mm512_setzero_si512(),
3760 _mm512_setzero_si512(),
3761 _mm512_setzero_si512(),
3762 _mm512_setzero_si512(),
3763 );
3764 let mut j = 0usize;
3765 while j + 256 <= n {
3766 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
3767 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
3768 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
3769 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
3770 j += 256;
3771 }
3772 while j + 64 <= n {
3773 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
3774 j += 64;
3775 }
3776 let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
3777 _mm512_add_epi32(a0, a1),
3778 _mm512_add_epi32(a2, a3),
3779 ));
3780 while j < n {
3782 total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
3783 j += 1;
3784 }
3785 total
3786 }
3787}
3788
3789#[cfg(target_arch = "x86_64")]
3795#[target_feature(enable = "avx2")]
3796unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
3797 unsafe {
3800 use core::arch::x86_64::*;
3801 let lomask = _mm_set1_epi8(0x0F);
3802 let eight = _mm256_set1_epi8(8);
3803 let ones = _mm256_set1_epi16(1);
3804 let mut acc = 0f32;
3805 for gi in 0..gpr {
3806 let g = g0 + gi;
3807 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
3808 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
3809 let lo = _mm_and_si128(b, lomask);
3810 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3811 let w = _mm256_sub_epi8(
3812 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3813 eight,
3814 );
3815 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3816 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
3817 let d = _mm256_madd_epi16(p16, ones);
3818 let hi128 = _mm256_extracti128_si256::<1>(d);
3819 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3820 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3821 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3822 acc += _mm_cvtsi128_si32(s32) as f32 * s;
3823 }
3824 acc
3825 }
3826}
3827
3828#[cfg(target_arch = "x86_64")]
3831#[target_feature(enable = "avx2")]
3832unsafe fn dot_q4_row_avx2_2(
3833 packed: &[u8],
3834 scales: &[u8],
3835 g0: usize,
3836 gpr: usize,
3837 xq1: &[i8],
3838 xq2: &[i8],
3839) -> (f32, f32) {
3840 unsafe {
3842 use core::arch::x86_64::*;
3843 let lomask = _mm_set1_epi8(0x0F);
3844 let eight = _mm256_set1_epi8(8);
3845 let ones = _mm256_set1_epi16(1);
3846 let (mut acc1, mut acc2) = (0f32, 0f32);
3847 #[inline(always)]
3848 unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
3849 unsafe {
3850 use core::arch::x86_64::*;
3851 let hi128 = _mm256_extracti128_si256::<1>(d);
3852 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3853 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3854 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3855 _mm_cvtsi128_si32(s32)
3856 }
3857 }
3858 for gi in 0..gpr {
3859 let g = g0 + gi;
3860 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
3861 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
3862 let lo = _mm_and_si128(b, lomask);
3863 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3864 let w = _mm256_sub_epi8(
3865 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3866 eight,
3867 );
3868 let aw = _mm256_abs_epi8(w);
3869 let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3870 let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3871 let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
3872 let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
3873 acc1 += hsum(d1) as f32 * s;
3874 acc2 += hsum(d2) as f32 * s;
3875 }
3876 (acc1, acc2)
3877 }
3878}
3879
3880#[cfg(target_arch = "x86_64")]
3882fn q8_range_avx2(
3883 q: &[u8],
3884 row_scale: &[f32],
3885 act: &SplitAct,
3886 cols: usize,
3887 out_addr: SendMut,
3888 start: usize,
3889 end: usize,
3890) {
3891 for o in start..end {
3892 let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
3893 unsafe { *out_addr.at(o) = v };
3895 }
3896}
3897
3898#[cfg(target_arch = "x86_64")]
3900#[allow(clippy::too_many_arguments)]
3901fn q8_range2_avx2(
3902 q: &[u8],
3903 row_scale: &[f32],
3904 a1: &SplitAct,
3905 a2: &SplitAct,
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) = row_dot_avx2(row, a1) * row_scale[o];
3917 *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
3918 }
3919 }
3920}
3921
3922#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
3928fn sdot_enabled() -> bool {
3929 use std::sync::OnceLock;
3930 static ON: OnceLock<bool> = OnceLock::new();
3931 *ON.get_or_init(|| {
3932 let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
3933 #[cfg(target_arch = "aarch64")]
3934 {
3935 want && std::arch::is_aarch64_feature_detected!("dotprod")
3936 }
3937 #[cfg(not(target_arch = "aarch64"))]
3938 {
3939 let _ = want;
3940 false
3941 }
3942 })
3943}
3944
3945struct SplitAct {
3950 xq: Vec<i8>,
3951 sx: f32,
3952 outliers: Vec<(usize, f32)>,
3953 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
3956 xsum: i32,
3957}
3958
3959thread_local! {
3960 static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
3963 const { std::cell::RefCell::new(Vec::new()) };
3964}
3965
3966impl Drop for SplitAct {
3967 fn drop(&mut self) {
3968 let buf = std::mem::take(&mut self.xq);
3969 if buf.capacity() > 0 {
3970 XQ_FREE.with(|f| {
3971 let mut f = f.borrow_mut();
3972 if f.len() < 16 {
3973 f.push(buf);
3974 }
3975 });
3976 }
3977 }
3978}
3979
3980fn split_act(x: &[f32]) -> SplitAct {
3981 let n = x.len();
3982 let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
3983 let thr = 8.0 * rms;
3984 let mut outliers: Vec<(usize, f32)> = Vec::new();
3988 let mut amax = 0f32;
3989 for (j, &v) in x.iter().enumerate() {
3990 let a = v.abs();
3991 if a > thr {
3992 outliers.push((j, v));
3993 } else if a > amax {
3994 amax = a;
3995 }
3996 }
3997 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
3998 let inv = 1.0 / sx;
3999 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
4000 xq.clear();
4001 xq.reserve(n);
4002 if outliers.is_empty() {
4003 xq.extend(x.iter().map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8));
4004 } else {
4005 xq.extend(x.iter().map(|&v| {
4007 if v.abs() > thr {
4008 0
4009 } else {
4010 (v * inv).round().clamp(-127.0, 127.0) as i8
4011 }
4012 }));
4013 }
4014 let xsum = xq.iter().map(|&v| v as i32).sum();
4015 SplitAct { xq, sx, outliers, xsum }
4016}
4017
4018#[cfg(target_arch = "aarch64")]
4021#[target_feature(enable = "neon,dotprod")]
4022unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
4023 unsafe {
4025 use core::arch::aarch64::*;
4026 use core::arch::asm;
4027 let wp = w.as_ptr() as *const i8;
4028 let n = w.len();
4029 let (mut a0, mut a1, mut a2, mut a3) =
4030 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
4031 let mut i = 0;
4032 while i + 64 <= n {
4033 let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
4034 let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
4035 let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
4036 let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
4037 asm!(
4038 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
4039 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
4040 "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
4041 "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
4042 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
4043 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
4044 w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
4045 options(pure, nomem, nostack),
4046 );
4047 i += 64;
4048 }
4049 while i + 16 <= n {
4050 let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
4051 asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
4052 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
4053 i += 16;
4054 }
4055 let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
4056 while i < n {
4057 s += (*wp.add(i)) as i32 * xq[i] as i32;
4058 i += 1;
4059 }
4060 s
4061}
4062}
4063
4064#[cfg(target_arch = "aarch64")]
4068#[target_feature(enable = "neon,dotprod")]
4069unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
4070 unsafe {
4072 use core::arch::aarch64::*;
4073 use core::arch::asm;
4074 let n = xq.len();
4075 let px = xq.as_ptr();
4076 let (p0, p1, p2, p3) = (
4077 w0.as_ptr() as *const i8,
4078 w1.as_ptr() as *const i8,
4079 w2.as_ptr() as *const i8,
4080 w3.as_ptr() as *const i8,
4081 );
4082 let (mut a0, mut a1, mut a2, mut a3) =
4083 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
4084 let mut i = 0;
4085 while i + 16 <= n {
4086 let x = vld1q_s8(px.add(i));
4087 let v0 = vld1q_s8(p0.add(i));
4088 let v1 = vld1q_s8(p1.add(i));
4089 let v2 = vld1q_s8(p2.add(i));
4090 let v3 = vld1q_s8(p3.add(i));
4091 asm!(
4092 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
4093 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
4094 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
4095 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
4096 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
4097 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
4098 options(pure, nomem, nostack),
4099 );
4100 i += 16;
4101 }
4102 let mut r = [vaddvq_s32(a0), vaddvq_s32(a1), vaddvq_s32(a2), vaddvq_s32(a3)];
4103 while i < n {
4104 let xi = *px.add(i) as i32;
4105 r[0] += (*p0.add(i)) as i32 * xi;
4106 r[1] += (*p1.add(i)) as i32 * xi;
4107 r[2] += (*p2.add(i)) as i32 * xi;
4108 r[3] += (*p3.add(i)) as i32 * xi;
4109 i += 1;
4110 }
4111 r
4112}
4113}
4114
4115#[cfg(target_arch = "aarch64")]
4122#[target_feature(enable = "neon,dotprod")]
4123unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
4124 unsafe {
4127 use core::arch::aarch64::*;
4128 use core::arch::asm;
4129 let n = xq.len();
4130 let px = xq.as_ptr();
4131 let pg = g.as_ptr() as *const i8;
4132 let (mut a0, mut a1, mut a2, mut a3) =
4133 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
4134 let mut i = 0;
4135 while i + 16 <= n {
4136 let x = vld1q_s8(px.add(i));
4137 let base = pg.add(4 * i);
4138 let v0 = vld1q_s8(base);
4139 let v1 = vld1q_s8(base.add(16));
4140 let v2 = vld1q_s8(base.add(32));
4141 let v3 = vld1q_s8(base.add(48));
4142 asm!(
4143 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
4144 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
4145 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
4146 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
4147 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
4148 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
4149 options(pure, nomem, nostack),
4150 );
4151 i += 16;
4152 }
4153 [vaddvq_s32(a0), vaddvq_s32(a1), vaddvq_s32(a2), vaddvq_s32(a3)]
4154 }
4155}
4156
4157#[cfg(target_arch = "aarch64")]
4163fn q8_range_sdot(
4164 q: &[u8],
4165 rep: &[u8],
4166 row_scale: &[f32],
4167 act: &SplitAct,
4168 cols: usize,
4169 out_addr: SendMut,
4170 start: usize,
4171 end: usize,
4172) {
4173 let mut o = start;
4174 if !rep.is_empty() {
4177 while o < end && o % 4 != 0 {
4178 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
4179 unsafe { *out_addr.at(o) = v };
4180 o += 1;
4181 }
4182 }
4183 while o + 4 <= end {
4184 let r = if rep.is_empty() {
4185 unsafe {
4186 dot_i8_sdot_4rows(
4187 &q[o * cols..(o + 1) * cols],
4188 &q[(o + 1) * cols..(o + 2) * cols],
4189 &q[(o + 2) * cols..(o + 3) * cols],
4190 &q[(o + 3) * cols..(o + 4) * cols],
4191 &act.xq,
4192 )
4193 }
4194 } else {
4195 unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
4196 };
4197 for k in 0..4 {
4198 let mut acc = r[k] as f32 * act.sx;
4199 for &(j, xv) in &act.outliers {
4200 acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
4201 }
4202 unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
4204 }
4205 o += 4;
4206 }
4207 while o < end {
4208 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
4209 unsafe { *out_addr.at(o) = v };
4210 o += 1;
4211 }
4212}
4213
4214#[cfg(target_arch = "aarch64")]
4217#[allow(clippy::too_many_arguments)]
4218fn q8_range2_sdot(
4219 q: &[u8],
4220 row_scale: &[f32],
4221 a1: &SplitAct,
4222 a2: &SplitAct,
4223 cols: usize,
4224 p1: SendMut,
4225 p2: SendMut,
4226 start: usize,
4227 end: usize,
4228) {
4229 for o in start..end {
4230 let row = &q[o * cols..(o + 1) * cols];
4231 unsafe {
4233 *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
4234 *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
4235 }
4236 }
4237}
4238
4239#[allow(clippy::too_many_arguments)]
4241fn q8_range2_f32(
4242 q: &[u8],
4243 row_scale: &[f32],
4244 x1: &[f32],
4245 x2: &[f32],
4246 cols: usize,
4247 p1: SendMut,
4248 p2: SendMut,
4249 start: usize,
4250 end: usize,
4251) {
4252 for o in start..end {
4253 let row = &q[o * cols..(o + 1) * cols];
4254 unsafe {
4256 *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
4257 *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
4258 }
4259 }
4260}
4261
4262fn q8_range_f32(
4264 q: &[u8],
4265 row_scale: &[f32],
4266 xs: &[f32],
4267 cols: usize,
4268 out_addr: SendMut,
4269 start: usize,
4270 end: usize,
4271) {
4272 for o in start..end {
4273 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
4274 unsafe { *out_addr.at(o) = v };
4276 }
4277}
4278
4279#[cfg(target_arch = "aarch64")]
4282#[inline]
4283fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
4284 let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
4285 for &(j, xv) in &act.outliers {
4286 acc += (row[j] as i8) as f32 * xv;
4287 }
4288 acc
4289}
4290
4291#[cfg(target_arch = "aarch64")]
4299#[target_feature(enable = "neon,dotprod")]
4300unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
4301 unsafe {
4304 use core::arch::aarch64::*;
4305 use core::arch::asm;
4306 let lomask = vdupq_n_u8(0x0F);
4307 let eight = vdupq_n_s8(8);
4308 let mut acc = 0f32;
4309 for gi in 0..gpr {
4310 let g = g0 + gi;
4311 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4312 let b = vld1q_u8(packed.as_ptr().add(g * 16));
4313 let lo = vandq_u8(b, lomask);
4314 let hi = vshrq_n_u8::<4>(b);
4315 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4316 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4317 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4318 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4319 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4320 asm!(
4321 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4322 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4323 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4324 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4325 options(pure, nomem, nostack),
4326 );
4327 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4328 }
4329 acc
4330 }
4331}
4332
4333#[cfg(target_arch = "aarch64")]
4338#[target_feature(enable = "neon,dotprod")]
4339unsafe fn dot_q4_row_sdot2(
4340 packed: &[u8],
4341 scales: &[u8],
4342 g0: usize,
4343 gpr: usize,
4344 xq1: &[i8],
4345 xq2: &[i8],
4346) -> (f32, f32) {
4347 unsafe {
4350 use core::arch::aarch64::*;
4351 use core::arch::asm;
4352 let lomask = vdupq_n_u8(0x0F);
4353 let eight = vdupq_n_s8(8);
4354 let (mut acc1, mut acc2) = (0f32, 0f32);
4355 for gi in 0..gpr {
4356 let g = g0 + gi;
4357 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4358 let b = vld1q_u8(packed.as_ptr().add(g * 16));
4359 let lo = vandq_u8(b, lomask);
4360 let hi = vshrq_n_u8::<4>(b);
4361 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4362 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4363 let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
4364 let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
4365 let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
4366 let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
4367 let (mut a0, mut a1, mut b0, mut b1) =
4368 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
4369 asm!(
4370 "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
4371 "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
4372 "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
4373 "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
4374 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4375 b0 = inout(vreg) b0, b1 = inout(vreg) b1,
4376 e0 = in(vreg) e0, e1 = in(vreg) e1,
4377 x10 = in(vreg) x10, x11 = in(vreg) x11,
4378 x20 = in(vreg) x20, x21 = in(vreg) x21,
4379 options(pure, nomem, nostack),
4380 );
4381 acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4382 acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
4383 }
4384 (acc1, acc2)
4385 }
4386}
4387
4388#[inline]
4393pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
4394 #[cfg(target_arch = "aarch64")]
4395 unsafe {
4396 return axpy_i8_f32_neon(acc, row, w);
4397 }
4398 #[cfg(target_arch = "x86_64")]
4399 if avx2_enabled() {
4400 return unsafe { axpy_i8_f32_avx2(acc, row, w) };
4401 }
4402 #[allow(unreachable_code)]
4403 {
4404 for (a, &b) in acc.iter_mut().zip(row) {
4405 *a += w * b as f32;
4406 }
4407 }
4408}
4409
4410#[cfg(target_arch = "x86_64")]
4412#[target_feature(enable = "avx2,fma")]
4413unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
4414 unsafe {
4416 use core::arch::x86_64::*;
4417 let n = acc.len().min(row.len());
4418 let ap = acc.as_mut_ptr();
4419 let rp = row.as_ptr();
4420 let wv = _mm256_set1_ps(w);
4421 let mut j = 0usize;
4422 while j + 16 <= n {
4423 let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
4424 let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
4425 let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
4426 let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
4427 let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
4428 _mm256_storeu_ps(ap.add(j), v0);
4429 _mm256_storeu_ps(ap.add(j + 8), v1);
4430 j += 16;
4431 }
4432 while j < n {
4433 *ap.add(j) += w * (*rp.add(j)) as f32;
4434 j += 1;
4435 }
4436 }
4437}
4438
4439#[cfg(target_arch = "aarch64")]
4440#[target_feature(enable = "neon")]
4441unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
4442 unsafe {
4444 use core::arch::aarch64::*;
4445 let n = acc.len().min(row.len());
4446 let ap = acc.as_mut_ptr();
4447 let rp = row.as_ptr();
4448 let wv = vdupq_n_f32(w);
4449 let mut j = 0usize;
4450 while j + 16 <= n {
4451 let rb = vld1q_s8(rp.add(j));
4452 let lo = vmovl_s8(vget_low_s8(rb));
4453 let hi = vmovl_s8(vget_high_s8(rb));
4454 for (off, half) in [(0, lo), (8, hi)] {
4455 let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
4456 let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
4457 let o = j + off;
4458 vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
4459 vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
4460 }
4461 j += 16;
4462 }
4463 while j < n {
4464 *ap.add(j) += w * (*rp.add(j)) as f32;
4465 j += 1;
4466 }
4467}
4468}
4469
4470#[inline]
4473pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
4474 #[cfg(target_arch = "aarch64")]
4475 unsafe {
4476 return dot_i8_f32_neon(w, x);
4477 }
4478 #[cfg(target_arch = "x86_64")]
4479 if avx2_enabled() {
4480 return unsafe { dot_i8_f32_avx2(w, x) };
4481 }
4482 #[allow(unreachable_code)]
4483 {
4484 let mut sum = 0.0f32;
4485 for (j, &b) in w.iter().enumerate() {
4486 sum += (b as i8) as f32 * x[j];
4487 }
4488 sum
4489 }
4490}
4491
4492#[inline]
4496fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
4497 #[cfg(target_arch = "aarch64")]
4498 unsafe {
4499 return dot_i8_col_f32_neon(w, x, col);
4500 }
4501 #[allow(unreachable_code)]
4502 {
4503 let mut sum = 0.0f32;
4504 for (j, &b) in w.iter().enumerate() {
4505 sum += (b as i8) as f32 * x[j] * col[j];
4506 }
4507 sum
4508 }
4509}
4510
4511#[cfg(target_arch = "aarch64")]
4512#[target_feature(enable = "neon")]
4513unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
4514 unsafe {
4516 use core::arch::aarch64::*;
4517 let n = x.len();
4518 let wp = w.as_ptr() as *const i8;
4519 let xp = x.as_ptr();
4520 let cp = col.as_ptr();
4521 let (mut a0, mut a1, mut a2, mut a3) =
4522 (vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4523 let mut j = 0usize;
4524 while j + 16 <= n {
4525 let wb = vld1q_s8(wp.add(j));
4526 let lo = vmovl_s8(vget_low_s8(wb));
4527 let hi = vmovl_s8(vget_high_s8(wb));
4528 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
4529 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
4530 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
4531 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
4532 a0 = vfmaq_f32(a0, w0, vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))));
4533 a1 = vfmaq_f32(a1, w1, vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))));
4534 a2 = vfmaq_f32(a2, w2, vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))));
4535 a3 = vfmaq_f32(a3, w3, vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))));
4536 j += 16;
4537 }
4538 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
4539 while j < n {
4540 sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
4541 j += 1;
4542 }
4543 sum
4544}
4545}
4546
4547#[cfg(target_arch = "aarch64")]
4548#[target_feature(enable = "neon")]
4549unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
4550 unsafe {
4552 use core::arch::aarch64::*;
4553 let n = x.len();
4554 let wp = w.as_ptr() as *const i8;
4555 let xp = x.as_ptr();
4556 let (mut a0, mut a1, mut a2, mut a3) =
4557 (vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4558 let mut j = 0usize;
4559 while j + 16 <= n {
4560 let wb = vld1q_s8(wp.add(j));
4561 let lo = vmovl_s8(vget_low_s8(wb));
4562 let hi = vmovl_s8(vget_high_s8(wb));
4563 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
4564 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
4565 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
4566 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
4567 a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
4568 a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
4569 a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
4570 a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
4571 j += 16;
4572 }
4573 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
4574 while j < n {
4575 sum += (*wp.add(j)) as f32 * *xp.add(j);
4576 j += 1;
4577 }
4578 sum
4579}
4580}
4581
4582#[allow(clippy::too_many_arguments)]
4583fn qmatvec(
4584 q: &[u8],
4585 rep: &[u8],
4586 row_scale: &[f32],
4587 xs: &[f32],
4588 rows: usize,
4589 cols: usize,
4590 out: &mut [f32],
4591 pool: Option<&Pool>,
4592) {
4593 debug_assert_eq!(out.len(), rows);
4594 #[cfg(not(target_arch = "aarch64"))]
4595 let _ = rep;
4596
4597 #[cfg(target_arch = "aarch64")]
4598 if sdot_enabled() {
4599 let act = split_act(xs);
4600 let out_addr = SendMut(out.as_mut_ptr());
4601 let run_range = |start: usize, end: usize| {
4602 q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
4603 };
4604 match pool {
4605 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
4606 _ => run_range(0, rows),
4607 }
4608 return;
4609 }
4610 #[cfg(target_arch = "x86_64")]
4613 if avx2_a8w8_enabled() {
4614 let act = split_act(xs);
4615 let out_addr = SendMut(out.as_mut_ptr());
4616 let run_range =
4617 |start: usize, end: usize| q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end);
4618 match pool {
4619 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
4620 _ => run_range(0, rows),
4621 }
4622 return;
4623 }
4624
4625 let row_dot = |o: usize| -> f32 { dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o] };
4626 match pool {
4627 Some(pool) if rows >= 256 => {
4628 let out_addr = SendMut(out.as_mut_ptr());
4629 pool.run(&move |widx, n| {
4630 let chunk = rows.div_ceil(n);
4631 let start = widx * chunk;
4632 let end = (start + chunk).min(rows);
4633 for o in start..end {
4634 unsafe { *out_addr.at(o) = row_dot(o) };
4636 }
4637 });
4638 }
4639 _ => {
4640 for (o, dst) in out.iter_mut().enumerate() {
4641 *dst = row_dot(o);
4642 }
4643 }
4644 }
4645}
4646
4647#[allow(clippy::too_many_arguments)]
4648fn qmatvec2(
4649 q: &[u8],
4650 row_scale: &[f32],
4651 x1: &[f32],
4652 x2: &[f32],
4653 rows: usize,
4654 cols: usize,
4655 o1: &mut [f32],
4656 o2: &mut [f32],
4657 pool: Option<&Pool>,
4658) {
4659 #[cfg(target_arch = "aarch64")]
4660 if sdot_enabled() {
4661 let a1s = split_act(x1);
4662 let a2s = split_act(x2);
4663 let p1 = SendMut(o1.as_mut_ptr());
4664 let p2 = SendMut(o2.as_mut_ptr());
4665 let run_range = |start: usize, end: usize| {
4666 q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
4667 };
4668 match pool {
4669 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
4670 _ => run_range(0, rows),
4671 }
4672 return;
4673 }
4674 #[cfg(target_arch = "x86_64")]
4675 if avx2_a8w8_enabled() {
4676 let a1s = split_act(x1);
4677 let a2s = split_act(x2);
4678 let p1 = SendMut(o1.as_mut_ptr());
4679 let p2 = SendMut(o2.as_mut_ptr());
4680 let run_range = |start: usize, end: usize| {
4681 q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
4682 };
4683 match pool {
4684 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
4685 _ => run_range(0, rows),
4686 }
4687 return;
4688 }
4689
4690 let row_dots = |o: usize| -> (f32, f32) {
4691 let row = &q[o * cols..(o + 1) * cols];
4692 (dot_i8_f32(row, x1) * row_scale[o], dot_i8_f32(row, x2) * row_scale[o])
4693 };
4694 match pool {
4695 Some(pool) if rows >= 256 => {
4696 let p1 = SendMut(o1.as_mut_ptr());
4697 let p2 = SendMut(o2.as_mut_ptr());
4698 pool.run(&move |widx, n| {
4699 let chunk = rows.div_ceil(n);
4700 let start = widx * chunk;
4701 let end = (start + chunk).min(rows);
4702 for o in start..end {
4703 let (s1, s2) = row_dots(o);
4704 unsafe {
4706 *p1.at(o) = s1;
4707 *p2.at(o) = s2;
4708 }
4709 }
4710 });
4711 }
4712 _ => {
4713 for o in 0..rows {
4714 let (s1, s2) = row_dots(o);
4715 o1[o] = s1;
4716 o2[o] = s2;
4717 }
4718 }
4719 }
4720}
4721
4722#[derive(Clone, Copy)]
4723struct SendMut(*mut f32);
4724unsafe impl Send for SendMut {}
4725unsafe impl Sync for SendMut {}
4726
4727impl SendMut {
4728 #[inline]
4729 fn at(self, i: usize) -> *mut f32 {
4730 unsafe { self.0.add(i) }
4731 }
4732}
4733
4734#[cfg(test)]
4735mod tests {
4736 use super::*;
4737
4738 #[test]
4739 fn f32_matvec_matches_matvec_rows_bitexact() {
4740 let (rows, cols) = (300, 40);
4741 let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
4742 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
4743 let qt = QTensor::from_f32(w.clone(), rows, cols);
4744
4745 let mut a = vec![0.0f32; rows];
4746 matvec_rows(None, &w, &x, &mut a);
4747 let mut b = vec![0.0f32; rows];
4748 qt.matvec(&x, &mut b, None);
4749 assert_eq!(a, b);
4750 }
4751
4752 #[test]
4753 fn sdot_kernel_exact_on_grid() {
4754 eprintln!("sdot_enabled = {}", sdot_enabled());
4759 let (rows, cols) = (9, 80); let w: Vec<u8> = (0..rows * cols)
4761 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
4762 .collect();
4763 let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
4764 let x: Vec<f32> = (0..cols)
4765 .map(|i| match i % 3 {
4766 0 => 1.0,
4767 1 => -1.0,
4768 _ => 0.0,
4769 })
4770 .collect();
4771 let mut a = vec![0.0f32; rows];
4772 qmatvec(&w, &[], &scales, &x, rows, cols, &mut a, None);
4773 for o in 0..rows {
4774 let mut acc = 0.0f32;
4775 for j in 0..cols {
4776 acc += (w[o * cols + j] as i8) as f32 * x[j];
4777 }
4778 let expect = acc * scales[o];
4779 assert!(
4780 (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
4781 "row {o}: {} vs {expect}",
4782 a[o]
4783 );
4784 }
4785 }
4786
4787 #[test]
4788 fn q1_kernels_match_exact_reference() {
4789 let (rows, cols) = (7, 96);
4791 let gpr = cols / GROUP_SIZE;
4792 let mut bytes = Vec::new();
4793 for t in 0..rows * gpr {
4794 let s = 0.01 + (t % 13) as f32 * 0.003;
4795 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
4796 for j in 0..4 {
4797 bytes.push(((t * 31 + j * 97) % 251) as u8);
4798 }
4799 }
4800 let x: Vec<f32> = (0..cols)
4802 .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
4803 .collect();
4804 let mut w = vec![0.0f32; rows * cols];
4806 cortiq_core::quant::dequant_q1(&bytes, &mut w);
4807 let mut expect = vec![0.0f32; rows];
4808 for o in 0..rows {
4809 expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
4810 }
4811 let mut got = vec![0.0f32; rows];
4812 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
4813 for o in 0..rows {
4814 assert!(
4815 (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
4816 "row {o}: {} vs {}",
4817 got[o],
4818 expect[o]
4819 );
4820 }
4821 let x2: Vec<f32> = x.iter().map(|v| -v).collect();
4823 let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
4824 q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
4825 assert_eq!(a1, got);
4826 let mut xs = x.clone();
4827 xs.extend_from_slice(&x2);
4828 let mut mm = vec![0.0f32; 2 * rows];
4829 q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
4830 assert_eq!(&mm[..rows], got.as_slice());
4831 assert_eq!(&mm[rows..], a2.as_slice());
4832 }
4833
4834 #[test]
4835 fn repack_is_bit_identical() {
4836 let (rows, cols) = (267, 96); let w: Vec<u8> = (0..rows * cols)
4842 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
4843 .collect();
4844 let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
4845 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
4846 let rep = q8_repack_layout(&w, rows, cols);
4847 for g in 0..rows / 4 {
4849 for c in 0..cols / 16 {
4850 for lane in 0..4 {
4851 assert_eq!(
4852 &rep[g * 4 * cols + c * 64 + lane * 16..g * 4 * cols + c * 64 + lane * 16 + 16],
4853 &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
4854 );
4855 }
4856 }
4857 }
4858 let mut a = vec![0.0f32; rows];
4859 qmatvec(&w, &[], &scales, &x, rows, cols, &mut a, None);
4860 let mut b = vec![0.0f32; rows];
4861 qmatvec(&w, &rep, &scales, &x, rows, cols, &mut b, None);
4862 assert_eq!(a, b, "full-range repack output diverged");
4863
4864 #[cfg(target_arch = "aarch64")]
4865 if sdot_enabled() {
4866 let act = split_act(&x);
4868 let mut c1 = vec![0.0f32; rows];
4869 let mut c2 = vec![0.0f32; rows];
4870 q8_range_sdot(&w, &[], &scales, &act, cols, SendMut(c1.as_mut_ptr()), 3, rows - 2);
4871 q8_range_sdot(&w, &rep, &scales, &act, cols, SendMut(c2.as_mut_ptr()), 3, rows - 2);
4872 assert_eq!(c1, c2, "unaligned-range repack output diverged");
4873 }
4874 }
4875
4876 #[test]
4877 fn sdot_a8w8_noise_is_bounded() {
4878 let (rows, cols) = (16, 512);
4882 let w: Vec<u8> = (0..rows * cols)
4883 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
4884 .collect();
4885 let scales = vec![0.01f32; rows];
4886 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
4887 let mut a = vec![0.0f32; rows];
4888 qmatvec(&w, &[], &scales, &x, rows, cols, &mut a, None);
4889 let (mut num, mut den) = (0f64, 0f64);
4890 for o in 0..rows {
4891 let mut acc = 0.0f32;
4892 for j in 0..cols {
4893 acc += (w[o * cols + j] as i8) as f32 * x[j];
4894 }
4895 let expect = acc * scales[o];
4896 num += ((a[o] - expect) as f64).powi(2);
4897 den += (expect as f64).powi(2);
4898 }
4899 let rel = (num / den.max(1e-12)).sqrt();
4900 assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
4901 }
4902
4903 #[test]
4904 fn i8_dot_neon_matches_scalar() {
4905 let n = 100;
4906 let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
4907 let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
4908 let mut scalar = 0.0f32;
4909 for j in 0..n {
4910 scalar += (w[j] as i8) as f32 * x[j];
4911 }
4912 let fast = dot_i8_f32(&w, &x);
4913 assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
4914 }
4915
4916 #[test]
4918 fn vbitmatvec_matches_full_dequant() {
4919 let (rows, cols) = (6, 64);
4920 let ng = cols / GROUP_SIZE;
4921 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
4923 let mut bytes = bits.clone();
4924 for g in 0..rows * ng {
4925 let s = 0.02 + 0.001 * g as f32;
4926 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
4927 }
4928 for r in 0..rows {
4929 let b = bits[r] as usize;
4930 let (mut acc, mut nb) = (0u64, 0usize);
4931 let mut rowbytes = Vec::new();
4932 for i in 0..cols {
4933 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
4934 acc = (acc << b) | v;
4935 nb += b;
4936 while nb >= 8 {
4937 nb -= 8;
4938 rowbytes.push(((acc >> nb) & 0xFF) as u8);
4939 }
4940 }
4941 if nb > 0 {
4942 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
4943 }
4944 bytes.extend_from_slice(&rowbytes);
4945 }
4946 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
4947
4948 let mut reference = vec![0f32; rows * cols];
4949 cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
4950 let mut expect = vec![0f32; rows];
4951 for r in 0..rows {
4952 expect[r] = reference[r * cols..(r + 1) * cols]
4953 .iter()
4954 .zip(&x)
4955 .map(|(w, xv)| w * xv)
4956 .sum();
4957 }
4958 let mut got = vec![0f32; rows];
4959 let offsets = vbit_row_offsets(&bytes, rows, cols);
4960 vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
4961 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
4965 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
4966 for r in 0..rows {
4967 assert!(
4968 (got[r] - expect[r]).abs() < tol * scale,
4969 "row {r}: {} vs {}",
4970 got[r],
4971 expect[r]
4972 );
4973 }
4974 }
4975
4976 #[test]
4981 #[cfg(target_arch = "x86_64")]
4982 fn vbit_matmat_blocked_matches_per_row() {
4983 let (rows, cols, b) = (64usize, 128usize, 9usize);
4984 let ng = cols / GROUP_SIZE;
4985 let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
4986 let mut bytes = bits.clone();
4987 for g in 0..rows * ng {
4988 let sc = 0.02 + 0.0005 * g as f32;
4989 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
4990 }
4991 for r in 0..rows {
4992 let bw = bits[r] as usize;
4993 let (mut acc, mut nb) = (0u64, 0usize);
4994 let mut rowbytes = Vec::new();
4995 for i in 0..cols {
4996 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
4997 acc = (acc << bw) | v;
4998 nb += bw;
4999 while nb >= 8 {
5000 nb -= 8;
5001 rowbytes.push(((acc >> nb) & 0xFF) as u8);
5002 }
5003 }
5004 if nb > 0 {
5005 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
5006 }
5007 bytes.extend_from_slice(&rowbytes);
5008 }
5009 let x: Vec<f32> =
5010 (0..b * cols).map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5).collect();
5011 let offsets = vbit_row_offsets(&bytes, rows, cols);
5012 let mut y_a = vec![0f32; b * rows];
5013 let mut y_b = vec![0f32; b * rows];
5014 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
5015 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
5016 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
5017 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
5018 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
5019 let max_d =
5020 y_a.iter().zip(&y_b).map(|(p, q)| (p - q).abs()).fold(0.0f32, f32::max);
5021 assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
5022 }
5023
5024 #[test]
5025 fn q4matvec_matches_full_dequant() {
5026 let (rows, cols) = (8, 64);
5027 let groups = rows * cols / GROUP_SIZE;
5028 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
5030 for i in 0..groups * 16 {
5031 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
5032 }
5033 for g in 0..groups {
5034 let s = 0.01 + 0.003 * g as f32;
5035 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
5036 }
5037 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
5038
5039 let mut reference = vec![0.0f32; rows * cols];
5040 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
5041 let mut expect = vec![0.0f32; rows];
5042 for r in 0..rows {
5043 expect[r] = reference[r * cols..(r + 1) * cols]
5044 .iter()
5045 .zip(&x)
5046 .map(|(w, xv)| w * xv)
5047 .sum();
5048 }
5049
5050 let mut got = vec![0.0f32; rows];
5051 q4matvec(&bytes, &x, rows, cols, &mut got, None);
5052 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
5056 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
5057 for r in 0..rows {
5058 assert!(
5059 (got[r] - expect[r]).abs() < tol * scale,
5060 "row {r}: {} vs {}",
5061 got[r],
5062 expect[r]
5063 );
5064 }
5065 }
5066
5067 #[test]
5070 fn vbitmatvec2_equals_two_singles() {
5071 let (rows, cols) = (6, 64);
5072 let ng = cols / GROUP_SIZE;
5073 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
5074 let mut bytes = bits.clone();
5075 for g in 0..rows * ng {
5076 let s = 0.02 + 0.001 * g as f32;
5077 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
5078 }
5079 for r in 0..rows {
5080 let b = bits[r] as usize;
5081 let (mut acc, mut nb) = (0u64, 0usize);
5082 let mut rowbytes = Vec::new();
5083 for i in 0..cols {
5084 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
5085 acc = (acc << b) | v;
5086 nb += b;
5087 while nb >= 8 {
5088 nb -= 8;
5089 rowbytes.push(((acc >> nb) & 0xFF) as u8);
5090 }
5091 }
5092 if nb > 0 {
5093 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
5094 }
5095 bytes.extend_from_slice(&rowbytes);
5096 }
5097 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
5098 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
5099 let offsets = vbit_row_offsets(&bytes, rows, cols);
5100
5101 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
5102 vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
5103 vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
5104 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
5105 vbitmatvec2(&bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
5106 assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
5107 assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
5108 }
5109
5110 #[test]
5112 fn q4matvec2_equals_two_singles() {
5113 let (rows, cols) = (8, 128);
5114 let groups = rows * cols / GROUP_SIZE;
5115 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
5116 for i in 0..groups * 16 {
5117 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
5118 }
5119 for g in 0..groups {
5120 let s = 0.01 + 0.003 * g as f32;
5121 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
5122 }
5123 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
5126 x1[9] = 250.0;
5127 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
5128
5129 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
5130 q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
5131 q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
5132 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
5133 q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
5134 assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
5135 assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
5136 }
5137
5138 #[test]
5141 fn matvec_many_equals_separate_matvecs() {
5142 use crate::pool::Pool;
5143 let (r1, r2, cols) = (300, 200, 64);
5144 let mk = |salt: usize, rows: usize| {
5145 QTensor::from_f32(
5146 (0..rows * cols).map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5).collect(),
5147 rows,
5148 cols,
5149 )
5150 };
5151 let (a, b) = (mk(1, r1), mk(5, r2));
5152 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
5153 let pool = Pool::new(3);
5154
5155 let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
5156 a.matvec(&x, &mut ea, Some(&pool));
5157 b.matvec(&x, &mut eb, Some(&pool));
5158 let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
5159 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
5160 assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
5161 assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
5162 }
5163
5164 #[test]
5167 fn batched_matmat_equals_per_position_matvec() {
5168 let (rows, cols, b) = (8, 64, 5);
5169 let groups = rows * cols / GROUP_SIZE;
5171 let mut q4 = Vec::new();
5172 for i in 0..groups * 16 {
5173 q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
5174 }
5175 for g in 0..groups {
5176 q4.extend_from_slice(
5177 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
5178 );
5179 }
5180 let ng = cols / GROUP_SIZE;
5182 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
5183 let mut vb = bits.clone();
5184 for g in 0..rows * ng {
5185 vb.extend_from_slice(
5186 &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
5187 );
5188 }
5189 for r in 0..rows {
5190 let bw = bits[r] as usize;
5191 let (mut acc, mut nb) = (0u64, 0usize);
5192 let mut rowbytes = Vec::new();
5193 for i in 0..cols {
5194 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
5195 acc = (acc << bw) | v;
5196 nb += bw;
5197 while nb >= 8 {
5198 nb -= 8;
5199 rowbytes.push(((acc >> nb) & 0xFF) as u8);
5200 }
5201 }
5202 if nb > 0 {
5203 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
5204 }
5205 vb.extend_from_slice(&rowbytes);
5206 }
5207 let offsets = vbit_row_offsets(&vb, rows, cols);
5208
5209 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
5210
5211 let mut got = vec![0f32; b * rows];
5213 q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
5214 for bi in 0..b {
5215 let mut expect = vec![0f32; rows];
5216 q4matvec(&q4, &xs[bi * cols..(bi + 1) * cols], rows, cols, &mut expect, None);
5217 assert_eq!(&got[bi * rows..(bi + 1) * rows], &expect[..], "q4 batch pos {bi}");
5218 }
5219
5220 let mut got = vec![0f32; b * rows];
5222 vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
5223 for bi in 0..b {
5224 let mut expect = vec![0f32; rows];
5225 vbitmatvec(
5226 &vb, &offsets, &xs[bi * cols..(bi + 1) * cols], rows, cols, &mut expect, None,
5227 );
5228 assert_eq!(&got[bi * rows..(bi + 1) * rows], &expect[..], "vbit batch pos {bi}");
5229 }
5230 }
5231
5232 #[test]
5236 fn q4_tiled_matches_q4_block_bitexact() {
5237 let (rows, cols, b) = (8usize, 128usize, 3usize);
5238 let groups = rows * cols / GROUP_SIZE;
5239 let mut split = Vec::with_capacity(groups * 18);
5240 for i in 0..groups * 16 {
5241 split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
5242 }
5243 for g in 0..groups {
5244 split.extend_from_slice(
5245 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
5246 );
5247 }
5248 let (packed, scales) = split.split_at(groups * 16);
5250 let mut tiled = Vec::with_capacity(groups * Q4_TILE);
5251 for g in 0..groups {
5252 tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
5253 tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
5254 }
5255
5256 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
5257 x1[9] = 250.0; let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
5259
5260 let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
5261 q4matvec(&split, &x1, rows, cols, &mut a, None);
5262 q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
5263 assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
5264
5265 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
5266 let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
5267 q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
5268 q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
5269 assert_eq!(a1, t1);
5270 assert_eq!(a2, t2);
5271
5272 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
5273 let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
5274 q4matmat(&split, &xs, b, rows, cols, &mut am, None);
5275 q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
5276 assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
5277 }
5278
5279 #[test]
5286 fn q4matvec_sdot_outlier_exact() {
5287 let (rows, cols) = (4, 128);
5288 let groups = rows * cols / GROUP_SIZE;
5289 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
5290 for i in 0..groups * 16 {
5291 bytes.push(((i * 11 + 5) % 256) as u8);
5292 }
5293 for g in 0..groups {
5294 let s = 0.02 + 0.002 * g as f32;
5295 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
5296 }
5297 let mut x: Vec<f32> = (0..cols)
5298 .map(|i| match i % 3 {
5299 0 => 1.0,
5300 1 => -1.0,
5301 _ => 0.0,
5302 })
5303 .collect();
5304 x[17] = 300.0; let mut reference = vec![0.0f32; rows * cols];
5307 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
5308 let mut expect = vec![0.0f32; rows];
5309 for r in 0..rows {
5310 expect[r] = reference[r * cols..(r + 1) * cols]
5311 .iter()
5312 .zip(&x)
5313 .map(|(w, xv)| w * xv)
5314 .sum();
5315 }
5316 let mut got = vec![0.0f32; rows];
5317 q4matvec(&bytes, &x, rows, cols, &mut got, None);
5318 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
5319 for r in 0..rows {
5320 assert!(
5321 (got[r] - expect[r]).abs() < 2e-3 * scale,
5322 "row {r}: {} vs {} (outlier term must be exact)",
5323 got[r],
5324 expect[r]
5325 );
5326 }
5327 }
5328}