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 #[cfg(target_arch = "aarch64")]
1417 if std::env::var("CMF_FORCE_NEON_GEMM").map(|v| v == "1").unwrap_or(false) {
1418 return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
1419 }
1420 unsafe {
1421 accel_blas::cblas_sgemm(
1422 101, 111, if b_rows_are_n { 112 } else { 111 },
1425 m as i32,
1426 n as i32,
1427 k as i32,
1428 alpha,
1429 a.as_ptr(),
1430 lda as i32,
1431 b_mat.as_ptr(),
1432 ldb as i32,
1433 0.0,
1434 c.as_mut_ptr(),
1435 ldc as i32,
1436 );
1437 }
1438}
1439
1440#[cfg(target_os = "macos")]
1447fn qmatmat_accel(
1448 q: &[u8],
1449 row_scale: &[f32],
1450 pre: &[std::borrow::Cow<'_, [f32]>],
1451 rows: usize,
1452 cols: usize,
1453 out: &mut [f32],
1454 pool: Option<&Pool>,
1455) {
1456 const TR: usize = 2048;
1461 let b = pre.len();
1462 thread_local! {
1463 static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
1464 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
1465 }
1466 XPANEL.with(|xp| {
1467 WTILE.with(|wt| {
1468 let mut xpanel = xp.borrow_mut();
1469 xpanel.clear();
1470 for x in pre {
1471 xpanel.extend_from_slice(x);
1472 }
1473 let mut wtile = wt.borrow_mut();
1474 wtile.resize(TR * cols, 0.0);
1475 let mut r0 = 0usize;
1476 while r0 < rows {
1477 let tr = TR.min(rows - r0);
1478 let wt_addr = SendMut(wtile.as_mut_ptr());
1480 let run = |start: usize, end: usize| {
1481 for r in start..end {
1482 let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
1483 let s = row_scale[r0 + r];
1484 let dst = unsafe {
1486 std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols)
1487 };
1488 for (d, &v) in dst.iter_mut().zip(row) {
1489 *d = (v as i8) as f32 * s;
1490 }
1491 }
1492 };
1493 dispatch_rows(pool, tr, &run);
1494 unsafe {
1496 accel_blas::cblas_sgemm(
1497 101, 111, 112, b as i32,
1501 tr as i32,
1502 cols as i32,
1503 1.0,
1504 xpanel.as_ptr(),
1505 cols as i32,
1506 wtile.as_ptr(),
1507 cols as i32,
1508 0.0,
1509 out.as_mut_ptr().add(r0),
1510 rows as i32,
1511 );
1512 }
1513 r0 += tr;
1514 }
1515 })
1516 });
1517}
1518
1519fn qmatmat(
1520 q: &[u8],
1521 row_scale: &[f32],
1522 pre: &[std::borrow::Cow<'_, [f32]>],
1523 rows: usize,
1524 cols: usize,
1525 out: &mut [f32],
1526 pool: Option<&Pool>,
1527) {
1528 let b = pre.len();
1529 debug_assert_eq!(out.len(), b * rows);
1530 #[cfg(target_os = "macos")]
1535 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
1536 qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
1537 return;
1538 }
1539 #[cfg(target_arch = "aarch64")]
1540 if sdot_enabled() {
1541 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
1542 let out_addr = SendMut(out.as_mut_ptr());
1543 let blocked_ok =
1546 std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true);
1547 let use_i8mm = i8mm_enabled();
1548 if blocked_ok {
1549 let run = |start: usize, end: usize| {
1550 let mut o = start;
1551 while o < end {
1552 if o + 2 <= end {
1553 let r0 = &q[o * cols..(o + 1) * cols];
1554 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
1555 let mut bi = 0usize;
1556 while bi + 4 <= acts.len() {
1557 let xs = [
1558 acts[bi].xq.as_slice(),
1559 acts[bi + 1].xq.as_slice(),
1560 acts[bi + 2].xq.as_slice(),
1561 acts[bi + 3].xq.as_slice(),
1562 ];
1563 let d = if use_i8mm {
1564 unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
1565 } else {
1566 unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
1567 };
1568 for (r, row) in [r0, r1].into_iter().enumerate() {
1569 for k in 0..4 {
1570 let act = &acts[bi + k];
1571 let mut v = d[r][k] as f32 * act.sx;
1572 for &(j, xv) in &act.outliers {
1573 v += (row[j] as i8) as f32 * xv;
1574 }
1575 unsafe {
1576 *out_addr.at((bi + k) * rows + o + r) =
1577 v * row_scale[o + r]
1578 };
1579 }
1580 }
1581 bi += 4;
1582 }
1583 while bi < acts.len() {
1584 for (r, row) in [r0, r1].into_iter().enumerate() {
1585 let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
1586 unsafe { *out_addr.at(bi * rows + o + r) = v };
1587 }
1588 bi += 1;
1589 }
1590 o += 2;
1591 } else {
1592 let row = &q[o * cols..(o + 1) * cols];
1593 for (bi, act) in acts.iter().enumerate() {
1594 let v = row_dot_sdot(row, act) * row_scale[o];
1595 unsafe { *out_addr.at(bi * rows + o) = v };
1596 }
1597 o += 1;
1598 }
1599 }
1600 };
1601 dispatch_rows(pool, rows, &run);
1602 return;
1603 }
1604 let run = |start: usize, end: usize| {
1605 for o in start..end {
1606 let row = &q[o * cols..(o + 1) * cols];
1607 for (bi, act) in acts.iter().enumerate() {
1608 let v = row_dot_sdot(row, act) * row_scale[o];
1609 unsafe { *out_addr.at(bi * rows + o) = v };
1610 }
1611 }
1612 };
1613 dispatch_rows(pool, rows, &run);
1614 return;
1615 }
1616 #[cfg(target_arch = "x86_64")]
1621 if avx2_a8w8_enabled() {
1622 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
1623 let out_addr = SendMut(out.as_mut_ptr());
1624 let blocked_ok = std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true);
1627 if !avx512vnni_enabled() && blocked_ok {
1628 let run = |start: usize, end: usize| {
1629 let mut o = start;
1630 while o < end {
1631 if o + 2 <= end {
1632 let r0 = &q[o * cols..(o + 1) * cols];
1633 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
1634 let mut bi = 0usize;
1635 while bi + 4 <= acts.len() {
1636 let xs = [
1637 acts[bi].xq.as_slice(),
1638 acts[bi + 1].xq.as_slice(),
1639 acts[bi + 2].xq.as_slice(),
1640 acts[bi + 3].xq.as_slice(),
1641 ];
1642 let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
1643 for (r, row) in [r0, r1].into_iter().enumerate() {
1644 for k in 0..4 {
1645 let act = &acts[bi + k];
1646 let mut v = d[r][k] as f32 * act.sx;
1647 for &(j, xv) in &act.outliers {
1648 v += (row[j] as i8) as f32 * xv;
1649 }
1650 unsafe {
1651 *out_addr.at((bi + k) * rows + o + r) =
1652 v * row_scale[o + r]
1653 };
1654 }
1655 }
1656 bi += 4;
1657 }
1658 while bi < acts.len() {
1659 for (r, row) in [r0, r1].into_iter().enumerate() {
1660 let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
1661 unsafe { *out_addr.at(bi * rows + o + r) = v };
1662 }
1663 bi += 1;
1664 }
1665 o += 2;
1666 } else {
1667 let row = &q[o * cols..(o + 1) * cols];
1668 for (bi, act) in acts.iter().enumerate() {
1669 let v = row_dot_avx2(row, act) * row_scale[o];
1670 unsafe { *out_addr.at(bi * rows + o) = v };
1671 }
1672 o += 1;
1673 }
1674 }
1675 };
1676 dispatch_rows(pool, rows, &run);
1677 return;
1678 }
1679 let run = |start: usize, end: usize| {
1680 for o in start..end {
1681 let row = &q[o * cols..(o + 1) * cols];
1682 for (bi, act) in acts.iter().enumerate() {
1683 let v = row_dot_avx2(row, act) * row_scale[o];
1684 unsafe { *out_addr.at(bi * rows + o) = v };
1685 }
1686 }
1687 };
1688 dispatch_rows(pool, rows, &run);
1689 return;
1690 }
1691 let out_addr = SendMut(out.as_mut_ptr());
1692 let run = |start: usize, end: usize| {
1693 for o in start..end {
1694 let row = &q[o * cols..(o + 1) * cols];
1695 for (bi, x) in pre.iter().enumerate() {
1696 let mut acc = 0f32;
1697 for j in 0..cols {
1698 acc += (row[j] as i8) as f32 * x[j];
1699 }
1700 unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
1701 }
1702 }
1703 };
1704 dispatch_rows(pool, rows, &run);
1705}
1706
1707fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
1710 match pool {
1711 Some(pool) if rows >= 256 => pool.run_rows(rows, run),
1712 _ => run(0, rows),
1713 }
1714}
1715
1716fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
1718 let groups = rows * cols / GROUP_SIZE;
1719 bytes.split_at(groups * 16)
1720}
1721
1722#[inline]
1727fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
1728 #[cfg(target_arch = "aarch64")]
1729 unsafe {
1730 return vbit_fill4_neon(data, buf);
1731 }
1732 #[cfg(target_arch = "x86_64")]
1733 if avx2_enabled() {
1734 return unsafe { vbit_fill4_avx2(data, buf) };
1735 }
1736 #[allow(unreachable_code)]
1737 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
1738 let u = unpack8::<4>(&data[blk * 4..]);
1739 for k in 0..8 {
1740 chunk[k] = (u[k] - 7) as i8 as u8;
1741 }
1742 }
1743}
1744
1745#[cfg(target_arch = "aarch64")]
1746#[target_feature(enable = "neon")]
1747unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
1748 unsafe {
1751 use core::arch::aarch64::*;
1752 let n = buf.len();
1753 let mask = vdupq_n_u8(0x0F);
1754 let seven = vdupq_n_s8(7);
1755 let mut g = 0usize;
1756 while g * 32 + 32 <= n {
1757 let b = vld1q_u8(data.as_ptr().add(g * 16));
1758 let hi = vshrq_n_u8::<4>(b);
1759 let lo = vandq_u8(b, mask);
1760 let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
1761 let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
1762 vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
1763 vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
1764 g += 1;
1765 }
1766 }
1767}
1768
1769#[cfg(target_arch = "x86_64")]
1770#[target_feature(enable = "avx2")]
1771unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
1772 unsafe {
1774 use core::arch::x86_64::*;
1775 let n = buf.len();
1776 let mask = _mm_set1_epi8(0x0F);
1777 let seven = _mm256_set1_epi8(7);
1778 let mut g = 0usize;
1779 while g * 32 + 32 <= n {
1780 let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
1781 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
1782 let lo = _mm_and_si128(b, mask);
1783 let z = _mm256_sub_epi8(
1784 _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
1785 seven,
1786 );
1787 _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
1788 g += 1;
1789 }
1790 }
1791}
1792
1793#[inline(always)]
1798fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
1799 let mut acc = 0u64;
1800 for i in 0..B {
1801 acc = (acc << 8) | data[i] as u64;
1802 }
1803 let mask = (1u64 << B) - 1;
1804 let mut out = [0i32; 8];
1805 for (k, o) in out.iter_mut().enumerate() {
1806 *o = ((acc >> ((7 - k) * B)) & mask) as i32;
1807 }
1808 out
1809}
1810
1811#[allow(clippy::too_many_arguments)]
1817fn vbitmatvec(
1818 bytes: &[u8],
1819 offsets: &[usize],
1820 x: &[f32],
1821 rows: usize,
1822 cols: usize,
1823 out: &mut [f32],
1824 pool: Option<&Pool>,
1825) {
1826 debug_assert_eq!(out.len(), rows);
1827 debug_assert_eq!(offsets.len(), rows + 1);
1828
1829 if a8w8_enabled() {
1833 let act = split_act(x);
1834 let out_addr = SendMut(out.as_mut_ptr());
1835 let run = move |start: usize, end: usize| {
1836 vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
1837 };
1838 dispatch_rows(pool, rows, &run);
1839 return;
1840 }
1841
1842 let out_addr = SendMut(out.as_mut_ptr());
1843 let run = move |start: usize, end: usize| {
1844 vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
1845 };
1846 dispatch_rows(pool, rows, &run);
1847}
1848
1849#[allow(clippy::too_many_arguments)]
1853fn vbit_range_a8w8(
1854 bytes: &[u8],
1855 offsets: &[usize],
1856 x: &[f32],
1857 act: &SplitAct,
1858 rows: usize,
1859 cols: usize,
1860 out: SendMut,
1861 start: usize,
1862 end: usize,
1863) {
1864 let ng = cols / GROUP_SIZE;
1865 let bits = &bytes[..rows];
1866 let sc_off = rows;
1867 let row_dot = |r: usize| -> f32 {
1868 let b = bits[r] as usize;
1869 let l = ((1i32 << (b - 1)) - 1) as i32;
1870 let mask = (1u64 << b) - 1;
1871 let data = &bytes[offsets[r]..offsets[r + 1]];
1872 if b == 8 {
1873 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
1875 let mut dot = 0f32;
1876 for g in 0..ng {
1877 let so = (r * ng + g) * 2;
1878 let sgf = f16_to_f32(u16::from_le_bytes([
1879 bytes[sc_off + so],
1880 bytes[sc_off + so + 1],
1881 ]));
1882 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
1883 let mut gd = 0f32;
1884 for &xv in xg.iter() {
1885 if nbits < 8 {
1886 acc = (acc << 8) | data[idx] as u64;
1887 idx += 1;
1888 nbits += 8;
1889 }
1890 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
1891 nbits -= 8;
1892 gd += (u - l) as f32 * xv;
1893 }
1894 dot += gd * sgf;
1895 }
1896 return dot;
1897 }
1898 thread_local! {
1902 static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
1903 const { std::cell::RefCell::new(Vec::new()) };
1904 }
1905 #[inline(always)]
1906 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
1907 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
1908 let u = unpack8::<B>(&data[blk * B..]);
1909 for k in 0..8 {
1910 chunk[k] = (u[k] - l) as i8 as u8;
1911 }
1912 }
1913 }
1914 let _ = mask;
1915 VBIT_SCRATCH.with(|scratch| {
1916 let mut buf = scratch.borrow_mut();
1917 buf.resize(cols, 0);
1918 match b {
1919 3 => fill::<3>(data, l, &mut buf),
1920 4 => vbit_fill4(data, &mut buf),
1921 5 => fill::<5>(data, l, &mut buf),
1922 6 => fill::<6>(data, l, &mut buf),
1923 _ => unreachable!(),
1924 }
1925 let mut dot = 0f32;
1926 for g in 0..ng {
1927 let so = (r * ng + g) * 2;
1928 let s = f16_to_f32(u16::from_le_bytes([
1929 bytes[sc_off + so],
1930 bytes[sc_off + so + 1],
1931 ]));
1932 let d = dot_i8_i8(
1933 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
1934 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
1935 ) as f32
1936 * act.sx;
1937 dot += d * s;
1938 }
1939 for &(j, xv) in &act.outliers {
1940 let so = (r * ng + j / GROUP_SIZE) * 2;
1941 let s = f16_to_f32(u16::from_le_bytes([
1942 bytes[sc_off + so],
1943 bytes[sc_off + so + 1],
1944 ]));
1945 dot += (buf[j] as i8) as f32 * s * xv;
1947 }
1948 dot
1949 })
1950 };
1951 for r in start..end {
1952 unsafe { *out.at(r) = row_dot(r) };
1954 }
1955}
1956
1957#[allow(clippy::too_many_arguments)]
1959fn vbit_range_f32(
1960 bytes: &[u8],
1961 offsets: &[usize],
1962 x: &[f32],
1963 rows: usize,
1964 cols: usize,
1965 out: SendMut,
1966 start: usize,
1967 end: usize,
1968) {
1969 let ng = cols / GROUP_SIZE;
1970 let bits = &bytes[..rows];
1971 let sc_off = rows;
1972 #[inline(always)]
1976 fn dot_row<const B: usize>(
1977 data: &[u8],
1978 bytes: &[u8],
1979 sc_off: usize,
1980 r: usize,
1981 ng: usize,
1982 x: &[f32],
1983 ) -> f32 {
1984 let l = ((1i32 << (B - 1)) - 1) as f32;
1985 let gbytes = GROUP_SIZE * B / 8;
1986 let mut dot = 0f32;
1987 for g in 0..ng {
1988 let so = (r * ng + g) * 2;
1989 let s = f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]));
1990 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
1991 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
1992 let mut gd = 0f32;
1993 for blk in 0..GROUP_SIZE / 8 {
1994 let u = unpack8::<B>(&gd0[blk * B..]);
1995 let xb = &xg[blk * 8..blk * 8 + 8];
1996 for k in 0..8 {
1997 gd += (u[k] as f32 - l) * xb[k];
1998 }
1999 }
2000 dot += gd * s;
2001 }
2002 dot
2003 }
2004 for r in start..end {
2005 let data = &bytes[offsets[r]..offsets[r + 1]];
2006 let v = match bits[r] {
2007 3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
2008 4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
2009 5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
2010 6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
2011 8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
2012 b => unreachable!("vbit bit-width {b} (validated at load)"),
2013 };
2014 unsafe { *out.at(r) = v };
2016 }
2017}
2018
2019#[allow(clippy::too_many_arguments)]
2024fn vbitmatvec2(
2025 bytes: &[u8],
2026 offsets: &[usize],
2027 x1: &[f32],
2028 x2: &[f32],
2029 rows: usize,
2030 cols: usize,
2031 o1: &mut [f32],
2032 o2: &mut [f32],
2033 pool: Option<&Pool>,
2034) {
2035 debug_assert_eq!(o1.len(), rows);
2036 debug_assert_eq!(o2.len(), rows);
2037
2038 if a8w8_enabled() {
2039 let a1 = split_act(x1);
2040 let a2 = split_act(x2);
2041 let p1 = SendMut(o1.as_mut_ptr());
2042 let p2 = SendMut(o2.as_mut_ptr());
2043 let run = move |start: usize, end: usize| {
2044 vbit_range2_a8w8(bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end)
2045 };
2046 dispatch_rows(pool, rows, &run);
2047 return;
2048 }
2049
2050 let p1 = SendMut(o1.as_mut_ptr());
2051 let p2 = SendMut(o2.as_mut_ptr());
2052 let run = move |start: usize, end: usize| {
2053 vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
2054 };
2055 dispatch_rows(pool, rows, &run);
2056}
2057
2058#[allow(clippy::too_many_arguments)]
2062fn vbit_range2_a8w8(
2063 bytes: &[u8],
2064 offsets: &[usize],
2065 x1: &[f32],
2066 x2: &[f32],
2067 a1: &SplitAct,
2068 a2: &SplitAct,
2069 rows: usize,
2070 cols: usize,
2071 p1: SendMut,
2072 p2: SendMut,
2073 start: usize,
2074 end: usize,
2075) {
2076 let ng = cols / GROUP_SIZE;
2077 let bits = &bytes[..rows];
2078 let sc_off = rows;
2079 let row_dots = |r: usize| -> (f32, f32) {
2080 let b = bits[r] as usize;
2081 let l = (1i32 << (b - 1)) - 1;
2082 let data = &bytes[offsets[r]..offsets[r + 1]];
2083 if b == 8 {
2084 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
2087 let (mut d1, mut d2) = (0f32, 0f32);
2088 for g in 0..ng {
2089 let so = (r * ng + g) * 2;
2090 let sgf = f16_to_f32(u16::from_le_bytes([
2091 bytes[sc_off + so],
2092 bytes[sc_off + so + 1],
2093 ]));
2094 let (mut g1, mut g2) = (0f32, 0f32);
2095 for k in 0..GROUP_SIZE {
2096 if nbits < 8 {
2097 acc = (acc << 8) | data[idx] as u64;
2098 idx += 1;
2099 nbits += 8;
2100 }
2101 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
2102 nbits -= 8;
2103 let w = (u - l) as f32;
2104 g1 += w * x1[g * GROUP_SIZE + k];
2105 g2 += w * x2[g * GROUP_SIZE + k];
2106 }
2107 d1 += g1 * sgf;
2108 d2 += g2 * sgf;
2109 }
2110 return (d1, d2);
2111 }
2112 thread_local! {
2113 static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
2114 const { std::cell::RefCell::new(Vec::new()) };
2115 }
2116 #[inline(always)]
2117 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
2118 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
2119 let u = unpack8::<B>(&data[blk * B..]);
2120 for k in 0..8 {
2121 chunk[k] = (u[k] - l) as i8 as u8;
2122 }
2123 }
2124 }
2125 VBIT_SCRATCH2.with(|scratch| {
2126 let mut buf = scratch.borrow_mut();
2127 buf.resize(cols, 0);
2128 match b {
2129 3 => fill::<3>(data, l, &mut buf),
2130 4 => vbit_fill4(data, &mut buf),
2131 5 => fill::<5>(data, l, &mut buf),
2132 6 => fill::<6>(data, l, &mut buf),
2133 _ => unreachable!(),
2134 }
2135 let (mut d1, mut d2) = (0f32, 0f32);
2136 for g in 0..ng {
2137 let so = (r * ng + g) * 2;
2138 let s = f16_to_f32(u16::from_le_bytes([
2139 bytes[sc_off + so],
2140 bytes[sc_off + so + 1],
2141 ]));
2142 let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2143 let v1 =
2144 dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
2145 let v2 =
2146 dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
2147 d1 += v1 * s;
2148 d2 += v2 * s;
2149 }
2150 for &(j, xv) in &a1.outliers {
2151 let so = (r * ng + j / GROUP_SIZE) * 2;
2152 let s = f16_to_f32(u16::from_le_bytes([
2153 bytes[sc_off + so],
2154 bytes[sc_off + so + 1],
2155 ]));
2156 d1 += (buf[j] as i8) as f32 * s * xv;
2157 }
2158 for &(j, xv) in &a2.outliers {
2159 let so = (r * ng + j / GROUP_SIZE) * 2;
2160 let s = f16_to_f32(u16::from_le_bytes([
2161 bytes[sc_off + so],
2162 bytes[sc_off + so + 1],
2163 ]));
2164 d2 += (buf[j] as i8) as f32 * s * xv;
2165 }
2166 (d1, d2)
2167 })
2168 };
2169 for r in start..end {
2170 let (v1, v2) = row_dots(r);
2171 unsafe {
2173 *p1.at(r) = v1;
2174 *p2.at(r) = v2;
2175 }
2176 }
2177}
2178
2179#[allow(clippy::too_many_arguments)]
2183fn vbit_range2_f32(
2184 bytes: &[u8],
2185 offsets: &[usize],
2186 x1: &[f32],
2187 x2: &[f32],
2188 rows: usize,
2189 cols: usize,
2190 p1: SendMut,
2191 p2: SendMut,
2192 start: usize,
2193 end: usize,
2194) {
2195 let ng = cols / GROUP_SIZE;
2196 let bits = &bytes[..rows];
2197 let sc_off = rows;
2198 #[inline(always)]
2199 #[allow(clippy::too_many_arguments)]
2200 fn dot_row2<const B: usize>(
2201 data: &[u8],
2202 bytes: &[u8],
2203 sc_off: usize,
2204 r: usize,
2205 ng: usize,
2206 x1: &[f32],
2207 x2: &[f32],
2208 ) -> (f32, f32) {
2209 let l = ((1i32 << (B - 1)) - 1) as f32;
2210 let gbytes = GROUP_SIZE * B / 8;
2211 let (mut d1, mut d2) = (0f32, 0f32);
2212 for g in 0..ng {
2213 let so = (r * ng + g) * 2;
2214 let s = f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]));
2215 let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2216 let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2217 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
2218 let (mut g1, mut g2) = (0f32, 0f32);
2219 for blk in 0..GROUP_SIZE / 8 {
2220 let u = unpack8::<B>(&gd0[blk * B..]);
2221 for k in 0..8 {
2222 let w = u[k] as f32 - l;
2223 g1 += w * x1g[blk * 8 + k];
2224 g2 += w * x2g[blk * 8 + k];
2225 }
2226 }
2227 d1 += g1 * s;
2228 d2 += g2 * s;
2229 }
2230 (d1, d2)
2231 }
2232 for r in start..end {
2233 let data = &bytes[offsets[r]..offsets[r + 1]];
2234 let (v1, v2) = match bits[r] {
2235 3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
2236 4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
2237 5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
2238 6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
2239 8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
2240 b => unreachable!("vbit bit-width {b} (validated at load)"),
2241 };
2242 unsafe {
2244 *p1.at(r) = v1;
2245 *p2.at(r) = v2;
2246 }
2247 }
2248}
2249
2250#[inline]
2257#[allow(unreachable_code)]
2258fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
2259 #[cfg(target_arch = "aarch64")]
2260 unsafe {
2261 return dot_q4t_row_sdot(bytes, r, gpr, xq);
2262 }
2263 #[cfg(target_arch = "x86_64")]
2264 unsafe {
2265 return dot_q4t_row_avx2(bytes, r, gpr, xq);
2266 }
2267 let mut acc = 0f32;
2268 for gi in 0..gpr {
2269 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
2270 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2271 let mut d = 0i32;
2272 for (k, &b) in tile[2..].iter().enumerate() {
2273 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
2274 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
2275 }
2276 acc += d as f32 * s;
2277 }
2278 acc
2279}
2280
2281#[cfg(target_arch = "aarch64")]
2282#[target_feature(enable = "neon,dotprod")]
2283unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
2284 unsafe {
2287 use core::arch::aarch64::*;
2288 use core::arch::asm;
2289 let lomask = vdupq_n_u8(0x0F);
2290 let eight = vdupq_n_s8(8);
2291 let mut acc = 0f32;
2292 for gi in 0..gpr {
2293 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
2294 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2295 let b = vld1q_u8(t.add(2));
2296 let lo = vandq_u8(b, lomask);
2297 let hi = vshrq_n_u8::<4>(b);
2298 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
2299 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
2300 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
2301 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
2302 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
2303 asm!(
2304 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
2305 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
2306 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
2307 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
2308 options(pure, nomem, nostack),
2309 );
2310 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
2311 }
2312 acc
2313 }
2314}
2315
2316#[cfg(target_arch = "x86_64")]
2317#[target_feature(enable = "avx2")]
2318unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
2319 unsafe {
2321 use core::arch::x86_64::*;
2322 let lomask = _mm_set1_epi8(0x0F);
2323 let eight = _mm256_set1_epi8(8);
2324 let ones = _mm256_set1_epi16(1);
2325 let mut acc = 0f32;
2326 for gi in 0..gpr {
2327 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
2328 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2329 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
2330 let lo = _mm_and_si128(b, lomask);
2331 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
2332 let w = _mm256_sub_epi8(
2333 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
2334 eight,
2335 );
2336 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2337 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
2338 let d = _mm256_madd_epi16(p16, ones);
2339 let hi128 = _mm256_extracti128_si256::<1>(d);
2340 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
2341 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2342 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2343 acc += _mm_cvtsi128_si32(s32) as f32 * s;
2344 }
2345 acc
2346 }
2347}
2348
2349#[cfg(target_arch = "x86_64")]
2354#[target_feature(enable = "avx2")]
2355unsafe fn dot_q4t_row_1x4_avx2(
2356 bytes: &[u8],
2357 r: usize,
2358 gpr: usize,
2359 xs: [&[i8]; 4],
2360) -> [f32; 4] {
2361 unsafe {
2363 use core::arch::x86_64::*;
2364 let lomask = _mm_set1_epi8(0x0F);
2365 let eight = _mm256_set1_epi8(8);
2366 let ones = _mm256_set1_epi16(1);
2367 let mut acc = [0f32; 4];
2368 for gi in 0..gpr {
2369 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
2370 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2371 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
2372 let lo = _mm_and_si128(bb, lomask);
2373 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
2374 let w = _mm256_sub_epi8(
2375 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
2376 eight,
2377 );
2378 let aw = _mm256_abs_epi8(w);
2379 for (k, xq) in xs.iter().enumerate() {
2380 let x =
2381 _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2382 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
2383 let d = _mm256_madd_epi16(p16, ones);
2384 let hi128 = _mm256_extracti128_si256::<1>(d);
2385 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
2386 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2387 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2388 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
2389 }
2390 }
2391 acc
2392 }
2393}
2394
2395#[inline]
2397fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
2398 let gi = j / GROUP_SIZE;
2399 let k = j % GROUP_SIZE;
2400 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
2401 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2402 let byte = tile[2 + k / 2];
2403 let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
2404 ((nib as i32 - 8) as f32, s)
2405}
2406
2407#[inline]
2410fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
2411 let mut acc = 0f32;
2412 for gi in 0..gpr {
2413 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
2414 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2415 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
2416 let mut ga = 0f32;
2417 for (k, &b) in tile[2..].iter().enumerate() {
2418 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
2419 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
2420 }
2421 acc += ga * s;
2422 }
2423 acc
2424}
2425
2426fn q4t_matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
2428 debug_assert_eq!(out.len(), rows);
2429 let gpr = cols / GROUP_SIZE;
2430 let out_addr = SendMut(out.as_mut_ptr());
2431 if a8w8_enabled() {
2432 let act = split_act(x);
2433 let run = move |start: usize, end: usize| {
2434 for r in start..end {
2435 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
2436 for &(j, xv) in &act.outliers {
2437 let (w, s) = q4t_outlier(bytes, r, gpr, j);
2438 acc += w * s * xv;
2439 }
2440 unsafe { *out_addr.at(r) = acc };
2442 }
2443 };
2444 dispatch_rows(pool, rows, &run);
2445 return;
2446 }
2447 let run = move |start: usize, end: usize| {
2448 for r in start..end {
2449 unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
2451 }
2452 };
2453 dispatch_rows(pool, rows, &run);
2454}
2455
2456#[allow(clippy::too_many_arguments)]
2458fn q4t_matvec2(
2459 bytes: &[u8],
2460 x1: &[f32],
2461 x2: &[f32],
2462 rows: usize,
2463 cols: usize,
2464 o1: &mut [f32],
2465 o2: &mut [f32],
2466 pool: Option<&Pool>,
2467) {
2468 let gpr = cols / GROUP_SIZE;
2469 let p1 = SendMut(o1.as_mut_ptr());
2470 let p2 = SendMut(o2.as_mut_ptr());
2471 if a8w8_enabled() {
2472 let a1 = split_act(x1);
2473 let a2 = split_act(x2);
2474 let run = move |start: usize, end: usize| {
2475 for r in start..end {
2476 let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
2477 let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
2478 for &(j, xv) in &a1.outliers {
2479 let (w, s) = q4t_outlier(bytes, r, gpr, j);
2480 v1 += w * s * xv;
2481 }
2482 for &(j, xv) in &a2.outliers {
2483 let (w, s) = q4t_outlier(bytes, r, gpr, j);
2484 v2 += w * s * xv;
2485 }
2486 unsafe {
2488 *p1.at(r) = v1;
2489 *p2.at(r) = v2;
2490 }
2491 }
2492 };
2493 dispatch_rows(pool, rows, &run);
2494 return;
2495 }
2496 let run = move |start: usize, end: usize| {
2497 for r in start..end {
2498 unsafe {
2500 *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
2501 *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
2502 }
2503 }
2504 };
2505 dispatch_rows(pool, rows, &run);
2506}
2507
2508#[allow(clippy::too_many_arguments)]
2510fn q4t_matmat(
2511 bytes: &[u8],
2512 xs_all: &[f32],
2513 b: usize,
2514 rows: usize,
2515 cols: usize,
2516 out: &mut [f32],
2517 pool: Option<&Pool>,
2518) {
2519 debug_assert_eq!(out.len(), b * rows);
2520 let gpr = cols / GROUP_SIZE;
2521 let out_addr = SendMut(out.as_mut_ptr());
2522 if a8w8_enabled() {
2523 let acts: Vec<SplitAct> =
2524 (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
2525 let acts = &acts;
2526 #[cfg(target_arch = "x86_64")]
2527 let blocked_ok = avx2_enabled()
2528 && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true);
2529 #[cfg(not(target_arch = "x86_64"))]
2530 let blocked_ok = false;
2531 let run = move |start: usize, end: usize| {
2532 for r in start..end {
2533 let mut bi = 0usize;
2534 #[cfg(target_arch = "x86_64")]
2535 if blocked_ok {
2536 while bi + 4 <= acts.len() {
2537 let xs = [
2538 acts[bi].xq.as_slice(),
2539 acts[bi + 1].xq.as_slice(),
2540 acts[bi + 2].xq.as_slice(),
2541 acts[bi + 3].xq.as_slice(),
2542 ];
2543 let d = unsafe { dot_q4t_row_1x4_avx2(bytes, r, gpr, xs) };
2544 for k in 0..4 {
2545 let act = &acts[bi + k];
2546 let mut acc = d[k] * act.sx;
2547 for &(j, xv) in &act.outliers {
2548 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
2549 acc += w * sc * xv;
2550 }
2551 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
2553 }
2554 bi += 4;
2555 }
2556 }
2557 let _ = blocked_ok;
2558 while bi < acts.len() {
2559 let act = &acts[bi];
2560 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
2561 for &(j, xv) in &act.outliers {
2562 let (w, s) = q4t_outlier(bytes, r, gpr, j);
2563 acc += w * s * xv;
2564 }
2565 unsafe { *out_addr.at(bi * rows + r) = acc };
2567 bi += 1;
2568 }
2569 }
2570 };
2571 dispatch_rows(pool, rows, &run);
2572 return;
2573 }
2574 let run = move |start: usize, end: usize| {
2575 for r in start..end {
2576 for bi in 0..b {
2577 let x = &xs_all[bi * cols..(bi + 1) * cols];
2578 unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
2580 }
2581 }
2582 };
2583 dispatch_rows(pool, rows, &run);
2584}
2585
2586fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
2595 (0..gpr)
2596 .map(|gi| {
2597 xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
2598 .iter()
2599 .map(|&v| v as i32)
2600 .sum()
2601 })
2602 .collect()
2603}
2604
2605#[inline]
2609#[allow(unreachable_code)]
2610#[cfg(target_arch = "x86_64")]
2615#[target_feature(enable = "avx2")]
2616unsafe fn dot_q1_row_avx2(
2617 bytes: &[u8],
2618 r: usize,
2619 gpr: usize,
2620 xq: &[i8],
2621 gsum: &[i32],
2622) -> f32 {
2623 unsafe {
2625 use core::arch::x86_64::*;
2626 let expand = _mm256_setr_epi8(
2628 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1,
2629 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3,
2630 );
2631 let bitsel = _mm256_setr_epi8(
2632 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128,
2633 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128,
2634 );
2635 let ones8 = _mm256_set1_epi8(1);
2636 let ones16 = _mm256_set1_epi16(1);
2637 let mut acc = 0f32;
2638 for gi in 0..gpr {
2639 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
2640 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2641 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
2642 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
2643 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
2644 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2645 let sel = _mm256_and_si256(x, mask);
2646 let p16 = _mm256_maddubs_epi16(ones8, sel);
2648 let d32 = _mm256_madd_epi16(p16, ones16);
2649 let hi128 = _mm256_extracti128_si256::<1>(d32);
2650 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
2651 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2652 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2653 let msum = _mm_cvtsi128_si32(s32);
2654 let d = 2 * msum - gsum[gi];
2657 acc += d as f32 * s;
2658 }
2659 acc
2660 }
2661}
2662
2663#[cfg(target_arch = "x86_64")]
2666#[target_feature(enable = "avx2")]
2667unsafe fn dot_q1_row_1x4_avx2(
2668 bytes: &[u8],
2669 r: usize,
2670 gpr: usize,
2671 xs: [&[i8]; 4],
2672 gsums: [&[i32]; 4],
2673) -> [f32; 4] {
2674 unsafe {
2676 use core::arch::x86_64::*;
2677 let expand = _mm256_setr_epi8(
2678 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1,
2679 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3,
2680 );
2681 let bitsel = _mm256_setr_epi8(
2682 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128,
2683 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128,
2684 );
2685 let ones8 = _mm256_set1_epi8(1);
2686 let ones16 = _mm256_set1_epi16(1);
2687 let mut acc = [0f32; 4];
2688 for gi in 0..gpr {
2689 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
2690 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2691 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
2692 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
2693 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
2694 for (k, xq) in xs.iter().enumerate() {
2695 let x =
2696 _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2697 let sel = _mm256_and_si256(x, mask);
2698 let p16 = _mm256_maddubs_epi16(ones8, sel);
2699 let d32 = _mm256_madd_epi16(p16, ones16);
2700 let hi128 = _mm256_extracti128_si256::<1>(d32);
2701 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
2702 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2703 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2704 let msum = _mm_cvtsi128_si32(s32);
2705 let d = 2 * msum - gsums[k][gi];
2706 acc[k] += d as f32 * s;
2707 }
2708 }
2709 acc
2710 }
2711}
2712
2713fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
2714 #[cfg(target_arch = "aarch64")]
2715 unsafe {
2716 return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
2717 }
2718 #[cfg(target_arch = "x86_64")]
2719 if avx2_enabled() {
2720 unsafe {
2721 return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
2722 }
2723 }
2724 let _ = gsum;
2725 let mut acc = 0f32;
2726 for gi in 0..gpr {
2727 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
2728 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2729 let mut d = 0i32;
2730 for (j, &b) in tile[2..].iter().enumerate() {
2731 for k in 0..8 {
2732 let w = ((b >> k) & 1) as i32 * 2 - 1;
2733 d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
2734 }
2735 }
2736 acc += d as f32 * s;
2737 }
2738 acc
2739}
2740
2741#[cfg(target_arch = "aarch64")]
2750#[target_feature(enable = "neon,dotprod")]
2751unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
2752 unsafe {
2755 use core::arch::aarch64::*;
2756 use core::arch::asm;
2757 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
2758 let m = vld1q_u8(MASKS.as_ptr());
2759 macro_rules! tile_dot {
2761 ($t:expr, $x:expr) => {{
2762 let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
2763 let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
2764 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
2765 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
2766 let x0 = vld1q_s8($x);
2767 let x1 = vld1q_s8($x.add(16));
2768 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
2769 asm!(
2770 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
2771 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
2772 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
2773 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
2774 options(pure, nomem, nostack),
2775 );
2776 vaddq_s32(a0, a1)
2777 }};
2778 }
2779 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
2788 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
2789 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
2790 const IW11: [u8; 16] = [10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11];
2791 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
2792 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
2793 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
2794 let isc = vld1_u8(ISC.as_ptr());
2795 macro_rules! tile_dot_tbl {
2797 ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
2798 let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
2799 let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
2800 let x0 = vld1q_s8($x);
2801 let x1 = vld1q_s8($x.add(16));
2802 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
2803 asm!(
2804 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
2805 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
2806 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
2807 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
2808 options(pure, nomem, nostack),
2809 );
2810 vaddq_s32(a0, a1)
2811 }};
2812 }
2813 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
2814 let row_base = r * gpr * Q1_TILE;
2815 let abs_end = bytes.len();
2816 let xp = xq.as_ptr();
2817 let gp = gsum.as_ptr();
2818 let mut accv = vdupq_n_f32(0.0);
2819 let mut gi = 0;
2820 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
2823 let t0 = base.add(gi * Q1_TILE);
2824 let ld_a = vld1q_u8(t0);
2825 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
2826 let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
2827 let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
2828 let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
2829 let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
2830 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
2832 let g = vld1q_s32(gp.add(gi));
2833 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
2834 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
2835 let scf: float32x4_t;
2836 asm!(
2837 "fcvtl {o:v}.4s, {i:v}.4h",
2838 o = out(vreg) scf, i = in(vreg) sc16,
2839 options(pure, nomem, nostack),
2840 );
2841 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
2842 gi += 4;
2843 }
2844 let mut acc = vaddvq_f32(accv);
2845 while gi < gpr {
2846 let t = base.add(gi * Q1_TILE);
2847 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2848 let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
2849 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
2850 gi += 1;
2851 }
2852 acc
2853 }
2854}
2855
2856#[cfg(target_arch = "aarch64")]
2861#[target_feature(enable = "neon,dotprod")]
2862unsafe fn dot_q1_row_1x4_sdot(
2863 bytes: &[u8],
2864 r: usize,
2865 gpr: usize,
2866 xs: [&[i8]; 4],
2867 gs: [&[i32]; 4],
2868) -> [f32; 4] {
2869 unsafe {
2871 use core::arch::aarch64::*;
2872 use core::arch::asm;
2873 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
2874 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
2875 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
2876 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
2877 const IW11: [u8; 16] = [10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11];
2878 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
2879 let m = vld1q_u8(MASKS.as_ptr());
2880 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
2881 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
2882 let isc = vld1_u8(ISC.as_ptr());
2883 macro_rules! sdot2 {
2884 ($w0:expr, $w1:expr, $x:expr) => {{
2885 let x0 = vld1q_s8($x);
2886 let x1 = vld1q_s8($x.add(16));
2887 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
2888 asm!(
2889 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
2890 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
2891 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
2892 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
2893 options(pure, nomem, nostack),
2894 );
2895 vaddq_s32(a0, a1)
2896 }};
2897 }
2898 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
2899 let row_base = r * gpr * Q1_TILE;
2900 let abs_end = bytes.len();
2901 let mut accv = [vdupq_n_f32(0.0); 4];
2902 let mut gi = 0;
2903 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
2904 let t0 = base.add(gi * Q1_TILE);
2905 let ld_a = vld1q_u8(t0);
2906 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
2907 let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
2909 let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
2910 let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
2911 let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
2912 let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
2913 let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
2914 let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
2915 let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
2916 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
2917 let scf: float32x4_t;
2918 asm!(
2919 "fcvtl {o:v}.4s, {i:v}.4h",
2920 o = out(vreg) scf, i = in(vreg) sc16,
2921 options(pure, nomem, nostack),
2922 );
2923 for k in 0..4 {
2924 let xp = xs[k].as_ptr();
2925 let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
2926 let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
2927 let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
2928 let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
2929 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
2930 let g = vld1q_s32(gs[k].as_ptr().add(gi));
2931 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
2932 accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
2933 }
2934 gi += 4;
2935 }
2936 let mut acc = [
2937 vaddvq_f32(accv[0]),
2938 vaddvq_f32(accv[1]),
2939 vaddvq_f32(accv[2]),
2940 vaddvq_f32(accv[3]),
2941 ];
2942 while gi < gpr {
2943 let t = base.add(gi * Q1_TILE);
2944 let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2945 let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
2946 let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
2947 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
2948 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
2949 for k in 0..4 {
2950 let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
2951 acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
2952 }
2953 gi += 1;
2954 }
2955 acc
2956 }
2957}
2958
2959#[inline]
2961fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
2962 let gi = j / GROUP_SIZE;
2963 let k = j % GROUP_SIZE;
2964 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
2965 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2966 let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
2967 ((bit as i32 * 2 - 1) as f32, s)
2968}
2969
2970#[inline]
2972fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
2973 let mut acc = 0f32;
2974 for gi in 0..gpr {
2975 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
2976 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2977 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
2978 let mut ga = 0f32;
2979 for (j, &b) in tile[2..].iter().enumerate() {
2980 for k in 0..8 {
2981 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
2982 }
2983 }
2984 acc += ga * s;
2985 }
2986 acc
2987}
2988
2989#[allow(clippy::too_many_arguments)]
2992fn q1_range_a8w8(
2993 bytes: &[u8],
2994 gpr: usize,
2995 act: &SplitAct,
2996 gsum: &[i32],
2997 out: SendMut,
2998 start: usize,
2999 end: usize,
3000) {
3001 for r in start..end {
3002 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
3003 for &(j, xv) in &act.outliers {
3004 let (w, s) = q1_outlier(bytes, r, gpr, j);
3005 acc += w * s * xv;
3006 }
3007 unsafe { *out.at(r) = acc };
3009 }
3010}
3011
3012fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
3014 for r in start..end {
3015 unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
3017 }
3018}
3019
3020fn q1_matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
3022 debug_assert_eq!(out.len(), rows);
3023 let gpr = cols / GROUP_SIZE;
3024 let out_addr = SendMut(out.as_mut_ptr());
3025 if a8w8_enabled() {
3026 let act = split_act(x);
3027 let gsum = q1_group_sums(&act.xq, gpr);
3028 let (act, gsum) = (&act, &gsum);
3029 let run =
3030 move |start: usize, end: usize| q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end);
3031 dispatch_rows(pool, rows, &run);
3032 return;
3033 }
3034 let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
3035 dispatch_rows(pool, rows, &run);
3036}
3037
3038#[allow(clippy::too_many_arguments)]
3040fn q1_matvec2(
3041 bytes: &[u8],
3042 x1: &[f32],
3043 x2: &[f32],
3044 rows: usize,
3045 cols: usize,
3046 o1: &mut [f32],
3047 o2: &mut [f32],
3048 pool: Option<&Pool>,
3049) {
3050 let gpr = cols / GROUP_SIZE;
3051 let p1 = SendMut(o1.as_mut_ptr());
3052 let p2 = SendMut(o2.as_mut_ptr());
3053 if a8w8_enabled() {
3054 let a1 = split_act(x1);
3055 let a2 = split_act(x2);
3056 let g1 = q1_group_sums(&a1.xq, gpr);
3057 let g2 = q1_group_sums(&a2.xq, gpr);
3058 let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
3059 let run = move |start: usize, end: usize| {
3060 for r in start..end {
3061 let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
3062 let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
3063 for &(j, xv) in &a1.outliers {
3064 let (w, s) = q1_outlier(bytes, r, gpr, j);
3065 v1 += w * s * xv;
3066 }
3067 for &(j, xv) in &a2.outliers {
3068 let (w, s) = q1_outlier(bytes, r, gpr, j);
3069 v2 += w * s * xv;
3070 }
3071 unsafe {
3073 *p1.at(r) = v1;
3074 *p2.at(r) = v2;
3075 }
3076 }
3077 };
3078 dispatch_rows(pool, rows, &run);
3079 return;
3080 }
3081 let run = move |start: usize, end: usize| {
3082 for r in start..end {
3083 unsafe {
3085 *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
3086 *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
3087 }
3088 }
3089 };
3090 dispatch_rows(pool, rows, &run);
3091}
3092
3093#[allow(clippy::too_many_arguments)]
3095fn q1_matmat(
3096 bytes: &[u8],
3097 xs_all: &[f32],
3098 b: usize,
3099 rows: usize,
3100 cols: usize,
3101 out: &mut [f32],
3102 pool: Option<&Pool>,
3103) {
3104 debug_assert_eq!(out.len(), b * rows);
3105 let gpr = cols / GROUP_SIZE;
3106 let out_addr = SendMut(out.as_mut_ptr());
3107 if a8w8_enabled() {
3108 let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
3109 .map(|bi| {
3110 let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
3111 let gsum = q1_group_sums(&act.xq, gpr);
3112 (act, gsum)
3113 })
3114 .collect();
3115 let acts = &acts;
3116 #[cfg(target_arch = "x86_64")]
3117 let blocked_ok = avx2_enabled()
3118 && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true);
3119 #[cfg(target_arch = "aarch64")]
3120 let blocked_ok = sdot_enabled()
3121 && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true);
3122 let run = move |start: usize, end: usize| {
3123 for r in start..end {
3124 let mut bi = 0usize;
3125 #[cfg(target_arch = "aarch64")]
3128 if blocked_ok {
3129 while bi + 4 <= acts.len() {
3130 let xs = [
3131 acts[bi].0.xq.as_slice(),
3132 acts[bi + 1].0.xq.as_slice(),
3133 acts[bi + 2].0.xq.as_slice(),
3134 acts[bi + 3].0.xq.as_slice(),
3135 ];
3136 let gs = [
3137 acts[bi].1.as_slice(),
3138 acts[bi + 1].1.as_slice(),
3139 acts[bi + 2].1.as_slice(),
3140 acts[bi + 3].1.as_slice(),
3141 ];
3142 let d = unsafe { dot_q1_row_1x4_sdot(bytes, r, gpr, xs, gs) };
3143 for k in 0..4 {
3144 let (act, _) = &acts[bi + k];
3145 let mut acc = d[k] * act.sx;
3146 for &(j, xv) in &act.outliers {
3147 let (w, sc) = q1_outlier(bytes, r, gpr, j);
3148 acc += w * sc * xv;
3149 }
3150 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
3152 }
3153 bi += 4;
3154 }
3155 }
3156 #[cfg(target_arch = "x86_64")]
3157 if blocked_ok {
3158 while bi + 4 <= acts.len() {
3159 let xs = [
3160 acts[bi].0.xq.as_slice(),
3161 acts[bi + 1].0.xq.as_slice(),
3162 acts[bi + 2].0.xq.as_slice(),
3163 acts[bi + 3].0.xq.as_slice(),
3164 ];
3165 let gs = [
3166 acts[bi].1.as_slice(),
3167 acts[bi + 1].1.as_slice(),
3168 acts[bi + 2].1.as_slice(),
3169 acts[bi + 3].1.as_slice(),
3170 ];
3171 let d = unsafe { dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs) };
3172 for k in 0..4 {
3173 let (act, _) = &acts[bi + k];
3174 let mut acc = d[k] * act.sx;
3175 for &(j, xv) in &act.outliers {
3176 let (w, sc) = q1_outlier(bytes, r, gpr, j);
3177 acc += w * sc * xv;
3178 }
3179 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
3181 }
3182 bi += 4;
3183 }
3184 }
3185 while bi < acts.len() {
3186 let (act, gsum) = &acts[bi];
3187 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
3188 for &(j, xv) in &act.outliers {
3189 let (w, s) = q1_outlier(bytes, r, gpr, j);
3190 acc += w * s * xv;
3191 }
3192 unsafe { *out_addr.at(bi * rows + r) = acc };
3194 bi += 1;
3195 }
3196 }
3197 };
3198 dispatch_rows(pool, rows, &run);
3199 return;
3200 }
3201 let run = move |start: usize, end: usize| {
3202 for r in start..end {
3203 for bi in 0..b {
3204 let x = &xs_all[bi * cols..(bi + 1) * cols];
3205 unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
3207 }
3208 }
3209 };
3210 dispatch_rows(pool, rows, &run);
3211}
3212
3213fn q4matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
3219 debug_assert_eq!(out.len(), rows);
3220 let (packed, scales) = q4_split(bytes, rows, cols);
3221 let gpr = cols / GROUP_SIZE;
3222 let out_addr = SendMut(out.as_mut_ptr());
3223
3224 if a8w8_enabled() {
3225 let act = split_act(x);
3226 let run = move |start: usize, end: usize| {
3227 q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
3228 };
3229 dispatch_rows(pool, rows, &run);
3230 return;
3231 }
3232
3233 let run = move |start: usize, end: usize| {
3234 q4_range_f32(packed, scales, gpr, x, out_addr, start, end)
3235 };
3236 dispatch_rows(pool, rows, &run);
3237}
3238
3239#[inline]
3242#[allow(unreachable_code)]
3243#[cfg(target_arch = "x86_64")]
3248#[target_feature(enable = "avx2")]
3249unsafe fn dot_q4b_row_1x4_avx2(
3250 buf: &[u8],
3251 scales: &[u8],
3252 g0: usize,
3253 gpr: usize,
3254 xs: [&[i8]; 4],
3255) -> [f32; 4] {
3256 unsafe {
3258 use core::arch::x86_64::*;
3259 let ones = _mm256_set1_epi16(1);
3260 let mut acc = [0f32; 4];
3261 for gi in 0..gpr {
3262 let s = f16_to_f32(u16::from_le_bytes([
3263 scales[(g0 + gi) * 2],
3264 scales[(g0 + gi) * 2 + 1],
3265 ]));
3266 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3267 let aw = _mm256_abs_epi8(w);
3268 for (k, xq) in xs.iter().enumerate() {
3269 let x =
3270 _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3271 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
3272 let d = _mm256_madd_epi16(p16, ones);
3273 let hi128 = _mm256_extracti128_si256::<1>(d);
3274 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3275 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3276 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3277 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
3278 }
3279 }
3280 acc
3281 }
3282}
3283
3284#[cfg(target_arch = "x86_64")]
3290#[target_feature(enable = "avx2")]
3291unsafe fn dot_q4b_row_1x4_sx_avx2(
3292 buf: &[u8],
3293 scales: &[u8],
3294 g0: usize,
3295 gpr: usize,
3296 xs: [&[i8]; 4],
3297 sxs: [f32; 4],
3298) -> [f32; 4] {
3299 unsafe {
3301 use core::arch::x86_64::*;
3302 let ones = _mm256_set1_epi16(1);
3303 let mut acc = [0f32; 4];
3304 for gi in 0..gpr {
3305 let s = f16_to_f32(u16::from_le_bytes([
3306 scales[(g0 + gi) * 2],
3307 scales[(g0 + gi) * 2 + 1],
3308 ]));
3309 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3310 let aw = _mm256_abs_epi8(w);
3311 for (k, xq) in xs.iter().enumerate() {
3312 let x =
3313 _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3314 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
3315 let d = _mm256_madd_epi16(p16, ones);
3316 let hi128 = _mm256_extracti128_si256::<1>(d);
3317 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3318 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3319 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3320 acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
3321 }
3322 }
3323 acc
3324 }
3325}
3326
3327fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
3328 #[cfg(target_arch = "aarch64")]
3329 unsafe {
3330 return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
3331 }
3332 #[cfg(target_arch = "x86_64")]
3333 unsafe {
3334 return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
3335 }
3336 let mut acc = 0f32;
3337 for gi in 0..gpr {
3338 let g = g0 + gi;
3339 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
3340 let mut d = 0i32;
3341 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
3342 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
3343 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
3344 }
3345 acc += d as f32 * s;
3346 }
3347 acc
3348}
3349
3350#[inline]
3352#[allow(unreachable_code)]
3353fn dot_q4_row_i8_2(
3354 packed: &[u8],
3355 scales: &[u8],
3356 g0: usize,
3357 gpr: usize,
3358 xq1: &[i8],
3359 xq2: &[i8],
3360) -> (f32, f32) {
3361 #[cfg(target_arch = "aarch64")]
3362 unsafe {
3363 return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
3364 }
3365 #[cfg(target_arch = "x86_64")]
3366 unsafe {
3367 return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
3368 }
3369 (
3370 dot_q4_row_i8(packed, scales, g0, gpr, xq1),
3371 dot_q4_row_i8(packed, scales, g0, gpr, xq2),
3372 )
3373}
3374
3375#[allow(clippy::too_many_arguments)]
3378fn q4_range_a8w8(
3379 packed: &[u8],
3380 scales: &[u8],
3381 gpr: usize,
3382 cols: usize,
3383 act: &SplitAct,
3384 out: SendMut,
3385 start: usize,
3386 end: usize,
3387) {
3388 for r in start..end {
3389 let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
3390 for &(j, xv) in &act.outliers {
3392 let flat = r * cols + j;
3393 let byte = packed[flat / 2];
3394 let nib = if flat & 1 == 0 { byte & 0x0F } else { byte >> 4 };
3395 let s = f16_to_f32(u16::from_le_bytes([
3396 scales[(flat / GROUP_SIZE) * 2],
3397 scales[(flat / GROUP_SIZE) * 2 + 1],
3398 ]));
3399 acc += ((nib as i32 - 8) as f32) * s * xv;
3400 }
3401 unsafe { *out.at(r) = acc };
3403 }
3404}
3405
3406#[allow(clippy::too_many_arguments)]
3409fn q4_range2_a8w8(
3410 packed: &[u8],
3411 scales: &[u8],
3412 gpr: usize,
3413 cols: usize,
3414 a1: &SplitAct,
3415 a2: &SplitAct,
3416 p1: SendMut,
3417 p2: SendMut,
3418 start: usize,
3419 end: usize,
3420) {
3421 for r in start..end {
3422 let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
3423 let mut acc1 = s1 * a1.sx;
3424 let mut acc2 = s2 * a2.sx;
3425 let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
3427 for &(j, xv) in outliers {
3428 let flat = r * cols + j;
3429 let byte = packed[flat / 2];
3430 let nib = if flat & 1 == 0 { byte & 0x0F } else { byte >> 4 };
3431 let s = f16_to_f32(u16::from_le_bytes([
3432 scales[(flat / GROUP_SIZE) * 2],
3433 scales[(flat / GROUP_SIZE) * 2 + 1],
3434 ]));
3435 *acc += ((nib as i32 - 8) as f32) * s * xv;
3436 }
3437 };
3438 fix(&a1.outliers, &mut acc1);
3439 fix(&a2.outliers, &mut acc2);
3440 unsafe {
3442 *p1.at(r) = acc1;
3443 *p2.at(r) = acc2;
3444 }
3445 }
3446}
3447
3448fn q4_range_f32(
3450 packed: &[u8],
3451 scales: &[u8],
3452 gpr: usize,
3453 x: &[f32],
3454 out: SendMut,
3455 start: usize,
3456 end: usize,
3457) {
3458 for r in start..end {
3459 let mut acc = 0f32;
3460 for gi in 0..gpr {
3461 let g = r * gpr + gi;
3462 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
3463 let pk = &packed[g * 16..(g + 1) * 16];
3464 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
3465 let mut ga = 0f32;
3466 for (k, &b) in pk.iter().enumerate() {
3467 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
3468 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
3469 }
3470 acc += ga * s;
3471 }
3472 unsafe { *out.at(r) = acc };
3474 }
3475}
3476
3477#[allow(clippy::too_many_arguments)]
3481fn q4matvec2(
3482 bytes: &[u8],
3483 x1: &[f32],
3484 x2: &[f32],
3485 rows: usize,
3486 cols: usize,
3487 o1: &mut [f32],
3488 o2: &mut [f32],
3489 pool: Option<&Pool>,
3490) {
3491 debug_assert_eq!(o1.len(), rows);
3492 debug_assert_eq!(o2.len(), rows);
3493 let (packed, scales) = q4_split(bytes, rows, cols);
3494 let gpr = cols / GROUP_SIZE;
3495
3496 if a8w8_enabled() {
3497 let a1 = split_act(x1);
3498 let a2 = split_act(x2);
3499 let p1 = SendMut(o1.as_mut_ptr());
3500 let p2 = SendMut(o2.as_mut_ptr());
3501 let run = move |start: usize, end: usize| {
3502 q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
3503 };
3504 dispatch_rows(pool, rows, &run);
3505 return;
3506 }
3507
3508 let p1 = SendMut(o1.as_mut_ptr());
3509 let p2 = SendMut(o2.as_mut_ptr());
3510 let run = move |start: usize, end: usize| {
3511 q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
3512 };
3513 dispatch_rows(pool, rows, &run);
3514}
3515
3516#[allow(clippy::too_many_arguments)]
3518fn q4_range2_f32(
3519 packed: &[u8],
3520 scales: &[u8],
3521 gpr: usize,
3522 x1: &[f32],
3523 x2: &[f32],
3524 p1: SendMut,
3525 p2: SendMut,
3526 start: usize,
3527 end: usize,
3528) {
3529 for r in start..end {
3530 let (mut acc1, mut acc2) = (0f32, 0f32);
3531 for gi in 0..gpr {
3532 let g = r * gpr + gi;
3533 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
3534 let pk = &packed[g * 16..(g + 1) * 16];
3535 let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
3536 let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
3537 let (mut g1, mut g2) = (0f32, 0f32);
3538 for (k, &b) in pk.iter().enumerate() {
3539 let wl = (b & 0x0F) as f32 - 8.0;
3540 let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
3541 g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
3542 g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
3543 }
3544 acc1 += g1 * s;
3545 acc2 += g2 * s;
3546 }
3547 unsafe {
3549 *p1.at(r) = acc1;
3550 *p2.at(r) = acc2;
3551 }
3552 }
3553}
3554
3555thread_local! {
3556 static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
3559 static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3560}
3561
3562#[allow(clippy::too_many_arguments)]
3568fn q4matmat(
3569 bytes: &[u8],
3570 xs_all: &[f32],
3571 b: usize,
3572 rows: usize,
3573 cols: usize,
3574 out: &mut [f32],
3575 pool: Option<&Pool>,
3576) {
3577 debug_assert_eq!(xs_all.len(), b * cols);
3578 debug_assert_eq!(out.len(), b * rows);
3579 let (packed, scales) = q4_split(bytes, rows, cols);
3580 let gpr = cols / GROUP_SIZE;
3581 let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
3582
3583 if a8w8_enabled() {
3584 let acts: Vec<SplitAct> =
3585 (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
3586 let acts = &acts;
3587 let out_addr = SendMut(out.as_mut_ptr());
3588 let run = move |start: usize, end: usize| {
3589 ROW_I8.with(|rb| {
3590 let mut buf = rb.borrow_mut();
3591 buf.resize(cols, 0);
3592 for r in start..end {
3593 for gi in 0..gpr {
3597 let g = r * gpr + gi;
3598 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
3599 buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
3600 buf[gi * GROUP_SIZE + k * 2 + 1] =
3601 (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
3602 }
3603 }
3604 let mut bi = 0usize;
3605 #[cfg(target_arch = "x86_64")]
3606 if avx2_enabled()
3607 && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true)
3608 {
3609 while bi + 4 <= acts.len() {
3610 let xs = [
3611 acts[bi].xq.as_slice(),
3612 acts[bi + 1].xq.as_slice(),
3613 acts[bi + 2].xq.as_slice(),
3614 acts[bi + 3].xq.as_slice(),
3615 ];
3616 let d = unsafe {
3617 dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
3618 };
3619 for k in 0..4 {
3620 let act = &acts[bi + k];
3621 let mut acc = d[k] * act.sx;
3622 for &(j, xv) in &act.outliers {
3623 acc += (buf[j] as i8) as f32
3624 * gscale((r * cols + j) / GROUP_SIZE)
3625 * xv;
3626 }
3627 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
3629 }
3630 bi += 4;
3631 }
3632 }
3633 while bi < acts.len() {
3634 let act = &acts[bi];
3635 let mut acc = 0f32;
3636 for gi in 0..gpr {
3637 let d = dot_i8_i8(
3638 &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
3639 &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
3640 );
3641 acc += d as f32 * gscale(r * gpr + gi);
3642 }
3643 acc *= act.sx;
3644 for &(j, xv) in &act.outliers {
3646 acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
3647 }
3648 unsafe { *out_addr.at(bi * rows + r) = acc };
3650 bi += 1;
3651 }
3652 }
3653 })
3654 };
3655 dispatch_rows(pool, rows, &run);
3656 return;
3657 }
3658
3659 let out_addr = SendMut(out.as_mut_ptr());
3660 let run = move |start: usize, end: usize| {
3661 ROW_F32.with(|rb| {
3662 let mut buf = rb.borrow_mut();
3663 buf.resize(cols, 0.0);
3664 for r in start..end {
3665 for gi in 0..gpr {
3668 let g = r * gpr + gi;
3669 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
3670 buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
3671 buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
3672 }
3673 }
3674 for bi in 0..b {
3675 let x = &xs_all[bi * cols..(bi + 1) * cols];
3676 let mut acc = 0f32;
3677 for gi in 0..gpr {
3678 let mut ga = 0f32;
3679 for k in 0..GROUP_SIZE / 2 {
3684 let e = gi * GROUP_SIZE + k * 2;
3685 ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
3686 }
3687 acc += ga * gscale(r * gpr + gi);
3688 }
3689 unsafe { *out_addr.at(bi * rows + r) = acc };
3691 }
3692 }
3693 })
3694 };
3695 dispatch_rows(pool, rows, &run);
3696}
3697
3698#[allow(clippy::too_many_arguments)]
3703fn vbitmatmat(
3704 bytes: &[u8],
3705 offsets: &[usize],
3706 xs_all: &[f32],
3707 b: usize,
3708 rows: usize,
3709 cols: usize,
3710 out: &mut [f32],
3711 pool: Option<&Pool>,
3712) {
3713 debug_assert_eq!(xs_all.len(), b * cols);
3714 debug_assert_eq!(out.len(), b * rows);
3715 debug_assert_eq!(offsets.len(), rows + 1);
3716 let ng = cols / GROUP_SIZE;
3717 let bits = &bytes[..rows];
3718 let sc_off = rows;
3719 let gscale = |r: usize, g: usize| {
3720 let so = (r * ng + g) * 2;
3721 f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]))
3722 };
3723
3724 let decode_f32 = |r: usize, dst: &mut [f32]| {
3726 let bw = bits[r] as usize;
3727 let l = ((1i32 << (bw - 1)) - 1) as f32;
3728 let data = &bytes[offsets[r]..offsets[r + 1]];
3729 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3730 for d in dst.iter_mut() {
3731 while nbits < bw {
3732 acc = (acc << 8) | data[idx] as u64;
3733 idx += 1;
3734 nbits += 8;
3735 }
3736 let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
3737 nbits -= bw;
3738 *d = u - l;
3739 }
3740 };
3741
3742 if a8w8_enabled() {
3743 let acts: Vec<SplitAct> =
3744 (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
3745 let acts = &acts;
3746 let out_addr = SendMut(out.as_mut_ptr());
3747 let run = move |start: usize, end: usize| {
3748 for r in start..end {
3749 let bw = bits[r] as usize;
3750 if bw == 8 {
3751 ROW_F32.with(|rb| {
3754 let mut buf = rb.borrow_mut();
3755 buf.resize(cols, 0.0);
3756 decode_f32(r, &mut buf);
3757 for bi in 0..b {
3758 let x = &xs_all[bi * cols..(bi + 1) * cols];
3759 let mut dot = 0f32;
3760 for g in 0..ng {
3761 let mut gd = 0f32;
3762 for k in 0..GROUP_SIZE {
3763 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
3764 }
3765 dot += gd * gscale(r, g);
3766 }
3767 unsafe { *out_addr.at(bi * rows + r) = dot };
3769 }
3770 });
3771 continue;
3772 }
3773 let l = (1i32 << (bw - 1)) - 1;
3774 let data = &bytes[offsets[r]..offsets[r + 1]];
3775 ROW_I8.with(|rb| {
3776 let mut buf = rb.borrow_mut();
3777 buf.resize(cols, 0);
3778 #[inline(always)]
3779 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3780 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3781 let u = unpack8::<B>(&data[blk * B..]);
3782 for k in 0..8 {
3783 chunk[k] = (u[k] - l) as i8 as u8;
3784 }
3785 }
3786 }
3787 match bw {
3788 3 => fill::<3>(data, l, &mut buf),
3789 4 => vbit_fill4(data, &mut buf),
3790 5 => fill::<5>(data, l, &mut buf),
3791 6 => fill::<6>(data, l, &mut buf),
3792 _ => unreachable!("vbit bit-width {bw} (validated at load)"),
3793 }
3794 let mut bi = 0usize;
3795 #[cfg(target_arch = "x86_64")]
3799 if avx2_enabled()
3800 && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true)
3801 {
3802 while bi + 4 <= acts.len() {
3803 let xs = [
3804 acts[bi].xq.as_slice(),
3805 acts[bi + 1].xq.as_slice(),
3806 acts[bi + 2].xq.as_slice(),
3807 acts[bi + 3].xq.as_slice(),
3808 ];
3809 let sxs = [
3810 acts[bi].sx,
3811 acts[bi + 1].sx,
3812 acts[bi + 2].sx,
3813 acts[bi + 3].sx,
3814 ];
3815 let d = unsafe {
3816 dot_q4b_row_1x4_sx_avx2(&buf, &bytes[sc_off..], r * ng, ng, xs, sxs)
3817 };
3818 for k in 0..4 {
3819 let act = &acts[bi + k];
3820 let mut dot = d[k];
3821 for &(j, xv) in &act.outliers {
3822 dot +=
3823 (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
3824 }
3825 unsafe { *out_addr.at((bi + k) * rows + r) = dot };
3827 }
3828 bi += 4;
3829 }
3830 }
3831 while bi < acts.len() {
3832 let act = &acts[bi];
3833 let mut dot = 0f32;
3834 for g in 0..ng {
3835 let d = dot_i8_i8(
3836 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3837 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3838 ) as f32
3839 * act.sx;
3840 dot += d * gscale(r, g);
3841 }
3842 for &(j, xv) in &act.outliers {
3843 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
3844 }
3845 unsafe { *out_addr.at(bi * rows + r) = dot };
3847 bi += 1;
3848 }
3849 });
3850 }
3851 };
3852 dispatch_rows(pool, rows, &run);
3853 return;
3854 }
3855
3856 let out_addr = SendMut(out.as_mut_ptr());
3857 let run = move |start: usize, end: usize| {
3858 ROW_F32.with(|rb| {
3859 let mut buf = rb.borrow_mut();
3860 buf.resize(cols, 0.0);
3861 for r in start..end {
3862 decode_f32(r, &mut buf);
3863 for bi in 0..b {
3864 let x = &xs_all[bi * cols..(bi + 1) * cols];
3865 let mut dot = 0f32;
3866 for g in 0..ng {
3867 let mut gd = 0f32;
3868 for k in 0..GROUP_SIZE {
3869 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
3870 }
3871 dot += gd * gscale(r, g);
3872 }
3873 unsafe { *out_addr.at(bi * rows + r) = dot };
3875 }
3876 }
3877 })
3878 };
3879 dispatch_rows(pool, rows, &run);
3880}
3881
3882pub(crate) fn gpu_batch_job<'a>(
3886 t: &'a QTensor,
3887 x: &[f32],
3888) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
3889 match t {
3890 QTensor::Mapped {
3891 model,
3892 idx,
3893 dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
3894 rows,
3895 cols,
3896 row_scale,
3897 col_field,
3898 ..
3899 } => Some((
3900 model.clone(),
3901 crate::gpu::BatchJob {
3902 idx: *idx,
3903 rows: *rows,
3904 cols: *cols,
3905 row_scale,
3906 xs: prescale(x, col_field, *dt).into_owned(),
3907 q1: false,
3908 },
3909 )),
3910 QTensor::Mapped {
3912 model,
3913 idx,
3914 dtype: TensorDtype::Q1,
3915 rows,
3916 cols,
3917 ..
3918 } => Some((
3919 model.clone(),
3920 crate::gpu::BatchJob {
3921 idx: *idx,
3922 rows: *rows,
3923 cols: *cols,
3924 row_scale: &[],
3925 xs: x.to_vec(),
3926 q1: true,
3927 },
3928 )),
3929 _ => None,
3930 }
3931}
3932
3933pub(crate) fn prescale<'a>(
3937 x: &'a [f32],
3938 col_field: &[f32],
3939 dtype: TensorDtype,
3940) -> std::borrow::Cow<'a, [f32]> {
3941 if dtype == TensorDtype::Q8_2f {
3942 x.iter().zip(col_field).map(|(a, c)| a * c).collect()
3943 } else {
3944 std::borrow::Cow::Borrowed(x)
3945 }
3946}
3947
3948#[cfg(target_arch = "x86_64")]
3953pub(crate) fn avx2_enabled() -> bool {
3954 use std::sync::OnceLock;
3955 static ON: OnceLock<bool> = OnceLock::new();
3956 *ON.get_or_init(|| {
3957 std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
3958 && std::arch::is_x86_feature_detected!("avx2")
3959 && std::arch::is_x86_feature_detected!("fma")
3960 })
3961}
3962
3963#[cfg(target_arch = "x86_64")]
3968fn avx2_a8w8_enabled() -> bool {
3969 use std::sync::OnceLock;
3970 static ON: OnceLock<bool> = OnceLock::new();
3971 *ON.get_or_init(|| {
3972 avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
3973 })
3974}
3975
3976#[inline]
3980pub(crate) fn a8w8_enabled() -> bool {
3981 #[cfg(target_arch = "aarch64")]
3982 {
3983 sdot_enabled()
3984 }
3985 #[cfg(target_arch = "x86_64")]
3986 {
3987 avx2_a8w8_enabled()
3988 }
3989 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
3990 {
3991 false
3992 }
3993}
3994
3995#[inline]
3998#[allow(unreachable_code)]
3999fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
4000 #[cfg(target_arch = "aarch64")]
4001 unsafe {
4002 return dot_i8_sdot(w, xq);
4003 }
4004 #[cfg(target_arch = "x86_64")]
4005 unsafe {
4006 if avx512vnni_enabled() {
4007 return dot_i8_i8_vnni(w, xq);
4008 }
4009 return dot_i8_i8_avx2(w, xq);
4010 }
4011 w.iter().zip(xq).map(|(&a, &b)| (a as i8) as i32 * b as i32).sum()
4012}
4013
4014#[cfg(target_arch = "x86_64")]
4018fn avx512vnni_enabled() -> bool {
4019 use std::sync::OnceLock;
4020 static ON: OnceLock<bool> = OnceLock::new();
4021 *ON.get_or_init(|| {
4022 std::env::var("CMF_AVX512").map(|v| v != "0").unwrap_or(true)
4023 && std::arch::is_x86_feature_detected!("avx512f")
4024 && std::arch::is_x86_feature_detected!("avx512bw")
4025 && std::arch::is_x86_feature_detected!("avx512vl")
4026 && std::arch::is_x86_feature_detected!("avx512vnni")
4027 })
4028}
4029
4030#[cfg(target_arch = "x86_64")]
4035#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4036unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
4037 unsafe {
4039 use core::arch::x86_64::*;
4040 let n = w.len();
4041 let mut j = 0usize;
4042 let mut total: i32;
4043 {
4048 #[inline(always)]
4049 unsafe fn step(
4050 w: *const u8,
4051 x: *const i8,
4052 acc: core::arch::x86_64::__m512i,
4053 ) -> core::arch::x86_64::__m512i {
4054 unsafe {
4055 use core::arch::x86_64::*;
4056 let wv = _mm512_loadu_si512(w as *const _);
4057 let xv = _mm512_loadu_si512(x as *const _);
4058 let aw = _mm512_abs_epi8(wv);
4059 let neg = _mm512_movepi8_mask(wv);
4060 let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
4061 _mm512_dpbusd_epi32(acc, aw, sx)
4062 }
4063 }
4064 let (mut a0, mut a1, mut a2, mut a3) = (
4065 _mm512_setzero_si512(),
4066 _mm512_setzero_si512(),
4067 _mm512_setzero_si512(),
4068 _mm512_setzero_si512(),
4069 );
4070 while j + 256 <= n {
4071 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
4072 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
4073 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
4074 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
4075 j += 256;
4076 }
4077 while j + 64 <= n {
4078 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
4079 j += 64;
4080 }
4081 let s01 = _mm512_add_epi32(a0, a1);
4082 let s23 = _mm512_add_epi32(a2, a3);
4083 total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
4084 }
4085 if j + 32 <= n {
4087 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
4088 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
4089 let d = _mm256_dpbusd_epi32(
4090 _mm256_setzero_si256(),
4091 _mm256_abs_epi8(wv),
4092 _mm256_sign_epi8(xv, wv),
4093 );
4094 let hi128 = _mm256_extracti128_si256::<1>(d);
4095 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4096 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4097 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4098 total += _mm_cvtsi128_si32(s32);
4099 j += 32;
4100 }
4101 while j < n {
4102 total += (w[j] as i8) as i32 * xq[j] as i32;
4103 j += 1;
4104 }
4105 total
4106 }
4107}
4108
4109#[cfg(target_arch = "x86_64")]
4111#[target_feature(enable = "avx2,fma")]
4112unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
4113 unsafe {
4115 use core::arch::x86_64::*;
4116 let n = x.len();
4117 let wp = w.as_ptr();
4118 let xp = x.as_ptr();
4119 let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
4120 let mut j = 0usize;
4121 while j + 16 <= n {
4122 let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
4123 let lo = _mm256_cvtepi8_epi32(wb);
4124 let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
4125 a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
4126 a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
4127 j += 16;
4128 }
4129 let acc = _mm256_add_ps(a0, a1);
4130 let hi128 = _mm256_extractf128_ps::<1>(acc);
4131 let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
4132 let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
4133 let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
4134 let mut sum = _mm_cvtss_f32(s32);
4135 while j < n {
4136 sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
4137 j += 1;
4138 }
4139 sum
4140 }
4141}
4142
4143#[cfg(target_arch = "x86_64")]
4148#[target_feature(enable = "avx2")]
4149unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
4150 unsafe {
4152 use core::arch::x86_64::*;
4153 let n = w.len();
4154 let ones = _mm256_set1_epi16(1);
4155 let mut acc = _mm256_setzero_si256();
4156 let mut j = 0usize;
4157 while j + 32 <= n {
4158 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
4159 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
4160 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
4161 acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
4162 j += 32;
4163 }
4164 let hi128 = _mm256_extracti128_si256::<1>(acc);
4165 let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
4166 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4167 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4168 let mut s = _mm_cvtsi128_si32(s32);
4169 while j < n {
4170 s += (w[j] as i8) as i32 * xq[j] as i32;
4171 j += 1;
4172 }
4173 s
4174 }
4175}
4176
4177#[cfg(target_arch = "aarch64")]
4181#[target_feature(enable = "neon,i8mm")]
4182unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
4183 unsafe {
4185 use core::arch::aarch64::*;
4186 use core::arch::asm;
4187 let n = w0.len();
4188 let w0p = w0.as_ptr() as *const i8;
4189 let w1p = w1.as_ptr() as *const i8;
4190 let mut acc01 = vdupq_n_s32(0);
4193 let mut acc23 = vdupq_n_s32(0);
4194 let mut i = 0usize;
4195 while i + 8 <= n {
4196 let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
4197 let xb01 =
4198 vcombine_s8(vld1_s8(xs[0].as_ptr().add(i)), vld1_s8(xs[1].as_ptr().add(i)));
4199 let xb23 =
4200 vcombine_s8(vld1_s8(xs[2].as_ptr().add(i)), vld1_s8(xs[3].as_ptr().add(i)));
4201 asm!(
4202 "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
4203 "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
4204 a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
4205 w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
4206 options(pure, nomem, nostack),
4207 );
4208 i += 8;
4209 }
4210 let mut out = [[0i32; 4]; 2];
4211 let a01: [i32; 4] = core::mem::transmute(acc01);
4212 let a23: [i32; 4] = core::mem::transmute(acc23);
4213 out[0][0] = a01[0];
4214 out[0][1] = a01[1];
4215 out[1][0] = a01[2];
4216 out[1][1] = a01[3];
4217 out[0][2] = a23[0];
4218 out[0][3] = a23[1];
4219 out[1][2] = a23[2];
4220 out[1][3] = a23[3];
4221 if i < n {
4222 for (k, x) in xs.iter().enumerate() {
4223 for j in i..n {
4224 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
4225 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
4226 }
4227 }
4228 }
4229 out
4230 }
4231}
4232
4233#[cfg(target_arch = "aarch64")]
4237#[target_feature(enable = "neon,dotprod")]
4238unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
4239 unsafe {
4241 use core::arch::aarch64::*;
4242 use core::arch::asm;
4243 let n = w0.len();
4244 let w0p = w0.as_ptr() as *const i8;
4245 let w1p = w1.as_ptr() as *const i8;
4246 let mut acc = [[vdupq_n_s32(0); 4]; 2];
4247 let mut i = 0usize;
4248 while i + 16 <= n {
4249 let wv0 = vld1q_s8(w0p.add(i));
4250 let wv1 = vld1q_s8(w1p.add(i));
4251 for (k, x) in xs.iter().enumerate() {
4252 let xv = vld1q_s8(x.as_ptr().add(i));
4253 let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
4254 asm!(
4255 "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
4256 "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
4257 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4258 w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
4259 options(pure, nomem, nostack),
4260 );
4261 acc[0][k] = a0;
4262 acc[1][k] = a1;
4263 }
4264 i += 16;
4265 }
4266 let mut out = [[0i32; 4]; 2];
4267 for r in 0..2 {
4268 for k in 0..4 {
4269 out[r][k] = vaddvq_s32(acc[r][k]);
4270 }
4271 }
4272 if i < n {
4273 for (k, x) in xs.iter().enumerate() {
4274 for j in i..n {
4275 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
4276 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
4277 }
4278 }
4279 }
4280 out
4281 }
4282}
4283
4284#[cfg(target_arch = "x86_64")]
4290#[target_feature(enable = "avx2")]
4291unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
4292 unsafe {
4294 use core::arch::x86_64::*;
4295 let n = w0.len();
4296 let ones = _mm256_set1_epi16(1);
4297 let mut acc = [[_mm256_setzero_si256(); 4]; 2];
4298 let mut j = 0usize;
4299 while j + 32 <= n {
4300 let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
4301 let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
4302 let aw0 = _mm256_abs_epi8(wv0);
4303 let aw1 = _mm256_abs_epi8(wv1);
4304 for (k, x) in xs.iter().enumerate() {
4305 let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
4306 let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
4307 acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
4308 let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
4309 acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
4310 }
4311 j += 32;
4312 }
4313 let mut out = [[0i32; 4]; 2];
4314 for r in 0..2 {
4315 for k in 0..4 {
4316 let a = acc[r][k];
4317 let hi128 = _mm256_extracti128_si256::<1>(a);
4318 let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
4319 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4320 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4321 out[r][k] = _mm_cvtsi128_si32(s32);
4322 }
4323 }
4324 if j < n {
4325 for (k, x) in xs.iter().enumerate() {
4326 for i in j..n {
4327 out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
4328 out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
4329 }
4330 }
4331 }
4332 out
4333 }
4334}
4335
4336#[cfg(target_arch = "x86_64")]
4341#[inline]
4342fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
4343 let dot = if avx512vnni_enabled() && row.len() >= 64 {
4344 (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
4345 } else {
4346 unsafe { dot_i8_i8_avx2(row, &act.xq) }
4347 };
4348 let mut acc = dot as f32 * act.sx;
4349 for &(j, xv) in &act.outliers {
4350 acc += (row[j] as i8) as f32 * xv;
4351 }
4352 acc
4353}
4354
4355#[cfg(target_arch = "x86_64")]
4359#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4360unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
4361 unsafe {
4363 use core::arch::x86_64::*;
4364 let n = w.len();
4365 let flip = _mm512_set1_epi8(-128); #[inline(always)]
4367 unsafe fn step(
4368 w: *const u8,
4369 x: *const i8,
4370 flip: core::arch::x86_64::__m512i,
4371 acc: core::arch::x86_64::__m512i,
4372 ) -> core::arch::x86_64::__m512i {
4373 unsafe {
4374 use core::arch::x86_64::*;
4375 let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
4376 _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
4377 }
4378 }
4379 let (mut a0, mut a1, mut a2, mut a3) = (
4380 _mm512_setzero_si512(),
4381 _mm512_setzero_si512(),
4382 _mm512_setzero_si512(),
4383 _mm512_setzero_si512(),
4384 );
4385 let mut j = 0usize;
4386 while j + 256 <= n {
4387 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
4388 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
4389 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
4390 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
4391 j += 256;
4392 }
4393 while j + 64 <= n {
4394 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
4395 j += 64;
4396 }
4397 let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
4398 _mm512_add_epi32(a0, a1),
4399 _mm512_add_epi32(a2, a3),
4400 ));
4401 while j < n {
4403 total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
4404 j += 1;
4405 }
4406 total
4407 }
4408}
4409
4410#[cfg(target_arch = "x86_64")]
4416#[target_feature(enable = "avx2")]
4417unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
4418 unsafe {
4421 use core::arch::x86_64::*;
4422 let lomask = _mm_set1_epi8(0x0F);
4423 let eight = _mm256_set1_epi8(8);
4424 let ones = _mm256_set1_epi16(1);
4425 let mut acc = 0f32;
4426 for gi in 0..gpr {
4427 let g = g0 + gi;
4428 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4429 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
4430 let lo = _mm_and_si128(b, lomask);
4431 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4432 let w = _mm256_sub_epi8(
4433 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4434 eight,
4435 );
4436 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4437 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4438 let d = _mm256_madd_epi16(p16, ones);
4439 let hi128 = _mm256_extracti128_si256::<1>(d);
4440 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4441 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4442 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4443 acc += _mm_cvtsi128_si32(s32) as f32 * s;
4444 }
4445 acc
4446 }
4447}
4448
4449#[cfg(target_arch = "x86_64")]
4452#[target_feature(enable = "avx2")]
4453unsafe fn dot_q4_row_avx2_2(
4454 packed: &[u8],
4455 scales: &[u8],
4456 g0: usize,
4457 gpr: usize,
4458 xq1: &[i8],
4459 xq2: &[i8],
4460) -> (f32, f32) {
4461 unsafe {
4463 use core::arch::x86_64::*;
4464 let lomask = _mm_set1_epi8(0x0F);
4465 let eight = _mm256_set1_epi8(8);
4466 let ones = _mm256_set1_epi16(1);
4467 let (mut acc1, mut acc2) = (0f32, 0f32);
4468 #[inline(always)]
4469 unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
4470 unsafe {
4471 use core::arch::x86_64::*;
4472 let hi128 = _mm256_extracti128_si256::<1>(d);
4473 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4474 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4475 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4476 _mm_cvtsi128_si32(s32)
4477 }
4478 }
4479 for gi in 0..gpr {
4480 let g = g0 + gi;
4481 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4482 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
4483 let lo = _mm_and_si128(b, lomask);
4484 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4485 let w = _mm256_sub_epi8(
4486 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4487 eight,
4488 );
4489 let aw = _mm256_abs_epi8(w);
4490 let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4491 let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4492 let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
4493 let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
4494 acc1 += hsum(d1) as f32 * s;
4495 acc2 += hsum(d2) as f32 * s;
4496 }
4497 (acc1, acc2)
4498 }
4499}
4500
4501#[cfg(target_arch = "x86_64")]
4503fn q8_range_avx2(
4504 q: &[u8],
4505 row_scale: &[f32],
4506 act: &SplitAct,
4507 cols: usize,
4508 out_addr: SendMut,
4509 start: usize,
4510 end: usize,
4511) {
4512 for o in start..end {
4513 let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
4514 unsafe { *out_addr.at(o) = v };
4516 }
4517}
4518
4519#[cfg(target_arch = "x86_64")]
4521#[allow(clippy::too_many_arguments)]
4522fn q8_range2_avx2(
4523 q: &[u8],
4524 row_scale: &[f32],
4525 a1: &SplitAct,
4526 a2: &SplitAct,
4527 cols: usize,
4528 p1: SendMut,
4529 p2: SendMut,
4530 start: usize,
4531 end: usize,
4532) {
4533 for o in start..end {
4534 let row = &q[o * cols..(o + 1) * cols];
4535 unsafe {
4537 *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
4538 *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
4539 }
4540 }
4541}
4542
4543#[cfg(target_arch = "aarch64")]
4554fn i8mm_enabled() -> bool {
4555 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4556 *ON.get_or_init(|| {
4557 std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
4558 && std::arch::is_aarch64_feature_detected!("i8mm")
4559 })
4560}
4561
4562#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
4566fn sdot_enabled() -> bool {
4567 use std::sync::OnceLock;
4568 static ON: OnceLock<bool> = OnceLock::new();
4569 *ON.get_or_init(|| {
4570 let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
4571 #[cfg(target_arch = "aarch64")]
4572 {
4573 want && std::arch::is_aarch64_feature_detected!("dotprod")
4574 }
4575 #[cfg(not(target_arch = "aarch64"))]
4576 {
4577 let _ = want;
4578 false
4579 }
4580 })
4581}
4582
4583struct SplitAct {
4588 xq: Vec<i8>,
4589 sx: f32,
4590 outliers: Vec<(usize, f32)>,
4591 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
4594 xsum: i32,
4595}
4596
4597thread_local! {
4598 static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
4601 const { std::cell::RefCell::new(Vec::new()) };
4602}
4603
4604impl Drop for SplitAct {
4605 fn drop(&mut self) {
4606 let buf = std::mem::take(&mut self.xq);
4607 if buf.capacity() > 0 {
4608 XQ_FREE.with(|f| {
4609 let mut f = f.borrow_mut();
4610 if f.len() < 16 {
4611 f.push(buf);
4612 }
4613 });
4614 }
4615 }
4616}
4617
4618fn split_act(x: &[f32]) -> SplitAct {
4619 let n = x.len();
4620 let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
4621 let thr = 8.0 * rms;
4622 let mut outliers: Vec<(usize, f32)> = Vec::new();
4626 let mut amax = 0f32;
4627 for (j, &v) in x.iter().enumerate() {
4628 let a = v.abs();
4629 if a > thr {
4630 outliers.push((j, v));
4631 } else if a > amax {
4632 amax = a;
4633 }
4634 }
4635 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
4636 let inv = 1.0 / sx;
4637 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
4638 xq.clear();
4639 xq.reserve(n);
4640 if outliers.is_empty() {
4641 xq.extend(x.iter().map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8));
4642 } else {
4643 xq.extend(x.iter().map(|&v| {
4645 if v.abs() > thr {
4646 0
4647 } else {
4648 (v * inv).round().clamp(-127.0, 127.0) as i8
4649 }
4650 }));
4651 }
4652 let xsum = xq.iter().map(|&v| v as i32).sum();
4653 SplitAct { xq, sx, outliers, xsum }
4654}
4655
4656#[cfg(target_arch = "aarch64")]
4659#[target_feature(enable = "neon,dotprod")]
4660unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
4661 unsafe {
4663 use core::arch::aarch64::*;
4664 use core::arch::asm;
4665 let wp = w.as_ptr() as *const i8;
4666 let n = w.len();
4667 let (mut a0, mut a1, mut a2, mut a3) =
4668 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
4669 let mut i = 0;
4670 while i + 64 <= n {
4671 let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
4672 let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
4673 let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
4674 let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
4675 asm!(
4676 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
4677 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
4678 "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
4679 "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
4680 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
4681 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
4682 w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
4683 options(pure, nomem, nostack),
4684 );
4685 i += 64;
4686 }
4687 while i + 16 <= n {
4688 let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
4689 asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
4690 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
4691 i += 16;
4692 }
4693 let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
4694 while i < n {
4695 s += (*wp.add(i)) as i32 * xq[i] as i32;
4696 i += 1;
4697 }
4698 s
4699}
4700}
4701
4702#[cfg(target_arch = "aarch64")]
4706#[target_feature(enable = "neon,dotprod")]
4707unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
4708 unsafe {
4710 use core::arch::aarch64::*;
4711 use core::arch::asm;
4712 let n = xq.len();
4713 let px = xq.as_ptr();
4714 let (p0, p1, p2, p3) = (
4715 w0.as_ptr() as *const i8,
4716 w1.as_ptr() as *const i8,
4717 w2.as_ptr() as *const i8,
4718 w3.as_ptr() as *const i8,
4719 );
4720 let (mut a0, mut a1, mut a2, mut a3) =
4721 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
4722 let mut i = 0;
4723 while i + 16 <= n {
4724 let x = vld1q_s8(px.add(i));
4725 let v0 = vld1q_s8(p0.add(i));
4726 let v1 = vld1q_s8(p1.add(i));
4727 let v2 = vld1q_s8(p2.add(i));
4728 let v3 = vld1q_s8(p3.add(i));
4729 asm!(
4730 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
4731 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
4732 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
4733 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
4734 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
4735 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
4736 options(pure, nomem, nostack),
4737 );
4738 i += 16;
4739 }
4740 let mut r = [vaddvq_s32(a0), vaddvq_s32(a1), vaddvq_s32(a2), vaddvq_s32(a3)];
4741 while i < n {
4742 let xi = *px.add(i) as i32;
4743 r[0] += (*p0.add(i)) as i32 * xi;
4744 r[1] += (*p1.add(i)) as i32 * xi;
4745 r[2] += (*p2.add(i)) as i32 * xi;
4746 r[3] += (*p3.add(i)) as i32 * xi;
4747 i += 1;
4748 }
4749 r
4750}
4751}
4752
4753#[cfg(target_arch = "aarch64")]
4760#[target_feature(enable = "neon,dotprod")]
4761unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
4762 unsafe {
4765 use core::arch::aarch64::*;
4766 use core::arch::asm;
4767 let n = xq.len();
4768 let px = xq.as_ptr();
4769 let pg = g.as_ptr() as *const i8;
4770 let (mut a0, mut a1, mut a2, mut a3) =
4771 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
4772 let mut i = 0;
4773 while i + 16 <= n {
4774 let x = vld1q_s8(px.add(i));
4775 let base = pg.add(4 * i);
4776 let v0 = vld1q_s8(base);
4777 let v1 = vld1q_s8(base.add(16));
4778 let v2 = vld1q_s8(base.add(32));
4779 let v3 = vld1q_s8(base.add(48));
4780 asm!(
4781 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
4782 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
4783 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
4784 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
4785 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
4786 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
4787 options(pure, nomem, nostack),
4788 );
4789 i += 16;
4790 }
4791 [vaddvq_s32(a0), vaddvq_s32(a1), vaddvq_s32(a2), vaddvq_s32(a3)]
4792 }
4793}
4794
4795#[cfg(target_arch = "aarch64")]
4801fn q8_range_sdot(
4802 q: &[u8],
4803 rep: &[u8],
4804 row_scale: &[f32],
4805 act: &SplitAct,
4806 cols: usize,
4807 out_addr: SendMut,
4808 start: usize,
4809 end: usize,
4810) {
4811 let mut o = start;
4812 if !rep.is_empty() {
4815 while o < end && o % 4 != 0 {
4816 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
4817 unsafe { *out_addr.at(o) = v };
4818 o += 1;
4819 }
4820 }
4821 while o + 4 <= end {
4822 let r = if rep.is_empty() {
4823 unsafe {
4824 dot_i8_sdot_4rows(
4825 &q[o * cols..(o + 1) * cols],
4826 &q[(o + 1) * cols..(o + 2) * cols],
4827 &q[(o + 2) * cols..(o + 3) * cols],
4828 &q[(o + 3) * cols..(o + 4) * cols],
4829 &act.xq,
4830 )
4831 }
4832 } else {
4833 unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
4834 };
4835 for k in 0..4 {
4836 let mut acc = r[k] as f32 * act.sx;
4837 for &(j, xv) in &act.outliers {
4838 acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
4839 }
4840 unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
4842 }
4843 o += 4;
4844 }
4845 while o < end {
4846 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
4847 unsafe { *out_addr.at(o) = v };
4848 o += 1;
4849 }
4850}
4851
4852#[cfg(target_arch = "aarch64")]
4855#[allow(clippy::too_many_arguments)]
4856fn q8_range2_sdot(
4857 q: &[u8],
4858 row_scale: &[f32],
4859 a1: &SplitAct,
4860 a2: &SplitAct,
4861 cols: usize,
4862 p1: SendMut,
4863 p2: SendMut,
4864 start: usize,
4865 end: usize,
4866) {
4867 for o in start..end {
4868 let row = &q[o * cols..(o + 1) * cols];
4869 unsafe {
4871 *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
4872 *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
4873 }
4874 }
4875}
4876
4877#[allow(clippy::too_many_arguments)]
4879fn q8_range2_f32(
4880 q: &[u8],
4881 row_scale: &[f32],
4882 x1: &[f32],
4883 x2: &[f32],
4884 cols: usize,
4885 p1: SendMut,
4886 p2: SendMut,
4887 start: usize,
4888 end: usize,
4889) {
4890 for o in start..end {
4891 let row = &q[o * cols..(o + 1) * cols];
4892 unsafe {
4894 *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
4895 *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
4896 }
4897 }
4898}
4899
4900fn q8_range_f32(
4902 q: &[u8],
4903 row_scale: &[f32],
4904 xs: &[f32],
4905 cols: usize,
4906 out_addr: SendMut,
4907 start: usize,
4908 end: usize,
4909) {
4910 for o in start..end {
4911 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
4912 unsafe { *out_addr.at(o) = v };
4914 }
4915}
4916
4917#[cfg(target_arch = "aarch64")]
4920#[inline]
4921fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
4922 let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
4923 for &(j, xv) in &act.outliers {
4924 acc += (row[j] as i8) as f32 * xv;
4925 }
4926 acc
4927}
4928
4929#[cfg(target_arch = "aarch64")]
4937#[target_feature(enable = "neon,dotprod")]
4938unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
4939 unsafe {
4942 use core::arch::aarch64::*;
4943 use core::arch::asm;
4944 let lomask = vdupq_n_u8(0x0F);
4945 let eight = vdupq_n_s8(8);
4946 let mut acc = 0f32;
4947 for gi in 0..gpr {
4948 let g = g0 + gi;
4949 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4950 let b = vld1q_u8(packed.as_ptr().add(g * 16));
4951 let lo = vandq_u8(b, lomask);
4952 let hi = vshrq_n_u8::<4>(b);
4953 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4954 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4955 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4956 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4957 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4958 asm!(
4959 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4960 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4961 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4962 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4963 options(pure, nomem, nostack),
4964 );
4965 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4966 }
4967 acc
4968 }
4969}
4970
4971#[cfg(target_arch = "aarch64")]
4976#[target_feature(enable = "neon,dotprod")]
4977unsafe fn dot_q4_row_sdot2(
4978 packed: &[u8],
4979 scales: &[u8],
4980 g0: usize,
4981 gpr: usize,
4982 xq1: &[i8],
4983 xq2: &[i8],
4984) -> (f32, f32) {
4985 unsafe {
4988 use core::arch::aarch64::*;
4989 use core::arch::asm;
4990 let lomask = vdupq_n_u8(0x0F);
4991 let eight = vdupq_n_s8(8);
4992 let (mut acc1, mut acc2) = (0f32, 0f32);
4993 for gi in 0..gpr {
4994 let g = g0 + gi;
4995 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4996 let b = vld1q_u8(packed.as_ptr().add(g * 16));
4997 let lo = vandq_u8(b, lomask);
4998 let hi = vshrq_n_u8::<4>(b);
4999 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5000 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5001 let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
5002 let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
5003 let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
5004 let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
5005 let (mut a0, mut a1, mut b0, mut b1) =
5006 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
5007 asm!(
5008 "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
5009 "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
5010 "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
5011 "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
5012 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5013 b0 = inout(vreg) b0, b1 = inout(vreg) b1,
5014 e0 = in(vreg) e0, e1 = in(vreg) e1,
5015 x10 = in(vreg) x10, x11 = in(vreg) x11,
5016 x20 = in(vreg) x20, x21 = in(vreg) x21,
5017 options(pure, nomem, nostack),
5018 );
5019 acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5020 acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
5021 }
5022 (acc1, acc2)
5023 }
5024}
5025
5026#[inline]
5031pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
5032 #[cfg(target_arch = "aarch64")]
5033 unsafe {
5034 return axpy_i8_f32_neon(acc, row, w);
5035 }
5036 #[cfg(target_arch = "x86_64")]
5037 if avx2_enabled() {
5038 return unsafe { axpy_i8_f32_avx2(acc, row, w) };
5039 }
5040 #[allow(unreachable_code)]
5041 {
5042 for (a, &b) in acc.iter_mut().zip(row) {
5043 *a += w * b as f32;
5044 }
5045 }
5046}
5047
5048#[cfg(target_arch = "x86_64")]
5050#[target_feature(enable = "avx2,fma")]
5051unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
5052 unsafe {
5054 use core::arch::x86_64::*;
5055 let n = acc.len().min(row.len());
5056 let ap = acc.as_mut_ptr();
5057 let rp = row.as_ptr();
5058 let wv = _mm256_set1_ps(w);
5059 let mut j = 0usize;
5060 while j + 16 <= n {
5061 let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
5062 let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
5063 let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
5064 let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
5065 let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
5066 _mm256_storeu_ps(ap.add(j), v0);
5067 _mm256_storeu_ps(ap.add(j + 8), v1);
5068 j += 16;
5069 }
5070 while j < n {
5071 *ap.add(j) += w * (*rp.add(j)) as f32;
5072 j += 1;
5073 }
5074 }
5075}
5076
5077#[cfg(target_arch = "aarch64")]
5078#[target_feature(enable = "neon")]
5079unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
5080 unsafe {
5082 use core::arch::aarch64::*;
5083 let n = acc.len().min(row.len());
5084 let ap = acc.as_mut_ptr();
5085 let rp = row.as_ptr();
5086 let wv = vdupq_n_f32(w);
5087 let mut j = 0usize;
5088 while j + 16 <= n {
5089 let rb = vld1q_s8(rp.add(j));
5090 let lo = vmovl_s8(vget_low_s8(rb));
5091 let hi = vmovl_s8(vget_high_s8(rb));
5092 for (off, half) in [(0, lo), (8, hi)] {
5093 let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
5094 let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
5095 let o = j + off;
5096 vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
5097 vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
5098 }
5099 j += 16;
5100 }
5101 while j < n {
5102 *ap.add(j) += w * (*rp.add(j)) as f32;
5103 j += 1;
5104 }
5105}
5106}
5107
5108#[inline]
5111pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
5112 #[cfg(target_arch = "aarch64")]
5113 unsafe {
5114 return dot_i8_f32_neon(w, x);
5115 }
5116 #[cfg(target_arch = "x86_64")]
5117 if avx2_enabled() {
5118 return unsafe { dot_i8_f32_avx2(w, x) };
5119 }
5120 #[allow(unreachable_code)]
5121 {
5122 let mut sum = 0.0f32;
5123 for (j, &b) in w.iter().enumerate() {
5124 sum += (b as i8) as f32 * x[j];
5125 }
5126 sum
5127 }
5128}
5129
5130#[inline]
5134fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
5135 #[cfg(target_arch = "aarch64")]
5136 unsafe {
5137 return dot_i8_col_f32_neon(w, x, col);
5138 }
5139 #[allow(unreachable_code)]
5140 {
5141 let mut sum = 0.0f32;
5142 for (j, &b) in w.iter().enumerate() {
5143 sum += (b as i8) as f32 * x[j] * col[j];
5144 }
5145 sum
5146 }
5147}
5148
5149#[cfg(target_arch = "aarch64")]
5150#[target_feature(enable = "neon")]
5151unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
5152 unsafe {
5154 use core::arch::aarch64::*;
5155 let n = x.len();
5156 let wp = w.as_ptr() as *const i8;
5157 let xp = x.as_ptr();
5158 let cp = col.as_ptr();
5159 let (mut a0, mut a1, mut a2, mut a3) =
5160 (vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0));
5161 let mut j = 0usize;
5162 while j + 16 <= n {
5163 let wb = vld1q_s8(wp.add(j));
5164 let lo = vmovl_s8(vget_low_s8(wb));
5165 let hi = vmovl_s8(vget_high_s8(wb));
5166 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
5167 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
5168 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
5169 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
5170 a0 = vfmaq_f32(a0, w0, vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))));
5171 a1 = vfmaq_f32(a1, w1, vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))));
5172 a2 = vfmaq_f32(a2, w2, vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))));
5173 a3 = vfmaq_f32(a3, w3, vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))));
5174 j += 16;
5175 }
5176 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
5177 while j < n {
5178 sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
5179 j += 1;
5180 }
5181 sum
5182}
5183}
5184
5185#[cfg(target_arch = "aarch64")]
5186#[target_feature(enable = "neon")]
5187unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
5188 unsafe {
5190 use core::arch::aarch64::*;
5191 let n = x.len();
5192 let wp = w.as_ptr() as *const i8;
5193 let xp = x.as_ptr();
5194 let (mut a0, mut a1, mut a2, mut a3) =
5195 (vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0));
5196 let mut j = 0usize;
5197 while j + 16 <= n {
5198 let wb = vld1q_s8(wp.add(j));
5199 let lo = vmovl_s8(vget_low_s8(wb));
5200 let hi = vmovl_s8(vget_high_s8(wb));
5201 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
5202 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
5203 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
5204 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
5205 a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
5206 a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
5207 a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
5208 a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
5209 j += 16;
5210 }
5211 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
5212 while j < n {
5213 sum += (*wp.add(j)) as f32 * *xp.add(j);
5214 j += 1;
5215 }
5216 sum
5217}
5218}
5219
5220#[allow(clippy::too_many_arguments)]
5221fn qmatvec(
5222 q: &[u8],
5223 rep: &[u8],
5224 row_scale: &[f32],
5225 xs: &[f32],
5226 rows: usize,
5227 cols: usize,
5228 out: &mut [f32],
5229 pool: Option<&Pool>,
5230) {
5231 debug_assert_eq!(out.len(), rows);
5232 #[cfg(not(target_arch = "aarch64"))]
5233 let _ = rep;
5234
5235 #[cfg(target_arch = "aarch64")]
5236 if sdot_enabled() {
5237 let act = split_act(xs);
5238 let out_addr = SendMut(out.as_mut_ptr());
5239 let run_range = |start: usize, end: usize| {
5240 q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
5241 };
5242 match pool {
5243 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
5244 _ => run_range(0, rows),
5245 }
5246 return;
5247 }
5248 #[cfg(target_arch = "x86_64")]
5251 if avx2_a8w8_enabled() {
5252 let act = split_act(xs);
5253 let out_addr = SendMut(out.as_mut_ptr());
5254 let run_range =
5255 |start: usize, end: usize| q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end);
5256 match pool {
5257 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
5258 _ => run_range(0, rows),
5259 }
5260 return;
5261 }
5262
5263 let row_dot = |o: usize| -> f32 { dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o] };
5264 match pool {
5265 Some(pool) if rows >= 256 => {
5266 let out_addr = SendMut(out.as_mut_ptr());
5267 pool.run(&move |widx, n| {
5268 let chunk = rows.div_ceil(n);
5269 let start = widx * chunk;
5270 let end = (start + chunk).min(rows);
5271 for o in start..end {
5272 unsafe { *out_addr.at(o) = row_dot(o) };
5274 }
5275 });
5276 }
5277 _ => {
5278 for (o, dst) in out.iter_mut().enumerate() {
5279 *dst = row_dot(o);
5280 }
5281 }
5282 }
5283}
5284
5285#[allow(clippy::too_many_arguments)]
5286fn qmatvec2(
5287 q: &[u8],
5288 row_scale: &[f32],
5289 x1: &[f32],
5290 x2: &[f32],
5291 rows: usize,
5292 cols: usize,
5293 o1: &mut [f32],
5294 o2: &mut [f32],
5295 pool: Option<&Pool>,
5296) {
5297 #[cfg(target_arch = "aarch64")]
5298 if sdot_enabled() {
5299 let a1s = split_act(x1);
5300 let a2s = split_act(x2);
5301 let p1 = SendMut(o1.as_mut_ptr());
5302 let p2 = SendMut(o2.as_mut_ptr());
5303 let run_range = |start: usize, end: usize| {
5304 q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
5305 };
5306 match pool {
5307 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
5308 _ => run_range(0, rows),
5309 }
5310 return;
5311 }
5312 #[cfg(target_arch = "x86_64")]
5313 if avx2_a8w8_enabled() {
5314 let a1s = split_act(x1);
5315 let a2s = split_act(x2);
5316 let p1 = SendMut(o1.as_mut_ptr());
5317 let p2 = SendMut(o2.as_mut_ptr());
5318 let run_range = |start: usize, end: usize| {
5319 q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
5320 };
5321 match pool {
5322 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
5323 _ => run_range(0, rows),
5324 }
5325 return;
5326 }
5327
5328 let row_dots = |o: usize| -> (f32, f32) {
5329 let row = &q[o * cols..(o + 1) * cols];
5330 (dot_i8_f32(row, x1) * row_scale[o], dot_i8_f32(row, x2) * row_scale[o])
5331 };
5332 match pool {
5333 Some(pool) if rows >= 256 => {
5334 let p1 = SendMut(o1.as_mut_ptr());
5335 let p2 = SendMut(o2.as_mut_ptr());
5336 pool.run(&move |widx, n| {
5337 let chunk = rows.div_ceil(n);
5338 let start = widx * chunk;
5339 let end = (start + chunk).min(rows);
5340 for o in start..end {
5341 let (s1, s2) = row_dots(o);
5342 unsafe {
5344 *p1.at(o) = s1;
5345 *p2.at(o) = s2;
5346 }
5347 }
5348 });
5349 }
5350 _ => {
5351 for o in 0..rows {
5352 let (s1, s2) = row_dots(o);
5353 o1[o] = s1;
5354 o2[o] = s2;
5355 }
5356 }
5357 }
5358}
5359
5360#[derive(Clone, Copy)]
5361struct SendMut(*mut f32);
5362unsafe impl Send for SendMut {}
5363unsafe impl Sync for SendMut {}
5364
5365impl SendMut {
5366 #[inline]
5367 fn at(self, i: usize) -> *mut f32 {
5368 unsafe { self.0.add(i) }
5369 }
5370}
5371
5372#[cfg(test)]
5373mod tests {
5374 use super::*;
5375
5376 #[test]
5377 fn f32_matvec_matches_matvec_rows_bitexact() {
5378 let (rows, cols) = (300, 40);
5379 let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
5380 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
5381 let qt = QTensor::from_f32(w.clone(), rows, cols);
5382
5383 let mut a = vec![0.0f32; rows];
5384 matvec_rows(None, &w, &x, &mut a);
5385 let mut b = vec![0.0f32; rows];
5386 qt.matvec(&x, &mut b, None);
5387 assert_eq!(a, b);
5388 }
5389
5390 #[test]
5391 fn sdot_kernel_exact_on_grid() {
5392 eprintln!("sdot_enabled = {}", sdot_enabled());
5397 let (rows, cols) = (9, 80); let w: Vec<u8> = (0..rows * cols)
5399 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
5400 .collect();
5401 let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
5402 let x: Vec<f32> = (0..cols)
5403 .map(|i| match i % 3 {
5404 0 => 1.0,
5405 1 => -1.0,
5406 _ => 0.0,
5407 })
5408 .collect();
5409 let mut a = vec![0.0f32; rows];
5410 qmatvec(&w, &[], &scales, &x, rows, cols, &mut a, None);
5411 for o in 0..rows {
5412 let mut acc = 0.0f32;
5413 for j in 0..cols {
5414 acc += (w[o * cols + j] as i8) as f32 * x[j];
5415 }
5416 let expect = acc * scales[o];
5417 assert!(
5418 (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
5419 "row {o}: {} vs {expect}",
5420 a[o]
5421 );
5422 }
5423 }
5424
5425 #[test]
5426 fn q1_tbl_fast_path_matches_reference() {
5427 let (rows, cols) = (5, 256);
5432 let gpr = cols / GROUP_SIZE;
5433 let mut bytes = Vec::new();
5434 for t in 0..rows * gpr {
5435 let s = 0.007 + (t % 11) as f32 * 0.004;
5436 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
5437 for j in 0..4 {
5438 bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
5439 }
5440 }
5441 let x: Vec<f32> = (0..cols)
5442 .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
5443 .collect();
5444 let mut w = vec![0.0f32; rows * cols];
5445 cortiq_core::quant::dequant_q1(&bytes, &mut w);
5446 let mut got = vec![0.0f32; rows];
5447 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
5448 for o in 0..rows {
5449 let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
5450 assert!(
5451 (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
5452 "row {o}: {} vs {expect}",
5453 got[o]
5454 );
5455 }
5456 let b = 5usize;
5459 let mut xs_all = Vec::new();
5460 for bi in 0..b {
5461 xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
5462 }
5463 let mut mm = vec![0.0f32; b * rows];
5464 q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
5465 for bi in 0..b {
5466 let mut single = vec![0.0f32; rows];
5467 q1_matvec(&bytes, &xs_all[bi * cols..(bi + 1) * cols], rows, cols, &mut single, None);
5468 assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
5469 }
5470 }
5471
5472 #[test]
5473 fn q1_kernels_match_exact_reference() {
5474 let (rows, cols) = (7, 96);
5476 let gpr = cols / GROUP_SIZE;
5477 let mut bytes = Vec::new();
5478 for t in 0..rows * gpr {
5479 let s = 0.01 + (t % 13) as f32 * 0.003;
5480 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
5481 for j in 0..4 {
5482 bytes.push(((t * 31 + j * 97) % 251) as u8);
5483 }
5484 }
5485 let x: Vec<f32> = (0..cols)
5487 .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
5488 .collect();
5489 let mut w = vec![0.0f32; rows * cols];
5491 cortiq_core::quant::dequant_q1(&bytes, &mut w);
5492 let mut expect = vec![0.0f32; rows];
5493 for o in 0..rows {
5494 expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
5495 }
5496 let mut got = vec![0.0f32; rows];
5497 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
5498 for o in 0..rows {
5499 assert!(
5500 (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
5501 "row {o}: {} vs {}",
5502 got[o],
5503 expect[o]
5504 );
5505 }
5506 let x2: Vec<f32> = x.iter().map(|v| -v).collect();
5508 let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
5509 q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
5510 assert_eq!(a1, got);
5511 let mut xs = x.clone();
5512 xs.extend_from_slice(&x2);
5513 let mut mm = vec![0.0f32; 2 * rows];
5514 q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
5515 assert_eq!(&mm[..rows], got.as_slice());
5516 assert_eq!(&mm[rows..], a2.as_slice());
5517 }
5518
5519 #[test]
5520 fn repack_is_bit_identical() {
5521 let (rows, cols) = (267, 96); let w: Vec<u8> = (0..rows * cols)
5527 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
5528 .collect();
5529 let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
5530 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
5531 let rep = q8_repack_layout(&w, rows, cols);
5532 for g in 0..rows / 4 {
5534 for c in 0..cols / 16 {
5535 for lane in 0..4 {
5536 assert_eq!(
5537 &rep[g * 4 * cols + c * 64 + lane * 16..g * 4 * cols + c * 64 + lane * 16 + 16],
5538 &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
5539 );
5540 }
5541 }
5542 }
5543 let mut a = vec![0.0f32; rows];
5544 qmatvec(&w, &[], &scales, &x, rows, cols, &mut a, None);
5545 let mut b = vec![0.0f32; rows];
5546 qmatvec(&w, &rep, &scales, &x, rows, cols, &mut b, None);
5547 assert_eq!(a, b, "full-range repack output diverged");
5548
5549 #[cfg(target_arch = "aarch64")]
5550 if sdot_enabled() {
5551 let act = split_act(&x);
5553 let mut c1 = vec![0.0f32; rows];
5554 let mut c2 = vec![0.0f32; rows];
5555 q8_range_sdot(&w, &[], &scales, &act, cols, SendMut(c1.as_mut_ptr()), 3, rows - 2);
5556 q8_range_sdot(&w, &rep, &scales, &act, cols, SendMut(c2.as_mut_ptr()), 3, rows - 2);
5557 assert_eq!(c1, c2, "unaligned-range repack output diverged");
5558 }
5559 }
5560
5561 #[test]
5562 fn sdot_a8w8_noise_is_bounded() {
5563 let (rows, cols) = (16, 512);
5567 let w: Vec<u8> = (0..rows * cols)
5568 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
5569 .collect();
5570 let scales = vec![0.01f32; rows];
5571 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
5572 let mut a = vec![0.0f32; rows];
5573 qmatvec(&w, &[], &scales, &x, rows, cols, &mut a, None);
5574 let (mut num, mut den) = (0f64, 0f64);
5575 for o in 0..rows {
5576 let mut acc = 0.0f32;
5577 for j in 0..cols {
5578 acc += (w[o * cols + j] as i8) as f32 * x[j];
5579 }
5580 let expect = acc * scales[o];
5581 num += ((a[o] - expect) as f64).powi(2);
5582 den += (expect as f64).powi(2);
5583 }
5584 let rel = (num / den.max(1e-12)).sqrt();
5585 assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
5586 }
5587
5588 #[test]
5589 fn i8_dot_neon_matches_scalar() {
5590 let n = 100;
5591 let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
5592 let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
5593 let mut scalar = 0.0f32;
5594 for j in 0..n {
5595 scalar += (w[j] as i8) as f32 * x[j];
5596 }
5597 let fast = dot_i8_f32(&w, &x);
5598 assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
5599 }
5600
5601 #[test]
5603 fn vbitmatvec_matches_full_dequant() {
5604 let (rows, cols) = (6, 64);
5605 let ng = cols / GROUP_SIZE;
5606 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
5608 let mut bytes = bits.clone();
5609 for g in 0..rows * ng {
5610 let s = 0.02 + 0.001 * g as f32;
5611 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
5612 }
5613 for r in 0..rows {
5614 let b = bits[r] as usize;
5615 let (mut acc, mut nb) = (0u64, 0usize);
5616 let mut rowbytes = Vec::new();
5617 for i in 0..cols {
5618 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
5619 acc = (acc << b) | v;
5620 nb += b;
5621 while nb >= 8 {
5622 nb -= 8;
5623 rowbytes.push(((acc >> nb) & 0xFF) as u8);
5624 }
5625 }
5626 if nb > 0 {
5627 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
5628 }
5629 bytes.extend_from_slice(&rowbytes);
5630 }
5631 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
5632
5633 let mut reference = vec![0f32; rows * cols];
5634 cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
5635 let mut expect = vec![0f32; rows];
5636 for r in 0..rows {
5637 expect[r] = reference[r * cols..(r + 1) * cols]
5638 .iter()
5639 .zip(&x)
5640 .map(|(w, xv)| w * xv)
5641 .sum();
5642 }
5643 let mut got = vec![0f32; rows];
5644 let offsets = vbit_row_offsets(&bytes, rows, cols);
5645 vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
5646 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
5650 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
5651 for r in 0..rows {
5652 assert!(
5653 (got[r] - expect[r]).abs() < tol * scale,
5654 "row {r}: {} vs {}",
5655 got[r],
5656 expect[r]
5657 );
5658 }
5659 }
5660
5661 #[test]
5666 #[cfg(target_arch = "x86_64")]
5667 fn vbit_matmat_blocked_matches_per_row() {
5668 let (rows, cols, b) = (64usize, 128usize, 9usize);
5669 let ng = cols / GROUP_SIZE;
5670 let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
5671 let mut bytes = bits.clone();
5672 for g in 0..rows * ng {
5673 let sc = 0.02 + 0.0005 * g as f32;
5674 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
5675 }
5676 for r in 0..rows {
5677 let bw = bits[r] as usize;
5678 let (mut acc, mut nb) = (0u64, 0usize);
5679 let mut rowbytes = Vec::new();
5680 for i in 0..cols {
5681 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
5682 acc = (acc << bw) | v;
5683 nb += bw;
5684 while nb >= 8 {
5685 nb -= 8;
5686 rowbytes.push(((acc >> nb) & 0xFF) as u8);
5687 }
5688 }
5689 if nb > 0 {
5690 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
5691 }
5692 bytes.extend_from_slice(&rowbytes);
5693 }
5694 let x: Vec<f32> =
5695 (0..b * cols).map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5).collect();
5696 let offsets = vbit_row_offsets(&bytes, rows, cols);
5697 let mut y_a = vec![0f32; b * rows];
5698 let mut y_b = vec![0f32; b * rows];
5699 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
5700 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
5701 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
5702 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
5703 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
5704 let max_d =
5705 y_a.iter().zip(&y_b).map(|(p, q)| (p - q).abs()).fold(0.0f32, f32::max);
5706 assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
5707 }
5708
5709 #[test]
5710 fn q4matvec_matches_full_dequant() {
5711 let (rows, cols) = (8, 64);
5712 let groups = rows * cols / GROUP_SIZE;
5713 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
5715 for i in 0..groups * 16 {
5716 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
5717 }
5718 for g in 0..groups {
5719 let s = 0.01 + 0.003 * g as f32;
5720 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
5721 }
5722 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
5723
5724 let mut reference = vec![0.0f32; rows * cols];
5725 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
5726 let mut expect = vec![0.0f32; rows];
5727 for r in 0..rows {
5728 expect[r] = reference[r * cols..(r + 1) * cols]
5729 .iter()
5730 .zip(&x)
5731 .map(|(w, xv)| w * xv)
5732 .sum();
5733 }
5734
5735 let mut got = vec![0.0f32; rows];
5736 q4matvec(&bytes, &x, rows, cols, &mut got, None);
5737 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
5741 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
5742 for r in 0..rows {
5743 assert!(
5744 (got[r] - expect[r]).abs() < tol * scale,
5745 "row {r}: {} vs {}",
5746 got[r],
5747 expect[r]
5748 );
5749 }
5750 }
5751
5752 #[test]
5755 fn vbitmatvec2_equals_two_singles() {
5756 let (rows, cols) = (6, 64);
5757 let ng = cols / GROUP_SIZE;
5758 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
5759 let mut bytes = bits.clone();
5760 for g in 0..rows * ng {
5761 let s = 0.02 + 0.001 * g as f32;
5762 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
5763 }
5764 for r in 0..rows {
5765 let b = bits[r] as usize;
5766 let (mut acc, mut nb) = (0u64, 0usize);
5767 let mut rowbytes = Vec::new();
5768 for i in 0..cols {
5769 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
5770 acc = (acc << b) | v;
5771 nb += b;
5772 while nb >= 8 {
5773 nb -= 8;
5774 rowbytes.push(((acc >> nb) & 0xFF) as u8);
5775 }
5776 }
5777 if nb > 0 {
5778 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
5779 }
5780 bytes.extend_from_slice(&rowbytes);
5781 }
5782 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
5783 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
5784 let offsets = vbit_row_offsets(&bytes, rows, cols);
5785
5786 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
5787 vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
5788 vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
5789 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
5790 vbitmatvec2(&bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
5791 assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
5792 assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
5793 }
5794
5795 #[test]
5797 fn q4matvec2_equals_two_singles() {
5798 let (rows, cols) = (8, 128);
5799 let groups = rows * cols / GROUP_SIZE;
5800 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
5801 for i in 0..groups * 16 {
5802 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
5803 }
5804 for g in 0..groups {
5805 let s = 0.01 + 0.003 * g as f32;
5806 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
5807 }
5808 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
5811 x1[9] = 250.0;
5812 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
5813
5814 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
5815 q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
5816 q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
5817 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
5818 q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
5819 assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
5820 assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
5821 }
5822
5823 #[test]
5826 fn matvec_many_equals_separate_matvecs() {
5827 use crate::pool::Pool;
5828 let (r1, r2, cols) = (300, 200, 64);
5829 let mk = |salt: usize, rows: usize| {
5830 QTensor::from_f32(
5831 (0..rows * cols).map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5).collect(),
5832 rows,
5833 cols,
5834 )
5835 };
5836 let (a, b) = (mk(1, r1), mk(5, r2));
5837 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
5838 let pool = Pool::new(3);
5839
5840 let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
5841 a.matvec(&x, &mut ea, Some(&pool));
5842 b.matvec(&x, &mut eb, Some(&pool));
5843 let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
5844 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
5845 assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
5846 assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
5847 }
5848
5849 #[test]
5852 fn batched_matmat_equals_per_position_matvec() {
5853 let (rows, cols, b) = (8, 64, 5);
5854 let groups = rows * cols / GROUP_SIZE;
5856 let mut q4 = Vec::new();
5857 for i in 0..groups * 16 {
5858 q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
5859 }
5860 for g in 0..groups {
5861 q4.extend_from_slice(
5862 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
5863 );
5864 }
5865 let ng = cols / GROUP_SIZE;
5867 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
5868 let mut vb = bits.clone();
5869 for g in 0..rows * ng {
5870 vb.extend_from_slice(
5871 &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
5872 );
5873 }
5874 for r in 0..rows {
5875 let bw = bits[r] as usize;
5876 let (mut acc, mut nb) = (0u64, 0usize);
5877 let mut rowbytes = Vec::new();
5878 for i in 0..cols {
5879 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
5880 acc = (acc << bw) | v;
5881 nb += bw;
5882 while nb >= 8 {
5883 nb -= 8;
5884 rowbytes.push(((acc >> nb) & 0xFF) as u8);
5885 }
5886 }
5887 if nb > 0 {
5888 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
5889 }
5890 vb.extend_from_slice(&rowbytes);
5891 }
5892 let offsets = vbit_row_offsets(&vb, rows, cols);
5893
5894 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
5895
5896 let mut got = vec![0f32; b * rows];
5898 q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
5899 for bi in 0..b {
5900 let mut expect = vec![0f32; rows];
5901 q4matvec(&q4, &xs[bi * cols..(bi + 1) * cols], rows, cols, &mut expect, None);
5902 assert_eq!(&got[bi * rows..(bi + 1) * rows], &expect[..], "q4 batch pos {bi}");
5903 }
5904
5905 let mut got = vec![0f32; b * rows];
5907 vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
5908 for bi in 0..b {
5909 let mut expect = vec![0f32; rows];
5910 vbitmatvec(
5911 &vb, &offsets, &xs[bi * cols..(bi + 1) * cols], rows, cols, &mut expect, None,
5912 );
5913 assert_eq!(&got[bi * rows..(bi + 1) * rows], &expect[..], "vbit batch pos {bi}");
5914 }
5915 }
5916
5917 #[test]
5921 fn q4_tiled_matches_q4_block_bitexact() {
5922 let (rows, cols, b) = (8usize, 128usize, 3usize);
5923 let groups = rows * cols / GROUP_SIZE;
5924 let mut split = Vec::with_capacity(groups * 18);
5925 for i in 0..groups * 16 {
5926 split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
5927 }
5928 for g in 0..groups {
5929 split.extend_from_slice(
5930 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
5931 );
5932 }
5933 let (packed, scales) = split.split_at(groups * 16);
5935 let mut tiled = Vec::with_capacity(groups * Q4_TILE);
5936 for g in 0..groups {
5937 tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
5938 tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
5939 }
5940
5941 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
5942 x1[9] = 250.0; let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
5944
5945 let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
5946 q4matvec(&split, &x1, rows, cols, &mut a, None);
5947 q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
5948 assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
5949
5950 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
5951 let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
5952 q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
5953 q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
5954 assert_eq!(a1, t1);
5955 assert_eq!(a2, t2);
5956
5957 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
5958 let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
5959 q4matmat(&split, &xs, b, rows, cols, &mut am, None);
5960 q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
5961 assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
5962 }
5963
5964 #[test]
5971 fn q4matvec_sdot_outlier_exact() {
5972 let (rows, cols) = (4, 128);
5973 let groups = rows * cols / GROUP_SIZE;
5974 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
5975 for i in 0..groups * 16 {
5976 bytes.push(((i * 11 + 5) % 256) as u8);
5977 }
5978 for g in 0..groups {
5979 let s = 0.02 + 0.002 * g as f32;
5980 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
5981 }
5982 let mut x: Vec<f32> = (0..cols)
5983 .map(|i| match i % 3 {
5984 0 => 1.0,
5985 1 => -1.0,
5986 _ => 0.0,
5987 })
5988 .collect();
5989 x[17] = 300.0; let mut reference = vec![0.0f32; rows * cols];
5992 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
5993 let mut expect = vec![0.0f32; rows];
5994 for r in 0..rows {
5995 expect[r] = reference[r * cols..(r + 1) * cols]
5996 .iter()
5997 .zip(&x)
5998 .map(|(w, xv)| w * xv)
5999 .sum();
6000 }
6001 let mut got = vec![0.0f32; rows];
6002 q4matvec(&bytes, &x, rows, cols, &mut got, None);
6003 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
6004 for r in 0..rows {
6005 assert!(
6006 (got[r] - expect[r]).abs() < 2e-3 * scale,
6007 "row {r}: {} vs {} (outlier term must be exact)",
6008 got[r],
6009 expect[r]
6010 );
6011 }
6012 }
6013}