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 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
2780 let xp = xq.as_ptr();
2781 let gp = gsum.as_ptr();
2782 let mut accv = vdupq_n_f32(0.0);
2783 let mut gi = 0;
2784 while gi + 4 <= gpr {
2785 let t0 = base.add(gi * Q1_TILE);
2786 let t1 = base.add((gi + 1) * Q1_TILE);
2787 let t2 = base.add((gi + 2) * Q1_TILE);
2788 let t3 = base.add((gi + 3) * Q1_TILE);
2789 let d0 = tile_dot!(t0, xp.add(gi * GROUP_SIZE));
2790 let d1 = tile_dot!(t1, xp.add((gi + 1) * GROUP_SIZE));
2791 let d2 = tile_dot!(t2, xp.add((gi + 2) * GROUP_SIZE));
2792 let d3 = tile_dot!(t3, xp.add((gi + 3) * GROUP_SIZE));
2793 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
2795 let g = vld1q_s32(gp.add(gi));
2796 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
2797 let sc = [
2798 f16_to_f32(u16::from_le_bytes([*t0, *t0.add(1)])),
2799 f16_to_f32(u16::from_le_bytes([*t1, *t1.add(1)])),
2800 f16_to_f32(u16::from_le_bytes([*t2, *t2.add(1)])),
2801 f16_to_f32(u16::from_le_bytes([*t3, *t3.add(1)])),
2802 ];
2803 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), vld1q_f32(sc.as_ptr()));
2804 gi += 4;
2805 }
2806 let mut acc = vaddvq_f32(accv);
2807 while gi < gpr {
2808 let t = base.add(gi * Q1_TILE);
2809 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2810 let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
2811 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
2812 gi += 1;
2813 }
2814 acc
2815 }
2816}
2817
2818#[inline]
2820fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
2821 let gi = j / GROUP_SIZE;
2822 let k = j % GROUP_SIZE;
2823 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
2824 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2825 let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
2826 ((bit as i32 * 2 - 1) as f32, s)
2827}
2828
2829#[inline]
2831fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
2832 let mut acc = 0f32;
2833 for gi in 0..gpr {
2834 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
2835 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2836 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
2837 let mut ga = 0f32;
2838 for (j, &b) in tile[2..].iter().enumerate() {
2839 for k in 0..8 {
2840 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
2841 }
2842 }
2843 acc += ga * s;
2844 }
2845 acc
2846}
2847
2848#[allow(clippy::too_many_arguments)]
2851fn q1_range_a8w8(
2852 bytes: &[u8],
2853 gpr: usize,
2854 act: &SplitAct,
2855 gsum: &[i32],
2856 out: SendMut,
2857 start: usize,
2858 end: usize,
2859) {
2860 for r in start..end {
2861 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
2862 for &(j, xv) in &act.outliers {
2863 let (w, s) = q1_outlier(bytes, r, gpr, j);
2864 acc += w * s * xv;
2865 }
2866 unsafe { *out.at(r) = acc };
2868 }
2869}
2870
2871fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
2873 for r in start..end {
2874 unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
2876 }
2877}
2878
2879fn q1_matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
2881 debug_assert_eq!(out.len(), rows);
2882 let gpr = cols / GROUP_SIZE;
2883 let out_addr = SendMut(out.as_mut_ptr());
2884 if a8w8_enabled() {
2885 let act = split_act(x);
2886 let gsum = q1_group_sums(&act.xq, gpr);
2887 let (act, gsum) = (&act, &gsum);
2888 let run =
2889 move |start: usize, end: usize| q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end);
2890 dispatch_rows(pool, rows, &run);
2891 return;
2892 }
2893 let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
2894 dispatch_rows(pool, rows, &run);
2895}
2896
2897#[allow(clippy::too_many_arguments)]
2899fn q1_matvec2(
2900 bytes: &[u8],
2901 x1: &[f32],
2902 x2: &[f32],
2903 rows: usize,
2904 cols: usize,
2905 o1: &mut [f32],
2906 o2: &mut [f32],
2907 pool: Option<&Pool>,
2908) {
2909 let gpr = cols / GROUP_SIZE;
2910 let p1 = SendMut(o1.as_mut_ptr());
2911 let p2 = SendMut(o2.as_mut_ptr());
2912 if a8w8_enabled() {
2913 let a1 = split_act(x1);
2914 let a2 = split_act(x2);
2915 let g1 = q1_group_sums(&a1.xq, gpr);
2916 let g2 = q1_group_sums(&a2.xq, gpr);
2917 let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
2918 let run = move |start: usize, end: usize| {
2919 for r in start..end {
2920 let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
2921 let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
2922 for &(j, xv) in &a1.outliers {
2923 let (w, s) = q1_outlier(bytes, r, gpr, j);
2924 v1 += w * s * xv;
2925 }
2926 for &(j, xv) in &a2.outliers {
2927 let (w, s) = q1_outlier(bytes, r, gpr, j);
2928 v2 += w * s * xv;
2929 }
2930 unsafe {
2932 *p1.at(r) = v1;
2933 *p2.at(r) = v2;
2934 }
2935 }
2936 };
2937 dispatch_rows(pool, rows, &run);
2938 return;
2939 }
2940 let run = move |start: usize, end: usize| {
2941 for r in start..end {
2942 unsafe {
2944 *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
2945 *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
2946 }
2947 }
2948 };
2949 dispatch_rows(pool, rows, &run);
2950}
2951
2952#[allow(clippy::too_many_arguments)]
2954fn q1_matmat(
2955 bytes: &[u8],
2956 xs_all: &[f32],
2957 b: usize,
2958 rows: usize,
2959 cols: usize,
2960 out: &mut [f32],
2961 pool: Option<&Pool>,
2962) {
2963 debug_assert_eq!(out.len(), b * rows);
2964 let gpr = cols / GROUP_SIZE;
2965 let out_addr = SendMut(out.as_mut_ptr());
2966 if a8w8_enabled() {
2967 let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
2968 .map(|bi| {
2969 let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
2970 let gsum = q1_group_sums(&act.xq, gpr);
2971 (act, gsum)
2972 })
2973 .collect();
2974 let acts = &acts;
2975 #[cfg(target_arch = "x86_64")]
2976 let blocked_ok = avx2_enabled()
2977 && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true);
2978 let run = move |start: usize, end: usize| {
2979 for r in start..end {
2980 let mut bi = 0usize;
2981 #[cfg(target_arch = "x86_64")]
2984 if blocked_ok {
2985 while bi + 4 <= acts.len() {
2986 let xs = [
2987 acts[bi].0.xq.as_slice(),
2988 acts[bi + 1].0.xq.as_slice(),
2989 acts[bi + 2].0.xq.as_slice(),
2990 acts[bi + 3].0.xq.as_slice(),
2991 ];
2992 let gs = [
2993 acts[bi].1.as_slice(),
2994 acts[bi + 1].1.as_slice(),
2995 acts[bi + 2].1.as_slice(),
2996 acts[bi + 3].1.as_slice(),
2997 ];
2998 let d = unsafe { dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs) };
2999 for k in 0..4 {
3000 let (act, _) = &acts[bi + k];
3001 let mut acc = d[k] * act.sx;
3002 for &(j, xv) in &act.outliers {
3003 let (w, sc) = q1_outlier(bytes, r, gpr, j);
3004 acc += w * sc * xv;
3005 }
3006 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
3008 }
3009 bi += 4;
3010 }
3011 }
3012 while bi < acts.len() {
3013 let (act, gsum) = &acts[bi];
3014 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
3015 for &(j, xv) in &act.outliers {
3016 let (w, s) = q1_outlier(bytes, r, gpr, j);
3017 acc += w * s * xv;
3018 }
3019 unsafe { *out_addr.at(bi * rows + r) = acc };
3021 bi += 1;
3022 }
3023 }
3024 };
3025 dispatch_rows(pool, rows, &run);
3026 return;
3027 }
3028 let run = move |start: usize, end: usize| {
3029 for r in start..end {
3030 for bi in 0..b {
3031 let x = &xs_all[bi * cols..(bi + 1) * cols];
3032 unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
3034 }
3035 }
3036 };
3037 dispatch_rows(pool, rows, &run);
3038}
3039
3040fn q4matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
3046 debug_assert_eq!(out.len(), rows);
3047 let (packed, scales) = q4_split(bytes, rows, cols);
3048 let gpr = cols / GROUP_SIZE;
3049 let out_addr = SendMut(out.as_mut_ptr());
3050
3051 if a8w8_enabled() {
3052 let act = split_act(x);
3053 let run = move |start: usize, end: usize| {
3054 q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
3055 };
3056 dispatch_rows(pool, rows, &run);
3057 return;
3058 }
3059
3060 let run = move |start: usize, end: usize| {
3061 q4_range_f32(packed, scales, gpr, x, out_addr, start, end)
3062 };
3063 dispatch_rows(pool, rows, &run);
3064}
3065
3066#[inline]
3069#[allow(unreachable_code)]
3070#[cfg(target_arch = "x86_64")]
3075#[target_feature(enable = "avx2")]
3076unsafe fn dot_q4b_row_1x4_avx2(
3077 buf: &[u8],
3078 scales: &[u8],
3079 g0: usize,
3080 gpr: usize,
3081 xs: [&[i8]; 4],
3082) -> [f32; 4] {
3083 unsafe {
3085 use core::arch::x86_64::*;
3086 let ones = _mm256_set1_epi16(1);
3087 let mut acc = [0f32; 4];
3088 for gi in 0..gpr {
3089 let s = f16_to_f32(u16::from_le_bytes([
3090 scales[(g0 + gi) * 2],
3091 scales[(g0 + gi) * 2 + 1],
3092 ]));
3093 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3094 let aw = _mm256_abs_epi8(w);
3095 for (k, xq) in xs.iter().enumerate() {
3096 let x =
3097 _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3098 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
3099 let d = _mm256_madd_epi16(p16, ones);
3100 let hi128 = _mm256_extracti128_si256::<1>(d);
3101 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3102 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3103 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3104 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
3105 }
3106 }
3107 acc
3108 }
3109}
3110
3111#[cfg(target_arch = "x86_64")]
3117#[target_feature(enable = "avx2")]
3118unsafe fn dot_q4b_row_1x4_sx_avx2(
3119 buf: &[u8],
3120 scales: &[u8],
3121 g0: usize,
3122 gpr: usize,
3123 xs: [&[i8]; 4],
3124 sxs: [f32; 4],
3125) -> [f32; 4] {
3126 unsafe {
3128 use core::arch::x86_64::*;
3129 let ones = _mm256_set1_epi16(1);
3130 let mut acc = [0f32; 4];
3131 for gi in 0..gpr {
3132 let s = f16_to_f32(u16::from_le_bytes([
3133 scales[(g0 + gi) * 2],
3134 scales[(g0 + gi) * 2 + 1],
3135 ]));
3136 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3137 let aw = _mm256_abs_epi8(w);
3138 for (k, xq) in xs.iter().enumerate() {
3139 let x =
3140 _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3141 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
3142 let d = _mm256_madd_epi16(p16, ones);
3143 let hi128 = _mm256_extracti128_si256::<1>(d);
3144 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3145 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3146 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3147 acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
3148 }
3149 }
3150 acc
3151 }
3152}
3153
3154fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
3155 #[cfg(target_arch = "aarch64")]
3156 unsafe {
3157 return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
3158 }
3159 #[cfg(target_arch = "x86_64")]
3160 unsafe {
3161 return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
3162 }
3163 let mut acc = 0f32;
3164 for gi in 0..gpr {
3165 let g = g0 + gi;
3166 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
3167 let mut d = 0i32;
3168 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
3169 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
3170 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
3171 }
3172 acc += d as f32 * s;
3173 }
3174 acc
3175}
3176
3177#[inline]
3179#[allow(unreachable_code)]
3180fn dot_q4_row_i8_2(
3181 packed: &[u8],
3182 scales: &[u8],
3183 g0: usize,
3184 gpr: usize,
3185 xq1: &[i8],
3186 xq2: &[i8],
3187) -> (f32, f32) {
3188 #[cfg(target_arch = "aarch64")]
3189 unsafe {
3190 return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
3191 }
3192 #[cfg(target_arch = "x86_64")]
3193 unsafe {
3194 return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
3195 }
3196 (
3197 dot_q4_row_i8(packed, scales, g0, gpr, xq1),
3198 dot_q4_row_i8(packed, scales, g0, gpr, xq2),
3199 )
3200}
3201
3202#[allow(clippy::too_many_arguments)]
3205fn q4_range_a8w8(
3206 packed: &[u8],
3207 scales: &[u8],
3208 gpr: usize,
3209 cols: usize,
3210 act: &SplitAct,
3211 out: SendMut,
3212 start: usize,
3213 end: usize,
3214) {
3215 for r in start..end {
3216 let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
3217 for &(j, xv) in &act.outliers {
3219 let flat = r * cols + j;
3220 let byte = packed[flat / 2];
3221 let nib = if flat & 1 == 0 { byte & 0x0F } else { byte >> 4 };
3222 let s = f16_to_f32(u16::from_le_bytes([
3223 scales[(flat / GROUP_SIZE) * 2],
3224 scales[(flat / GROUP_SIZE) * 2 + 1],
3225 ]));
3226 acc += ((nib as i32 - 8) as f32) * s * xv;
3227 }
3228 unsafe { *out.at(r) = acc };
3230 }
3231}
3232
3233#[allow(clippy::too_many_arguments)]
3236fn q4_range2_a8w8(
3237 packed: &[u8],
3238 scales: &[u8],
3239 gpr: usize,
3240 cols: usize,
3241 a1: &SplitAct,
3242 a2: &SplitAct,
3243 p1: SendMut,
3244 p2: SendMut,
3245 start: usize,
3246 end: usize,
3247) {
3248 for r in start..end {
3249 let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
3250 let mut acc1 = s1 * a1.sx;
3251 let mut acc2 = s2 * a2.sx;
3252 let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
3254 for &(j, xv) in outliers {
3255 let flat = r * cols + j;
3256 let byte = packed[flat / 2];
3257 let nib = if flat & 1 == 0 { byte & 0x0F } else { byte >> 4 };
3258 let s = f16_to_f32(u16::from_le_bytes([
3259 scales[(flat / GROUP_SIZE) * 2],
3260 scales[(flat / GROUP_SIZE) * 2 + 1],
3261 ]));
3262 *acc += ((nib as i32 - 8) as f32) * s * xv;
3263 }
3264 };
3265 fix(&a1.outliers, &mut acc1);
3266 fix(&a2.outliers, &mut acc2);
3267 unsafe {
3269 *p1.at(r) = acc1;
3270 *p2.at(r) = acc2;
3271 }
3272 }
3273}
3274
3275fn q4_range_f32(
3277 packed: &[u8],
3278 scales: &[u8],
3279 gpr: usize,
3280 x: &[f32],
3281 out: SendMut,
3282 start: usize,
3283 end: usize,
3284) {
3285 for r in start..end {
3286 let mut acc = 0f32;
3287 for gi in 0..gpr {
3288 let g = r * gpr + gi;
3289 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
3290 let pk = &packed[g * 16..(g + 1) * 16];
3291 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
3292 let mut ga = 0f32;
3293 for (k, &b) in pk.iter().enumerate() {
3294 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
3295 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
3296 }
3297 acc += ga * s;
3298 }
3299 unsafe { *out.at(r) = acc };
3301 }
3302}
3303
3304#[allow(clippy::too_many_arguments)]
3308fn q4matvec2(
3309 bytes: &[u8],
3310 x1: &[f32],
3311 x2: &[f32],
3312 rows: usize,
3313 cols: usize,
3314 o1: &mut [f32],
3315 o2: &mut [f32],
3316 pool: Option<&Pool>,
3317) {
3318 debug_assert_eq!(o1.len(), rows);
3319 debug_assert_eq!(o2.len(), rows);
3320 let (packed, scales) = q4_split(bytes, rows, cols);
3321 let gpr = cols / GROUP_SIZE;
3322
3323 if a8w8_enabled() {
3324 let a1 = split_act(x1);
3325 let a2 = split_act(x2);
3326 let p1 = SendMut(o1.as_mut_ptr());
3327 let p2 = SendMut(o2.as_mut_ptr());
3328 let run = move |start: usize, end: usize| {
3329 q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
3330 };
3331 dispatch_rows(pool, rows, &run);
3332 return;
3333 }
3334
3335 let p1 = SendMut(o1.as_mut_ptr());
3336 let p2 = SendMut(o2.as_mut_ptr());
3337 let run = move |start: usize, end: usize| {
3338 q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
3339 };
3340 dispatch_rows(pool, rows, &run);
3341}
3342
3343#[allow(clippy::too_many_arguments)]
3345fn q4_range2_f32(
3346 packed: &[u8],
3347 scales: &[u8],
3348 gpr: usize,
3349 x1: &[f32],
3350 x2: &[f32],
3351 p1: SendMut,
3352 p2: SendMut,
3353 start: usize,
3354 end: usize,
3355) {
3356 for r in start..end {
3357 let (mut acc1, mut acc2) = (0f32, 0f32);
3358 for gi in 0..gpr {
3359 let g = r * gpr + gi;
3360 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
3361 let pk = &packed[g * 16..(g + 1) * 16];
3362 let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
3363 let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
3364 let (mut g1, mut g2) = (0f32, 0f32);
3365 for (k, &b) in pk.iter().enumerate() {
3366 let wl = (b & 0x0F) as f32 - 8.0;
3367 let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
3368 g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
3369 g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
3370 }
3371 acc1 += g1 * s;
3372 acc2 += g2 * s;
3373 }
3374 unsafe {
3376 *p1.at(r) = acc1;
3377 *p2.at(r) = acc2;
3378 }
3379 }
3380}
3381
3382thread_local! {
3383 static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
3386 static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3387}
3388
3389#[allow(clippy::too_many_arguments)]
3395fn q4matmat(
3396 bytes: &[u8],
3397 xs_all: &[f32],
3398 b: usize,
3399 rows: usize,
3400 cols: usize,
3401 out: &mut [f32],
3402 pool: Option<&Pool>,
3403) {
3404 debug_assert_eq!(xs_all.len(), b * cols);
3405 debug_assert_eq!(out.len(), b * rows);
3406 let (packed, scales) = q4_split(bytes, rows, cols);
3407 let gpr = cols / GROUP_SIZE;
3408 let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
3409
3410 if a8w8_enabled() {
3411 let acts: Vec<SplitAct> =
3412 (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
3413 let acts = &acts;
3414 let out_addr = SendMut(out.as_mut_ptr());
3415 let run = move |start: usize, end: usize| {
3416 ROW_I8.with(|rb| {
3417 let mut buf = rb.borrow_mut();
3418 buf.resize(cols, 0);
3419 for r in start..end {
3420 for gi in 0..gpr {
3424 let g = r * gpr + gi;
3425 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
3426 buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
3427 buf[gi * GROUP_SIZE + k * 2 + 1] =
3428 (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
3429 }
3430 }
3431 let mut bi = 0usize;
3432 #[cfg(target_arch = "x86_64")]
3433 if avx2_enabled()
3434 && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true)
3435 {
3436 while bi + 4 <= acts.len() {
3437 let xs = [
3438 acts[bi].xq.as_slice(),
3439 acts[bi + 1].xq.as_slice(),
3440 acts[bi + 2].xq.as_slice(),
3441 acts[bi + 3].xq.as_slice(),
3442 ];
3443 let d = unsafe {
3444 dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
3445 };
3446 for k in 0..4 {
3447 let act = &acts[bi + k];
3448 let mut acc = d[k] * act.sx;
3449 for &(j, xv) in &act.outliers {
3450 acc += (buf[j] as i8) as f32
3451 * gscale((r * cols + j) / GROUP_SIZE)
3452 * xv;
3453 }
3454 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
3456 }
3457 bi += 4;
3458 }
3459 }
3460 while bi < acts.len() {
3461 let act = &acts[bi];
3462 let mut acc = 0f32;
3463 for gi in 0..gpr {
3464 let d = dot_i8_i8(
3465 &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
3466 &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
3467 );
3468 acc += d as f32 * gscale(r * gpr + gi);
3469 }
3470 acc *= act.sx;
3471 for &(j, xv) in &act.outliers {
3473 acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
3474 }
3475 unsafe { *out_addr.at(bi * rows + r) = acc };
3477 bi += 1;
3478 }
3479 }
3480 })
3481 };
3482 dispatch_rows(pool, rows, &run);
3483 return;
3484 }
3485
3486 let out_addr = SendMut(out.as_mut_ptr());
3487 let run = move |start: usize, end: usize| {
3488 ROW_F32.with(|rb| {
3489 let mut buf = rb.borrow_mut();
3490 buf.resize(cols, 0.0);
3491 for r in start..end {
3492 for gi in 0..gpr {
3495 let g = r * gpr + gi;
3496 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
3497 buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
3498 buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
3499 }
3500 }
3501 for bi in 0..b {
3502 let x = &xs_all[bi * cols..(bi + 1) * cols];
3503 let mut acc = 0f32;
3504 for gi in 0..gpr {
3505 let mut ga = 0f32;
3506 for k in 0..GROUP_SIZE / 2 {
3511 let e = gi * GROUP_SIZE + k * 2;
3512 ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
3513 }
3514 acc += ga * gscale(r * gpr + gi);
3515 }
3516 unsafe { *out_addr.at(bi * rows + r) = acc };
3518 }
3519 }
3520 })
3521 };
3522 dispatch_rows(pool, rows, &run);
3523}
3524
3525#[allow(clippy::too_many_arguments)]
3530fn vbitmatmat(
3531 bytes: &[u8],
3532 offsets: &[usize],
3533 xs_all: &[f32],
3534 b: usize,
3535 rows: usize,
3536 cols: usize,
3537 out: &mut [f32],
3538 pool: Option<&Pool>,
3539) {
3540 debug_assert_eq!(xs_all.len(), b * cols);
3541 debug_assert_eq!(out.len(), b * rows);
3542 debug_assert_eq!(offsets.len(), rows + 1);
3543 let ng = cols / GROUP_SIZE;
3544 let bits = &bytes[..rows];
3545 let sc_off = rows;
3546 let gscale = |r: usize, g: usize| {
3547 let so = (r * ng + g) * 2;
3548 f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]))
3549 };
3550
3551 let decode_f32 = |r: usize, dst: &mut [f32]| {
3553 let bw = bits[r] as usize;
3554 let l = ((1i32 << (bw - 1)) - 1) as f32;
3555 let data = &bytes[offsets[r]..offsets[r + 1]];
3556 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3557 for d in dst.iter_mut() {
3558 while nbits < bw {
3559 acc = (acc << 8) | data[idx] as u64;
3560 idx += 1;
3561 nbits += 8;
3562 }
3563 let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
3564 nbits -= bw;
3565 *d = u - l;
3566 }
3567 };
3568
3569 if a8w8_enabled() {
3570 let acts: Vec<SplitAct> =
3571 (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
3572 let acts = &acts;
3573 let out_addr = SendMut(out.as_mut_ptr());
3574 let run = move |start: usize, end: usize| {
3575 for r in start..end {
3576 let bw = bits[r] as usize;
3577 if bw == 8 {
3578 ROW_F32.with(|rb| {
3581 let mut buf = rb.borrow_mut();
3582 buf.resize(cols, 0.0);
3583 decode_f32(r, &mut buf);
3584 for bi in 0..b {
3585 let x = &xs_all[bi * cols..(bi + 1) * cols];
3586 let mut dot = 0f32;
3587 for g in 0..ng {
3588 let mut gd = 0f32;
3589 for k in 0..GROUP_SIZE {
3590 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
3591 }
3592 dot += gd * gscale(r, g);
3593 }
3594 unsafe { *out_addr.at(bi * rows + r) = dot };
3596 }
3597 });
3598 continue;
3599 }
3600 let l = (1i32 << (bw - 1)) - 1;
3601 let data = &bytes[offsets[r]..offsets[r + 1]];
3602 ROW_I8.with(|rb| {
3603 let mut buf = rb.borrow_mut();
3604 buf.resize(cols, 0);
3605 #[inline(always)]
3606 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3607 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3608 let u = unpack8::<B>(&data[blk * B..]);
3609 for k in 0..8 {
3610 chunk[k] = (u[k] - l) as i8 as u8;
3611 }
3612 }
3613 }
3614 match bw {
3615 3 => fill::<3>(data, l, &mut buf),
3616 4 => vbit_fill4(data, &mut buf),
3617 5 => fill::<5>(data, l, &mut buf),
3618 6 => fill::<6>(data, l, &mut buf),
3619 _ => unreachable!("vbit bit-width {bw} (validated at load)"),
3620 }
3621 let mut bi = 0usize;
3622 #[cfg(target_arch = "x86_64")]
3626 if avx2_enabled()
3627 && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true)
3628 {
3629 while bi + 4 <= acts.len() {
3630 let xs = [
3631 acts[bi].xq.as_slice(),
3632 acts[bi + 1].xq.as_slice(),
3633 acts[bi + 2].xq.as_slice(),
3634 acts[bi + 3].xq.as_slice(),
3635 ];
3636 let sxs = [
3637 acts[bi].sx,
3638 acts[bi + 1].sx,
3639 acts[bi + 2].sx,
3640 acts[bi + 3].sx,
3641 ];
3642 let d = unsafe {
3643 dot_q4b_row_1x4_sx_avx2(&buf, &bytes[sc_off..], r * ng, ng, xs, sxs)
3644 };
3645 for k in 0..4 {
3646 let act = &acts[bi + k];
3647 let mut dot = d[k];
3648 for &(j, xv) in &act.outliers {
3649 dot +=
3650 (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
3651 }
3652 unsafe { *out_addr.at((bi + k) * rows + r) = dot };
3654 }
3655 bi += 4;
3656 }
3657 }
3658 while bi < acts.len() {
3659 let act = &acts[bi];
3660 let mut dot = 0f32;
3661 for g in 0..ng {
3662 let d = dot_i8_i8(
3663 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3664 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3665 ) as f32
3666 * act.sx;
3667 dot += d * gscale(r, g);
3668 }
3669 for &(j, xv) in &act.outliers {
3670 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
3671 }
3672 unsafe { *out_addr.at(bi * rows + r) = dot };
3674 bi += 1;
3675 }
3676 });
3677 }
3678 };
3679 dispatch_rows(pool, rows, &run);
3680 return;
3681 }
3682
3683 let out_addr = SendMut(out.as_mut_ptr());
3684 let run = move |start: usize, end: usize| {
3685 ROW_F32.with(|rb| {
3686 let mut buf = rb.borrow_mut();
3687 buf.resize(cols, 0.0);
3688 for r in start..end {
3689 decode_f32(r, &mut buf);
3690 for bi in 0..b {
3691 let x = &xs_all[bi * cols..(bi + 1) * cols];
3692 let mut dot = 0f32;
3693 for g in 0..ng {
3694 let mut gd = 0f32;
3695 for k in 0..GROUP_SIZE {
3696 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
3697 }
3698 dot += gd * gscale(r, g);
3699 }
3700 unsafe { *out_addr.at(bi * rows + r) = dot };
3702 }
3703 }
3704 })
3705 };
3706 dispatch_rows(pool, rows, &run);
3707}
3708
3709pub(crate) fn gpu_batch_job<'a>(
3713 t: &'a QTensor,
3714 x: &[f32],
3715) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
3716 match t {
3717 QTensor::Mapped {
3718 model,
3719 idx,
3720 dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
3721 rows,
3722 cols,
3723 row_scale,
3724 col_field,
3725 ..
3726 } => Some((
3727 model.clone(),
3728 crate::gpu::BatchJob {
3729 idx: *idx,
3730 rows: *rows,
3731 cols: *cols,
3732 row_scale,
3733 xs: prescale(x, col_field, *dt).into_owned(),
3734 q1: false,
3735 },
3736 )),
3737 QTensor::Mapped {
3739 model,
3740 idx,
3741 dtype: TensorDtype::Q1,
3742 rows,
3743 cols,
3744 ..
3745 } => Some((
3746 model.clone(),
3747 crate::gpu::BatchJob {
3748 idx: *idx,
3749 rows: *rows,
3750 cols: *cols,
3751 row_scale: &[],
3752 xs: x.to_vec(),
3753 q1: true,
3754 },
3755 )),
3756 _ => None,
3757 }
3758}
3759
3760pub(crate) fn prescale<'a>(
3764 x: &'a [f32],
3765 col_field: &[f32],
3766 dtype: TensorDtype,
3767) -> std::borrow::Cow<'a, [f32]> {
3768 if dtype == TensorDtype::Q8_2f {
3769 x.iter().zip(col_field).map(|(a, c)| a * c).collect()
3770 } else {
3771 std::borrow::Cow::Borrowed(x)
3772 }
3773}
3774
3775#[cfg(target_arch = "x86_64")]
3780pub(crate) fn avx2_enabled() -> bool {
3781 use std::sync::OnceLock;
3782 static ON: OnceLock<bool> = OnceLock::new();
3783 *ON.get_or_init(|| {
3784 std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
3785 && std::arch::is_x86_feature_detected!("avx2")
3786 && std::arch::is_x86_feature_detected!("fma")
3787 })
3788}
3789
3790#[cfg(target_arch = "x86_64")]
3795fn avx2_a8w8_enabled() -> bool {
3796 use std::sync::OnceLock;
3797 static ON: OnceLock<bool> = OnceLock::new();
3798 *ON.get_or_init(|| {
3799 avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
3800 })
3801}
3802
3803#[inline]
3807pub(crate) fn a8w8_enabled() -> bool {
3808 #[cfg(target_arch = "aarch64")]
3809 {
3810 sdot_enabled()
3811 }
3812 #[cfg(target_arch = "x86_64")]
3813 {
3814 avx2_a8w8_enabled()
3815 }
3816 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
3817 {
3818 false
3819 }
3820}
3821
3822#[inline]
3825#[allow(unreachable_code)]
3826fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
3827 #[cfg(target_arch = "aarch64")]
3828 unsafe {
3829 return dot_i8_sdot(w, xq);
3830 }
3831 #[cfg(target_arch = "x86_64")]
3832 unsafe {
3833 if avx512vnni_enabled() {
3834 return dot_i8_i8_vnni(w, xq);
3835 }
3836 return dot_i8_i8_avx2(w, xq);
3837 }
3838 w.iter().zip(xq).map(|(&a, &b)| (a as i8) as i32 * b as i32).sum()
3839}
3840
3841#[cfg(target_arch = "x86_64")]
3845fn avx512vnni_enabled() -> bool {
3846 use std::sync::OnceLock;
3847 static ON: OnceLock<bool> = OnceLock::new();
3848 *ON.get_or_init(|| {
3849 std::env::var("CMF_AVX512").map(|v| v != "0").unwrap_or(true)
3850 && std::arch::is_x86_feature_detected!("avx512f")
3851 && std::arch::is_x86_feature_detected!("avx512bw")
3852 && std::arch::is_x86_feature_detected!("avx512vl")
3853 && std::arch::is_x86_feature_detected!("avx512vnni")
3854 })
3855}
3856
3857#[cfg(target_arch = "x86_64")]
3862#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
3863unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
3864 unsafe {
3866 use core::arch::x86_64::*;
3867 let n = w.len();
3868 let mut j = 0usize;
3869 let mut total: i32;
3870 {
3875 #[inline(always)]
3876 unsafe fn step(
3877 w: *const u8,
3878 x: *const i8,
3879 acc: core::arch::x86_64::__m512i,
3880 ) -> core::arch::x86_64::__m512i {
3881 unsafe {
3882 use core::arch::x86_64::*;
3883 let wv = _mm512_loadu_si512(w as *const _);
3884 let xv = _mm512_loadu_si512(x as *const _);
3885 let aw = _mm512_abs_epi8(wv);
3886 let neg = _mm512_movepi8_mask(wv);
3887 let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
3888 _mm512_dpbusd_epi32(acc, aw, sx)
3889 }
3890 }
3891 let (mut a0, mut a1, mut a2, mut a3) = (
3892 _mm512_setzero_si512(),
3893 _mm512_setzero_si512(),
3894 _mm512_setzero_si512(),
3895 _mm512_setzero_si512(),
3896 );
3897 while j + 256 <= n {
3898 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
3899 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
3900 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
3901 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
3902 j += 256;
3903 }
3904 while j + 64 <= n {
3905 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
3906 j += 64;
3907 }
3908 let s01 = _mm512_add_epi32(a0, a1);
3909 let s23 = _mm512_add_epi32(a2, a3);
3910 total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
3911 }
3912 if j + 32 <= n {
3914 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
3915 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
3916 let d = _mm256_dpbusd_epi32(
3917 _mm256_setzero_si256(),
3918 _mm256_abs_epi8(wv),
3919 _mm256_sign_epi8(xv, wv),
3920 );
3921 let hi128 = _mm256_extracti128_si256::<1>(d);
3922 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3923 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3924 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3925 total += _mm_cvtsi128_si32(s32);
3926 j += 32;
3927 }
3928 while j < n {
3929 total += (w[j] as i8) as i32 * xq[j] as i32;
3930 j += 1;
3931 }
3932 total
3933 }
3934}
3935
3936#[cfg(target_arch = "x86_64")]
3938#[target_feature(enable = "avx2,fma")]
3939unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
3940 unsafe {
3942 use core::arch::x86_64::*;
3943 let n = x.len();
3944 let wp = w.as_ptr();
3945 let xp = x.as_ptr();
3946 let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
3947 let mut j = 0usize;
3948 while j + 16 <= n {
3949 let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
3950 let lo = _mm256_cvtepi8_epi32(wb);
3951 let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
3952 a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
3953 a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
3954 j += 16;
3955 }
3956 let acc = _mm256_add_ps(a0, a1);
3957 let hi128 = _mm256_extractf128_ps::<1>(acc);
3958 let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
3959 let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
3960 let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
3961 let mut sum = _mm_cvtss_f32(s32);
3962 while j < n {
3963 sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
3964 j += 1;
3965 }
3966 sum
3967 }
3968}
3969
3970#[cfg(target_arch = "x86_64")]
3975#[target_feature(enable = "avx2")]
3976unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
3977 unsafe {
3979 use core::arch::x86_64::*;
3980 let n = w.len();
3981 let ones = _mm256_set1_epi16(1);
3982 let mut acc = _mm256_setzero_si256();
3983 let mut j = 0usize;
3984 while j + 32 <= n {
3985 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
3986 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
3987 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
3988 acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
3989 j += 32;
3990 }
3991 let hi128 = _mm256_extracti128_si256::<1>(acc);
3992 let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
3993 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3994 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3995 let mut s = _mm_cvtsi128_si32(s32);
3996 while j < n {
3997 s += (w[j] as i8) as i32 * xq[j] as i32;
3998 j += 1;
3999 }
4000 s
4001 }
4002}
4003
4004#[cfg(target_arch = "aarch64")]
4008#[target_feature(enable = "neon,i8mm")]
4009unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
4010 unsafe {
4012 use core::arch::aarch64::*;
4013 use core::arch::asm;
4014 let n = w0.len();
4015 let w0p = w0.as_ptr() as *const i8;
4016 let w1p = w1.as_ptr() as *const i8;
4017 let mut acc01 = vdupq_n_s32(0);
4020 let mut acc23 = vdupq_n_s32(0);
4021 let mut i = 0usize;
4022 while i + 8 <= n {
4023 let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
4024 let xb01 =
4025 vcombine_s8(vld1_s8(xs[0].as_ptr().add(i)), vld1_s8(xs[1].as_ptr().add(i)));
4026 let xb23 =
4027 vcombine_s8(vld1_s8(xs[2].as_ptr().add(i)), vld1_s8(xs[3].as_ptr().add(i)));
4028 asm!(
4029 "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
4030 "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
4031 a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
4032 w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
4033 options(pure, nomem, nostack),
4034 );
4035 i += 8;
4036 }
4037 let mut out = [[0i32; 4]; 2];
4038 let a01: [i32; 4] = core::mem::transmute(acc01);
4039 let a23: [i32; 4] = core::mem::transmute(acc23);
4040 out[0][0] = a01[0];
4041 out[0][1] = a01[1];
4042 out[1][0] = a01[2];
4043 out[1][1] = a01[3];
4044 out[0][2] = a23[0];
4045 out[0][3] = a23[1];
4046 out[1][2] = a23[2];
4047 out[1][3] = a23[3];
4048 if i < n {
4049 for (k, x) in xs.iter().enumerate() {
4050 for j in i..n {
4051 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
4052 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
4053 }
4054 }
4055 }
4056 out
4057 }
4058}
4059
4060#[cfg(target_arch = "aarch64")]
4064#[target_feature(enable = "neon,dotprod")]
4065unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
4066 unsafe {
4068 use core::arch::aarch64::*;
4069 use core::arch::asm;
4070 let n = w0.len();
4071 let w0p = w0.as_ptr() as *const i8;
4072 let w1p = w1.as_ptr() as *const i8;
4073 let mut acc = [[vdupq_n_s32(0); 4]; 2];
4074 let mut i = 0usize;
4075 while i + 16 <= n {
4076 let wv0 = vld1q_s8(w0p.add(i));
4077 let wv1 = vld1q_s8(w1p.add(i));
4078 for (k, x) in xs.iter().enumerate() {
4079 let xv = vld1q_s8(x.as_ptr().add(i));
4080 let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
4081 asm!(
4082 "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
4083 "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
4084 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4085 w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
4086 options(pure, nomem, nostack),
4087 );
4088 acc[0][k] = a0;
4089 acc[1][k] = a1;
4090 }
4091 i += 16;
4092 }
4093 let mut out = [[0i32; 4]; 2];
4094 for r in 0..2 {
4095 for k in 0..4 {
4096 out[r][k] = vaddvq_s32(acc[r][k]);
4097 }
4098 }
4099 if i < n {
4100 for (k, x) in xs.iter().enumerate() {
4101 for j in i..n {
4102 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
4103 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
4104 }
4105 }
4106 }
4107 out
4108 }
4109}
4110
4111#[cfg(target_arch = "x86_64")]
4117#[target_feature(enable = "avx2")]
4118unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
4119 unsafe {
4121 use core::arch::x86_64::*;
4122 let n = w0.len();
4123 let ones = _mm256_set1_epi16(1);
4124 let mut acc = [[_mm256_setzero_si256(); 4]; 2];
4125 let mut j = 0usize;
4126 while j + 32 <= n {
4127 let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
4128 let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
4129 let aw0 = _mm256_abs_epi8(wv0);
4130 let aw1 = _mm256_abs_epi8(wv1);
4131 for (k, x) in xs.iter().enumerate() {
4132 let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
4133 let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
4134 acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
4135 let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
4136 acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
4137 }
4138 j += 32;
4139 }
4140 let mut out = [[0i32; 4]; 2];
4141 for r in 0..2 {
4142 for k in 0..4 {
4143 let a = acc[r][k];
4144 let hi128 = _mm256_extracti128_si256::<1>(a);
4145 let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
4146 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4147 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4148 out[r][k] = _mm_cvtsi128_si32(s32);
4149 }
4150 }
4151 if j < n {
4152 for (k, x) in xs.iter().enumerate() {
4153 for i in j..n {
4154 out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
4155 out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
4156 }
4157 }
4158 }
4159 out
4160 }
4161}
4162
4163#[cfg(target_arch = "x86_64")]
4168#[inline]
4169fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
4170 let dot = if avx512vnni_enabled() && row.len() >= 64 {
4171 (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
4172 } else {
4173 unsafe { dot_i8_i8_avx2(row, &act.xq) }
4174 };
4175 let mut acc = dot as f32 * act.sx;
4176 for &(j, xv) in &act.outliers {
4177 acc += (row[j] as i8) as f32 * xv;
4178 }
4179 acc
4180}
4181
4182#[cfg(target_arch = "x86_64")]
4186#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4187unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
4188 unsafe {
4190 use core::arch::x86_64::*;
4191 let n = w.len();
4192 let flip = _mm512_set1_epi8(-128); #[inline(always)]
4194 unsafe fn step(
4195 w: *const u8,
4196 x: *const i8,
4197 flip: core::arch::x86_64::__m512i,
4198 acc: core::arch::x86_64::__m512i,
4199 ) -> core::arch::x86_64::__m512i {
4200 unsafe {
4201 use core::arch::x86_64::*;
4202 let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
4203 _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
4204 }
4205 }
4206 let (mut a0, mut a1, mut a2, mut a3) = (
4207 _mm512_setzero_si512(),
4208 _mm512_setzero_si512(),
4209 _mm512_setzero_si512(),
4210 _mm512_setzero_si512(),
4211 );
4212 let mut j = 0usize;
4213 while j + 256 <= n {
4214 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
4215 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
4216 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
4217 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
4218 j += 256;
4219 }
4220 while j + 64 <= n {
4221 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
4222 j += 64;
4223 }
4224 let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
4225 _mm512_add_epi32(a0, a1),
4226 _mm512_add_epi32(a2, a3),
4227 ));
4228 while j < n {
4230 total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
4231 j += 1;
4232 }
4233 total
4234 }
4235}
4236
4237#[cfg(target_arch = "x86_64")]
4243#[target_feature(enable = "avx2")]
4244unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
4245 unsafe {
4248 use core::arch::x86_64::*;
4249 let lomask = _mm_set1_epi8(0x0F);
4250 let eight = _mm256_set1_epi8(8);
4251 let ones = _mm256_set1_epi16(1);
4252 let mut acc = 0f32;
4253 for gi in 0..gpr {
4254 let g = g0 + gi;
4255 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4256 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
4257 let lo = _mm_and_si128(b, lomask);
4258 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4259 let w = _mm256_sub_epi8(
4260 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4261 eight,
4262 );
4263 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4264 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4265 let d = _mm256_madd_epi16(p16, ones);
4266 let hi128 = _mm256_extracti128_si256::<1>(d);
4267 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4268 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4269 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4270 acc += _mm_cvtsi128_si32(s32) as f32 * s;
4271 }
4272 acc
4273 }
4274}
4275
4276#[cfg(target_arch = "x86_64")]
4279#[target_feature(enable = "avx2")]
4280unsafe fn dot_q4_row_avx2_2(
4281 packed: &[u8],
4282 scales: &[u8],
4283 g0: usize,
4284 gpr: usize,
4285 xq1: &[i8],
4286 xq2: &[i8],
4287) -> (f32, f32) {
4288 unsafe {
4290 use core::arch::x86_64::*;
4291 let lomask = _mm_set1_epi8(0x0F);
4292 let eight = _mm256_set1_epi8(8);
4293 let ones = _mm256_set1_epi16(1);
4294 let (mut acc1, mut acc2) = (0f32, 0f32);
4295 #[inline(always)]
4296 unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
4297 unsafe {
4298 use core::arch::x86_64::*;
4299 let hi128 = _mm256_extracti128_si256::<1>(d);
4300 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4301 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4302 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4303 _mm_cvtsi128_si32(s32)
4304 }
4305 }
4306 for gi in 0..gpr {
4307 let g = g0 + gi;
4308 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4309 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
4310 let lo = _mm_and_si128(b, lomask);
4311 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4312 let w = _mm256_sub_epi8(
4313 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4314 eight,
4315 );
4316 let aw = _mm256_abs_epi8(w);
4317 let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4318 let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4319 let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
4320 let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
4321 acc1 += hsum(d1) as f32 * s;
4322 acc2 += hsum(d2) as f32 * s;
4323 }
4324 (acc1, acc2)
4325 }
4326}
4327
4328#[cfg(target_arch = "x86_64")]
4330fn q8_range_avx2(
4331 q: &[u8],
4332 row_scale: &[f32],
4333 act: &SplitAct,
4334 cols: usize,
4335 out_addr: SendMut,
4336 start: usize,
4337 end: usize,
4338) {
4339 for o in start..end {
4340 let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
4341 unsafe { *out_addr.at(o) = v };
4343 }
4344}
4345
4346#[cfg(target_arch = "x86_64")]
4348#[allow(clippy::too_many_arguments)]
4349fn q8_range2_avx2(
4350 q: &[u8],
4351 row_scale: &[f32],
4352 a1: &SplitAct,
4353 a2: &SplitAct,
4354 cols: usize,
4355 p1: SendMut,
4356 p2: SendMut,
4357 start: usize,
4358 end: usize,
4359) {
4360 for o in start..end {
4361 let row = &q[o * cols..(o + 1) * cols];
4362 unsafe {
4364 *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
4365 *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
4366 }
4367 }
4368}
4369
4370#[cfg(target_arch = "aarch64")]
4381fn i8mm_enabled() -> bool {
4382 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4383 *ON.get_or_init(|| {
4384 std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
4385 && std::arch::is_aarch64_feature_detected!("i8mm")
4386 })
4387}
4388
4389#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
4393fn sdot_enabled() -> bool {
4394 use std::sync::OnceLock;
4395 static ON: OnceLock<bool> = OnceLock::new();
4396 *ON.get_or_init(|| {
4397 let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
4398 #[cfg(target_arch = "aarch64")]
4399 {
4400 want && std::arch::is_aarch64_feature_detected!("dotprod")
4401 }
4402 #[cfg(not(target_arch = "aarch64"))]
4403 {
4404 let _ = want;
4405 false
4406 }
4407 })
4408}
4409
4410struct SplitAct {
4415 xq: Vec<i8>,
4416 sx: f32,
4417 outliers: Vec<(usize, f32)>,
4418 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
4421 xsum: i32,
4422}
4423
4424thread_local! {
4425 static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
4428 const { std::cell::RefCell::new(Vec::new()) };
4429}
4430
4431impl Drop for SplitAct {
4432 fn drop(&mut self) {
4433 let buf = std::mem::take(&mut self.xq);
4434 if buf.capacity() > 0 {
4435 XQ_FREE.with(|f| {
4436 let mut f = f.borrow_mut();
4437 if f.len() < 16 {
4438 f.push(buf);
4439 }
4440 });
4441 }
4442 }
4443}
4444
4445fn split_act(x: &[f32]) -> SplitAct {
4446 let n = x.len();
4447 let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
4448 let thr = 8.0 * rms;
4449 let mut outliers: Vec<(usize, f32)> = Vec::new();
4453 let mut amax = 0f32;
4454 for (j, &v) in x.iter().enumerate() {
4455 let a = v.abs();
4456 if a > thr {
4457 outliers.push((j, v));
4458 } else if a > amax {
4459 amax = a;
4460 }
4461 }
4462 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
4463 let inv = 1.0 / sx;
4464 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
4465 xq.clear();
4466 xq.reserve(n);
4467 if outliers.is_empty() {
4468 xq.extend(x.iter().map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8));
4469 } else {
4470 xq.extend(x.iter().map(|&v| {
4472 if v.abs() > thr {
4473 0
4474 } else {
4475 (v * inv).round().clamp(-127.0, 127.0) as i8
4476 }
4477 }));
4478 }
4479 let xsum = xq.iter().map(|&v| v as i32).sum();
4480 SplitAct { xq, sx, outliers, xsum }
4481}
4482
4483#[cfg(target_arch = "aarch64")]
4486#[target_feature(enable = "neon,dotprod")]
4487unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
4488 unsafe {
4490 use core::arch::aarch64::*;
4491 use core::arch::asm;
4492 let wp = w.as_ptr() as *const i8;
4493 let n = w.len();
4494 let (mut a0, mut a1, mut a2, mut a3) =
4495 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
4496 let mut i = 0;
4497 while i + 64 <= n {
4498 let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
4499 let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
4500 let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
4501 let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
4502 asm!(
4503 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
4504 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
4505 "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
4506 "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
4507 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
4508 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
4509 w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
4510 options(pure, nomem, nostack),
4511 );
4512 i += 64;
4513 }
4514 while i + 16 <= n {
4515 let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
4516 asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
4517 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
4518 i += 16;
4519 }
4520 let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
4521 while i < n {
4522 s += (*wp.add(i)) as i32 * xq[i] as i32;
4523 i += 1;
4524 }
4525 s
4526}
4527}
4528
4529#[cfg(target_arch = "aarch64")]
4533#[target_feature(enable = "neon,dotprod")]
4534unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
4535 unsafe {
4537 use core::arch::aarch64::*;
4538 use core::arch::asm;
4539 let n = xq.len();
4540 let px = xq.as_ptr();
4541 let (p0, p1, p2, p3) = (
4542 w0.as_ptr() as *const i8,
4543 w1.as_ptr() as *const i8,
4544 w2.as_ptr() as *const i8,
4545 w3.as_ptr() as *const i8,
4546 );
4547 let (mut a0, mut a1, mut a2, mut a3) =
4548 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
4549 let mut i = 0;
4550 while i + 16 <= n {
4551 let x = vld1q_s8(px.add(i));
4552 let v0 = vld1q_s8(p0.add(i));
4553 let v1 = vld1q_s8(p1.add(i));
4554 let v2 = vld1q_s8(p2.add(i));
4555 let v3 = vld1q_s8(p3.add(i));
4556 asm!(
4557 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
4558 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
4559 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
4560 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
4561 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
4562 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
4563 options(pure, nomem, nostack),
4564 );
4565 i += 16;
4566 }
4567 let mut r = [vaddvq_s32(a0), vaddvq_s32(a1), vaddvq_s32(a2), vaddvq_s32(a3)];
4568 while i < n {
4569 let xi = *px.add(i) as i32;
4570 r[0] += (*p0.add(i)) as i32 * xi;
4571 r[1] += (*p1.add(i)) as i32 * xi;
4572 r[2] += (*p2.add(i)) as i32 * xi;
4573 r[3] += (*p3.add(i)) as i32 * xi;
4574 i += 1;
4575 }
4576 r
4577}
4578}
4579
4580#[cfg(target_arch = "aarch64")]
4587#[target_feature(enable = "neon,dotprod")]
4588unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
4589 unsafe {
4592 use core::arch::aarch64::*;
4593 use core::arch::asm;
4594 let n = xq.len();
4595 let px = xq.as_ptr();
4596 let pg = g.as_ptr() as *const i8;
4597 let (mut a0, mut a1, mut a2, mut a3) =
4598 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
4599 let mut i = 0;
4600 while i + 16 <= n {
4601 let x = vld1q_s8(px.add(i));
4602 let base = pg.add(4 * i);
4603 let v0 = vld1q_s8(base);
4604 let v1 = vld1q_s8(base.add(16));
4605 let v2 = vld1q_s8(base.add(32));
4606 let v3 = vld1q_s8(base.add(48));
4607 asm!(
4608 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
4609 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
4610 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
4611 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
4612 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
4613 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
4614 options(pure, nomem, nostack),
4615 );
4616 i += 16;
4617 }
4618 [vaddvq_s32(a0), vaddvq_s32(a1), vaddvq_s32(a2), vaddvq_s32(a3)]
4619 }
4620}
4621
4622#[cfg(target_arch = "aarch64")]
4628fn q8_range_sdot(
4629 q: &[u8],
4630 rep: &[u8],
4631 row_scale: &[f32],
4632 act: &SplitAct,
4633 cols: usize,
4634 out_addr: SendMut,
4635 start: usize,
4636 end: usize,
4637) {
4638 let mut o = start;
4639 if !rep.is_empty() {
4642 while o < end && o % 4 != 0 {
4643 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
4644 unsafe { *out_addr.at(o) = v };
4645 o += 1;
4646 }
4647 }
4648 while o + 4 <= end {
4649 let r = if rep.is_empty() {
4650 unsafe {
4651 dot_i8_sdot_4rows(
4652 &q[o * cols..(o + 1) * cols],
4653 &q[(o + 1) * cols..(o + 2) * cols],
4654 &q[(o + 2) * cols..(o + 3) * cols],
4655 &q[(o + 3) * cols..(o + 4) * cols],
4656 &act.xq,
4657 )
4658 }
4659 } else {
4660 unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
4661 };
4662 for k in 0..4 {
4663 let mut acc = r[k] as f32 * act.sx;
4664 for &(j, xv) in &act.outliers {
4665 acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
4666 }
4667 unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
4669 }
4670 o += 4;
4671 }
4672 while o < end {
4673 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
4674 unsafe { *out_addr.at(o) = v };
4675 o += 1;
4676 }
4677}
4678
4679#[cfg(target_arch = "aarch64")]
4682#[allow(clippy::too_many_arguments)]
4683fn q8_range2_sdot(
4684 q: &[u8],
4685 row_scale: &[f32],
4686 a1: &SplitAct,
4687 a2: &SplitAct,
4688 cols: usize,
4689 p1: SendMut,
4690 p2: SendMut,
4691 start: usize,
4692 end: usize,
4693) {
4694 for o in start..end {
4695 let row = &q[o * cols..(o + 1) * cols];
4696 unsafe {
4698 *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
4699 *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
4700 }
4701 }
4702}
4703
4704#[allow(clippy::too_many_arguments)]
4706fn q8_range2_f32(
4707 q: &[u8],
4708 row_scale: &[f32],
4709 x1: &[f32],
4710 x2: &[f32],
4711 cols: usize,
4712 p1: SendMut,
4713 p2: SendMut,
4714 start: usize,
4715 end: usize,
4716) {
4717 for o in start..end {
4718 let row = &q[o * cols..(o + 1) * cols];
4719 unsafe {
4721 *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
4722 *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
4723 }
4724 }
4725}
4726
4727fn q8_range_f32(
4729 q: &[u8],
4730 row_scale: &[f32],
4731 xs: &[f32],
4732 cols: usize,
4733 out_addr: SendMut,
4734 start: usize,
4735 end: usize,
4736) {
4737 for o in start..end {
4738 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
4739 unsafe { *out_addr.at(o) = v };
4741 }
4742}
4743
4744#[cfg(target_arch = "aarch64")]
4747#[inline]
4748fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
4749 let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
4750 for &(j, xv) in &act.outliers {
4751 acc += (row[j] as i8) as f32 * xv;
4752 }
4753 acc
4754}
4755
4756#[cfg(target_arch = "aarch64")]
4764#[target_feature(enable = "neon,dotprod")]
4765unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
4766 unsafe {
4769 use core::arch::aarch64::*;
4770 use core::arch::asm;
4771 let lomask = vdupq_n_u8(0x0F);
4772 let eight = vdupq_n_s8(8);
4773 let mut acc = 0f32;
4774 for gi in 0..gpr {
4775 let g = g0 + gi;
4776 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4777 let b = vld1q_u8(packed.as_ptr().add(g * 16));
4778 let lo = vandq_u8(b, lomask);
4779 let hi = vshrq_n_u8::<4>(b);
4780 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4781 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4782 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4783 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4784 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4785 asm!(
4786 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4787 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4788 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4789 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4790 options(pure, nomem, nostack),
4791 );
4792 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4793 }
4794 acc
4795 }
4796}
4797
4798#[cfg(target_arch = "aarch64")]
4803#[target_feature(enable = "neon,dotprod")]
4804unsafe fn dot_q4_row_sdot2(
4805 packed: &[u8],
4806 scales: &[u8],
4807 g0: usize,
4808 gpr: usize,
4809 xq1: &[i8],
4810 xq2: &[i8],
4811) -> (f32, f32) {
4812 unsafe {
4815 use core::arch::aarch64::*;
4816 use core::arch::asm;
4817 let lomask = vdupq_n_u8(0x0F);
4818 let eight = vdupq_n_s8(8);
4819 let (mut acc1, mut acc2) = (0f32, 0f32);
4820 for gi in 0..gpr {
4821 let g = g0 + gi;
4822 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4823 let b = vld1q_u8(packed.as_ptr().add(g * 16));
4824 let lo = vandq_u8(b, lomask);
4825 let hi = vshrq_n_u8::<4>(b);
4826 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4827 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4828 let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
4829 let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
4830 let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
4831 let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
4832 let (mut a0, mut a1, mut b0, mut b1) =
4833 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
4834 asm!(
4835 "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
4836 "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
4837 "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
4838 "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
4839 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4840 b0 = inout(vreg) b0, b1 = inout(vreg) b1,
4841 e0 = in(vreg) e0, e1 = in(vreg) e1,
4842 x10 = in(vreg) x10, x11 = in(vreg) x11,
4843 x20 = in(vreg) x20, x21 = in(vreg) x21,
4844 options(pure, nomem, nostack),
4845 );
4846 acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4847 acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
4848 }
4849 (acc1, acc2)
4850 }
4851}
4852
4853#[inline]
4858pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
4859 #[cfg(target_arch = "aarch64")]
4860 unsafe {
4861 return axpy_i8_f32_neon(acc, row, w);
4862 }
4863 #[cfg(target_arch = "x86_64")]
4864 if avx2_enabled() {
4865 return unsafe { axpy_i8_f32_avx2(acc, row, w) };
4866 }
4867 #[allow(unreachable_code)]
4868 {
4869 for (a, &b) in acc.iter_mut().zip(row) {
4870 *a += w * b as f32;
4871 }
4872 }
4873}
4874
4875#[cfg(target_arch = "x86_64")]
4877#[target_feature(enable = "avx2,fma")]
4878unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
4879 unsafe {
4881 use core::arch::x86_64::*;
4882 let n = acc.len().min(row.len());
4883 let ap = acc.as_mut_ptr();
4884 let rp = row.as_ptr();
4885 let wv = _mm256_set1_ps(w);
4886 let mut j = 0usize;
4887 while j + 16 <= n {
4888 let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
4889 let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
4890 let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
4891 let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
4892 let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
4893 _mm256_storeu_ps(ap.add(j), v0);
4894 _mm256_storeu_ps(ap.add(j + 8), v1);
4895 j += 16;
4896 }
4897 while j < n {
4898 *ap.add(j) += w * (*rp.add(j)) as f32;
4899 j += 1;
4900 }
4901 }
4902}
4903
4904#[cfg(target_arch = "aarch64")]
4905#[target_feature(enable = "neon")]
4906unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
4907 unsafe {
4909 use core::arch::aarch64::*;
4910 let n = acc.len().min(row.len());
4911 let ap = acc.as_mut_ptr();
4912 let rp = row.as_ptr();
4913 let wv = vdupq_n_f32(w);
4914 let mut j = 0usize;
4915 while j + 16 <= n {
4916 let rb = vld1q_s8(rp.add(j));
4917 let lo = vmovl_s8(vget_low_s8(rb));
4918 let hi = vmovl_s8(vget_high_s8(rb));
4919 for (off, half) in [(0, lo), (8, hi)] {
4920 let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
4921 let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
4922 let o = j + off;
4923 vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
4924 vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
4925 }
4926 j += 16;
4927 }
4928 while j < n {
4929 *ap.add(j) += w * (*rp.add(j)) as f32;
4930 j += 1;
4931 }
4932}
4933}
4934
4935#[inline]
4938pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
4939 #[cfg(target_arch = "aarch64")]
4940 unsafe {
4941 return dot_i8_f32_neon(w, x);
4942 }
4943 #[cfg(target_arch = "x86_64")]
4944 if avx2_enabled() {
4945 return unsafe { dot_i8_f32_avx2(w, x) };
4946 }
4947 #[allow(unreachable_code)]
4948 {
4949 let mut sum = 0.0f32;
4950 for (j, &b) in w.iter().enumerate() {
4951 sum += (b as i8) as f32 * x[j];
4952 }
4953 sum
4954 }
4955}
4956
4957#[inline]
4961fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
4962 #[cfg(target_arch = "aarch64")]
4963 unsafe {
4964 return dot_i8_col_f32_neon(w, x, col);
4965 }
4966 #[allow(unreachable_code)]
4967 {
4968 let mut sum = 0.0f32;
4969 for (j, &b) in w.iter().enumerate() {
4970 sum += (b as i8) as f32 * x[j] * col[j];
4971 }
4972 sum
4973 }
4974}
4975
4976#[cfg(target_arch = "aarch64")]
4977#[target_feature(enable = "neon")]
4978unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
4979 unsafe {
4981 use core::arch::aarch64::*;
4982 let n = x.len();
4983 let wp = w.as_ptr() as *const i8;
4984 let xp = x.as_ptr();
4985 let cp = col.as_ptr();
4986 let (mut a0, mut a1, mut a2, mut a3) =
4987 (vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4988 let mut j = 0usize;
4989 while j + 16 <= n {
4990 let wb = vld1q_s8(wp.add(j));
4991 let lo = vmovl_s8(vget_low_s8(wb));
4992 let hi = vmovl_s8(vget_high_s8(wb));
4993 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
4994 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
4995 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
4996 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
4997 a0 = vfmaq_f32(a0, w0, vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))));
4998 a1 = vfmaq_f32(a1, w1, vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))));
4999 a2 = vfmaq_f32(a2, w2, vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))));
5000 a3 = vfmaq_f32(a3, w3, vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))));
5001 j += 16;
5002 }
5003 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
5004 while j < n {
5005 sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
5006 j += 1;
5007 }
5008 sum
5009}
5010}
5011
5012#[cfg(target_arch = "aarch64")]
5013#[target_feature(enable = "neon")]
5014unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
5015 unsafe {
5017 use core::arch::aarch64::*;
5018 let n = x.len();
5019 let wp = w.as_ptr() as *const i8;
5020 let xp = x.as_ptr();
5021 let (mut a0, mut a1, mut a2, mut a3) =
5022 (vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0));
5023 let mut j = 0usize;
5024 while j + 16 <= n {
5025 let wb = vld1q_s8(wp.add(j));
5026 let lo = vmovl_s8(vget_low_s8(wb));
5027 let hi = vmovl_s8(vget_high_s8(wb));
5028 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
5029 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
5030 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
5031 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
5032 a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
5033 a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
5034 a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
5035 a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
5036 j += 16;
5037 }
5038 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
5039 while j < n {
5040 sum += (*wp.add(j)) as f32 * *xp.add(j);
5041 j += 1;
5042 }
5043 sum
5044}
5045}
5046
5047#[allow(clippy::too_many_arguments)]
5048fn qmatvec(
5049 q: &[u8],
5050 rep: &[u8],
5051 row_scale: &[f32],
5052 xs: &[f32],
5053 rows: usize,
5054 cols: usize,
5055 out: &mut [f32],
5056 pool: Option<&Pool>,
5057) {
5058 debug_assert_eq!(out.len(), rows);
5059 #[cfg(not(target_arch = "aarch64"))]
5060 let _ = rep;
5061
5062 #[cfg(target_arch = "aarch64")]
5063 if sdot_enabled() {
5064 let act = split_act(xs);
5065 let out_addr = SendMut(out.as_mut_ptr());
5066 let run_range = |start: usize, end: usize| {
5067 q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
5068 };
5069 match pool {
5070 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
5071 _ => run_range(0, rows),
5072 }
5073 return;
5074 }
5075 #[cfg(target_arch = "x86_64")]
5078 if avx2_a8w8_enabled() {
5079 let act = split_act(xs);
5080 let out_addr = SendMut(out.as_mut_ptr());
5081 let run_range =
5082 |start: usize, end: usize| q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end);
5083 match pool {
5084 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
5085 _ => run_range(0, rows),
5086 }
5087 return;
5088 }
5089
5090 let row_dot = |o: usize| -> f32 { dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o] };
5091 match pool {
5092 Some(pool) if rows >= 256 => {
5093 let out_addr = SendMut(out.as_mut_ptr());
5094 pool.run(&move |widx, n| {
5095 let chunk = rows.div_ceil(n);
5096 let start = widx * chunk;
5097 let end = (start + chunk).min(rows);
5098 for o in start..end {
5099 unsafe { *out_addr.at(o) = row_dot(o) };
5101 }
5102 });
5103 }
5104 _ => {
5105 for (o, dst) in out.iter_mut().enumerate() {
5106 *dst = row_dot(o);
5107 }
5108 }
5109 }
5110}
5111
5112#[allow(clippy::too_many_arguments)]
5113fn qmatvec2(
5114 q: &[u8],
5115 row_scale: &[f32],
5116 x1: &[f32],
5117 x2: &[f32],
5118 rows: usize,
5119 cols: usize,
5120 o1: &mut [f32],
5121 o2: &mut [f32],
5122 pool: Option<&Pool>,
5123) {
5124 #[cfg(target_arch = "aarch64")]
5125 if sdot_enabled() {
5126 let a1s = split_act(x1);
5127 let a2s = split_act(x2);
5128 let p1 = SendMut(o1.as_mut_ptr());
5129 let p2 = SendMut(o2.as_mut_ptr());
5130 let run_range = |start: usize, end: usize| {
5131 q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
5132 };
5133 match pool {
5134 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
5135 _ => run_range(0, rows),
5136 }
5137 return;
5138 }
5139 #[cfg(target_arch = "x86_64")]
5140 if avx2_a8w8_enabled() {
5141 let a1s = split_act(x1);
5142 let a2s = split_act(x2);
5143 let p1 = SendMut(o1.as_mut_ptr());
5144 let p2 = SendMut(o2.as_mut_ptr());
5145 let run_range = |start: usize, end: usize| {
5146 q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
5147 };
5148 match pool {
5149 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
5150 _ => run_range(0, rows),
5151 }
5152 return;
5153 }
5154
5155 let row_dots = |o: usize| -> (f32, f32) {
5156 let row = &q[o * cols..(o + 1) * cols];
5157 (dot_i8_f32(row, x1) * row_scale[o], dot_i8_f32(row, x2) * row_scale[o])
5158 };
5159 match pool {
5160 Some(pool) if rows >= 256 => {
5161 let p1 = SendMut(o1.as_mut_ptr());
5162 let p2 = SendMut(o2.as_mut_ptr());
5163 pool.run(&move |widx, n| {
5164 let chunk = rows.div_ceil(n);
5165 let start = widx * chunk;
5166 let end = (start + chunk).min(rows);
5167 for o in start..end {
5168 let (s1, s2) = row_dots(o);
5169 unsafe {
5171 *p1.at(o) = s1;
5172 *p2.at(o) = s2;
5173 }
5174 }
5175 });
5176 }
5177 _ => {
5178 for o in 0..rows {
5179 let (s1, s2) = row_dots(o);
5180 o1[o] = s1;
5181 o2[o] = s2;
5182 }
5183 }
5184 }
5185}
5186
5187#[derive(Clone, Copy)]
5188struct SendMut(*mut f32);
5189unsafe impl Send for SendMut {}
5190unsafe impl Sync for SendMut {}
5191
5192impl SendMut {
5193 #[inline]
5194 fn at(self, i: usize) -> *mut f32 {
5195 unsafe { self.0.add(i) }
5196 }
5197}
5198
5199#[cfg(test)]
5200mod tests {
5201 use super::*;
5202
5203 #[test]
5204 fn f32_matvec_matches_matvec_rows_bitexact() {
5205 let (rows, cols) = (300, 40);
5206 let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
5207 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
5208 let qt = QTensor::from_f32(w.clone(), rows, cols);
5209
5210 let mut a = vec![0.0f32; rows];
5211 matvec_rows(None, &w, &x, &mut a);
5212 let mut b = vec![0.0f32; rows];
5213 qt.matvec(&x, &mut b, None);
5214 assert_eq!(a, b);
5215 }
5216
5217 #[test]
5218 fn sdot_kernel_exact_on_grid() {
5219 eprintln!("sdot_enabled = {}", sdot_enabled());
5224 let (rows, cols) = (9, 80); let w: Vec<u8> = (0..rows * cols)
5226 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
5227 .collect();
5228 let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
5229 let x: Vec<f32> = (0..cols)
5230 .map(|i| match i % 3 {
5231 0 => 1.0,
5232 1 => -1.0,
5233 _ => 0.0,
5234 })
5235 .collect();
5236 let mut a = vec![0.0f32; rows];
5237 qmatvec(&w, &[], &scales, &x, rows, cols, &mut a, None);
5238 for o in 0..rows {
5239 let mut acc = 0.0f32;
5240 for j in 0..cols {
5241 acc += (w[o * cols + j] as i8) as f32 * x[j];
5242 }
5243 let expect = acc * scales[o];
5244 assert!(
5245 (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
5246 "row {o}: {} vs {expect}",
5247 a[o]
5248 );
5249 }
5250 }
5251
5252 #[test]
5253 fn q1_kernels_match_exact_reference() {
5254 let (rows, cols) = (7, 96);
5256 let gpr = cols / GROUP_SIZE;
5257 let mut bytes = Vec::new();
5258 for t in 0..rows * gpr {
5259 let s = 0.01 + (t % 13) as f32 * 0.003;
5260 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
5261 for j in 0..4 {
5262 bytes.push(((t * 31 + j * 97) % 251) as u8);
5263 }
5264 }
5265 let x: Vec<f32> = (0..cols)
5267 .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
5268 .collect();
5269 let mut w = vec![0.0f32; rows * cols];
5271 cortiq_core::quant::dequant_q1(&bytes, &mut w);
5272 let mut expect = vec![0.0f32; rows];
5273 for o in 0..rows {
5274 expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
5275 }
5276 let mut got = vec![0.0f32; rows];
5277 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
5278 for o in 0..rows {
5279 assert!(
5280 (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
5281 "row {o}: {} vs {}",
5282 got[o],
5283 expect[o]
5284 );
5285 }
5286 let x2: Vec<f32> = x.iter().map(|v| -v).collect();
5288 let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
5289 q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
5290 assert_eq!(a1, got);
5291 let mut xs = x.clone();
5292 xs.extend_from_slice(&x2);
5293 let mut mm = vec![0.0f32; 2 * rows];
5294 q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
5295 assert_eq!(&mm[..rows], got.as_slice());
5296 assert_eq!(&mm[rows..], a2.as_slice());
5297 }
5298
5299 #[test]
5300 fn repack_is_bit_identical() {
5301 let (rows, cols) = (267, 96); let w: Vec<u8> = (0..rows * cols)
5307 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
5308 .collect();
5309 let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
5310 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
5311 let rep = q8_repack_layout(&w, rows, cols);
5312 for g in 0..rows / 4 {
5314 for c in 0..cols / 16 {
5315 for lane in 0..4 {
5316 assert_eq!(
5317 &rep[g * 4 * cols + c * 64 + lane * 16..g * 4 * cols + c * 64 + lane * 16 + 16],
5318 &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
5319 );
5320 }
5321 }
5322 }
5323 let mut a = vec![0.0f32; rows];
5324 qmatvec(&w, &[], &scales, &x, rows, cols, &mut a, None);
5325 let mut b = vec![0.0f32; rows];
5326 qmatvec(&w, &rep, &scales, &x, rows, cols, &mut b, None);
5327 assert_eq!(a, b, "full-range repack output diverged");
5328
5329 #[cfg(target_arch = "aarch64")]
5330 if sdot_enabled() {
5331 let act = split_act(&x);
5333 let mut c1 = vec![0.0f32; rows];
5334 let mut c2 = vec![0.0f32; rows];
5335 q8_range_sdot(&w, &[], &scales, &act, cols, SendMut(c1.as_mut_ptr()), 3, rows - 2);
5336 q8_range_sdot(&w, &rep, &scales, &act, cols, SendMut(c2.as_mut_ptr()), 3, rows - 2);
5337 assert_eq!(c1, c2, "unaligned-range repack output diverged");
5338 }
5339 }
5340
5341 #[test]
5342 fn sdot_a8w8_noise_is_bounded() {
5343 let (rows, cols) = (16, 512);
5347 let w: Vec<u8> = (0..rows * cols)
5348 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
5349 .collect();
5350 let scales = vec![0.01f32; rows];
5351 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
5352 let mut a = vec![0.0f32; rows];
5353 qmatvec(&w, &[], &scales, &x, rows, cols, &mut a, None);
5354 let (mut num, mut den) = (0f64, 0f64);
5355 for o in 0..rows {
5356 let mut acc = 0.0f32;
5357 for j in 0..cols {
5358 acc += (w[o * cols + j] as i8) as f32 * x[j];
5359 }
5360 let expect = acc * scales[o];
5361 num += ((a[o] - expect) as f64).powi(2);
5362 den += (expect as f64).powi(2);
5363 }
5364 let rel = (num / den.max(1e-12)).sqrt();
5365 assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
5366 }
5367
5368 #[test]
5369 fn i8_dot_neon_matches_scalar() {
5370 let n = 100;
5371 let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
5372 let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
5373 let mut scalar = 0.0f32;
5374 for j in 0..n {
5375 scalar += (w[j] as i8) as f32 * x[j];
5376 }
5377 let fast = dot_i8_f32(&w, &x);
5378 assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
5379 }
5380
5381 #[test]
5383 fn vbitmatvec_matches_full_dequant() {
5384 let (rows, cols) = (6, 64);
5385 let ng = cols / GROUP_SIZE;
5386 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
5388 let mut bytes = bits.clone();
5389 for g in 0..rows * ng {
5390 let s = 0.02 + 0.001 * g as f32;
5391 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
5392 }
5393 for r in 0..rows {
5394 let b = bits[r] as usize;
5395 let (mut acc, mut nb) = (0u64, 0usize);
5396 let mut rowbytes = Vec::new();
5397 for i in 0..cols {
5398 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
5399 acc = (acc << b) | v;
5400 nb += b;
5401 while nb >= 8 {
5402 nb -= 8;
5403 rowbytes.push(((acc >> nb) & 0xFF) as u8);
5404 }
5405 }
5406 if nb > 0 {
5407 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
5408 }
5409 bytes.extend_from_slice(&rowbytes);
5410 }
5411 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
5412
5413 let mut reference = vec![0f32; rows * cols];
5414 cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
5415 let mut expect = vec![0f32; rows];
5416 for r in 0..rows {
5417 expect[r] = reference[r * cols..(r + 1) * cols]
5418 .iter()
5419 .zip(&x)
5420 .map(|(w, xv)| w * xv)
5421 .sum();
5422 }
5423 let mut got = vec![0f32; rows];
5424 let offsets = vbit_row_offsets(&bytes, rows, cols);
5425 vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
5426 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
5430 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
5431 for r in 0..rows {
5432 assert!(
5433 (got[r] - expect[r]).abs() < tol * scale,
5434 "row {r}: {} vs {}",
5435 got[r],
5436 expect[r]
5437 );
5438 }
5439 }
5440
5441 #[test]
5446 #[cfg(target_arch = "x86_64")]
5447 fn vbit_matmat_blocked_matches_per_row() {
5448 let (rows, cols, b) = (64usize, 128usize, 9usize);
5449 let ng = cols / GROUP_SIZE;
5450 let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
5451 let mut bytes = bits.clone();
5452 for g in 0..rows * ng {
5453 let sc = 0.02 + 0.0005 * g as f32;
5454 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
5455 }
5456 for r in 0..rows {
5457 let bw = bits[r] as usize;
5458 let (mut acc, mut nb) = (0u64, 0usize);
5459 let mut rowbytes = Vec::new();
5460 for i in 0..cols {
5461 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
5462 acc = (acc << bw) | v;
5463 nb += bw;
5464 while nb >= 8 {
5465 nb -= 8;
5466 rowbytes.push(((acc >> nb) & 0xFF) as u8);
5467 }
5468 }
5469 if nb > 0 {
5470 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
5471 }
5472 bytes.extend_from_slice(&rowbytes);
5473 }
5474 let x: Vec<f32> =
5475 (0..b * cols).map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5).collect();
5476 let offsets = vbit_row_offsets(&bytes, rows, cols);
5477 let mut y_a = vec![0f32; b * rows];
5478 let mut y_b = vec![0f32; b * rows];
5479 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
5480 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
5481 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
5482 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
5483 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
5484 let max_d =
5485 y_a.iter().zip(&y_b).map(|(p, q)| (p - q).abs()).fold(0.0f32, f32::max);
5486 assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
5487 }
5488
5489 #[test]
5490 fn q4matvec_matches_full_dequant() {
5491 let (rows, cols) = (8, 64);
5492 let groups = rows * cols / GROUP_SIZE;
5493 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
5495 for i in 0..groups * 16 {
5496 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
5497 }
5498 for g in 0..groups {
5499 let s = 0.01 + 0.003 * g as f32;
5500 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
5501 }
5502 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
5503
5504 let mut reference = vec![0.0f32; rows * cols];
5505 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
5506 let mut expect = vec![0.0f32; rows];
5507 for r in 0..rows {
5508 expect[r] = reference[r * cols..(r + 1) * cols]
5509 .iter()
5510 .zip(&x)
5511 .map(|(w, xv)| w * xv)
5512 .sum();
5513 }
5514
5515 let mut got = vec![0.0f32; rows];
5516 q4matvec(&bytes, &x, rows, cols, &mut got, None);
5517 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
5521 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
5522 for r in 0..rows {
5523 assert!(
5524 (got[r] - expect[r]).abs() < tol * scale,
5525 "row {r}: {} vs {}",
5526 got[r],
5527 expect[r]
5528 );
5529 }
5530 }
5531
5532 #[test]
5535 fn vbitmatvec2_equals_two_singles() {
5536 let (rows, cols) = (6, 64);
5537 let ng = cols / GROUP_SIZE;
5538 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
5539 let mut bytes = bits.clone();
5540 for g in 0..rows * ng {
5541 let s = 0.02 + 0.001 * g as f32;
5542 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
5543 }
5544 for r in 0..rows {
5545 let b = bits[r] as usize;
5546 let (mut acc, mut nb) = (0u64, 0usize);
5547 let mut rowbytes = Vec::new();
5548 for i in 0..cols {
5549 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
5550 acc = (acc << b) | v;
5551 nb += b;
5552 while nb >= 8 {
5553 nb -= 8;
5554 rowbytes.push(((acc >> nb) & 0xFF) as u8);
5555 }
5556 }
5557 if nb > 0 {
5558 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
5559 }
5560 bytes.extend_from_slice(&rowbytes);
5561 }
5562 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
5563 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
5564 let offsets = vbit_row_offsets(&bytes, rows, cols);
5565
5566 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
5567 vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
5568 vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
5569 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
5570 vbitmatvec2(&bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
5571 assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
5572 assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
5573 }
5574
5575 #[test]
5577 fn q4matvec2_equals_two_singles() {
5578 let (rows, cols) = (8, 128);
5579 let groups = rows * cols / GROUP_SIZE;
5580 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
5581 for i in 0..groups * 16 {
5582 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
5583 }
5584 for g in 0..groups {
5585 let s = 0.01 + 0.003 * g as f32;
5586 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
5587 }
5588 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
5591 x1[9] = 250.0;
5592 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
5593
5594 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
5595 q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
5596 q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
5597 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
5598 q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
5599 assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
5600 assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
5601 }
5602
5603 #[test]
5606 fn matvec_many_equals_separate_matvecs() {
5607 use crate::pool::Pool;
5608 let (r1, r2, cols) = (300, 200, 64);
5609 let mk = |salt: usize, rows: usize| {
5610 QTensor::from_f32(
5611 (0..rows * cols).map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5).collect(),
5612 rows,
5613 cols,
5614 )
5615 };
5616 let (a, b) = (mk(1, r1), mk(5, r2));
5617 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
5618 let pool = Pool::new(3);
5619
5620 let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
5621 a.matvec(&x, &mut ea, Some(&pool));
5622 b.matvec(&x, &mut eb, Some(&pool));
5623 let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
5624 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
5625 assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
5626 assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
5627 }
5628
5629 #[test]
5632 fn batched_matmat_equals_per_position_matvec() {
5633 let (rows, cols, b) = (8, 64, 5);
5634 let groups = rows * cols / GROUP_SIZE;
5636 let mut q4 = Vec::new();
5637 for i in 0..groups * 16 {
5638 q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
5639 }
5640 for g in 0..groups {
5641 q4.extend_from_slice(
5642 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
5643 );
5644 }
5645 let ng = cols / GROUP_SIZE;
5647 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
5648 let mut vb = bits.clone();
5649 for g in 0..rows * ng {
5650 vb.extend_from_slice(
5651 &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
5652 );
5653 }
5654 for r in 0..rows {
5655 let bw = bits[r] as usize;
5656 let (mut acc, mut nb) = (0u64, 0usize);
5657 let mut rowbytes = Vec::new();
5658 for i in 0..cols {
5659 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
5660 acc = (acc << bw) | v;
5661 nb += bw;
5662 while nb >= 8 {
5663 nb -= 8;
5664 rowbytes.push(((acc >> nb) & 0xFF) as u8);
5665 }
5666 }
5667 if nb > 0 {
5668 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
5669 }
5670 vb.extend_from_slice(&rowbytes);
5671 }
5672 let offsets = vbit_row_offsets(&vb, rows, cols);
5673
5674 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
5675
5676 let mut got = vec![0f32; b * rows];
5678 q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
5679 for bi in 0..b {
5680 let mut expect = vec![0f32; rows];
5681 q4matvec(&q4, &xs[bi * cols..(bi + 1) * cols], rows, cols, &mut expect, None);
5682 assert_eq!(&got[bi * rows..(bi + 1) * rows], &expect[..], "q4 batch pos {bi}");
5683 }
5684
5685 let mut got = vec![0f32; b * rows];
5687 vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
5688 for bi in 0..b {
5689 let mut expect = vec![0f32; rows];
5690 vbitmatvec(
5691 &vb, &offsets, &xs[bi * cols..(bi + 1) * cols], rows, cols, &mut expect, None,
5692 );
5693 assert_eq!(&got[bi * rows..(bi + 1) * rows], &expect[..], "vbit batch pos {bi}");
5694 }
5695 }
5696
5697 #[test]
5701 fn q4_tiled_matches_q4_block_bitexact() {
5702 let (rows, cols, b) = (8usize, 128usize, 3usize);
5703 let groups = rows * cols / GROUP_SIZE;
5704 let mut split = Vec::with_capacity(groups * 18);
5705 for i in 0..groups * 16 {
5706 split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
5707 }
5708 for g in 0..groups {
5709 split.extend_from_slice(
5710 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
5711 );
5712 }
5713 let (packed, scales) = split.split_at(groups * 16);
5715 let mut tiled = Vec::with_capacity(groups * Q4_TILE);
5716 for g in 0..groups {
5717 tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
5718 tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
5719 }
5720
5721 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
5722 x1[9] = 250.0; let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
5724
5725 let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
5726 q4matvec(&split, &x1, rows, cols, &mut a, None);
5727 q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
5728 assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
5729
5730 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
5731 let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
5732 q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
5733 q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
5734 assert_eq!(a1, t1);
5735 assert_eq!(a2, t2);
5736
5737 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
5738 let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
5739 q4matmat(&split, &xs, b, rows, cols, &mut am, None);
5740 q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
5741 assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
5742 }
5743
5744 #[test]
5751 fn q4matvec_sdot_outlier_exact() {
5752 let (rows, cols) = (4, 128);
5753 let groups = rows * cols / GROUP_SIZE;
5754 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
5755 for i in 0..groups * 16 {
5756 bytes.push(((i * 11 + 5) % 256) as u8);
5757 }
5758 for g in 0..groups {
5759 let s = 0.02 + 0.002 * g as f32;
5760 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
5761 }
5762 let mut x: Vec<f32> = (0..cols)
5763 .map(|i| match i % 3 {
5764 0 => 1.0,
5765 1 => -1.0,
5766 _ => 0.0,
5767 })
5768 .collect();
5769 x[17] = 300.0; let mut reference = vec![0.0f32; rows * cols];
5772 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
5773 let mut expect = vec![0.0f32; rows];
5774 for r in 0..rows {
5775 expect[r] = reference[r * cols..(r + 1) * cols]
5776 .iter()
5777 .zip(&x)
5778 .map(|(w, xv)| w * xv)
5779 .sum();
5780 }
5781 let mut got = vec![0.0f32; rows];
5782 q4matvec(&bytes, &x, rows, cols, &mut got, None);
5783 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
5784 for r in 0..rows {
5785 assert!(
5786 (got[r] - expect[r]).abs() < 2e-3 * scale,
5787 "row {r}: {} vs {} (outlier term must be exact)",
5788 got[r],
5789 expect[r]
5790 );
5791 }
5792 }
5793}