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(cfg!(target_os = "android"))
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 TensorDtype::Q1T if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
247 model: model.clone(),
248 idx,
249 dtype: entry.dtype,
250 rows,
251 cols,
252 row_scale: Vec::new(),
253 col_field: Vec::new(),
254 vbit_offsets: Vec::new(),
255 repack: Vec::new(),
256 }),
257 _ => {
259 let mut data = vec![0.0f32; rows * cols];
260 cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
261 Ok(Self::from_f32(data, rows, cols))
262 }
263 }
264 }
265
266 pub(crate) fn is_q1(&self) -> bool {
269 matches!(self, Self::Mapped { dtype: TensorDtype::Q1, .. })
270 }
271
272 pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
275 match self {
276 Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
277 _ => None,
278 }
279 }
280
281 pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
288 match self {
289 Self::Mapped {
290 idx,
291 dtype: TensorDtype::Q1 | TensorDtype::Q1T | TensorDtype::Q4Block | TensorDtype::Q8Row,
292 rows,
293 cols,
294 ..
295 } => Some((*idx, *rows, *cols)),
296 _ => None,
297 }
298 }
299
300 pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
305 match self {
306 Self::Mapped {
307 idx,
308 dtype: TensorDtype::Q8Row,
309 rows,
310 cols,
311 row_scale,
312 col_field,
313 ..
314 } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
315 _ => None,
316 }
317 }
318
319 pub fn rows(&self) -> usize {
320 match self {
321 Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
322 }
323 }
324
325 pub fn cols(&self) -> usize {
326 match self {
327 Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
328 }
329 }
330
331 pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
334 match self {
335 Self::Mapped { model, idx, dtype: TensorDtype::Q1, .. } => Some((model, *idx)),
336 _ => None,
337 }
338 }
339
340 pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
344 match self {
345 Self::Mapped { model, idx, dtype: TensorDtype::Q8Row, row_scale, .. } => {
346 Some((model, *idx, 0, row_scale.as_slice()))
347 }
348 Self::Mapped { model, idx, dtype: TensorDtype::Q1, .. } => Some((model, *idx, 1, &[])),
349 Self::Mapped { model, idx, dtype: TensorDtype::Q4Tiled, .. } => Some((model, *idx, 2, &[])),
350 Self::Mapped { model, idx, dtype: TensorDtype::Q1T, .. } => Some((model, *idx, 3, &[])),
351 _ => None,
352 }
353 }
354
355 pub fn as_f32(&self) -> Option<&[f32]> {
358 match self {
359 Self::F32 { data, .. } => Some(data),
360 Self::Mapped { .. } => None,
361 }
362 }
363
364 fn quant_bytes(&self) -> &[u8] {
365 match self {
366 Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
367 Self::F32 { .. } => unreachable!("quant_bytes on F32"),
368 }
369 }
370
371 pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
373 let cols = self.cols();
374 debug_assert_eq!(dst.len(), cols);
375 match self {
376 Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
377 Self::Mapped {
378 dtype,
379 row_scale,
380 col_field,
381 vbit_offsets,
382 ..
383 } => {
384 if *dtype == TensorDtype::Q4Tiled {
385 let bytes = self.quant_bytes();
386 let gpr = cols / GROUP_SIZE;
387 for gi in 0..gpr {
388 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
389 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
390 for (k, &b) in tile[2..].iter().enumerate() {
391 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
392 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
393 }
394 }
395 return;
396 }
397 if *dtype == TensorDtype::Q4Block {
398 let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
399 let gpr = cols / GROUP_SIZE;
400 for gi in 0..gpr {
401 let g = r * gpr + gi;
402 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
403 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
404 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
405 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
406 }
407 }
408 return;
409 }
410 if *dtype == TensorDtype::Q1 {
411 let bytes = self.quant_bytes();
412 let gpr = cols / GROUP_SIZE;
413 for gi in 0..gpr {
414 let tile =
415 &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
416 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
417 for (j, &b) in tile[2..].iter().enumerate() {
418 for k in 0..8 {
419 dst[gi * GROUP_SIZE + j * 8 + k] =
420 (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
421 }
422 }
423 }
424 return;
425 }
426 if *dtype == TensorDtype::Q1T {
427 let bytes = self.quant_bytes();
428 let gpr = cols / GROUP_SIZE;
429 let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
430 for gi in 0..gpr {
431 let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
432 let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
433 let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
434 for k in 0..GROUP_SIZE {
435 dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k) {
436 1 => s,
437 2 => -s,
438 _ => 0.0,
439 };
440 }
441 }
442 let rows = self.rows();
444 let entries = base_len + (rows + 1) * 4;
445 if entries <= bytes.len() {
446 let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
447 let r0 = u32::from_le_bytes([ptrs[r * 4], ptrs[r * 4 + 1], ptrs[r * 4 + 2], ptrs[r * 4 + 3]]) as usize;
448 let r1 = u32::from_le_bytes([ptrs[(r + 1) * 4], ptrs[(r + 1) * 4 + 1], ptrs[(r + 1) * 4 + 2], ptrs[(r + 1) * 4 + 3]]) as usize;
449 let off = entries + r0 * 4;
450 for i in 0..r1 - r0 {
451 let item = &bytes[off + i * 4..off + i * 4 + 4];
452 let c = u16::from_le_bytes([item[0], item[1]]) as usize;
453 let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([item[2], item[3]]));
454 if c < cols {
455 dst[c] = v;
456 }
457 }
458 }
459 return;
460 }
461 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
462 let bytes = self.quant_bytes();
463 let rows = self.rows();
464 let ng = cols / GROUP_SIZE;
465 let bits = &bytes[..rows];
466 let sc_off = rows;
467 let off = vbit_offsets[r];
470 let b = bits[r] as usize;
471 let l = ((1usize << (b - 1)) - 1) as f32;
472 let data = &bytes[off..];
473 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
474 for (i, d) in dst.iter_mut().enumerate() {
475 while nbits < b {
476 acc = (acc << 8) | data[idx] as u64;
477 idx += 1;
478 nbits += 8;
479 }
480 let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
481 nbits -= b;
482 let so = (r * ng + i / GROUP_SIZE) * 2;
483 let sv = f16_to_f32(u16::from_le_bytes([
484 bytes[sc_off + so],
485 bytes[sc_off + so + 1],
486 ]));
487 *d = (u - l) * sv;
488 }
489 return;
490 }
491 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
492 let s = row_scale[r];
493 match dtype {
494 TensorDtype::Q8Row => {
495 for (d, &b) in dst.iter_mut().zip(q) {
496 *d = (b as i8) as f32 * s;
497 }
498 }
499 TensorDtype::Q8_2f => {
500 for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
501 *d = (b as i8) as f32 * s * col_field[i];
502 }
503 }
504 _ => unreachable!(),
505 }
506 }
507 }
508 }
509
510 pub fn sparse_col_ok(&self) -> bool {
515 match self {
516 Self::F32 { .. } => true,
517 Self::Mapped { dtype, .. } => {
518 matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
519 }
520 }
521 }
522
523 pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
527 let inter = self.cols();
528 let hidden = self.rows();
529 debug_assert_eq!(out.len(), hidden);
530 match self {
531 Self::F32 { data, .. } => {
532 for (k, o) in out.iter_mut().enumerate() {
533 *o += w * data[k * inter + c];
534 }
535 }
536 Self::Mapped {
537 dtype, row_scale, col_field, ..
538 } => {
539 let q = self.quant_bytes();
540 let colf = if *dtype == TensorDtype::Q8_2f {
541 col_field[c]
542 } else {
543 1.0
544 };
545 let wc = w * colf;
546 for (k, o) in out.iter_mut().enumerate() {
547 let b = q[k * inter + c] as i8 as f32;
548 *o += wc * b * row_scale[k];
549 }
550 }
551 }
552 }
553
554 pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
558 let cols = self.cols();
559 match self {
560 Self::F32 { data, .. } => {
561 let row = &data[r * cols..(r + 1) * cols];
562 row.iter().zip(x).map(|(w, v)| w * v).sum()
563 }
564 Self::Mapped { dtype, row_scale, col_field, .. } => match dtype {
565 TensorDtype::Q8Row => {
566 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
567 dot_i8_f32(q, x) * row_scale[r]
568 }
569 TensorDtype::Q8_2f => {
570 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
571 dot_i8_col_f32(q, x, col_field) * row_scale[r]
572 }
573 _ => {
574 self.row_f32(r, scratch);
575 scratch.iter().zip(x).map(|(w, v)| w * v).sum()
576 }
577 },
578 }
579 }
580
581 pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
584 match self {
585 Self::F32 { data, .. } => matvec_rows(pool, data, x, out),
586 Self::Mapped {
587 model,
588 idx,
589 dtype,
590 rows,
591 cols,
592 row_scale,
593 col_field,
594 vbit_offsets,
595 repack,
596 } => {
597 let _ = (model, idx);
598 if *dtype == TensorDtype::Q4Block {
599 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
603 let t0 = std::time::Instant::now();
604 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
605 crate::gpu::ProbeArm::Gpu => {
606 if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
607 crate::gpu::probe_record(
608 crate::gpu::OpClass::Matvec,
609 true,
610 t0.elapsed(),
611 );
612 return;
613 }
614 }
615 crate::gpu::ProbeArm::CpuTimed => {
616 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
617 crate::gpu::probe_record(
618 crate::gpu::OpClass::Matvec,
619 false,
620 t0.elapsed(),
621 );
622 return;
623 }
624 crate::gpu::ProbeArm::Cpu => {}
625 }
626 }
627 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
628 return;
629 }
630 if *dtype == TensorDtype::Q4Tiled {
631 q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
632 return;
633 }
634 if *dtype == TensorDtype::Q1 {
635 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
640 let t0 = std::time::Instant::now();
641 let arm = if crate::gpu::q1_force() {
642 crate::gpu::ProbeArm::Gpu
643 } else {
644 crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
645 };
646 match arm {
647 crate::gpu::ProbeArm::Gpu => {
648 if crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out) {
649 crate::gpu::probe_record(
650 crate::gpu::OpClass::Matvec,
651 true,
652 t0.elapsed(),
653 );
654 return;
655 }
656 }
657 crate::gpu::ProbeArm::CpuTimed => {
658 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
659 crate::gpu::probe_record(
660 crate::gpu::OpClass::Matvec,
661 false,
662 t0.elapsed(),
663 );
664 return;
665 }
666 crate::gpu::ProbeArm::Cpu => {}
667 }
668 }
669 q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
670 return;
671 }
672 if *dtype == TensorDtype::Q1T {
673 if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
677 let t0 = std::time::Instant::now();
678 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
679 crate::gpu::ProbeArm::Gpu => {
680 if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
681 q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
682 crate::gpu::probe_record(
683 crate::gpu::OpClass::Matvec,
684 true,
685 t0.elapsed(),
686 );
687 return;
688 }
689 }
690 crate::gpu::ProbeArm::CpuTimed => {
691 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
692 crate::gpu::probe_record(
693 crate::gpu::OpClass::Matvec,
694 false,
695 t0.elapsed(),
696 );
697 return;
698 }
699 crate::gpu::ProbeArm::Cpu => {}
700 }
701 }
702 q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
703 return;
704 }
705 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
706 vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
707 return;
708 }
709 let xs = prescale(x, col_field, *dtype);
710 if *rows >= crate::gpu::min_rows()
715 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
716 && std::env::var("CMF_GPU_LMHEAD").map(|v| v != "0").unwrap_or(true)
717 && crate::gpu::enabled_here()
718 {
719 let t0 = std::time::Instant::now();
722 match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
723 crate::gpu::ProbeArm::Gpu => {}
724 crate::gpu::ProbeArm::CpuTimed => {
725 qmatvec(self.quant_bytes(), repack, row_scale, x, col_field, *dtype, *rows, *cols, out, pool);
726 crate::gpu::probe_record(
727 crate::gpu::OpClass::Matvec,
728 false,
729 t0.elapsed(),
730 );
731 return;
732 }
733 crate::gpu::ProbeArm::Cpu => {
734 qmatvec(self.quant_bytes(), repack, row_scale, x, col_field, *dtype, *rows, *cols, out, pool);
735 return;
736 }
737 }
738 let frac = std::env::var("CMF_GPU_SPLIT")
739 .ok()
740 .and_then(|v| v.parse::<f32>().ok())
741 .unwrap_or(0.5)
742 .clamp(0.0, 1.0);
743 let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
744 let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
745 let bytes = self.quant_bytes();
746 let ok = std::thread::scope(|sc| {
747 let g = sc.spawn(|| {
748 crate::gpu::q8_matvec_range(
749 model,
750 *idx,
751 cpu_rows,
752 &row_scale[cpu_rows..],
753 &xs,
754 *rows - cpu_rows,
755 *cols,
756 out_gpu,
757 )
758 });
759 if cpu_rows > 0 {
760 let rep_cpu = if repack.is_empty() {
763 &[][..]
764 } else {
765 &repack[..(cpu_rows / 4) * 4 * *cols]
766 };
767 qmatvec(
768 &bytes[..cpu_rows * *cols],
769 rep_cpu,
770 &row_scale[..cpu_rows],
771 x,
772 col_field,
773 *dtype,
774 cpu_rows,
775 *cols,
776 out_cpu,
777 pool,
778 );
779 }
780 g.join().unwrap_or(false)
781 });
782 if ok {
783 crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
784 return;
785 }
786 qmatvec(
789 &bytes[cpu_rows * *cols..(*rows) * *cols],
790 &[],
791 &row_scale[cpu_rows..],
792 x,
793 col_field,
794 *dtype,
795 *rows - cpu_rows,
796 *cols,
797 out_gpu,
798 pool,
799 );
800 return;
801 }
802 qmatvec(self.quant_bytes(), repack, row_scale, x, col_field, *dtype, *rows, *cols, out, pool);
803 }
804 }
805 }
806
807 pub fn matvec2(&self, x1: &[f32], x2: &[f32], o1: &mut [f32], o2: &mut [f32], pool: Option<&Pool>) {
809 match self {
810 Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
811 Self::Mapped {
812 dtype,
813 rows,
814 cols,
815 row_scale,
816 col_field,
817 vbit_offsets,
818 ..
819 } => {
820 if *dtype == TensorDtype::Q4Block {
821 q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
822 return;
823 }
824 if *dtype == TensorDtype::Q4Tiled {
825 q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
826 return;
827 }
828 if *dtype == TensorDtype::Q1 {
829 q1_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
830 return;
831 }
832 if *dtype == TensorDtype::Q1T {
833 q1t_matvec(self.quant_bytes(), x1, *rows, *cols, o1, pool);
838 q1t_matvec(self.quant_bytes(), x2, *rows, *cols, o2, pool);
839 return;
840 }
841 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
842 vbitmatvec2(self.quant_bytes(), vbit_offsets, x1, x2, *rows, *cols, o1, o2, pool);
843 return;
844 }
845 qmatvec2(self.quant_bytes(), row_scale, x1, x2, col_field, *dtype, *rows, *cols, o1, o2, pool);
846 }
847 }
848 }
849}
850
851impl QTensor {
852 pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
857 let cols = self.cols();
858 let rows = self.rows();
859 debug_assert_eq!(xs_all.len(), b * cols);
860 debug_assert_eq!(out.len(), b * rows);
861 if crate::gptq_capture::capturing() {
865 if let Self::Mapped { model, idx, .. } = self {
866 crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
867 }
868 }
869 match self {
870 Self::F32 { data, .. } => {
871 let out_addr = SendMut(out.as_mut_ptr());
872 let run = |start: usize, end: usize| {
873 for o in start..end {
874 let row = &data[o * cols..(o + 1) * cols];
875 for bi in 0..b {
876 let x = &xs_all[bi * cols..(bi + 1) * cols];
877 let mut acc = 0f32;
878 for j in 0..cols {
879 acc += row[j] * x[j];
880 }
881 unsafe { *out_addr.at(bi * rows + o) = acc };
882 }
883 }
884 };
885 dispatch_rows(pool, rows, &run);
886 }
887 Self::Mapped {
888 dtype,
889 row_scale,
890 col_field,
891 vbit_offsets,
892 ..
893 } => {
894 if *dtype == TensorDtype::Q4Block {
895 q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
896 return;
897 }
898 if *dtype == TensorDtype::Q4Tiled {
899 q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
900 return;
901 }
902 if *dtype == TensorDtype::Q1 {
903 if b >= 32
906 && b * rows * cols >= 128_000_000
907 && cols % 64 == 0
908 && crate::gpu::enabled_here()
909 {
910 if let Self::Mapped { model, idx, .. } = self {
911 let t0 = std::time::Instant::now();
912 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
913 crate::gpu::ProbeArm::Gpu => {
914 if crate::gpu::q1_matmat(model, *idx, xs_all, b, rows, cols, out) {
915 crate::gpu::probe_record(
916 crate::gpu::OpClass::Matmat,
917 true,
918 t0.elapsed(),
919 );
920 return;
921 }
922 }
923 crate::gpu::ProbeArm::CpuTimed => {
924 q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
925 crate::gpu::probe_record(
926 crate::gpu::OpClass::Matmat,
927 false,
928 t0.elapsed(),
929 );
930 return;
931 }
932 crate::gpu::ProbeArm::Cpu => {}
933 }
934 }
935 }
936 q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
937 return;
938 }
939 if *dtype == TensorDtype::Q1T {
940 if b >= 32
943 && b * rows * cols >= 128_000_000
944 && crate::gpu::enabled_here()
945 {
946 if let Self::Mapped { model, idx, .. } = self {
947 let t0 = std::time::Instant::now();
948 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
949 crate::gpu::ProbeArm::Gpu => {
950 if crate::gpu::q1t_matmat(model, *idx, xs_all, b, rows, cols, out) {
951 crate::gpu::probe_record(
952 crate::gpu::OpClass::Matmat,
953 true,
954 t0.elapsed(),
955 );
956 return;
957 }
958 }
959 crate::gpu::ProbeArm::CpuTimed => {
960 q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
961 crate::gpu::probe_record(
962 crate::gpu::OpClass::Matmat,
963 false,
964 t0.elapsed(),
965 );
966 return;
967 }
968 crate::gpu::ProbeArm::Cpu => {}
969 }
970 }
971 }
972 q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
973 return;
974 }
975 if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
976 vbitmatmat(self.quant_bytes(), vbit_offsets, xs_all, b, rows, cols, out, pool);
977 return;
978 }
979 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
980 .map(|bi| {
981 prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype)
982 })
983 .collect();
984 if b >= 8
990 && b * rows * cols >= 128_000_000
991 && crate::gpu::enabled_here()
992 {
993 if let Self::Mapped { model, idx, .. } = self {
994 let t0 = std::time::Instant::now();
995 match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
996 crate::gpu::ProbeArm::Gpu
997 if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
998 && !crate::gpu::q8_resident_or_upload(model, *idx) =>
999 {
1000 let q = self.quant_bytes();
1004 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1005 return;
1006 }
1007 crate::gpu::ProbeArm::Gpu => {
1008 let flat: Vec<f32> =
1009 pre.iter().flat_map(|v| v.iter().copied()).collect();
1010 if crate::gpu::q8_matmat(
1011 model, *idx, row_scale, &flat, b, rows, cols, out)
1012 {
1013 crate::gpu::probe_record(
1014 crate::gpu::OpClass::Matmat,
1015 true,
1016 t0.elapsed(),
1017 );
1018 return;
1019 }
1020 }
1021 crate::gpu::ProbeArm::CpuTimed => {
1022 let q = self.quant_bytes();
1023 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1024 crate::gpu::probe_record(
1025 crate::gpu::OpClass::Matmat,
1026 false,
1027 t0.elapsed(),
1028 );
1029 return;
1030 }
1031 crate::gpu::ProbeArm::Cpu => {}
1032 }
1033 }
1034 }
1035 let q = self.quant_bytes();
1036 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1037 }
1038 }
1039 }
1040}
1041
1042impl QTensor {
1043 pub fn matvec_many<const N: usize>(
1050 ts: [&QTensor; N],
1051 x: &[f32],
1052 mut outs: [&mut [f32]; N],
1053 pool: Option<&Pool>,
1054 ) {
1055 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
1056 let uniform_q8 = ts.iter().all(|t| {
1057 matches!(
1058 t,
1059 Self::Mapped { dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f, .. }
1060 )
1061 });
1062 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
1063 let uniform_q4 = ts
1064 .iter()
1065 .all(|t| matches!(t, Self::Mapped { dtype: TensorDtype::Q4Block, .. }));
1066 let uniform_vbit = ts
1067 .iter()
1068 .all(|t| matches!(
1069 t,
1070 Self::Mapped { dtype: TensorDtype::Vbit | TensorDtype::VbitRo, .. }
1071 ));
1072 let uniform_q1 = ts
1073 .iter()
1074 .all(|t| matches!(t, Self::Mapped { dtype: TensorDtype::Q1, .. }));
1075 let Some(pool) = pool else {
1076 for (t, o) in ts.iter().zip(outs.iter_mut()) {
1077 t.matvec(x, o, None);
1078 }
1079 return;
1080 };
1081 if total_rows < 256
1082 || !(uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit || uniform_q1)
1083 {
1084 for (t, o) in ts.iter().zip(outs.iter_mut()) {
1085 t.matvec(x, o, Some(pool));
1086 }
1087 return;
1088 }
1089
1090 if uniform_q1 {
1091 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1094 if a8w8_enabled() {
1095 let act = split_act(x);
1096 let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
1097 let (act, gsum) = (&act, &gsum);
1098 let closures: [_; N] = std::array::from_fn(|i| {
1099 let (bytes, gpr, out) =
1100 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1101 move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
1102 });
1103 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1104 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1105 pool.run_many(&parts);
1106 } else {
1107 let closures: [_; N] = std::array::from_fn(|i| {
1108 let (bytes, gpr, out) =
1109 (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1110 move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
1111 });
1112 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1113 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1114 pool.run_many(&parts);
1115 }
1116 return;
1117 }
1118
1119 if uniform_q4 || uniform_vbit {
1120 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1121 if a8w8_enabled() {
1123 let act = split_act(x);
1124 let act = &act;
1125 if uniform_q4 {
1126 let closures: [_; N] = std::array::from_fn(|i| {
1127 let (packed, scales) =
1128 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1129 let (gpr, cols, out) =
1130 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
1131 move |s: usize, e: usize| {
1132 q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
1133 }
1134 });
1135 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1136 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1137 pool.run_many(&parts);
1138 } else {
1139 let closures: [_; N] = std::array::from_fn(|i| {
1140 let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
1141 let (bytes, rows, cols, out) =
1142 (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), outs_addr[i]);
1143 move |s: usize, e: usize| {
1144 vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
1145 }
1146 });
1147 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1148 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1149 pool.run_many(&parts);
1150 }
1151 return;
1152 }
1153 if uniform_q4 {
1154 let closures: [_; N] = std::array::from_fn(|i| {
1155 let (packed, scales) =
1156 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1157 let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1158 move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
1159 });
1160 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1161 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1162 pool.run_many(&parts);
1163 } else {
1164 let closures: [_; N] = std::array::from_fn(|i| {
1165 let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
1166 let (bytes, rows, cols, out) =
1167 (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), outs_addr[i]);
1168 move |s: usize, e: usize| {
1169 vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
1170 }
1171 });
1172 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1173 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1174 pool.run_many(&parts);
1175 }
1176 return;
1177 }
1178
1179 if uniform_f32 {
1180 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1181 let closures: [_; N] = std::array::from_fn(|i| {
1182 let Self::F32 { data, cols, .. } = ts[i] else { unreachable!() };
1183 let out = outs_addr[i];
1184 move |start: usize, end: usize| {
1185 for o in start..end {
1186 let row = &data[o * cols..(o + 1) * cols];
1187 let mut sum = 0.0f32;
1188 for j in 0..*cols {
1189 sum += row[j] * x[j];
1190 }
1191 unsafe { *out.at(o) = sum };
1193 }
1194 }
1195 });
1196 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1197 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1198 pool.run_many(&parts);
1199 return;
1200 }
1201
1202 struct Ctx<'a> {
1205 bytes: &'a [u8],
1206 #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
1207 rep: &'a [u8],
1208 row_scale: &'a [f32],
1209 cols: usize,
1210 xs: std::borrow::Cow<'a, [f32]>,
1211 }
1212 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
1213 let Self::Mapped { dtype, cols, row_scale, col_field, repack, .. } = ts[i] else {
1214 unreachable!()
1215 };
1216 Ctx {
1217 bytes: ts[i].quant_bytes(),
1218 rep: repack,
1219 row_scale,
1220 cols: *cols,
1221 xs: prescale(x, col_field, *dtype),
1222 }
1223 });
1224 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1225 #[cfg(target_arch = "aarch64")]
1226 if sdot_enabled() {
1227 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
1228 let closures: [_; N] = std::array::from_fn(|i| {
1229 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
1230 move |start: usize, end: usize| {
1231 q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
1232 }
1233 });
1234 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1235 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1236 pool.run_many(&parts);
1237 return;
1238 }
1239 #[cfg(target_arch = "x86_64")]
1240 if avx2_a8w8_enabled() {
1241 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
1242 let closures: [_; N] = std::array::from_fn(|i| {
1243 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
1244 move |start: usize, end: usize| {
1245 q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
1246 }
1247 });
1248 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1249 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1250 pool.run_many(&parts);
1251 return;
1252 }
1253 let closures: [_; N] = std::array::from_fn(|i| {
1254 let (c, out) = (&ctxs[i], outs_addr[i]);
1255 move |start: usize, end: usize| {
1256 q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
1257 }
1258 });
1259 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1260 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1261 pool.run_many(&parts);
1262 }
1263}
1264
1265impl QTensor {
1266 #[allow(clippy::needless_range_loop)]
1271 pub fn matvec2_many<const N: usize>(
1272 ts: [&QTensor; N],
1273 x1: &[f32],
1274 x2: &[f32],
1275 mut o1s: [&mut [f32]; N],
1276 mut o2s: [&mut [f32]; N],
1277 pool: Option<&Pool>,
1278 ) {
1279 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
1280 let uniform_q8 = ts.iter().all(|t| {
1281 matches!(
1282 t,
1283 Self::Mapped { dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f, .. }
1284 )
1285 });
1286 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
1287 let uniform_q4 = ts
1288 .iter()
1289 .all(|t| matches!(t, Self::Mapped { dtype: TensorDtype::Q4Block, .. }));
1290 let uniform_vbit = ts
1291 .iter()
1292 .all(|t| matches!(
1293 t,
1294 Self::Mapped { dtype: TensorDtype::Vbit | TensorDtype::VbitRo, .. }
1295 ));
1296 let fusable = pool.is_some()
1297 && total_rows >= 256
1298 && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
1299 if !fusable {
1300 for i in 0..N {
1301 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
1302 }
1303 return;
1304 }
1305 let pool = pool.unwrap();
1306
1307 if uniform_q4 || uniform_vbit {
1308 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1309 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1310 if a8w8_enabled() {
1312 let a1 = split_act(x1);
1313 let a2 = split_act(x2);
1314 let (a1, a2) = (&a1, &a2);
1315 if uniform_q4 {
1316 let closures: [_; N] = std::array::from_fn(|i| {
1317 let (packed, scales) =
1318 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1319 let (gpr, cols, o1, o2) =
1320 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
1321 move |s: usize, e: usize| {
1322 q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
1323 }
1324 });
1325 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1326 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1327 pool.run_many(&parts);
1328 } else {
1329 let closures: [_; N] = std::array::from_fn(|i| {
1330 let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
1331 let (bytes, rows, cols, o1, o2) =
1332 (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), p1[i], p2[i]);
1333 move |s: usize, e: usize| {
1334 vbit_range2_a8w8(
1335 bytes, vbit_offsets, x1, x2, a1, a2, rows, cols, o1, o2, s, e,
1336 )
1337 }
1338 });
1339 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1340 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1341 pool.run_many(&parts);
1342 }
1343 return;
1344 }
1345 if uniform_q4 {
1346 let closures: [_; N] = std::array::from_fn(|i| {
1347 let (packed, scales) =
1348 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1349 let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
1350 move |s: usize, e: usize| {
1351 q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
1352 }
1353 });
1354 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1355 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1356 pool.run_many(&parts);
1357 } else {
1358 let closures: [_; N] = std::array::from_fn(|i| {
1359 let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
1360 let (bytes, rows, cols, o1, o2) =
1361 (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), p1[i], p2[i]);
1362 move |s: usize, e: usize| {
1363 vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
1364 }
1365 });
1366 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1367 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1368 pool.run_many(&parts);
1369 }
1370 return;
1371 }
1372
1373 if uniform_f32 {
1374 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1375 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1376 let closures: [_; N] = std::array::from_fn(|i| {
1377 let Self::F32 { data, cols, .. } = ts[i] else { unreachable!() };
1378 let (o1, o2) = (p1[i], p2[i]);
1379 move |start: usize, end: usize| {
1380 for o in start..end {
1381 let row = &data[o * cols..(o + 1) * cols];
1382 let (mut s1, mut s2) = (0.0f32, 0.0f32);
1383 for j in 0..*cols {
1384 s1 += row[j] * x1[j];
1385 s2 += row[j] * x2[j];
1386 }
1387 unsafe {
1389 *o1.at(o) = s1;
1390 *o2.at(o) = s2;
1391 }
1392 }
1393 }
1394 });
1395 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1396 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1397 pool.run_many(&parts);
1398 return;
1399 }
1400
1401 struct Ctx<'a> {
1402 bytes: &'a [u8],
1403 row_scale: &'a [f32],
1404 cols: usize,
1405 xs1: std::borrow::Cow<'a, [f32]>,
1406 xs2: std::borrow::Cow<'a, [f32]>,
1407 }
1408 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
1409 let Self::Mapped { dtype, cols, row_scale, col_field, .. } = ts[i] else {
1410 unreachable!()
1411 };
1412 Ctx {
1413 bytes: ts[i].quant_bytes(),
1414 row_scale,
1415 cols: *cols,
1416 xs1: prescale(x1, col_field, *dtype),
1417 xs2: prescale(x2, col_field, *dtype),
1418 }
1419 });
1420 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
1421 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
1422 #[cfg(target_arch = "aarch64")]
1423 if sdot_enabled() {
1424 let acts: [(SplitAct, SplitAct); N] =
1425 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
1426 let closures: [_; N] = std::array::from_fn(|i| {
1427 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
1428 move |start: usize, end: usize| {
1429 q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
1430 }
1431 });
1432 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1433 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1434 pool.run_many(&parts);
1435 return;
1436 }
1437 #[cfg(target_arch = "x86_64")]
1438 if avx2_a8w8_enabled() {
1439 let acts: [(SplitAct, SplitAct); N] =
1440 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
1441 let closures: [_; N] = std::array::from_fn(|i| {
1442 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
1443 move |start: usize, end: usize| {
1444 q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
1445 }
1446 });
1447 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1448 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1449 pool.run_many(&parts);
1450 return;
1451 }
1452 let closures: [_; N] = std::array::from_fn(|i| {
1453 let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
1454 move |start: usize, end: usize| {
1455 q8_range2_f32(c.bytes, c.row_scale, &c.xs1, &c.xs2, c.cols, o1, o2, start, end)
1456 }
1457 });
1458 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1459 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1460 pool.run_many(&parts);
1461 }
1462}
1463
1464#[cfg(target_os = "macos")]
1469mod accel_blas {
1470 #[link(name = "Accelerate", kind = "framework")]
1471 unsafe extern "C" {
1472 pub fn cblas_sgemm(
1473 order: i32,
1474 trans_a: i32,
1475 trans_b: i32,
1476 m: i32,
1477 n: i32,
1478 k: i32,
1479 alpha: f32,
1480 a: *const f32,
1481 lda: i32,
1482 b: *const f32,
1483 ldb: i32,
1484 beta: f32,
1485 c: *mut f32,
1486 ldc: i32,
1487 );
1488 }
1489}
1490
1491#[cfg(target_os = "macos")]
1492pub(crate) fn accel_gemm_enabled() -> bool {
1493 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
1494 *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
1495}
1496
1497#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
1500pub(crate) fn accel_gemm_enabled() -> bool {
1501 true
1502}
1503
1504#[cfg(target_arch = "aarch64")]
1511#[allow(clippy::too_many_arguments)]
1512pub(crate) fn neon_gemm_rm(
1513 m: usize,
1514 n: usize,
1515 k: usize,
1516 alpha: f32,
1517 a: &[f32],
1518 lda: usize,
1519 b_mat: &[f32],
1520 ldb: usize,
1521 b_rows_are_n: bool,
1522 c: &mut [f32],
1523 ldc: usize,
1524) {
1525 debug_assert!(a.len() >= (m - 1) * lda + k);
1526 debug_assert!(c.len() >= (m - 1) * ldc + n);
1527 unsafe {
1529 use core::arch::aarch64::*;
1530 let mut i = 0usize;
1531 while i < m {
1532 let mi = (m - i).min(4);
1533 let mut j = 0usize;
1534 while j < n {
1535 let nj = (n - j).min(8);
1536 if mi == 4 && nj == 8 {
1537 let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
1538 let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
1539 let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
1540 let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
1541 for p in 0..k {
1542 let (b0, b1) = if b_rows_are_n {
1543 let base = b_mat.as_ptr().add(j * ldb + p);
1546 let g = |o: usize| *base.add(o * ldb);
1547 (
1548 [g(0), g(1), g(2), g(3)],
1549 [g(4), g(5), g(6), g(7)],
1550 )
1551 } else {
1552 let base = b_mat.as_ptr().add(p * ldb + j);
1553 (
1554 [*base, *base.add(1), *base.add(2), *base.add(3)],
1555 [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
1556 )
1557 };
1558 let bv0 = vld1q_f32(b0.as_ptr());
1559 let bv1 = vld1q_f32(b1.as_ptr());
1560 let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
1561 let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
1562 let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
1563 let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
1564 c0a = vfmaq_f32(c0a, a0, bv0);
1565 c0b = vfmaq_f32(c0b, a0, bv1);
1566 c1a = vfmaq_f32(c1a, a1, bv0);
1567 c1b = vfmaq_f32(c1b, a1, bv1);
1568 c2a = vfmaq_f32(c2a, a2, bv0);
1569 c2b = vfmaq_f32(c2b, a2, bv1);
1570 c3a = vfmaq_f32(c3a, a3, bv0);
1571 c3b = vfmaq_f32(c3b, a3, bv1);
1572 }
1573 let al = vdupq_n_f32(alpha);
1574 for (r, (ca, cb)) in
1575 [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)].iter().enumerate()
1576 {
1577 let dst = c.as_mut_ptr().add((i + r) * ldc + j);
1578 vst1q_f32(dst, vmulq_f32(*ca, al));
1579 vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
1580 }
1581 } else {
1582 for r in 0..mi {
1583 for q in 0..nj {
1584 let mut acc = 0f32;
1585 for p in 0..k {
1586 let bv = if b_rows_are_n {
1587 b_mat[(j + q) * ldb + p]
1588 } else {
1589 b_mat[p * ldb + j + q]
1590 };
1591 acc += a[(i + r) * lda + p] * bv;
1592 }
1593 c[(i + r) * ldc + j + q] = acc * alpha;
1594 }
1595 }
1596 }
1597 j += nj;
1598 }
1599 i += mi;
1600 }
1601 }
1602}
1603
1604#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
1606#[allow(clippy::too_many_arguments)]
1607pub(crate) fn sgemm_rm(
1608 m: usize,
1609 n: usize,
1610 k: usize,
1611 alpha: f32,
1612 a: &[f32],
1613 lda: usize,
1614 b_mat: &[f32],
1615 ldb: usize,
1616 b_rows_are_n: bool,
1617 c: &mut [f32],
1618 ldc: usize,
1619) {
1620 neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
1621}
1622
1623#[cfg(target_os = "macos")]
1626#[allow(clippy::too_many_arguments)]
1627pub(crate) fn sgemm_rm(
1628 m: usize,
1629 n: usize,
1630 k: usize,
1631 alpha: f32,
1632 a: &[f32],
1633 lda: usize,
1634 b_mat: &[f32],
1635 ldb: usize,
1636 b_rows_are_n: bool,
1637 c: &mut [f32],
1638 ldc: usize,
1639) {
1640 debug_assert!(a.len() >= (m - 1) * lda + k);
1641 debug_assert!(c.len() >= (m - 1) * ldc + n);
1642 #[cfg(target_arch = "aarch64")]
1647 if std::env::var("CMF_FORCE_NEON_GEMM").map(|v| v == "1").unwrap_or(false) {
1648 return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
1649 }
1650 unsafe {
1651 accel_blas::cblas_sgemm(
1652 101, 111, if b_rows_are_n { 112 } else { 111 },
1655 m as i32,
1656 n as i32,
1657 k as i32,
1658 alpha,
1659 a.as_ptr(),
1660 lda as i32,
1661 b_mat.as_ptr(),
1662 ldb as i32,
1663 0.0,
1664 c.as_mut_ptr(),
1665 ldc as i32,
1666 );
1667 }
1668}
1669
1670#[cfg(target_os = "macos")]
1677fn qmatmat_accel(
1678 q: &[u8],
1679 row_scale: &[f32],
1680 pre: &[std::borrow::Cow<'_, [f32]>],
1681 rows: usize,
1682 cols: usize,
1683 out: &mut [f32],
1684 pool: Option<&Pool>,
1685) {
1686 const TR: usize = 2048;
1691 let b = pre.len();
1692 thread_local! {
1693 static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
1694 static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
1695 }
1696 XPANEL.with(|xp| {
1697 WTILE.with(|wt| {
1698 let mut xpanel = xp.borrow_mut();
1699 xpanel.clear();
1700 for x in pre {
1701 xpanel.extend_from_slice(x);
1702 }
1703 let mut wtile = wt.borrow_mut();
1704 wtile.resize(TR * cols, 0.0);
1705 let mut r0 = 0usize;
1706 while r0 < rows {
1707 let tr = TR.min(rows - r0);
1708 let wt_addr = SendMut(wtile.as_mut_ptr());
1710 let run = |start: usize, end: usize| {
1711 for r in start..end {
1712 let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
1713 let s = row_scale[r0 + r];
1714 let dst = unsafe {
1716 std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols)
1717 };
1718 for (d, &v) in dst.iter_mut().zip(row) {
1719 *d = (v as i8) as f32 * s;
1720 }
1721 }
1722 };
1723 dispatch_rows(pool, tr, &run);
1724 unsafe {
1726 accel_blas::cblas_sgemm(
1727 101, 111, 112, b as i32,
1731 tr as i32,
1732 cols as i32,
1733 1.0,
1734 xpanel.as_ptr(),
1735 cols as i32,
1736 wtile.as_ptr(),
1737 cols as i32,
1738 0.0,
1739 out.as_mut_ptr().add(r0),
1740 rows as i32,
1741 );
1742 }
1743 r0 += tr;
1744 }
1745 })
1746 });
1747}
1748
1749fn qmatmat(
1750 q: &[u8],
1751 row_scale: &[f32],
1752 pre: &[std::borrow::Cow<'_, [f32]>],
1753 rows: usize,
1754 cols: usize,
1755 out: &mut [f32],
1756 pool: Option<&Pool>,
1757) {
1758 let b = pre.len();
1759 debug_assert_eq!(out.len(), b * rows);
1760 #[cfg(target_os = "macos")]
1765 if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
1766 qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
1767 return;
1768 }
1769 #[cfg(target_arch = "aarch64")]
1770 if sdot_enabled() {
1771 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
1772 let out_addr = SendMut(out.as_mut_ptr());
1773 let blocked_ok =
1776 std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true);
1777 let use_i8mm = i8mm_enabled();
1778 if blocked_ok {
1779 let run = |start: usize, end: usize| {
1780 let mut o = start;
1781 while o < end {
1782 if o + 2 <= end {
1783 let r0 = &q[o * cols..(o + 1) * cols];
1784 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
1785 let mut bi = 0usize;
1786 while bi + 4 <= acts.len() {
1787 let xs = [
1788 acts[bi].xq.as_slice(),
1789 acts[bi + 1].xq.as_slice(),
1790 acts[bi + 2].xq.as_slice(),
1791 acts[bi + 3].xq.as_slice(),
1792 ];
1793 let d = if use_i8mm {
1794 unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
1795 } else {
1796 unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
1797 };
1798 for (r, row) in [r0, r1].into_iter().enumerate() {
1799 for k in 0..4 {
1800 let act = &acts[bi + k];
1801 let mut v = d[r][k] as f32 * act.sx;
1802 for &(j, xv) in &act.outliers {
1803 v += (row[j] as i8) as f32 * xv;
1804 }
1805 unsafe {
1806 *out_addr.at((bi + k) * rows + o + r) =
1807 v * row_scale[o + r]
1808 };
1809 }
1810 }
1811 bi += 4;
1812 }
1813 while bi < acts.len() {
1814 for (r, row) in [r0, r1].into_iter().enumerate() {
1815 let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
1816 unsafe { *out_addr.at(bi * rows + o + r) = v };
1817 }
1818 bi += 1;
1819 }
1820 o += 2;
1821 } else {
1822 let row = &q[o * cols..(o + 1) * cols];
1823 for (bi, act) in acts.iter().enumerate() {
1824 let v = row_dot_sdot(row, act) * row_scale[o];
1825 unsafe { *out_addr.at(bi * rows + o) = v };
1826 }
1827 o += 1;
1828 }
1829 }
1830 };
1831 dispatch_rows(pool, rows, &run);
1832 return;
1833 }
1834 let run = |start: usize, end: usize| {
1835 for o in start..end {
1836 let row = &q[o * cols..(o + 1) * cols];
1837 for (bi, act) in acts.iter().enumerate() {
1838 let v = row_dot_sdot(row, act) * row_scale[o];
1839 unsafe { *out_addr.at(bi * rows + o) = v };
1840 }
1841 }
1842 };
1843 dispatch_rows(pool, rows, &run);
1844 return;
1845 }
1846 #[cfg(target_arch = "x86_64")]
1851 if avx2_a8w8_enabled() {
1852 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
1853 let out_addr = SendMut(out.as_mut_ptr());
1854 let blocked_ok = std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true);
1857 if !avx512vnni_enabled() && blocked_ok {
1858 let run = |start: usize, end: usize| {
1859 let mut o = start;
1860 while o < end {
1861 if o + 2 <= end {
1862 let r0 = &q[o * cols..(o + 1) * cols];
1863 let r1 = &q[(o + 1) * cols..(o + 2) * cols];
1864 let mut bi = 0usize;
1865 while bi + 4 <= acts.len() {
1866 let xs = [
1867 acts[bi].xq.as_slice(),
1868 acts[bi + 1].xq.as_slice(),
1869 acts[bi + 2].xq.as_slice(),
1870 acts[bi + 3].xq.as_slice(),
1871 ];
1872 let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
1873 for (r, row) in [r0, r1].into_iter().enumerate() {
1874 for k in 0..4 {
1875 let act = &acts[bi + k];
1876 let mut v = d[r][k] as f32 * act.sx;
1877 for &(j, xv) in &act.outliers {
1878 v += (row[j] as i8) as f32 * xv;
1879 }
1880 unsafe {
1881 *out_addr.at((bi + k) * rows + o + r) =
1882 v * row_scale[o + r]
1883 };
1884 }
1885 }
1886 bi += 4;
1887 }
1888 while bi < acts.len() {
1889 for (r, row) in [r0, r1].into_iter().enumerate() {
1890 let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
1891 unsafe { *out_addr.at(bi * rows + o + r) = v };
1892 }
1893 bi += 1;
1894 }
1895 o += 2;
1896 } else {
1897 let row = &q[o * cols..(o + 1) * cols];
1898 for (bi, act) in acts.iter().enumerate() {
1899 let v = row_dot_avx2(row, act) * row_scale[o];
1900 unsafe { *out_addr.at(bi * rows + o) = v };
1901 }
1902 o += 1;
1903 }
1904 }
1905 };
1906 dispatch_rows(pool, rows, &run);
1907 return;
1908 }
1909 let run = |start: usize, end: usize| {
1910 for o in start..end {
1911 let row = &q[o * cols..(o + 1) * cols];
1912 for (bi, act) in acts.iter().enumerate() {
1913 let v = row_dot_avx2(row, act) * row_scale[o];
1914 unsafe { *out_addr.at(bi * rows + o) = v };
1915 }
1916 }
1917 };
1918 dispatch_rows(pool, rows, &run);
1919 return;
1920 }
1921 let out_addr = SendMut(out.as_mut_ptr());
1922 let run = |start: usize, end: usize| {
1923 for o in start..end {
1924 let row = &q[o * cols..(o + 1) * cols];
1925 for (bi, x) in pre.iter().enumerate() {
1926 let mut acc = 0f32;
1927 for j in 0..cols {
1928 acc += (row[j] as i8) as f32 * x[j];
1929 }
1930 unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
1931 }
1932 }
1933 };
1934 dispatch_rows(pool, rows, &run);
1935}
1936
1937fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
1940 match pool {
1941 Some(pool) if rows >= 256 => pool.run_rows(rows, run),
1942 _ => run(0, rows),
1943 }
1944}
1945
1946fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
1948 let groups = rows * cols / GROUP_SIZE;
1949 bytes.split_at(groups * 16)
1950}
1951
1952#[inline]
1957fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
1958 #[cfg(target_arch = "aarch64")]
1959 unsafe {
1960 return vbit_fill4_neon(data, buf);
1961 }
1962 #[cfg(target_arch = "x86_64")]
1963 if avx2_enabled() {
1964 return unsafe { vbit_fill4_avx2(data, buf) };
1965 }
1966 #[allow(unreachable_code)]
1967 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
1968 let u = unpack8::<4>(&data[blk * 4..]);
1969 for k in 0..8 {
1970 chunk[k] = (u[k] - 7) as i8 as u8;
1971 }
1972 }
1973}
1974
1975#[cfg(target_arch = "aarch64")]
1976#[target_feature(enable = "neon")]
1977unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
1978 unsafe {
1981 use core::arch::aarch64::*;
1982 let n = buf.len();
1983 let mask = vdupq_n_u8(0x0F);
1984 let seven = vdupq_n_s8(7);
1985 let mut g = 0usize;
1986 while g * 32 + 32 <= n {
1987 let b = vld1q_u8(data.as_ptr().add(g * 16));
1988 let hi = vshrq_n_u8::<4>(b);
1989 let lo = vandq_u8(b, mask);
1990 let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
1991 let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
1992 vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
1993 vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
1994 g += 1;
1995 }
1996 }
1997}
1998
1999#[cfg(target_arch = "x86_64")]
2000#[target_feature(enable = "avx2")]
2001unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
2002 unsafe {
2004 use core::arch::x86_64::*;
2005 let n = buf.len();
2006 let mask = _mm_set1_epi8(0x0F);
2007 let seven = _mm256_set1_epi8(7);
2008 let mut g = 0usize;
2009 while g * 32 + 32 <= n {
2010 let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
2011 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
2012 let lo = _mm_and_si128(b, mask);
2013 let z = _mm256_sub_epi8(
2014 _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
2015 seven,
2016 );
2017 _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
2018 g += 1;
2019 }
2020 }
2021}
2022
2023#[inline(always)]
2028fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
2029 let mut acc = 0u64;
2030 for i in 0..B {
2031 acc = (acc << 8) | data[i] as u64;
2032 }
2033 let mask = (1u64 << B) - 1;
2034 let mut out = [0i32; 8];
2035 for (k, o) in out.iter_mut().enumerate() {
2036 *o = ((acc >> ((7 - k) * B)) & mask) as i32;
2037 }
2038 out
2039}
2040
2041#[allow(clippy::too_many_arguments)]
2047fn vbitmatvec(
2048 bytes: &[u8],
2049 offsets: &[usize],
2050 x: &[f32],
2051 rows: usize,
2052 cols: usize,
2053 out: &mut [f32],
2054 pool: Option<&Pool>,
2055) {
2056 debug_assert_eq!(out.len(), rows);
2057 debug_assert_eq!(offsets.len(), rows + 1);
2058
2059 if a8w8_enabled() {
2063 let act = split_act(x);
2064 let out_addr = SendMut(out.as_mut_ptr());
2065 let run = move |start: usize, end: usize| {
2066 vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
2067 };
2068 dispatch_rows(pool, rows, &run);
2069 return;
2070 }
2071
2072 let out_addr = SendMut(out.as_mut_ptr());
2073 let run = move |start: usize, end: usize| {
2074 vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
2075 };
2076 dispatch_rows(pool, rows, &run);
2077}
2078
2079#[allow(clippy::too_many_arguments)]
2083fn vbit_range_a8w8(
2084 bytes: &[u8],
2085 offsets: &[usize],
2086 x: &[f32],
2087 act: &SplitAct,
2088 rows: usize,
2089 cols: usize,
2090 out: SendMut,
2091 start: usize,
2092 end: usize,
2093) {
2094 let ng = cols / GROUP_SIZE;
2095 let bits = &bytes[..rows];
2096 let sc_off = rows;
2097 let row_dot = |r: usize| -> f32 {
2098 let b = bits[r] as usize;
2099 let l = ((1i32 << (b - 1)) - 1) as i32;
2100 let mask = (1u64 << b) - 1;
2101 let data = &bytes[offsets[r]..offsets[r + 1]];
2102 if b == 8 {
2103 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
2105 let mut dot = 0f32;
2106 for g in 0..ng {
2107 let so = (r * ng + g) * 2;
2108 let sgf = f16_to_f32(u16::from_le_bytes([
2109 bytes[sc_off + so],
2110 bytes[sc_off + so + 1],
2111 ]));
2112 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2113 let mut gd = 0f32;
2114 for &xv in xg.iter() {
2115 if nbits < 8 {
2116 acc = (acc << 8) | data[idx] as u64;
2117 idx += 1;
2118 nbits += 8;
2119 }
2120 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
2121 nbits -= 8;
2122 gd += (u - l) as f32 * xv;
2123 }
2124 dot += gd * sgf;
2125 }
2126 return dot;
2127 }
2128 thread_local! {
2132 static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
2133 const { std::cell::RefCell::new(Vec::new()) };
2134 }
2135 #[inline(always)]
2136 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
2137 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
2138 let u = unpack8::<B>(&data[blk * B..]);
2139 for k in 0..8 {
2140 chunk[k] = (u[k] - l) as i8 as u8;
2141 }
2142 }
2143 }
2144 let _ = mask;
2145 VBIT_SCRATCH.with(|scratch| {
2146 let mut buf = scratch.borrow_mut();
2147 buf.resize(cols, 0);
2148 match b {
2149 3 => fill::<3>(data, l, &mut buf),
2150 4 => vbit_fill4(data, &mut buf),
2151 5 => fill::<5>(data, l, &mut buf),
2152 6 => fill::<6>(data, l, &mut buf),
2153 _ => unreachable!(),
2154 }
2155 let mut dot = 0f32;
2156 for g in 0..ng {
2157 let so = (r * ng + g) * 2;
2158 let s = f16_to_f32(u16::from_le_bytes([
2159 bytes[sc_off + so],
2160 bytes[sc_off + so + 1],
2161 ]));
2162 let d = dot_i8_i8(
2163 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
2164 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
2165 ) as f32
2166 * act.sx;
2167 dot += d * s;
2168 }
2169 for &(j, xv) in &act.outliers {
2170 let so = (r * ng + j / GROUP_SIZE) * 2;
2171 let s = f16_to_f32(u16::from_le_bytes([
2172 bytes[sc_off + so],
2173 bytes[sc_off + so + 1],
2174 ]));
2175 dot += (buf[j] as i8) as f32 * s * xv;
2177 }
2178 dot
2179 })
2180 };
2181 for r in start..end {
2182 unsafe { *out.at(r) = row_dot(r) };
2184 }
2185}
2186
2187#[allow(clippy::too_many_arguments)]
2189fn vbit_range_f32(
2190 bytes: &[u8],
2191 offsets: &[usize],
2192 x: &[f32],
2193 rows: usize,
2194 cols: usize,
2195 out: SendMut,
2196 start: usize,
2197 end: usize,
2198) {
2199 let ng = cols / GROUP_SIZE;
2200 let bits = &bytes[..rows];
2201 let sc_off = rows;
2202 #[inline(always)]
2206 fn dot_row<const B: usize>(
2207 data: &[u8],
2208 bytes: &[u8],
2209 sc_off: usize,
2210 r: usize,
2211 ng: usize,
2212 x: &[f32],
2213 ) -> f32 {
2214 let l = ((1i32 << (B - 1)) - 1) as f32;
2215 let gbytes = GROUP_SIZE * B / 8;
2216 let mut dot = 0f32;
2217 for g in 0..ng {
2218 let so = (r * ng + g) * 2;
2219 let s = f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]));
2220 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2221 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
2222 let mut gd = 0f32;
2223 for blk in 0..GROUP_SIZE / 8 {
2224 let u = unpack8::<B>(&gd0[blk * B..]);
2225 let xb = &xg[blk * 8..blk * 8 + 8];
2226 for k in 0..8 {
2227 gd += (u[k] as f32 - l) * xb[k];
2228 }
2229 }
2230 dot += gd * s;
2231 }
2232 dot
2233 }
2234 for r in start..end {
2235 let data = &bytes[offsets[r]..offsets[r + 1]];
2236 let v = match bits[r] {
2237 3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
2238 4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
2239 5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
2240 6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
2241 8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
2242 b => unreachable!("vbit bit-width {b} (validated at load)"),
2243 };
2244 unsafe { *out.at(r) = v };
2246 }
2247}
2248
2249#[allow(clippy::too_many_arguments)]
2254fn vbitmatvec2(
2255 bytes: &[u8],
2256 offsets: &[usize],
2257 x1: &[f32],
2258 x2: &[f32],
2259 rows: usize,
2260 cols: usize,
2261 o1: &mut [f32],
2262 o2: &mut [f32],
2263 pool: Option<&Pool>,
2264) {
2265 debug_assert_eq!(o1.len(), rows);
2266 debug_assert_eq!(o2.len(), rows);
2267
2268 if a8w8_enabled() {
2269 let a1 = split_act(x1);
2270 let a2 = split_act(x2);
2271 let p1 = SendMut(o1.as_mut_ptr());
2272 let p2 = SendMut(o2.as_mut_ptr());
2273 let run = move |start: usize, end: usize| {
2274 vbit_range2_a8w8(bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end)
2275 };
2276 dispatch_rows(pool, rows, &run);
2277 return;
2278 }
2279
2280 let p1 = SendMut(o1.as_mut_ptr());
2281 let p2 = SendMut(o2.as_mut_ptr());
2282 let run = move |start: usize, end: usize| {
2283 vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
2284 };
2285 dispatch_rows(pool, rows, &run);
2286}
2287
2288#[allow(clippy::too_many_arguments)]
2292fn vbit_range2_a8w8(
2293 bytes: &[u8],
2294 offsets: &[usize],
2295 x1: &[f32],
2296 x2: &[f32],
2297 a1: &SplitAct,
2298 a2: &SplitAct,
2299 rows: usize,
2300 cols: usize,
2301 p1: SendMut,
2302 p2: SendMut,
2303 start: usize,
2304 end: usize,
2305) {
2306 let ng = cols / GROUP_SIZE;
2307 let bits = &bytes[..rows];
2308 let sc_off = rows;
2309 let row_dots = |r: usize| -> (f32, f32) {
2310 let b = bits[r] as usize;
2311 let l = (1i32 << (b - 1)) - 1;
2312 let data = &bytes[offsets[r]..offsets[r + 1]];
2313 if b == 8 {
2314 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
2317 let (mut d1, mut d2) = (0f32, 0f32);
2318 for g in 0..ng {
2319 let so = (r * ng + g) * 2;
2320 let sgf = f16_to_f32(u16::from_le_bytes([
2321 bytes[sc_off + so],
2322 bytes[sc_off + so + 1],
2323 ]));
2324 let (mut g1, mut g2) = (0f32, 0f32);
2325 for k in 0..GROUP_SIZE {
2326 if nbits < 8 {
2327 acc = (acc << 8) | data[idx] as u64;
2328 idx += 1;
2329 nbits += 8;
2330 }
2331 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
2332 nbits -= 8;
2333 let w = (u - l) as f32;
2334 g1 += w * x1[g * GROUP_SIZE + k];
2335 g2 += w * x2[g * GROUP_SIZE + k];
2336 }
2337 d1 += g1 * sgf;
2338 d2 += g2 * sgf;
2339 }
2340 return (d1, d2);
2341 }
2342 thread_local! {
2343 static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
2344 const { std::cell::RefCell::new(Vec::new()) };
2345 }
2346 #[inline(always)]
2347 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
2348 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
2349 let u = unpack8::<B>(&data[blk * B..]);
2350 for k in 0..8 {
2351 chunk[k] = (u[k] - l) as i8 as u8;
2352 }
2353 }
2354 }
2355 VBIT_SCRATCH2.with(|scratch| {
2356 let mut buf = scratch.borrow_mut();
2357 buf.resize(cols, 0);
2358 match b {
2359 3 => fill::<3>(data, l, &mut buf),
2360 4 => vbit_fill4(data, &mut buf),
2361 5 => fill::<5>(data, l, &mut buf),
2362 6 => fill::<6>(data, l, &mut buf),
2363 _ => unreachable!(),
2364 }
2365 let (mut d1, mut d2) = (0f32, 0f32);
2366 for g in 0..ng {
2367 let so = (r * ng + g) * 2;
2368 let s = f16_to_f32(u16::from_le_bytes([
2369 bytes[sc_off + so],
2370 bytes[sc_off + so + 1],
2371 ]));
2372 let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2373 let v1 =
2374 dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
2375 let v2 =
2376 dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
2377 d1 += v1 * s;
2378 d2 += v2 * s;
2379 }
2380 for &(j, xv) in &a1.outliers {
2381 let so = (r * ng + j / GROUP_SIZE) * 2;
2382 let s = f16_to_f32(u16::from_le_bytes([
2383 bytes[sc_off + so],
2384 bytes[sc_off + so + 1],
2385 ]));
2386 d1 += (buf[j] as i8) as f32 * s * xv;
2387 }
2388 for &(j, xv) in &a2.outliers {
2389 let so = (r * ng + j / GROUP_SIZE) * 2;
2390 let s = f16_to_f32(u16::from_le_bytes([
2391 bytes[sc_off + so],
2392 bytes[sc_off + so + 1],
2393 ]));
2394 d2 += (buf[j] as i8) as f32 * s * xv;
2395 }
2396 (d1, d2)
2397 })
2398 };
2399 for r in start..end {
2400 let (v1, v2) = row_dots(r);
2401 unsafe {
2403 *p1.at(r) = v1;
2404 *p2.at(r) = v2;
2405 }
2406 }
2407}
2408
2409#[allow(clippy::too_many_arguments)]
2413fn vbit_range2_f32(
2414 bytes: &[u8],
2415 offsets: &[usize],
2416 x1: &[f32],
2417 x2: &[f32],
2418 rows: usize,
2419 cols: usize,
2420 p1: SendMut,
2421 p2: SendMut,
2422 start: usize,
2423 end: usize,
2424) {
2425 let ng = cols / GROUP_SIZE;
2426 let bits = &bytes[..rows];
2427 let sc_off = rows;
2428 #[inline(always)]
2429 #[allow(clippy::too_many_arguments)]
2430 fn dot_row2<const B: usize>(
2431 data: &[u8],
2432 bytes: &[u8],
2433 sc_off: usize,
2434 r: usize,
2435 ng: usize,
2436 x1: &[f32],
2437 x2: &[f32],
2438 ) -> (f32, f32) {
2439 let l = ((1i32 << (B - 1)) - 1) as f32;
2440 let gbytes = GROUP_SIZE * B / 8;
2441 let (mut d1, mut d2) = (0f32, 0f32);
2442 for g in 0..ng {
2443 let so = (r * ng + g) * 2;
2444 let s = f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]));
2445 let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2446 let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
2447 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
2448 let (mut g1, mut g2) = (0f32, 0f32);
2449 for blk in 0..GROUP_SIZE / 8 {
2450 let u = unpack8::<B>(&gd0[blk * B..]);
2451 for k in 0..8 {
2452 let w = u[k] as f32 - l;
2453 g1 += w * x1g[blk * 8 + k];
2454 g2 += w * x2g[blk * 8 + k];
2455 }
2456 }
2457 d1 += g1 * s;
2458 d2 += g2 * s;
2459 }
2460 (d1, d2)
2461 }
2462 for r in start..end {
2463 let data = &bytes[offsets[r]..offsets[r + 1]];
2464 let (v1, v2) = match bits[r] {
2465 3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
2466 4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
2467 5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
2468 6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
2469 8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
2470 b => unreachable!("vbit bit-width {b} (validated at load)"),
2471 };
2472 unsafe {
2474 *p1.at(r) = v1;
2475 *p2.at(r) = v2;
2476 }
2477 }
2478}
2479
2480#[inline]
2487#[allow(unreachable_code)]
2488fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
2489 #[cfg(target_arch = "aarch64")]
2490 unsafe {
2491 return dot_q4t_row_sdot(bytes, r, gpr, xq);
2492 }
2493 #[cfg(target_arch = "x86_64")]
2494 unsafe {
2495 return dot_q4t_row_avx2(bytes, r, gpr, xq);
2496 }
2497 let mut acc = 0f32;
2498 for gi in 0..gpr {
2499 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
2500 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2501 let mut d = 0i32;
2502 for (k, &b) in tile[2..].iter().enumerate() {
2503 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
2504 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
2505 }
2506 acc += d as f32 * s;
2507 }
2508 acc
2509}
2510
2511#[cfg(target_arch = "aarch64")]
2512#[target_feature(enable = "neon,dotprod")]
2513unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
2514 unsafe {
2517 use core::arch::aarch64::*;
2518 use core::arch::asm;
2519 let lomask = vdupq_n_u8(0x0F);
2520 let eight = vdupq_n_s8(8);
2521 let mut acc = 0f32;
2522 for gi in 0..gpr {
2523 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
2524 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2525 let b = vld1q_u8(t.add(2));
2526 let lo = vandq_u8(b, lomask);
2527 let hi = vshrq_n_u8::<4>(b);
2528 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
2529 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
2530 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
2531 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
2532 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
2533 asm!(
2534 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
2535 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
2536 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
2537 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
2538 options(pure, nomem, nostack),
2539 );
2540 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
2541 }
2542 acc
2543 }
2544}
2545
2546#[cfg(target_arch = "x86_64")]
2547#[target_feature(enable = "avx2")]
2548unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
2549 unsafe {
2551 use core::arch::x86_64::*;
2552 let lomask = _mm_set1_epi8(0x0F);
2553 let eight = _mm256_set1_epi8(8);
2554 let ones = _mm256_set1_epi16(1);
2555 let mut acc = 0f32;
2556 for gi in 0..gpr {
2557 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
2558 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2559 let b = _mm_loadu_si128(t.add(2) as *const __m128i);
2560 let lo = _mm_and_si128(b, lomask);
2561 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
2562 let w = _mm256_sub_epi8(
2563 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
2564 eight,
2565 );
2566 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2567 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
2568 let d = _mm256_madd_epi16(p16, ones);
2569 let hi128 = _mm256_extracti128_si256::<1>(d);
2570 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
2571 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2572 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2573 acc += _mm_cvtsi128_si32(s32) as f32 * s;
2574 }
2575 acc
2576 }
2577}
2578
2579#[cfg(target_arch = "x86_64")]
2584#[target_feature(enable = "avx2")]
2585unsafe fn dot_q4t_row_1x4_avx2(
2586 bytes: &[u8],
2587 r: usize,
2588 gpr: usize,
2589 xs: [&[i8]; 4],
2590) -> [f32; 4] {
2591 unsafe {
2593 use core::arch::x86_64::*;
2594 let lomask = _mm_set1_epi8(0x0F);
2595 let eight = _mm256_set1_epi8(8);
2596 let ones = _mm256_set1_epi16(1);
2597 let mut acc = [0f32; 4];
2598 for gi in 0..gpr {
2599 let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
2600 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2601 let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
2602 let lo = _mm_and_si128(bb, lomask);
2603 let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
2604 let w = _mm256_sub_epi8(
2605 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
2606 eight,
2607 );
2608 let aw = _mm256_abs_epi8(w);
2609 for (k, xq) in xs.iter().enumerate() {
2610 let x =
2611 _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2612 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
2613 let d = _mm256_madd_epi16(p16, ones);
2614 let hi128 = _mm256_extracti128_si256::<1>(d);
2615 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
2616 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2617 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2618 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
2619 }
2620 }
2621 acc
2622 }
2623}
2624
2625#[inline]
2627fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
2628 let gi = j / GROUP_SIZE;
2629 let k = j % GROUP_SIZE;
2630 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
2631 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2632 let byte = tile[2 + k / 2];
2633 let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
2634 ((nib as i32 - 8) as f32, s)
2635}
2636
2637#[inline]
2640fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
2641 let mut acc = 0f32;
2642 for gi in 0..gpr {
2643 let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
2644 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2645 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
2646 let mut ga = 0f32;
2647 for (k, &b) in tile[2..].iter().enumerate() {
2648 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
2649 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
2650 }
2651 acc += ga * s;
2652 }
2653 acc
2654}
2655
2656fn q4t_matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
2658 debug_assert_eq!(out.len(), rows);
2659 let gpr = cols / GROUP_SIZE;
2660 let out_addr = SendMut(out.as_mut_ptr());
2661 if a8w8_enabled() {
2662 let act = split_act(x);
2663 let run = move |start: usize, end: usize| {
2664 for r in start..end {
2665 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
2666 for &(j, xv) in &act.outliers {
2667 let (w, s) = q4t_outlier(bytes, r, gpr, j);
2668 acc += w * s * xv;
2669 }
2670 unsafe { *out_addr.at(r) = acc };
2672 }
2673 };
2674 dispatch_rows(pool, rows, &run);
2675 return;
2676 }
2677 let run = move |start: usize, end: usize| {
2678 for r in start..end {
2679 unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
2681 }
2682 };
2683 dispatch_rows(pool, rows, &run);
2684}
2685
2686#[allow(clippy::too_many_arguments)]
2688fn q4t_matvec2(
2689 bytes: &[u8],
2690 x1: &[f32],
2691 x2: &[f32],
2692 rows: usize,
2693 cols: usize,
2694 o1: &mut [f32],
2695 o2: &mut [f32],
2696 pool: Option<&Pool>,
2697) {
2698 let gpr = cols / GROUP_SIZE;
2699 let p1 = SendMut(o1.as_mut_ptr());
2700 let p2 = SendMut(o2.as_mut_ptr());
2701 if a8w8_enabled() {
2702 let a1 = split_act(x1);
2703 let a2 = split_act(x2);
2704 let run = move |start: usize, end: usize| {
2705 for r in start..end {
2706 let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
2707 let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
2708 for &(j, xv) in &a1.outliers {
2709 let (w, s) = q4t_outlier(bytes, r, gpr, j);
2710 v1 += w * s * xv;
2711 }
2712 for &(j, xv) in &a2.outliers {
2713 let (w, s) = q4t_outlier(bytes, r, gpr, j);
2714 v2 += w * s * xv;
2715 }
2716 unsafe {
2718 *p1.at(r) = v1;
2719 *p2.at(r) = v2;
2720 }
2721 }
2722 };
2723 dispatch_rows(pool, rows, &run);
2724 return;
2725 }
2726 let run = move |start: usize, end: usize| {
2727 for r in start..end {
2728 unsafe {
2730 *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
2731 *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
2732 }
2733 }
2734 };
2735 dispatch_rows(pool, rows, &run);
2736}
2737
2738#[allow(clippy::too_many_arguments)]
2740fn q4t_matmat(
2741 bytes: &[u8],
2742 xs_all: &[f32],
2743 b: usize,
2744 rows: usize,
2745 cols: usize,
2746 out: &mut [f32],
2747 pool: Option<&Pool>,
2748) {
2749 debug_assert_eq!(out.len(), b * rows);
2750 let gpr = cols / GROUP_SIZE;
2751 let out_addr = SendMut(out.as_mut_ptr());
2752 if a8w8_enabled() {
2753 let acts: Vec<SplitAct> =
2754 (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
2755 let acts = &acts;
2756 #[cfg(target_arch = "x86_64")]
2757 let blocked_ok = avx2_enabled()
2758 && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true);
2759 #[cfg(not(target_arch = "x86_64"))]
2760 let blocked_ok = false;
2761 let run = move |start: usize, end: usize| {
2762 for r in start..end {
2763 let mut bi = 0usize;
2764 #[cfg(target_arch = "x86_64")]
2765 if blocked_ok {
2766 while bi + 4 <= acts.len() {
2767 let xs = [
2768 acts[bi].xq.as_slice(),
2769 acts[bi + 1].xq.as_slice(),
2770 acts[bi + 2].xq.as_slice(),
2771 acts[bi + 3].xq.as_slice(),
2772 ];
2773 let d = unsafe { dot_q4t_row_1x4_avx2(bytes, r, gpr, xs) };
2774 for k in 0..4 {
2775 let act = &acts[bi + k];
2776 let mut acc = d[k] * act.sx;
2777 for &(j, xv) in &act.outliers {
2778 let (w, sc) = q4t_outlier(bytes, r, gpr, j);
2779 acc += w * sc * xv;
2780 }
2781 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
2783 }
2784 bi += 4;
2785 }
2786 }
2787 let _ = blocked_ok;
2788 while bi < acts.len() {
2789 let act = &acts[bi];
2790 let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
2791 for &(j, xv) in &act.outliers {
2792 let (w, s) = q4t_outlier(bytes, r, gpr, j);
2793 acc += w * s * xv;
2794 }
2795 unsafe { *out_addr.at(bi * rows + r) = acc };
2797 bi += 1;
2798 }
2799 }
2800 };
2801 dispatch_rows(pool, rows, &run);
2802 return;
2803 }
2804 let run = move |start: usize, end: usize| {
2805 for r in start..end {
2806 for bi in 0..b {
2807 let x = &xs_all[bi * cols..(bi + 1) * cols];
2808 unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
2810 }
2811 }
2812 };
2813 dispatch_rows(pool, rows, &run);
2814}
2815
2816fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
2825 (0..gpr)
2826 .map(|gi| {
2827 xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
2828 .iter()
2829 .map(|&v| v as i32)
2830 .sum()
2831 })
2832 .collect()
2833}
2834
2835#[inline]
2839#[allow(unreachable_code)]
2840#[cfg(target_arch = "x86_64")]
2845#[target_feature(enable = "avx2")]
2846unsafe fn dot_q1_row_avx2(
2847 bytes: &[u8],
2848 r: usize,
2849 gpr: usize,
2850 xq: &[i8],
2851 gsum: &[i32],
2852) -> f32 {
2853 unsafe {
2855 use core::arch::x86_64::*;
2856 let expand = _mm256_setr_epi8(
2858 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1,
2859 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3,
2860 );
2861 let bitsel = _mm256_setr_epi8(
2862 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128,
2863 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128,
2864 );
2865 let ones8 = _mm256_set1_epi8(1);
2866 let ones16 = _mm256_set1_epi16(1);
2867 let mut acc = 0f32;
2868 for gi in 0..gpr {
2869 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
2870 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2871 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
2872 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
2873 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
2874 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2875 let sel = _mm256_and_si256(x, mask);
2876 let p16 = _mm256_maddubs_epi16(ones8, sel);
2878 let d32 = _mm256_madd_epi16(p16, ones16);
2879 let hi128 = _mm256_extracti128_si256::<1>(d32);
2880 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
2881 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2882 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2883 let msum = _mm_cvtsi128_si32(s32);
2884 let d = 2 * msum - gsum[gi];
2887 acc += d as f32 * s;
2888 }
2889 acc
2890 }
2891}
2892
2893#[cfg(target_arch = "x86_64")]
2896#[target_feature(enable = "avx2")]
2897unsafe fn dot_q1_row_1x4_avx2(
2898 bytes: &[u8],
2899 r: usize,
2900 gpr: usize,
2901 xs: [&[i8]; 4],
2902 gsums: [&[i32]; 4],
2903) -> [f32; 4] {
2904 unsafe {
2906 use core::arch::x86_64::*;
2907 let expand = _mm256_setr_epi8(
2908 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1,
2909 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3,
2910 );
2911 let bitsel = _mm256_setr_epi8(
2912 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128,
2913 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128,
2914 );
2915 let ones8 = _mm256_set1_epi8(1);
2916 let ones16 = _mm256_set1_epi16(1);
2917 let mut acc = [0f32; 4];
2918 for gi in 0..gpr {
2919 let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
2920 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
2921 let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
2922 let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
2923 let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
2924 for (k, xq) in xs.iter().enumerate() {
2925 let x =
2926 _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2927 let sel = _mm256_and_si256(x, mask);
2928 let p16 = _mm256_maddubs_epi16(ones8, sel);
2929 let d32 = _mm256_madd_epi16(p16, ones16);
2930 let hi128 = _mm256_extracti128_si256::<1>(d32);
2931 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
2932 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2933 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2934 let msum = _mm_cvtsi128_si32(s32);
2935 let d = 2 * msum - gsums[k][gi];
2936 acc[k] += d as f32 * s;
2937 }
2938 }
2939 acc
2940 }
2941}
2942
2943fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
2944 #[cfg(target_arch = "aarch64")]
2945 unsafe {
2946 return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
2947 }
2948 #[cfg(target_arch = "x86_64")]
2949 if avx2_enabled() {
2950 unsafe {
2951 return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
2952 }
2953 }
2954 let _ = gsum;
2955 let mut acc = 0f32;
2956 for gi in 0..gpr {
2957 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
2958 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
2959 let mut d = 0i32;
2960 for (j, &b) in tile[2..].iter().enumerate() {
2961 for k in 0..8 {
2962 let w = ((b >> k) & 1) as i32 * 2 - 1;
2963 d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
2964 }
2965 }
2966 acc += d as f32 * s;
2967 }
2968 acc
2969}
2970
2971#[cfg(target_arch = "aarch64")]
2980#[target_feature(enable = "neon,dotprod")]
2981unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
2982 unsafe {
2985 use core::arch::aarch64::*;
2986 use core::arch::asm;
2987 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
2988 let m = vld1q_u8(MASKS.as_ptr());
2989 macro_rules! tile_dot {
2991 ($t:expr, $x:expr) => {{
2992 let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
2993 let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
2994 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
2995 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
2996 let x0 = vld1q_s8($x);
2997 let x1 = vld1q_s8($x.add(16));
2998 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
2999 asm!(
3000 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
3001 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
3002 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3003 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
3004 options(pure, nomem, nostack),
3005 );
3006 vaddq_s32(a0, a1)
3007 }};
3008 }
3009 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
3018 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
3019 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
3020 const IW11: [u8; 16] = [10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11];
3021 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
3022 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
3023 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
3024 let isc = vld1_u8(ISC.as_ptr());
3025 macro_rules! tile_dot_tbl {
3027 ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
3028 let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
3029 let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
3030 let x0 = vld1q_s8($x);
3031 let x1 = vld1q_s8($x.add(16));
3032 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3033 asm!(
3034 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
3035 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
3036 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3037 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
3038 options(pure, nomem, nostack),
3039 );
3040 vaddq_s32(a0, a1)
3041 }};
3042 }
3043 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
3044 let row_base = r * gpr * Q1_TILE;
3045 let abs_end = bytes.len();
3046 let xp = xq.as_ptr();
3047 let gp = gsum.as_ptr();
3048 let mut accv = vdupq_n_f32(0.0);
3049 let mut gi = 0;
3050 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
3053 let t0 = base.add(gi * Q1_TILE);
3054 let ld_a = vld1q_u8(t0);
3055 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
3056 let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
3057 let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
3058 let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
3059 let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
3060 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
3062 let g = vld1q_s32(gp.add(gi));
3063 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
3064 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
3065 let scf: float32x4_t;
3066 asm!(
3067 "fcvtl {o:v}.4s, {i:v}.4h",
3068 o = out(vreg) scf, i = in(vreg) sc16,
3069 options(pure, nomem, nostack),
3070 );
3071 accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
3072 gi += 4;
3073 }
3074 let mut acc = vaddvq_f32(accv);
3075 while gi < gpr {
3076 let t = base.add(gi * Q1_TILE);
3077 let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3078 let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
3079 acc += (-(2 * d + *gp.add(gi))) as f32 * s;
3080 gi += 1;
3081 }
3082 acc
3083 }
3084}
3085
3086#[cfg(target_arch = "aarch64")]
3091#[target_feature(enable = "neon,dotprod")]
3092unsafe fn dot_q1_row_1x4_sdot(
3093 bytes: &[u8],
3094 r: usize,
3095 gpr: usize,
3096 xs: [&[i8]; 4],
3097 gs: [&[i32]; 4],
3098) -> [f32; 4] {
3099 unsafe {
3101 use core::arch::aarch64::*;
3102 use core::arch::asm;
3103 const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
3104 const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
3105 const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
3106 const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
3107 const IW11: [u8; 16] = [10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11];
3108 const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
3109 let m = vld1q_u8(MASKS.as_ptr());
3110 let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
3111 let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
3112 let isc = vld1_u8(ISC.as_ptr());
3113 macro_rules! sdot2 {
3114 ($w0:expr, $w1:expr, $x:expr) => {{
3115 let x0 = vld1q_s8($x);
3116 let x1 = vld1q_s8($x.add(16));
3117 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3118 asm!(
3119 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
3120 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
3121 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3122 w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
3123 options(pure, nomem, nostack),
3124 );
3125 vaddq_s32(a0, a1)
3126 }};
3127 }
3128 let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
3129 let row_base = r * gpr * Q1_TILE;
3130 let abs_end = bytes.len();
3131 let mut accv = [vdupq_n_f32(0.0); 4];
3132 let mut gi = 0;
3133 while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
3134 let t0 = base.add(gi * Q1_TILE);
3135 let ld_a = vld1q_u8(t0);
3136 let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
3137 let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
3139 let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
3140 let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
3141 let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
3142 let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
3143 let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
3144 let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
3145 let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
3146 let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
3147 let scf: float32x4_t;
3148 asm!(
3149 "fcvtl {o:v}.4s, {i:v}.4h",
3150 o = out(vreg) scf, i = in(vreg) sc16,
3151 options(pure, nomem, nostack),
3152 );
3153 for k in 0..4 {
3154 let xp = xs[k].as_ptr();
3155 let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
3156 let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
3157 let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
3158 let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
3159 let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
3160 let g = vld1q_s32(gs[k].as_ptr().add(gi));
3161 let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
3162 accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
3163 }
3164 gi += 4;
3165 }
3166 let mut acc = [
3167 vaddvq_f32(accv[0]),
3168 vaddvq_f32(accv[1]),
3169 vaddvq_f32(accv[2]),
3170 vaddvq_f32(accv[3]),
3171 ];
3172 while gi < gpr {
3173 let t = base.add(gi * Q1_TILE);
3174 let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3175 let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
3176 let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
3177 let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
3178 let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
3179 for k in 0..4 {
3180 let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
3181 acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
3182 }
3183 gi += 1;
3184 }
3185 acc
3186 }
3187}
3188
3189#[inline]
3191fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
3192 let gi = j / GROUP_SIZE;
3193 let k = j % GROUP_SIZE;
3194 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
3195 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3196 let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
3197 ((bit as i32 * 2 - 1) as f32, s)
3198}
3199
3200#[inline]
3202fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
3203 let mut acc = 0f32;
3204 for gi in 0..gpr {
3205 let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
3206 let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3207 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
3208 let mut ga = 0f32;
3209 for (j, &b) in tile[2..].iter().enumerate() {
3210 for k in 0..8 {
3211 ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
3212 }
3213 }
3214 acc += ga * s;
3215 }
3216 acc
3217}
3218
3219#[allow(clippy::too_many_arguments)]
3222fn q1_range_a8w8(
3223 bytes: &[u8],
3224 gpr: usize,
3225 act: &SplitAct,
3226 gsum: &[i32],
3227 out: SendMut,
3228 start: usize,
3229 end: usize,
3230) {
3231 for r in start..end {
3232 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
3233 for &(j, xv) in &act.outliers {
3234 let (w, s) = q1_outlier(bytes, r, gpr, j);
3235 acc += w * s * xv;
3236 }
3237 unsafe { *out.at(r) = acc };
3239 }
3240}
3241
3242fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
3244 for r in start..end {
3245 unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
3247 }
3248}
3249
3250fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
3255 let entries = base_len + (rows + 1) * 4;
3256 (base_len, entries, entries <= bytes.len())
3257}
3258
3259#[inline]
3261fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
3262 let o = rp_off + r * 4;
3263 u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
3264}
3265
3266const SIGN5: [[f32; 5]; 256] = {
3270 let mut lut = [[0.0f32; 5]; 256];
3271 let pow3 = [1u16, 3, 9, 27, 81];
3272 let mut byte = 0usize;
3273 while byte < 256 {
3274 let mut i = 0usize;
3275 while i < 5 {
3276 let code = (byte as u16 / pow3[i]) % 3;
3277 lut[byte][i] = if code == 1 {
3278 1.0
3279 } else if code == 2 {
3280 -1.0
3281 } else {
3282 0.0
3283 };
3284 i += 1;
3285 }
3286 byte += 1;
3287 }
3288 lut
3289};
3290
3291const SIGN5_I8: [[i8; 5]; 256] = {
3293 let mut lut = [[0i8; 5]; 256];
3294 let pow3 = [1u16, 3, 9, 27, 81];
3295 let mut byte = 0usize;
3296 while byte < 256 {
3297 let mut i = 0usize;
3298 while i < 5 {
3299 let code = (byte as u16 / pow3[i]) % 3;
3300 lut[byte][i] = if code == 1 {
3301 1
3302 } else if code == 2 {
3303 -1
3304 } else {
3305 0
3306 };
3307 i += 1;
3308 }
3309 byte += 1;
3310 }
3311 lut
3312};
3313
3314const SIGN5_U64: [u64; 256] = {
3320 let mut lut = [0u64; 256];
3321 let pow3 = [1u16, 3, 9, 27, 81];
3322 let mut byte = 0usize;
3323 while byte < 256 {
3324 let mut v = 0u64;
3325 let mut i = 0usize;
3326 while i < 5 {
3327 let code = (byte as u16 / pow3[i]) % 3;
3328 let s: u8 = if code == 1 {
3329 1
3330 } else if code == 2 {
3331 0xFF
3332 } else {
3333 0
3334 };
3335 v |= (s as u64) << (i * 8);
3336 i += 1;
3337 }
3338 lut[byte] = v;
3339 byte += 1;
3340 }
3341 lut
3342};
3343
3344#[inline]
3349fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
3350 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3351 let off = (r * gpr + j / GROUP_SIZE) * TILE;
3352 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
3353 let within = j % GROUP_SIZE;
3354 SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
3355}
3356
3357#[cfg(target_arch = "aarch64")]
3360#[target_feature(enable = "neon,dotprod")]
3361#[inline]
3362unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
3363 unsafe {
3365 use core::arch::aarch64::*;
3366 use core::arch::asm;
3367 let w0 = vld1q_s8(w);
3368 let w1 = vld1q_s8(w.add(16));
3369 let x0 = vld1q_s8(x);
3370 let x1 = vld1q_s8(x.add(16));
3371 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3372 asm!(
3373 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
3374 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
3375 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3376 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
3377 options(pure, nomem, nostack),
3378 );
3379 vaddvq_s32(vaddq_s32(a0, a1))
3380 }
3381}
3382
3383#[cfg(target_arch = "x86_64")]
3386#[target_feature(enable = "avx2")]
3387#[inline]
3388unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
3389 unsafe {
3391 use core::arch::x86_64::*;
3392 let wv = _mm256_loadu_si256(w as *const __m256i);
3393 let xv = _mm256_loadu_si256(x as *const __m256i);
3394 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
3395 let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
3396 let hi128 = _mm256_extracti128_si256::<1>(d);
3397 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3398 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3399 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3400 _mm_cvtsi128_si32(s32)
3401 }
3402}
3403
3404#[inline]
3409fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
3410 debug_assert!(dst.len() >= 40);
3411 unsafe {
3414 let p = dst.as_mut_ptr();
3415 for bi in 0..7 {
3416 core::ptr::write_unaligned(
3417 p.add(bi * 5) as *mut u64,
3418 SIGN5_U64[*codes.add(bi) as usize],
3419 );
3420 }
3421 }
3422}
3423
3424#[inline]
3429fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
3430 #[cfg(target_arch = "aarch64")]
3431 unsafe {
3432 return sdot32_i8(w, x);
3433 }
3434 #[cfg(target_arch = "x86_64")]
3435 unsafe {
3436 return i8dot32_avx2(w, x);
3437 }
3438 #[allow(unreachable_code)]
3439 unsafe {
3440 let mut s = 0i32;
3441 for k in 0..GROUP_SIZE {
3442 s += *w.add(k) as i32 * *x.add(k) as i32;
3443 }
3444 s
3445 }
3446}
3447
3448#[cfg(target_arch = "aarch64")]
3452#[target_feature(enable = "neon,dotprod")]
3453unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3454 unsafe {
3456 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3457 let mut acc = 0f32;
3458 let mut sg = [0i8; GROUP_SIZE + 8]; for gi in 0..gpr {
3460 let off = (r * gpr + gi) * TILE;
3461 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
3462 q1t_unpack_group_i8(bytes.as_ptr().add(off + 2), &mut sg);
3463 acc += sdot32_i8(sg.as_ptr(), xq.as_ptr().add(gi * GROUP_SIZE)) as f32 * s;
3464 }
3465 acc
3466 }
3467}
3468
3469#[cfg(target_arch = "x86_64")]
3471#[target_feature(enable = "avx2")]
3472unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3473 unsafe {
3475 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3476 let mut acc = 0f32;
3477 let mut sg = [0i8; GROUP_SIZE + 8]; for gi in 0..gpr {
3479 let off = (r * gpr + gi) * TILE;
3480 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
3481 q1t_unpack_group_i8(bytes.as_ptr().add(off + 2), &mut sg);
3482 acc += i8dot32_avx2(sg.as_ptr(), xq.as_ptr().add(gi * GROUP_SIZE)) as f32 * s;
3483 }
3484 acc
3485 }
3486}
3487
3488#[inline]
3492fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3493 #[cfg(target_arch = "aarch64")]
3494 unsafe {
3495 return q1t_dot_row_sdot(bytes, r, gpr, xq);
3496 }
3497 #[cfg(target_arch = "x86_64")]
3498 unsafe {
3499 return q1t_dot_row_avx2(bytes, r, gpr, xq);
3500 }
3501 #[allow(unreachable_code)]
3502 {
3503 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3504 let mut acc = 0f32;
3505 let mut sg = [0i8; GROUP_SIZE + 8]; for gi in 0..gpr {
3507 let off = (r * gpr + gi) * TILE;
3508 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
3509 q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
3510 let mut d = 0i32;
3511 for k in 0..GROUP_SIZE {
3512 d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
3513 }
3514 acc += d as f32 * s;
3515 }
3516 acc
3517 }
3518}
3519
3520fn q1t_row_outlier_correction(
3527 bytes: &[u8],
3528 r: usize,
3529 rp_off: usize,
3530 entries_off: usize,
3531 has_ov: bool,
3532 x: &[f32],
3533) -> f32 {
3534 if !has_ov {
3535 return 0.0;
3536 }
3537 let (c0, c1) = (q1t_rowptr(bytes, rp_off, r), q1t_rowptr(bytes, rp_off, r + 1));
3538 let mut corr = 0f32;
3539 for p in c0..c1 {
3540 let e = entries_off + p * 4;
3541 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
3542 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
3543 corr += val * x[col];
3544 }
3545 corr
3546}
3547
3548fn q1t_dequant_row(
3552 bytes: &[u8],
3553 r: usize,
3554 gpr: usize,
3555 rp_off: usize,
3556 entries_off: usize,
3557 has_ov: bool,
3558 buf: &mut [f32],
3559) {
3560 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3561 for g in 0..gpr {
3562 let off = (r * gpr + g) * TILE;
3563 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
3564 let codes = &bytes[off + 2..off + TILE];
3565 let bc = g * GROUP_SIZE;
3566 for bi in 0..6 {
3568 let lut = &SIGN5[codes[bi] as usize];
3569 let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
3570 for i in 0..5 {
3571 d[i] = lut[i] * s;
3572 }
3573 }
3574 let lut = &SIGN5[codes[6] as usize];
3575 buf[bc + 30] = lut[0] * s;
3576 buf[bc + 31] = lut[1] * s;
3577 }
3578 if !has_ov {
3579 return;
3580 }
3581 let (c0, c1) = (q1t_rowptr(bytes, rp_off, r), q1t_rowptr(bytes, rp_off, r + 1));
3582 for p in c0..c1 {
3583 let e = entries_off + p * 4;
3584 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
3585 buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
3586 }
3587}
3588
3589fn q1t_add_overlay(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
3593 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3594 let gpr = cols / GROUP_SIZE;
3595 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
3596 if !has_ov {
3597 return;
3598 }
3599 let out_addr = SendMut(out.as_mut_ptr());
3600 let run = move |start: usize, end: usize| {
3601 for r in start..end {
3602 let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
3603 unsafe { *out_addr.at(r) += corr };
3605 }
3606 };
3607 dispatch_rows(pool, rows, &run);
3608}
3609
3610fn q1t_matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
3614 debug_assert_eq!(out.len(), rows);
3615 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3616 let gpr = cols / GROUP_SIZE;
3617 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
3618 let out_addr = SendMut(out.as_mut_ptr());
3619 if a8w8_enabled() {
3623 let act = split_act(x);
3624 let act = &act;
3625 let run = move |start: usize, end: usize| {
3626 for r in start..end {
3627 let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
3628 for &(j, xv) in &act.outliers {
3629 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
3630 }
3631 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
3632 unsafe { *out_addr.at(r) = acc };
3634 }
3635 };
3636 dispatch_rows(pool, rows, &run);
3637 return;
3638 }
3639 let run = move |start: usize, end: usize| {
3640 let mut sg = [0f32; GROUP_SIZE];
3644 for r in start..end {
3645 let mut acc = 0f32;
3646 for g in 0..gpr {
3647 let off = (r * gpr + g) * TILE;
3648 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
3649 let codes = &bytes[off + 2..off + TILE];
3650 let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
3651 for bi in 0..6 {
3652 sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
3653 }
3654 let lut = &SIGN5[codes[6] as usize];
3655 sg[30] = lut[0];
3656 sg[31] = lut[1];
3657 let mut gsum = 0f32;
3658 for k in 0..GROUP_SIZE {
3659 gsum += sg[k] * xg[k];
3660 }
3661 acc += s * gsum;
3662 }
3663 acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
3664 unsafe { *out_addr.at(r) = acc };
3665 }
3666 };
3667 dispatch_rows(pool, rows, &run);
3668}
3669
3670fn q1t_matmat(
3673 bytes: &[u8],
3674 xs: &[f32],
3675 b: usize,
3676 rows: usize,
3677 cols: usize,
3678 out: &mut [f32],
3679 pool: Option<&Pool>,
3680) {
3681 debug_assert_eq!(out.len(), b * rows);
3682 const TILE: usize = cortiq_core::quant::Q1T_TILE;
3683 let gpr = cols / GROUP_SIZE;
3684 let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
3685 let out_addr = SendMut(out.as_mut_ptr());
3686 if a8w8_enabled() {
3690 let acts: Vec<SplitAct> = (0..b)
3691 .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
3692 .collect();
3693 let acts = &acts;
3694 let run = move |start: usize, end: usize| {
3695 let mut sg = vec![0i8; cols + 8]; let mut sc = vec![0f32; gpr]; let mut accs = vec![0f32; b]; for r in start..end {
3699 for g in 0..gpr {
3700 let off = (r * gpr + g) * TILE;
3701 sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
3702 q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg[g * GROUP_SIZE..]);
3703 }
3704 for bi in 0..b {
3705 let act = &acts[bi];
3706 let mut isum = 0f32;
3707 for g in 0..gpr {
3708 let d = q1t_i8dot32(
3709 sg.as_ptr().wrapping_add(g * GROUP_SIZE),
3710 act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
3711 );
3712 isum += d as f32 * sc[g];
3713 }
3714 let mut acc = isum * act.sx;
3715 for &(j, xv) in &act.outliers {
3716 acc += q1t_base_weight(bytes, r, gpr, j) * xv;
3717 }
3718 accs[bi] = acc;
3719 }
3720 if has_ov {
3724 let (c0, c1) =
3725 (q1t_rowptr(bytes, rp_off, r), q1t_rowptr(bytes, rp_off, r + 1));
3726 for p in c0..c1 {
3727 let e = ent_off + p * 4;
3728 let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
3729 let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
3730 for bi in 0..b {
3731 accs[bi] += val * xs[bi * cols + col];
3732 }
3733 }
3734 }
3735 for bi in 0..b {
3736 unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
3737 }
3738 }
3739 };
3740 dispatch_rows(pool, rows, &run);
3741 return;
3742 }
3743 let run = move |start: usize, end: usize| {
3744 let mut buf = vec![0f32; cols];
3745 for r in start..end {
3746 q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
3747 for bi in 0..b {
3748 let xr = &xs[bi * cols..(bi + 1) * cols];
3749 let mut acc = 0f32;
3750 for j in 0..cols {
3751 acc += buf[j] * xr[j];
3752 }
3753 unsafe { *out_addr.at(bi * rows + r) = acc };
3754 }
3755 }
3756 };
3757 dispatch_rows(pool, rows, &run);
3758}
3759
3760fn q1_matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
3761 debug_assert_eq!(out.len(), rows);
3762 let gpr = cols / GROUP_SIZE;
3763 let out_addr = SendMut(out.as_mut_ptr());
3764 if a8w8_enabled() {
3765 let act = split_act(x);
3766 let gsum = q1_group_sums(&act.xq, gpr);
3767 let (act, gsum) = (&act, &gsum);
3768 let run =
3769 move |start: usize, end: usize| q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end);
3770 dispatch_rows(pool, rows, &run);
3771 return;
3772 }
3773 let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
3774 dispatch_rows(pool, rows, &run);
3775}
3776
3777#[allow(clippy::too_many_arguments)]
3779fn q1_matvec2(
3780 bytes: &[u8],
3781 x1: &[f32],
3782 x2: &[f32],
3783 rows: usize,
3784 cols: usize,
3785 o1: &mut [f32],
3786 o2: &mut [f32],
3787 pool: Option<&Pool>,
3788) {
3789 let gpr = cols / GROUP_SIZE;
3790 let p1 = SendMut(o1.as_mut_ptr());
3791 let p2 = SendMut(o2.as_mut_ptr());
3792 if a8w8_enabled() {
3793 let a1 = split_act(x1);
3794 let a2 = split_act(x2);
3795 let g1 = q1_group_sums(&a1.xq, gpr);
3796 let g2 = q1_group_sums(&a2.xq, gpr);
3797 let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
3798 let run = move |start: usize, end: usize| {
3799 for r in start..end {
3800 let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
3801 let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
3802 for &(j, xv) in &a1.outliers {
3803 let (w, s) = q1_outlier(bytes, r, gpr, j);
3804 v1 += w * s * xv;
3805 }
3806 for &(j, xv) in &a2.outliers {
3807 let (w, s) = q1_outlier(bytes, r, gpr, j);
3808 v2 += w * s * xv;
3809 }
3810 unsafe {
3812 *p1.at(r) = v1;
3813 *p2.at(r) = v2;
3814 }
3815 }
3816 };
3817 dispatch_rows(pool, rows, &run);
3818 return;
3819 }
3820 let run = move |start: usize, end: usize| {
3821 for r in start..end {
3822 unsafe {
3824 *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
3825 *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
3826 }
3827 }
3828 };
3829 dispatch_rows(pool, rows, &run);
3830}
3831
3832#[allow(clippy::too_many_arguments)]
3834fn q1_matmat(
3835 bytes: &[u8],
3836 xs_all: &[f32],
3837 b: usize,
3838 rows: usize,
3839 cols: usize,
3840 out: &mut [f32],
3841 pool: Option<&Pool>,
3842) {
3843 debug_assert_eq!(out.len(), b * rows);
3844 let gpr = cols / GROUP_SIZE;
3845 let out_addr = SendMut(out.as_mut_ptr());
3846 if a8w8_enabled() {
3847 let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
3848 .map(|bi| {
3849 let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
3850 let gsum = q1_group_sums(&act.xq, gpr);
3851 (act, gsum)
3852 })
3853 .collect();
3854 let acts = &acts;
3855 #[cfg(target_arch = "x86_64")]
3856 let blocked_ok = avx2_enabled()
3857 && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true);
3858 #[cfg(target_arch = "aarch64")]
3859 let blocked_ok = sdot_enabled()
3860 && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true);
3861 let run = move |start: usize, end: usize| {
3862 for r in start..end {
3863 let mut bi = 0usize;
3864 #[cfg(target_arch = "aarch64")]
3867 if blocked_ok {
3868 while bi + 4 <= acts.len() {
3869 let xs = [
3870 acts[bi].0.xq.as_slice(),
3871 acts[bi + 1].0.xq.as_slice(),
3872 acts[bi + 2].0.xq.as_slice(),
3873 acts[bi + 3].0.xq.as_slice(),
3874 ];
3875 let gs = [
3876 acts[bi].1.as_slice(),
3877 acts[bi + 1].1.as_slice(),
3878 acts[bi + 2].1.as_slice(),
3879 acts[bi + 3].1.as_slice(),
3880 ];
3881 let d = unsafe { dot_q1_row_1x4_sdot(bytes, r, gpr, xs, gs) };
3882 for k in 0..4 {
3883 let (act, _) = &acts[bi + k];
3884 let mut acc = d[k] * act.sx;
3885 for &(j, xv) in &act.outliers {
3886 let (w, sc) = q1_outlier(bytes, r, gpr, j);
3887 acc += w * sc * xv;
3888 }
3889 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
3891 }
3892 bi += 4;
3893 }
3894 }
3895 #[cfg(target_arch = "x86_64")]
3896 if blocked_ok {
3897 while bi + 4 <= acts.len() {
3898 let xs = [
3899 acts[bi].0.xq.as_slice(),
3900 acts[bi + 1].0.xq.as_slice(),
3901 acts[bi + 2].0.xq.as_slice(),
3902 acts[bi + 3].0.xq.as_slice(),
3903 ];
3904 let gs = [
3905 acts[bi].1.as_slice(),
3906 acts[bi + 1].1.as_slice(),
3907 acts[bi + 2].1.as_slice(),
3908 acts[bi + 3].1.as_slice(),
3909 ];
3910 let d = unsafe { dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs) };
3911 for k in 0..4 {
3912 let (act, _) = &acts[bi + k];
3913 let mut acc = d[k] * act.sx;
3914 for &(j, xv) in &act.outliers {
3915 let (w, sc) = q1_outlier(bytes, r, gpr, j);
3916 acc += w * sc * xv;
3917 }
3918 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
3920 }
3921 bi += 4;
3922 }
3923 }
3924 while bi < acts.len() {
3925 let (act, gsum) = &acts[bi];
3926 let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
3927 for &(j, xv) in &act.outliers {
3928 let (w, s) = q1_outlier(bytes, r, gpr, j);
3929 acc += w * s * xv;
3930 }
3931 unsafe { *out_addr.at(bi * rows + r) = acc };
3933 bi += 1;
3934 }
3935 }
3936 };
3937 dispatch_rows(pool, rows, &run);
3938 return;
3939 }
3940 let run = move |start: usize, end: usize| {
3941 for r in start..end {
3942 for bi in 0..b {
3943 let x = &xs_all[bi * cols..(bi + 1) * cols];
3944 unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
3946 }
3947 }
3948 };
3949 dispatch_rows(pool, rows, &run);
3950}
3951
3952fn q4matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
3958 debug_assert_eq!(out.len(), rows);
3959 let (packed, scales) = q4_split(bytes, rows, cols);
3960 let gpr = cols / GROUP_SIZE;
3961 let out_addr = SendMut(out.as_mut_ptr());
3962
3963 if a8w8_enabled() {
3964 let act = split_act(x);
3965 let run = move |start: usize, end: usize| {
3966 q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
3967 };
3968 dispatch_rows(pool, rows, &run);
3969 return;
3970 }
3971
3972 let run = move |start: usize, end: usize| {
3973 q4_range_f32(packed, scales, gpr, x, out_addr, start, end)
3974 };
3975 dispatch_rows(pool, rows, &run);
3976}
3977
3978#[inline]
3981#[allow(unreachable_code)]
3982#[cfg(target_arch = "x86_64")]
3987#[target_feature(enable = "avx2")]
3988unsafe fn dot_q4b_row_1x4_avx2(
3989 buf: &[u8],
3990 scales: &[u8],
3991 g0: usize,
3992 gpr: usize,
3993 xs: [&[i8]; 4],
3994) -> [f32; 4] {
3995 unsafe {
3997 use core::arch::x86_64::*;
3998 let ones = _mm256_set1_epi16(1);
3999 let mut acc = [0f32; 4];
4000 for gi in 0..gpr {
4001 let s = f16_to_f32(u16::from_le_bytes([
4002 scales[(g0 + gi) * 2],
4003 scales[(g0 + gi) * 2 + 1],
4004 ]));
4005 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4006 let aw = _mm256_abs_epi8(w);
4007 for (k, xq) in xs.iter().enumerate() {
4008 let x =
4009 _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4010 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
4011 let d = _mm256_madd_epi16(p16, ones);
4012 let hi128 = _mm256_extracti128_si256::<1>(d);
4013 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4014 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4015 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4016 acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
4017 }
4018 }
4019 acc
4020 }
4021}
4022
4023#[cfg(target_arch = "x86_64")]
4029#[target_feature(enable = "avx2")]
4030unsafe fn dot_q4b_row_1x4_sx_avx2(
4031 buf: &[u8],
4032 scales: &[u8],
4033 g0: usize,
4034 gpr: usize,
4035 xs: [&[i8]; 4],
4036 sxs: [f32; 4],
4037) -> [f32; 4] {
4038 unsafe {
4040 use core::arch::x86_64::*;
4041 let ones = _mm256_set1_epi16(1);
4042 let mut acc = [0f32; 4];
4043 for gi in 0..gpr {
4044 let s = f16_to_f32(u16::from_le_bytes([
4045 scales[(g0 + gi) * 2],
4046 scales[(g0 + gi) * 2 + 1],
4047 ]));
4048 let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4049 let aw = _mm256_abs_epi8(w);
4050 for (k, xq) in xs.iter().enumerate() {
4051 let x =
4052 _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4053 let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
4054 let d = _mm256_madd_epi16(p16, ones);
4055 let hi128 = _mm256_extracti128_si256::<1>(d);
4056 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4057 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4058 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4059 acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
4060 }
4061 }
4062 acc
4063 }
4064}
4065
4066fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
4067 #[cfg(target_arch = "aarch64")]
4068 unsafe {
4069 return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
4070 }
4071 #[cfg(target_arch = "x86_64")]
4072 unsafe {
4073 return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
4074 }
4075 let mut acc = 0f32;
4076 for gi in 0..gpr {
4077 let g = g0 + gi;
4078 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4079 let mut d = 0i32;
4080 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
4081 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
4082 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
4083 }
4084 acc += d as f32 * s;
4085 }
4086 acc
4087}
4088
4089#[inline]
4091#[allow(unreachable_code)]
4092fn dot_q4_row_i8_2(
4093 packed: &[u8],
4094 scales: &[u8],
4095 g0: usize,
4096 gpr: usize,
4097 xq1: &[i8],
4098 xq2: &[i8],
4099) -> (f32, f32) {
4100 #[cfg(target_arch = "aarch64")]
4101 unsafe {
4102 return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
4103 }
4104 #[cfg(target_arch = "x86_64")]
4105 unsafe {
4106 return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
4107 }
4108 (
4109 dot_q4_row_i8(packed, scales, g0, gpr, xq1),
4110 dot_q4_row_i8(packed, scales, g0, gpr, xq2),
4111 )
4112}
4113
4114#[allow(clippy::too_many_arguments)]
4117fn q4_range_a8w8(
4118 packed: &[u8],
4119 scales: &[u8],
4120 gpr: usize,
4121 cols: usize,
4122 act: &SplitAct,
4123 out: SendMut,
4124 start: usize,
4125 end: usize,
4126) {
4127 for r in start..end {
4128 let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
4129 for &(j, xv) in &act.outliers {
4131 let flat = r * cols + j;
4132 let byte = packed[flat / 2];
4133 let nib = if flat & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4134 let s = f16_to_f32(u16::from_le_bytes([
4135 scales[(flat / GROUP_SIZE) * 2],
4136 scales[(flat / GROUP_SIZE) * 2 + 1],
4137 ]));
4138 acc += ((nib as i32 - 8) as f32) * s * xv;
4139 }
4140 unsafe { *out.at(r) = acc };
4142 }
4143}
4144
4145#[allow(clippy::too_many_arguments)]
4148fn q4_range2_a8w8(
4149 packed: &[u8],
4150 scales: &[u8],
4151 gpr: usize,
4152 cols: usize,
4153 a1: &SplitAct,
4154 a2: &SplitAct,
4155 p1: SendMut,
4156 p2: SendMut,
4157 start: usize,
4158 end: usize,
4159) {
4160 for r in start..end {
4161 let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
4162 let mut acc1 = s1 * a1.sx;
4163 let mut acc2 = s2 * a2.sx;
4164 let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
4166 for &(j, xv) in outliers {
4167 let flat = r * cols + j;
4168 let byte = packed[flat / 2];
4169 let nib = if flat & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4170 let s = f16_to_f32(u16::from_le_bytes([
4171 scales[(flat / GROUP_SIZE) * 2],
4172 scales[(flat / GROUP_SIZE) * 2 + 1],
4173 ]));
4174 *acc += ((nib as i32 - 8) as f32) * s * xv;
4175 }
4176 };
4177 fix(&a1.outliers, &mut acc1);
4178 fix(&a2.outliers, &mut acc2);
4179 unsafe {
4181 *p1.at(r) = acc1;
4182 *p2.at(r) = acc2;
4183 }
4184 }
4185}
4186
4187fn q4_range_f32(
4189 packed: &[u8],
4190 scales: &[u8],
4191 gpr: usize,
4192 x: &[f32],
4193 out: SendMut,
4194 start: usize,
4195 end: usize,
4196) {
4197 for r in start..end {
4198 let mut acc = 0f32;
4199 for gi in 0..gpr {
4200 let g = r * gpr + gi;
4201 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4202 let pk = &packed[g * 16..(g + 1) * 16];
4203 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4204 let mut ga = 0f32;
4205 for (k, &b) in pk.iter().enumerate() {
4206 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4207 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4208 }
4209 acc += ga * s;
4210 }
4211 unsafe { *out.at(r) = acc };
4213 }
4214}
4215
4216#[allow(clippy::too_many_arguments)]
4220fn q4matvec2(
4221 bytes: &[u8],
4222 x1: &[f32],
4223 x2: &[f32],
4224 rows: usize,
4225 cols: usize,
4226 o1: &mut [f32],
4227 o2: &mut [f32],
4228 pool: Option<&Pool>,
4229) {
4230 debug_assert_eq!(o1.len(), rows);
4231 debug_assert_eq!(o2.len(), rows);
4232 let (packed, scales) = q4_split(bytes, rows, cols);
4233 let gpr = cols / GROUP_SIZE;
4234
4235 if a8w8_enabled() {
4236 let a1 = split_act(x1);
4237 let a2 = split_act(x2);
4238 let p1 = SendMut(o1.as_mut_ptr());
4239 let p2 = SendMut(o2.as_mut_ptr());
4240 let run = move |start: usize, end: usize| {
4241 q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
4242 };
4243 dispatch_rows(pool, rows, &run);
4244 return;
4245 }
4246
4247 let p1 = SendMut(o1.as_mut_ptr());
4248 let p2 = SendMut(o2.as_mut_ptr());
4249 let run = move |start: usize, end: usize| {
4250 q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
4251 };
4252 dispatch_rows(pool, rows, &run);
4253}
4254
4255#[allow(clippy::too_many_arguments)]
4257fn q4_range2_f32(
4258 packed: &[u8],
4259 scales: &[u8],
4260 gpr: usize,
4261 x1: &[f32],
4262 x2: &[f32],
4263 p1: SendMut,
4264 p2: SendMut,
4265 start: usize,
4266 end: usize,
4267) {
4268 for r in start..end {
4269 let (mut acc1, mut acc2) = (0f32, 0f32);
4270 for gi in 0..gpr {
4271 let g = r * gpr + gi;
4272 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4273 let pk = &packed[g * 16..(g + 1) * 16];
4274 let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4275 let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4276 let (mut g1, mut g2) = (0f32, 0f32);
4277 for (k, &b) in pk.iter().enumerate() {
4278 let wl = (b & 0x0F) as f32 - 8.0;
4279 let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
4280 g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
4281 g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
4282 }
4283 acc1 += g1 * s;
4284 acc2 += g2 * s;
4285 }
4286 unsafe {
4288 *p1.at(r) = acc1;
4289 *p2.at(r) = acc2;
4290 }
4291 }
4292}
4293
4294thread_local! {
4295 static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
4298 static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4299}
4300
4301#[allow(clippy::too_many_arguments)]
4307fn q4matmat(
4308 bytes: &[u8],
4309 xs_all: &[f32],
4310 b: usize,
4311 rows: usize,
4312 cols: usize,
4313 out: &mut [f32],
4314 pool: Option<&Pool>,
4315) {
4316 debug_assert_eq!(xs_all.len(), b * cols);
4317 debug_assert_eq!(out.len(), b * rows);
4318 let (packed, scales) = q4_split(bytes, rows, cols);
4319 let gpr = cols / GROUP_SIZE;
4320 let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
4321
4322 if a8w8_enabled() {
4323 let acts: Vec<SplitAct> =
4324 (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
4325 let acts = &acts;
4326 let out_addr = SendMut(out.as_mut_ptr());
4327 let run = move |start: usize, end: usize| {
4328 ROW_I8.with(|rb| {
4329 let mut buf = rb.borrow_mut();
4330 buf.resize(cols, 0);
4331 for r in start..end {
4332 for gi in 0..gpr {
4336 let g = r * gpr + gi;
4337 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
4338 buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
4339 buf[gi * GROUP_SIZE + k * 2 + 1] =
4340 (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
4341 }
4342 }
4343 let mut bi = 0usize;
4344 #[cfg(target_arch = "x86_64")]
4345 if avx2_enabled()
4346 && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true)
4347 {
4348 while bi + 4 <= acts.len() {
4349 let xs = [
4350 acts[bi].xq.as_slice(),
4351 acts[bi + 1].xq.as_slice(),
4352 acts[bi + 2].xq.as_slice(),
4353 acts[bi + 3].xq.as_slice(),
4354 ];
4355 let d = unsafe {
4356 dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
4357 };
4358 for k in 0..4 {
4359 let act = &acts[bi + k];
4360 let mut acc = d[k] * act.sx;
4361 for &(j, xv) in &act.outliers {
4362 acc += (buf[j] as i8) as f32
4363 * gscale((r * cols + j) / GROUP_SIZE)
4364 * xv;
4365 }
4366 unsafe { *out_addr.at((bi + k) * rows + r) = acc };
4368 }
4369 bi += 4;
4370 }
4371 }
4372 while bi < acts.len() {
4373 let act = &acts[bi];
4374 let mut acc = 0f32;
4375 for gi in 0..gpr {
4376 let d = dot_i8_i8(
4377 &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
4378 &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
4379 );
4380 acc += d as f32 * gscale(r * gpr + gi);
4381 }
4382 acc *= act.sx;
4383 for &(j, xv) in &act.outliers {
4385 acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
4386 }
4387 unsafe { *out_addr.at(bi * rows + r) = acc };
4389 bi += 1;
4390 }
4391 }
4392 })
4393 };
4394 dispatch_rows(pool, rows, &run);
4395 return;
4396 }
4397
4398 let out_addr = SendMut(out.as_mut_ptr());
4399 let run = move |start: usize, end: usize| {
4400 ROW_F32.with(|rb| {
4401 let mut buf = rb.borrow_mut();
4402 buf.resize(cols, 0.0);
4403 for r in start..end {
4404 for gi in 0..gpr {
4407 let g = r * gpr + gi;
4408 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
4409 buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
4410 buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
4411 }
4412 }
4413 for bi in 0..b {
4414 let x = &xs_all[bi * cols..(bi + 1) * cols];
4415 let mut acc = 0f32;
4416 for gi in 0..gpr {
4417 let mut ga = 0f32;
4418 for k in 0..GROUP_SIZE / 2 {
4423 let e = gi * GROUP_SIZE + k * 2;
4424 ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
4425 }
4426 acc += ga * gscale(r * gpr + gi);
4427 }
4428 unsafe { *out_addr.at(bi * rows + r) = acc };
4430 }
4431 }
4432 })
4433 };
4434 dispatch_rows(pool, rows, &run);
4435}
4436
4437#[allow(clippy::too_many_arguments)]
4442fn vbitmatmat(
4443 bytes: &[u8],
4444 offsets: &[usize],
4445 xs_all: &[f32],
4446 b: usize,
4447 rows: usize,
4448 cols: usize,
4449 out: &mut [f32],
4450 pool: Option<&Pool>,
4451) {
4452 debug_assert_eq!(xs_all.len(), b * cols);
4453 debug_assert_eq!(out.len(), b * rows);
4454 debug_assert_eq!(offsets.len(), rows + 1);
4455 let ng = cols / GROUP_SIZE;
4456 let bits = &bytes[..rows];
4457 let sc_off = rows;
4458 let gscale = |r: usize, g: usize| {
4459 let so = (r * ng + g) * 2;
4460 f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]))
4461 };
4462
4463 let decode_f32 = |r: usize, dst: &mut [f32]| {
4465 let bw = bits[r] as usize;
4466 let l = ((1i32 << (bw - 1)) - 1) as f32;
4467 let data = &bytes[offsets[r]..offsets[r + 1]];
4468 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
4469 for d in dst.iter_mut() {
4470 while nbits < bw {
4471 acc = (acc << 8) | data[idx] as u64;
4472 idx += 1;
4473 nbits += 8;
4474 }
4475 let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
4476 nbits -= bw;
4477 *d = u - l;
4478 }
4479 };
4480
4481 if a8w8_enabled() {
4482 let acts: Vec<SplitAct> =
4483 (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
4484 let acts = &acts;
4485 let out_addr = SendMut(out.as_mut_ptr());
4486 let run = move |start: usize, end: usize| {
4487 for r in start..end {
4488 let bw = bits[r] as usize;
4489 if bw == 8 {
4490 ROW_F32.with(|rb| {
4493 let mut buf = rb.borrow_mut();
4494 buf.resize(cols, 0.0);
4495 decode_f32(r, &mut buf);
4496 for bi in 0..b {
4497 let x = &xs_all[bi * cols..(bi + 1) * cols];
4498 let mut dot = 0f32;
4499 for g in 0..ng {
4500 let mut gd = 0f32;
4501 for k in 0..GROUP_SIZE {
4502 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
4503 }
4504 dot += gd * gscale(r, g);
4505 }
4506 unsafe { *out_addr.at(bi * rows + r) = dot };
4508 }
4509 });
4510 continue;
4511 }
4512 let l = (1i32 << (bw - 1)) - 1;
4513 let data = &bytes[offsets[r]..offsets[r + 1]];
4514 ROW_I8.with(|rb| {
4515 let mut buf = rb.borrow_mut();
4516 buf.resize(cols, 0);
4517 #[inline(always)]
4518 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
4519 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4520 let u = unpack8::<B>(&data[blk * B..]);
4521 for k in 0..8 {
4522 chunk[k] = (u[k] - l) as i8 as u8;
4523 }
4524 }
4525 }
4526 match bw {
4527 3 => fill::<3>(data, l, &mut buf),
4528 4 => vbit_fill4(data, &mut buf),
4529 5 => fill::<5>(data, l, &mut buf),
4530 6 => fill::<6>(data, l, &mut buf),
4531 _ => unreachable!("vbit bit-width {bw} (validated at load)"),
4532 }
4533 let mut bi = 0usize;
4534 #[cfg(target_arch = "x86_64")]
4538 if avx2_enabled()
4539 && std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true)
4540 {
4541 while bi + 4 <= acts.len() {
4542 let xs = [
4543 acts[bi].xq.as_slice(),
4544 acts[bi + 1].xq.as_slice(),
4545 acts[bi + 2].xq.as_slice(),
4546 acts[bi + 3].xq.as_slice(),
4547 ];
4548 let sxs = [
4549 acts[bi].sx,
4550 acts[bi + 1].sx,
4551 acts[bi + 2].sx,
4552 acts[bi + 3].sx,
4553 ];
4554 let d = unsafe {
4555 dot_q4b_row_1x4_sx_avx2(&buf, &bytes[sc_off..], r * ng, ng, xs, sxs)
4556 };
4557 for k in 0..4 {
4558 let act = &acts[bi + k];
4559 let mut dot = d[k];
4560 for &(j, xv) in &act.outliers {
4561 dot +=
4562 (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
4563 }
4564 unsafe { *out_addr.at((bi + k) * rows + r) = dot };
4566 }
4567 bi += 4;
4568 }
4569 }
4570 while bi < acts.len() {
4571 let act = &acts[bi];
4572 let mut dot = 0f32;
4573 for g in 0..ng {
4574 let d = dot_i8_i8(
4575 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4576 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4577 ) as f32
4578 * act.sx;
4579 dot += d * gscale(r, g);
4580 }
4581 for &(j, xv) in &act.outliers {
4582 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
4583 }
4584 unsafe { *out_addr.at(bi * rows + r) = dot };
4586 bi += 1;
4587 }
4588 });
4589 }
4590 };
4591 dispatch_rows(pool, rows, &run);
4592 return;
4593 }
4594
4595 let out_addr = SendMut(out.as_mut_ptr());
4596 let run = move |start: usize, end: usize| {
4597 ROW_F32.with(|rb| {
4598 let mut buf = rb.borrow_mut();
4599 buf.resize(cols, 0.0);
4600 for r in start..end {
4601 decode_f32(r, &mut buf);
4602 for bi in 0..b {
4603 let x = &xs_all[bi * cols..(bi + 1) * cols];
4604 let mut dot = 0f32;
4605 for g in 0..ng {
4606 let mut gd = 0f32;
4607 for k in 0..GROUP_SIZE {
4608 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
4609 }
4610 dot += gd * gscale(r, g);
4611 }
4612 unsafe { *out_addr.at(bi * rows + r) = dot };
4614 }
4615 }
4616 })
4617 };
4618 dispatch_rows(pool, rows, &run);
4619}
4620
4621pub(crate) fn gpu_batch_job<'a>(
4625 t: &'a QTensor,
4626 x: &[f32],
4627) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
4628 match t {
4629 QTensor::Mapped {
4630 model,
4631 idx,
4632 dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
4633 rows,
4634 cols,
4635 row_scale,
4636 col_field,
4637 ..
4638 } => Some((
4639 model.clone(),
4640 crate::gpu::BatchJob {
4641 idx: *idx,
4642 rows: *rows,
4643 cols: *cols,
4644 row_scale,
4645 xs: prescale(x, col_field, *dt).into_owned(),
4646 q1: false,
4647 },
4648 )),
4649 QTensor::Mapped {
4651 model,
4652 idx,
4653 dtype: TensorDtype::Q1,
4654 rows,
4655 cols,
4656 ..
4657 } => Some((
4658 model.clone(),
4659 crate::gpu::BatchJob {
4660 idx: *idx,
4661 rows: *rows,
4662 cols: *cols,
4663 row_scale: &[],
4664 xs: x.to_vec(),
4665 q1: true,
4666 },
4667 )),
4668 _ => None,
4669 }
4670}
4671
4672thread_local! {
4673 static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = std::cell::RefCell::new(Vec::new());
4674 static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = std::cell::RefCell::new(Vec::new());
4675}
4676
4677pub(crate) fn prescale<'a>(
4678 x: &'a [f32],
4679 col_field: &[f32],
4680 dtype: TensorDtype,
4681) -> std::borrow::Cow<'a, [f32]> {
4682 if dtype == TensorDtype::Q8_2f {
4683 x.iter().zip(col_field).map(|(a, c)| a * c).collect()
4684 } else {
4685 std::borrow::Cow::Borrowed(x)
4686 }
4687}
4688
4689pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
4692 x: &[f32],
4693 col_field: &[f32],
4694 dtype: TensorDtype,
4695 buf_id: u8,
4696 f: F,
4697) -> R {
4698 if dtype == TensorDtype::Q8_2f {
4699 if buf_id == 1 {
4700 PRESCALE_BUF1.with(|b| {
4701 let mut buf = b.borrow_mut();
4702 buf.clear();
4703 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
4704 f(&buf)
4705 })
4706 } else {
4707 PRESCALE_BUF2.with(|b| {
4708 let mut buf = b.borrow_mut();
4709 buf.clear();
4710 buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
4711 f(&buf)
4712 })
4713 }
4714 } else {
4715 f(x)
4716 }
4717}
4718
4719#[cfg(target_arch = "x86_64")]
4724pub(crate) fn avx2_enabled() -> bool {
4725 use std::sync::OnceLock;
4726 static ON: OnceLock<bool> = OnceLock::new();
4727 *ON.get_or_init(|| {
4728 std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
4729 && std::arch::is_x86_feature_detected!("avx2")
4730 && std::arch::is_x86_feature_detected!("fma")
4731 })
4732}
4733
4734#[cfg(target_arch = "x86_64")]
4739fn avx2_a8w8_enabled() -> bool {
4740 use std::sync::OnceLock;
4741 static ON: OnceLock<bool> = OnceLock::new();
4742 *ON.get_or_init(|| {
4743 avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
4744 })
4745}
4746
4747#[inline]
4751pub(crate) fn a8w8_enabled() -> bool {
4752 #[cfg(target_arch = "aarch64")]
4753 {
4754 sdot_enabled()
4755 }
4756 #[cfg(target_arch = "x86_64")]
4757 {
4758 avx2_a8w8_enabled()
4759 }
4760 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
4761 {
4762 false
4763 }
4764}
4765
4766#[inline]
4769#[allow(unreachable_code)]
4770fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
4771 #[cfg(target_arch = "aarch64")]
4772 unsafe {
4773 return dot_i8_sdot(w, xq);
4774 }
4775 #[cfg(target_arch = "x86_64")]
4776 unsafe {
4777 if avx512vnni_enabled() {
4778 return dot_i8_i8_vnni(w, xq);
4779 }
4780 return dot_i8_i8_avx2(w, xq);
4781 }
4782 w.iter().zip(xq).map(|(&a, &b)| (a as i8) as i32 * b as i32).sum()
4783}
4784
4785#[cfg(target_arch = "x86_64")]
4789fn avx512vnni_enabled() -> bool {
4790 use std::sync::OnceLock;
4791 static ON: OnceLock<bool> = OnceLock::new();
4792 *ON.get_or_init(|| {
4793 std::env::var("CMF_AVX512").map(|v| v != "0").unwrap_or(true)
4794 && std::arch::is_x86_feature_detected!("avx512f")
4795 && std::arch::is_x86_feature_detected!("avx512bw")
4796 && std::arch::is_x86_feature_detected!("avx512vl")
4797 && std::arch::is_x86_feature_detected!("avx512vnni")
4798 })
4799}
4800
4801#[cfg(target_arch = "x86_64")]
4806#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4807unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
4808 unsafe {
4810 use core::arch::x86_64::*;
4811 let n = w.len();
4812 let mut j = 0usize;
4813 let mut total: i32;
4814 {
4819 #[inline(always)]
4820 unsafe fn step(
4821 w: *const u8,
4822 x: *const i8,
4823 acc: core::arch::x86_64::__m512i,
4824 ) -> core::arch::x86_64::__m512i {
4825 unsafe {
4826 use core::arch::x86_64::*;
4827 let wv = _mm512_loadu_si512(w as *const _);
4828 let xv = _mm512_loadu_si512(x as *const _);
4829 let aw = _mm512_abs_epi8(wv);
4830 let neg = _mm512_movepi8_mask(wv);
4831 let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
4832 _mm512_dpbusd_epi32(acc, aw, sx)
4833 }
4834 }
4835 let (mut a0, mut a1, mut a2, mut a3) = (
4836 _mm512_setzero_si512(),
4837 _mm512_setzero_si512(),
4838 _mm512_setzero_si512(),
4839 _mm512_setzero_si512(),
4840 );
4841 while j + 256 <= n {
4842 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
4843 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
4844 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
4845 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
4846 j += 256;
4847 }
4848 while j + 64 <= n {
4849 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
4850 j += 64;
4851 }
4852 let s01 = _mm512_add_epi32(a0, a1);
4853 let s23 = _mm512_add_epi32(a2, a3);
4854 total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
4855 }
4856 if j + 32 <= n {
4858 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
4859 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
4860 let d = _mm256_dpbusd_epi32(
4861 _mm256_setzero_si256(),
4862 _mm256_abs_epi8(wv),
4863 _mm256_sign_epi8(xv, wv),
4864 );
4865 let hi128 = _mm256_extracti128_si256::<1>(d);
4866 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4867 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4868 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4869 total += _mm_cvtsi128_si32(s32);
4870 j += 32;
4871 }
4872 while j < n {
4873 total += (w[j] as i8) as i32 * xq[j] as i32;
4874 j += 1;
4875 }
4876 total
4877 }
4878}
4879
4880#[cfg(target_arch = "x86_64")]
4882#[target_feature(enable = "avx2,fma")]
4883unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
4884 unsafe {
4886 use core::arch::x86_64::*;
4887 let n = x.len();
4888 let wp = w.as_ptr();
4889 let xp = x.as_ptr();
4890 let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
4891 let mut j = 0usize;
4892 while j + 16 <= n {
4893 let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
4894 let lo = _mm256_cvtepi8_epi32(wb);
4895 let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
4896 a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
4897 a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
4898 j += 16;
4899 }
4900 let acc = _mm256_add_ps(a0, a1);
4901 let hi128 = _mm256_extractf128_ps::<1>(acc);
4902 let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
4903 let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
4904 let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
4905 let mut sum = _mm_cvtss_f32(s32);
4906 while j < n {
4907 sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
4908 j += 1;
4909 }
4910 sum
4911 }
4912}
4913
4914#[cfg(target_arch = "x86_64")]
4919#[target_feature(enable = "avx2")]
4920unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
4921 unsafe {
4923 use core::arch::x86_64::*;
4924 let n = w.len();
4925 let ones = _mm256_set1_epi16(1);
4926 let mut acc = _mm256_setzero_si256();
4927 let mut j = 0usize;
4928 while j + 32 <= n {
4929 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
4930 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
4931 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
4932 acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
4933 j += 32;
4934 }
4935 let hi128 = _mm256_extracti128_si256::<1>(acc);
4936 let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
4937 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4938 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4939 let mut s = _mm_cvtsi128_si32(s32);
4940 while j < n {
4941 s += (w[j] as i8) as i32 * xq[j] as i32;
4942 j += 1;
4943 }
4944 s
4945 }
4946}
4947
4948#[cfg(target_arch = "aarch64")]
4952#[target_feature(enable = "neon,i8mm")]
4953unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
4954 unsafe {
4956 use core::arch::aarch64::*;
4957 use core::arch::asm;
4958 let n = w0.len();
4959 let w0p = w0.as_ptr() as *const i8;
4960 let w1p = w1.as_ptr() as *const i8;
4961 let mut acc01 = vdupq_n_s32(0);
4964 let mut acc23 = vdupq_n_s32(0);
4965 let mut i = 0usize;
4966 while i + 8 <= n {
4967 let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
4968 let xb01 =
4969 vcombine_s8(vld1_s8(xs[0].as_ptr().add(i)), vld1_s8(xs[1].as_ptr().add(i)));
4970 let xb23 =
4971 vcombine_s8(vld1_s8(xs[2].as_ptr().add(i)), vld1_s8(xs[3].as_ptr().add(i)));
4972 asm!(
4973 "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
4974 "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
4975 a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
4976 w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
4977 options(pure, nomem, nostack),
4978 );
4979 i += 8;
4980 }
4981 let mut out = [[0i32; 4]; 2];
4982 let a01: [i32; 4] = core::mem::transmute(acc01);
4983 let a23: [i32; 4] = core::mem::transmute(acc23);
4984 out[0][0] = a01[0];
4985 out[0][1] = a01[1];
4986 out[1][0] = a01[2];
4987 out[1][1] = a01[3];
4988 out[0][2] = a23[0];
4989 out[0][3] = a23[1];
4990 out[1][2] = a23[2];
4991 out[1][3] = a23[3];
4992 if i < n {
4993 for (k, x) in xs.iter().enumerate() {
4994 for j in i..n {
4995 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
4996 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
4997 }
4998 }
4999 }
5000 out
5001 }
5002}
5003
5004#[cfg(target_arch = "aarch64")]
5008#[target_feature(enable = "neon,dotprod")]
5009unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
5010 unsafe {
5012 use core::arch::aarch64::*;
5013 use core::arch::asm;
5014 let n = w0.len();
5015 let w0p = w0.as_ptr() as *const i8;
5016 let w1p = w1.as_ptr() as *const i8;
5017 let mut acc = [[vdupq_n_s32(0); 4]; 2];
5018 let mut i = 0usize;
5019 while i + 16 <= n {
5020 let wv0 = vld1q_s8(w0p.add(i));
5021 let wv1 = vld1q_s8(w1p.add(i));
5022 for (k, x) in xs.iter().enumerate() {
5023 let xv = vld1q_s8(x.as_ptr().add(i));
5024 let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
5025 asm!(
5026 "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
5027 "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
5028 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5029 w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
5030 options(pure, nomem, nostack),
5031 );
5032 acc[0][k] = a0;
5033 acc[1][k] = a1;
5034 }
5035 i += 16;
5036 }
5037 let mut out = [[0i32; 4]; 2];
5038 for r in 0..2 {
5039 for k in 0..4 {
5040 out[r][k] = vaddvq_s32(acc[r][k]);
5041 }
5042 }
5043 if i < n {
5044 for (k, x) in xs.iter().enumerate() {
5045 for j in i..n {
5046 out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
5047 out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
5048 }
5049 }
5050 }
5051 out
5052 }
5053}
5054
5055#[cfg(target_arch = "x86_64")]
5061#[target_feature(enable = "avx2")]
5062unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
5063 unsafe {
5065 use core::arch::x86_64::*;
5066 let n = w0.len();
5067 let ones = _mm256_set1_epi16(1);
5068 let mut acc = [[_mm256_setzero_si256(); 4]; 2];
5069 let mut j = 0usize;
5070 while j + 32 <= n {
5071 let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
5072 let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
5073 let aw0 = _mm256_abs_epi8(wv0);
5074 let aw1 = _mm256_abs_epi8(wv1);
5075 for (k, x) in xs.iter().enumerate() {
5076 let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
5077 let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
5078 acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
5079 let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
5080 acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
5081 }
5082 j += 32;
5083 }
5084 let mut out = [[0i32; 4]; 2];
5085 for r in 0..2 {
5086 for k in 0..4 {
5087 let a = acc[r][k];
5088 let hi128 = _mm256_extracti128_si256::<1>(a);
5089 let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
5090 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5091 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5092 out[r][k] = _mm_cvtsi128_si32(s32);
5093 }
5094 }
5095 if j < n {
5096 for (k, x) in xs.iter().enumerate() {
5097 for i in j..n {
5098 out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
5099 out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
5100 }
5101 }
5102 }
5103 out
5104 }
5105}
5106
5107#[cfg(target_arch = "x86_64")]
5112#[inline]
5113fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
5114 let dot = if avx512vnni_enabled() && row.len() >= 64 {
5115 (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
5116 } else {
5117 unsafe { dot_i8_i8_avx2(row, &act.xq) }
5118 };
5119 let mut acc = dot as f32 * act.sx;
5120 for &(j, xv) in &act.outliers {
5121 acc += (row[j] as i8) as f32 * xv;
5122 }
5123 acc
5124}
5125
5126#[cfg(target_arch = "x86_64")]
5130#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5131unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
5132 unsafe {
5134 use core::arch::x86_64::*;
5135 let n = w.len();
5136 let flip = _mm512_set1_epi8(-128); #[inline(always)]
5138 unsafe fn step(
5139 w: *const u8,
5140 x: *const i8,
5141 flip: core::arch::x86_64::__m512i,
5142 acc: core::arch::x86_64::__m512i,
5143 ) -> core::arch::x86_64::__m512i {
5144 unsafe {
5145 use core::arch::x86_64::*;
5146 let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
5147 _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
5148 }
5149 }
5150 let (mut a0, mut a1, mut a2, mut a3) = (
5151 _mm512_setzero_si512(),
5152 _mm512_setzero_si512(),
5153 _mm512_setzero_si512(),
5154 _mm512_setzero_si512(),
5155 );
5156 let mut j = 0usize;
5157 while j + 256 <= n {
5158 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
5159 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
5160 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
5161 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
5162 j += 256;
5163 }
5164 while j + 64 <= n {
5165 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
5166 j += 64;
5167 }
5168 let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
5169 _mm512_add_epi32(a0, a1),
5170 _mm512_add_epi32(a2, a3),
5171 ));
5172 while j < n {
5174 total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
5175 j += 1;
5176 }
5177 total
5178 }
5179}
5180
5181#[cfg(target_arch = "x86_64")]
5187#[target_feature(enable = "avx2")]
5188unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
5189 unsafe {
5192 use core::arch::x86_64::*;
5193 let lomask = _mm_set1_epi8(0x0F);
5194 let eight = _mm256_set1_epi8(8);
5195 let ones = _mm256_set1_epi16(1);
5196 let mut acc = 0f32;
5197 for gi in 0..gpr {
5198 let g = g0 + gi;
5199 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
5200 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
5201 let lo = _mm_and_si128(b, lomask);
5202 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5203 let w = _mm256_sub_epi8(
5204 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5205 eight,
5206 );
5207 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5208 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5209 let d = _mm256_madd_epi16(p16, ones);
5210 let hi128 = _mm256_extracti128_si256::<1>(d);
5211 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5212 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5213 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5214 acc += _mm_cvtsi128_si32(s32) as f32 * s;
5215 }
5216 acc
5217 }
5218}
5219
5220#[cfg(target_arch = "x86_64")]
5223#[target_feature(enable = "avx2")]
5224unsafe fn dot_q4_row_avx2_2(
5225 packed: &[u8],
5226 scales: &[u8],
5227 g0: usize,
5228 gpr: usize,
5229 xq1: &[i8],
5230 xq2: &[i8],
5231) -> (f32, f32) {
5232 unsafe {
5234 use core::arch::x86_64::*;
5235 let lomask = _mm_set1_epi8(0x0F);
5236 let eight = _mm256_set1_epi8(8);
5237 let ones = _mm256_set1_epi16(1);
5238 let (mut acc1, mut acc2) = (0f32, 0f32);
5239 #[inline(always)]
5240 unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
5241 unsafe {
5242 use core::arch::x86_64::*;
5243 let hi128 = _mm256_extracti128_si256::<1>(d);
5244 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5245 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5246 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5247 _mm_cvtsi128_si32(s32)
5248 }
5249 }
5250 for gi in 0..gpr {
5251 let g = g0 + gi;
5252 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
5253 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
5254 let lo = _mm_and_si128(b, lomask);
5255 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5256 let w = _mm256_sub_epi8(
5257 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5258 eight,
5259 );
5260 let aw = _mm256_abs_epi8(w);
5261 let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5262 let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5263 let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
5264 let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
5265 acc1 += hsum(d1) as f32 * s;
5266 acc2 += hsum(d2) as f32 * s;
5267 }
5268 (acc1, acc2)
5269 }
5270}
5271
5272#[cfg(target_arch = "x86_64")]
5274fn q8_range_avx2(
5275 q: &[u8],
5276 row_scale: &[f32],
5277 act: &SplitAct,
5278 cols: usize,
5279 out_addr: SendMut,
5280 start: usize,
5281 end: usize,
5282) {
5283 for o in start..end {
5284 let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
5285 unsafe { *out_addr.at(o) = v };
5287 }
5288}
5289
5290#[cfg(target_arch = "x86_64")]
5292#[allow(clippy::too_many_arguments)]
5293fn q8_range2_avx2(
5294 q: &[u8],
5295 row_scale: &[f32],
5296 a1: &SplitAct,
5297 a2: &SplitAct,
5298 cols: usize,
5299 p1: SendMut,
5300 p2: SendMut,
5301 start: usize,
5302 end: usize,
5303) {
5304 for o in start..end {
5305 let row = &q[o * cols..(o + 1) * cols];
5306 unsafe {
5308 *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
5309 *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
5310 }
5311 }
5312}
5313
5314#[cfg(target_arch = "aarch64")]
5325fn i8mm_enabled() -> bool {
5326 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5327 *ON.get_or_init(|| {
5328 std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
5329 && std::arch::is_aarch64_feature_detected!("i8mm")
5330 })
5331}
5332
5333#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
5337fn sdot_enabled() -> bool {
5338 use std::sync::OnceLock;
5339 static ON: OnceLock<bool> = OnceLock::new();
5340 *ON.get_or_init(|| {
5341 let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
5342 if !want { return false; }
5343
5344 #[cfg(target_arch = "aarch64")]
5345 {
5346 if std::arch::is_aarch64_feature_detected!("dotprod") {
5347 return true;
5348 }
5349 #[cfg(target_os = "android")]
5350 {
5351 if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
5352 if cpuinfo.lines().any(|l| (l.starts_with("Features") || l.starts_with("features")) && l.contains("asimddp")) {
5353 return true;
5354 }
5355 }
5356 }
5357 false
5358 }
5359 #[cfg(not(target_arch = "aarch64"))]
5360 {
5361 false
5362 }
5363 })
5364}
5365
5366struct SplitAct {
5371 xq: Vec<i8>,
5372 sx: f32,
5373 outliers: Vec<(usize, f32)>,
5374 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
5377 xsum: i32,
5378}
5379
5380thread_local! {
5381 static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
5384 const { std::cell::RefCell::new(Vec::new()) };
5385}
5386
5387impl Drop for SplitAct {
5388 fn drop(&mut self) {
5389 let buf = std::mem::take(&mut self.xq);
5390 if buf.capacity() > 0 {
5391 XQ_FREE.with(|f| {
5392 let mut f = f.borrow_mut();
5393 if f.len() < 16 {
5394 f.push(buf);
5395 }
5396 });
5397 }
5398 }
5399}
5400
5401fn split_act(x: &[f32]) -> SplitAct {
5402 let n = x.len();
5403 let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
5404 let thr = 8.0 * rms;
5405 let mut outliers: Vec<(usize, f32)> = Vec::new();
5409 let mut amax = 0f32;
5410 for (j, &v) in x.iter().enumerate() {
5411 let a = v.abs();
5412 if a > thr {
5413 outliers.push((j, v));
5414 } else if a > amax {
5415 amax = a;
5416 }
5417 }
5418 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
5419 let inv = 1.0 / sx;
5420 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
5421 xq.clear();
5422 xq.reserve(n);
5423 if outliers.is_empty() {
5424 xq.extend(x.iter().map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8));
5425 } else {
5426 xq.extend(x.iter().map(|&v| {
5428 if v.abs() > thr {
5429 0
5430 } else {
5431 (v * inv).round().clamp(-127.0, 127.0) as i8
5432 }
5433 }));
5434 }
5435 let xsum = xq.iter().map(|&v| v as i32).sum();
5436 SplitAct { xq, sx, outliers, xsum }
5437}
5438
5439fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
5440 let n = x.len();
5441 let rms = (x.iter().zip(col).map(|(&a, &c)| { let v = a * c; (v * v) as f64 })
5442 .sum::<f64>() / n.max(1) as f64).sqrt() as f32;
5443 let thr = 8.0 * rms;
5444
5445 let mut outliers = Vec::new();
5446 let mut amax = 0f32;
5447 for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
5448 let v = a * c;
5449 let s = v.abs();
5450 if s > thr {
5451 outliers.push((j, v));
5452 } else if s > amax {
5453 amax = s;
5454 }
5455 }
5456
5457 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
5458 let inv = 1.0 / sx;
5459 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
5460 xq.clear();
5461 xq.reserve(n);
5462 if outliers.is_empty() {
5463 xq.extend(x.iter().zip(col).map(|(&a, &c)| ((a * c) * inv).round().clamp(-127.0, 127.0) as i8));
5464 } else {
5465 xq.extend(x.iter().zip(col).map(|(&a, &c)| {
5466 let v = a * c;
5467 if v.abs() > thr { 0 } else { (v * inv).round().clamp(-127.0, 127.0) as i8 }
5468 }));
5469 }
5470 let xsum = xq.iter().map(|&v| v as i32).sum();
5471 SplitAct { xq, sx, outliers, xsum }
5472}
5473
5474#[cfg(target_arch = "aarch64")]
5477#[target_feature(enable = "neon,dotprod")]
5478unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
5479 unsafe {
5481 use core::arch::aarch64::*;
5482 use core::arch::asm;
5483 let wp = w.as_ptr() as *const i8;
5484 let n = w.len();
5485 let (mut a0, mut a1, mut a2, mut a3) =
5486 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
5487 let mut i = 0;
5488 while i + 64 <= n {
5489 let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
5490 let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
5491 let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
5492 let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
5493 asm!(
5494 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
5495 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
5496 "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
5497 "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
5498 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
5499 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
5500 w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
5501 options(pure, nomem, nostack),
5502 );
5503 i += 64;
5504 }
5505 while i + 16 <= n {
5506 let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
5507 asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
5508 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
5509 i += 16;
5510 }
5511 let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
5512 while i < n {
5513 s += (*wp.add(i)) as i32 * xq[i] as i32;
5514 i += 1;
5515 }
5516 s
5517}
5518}
5519
5520#[cfg(target_arch = "aarch64")]
5524#[target_feature(enable = "neon,dotprod")]
5525unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
5526 unsafe {
5528 use core::arch::aarch64::*;
5529 use core::arch::asm;
5530 let n = xq.len();
5531 let px = xq.as_ptr();
5532 let (p0, p1, p2, p3) = (
5533 w0.as_ptr() as *const i8,
5534 w1.as_ptr() as *const i8,
5535 w2.as_ptr() as *const i8,
5536 w3.as_ptr() as *const i8,
5537 );
5538 let (mut a0, mut a1, mut a2, mut a3) =
5539 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
5540 let mut i = 0;
5541 while i + 16 <= n {
5542 let x = vld1q_s8(px.add(i));
5543 let v0 = vld1q_s8(p0.add(i));
5544 let v1 = vld1q_s8(p1.add(i));
5545 let v2 = vld1q_s8(p2.add(i));
5546 let v3 = vld1q_s8(p3.add(i));
5547 asm!(
5548 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
5549 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
5550 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
5551 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
5552 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
5553 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
5554 options(pure, nomem, nostack),
5555 );
5556 i += 16;
5557 }
5558 let mut r = [vaddvq_s32(a0), vaddvq_s32(a1), vaddvq_s32(a2), vaddvq_s32(a3)];
5559 while i < n {
5560 let xi = *px.add(i) as i32;
5561 r[0] += (*p0.add(i)) as i32 * xi;
5562 r[1] += (*p1.add(i)) as i32 * xi;
5563 r[2] += (*p2.add(i)) as i32 * xi;
5564 r[3] += (*p3.add(i)) as i32 * xi;
5565 i += 1;
5566 }
5567 r
5568}
5569}
5570
5571#[cfg(target_arch = "aarch64")]
5578#[target_feature(enable = "neon,dotprod")]
5579unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
5580 unsafe {
5583 use core::arch::aarch64::*;
5584 use core::arch::asm;
5585 let n = xq.len();
5586 let px = xq.as_ptr();
5587 let pg = g.as_ptr() as *const i8;
5588 let (mut a0, mut a1, mut a2, mut a3) =
5589 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
5590 let mut i = 0;
5591 while i + 16 <= n {
5592 let x = vld1q_s8(px.add(i));
5593 let base = pg.add(4 * i);
5594 let v0 = vld1q_s8(base);
5595 let v1 = vld1q_s8(base.add(16));
5596 let v2 = vld1q_s8(base.add(32));
5597 let v3 = vld1q_s8(base.add(48));
5598 asm!(
5599 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
5600 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
5601 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
5602 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
5603 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
5604 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
5605 options(pure, nomem, nostack),
5606 );
5607 i += 16;
5608 }
5609 [vaddvq_s32(a0), vaddvq_s32(a1), vaddvq_s32(a2), vaddvq_s32(a3)]
5610 }
5611}
5612
5613#[cfg(target_arch = "aarch64")]
5619fn q8_range_sdot(
5620 q: &[u8],
5621 rep: &[u8],
5622 row_scale: &[f32],
5623 act: &SplitAct,
5624 cols: usize,
5625 out_addr: SendMut,
5626 start: usize,
5627 end: usize,
5628) {
5629 let mut o = start;
5630 if !rep.is_empty() {
5633 while o < end && o % 4 != 0 {
5634 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
5635 unsafe { *out_addr.at(o) = v };
5636 o += 1;
5637 }
5638 }
5639 while o + 4 <= end {
5640 let r = if rep.is_empty() {
5641 unsafe {
5642 dot_i8_sdot_4rows(
5643 &q[o * cols..(o + 1) * cols],
5644 &q[(o + 1) * cols..(o + 2) * cols],
5645 &q[(o + 2) * cols..(o + 3) * cols],
5646 &q[(o + 3) * cols..(o + 4) * cols],
5647 &act.xq,
5648 )
5649 }
5650 } else {
5651 unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
5652 };
5653 for k in 0..4 {
5654 let mut acc = r[k] as f32 * act.sx;
5655 for &(j, xv) in &act.outliers {
5656 acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
5657 }
5658 unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
5660 }
5661 o += 4;
5662 }
5663 while o < end {
5664 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
5665 unsafe { *out_addr.at(o) = v };
5666 o += 1;
5667 }
5668}
5669
5670#[cfg(target_arch = "aarch64")]
5673#[allow(clippy::too_many_arguments)]
5674fn q8_range2_sdot(
5675 q: &[u8],
5676 row_scale: &[f32],
5677 a1: &SplitAct,
5678 a2: &SplitAct,
5679 cols: usize,
5680 p1: SendMut,
5681 p2: SendMut,
5682 start: usize,
5683 end: usize,
5684) {
5685 for o in start..end {
5686 let row = &q[o * cols..(o + 1) * cols];
5687 unsafe {
5689 *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
5690 *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
5691 }
5692 }
5693}
5694
5695#[allow(clippy::too_many_arguments)]
5697fn q8_range2_f32(
5698 q: &[u8],
5699 row_scale: &[f32],
5700 x1: &[f32],
5701 x2: &[f32],
5702 cols: usize,
5703 p1: SendMut,
5704 p2: SendMut,
5705 start: usize,
5706 end: usize,
5707) {
5708 for o in start..end {
5709 let row = &q[o * cols..(o + 1) * cols];
5710 unsafe {
5712 *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
5713 *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
5714 }
5715 }
5716}
5717
5718fn q8_range_f32(
5720 q: &[u8],
5721 row_scale: &[f32],
5722 xs: &[f32],
5723 cols: usize,
5724 out_addr: SendMut,
5725 start: usize,
5726 end: usize,
5727) {
5728 for o in start..end {
5729 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
5730 unsafe { *out_addr.at(o) = v };
5732 }
5733}
5734
5735#[cfg(target_arch = "aarch64")]
5738#[inline]
5739fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
5740 let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
5741 for &(j, xv) in &act.outliers {
5742 acc += (row[j] as i8) as f32 * xv;
5743 }
5744 acc
5745}
5746
5747#[cfg(target_arch = "aarch64")]
5755#[target_feature(enable = "neon,dotprod")]
5756unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
5757 unsafe {
5760 use core::arch::aarch64::*;
5761 use core::arch::asm;
5762 let lomask = vdupq_n_u8(0x0F);
5763 let eight = vdupq_n_s8(8);
5764 let mut acc = 0f32;
5765 for gi in 0..gpr {
5766 let g = g0 + gi;
5767 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
5768 let b = vld1q_u8(packed.as_ptr().add(g * 16));
5769 let lo = vandq_u8(b, lomask);
5770 let hi = vshrq_n_u8::<4>(b);
5771 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5772 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5773 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5774 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5775 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5776 asm!(
5777 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5778 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5779 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5780 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5781 options(pure, nomem, nostack),
5782 );
5783 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5784 }
5785 acc
5786 }
5787}
5788
5789#[cfg(target_arch = "aarch64")]
5794#[target_feature(enable = "neon,dotprod")]
5795unsafe fn dot_q4_row_sdot2(
5796 packed: &[u8],
5797 scales: &[u8],
5798 g0: usize,
5799 gpr: usize,
5800 xq1: &[i8],
5801 xq2: &[i8],
5802) -> (f32, f32) {
5803 unsafe {
5806 use core::arch::aarch64::*;
5807 use core::arch::asm;
5808 let lomask = vdupq_n_u8(0x0F);
5809 let eight = vdupq_n_s8(8);
5810 let (mut acc1, mut acc2) = (0f32, 0f32);
5811 for gi in 0..gpr {
5812 let g = g0 + gi;
5813 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
5814 let b = vld1q_u8(packed.as_ptr().add(g * 16));
5815 let lo = vandq_u8(b, lomask);
5816 let hi = vshrq_n_u8::<4>(b);
5817 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5818 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5819 let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
5820 let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
5821 let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
5822 let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
5823 let (mut a0, mut a1, mut b0, mut b1) =
5824 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
5825 asm!(
5826 "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
5827 "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
5828 "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
5829 "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
5830 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5831 b0 = inout(vreg) b0, b1 = inout(vreg) b1,
5832 e0 = in(vreg) e0, e1 = in(vreg) e1,
5833 x10 = in(vreg) x10, x11 = in(vreg) x11,
5834 x20 = in(vreg) x20, x21 = in(vreg) x21,
5835 options(pure, nomem, nostack),
5836 );
5837 acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5838 acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
5839 }
5840 (acc1, acc2)
5841 }
5842}
5843
5844#[inline]
5849pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
5850 #[cfg(target_arch = "aarch64")]
5851 unsafe {
5852 return axpy_i8_f32_neon(acc, row, w);
5853 }
5854 #[cfg(target_arch = "x86_64")]
5855 if avx2_enabled() {
5856 return unsafe { axpy_i8_f32_avx2(acc, row, w) };
5857 }
5858 #[allow(unreachable_code)]
5859 {
5860 for (a, &b) in acc.iter_mut().zip(row) {
5861 *a += w * b as f32;
5862 }
5863 }
5864}
5865
5866#[cfg(target_arch = "x86_64")]
5868#[target_feature(enable = "avx2,fma")]
5869unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
5870 unsafe {
5872 use core::arch::x86_64::*;
5873 let n = acc.len().min(row.len());
5874 let ap = acc.as_mut_ptr();
5875 let rp = row.as_ptr();
5876 let wv = _mm256_set1_ps(w);
5877 let mut j = 0usize;
5878 while j + 16 <= n {
5879 let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
5880 let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
5881 let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
5882 let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
5883 let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
5884 _mm256_storeu_ps(ap.add(j), v0);
5885 _mm256_storeu_ps(ap.add(j + 8), v1);
5886 j += 16;
5887 }
5888 while j < n {
5889 *ap.add(j) += w * (*rp.add(j)) as f32;
5890 j += 1;
5891 }
5892 }
5893}
5894
5895#[cfg(target_arch = "aarch64")]
5896#[target_feature(enable = "neon")]
5897unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
5898 unsafe {
5900 use core::arch::aarch64::*;
5901 let n = acc.len().min(row.len());
5902 let ap = acc.as_mut_ptr();
5903 let rp = row.as_ptr();
5904 let wv = vdupq_n_f32(w);
5905 let mut j = 0usize;
5906 while j + 16 <= n {
5907 let rb = vld1q_s8(rp.add(j));
5908 let lo = vmovl_s8(vget_low_s8(rb));
5909 let hi = vmovl_s8(vget_high_s8(rb));
5910 for (off, half) in [(0, lo), (8, hi)] {
5911 let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
5912 let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
5913 let o = j + off;
5914 vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
5915 vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
5916 }
5917 j += 16;
5918 }
5919 while j < n {
5920 *ap.add(j) += w * (*rp.add(j)) as f32;
5921 j += 1;
5922 }
5923}
5924}
5925
5926#[inline]
5929pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
5930 #[cfg(target_arch = "aarch64")]
5931 unsafe {
5932 return dot_i8_f32_neon(w, x);
5933 }
5934 #[cfg(target_arch = "x86_64")]
5935 if avx2_enabled() {
5936 return unsafe { dot_i8_f32_avx2(w, x) };
5937 }
5938 #[allow(unreachable_code)]
5939 {
5940 let mut sum = 0.0f32;
5941 for (j, &b) in w.iter().enumerate() {
5942 sum += (b as i8) as f32 * x[j];
5943 }
5944 sum
5945 }
5946}
5947
5948#[inline]
5952fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
5953 #[cfg(target_arch = "aarch64")]
5954 unsafe {
5955 return dot_i8_col_f32_neon(w, x, col);
5956 }
5957 #[allow(unreachable_code)]
5958 {
5959 let mut sum = 0.0f32;
5960 for (j, &b) in w.iter().enumerate() {
5961 sum += (b as i8) as f32 * x[j] * col[j];
5962 }
5963 sum
5964 }
5965}
5966
5967#[cfg(target_arch = "aarch64")]
5968#[target_feature(enable = "neon")]
5969unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
5970 unsafe {
5972 use core::arch::aarch64::*;
5973 let n = x.len();
5974 let wp = w.as_ptr() as *const i8;
5975 let xp = x.as_ptr();
5976 let cp = col.as_ptr();
5977 let (mut a0, mut a1, mut a2, mut a3) =
5978 (vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0));
5979 let mut j = 0usize;
5980 while j + 16 <= n {
5981 let wb = vld1q_s8(wp.add(j));
5982 let lo = vmovl_s8(vget_low_s8(wb));
5983 let hi = vmovl_s8(vget_high_s8(wb));
5984 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
5985 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
5986 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
5987 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
5988 a0 = vfmaq_f32(a0, w0, vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))));
5989 a1 = vfmaq_f32(a1, w1, vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))));
5990 a2 = vfmaq_f32(a2, w2, vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))));
5991 a3 = vfmaq_f32(a3, w3, vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))));
5992 j += 16;
5993 }
5994 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
5995 while j < n {
5996 sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
5997 j += 1;
5998 }
5999 sum
6000}
6001}
6002
6003#[cfg(target_arch = "aarch64")]
6004#[target_feature(enable = "neon")]
6005unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
6006 unsafe {
6008 use core::arch::aarch64::*;
6009 let n = x.len();
6010 let wp = w.as_ptr() as *const i8;
6011 let xp = x.as_ptr();
6012 let (mut a0, mut a1, mut a2, mut a3) =
6013 (vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0));
6014 let mut j = 0usize;
6015 while j + 16 <= n {
6016 let wb = vld1q_s8(wp.add(j));
6017 let lo = vmovl_s8(vget_low_s8(wb));
6018 let hi = vmovl_s8(vget_high_s8(wb));
6019 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
6020 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
6021 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
6022 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
6023 a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
6024 a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
6025 a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
6026 a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
6027 j += 16;
6028 }
6029 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
6030 while j < n {
6031 sum += (*wp.add(j)) as f32 * *xp.add(j);
6032 j += 1;
6033 }
6034 sum
6035}
6036}
6037
6038#[allow(clippy::too_many_arguments)]
6039fn qmatvec(
6040 q: &[u8],
6041 rep: &[u8],
6042 row_scale: &[f32],
6043 x: &[f32],
6044 col_field: &[f32],
6045 dtype: TensorDtype,
6046 rows: usize,
6047 cols: usize,
6048 out: &mut [f32],
6049 pool: Option<&Pool>,
6050) {
6051 debug_assert_eq!(out.len(), rows);
6052 #[cfg(not(target_arch = "aarch64"))]
6053 let _ = rep;
6054
6055 #[cfg(target_arch = "aarch64")]
6056 if sdot_enabled() {
6057 let act = if dtype == TensorDtype::Q8_2f {
6058 split_act_q8_2f(x, col_field)
6059 } else {
6060 split_act(x)
6061 };
6062 let out_addr = SendMut(out.as_mut_ptr());
6063 let run_range = |start: usize, end: usize| {
6064 q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
6065 };
6066 match pool {
6067 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
6068 _ => run_range(0, rows),
6069 }
6070 return;
6071 }
6072 #[cfg(target_arch = "x86_64")]
6075 if avx2_a8w8_enabled() {
6076 let act = if dtype == TensorDtype::Q8_2f {
6077 split_act_q8_2f(x, col_field)
6078 } else {
6079 split_act(x)
6080 };
6081 let out_addr = SendMut(out.as_mut_ptr());
6082 let run_range =
6083 |start: usize, end: usize| q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end);
6084 match pool {
6085 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
6086 _ => run_range(0, rows),
6087 }
6088 return;
6089 }
6090
6091 prescale_with(x, col_field, dtype, 1, |xs| {
6092 let out_addr = SendMut(out.as_mut_ptr());
6093 let run_range = move |start: usize, end: usize| {
6094 for o in start..end {
6095 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
6096 unsafe { *out_addr.at(o) = v };
6098 }
6099 };
6100 match pool {
6101 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
6102 _ => run_range(0, rows),
6103 }
6104 });
6105}
6106
6107#[allow(clippy::too_many_arguments)]
6108fn qmatvec2(
6109 q: &[u8],
6110 row_scale: &[f32],
6111 x1: &[f32],
6112 x2: &[f32],
6113 col_field: &[f32],
6114 dtype: TensorDtype,
6115 rows: usize,
6116 cols: usize,
6117 o1: &mut [f32],
6118 o2: &mut [f32],
6119 pool: Option<&Pool>,
6120) {
6121 #[cfg(target_arch = "aarch64")]
6122 if sdot_enabled() {
6123 let a1s = if dtype == TensorDtype::Q8_2f { split_act_q8_2f(x1, col_field) } else { split_act(x1) };
6124 let a2s = if dtype == TensorDtype::Q8_2f { split_act_q8_2f(x2, col_field) } else { split_act(x2) };
6125 let p1 = SendMut(o1.as_mut_ptr());
6126 let p2 = SendMut(o2.as_mut_ptr());
6127 let run_range = |start: usize, end: usize| {
6128 q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
6129 };
6130 match pool {
6131 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
6132 _ => run_range(0, rows),
6133 }
6134 return;
6135 }
6136 #[cfg(target_arch = "x86_64")]
6137 if avx2_a8w8_enabled() {
6138 let a1s = if dtype == TensorDtype::Q8_2f { split_act_q8_2f(x1, col_field) } else { split_act(x1) };
6139 let a2s = if dtype == TensorDtype::Q8_2f { split_act_q8_2f(x2, col_field) } else { split_act(x2) };
6140 let p1 = SendMut(o1.as_mut_ptr());
6141 let p2 = SendMut(o2.as_mut_ptr());
6142 let run_range = |start: usize, end: usize| {
6143 q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
6144 };
6145 match pool {
6146 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
6147 _ => run_range(0, rows),
6148 }
6149 return;
6150 }
6151
6152 prescale_with(x1, col_field, dtype, 1, |x1s| {
6153 prescale_with(x2, col_field, dtype, 2, |x2s| {
6154 let p1 = SendMut(o1.as_mut_ptr());
6155 let p2 = SendMut(o2.as_mut_ptr());
6156 let run_range = move |start: usize, end: usize| {
6157 for o in start..end {
6158 let row = &q[o * cols..(o + 1) * cols];
6159 let s1 = dot_i8_f32(row, x1s) * row_scale[o];
6160 let s2 = dot_i8_f32(row, x2s) * row_scale[o];
6161 unsafe {
6163 *p1.at(o) = s1;
6164 *p2.at(o) = s2;
6165 }
6166 }
6167 };
6168 match pool {
6169 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
6170 _ => run_range(0, rows),
6171 }
6172 });
6173 });
6174}
6175
6176#[derive(Clone, Copy)]
6177struct SendMut(*mut f32);
6178unsafe impl Send for SendMut {}
6179unsafe impl Sync for SendMut {}
6180
6181impl SendMut {
6182 #[inline]
6183 fn at(self, i: usize) -> *mut f32 {
6184 unsafe { self.0.add(i) }
6185 }
6186}
6187
6188#[cfg(test)]
6189mod tests {
6190 use super::*;
6191
6192 #[test]
6193 fn f32_matvec_matches_matvec_rows_bitexact() {
6194 let (rows, cols) = (300, 40);
6195 let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
6196 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
6197 let qt = QTensor::from_f32(w.clone(), rows, cols);
6198
6199 let mut a = vec![0.0f32; rows];
6200 matvec_rows(None, &w, &x, &mut a);
6201 let mut b = vec![0.0f32; rows];
6202 qt.matvec(&x, &mut b, None);
6203 assert_eq!(a, b);
6204 }
6205
6206 #[test]
6207 fn sdot_kernel_exact_on_grid() {
6208 eprintln!("sdot_enabled = {}", sdot_enabled());
6213 let (rows, cols) = (9, 80); let w: Vec<u8> = (0..rows * cols)
6215 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
6216 .collect();
6217 let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
6218 let x: Vec<f32> = (0..cols)
6219 .map(|i| match i % 3 {
6220 0 => 1.0,
6221 1 => -1.0,
6222 _ => 0.0,
6223 })
6224 .collect();
6225 let mut a = vec![0.0f32; rows];
6226 qmatvec(&w, &[], &scales, &x, &[], TensorDtype::Q8Row, rows, cols, &mut a, None);
6227 for o in 0..rows {
6228 let mut acc = 0.0f32;
6229 for j in 0..cols {
6230 acc += (w[o * cols + j] as i8) as f32 * x[j];
6231 }
6232 let expect = acc * scales[o];
6233 assert!(
6234 (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
6235 "row {o}: {} vs {expect}",
6236 a[o]
6237 );
6238 }
6239 }
6240
6241 #[test]
6242 fn q1_tbl_fast_path_matches_reference() {
6243 let (rows, cols) = (5, 256);
6248 let gpr = cols / GROUP_SIZE;
6249 let mut bytes = Vec::new();
6250 for t in 0..rows * gpr {
6251 let s = 0.007 + (t % 11) as f32 * 0.004;
6252 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
6253 for j in 0..4 {
6254 bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
6255 }
6256 }
6257 let x: Vec<f32> = (0..cols)
6258 .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
6259 .collect();
6260 let mut w = vec![0.0f32; rows * cols];
6261 cortiq_core::quant::dequant_q1(&bytes, &mut w);
6262 let mut got = vec![0.0f32; rows];
6263 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
6264 for o in 0..rows {
6265 let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
6266 assert!(
6267 (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
6268 "row {o}: {} vs {expect}",
6269 got[o]
6270 );
6271 }
6272 let b = 5usize;
6275 let mut xs_all = Vec::new();
6276 for bi in 0..b {
6277 xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
6278 }
6279 let mut mm = vec![0.0f32; b * rows];
6280 q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
6281 for bi in 0..b {
6282 let mut single = vec![0.0f32; rows];
6283 q1_matvec(&bytes, &xs_all[bi * cols..(bi + 1) * cols], rows, cols, &mut single, None);
6284 assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
6285 }
6286 }
6287
6288 #[test]
6289 fn q1_kernels_match_exact_reference() {
6290 let (rows, cols) = (7, 96);
6292 let gpr = cols / GROUP_SIZE;
6293 let mut bytes = Vec::new();
6294 for t in 0..rows * gpr {
6295 let s = 0.01 + (t % 13) as f32 * 0.003;
6296 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
6297 for j in 0..4 {
6298 bytes.push(((t * 31 + j * 97) % 251) as u8);
6299 }
6300 }
6301 let x: Vec<f32> = (0..cols)
6303 .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
6304 .collect();
6305 let mut w = vec![0.0f32; rows * cols];
6307 cortiq_core::quant::dequant_q1(&bytes, &mut w);
6308 let mut expect = vec![0.0f32; rows];
6309 for o in 0..rows {
6310 expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
6311 }
6312 let mut got = vec![0.0f32; rows];
6313 q1_matvec(&bytes, &x, rows, cols, &mut got, None);
6314 for o in 0..rows {
6315 assert!(
6316 (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
6317 "row {o}: {} vs {}",
6318 got[o],
6319 expect[o]
6320 );
6321 }
6322 let x2: Vec<f32> = x.iter().map(|v| -v).collect();
6324 let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
6325 q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
6326 assert_eq!(a1, got);
6327 let mut xs = x.clone();
6328 xs.extend_from_slice(&x2);
6329 let mut mm = vec![0.0f32; 2 * rows];
6330 q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
6331 assert_eq!(&mm[..rows], got.as_slice());
6332 assert_eq!(&mm[rows..], a2.as_slice());
6333 }
6334
6335 #[test]
6336 fn repack_is_bit_identical() {
6337 let (rows, cols) = (267, 96); let w: Vec<u8> = (0..rows * cols)
6343 .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
6344 .collect();
6345 let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
6346 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
6347 let rep = q8_repack_layout(&w, rows, cols);
6348 for g in 0..rows / 4 {
6350 for c in 0..cols / 16 {
6351 for lane in 0..4 {
6352 assert_eq!(
6353 &rep[g * 4 * cols + c * 64 + lane * 16..g * 4 * cols + c * 64 + lane * 16 + 16],
6354 &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
6355 );
6356 }
6357 }
6358 }
6359 let mut a = vec![0.0f32; rows];
6360 qmatvec(&w, &[], &scales, &x, &[], TensorDtype::Q8Row, rows, cols, &mut a, None);
6361 let mut b = vec![0.0f32; rows];
6362 qmatvec(&w, &rep, &scales, &x, &[], TensorDtype::Q8Row, rows, cols, &mut b, None);
6363 assert_eq!(a, b, "full-range repack output diverged");
6364
6365 #[cfg(target_arch = "aarch64")]
6366 if sdot_enabled() {
6367 let act = split_act(&x);
6369 let mut c1 = vec![0.0f32; rows];
6370 let mut c2 = vec![0.0f32; rows];
6371 q8_range_sdot(&w, &[], &scales, &act, cols, SendMut(c1.as_mut_ptr()), 3, rows - 2);
6372 q8_range_sdot(&w, &rep, &scales, &act, cols, SendMut(c2.as_mut_ptr()), 3, rows - 2);
6373 assert_eq!(c1, c2, "unaligned-range repack output diverged");
6374 }
6375 }
6376
6377 #[test]
6378 fn sdot_a8w8_noise_is_bounded() {
6379 let (rows, cols) = (16, 512);
6383 let w: Vec<u8> = (0..rows * cols)
6384 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
6385 .collect();
6386 let scales = vec![0.01f32; rows];
6387 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
6388 let mut a = vec![0.0f32; rows];
6389 qmatvec(&w, &[], &scales, &x, &[], TensorDtype::Q8Row, rows, cols, &mut a, None);
6390 let (mut num, mut den) = (0f64, 0f64);
6391 for o in 0..rows {
6392 let mut acc = 0.0f32;
6393 for j in 0..cols {
6394 acc += (w[o * cols + j] as i8) as f32 * x[j];
6395 }
6396 let expect = acc * scales[o];
6397 num += ((a[o] - expect) as f64).powi(2);
6398 den += (expect as f64).powi(2);
6399 }
6400 let rel = (num / den.max(1e-12)).sqrt();
6401 assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
6402 }
6403
6404 #[test]
6405 fn i8_dot_neon_matches_scalar() {
6406 let n = 100;
6407 let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
6408 let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
6409 let mut scalar = 0.0f32;
6410 for j in 0..n {
6411 scalar += (w[j] as i8) as f32 * x[j];
6412 }
6413 let fast = dot_i8_f32(&w, &x);
6414 assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
6415 }
6416
6417 #[test]
6419 fn vbitmatvec_matches_full_dequant() {
6420 let (rows, cols) = (6, 64);
6421 let ng = cols / GROUP_SIZE;
6422 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
6424 let mut bytes = bits.clone();
6425 for g in 0..rows * ng {
6426 let s = 0.02 + 0.001 * g as f32;
6427 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
6428 }
6429 for r in 0..rows {
6430 let b = bits[r] as usize;
6431 let (mut acc, mut nb) = (0u64, 0usize);
6432 let mut rowbytes = Vec::new();
6433 for i in 0..cols {
6434 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
6435 acc = (acc << b) | v;
6436 nb += b;
6437 while nb >= 8 {
6438 nb -= 8;
6439 rowbytes.push(((acc >> nb) & 0xFF) as u8);
6440 }
6441 }
6442 if nb > 0 {
6443 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
6444 }
6445 bytes.extend_from_slice(&rowbytes);
6446 }
6447 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
6448
6449 let mut reference = vec![0f32; rows * cols];
6450 cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
6451 let mut expect = vec![0f32; rows];
6452 for r in 0..rows {
6453 expect[r] = reference[r * cols..(r + 1) * cols]
6454 .iter()
6455 .zip(&x)
6456 .map(|(w, xv)| w * xv)
6457 .sum();
6458 }
6459 let mut got = vec![0f32; rows];
6460 let offsets = vbit_row_offsets(&bytes, rows, cols);
6461 vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
6462 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
6466 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
6467 for r in 0..rows {
6468 assert!(
6469 (got[r] - expect[r]).abs() < tol * scale,
6470 "row {r}: {} vs {}",
6471 got[r],
6472 expect[r]
6473 );
6474 }
6475 }
6476
6477 #[test]
6482 #[cfg(target_arch = "x86_64")]
6483 fn vbit_matmat_blocked_matches_per_row() {
6484 let (rows, cols, b) = (64usize, 128usize, 9usize);
6485 let ng = cols / GROUP_SIZE;
6486 let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
6487 let mut bytes = bits.clone();
6488 for g in 0..rows * ng {
6489 let sc = 0.02 + 0.0005 * g as f32;
6490 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
6491 }
6492 for r in 0..rows {
6493 let bw = bits[r] as usize;
6494 let (mut acc, mut nb) = (0u64, 0usize);
6495 let mut rowbytes = Vec::new();
6496 for i in 0..cols {
6497 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
6498 acc = (acc << bw) | v;
6499 nb += bw;
6500 while nb >= 8 {
6501 nb -= 8;
6502 rowbytes.push(((acc >> nb) & 0xFF) as u8);
6503 }
6504 }
6505 if nb > 0 {
6506 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
6507 }
6508 bytes.extend_from_slice(&rowbytes);
6509 }
6510 let x: Vec<f32> =
6511 (0..b * cols).map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5).collect();
6512 let offsets = vbit_row_offsets(&bytes, rows, cols);
6513 let mut y_a = vec![0f32; b * rows];
6514 let mut y_b = vec![0f32; b * rows];
6515 unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
6516 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
6517 unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
6518 vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
6519 unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
6520 let max_d =
6521 y_a.iter().zip(&y_b).map(|(p, q)| (p - q).abs()).fold(0.0f32, f32::max);
6522 assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
6523 }
6524
6525 #[test]
6526 fn q4matvec_matches_full_dequant() {
6527 let (rows, cols) = (8, 64);
6528 let groups = rows * cols / GROUP_SIZE;
6529 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
6531 for i in 0..groups * 16 {
6532 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
6533 }
6534 for g in 0..groups {
6535 let s = 0.01 + 0.003 * g as f32;
6536 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
6537 }
6538 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
6539
6540 let mut reference = vec![0.0f32; rows * cols];
6541 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
6542 let mut expect = vec![0.0f32; rows];
6543 for r in 0..rows {
6544 expect[r] = reference[r * cols..(r + 1) * cols]
6545 .iter()
6546 .zip(&x)
6547 .map(|(w, xv)| w * xv)
6548 .sum();
6549 }
6550
6551 let mut got = vec![0.0f32; rows];
6552 q4matvec(&bytes, &x, rows, cols, &mut got, None);
6553 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
6557 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
6558 for r in 0..rows {
6559 assert!(
6560 (got[r] - expect[r]).abs() < tol * scale,
6561 "row {r}: {} vs {}",
6562 got[r],
6563 expect[r]
6564 );
6565 }
6566 }
6567
6568 #[test]
6571 fn vbitmatvec2_equals_two_singles() {
6572 let (rows, cols) = (6, 64);
6573 let ng = cols / GROUP_SIZE;
6574 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
6575 let mut bytes = bits.clone();
6576 for g in 0..rows * ng {
6577 let s = 0.02 + 0.001 * g as f32;
6578 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
6579 }
6580 for r in 0..rows {
6581 let b = bits[r] as usize;
6582 let (mut acc, mut nb) = (0u64, 0usize);
6583 let mut rowbytes = Vec::new();
6584 for i in 0..cols {
6585 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
6586 acc = (acc << b) | v;
6587 nb += b;
6588 while nb >= 8 {
6589 nb -= 8;
6590 rowbytes.push(((acc >> nb) & 0xFF) as u8);
6591 }
6592 }
6593 if nb > 0 {
6594 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
6595 }
6596 bytes.extend_from_slice(&rowbytes);
6597 }
6598 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
6599 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
6600 let offsets = vbit_row_offsets(&bytes, rows, cols);
6601
6602 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
6603 vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
6604 vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
6605 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
6606 vbitmatvec2(&bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
6607 assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
6608 assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
6609 }
6610
6611 #[test]
6613 fn q4matvec2_equals_two_singles() {
6614 let (rows, cols) = (8, 128);
6615 let groups = rows * cols / GROUP_SIZE;
6616 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
6617 for i in 0..groups * 16 {
6618 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
6619 }
6620 for g in 0..groups {
6621 let s = 0.01 + 0.003 * g as f32;
6622 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
6623 }
6624 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
6627 x1[9] = 250.0;
6628 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
6629
6630 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
6631 q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
6632 q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
6633 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
6634 q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
6635 assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
6636 assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
6637 }
6638
6639 #[test]
6642 fn matvec_many_equals_separate_matvecs() {
6643 use crate::pool::Pool;
6644 let (r1, r2, cols) = (300, 200, 64);
6645 let mk = |salt: usize, rows: usize| {
6646 QTensor::from_f32(
6647 (0..rows * cols).map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5).collect(),
6648 rows,
6649 cols,
6650 )
6651 };
6652 let (a, b) = (mk(1, r1), mk(5, r2));
6653 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
6654 let pool = Pool::new(3);
6655
6656 let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
6657 a.matvec(&x, &mut ea, Some(&pool));
6658 b.matvec(&x, &mut eb, Some(&pool));
6659 let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
6660 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
6661 assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
6662 assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
6663 }
6664
6665 #[test]
6668 fn batched_matmat_equals_per_position_matvec() {
6669 let (rows, cols, b) = (8, 64, 5);
6670 let groups = rows * cols / GROUP_SIZE;
6672 let mut q4 = Vec::new();
6673 for i in 0..groups * 16 {
6674 q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
6675 }
6676 for g in 0..groups {
6677 q4.extend_from_slice(
6678 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
6679 );
6680 }
6681 let ng = cols / GROUP_SIZE;
6683 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
6684 let mut vb = bits.clone();
6685 for g in 0..rows * ng {
6686 vb.extend_from_slice(
6687 &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
6688 );
6689 }
6690 for r in 0..rows {
6691 let bw = bits[r] as usize;
6692 let (mut acc, mut nb) = (0u64, 0usize);
6693 let mut rowbytes = Vec::new();
6694 for i in 0..cols {
6695 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
6696 acc = (acc << bw) | v;
6697 nb += bw;
6698 while nb >= 8 {
6699 nb -= 8;
6700 rowbytes.push(((acc >> nb) & 0xFF) as u8);
6701 }
6702 }
6703 if nb > 0 {
6704 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
6705 }
6706 vb.extend_from_slice(&rowbytes);
6707 }
6708 let offsets = vbit_row_offsets(&vb, rows, cols);
6709
6710 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
6711
6712 let mut got = vec![0f32; b * rows];
6714 q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
6715 for bi in 0..b {
6716 let mut expect = vec![0f32; rows];
6717 q4matvec(&q4, &xs[bi * cols..(bi + 1) * cols], rows, cols, &mut expect, None);
6718 assert_eq!(&got[bi * rows..(bi + 1) * rows], &expect[..], "q4 batch pos {bi}");
6719 }
6720
6721 let mut got = vec![0f32; b * rows];
6723 vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
6724 for bi in 0..b {
6725 let mut expect = vec![0f32; rows];
6726 vbitmatvec(
6727 &vb, &offsets, &xs[bi * cols..(bi + 1) * cols], rows, cols, &mut expect, None,
6728 );
6729 assert_eq!(&got[bi * rows..(bi + 1) * rows], &expect[..], "vbit batch pos {bi}");
6730 }
6731 }
6732
6733 #[test]
6737 fn q4_tiled_matches_q4_block_bitexact() {
6738 let (rows, cols, b) = (8usize, 128usize, 3usize);
6739 let groups = rows * cols / GROUP_SIZE;
6740 let mut split = Vec::with_capacity(groups * 18);
6741 for i in 0..groups * 16 {
6742 split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
6743 }
6744 for g in 0..groups {
6745 split.extend_from_slice(
6746 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
6747 );
6748 }
6749 let (packed, scales) = split.split_at(groups * 16);
6751 let mut tiled = Vec::with_capacity(groups * Q4_TILE);
6752 for g in 0..groups {
6753 tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
6754 tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
6755 }
6756
6757 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
6758 x1[9] = 250.0; let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
6760
6761 let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
6762 q4matvec(&split, &x1, rows, cols, &mut a, None);
6763 q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
6764 assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
6765
6766 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
6767 let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
6768 q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
6769 q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
6770 assert_eq!(a1, t1);
6771 assert_eq!(a2, t2);
6772
6773 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
6774 let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
6775 q4matmat(&split, &xs, b, rows, cols, &mut am, None);
6776 q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
6777 assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
6778 }
6779
6780 #[test]
6787 fn q4matvec_sdot_outlier_exact() {
6788 let (rows, cols) = (4, 128);
6789 let groups = rows * cols / GROUP_SIZE;
6790 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
6791 for i in 0..groups * 16 {
6792 bytes.push(((i * 11 + 5) % 256) as u8);
6793 }
6794 for g in 0..groups {
6795 let s = 0.02 + 0.002 * g as f32;
6796 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
6797 }
6798 let mut x: Vec<f32> = (0..cols)
6799 .map(|i| match i % 3 {
6800 0 => 1.0,
6801 1 => -1.0,
6802 _ => 0.0,
6803 })
6804 .collect();
6805 x[17] = 300.0; let mut reference = vec![0.0f32; rows * cols];
6808 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
6809 let mut expect = vec![0.0f32; rows];
6810 for r in 0..rows {
6811 expect[r] = reference[r * cols..(r + 1) * cols]
6812 .iter()
6813 .zip(&x)
6814 .map(|(w, xv)| w * xv)
6815 .sum();
6816 }
6817 let mut got = vec![0.0f32; rows];
6818 q4matvec(&bytes, &x, rows, cols, &mut got, None);
6819 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
6820 for r in 0..rows {
6821 assert!(
6822 (got[r] - expect[r]).abs() < 2e-3 * scale,
6823 "row {r}: {} vs {} (outlier term must be exact)",
6824 got[r],
6825 expect[r]
6826 );
6827 }
6828 }
6829
6830 #[test]
6834 fn q1t_matvec_matches_reference() {
6835 use cortiq_core::quant::{dequant_q1t, f32_to_f16};
6836 let (rows, cols) = (3usize, 64usize); let gpr = cols / GROUP_SIZE;
6838 let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
6839 let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
6841 let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
6842 let mut bytes = Vec::new();
6843 for r in 0..rows {
6844 for g in 0..gpr {
6845 bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
6846 let mut c = [0u8; 7];
6847 for k in 0..GROUP_SIZE {
6848 let code = if is_out(r * cols + g * GROUP_SIZE + k) {
6850 0
6851 } else {
6852 ((k + r * 3 + g) % 3) as u8 };
6854 cortiq_core::quant::q1t_pack(&mut c, k, code);
6855 }
6856 bytes.extend_from_slice(&c);
6857 }
6858 }
6859 let mut row_ptr = vec![0u32; rows + 1];
6862 for &(idx, _) in &outliers {
6863 row_ptr[idx as usize / cols + 1] += 1;
6864 }
6865 for r in 0..rows {
6866 row_ptr[r + 1] += row_ptr[r];
6867 }
6868 for &p in &row_ptr {
6869 bytes.extend_from_slice(&p.to_le_bytes());
6870 }
6871 for &(idx, v) in &outliers {
6872 bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
6873 bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
6874 }
6875
6876 let mut refw = vec![0f32; rows * cols];
6877 dequant_q1t(&bytes, rows, cols, &mut refw);
6878 let x: Vec<f32> = (0..cols).map(|j| if j % 3 == 0 { 1.0 } else { -1.0 }).collect();
6881 let mut expect = vec![0f32; rows];
6882 for r in 0..rows {
6883 let mut a = 0.0f32;
6884 for j in 0..cols {
6885 a += refw[r * cols + j] * x[j];
6886 }
6887 expect[r] = a;
6888 }
6889 let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
6890 let mut got = vec![0f32; rows];
6891 q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
6892 for r in 0..rows {
6893 assert!((got[r] - expect[r]).abs() < tol(expect[r]), "row {r}: {} vs {}", got[r], expect[r]);
6894 }
6895 let x2: Vec<f32> = x.iter().chain(x.iter().map(|v| v)).copied().collect();
6897 let mut gm = vec![0f32; 2 * rows];
6898 q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
6899 for r in 0..rows {
6900 assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
6901 assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
6902 }
6903 let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
6908 q1t_matvec(&bytes, &x, rows, cols, &mut p1, None);
6909 q1t_matvec(&bytes, &x, rows, cols, &mut p2, None);
6910 for r in 0..rows {
6911 assert!((p1[r] - expect[r]).abs() < tol(expect[r]) && (p2[r] - expect[r]).abs() < tol(expect[r]));
6912 }
6913 }
6914
6915 #[test]
6919 #[ignore]
6920 fn q1t_matvec_speed() {
6921 use cortiq_core::quant::{f32_to_f16, q1t_code, q1t_pack, Q1T_TILE};
6922 use std::time::Instant;
6923 let (rows, cols) = (8192usize, 4096usize); let gpr = cols / GROUP_SIZE;
6925 let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
6926 for r in 0..rows {
6927 for g in 0..gpr {
6928 let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
6929 bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
6930 let mut c = [0u8; 7];
6931 for k in 0..GROUP_SIZE {
6932 q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
6933 }
6934 bytes.extend_from_slice(&c);
6935 }
6936 }
6937 let (n, stride) = (rows * cols, 40usize); let mut row_ptr = vec![0u32; rows + 1];
6939 let mut idx = 0usize;
6940 while idx < n {
6941 row_ptr[idx / cols + 1] += 1;
6942 idx += stride;
6943 }
6944 for r in 0..rows {
6945 row_ptr[r + 1] += row_ptr[r];
6946 }
6947 for &p in &row_ptr {
6948 bytes.extend_from_slice(&p.to_le_bytes());
6949 }
6950 let mut idx = 0usize;
6951 while idx < n {
6952 bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
6953 bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
6954 idx += stride;
6955 }
6956 let x: Vec<f32> = (0..cols).map(|j| if j % 3 == 0 { 1.0 } else { -1.0 }).collect();
6959 let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
6960
6961 let slow = |out: &mut [f32]| {
6963 let mut buf = vec![0f32; cols];
6964 for r in 0..rows {
6965 for g in 0..gpr {
6966 let off = (r * gpr + g) * Q1T_TILE;
6967 let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6968 let codes = &bytes[off + 2..off + Q1T_TILE];
6969 for k in 0..GROUP_SIZE {
6970 buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
6971 1 => s,
6972 2 => -s,
6973 _ => 0.0,
6974 };
6975 }
6976 }
6977 out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
6978 + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
6979 }
6980 };
6981 let iters = 5;
6982 let mut a = vec![0f32; rows];
6983 slow(&mut a); let t = Instant::now();
6985 for _ in 0..iters {
6986 slow(&mut a);
6987 }
6988 let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
6989
6990 let mut b = vec![0f32; rows];
6991 q1t_matvec(&bytes, &x, rows, cols, &mut b, None); let t = Instant::now();
6993 for _ in 0..iters {
6994 q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
6995 }
6996 let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
6997
6998 for r in 0..rows {
6999 assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
7000 }
7001 println!(
7002 "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms fused-LUT {fast_ms:.2} ms => {:.2}x",
7003 slow_ms / fast_ms
7004 );
7005 }
7006}