1use crate::pool::{matvec_rows, matvec_rows2, Pool};
15use cortiq_core::quant::{f16_to_f32, GROUP_SIZE};
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 },
41}
42
43fn vbit_row_offsets(bytes: &[u8], rows: usize, cols: usize) -> Vec<usize> {
46 let ng = cols / GROUP_SIZE;
47 let bits = &bytes[..rows];
48 let mut offsets = Vec::with_capacity(rows + 1);
49 let mut off = rows + rows * ng * 2;
50 for r in 0..rows {
51 offsets.push(off);
52 off += (cols * bits[r] as usize + 7) / 8;
53 }
54 offsets.push(off);
55 offsets
56}
57
58impl QTensor {
59 pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
60 debug_assert_eq!(data.len(), rows * cols);
61 Self::F32 { data, rows, cols }
62 }
63
64 pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
67 let idx = model
70 .tensor_index(name)
71 .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
72 let entry = &model.tensors[idx];
73 if entry.shape.len() != 2 {
74 return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
75 }
76 let (rows, cols) = (entry.shape[0], entry.shape[1]);
77 let bytes = model.entry_bytes(entry);
78
79 match entry.dtype {
80 TensorDtype::Q8Row | TensorDtype::Q8_2f => {
81 let n = rows * cols;
82 let scales_off = n;
83 let row_scale: Vec<f32> = (0..rows)
84 .map(|o| {
85 f16_to_f32(u16::from_le_bytes([
86 bytes[scales_off + o * 2],
87 bytes[scales_off + o * 2 + 1],
88 ]))
89 })
90 .collect();
91 let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
92 let col_off = n + rows * 2;
93 (0..cols)
94 .map(|i| {
95 f16_to_f32(u16::from_le_bytes([
96 bytes[col_off + i * 2],
97 bytes[col_off + i * 2 + 1],
98 ]))
99 })
100 .collect()
101 } else {
102 Vec::new()
103 };
104 Ok(Self::Mapped {
105 model: model.clone(),
106 idx,
107 dtype: entry.dtype,
108 rows,
109 cols,
110 row_scale,
111 col_field,
112 vbit_offsets: Vec::new(),
113 })
114 }
115 TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
117 model: model.clone(),
118 idx,
119 dtype: entry.dtype,
120 rows,
121 cols,
122 row_scale: Vec::new(),
123 col_field: Vec::new(),
124 vbit_offsets: vbit_row_offsets(bytes, rows, cols),
125 }),
126 TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
129 model: model.clone(),
130 idx,
131 dtype: entry.dtype,
132 rows,
133 cols,
134 row_scale: Vec::new(),
135 col_field: Vec::new(),
136 vbit_offsets: Vec::new(),
137 }),
138 _ => {
140 let mut data = vec![0.0f32; rows * cols];
141 cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
142 Ok(Self::from_f32(data, rows, cols))
143 }
144 }
145 }
146
147 pub fn rows(&self) -> usize {
148 match self {
149 Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
150 }
151 }
152
153 pub fn cols(&self) -> usize {
154 match self {
155 Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
156 }
157 }
158
159 pub fn as_f32(&self) -> Option<&[f32]> {
162 match self {
163 Self::F32 { data, .. } => Some(data),
164 Self::Mapped { .. } => None,
165 }
166 }
167
168 fn quant_bytes(&self) -> &[u8] {
169 match self {
170 Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
171 Self::F32 { .. } => unreachable!("quant_bytes on F32"),
172 }
173 }
174
175 pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
177 let cols = self.cols();
178 debug_assert_eq!(dst.len(), cols);
179 match self {
180 Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
181 Self::Mapped {
182 dtype,
183 row_scale,
184 col_field,
185 vbit_offsets,
186 ..
187 } => {
188 if *dtype == TensorDtype::Q4Block {
189 let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
190 let gpr = cols / GROUP_SIZE;
191 for gi in 0..gpr {
192 let g = r * gpr + gi;
193 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
194 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
195 dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
196 dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
197 }
198 }
199 return;
200 }
201 if *dtype == TensorDtype::Vbit {
202 let bytes = self.quant_bytes();
203 let rows = self.rows();
204 let ng = cols / GROUP_SIZE;
205 let bits = &bytes[..rows];
206 let sc_off = rows;
207 let off = vbit_offsets[r];
210 let b = bits[r] as usize;
211 let l = ((1usize << (b - 1)) - 1) as f32;
212 let data = &bytes[off..];
213 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
214 for (i, d) in dst.iter_mut().enumerate() {
215 while nbits < b {
216 acc = (acc << 8) | data[idx] as u64;
217 idx += 1;
218 nbits += 8;
219 }
220 let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
221 nbits -= b;
222 let so = (r * ng + i / GROUP_SIZE) * 2;
223 let sv = f16_to_f32(u16::from_le_bytes([
224 bytes[sc_off + so],
225 bytes[sc_off + so + 1],
226 ]));
227 *d = (u - l) * sv;
228 }
229 return;
230 }
231 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
232 let s = row_scale[r];
233 match dtype {
234 TensorDtype::Q8Row => {
235 for (d, &b) in dst.iter_mut().zip(q) {
236 *d = (b as i8) as f32 * s;
237 }
238 }
239 TensorDtype::Q8_2f => {
240 for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
241 *d = (b as i8) as f32 * s * col_field[i];
242 }
243 }
244 _ => unreachable!(),
245 }
246 }
247 }
248 }
249
250 pub fn sparse_col_ok(&self) -> bool {
255 match self {
256 Self::F32 { .. } => true,
257 Self::Mapped { dtype, .. } => {
258 matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
259 }
260 }
261 }
262
263 pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
267 let inter = self.cols();
268 let hidden = self.rows();
269 debug_assert_eq!(out.len(), hidden);
270 match self {
271 Self::F32 { data, .. } => {
272 for (k, o) in out.iter_mut().enumerate() {
273 *o += w * data[k * inter + c];
274 }
275 }
276 Self::Mapped {
277 dtype, row_scale, col_field, ..
278 } => {
279 let q = self.quant_bytes();
280 let colf = if *dtype == TensorDtype::Q8_2f {
281 col_field[c]
282 } else {
283 1.0
284 };
285 let wc = w * colf;
286 for (k, o) in out.iter_mut().enumerate() {
287 let b = q[k * inter + c] as i8 as f32;
288 *o += wc * b * row_scale[k];
289 }
290 }
291 }
292 }
293
294 pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
298 let cols = self.cols();
299 match self {
300 Self::F32 { data, .. } => {
301 let row = &data[r * cols..(r + 1) * cols];
302 row.iter().zip(x).map(|(w, v)| w * v).sum()
303 }
304 Self::Mapped { dtype, row_scale, col_field, .. } => match dtype {
305 TensorDtype::Q8Row => {
306 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
307 dot_i8_f32(q, x) * row_scale[r]
308 }
309 TensorDtype::Q8_2f => {
310 let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
311 dot_i8_col_f32(q, x, col_field) * row_scale[r]
312 }
313 _ => {
314 self.row_f32(r, scratch);
315 scratch.iter().zip(x).map(|(w, v)| w * v).sum()
316 }
317 },
318 }
319 }
320
321 pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
324 match self {
325 Self::F32 { data, .. } => matvec_rows(pool, data, x, out),
326 Self::Mapped {
327 model,
328 idx,
329 dtype,
330 rows,
331 cols,
332 row_scale,
333 col_field,
334 vbit_offsets,
335 } => {
336 let _ = (model, idx);
337 if *dtype == TensorDtype::Q4Block {
338 q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
339 return;
340 }
341 if *dtype == TensorDtype::Vbit {
342 vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
343 return;
344 }
345 let xs = prescale(x, col_field, *dtype);
346 if *rows >= crate::gpu::min_rows()
351 && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
352 && std::env::var("CMF_GPU_LMHEAD").map(|v| v != "0").unwrap_or(true)
353 && crate::gpu::enabled_here()
354 {
355 let frac = std::env::var("CMF_GPU_SPLIT")
356 .ok()
357 .and_then(|v| v.parse::<f32>().ok())
358 .unwrap_or(0.5)
359 .clamp(0.0, 1.0);
360 let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
361 let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
362 let bytes = self.quant_bytes();
363 let ok = std::thread::scope(|sc| {
364 let g = sc.spawn(|| {
365 crate::gpu::q8_matvec_range(
366 model,
367 *idx,
368 cpu_rows,
369 &row_scale[cpu_rows..],
370 &xs,
371 *rows - cpu_rows,
372 *cols,
373 out_gpu,
374 )
375 });
376 if cpu_rows > 0 {
377 qmatvec(
378 &bytes[..cpu_rows * *cols],
379 &row_scale[..cpu_rows],
380 &xs,
381 cpu_rows,
382 *cols,
383 out_cpu,
384 pool,
385 );
386 }
387 g.join().unwrap_or(false)
388 });
389 if ok {
390 return;
391 }
392 qmatvec(
394 &bytes[cpu_rows * *cols..(*rows) * *cols],
395 &row_scale[cpu_rows..],
396 &xs,
397 *rows - cpu_rows,
398 *cols,
399 out_gpu,
400 pool,
401 );
402 return;
403 }
404 qmatvec(self.quant_bytes(), row_scale, &xs, *rows, *cols, out, pool);
405 }
406 }
407 }
408
409 pub fn matvec2(&self, x1: &[f32], x2: &[f32], o1: &mut [f32], o2: &mut [f32], pool: Option<&Pool>) {
411 match self {
412 Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
413 Self::Mapped {
414 dtype,
415 rows,
416 cols,
417 row_scale,
418 col_field,
419 vbit_offsets,
420 ..
421 } => {
422 if *dtype == TensorDtype::Q4Block {
423 q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
424 return;
425 }
426 if *dtype == TensorDtype::Vbit {
427 vbitmatvec2(self.quant_bytes(), vbit_offsets, x1, x2, *rows, *cols, o1, o2, pool);
428 return;
429 }
430 let x1s = prescale(x1, col_field, *dtype);
431 let x2s = prescale(x2, col_field, *dtype);
432 qmatvec2(self.quant_bytes(), row_scale, &x1s, &x2s, *rows, *cols, o1, o2, pool);
433 }
434 }
435 }
436}
437
438impl QTensor {
439 pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
444 let cols = self.cols();
445 let rows = self.rows();
446 debug_assert_eq!(xs_all.len(), b * cols);
447 debug_assert_eq!(out.len(), b * rows);
448 match self {
449 Self::F32 { data, .. } => {
450 let out_addr = SendMut(out.as_mut_ptr());
451 let run = |start: usize, end: usize| {
452 for o in start..end {
453 let row = &data[o * cols..(o + 1) * cols];
454 for bi in 0..b {
455 let x = &xs_all[bi * cols..(bi + 1) * cols];
456 let mut acc = 0f32;
457 for j in 0..cols {
458 acc += row[j] * x[j];
459 }
460 unsafe { *out_addr.at(bi * rows + o) = acc };
461 }
462 }
463 };
464 dispatch_rows(pool, rows, &run);
465 }
466 Self::Mapped {
467 dtype,
468 row_scale,
469 col_field,
470 vbit_offsets,
471 ..
472 } => {
473 if *dtype == TensorDtype::Q4Block {
474 q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
475 return;
476 }
477 if *dtype == TensorDtype::Vbit {
478 vbitmatmat(self.quant_bytes(), vbit_offsets, xs_all, b, rows, cols, out, pool);
479 return;
480 }
481 let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
482 .map(|bi| {
483 prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype)
484 })
485 .collect();
486 if b >= 8
489 && b * rows * cols >= 128_000_000
490 && crate::gpu::enabled_here()
491 {
492 if let Self::Mapped { model, idx, .. } = self {
493 let flat: Vec<f32> =
494 pre.iter().flat_map(|v| v.iter().copied()).collect();
495 if crate::gpu::q8_matmat(
496 model, *idx, row_scale, &flat, b, rows, cols, out)
497 {
498 return;
499 }
500 }
501 }
502 let q = self.quant_bytes();
503 qmatmat(q, row_scale, &pre, rows, cols, out, pool);
504 }
505 }
506 }
507}
508
509impl QTensor {
510 pub fn matvec_many<const N: usize>(
517 ts: [&QTensor; N],
518 x: &[f32],
519 mut outs: [&mut [f32]; N],
520 pool: Option<&Pool>,
521 ) {
522 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
523 let uniform_q8 = ts.iter().all(|t| {
524 matches!(
525 t,
526 Self::Mapped { dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f, .. }
527 )
528 });
529 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
530 let uniform_q4 = ts
531 .iter()
532 .all(|t| matches!(t, Self::Mapped { dtype: TensorDtype::Q4Block, .. }));
533 let uniform_vbit = ts
534 .iter()
535 .all(|t| matches!(t, Self::Mapped { dtype: TensorDtype::Vbit, .. }));
536 let Some(pool) = pool else {
537 for (t, o) in ts.iter().zip(outs.iter_mut()) {
538 t.matvec(x, o, None);
539 }
540 return;
541 };
542 if total_rows < 256 || !(uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit) {
543 for (t, o) in ts.iter().zip(outs.iter_mut()) {
544 t.matvec(x, o, Some(pool));
545 }
546 return;
547 }
548
549 if uniform_q4 || uniform_vbit {
550 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
551 if a8w8_enabled() {
553 let act = split_act(x);
554 let act = &act;
555 if uniform_q4 {
556 let closures: [_; N] = std::array::from_fn(|i| {
557 let (packed, scales) =
558 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
559 let (gpr, cols, out) =
560 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
561 move |s: usize, e: usize| {
562 q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
563 }
564 });
565 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
566 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
567 pool.run_many(&parts);
568 } else {
569 let closures: [_; N] = std::array::from_fn(|i| {
570 let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
571 let (bytes, rows, cols, out) =
572 (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), outs_addr[i]);
573 move |s: usize, e: usize| {
574 vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
575 }
576 });
577 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
578 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
579 pool.run_many(&parts);
580 }
581 return;
582 }
583 if uniform_q4 {
584 let closures: [_; N] = std::array::from_fn(|i| {
585 let (packed, scales) =
586 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
587 let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
588 move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
589 });
590 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
591 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
592 pool.run_many(&parts);
593 } else {
594 let closures: [_; N] = std::array::from_fn(|i| {
595 let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
596 let (bytes, rows, cols, out) =
597 (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), outs_addr[i]);
598 move |s: usize, e: usize| {
599 vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
600 }
601 });
602 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
603 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
604 pool.run_many(&parts);
605 }
606 return;
607 }
608
609 if uniform_f32 {
610 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
611 let closures: [_; N] = std::array::from_fn(|i| {
612 let Self::F32 { data, cols, .. } = ts[i] else { unreachable!() };
613 let out = outs_addr[i];
614 move |start: usize, end: usize| {
615 for o in start..end {
616 let row = &data[o * cols..(o + 1) * cols];
617 let mut sum = 0.0f32;
618 for j in 0..*cols {
619 sum += row[j] * x[j];
620 }
621 unsafe { *out.at(o) = sum };
623 }
624 }
625 });
626 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
627 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
628 pool.run_many(&parts);
629 return;
630 }
631
632 struct Ctx<'a> {
635 bytes: &'a [u8],
636 row_scale: &'a [f32],
637 cols: usize,
638 xs: std::borrow::Cow<'a, [f32]>,
639 }
640 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
641 let Self::Mapped { dtype, cols, row_scale, col_field, .. } = ts[i] else {
642 unreachable!()
643 };
644 Ctx {
645 bytes: ts[i].quant_bytes(),
646 row_scale,
647 cols: *cols,
648 xs: prescale(x, col_field, *dtype),
649 }
650 });
651 let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
652 #[cfg(target_arch = "aarch64")]
653 if sdot_enabled() {
654 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
655 let closures: [_; N] = std::array::from_fn(|i| {
656 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
657 move |start: usize, end: usize| {
658 q8_range_sdot(c.bytes, c.row_scale, act, c.cols, out, start, end)
659 }
660 });
661 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
662 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
663 pool.run_many(&parts);
664 return;
665 }
666 #[cfg(target_arch = "x86_64")]
667 if avx2_a8w8_enabled() {
668 let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
669 let closures: [_; N] = std::array::from_fn(|i| {
670 let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
671 move |start: usize, end: usize| {
672 q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
673 }
674 });
675 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
676 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
677 pool.run_many(&parts);
678 return;
679 }
680 let closures: [_; N] = std::array::from_fn(|i| {
681 let (c, out) = (&ctxs[i], outs_addr[i]);
682 move |start: usize, end: usize| {
683 q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
684 }
685 });
686 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
687 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
688 pool.run_many(&parts);
689 }
690}
691
692impl QTensor {
693 #[allow(clippy::needless_range_loop)]
698 pub fn matvec2_many<const N: usize>(
699 ts: [&QTensor; N],
700 x1: &[f32],
701 x2: &[f32],
702 mut o1s: [&mut [f32]; N],
703 mut o2s: [&mut [f32]; N],
704 pool: Option<&Pool>,
705 ) {
706 let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
707 let uniform_q8 = ts.iter().all(|t| {
708 matches!(
709 t,
710 Self::Mapped { dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f, .. }
711 )
712 });
713 let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
714 let uniform_q4 = ts
715 .iter()
716 .all(|t| matches!(t, Self::Mapped { dtype: TensorDtype::Q4Block, .. }));
717 let uniform_vbit = ts
718 .iter()
719 .all(|t| matches!(t, Self::Mapped { dtype: TensorDtype::Vbit, .. }));
720 let fusable = pool.is_some()
721 && total_rows >= 256
722 && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
723 if !fusable {
724 for i in 0..N {
725 ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
726 }
727 return;
728 }
729 let pool = pool.unwrap();
730
731 if uniform_q4 || uniform_vbit {
732 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
733 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
734 if a8w8_enabled() {
736 let a1 = split_act(x1);
737 let a2 = split_act(x2);
738 let (a1, a2) = (&a1, &a2);
739 if uniform_q4 {
740 let closures: [_; N] = std::array::from_fn(|i| {
741 let (packed, scales) =
742 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
743 let (gpr, cols, o1, o2) =
744 (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
745 move |s: usize, e: usize| {
746 q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
747 }
748 });
749 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
750 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
751 pool.run_many(&parts);
752 } else {
753 let closures: [_; N] = std::array::from_fn(|i| {
754 let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
755 let (bytes, rows, cols, o1, o2) =
756 (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), p1[i], p2[i]);
757 move |s: usize, e: usize| {
758 vbit_range2_a8w8(
759 bytes, vbit_offsets, x1, x2, a1, a2, rows, cols, o1, o2, s, e,
760 )
761 }
762 });
763 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
764 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
765 pool.run_many(&parts);
766 }
767 return;
768 }
769 if uniform_q4 {
770 let closures: [_; N] = std::array::from_fn(|i| {
771 let (packed, scales) =
772 q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
773 let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
774 move |s: usize, e: usize| {
775 q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
776 }
777 });
778 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
779 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
780 pool.run_many(&parts);
781 } else {
782 let closures: [_; N] = std::array::from_fn(|i| {
783 let Self::Mapped { vbit_offsets, .. } = ts[i] else { unreachable!() };
784 let (bytes, rows, cols, o1, o2) =
785 (ts[i].quant_bytes(), ts[i].rows(), ts[i].cols(), p1[i], p2[i]);
786 move |s: usize, e: usize| {
787 vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
788 }
789 });
790 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
791 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
792 pool.run_many(&parts);
793 }
794 return;
795 }
796
797 if uniform_f32 {
798 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
799 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
800 let closures: [_; N] = std::array::from_fn(|i| {
801 let Self::F32 { data, cols, .. } = ts[i] else { unreachable!() };
802 let (o1, o2) = (p1[i], p2[i]);
803 move |start: usize, end: usize| {
804 for o in start..end {
805 let row = &data[o * cols..(o + 1) * cols];
806 let (mut s1, mut s2) = (0.0f32, 0.0f32);
807 for j in 0..*cols {
808 s1 += row[j] * x1[j];
809 s2 += row[j] * x2[j];
810 }
811 unsafe {
813 *o1.at(o) = s1;
814 *o2.at(o) = s2;
815 }
816 }
817 }
818 });
819 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
820 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
821 pool.run_many(&parts);
822 return;
823 }
824
825 struct Ctx<'a> {
826 bytes: &'a [u8],
827 row_scale: &'a [f32],
828 cols: usize,
829 xs1: std::borrow::Cow<'a, [f32]>,
830 xs2: std::borrow::Cow<'a, [f32]>,
831 }
832 let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
833 let Self::Mapped { dtype, cols, row_scale, col_field, .. } = ts[i] else {
834 unreachable!()
835 };
836 Ctx {
837 bytes: ts[i].quant_bytes(),
838 row_scale,
839 cols: *cols,
840 xs1: prescale(x1, col_field, *dtype),
841 xs2: prescale(x2, col_field, *dtype),
842 }
843 });
844 let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
845 let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
846 #[cfg(target_arch = "aarch64")]
847 if sdot_enabled() {
848 let acts: [(SplitAct, SplitAct); N] =
849 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
850 let closures: [_; N] = std::array::from_fn(|i| {
851 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
852 move |start: usize, end: usize| {
853 q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
854 }
855 });
856 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
857 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
858 pool.run_many(&parts);
859 return;
860 }
861 #[cfg(target_arch = "x86_64")]
862 if avx2_a8w8_enabled() {
863 let acts: [(SplitAct, SplitAct); N] =
864 std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
865 let closures: [_; N] = std::array::from_fn(|i| {
866 let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
867 move |start: usize, end: usize| {
868 q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
869 }
870 });
871 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
872 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
873 pool.run_many(&parts);
874 return;
875 }
876 let closures: [_; N] = std::array::from_fn(|i| {
877 let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
878 move |start: usize, end: usize| {
879 q8_range2_f32(c.bytes, c.row_scale, &c.xs1, &c.xs2, c.cols, o1, o2, start, end)
880 }
881 });
882 let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
883 std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
884 pool.run_many(&parts);
885 }
886}
887
888fn qmatmat(
891 q: &[u8],
892 row_scale: &[f32],
893 pre: &[std::borrow::Cow<'_, [f32]>],
894 rows: usize,
895 cols: usize,
896 out: &mut [f32],
897 pool: Option<&Pool>,
898) {
899 let b = pre.len();
900 debug_assert_eq!(out.len(), b * rows);
901 #[cfg(target_arch = "aarch64")]
902 if sdot_enabled() {
903 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
904 let out_addr = SendMut(out.as_mut_ptr());
905 let run = |start: usize, end: usize| {
906 for o in start..end {
907 let row = &q[o * cols..(o + 1) * cols];
908 for (bi, act) in acts.iter().enumerate() {
909 let v = row_dot_sdot(row, act) * row_scale[o];
910 unsafe { *out_addr.at(bi * rows + o) = v };
911 }
912 }
913 };
914 dispatch_rows(pool, rows, &run);
915 return;
916 }
917 #[cfg(target_arch = "x86_64")]
921 if avx2_a8w8_enabled() {
922 let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
923 let out_addr = SendMut(out.as_mut_ptr());
924 let run = |start: usize, end: usize| {
925 for o in start..end {
926 let row = &q[o * cols..(o + 1) * cols];
927 for (bi, act) in acts.iter().enumerate() {
928 let v = row_dot_avx2(row, act) * row_scale[o];
929 unsafe { *out_addr.at(bi * rows + o) = v };
930 }
931 }
932 };
933 dispatch_rows(pool, rows, &run);
934 return;
935 }
936 let out_addr = SendMut(out.as_mut_ptr());
937 let run = |start: usize, end: usize| {
938 for o in start..end {
939 let row = &q[o * cols..(o + 1) * cols];
940 for (bi, x) in pre.iter().enumerate() {
941 let mut acc = 0f32;
942 for j in 0..cols {
943 acc += (row[j] as i8) as f32 * x[j];
944 }
945 unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
946 }
947 }
948 };
949 dispatch_rows(pool, rows, &run);
950}
951
952fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
955 match pool {
956 Some(pool) if rows >= 256 => pool.run_rows(rows, run),
957 _ => run(0, rows),
958 }
959}
960
961fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
963 let groups = rows * cols / GROUP_SIZE;
964 bytes.split_at(groups * 16)
965}
966
967#[inline]
972fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
973 #[cfg(target_arch = "aarch64")]
974 unsafe {
975 return vbit_fill4_neon(data, buf);
976 }
977 #[cfg(target_arch = "x86_64")]
978 if avx2_enabled() {
979 return unsafe { vbit_fill4_avx2(data, buf) };
980 }
981 #[allow(unreachable_code)]
982 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
983 let u = unpack8::<4>(&data[blk * 4..]);
984 for k in 0..8 {
985 chunk[k] = (u[k] - 7) as i8 as u8;
986 }
987 }
988}
989
990#[cfg(target_arch = "aarch64")]
991#[target_feature(enable = "neon")]
992unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
993 unsafe {
996 use core::arch::aarch64::*;
997 let n = buf.len();
998 let mask = vdupq_n_u8(0x0F);
999 let seven = vdupq_n_s8(7);
1000 let mut g = 0usize;
1001 while g * 32 + 32 <= n {
1002 let b = vld1q_u8(data.as_ptr().add(g * 16));
1003 let hi = vshrq_n_u8::<4>(b);
1004 let lo = vandq_u8(b, mask);
1005 let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
1006 let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
1007 vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
1008 vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
1009 g += 1;
1010 }
1011 }
1012}
1013
1014#[cfg(target_arch = "x86_64")]
1015#[target_feature(enable = "avx2")]
1016unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
1017 unsafe {
1019 use core::arch::x86_64::*;
1020 let n = buf.len();
1021 let mask = _mm_set1_epi8(0x0F);
1022 let seven = _mm256_set1_epi8(7);
1023 let mut g = 0usize;
1024 while g * 32 + 32 <= n {
1025 let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
1026 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
1027 let lo = _mm_and_si128(b, mask);
1028 let z = _mm256_sub_epi8(
1029 _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
1030 seven,
1031 );
1032 _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
1033 g += 1;
1034 }
1035 }
1036}
1037
1038#[inline(always)]
1043fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
1044 let mut acc = 0u64;
1045 for i in 0..B {
1046 acc = (acc << 8) | data[i] as u64;
1047 }
1048 let mask = (1u64 << B) - 1;
1049 let mut out = [0i32; 8];
1050 for (k, o) in out.iter_mut().enumerate() {
1051 *o = ((acc >> ((7 - k) * B)) & mask) as i32;
1052 }
1053 out
1054}
1055
1056#[allow(clippy::too_many_arguments)]
1062fn vbitmatvec(
1063 bytes: &[u8],
1064 offsets: &[usize],
1065 x: &[f32],
1066 rows: usize,
1067 cols: usize,
1068 out: &mut [f32],
1069 pool: Option<&Pool>,
1070) {
1071 debug_assert_eq!(out.len(), rows);
1072 debug_assert_eq!(offsets.len(), rows + 1);
1073
1074 if a8w8_enabled() {
1078 let act = split_act(x);
1079 let out_addr = SendMut(out.as_mut_ptr());
1080 let run = move |start: usize, end: usize| {
1081 vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
1082 };
1083 dispatch_rows(pool, rows, &run);
1084 return;
1085 }
1086
1087 let out_addr = SendMut(out.as_mut_ptr());
1088 let run = move |start: usize, end: usize| {
1089 vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
1090 };
1091 dispatch_rows(pool, rows, &run);
1092}
1093
1094#[allow(clippy::too_many_arguments)]
1098fn vbit_range_a8w8(
1099 bytes: &[u8],
1100 offsets: &[usize],
1101 x: &[f32],
1102 act: &SplitAct,
1103 rows: usize,
1104 cols: usize,
1105 out: SendMut,
1106 start: usize,
1107 end: usize,
1108) {
1109 let ng = cols / GROUP_SIZE;
1110 let bits = &bytes[..rows];
1111 let sc_off = rows;
1112 let row_dot = |r: usize| -> f32 {
1113 let b = bits[r] as usize;
1114 let l = ((1i32 << (b - 1)) - 1) as i32;
1115 let mask = (1u64 << b) - 1;
1116 let data = &bytes[offsets[r]..offsets[r + 1]];
1117 if b == 8 {
1118 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
1120 let mut dot = 0f32;
1121 for g in 0..ng {
1122 let so = (r * ng + g) * 2;
1123 let sgf = f16_to_f32(u16::from_le_bytes([
1124 bytes[sc_off + so],
1125 bytes[sc_off + so + 1],
1126 ]));
1127 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
1128 let mut gd = 0f32;
1129 for &xv in xg.iter() {
1130 if nbits < 8 {
1131 acc = (acc << 8) | data[idx] as u64;
1132 idx += 1;
1133 nbits += 8;
1134 }
1135 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
1136 nbits -= 8;
1137 gd += (u - l) as f32 * xv;
1138 }
1139 dot += gd * sgf;
1140 }
1141 return dot;
1142 }
1143 thread_local! {
1147 static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
1148 const { std::cell::RefCell::new(Vec::new()) };
1149 }
1150 #[inline(always)]
1151 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
1152 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
1153 let u = unpack8::<B>(&data[blk * B..]);
1154 for k in 0..8 {
1155 chunk[k] = (u[k] - l) as i8 as u8;
1156 }
1157 }
1158 }
1159 let _ = mask;
1160 VBIT_SCRATCH.with(|scratch| {
1161 let mut buf = scratch.borrow_mut();
1162 buf.resize(cols, 0);
1163 match b {
1164 3 => fill::<3>(data, l, &mut buf),
1165 4 => vbit_fill4(data, &mut buf),
1166 5 => fill::<5>(data, l, &mut buf),
1167 6 => fill::<6>(data, l, &mut buf),
1168 _ => unreachable!(),
1169 }
1170 let mut dot = 0f32;
1171 for g in 0..ng {
1172 let so = (r * ng + g) * 2;
1173 let s = f16_to_f32(u16::from_le_bytes([
1174 bytes[sc_off + so],
1175 bytes[sc_off + so + 1],
1176 ]));
1177 let d = dot_i8_i8(
1178 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
1179 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
1180 ) as f32
1181 * act.sx;
1182 dot += d * s;
1183 }
1184 for &(j, xv) in &act.outliers {
1185 let so = (r * ng + j / GROUP_SIZE) * 2;
1186 let s = f16_to_f32(u16::from_le_bytes([
1187 bytes[sc_off + so],
1188 bytes[sc_off + so + 1],
1189 ]));
1190 dot += (buf[j] as i8) as f32 * s * xv;
1192 }
1193 dot
1194 })
1195 };
1196 for r in start..end {
1197 unsafe { *out.at(r) = row_dot(r) };
1199 }
1200}
1201
1202#[allow(clippy::too_many_arguments)]
1204fn vbit_range_f32(
1205 bytes: &[u8],
1206 offsets: &[usize],
1207 x: &[f32],
1208 rows: usize,
1209 cols: usize,
1210 out: SendMut,
1211 start: usize,
1212 end: usize,
1213) {
1214 let ng = cols / GROUP_SIZE;
1215 let bits = &bytes[..rows];
1216 let sc_off = rows;
1217 #[inline(always)]
1221 fn dot_row<const B: usize>(
1222 data: &[u8],
1223 bytes: &[u8],
1224 sc_off: usize,
1225 r: usize,
1226 ng: usize,
1227 x: &[f32],
1228 ) -> f32 {
1229 let l = ((1i32 << (B - 1)) - 1) as f32;
1230 let gbytes = GROUP_SIZE * B / 8;
1231 let mut dot = 0f32;
1232 for g in 0..ng {
1233 let so = (r * ng + g) * 2;
1234 let s = f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]));
1235 let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
1236 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
1237 let mut gd = 0f32;
1238 for blk in 0..GROUP_SIZE / 8 {
1239 let u = unpack8::<B>(&gd0[blk * B..]);
1240 let xb = &xg[blk * 8..blk * 8 + 8];
1241 for k in 0..8 {
1242 gd += (u[k] as f32 - l) * xb[k];
1243 }
1244 }
1245 dot += gd * s;
1246 }
1247 dot
1248 }
1249 for r in start..end {
1250 let data = &bytes[offsets[r]..offsets[r + 1]];
1251 let v = match bits[r] {
1252 3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
1253 4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
1254 5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
1255 6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
1256 8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
1257 b => unreachable!("vbit bit-width {b} (validated at load)"),
1258 };
1259 unsafe { *out.at(r) = v };
1261 }
1262}
1263
1264#[allow(clippy::too_many_arguments)]
1269fn vbitmatvec2(
1270 bytes: &[u8],
1271 offsets: &[usize],
1272 x1: &[f32],
1273 x2: &[f32],
1274 rows: usize,
1275 cols: usize,
1276 o1: &mut [f32],
1277 o2: &mut [f32],
1278 pool: Option<&Pool>,
1279) {
1280 debug_assert_eq!(o1.len(), rows);
1281 debug_assert_eq!(o2.len(), rows);
1282
1283 if a8w8_enabled() {
1284 let a1 = split_act(x1);
1285 let a2 = split_act(x2);
1286 let p1 = SendMut(o1.as_mut_ptr());
1287 let p2 = SendMut(o2.as_mut_ptr());
1288 let run = move |start: usize, end: usize| {
1289 vbit_range2_a8w8(bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end)
1290 };
1291 dispatch_rows(pool, rows, &run);
1292 return;
1293 }
1294
1295 let p1 = SendMut(o1.as_mut_ptr());
1296 let p2 = SendMut(o2.as_mut_ptr());
1297 let run = move |start: usize, end: usize| {
1298 vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
1299 };
1300 dispatch_rows(pool, rows, &run);
1301}
1302
1303#[allow(clippy::too_many_arguments)]
1307fn vbit_range2_a8w8(
1308 bytes: &[u8],
1309 offsets: &[usize],
1310 x1: &[f32],
1311 x2: &[f32],
1312 a1: &SplitAct,
1313 a2: &SplitAct,
1314 rows: usize,
1315 cols: usize,
1316 p1: SendMut,
1317 p2: SendMut,
1318 start: usize,
1319 end: usize,
1320) {
1321 let ng = cols / GROUP_SIZE;
1322 let bits = &bytes[..rows];
1323 let sc_off = rows;
1324 let row_dots = |r: usize| -> (f32, f32) {
1325 let b = bits[r] as usize;
1326 let l = (1i32 << (b - 1)) - 1;
1327 let data = &bytes[offsets[r]..offsets[r + 1]];
1328 if b == 8 {
1329 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
1332 let (mut d1, mut d2) = (0f32, 0f32);
1333 for g in 0..ng {
1334 let so = (r * ng + g) * 2;
1335 let sgf = f16_to_f32(u16::from_le_bytes([
1336 bytes[sc_off + so],
1337 bytes[sc_off + so + 1],
1338 ]));
1339 let (mut g1, mut g2) = (0f32, 0f32);
1340 for k in 0..GROUP_SIZE {
1341 if nbits < 8 {
1342 acc = (acc << 8) | data[idx] as u64;
1343 idx += 1;
1344 nbits += 8;
1345 }
1346 let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
1347 nbits -= 8;
1348 let w = (u - l) as f32;
1349 g1 += w * x1[g * GROUP_SIZE + k];
1350 g2 += w * x2[g * GROUP_SIZE + k];
1351 }
1352 d1 += g1 * sgf;
1353 d2 += g2 * sgf;
1354 }
1355 return (d1, d2);
1356 }
1357 thread_local! {
1358 static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
1359 const { std::cell::RefCell::new(Vec::new()) };
1360 }
1361 #[inline(always)]
1362 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
1363 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
1364 let u = unpack8::<B>(&data[blk * B..]);
1365 for k in 0..8 {
1366 chunk[k] = (u[k] - l) as i8 as u8;
1367 }
1368 }
1369 }
1370 VBIT_SCRATCH2.with(|scratch| {
1371 let mut buf = scratch.borrow_mut();
1372 buf.resize(cols, 0);
1373 match b {
1374 3 => fill::<3>(data, l, &mut buf),
1375 4 => vbit_fill4(data, &mut buf),
1376 5 => fill::<5>(data, l, &mut buf),
1377 6 => fill::<6>(data, l, &mut buf),
1378 _ => unreachable!(),
1379 }
1380 let (mut d1, mut d2) = (0f32, 0f32);
1381 for g in 0..ng {
1382 let so = (r * ng + g) * 2;
1383 let s = f16_to_f32(u16::from_le_bytes([
1384 bytes[sc_off + so],
1385 bytes[sc_off + so + 1],
1386 ]));
1387 let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
1388 let v1 =
1389 dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
1390 let v2 =
1391 dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
1392 d1 += v1 * s;
1393 d2 += v2 * s;
1394 }
1395 for &(j, xv) in &a1.outliers {
1396 let so = (r * ng + j / GROUP_SIZE) * 2;
1397 let s = f16_to_f32(u16::from_le_bytes([
1398 bytes[sc_off + so],
1399 bytes[sc_off + so + 1],
1400 ]));
1401 d1 += (buf[j] as i8) as f32 * s * xv;
1402 }
1403 for &(j, xv) in &a2.outliers {
1404 let so = (r * ng + j / GROUP_SIZE) * 2;
1405 let s = f16_to_f32(u16::from_le_bytes([
1406 bytes[sc_off + so],
1407 bytes[sc_off + so + 1],
1408 ]));
1409 d2 += (buf[j] as i8) as f32 * s * xv;
1410 }
1411 (d1, d2)
1412 })
1413 };
1414 for r in start..end {
1415 let (v1, v2) = row_dots(r);
1416 unsafe {
1418 *p1.at(r) = v1;
1419 *p2.at(r) = v2;
1420 }
1421 }
1422}
1423
1424#[allow(clippy::too_many_arguments)]
1428fn vbit_range2_f32(
1429 bytes: &[u8],
1430 offsets: &[usize],
1431 x1: &[f32],
1432 x2: &[f32],
1433 rows: usize,
1434 cols: usize,
1435 p1: SendMut,
1436 p2: SendMut,
1437 start: usize,
1438 end: usize,
1439) {
1440 let ng = cols / GROUP_SIZE;
1441 let bits = &bytes[..rows];
1442 let sc_off = rows;
1443 #[inline(always)]
1444 #[allow(clippy::too_many_arguments)]
1445 fn dot_row2<const B: usize>(
1446 data: &[u8],
1447 bytes: &[u8],
1448 sc_off: usize,
1449 r: usize,
1450 ng: usize,
1451 x1: &[f32],
1452 x2: &[f32],
1453 ) -> (f32, f32) {
1454 let l = ((1i32 << (B - 1)) - 1) as f32;
1455 let gbytes = GROUP_SIZE * B / 8;
1456 let (mut d1, mut d2) = (0f32, 0f32);
1457 for g in 0..ng {
1458 let so = (r * ng + g) * 2;
1459 let s = f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]));
1460 let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
1461 let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
1462 let gd0 = &data[g * gbytes..(g + 1) * gbytes];
1463 let (mut g1, mut g2) = (0f32, 0f32);
1464 for blk in 0..GROUP_SIZE / 8 {
1465 let u = unpack8::<B>(&gd0[blk * B..]);
1466 for k in 0..8 {
1467 let w = u[k] as f32 - l;
1468 g1 += w * x1g[blk * 8 + k];
1469 g2 += w * x2g[blk * 8 + k];
1470 }
1471 }
1472 d1 += g1 * s;
1473 d2 += g2 * s;
1474 }
1475 (d1, d2)
1476 }
1477 for r in start..end {
1478 let data = &bytes[offsets[r]..offsets[r + 1]];
1479 let (v1, v2) = match bits[r] {
1480 3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
1481 4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
1482 5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
1483 6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
1484 8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
1485 b => unreachable!("vbit bit-width {b} (validated at load)"),
1486 };
1487 unsafe {
1489 *p1.at(r) = v1;
1490 *p2.at(r) = v2;
1491 }
1492 }
1493}
1494
1495fn q4matvec(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
1501 debug_assert_eq!(out.len(), rows);
1502 let (packed, scales) = q4_split(bytes, rows, cols);
1503 let gpr = cols / GROUP_SIZE;
1504 let out_addr = SendMut(out.as_mut_ptr());
1505
1506 if a8w8_enabled() {
1507 let act = split_act(x);
1508 let run = move |start: usize, end: usize| {
1509 q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
1510 };
1511 dispatch_rows(pool, rows, &run);
1512 return;
1513 }
1514
1515 let run = move |start: usize, end: usize| {
1516 q4_range_f32(packed, scales, gpr, x, out_addr, start, end)
1517 };
1518 dispatch_rows(pool, rows, &run);
1519}
1520
1521#[inline]
1524#[allow(unreachable_code)]
1525fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
1526 #[cfg(target_arch = "aarch64")]
1527 unsafe {
1528 return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
1529 }
1530 #[cfg(target_arch = "x86_64")]
1531 unsafe {
1532 return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
1533 }
1534 let mut acc = 0f32;
1535 for gi in 0..gpr {
1536 let g = g0 + gi;
1537 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
1538 let mut d = 0i32;
1539 for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
1540 d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
1541 + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
1542 }
1543 acc += d as f32 * s;
1544 }
1545 acc
1546}
1547
1548#[inline]
1550#[allow(unreachable_code)]
1551fn dot_q4_row_i8_2(
1552 packed: &[u8],
1553 scales: &[u8],
1554 g0: usize,
1555 gpr: usize,
1556 xq1: &[i8],
1557 xq2: &[i8],
1558) -> (f32, f32) {
1559 #[cfg(target_arch = "aarch64")]
1560 unsafe {
1561 return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
1562 }
1563 #[cfg(target_arch = "x86_64")]
1564 unsafe {
1565 return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
1566 }
1567 (
1568 dot_q4_row_i8(packed, scales, g0, gpr, xq1),
1569 dot_q4_row_i8(packed, scales, g0, gpr, xq2),
1570 )
1571}
1572
1573#[allow(clippy::too_many_arguments)]
1576fn q4_range_a8w8(
1577 packed: &[u8],
1578 scales: &[u8],
1579 gpr: usize,
1580 cols: usize,
1581 act: &SplitAct,
1582 out: SendMut,
1583 start: usize,
1584 end: usize,
1585) {
1586 for r in start..end {
1587 let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
1588 for &(j, xv) in &act.outliers {
1590 let flat = r * cols + j;
1591 let byte = packed[flat / 2];
1592 let nib = if flat & 1 == 0 { byte & 0x0F } else { byte >> 4 };
1593 let s = f16_to_f32(u16::from_le_bytes([
1594 scales[(flat / GROUP_SIZE) * 2],
1595 scales[(flat / GROUP_SIZE) * 2 + 1],
1596 ]));
1597 acc += ((nib as i32 - 8) as f32) * s * xv;
1598 }
1599 unsafe { *out.at(r) = acc };
1601 }
1602}
1603
1604#[allow(clippy::too_many_arguments)]
1607fn q4_range2_a8w8(
1608 packed: &[u8],
1609 scales: &[u8],
1610 gpr: usize,
1611 cols: usize,
1612 a1: &SplitAct,
1613 a2: &SplitAct,
1614 p1: SendMut,
1615 p2: SendMut,
1616 start: usize,
1617 end: usize,
1618) {
1619 for r in start..end {
1620 let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
1621 let mut acc1 = s1 * a1.sx;
1622 let mut acc2 = s2 * a2.sx;
1623 let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
1625 for &(j, xv) in outliers {
1626 let flat = r * cols + j;
1627 let byte = packed[flat / 2];
1628 let nib = if flat & 1 == 0 { byte & 0x0F } else { byte >> 4 };
1629 let s = f16_to_f32(u16::from_le_bytes([
1630 scales[(flat / GROUP_SIZE) * 2],
1631 scales[(flat / GROUP_SIZE) * 2 + 1],
1632 ]));
1633 *acc += ((nib as i32 - 8) as f32) * s * xv;
1634 }
1635 };
1636 fix(&a1.outliers, &mut acc1);
1637 fix(&a2.outliers, &mut acc2);
1638 unsafe {
1640 *p1.at(r) = acc1;
1641 *p2.at(r) = acc2;
1642 }
1643 }
1644}
1645
1646fn q4_range_f32(
1648 packed: &[u8],
1649 scales: &[u8],
1650 gpr: usize,
1651 x: &[f32],
1652 out: SendMut,
1653 start: usize,
1654 end: usize,
1655) {
1656 for r in start..end {
1657 let mut acc = 0f32;
1658 for gi in 0..gpr {
1659 let g = r * gpr + gi;
1660 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
1661 let pk = &packed[g * 16..(g + 1) * 16];
1662 let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
1663 let mut ga = 0f32;
1664 for (k, &b) in pk.iter().enumerate() {
1665 ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
1666 + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
1667 }
1668 acc += ga * s;
1669 }
1670 unsafe { *out.at(r) = acc };
1672 }
1673}
1674
1675#[allow(clippy::too_many_arguments)]
1679fn q4matvec2(
1680 bytes: &[u8],
1681 x1: &[f32],
1682 x2: &[f32],
1683 rows: usize,
1684 cols: usize,
1685 o1: &mut [f32],
1686 o2: &mut [f32],
1687 pool: Option<&Pool>,
1688) {
1689 debug_assert_eq!(o1.len(), rows);
1690 debug_assert_eq!(o2.len(), rows);
1691 let (packed, scales) = q4_split(bytes, rows, cols);
1692 let gpr = cols / GROUP_SIZE;
1693
1694 if a8w8_enabled() {
1695 let a1 = split_act(x1);
1696 let a2 = split_act(x2);
1697 let p1 = SendMut(o1.as_mut_ptr());
1698 let p2 = SendMut(o2.as_mut_ptr());
1699 let run = move |start: usize, end: usize| {
1700 q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
1701 };
1702 dispatch_rows(pool, rows, &run);
1703 return;
1704 }
1705
1706 let p1 = SendMut(o1.as_mut_ptr());
1707 let p2 = SendMut(o2.as_mut_ptr());
1708 let run = move |start: usize, end: usize| {
1709 q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
1710 };
1711 dispatch_rows(pool, rows, &run);
1712}
1713
1714#[allow(clippy::too_many_arguments)]
1716fn q4_range2_f32(
1717 packed: &[u8],
1718 scales: &[u8],
1719 gpr: usize,
1720 x1: &[f32],
1721 x2: &[f32],
1722 p1: SendMut,
1723 p2: SendMut,
1724 start: usize,
1725 end: usize,
1726) {
1727 for r in start..end {
1728 let (mut acc1, mut acc2) = (0f32, 0f32);
1729 for gi in 0..gpr {
1730 let g = r * gpr + gi;
1731 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
1732 let pk = &packed[g * 16..(g + 1) * 16];
1733 let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
1734 let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
1735 let (mut g1, mut g2) = (0f32, 0f32);
1736 for (k, &b) in pk.iter().enumerate() {
1737 let wl = (b & 0x0F) as f32 - 8.0;
1738 let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
1739 g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
1740 g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
1741 }
1742 acc1 += g1 * s;
1743 acc2 += g2 * s;
1744 }
1745 unsafe {
1747 *p1.at(r) = acc1;
1748 *p2.at(r) = acc2;
1749 }
1750 }
1751}
1752
1753thread_local! {
1754 static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
1757 static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
1758}
1759
1760#[allow(clippy::too_many_arguments)]
1766fn q4matmat(
1767 bytes: &[u8],
1768 xs_all: &[f32],
1769 b: usize,
1770 rows: usize,
1771 cols: usize,
1772 out: &mut [f32],
1773 pool: Option<&Pool>,
1774) {
1775 debug_assert_eq!(xs_all.len(), b * cols);
1776 debug_assert_eq!(out.len(), b * rows);
1777 let (packed, scales) = q4_split(bytes, rows, cols);
1778 let gpr = cols / GROUP_SIZE;
1779 let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
1780
1781 if a8w8_enabled() {
1782 let acts: Vec<SplitAct> =
1783 (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
1784 let acts = &acts;
1785 let out_addr = SendMut(out.as_mut_ptr());
1786 let run = move |start: usize, end: usize| {
1787 ROW_I8.with(|rb| {
1788 let mut buf = rb.borrow_mut();
1789 buf.resize(cols, 0);
1790 for r in start..end {
1791 for gi in 0..gpr {
1795 let g = r * gpr + gi;
1796 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
1797 buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
1798 buf[gi * GROUP_SIZE + k * 2 + 1] =
1799 (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
1800 }
1801 }
1802 for (bi, act) in acts.iter().enumerate() {
1803 let mut acc = 0f32;
1804 for gi in 0..gpr {
1805 let d = dot_i8_i8(
1806 &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
1807 &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
1808 );
1809 acc += d as f32 * gscale(r * gpr + gi);
1810 }
1811 acc *= act.sx;
1812 for &(j, xv) in &act.outliers {
1814 acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
1815 }
1816 unsafe { *out_addr.at(bi * rows + r) = acc };
1818 }
1819 }
1820 })
1821 };
1822 dispatch_rows(pool, rows, &run);
1823 return;
1824 }
1825
1826 let out_addr = SendMut(out.as_mut_ptr());
1827 let run = move |start: usize, end: usize| {
1828 ROW_F32.with(|rb| {
1829 let mut buf = rb.borrow_mut();
1830 buf.resize(cols, 0.0);
1831 for r in start..end {
1832 for gi in 0..gpr {
1835 let g = r * gpr + gi;
1836 for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
1837 buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
1838 buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
1839 }
1840 }
1841 for bi in 0..b {
1842 let x = &xs_all[bi * cols..(bi + 1) * cols];
1843 let mut acc = 0f32;
1844 for gi in 0..gpr {
1845 let mut ga = 0f32;
1846 for k in 0..GROUP_SIZE / 2 {
1851 let e = gi * GROUP_SIZE + k * 2;
1852 ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
1853 }
1854 acc += ga * gscale(r * gpr + gi);
1855 }
1856 unsafe { *out_addr.at(bi * rows + r) = acc };
1858 }
1859 }
1860 })
1861 };
1862 dispatch_rows(pool, rows, &run);
1863}
1864
1865#[allow(clippy::too_many_arguments)]
1870fn vbitmatmat(
1871 bytes: &[u8],
1872 offsets: &[usize],
1873 xs_all: &[f32],
1874 b: usize,
1875 rows: usize,
1876 cols: usize,
1877 out: &mut [f32],
1878 pool: Option<&Pool>,
1879) {
1880 debug_assert_eq!(xs_all.len(), b * cols);
1881 debug_assert_eq!(out.len(), b * rows);
1882 debug_assert_eq!(offsets.len(), rows + 1);
1883 let ng = cols / GROUP_SIZE;
1884 let bits = &bytes[..rows];
1885 let sc_off = rows;
1886 let gscale = |r: usize, g: usize| {
1887 let so = (r * ng + g) * 2;
1888 f16_to_f32(u16::from_le_bytes([bytes[sc_off + so], bytes[sc_off + so + 1]]))
1889 };
1890
1891 let decode_f32 = |r: usize, dst: &mut [f32]| {
1893 let bw = bits[r] as usize;
1894 let l = ((1i32 << (bw - 1)) - 1) as f32;
1895 let data = &bytes[offsets[r]..offsets[r + 1]];
1896 let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
1897 for d in dst.iter_mut() {
1898 while nbits < bw {
1899 acc = (acc << 8) | data[idx] as u64;
1900 idx += 1;
1901 nbits += 8;
1902 }
1903 let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
1904 nbits -= bw;
1905 *d = u - l;
1906 }
1907 };
1908
1909 if a8w8_enabled() {
1910 let acts: Vec<SplitAct> =
1911 (0..b).map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols])).collect();
1912 let acts = &acts;
1913 let out_addr = SendMut(out.as_mut_ptr());
1914 let run = move |start: usize, end: usize| {
1915 for r in start..end {
1916 let bw = bits[r] as usize;
1917 if bw == 8 {
1918 ROW_F32.with(|rb| {
1921 let mut buf = rb.borrow_mut();
1922 buf.resize(cols, 0.0);
1923 decode_f32(r, &mut buf);
1924 for bi in 0..b {
1925 let x = &xs_all[bi * cols..(bi + 1) * cols];
1926 let mut dot = 0f32;
1927 for g in 0..ng {
1928 let mut gd = 0f32;
1929 for k in 0..GROUP_SIZE {
1930 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
1931 }
1932 dot += gd * gscale(r, g);
1933 }
1934 unsafe { *out_addr.at(bi * rows + r) = dot };
1936 }
1937 });
1938 continue;
1939 }
1940 let l = (1i32 << (bw - 1)) - 1;
1941 let data = &bytes[offsets[r]..offsets[r + 1]];
1942 ROW_I8.with(|rb| {
1943 let mut buf = rb.borrow_mut();
1944 buf.resize(cols, 0);
1945 #[inline(always)]
1946 fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
1947 for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
1948 let u = unpack8::<B>(&data[blk * B..]);
1949 for k in 0..8 {
1950 chunk[k] = (u[k] - l) as i8 as u8;
1951 }
1952 }
1953 }
1954 match bw {
1955 3 => fill::<3>(data, l, &mut buf),
1956 4 => vbit_fill4(data, &mut buf),
1957 5 => fill::<5>(data, l, &mut buf),
1958 6 => fill::<6>(data, l, &mut buf),
1959 _ => unreachable!("vbit bit-width {bw} (validated at load)"),
1960 }
1961 for (bi, act) in acts.iter().enumerate() {
1962 let mut dot = 0f32;
1963 for g in 0..ng {
1964 let d = dot_i8_i8(
1965 &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
1966 &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
1967 ) as f32
1968 * act.sx;
1969 dot += d * gscale(r, g);
1970 }
1971 for &(j, xv) in &act.outliers {
1972 dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
1973 }
1974 unsafe { *out_addr.at(bi * rows + r) = dot };
1976 }
1977 });
1978 }
1979 };
1980 dispatch_rows(pool, rows, &run);
1981 return;
1982 }
1983
1984 let out_addr = SendMut(out.as_mut_ptr());
1985 let run = move |start: usize, end: usize| {
1986 ROW_F32.with(|rb| {
1987 let mut buf = rb.borrow_mut();
1988 buf.resize(cols, 0.0);
1989 for r in start..end {
1990 decode_f32(r, &mut buf);
1991 for bi in 0..b {
1992 let x = &xs_all[bi * cols..(bi + 1) * cols];
1993 let mut dot = 0f32;
1994 for g in 0..ng {
1995 let mut gd = 0f32;
1996 for k in 0..GROUP_SIZE {
1997 gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
1998 }
1999 dot += gd * gscale(r, g);
2000 }
2001 unsafe { *out_addr.at(bi * rows + r) = dot };
2003 }
2004 }
2005 })
2006 };
2007 dispatch_rows(pool, rows, &run);
2008}
2009
2010pub(crate) fn prescale<'a>(
2014 x: &'a [f32],
2015 col_field: &[f32],
2016 dtype: TensorDtype,
2017) -> std::borrow::Cow<'a, [f32]> {
2018 if dtype == TensorDtype::Q8_2f {
2019 x.iter().zip(col_field).map(|(a, c)| a * c).collect()
2020 } else {
2021 std::borrow::Cow::Borrowed(x)
2022 }
2023}
2024
2025#[cfg(target_arch = "x86_64")]
2030pub(crate) fn avx2_enabled() -> bool {
2031 use std::sync::OnceLock;
2032 static ON: OnceLock<bool> = OnceLock::new();
2033 *ON.get_or_init(|| {
2034 std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
2035 && std::arch::is_x86_feature_detected!("avx2")
2036 && std::arch::is_x86_feature_detected!("fma")
2037 })
2038}
2039
2040#[cfg(target_arch = "x86_64")]
2045fn avx2_a8w8_enabled() -> bool {
2046 use std::sync::OnceLock;
2047 static ON: OnceLock<bool> = OnceLock::new();
2048 *ON.get_or_init(|| {
2049 avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
2050 })
2051}
2052
2053#[inline]
2057pub(crate) fn a8w8_enabled() -> bool {
2058 #[cfg(target_arch = "aarch64")]
2059 {
2060 sdot_enabled()
2061 }
2062 #[cfg(target_arch = "x86_64")]
2063 {
2064 avx2_a8w8_enabled()
2065 }
2066 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
2067 {
2068 false
2069 }
2070}
2071
2072#[inline]
2075#[allow(unreachable_code)]
2076fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
2077 #[cfg(target_arch = "aarch64")]
2078 unsafe {
2079 return dot_i8_sdot(w, xq);
2080 }
2081 #[cfg(target_arch = "x86_64")]
2082 unsafe {
2083 if avx512vnni_enabled() {
2084 return dot_i8_i8_vnni(w, xq);
2085 }
2086 return dot_i8_i8_avx2(w, xq);
2087 }
2088 w.iter().zip(xq).map(|(&a, &b)| (a as i8) as i32 * b as i32).sum()
2089}
2090
2091#[cfg(target_arch = "x86_64")]
2095fn avx512vnni_enabled() -> bool {
2096 use std::sync::OnceLock;
2097 static ON: OnceLock<bool> = OnceLock::new();
2098 *ON.get_or_init(|| {
2099 std::env::var("CMF_AVX512").map(|v| v != "0").unwrap_or(true)
2100 && std::arch::is_x86_feature_detected!("avx512f")
2101 && std::arch::is_x86_feature_detected!("avx512bw")
2102 && std::arch::is_x86_feature_detected!("avx512vl")
2103 && std::arch::is_x86_feature_detected!("avx512vnni")
2104 })
2105}
2106
2107#[cfg(target_arch = "x86_64")]
2112#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
2113unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
2114 unsafe {
2116 use core::arch::x86_64::*;
2117 let n = w.len();
2118 let mut j = 0usize;
2119 let mut total: i32;
2120 {
2125 #[inline(always)]
2126 unsafe fn step(
2127 w: *const u8,
2128 x: *const i8,
2129 acc: core::arch::x86_64::__m512i,
2130 ) -> core::arch::x86_64::__m512i {
2131 unsafe {
2132 use core::arch::x86_64::*;
2133 let wv = _mm512_loadu_si512(w as *const _);
2134 let xv = _mm512_loadu_si512(x as *const _);
2135 let aw = _mm512_abs_epi8(wv);
2136 let neg = _mm512_movepi8_mask(wv);
2137 let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
2138 _mm512_dpbusd_epi32(acc, aw, sx)
2139 }
2140 }
2141 let (mut a0, mut a1, mut a2, mut a3) = (
2142 _mm512_setzero_si512(),
2143 _mm512_setzero_si512(),
2144 _mm512_setzero_si512(),
2145 _mm512_setzero_si512(),
2146 );
2147 while j + 256 <= n {
2148 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
2149 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
2150 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
2151 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
2152 j += 256;
2153 }
2154 while j + 64 <= n {
2155 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
2156 j += 64;
2157 }
2158 let s01 = _mm512_add_epi32(a0, a1);
2159 let s23 = _mm512_add_epi32(a2, a3);
2160 total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
2161 }
2162 if j + 32 <= n {
2164 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
2165 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
2166 let d = _mm256_dpbusd_epi32(
2167 _mm256_setzero_si256(),
2168 _mm256_abs_epi8(wv),
2169 _mm256_sign_epi8(xv, wv),
2170 );
2171 let hi128 = _mm256_extracti128_si256::<1>(d);
2172 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
2173 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2174 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2175 total += _mm_cvtsi128_si32(s32);
2176 j += 32;
2177 }
2178 while j < n {
2179 total += (w[j] as i8) as i32 * xq[j] as i32;
2180 j += 1;
2181 }
2182 total
2183 }
2184}
2185
2186#[cfg(target_arch = "x86_64")]
2188#[target_feature(enable = "avx2,fma")]
2189unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
2190 unsafe {
2192 use core::arch::x86_64::*;
2193 let n = x.len();
2194 let wp = w.as_ptr();
2195 let xp = x.as_ptr();
2196 let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
2197 let mut j = 0usize;
2198 while j + 16 <= n {
2199 let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
2200 let lo = _mm256_cvtepi8_epi32(wb);
2201 let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
2202 a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
2203 a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
2204 j += 16;
2205 }
2206 let acc = _mm256_add_ps(a0, a1);
2207 let hi128 = _mm256_extractf128_ps::<1>(acc);
2208 let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
2209 let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
2210 let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
2211 let mut sum = _mm_cvtss_f32(s32);
2212 while j < n {
2213 sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
2214 j += 1;
2215 }
2216 sum
2217 }
2218}
2219
2220#[cfg(target_arch = "x86_64")]
2225#[target_feature(enable = "avx2")]
2226unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
2227 unsafe {
2229 use core::arch::x86_64::*;
2230 let n = w.len();
2231 let ones = _mm256_set1_epi16(1);
2232 let mut acc = _mm256_setzero_si256();
2233 let mut j = 0usize;
2234 while j + 32 <= n {
2235 let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
2236 let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
2237 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
2238 acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
2239 j += 32;
2240 }
2241 let hi128 = _mm256_extracti128_si256::<1>(acc);
2242 let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
2243 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2244 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2245 let mut s = _mm_cvtsi128_si32(s32);
2246 while j < n {
2247 s += (w[j] as i8) as i32 * xq[j] as i32;
2248 j += 1;
2249 }
2250 s
2251 }
2252}
2253
2254#[cfg(target_arch = "x86_64")]
2259#[inline]
2260fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
2261 let dot = if avx512vnni_enabled() && row.len() >= 64 {
2262 (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
2263 } else {
2264 unsafe { dot_i8_i8_avx2(row, &act.xq) }
2265 };
2266 let mut acc = dot as f32 * act.sx;
2267 for &(j, xv) in &act.outliers {
2268 acc += (row[j] as i8) as f32 * xv;
2269 }
2270 acc
2271}
2272
2273#[cfg(target_arch = "x86_64")]
2277#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
2278unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
2279 unsafe {
2281 use core::arch::x86_64::*;
2282 let n = w.len();
2283 let flip = _mm512_set1_epi8(-128); #[inline(always)]
2285 unsafe fn step(
2286 w: *const u8,
2287 x: *const i8,
2288 flip: core::arch::x86_64::__m512i,
2289 acc: core::arch::x86_64::__m512i,
2290 ) -> core::arch::x86_64::__m512i {
2291 unsafe {
2292 use core::arch::x86_64::*;
2293 let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
2294 _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
2295 }
2296 }
2297 let (mut a0, mut a1, mut a2, mut a3) = (
2298 _mm512_setzero_si512(),
2299 _mm512_setzero_si512(),
2300 _mm512_setzero_si512(),
2301 _mm512_setzero_si512(),
2302 );
2303 let mut j = 0usize;
2304 while j + 256 <= n {
2305 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
2306 a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
2307 a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
2308 a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
2309 j += 256;
2310 }
2311 while j + 64 <= n {
2312 a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
2313 j += 64;
2314 }
2315 let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
2316 _mm512_add_epi32(a0, a1),
2317 _mm512_add_epi32(a2, a3),
2318 ));
2319 while j < n {
2321 total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
2322 j += 1;
2323 }
2324 total
2325 }
2326}
2327
2328#[cfg(target_arch = "x86_64")]
2334#[target_feature(enable = "avx2")]
2335unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
2336 unsafe {
2339 use core::arch::x86_64::*;
2340 let lomask = _mm_set1_epi8(0x0F);
2341 let eight = _mm256_set1_epi8(8);
2342 let ones = _mm256_set1_epi16(1);
2343 let mut acc = 0f32;
2344 for gi in 0..gpr {
2345 let g = g0 + gi;
2346 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
2347 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
2348 let lo = _mm_and_si128(b, lomask);
2349 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
2350 let w = _mm256_sub_epi8(
2351 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
2352 eight,
2353 );
2354 let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2355 let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
2356 let d = _mm256_madd_epi16(p16, ones);
2357 let hi128 = _mm256_extracti128_si256::<1>(d);
2358 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
2359 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2360 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2361 acc += _mm_cvtsi128_si32(s32) as f32 * s;
2362 }
2363 acc
2364 }
2365}
2366
2367#[cfg(target_arch = "x86_64")]
2370#[target_feature(enable = "avx2")]
2371unsafe fn dot_q4_row_avx2_2(
2372 packed: &[u8],
2373 scales: &[u8],
2374 g0: usize,
2375 gpr: usize,
2376 xq1: &[i8],
2377 xq2: &[i8],
2378) -> (f32, f32) {
2379 unsafe {
2381 use core::arch::x86_64::*;
2382 let lomask = _mm_set1_epi8(0x0F);
2383 let eight = _mm256_set1_epi8(8);
2384 let ones = _mm256_set1_epi16(1);
2385 let (mut acc1, mut acc2) = (0f32, 0f32);
2386 #[inline(always)]
2387 unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
2388 unsafe {
2389 use core::arch::x86_64::*;
2390 let hi128 = _mm256_extracti128_si256::<1>(d);
2391 let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
2392 let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
2393 let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
2394 _mm_cvtsi128_si32(s32)
2395 }
2396 }
2397 for gi in 0..gpr {
2398 let g = g0 + gi;
2399 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
2400 let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
2401 let lo = _mm_and_si128(b, lomask);
2402 let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
2403 let w = _mm256_sub_epi8(
2404 _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
2405 eight,
2406 );
2407 let aw = _mm256_abs_epi8(w);
2408 let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2409 let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
2410 let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
2411 let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
2412 acc1 += hsum(d1) as f32 * s;
2413 acc2 += hsum(d2) as f32 * s;
2414 }
2415 (acc1, acc2)
2416 }
2417}
2418
2419#[cfg(target_arch = "x86_64")]
2421fn q8_range_avx2(
2422 q: &[u8],
2423 row_scale: &[f32],
2424 act: &SplitAct,
2425 cols: usize,
2426 out_addr: SendMut,
2427 start: usize,
2428 end: usize,
2429) {
2430 for o in start..end {
2431 let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
2432 unsafe { *out_addr.at(o) = v };
2434 }
2435}
2436
2437#[cfg(target_arch = "x86_64")]
2439#[allow(clippy::too_many_arguments)]
2440fn q8_range2_avx2(
2441 q: &[u8],
2442 row_scale: &[f32],
2443 a1: &SplitAct,
2444 a2: &SplitAct,
2445 cols: usize,
2446 p1: SendMut,
2447 p2: SendMut,
2448 start: usize,
2449 end: usize,
2450) {
2451 for o in start..end {
2452 let row = &q[o * cols..(o + 1) * cols];
2453 unsafe {
2455 *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
2456 *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
2457 }
2458 }
2459}
2460
2461#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
2467fn sdot_enabled() -> bool {
2468 use std::sync::OnceLock;
2469 static ON: OnceLock<bool> = OnceLock::new();
2470 *ON.get_or_init(|| {
2471 let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
2472 #[cfg(target_arch = "aarch64")]
2473 {
2474 want && std::arch::is_aarch64_feature_detected!("dotprod")
2475 }
2476 #[cfg(not(target_arch = "aarch64"))]
2477 {
2478 let _ = want;
2479 false
2480 }
2481 })
2482}
2483
2484struct SplitAct {
2489 xq: Vec<i8>,
2490 sx: f32,
2491 outliers: Vec<(usize, f32)>,
2492 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
2495 xsum: i32,
2496}
2497
2498thread_local! {
2499 static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
2502 const { std::cell::RefCell::new(Vec::new()) };
2503}
2504
2505impl Drop for SplitAct {
2506 fn drop(&mut self) {
2507 let buf = std::mem::take(&mut self.xq);
2508 if buf.capacity() > 0 {
2509 XQ_FREE.with(|f| {
2510 let mut f = f.borrow_mut();
2511 if f.len() < 16 {
2512 f.push(buf);
2513 }
2514 });
2515 }
2516 }
2517}
2518
2519fn split_act(x: &[f32]) -> SplitAct {
2520 let n = x.len();
2521 let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
2522 let thr = 8.0 * rms;
2523 let mut outliers: Vec<(usize, f32)> = Vec::new();
2527 let mut amax = 0f32;
2528 for (j, &v) in x.iter().enumerate() {
2529 let a = v.abs();
2530 if a > thr {
2531 outliers.push((j, v));
2532 } else if a > amax {
2533 amax = a;
2534 }
2535 }
2536 let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
2537 let inv = 1.0 / sx;
2538 let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
2539 xq.clear();
2540 xq.reserve(n);
2541 if outliers.is_empty() {
2542 xq.extend(x.iter().map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8));
2543 } else {
2544 xq.extend(x.iter().map(|&v| {
2546 if v.abs() > thr {
2547 0
2548 } else {
2549 (v * inv).round().clamp(-127.0, 127.0) as i8
2550 }
2551 }));
2552 }
2553 let xsum = xq.iter().map(|&v| v as i32).sum();
2554 SplitAct { xq, sx, outliers, xsum }
2555}
2556
2557#[cfg(target_arch = "aarch64")]
2560#[target_feature(enable = "neon,dotprod")]
2561unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
2562 unsafe {
2564 use core::arch::aarch64::*;
2565 use core::arch::asm;
2566 let wp = w.as_ptr() as *const i8;
2567 let n = w.len();
2568 let (mut a0, mut a1, mut a2, mut a3) =
2569 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
2570 let mut i = 0;
2571 while i + 64 <= n {
2572 let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
2573 let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
2574 let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
2575 let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
2576 asm!(
2577 "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
2578 "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
2579 "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
2580 "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
2581 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
2582 w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
2583 w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
2584 options(pure, nomem, nostack),
2585 );
2586 i += 64;
2587 }
2588 while i + 16 <= n {
2589 let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
2590 asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
2591 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
2592 i += 16;
2593 }
2594 let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
2595 while i < n {
2596 s += (*wp.add(i)) as i32 * xq[i] as i32;
2597 i += 1;
2598 }
2599 s
2600}
2601}
2602
2603#[cfg(target_arch = "aarch64")]
2607#[target_feature(enable = "neon,dotprod")]
2608unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
2609 unsafe {
2611 use core::arch::aarch64::*;
2612 use core::arch::asm;
2613 let n = xq.len();
2614 let px = xq.as_ptr();
2615 let (p0, p1, p2, p3) = (
2616 w0.as_ptr() as *const i8,
2617 w1.as_ptr() as *const i8,
2618 w2.as_ptr() as *const i8,
2619 w3.as_ptr() as *const i8,
2620 );
2621 let (mut a0, mut a1, mut a2, mut a3) =
2622 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
2623 let mut i = 0;
2624 while i + 16 <= n {
2625 let x = vld1q_s8(px.add(i));
2626 let v0 = vld1q_s8(p0.add(i));
2627 let v1 = vld1q_s8(p1.add(i));
2628 let v2 = vld1q_s8(p2.add(i));
2629 let v3 = vld1q_s8(p3.add(i));
2630 asm!(
2631 "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
2632 "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
2633 "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
2634 "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
2635 a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
2636 v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
2637 options(pure, nomem, nostack),
2638 );
2639 i += 16;
2640 }
2641 let mut r = [vaddvq_s32(a0), vaddvq_s32(a1), vaddvq_s32(a2), vaddvq_s32(a3)];
2642 while i < n {
2643 let xi = *px.add(i) as i32;
2644 r[0] += (*p0.add(i)) as i32 * xi;
2645 r[1] += (*p1.add(i)) as i32 * xi;
2646 r[2] += (*p2.add(i)) as i32 * xi;
2647 r[3] += (*p3.add(i)) as i32 * xi;
2648 i += 1;
2649 }
2650 r
2651}
2652}
2653
2654#[cfg(target_arch = "aarch64")]
2658fn q8_range_sdot(
2659 q: &[u8],
2660 row_scale: &[f32],
2661 act: &SplitAct,
2662 cols: usize,
2663 out_addr: SendMut,
2664 start: usize,
2665 end: usize,
2666) {
2667 let mut o = start;
2668 while o + 4 <= end {
2669 let r = unsafe {
2670 dot_i8_sdot_4rows(
2671 &q[o * cols..(o + 1) * cols],
2672 &q[(o + 1) * cols..(o + 2) * cols],
2673 &q[(o + 2) * cols..(o + 3) * cols],
2674 &q[(o + 3) * cols..(o + 4) * cols],
2675 &act.xq,
2676 )
2677 };
2678 for k in 0..4 {
2679 let mut acc = r[k] as f32 * act.sx;
2680 for &(j, xv) in &act.outliers {
2681 acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
2682 }
2683 unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
2685 }
2686 o += 4;
2687 }
2688 while o < end {
2689 let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
2690 unsafe { *out_addr.at(o) = v };
2691 o += 1;
2692 }
2693}
2694
2695#[cfg(target_arch = "aarch64")]
2698#[allow(clippy::too_many_arguments)]
2699fn q8_range2_sdot(
2700 q: &[u8],
2701 row_scale: &[f32],
2702 a1: &SplitAct,
2703 a2: &SplitAct,
2704 cols: usize,
2705 p1: SendMut,
2706 p2: SendMut,
2707 start: usize,
2708 end: usize,
2709) {
2710 for o in start..end {
2711 let row = &q[o * cols..(o + 1) * cols];
2712 unsafe {
2714 *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
2715 *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
2716 }
2717 }
2718}
2719
2720#[allow(clippy::too_many_arguments)]
2722fn q8_range2_f32(
2723 q: &[u8],
2724 row_scale: &[f32],
2725 x1: &[f32],
2726 x2: &[f32],
2727 cols: usize,
2728 p1: SendMut,
2729 p2: SendMut,
2730 start: usize,
2731 end: usize,
2732) {
2733 for o in start..end {
2734 let row = &q[o * cols..(o + 1) * cols];
2735 unsafe {
2737 *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
2738 *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
2739 }
2740 }
2741}
2742
2743fn q8_range_f32(
2745 q: &[u8],
2746 row_scale: &[f32],
2747 xs: &[f32],
2748 cols: usize,
2749 out_addr: SendMut,
2750 start: usize,
2751 end: usize,
2752) {
2753 for o in start..end {
2754 let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
2755 unsafe { *out_addr.at(o) = v };
2757 }
2758}
2759
2760#[cfg(target_arch = "aarch64")]
2763#[inline]
2764fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
2765 let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
2766 for &(j, xv) in &act.outliers {
2767 acc += (row[j] as i8) as f32 * xv;
2768 }
2769 acc
2770}
2771
2772#[cfg(target_arch = "aarch64")]
2780#[target_feature(enable = "neon,dotprod")]
2781unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
2782 unsafe {
2785 use core::arch::aarch64::*;
2786 use core::arch::asm;
2787 let lomask = vdupq_n_u8(0x0F);
2788 let eight = vdupq_n_s8(8);
2789 let mut acc = 0f32;
2790 for gi in 0..gpr {
2791 let g = g0 + gi;
2792 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
2793 let b = vld1q_u8(packed.as_ptr().add(g * 16));
2794 let lo = vandq_u8(b, lomask);
2795 let hi = vshrq_n_u8::<4>(b);
2796 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
2797 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
2798 let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
2799 let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
2800 let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
2801 asm!(
2802 "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
2803 "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
2804 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
2805 e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
2806 options(pure, nomem, nostack),
2807 );
2808 acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
2809 }
2810 acc
2811 }
2812}
2813
2814#[cfg(target_arch = "aarch64")]
2819#[target_feature(enable = "neon,dotprod")]
2820unsafe fn dot_q4_row_sdot2(
2821 packed: &[u8],
2822 scales: &[u8],
2823 g0: usize,
2824 gpr: usize,
2825 xq1: &[i8],
2826 xq2: &[i8],
2827) -> (f32, f32) {
2828 unsafe {
2831 use core::arch::aarch64::*;
2832 use core::arch::asm;
2833 let lomask = vdupq_n_u8(0x0F);
2834 let eight = vdupq_n_s8(8);
2835 let (mut acc1, mut acc2) = (0f32, 0f32);
2836 for gi in 0..gpr {
2837 let g = g0 + gi;
2838 let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
2839 let b = vld1q_u8(packed.as_ptr().add(g * 16));
2840 let lo = vandq_u8(b, lomask);
2841 let hi = vshrq_n_u8::<4>(b);
2842 let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
2843 let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
2844 let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
2845 let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
2846 let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
2847 let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
2848 let (mut a0, mut a1, mut b0, mut b1) =
2849 (vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0));
2850 asm!(
2851 "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
2852 "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
2853 "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
2854 "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
2855 a0 = inout(vreg) a0, a1 = inout(vreg) a1,
2856 b0 = inout(vreg) b0, b1 = inout(vreg) b1,
2857 e0 = in(vreg) e0, e1 = in(vreg) e1,
2858 x10 = in(vreg) x10, x11 = in(vreg) x11,
2859 x20 = in(vreg) x20, x21 = in(vreg) x21,
2860 options(pure, nomem, nostack),
2861 );
2862 acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
2863 acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
2864 }
2865 (acc1, acc2)
2866 }
2867}
2868
2869#[inline]
2874pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
2875 #[cfg(target_arch = "aarch64")]
2876 unsafe {
2877 return axpy_i8_f32_neon(acc, row, w);
2878 }
2879 #[cfg(target_arch = "x86_64")]
2880 if avx2_enabled() {
2881 return unsafe { axpy_i8_f32_avx2(acc, row, w) };
2882 }
2883 #[allow(unreachable_code)]
2884 {
2885 for (a, &b) in acc.iter_mut().zip(row) {
2886 *a += w * b as f32;
2887 }
2888 }
2889}
2890
2891#[cfg(target_arch = "x86_64")]
2893#[target_feature(enable = "avx2,fma")]
2894unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
2895 unsafe {
2897 use core::arch::x86_64::*;
2898 let n = acc.len().min(row.len());
2899 let ap = acc.as_mut_ptr();
2900 let rp = row.as_ptr();
2901 let wv = _mm256_set1_ps(w);
2902 let mut j = 0usize;
2903 while j + 16 <= n {
2904 let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
2905 let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
2906 let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
2907 let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
2908 let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
2909 _mm256_storeu_ps(ap.add(j), v0);
2910 _mm256_storeu_ps(ap.add(j + 8), v1);
2911 j += 16;
2912 }
2913 while j < n {
2914 *ap.add(j) += w * (*rp.add(j)) as f32;
2915 j += 1;
2916 }
2917 }
2918}
2919
2920#[cfg(target_arch = "aarch64")]
2921#[target_feature(enable = "neon")]
2922unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
2923 unsafe {
2925 use core::arch::aarch64::*;
2926 let n = acc.len().min(row.len());
2927 let ap = acc.as_mut_ptr();
2928 let rp = row.as_ptr();
2929 let wv = vdupq_n_f32(w);
2930 let mut j = 0usize;
2931 while j + 16 <= n {
2932 let rb = vld1q_s8(rp.add(j));
2933 let lo = vmovl_s8(vget_low_s8(rb));
2934 let hi = vmovl_s8(vget_high_s8(rb));
2935 for (off, half) in [(0, lo), (8, hi)] {
2936 let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
2937 let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
2938 let o = j + off;
2939 vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
2940 vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
2941 }
2942 j += 16;
2943 }
2944 while j < n {
2945 *ap.add(j) += w * (*rp.add(j)) as f32;
2946 j += 1;
2947 }
2948}
2949}
2950
2951#[inline]
2954pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
2955 #[cfg(target_arch = "aarch64")]
2956 unsafe {
2957 return dot_i8_f32_neon(w, x);
2958 }
2959 #[cfg(target_arch = "x86_64")]
2960 if avx2_enabled() {
2961 return unsafe { dot_i8_f32_avx2(w, x) };
2962 }
2963 #[allow(unreachable_code)]
2964 {
2965 let mut sum = 0.0f32;
2966 for (j, &b) in w.iter().enumerate() {
2967 sum += (b as i8) as f32 * x[j];
2968 }
2969 sum
2970 }
2971}
2972
2973#[inline]
2977fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
2978 #[cfg(target_arch = "aarch64")]
2979 unsafe {
2980 return dot_i8_col_f32_neon(w, x, col);
2981 }
2982 #[allow(unreachable_code)]
2983 {
2984 let mut sum = 0.0f32;
2985 for (j, &b) in w.iter().enumerate() {
2986 sum += (b as i8) as f32 * x[j] * col[j];
2987 }
2988 sum
2989 }
2990}
2991
2992#[cfg(target_arch = "aarch64")]
2993#[target_feature(enable = "neon")]
2994unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
2995 unsafe {
2997 use core::arch::aarch64::*;
2998 let n = x.len();
2999 let wp = w.as_ptr() as *const i8;
3000 let xp = x.as_ptr();
3001 let cp = col.as_ptr();
3002 let (mut a0, mut a1, mut a2, mut a3) =
3003 (vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3004 let mut j = 0usize;
3005 while j + 16 <= n {
3006 let wb = vld1q_s8(wp.add(j));
3007 let lo = vmovl_s8(vget_low_s8(wb));
3008 let hi = vmovl_s8(vget_high_s8(wb));
3009 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
3010 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
3011 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
3012 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
3013 a0 = vfmaq_f32(a0, w0, vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))));
3014 a1 = vfmaq_f32(a1, w1, vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))));
3015 a2 = vfmaq_f32(a2, w2, vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))));
3016 a3 = vfmaq_f32(a3, w3, vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))));
3017 j += 16;
3018 }
3019 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
3020 while j < n {
3021 sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
3022 j += 1;
3023 }
3024 sum
3025}
3026}
3027
3028#[cfg(target_arch = "aarch64")]
3029#[target_feature(enable = "neon")]
3030unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
3031 unsafe {
3033 use core::arch::aarch64::*;
3034 let n = x.len();
3035 let wp = w.as_ptr() as *const i8;
3036 let xp = x.as_ptr();
3037 let (mut a0, mut a1, mut a2, mut a3) =
3038 (vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3039 let mut j = 0usize;
3040 while j + 16 <= n {
3041 let wb = vld1q_s8(wp.add(j));
3042 let lo = vmovl_s8(vget_low_s8(wb));
3043 let hi = vmovl_s8(vget_high_s8(wb));
3044 let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
3045 let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
3046 let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
3047 let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
3048 a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
3049 a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
3050 a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
3051 a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
3052 j += 16;
3053 }
3054 let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
3055 while j < n {
3056 sum += (*wp.add(j)) as f32 * *xp.add(j);
3057 j += 1;
3058 }
3059 sum
3060}
3061}
3062
3063#[allow(clippy::too_many_arguments)]
3064fn qmatvec(
3065 q: &[u8],
3066 row_scale: &[f32],
3067 xs: &[f32],
3068 rows: usize,
3069 cols: usize,
3070 out: &mut [f32],
3071 pool: Option<&Pool>,
3072) {
3073 debug_assert_eq!(out.len(), rows);
3074
3075 #[cfg(target_arch = "aarch64")]
3076 if sdot_enabled() {
3077 let act = split_act(xs);
3078 let out_addr = SendMut(out.as_mut_ptr());
3079 let run_range =
3080 |start: usize, end: usize| q8_range_sdot(q, row_scale, &act, cols, out_addr, start, end);
3081 match pool {
3082 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
3083 _ => run_range(0, rows),
3084 }
3085 return;
3086 }
3087 #[cfg(target_arch = "x86_64")]
3090 if avx2_a8w8_enabled() {
3091 let act = split_act(xs);
3092 let out_addr = SendMut(out.as_mut_ptr());
3093 let run_range =
3094 |start: usize, end: usize| q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end);
3095 match pool {
3096 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
3097 _ => run_range(0, rows),
3098 }
3099 return;
3100 }
3101
3102 let row_dot = |o: usize| -> f32 { dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o] };
3103 match pool {
3104 Some(pool) if rows >= 256 => {
3105 let out_addr = SendMut(out.as_mut_ptr());
3106 pool.run(&move |widx, n| {
3107 let chunk = rows.div_ceil(n);
3108 let start = widx * chunk;
3109 let end = (start + chunk).min(rows);
3110 for o in start..end {
3111 unsafe { *out_addr.at(o) = row_dot(o) };
3113 }
3114 });
3115 }
3116 _ => {
3117 for (o, dst) in out.iter_mut().enumerate() {
3118 *dst = row_dot(o);
3119 }
3120 }
3121 }
3122}
3123
3124#[allow(clippy::too_many_arguments)]
3125fn qmatvec2(
3126 q: &[u8],
3127 row_scale: &[f32],
3128 x1: &[f32],
3129 x2: &[f32],
3130 rows: usize,
3131 cols: usize,
3132 o1: &mut [f32],
3133 o2: &mut [f32],
3134 pool: Option<&Pool>,
3135) {
3136 #[cfg(target_arch = "aarch64")]
3137 if sdot_enabled() {
3138 let a1s = split_act(x1);
3139 let a2s = split_act(x2);
3140 let p1 = SendMut(o1.as_mut_ptr());
3141 let p2 = SendMut(o2.as_mut_ptr());
3142 let run_range = |start: usize, end: usize| {
3143 q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
3144 };
3145 match pool {
3146 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
3147 _ => run_range(0, rows),
3148 }
3149 return;
3150 }
3151 #[cfg(target_arch = "x86_64")]
3152 if avx2_a8w8_enabled() {
3153 let a1s = split_act(x1);
3154 let a2s = split_act(x2);
3155 let p1 = SendMut(o1.as_mut_ptr());
3156 let p2 = SendMut(o2.as_mut_ptr());
3157 let run_range = |start: usize, end: usize| {
3158 q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
3159 };
3160 match pool {
3161 Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
3162 _ => run_range(0, rows),
3163 }
3164 return;
3165 }
3166
3167 let row_dots = |o: usize| -> (f32, f32) {
3168 let row = &q[o * cols..(o + 1) * cols];
3169 (dot_i8_f32(row, x1) * row_scale[o], dot_i8_f32(row, x2) * row_scale[o])
3170 };
3171 match pool {
3172 Some(pool) if rows >= 256 => {
3173 let p1 = SendMut(o1.as_mut_ptr());
3174 let p2 = SendMut(o2.as_mut_ptr());
3175 pool.run(&move |widx, n| {
3176 let chunk = rows.div_ceil(n);
3177 let start = widx * chunk;
3178 let end = (start + chunk).min(rows);
3179 for o in start..end {
3180 let (s1, s2) = row_dots(o);
3181 unsafe {
3183 *p1.at(o) = s1;
3184 *p2.at(o) = s2;
3185 }
3186 }
3187 });
3188 }
3189 _ => {
3190 for o in 0..rows {
3191 let (s1, s2) = row_dots(o);
3192 o1[o] = s1;
3193 o2[o] = s2;
3194 }
3195 }
3196 }
3197}
3198
3199#[derive(Clone, Copy)]
3200struct SendMut(*mut f32);
3201unsafe impl Send for SendMut {}
3202unsafe impl Sync for SendMut {}
3203
3204impl SendMut {
3205 #[inline]
3206 fn at(self, i: usize) -> *mut f32 {
3207 unsafe { self.0.add(i) }
3208 }
3209}
3210
3211#[cfg(test)]
3212mod tests {
3213 use super::*;
3214
3215 #[test]
3216 fn f32_matvec_matches_matvec_rows_bitexact() {
3217 let (rows, cols) = (300, 40);
3218 let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
3219 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
3220 let qt = QTensor::from_f32(w.clone(), rows, cols);
3221
3222 let mut a = vec![0.0f32; rows];
3223 matvec_rows(None, &w, &x, &mut a);
3224 let mut b = vec![0.0f32; rows];
3225 qt.matvec(&x, &mut b, None);
3226 assert_eq!(a, b);
3227 }
3228
3229 #[test]
3230 fn sdot_kernel_exact_on_grid() {
3231 eprintln!("sdot_enabled = {}", sdot_enabled());
3236 let (rows, cols) = (9, 80); let w: Vec<u8> = (0..rows * cols)
3238 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
3239 .collect();
3240 let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
3241 let x: Vec<f32> = (0..cols)
3242 .map(|i| match i % 3 {
3243 0 => 1.0,
3244 1 => -1.0,
3245 _ => 0.0,
3246 })
3247 .collect();
3248 let mut a = vec![0.0f32; rows];
3249 qmatvec(&w, &scales, &x, rows, cols, &mut a, None);
3250 for o in 0..rows {
3251 let mut acc = 0.0f32;
3252 for j in 0..cols {
3253 acc += (w[o * cols + j] as i8) as f32 * x[j];
3254 }
3255 let expect = acc * scales[o];
3256 assert!(
3257 (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
3258 "row {o}: {} vs {expect}",
3259 a[o]
3260 );
3261 }
3262 }
3263
3264 #[test]
3265 fn sdot_a8w8_noise_is_bounded() {
3266 let (rows, cols) = (16, 512);
3270 let w: Vec<u8> = (0..rows * cols)
3271 .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
3272 .collect();
3273 let scales = vec![0.01f32; rows];
3274 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
3275 let mut a = vec![0.0f32; rows];
3276 qmatvec(&w, &scales, &x, rows, cols, &mut a, None);
3277 let (mut num, mut den) = (0f64, 0f64);
3278 for o in 0..rows {
3279 let mut acc = 0.0f32;
3280 for j in 0..cols {
3281 acc += (w[o * cols + j] as i8) as f32 * x[j];
3282 }
3283 let expect = acc * scales[o];
3284 num += ((a[o] - expect) as f64).powi(2);
3285 den += (expect as f64).powi(2);
3286 }
3287 let rel = (num / den.max(1e-12)).sqrt();
3288 assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
3289 }
3290
3291 #[test]
3292 fn i8_dot_neon_matches_scalar() {
3293 let n = 100;
3294 let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
3295 let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
3296 let mut scalar = 0.0f32;
3297 for j in 0..n {
3298 scalar += (w[j] as i8) as f32 * x[j];
3299 }
3300 let fast = dot_i8_f32(&w, &x);
3301 assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
3302 }
3303
3304 #[test]
3306 fn vbitmatvec_matches_full_dequant() {
3307 let (rows, cols) = (6, 64);
3308 let ng = cols / GROUP_SIZE;
3309 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
3311 let mut bytes = bits.clone();
3312 for g in 0..rows * ng {
3313 let s = 0.02 + 0.001 * g as f32;
3314 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
3315 }
3316 for r in 0..rows {
3317 let b = bits[r] as usize;
3318 let (mut acc, mut nb) = (0u64, 0usize);
3319 let mut rowbytes = Vec::new();
3320 for i in 0..cols {
3321 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
3322 acc = (acc << b) | v;
3323 nb += b;
3324 while nb >= 8 {
3325 nb -= 8;
3326 rowbytes.push(((acc >> nb) & 0xFF) as u8);
3327 }
3328 }
3329 if nb > 0 {
3330 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
3331 }
3332 bytes.extend_from_slice(&rowbytes);
3333 }
3334 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
3335
3336 let mut reference = vec![0f32; rows * cols];
3337 cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
3338 let mut expect = vec![0f32; rows];
3339 for r in 0..rows {
3340 expect[r] = reference[r * cols..(r + 1) * cols]
3341 .iter()
3342 .zip(&x)
3343 .map(|(w, xv)| w * xv)
3344 .sum();
3345 }
3346 let mut got = vec![0f32; rows];
3347 let offsets = vbit_row_offsets(&bytes, rows, cols);
3348 vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
3349 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
3353 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
3354 for r in 0..rows {
3355 assert!(
3356 (got[r] - expect[r]).abs() < tol * scale,
3357 "row {r}: {} vs {}",
3358 got[r],
3359 expect[r]
3360 );
3361 }
3362 }
3363
3364 #[test]
3367 fn q4matvec_matches_full_dequant() {
3368 let (rows, cols) = (8, 64);
3369 let groups = rows * cols / GROUP_SIZE;
3370 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
3372 for i in 0..groups * 16 {
3373 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
3374 }
3375 for g in 0..groups {
3376 let s = 0.01 + 0.003 * g as f32;
3377 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
3378 }
3379 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
3380
3381 let mut reference = vec![0.0f32; rows * cols];
3382 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
3383 let mut expect = vec![0.0f32; rows];
3384 for r in 0..rows {
3385 expect[r] = reference[r * cols..(r + 1) * cols]
3386 .iter()
3387 .zip(&x)
3388 .map(|(w, xv)| w * xv)
3389 .sum();
3390 }
3391
3392 let mut got = vec![0.0f32; rows];
3393 q4matvec(&bytes, &x, rows, cols, &mut got, None);
3394 let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
3398 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
3399 for r in 0..rows {
3400 assert!(
3401 (got[r] - expect[r]).abs() < tol * scale,
3402 "row {r}: {} vs {}",
3403 got[r],
3404 expect[r]
3405 );
3406 }
3407 }
3408
3409 #[test]
3412 fn vbitmatvec2_equals_two_singles() {
3413 let (rows, cols) = (6, 64);
3414 let ng = cols / GROUP_SIZE;
3415 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
3416 let mut bytes = bits.clone();
3417 for g in 0..rows * ng {
3418 let s = 0.02 + 0.001 * g as f32;
3419 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
3420 }
3421 for r in 0..rows {
3422 let b = bits[r] as usize;
3423 let (mut acc, mut nb) = (0u64, 0usize);
3424 let mut rowbytes = Vec::new();
3425 for i in 0..cols {
3426 let v = ((i * 7 + r * 13) % (1 << b)) as u64;
3427 acc = (acc << b) | v;
3428 nb += b;
3429 while nb >= 8 {
3430 nb -= 8;
3431 rowbytes.push(((acc >> nb) & 0xFF) as u8);
3432 }
3433 }
3434 if nb > 0 {
3435 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
3436 }
3437 bytes.extend_from_slice(&rowbytes);
3438 }
3439 let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
3440 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
3441 let offsets = vbit_row_offsets(&bytes, rows, cols);
3442
3443 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
3444 vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
3445 vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
3446 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
3447 vbitmatvec2(&bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
3448 assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
3449 assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
3450 }
3451
3452 #[test]
3454 fn q4matvec2_equals_two_singles() {
3455 let (rows, cols) = (8, 128);
3456 let groups = rows * cols / GROUP_SIZE;
3457 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
3458 for i in 0..groups * 16 {
3459 bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
3460 }
3461 for g in 0..groups {
3462 let s = 0.01 + 0.003 * g as f32;
3463 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
3464 }
3465 let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
3468 x1[9] = 250.0;
3469 let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
3470
3471 let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
3472 q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
3473 q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
3474 let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
3475 q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
3476 assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
3477 assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
3478 }
3479
3480 #[test]
3483 fn matvec_many_equals_separate_matvecs() {
3484 use crate::pool::Pool;
3485 let (r1, r2, cols) = (300, 200, 64);
3486 let mk = |salt: usize, rows: usize| {
3487 QTensor::from_f32(
3488 (0..rows * cols).map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5).collect(),
3489 rows,
3490 cols,
3491 )
3492 };
3493 let (a, b) = (mk(1, r1), mk(5, r2));
3494 let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
3495 let pool = Pool::new(3);
3496
3497 let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
3498 a.matvec(&x, &mut ea, Some(&pool));
3499 b.matvec(&x, &mut eb, Some(&pool));
3500 let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
3501 QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
3502 assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
3503 assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
3504 }
3505
3506 #[test]
3509 fn batched_matmat_equals_per_position_matvec() {
3510 let (rows, cols, b) = (8, 64, 5);
3511 let groups = rows * cols / GROUP_SIZE;
3513 let mut q4 = Vec::new();
3514 for i in 0..groups * 16 {
3515 q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
3516 }
3517 for g in 0..groups {
3518 q4.extend_from_slice(
3519 &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
3520 );
3521 }
3522 let ng = cols / GROUP_SIZE;
3524 let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
3525 let mut vb = bits.clone();
3526 for g in 0..rows * ng {
3527 vb.extend_from_slice(
3528 &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
3529 );
3530 }
3531 for r in 0..rows {
3532 let bw = bits[r] as usize;
3533 let (mut acc, mut nb) = (0u64, 0usize);
3534 let mut rowbytes = Vec::new();
3535 for i in 0..cols {
3536 let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
3537 acc = (acc << bw) | v;
3538 nb += bw;
3539 while nb >= 8 {
3540 nb -= 8;
3541 rowbytes.push(((acc >> nb) & 0xFF) as u8);
3542 }
3543 }
3544 if nb > 0 {
3545 rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
3546 }
3547 vb.extend_from_slice(&rowbytes);
3548 }
3549 let offsets = vbit_row_offsets(&vb, rows, cols);
3550
3551 let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
3552
3553 let mut got = vec![0f32; b * rows];
3555 q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
3556 for bi in 0..b {
3557 let mut expect = vec![0f32; rows];
3558 q4matvec(&q4, &xs[bi * cols..(bi + 1) * cols], rows, cols, &mut expect, None);
3559 assert_eq!(&got[bi * rows..(bi + 1) * rows], &expect[..], "q4 batch pos {bi}");
3560 }
3561
3562 let mut got = vec![0f32; b * rows];
3564 vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
3565 for bi in 0..b {
3566 let mut expect = vec![0f32; rows];
3567 vbitmatvec(
3568 &vb, &offsets, &xs[bi * cols..(bi + 1) * cols], rows, cols, &mut expect, None,
3569 );
3570 assert_eq!(&got[bi * rows..(bi + 1) * rows], &expect[..], "vbit batch pos {bi}");
3571 }
3572 }
3573
3574 #[test]
3581 fn q4matvec_sdot_outlier_exact() {
3582 let (rows, cols) = (4, 128);
3583 let groups = rows * cols / GROUP_SIZE;
3584 let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
3585 for i in 0..groups * 16 {
3586 bytes.push(((i * 11 + 5) % 256) as u8);
3587 }
3588 for g in 0..groups {
3589 let s = 0.02 + 0.002 * g as f32;
3590 bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
3591 }
3592 let mut x: Vec<f32> = (0..cols)
3593 .map(|i| match i % 3 {
3594 0 => 1.0,
3595 1 => -1.0,
3596 _ => 0.0,
3597 })
3598 .collect();
3599 x[17] = 300.0; let mut reference = vec![0.0f32; rows * cols];
3602 cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
3603 let mut expect = vec![0.0f32; rows];
3604 for r in 0..rows {
3605 expect[r] = reference[r * cols..(r + 1) * cols]
3606 .iter()
3607 .zip(&x)
3608 .map(|(w, xv)| w * xv)
3609 .sum();
3610 }
3611 let mut got = vec![0.0f32; rows];
3612 q4matvec(&bytes, &x, rows, cols, &mut got, None);
3613 let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
3614 for r in 0..rows {
3615 assert!(
3616 (got[r] - expect[r]).abs() < 2e-3 * scale,
3617 "row {r}: {} vs {} (outlier term must be exact)",
3618 got[r],
3619 expect[r]
3620 );
3621 }
3622 }
3623}