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