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