1use crate::pool::{matvec_rows, matvec_rows2, Pool};
15use cortiq_core::quant::{f16_to_f32, GROUP_SIZE, Q1_TILE, Q4_TILE};
16use cortiq_core::{CmfModel, TensorDtype};
17use std::sync::Arc;
18
19pub enum QTensor {
20 F32 {
21 data: Vec<f32>,
22 rows: usize,
23 cols: usize,
24 },
25 Mapped {
26 model: Arc<CmfModel>,
27 idx: usize,
29 dtype: TensorDtype,
30 rows: usize,
31 cols: usize,
32 row_scale: Vec<f32>,
34 col_field: Vec<f32>,
36 vbit_offsets: Vec<usize>,
40 repack: Vec<u8>,
48 },
49}
50
51fn repack_enabled() -> bool {
58 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
59 *ON.get_or_init(|| {
60 std::env::var("CMF_REPACK").map(|v| v == "1").unwrap_or(false)
61 })
62}
63
64fn q8_repack(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
68 #[cfg(target_arch = "aarch64")]
69 let arch_ok = sdot_enabled();
70 #[cfg(not(target_arch = "aarch64"))]
71 let arch_ok = false;
72 if !arch_ok || !repack_enabled() || rows < 256 || cols % 16 != 0 {
73 return Vec::new();
74 }
75 q8_repack_layout(bytes, rows, cols)
76}
77
78fn q8_repack_layout(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
81 let groups = rows / 4;
82 let mut rep = vec![0u8; groups * 4 * cols];
83 for g in 0..groups {
84 let dst = &mut rep[g * 4 * cols..(g + 1) * 4 * cols];
85 for c in 0..cols / 16 {
86 for lane in 0..4 {
87 let src = (g * 4 + lane) * cols + c * 16;
88 dst[c * 64 + lane * 16..c * 64 + lane * 16 + 16]
89 .copy_from_slice(&bytes[src..src + 16]);
90 }
91 }
92 }
93 rep
94}
95
96fn vbit_row_offsets(bytes: &[u8], rows: usize, cols: usize) -> Vec<usize> {
99 let ng = cols / GROUP_SIZE;
100 let bits = &bytes[..rows];
101 let mut offsets = Vec::with_capacity(rows + 1);
102 let mut off = rows + rows * ng * 2;
103 for r in 0..rows {
104 offsets.push(off);
105 off += (cols * bits[r] as usize + 7) / 8;
106 }
107 offsets.push(off);
108 offsets
109}
110
111impl QTensor {
112 pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
113 debug_assert_eq!(data.len(), rows * cols);
114 Self::F32 { data, rows, cols }
115 }
116
117 pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
120 let idx = model
123 .tensor_index(name)
124 .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
125 let entry = &model.tensors[idx];
126 if entry.shape.len() != 2 {
127 return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
128 }
129 let (rows, cols) = (entry.shape[0], entry.shape[1]);
130 let bytes = model.entry_bytes(entry);
131
132 match entry.dtype {
133 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
134 let n = rows * cols;
135 let scales_off = n;
136 let row_scale: Vec<f32> = (0..rows)
137 .map(|o| {
138 f16_to_f32(u16::from_le_bytes([
139 bytes[scales_off + o * 2],
140 bytes[scales_off + o * 2 + 1],
141 ]))
142 })
143 .collect();
144 let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
145 let col_off = n + rows * 2;
146 (0..cols)
147 .map(|i| {
148 f16_to_f32(u16::from_le_bytes([
149 bytes[col_off + i * 2],
150 bytes[col_off + i * 2 + 1],
151 ]))
152 })
153 .collect()
154 } else {
155 Vec::new()
156 };
157 Ok(Self::Mapped {
158 model: model.clone(),
159 idx,
160 dtype: entry.dtype,
161 rows,
162 cols,
163 row_scale,
164 col_field,
165 vbit_offsets: Vec::new(),
166 repack: q8_repack(bytes, rows, cols),
167 })
168 }
169 TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
171 model: model.clone(),
172 idx,
173 dtype: entry.dtype,
174 rows,
175 cols,
176 row_scale: Vec::new(),
177 col_field: Vec::new(),
178 vbit_offsets: vbit_row_offsets(bytes, rows, cols),
179 repack: Vec::new(),
180 }),
181 TensorDtype::VbitRo if cols % GROUP_SIZE == 0 => {
185 let (_, off_off, packed_off) =
186 cortiq_core::quant::vbit_ro_sections(rows, cols);
187 let offsets: Vec<usize> = (0..=rows)
188 .map(|r| {
189 packed_off + cortiq_core::quant::vbit_ro_offset(bytes, off_off, r)
190 })
191 .collect();
192 Ok(Self::Mapped {
193 model: model.clone(),
194 idx,
195 dtype: entry.dtype,
196 rows,
197 cols,
198 row_scale: Vec::new(),
199 col_field: Vec::new(),
200 vbit_offsets: offsets,
201 repack: Vec::new(),
202 })
203 }
204 TensorDtype::Q4Tiled if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
210 model: model.clone(),
211 idx,
212 dtype: entry.dtype,
213 rows,
214 cols,
215 row_scale: Vec::new(),
216 col_field: Vec::new(),
217 vbit_offsets: Vec::new(),
218 repack: Vec::new(),
219 }),
220 TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
221 model: model.clone(),
222 idx,
223 dtype: entry.dtype,
224 rows,
225 cols,
226 row_scale: Vec::new(),
227 col_field: Vec::new(),
228 vbit_offsets: Vec::new(),
229 repack: Vec::new(),
230 }),
231 TensorDtype::Q1 if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
233 model: model.clone(),
234 idx,
235 dtype: entry.dtype,
236 rows,
237 cols,
238 row_scale: Vec::new(),
239 col_field: Vec::new(),
240 vbit_offsets: Vec::new(),
241 repack: Vec::new(),
242 }),
243 _ => {
245 let mut data = vec![0.0f32; rows * cols];
246 cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
247 Ok(Self::from_f32(data, rows, cols))
248 }
249 }
250 }
251
252 pub(crate) fn is_q1(&self) -> bool {
255 matches!(self, Self::Mapped { dtype: TensorDtype::Q1, .. })
256 }
257
258 pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
261 match self {
262 Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
263 _ => None,
264 }
265 }
266
267 pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
270 match self {
271 Self::Mapped { idx, dtype: TensorDtype::Q1, rows, cols, .. } => {
272 Some((*idx, *rows, *cols))
273 }
274 _ => None,
275 }
276 }
277
278 pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
283 match self {
284 Self::Mapped {
285 idx,
286 dtype: TensorDtype::Q8Row,
287 rows,
288 cols,
289 row_scale,
290 col_field,
291 ..
292 } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
293 _ => None,
294 }
295 }
296
297 pub fn rows(&self) -> usize {
298 match self {
299 Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
300 }
301 }
302
303 pub fn cols(&self) -> usize {
304 match self {
305 Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
306 }
307 }
308
309 pub fn as_f32(&self) -> Option<&[f32]> {
312 match self {
313 Self::F32 { data, .. } => Some(data),
314 Self::Mapped { .. } => None,
315 }
316 }
317
318 fn quant_bytes(&self) -> &[u8] {
319 match self {
320 Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
321 Self::F32 { .. } => unreachable!("quant_bytes on F32"),
322 }
323 }
324
325 pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
327 let cols = self.cols();
328 debug_assert_eq!(dst.len(), cols);
329 match self {
330 Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
331 Self::Mapped {
332 dtype,
333 row_scale,
334 col_field,
335 vbit_offsets,
336 ..
337 } => {
338 if *dtype == TensorDtype::Q4Tiled {
339 let bytes = self.quant_bytes();
340 let gpr = cols / GROUP_SIZE;
341 for gi in 0..gpr {
342 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
343 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
344 for (k, &b) in tile[2..].iter().enumerate() {
345 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
346 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
347 }
348 }
349 return;
350 }
351 if *dtype == TensorDtype::Q4Block {
352 let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
353 let gpr = cols / GROUP_SIZE;
354 for gi in 0..gpr {
355 let g = r * gpr + gi;
356 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
357 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
358 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
359 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
360 }
361 }
362 return;
363 }
364 if *dtype == TensorDtype::Q1 {
365 let bytes = self.quant_bytes();
366 let gpr = cols / GROUP_SIZE;
367 for gi in 0..gpr {
368 let tile =
369 &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
370 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
371 for (j, &b) in tile[2..].iter().enumerate() {
372 for k in 0..8 {
373 dst[gi * GROUP_SIZE + j * 8 + k] =
374 (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
375 }
376 }
377 }
378 return;
379 }
380 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
381 let bytes = self.quant_bytes();
382 let rows = self.rows();
383 let ng = cols / GROUP_SIZE;
384 let bits = &bytes[..rows];
385 let sc_off = rows;
386 let off = vbit_offsets[r];
389 let b = bits[r] as usize;
390 let l = ((1usize << (b - 1)) - 1) as f32;
391 let data = &bytes[off..];
392 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
393 for (i, d) in dst.iter_mut().enumerate() {
394 while nbits < b {
395 acc = (acc << 8) | data[idx] as u64;
396 idx += 1;
397 nbits += 8;
398 }
399 let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
400 nbits -= b;
401 let so = (r * ng + i / GROUP_SIZE) * 2;
402 let sv = f16_to_f32(u16::from_le_bytes([
403 bytes[sc_off + so],
404 bytes[sc_off + so + 1],
405 ]));
406 *d = (u - l) * sv;
407 }
408 return;
409 }
410 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
411 let s = row_scale[r];
412 match dtype {
413 TensorDtype::Q8Row => {
414 for (d, &b) in dst.iter_mut().zip(q) {
415 *d = (b as i8) as f32 * s;
416 }
417 }
418 TensorDtype::Q8_2f => {
419 for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
420 *d = (b as i8) as f32 * s * col_field[i];
421 }
422 }
423 _ => unreachable!(),
424 }
425 }
426 }
427 }
428
429 pub fn sparse_col_ok(&self) -> bool {
434 match self {
435 Self::F32 { .. } => true,
436 Self::Mapped { dtype, .. } => {
437 matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
438 }
439 }
440 }
441
442 pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
446 let inter = self.cols();
447 let hidden = self.rows();
448 debug_assert_eq!(out.len(), hidden);
449 match self {
450 Self::F32 { data, .. } => {
451 for (k, o) in out.iter_mut().enumerate() {
452 *o += w * data[k * inter + c];
453 }
454 }
455 Self::Mapped {
456 dtype, row_scale, col_field, ..
457 } => {
458 let q = self.quant_bytes();
459 let colf = if *dtype == TensorDtype::Q8_2f {
460 col_field[c]
461 } else {
462 1.0
463 };
464 let wc = w * colf;
465 for (k, o) in out.iter_mut().enumerate() {
466 let b = q[k * inter + c] as i8 as f32;
467 *o += wc * b * row_scale[k];
468 }
469 }
470 }
471 }
472
473 pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
477 let cols = self.cols();
478 match self {
479 Self::F32 { data, .. } => {
480 let row = &data[r * cols..(r + 1) * cols];
481 row.iter().zip(x).map(|(w, v)| w * v).sum()
482 }
483 Self::Mapped { dtype, row_scale, col_field, .. } => match dtype {
484 TensorDtype::Q8Row => {
485 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
486 dot_i8_f32(q, x) * row_scale[r]
487 }
488 TensorDtype::Q8_2f => {
489 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
490 dot_i8_col_f32(q, x, col_field) * row_scale[r]
491 }
492 _ => {
493 self.row_f32(r, scratch);
494 scratch.iter().zip(x).map(|(w, v)| w * v).sum()
495 }
496 },
497 }
498 }
499
500 pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
503 match self {
504 Self::F32 { data, .. } => matvec_rows(pool, data, x, out),
505 Self::Mapped {
506 model,
507 idx,
508 dtype,
509 rows,
510 cols,
511 row_scale,
512 col_field,
513 vbit_offsets,
514 repack,
515 } => {
516 let _ = (model, idx);
517 if *dtype == TensorDtype::Q4Block {
518 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
519 return;
520 }
521 if *dtype == TensorDtype::Q4Tiled {
522 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
523 return;
524 }
525 if *dtype == TensorDtype::Q1 {
526 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
531 let t0 = std::time::Instant::now();
532 let arm = if crate::gpu::q1_force() {
533 crate::gpu::ProbeArm::Gpu
534 } else {
535 crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
536 };
537 match arm {
538 crate::gpu::ProbeArm::Gpu => {
539 if crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out) {
540 crate::gpu::probe_record(
541 crate::gpu::OpClass::Matvec,
542 true,
543 t0.elapsed(),
544 );
545 return;
546 }
547 }
548 crate::gpu::ProbeArm::CpuTimed => {
549 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
550 crate::gpu::probe_record(
551 crate::gpu::OpClass::Matvec,
552 false,
553 t0.elapsed(),
554 );
555 return;
556 }
557 crate::gpu::ProbeArm::Cpu => {}
558 }
559 }
560 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
561 return;
562 }
563 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
564 vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
565 return;
566 }
567 let xs = prescale(x, col_field, *dtype);
568 if *rows >= crate::gpu::min_rows()
573 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
574 && std::env::var("CMF_GPU_LMHEAD").map(|v| v != "0").unwrap_or(true)
575 && crate::gpu::enabled_here()
576 {
577 let t0 = std::time::Instant::now();
580 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
581 crate::gpu::ProbeArm::Gpu => {}
582 crate::gpu::ProbeArm::CpuTimed => {
583 qmatvec(self.quant_bytes(), repack, row_scale, &xs, *rows, *cols, out, pool);
584 crate::gpu::probe_record(
585 crate::gpu::OpClass::Matvec,
586 false,
587 t0.elapsed(),
588 );
589 return;
590 }
591 crate::gpu::ProbeArm::Cpu => {
592 qmatvec(self.quant_bytes(), repack, row_scale, &xs, *rows, *cols, out, pool);
593 return;
594 }
595 }
596 let frac = std::env::var("CMF_GPU_SPLIT")
597 .ok()
598 .and_then(|v| v.parse::<f32>().ok())
599 .unwrap_or(0.5)
600 .clamp(0.0, 1.0);
601 let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
602 let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
603 let bytes = self.quant_bytes();
604 let ok = std::thread::scope(|sc| {
605 let g = sc.spawn(|| {
606 crate::gpu::q8_matvec_range(
607 model,
608 *idx,
609 cpu_rows,
610 &row_scale[cpu_rows..],
611 &xs,
612 *rows - cpu_rows,
613 *cols,
614 out_gpu,
615 )
616 });
617 if cpu_rows > 0 {
618 let rep_cpu = if repack.is_empty() {
621 &[][..]
622 } else {
623 &repack[..(cpu_rows / 4) * 4 * *cols]
624 };
625 qmatvec(
626 &bytes[..cpu_rows * *cols],
627 rep_cpu,
628 &row_scale[..cpu_rows],
629 &xs,
630 cpu_rows,
631 *cols,
632 out_cpu,
633 pool,
634 );
635 }
636 g.join().unwrap_or(false)
637 });
638 if ok {
639 crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
640 return;
641 }
642 qmatvec(
645 &bytes[cpu_rows * *cols..(*rows) * *cols],
646 &[],
647 &row_scale[cpu_rows..],
648 &xs,
649 *rows - cpu_rows,
650 *cols,
651 out_gpu,
652 pool,
653 );
654 return;
655 }
656 qmatvec(self.quant_bytes(), repack, row_scale, &xs, *rows, *cols, out, pool);
657 }
658 }
659 }
660
661 pub fn matvec2(&self, x1: &[f32], x2: &[f32], o1: &mut [f32], o2: &mut [f32], pool: Option<&Pool>) {
663 match self {
664 Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
665 Self::Mapped {
666 dtype,
667 rows,
668 cols,
669 row_scale,
670 col_field,
671 vbit_offsets,
672 ..
673 } => {
674 if *dtype == TensorDtype::Q4Block {
675 q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
676 return;
677 }
678 if *dtype == TensorDtype::Q4Tiled {
679 q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
680 return;
681 }
682 if *dtype == TensorDtype::Q1 {
683 q1_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
684 return;
685 }
686 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
687 vbitmatvec2(self.quant_bytes(), vbit_offsets, x1, x2, *rows, *cols, o1, o2, pool);
688 return;
689 }
690 let x1s = prescale(x1, col_field, *dtype);
691 let x2s = prescale(x2, col_field, *dtype);
692 qmatvec2(self.quant_bytes(), row_scale, &x1s, &x2s, *rows, *cols, o1, o2, pool);
693 }
694 }
695 }
696}
697
698impl QTensor {
699 pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
704 let cols = self.cols();
705 let rows = self.rows();
706 debug_assert_eq!(xs_all.len(), b * cols);
707 debug_assert_eq!(out.len(), b * rows);
708 match self {
709 Self::F32 { data, .. } => {
710 let out_addr = SendMut(out.as_mut_ptr());
711 let run = |start: usize, end: usize| {
712 for o in start..end {
713 let row = &data[o * cols..(o + 1) * cols];
714 for bi in 0..b {
715 let x = &xs_all[bi * cols..(bi + 1) * cols];
716 let mut acc = 0f32;
717 for j in 0..cols {
718 acc += row[j] * x[j];
719 }
720 unsafe { *out_addr.at(bi * rows + o) = acc };
721 }
722 }
723 };
724 dispatch_rows(pool, rows, &run);
725 }
726 Self::Mapped {
727 dtype,
728 row_scale,
729 col_field,
730 vbit_offsets,
731 ..
732 } => {
733 if *dtype == TensorDtype::Q4Block {
734 q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
735 return;
736 }
737 if *dtype == TensorDtype::Q4Tiled {
738 q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
739 return;
740 }
741 if *dtype == TensorDtype::Q1 {
742 q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
743 return;
744 }
745 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
746 vbitmatmat(self.quant_bytes(), vbit_offsets, xs_all, b, rows, cols, out, pool);
747 return;
748 }
749 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
750 .map(|bi| {
751 prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype)
752 })
753 .collect();
754 if b >= 8
760 && b * rows * cols >= 128_000_000
761 && crate::gpu::enabled_here()
762 {
763 if let Self::Mapped { model, idx, .. } = self {
764 let t0 = std::time::Instant::now();
765 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
766 crate::gpu::ProbeArm::Gpu
767 if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
768 && !crate::gpu::q8_resident_or_upload(model, *idx) =>
769 {
770 let q = self.quant_bytes();
774 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
775 return;
776 }
777 crate::gpu::ProbeArm::Gpu => {
778 let flat: Vec<f32> =
779 pre.iter().flat_map(|v| v.iter().copied()).collect();
780 if crate::gpu::q8_matmat(
781 model, *idx, row_scale, &flat, b, rows, cols, out)
782 {
783 crate::gpu::probe_record(
784 crate::gpu::OpClass::Matmat,
785 true,
786 t0.elapsed(),
787 );
788 return;
789 }
790 }
791 crate::gpu::ProbeArm::CpuTimed => {
792 let q = self.quant_bytes();
793 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
794 crate::gpu::probe_record(
795 crate::gpu::OpClass::Matmat,
796 false,
797 t0.elapsed(),
798 );
799 return;
800 }
801 crate::gpu::ProbeArm::Cpu => {}
802 }
803 }
804 }
805 let q = self.quant_bytes();
806 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
807 }
808 }
809 }
810}
811
812impl QTensor {
813 pub fn matvec_many<const N: usize>(
820 ts: [&QTensor; N],
821 x: &[f32],
822 mut outs: [&mut [f32]; N],
823 pool: Option<&Pool>,
824 ) {
825 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
826 let uniform_q8 = ts.iter().all(|t| {
827 matches!(
828 t,
829 Self::Mapped { dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f, .. }
830 )
831 });
832 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
833 let uniform_q4 = ts
834 .iter()
835 .all(|t| matches!(t, Self::Mapped { dtype: TensorDtype::Q4Block, .. }));
836 let uniform_vbit = ts
837 .iter()
838 .all(|t| matches!(
839 t,
840 Self::Mapped { dtype: TensorDtype::Vbit | TensorDtype::VbitRo, .. }
841 ));
842 let uniform_q1 = ts
843 .iter()
844 .all(|t| matches!(t, Self::Mapped { dtype: TensorDtype::Q1, .. }));
845 let Some(pool) = pool else {
846 for (t, o) in ts.iter().zip(outs.iter_mut()) {
847 t.matvec(x, o, None);
848 }
849 return;
850 };
851 if total_rows < 256
852 || !(uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit || uniform_q1)
853 {
854 for (t, o) in ts.iter().zip(outs.iter_mut()) {
855 t.matvec(x, o, Some(pool));
856 }
857 return;
858 }
859
860 if uniform_q1 {
861 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
864 if a8w8_enabled() {
865 let act = split_act(x);
866 let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
867 let (act, gsum) = (&act, &gsum);
868 let closures: [_; N] = std::array::from_fn(|i| {
869 let (bytes, gpr, out) =
870 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
871 move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
872 });
873 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
874 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
875 pool.run_many(&parts);
876 } else {
877 let closures: [_; N] = std::array::from_fn(|i| {
878 let (bytes, gpr, out) =
879 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
880 move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
881 });
882 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
883 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
884 pool.run_many(&parts);
885 }
886 return;
887 }
888
889 if uniform_q4 || uniform_vbit {
890 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
891 if a8w8_enabled() {
893 let act = split_act(x);
894 let act = &act;
895 if uniform_q4 {
896 let closures: [_; N] = std::array::from_fn(|i| {
897 let (packed, scales) =
898 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
899 let (gpr, cols, out) =
900 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
901 move |s: usize, e: usize| {
902 q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
903 }
904 });
905 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
906 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
907 pool.run_many(&parts);
908 } else {
909 let closures: [_; N] = std::array::from_fn(|i| {
910 let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
911 let (bytes, rows, cols, out) =
912 (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), outs_addr[i]);
913 move |s: usize, e: usize| {
914 vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
915 }
916 });
917 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
918 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
919 pool.run_many(&parts);
920 }
921 return;
922 }
923 if uniform_q4 {
924 let closures: [_; N] = std::array::from_fn(|i| {
925 let (packed, scales) =
926 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
927 let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
928 move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
929 });
930 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
931 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
932 pool.run_many(&parts);
933 } else {
934 let closures: [_; N] = std::array::from_fn(|i| {
935 let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
936 let (bytes, rows, cols, out) =
937 (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), outs_addr[i]);
938 move |s: usize, e: usize| {
939 vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
940 }
941 });
942 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
943 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
944 pool.run_many(&parts);
945 }
946 return;
947 }
948
949 if uniform_f32 {
950 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
951 let closures: [_; N] = std::array::from_fn(|i| {
952 let Self::F32 { data, cols, .. } = ts[i] else { unreachable!() };
953 let out = outs_addr[i];
954 move |start: usize, end: usize| {
955 for o in start..end {
956 let row = &data[o * cols..(o + 1) * cols];
957 let mut sum = 0.0f32;
958 for j in 0..*cols {
959 sum += row[j] * x[j];
960 }
961 unsafe { *out.at(o) = sum };
963 }
964 }
965 });
966 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
967 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
968 pool.run_many(&parts);
969 return;
970 }
971
972 struct Ctx<'a> {
975 bytes: &'a [u8],
976 #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
977 rep: &'a [u8],
978 row_scale: &'a [f32],
979 cols: usize,
980 xs: std::borrow::Cow<'a, [f32]>,
981 }
982 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
983 let Self::Mapped { dtype, cols, row_scale, col_field, repack, .. } = ts[i] else {
984 unreachable!()
985 };
986 Ctx {
987 bytes: ts[i].quant_bytes(),
988 rep: repack,
989 row_scale,
990 cols: *cols,
991 xs: prescale(x, col_field, *dtype),
992 }
993 });
994 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
995 #[cfg(target_arch = "aarch64")]
996 if sdot_enabled() {
997 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
998 let closures: [_; N] = std::array::from_fn(|i| {
999 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
1000 move |start: usize, end: usize| {
1001 q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
1002 }
1003 });
1004 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1005 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1006 pool.run_many(&parts);
1007 return;
1008 }
1009 #[cfg(target_arch = "x86_64")]
1010 if avx2_a8w8_enabled() {
1011 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
1012 let closures: [_; N] = std::array::from_fn(|i| {
1013 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
1014 move |start: usize, end: usize| {
1015 q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
1016 }
1017 });
1018 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1019 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1020 pool.run_many(&parts);
1021 return;
1022 }
1023 let closures: [_; N] = std::array::from_fn(|i| {
1024 let (c, out) = (&ctxs[i], outs_addr[i]);
1025 move |start: usize, end: usize| {
1026 q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
1027 }
1028 });
1029 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1030 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1031 pool.run_many(&parts);
1032 }
1033}
1034
1035impl QTensor {
1036 #[allow(clippy::needless_range_loop)]
1041 pub fn matvec2_many<const N: usize>(
1042 ts: [&QTensor; N],
1043 x1: &[f32],
1044 x2: &[f32],
1045 mut o1s: [&mut [f32]; N],
1046 mut o2s: [&mut [f32]; N],
1047 pool: Option<&Pool>,
1048 ) {
1049 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
1050 let uniform_q8 = ts.iter().all(|t| {
1051 matches!(
1052 t,
1053 Self::Mapped { dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f, .. }
1054 )
1055 });
1056 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
1057 let uniform_q4 = ts
1058 .iter()
1059 .all(|t| matches!(t, Self::Mapped { dtype: TensorDtype::Q4Block, .. }));
1060 let uniform_vbit = ts
1061 .iter()
1062 .all(|t| matches!(
1063 t,
1064 Self::Mapped { dtype: TensorDtype::Vbit | TensorDtype::VbitRo, .. }
1065 ));
1066 let fusable = pool.is_some()
1067 && total_rows >= 256
1068 && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
1069 if !fusable {
1070 for i in 0..N {
1071 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
1072 }
1073 return;
1074 }
1075 let pool = pool.unwrap();
1076
1077 if uniform_q4 || uniform_vbit {
1078 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1079 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1080 if a8w8_enabled() {
1082 let a1 = split_act(x1);
1083 let a2 = split_act(x2);
1084 let (a1, a2) = (&a1, &a2);
1085 if uniform_q4 {
1086 let closures: [_; N] = std::array::from_fn(|i| {
1087 let (packed, scales) =
1088 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1089 let (gpr, cols, o1, o2) =
1090 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
1091 move |s: usize, e: usize| {
1092 q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
1093 }
1094 });
1095 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1096 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1097 pool.run_many(&parts);
1098 } else {
1099 let closures: [_; N] = std::array::from_fn(|i| {
1100 let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
1101 let (bytes, rows, cols, o1, o2) =
1102 (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), p1[i], p2[i]);
1103 move |s: usize, e: usize| {
1104 vbit_range2_a8w8(
1105 bytes, vbit_offsets, x1, x2, a1, a2, rows, cols, o1, o2, s, e,
1106 )
1107 }
1108 });
1109 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1110 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1111 pool.run_many(&parts);
1112 }
1113 return;
1114 }
1115 if uniform_q4 {
1116 let closures: [_; N] = std::array::from_fn(|i| {
1117 let (packed, scales) =
1118 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1119 let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
1120 move |s: usize, e: usize| {
1121 q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
1122 }
1123 });
1124 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1125 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1126 pool.run_many(&parts);
1127 } else {
1128 let closures: [_; N] = std::array::from_fn(|i| {
1129 let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
1130 let (bytes, rows, cols, o1, o2) =
1131 (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), p1[i], p2[i]);
1132 move |s: usize, e: usize| {
1133 vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
1134 }
1135 });
1136 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1137 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1138 pool.run_many(&parts);
1139 }
1140 return;
1141 }
1142
1143 if uniform_f32 {
1144 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1145 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1146 let closures: [_; N] = std::array::from_fn(|i| {
1147 let Self::F32 { data, cols, .. } = ts[i] else { unreachable!() };
1148 let (o1, o2) = (p1[i], p2[i]);
1149 move |start: usize, end: usize| {
1150 for o in start..end {
1151 let row = &data[o * cols..(o + 1) * cols];
1152 let (mut s1, mut s2) = (0.0f32, 0.0f32);
1153 for j in 0..*cols {
1154 s1 += row[j] * x1[j];
1155 s2 += row[j] * x2[j];
1156 }
1157 unsafe {
1159 *o1.at(o) = s1;
1160 *o2.at(o) = s2;
1161 }
1162 }
1163 }
1164 });
1165 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1166 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1167 pool.run_many(&parts);
1168 return;
1169 }
1170
1171 struct Ctx<'a> {
1172 bytes: &'a [u8],
1173 row_scale: &'a [f32],
1174 cols: usize,
1175 xs1: std::borrow::Cow<'a, [f32]>,
1176 xs2: std::borrow::Cow<'a, [f32]>,
1177 }
1178 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
1179 let Self::Mapped { dtype, cols, row_scale, col_field, .. } = ts[i] else {
1180 unreachable!()
1181 };
1182 Ctx {
1183 bytes: ts[i].quant_bytes(),
1184 row_scale,
1185 cols: *cols,
1186 xs1: prescale(x1, col_field, *dtype),
1187 xs2: prescale(x2, col_field, *dtype),
1188 }
1189 });
1190 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1191 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1192 #[cfg(target_arch = "aarch64")]
1193 if sdot_enabled() {
1194 let acts: [(SplitAct, SplitAct); N] =
1195 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
1196 let closures: [_; N] = std::array::from_fn(|i| {
1197 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
1198 move |start: usize, end: usize| {
1199 q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
1200 }
1201 });
1202 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1203 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1204 pool.run_many(&parts);
1205 return;
1206 }
1207 #[cfg(target_arch = "x86_64")]
1208 if avx2_a8w8_enabled() {
1209 let acts: [(SplitAct, SplitAct); N] =
1210 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
1211 let closures: [_; N] = std::array::from_fn(|i| {
1212 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
1213 move |start: usize, end: usize| {
1214 q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
1215 }
1216 });
1217 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1218 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1219 pool.run_many(&parts);
1220 return;
1221 }
1222 let closures: [_; N] = std::array::from_fn(|i| {
1223 let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
1224 move |start: usize, end: usize| {
1225 q8_range2_f32(c.bytes, c.row_scale, &c.xs1, &c.xs2, c.cols, o1, o2, start, end)
1226 }
1227 });
1228 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1229 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1230 pool.run_many(&parts);
1231 }
1232}
1233
1234#[cfg(target_os = "macos")]
1239mod accel_blas {
1240 #[link(name = "Accelerate", kind = "framework")]
1241 unsafe extern "C" {
1242 pub fn cblas_sgemm(
1243 order: i32,
1244 trans_a: i32,
1245 trans_b: i32,
1246 m: i32,
1247 n: i32,
1248 k: i32,
1249 alpha: f32,
1250 a: *const f32,
1251 lda: i32,
1252 b: *const f32,
1253 ldb: i32,
1254 beta: f32,
1255 c: *mut f32,
1256 ldc: i32,
1257 );
1258 }
1259}
1260
1261#[cfg(target_os = "macos")]
1262pub(crate) fn accel_gemm_enabled() -> bool {
1263 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1264 *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
1265}
1266
1267#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
1270pub(crate) fn accel_gemm_enabled() -> bool {
1271 true
1272}
1273
1274#[cfg(target_arch = "aarch64")]
1281#[allow(clippy::too_many_arguments)]
1282pub(crate) fn neon_gemm_rm(
1283 m: usize,
1284 n: usize,
1285 k: usize,
1286 alpha: f32,
1287 a: &[f32],
1288 lda: usize,
1289 b_mat: &[f32],
1290 ldb: usize,
1291 b_rows_are_n: bool,
1292 c: &mut [f32],
1293 ldc: usize,
1294) {
1295 debug_assert!(a.len() >= (m - 1) * lda + k);
1296 debug_assert!(c.len() >= (m - 1) * ldc + n);
1297 unsafe {
1299 use core::arch::aarch64::*;
1300 let mut i = 0usize;
1301 while i < m {
1302 let mi = (m - i).min(4);
1303 let mut j = 0usize;
1304 while j < n {
1305 let nj = (n - j).min(8);
1306 if mi == 4 && nj == 8 {
1307 let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
1308 let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
1309 let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
1310 let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
1311 for p in 0..k {
1312 let (b0, b1) = if b_rows_are_n {
1313 let base = b_mat.as_ptr().add(j * ldb + p);
1316 let g = |o: usize| *base.add(o * ldb);
1317 (
1318 [g(0), g(1), g(2), g(3)],
1319 [g(4), g(5), g(6), g(7)],
1320 )
1321 } else {
1322 let base = b_mat.as_ptr().add(p * ldb + j);
1323 (
1324 [*base, *base.add(1), *base.add(2), *base.add(3)],
1325 [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
1326 )
1327 };
1328 let bv0 = vld1q_f32(b0.as_ptr());
1329 let bv1 = vld1q_f32(b1.as_ptr());
1330 let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
1331 let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
1332 let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
1333 let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
1334 c0a = vfmaq_f32(c0a, a0, bv0);
1335 c0b = vfmaq_f32(c0b, a0, bv1);
1336 c1a = vfmaq_f32(c1a, a1, bv0);
1337 c1b = vfmaq_f32(c1b, a1, bv1);
1338 c2a = vfmaq_f32(c2a, a2, bv0);
1339 c2b = vfmaq_f32(c2b, a2, bv1);
1340 c3a = vfmaq_f32(c3a, a3, bv0);
1341 c3b = vfmaq_f32(c3b, a3, bv1);
1342 }
1343 let al = vdupq_n_f32(alpha);
1344 for (r, (ca, cb)) in
1345 [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)].iter().enumerate()
1346 {
1347 let dst = c.as_mut_ptr().add((i + r) * ldc + j);
1348 vst1q_f32(dst, vmulq_f32(*ca, al));
1349 vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
1350 }
1351 } else {
1352 for r in 0..mi {
1353 for q in 0..nj {
1354 let mut acc = 0f32;
1355 for p in 0..k {
1356 let bv = if b_rows_are_n {
1357 b_mat[(j + q) * ldb + p]
1358 } else {
1359 b_mat[p * ldb + j + q]
1360 };
1361 acc += a[(i + r) * lda + p] * bv;
1362 }
1363 c[(i + r) * ldc + j + q] = acc * alpha;
1364 }
1365 }
1366 }
1367 j += nj;
1368 }
1369 i += mi;
1370 }
1371 }
1372}
1373
1374#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
1376#[allow(clippy::too_many_arguments)]
1377pub(crate) fn sgemm_rm(
1378 m: usize,
1379 n: usize,
1380 k: usize,
1381 alpha: f32,
1382 a: &[f32],
1383 lda: usize,
1384 b_mat: &[f32],
1385 ldb: usize,
1386 b_rows_are_n: bool,
1387 c: &mut [f32],
1388 ldc: usize,
1389) {
1390 neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
1391}
1392
1393#[cfg(target_os = "macos")]
1396#[allow(clippy::too_many_arguments)]
1397pub(crate) fn sgemm_rm(
1398 m: usize,
1399 n: usize,
1400 k: usize,
1401 alpha: f32,
1402 a: &[f32],
1403 lda: usize,
1404 b_mat: &[f32],
1405 ldb: usize,
1406 b_rows_are_n: bool,
1407 c: &mut [f32],
1408 ldc: usize,
1409) {
1410 debug_assert!(a.len() >= (m - 1) * lda + k);
1411 debug_assert!(c.len() >= (m - 1) * ldc + n);
1412 if std::env::var("CMF_FORCE_NEON_GEMM").map(|v| v == "1").unwrap_or(false) {
1416 return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
1417 }
1418 unsafe {
1419 accel_blas::cblas_sgemm(
1420 101, 111, if b_rows_are_n { 112 } else { 111 },
1423 m as i32,
1424 n as i32,
1425 k as i32,
1426 alpha,
1427 a.as_ptr(),
1428 lda as i32,
1429 b_mat.as_ptr(),
1430 ldb as i32,
1431 0.0,
1432 c.as_mut_ptr(),
1433 ldc as i32,
1434 );
1435 }
1436}
1437
1438#[cfg(target_os = "macos")]
1445fn qmatmat_accel(
1446 q: &[u8],
1447 row_scale: &[f32],
1448 pre: &[std::borrow::Cow<'_, [f32]>],
1449 rows: usize,
1450 cols: usize,
1451 out: &mut [f32],
1452 pool: Option<&Pool>,
1453) {
1454 const TR: usize = 2048;
1459 let b = pre.len();
1460 thread_local! {
1461 static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
1462 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
1463 }
1464 XPANEL.with(|xp| {
1465 WTILE.with(|wt| {
1466 let mut xpanel = xp.borrow_mut();
1467 xpanel.clear();
1468 for x in pre {
1469 xpanel.extend_from_slice(x);
1470 }
1471 let mut wtile = wt.borrow_mut();
1472 wtile.resize(TR * cols, 0.0);
1473 let mut r0 = 0usize;
1474 while r0 < rows {
1475 let tr = TR.min(rows - r0);
1476 let wt_addr = SendMut(wtile.as_mut_ptr());
1478 let run = |start: usize, end: usize| {
1479 for r in start..end {
1480 let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
1481 let s = row_scale[r0 + r];
1482 let dst = unsafe {
1484 std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols)
1485 };
1486 for (d, &v) in dst.iter_mut().zip(row) {
1487 *d = (v as i8) as f32 * s;
1488 }
1489 }
1490 };
1491 dispatch_rows(pool, tr, &run);
1492 unsafe {
1494 accel_blas::cblas_sgemm(
1495 101, 111, 112, b as i32,
1499 tr as i32,
1500 cols as i32,
1501 1.0,
1502 xpanel.as_ptr(),
1503 cols as i32,
1504 wtile.as_ptr(),
1505 cols as i32,
1506 0.0,
1507 out.as_mut_ptr().add(r0),
1508 rows as i32,
1509 );
1510 }
1511 r0 += tr;
1512 }
1513 })
1514 });
1515}
1516
1517fn qmatmat(
1518 q: &[u8],
1519 row_scale: &[f32],
1520 pre: &[std::borrow::Cow<'_, [f32]>],
1521 rows: usize,
1522 cols: usize,
1523 out: &mut [f32],
1524 pool: Option<&Pool>,
1525) {
1526 let b = pre.len();
1527 debug_assert_eq!(out.len(), b * rows);
1528 #[cfg(target_os = "macos")]
1533 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
1534 qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
1535 return;
1536 }
1537 #[cfg(target_arch = "aarch64")]
1538 if sdot_enabled() {
1539 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
1540 let out_addr = SendMut(out.as_mut_ptr());
1541 let blocked_ok =
1544 std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true);
1545 if blocked_ok {
1546 let run = |start: usize, end: usize| {
1547 let mut o = start;
1548 while o < end {
1549 if o + 2 <= end {
1550 let r0 = &q[o * cols..(o + 1) * cols];
1551 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
1552 let mut bi = 0usize;
1553 while bi + 4 <= acts.len() {
1554 let xs = [
1555 acts[bi].xq.as_slice(),
1556 acts[bi + 1].xq.as_slice(),
1557 acts[bi + 2].xq.as_slice(),
1558 acts[bi + 3].xq.as_slice(),
1559 ];
1560 let d = unsafe { dot_i8_sdot_2x4(r0, r1, xs) };
1561 for (r, row) in [r0, r1].into_iter().enumerate() {
1562 for k in 0..4 {
1563 let act = &acts[bi + k];
1564 let mut v = d[r][k] as f32 * act.sx;
1565 for &(j, xv) in &act.outliers {
1566 v += (row[j] as i8) as f32 * xv;
1567 }
1568 unsafe {
1569 *out_addr.at((bi + k) * rows + o + r) =
1570 v * row_scale[o + r]
1571 };
1572 }
1573 }
1574 bi += 4;
1575 }
1576 while bi < acts.len() {
1577 for (r, row) in [r0, r1].into_iter().enumerate() {
1578 let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
1579 unsafe { *out_addr.at(bi * rows + o + r) = v };
1580 }
1581 bi += 1;
1582 }
1583 o += 2;
1584 } else {
1585 let row = &q[o * cols..(o + 1) * cols];
1586 for (bi, act) in acts.iter().enumerate() {
1587 let v = row_dot_sdot(row, act) * row_scale[o];
1588 unsafe { *out_addr.at(bi * rows + o) = v };
1589 }
1590 o += 1;
1591 }
1592 }
1593 };
1594 dispatch_rows(pool, rows, &run);
1595 return;
1596 }
1597 let run = |start: usize, end: usize| {
1598 for o in start..end {
1599 let row = &q[o * cols..(o + 1) * cols];
1600 for (bi, act) in acts.iter().enumerate() {
1601 let v = row_dot_sdot(row, act) * row_scale[o];
1602 unsafe { *out_addr.at(bi * rows + o) = v };
1603 }
1604 }
1605 };
1606 dispatch_rows(pool, rows, &run);
1607 return;
1608 }
1609 #[cfg(target_arch = "x86_64")]
1614 if avx2_a8w8_enabled() {
1615 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
1616 let out_addr = SendMut(out.as_mut_ptr());
1617 let blocked_ok = std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true);
1620 if !avx512vnni_enabled() && blocked_ok {
1621 let run = |start: usize, end: usize| {
1622 let mut o = start;
1623 while o < end {
1624 if o + 2 <= end {
1625 let r0 = &q[o * cols..(o + 1) * cols];
1626 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
1627 let mut bi = 0usize;
1628 while bi + 4 <= acts.len() {
1629 let xs = [
1630 acts[bi].xq.as_slice(),
1631 acts[bi + 1].xq.as_slice(),
1632 acts[bi + 2].xq.as_slice(),
1633 acts[bi + 3].xq.as_slice(),
1634 ];
1635 let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
1636 for (r, row) in [r0, r1].into_iter().enumerate() {
1637 for k in 0..4 {
1638 let act = &acts[bi + k];
1639 let mut v = d[r][k] as f32 * act.sx;
1640 for &(j, xv) in &act.outliers {
1641 v += (row[j] as i8) as f32 * xv;
1642 }
1643 unsafe {
1644 *out_addr.at((bi + k) * rows + o + r) =
1645 v * row_scale[o + r]
1646 };
1647 }
1648 }
1649 bi += 4;
1650 }
1651 while bi < acts.len() {
1652 for (r, row) in [r0, r1].into_iter().enumerate() {
1653 let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
1654 unsafe { *out_addr.at(bi * rows + o + r) = v };
1655 }
1656 bi += 1;
1657 }
1658 o += 2;
1659 } else {
1660 let row = &q[o * cols..(o + 1) * cols];
1661 for (bi, act) in acts.iter().enumerate() {
1662 let v = row_dot_avx2(row, act) * row_scale[o];
1663 unsafe { *out_addr.at(bi * rows + o) = v };
1664 }
1665 o += 1;
1666 }
1667 }
1668 };
1669 dispatch_rows(pool, rows, &run);
1670 return;
1671 }
1672 let run = |start: usize, end: usize| {
1673 for o in start..end {
1674 let row = &q[o * cols..(o + 1) * cols];
1675 for (bi, act) in acts.iter().enumerate() {
1676 let v = row_dot_avx2(row, act) * row_scale[o];
1677 unsafe { *out_addr.at(bi * rows + o) = v };
1678 }
1679 }
1680 };
1681 dispatch_rows(pool, rows, &run);
1682 return;
1683 }
1684 let out_addr = SendMut(out.as_mut_ptr());
1685 let run = |start: usize, end: usize| {
1686 for o in start..end {
1687 let row = &q[o * cols..(o + 1) * cols];
1688 for (bi, x) in pre.iter().enumerate() {
1689 let mut acc = 0f32;
1690 for j in 0..cols {
1691 acc += (row[j] as i8) as f32 * x[j];
1692 }
1693 unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
1694 }
1695 }
1696 };
1697 dispatch_rows(pool, rows, &run);
1698}
1699
1700fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
1703 match pool {
1704 Some(pool) if rows >= 256 => pool.run_rows(rows, run),
1705 _ => run(0, rows),
1706 }
1707}
1708
1709fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
1711 let groups = rows * cols / GROUP_SIZE;
1712 bytes.split_at(groups * 16)
1713}
1714
1715#[inline]
1720fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
1721 #[cfg(target_arch = "aarch64")]
1722 unsafe {
1723 return vbit_fill4_neon(data, buf);
1724 }
1725 #[cfg(target_arch = "x86_64")]
1726 if avx2_enabled() {
1727 return unsafe { vbit_fill4_avx2(data, buf) };
1728 }
1729 #[allow(unreachable_code)]
1730 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
1731 let u = unpack8::<4>(&data[blk * 4..]);
1732 for k in 0..8 {
1733 chunk[k] = (u[k] - 7) as i8 as u8;
1734 }
1735 }
1736}
1737
1738#[cfg(target_arch = "aarch64")]
1739#[target_feature(enable = "neon")]
1740unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
1741 unsafe {
1744 use core::arch::aarch64::*;
1745 let n = buf.len();
1746 let mask = vdupq_n_u8(0x0F);
1747 let seven = vdupq_n_s8(7);
1748 let mut g = 0usize;
1749 while g * 32 + 32 <= n {
1750 let b = vld1q_u8(data.as_ptr().add(g * 16));
1751 let hi = vshrq_n_u8::<4>(b);
1752 let lo = vandq_u8(b, mask);
1753 let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
1754 let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
1755 vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
1756 vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
1757 g += 1;
1758 }
1759 }
1760}
1761
1762#[cfg(target_arch = "x86_64")]
1763#[target_feature(enable = "avx2")]
1764unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
1765 unsafe {
1767 use core::arch::x86_64::*;
1768 let n = buf.len();
1769 let mask = _mm_set1_epi8(0x0F);
1770 let seven = _mm256_set1_epi8(7);
1771 let mut g = 0usize;
1772 while g * 32 + 32 <= n {
1773 let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
1774 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
1775 let lo = _mm_and_si128(b, mask);
1776 let z = _mm256_sub_epi8(
1777 _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
1778 seven,
1779 );
1780 _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
1781 g += 1;
1782 }
1783 }
1784}
1785
1786#[inline(always)]
1791fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
1792 let mut acc = 0u64;
1793 for i in 0..B {
1794 acc = (acc << 8) | data[i] as u64;
1795 }
1796 let mask = (1u64 << B) - 1;
1797 let mut out = [0i32; 8];
1798 for (k, o) in out.iter_mut().enumerate() {
1799 *o = ((acc >> ((7 - k) * B)) & mask) as i32;
1800 }
1801 out
1802}
1803
1804#[allow(clippy::too_many_arguments)]
1810fn vbitmatvec(
1811 bytes: &[u8],
1812 offsets: &[usize],
1813 x: &[f32],
1814 rows: usize,
1815 cols: usize,
1816 out: &mut [f32],
1817 pool: Option<&Pool>,
1818) {
1819 debug_assert_eq!(out.len(), rows);
1820 debug_assert_eq!(offsets.len(), rows + 1);
1821
1822 if a8w8_enabled() {
1826 let act = split_act(x);
1827 let out_addr = SendMut(out.as_mut_ptr());
1828 let run = move |start: usize, end: usize| {
1829 vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
1830 };
1831 dispatch_rows(pool, rows, &run);
1832 return;
1833 }
1834
1835 let out_addr = SendMut(out.as_mut_ptr());
1836 let run = move |start: usize, end: usize| {
1837 vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
1838 };
1839 dispatch_rows(pool, rows, &run);
1840}
1841
1842#[allow(clippy::too_many_arguments)]
1846fn vbit_range_a8w8(
1847 bytes: &[u8],
1848 offsets: &[usize],
1849 x: &[f32],
1850 act: &SplitAct,
1851 rows: usize,
1852 cols: usize,
1853 out: SendMut,
1854 start: usize,
1855 end: usize,
1856) {
1857 let ng = cols / GROUP_SIZE;
1858 let bits = &bytes[..rows];
1859 let sc_off = rows;
1860 let row_dot = |r: usize| -> f32 {
1861 let b = bits[r] as usize;
1862 let l = ((1i32 << (b - 1)) - 1) as i32;
1863 let mask = (1u64 << b) - 1;
1864 let data = &bytes[offsets[r]..offsets[r + 1]];
1865 if b == 8 {
1866 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
1868 let mut dot = 0f32;
1869 for g in 0..ng {
1870 let so = (r * ng + g) * 2;
1871 let sgf = f16_to_f32(u16::from_le_bytes([
1872 bytes[sc_off + so],
1873 bytes[sc_off + so + 1],
1874 ]));
1875 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
1876 let mut gd = 0f32;
1877 for &xv in xg.iter() {
1878 if nbits < 8 {
1879 acc = (acc << 8) | data[idx] as u64;
1880 idx += 1;
1881 nbits += 8;
1882 }
1883 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
1884 nbits -= 8;
1885 gd += (u - l) as f32 * xv;
1886 }
1887 dot += gd * sgf;
1888 }
1889 return dot;
1890 }
1891 thread_local! {
1895 static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
1896 const { std::cell::RefCell::new(Vec::new()) };
1897 }
1898 #[inline(always)]
1899 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
1900 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
1901 let u = unpack8::<B>(&data[blk * B..]);
1902 for k in 0..8 {
1903 chunk[k] = (u[k] - l) as i8 as u8;
1904 }
1905 }
1906 }
1907 let _ = mask;
1908 VBIT_SCRATCH.with(|scratch| {
1909 let mut buf = scratch.borrow_mut();
1910 buf.resize(cols, 0);
1911 match b {
1912 3 => fill::<3>(data, l, &mut buf),
1913 4 => vbit_fill4(data, &mut buf),
1914 5 => fill::<5>(data, l, &mut buf),
1915 6 => fill::<6>(data, l, &mut buf),
1916 _ => unreachable!(),
1917 }
1918 let mut dot = 0f32;
1919 for g in 0..ng {
1920 let so = (r * ng + g) * 2;
1921 let s = f16_to_f32(u16::from_le_bytes([
1922 bytes[sc_off + so],
1923 bytes[sc_off + so + 1],
1924 ]));
1925 let d = dot_i8_i8(
1926 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
1927 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
1928 ) as f32
1929 * act.sx;
1930 dot += d * s;
1931 }
1932 for &(j, xv) in &act.outliers {
1933 let so = (r * ng + j / GROUP_SIZE) * 2;
1934 let s = f16_to_f32(u16::from_le_bytes([
1935 bytes[sc_off + so],
1936 bytes[sc_off + so + 1],
1937 ]));
1938 dot += (buf[j] as i8) as f32 * s * xv;
1940 }
1941 dot
1942 })
1943 };
1944 for r in start..end {
1945 unsafe { *out.at(r) = row_dot(r) };
1947 }
1948}
1949
1950#[allow(clippy::too_many_arguments)]
1952fn vbit_range_f32(
1953 bytes: &[u8],
1954 offsets: &[usize],
1955 x: &[f32],
1956 rows: usize,
1957 cols: usize,
1958 out: SendMut,
1959 start: usize,
1960 end: usize,
1961) {
1962 let ng = cols / GROUP_SIZE;
1963 let bits = &bytes[..rows];
1964 let sc_off = rows;
1965 #[inline(always)]
1969 fn dot_row<const B: usize>(
1970 data: &[u8],
1971 bytes: &[u8],
1972 sc_off: usize,
1973 r: usize,
1974 ng: usize,
1975 x: &[f32],
1976 ) -> f32 {
1977 let l = ((1i32 << (B - 1)) - 1) as f32;
1978 let gbytes = GROUP_SIZE * B / 8;
1979 let mut dot = 0f32;
1980 for g in 0..ng {
1981 let so = (r * ng + g) * 2;
1982 let s = f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]));
1983 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
1984 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
1985 let mut gd = 0f32;
1986 for blk in 0..GROUP_SIZE / 8 {
1987 let u = unpack8::<B>(&gd0[blk * B..]);
1988 let xb = &xg[blk * 8..blk * 8 + 8];
1989 for k in 0..8 {
1990 gd += (u[k] as f32 - l) * xb[k];
1991 }
1992 }
1993 dot += gd * s;
1994 }
1995 dot
1996 }
1997 for r in start..end {
1998 let data = &bytes[offsets[r]..offsets[r + 1]];
1999 let v = match bits[r] {
2000 3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
2001 4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
2002 5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
2003 6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
2004 8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
2005 b => unreachable!("vbit bit-width {b} (validated at load)"),
2006 };
2007 unsafe { *out.at(r) = v };
2009 }
2010}
2011
2012#[allow(clippy::too_many_arguments)]
2017fn vbitmatvec2(
2018 bytes: &[u8],
2019 offsets: &[usize],
2020 x1: &[f32],
2021 x2: &[f32],
2022 rows: usize,
2023 cols: usize,
2024 o1: &mut [f32],
2025 o2: &mut [f32],
2026 pool: Option<&Pool>,
2027) {
2028 debug_assert_eq!(o1.len(), rows);
2029 debug_assert_eq!(o2.len(), rows);
2030
2031 if a8w8_enabled() {
2032 let a1 = split_act(x1);
2033 let a2 = split_act(x2);
2034 let p1 = SendMut(o1.as_mut_ptr());
2035 let p2 = SendMut(o2.as_mut_ptr());
2036 let run = move |start: usize, end: usize| {
2037 vbit_range2_a8w8(bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end)
2038 };
2039 dispatch_rows(pool, rows, &run);
2040 return;
2041 }
2042
2043 let p1 = SendMut(o1.as_mut_ptr());
2044 let p2 = SendMut(o2.as_mut_ptr());
2045 let run = move |start: usize, end: usize| {
2046 vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
2047 };
2048 dispatch_rows(pool, rows, &run);
2049}
2050
2051#[allow(clippy::too_many_arguments)]
2055fn vbit_range2_a8w8(
2056 bytes: &[u8],
2057 offsets: &[usize],
2058 x1: &[f32],
2059 x2: &[f32],
2060 a1: &SplitAct,
2061 a2: &SplitAct,
2062 rows: usize,
2063 cols: usize,
2064 p1: SendMut,
2065 p2: SendMut,
2066 start: usize,
2067 end: usize,
2068) {
2069 let ng = cols / GROUP_SIZE;
2070 let bits = &bytes[..rows];
2071 let sc_off = rows;
2072 let row_dots = |r: usize| -> (f32, f32) {
2073 let b = bits[r] as usize;
2074 let l = (1i32 << (b - 1)) - 1;
2075 let data = &bytes[offsets[r]..offsets[r + 1]];
2076 if b == 8 {
2077 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
2080 let (mut d1, mut d2) = (0f32, 0f32);
2081 for g in 0..ng {
2082 let so = (r * ng + g) * 2;
2083 let sgf = f16_to_f32(u16::from_le_bytes([
2084 bytes[sc_off + so],
2085 bytes[sc_off + so + 1],
2086 ]));
2087 let (mut g1, mut g2) = (0f32, 0f32);
2088 for k in 0..GROUP_SIZE {
2089 if nbits < 8 {
2090 acc = (acc << 8) | data[idx] as u64;
2091 idx += 1;
2092 nbits += 8;
2093 }
2094 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
2095 nbits -= 8;
2096 let w = (u - l) as f32;
2097 g1 += w * x1[g * GROUP_SIZE + k];
2098 g2 += w * x2[g * GROUP_SIZE + k];
2099 }
2100 d1 += g1 * sgf;
2101 d2 += g2 * sgf;
2102 }
2103 return (d1, d2);
2104 }
2105 thread_local! {
2106 static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
2107 const { std::cell::RefCell::new(Vec::new()) };
2108 }
2109 #[inline(always)]
2110 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
2111 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
2112 let u = unpack8::<B>(&data[blk * B..]);
2113 for k in 0..8 {
2114 chunk[k] = (u[k] - l) as i8 as u8;
2115 }
2116 }
2117 }
2118 VBIT_SCRATCH2.with(|scratch| {
2119 let mut buf = scratch.borrow_mut();
2120 buf.resize(cols, 0);
2121 match b {
2122 3 => fill::<3>(data, l, &mut buf),
2123 4 => vbit_fill4(data, &mut buf),
2124 5 => fill::<5>(data, l, &mut buf),
2125 6 => fill::<6>(data, l, &mut buf),
2126 _ => unreachable!(),
2127 }
2128 let (mut d1, mut d2) = (0f32, 0f32);
2129 for g in 0..ng {
2130 let so = (r * ng + g) * 2;
2131 let s = f16_to_f32(u16::from_le_bytes([
2132 bytes[sc_off + so],
2133 bytes[sc_off + so + 1],
2134 ]));
2135 let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2136 let v1 =
2137 dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
2138 let v2 =
2139 dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
2140 d1 += v1 * s;
2141 d2 += v2 * s;
2142 }
2143 for &(j, xv) in &a1.outliers {
2144 let so = (r * ng + j / GROUP_SIZE) * 2;
2145 let s = f16_to_f32(u16::from_le_bytes([
2146 bytes[sc_off + so],
2147 bytes[sc_off + so + 1],
2148 ]));
2149 d1 += (buf[j] as i8) as f32 * s * xv;
2150 }
2151 for &(j, xv) in &a2.outliers {
2152 let so = (r * ng + j / GROUP_SIZE) * 2;
2153 let s = f16_to_f32(u16::from_le_bytes([
2154 bytes[sc_off + so],
2155 bytes[sc_off + so + 1],
2156 ]));
2157 d2 += (buf[j] as i8) as f32 * s * xv;
2158 }
2159 (d1, d2)
2160 })
2161 };
2162 for r in start..end {
2163 let (v1, v2) = row_dots(r);
2164 unsafe {
2166 *p1.at(r) = v1;
2167 *p2.at(r) = v2;
2168 }
2169 }
2170}
2171
2172#[allow(clippy::too_many_arguments)]
2176fn vbit_range2_f32(
2177 bytes: &[u8],
2178 offsets: &[usize],
2179 x1: &[f32],
2180 x2: &[f32],
2181 rows: usize,
2182 cols: usize,
2183 p1: SendMut,
2184 p2: SendMut,
2185 start: usize,
2186 end: usize,
2187) {
2188 let ng = cols / GROUP_SIZE;
2189 let bits = &bytes[..rows];
2190 let sc_off = rows;
2191 #[inline(always)]
2192 #[allow(clippy::too_many_arguments)]
2193 fn dot_row2<const B: usize>(
2194 data: &[u8],
2195 bytes: &[u8],
2196 sc_off: usize,
2197 r: usize,
2198 ng: usize,
2199 x1: &[f32],
2200 x2: &[f32],
2201 ) -> (f32, f32) {
2202 let l = ((1i32 << (B - 1)) - 1) as f32;
2203 let gbytes = GROUP_SIZE * B / 8;
2204 let (mut d1, mut d2) = (0f32, 0f32);
2205 for g in 0..ng {
2206 let so = (r * ng + g) * 2;
2207 let s = f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]));
2208 let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2209 let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2210 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
2211 let (mut g1, mut g2) = (0f32, 0f32);
2212 for blk in 0..GROUP_SIZE / 8 {
2213 let u = unpack8::<B>(&gd0[blk * B..]);
2214 for k in 0..8 {
2215 let w = u[k] as f32 - l;
2216 g1 += w * x1g[blk * 8 + k];
2217 g2 += w * x2g[blk * 8 + k];
2218 }
2219 }
2220 d1 += g1 * s;
2221 d2 += g2 * s;
2222 }
2223 (d1, d2)
2224 }
2225 for r in start..end {
2226 let data = &bytes[offsets[r]..offsets[r + 1]];
2227 let (v1, v2) = match bits[r] {
2228 3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
2229 4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
2230 5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
2231 6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
2232 8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
2233 b => unreachable!("vbit bit-width {b} (validated at load)"),
2234 };
2235 unsafe {
2237 *p1.at(r) = v1;
2238 *p2.at(r) = v2;
2239 }
2240 }
2241}
2242
2243#[inline]
2250#[allow(unreachable_code)]
2251fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
2252 #[cfg(target_arch = "aarch64")]
2253 unsafe {
2254 return dot_q4t_row_sdot(bytes, r, gpr, xq);
2255 }
2256 #[cfg(target_arch = "x86_64")]
2257 unsafe {
2258 return dot_q4t_row_avx2(bytes, r, gpr, xq);
2259 }
2260 let mut acc = 0f32;
2261 for gi in 0..gpr {
2262 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
2263 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2264 let mut d = 0i32;
2265 for (k, &b) in tile[2..].iter().enumerate() {
2266 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
2267 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
2268 }
2269 acc += d as f32 * s;
2270 }
2271 acc
2272}
2273
2274#[cfg(target_arch = "aarch64")]
2275#[target_feature(enable = "neon,dotprod")]
2276unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
2277 unsafe {
2280 use core::arch::aarch64::*;
2281 use core::arch::asm;
2282 let lomask = vdupq_n_u8(0x0F);
2283 let eight = vdupq_n_s8(8);
2284 let mut acc = 0f32;
2285 for gi in 0..gpr {
2286 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
2287 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2288 let b = vld1q_u8(t.add(2));
2289 let lo = vandq_u8(b, lomask);
2290 let hi = vshrq_n_u8::<4>(b);
2291 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
2292 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
2293 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
2294 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
2295 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
2296 asm!(
2297 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
2298 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
2299 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
2300 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
2301 options(pure, nomem, nostack),
2302 );
2303 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
2304 }
2305 acc
2306 }
2307}
2308
2309#[cfg(target_arch = "x86_64")]
2310#[target_feature(enable = "avx2")]
2311unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
2312 unsafe {
2314 use core::arch::x86_64::*;
2315 let lomask = _mm_set1_epi8(0x0F);
2316 let eight = _mm256_set1_epi8(8);
2317 let ones = _mm256_set1_epi16(1);
2318 let mut acc = 0f32;
2319 for gi in 0..gpr {
2320 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
2321 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2322 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
2323 let lo = _mm_and_si128(b, lomask);
2324 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
2325 let w = _mm256_sub_epi8(
2326 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
2327 eight,
2328 );
2329 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2330 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
2331 let d = _mm256_madd_epi16(p16, ones);
2332 let hi128 = _mm256_extracti128_si256::<1>(d);
2333 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
2334 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2335 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2336 acc += _mm_cvtsi128_si32(s32) as f32 * s;
2337 }
2338 acc
2339 }
2340}
2341
2342#[cfg(target_arch = "x86_64")]
2347#[target_feature(enable = "avx2")]
2348unsafe fn dot_q4t_row_1x4_avx2(
2349 bytes: &[u8],
2350 r: usize,
2351 gpr: usize,
2352 xs: [&[i8]; 4],
2353) -> [f32; 4] {
2354 unsafe {
2356 use core::arch::x86_64::*;
2357 let lomask = _mm_set1_epi8(0x0F);
2358 let eight = _mm256_set1_epi8(8);
2359 let ones = _mm256_set1_epi16(1);
2360 let mut acc = [0f32; 4];
2361 for gi in 0..gpr {
2362 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
2363 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2364 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
2365 let lo = _mm_and_si128(bb, lomask);
2366 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
2367 let w = _mm256_sub_epi8(
2368 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
2369 eight,
2370 );
2371 let aw = _mm256_abs_epi8(w);
2372 for (k, xq) in xs.iter().enumerate() {
2373 let x =
2374 _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2375 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
2376 let d = _mm256_madd_epi16(p16, ones);
2377 let hi128 = _mm256_extracti128_si256::<1>(d);
2378 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
2379 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2380 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2381 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
2382 }
2383 }
2384 acc
2385 }
2386}
2387
2388#[inline]
2390fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
2391 let gi = j / GROUP_SIZE;
2392 let k = j % GROUP_SIZE;
2393 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
2394 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2395 let byte = tile[2 + k / 2];
2396 let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
2397 ((nib as i32 - 8) as f32, s)
2398}
2399
2400#[inline]
2403fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
2404 let mut acc = 0f32;
2405 for gi in 0..gpr {
2406 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
2407 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2408 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
2409 let mut ga = 0f32;
2410 for (k, &b) in tile[2..].iter().enumerate() {
2411 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
2412 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
2413 }
2414 acc += ga * s;
2415 }
2416 acc
2417}
2418
2419fn q4t_matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
2421 debug_assert_eq!(out.len(), rows);
2422 let gpr = cols / GROUP_SIZE;
2423 let out_addr = SendMut(out.as_mut_ptr());
2424 if a8w8_enabled() {
2425 let act = split_act(x);
2426 let run = move |start: usize, end: usize| {
2427 for r in start..end {
2428 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
2429 for &(j, xv) in &act.outliers {
2430 let (w, s) = q4t_outlier(bytes, r, gpr, j);
2431 acc += w * s * xv;
2432 }
2433 unsafe { *out_addr.at(r) = acc };
2435 }
2436 };
2437 dispatch_rows(pool, rows, &run);
2438 return;
2439 }
2440 let run = move |start: usize, end: usize| {
2441 for r in start..end {
2442 unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
2444 }
2445 };
2446 dispatch_rows(pool, rows, &run);
2447}
2448
2449#[allow(clippy::too_many_arguments)]
2451fn q4t_matvec2(
2452 bytes: &[u8],
2453 x1: &[f32],
2454 x2: &[f32],
2455 rows: usize,
2456 cols: usize,
2457 o1: &mut [f32],
2458 o2: &mut [f32],
2459 pool: Option<&Pool>,
2460) {
2461 let gpr = cols / GROUP_SIZE;
2462 let p1 = SendMut(o1.as_mut_ptr());
2463 let p2 = SendMut(o2.as_mut_ptr());
2464 if a8w8_enabled() {
2465 let a1 = split_act(x1);
2466 let a2 = split_act(x2);
2467 let run = move |start: usize, end: usize| {
2468 for r in start..end {
2469 let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
2470 let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
2471 for &(j, xv) in &a1.outliers {
2472 let (w, s) = q4t_outlier(bytes, r, gpr, j);
2473 v1 += w * s * xv;
2474 }
2475 for &(j, xv) in &a2.outliers {
2476 let (w, s) = q4t_outlier(bytes, r, gpr, j);
2477 v2 += w * s * xv;
2478 }
2479 unsafe {
2481 *p1.at(r) = v1;
2482 *p2.at(r) = v2;
2483 }
2484 }
2485 };
2486 dispatch_rows(pool, rows, &run);
2487 return;
2488 }
2489 let run = move |start: usize, end: usize| {
2490 for r in start..end {
2491 unsafe {
2493 *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
2494 *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
2495 }
2496 }
2497 };
2498 dispatch_rows(pool, rows, &run);
2499}
2500
2501#[allow(clippy::too_many_arguments)]
2503fn q4t_matmat(
2504 bytes: &[u8],
2505 xs_all: &[f32],
2506 b: usize,
2507 rows: usize,
2508 cols: usize,
2509 out: &mut [f32],
2510 pool: Option<&Pool>,
2511) {
2512 debug_assert_eq!(out.len(), b * rows);
2513 let gpr = cols / GROUP_SIZE;
2514 let out_addr = SendMut(out.as_mut_ptr());
2515 if a8w8_enabled() {
2516 let acts: Vec<SplitAct> =
2517 (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
2518 let acts = &acts;
2519 #[cfg(target_arch = "x86_64")]
2520 let blocked_ok = avx2_enabled()
2521 && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true);
2522 #[cfg(not(target_arch = "x86_64"))]
2523 let blocked_ok = false;
2524 let run = move |start: usize, end: usize| {
2525 for r in start..end {
2526 let mut bi = 0usize;
2527 #[cfg(target_arch = "x86_64")]
2528 if blocked_ok {
2529 while bi + 4 <= acts.len() {
2530 let xs = [
2531 acts[bi].xq.as_slice(),
2532 acts[bi + 1].xq.as_slice(),
2533 acts[bi + 2].xq.as_slice(),
2534 acts[bi + 3].xq.as_slice(),
2535 ];
2536 let d = unsafe { dot_q4t_row_1x4_avx2(bytes, r, gpr, xs) };
2537 for k in 0..4 {
2538 let act = &acts[bi + k];
2539 let mut acc = d[k] * act.sx;
2540 for &(j, xv) in &act.outliers {
2541 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
2542 acc += w * sc * xv;
2543 }
2544 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
2546 }
2547 bi += 4;
2548 }
2549 }
2550 let _ = blocked_ok;
2551 while bi < acts.len() {
2552 let act = &acts[bi];
2553 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
2554 for &(j, xv) in &act.outliers {
2555 let (w, s) = q4t_outlier(bytes, r, gpr, j);
2556 acc += w * s * xv;
2557 }
2558 unsafe { *out_addr.at(bi * rows + r) = acc };
2560 bi += 1;
2561 }
2562 }
2563 };
2564 dispatch_rows(pool, rows, &run);
2565 return;
2566 }
2567 let run = move |start: usize, end: usize| {
2568 for r in start..end {
2569 for bi in 0..b {
2570 let x = &xs_all[bi * cols..(bi + 1) * cols];
2571 unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
2573 }
2574 }
2575 };
2576 dispatch_rows(pool, rows, &run);
2577}
2578
2579fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
2588 (0..gpr)
2589 .map(|gi| {
2590 xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
2591 .iter()
2592 .map(|&v| v as i32)
2593 .sum()
2594 })
2595 .collect()
2596}
2597
2598#[inline]
2602#[allow(unreachable_code)]
2603#[cfg(target_arch = "x86_64")]
2608#[target_feature(enable = "avx2")]
2609unsafe fn dot_q1_row_avx2(
2610 bytes: &[u8],
2611 r: usize,
2612 gpr: usize,
2613 xq: &[i8],
2614 gsum: &[i32],
2615) -> f32 {
2616 unsafe {
2618 use core::arch::x86_64::*;
2619 let expand = _mm256_setr_epi8(
2621 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1,
2622 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3,
2623 );
2624 let bitsel = _mm256_setr_epi8(
2625 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128,
2626 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128,
2627 );
2628 let ones8 = _mm256_set1_epi8(1);
2629 let ones16 = _mm256_set1_epi16(1);
2630 let mut acc = 0f32;
2631 for gi in 0..gpr {
2632 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
2633 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2634 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
2635 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
2636 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
2637 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2638 let sel = _mm256_and_si256(x, mask);
2639 let p16 = _mm256_maddubs_epi16(ones8, sel);
2641 let d32 = _mm256_madd_epi16(p16, ones16);
2642 let hi128 = _mm256_extracti128_si256::<1>(d32);
2643 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
2644 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2645 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2646 let msum = _mm_cvtsi128_si32(s32);
2647 let d = 2 * msum - gsum[gi];
2650 acc += d as f32 * s;
2651 }
2652 acc
2653 }
2654}
2655
2656#[cfg(target_arch = "x86_64")]
2659#[target_feature(enable = "avx2")]
2660unsafe fn dot_q1_row_1x4_avx2(
2661 bytes: &[u8],
2662 r: usize,
2663 gpr: usize,
2664 xs: [&[i8]; 4],
2665 gsums: [&[i32]; 4],
2666) -> [f32; 4] {
2667 unsafe {
2669 use core::arch::x86_64::*;
2670 let expand = _mm256_setr_epi8(
2671 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1,
2672 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3,
2673 );
2674 let bitsel = _mm256_setr_epi8(
2675 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128,
2676 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128,
2677 );
2678 let ones8 = _mm256_set1_epi8(1);
2679 let ones16 = _mm256_set1_epi16(1);
2680 let mut acc = [0f32; 4];
2681 for gi in 0..gpr {
2682 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
2683 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2684 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
2685 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
2686 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
2687 for (k, xq) in xs.iter().enumerate() {
2688 let x =
2689 _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2690 let sel = _mm256_and_si256(x, mask);
2691 let p16 = _mm256_maddubs_epi16(ones8, sel);
2692 let d32 = _mm256_madd_epi16(p16, ones16);
2693 let hi128 = _mm256_extracti128_si256::<1>(d32);
2694 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
2695 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2696 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2697 let msum = _mm_cvtsi128_si32(s32);
2698 let d = 2 * msum - gsums[k][gi];
2699 acc[k] += d as f32 * s;
2700 }
2701 }
2702 acc
2703 }
2704}
2705
2706fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
2707 #[cfg(target_arch = "aarch64")]
2708 unsafe {
2709 return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
2710 }
2711 #[cfg(target_arch = "x86_64")]
2712 if avx2_enabled() {
2713 unsafe {
2714 return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
2715 }
2716 }
2717 let _ = gsum;
2718 let mut acc = 0f32;
2719 for gi in 0..gpr {
2720 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
2721 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2722 let mut d = 0i32;
2723 for (j, &b) in tile[2..].iter().enumerate() {
2724 for k in 0..8 {
2725 let w = ((b >> k) & 1) as i32 * 2 - 1;
2726 d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
2727 }
2728 }
2729 acc += d as f32 * s;
2730 }
2731 acc
2732}
2733
2734#[cfg(target_arch = "aarch64")]
2743#[target_feature(enable = "neon,dotprod")]
2744unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
2745 unsafe {
2748 use core::arch::aarch64::*;
2749 use core::arch::asm;
2750 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
2751 let m = vld1q_u8(MASKS.as_ptr());
2752 macro_rules! tile_dot {
2754 ($t:expr, $x:expr) => {{
2755 let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
2756 let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
2757 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
2758 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
2759 let x0 = vld1q_s8($x);
2760 let x1 = vld1q_s8($x.add(16));
2761 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
2762 asm!(
2763 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
2764 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
2765 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
2766 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
2767 options(pure, nomem, nostack),
2768 );
2769 vaddq_s32(a0, a1)
2770 }};
2771 }
2772 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
2773 let xp = xq.as_ptr();
2774 let gp = gsum.as_ptr();
2775 let mut accv = vdupq_n_f32(0.0);
2776 let mut gi = 0;
2777 while gi + 4 <= gpr {
2778 let t0 = base.add(gi * Q1_TILE);
2779 let t1 = base.add((gi + 1) * Q1_TILE);
2780 let t2 = base.add((gi + 2) * Q1_TILE);
2781 let t3 = base.add((gi + 3) * Q1_TILE);
2782 let d0 = tile_dot!(t0, xp.add(gi * GROUP_SIZE));
2783 let d1 = tile_dot!(t1, xp.add((gi + 1) * GROUP_SIZE));
2784 let d2 = tile_dot!(t2, xp.add((gi + 2) * GROUP_SIZE));
2785 let d3 = tile_dot!(t3, xp.add((gi + 3) * GROUP_SIZE));
2786 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
2788 let g = vld1q_s32(gp.add(gi));
2789 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
2790 let sc = [
2791 f16_to_f32(u16::from_le_bytes([*t0, *t0.add(1)])),
2792 f16_to_f32(u16::from_le_bytes([*t1, *t1.add(1)])),
2793 f16_to_f32(u16::from_le_bytes([*t2, *t2.add(1)])),
2794 f16_to_f32(u16::from_le_bytes([*t3, *t3.add(1)])),
2795 ];
2796 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), vld1q_f32(sc.as_ptr()));
2797 gi += 4;
2798 }
2799 let mut acc = vaddvq_f32(accv);
2800 while gi < gpr {
2801 let t = base.add(gi * Q1_TILE);
2802 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2803 let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
2804 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
2805 gi += 1;
2806 }
2807 acc
2808 }
2809}
2810
2811#[inline]
2813fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
2814 let gi = j / GROUP_SIZE;
2815 let k = j % GROUP_SIZE;
2816 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
2817 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2818 let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
2819 ((bit as i32 * 2 - 1) as f32, s)
2820}
2821
2822#[inline]
2824fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
2825 let mut acc = 0f32;
2826 for gi in 0..gpr {
2827 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
2828 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2829 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
2830 let mut ga = 0f32;
2831 for (j, &b) in tile[2..].iter().enumerate() {
2832 for k in 0..8 {
2833 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
2834 }
2835 }
2836 acc += ga * s;
2837 }
2838 acc
2839}
2840
2841#[allow(clippy::too_many_arguments)]
2844fn q1_range_a8w8(
2845 bytes: &[u8],
2846 gpr: usize,
2847 act: &SplitAct,
2848 gsum: &[i32],
2849 out: SendMut,
2850 start: usize,
2851 end: usize,
2852) {
2853 for r in start..end {
2854 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
2855 for &(j, xv) in &act.outliers {
2856 let (w, s) = q1_outlier(bytes, r, gpr, j);
2857 acc += w * s * xv;
2858 }
2859 unsafe { *out.at(r) = acc };
2861 }
2862}
2863
2864fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
2866 for r in start..end {
2867 unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
2869 }
2870}
2871
2872fn q1_matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
2874 debug_assert_eq!(out.len(), rows);
2875 let gpr = cols / GROUP_SIZE;
2876 let out_addr = SendMut(out.as_mut_ptr());
2877 if a8w8_enabled() {
2878 let act = split_act(x);
2879 let gsum = q1_group_sums(&act.xq, gpr);
2880 let (act, gsum) = (&act, &gsum);
2881 let run =
2882 move |start: usize, end: usize| q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end);
2883 dispatch_rows(pool, rows, &run);
2884 return;
2885 }
2886 let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
2887 dispatch_rows(pool, rows, &run);
2888}
2889
2890#[allow(clippy::too_many_arguments)]
2892fn q1_matvec2(
2893 bytes: &[u8],
2894 x1: &[f32],
2895 x2: &[f32],
2896 rows: usize,
2897 cols: usize,
2898 o1: &mut [f32],
2899 o2: &mut [f32],
2900 pool: Option<&Pool>,
2901) {
2902 let gpr = cols / GROUP_SIZE;
2903 let p1 = SendMut(o1.as_mut_ptr());
2904 let p2 = SendMut(o2.as_mut_ptr());
2905 if a8w8_enabled() {
2906 let a1 = split_act(x1);
2907 let a2 = split_act(x2);
2908 let g1 = q1_group_sums(&a1.xq, gpr);
2909 let g2 = q1_group_sums(&a2.xq, gpr);
2910 let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
2911 let run = move |start: usize, end: usize| {
2912 for r in start..end {
2913 let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
2914 let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
2915 for &(j, xv) in &a1.outliers {
2916 let (w, s) = q1_outlier(bytes, r, gpr, j);
2917 v1 += w * s * xv;
2918 }
2919 for &(j, xv) in &a2.outliers {
2920 let (w, s) = q1_outlier(bytes, r, gpr, j);
2921 v2 += w * s * xv;
2922 }
2923 unsafe {
2925 *p1.at(r) = v1;
2926 *p2.at(r) = v2;
2927 }
2928 }
2929 };
2930 dispatch_rows(pool, rows, &run);
2931 return;
2932 }
2933 let run = move |start: usize, end: usize| {
2934 for r in start..end {
2935 unsafe {
2937 *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
2938 *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
2939 }
2940 }
2941 };
2942 dispatch_rows(pool, rows, &run);
2943}
2944
2945#[allow(clippy::too_many_arguments)]
2947fn q1_matmat(
2948 bytes: &[u8],
2949 xs_all: &[f32],
2950 b: usize,
2951 rows: usize,
2952 cols: usize,
2953 out: &mut [f32],
2954 pool: Option<&Pool>,
2955) {
2956 debug_assert_eq!(out.len(), b * rows);
2957 let gpr = cols / GROUP_SIZE;
2958 let out_addr = SendMut(out.as_mut_ptr());
2959 if a8w8_enabled() {
2960 let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
2961 .map(|bi| {
2962 let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
2963 let gsum = q1_group_sums(&act.xq, gpr);
2964 (act, gsum)
2965 })
2966 .collect();
2967 let acts = &acts;
2968 #[cfg(target_arch = "x86_64")]
2969 let blocked_ok = avx2_enabled()
2970 && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true);
2971 let run = move |start: usize, end: usize| {
2972 for r in start..end {
2973 let mut bi = 0usize;
2974 #[cfg(target_arch = "x86_64")]
2977 if blocked_ok {
2978 while bi + 4 <= acts.len() {
2979 let xs = [
2980 acts[bi].0.xq.as_slice(),
2981 acts[bi + 1].0.xq.as_slice(),
2982 acts[bi + 2].0.xq.as_slice(),
2983 acts[bi + 3].0.xq.as_slice(),
2984 ];
2985 let gs = [
2986 acts[bi].1.as_slice(),
2987 acts[bi + 1].1.as_slice(),
2988 acts[bi + 2].1.as_slice(),
2989 acts[bi + 3].1.as_slice(),
2990 ];
2991 let d = unsafe { dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs) };
2992 for k in 0..4 {
2993 let (act, _) = &acts[bi + k];
2994 let mut acc = d[k] * act.sx;
2995 for &(j, xv) in &act.outliers {
2996 let (w, sc) = q1_outlier(bytes, r, gpr, j);
2997 acc += w * sc * xv;
2998 }
2999 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
3001 }
3002 bi += 4;
3003 }
3004 }
3005 while bi < acts.len() {
3006 let (act, gsum) = &acts[bi];
3007 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
3008 for &(j, xv) in &act.outliers {
3009 let (w, s) = q1_outlier(bytes, r, gpr, j);
3010 acc += w * s * xv;
3011 }
3012 unsafe { *out_addr.at(bi * rows + r) = acc };
3014 bi += 1;
3015 }
3016 }
3017 };
3018 dispatch_rows(pool, rows, &run);
3019 return;
3020 }
3021 let run = move |start: usize, end: usize| {
3022 for r in start..end {
3023 for bi in 0..b {
3024 let x = &xs_all[bi * cols..(bi + 1) * cols];
3025 unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
3027 }
3028 }
3029 };
3030 dispatch_rows(pool, rows, &run);
3031}
3032
3033fn q4matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
3039 debug_assert_eq!(out.len(), rows);
3040 let (packed, scales) = q4_split(bytes, rows, cols);
3041 let gpr = cols / GROUP_SIZE;
3042 let out_addr = SendMut(out.as_mut_ptr());
3043
3044 if a8w8_enabled() {
3045 let act = split_act(x);
3046 let run = move |start: usize, end: usize| {
3047 q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
3048 };
3049 dispatch_rows(pool, rows, &run);
3050 return;
3051 }
3052
3053 let run = move |start: usize, end: usize| {
3054 q4_range_f32(packed, scales, gpr, x, out_addr, start, end)
3055 };
3056 dispatch_rows(pool, rows, &run);
3057}
3058
3059#[inline]
3062#[allow(unreachable_code)]
3063#[cfg(target_arch = "x86_64")]
3068#[target_feature(enable = "avx2")]
3069unsafe fn dot_q4b_row_1x4_avx2(
3070 buf: &[u8],
3071 scales: &[u8],
3072 g0: usize,
3073 gpr: usize,
3074 xs: [&[i8]; 4],
3075) -> [f32; 4] {
3076 unsafe {
3078 use core::arch::x86_64::*;
3079 let ones = _mm256_set1_epi16(1);
3080 let mut acc = [0f32; 4];
3081 for gi in 0..gpr {
3082 let s = f16_to_f32(u16::from_le_bytes([
3083 scales[(g0 + gi) * 2],
3084 scales[(g0 + gi) * 2 + 1],
3085 ]));
3086 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3087 let aw = _mm256_abs_epi8(w);
3088 for (k, xq) in xs.iter().enumerate() {
3089 let x =
3090 _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3091 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
3092 let d = _mm256_madd_epi16(p16, ones);
3093 let hi128 = _mm256_extracti128_si256::<1>(d);
3094 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3095 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3096 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3097 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
3098 }
3099 }
3100 acc
3101 }
3102}
3103
3104#[cfg(target_arch = "x86_64")]
3110#[target_feature(enable = "avx2")]
3111unsafe fn dot_q4b_row_1x4_sx_avx2(
3112 buf: &[u8],
3113 scales: &[u8],
3114 g0: usize,
3115 gpr: usize,
3116 xs: [&[i8]; 4],
3117 sxs: [f32; 4],
3118) -> [f32; 4] {
3119 unsafe {
3121 use core::arch::x86_64::*;
3122 let ones = _mm256_set1_epi16(1);
3123 let mut acc = [0f32; 4];
3124 for gi in 0..gpr {
3125 let s = f16_to_f32(u16::from_le_bytes([
3126 scales[(g0 + gi) * 2],
3127 scales[(g0 + gi) * 2 + 1],
3128 ]));
3129 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3130 let aw = _mm256_abs_epi8(w);
3131 for (k, xq) in xs.iter().enumerate() {
3132 let x =
3133 _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3134 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
3135 let d = _mm256_madd_epi16(p16, ones);
3136 let hi128 = _mm256_extracti128_si256::<1>(d);
3137 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3138 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3139 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3140 acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
3141 }
3142 }
3143 acc
3144 }
3145}
3146
3147fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
3148 #[cfg(target_arch = "aarch64")]
3149 unsafe {
3150 return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
3151 }
3152 #[cfg(target_arch = "x86_64")]
3153 unsafe {
3154 return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
3155 }
3156 let mut acc = 0f32;
3157 for gi in 0..gpr {
3158 let g = g0 + gi;
3159 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
3160 let mut d = 0i32;
3161 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
3162 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
3163 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
3164 }
3165 acc += d as f32 * s;
3166 }
3167 acc
3168}
3169
3170#[inline]
3172#[allow(unreachable_code)]
3173fn dot_q4_row_i8_2(
3174 packed: &[u8],
3175 scales: &[u8],
3176 g0: usize,
3177 gpr: usize,
3178 xq1: &[i8],
3179 xq2: &[i8],
3180) -> (f32, f32) {
3181 #[cfg(target_arch = "aarch64")]
3182 unsafe {
3183 return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
3184 }
3185 #[cfg(target_arch = "x86_64")]
3186 unsafe {
3187 return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
3188 }
3189 (
3190 dot_q4_row_i8(packed, scales, g0, gpr, xq1),
3191 dot_q4_row_i8(packed, scales, g0, gpr, xq2),
3192 )
3193}
3194
3195#[allow(clippy::too_many_arguments)]
3198fn q4_range_a8w8(
3199 packed: &[u8],
3200 scales: &[u8],
3201 gpr: usize,
3202 cols: usize,
3203 act: &SplitAct,
3204 out: SendMut,
3205 start: usize,
3206 end: usize,
3207) {
3208 for r in start..end {
3209 let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
3210 for &(j, xv) in &act.outliers {
3212 let flat = r * cols + j;
3213 let byte = packed[flat / 2];
3214 let nib = if flat & 1 == 0 { byte & 0x0F } else { byte >> 4 };
3215 let s = f16_to_f32(u16::from_le_bytes([
3216 scales[(flat / GROUP_SIZE) * 2],
3217 scales[(flat / GROUP_SIZE) * 2 + 1],
3218 ]));
3219 acc += ((nib as i32 - 8) as f32) * s * xv;
3220 }
3221 unsafe { *out.at(r) = acc };
3223 }
3224}
3225
3226#[allow(clippy::too_many_arguments)]
3229fn q4_range2_a8w8(
3230 packed: &[u8],
3231 scales: &[u8],
3232 gpr: usize,
3233 cols: usize,
3234 a1: &SplitAct,
3235 a2: &SplitAct,
3236 p1: SendMut,
3237 p2: SendMut,
3238 start: usize,
3239 end: usize,
3240) {
3241 for r in start..end {
3242 let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
3243 let mut acc1 = s1 * a1.sx;
3244 let mut acc2 = s2 * a2.sx;
3245 let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
3247 for &(j, xv) in outliers {
3248 let flat = r * cols + j;
3249 let byte = packed[flat / 2];
3250 let nib = if flat & 1 == 0 { byte & 0x0F } else { byte >> 4 };
3251 let s = f16_to_f32(u16::from_le_bytes([
3252 scales[(flat / GROUP_SIZE) * 2],
3253 scales[(flat / GROUP_SIZE) * 2 + 1],
3254 ]));
3255 *acc += ((nib as i32 - 8) as f32) * s * xv;
3256 }
3257 };
3258 fix(&a1.outliers, &mut acc1);
3259 fix(&a2.outliers, &mut acc2);
3260 unsafe {
3262 *p1.at(r) = acc1;
3263 *p2.at(r) = acc2;
3264 }
3265 }
3266}
3267
3268fn q4_range_f32(
3270 packed: &[u8],
3271 scales: &[u8],
3272 gpr: usize,
3273 x: &[f32],
3274 out: SendMut,
3275 start: usize,
3276 end: usize,
3277) {
3278 for r in start..end {
3279 let mut acc = 0f32;
3280 for gi in 0..gpr {
3281 let g = r * gpr + gi;
3282 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
3283 let pk = &packed[g * 16..(g + 1) * 16];
3284 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
3285 let mut ga = 0f32;
3286 for (k, &b) in pk.iter().enumerate() {
3287 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
3288 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
3289 }
3290 acc += ga * s;
3291 }
3292 unsafe { *out.at(r) = acc };
3294 }
3295}
3296
3297#[allow(clippy::too_many_arguments)]
3301fn q4matvec2(
3302 bytes: &[u8],
3303 x1: &[f32],
3304 x2: &[f32],
3305 rows: usize,
3306 cols: usize,
3307 o1: &mut [f32],
3308 o2: &mut [f32],
3309 pool: Option<&Pool>,
3310) {
3311 debug_assert_eq!(o1.len(), rows);
3312 debug_assert_eq!(o2.len(), rows);
3313 let (packed, scales) = q4_split(bytes, rows, cols);
3314 let gpr = cols / GROUP_SIZE;
3315
3316 if a8w8_enabled() {
3317 let a1 = split_act(x1);
3318 let a2 = split_act(x2);
3319 let p1 = SendMut(o1.as_mut_ptr());
3320 let p2 = SendMut(o2.as_mut_ptr());
3321 let run = move |start: usize, end: usize| {
3322 q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
3323 };
3324 dispatch_rows(pool, rows, &run);
3325 return;
3326 }
3327
3328 let p1 = SendMut(o1.as_mut_ptr());
3329 let p2 = SendMut(o2.as_mut_ptr());
3330 let run = move |start: usize, end: usize| {
3331 q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
3332 };
3333 dispatch_rows(pool, rows, &run);
3334}
3335
3336#[allow(clippy::too_many_arguments)]
3338fn q4_range2_f32(
3339 packed: &[u8],
3340 scales: &[u8],
3341 gpr: usize,
3342 x1: &[f32],
3343 x2: &[f32],
3344 p1: SendMut,
3345 p2: SendMut,
3346 start: usize,
3347 end: usize,
3348) {
3349 for r in start..end {
3350 let (mut acc1, mut acc2) = (0f32, 0f32);
3351 for gi in 0..gpr {
3352 let g = r * gpr + gi;
3353 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
3354 let pk = &packed[g * 16..(g + 1) * 16];
3355 let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
3356 let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
3357 let (mut g1, mut g2) = (0f32, 0f32);
3358 for (k, &b) in pk.iter().enumerate() {
3359 let wl = (b & 0x0F) as f32 - 8.0;
3360 let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
3361 g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
3362 g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
3363 }
3364 acc1 += g1 * s;
3365 acc2 += g2 * s;
3366 }
3367 unsafe {
3369 *p1.at(r) = acc1;
3370 *p2.at(r) = acc2;
3371 }
3372 }
3373}
3374
3375thread_local! {
3376 static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
3379 static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3380}
3381
3382#[allow(clippy::too_many_arguments)]
3388fn q4matmat(
3389 bytes: &[u8],
3390 xs_all: &[f32],
3391 b: usize,
3392 rows: usize,
3393 cols: usize,
3394 out: &mut [f32],
3395 pool: Option<&Pool>,
3396) {
3397 debug_assert_eq!(xs_all.len(), b * cols);
3398 debug_assert_eq!(out.len(), b * rows);
3399 let (packed, scales) = q4_split(bytes, rows, cols);
3400 let gpr = cols / GROUP_SIZE;
3401 let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
3402
3403 if a8w8_enabled() {
3404 let acts: Vec<SplitAct> =
3405 (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
3406 let acts = &acts;
3407 let out_addr = SendMut(out.as_mut_ptr());
3408 let run = move |start: usize, end: usize| {
3409 ROW_I8.with(|rb| {
3410 let mut buf = rb.borrow_mut();
3411 buf.resize(cols, 0);
3412 for r in start..end {
3413 for gi in 0..gpr {
3417 let g = r * gpr + gi;
3418 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
3419 buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
3420 buf[gi * GROUP_SIZE + k * 2 + 1] =
3421 (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
3422 }
3423 }
3424 let mut bi = 0usize;
3425 #[cfg(target_arch = "x86_64")]
3426 if avx2_enabled()
3427 && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true)
3428 {
3429 while bi + 4 <= acts.len() {
3430 let xs = [
3431 acts[bi].xq.as_slice(),
3432 acts[bi + 1].xq.as_slice(),
3433 acts[bi + 2].xq.as_slice(),
3434 acts[bi + 3].xq.as_slice(),
3435 ];
3436 let d = unsafe {
3437 dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
3438 };
3439 for k in 0..4 {
3440 let act = &acts[bi + k];
3441 let mut acc = d[k] * act.sx;
3442 for &(j, xv) in &act.outliers {
3443 acc += (buf[j] as i8) as f32
3444 * gscale((r * cols + j) / GROUP_SIZE)
3445 * xv;
3446 }
3447 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
3449 }
3450 bi += 4;
3451 }
3452 }
3453 while bi < acts.len() {
3454 let act = &acts[bi];
3455 let mut acc = 0f32;
3456 for gi in 0..gpr {
3457 let d = dot_i8_i8(
3458 &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
3459 &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
3460 );
3461 acc += d as f32 * gscale(r * gpr + gi);
3462 }
3463 acc *= act.sx;
3464 for &(j, xv) in &act.outliers {
3466 acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
3467 }
3468 unsafe { *out_addr.at(bi * rows + r) = acc };
3470 bi += 1;
3471 }
3472 }
3473 })
3474 };
3475 dispatch_rows(pool, rows, &run);
3476 return;
3477 }
3478
3479 let out_addr = SendMut(out.as_mut_ptr());
3480 let run = move |start: usize, end: usize| {
3481 ROW_F32.with(|rb| {
3482 let mut buf = rb.borrow_mut();
3483 buf.resize(cols, 0.0);
3484 for r in start..end {
3485 for gi in 0..gpr {
3488 let g = r * gpr + gi;
3489 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
3490 buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
3491 buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
3492 }
3493 }
3494 for bi in 0..b {
3495 let x = &xs_all[bi * cols..(bi + 1) * cols];
3496 let mut acc = 0f32;
3497 for gi in 0..gpr {
3498 let mut ga = 0f32;
3499 for k in 0..GROUP_SIZE / 2 {
3504 let e = gi * GROUP_SIZE + k * 2;
3505 ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
3506 }
3507 acc += ga * gscale(r * gpr + gi);
3508 }
3509 unsafe { *out_addr.at(bi * rows + r) = acc };
3511 }
3512 }
3513 })
3514 };
3515 dispatch_rows(pool, rows, &run);
3516}
3517
3518#[allow(clippy::too_many_arguments)]
3523fn vbitmatmat(
3524 bytes: &[u8],
3525 offsets: &[usize],
3526 xs_all: &[f32],
3527 b: usize,
3528 rows: usize,
3529 cols: usize,
3530 out: &mut [f32],
3531 pool: Option<&Pool>,
3532) {
3533 debug_assert_eq!(xs_all.len(), b * cols);
3534 debug_assert_eq!(out.len(), b * rows);
3535 debug_assert_eq!(offsets.len(), rows + 1);
3536 let ng = cols / GROUP_SIZE;
3537 let bits = &bytes[..rows];
3538 let sc_off = rows;
3539 let gscale = |r: usize, g: usize| {
3540 let so = (r * ng + g) * 2;
3541 f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]))
3542 };
3543
3544 let decode_f32 = |r: usize, dst: &mut [f32]| {
3546 let bw = bits[r] as usize;
3547 let l = ((1i32 << (bw - 1)) - 1) as f32;
3548 let data = &bytes[offsets[r]..offsets[r + 1]];
3549 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3550 for d in dst.iter_mut() {
3551 while nbits < bw {
3552 acc = (acc << 8) | data[idx] as u64;
3553 idx += 1;
3554 nbits += 8;
3555 }
3556 let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
3557 nbits -= bw;
3558 *d = u - l;
3559 }
3560 };
3561
3562 if a8w8_enabled() {
3563 let acts: Vec<SplitAct> =
3564 (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
3565 let acts = &acts;
3566 let out_addr = SendMut(out.as_mut_ptr());
3567 let run = move |start: usize, end: usize| {
3568 for r in start..end {
3569 let bw = bits[r] as usize;
3570 if bw == 8 {
3571 ROW_F32.with(|rb| {
3574 let mut buf = rb.borrow_mut();
3575 buf.resize(cols, 0.0);
3576 decode_f32(r, &mut buf);
3577 for bi in 0..b {
3578 let x = &xs_all[bi * cols..(bi + 1) * cols];
3579 let mut dot = 0f32;
3580 for g in 0..ng {
3581 let mut gd = 0f32;
3582 for k in 0..GROUP_SIZE {
3583 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
3584 }
3585 dot += gd * gscale(r, g);
3586 }
3587 unsafe { *out_addr.at(bi * rows + r) = dot };
3589 }
3590 });
3591 continue;
3592 }
3593 let l = (1i32 << (bw - 1)) - 1;
3594 let data = &bytes[offsets[r]..offsets[r + 1]];
3595 ROW_I8.with(|rb| {
3596 let mut buf = rb.borrow_mut();
3597 buf.resize(cols, 0);
3598 #[inline(always)]
3599 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3600 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3601 let u = unpack8::<B>(&data[blk * B..]);
3602 for k in 0..8 {
3603 chunk[k] = (u[k] - l) as i8 as u8;
3604 }
3605 }
3606 }
3607 match bw {
3608 3 => fill::<3>(data, l, &mut buf),
3609 4 => vbit_fill4(data, &mut buf),
3610 5 => fill::<5>(data, l, &mut buf),
3611 6 => fill::<6>(data, l, &mut buf),
3612 _ => unreachable!("vbit bit-width {bw} (validated at load)"),
3613 }
3614 let mut bi = 0usize;
3615 #[cfg(target_arch = "x86_64")]
3619 if avx2_enabled()
3620 && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true)
3621 {
3622 while bi + 4 <= acts.len() {
3623 let xs = [
3624 acts[bi].xq.as_slice(),
3625 acts[bi + 1].xq.as_slice(),
3626 acts[bi + 2].xq.as_slice(),
3627 acts[bi + 3].xq.as_slice(),
3628 ];
3629 let sxs = [
3630 acts[bi].sx,
3631 acts[bi + 1].sx,
3632 acts[bi + 2].sx,
3633 acts[bi + 3].sx,
3634 ];
3635 let d = unsafe {
3636 dot_q4b_row_1x4_sx_avx2(&buf, &bytes[sc_off..], r * ng, ng, xs, sxs)
3637 };
3638 for k in 0..4 {
3639 let act = &acts[bi + k];
3640 let mut dot = d[k];
3641 for &(j, xv) in &act.outliers {
3642 dot +=
3643 (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
3644 }
3645 unsafe { *out_addr.at((bi + k) * rows + r) = dot };
3647 }
3648 bi += 4;
3649 }
3650 }
3651 while bi < acts.len() {
3652 let act = &acts[bi];
3653 let mut dot = 0f32;
3654 for g in 0..ng {
3655 let d = dot_i8_i8(
3656 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3657 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3658 ) as f32
3659 * act.sx;
3660 dot += d * gscale(r, g);
3661 }
3662 for &(j, xv) in &act.outliers {
3663 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
3664 }
3665 unsafe { *out_addr.at(bi * rows + r) = dot };
3667 bi += 1;
3668 }
3669 });
3670 }
3671 };
3672 dispatch_rows(pool, rows, &run);
3673 return;
3674 }
3675
3676 let out_addr = SendMut(out.as_mut_ptr());
3677 let run = move |start: usize, end: usize| {
3678 ROW_F32.with(|rb| {
3679 let mut buf = rb.borrow_mut();
3680 buf.resize(cols, 0.0);
3681 for r in start..end {
3682 decode_f32(r, &mut buf);
3683 for bi in 0..b {
3684 let x = &xs_all[bi * cols..(bi + 1) * cols];
3685 let mut dot = 0f32;
3686 for g in 0..ng {
3687 let mut gd = 0f32;
3688 for k in 0..GROUP_SIZE {
3689 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
3690 }
3691 dot += gd * gscale(r, g);
3692 }
3693 unsafe { *out_addr.at(bi * rows + r) = dot };
3695 }
3696 }
3697 })
3698 };
3699 dispatch_rows(pool, rows, &run);
3700}
3701
3702pub(crate) fn gpu_batch_job<'a>(
3706 t: &'a QTensor,
3707 x: &[f32],
3708) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
3709 match t {
3710 QTensor::Mapped {
3711 model,
3712 idx,
3713 dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
3714 rows,
3715 cols,
3716 row_scale,
3717 col_field,
3718 ..
3719 } => Some((
3720 model.clone(),
3721 crate::gpu::BatchJob {
3722 idx: *idx,
3723 rows: *rows,
3724 cols: *cols,
3725 row_scale,
3726 xs: prescale(x, col_field, *dt).into_owned(),
3727 q1: false,
3728 },
3729 )),
3730 QTensor::Mapped {
3732 model,
3733 idx,
3734 dtype: TensorDtype::Q1,
3735 rows,
3736 cols,
3737 ..
3738 } => Some((
3739 model.clone(),
3740 crate::gpu::BatchJob {
3741 idx: *idx,
3742 rows: *rows,
3743 cols: *cols,
3744 row_scale: &[],
3745 xs: x.to_vec(),
3746 q1: true,
3747 },
3748 )),
3749 _ => None,
3750 }
3751}
3752
3753pub(crate) fn prescale<'a>(
3757 x: &'a [f32],
3758 col_field: &[f32],
3759 dtype: TensorDtype,
3760) -> std::borrow::Cow<'a, [f32]> {
3761 if dtype == TensorDtype::Q8_2f {
3762 x.iter().zip(col_field).map(|(a, c)| a * c).collect()
3763 } else {
3764 std::borrow::Cow::Borrowed(x)
3765 }
3766}
3767
3768#[cfg(target_arch = "x86_64")]
3773pub(crate) fn avx2_enabled() -> bool {
3774 use std::sync::OnceLock;
3775 static ON: OnceLock<bool> = OnceLock::new();
3776 *ON.get_or_init(|| {
3777 std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
3778 && std::arch::is_x86_feature_detected!("avx2")
3779 && std::arch::is_x86_feature_detected!("fma")
3780 })
3781}
3782
3783#[cfg(target_arch = "x86_64")]
3788fn avx2_a8w8_enabled() -> bool {
3789 use std::sync::OnceLock;
3790 static ON: OnceLock<bool> = OnceLock::new();
3791 *ON.get_or_init(|| {
3792 avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
3793 })
3794}
3795
3796#[inline]
3800pub(crate) fn a8w8_enabled() -> bool {
3801 #[cfg(target_arch = "aarch64")]
3802 {
3803 sdot_enabled()
3804 }
3805 #[cfg(target_arch = "x86_64")]
3806 {
3807 avx2_a8w8_enabled()
3808 }
3809 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
3810 {
3811 false
3812 }
3813}
3814
3815#[inline]
3818#[allow(unreachable_code)]
3819fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
3820 #[cfg(target_arch = "aarch64")]
3821 unsafe {
3822 return dot_i8_sdot(w, xq);
3823 }
3824 #[cfg(target_arch = "x86_64")]
3825 unsafe {
3826 if avx512vnni_enabled() {
3827 return dot_i8_i8_vnni(w, xq);
3828 }
3829 return dot_i8_i8_avx2(w, xq);
3830 }
3831 w.iter().zip(xq).map(|(&a, &b)| (a as i8) as i32 * b as i32).sum()
3832}
3833
3834#[cfg(target_arch = "x86_64")]
3838fn avx512vnni_enabled() -> bool {
3839 use std::sync::OnceLock;
3840 static ON: OnceLock<bool> = OnceLock::new();
3841 *ON.get_or_init(|| {
3842 std::env::var("CMF_AVX512").map(|v| v != "0").unwrap_or(true)
3843 && std::arch::is_x86_feature_detected!("avx512f")
3844 && std::arch::is_x86_feature_detected!("avx512bw")
3845 && std::arch::is_x86_feature_detected!("avx512vl")
3846 && std::arch::is_x86_feature_detected!("avx512vnni")
3847 })
3848}
3849
3850#[cfg(target_arch = "x86_64")]
3855#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
3856unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
3857 unsafe {
3859 use core::arch::x86_64::*;
3860 let n = w.len();
3861 let mut j = 0usize;
3862 let mut total: i32;
3863 {
3868 #[inline(always)]
3869 unsafe fn step(
3870 w: *const u8,
3871 x: *const i8,
3872 acc: core::arch::x86_64::__m512i,
3873 ) -> core::arch::x86_64::__m512i {
3874 unsafe {
3875 use core::arch::x86_64::*;
3876 let wv = _mm512_loadu_si512(w as *const _);
3877 let xv = _mm512_loadu_si512(x as *const _);
3878 let aw = _mm512_abs_epi8(wv);
3879 let neg = _mm512_movepi8_mask(wv);
3880 let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
3881 _mm512_dpbusd_epi32(acc, aw, sx)
3882 }
3883 }
3884 let (mut a0, mut a1, mut a2, mut a3) = (
3885 _mm512_setzero_si512(),
3886 _mm512_setzero_si512(),
3887 _mm512_setzero_si512(),
3888 _mm512_setzero_si512(),
3889 );
3890 while j + 256 <= n {
3891 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
3892 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
3893 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
3894 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
3895 j += 256;
3896 }
3897 while j + 64 <= n {
3898 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
3899 j += 64;
3900 }
3901 let s01 = _mm512_add_epi32(a0, a1);
3902 let s23 = _mm512_add_epi32(a2, a3);
3903 total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
3904 }
3905 if j + 32 <= n {
3907 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
3908 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
3909 let d = _mm256_dpbusd_epi32(
3910 _mm256_setzero_si256(),
3911 _mm256_abs_epi8(wv),
3912 _mm256_sign_epi8(xv, wv),
3913 );
3914 let hi128 = _mm256_extracti128_si256::<1>(d);
3915 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3916 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3917 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3918 total += _mm_cvtsi128_si32(s32);
3919 j += 32;
3920 }
3921 while j < n {
3922 total += (w[j] as i8) as i32 * xq[j] as i32;
3923 j += 1;
3924 }
3925 total
3926 }
3927}
3928
3929#[cfg(target_arch = "x86_64")]
3931#[target_feature(enable = "avx2,fma")]
3932unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
3933 unsafe {
3935 use core::arch::x86_64::*;
3936 let n = x.len();
3937 let wp = w.as_ptr();
3938 let xp = x.as_ptr();
3939 let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
3940 let mut j = 0usize;
3941 while j + 16 <= n {
3942 let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
3943 let lo = _mm256_cvtepi8_epi32(wb);
3944 let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
3945 a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
3946 a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
3947 j += 16;
3948 }
3949 let acc = _mm256_add_ps(a0, a1);
3950 let hi128 = _mm256_extractf128_ps::<1>(acc);
3951 let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
3952 let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
3953 let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
3954 let mut sum = _mm_cvtss_f32(s32);
3955 while j < n {
3956 sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
3957 j += 1;
3958 }
3959 sum
3960 }
3961}
3962
3963#[cfg(target_arch = "x86_64")]
3968#[target_feature(enable = "avx2")]
3969unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
3970 unsafe {
3972 use core::arch::x86_64::*;
3973 let n = w.len();
3974 let ones = _mm256_set1_epi16(1);
3975 let mut acc = _mm256_setzero_si256();
3976 let mut j = 0usize;
3977 while j + 32 <= n {
3978 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
3979 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
3980 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
3981 acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
3982 j += 32;
3983 }
3984 let hi128 = _mm256_extracti128_si256::<1>(acc);
3985 let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
3986 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3987 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3988 let mut s = _mm_cvtsi128_si32(s32);
3989 while j < n {
3990 s += (w[j] as i8) as i32 * xq[j] as i32;
3991 j += 1;
3992 }
3993 s
3994 }
3995}
3996
3997#[cfg(target_arch = "aarch64")]
4001#[target_feature(enable = "neon,dotprod")]
4002unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
4003 unsafe {
4005 use core::arch::aarch64::*;
4006 use core::arch::asm;
4007 let n = w0.len();
4008 let w0p = w0.as_ptr() as *const i8;
4009 let w1p = w1.as_ptr() as *const i8;
4010 let mut acc = [[vdupq_n_s32(0); 4]; 2];
4011 let mut i = 0usize;
4012 while i + 16 <= n {
4013 let wv0 = vld1q_s8(w0p.add(i));
4014 let wv1 = vld1q_s8(w1p.add(i));
4015 for (k, x) in xs.iter().enumerate() {
4016 let xv = vld1q_s8(x.as_ptr().add(i));
4017 let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
4018 asm!(
4019 "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
4020 "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
4021 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4022 w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
4023 options(pure, nomem, nostack),
4024 );
4025 acc[0][k] = a0;
4026 acc[1][k] = a1;
4027 }
4028 i += 16;
4029 }
4030 let mut out = [[0i32; 4]; 2];
4031 for r in 0..2 {
4032 for k in 0..4 {
4033 out[r][k] = vaddvq_s32(acc[r][k]);
4034 }
4035 }
4036 if i < n {
4037 for (k, x) in xs.iter().enumerate() {
4038 for j in i..n {
4039 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
4040 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
4041 }
4042 }
4043 }
4044 out
4045 }
4046}
4047
4048#[cfg(target_arch = "x86_64")]
4054#[target_feature(enable = "avx2")]
4055unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
4056 unsafe {
4058 use core::arch::x86_64::*;
4059 let n = w0.len();
4060 let ones = _mm256_set1_epi16(1);
4061 let mut acc = [[_mm256_setzero_si256(); 4]; 2];
4062 let mut j = 0usize;
4063 while j + 32 <= n {
4064 let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
4065 let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
4066 let aw0 = _mm256_abs_epi8(wv0);
4067 let aw1 = _mm256_abs_epi8(wv1);
4068 for (k, x) in xs.iter().enumerate() {
4069 let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
4070 let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
4071 acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
4072 let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
4073 acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
4074 }
4075 j += 32;
4076 }
4077 let mut out = [[0i32; 4]; 2];
4078 for r in 0..2 {
4079 for k in 0..4 {
4080 let a = acc[r][k];
4081 let hi128 = _mm256_extracti128_si256::<1>(a);
4082 let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
4083 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4084 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4085 out[r][k] = _mm_cvtsi128_si32(s32);
4086 }
4087 }
4088 if j < n {
4089 for (k, x) in xs.iter().enumerate() {
4090 for i in j..n {
4091 out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
4092 out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
4093 }
4094 }
4095 }
4096 out
4097 }
4098}
4099
4100#[cfg(target_arch = "x86_64")]
4105#[inline]
4106fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
4107 let dot = if avx512vnni_enabled() && row.len() >= 64 {
4108 (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
4109 } else {
4110 unsafe { dot_i8_i8_avx2(row, &act.xq) }
4111 };
4112 let mut acc = dot as f32 * act.sx;
4113 for &(j, xv) in &act.outliers {
4114 acc += (row[j] as i8) as f32 * xv;
4115 }
4116 acc
4117}
4118
4119#[cfg(target_arch = "x86_64")]
4123#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4124unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
4125 unsafe {
4127 use core::arch::x86_64::*;
4128 let n = w.len();
4129 let flip = _mm512_set1_epi8(-128); #[inline(always)]
4131 unsafe fn step(
4132 w: *const u8,
4133 x: *const i8,
4134 flip: core::arch::x86_64::__m512i,
4135 acc: core::arch::x86_64::__m512i,
4136 ) -> core::arch::x86_64::__m512i {
4137 unsafe {
4138 use core::arch::x86_64::*;
4139 let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
4140 _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
4141 }
4142 }
4143 let (mut a0, mut a1, mut a2, mut a3) = (
4144 _mm512_setzero_si512(),
4145 _mm512_setzero_si512(),
4146 _mm512_setzero_si512(),
4147 _mm512_setzero_si512(),
4148 );
4149 let mut j = 0usize;
4150 while j + 256 <= n {
4151 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
4152 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
4153 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
4154 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
4155 j += 256;
4156 }
4157 while j + 64 <= n {
4158 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
4159 j += 64;
4160 }
4161 let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
4162 _mm512_add_epi32(a0, a1),
4163 _mm512_add_epi32(a2, a3),
4164 ));
4165 while j < n {
4167 total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
4168 j += 1;
4169 }
4170 total
4171 }
4172}
4173
4174#[cfg(target_arch = "x86_64")]
4180#[target_feature(enable = "avx2")]
4181unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
4182 unsafe {
4185 use core::arch::x86_64::*;
4186 let lomask = _mm_set1_epi8(0x0F);
4187 let eight = _mm256_set1_epi8(8);
4188 let ones = _mm256_set1_epi16(1);
4189 let mut acc = 0f32;
4190 for gi in 0..gpr {
4191 let g = g0 + gi;
4192 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4193 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
4194 let lo = _mm_and_si128(b, lomask);
4195 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4196 let w = _mm256_sub_epi8(
4197 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4198 eight,
4199 );
4200 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4201 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4202 let d = _mm256_madd_epi16(p16, ones);
4203 let hi128 = _mm256_extracti128_si256::<1>(d);
4204 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4205 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4206 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4207 acc += _mm_cvtsi128_si32(s32) as f32 * s;
4208 }
4209 acc
4210 }
4211}
4212
4213#[cfg(target_arch = "x86_64")]
4216#[target_feature(enable = "avx2")]
4217unsafe fn dot_q4_row_avx2_2(
4218 packed: &[u8],
4219 scales: &[u8],
4220 g0: usize,
4221 gpr: usize,
4222 xq1: &[i8],
4223 xq2: &[i8],
4224) -> (f32, f32) {
4225 unsafe {
4227 use core::arch::x86_64::*;
4228 let lomask = _mm_set1_epi8(0x0F);
4229 let eight = _mm256_set1_epi8(8);
4230 let ones = _mm256_set1_epi16(1);
4231 let (mut acc1, mut acc2) = (0f32, 0f32);
4232 #[inline(always)]
4233 unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
4234 unsafe {
4235 use core::arch::x86_64::*;
4236 let hi128 = _mm256_extracti128_si256::<1>(d);
4237 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4238 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4239 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4240 _mm_cvtsi128_si32(s32)
4241 }
4242 }
4243 for gi in 0..gpr {
4244 let g = g0 + gi;
4245 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4246 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
4247 let lo = _mm_and_si128(b, lomask);
4248 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4249 let w = _mm256_sub_epi8(
4250 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4251 eight,
4252 );
4253 let aw = _mm256_abs_epi8(w);
4254 let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4255 let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4256 let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
4257 let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
4258 acc1 += hsum(d1) as f32 * s;
4259 acc2 += hsum(d2) as f32 * s;
4260 }
4261 (acc1, acc2)
4262 }
4263}
4264
4265#[cfg(target_arch = "x86_64")]
4267fn q8_range_avx2(
4268 q: &[u8],
4269 row_scale: &[f32],
4270 act: &SplitAct,
4271 cols: usize,
4272 out_addr: SendMut,
4273 start: usize,
4274 end: usize,
4275) {
4276 for o in start..end {
4277 let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
4278 unsafe { *out_addr.at(o) = v };
4280 }
4281}
4282
4283#[cfg(target_arch = "x86_64")]
4285#[allow(clippy::too_many_arguments)]
4286fn q8_range2_avx2(
4287 q: &[u8],
4288 row_scale: &[f32],
4289 a1: &SplitAct,
4290 a2: &SplitAct,
4291 cols: usize,
4292 p1: SendMut,
4293 p2: SendMut,
4294 start: usize,
4295 end: usize,
4296) {
4297 for o in start..end {
4298 let row = &q[o * cols..(o + 1) * cols];
4299 unsafe {
4301 *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
4302 *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
4303 }
4304 }
4305}
4306
4307#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
4313fn sdot_enabled() -> bool {
4314 use std::sync::OnceLock;
4315 static ON: OnceLock<bool> = OnceLock::new();
4316 *ON.get_or_init(|| {
4317 let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
4318 #[cfg(target_arch = "aarch64")]
4319 {
4320 want && std::arch::is_aarch64_feature_detected!("dotprod")
4321 }
4322 #[cfg(not(target_arch = "aarch64"))]
4323 {
4324 let _ = want;
4325 false
4326 }
4327 })
4328}
4329
4330struct SplitAct {
4335 xq: Vec<i8>,
4336 sx: f32,
4337 outliers: Vec<(usize, f32)>,
4338 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
4341 xsum: i32,
4342}
4343
4344thread_local! {
4345 static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
4348 const { std::cell::RefCell::new(Vec::new()) };
4349}
4350
4351impl Drop for SplitAct {
4352 fn drop(&mut self) {
4353 let buf = std::mem::take(&mut self.xq);
4354 if buf.capacity() > 0 {
4355 XQ_FREE.with(|f| {
4356 let mut f = f.borrow_mut();
4357 if f.len() < 16 {
4358 f.push(buf);
4359 }
4360 });
4361 }
4362 }
4363}
4364
4365fn split_act(x: &[f32]) -> SplitAct {
4366 let n = x.len();
4367 let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
4368 let thr = 8.0 * rms;
4369 let mut outliers: Vec<(usize, f32)> = Vec::new();
4373 let mut amax = 0f32;
4374 for (j, &v) in x.iter().enumerate() {
4375 let a = v.abs();
4376 if a > thr {
4377 outliers.push((j, v));
4378 } else if a > amax {
4379 amax = a;
4380 }
4381 }
4382 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
4383 let inv = 1.0 / sx;
4384 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
4385 xq.clear();
4386 xq.reserve(n);
4387 if outliers.is_empty() {
4388 xq.extend(x.iter().map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8));
4389 } else {
4390 xq.extend(x.iter().map(|&v| {
4392 if v.abs() > thr {
4393 0
4394 } else {
4395 (v * inv).round().clamp(-127.0, 127.0) as i8
4396 }
4397 }));
4398 }
4399 let xsum = xq.iter().map(|&v| v as i32).sum();
4400 SplitAct { xq, sx, outliers, xsum }
4401}
4402
4403#[cfg(target_arch = "aarch64")]
4406#[target_feature(enable = "neon,dotprod")]
4407unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
4408 unsafe {
4410 use core::arch::aarch64::*;
4411 use core::arch::asm;
4412 let wp = w.as_ptr() as *const i8;
4413 let n = w.len();
4414 let (mut a0, mut a1, mut a2, mut a3) =
4415 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
4416 let mut i = 0;
4417 while i + 64 <= n {
4418 let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
4419 let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
4420 let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
4421 let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
4422 asm!(
4423 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
4424 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
4425 "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
4426 "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
4427 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
4428 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
4429 w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
4430 options(pure, nomem, nostack),
4431 );
4432 i += 64;
4433 }
4434 while i + 16 <= n {
4435 let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
4436 asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
4437 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
4438 i += 16;
4439 }
4440 let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
4441 while i < n {
4442 s += (*wp.add(i)) as i32 * xq[i] as i32;
4443 i += 1;
4444 }
4445 s
4446}
4447}
4448
4449#[cfg(target_arch = "aarch64")]
4453#[target_feature(enable = "neon,dotprod")]
4454unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
4455 unsafe {
4457 use core::arch::aarch64::*;
4458 use core::arch::asm;
4459 let n = xq.len();
4460 let px = xq.as_ptr();
4461 let (p0, p1, p2, p3) = (
4462 w0.as_ptr() as *const i8,
4463 w1.as_ptr() as *const i8,
4464 w2.as_ptr() as *const i8,
4465 w3.as_ptr() as *const i8,
4466 );
4467 let (mut a0, mut a1, mut a2, mut a3) =
4468 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
4469 let mut i = 0;
4470 while i + 16 <= n {
4471 let x = vld1q_s8(px.add(i));
4472 let v0 = vld1q_s8(p0.add(i));
4473 let v1 = vld1q_s8(p1.add(i));
4474 let v2 = vld1q_s8(p2.add(i));
4475 let v3 = vld1q_s8(p3.add(i));
4476 asm!(
4477 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
4478 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
4479 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
4480 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
4481 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
4482 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
4483 options(pure, nomem, nostack),
4484 );
4485 i += 16;
4486 }
4487 let mut r = [vaddvq_s32(a0), vaddvq_s32(a1), vaddvq_s32(a2), vaddvq_s32(a3)];
4488 while i < n {
4489 let xi = *px.add(i) as i32;
4490 r[0] += (*p0.add(i)) as i32 * xi;
4491 r[1] += (*p1.add(i)) as i32 * xi;
4492 r[2] += (*p2.add(i)) as i32 * xi;
4493 r[3] += (*p3.add(i)) as i32 * xi;
4494 i += 1;
4495 }
4496 r
4497}
4498}
4499
4500#[cfg(target_arch = "aarch64")]
4507#[target_feature(enable = "neon,dotprod")]
4508unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
4509 unsafe {
4512 use core::arch::aarch64::*;
4513 use core::arch::asm;
4514 let n = xq.len();
4515 let px = xq.as_ptr();
4516 let pg = g.as_ptr() as *const i8;
4517 let (mut a0, mut a1, mut a2, mut a3) =
4518 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
4519 let mut i = 0;
4520 while i + 16 <= n {
4521 let x = vld1q_s8(px.add(i));
4522 let base = pg.add(4 * i);
4523 let v0 = vld1q_s8(base);
4524 let v1 = vld1q_s8(base.add(16));
4525 let v2 = vld1q_s8(base.add(32));
4526 let v3 = vld1q_s8(base.add(48));
4527 asm!(
4528 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
4529 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
4530 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
4531 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
4532 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
4533 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
4534 options(pure, nomem, nostack),
4535 );
4536 i += 16;
4537 }
4538 [vaddvq_s32(a0), vaddvq_s32(a1), vaddvq_s32(a2), vaddvq_s32(a3)]
4539 }
4540}
4541
4542#[cfg(target_arch = "aarch64")]
4548fn q8_range_sdot(
4549 q: &[u8],
4550 rep: &[u8],
4551 row_scale: &[f32],
4552 act: &SplitAct,
4553 cols: usize,
4554 out_addr: SendMut,
4555 start: usize,
4556 end: usize,
4557) {
4558 let mut o = start;
4559 if !rep.is_empty() {
4562 while o < end && o % 4 != 0 {
4563 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
4564 unsafe { *out_addr.at(o) = v };
4565 o += 1;
4566 }
4567 }
4568 while o + 4 <= end {
4569 let r = if rep.is_empty() {
4570 unsafe {
4571 dot_i8_sdot_4rows(
4572 &q[o * cols..(o + 1) * cols],
4573 &q[(o + 1) * cols..(o + 2) * cols],
4574 &q[(o + 2) * cols..(o + 3) * cols],
4575 &q[(o + 3) * cols..(o + 4) * cols],
4576 &act.xq,
4577 )
4578 }
4579 } else {
4580 unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
4581 };
4582 for k in 0..4 {
4583 let mut acc = r[k] as f32 * act.sx;
4584 for &(j, xv) in &act.outliers {
4585 acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
4586 }
4587 unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
4589 }
4590 o += 4;
4591 }
4592 while o < end {
4593 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
4594 unsafe { *out_addr.at(o) = v };
4595 o += 1;
4596 }
4597}
4598
4599#[cfg(target_arch = "aarch64")]
4602#[allow(clippy::too_many_arguments)]
4603fn q8_range2_sdot(
4604 q: &[u8],
4605 row_scale: &[f32],
4606 a1: &SplitAct,
4607 a2: &SplitAct,
4608 cols: usize,
4609 p1: SendMut,
4610 p2: SendMut,
4611 start: usize,
4612 end: usize,
4613) {
4614 for o in start..end {
4615 let row = &q[o * cols..(o + 1) * cols];
4616 unsafe {
4618 *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
4619 *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
4620 }
4621 }
4622}
4623
4624#[allow(clippy::too_many_arguments)]
4626fn q8_range2_f32(
4627 q: &[u8],
4628 row_scale: &[f32],
4629 x1: &[f32],
4630 x2: &[f32],
4631 cols: usize,
4632 p1: SendMut,
4633 p2: SendMut,
4634 start: usize,
4635 end: usize,
4636) {
4637 for o in start..end {
4638 let row = &q[o * cols..(o + 1) * cols];
4639 unsafe {
4641 *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
4642 *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
4643 }
4644 }
4645}
4646
4647fn q8_range_f32(
4649 q: &[u8],
4650 row_scale: &[f32],
4651 xs: &[f32],
4652 cols: usize,
4653 out_addr: SendMut,
4654 start: usize,
4655 end: usize,
4656) {
4657 for o in start..end {
4658 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
4659 unsafe { *out_addr.at(o) = v };
4661 }
4662}
4663
4664#[cfg(target_arch = "aarch64")]
4667#[inline]
4668fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
4669 let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
4670 for &(j, xv) in &act.outliers {
4671 acc += (row[j] as i8) as f32 * xv;
4672 }
4673 acc
4674}
4675
4676#[cfg(target_arch = "aarch64")]
4684#[target_feature(enable = "neon,dotprod")]
4685unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
4686 unsafe {
4689 use core::arch::aarch64::*;
4690 use core::arch::asm;
4691 let lomask = vdupq_n_u8(0x0F);
4692 let eight = vdupq_n_s8(8);
4693 let mut acc = 0f32;
4694 for gi in 0..gpr {
4695 let g = g0 + gi;
4696 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4697 let b = vld1q_u8(packed.as_ptr().add(g * 16));
4698 let lo = vandq_u8(b, lomask);
4699 let hi = vshrq_n_u8::<4>(b);
4700 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4701 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4702 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4703 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4704 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4705 asm!(
4706 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4707 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4708 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4709 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4710 options(pure, nomem, nostack),
4711 );
4712 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4713 }
4714 acc
4715 }
4716}
4717
4718#[cfg(target_arch = "aarch64")]
4723#[target_feature(enable = "neon,dotprod")]
4724unsafe fn dot_q4_row_sdot2(
4725 packed: &[u8],
4726 scales: &[u8],
4727 g0: usize,
4728 gpr: usize,
4729 xq1: &[i8],
4730 xq2: &[i8],
4731) -> (f32, f32) {
4732 unsafe {
4735 use core::arch::aarch64::*;
4736 use core::arch::asm;
4737 let lomask = vdupq_n_u8(0x0F);
4738 let eight = vdupq_n_s8(8);
4739 let (mut acc1, mut acc2) = (0f32, 0f32);
4740 for gi in 0..gpr {
4741 let g = g0 + gi;
4742 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4743 let b = vld1q_u8(packed.as_ptr().add(g * 16));
4744 let lo = vandq_u8(b, lomask);
4745 let hi = vshrq_n_u8::<4>(b);
4746 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4747 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4748 let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
4749 let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
4750 let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
4751 let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
4752 let (mut a0, mut a1, mut b0, mut b1) =
4753 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
4754 asm!(
4755 "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
4756 "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
4757 "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
4758 "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
4759 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4760 b0 = inout(vreg) b0, b1 = inout(vreg) b1,
4761 e0 = in(vreg) e0, e1 = in(vreg) e1,
4762 x10 = in(vreg) x10, x11 = in(vreg) x11,
4763 x20 = in(vreg) x20, x21 = in(vreg) x21,
4764 options(pure, nomem, nostack),
4765 );
4766 acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4767 acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
4768 }
4769 (acc1, acc2)
4770 }
4771}
4772
4773#[inline]
4778pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
4779 #[cfg(target_arch = "aarch64")]
4780 unsafe {
4781 return axpy_i8_f32_neon(acc, row, w);
4782 }
4783 #[cfg(target_arch = "x86_64")]
4784 if avx2_enabled() {
4785 return unsafe { axpy_i8_f32_avx2(acc, row, w) };
4786 }
4787 #[allow(unreachable_code)]
4788 {
4789 for (a, &b) in acc.iter_mut().zip(row) {
4790 *a += w * b as f32;
4791 }
4792 }
4793}
4794
4795#[cfg(target_arch = "x86_64")]
4797#[target_feature(enable = "avx2,fma")]
4798unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
4799 unsafe {
4801 use core::arch::x86_64::*;
4802 let n = acc.len().min(row.len());
4803 let ap = acc.as_mut_ptr();
4804 let rp = row.as_ptr();
4805 let wv = _mm256_set1_ps(w);
4806 let mut j = 0usize;
4807 while j + 16 <= n {
4808 let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
4809 let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
4810 let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
4811 let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
4812 let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
4813 _mm256_storeu_ps(ap.add(j), v0);
4814 _mm256_storeu_ps(ap.add(j + 8), v1);
4815 j += 16;
4816 }
4817 while j < n {
4818 *ap.add(j) += w * (*rp.add(j)) as f32;
4819 j += 1;
4820 }
4821 }
4822}
4823
4824#[cfg(target_arch = "aarch64")]
4825#[target_feature(enable = "neon")]
4826unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
4827 unsafe {
4829 use core::arch::aarch64::*;
4830 let n = acc.len().min(row.len());
4831 let ap = acc.as_mut_ptr();
4832 let rp = row.as_ptr();
4833 let wv = vdupq_n_f32(w);
4834 let mut j = 0usize;
4835 while j + 16 <= n {
4836 let rb = vld1q_s8(rp.add(j));
4837 let lo = vmovl_s8(vget_low_s8(rb));
4838 let hi = vmovl_s8(vget_high_s8(rb));
4839 for (off, half) in [(0, lo), (8, hi)] {
4840 let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
4841 let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
4842 let o = j + off;
4843 vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
4844 vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
4845 }
4846 j += 16;
4847 }
4848 while j < n {
4849 *ap.add(j) += w * (*rp.add(j)) as f32;
4850 j += 1;
4851 }
4852}
4853}
4854
4855#[inline]
4858pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
4859 #[cfg(target_arch = "aarch64")]
4860 unsafe {
4861 return dot_i8_f32_neon(w, x);
4862 }
4863 #[cfg(target_arch = "x86_64")]
4864 if avx2_enabled() {
4865 return unsafe { dot_i8_f32_avx2(w, x) };
4866 }
4867 #[allow(unreachable_code)]
4868 {
4869 let mut sum = 0.0f32;
4870 for (j, &b) in w.iter().enumerate() {
4871 sum += (b as i8) as f32 * x[j];
4872 }
4873 sum
4874 }
4875}
4876
4877#[inline]
4881fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
4882 #[cfg(target_arch = "aarch64")]
4883 unsafe {
4884 return dot_i8_col_f32_neon(w, x, col);
4885 }
4886 #[allow(unreachable_code)]
4887 {
4888 let mut sum = 0.0f32;
4889 for (j, &b) in w.iter().enumerate() {
4890 sum += (b as i8) as f32 * x[j] * col[j];
4891 }
4892 sum
4893 }
4894}
4895
4896#[cfg(target_arch = "aarch64")]
4897#[target_feature(enable = "neon")]
4898unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
4899 unsafe {
4901 use core::arch::aarch64::*;
4902 let n = x.len();
4903 let wp = w.as_ptr() as *const i8;
4904 let xp = x.as_ptr();
4905 let cp = col.as_ptr();
4906 let (mut a0, mut a1, mut a2, mut a3) =
4907 (vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4908 let mut j = 0usize;
4909 while j + 16 <= n {
4910 let wb = vld1q_s8(wp.add(j));
4911 let lo = vmovl_s8(vget_low_s8(wb));
4912 let hi = vmovl_s8(vget_high_s8(wb));
4913 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
4914 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
4915 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
4916 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
4917 a0 = vfmaq_f32(a0, w0, vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))));
4918 a1 = vfmaq_f32(a1, w1, vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))));
4919 a2 = vfmaq_f32(a2, w2, vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))));
4920 a3 = vfmaq_f32(a3, w3, vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))));
4921 j += 16;
4922 }
4923 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
4924 while j < n {
4925 sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
4926 j += 1;
4927 }
4928 sum
4929}
4930}
4931
4932#[cfg(target_arch = "aarch64")]
4933#[target_feature(enable = "neon")]
4934unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
4935 unsafe {
4937 use core::arch::aarch64::*;
4938 let n = x.len();
4939 let wp = w.as_ptr() as *const i8;
4940 let xp = x.as_ptr();
4941 let (mut a0, mut a1, mut a2, mut a3) =
4942 (vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4943 let mut j = 0usize;
4944 while j + 16 <= n {
4945 let wb = vld1q_s8(wp.add(j));
4946 let lo = vmovl_s8(vget_low_s8(wb));
4947 let hi = vmovl_s8(vget_high_s8(wb));
4948 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
4949 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
4950 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
4951 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
4952 a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
4953 a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
4954 a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
4955 a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
4956 j += 16;
4957 }
4958 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
4959 while j < n {
4960 sum += (*wp.add(j)) as f32 * *xp.add(j);
4961 j += 1;
4962 }
4963 sum
4964}
4965}
4966
4967#[allow(clippy::too_many_arguments)]
4968fn qmatvec(
4969 q: &[u8],
4970 rep: &[u8],
4971 row_scale: &[f32],
4972 xs: &[f32],
4973 rows: usize,
4974 cols: usize,
4975 out: &mut [f32],
4976 pool: Option<&Pool>,
4977) {
4978 debug_assert_eq!(out.len(), rows);
4979 #[cfg(not(target_arch = "aarch64"))]
4980 let _ = rep;
4981
4982 #[cfg(target_arch = "aarch64")]
4983 if sdot_enabled() {
4984 let act = split_act(xs);
4985 let out_addr = SendMut(out.as_mut_ptr());
4986 let run_range = |start: usize, end: usize| {
4987 q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
4988 };
4989 match pool {
4990 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
4991 _ => run_range(0, rows),
4992 }
4993 return;
4994 }
4995 #[cfg(target_arch = "x86_64")]
4998 if avx2_a8w8_enabled() {
4999 let act = split_act(xs);
5000 let out_addr = SendMut(out.as_mut_ptr());
5001 let run_range =
5002 |start: usize, end: usize| q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end);
5003 match pool {
5004 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
5005 _ => run_range(0, rows),
5006 }
5007 return;
5008 }
5009
5010 let row_dot = |o: usize| -> f32 { dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o] };
5011 match pool {
5012 Some(pool) if rows >= 256 => {
5013 let out_addr = SendMut(out.as_mut_ptr());
5014 pool.run(&move |widx, n| {
5015 let chunk = rows.div_ceil(n);
5016 let start = widx * chunk;
5017 let end = (start + chunk).min(rows);
5018 for o in start..end {
5019 unsafe { *out_addr.at(o) = row_dot(o) };
5021 }
5022 });
5023 }
5024 _ => {
5025 for (o, dst) in out.iter_mut().enumerate() {
5026 *dst = row_dot(o);
5027 }
5028 }
5029 }
5030}
5031
5032#[allow(clippy::too_many_arguments)]
5033fn qmatvec2(
5034 q: &[u8],
5035 row_scale: &[f32],
5036 x1: &[f32],
5037 x2: &[f32],
5038 rows: usize,
5039 cols: usize,
5040 o1: &mut [f32],
5041 o2: &mut [f32],
5042 pool: Option<&Pool>,
5043) {
5044 #[cfg(target_arch = "aarch64")]
5045 if sdot_enabled() {
5046 let a1s = split_act(x1);
5047 let a2s = split_act(x2);
5048 let p1 = SendMut(o1.as_mut_ptr());
5049 let p2 = SendMut(o2.as_mut_ptr());
5050 let run_range = |start: usize, end: usize| {
5051 q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
5052 };
5053 match pool {
5054 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
5055 _ => run_range(0, rows),
5056 }
5057 return;
5058 }
5059 #[cfg(target_arch = "x86_64")]
5060 if avx2_a8w8_enabled() {
5061 let a1s = split_act(x1);
5062 let a2s = split_act(x2);
5063 let p1 = SendMut(o1.as_mut_ptr());
5064 let p2 = SendMut(o2.as_mut_ptr());
5065 let run_range = |start: usize, end: usize| {
5066 q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
5067 };
5068 match pool {
5069 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
5070 _ => run_range(0, rows),
5071 }
5072 return;
5073 }
5074
5075 let row_dots = |o: usize| -> (f32, f32) {
5076 let row = &q[o * cols..(o + 1) * cols];
5077 (dot_i8_f32(row, x1) * row_scale[o], dot_i8_f32(row, x2) * row_scale[o])
5078 };
5079 match pool {
5080 Some(pool) if rows >= 256 => {
5081 let p1 = SendMut(o1.as_mut_ptr());
5082 let p2 = SendMut(o2.as_mut_ptr());
5083 pool.run(&move |widx, n| {
5084 let chunk = rows.div_ceil(n);
5085 let start = widx * chunk;
5086 let end = (start + chunk).min(rows);
5087 for o in start..end {
5088 let (s1, s2) = row_dots(o);
5089 unsafe {
5091 *p1.at(o) = s1;
5092 *p2.at(o) = s2;
5093 }
5094 }
5095 });
5096 }
5097 _ => {
5098 for o in 0..rows {
5099 let (s1, s2) = row_dots(o);
5100 o1[o] = s1;
5101 o2[o] = s2;
5102 }
5103 }
5104 }
5105}
5106
5107#[derive(Clone, Copy)]
5108struct SendMut(*mut f32);
5109unsafe impl Send for SendMut {}
5110unsafe impl Sync for SendMut {}
5111
5112impl SendMut {
5113 #[inline]
5114 fn at(self, i: usize) -> *mut f32 {
5115 unsafe { self.0.add(i) }
5116 }
5117}
5118
5119#[cfg(test)]
5120mod tests {
5121 use super::*;
5122
5123 #[test]
5124 fn f32_matvec_matches_matvec_rows_bitexact() {
5125 let (rows, cols) = (300, 40);
5126 let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
5127 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
5128 let qt = QTensor::from_f32(w.clone(), rows, cols);
5129
5130 let mut a = vec![0.0f32; rows];
5131 matvec_rows(None, &w, &x, &mut a);
5132 let mut b = vec![0.0f32; rows];
5133 qt.matvec(&x, &mut b, None);
5134 assert_eq!(a, b);
5135 }
5136
5137 #[test]
5138 fn sdot_kernel_exact_on_grid() {
5139 eprintln!("sdot_enabled = {}", sdot_enabled());
5144 let (rows, cols) = (9, 80); let w: Vec<u8> = (0..rows * cols)
5146 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
5147 .collect();
5148 let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
5149 let x: Vec<f32> = (0..cols)
5150 .map(|i| match i % 3 {
5151 0 => 1.0,
5152 1 => -1.0,
5153 _ => 0.0,
5154 })
5155 .collect();
5156 let mut a = vec![0.0f32; rows];
5157 qmatvec(&w, &[], &scales, &x, rows, cols, &mut a, None);
5158 for o in 0..rows {
5159 let mut acc = 0.0f32;
5160 for j in 0..cols {
5161 acc += (w[o * cols + j] as i8) as f32 * x[j];
5162 }
5163 let expect = acc * scales[o];
5164 assert!(
5165 (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
5166 "row {o}: {} vs {expect}",
5167 a[o]
5168 );
5169 }
5170 }
5171
5172 #[test]
5173 fn q1_kernels_match_exact_reference() {
5174 let (rows, cols) = (7, 96);
5176 let gpr = cols / GROUP_SIZE;
5177 let mut bytes = Vec::new();
5178 for t in 0..rows * gpr {
5179 let s = 0.01 + (t % 13) as f32 * 0.003;
5180 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
5181 for j in 0..4 {
5182 bytes.push(((t * 31 + j * 97) % 251) as u8);
5183 }
5184 }
5185 let x: Vec<f32> = (0..cols)
5187 .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
5188 .collect();
5189 let mut w = vec![0.0f32; rows * cols];
5191 cortiq_core::quant::dequant_q1(&bytes, &mut w);
5192 let mut expect = vec![0.0f32; rows];
5193 for o in 0..rows {
5194 expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
5195 }
5196 let mut got = vec![0.0f32; rows];
5197 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
5198 for o in 0..rows {
5199 assert!(
5200 (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
5201 "row {o}: {} vs {}",
5202 got[o],
5203 expect[o]
5204 );
5205 }
5206 let x2: Vec<f32> = x.iter().map(|v| -v).collect();
5208 let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
5209 q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
5210 assert_eq!(a1, got);
5211 let mut xs = x.clone();
5212 xs.extend_from_slice(&x2);
5213 let mut mm = vec![0.0f32; 2 * rows];
5214 q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
5215 assert_eq!(&mm[..rows], got.as_slice());
5216 assert_eq!(&mm[rows..], a2.as_slice());
5217 }
5218
5219 #[test]
5220 fn repack_is_bit_identical() {
5221 let (rows, cols) = (267, 96); let w: Vec<u8> = (0..rows * cols)
5227 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
5228 .collect();
5229 let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
5230 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
5231 let rep = q8_repack_layout(&w, rows, cols);
5232 for g in 0..rows / 4 {
5234 for c in 0..cols / 16 {
5235 for lane in 0..4 {
5236 assert_eq!(
5237 &rep[g * 4 * cols + c * 64 + lane * 16..g * 4 * cols + c * 64 + lane * 16 + 16],
5238 &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
5239 );
5240 }
5241 }
5242 }
5243 let mut a = vec![0.0f32; rows];
5244 qmatvec(&w, &[], &scales, &x, rows, cols, &mut a, None);
5245 let mut b = vec![0.0f32; rows];
5246 qmatvec(&w, &rep, &scales, &x, rows, cols, &mut b, None);
5247 assert_eq!(a, b, "full-range repack output diverged");
5248
5249 #[cfg(target_arch = "aarch64")]
5250 if sdot_enabled() {
5251 let act = split_act(&x);
5253 let mut c1 = vec![0.0f32; rows];
5254 let mut c2 = vec![0.0f32; rows];
5255 q8_range_sdot(&w, &[], &scales, &act, cols, SendMut(c1.as_mut_ptr()), 3, rows - 2);
5256 q8_range_sdot(&w, &rep, &scales, &act, cols, SendMut(c2.as_mut_ptr()), 3, rows - 2);
5257 assert_eq!(c1, c2, "unaligned-range repack output diverged");
5258 }
5259 }
5260
5261 #[test]
5262 fn sdot_a8w8_noise_is_bounded() {
5263 let (rows, cols) = (16, 512);
5267 let w: Vec<u8> = (0..rows * cols)
5268 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
5269 .collect();
5270 let scales = vec![0.01f32; rows];
5271 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
5272 let mut a = vec![0.0f32; rows];
5273 qmatvec(&w, &[], &scales, &x, rows, cols, &mut a, None);
5274 let (mut num, mut den) = (0f64, 0f64);
5275 for o in 0..rows {
5276 let mut acc = 0.0f32;
5277 for j in 0..cols {
5278 acc += (w[o * cols + j] as i8) as f32 * x[j];
5279 }
5280 let expect = acc * scales[o];
5281 num += ((a[o] - expect) as f64).powi(2);
5282 den += (expect as f64).powi(2);
5283 }
5284 let rel = (num / den.max(1e-12)).sqrt();
5285 assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
5286 }
5287
5288 #[test]
5289 fn i8_dot_neon_matches_scalar() {
5290 let n = 100;
5291 let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
5292 let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
5293 let mut scalar = 0.0f32;
5294 for j in 0..n {
5295 scalar += (w[j] as i8) as f32 * x[j];
5296 }
5297 let fast = dot_i8_f32(&w, &x);
5298 assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
5299 }
5300
5301 #[test]
5303 fn vbitmatvec_matches_full_dequant() {
5304 let (rows, cols) = (6, 64);
5305 let ng = cols / GROUP_SIZE;
5306 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
5308 let mut bytes = bits.clone();
5309 for g in 0..rows * ng {
5310 let s = 0.02 + 0.001 * g as f32;
5311 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
5312 }
5313 for r in 0..rows {
5314 let b = bits[r] as usize;
5315 let (mut acc, mut nb) = (0u64, 0usize);
5316 let mut rowbytes = Vec::new();
5317 for i in 0..cols {
5318 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
5319 acc = (acc << b) | v;
5320 nb += b;
5321 while nb >= 8 {
5322 nb -= 8;
5323 rowbytes.push(((acc >> nb) & 0xFF) as u8);
5324 }
5325 }
5326 if nb > 0 {
5327 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
5328 }
5329 bytes.extend_from_slice(&rowbytes);
5330 }
5331 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
5332
5333 let mut reference = vec![0f32; rows * cols];
5334 cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
5335 let mut expect = vec![0f32; rows];
5336 for r in 0..rows {
5337 expect[r] = reference[r * cols..(r + 1) * cols]
5338 .iter()
5339 .zip(&x)
5340 .map(|(w, xv)| w * xv)
5341 .sum();
5342 }
5343 let mut got = vec![0f32; rows];
5344 let offsets = vbit_row_offsets(&bytes, rows, cols);
5345 vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
5346 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
5350 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
5351 for r in 0..rows {
5352 assert!(
5353 (got[r] - expect[r]).abs() < tol * scale,
5354 "row {r}: {} vs {}",
5355 got[r],
5356 expect[r]
5357 );
5358 }
5359 }
5360
5361 #[test]
5366 #[cfg(target_arch = "x86_64")]
5367 fn vbit_matmat_blocked_matches_per_row() {
5368 let (rows, cols, b) = (64usize, 128usize, 9usize);
5369 let ng = cols / GROUP_SIZE;
5370 let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
5371 let mut bytes = bits.clone();
5372 for g in 0..rows * ng {
5373 let sc = 0.02 + 0.0005 * g as f32;
5374 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
5375 }
5376 for r in 0..rows {
5377 let bw = bits[r] as usize;
5378 let (mut acc, mut nb) = (0u64, 0usize);
5379 let mut rowbytes = Vec::new();
5380 for i in 0..cols {
5381 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
5382 acc = (acc << bw) | v;
5383 nb += bw;
5384 while nb >= 8 {
5385 nb -= 8;
5386 rowbytes.push(((acc >> nb) & 0xFF) as u8);
5387 }
5388 }
5389 if nb > 0 {
5390 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
5391 }
5392 bytes.extend_from_slice(&rowbytes);
5393 }
5394 let x: Vec<f32> =
5395 (0..b * cols).map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5).collect();
5396 let offsets = vbit_row_offsets(&bytes, rows, cols);
5397 let mut y_a = vec![0f32; b * rows];
5398 let mut y_b = vec![0f32; b * rows];
5399 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
5400 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
5401 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
5402 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
5403 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
5404 let max_d =
5405 y_a.iter().zip(&y_b).map(|(p, q)| (p - q).abs()).fold(0.0f32, f32::max);
5406 assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
5407 }
5408
5409 #[test]
5410 fn q4matvec_matches_full_dequant() {
5411 let (rows, cols) = (8, 64);
5412 let groups = rows * cols / GROUP_SIZE;
5413 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
5415 for i in 0..groups * 16 {
5416 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
5417 }
5418 for g in 0..groups {
5419 let s = 0.01 + 0.003 * g as f32;
5420 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
5421 }
5422 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
5423
5424 let mut reference = vec![0.0f32; rows * cols];
5425 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
5426 let mut expect = vec![0.0f32; rows];
5427 for r in 0..rows {
5428 expect[r] = reference[r * cols..(r + 1) * cols]
5429 .iter()
5430 .zip(&x)
5431 .map(|(w, xv)| w * xv)
5432 .sum();
5433 }
5434
5435 let mut got = vec![0.0f32; rows];
5436 q4matvec(&bytes, &x, rows, cols, &mut got, None);
5437 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
5441 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
5442 for r in 0..rows {
5443 assert!(
5444 (got[r] - expect[r]).abs() < tol * scale,
5445 "row {r}: {} vs {}",
5446 got[r],
5447 expect[r]
5448 );
5449 }
5450 }
5451
5452 #[test]
5455 fn vbitmatvec2_equals_two_singles() {
5456 let (rows, cols) = (6, 64);
5457 let ng = cols / GROUP_SIZE;
5458 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
5459 let mut bytes = bits.clone();
5460 for g in 0..rows * ng {
5461 let s = 0.02 + 0.001 * g as f32;
5462 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
5463 }
5464 for r in 0..rows {
5465 let b = bits[r] as usize;
5466 let (mut acc, mut nb) = (0u64, 0usize);
5467 let mut rowbytes = Vec::new();
5468 for i in 0..cols {
5469 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
5470 acc = (acc << b) | v;
5471 nb += b;
5472 while nb >= 8 {
5473 nb -= 8;
5474 rowbytes.push(((acc >> nb) & 0xFF) as u8);
5475 }
5476 }
5477 if nb > 0 {
5478 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
5479 }
5480 bytes.extend_from_slice(&rowbytes);
5481 }
5482 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
5483 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
5484 let offsets = vbit_row_offsets(&bytes, rows, cols);
5485
5486 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
5487 vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
5488 vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
5489 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
5490 vbitmatvec2(&bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
5491 assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
5492 assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
5493 }
5494
5495 #[test]
5497 fn q4matvec2_equals_two_singles() {
5498 let (rows, cols) = (8, 128);
5499 let groups = rows * cols / GROUP_SIZE;
5500 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
5501 for i in 0..groups * 16 {
5502 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
5503 }
5504 for g in 0..groups {
5505 let s = 0.01 + 0.003 * g as f32;
5506 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
5507 }
5508 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
5511 x1[9] = 250.0;
5512 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
5513
5514 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
5515 q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
5516 q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
5517 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
5518 q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
5519 assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
5520 assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
5521 }
5522
5523 #[test]
5526 fn matvec_many_equals_separate_matvecs() {
5527 use crate::pool::Pool;
5528 let (r1, r2, cols) = (300, 200, 64);
5529 let mk = |salt: usize, rows: usize| {
5530 QTensor::from_f32(
5531 (0..rows * cols).map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5).collect(),
5532 rows,
5533 cols,
5534 )
5535 };
5536 let (a, b) = (mk(1, r1), mk(5, r2));
5537 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
5538 let pool = Pool::new(3);
5539
5540 let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
5541 a.matvec(&x, &mut ea, Some(&pool));
5542 b.matvec(&x, &mut eb, Some(&pool));
5543 let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
5544 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
5545 assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
5546 assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
5547 }
5548
5549 #[test]
5552 fn batched_matmat_equals_per_position_matvec() {
5553 let (rows, cols, b) = (8, 64, 5);
5554 let groups = rows * cols / GROUP_SIZE;
5556 let mut q4 = Vec::new();
5557 for i in 0..groups * 16 {
5558 q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
5559 }
5560 for g in 0..groups {
5561 q4.extend_from_slice(
5562 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
5563 );
5564 }
5565 let ng = cols / GROUP_SIZE;
5567 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
5568 let mut vb = bits.clone();
5569 for g in 0..rows * ng {
5570 vb.extend_from_slice(
5571 &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
5572 );
5573 }
5574 for r in 0..rows {
5575 let bw = bits[r] as usize;
5576 let (mut acc, mut nb) = (0u64, 0usize);
5577 let mut rowbytes = Vec::new();
5578 for i in 0..cols {
5579 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
5580 acc = (acc << bw) | v;
5581 nb += bw;
5582 while nb >= 8 {
5583 nb -= 8;
5584 rowbytes.push(((acc >> nb) & 0xFF) as u8);
5585 }
5586 }
5587 if nb > 0 {
5588 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
5589 }
5590 vb.extend_from_slice(&rowbytes);
5591 }
5592 let offsets = vbit_row_offsets(&vb, rows, cols);
5593
5594 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
5595
5596 let mut got = vec![0f32; b * rows];
5598 q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
5599 for bi in 0..b {
5600 let mut expect = vec![0f32; rows];
5601 q4matvec(&q4, &xs[bi * cols..(bi + 1) * cols], rows, cols, &mut expect, None);
5602 assert_eq!(&got[bi * rows..(bi + 1) * rows], &expect[..], "q4 batch pos {bi}");
5603 }
5604
5605 let mut got = vec![0f32; b * rows];
5607 vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
5608 for bi in 0..b {
5609 let mut expect = vec![0f32; rows];
5610 vbitmatvec(
5611 &vb, &offsets, &xs[bi * cols..(bi + 1) * cols], rows, cols, &mut expect, None,
5612 );
5613 assert_eq!(&got[bi * rows..(bi + 1) * rows], &expect[..], "vbit batch pos {bi}");
5614 }
5615 }
5616
5617 #[test]
5621 fn q4_tiled_matches_q4_block_bitexact() {
5622 let (rows, cols, b) = (8usize, 128usize, 3usize);
5623 let groups = rows * cols / GROUP_SIZE;
5624 let mut split = Vec::with_capacity(groups * 18);
5625 for i in 0..groups * 16 {
5626 split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
5627 }
5628 for g in 0..groups {
5629 split.extend_from_slice(
5630 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
5631 );
5632 }
5633 let (packed, scales) = split.split_at(groups * 16);
5635 let mut tiled = Vec::with_capacity(groups * Q4_TILE);
5636 for g in 0..groups {
5637 tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
5638 tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
5639 }
5640
5641 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
5642 x1[9] = 250.0; let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
5644
5645 let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
5646 q4matvec(&split, &x1, rows, cols, &mut a, None);
5647 q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
5648 assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
5649
5650 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
5651 let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
5652 q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
5653 q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
5654 assert_eq!(a1, t1);
5655 assert_eq!(a2, t2);
5656
5657 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
5658 let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
5659 q4matmat(&split, &xs, b, rows, cols, &mut am, None);
5660 q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
5661 assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
5662 }
5663
5664 #[test]
5671 fn q4matvec_sdot_outlier_exact() {
5672 let (rows, cols) = (4, 128);
5673 let groups = rows * cols / GROUP_SIZE;
5674 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
5675 for i in 0..groups * 16 {
5676 bytes.push(((i * 11 + 5) % 256) as u8);
5677 }
5678 for g in 0..groups {
5679 let s = 0.02 + 0.002 * g as f32;
5680 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
5681 }
5682 let mut x: Vec<f32> = (0..cols)
5683 .map(|i| match i % 3 {
5684 0 => 1.0,
5685 1 => -1.0,
5686 _ => 0.0,
5687 })
5688 .collect();
5689 x[17] = 300.0; let mut reference = vec![0.0f32; rows * cols];
5692 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
5693 let mut expect = vec![0.0f32; rows];
5694 for r in 0..rows {
5695 expect[r] = reference[r * cols..(r + 1) * cols]
5696 .iter()
5697 .zip(&x)
5698 .map(|(w, xv)| w * xv)
5699 .sum();
5700 }
5701 let mut got = vec![0.0f32; rows];
5702 q4matvec(&bytes, &x, rows, cols, &mut got, None);
5703 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
5704 for r in 0..rows {
5705 assert!(
5706 (got[r] - expect[r]).abs() < 2e-3 * scale,
5707 "row {r}: {} vs {} (outlier term must be exact)",
5708 got[r],
5709 expect[r]
5710 );
5711 }
5712 }
5713}