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