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