use std::arch::aarch64::*;
use std::arch::asm;
use super::{
InstructionTier, KernelArm, KernelBackendTag, KernelFeatures, KernelTable,
MAX_DOT_I8_DIMENSION, scalar,
};
const I8_LANES: usize = 16;
const I8_UNROLL: usize = 4;
const I8_BLOCK: usize = I8_LANES * I8_UNROLL;
const F32_LANES: usize = 4;
const FLOAT_UNROLL: usize = 4; const FLOAT_ACCUMULATORS: usize = 4; const FLOAT_BLOCK: usize = F32_LANES * FLOAT_UNROLL;
const _: [(); FLOAT_UNROLL] = [(); FLOAT_ACCUMULATORS];
const HAMMING_LANES: usize = 16;
const HAMMING_UNROLL: usize = 4; const HAMMING_BLOCK: usize = HAMMING_LANES * HAMMING_UNROLL;
const PACKED_LANES: usize = 16;
const BIT4_DIMENSIONS_PER_BLOCK: usize = PACKED_LANES * 2;
const BIT4_RAW_DOTPROD_BLOCK: usize = BIT4_DIMENSIONS_PER_BLOCK * 4;
const BIT4_PREPARED_DOTPROD_BLOCK: usize = BIT4_DIMENSIONS_PER_BLOCK * 4;
const WIDE_STREAM_LANES: usize = 16;
const WIDE_STREAM_UNROLL: usize = 4;
const WIDE_STREAM_BLOCK: usize = WIDE_STREAM_LANES * WIDE_STREAM_UNROLL;
pub(super) fn table(features: KernelFeatures) -> KernelTable {
if features.dotprod {
dotprod_table(features)
} else {
widen_table(features)
}
}
pub(super) fn dotprod_table(features: KernelFeatures) -> KernelTable {
KernelTable {
backend: KernelBackendTag::NeonDotprodU4,
arm: KernelArm::Neon,
tier: InstructionTier::NeonDotprod,
dot_i8: dot_i8_dotprod,
dot_f32,
dot_f16: f16_kernel(features),
hamming_u1,
dot_i8_batch: dot_i8_batch_dotprod,
hamming_u1_batch,
dot_bit4: dot_bit4_dotprod,
dot_bit4_prepared: dot_bit4_prepared_dotprod,
dot_bit4_batch: dot_bit4_batch_dotprod,
score_bit4_prepared_batch: score_bit4_prepared_batch_dotprod,
score_bit4_ptrs: score_bit4_ptrs_dotprod,
}
}
#[cfg(any(test, feature = "benchmark-kernels"))]
pub(super) fn i8mm_table(features: KernelFeatures) -> KernelTable {
KernelTable {
backend: KernelBackendTag::NeonI8mm,
arm: KernelArm::Neon,
tier: InstructionTier::NeonI8mmReserved,
dot_i8: dot_i8_dotprod,
dot_f32,
dot_f16: f16_kernel(features),
hamming_u1,
dot_i8_batch: dot_i8_batch_i8mm,
hamming_u1_batch,
dot_bit4: dot_bit4_dotprod,
dot_bit4_prepared: dot_bit4_prepared_dotprod,
dot_bit4_batch: dot_bit4_batch_dotprod,
score_bit4_prepared_batch: score_bit4_prepared_batch_dotprod,
score_bit4_ptrs: score_bit4_ptrs_dotprod,
}
}
#[cfg(any(test, feature = "benchmark-kernels"))]
pub(super) fn dotprod_u2_table(features: KernelFeatures) -> KernelTable {
dotprod_shape_table(
features,
KernelBackendTag::NeonDotprodU2,
dot_i8_dotprod_u2,
dot_i8_batch_dotprod_u2,
)
}
#[cfg(any(test, feature = "benchmark-kernels"))]
pub(super) fn dotprod_u6_table(features: KernelFeatures) -> KernelTable {
dotprod_shape_table(
features,
KernelBackendTag::NeonDotprodU6,
dot_i8_dotprod_u6,
dot_i8_batch_dotprod_u6,
)
}
#[cfg(any(test, feature = "benchmark-kernels"))]
pub(super) fn dotprod_u8_table(features: KernelFeatures) -> KernelTable {
dotprod_shape_table(
features,
KernelBackendTag::NeonDotprodU8,
dot_i8_dotprod_u8,
dot_i8_batch_dotprod_u8,
)
}
#[cfg(any(test, feature = "benchmark-kernels"))]
pub(super) fn dotprod_prefetch_table(features: KernelFeatures) -> KernelTable {
dotprod_shape_table(
features,
KernelBackendTag::NeonDotprodU4Prefetch,
dot_i8_dotprod,
dot_i8_batch_dotprod_prefetch,
)
}
#[cfg(any(test, feature = "benchmark-kernels"))]
fn dotprod_shape_table(
features: KernelFeatures,
backend: KernelBackendTag,
dot_i8: super::DotI8Fn,
dot_i8_batch: super::DotI8BatchFn,
) -> KernelTable {
KernelTable {
backend,
arm: KernelArm::Neon,
tier: InstructionTier::NeonDotprod,
dot_i8,
dot_f32,
dot_f16: f16_kernel(features),
hamming_u1,
dot_i8_batch,
hamming_u1_batch,
dot_bit4: dot_bit4_dotprod,
dot_bit4_prepared: dot_bit4_prepared_dotprod,
dot_bit4_batch: dot_bit4_batch_dotprod,
score_bit4_prepared_batch: score_bit4_prepared_batch_dotprod,
score_bit4_ptrs: score_bit4_ptrs_dotprod,
}
}
pub(super) fn widen_table(features: KernelFeatures) -> KernelTable {
KernelTable {
backend: KernelBackendTag::NeonWiden,
arm: KernelArm::Neon,
tier: InstructionTier::NeonWiden,
dot_i8: dot_i8_widen,
dot_f32,
dot_f16: f16_kernel(features),
hamming_u1,
dot_i8_batch: dot_i8_batch_widen,
hamming_u1_batch,
dot_bit4: dot_bit4_widen,
dot_bit4_prepared: dot_bit4_prepared_widen,
dot_bit4_batch: dot_bit4_batch_widen,
score_bit4_prepared_batch: scalar::score_bit4_prepared_batch,
score_bit4_ptrs: scalar::score_bit4_ptrs,
}
}
#[inline]
fn dot_i8_dotprod(a: &[i8], b: &[i8]) -> i32 {
debug_assert_eq!(a.len(), b.len(), "kernel lengths must be pre-validated");
debug_assert!(
a.len() <= MAX_DOT_I8_DIMENSION,
"i8 dimension must be pre-validated"
);
unsafe { dot_i8_dotprod_inner(a, b) }
}
#[target_feature(enable = "dotprod")]
unsafe fn dot_i8_dotprod_inner(a: &[i8], b: &[i8]) -> i32 {
unsafe {
let mut acc0 = vdupq_n_s32(0);
let mut acc1 = vdupq_n_s32(0);
let mut acc2 = vdupq_n_s32(0);
let mut acc3 = vdupq_n_s32(0);
let processed = a.len() / I8_BLOCK * I8_BLOCK;
let mut base = 0_usize;
while base < processed {
let a_ptr = a.as_ptr().wrapping_add(base);
let b_ptr = b.as_ptr().wrapping_add(base);
acc0 = dotprod_mac(acc0, vld1q_s8(a_ptr), vld1q_s8(b_ptr));
acc1 = dotprod_mac(
acc1,
vld1q_s8(a_ptr.wrapping_add(I8_LANES)),
vld1q_s8(b_ptr.wrapping_add(I8_LANES)),
);
acc2 = dotprod_mac(
acc2,
vld1q_s8(a_ptr.wrapping_add(I8_LANES * 2)),
vld1q_s8(b_ptr.wrapping_add(I8_LANES * 2)),
);
acc3 = dotprod_mac(
acc3,
vld1q_s8(a_ptr.wrapping_add(I8_LANES * 3)),
vld1q_s8(b_ptr.wrapping_add(I8_LANES * 3)),
);
base += I8_BLOCK;
}
let vectors = vaddq_s32(vaddq_s32(acc0, acc1), vaddq_s32(acc2, acc3));
let sum = vaddvq_s32(vectors);
add_i8_tail(sum, a, b, processed)
}
}
#[cfg(any(test, feature = "benchmark-kernels"))]
fn dot_i8_dotprod_u2(a: &[i8], b: &[i8]) -> i32 {
dot_i8_dotprod_shape::<2>(a, b)
}
#[cfg(any(test, feature = "benchmark-kernels"))]
fn dot_i8_dotprod_u6(a: &[i8], b: &[i8]) -> i32 {
dot_i8_dotprod_shape::<6>(a, b)
}
#[cfg(any(test, feature = "benchmark-kernels"))]
fn dot_i8_dotprod_u8(a: &[i8], b: &[i8]) -> i32 {
dot_i8_dotprod_shape::<8>(a, b)
}
#[cfg(any(test, feature = "benchmark-kernels"))]
fn dot_i8_dotprod_shape<const ACCUMULATORS: usize>(a: &[i8], b: &[i8]) -> i32 {
debug_assert_eq!(a.len(), b.len(), "kernel lengths must be pre-validated");
debug_assert!(
a.len() <= MAX_DOT_I8_DIMENSION,
"i8 dimension must be pre-validated"
);
unsafe { dot_i8_dotprod_shape_inner::<ACCUMULATORS>(a, b) }
}
#[target_feature(enable = "dotprod")]
#[cfg(any(test, feature = "benchmark-kernels"))]
unsafe fn dot_i8_dotprod_shape_inner<const ACCUMULATORS: usize>(a: &[i8], b: &[i8]) -> i32 {
debug_assert!(ACCUMULATORS > 0);
let block = I8_LANES * ACCUMULATORS;
let processed = a.len() / block * block;
let mut accumulators = [vdupq_n_s32(0); ACCUMULATORS];
let mut base = 0_usize;
while base < processed {
for (index, accumulator) in accumulators.iter_mut().enumerate() {
let offset = base + index * I8_LANES;
let left = unsafe { vld1q_s8(a.as_ptr().add(offset)) };
let right = unsafe { vld1q_s8(b.as_ptr().add(offset)) };
*accumulator = unsafe { dotprod_mac(*accumulator, left, right) };
}
base += block;
}
let vectors = accumulators
.into_iter()
.fold(vdupq_n_s32(0), |sum, accumulator| {
vaddq_s32(sum, accumulator)
});
add_i8_tail(vaddvq_s32(vectors), a, b, processed)
}
#[target_feature(enable = "dotprod")]
unsafe fn dotprod_mac(mut acc: int32x4_t, a: int8x16_t, b: int8x16_t) -> int32x4_t {
unsafe {
asm!(
"sdot {acc:v}.4s, {a:v}.16b, {b:v}.16b",
acc = inout(vreg) acc,
a = in(vreg) a,
b = in(vreg) b,
options(pure, nomem, nostack)
);
}
acc
}
#[inline]
fn dot_i8_widen(a: &[i8], b: &[i8]) -> i32 {
debug_assert_eq!(a.len(), b.len(), "kernel lengths must be pre-validated");
debug_assert!(
a.len() <= MAX_DOT_I8_DIMENSION,
"i8 dimension must be pre-validated"
);
unsafe { dot_i8_widen_inner(a, b) }
}
#[target_feature(enable = "neon")]
unsafe fn dot_i8_widen_inner(a: &[i8], b: &[i8]) -> i32 {
unsafe {
let mut acc0 = vdupq_n_s32(0);
let mut acc1 = vdupq_n_s32(0);
let mut acc2 = vdupq_n_s32(0);
let mut acc3 = vdupq_n_s32(0);
let processed = a.len() / I8_BLOCK * I8_BLOCK;
let mut base = 0_usize;
while base < processed {
let a_ptr = a.as_ptr().wrapping_add(base);
let b_ptr = b.as_ptr().wrapping_add(base);
acc0 = widen_mac(acc0, vld1q_s8(a_ptr), vld1q_s8(b_ptr));
acc1 = widen_mac(
acc1,
vld1q_s8(a_ptr.wrapping_add(I8_LANES)),
vld1q_s8(b_ptr.wrapping_add(I8_LANES)),
);
acc2 = widen_mac(
acc2,
vld1q_s8(a_ptr.wrapping_add(I8_LANES * 2)),
vld1q_s8(b_ptr.wrapping_add(I8_LANES * 2)),
);
acc3 = widen_mac(
acc3,
vld1q_s8(a_ptr.wrapping_add(I8_LANES * 3)),
vld1q_s8(b_ptr.wrapping_add(I8_LANES * 3)),
);
base += I8_BLOCK;
}
let vectors = vaddq_s32(vaddq_s32(acc0, acc1), vaddq_s32(acc2, acc3));
let sum = vaddvq_s32(vectors);
add_i8_tail(sum, a, b, processed)
}
}
#[target_feature(enable = "neon")]
unsafe fn widen_mac(acc: int32x4_t, a: int8x16_t, b: int8x16_t) -> int32x4_t {
let low_products = vmull_s8(vget_low_s8(a), vget_low_s8(b));
let high_products = vmull_s8(vget_high_s8(a), vget_high_s8(b));
vaddq_s32(
acc,
vaddq_s32(vpaddlq_s16(low_products), vpaddlq_s16(high_products)),
)
}
fn add_i8_tail(sum: i32, a: &[i8], b: &[i8], processed: usize) -> i32 {
let Some(a_tail) = a.get(processed..) else {
return sum;
};
let Some(b_tail) = b.get(processed..) else {
return sum;
};
sum + a_tail
.iter()
.zip(b_tail)
.map(|(&left, &right)| i32::from(left) * i32::from(right))
.sum::<i32>()
}
#[inline]
fn dot_bit4_dotprod(q: &[i8], codes: &[u8]) -> i32 {
packed_shape_debug_assertions(q, codes, 2);
unsafe { dot_bit4_dotprod_inner(q, codes) }
}
#[target_feature(enable = "dotprod")]
unsafe fn dot_bit4_dotprod_inner(q: &[i8], codes: &[u8]) -> i32 {
unsafe {
let processed = q.len() / BIT4_RAW_DOTPROD_BLOCK * BIT4_RAW_DOTPROD_BLOCK;
let mut acc0 = vdupq_n_s32(0);
let mut acc1 = vdupq_n_s32(0);
let mut acc2 = vdupq_n_s32(0);
let mut acc3 = vdupq_n_s32(0);
let mut acc4 = vdupq_n_s32(0);
let mut acc5 = vdupq_n_s32(0);
let mut acc6 = vdupq_n_s32(0);
let mut acc7 = vdupq_n_s32(0);
let mut base = 0_usize;
while base < processed {
let code_ptr = codes.as_ptr().wrapping_add(base / 2);
let (values0, values1) = unpack_bit4_block(code_ptr);
let (values2, values3) = unpack_bit4_block(code_ptr.wrapping_add(PACKED_LANES));
let (values4, values5) = unpack_bit4_block(code_ptr.wrapping_add(PACKED_LANES * 2));
let (values6, values7) = unpack_bit4_block(code_ptr.wrapping_add(PACKED_LANES * 3));
let query_ptr = q.as_ptr().wrapping_add(base);
acc0 = dotprod_mac(acc0, vld1q_s8(query_ptr), values0);
acc1 = dotprod_mac(acc1, vld1q_s8(query_ptr.wrapping_add(I8_LANES)), values1);
acc2 = dotprod_mac(
acc2,
vld1q_s8(query_ptr.wrapping_add(I8_LANES * 2)),
values2,
);
acc3 = dotprod_mac(
acc3,
vld1q_s8(query_ptr.wrapping_add(I8_LANES * 3)),
values3,
);
acc4 = dotprod_mac(
acc4,
vld1q_s8(query_ptr.wrapping_add(I8_LANES * 4)),
values4,
);
acc5 = dotprod_mac(
acc5,
vld1q_s8(query_ptr.wrapping_add(I8_LANES * 5)),
values5,
);
acc6 = dotprod_mac(
acc6,
vld1q_s8(query_ptr.wrapping_add(I8_LANES * 6)),
values6,
);
acc7 = dotprod_mac(
acc7,
vld1q_s8(query_ptr.wrapping_add(I8_LANES * 7)),
values7,
);
base += BIT4_RAW_DOTPROD_BLOCK;
}
let vectors = vaddq_s32(
vaddq_s32(vaddq_s32(acc0, acc1), vaddq_s32(acc2, acc3)),
vaddq_s32(vaddq_s32(acc4, acc5), vaddq_s32(acc6, acc7)),
);
packed_bit4_tail(vaddvq_s32(vectors), q, codes, processed)
}
}
#[inline]
fn dot_bit4_prepared_dotprod(q: &[i8], query_sum: i32, codes: &[u8]) -> i32 {
packed_shape_debug_assertions(q, codes, 2);
unsafe { dot_bit4_prepared_dotprod_inner(q, query_sum, codes) }
}
#[target_feature(enable = "dotprod")]
unsafe fn dot_bit4_prepared_dotprod_inner(q: &[i8], query_sum: i32, codes: &[u8]) -> i32 {
unsafe {
let processed = q.len() / BIT4_PREPARED_DOTPROD_BLOCK * BIT4_PREPARED_DOTPROD_BLOCK;
let mut acc0 = vdupq_n_s32(0);
let mut acc1 = vdupq_n_s32(0);
let mut acc2 = vdupq_n_s32(0);
let mut acc3 = vdupq_n_s32(0);
let mut acc4 = vdupq_n_s32(0);
let mut acc5 = vdupq_n_s32(0);
let mut acc6 = vdupq_n_s32(0);
let mut acc7 = vdupq_n_s32(0);
let mut base = 0_usize;
while base < processed {
let code_ptr = codes.as_ptr().wrapping_add(base / 2);
let query_ptr = q.as_ptr().wrapping_add(base);
let (values0, values1) = unpack_bit4_split_block(code_ptr);
let (values2, values3) = unpack_bit4_split_block(code_ptr.wrapping_add(PACKED_LANES));
let (values4, values5) =
unpack_bit4_split_block(code_ptr.wrapping_add(PACKED_LANES * 2));
let (values6, values7) =
unpack_bit4_split_block(code_ptr.wrapping_add(PACKED_LANES * 3));
acc0 = dotprod_mac(acc0, vld1q_s8(query_ptr), values0);
acc1 = dotprod_mac(acc1, vld1q_s8(query_ptr.wrapping_add(I8_LANES)), values1);
acc2 = dotprod_mac(
acc2,
vld1q_s8(query_ptr.wrapping_add(I8_LANES * 2)),
values2,
);
acc3 = dotprod_mac(
acc3,
vld1q_s8(query_ptr.wrapping_add(I8_LANES * 3)),
values3,
);
acc4 = dotprod_mac(
acc4,
vld1q_s8(query_ptr.wrapping_add(I8_LANES * 4)),
values4,
);
acc5 = dotprod_mac(
acc5,
vld1q_s8(query_ptr.wrapping_add(I8_LANES * 5)),
values5,
);
acc6 = dotprod_mac(
acc6,
vld1q_s8(query_ptr.wrapping_add(I8_LANES * 6)),
values6,
);
acc7 = dotprod_mac(
acc7,
vld1q_s8(query_ptr.wrapping_add(I8_LANES * 7)),
values7,
);
base += BIT4_PREPARED_DOTPROD_BLOCK;
}
let vectors = vaddq_s32(
vaddq_s32(vaddq_s32(acc0, acc1), vaddq_s32(acc2, acc3)),
vaddq_s32(vaddq_s32(acc4, acc5), vaddq_s32(acc6, acc7)),
);
packed_bit4_prepared_tail(vaddvq_s32(vectors), q, query_sum, codes, processed)
}
}
#[inline]
fn dot_bit4_widen(q: &[i8], codes: &[u8]) -> i32 {
packed_shape_debug_assertions(q, codes, 2);
unsafe { dot_bit4_widen_inner(q, codes) }
}
#[target_feature(enable = "neon")]
unsafe fn dot_bit4_widen_inner(q: &[i8], codes: &[u8]) -> i32 {
unsafe {
let processed = q.len() / BIT4_DIMENSIONS_PER_BLOCK * BIT4_DIMENSIONS_PER_BLOCK;
let mut acc = vdupq_n_s32(0);
let mut base = 0_usize;
while base < processed {
let code_ptr = codes.as_ptr().wrapping_add(base / 2);
let (values0, values1) = unpack_bit4_block(code_ptr);
let query_ptr = q.as_ptr().wrapping_add(base);
acc = widen_mac(acc, vld1q_s8(query_ptr), values0);
acc = widen_mac(acc, vld1q_s8(query_ptr.wrapping_add(I8_LANES)), values1);
base += BIT4_DIMENSIONS_PER_BLOCK;
}
packed_bit4_tail(vaddvq_s32(acc), q, codes, processed)
}
}
#[inline]
fn dot_bit4_prepared_widen(q: &[i8], query_sum: i32, codes: &[u8]) -> i32 {
packed_shape_debug_assertions(q, codes, 2);
unsafe { dot_bit4_prepared_widen_inner(q, query_sum, codes) }
}
#[target_feature(enable = "neon")]
unsafe fn dot_bit4_prepared_widen_inner(q: &[i8], query_sum: i32, codes: &[u8]) -> i32 {
unsafe {
let processed = q.len() / BIT4_DIMENSIONS_PER_BLOCK * BIT4_DIMENSIONS_PER_BLOCK;
let mut acc0 = vdupq_n_s32(0);
let mut acc1 = vdupq_n_s32(0);
let mut base = 0_usize;
while base < processed {
let code_ptr = codes.as_ptr().wrapping_add(base / 2);
let (values0, values1) = unpack_bit4_split_block(code_ptr);
let query_ptr = q.as_ptr().wrapping_add(base);
acc0 = widen_mac(acc0, vld1q_s8(query_ptr), values0);
acc1 = widen_mac(acc1, vld1q_s8(query_ptr.wrapping_add(I8_LANES)), values1);
base += BIT4_DIMENSIONS_PER_BLOCK;
}
packed_bit4_prepared_tail(
vaddvq_s32(vaddq_s32(acc0, acc1)),
q,
query_sum,
codes,
processed,
)
}
}
#[target_feature(enable = "neon")]
unsafe fn unpack_bit4_block(code_ptr: *const u8) -> (int8x16_t, int8x16_t) {
unsafe {
let packed = vld1q_u8(code_ptr);
let mask = vdupq_n_u8(0x0f);
let high = vshrq_n_u8(packed, 4);
let low = vandq_u8(packed, mask);
(
bit4_grid(vzip1q_u8(high, low)),
bit4_grid(vzip2q_u8(high, low)),
)
}
}
#[target_feature(enable = "neon")]
unsafe fn unpack_bit4_split_block(code_ptr: *const u8) -> (int8x16_t, int8x16_t) {
unsafe {
let packed = vld1q_u8(code_ptr);
let mask = vdupq_n_u8(0x0f);
(
vreinterpretq_s8_u8(vshrq_n_u8(packed, 4)),
vreinterpretq_s8_u8(vandq_u8(packed, mask)),
)
}
}
#[target_feature(enable = "neon")]
unsafe fn bit4_grid(fields: uint8x16_t) -> int8x16_t {
vreinterpretq_s8_u8(vsubq_u8(vshlq_n_u8(fields, 1), vdupq_n_u8(15)))
}
fn packed_shape_debug_assertions(q: &[i8], codes: &[u8], fields_per_byte: usize) {
debug_assert_eq!(codes.len(), q.len().div_ceil(fields_per_byte));
debug_assert!(q.len() <= MAX_DOT_I8_DIMENSION);
}
fn packed_bit4_tail(sum: i32, q: &[i8], codes: &[u8], processed: usize) -> i32 {
let Some(q_tail) = q.get(processed..) else {
return sum;
};
let Some(code_tail) = codes.get(processed / 2..) else {
return sum;
};
sum + scalar::dot_bit4(q_tail, code_tail)
}
fn packed_bit4_prepared_tail(
sum: i32,
q: &[i8],
query_sum: i32,
codes: &[u8],
processed: usize,
) -> i32 {
let Some(q_tail) = q.get(processed..) else {
return sum;
};
let Some(code_tail) = codes.get(processed / 2..) else {
return sum;
};
let unsigned_sum = sum + scalar::dot_bit4_prepared_unsigned(q_tail, code_tail);
2 * unsigned_sum - 15 * query_sum
}
#[inline]
fn dot_f32(a: &[f32], b: &[f32]) -> f32 {
debug_assert_eq!(a.len(), b.len(), "kernel lengths must be pre-validated");
unsafe { dot_f32_inner(a, b) }
}
#[target_feature(enable = "neon")]
unsafe fn dot_f32_inner(a: &[f32], b: &[f32]) -> f32 {
unsafe {
let mut acc0 = vdupq_n_f32(0.0);
let mut acc1 = vdupq_n_f32(0.0);
let mut acc2 = vdupq_n_f32(0.0);
let mut acc3 = vdupq_n_f32(0.0);
let processed = a.len() / FLOAT_BLOCK * FLOAT_BLOCK;
let mut base = 0_usize;
while base < processed {
let a_ptr = a.as_ptr().wrapping_add(base);
let b_ptr = b.as_ptr().wrapping_add(base);
acc0 = vfmaq_f32(acc0, vld1q_f32(a_ptr), vld1q_f32(b_ptr));
acc1 = vfmaq_f32(
acc1,
vld1q_f32(a_ptr.wrapping_add(F32_LANES)),
vld1q_f32(b_ptr.wrapping_add(F32_LANES)),
);
acc2 = vfmaq_f32(
acc2,
vld1q_f32(a_ptr.wrapping_add(F32_LANES * 2)),
vld1q_f32(b_ptr.wrapping_add(F32_LANES * 2)),
);
acc3 = vfmaq_f32(
acc3,
vld1q_f32(a_ptr.wrapping_add(F32_LANES * 3)),
vld1q_f32(b_ptr.wrapping_add(F32_LANES * 3)),
);
base += FLOAT_BLOCK;
}
let vectors = vaddq_f32(vaddq_f32(acc0, acc1), vaddq_f32(acc2, acc3));
add_f32_tail(vaddvq_f32(vectors), a, b, processed)
}
}
#[inline]
fn f16_kernel(features: KernelFeatures) -> super::DotF16Fn {
if features.fp16 {
dot_f16_fp16
} else {
dot_f16_fallback
}
}
#[inline]
fn dot_f16_fp16(a: &[u16], b: &[u16]) -> f32 {
debug_assert_eq!(a.len(), b.len(), "kernel lengths must be pre-validated");
unsafe { dot_f16_fp16_inner(a, b) }
}
#[target_feature(enable = "neon,fp16")]
unsafe fn dot_f16_fp16_inner(a: &[u16], b: &[u16]) -> f32 {
unsafe {
let mut acc0 = vdupq_n_f32(0.0);
let mut acc1 = vdupq_n_f32(0.0);
let mut acc2 = vdupq_n_f32(0.0);
let mut acc3 = vdupq_n_f32(0.0);
let processed = a.len() / FLOAT_BLOCK * FLOAT_BLOCK;
let mut base = 0_usize;
while base < processed {
let a_ptr = a.as_ptr().wrapping_add(base);
let b_ptr = b.as_ptr().wrapping_add(base);
acc0 = f16_mac(acc0, a_ptr, b_ptr);
acc1 = f16_mac(
acc1,
a_ptr.wrapping_add(F32_LANES),
b_ptr.wrapping_add(F32_LANES),
);
acc2 = f16_mac(
acc2,
a_ptr.wrapping_add(F32_LANES * 2),
b_ptr.wrapping_add(F32_LANES * 2),
);
acc3 = f16_mac(
acc3,
a_ptr.wrapping_add(F32_LANES * 3),
b_ptr.wrapping_add(F32_LANES * 3),
);
base += FLOAT_BLOCK;
}
let vectors = vaddq_f32(vaddq_f32(acc0, acc1), vaddq_f32(acc2, acc3));
add_f16_tail(vaddvq_f32(vectors), a, b, processed)
}
}
#[target_feature(enable = "neon,fp16")]
unsafe fn f16_mac(acc: float32x4_t, a: *const u16, b: *const u16) -> float32x4_t {
unsafe { vfmaq_f32(acc, f16x4_to_f32(vld1_u16(a)), f16x4_to_f32(vld1_u16(b))) }
}
#[target_feature(enable = "neon,fp16")]
unsafe fn f16x4_to_f32(input: uint16x4_t) -> float32x4_t {
let output: float32x4_t;
unsafe {
asm!(
"fcvtl {output:v}.4s, {input:v}.4h",
output = out(vreg) output,
input = in(vreg) input,
options(pure, nomem, nostack)
);
}
output
}
#[inline]
fn dot_f16_fallback(a: &[u16], b: &[u16]) -> f32 {
debug_assert_eq!(a.len(), b.len(), "kernel lengths must be pre-validated");
unsafe { dot_f16_fallback_inner(a, b) }
}
#[target_feature(enable = "neon")]
unsafe fn dot_f16_fallback_inner(a: &[u16], b: &[u16]) -> f32 {
unsafe {
let mut acc0 = vdupq_n_f32(0.0);
let mut acc1 = vdupq_n_f32(0.0);
let mut acc2 = vdupq_n_f32(0.0);
let mut acc3 = vdupq_n_f32(0.0);
let processed = a.len() / FLOAT_BLOCK * FLOAT_BLOCK;
let mut base = 0_usize;
while base < processed {
let a_block = convert_f16_block(a, base);
let b_block = convert_f16_block(b, base);
acc0 = vfmaq_f32(
acc0,
vld1q_f32(a_block.as_ptr()),
vld1q_f32(b_block.as_ptr()),
);
acc1 = vfmaq_f32(
acc1,
vld1q_f32(a_block.as_ptr().wrapping_add(F32_LANES)),
vld1q_f32(b_block.as_ptr().wrapping_add(F32_LANES)),
);
acc2 = vfmaq_f32(
acc2,
vld1q_f32(a_block.as_ptr().wrapping_add(F32_LANES * 2)),
vld1q_f32(b_block.as_ptr().wrapping_add(F32_LANES * 2)),
);
acc3 = vfmaq_f32(
acc3,
vld1q_f32(a_block.as_ptr().wrapping_add(F32_LANES * 3)),
vld1q_f32(b_block.as_ptr().wrapping_add(F32_LANES * 3)),
);
base += FLOAT_BLOCK;
}
let vectors = vaddq_f32(vaddq_f32(acc0, acc1), vaddq_f32(acc2, acc3));
add_f16_tail(vaddvq_f32(vectors), a, b, processed)
}
}
fn convert_f16_block(values: &[u16], base: usize) -> [f32; FLOAT_BLOCK] {
let mut converted = [0.0_f32; FLOAT_BLOCK];
let Some(block) = values.get(base..base + FLOAT_BLOCK) else {
return converted;
};
for (output, &bits) in converted.iter_mut().zip(block) {
*output = scalar::f16_to_f32(bits);
}
converted
}
fn add_f32_tail(sum: f32, a: &[f32], b: &[f32], processed: usize) -> f32 {
let Some(a_tail) = a.get(processed..) else {
return sum;
};
let Some(b_tail) = b.get(processed..) else {
return sum;
};
a_tail
.iter()
.zip(b_tail)
.fold(sum, |acc, (&left, &right)| acc + left * right)
}
fn add_f16_tail(sum: f32, a: &[u16], b: &[u16], processed: usize) -> f32 {
let Some(a_tail) = a.get(processed..) else {
return sum;
};
let Some(b_tail) = b.get(processed..) else {
return sum;
};
a_tail.iter().zip(b_tail).fold(sum, |acc, (&left, &right)| {
acc + scalar::f16_to_f32(left) * scalar::f16_to_f32(right)
})
}
pub(super) fn wide_stream_checksum(bytes: &[u8]) -> u64 {
unsafe { wide_stream_checksum_inner(bytes) }
}
#[target_feature(enable = "neon")]
unsafe fn wide_stream_checksum_inner(bytes: &[u8]) -> u64 {
unsafe {
let mut acc0 = vdupq_n_u64(0);
let mut acc1 = vdupq_n_u64(0);
let mut acc2 = vdupq_n_u64(0);
let mut acc3 = vdupq_n_u64(0);
let processed = bytes.len() / WIDE_STREAM_BLOCK * WIDE_STREAM_BLOCK;
let mut base = 0_usize;
while base < processed {
let ptr = bytes.as_ptr().wrapping_add(base);
acc0 = vaddq_u64(acc0, vreinterpretq_u64_u8(vld1q_u8(ptr)));
acc1 = vaddq_u64(
acc1,
vreinterpretq_u64_u8(vld1q_u8(ptr.wrapping_add(WIDE_STREAM_LANES))),
);
acc2 = vaddq_u64(
acc2,
vreinterpretq_u64_u8(vld1q_u8(ptr.wrapping_add(WIDE_STREAM_LANES * 2))),
);
acc3 = vaddq_u64(
acc3,
vreinterpretq_u64_u8(vld1q_u8(ptr.wrapping_add(WIDE_STREAM_LANES * 3))),
);
base += WIDE_STREAM_BLOCK;
}
let vectors = vaddq_u64(vaddq_u64(acc0, acc1), vaddq_u64(acc2, acc3));
bytes.get(processed..).map_or(vaddvq_u64(vectors), |tail| {
tail.iter().fold(vaddvq_u64(vectors), |checksum, &byte| {
checksum.wrapping_add(u64::from(byte))
})
})
}
}
#[inline]
fn hamming_u1(a: &[u8], b: &[u8]) -> u32 {
debug_assert_eq!(a.len(), b.len(), "kernel lengths must be pre-validated");
unsafe { hamming_u1_inner(a, b) }
}
#[target_feature(enable = "neon")]
#[inline]
unsafe fn hamming_u1_inner(a: &[u8], b: &[u8]) -> u32 {
unsafe {
let block_processed = a.len() / HAMMING_BLOCK * HAMMING_BLOCK;
let vector_processed = a.len() / HAMMING_LANES * HAMMING_LANES;
let mut sum = 0_u32;
let mut base = 0_usize;
while base < block_processed {
let a_ptr = a.as_ptr().wrapping_add(base);
let b_ptr = b.as_ptr().wrapping_add(base);
let count0 = vcntq_u8(veorq_u8(vld1q_u8(a_ptr), vld1q_u8(b_ptr)));
let count1 = vcntq_u8(veorq_u8(
vld1q_u8(a_ptr.wrapping_add(HAMMING_LANES)),
vld1q_u8(b_ptr.wrapping_add(HAMMING_LANES)),
));
let count2 = vcntq_u8(veorq_u8(
vld1q_u8(a_ptr.wrapping_add(HAMMING_LANES * 2)),
vld1q_u8(b_ptr.wrapping_add(HAMMING_LANES * 2)),
));
let count3 = vcntq_u8(veorq_u8(
vld1q_u8(a_ptr.wrapping_add(HAMMING_LANES * 3)),
vld1q_u8(b_ptr.wrapping_add(HAMMING_LANES * 3)),
));
let counts = vaddq_u8(vaddq_u8(count0, count1), vaddq_u8(count2, count3));
sum += u32::from(vaddlvq_u8(counts));
base += HAMMING_BLOCK;
}
while base < vector_processed {
let count = vcntq_u8(veorq_u8(
vld1q_u8(a.as_ptr().wrapping_add(base)),
vld1q_u8(b.as_ptr().wrapping_add(base)),
));
sum += u32::from(vaddvq_u8(count));
base += HAMMING_LANES;
}
if let (Some(a_tail), Some(b_tail)) = (a.get(vector_processed..), b.get(vector_processed..))
{
sum += scalar::hamming_u1(a_tail, b_tail);
}
sum
}
}
fn dot_i8_batch_dotprod(q: &[i8], rows: &[i8], d: usize, out: &mut [i32]) {
batch_i8(q, rows, d, out, dot_i8_dotprod);
}
fn dot_i8_batch_widen(q: &[i8], rows: &[i8], d: usize, out: &mut [i32]) {
batch_i8(q, rows, d, out, dot_i8_widen);
}
#[cfg(any(test, feature = "benchmark-kernels"))]
fn dot_i8_batch_dotprod_u2(q: &[i8], rows: &[i8], d: usize, out: &mut [i32]) {
batch_i8(q, rows, d, out, dot_i8_dotprod_u2);
}
#[cfg(any(test, feature = "benchmark-kernels"))]
fn dot_i8_batch_dotprod_u6(q: &[i8], rows: &[i8], d: usize, out: &mut [i32]) {
batch_i8(q, rows, d, out, dot_i8_dotprod_u6);
}
#[cfg(any(test, feature = "benchmark-kernels"))]
fn dot_i8_batch_dotprod_u8(q: &[i8], rows: &[i8], d: usize, out: &mut [i32]) {
batch_i8(q, rows, d, out, dot_i8_dotprod_u8);
}
#[cfg(any(test, feature = "benchmark-kernels"))]
fn dot_i8_batch_dotprod_prefetch(q: &[i8], rows: &[i8], d: usize, out: &mut [i32]) {
debug_assert_eq!(q.len(), d, "query dimension must be pre-validated");
debug_assert!(
d <= MAX_DOT_I8_DIMENSION,
"i8 dimension must be pre-validated"
);
debug_assert_eq!(
rows.len(),
d.saturating_mul(out.len()),
"batch shape must be pre-validated"
);
if d == 0 {
out.fill(0);
return;
}
unsafe { dot_i8_batch_dotprod_prefetch_inner(q, rows, d, out) };
}
#[target_feature(enable = "dotprod")]
#[cfg(any(test, feature = "benchmark-kernels"))]
unsafe fn dot_i8_batch_dotprod_prefetch_inner(q: &[i8], rows: &[i8], d: usize, out: &mut [i32]) {
let row_count = out.len();
for (row_index, (row, result)) in rows.chunks_exact(d).zip(out.iter_mut()).enumerate() {
if row_index + 1 < row_count {
let next = rows.as_ptr().wrapping_add((row_index + 1) * d);
unsafe {
asm!(
"prfm pldl1keep, [{address}]",
address = in(reg) next,
options(readonly, nostack)
);
}
}
*result = unsafe { dot_i8_dotprod_inner(q, row) };
}
}
#[cfg(any(test, feature = "benchmark-kernels"))]
fn dot_i8_batch_i8mm(q: &[i8], rows: &[i8], d: usize, out: &mut [i32]) {
debug_assert_eq!(q.len(), d, "query dimension must be pre-validated");
debug_assert!(
d <= MAX_DOT_I8_DIMENSION,
"i8 dimension must be pre-validated"
);
debug_assert_eq!(
rows.len(),
d.saturating_mul(out.len()),
"batch shape must be pre-validated"
);
if d == 0 {
out.fill(0);
return;
}
unsafe { dot_i8_batch_i8mm_inner(q, rows, d, out) };
}
#[target_feature(enable = "i8mm,dotprod")]
#[cfg(any(test, feature = "benchmark-kernels"))]
unsafe fn dot_i8_batch_i8mm_inner(q: &[i8], rows: &[i8], d: usize, out: &mut [i32]) {
let paired_rows = out.len() / 2 * 2;
let processed = d / I8_LANES * I8_LANES;
let mut row_index = 0_usize;
while row_index < paired_rows {
let row0_start = row_index * d;
let row1_start = row0_start + d;
let row0 = unsafe { rows.get_unchecked(row0_start..row0_start + d) };
let row1 = unsafe { rows.get_unchecked(row1_start..row1_start + d) };
let mut acc = vdupq_n_s32(0);
let mut base = 0_usize;
while base < processed {
let query = unsafe { vld1q_s8(q.as_ptr().add(base)) };
let left = unsafe { vld1q_s8(row0.as_ptr().add(base)) };
let right = unsafe { vld1q_s8(row1.as_ptr().add(base)) };
let query_low = vget_low_s8(query);
let query_high = vget_high_s8(query);
let duplicated_low = vcombine_s8(query_low, query_low);
let duplicated_high = vcombine_s8(query_high, query_high);
let rows_low = vcombine_s8(vget_low_s8(left), vget_low_s8(right));
let rows_high = vcombine_s8(vget_high_s8(left), vget_high_s8(right));
acc = unsafe { i8mm_mac(acc, duplicated_low, rows_low) };
acc = unsafe { i8mm_mac(acc, duplicated_high, rows_high) };
base += I8_LANES;
}
let sum0 = vgetq_lane_s32::<0>(acc);
let sum1 = vgetq_lane_s32::<1>(acc);
let tail0 = add_i8_tail(sum0, q, row0, processed);
let tail1 = add_i8_tail(sum1, q, row1, processed);
unsafe {
*out.get_unchecked_mut(row_index) = tail0;
*out.get_unchecked_mut(row_index + 1) = tail1;
}
row_index += 2;
}
if paired_rows < out.len() {
let start = paired_rows * d;
let row = unsafe { rows.get_unchecked(start..start + d) };
*unsafe { out.get_unchecked_mut(paired_rows) } = unsafe { dot_i8_dotprod_inner(q, row) };
}
}
#[target_feature(enable = "i8mm")]
#[cfg(any(test, feature = "benchmark-kernels"))]
unsafe fn i8mm_mac(mut acc: int32x4_t, left: int8x16_t, right: int8x16_t) -> int32x4_t {
unsafe {
asm!(
"smmla {acc:v}.4s, {left:v}.16b, {right:v}.16b",
acc = inout(vreg) acc,
left = in(vreg) left,
right = in(vreg) right,
options(pure, nomem, nostack)
);
}
acc
}
fn dot_bit4_batch_dotprod(q: &[i8], rows: &[u8], d: usize, out: &mut [i32]) {
debug_assert_eq!(q.len(), d, "query dimension must be pre-validated");
debug_assert!(d <= MAX_DOT_I8_DIMENSION);
let row_bytes = d.div_ceil(2);
debug_assert_eq!(rows.len(), row_bytes.saturating_mul(out.len()));
if row_bytes == 0 {
out.fill(0);
return;
}
let query_sum = q.iter().map(|&code| i32::from(code)).sum();
unsafe { dot_bit4_batch_dotprod_r4_inner(q, query_sum, rows, row_bytes, out) };
}
#[target_feature(enable = "dotprod")]
unsafe fn dot_bit4_batch_dotprod_r4_inner(
q: &[i8],
query_sum: i32,
rows: &[u8],
row_bytes: usize,
out: &mut [i32],
) {
unsafe {
const BLOCK: usize = BIT4_DIMENSIONS_PER_BLOCK * 2;
let processed = q.len() / BLOCK * BLOCK;
let processed_query_sum = if processed == q.len() {
query_sum
} else {
q.get(..processed)
.unwrap_or_default()
.iter()
.map(|&code| i32::from(code))
.sum()
};
let grouped_rows = out.len() / 4 * 4;
let mut row_index = 0_usize;
while row_index < grouped_rows {
let row0_ptr = rows.as_ptr().add(row_index * row_bytes);
let row1_ptr = row0_ptr.add(row_bytes);
let row2_ptr = row1_ptr.add(row_bytes);
let row3_ptr = row2_ptr.add(row_bytes);
let mut row0_acc0 = vdupq_n_s32(0);
let mut row0_acc1 = vdupq_n_s32(0);
let mut row1_acc0 = vdupq_n_s32(0);
let mut row1_acc1 = vdupq_n_s32(0);
let mut row2_acc0 = vdupq_n_s32(0);
let mut row2_acc1 = vdupq_n_s32(0);
let mut row3_acc0 = vdupq_n_s32(0);
let mut row3_acc1 = vdupq_n_s32(0);
let mut base = 0_usize;
while base < processed {
let query_ptr = q.as_ptr().add(base);
let coordinate0 = vld1q_s8(query_ptr);
let coordinate1 = vld1q_s8(query_ptr.add(I8_LANES));
let coordinate2 = vld1q_s8(query_ptr.add(I8_LANES * 2));
let coordinate3 = vld1q_s8(query_ptr.add(I8_LANES * 3));
let query_even0 = vuzp1q_s8(coordinate0, coordinate1);
let query_odd0 = vuzp2q_s8(coordinate0, coordinate1);
let query_even1 = vuzp1q_s8(coordinate2, coordinate3);
let query_odd1 = vuzp2q_s8(coordinate2, coordinate3);
let code_offset = base / 2;
row0_acc0 = bit4_split_mac(
row0_acc0,
query_even0,
query_odd0,
row0_ptr.add(code_offset),
);
row1_acc0 = bit4_split_mac(
row1_acc0,
query_even0,
query_odd0,
row1_ptr.add(code_offset),
);
row2_acc0 = bit4_split_mac(
row2_acc0,
query_even0,
query_odd0,
row2_ptr.add(code_offset),
);
row3_acc0 = bit4_split_mac(
row3_acc0,
query_even0,
query_odd0,
row3_ptr.add(code_offset),
);
row0_acc1 = bit4_split_mac(
row0_acc1,
query_even1,
query_odd1,
row0_ptr.add(code_offset + PACKED_LANES),
);
row1_acc1 = bit4_split_mac(
row1_acc1,
query_even1,
query_odd1,
row1_ptr.add(code_offset + PACKED_LANES),
);
row2_acc1 = bit4_split_mac(
row2_acc1,
query_even1,
query_odd1,
row2_ptr.add(code_offset + PACKED_LANES),
);
row3_acc1 = bit4_split_mac(
row3_acc1,
query_even1,
query_odd1,
row3_ptr.add(code_offset + PACKED_LANES),
);
base += BLOCK;
}
let pairs01 = vpaddq_s32(
vaddq_s32(row0_acc0, row0_acc1),
vaddq_s32(row1_acc0, row1_acc1),
);
let pairs23 = vpaddq_s32(
vaddq_s32(row2_acc0, row2_acc1),
vaddq_s32(row3_acc0, row3_acc1),
);
let unsigned_sums = vpaddq_s32(pairs01, pairs23);
let integer_dots = if processed == q.len() {
vsubq_s32(vshlq_n_s32(unsigned_sums, 1), vdupq_n_s32(15 * query_sum))
} else {
let mut sums = [0_i32; 4];
vst1q_s32(sums.as_mut_ptr(), unsigned_sums);
let row0 = std::slice::from_raw_parts(row0_ptr, row_bytes);
let row1 = std::slice::from_raw_parts(row1_ptr, row_bytes);
let row2 = std::slice::from_raw_parts(row2_ptr, row_bytes);
let row3 = std::slice::from_raw_parts(row3_ptr, row_bytes);
sums[0] =
packed_bit4_coordinate_tail(sums[0], q, row0, processed, processed_query_sum);
sums[1] =
packed_bit4_coordinate_tail(sums[1], q, row1, processed, processed_query_sum);
sums[2] =
packed_bit4_coordinate_tail(sums[2], q, row2, processed, processed_query_sum);
sums[3] =
packed_bit4_coordinate_tail(sums[3], q, row3, processed, processed_query_sum);
vld1q_s32(sums.as_ptr())
};
vst1q_s32(out.as_mut_ptr().add(row_index), integer_dots);
row_index += 4;
}
while row_index < out.len() {
let row_ptr = rows.as_ptr().add(row_index * row_bytes);
let row = std::slice::from_raw_parts(row_ptr, row_bytes);
*out.get_unchecked_mut(row_index) = dot_bit4_dotprod_inner(q, row);
row_index += 1;
}
}
}
fn packed_bit4_coordinate_tail(
unsigned_sum: i32,
q: &[i8],
codes: &[u8],
processed: usize,
processed_query_sum: i32,
) -> i32 {
let full_blocks = 2 * unsigned_sum - 15 * processed_query_sum;
let Some(q_tail) = q.get(processed..) else {
return full_blocks;
};
let Some(code_tail) = codes.get(processed / 2..) else {
return full_blocks;
};
full_blocks + scalar::dot_bit4(q_tail, code_tail)
}
fn score_bit4_prepared_batch_dotprod(
q: &[i8],
query_sum: i32,
query_scale_half: f64,
rows: &[u8],
d: usize,
factors: &[crate::quant::Bit4Factors],
out: &mut [f32],
) {
debug_assert_eq!(q.len(), d, "query dimension must be pre-validated");
debug_assert!(d <= MAX_DOT_I8_DIMENSION);
debug_assert_eq!(factors.len(), out.len());
let row_bytes = d.div_ceil(2);
debug_assert_eq!(rows.len(), row_bytes.saturating_mul(out.len()));
if row_bytes == 0 {
out.fill(0.0);
return;
}
unsafe {
score_bit4_prepared_batch_dotprod_r4_inner(
q,
query_sum,
query_scale_half,
rows,
row_bytes,
factors,
out,
);
}
}
unsafe fn score_bit4_ptrs_dotprod(
q: &[i8],
query_sum: i32,
query_scale_half: f64,
rows: &[*const u8; 4],
row_bytes: usize,
factors: &[crate::quant::Bit4Factors; 4],
out: &mut [f32; 4],
) {
debug_assert!(q.len() <= MAX_DOT_I8_DIMENSION);
debug_assert_eq!(row_bytes, q.len().div_ceil(2));
if row_bytes == 0 {
out.fill(0.0);
return;
}
unsafe {
score_bit4_prepared_ptrs_dotprod_r4_inner(
q,
query_sum,
query_scale_half,
rows,
row_bytes,
factors,
out.as_mut_ptr(),
);
}
}
#[target_feature(enable = "dotprod")]
unsafe fn score_bit4_prepared_batch_dotprod_r4_inner(
q: &[i8],
query_sum: i32,
query_scale_half: f64,
rows: &[u8],
row_bytes: usize,
factors: &[crate::quant::Bit4Factors],
out: &mut [f32],
) {
unsafe {
let grouped_rows = out.len() / 4 * 4;
let mut row_index = 0_usize;
while row_index < grouped_rows {
let row0_ptr = rows.as_ptr().add(row_index * row_bytes);
let pointers = [
row0_ptr,
row0_ptr.add(row_bytes),
row0_ptr.add(row_bytes * 2),
row0_ptr.add(row_bytes * 3),
];
let group_factors = [
*factors.get_unchecked(row_index),
*factors.get_unchecked(row_index + 1),
*factors.get_unchecked(row_index + 2),
*factors.get_unchecked(row_index + 3),
];
score_bit4_prepared_ptrs_dotprod_r4_inner(
q,
query_sum,
query_scale_half,
&pointers,
row_bytes,
&group_factors,
out.as_mut_ptr().add(row_index),
);
row_index += 4;
}
while row_index < out.len() {
let row_ptr = rows.as_ptr().add(row_index * row_bytes);
let row = std::slice::from_raw_parts(row_ptr, row_bytes);
let integer_dot = dot_bit4_prepared_dotprod_inner(q, query_sum, row);
*out.get_unchecked_mut(row_index) = scalar::bit4_score(
integer_dot,
*factors.get_unchecked(row_index),
query_scale_half,
);
row_index += 1;
}
}
}
#[target_feature(enable = "dotprod")]
unsafe fn score_bit4_prepared_ptrs_dotprod_r4_inner(
q: &[i8],
query_sum: i32,
query_scale_half: f64,
rows: &[*const u8; 4],
row_bytes: usize,
factors: &[crate::quant::Bit4Factors; 4],
out: *mut f32,
) {
unsafe {
const BLOCK: usize = BIT4_DIMENSIONS_PER_BLOCK * 2;
let processed = q.len() / BLOCK * BLOCK;
let row0_ptr = *rows.get_unchecked(0);
let row1_ptr = *rows.get_unchecked(1);
let row2_ptr = *rows.get_unchecked(2);
let row3_ptr = *rows.get_unchecked(3);
let mut row0_acc0 = vdupq_n_s32(0);
let mut row0_acc1 = vdupq_n_s32(0);
let mut row1_acc0 = vdupq_n_s32(0);
let mut row1_acc1 = vdupq_n_s32(0);
let mut row2_acc0 = vdupq_n_s32(0);
let mut row2_acc1 = vdupq_n_s32(0);
let mut row3_acc0 = vdupq_n_s32(0);
let mut row3_acc1 = vdupq_n_s32(0);
let mut base = 0_usize;
while base < processed {
let query_ptr = q.as_ptr().add(base);
let query0 = vld1q_s8(query_ptr);
let query1 = vld1q_s8(query_ptr.add(I8_LANES));
let query2 = vld1q_s8(query_ptr.add(I8_LANES * 2));
let query3 = vld1q_s8(query_ptr.add(I8_LANES * 3));
let code_offset = base / 2;
row0_acc0 = bit4_split_mac(row0_acc0, query0, query1, row0_ptr.add(code_offset));
row1_acc0 = bit4_split_mac(row1_acc0, query0, query1, row1_ptr.add(code_offset));
row2_acc0 = bit4_split_mac(row2_acc0, query0, query1, row2_ptr.add(code_offset));
row3_acc0 = bit4_split_mac(row3_acc0, query0, query1, row3_ptr.add(code_offset));
row0_acc1 = bit4_split_mac(
row0_acc1,
query2,
query3,
row0_ptr.add(code_offset + PACKED_LANES),
);
row1_acc1 = bit4_split_mac(
row1_acc1,
query2,
query3,
row1_ptr.add(code_offset + PACKED_LANES),
);
row2_acc1 = bit4_split_mac(
row2_acc1,
query2,
query3,
row2_ptr.add(code_offset + PACKED_LANES),
);
row3_acc1 = bit4_split_mac(
row3_acc1,
query2,
query3,
row3_ptr.add(code_offset + PACKED_LANES),
);
base += BLOCK;
}
let row0 = vaddq_s32(row0_acc0, row0_acc1);
let row1 = vaddq_s32(row1_acc0, row1_acc1);
let row2 = vaddq_s32(row2_acc0, row2_acc1);
let row3 = vaddq_s32(row3_acc0, row3_acc1);
let pairs01 = vpaddq_s32(row0, row1);
let pairs23 = vpaddq_s32(row2, row3);
let unsigned_sums = vpaddq_s32(pairs01, pairs23);
let integer_dots = if processed == q.len() {
vsubq_s32(vshlq_n_s32(unsigned_sums, 1), vdupq_n_s32(15 * query_sum))
} else {
let mut sums = [0_i32; 4];
vst1q_s32(sums.as_mut_ptr(), unsigned_sums);
let row0_slice = std::slice::from_raw_parts(row0_ptr, row_bytes);
let row1_slice = std::slice::from_raw_parts(row1_ptr, row_bytes);
let row2_slice = std::slice::from_raw_parts(row2_ptr, row_bytes);
let row3_slice = std::slice::from_raw_parts(row3_ptr, row_bytes);
sums[0] = packed_bit4_prepared_tail(sums[0], q, query_sum, row0_slice, processed);
sums[1] = packed_bit4_prepared_tail(sums[1], q, query_sum, row1_slice, processed);
sums[2] = packed_bit4_prepared_tail(sums[2], q, query_sum, row2_slice, processed);
sums[3] = packed_bit4_prepared_tail(sums[3], q, query_sum, row3_slice, processed);
vld1q_s32(sums.as_ptr())
};
store_bit4_score_pair(
integer_dots,
*factors.get_unchecked(0),
*factors.get_unchecked(1),
query_scale_half,
out,
);
store_bit4_score_pair(
vextq_s32(integer_dots, integer_dots, 2),
*factors.get_unchecked(2),
*factors.get_unchecked(3),
query_scale_half,
out.add(2),
);
}
}
#[target_feature(enable = "dotprod")]
unsafe fn bit4_split_mac(
mut acc: int32x4_t,
query_high: int8x16_t,
query_low: int8x16_t,
code_ptr: *const u8,
) -> int32x4_t {
unsafe {
let packed = vld1q_u8(code_ptr);
let high = vreinterpretq_s8_u8(vshrq_n_u8(packed, 4));
let low = vreinterpretq_s8_u8(vandq_u8(packed, vdupq_n_u8(0x0f)));
acc = dotprod_mac(acc, query_high, high);
dotprod_mac(acc, query_low, low)
}
}
#[inline]
#[target_feature(enable = "neon")]
unsafe fn store_bit4_score_pair(
integer_dots: int32x4_t,
factor0: crate::quant::Bit4Factors,
factor1: crate::quant::Bit4Factors,
query_scale_half: f64,
output: *mut f32,
) {
unsafe {
let (scale0, correction0) = factor0.scoring_parts();
let (scale1, correction1) = factor1.scoring_parts();
let scales = [scale0, scale1];
let corrections = [correction0, correction1];
let scale_f32 = vld1_f32(scales.as_ptr());
let correction_f32 = vld1_f32(corrections.as_ptr());
let factor_correction = vmulq_f64(vcvt_f64_f32(scale_f32), vcvt_f64_f32(correction_f32));
let dots_i64 = vmovl_s32(vget_low_s32(integer_dots));
let dots_f64 = vcvtq_f64_s64(dots_i64);
let scaled = vmulq_f64(vmulq_n_f64(factor_correction, query_scale_half), dots_f64);
let scores = vcvt_f32_f64(scaled);
let zero_mask = vceq_f32(scale_f32, vdup_n_f32(0.0));
let exact_scores = vbsl_u32(
zero_mask,
vreinterpret_u32_f32(vdup_n_f32(0.0)),
vreinterpret_u32_f32(scores),
);
vst1_f32(output, vreinterpret_f32_u32(exact_scores));
}
}
fn dot_bit4_batch_widen(q: &[i8], rows: &[u8], d: usize, out: &mut [i32]) {
batch_packed(q, rows, d, out, 2, dot_bit4_widen);
}
fn batch_i8(q: &[i8], rows: &[i8], d: usize, out: &mut [i32], kernel: fn(&[i8], &[i8]) -> i32) {
debug_assert_eq!(q.len(), d, "query dimension must be pre-validated");
debug_assert!(
d <= MAX_DOT_I8_DIMENSION,
"i8 dimension must be pre-validated"
);
debug_assert_eq!(
rows.len(),
d.saturating_mul(out.len()),
"batch shape must be pre-validated"
);
if d == 0 {
out.fill(0);
return;
}
for (row, result) in rows.chunks_exact(d).zip(out.iter_mut()) {
*result = kernel(q, row);
}
}
fn batch_packed(
q: &[i8],
rows: &[u8],
d: usize,
out: &mut [i32],
fields_per_byte: usize,
kernel: fn(&[i8], &[u8]) -> i32,
) {
debug_assert_eq!(q.len(), d, "query dimension must be pre-validated");
debug_assert!(d <= MAX_DOT_I8_DIMENSION);
let row_bytes = d.div_ceil(fields_per_byte);
debug_assert_eq!(
rows.len(),
row_bytes.saturating_mul(out.len()),
"batch shape must be pre-validated"
);
if row_bytes == 0 {
out.fill(0);
return;
}
for (row, result) in rows.chunks_exact(row_bytes).zip(out.iter_mut()) {
*result = kernel(q, row);
}
}
fn hamming_u1_batch(q: &[u8], rows: &[u8], d_bytes: usize, out: &mut [u32]) {
debug_assert_eq!(q.len(), d_bytes, "query dimension must be pre-validated");
debug_assert_eq!(
rows.len(),
d_bytes.saturating_mul(out.len()),
"batch shape must be pre-validated"
);
if d_bytes == 0 {
out.fill(0);
return;
}
unsafe { hamming_u1_batch_inner(q, rows, d_bytes, out) };
}
#[target_feature(enable = "neon")]
unsafe fn hamming_u1_batch_inner(q: &[u8], rows: &[u8], d_bytes: usize, out: &mut [u32]) {
for (row, result) in rows.chunks_exact(d_bytes).zip(out.iter_mut()) {
*result = unsafe { hamming_u1_inner(q, row) };
}
}
#[cfg(any(target_os = "macos", target_os = "ios"))]
pub(super) fn darwin_optional_feature(name: &[u8]) -> bool {
let Ok(name) = std::ffi::CStr::from_bytes_with_nul(name) else {
return false;
};
let mut value = 0_i32;
let mut value_len = size_of::<i32>();
let status = unsafe {
libc::sysctlbyname(
name.as_ptr(),
std::ptr::addr_of_mut!(value).cast(),
std::ptr::addr_of_mut!(value_len),
std::ptr::null_mut(),
0,
)
};
status == 0 && value_len == size_of::<i32>() && value != 0
}