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