diff --git a/jvector-native/src/main/java/io/github/jbellis/jvector/vector/NativeVectorUtilSupport.java b/jvector-native/src/main/java/io/github/jbellis/jvector/vector/NativeVectorUtilSupport.java index 4b627a244..5154747da 100644 --- a/jvector-native/src/main/java/io/github/jbellis/jvector/vector/NativeVectorUtilSupport.java +++ b/jvector-native/src/main/java/io/github/jbellis/jvector/vector/NativeVectorUtilSupport.java @@ -114,6 +114,11 @@ public void addInPlace(VectorFloat v1, float value) { NativeSimdOps.add_scalar_in_place_f32(((MemorySegmentVectorFloat) v1).get(), value, v1.length()); } + @Override + public void scale(VectorFloat vector, float multiplier) { + NativeSimdOps.scale_f32(((MemorySegmentVectorFloat) vector).get(), multiplier, vector.length()); + } + @Override public void subInPlace(VectorFloat v1, VectorFloat v2) { NativeSimdOps.sub_in_place_f32(((MemorySegmentVectorFloat) v1).get(), ((MemorySegmentVectorFloat) v2).get(), v1.length()); @@ -124,16 +129,61 @@ public void subInPlace(VectorFloat vector, float value) { NativeSimdOps.sub_scalar_in_place_f32(((MemorySegmentVectorFloat) vector).get(), value, vector.length()); } + @Override + public VectorFloat sub(VectorFloat a, VectorFloat b) { + return sub(a, 0, b, 0, a.length()); + } + + @Override + public VectorFloat sub(VectorFloat a, float value) { + return sub(a, 0, value, a.length()); + } + + @Override + public VectorFloat sub(VectorFloat a, int aOffset, VectorFloat b, int bOffset, int length) { + VectorFloat res = VectorizationProvider.getInstance().getVectorTypeSupport().createFloatVector(length); + NativeSimdOps.sub_f32( + ((MemorySegmentVectorFloat) a).get(), aOffset, + ((MemorySegmentVectorFloat) b).get(), bOffset, + ((MemorySegmentVectorFloat) res).get(), length); + return res; + } + + @Override + public VectorFloat sub(VectorFloat a, int aOffset, float value, int length) { + VectorFloat res = VectorizationProvider.getInstance().getVectorTypeSupport().createFloatVector(length); + NativeSimdOps.sub_scalar_f32( + ((MemorySegmentVectorFloat) a).get(), aOffset, + value, + ((MemorySegmentVectorFloat) res).get(), length); + return res; + } + @Override public float max(VectorFloat v) { return NativeSimdOps.max_f32(((MemorySegmentVectorFloat) v).get(), v.length()); } + @Override + public float min(VectorFloat v) { + return NativeSimdOps.min_f32(((MemorySegmentVectorFloat) v).get(), v.length()); + } + @Override public void minInPlace(VectorFloat v1, VectorFloat v2) { NativeSimdOps.min_in_place_f32(((MemorySegmentVectorFloat) v1).get(), ((MemorySegmentVectorFloat) v2).get(), v1.length()); } + @Override + public float sum(VectorFloat vector) { + return NativeSimdOps.sum_f32(((MemorySegmentVectorFloat) vector).get(), vector.length()); + } + + @Override + public int hammingDistance(long[] a, long[] b) { + return NativeSimdOps.hamming_distance(MemorySegment.ofArray(a), MemorySegment.ofArray(b), (long) a.length); + } + @Override public float assembleAndSumPQ( VectorFloat codebookPartialSums, diff --git a/jvector-native/src/main/java/io/github/jbellis/jvector/vector/cnative/NativeSimdOps.java b/jvector-native/src/main/java/io/github/jbellis/jvector/vector/cnative/NativeSimdOps.java index d822468a6..61b9c5ae8 100644 --- a/jvector-native/src/main/java/io/github/jbellis/jvector/vector/cnative/NativeSimdOps.java +++ b/jvector-native/src/main/java/io/github/jbellis/jvector/vector/cnative/NativeSimdOps.java @@ -1571,6 +1571,188 @@ public static void sub_scalar_in_place_f32(MemorySegment v1, float value, long l } } + private static class sub_f32 { + public static final FunctionDescriptor DESC = FunctionDescriptor.ofVoid( + NativeSimdOps.C_POINTER, + NativeSimdOps.C_LONG, + NativeSimdOps.C_POINTER, + NativeSimdOps.C_LONG, + NativeSimdOps.C_POINTER, + NativeSimdOps.C_LONG + ); + + public static final MemorySegment ADDR = NativeSimdOps.findOrThrow("sub_f32"); + + public static final MethodHandle HANDLE = Linker.nativeLinker().downcallHandle(ADDR, DESC, Linker.Option.critical(true)); + } + + /** + * Function descriptor for: + * {@snippet lang=c : + * void sub_f32(const float *a, size_t aoffset, const float *b, size_t boffset, float *res, size_t length) + * } + */ + public static FunctionDescriptor sub_f32$descriptor() { + return sub_f32.DESC; + } + + /** + * Downcall method handle for: + * {@snippet lang=c : + * void sub_f32(const float *a, size_t aoffset, const float *b, size_t boffset, float *res, size_t length) + * } + */ + public static MethodHandle sub_f32$handle() { + return sub_f32.HANDLE; + } + + /** + * Address for: + * {@snippet lang=c : + * void sub_f32(const float *a, size_t aoffset, const float *b, size_t boffset, float *res, size_t length) + * } + */ + public static MemorySegment sub_f32$address() { + return sub_f32.ADDR; + } + + /** + * {@snippet lang=c : + * void sub_f32(const float *a, size_t aoffset, const float *b, size_t boffset, float *res, size_t length) + * } + */ + public static void sub_f32(MemorySegment a, long aoffset, MemorySegment b, long boffset, MemorySegment res, long length) { + var mh$ = sub_f32.HANDLE; + try { + if (TRACE_DOWNCALLS) { + traceDowncall("sub_f32", a, aoffset, b, boffset, res, length); + } + mh$.invokeExact(a, aoffset, b, boffset, res, length); + } catch (Throwable ex$) { + throw new AssertionError("should not reach here", ex$); + } + } + + private static class sub_scalar_f32 { + public static final FunctionDescriptor DESC = FunctionDescriptor.ofVoid( + NativeSimdOps.C_POINTER, + NativeSimdOps.C_LONG, + NativeSimdOps.C_FLOAT, + NativeSimdOps.C_POINTER, + NativeSimdOps.C_LONG + ); + + public static final MemorySegment ADDR = NativeSimdOps.findOrThrow("sub_scalar_f32"); + + public static final MethodHandle HANDLE = Linker.nativeLinker().downcallHandle(ADDR, DESC, Linker.Option.critical(true)); + } + + /** + * Function descriptor for: + * {@snippet lang=c : + * void sub_scalar_f32(const float *a, size_t aoffset, float value, float *res, size_t length) + * } + */ + public static FunctionDescriptor sub_scalar_f32$descriptor() { + return sub_scalar_f32.DESC; + } + + /** + * Downcall method handle for: + * {@snippet lang=c : + * void sub_scalar_f32(const float *a, size_t aoffset, float value, float *res, size_t length) + * } + */ + public static MethodHandle sub_scalar_f32$handle() { + return sub_scalar_f32.HANDLE; + } + + /** + * Address for: + * {@snippet lang=c : + * void sub_scalar_f32(const float *a, size_t aoffset, float value, float *res, size_t length) + * } + */ + public static MemorySegment sub_scalar_f32$address() { + return sub_scalar_f32.ADDR; + } + + /** + * {@snippet lang=c : + * void sub_scalar_f32(const float *a, size_t aoffset, float value, float *res, size_t length) + * } + */ + public static void sub_scalar_f32(MemorySegment a, long aoffset, float value, MemorySegment res, long length) { + var mh$ = sub_scalar_f32.HANDLE; + try { + if (TRACE_DOWNCALLS) { + traceDowncall("sub_scalar_f32", a, aoffset, value, res, length); + } + mh$.invokeExact(a, aoffset, value, res, length); + } catch (Throwable ex$) { + throw new AssertionError("should not reach here", ex$); + } + } + + private static class scale_f32 { + public static final FunctionDescriptor DESC = FunctionDescriptor.ofVoid( + NativeSimdOps.C_POINTER, + NativeSimdOps.C_FLOAT, + NativeSimdOps.C_LONG + ); + + public static final MemorySegment ADDR = NativeSimdOps.findOrThrow("scale_f32"); + + public static final MethodHandle HANDLE = Linker.nativeLinker().downcallHandle(ADDR, DESC, Linker.Option.critical(true)); + } + + /** + * Function descriptor for: + * {@snippet lang=c : + * void scale_f32(float *v, float multiplier, size_t length) + * } + */ + public static FunctionDescriptor scale_f32$descriptor() { + return scale_f32.DESC; + } + + /** + * Downcall method handle for: + * {@snippet lang=c : + * void scale_f32(float *v, float multiplier, size_t length) + * } + */ + public static MethodHandle scale_f32$handle() { + return scale_f32.HANDLE; + } + + /** + * Address for: + * {@snippet lang=c : + * void scale_f32(float *v, float multiplier, size_t length) + * } + */ + public static MemorySegment scale_f32$address() { + return scale_f32.ADDR; + } + + /** + * {@snippet lang=c : + * void scale_f32(float *v, float multiplier, size_t length) + * } + */ + public static void scale_f32(MemorySegment v, float multiplier, long length) { + var mh$ = scale_f32.HANDLE; + try { + if (TRACE_DOWNCALLS) { + traceDowncall("scale_f32", v, multiplier, length); + } + mh$.invokeExact(v, multiplier, length); + } catch (Throwable ex$) { + throw new AssertionError("should not reach here", ex$); + } + } + private static class max_f32 { public static final FunctionDescriptor DESC = FunctionDescriptor.of( NativeSimdOps.C_FLOAT, @@ -1630,6 +1812,65 @@ public static float max_f32(MemorySegment v, long length) { } } + private static class min_f32 { + public static final FunctionDescriptor DESC = FunctionDescriptor.of( + NativeSimdOps.C_FLOAT, + NativeSimdOps.C_POINTER, + NativeSimdOps.C_LONG + ); + + public static final MemorySegment ADDR = NativeSimdOps.findOrThrow("min_f32"); + + public static final MethodHandle HANDLE = Linker.nativeLinker().downcallHandle(ADDR, DESC, Linker.Option.critical(true)); + } + + /** + * Function descriptor for: + * {@snippet lang=c : + * float min_f32(const float *v, size_t length) + * } + */ + public static FunctionDescriptor min_f32$descriptor() { + return min_f32.DESC; + } + + /** + * Downcall method handle for: + * {@snippet lang=c : + * float min_f32(const float *v, size_t length) + * } + */ + public static MethodHandle min_f32$handle() { + return min_f32.HANDLE; + } + + /** + * Address for: + * {@snippet lang=c : + * float min_f32(const float *v, size_t length) + * } + */ + public static MemorySegment min_f32$address() { + return min_f32.ADDR; + } + + /** + * {@snippet lang=c : + * float min_f32(const float *v, size_t length) + * } + */ + public static float min_f32(MemorySegment v, long length) { + var mh$ = min_f32.HANDLE; + try { + if (TRACE_DOWNCALLS) { + traceDowncall("min_f32", v, length); + } + return (float)mh$.invokeExact(v, length); + } catch (Throwable ex$) { + throw new AssertionError("should not reach here", ex$); + } + } + private static class min_in_place_f32 { public static final FunctionDescriptor DESC = FunctionDescriptor.ofVoid( NativeSimdOps.C_POINTER, @@ -1689,6 +1930,65 @@ public static void min_in_place_f32(MemorySegment v1, MemorySegment v2, long len } } + private static class sum_f32 { + public static final FunctionDescriptor DESC = FunctionDescriptor.of( + NativeSimdOps.C_FLOAT, + NativeSimdOps.C_POINTER, + NativeSimdOps.C_LONG + ); + + public static final MemorySegment ADDR = NativeSimdOps.findOrThrow("sum_f32"); + + public static final MethodHandle HANDLE = Linker.nativeLinker().downcallHandle(ADDR, DESC, Linker.Option.critical(true)); + } + + /** + * Function descriptor for: + * {@snippet lang=c : + * float sum_f32(const float *v, size_t length) + * } + */ + public static FunctionDescriptor sum_f32$descriptor() { + return sum_f32.DESC; + } + + /** + * Downcall method handle for: + * {@snippet lang=c : + * float sum_f32(const float *v, size_t length) + * } + */ + public static MethodHandle sum_f32$handle() { + return sum_f32.HANDLE; + } + + /** + * Address for: + * {@snippet lang=c : + * float sum_f32(const float *v, size_t length) + * } + */ + public static MemorySegment sum_f32$address() { + return sum_f32.ADDR; + } + + /** + * {@snippet lang=c : + * float sum_f32(const float *v, size_t length) + * } + */ + public static float sum_f32(MemorySegment v, long length) { + var mh$ = sum_f32.HANDLE; + try { + if (TRACE_DOWNCALLS) { + traceDowncall("sum_f32", v, length); + } + return (float)mh$.invokeExact(v, length); + } catch (Throwable ex$) { + throw new AssertionError("should not reach here", ex$); + } + } + private static class assemble_and_sum_f32 { public static final FunctionDescriptor DESC = FunctionDescriptor.of( NativeSimdOps.C_FLOAT, @@ -2506,6 +2806,66 @@ public static void nvq_shuffle_query_in_place_8bit(MemorySegment vector, long le } } + private static class hamming_distance { + public static final FunctionDescriptor DESC = FunctionDescriptor.of( + NativeSimdOps.C_INT, + NativeSimdOps.C_POINTER, + NativeSimdOps.C_POINTER, + NativeSimdOps.C_LONG + ); + + public static final MemorySegment ADDR = NativeSimdOps.findOrThrow("hamming_distance"); + + public static final MethodHandle HANDLE = Linker.nativeLinker().downcallHandle(ADDR, DESC, Linker.Option.critical(true)); + } + + /** + * Function descriptor for: + * {@snippet lang=c : + * int32_t hamming_distance(const uint64_t *a, const uint64_t *b, size_t length) + * } + */ + public static FunctionDescriptor hamming_distance$descriptor() { + return hamming_distance.DESC; + } + + /** + * Downcall method handle for: + * {@snippet lang=c : + * int32_t hamming_distance(const uint64_t *a, const uint64_t *b, size_t length) + * } + */ + public static MethodHandle hamming_distance$handle() { + return hamming_distance.HANDLE; + } + + /** + * Address for: + * {@snippet lang=c : + * int32_t hamming_distance(const uint64_t *a, const uint64_t *b, size_t length) + * } + */ + public static MemorySegment hamming_distance$address() { + return hamming_distance.ADDR; + } + + /** + * {@snippet lang=c : + * int32_t hamming_distance(const uint64_t *a, const uint64_t *b, size_t length) + * } + */ + public static int hamming_distance(MemorySegment a, MemorySegment b, long length) { + var mh$ = hamming_distance.HANDLE; + try { + if (TRACE_DOWNCALLS) { + traceDowncall("hamming_distance", a, b, length); + } + return (int)mh$.invokeExact(a, b, length); + } catch (Throwable ex$) { + throw new AssertionError("should not reach here", ex$); + } + } + private static class jvector_simd_get_active_isa { public static final FunctionDescriptor DESC = FunctionDescriptor.of( NativeSimdOps.C_POINTER ); diff --git a/jvector-native/src/main/native/src/jvector_simd_kernel_list.h b/jvector-native/src/main/native/src/jvector_simd_kernel_list.h index 63e7baa4d..8964f723a 100644 --- a/jvector-native/src/main/native/src/jvector_simd_kernel_list.h +++ b/jvector-native/src/main/native/src/jvector_simd_kernel_list.h @@ -37,13 +37,18 @@ KERNEL_ENTRY(float, cosine_f32, (const float *a, size_t aoffset, const float *b, size_t boffset, size_t length), (a, aoffset, b, boffset, length)) \ KERNEL_ENTRY(float, dot_product_f32, (const float *a, size_t aoffset, const float *b, size_t boffset, size_t length), (a, aoffset, b, boffset, length)) \ KERNEL_ENTRY(float, euclidean_f32, (const float *a, size_t aoffset, const float *b, size_t boffset, size_t length), (a, aoffset, b, boffset, length)) \ - /* Element-wise in-place arithmetic */ \ + /* Element-wise arithmetic and reductions */ \ KERNEL_ENTRY(void, add_in_place_f32, (float *v1, const float *v2, size_t length), (v1, v2, length)) \ KERNEL_ENTRY(void, add_scalar_in_place_f32, (float *v1, float value, size_t length), (v1, value, length)) \ KERNEL_ENTRY(void, sub_in_place_f32, (float *v1, const float *v2, size_t length), (v1, v2, length)) \ KERNEL_ENTRY(void, sub_scalar_in_place_f32, (float *v1, float value, size_t length), (v1, value, length)) \ + KERNEL_ENTRY(void, sub_f32, (const float *a, size_t aoffset, const float *b, size_t boffset, float *res, size_t length), (a, aoffset, b, boffset, res, length)) \ + KERNEL_ENTRY(void, sub_scalar_f32, (const float *a, size_t aoffset, float value, float *res, size_t length), (a, aoffset, value, res, length)) \ + KERNEL_ENTRY(void, scale_f32, (float *v, float multiplier, size_t length), (v, multiplier, length)) \ KERNEL_ENTRY(float, max_f32, (const float *v, size_t length), (v, length)) \ + KERNEL_ENTRY(float, min_f32, (const float *v, size_t length), (v, length)) \ KERNEL_ENTRY(void, min_in_place_f32, (float *v1, const float *v2, size_t length), (v1, v2, length)) \ + KERNEL_ENTRY(float, sum_f32, (const float *v, size_t length), (v, length)) \ /* PQ kernels */ \ KERNEL_ENTRY(float, assemble_and_sum_f32, (const float *data, int dataBase, const unsigned char *baseOffsets, int baseOffsetsOffset, size_t baseOffsetsLength), (data, dataBase, baseOffsets, baseOffsetsOffset, baseOffsetsLength)) \ KERNEL_ENTRY(float, assemble_and_sum_pq_f32, (const float *data, size_t subspaceCount, const unsigned char *baseOffsets1, int baseOffsetsOffset1, const unsigned char *baseOffsets2, int baseOffsetsOffset2, int clusterCount), (data, subspaceCount, baseOffsets1, baseOffsetsOffset1, baseOffsets2, baseOffsetsOffset2, clusterCount)) \ @@ -58,7 +63,9 @@ KERNEL_ENTRY(float, nvq_square_l2_distance_8bit, (const float *vector, const unsigned char *quantized, size_t length, float alpha, float x0, float minValue, float maxValue), (vector, quantized, length, alpha, x0, minValue, maxValue)) \ KERNEL_ENTRY(float, nvq_dot_product_8bit, (const float *vector, const unsigned char *quantized, size_t length, float alpha, float x0, float minValue, float maxValue), (vector, quantized, length, alpha, x0, minValue, maxValue)) \ KERNEL_ENTRY(int64_t, nvq_cosine_8bit_packed, (const float *vector, const unsigned char *quantized, size_t length, float alpha, float x0, float minValue, float maxValue, const float *centroid), (vector, quantized, length, alpha, x0, minValue, maxValue, centroid)) \ - KERNEL_ENTRY(void, nvq_shuffle_query_in_place_8bit, (float *vector, size_t length), (vector, length)) + KERNEL_ENTRY(void, nvq_shuffle_query_in_place_8bit, (float *vector, size_t length), (vector, length)) \ + /* Hamming distance */ \ + KERNEL_ENTRY(int32_t, hamming_distance, (const uint64_t *a, const uint64_t *b, size_t length), (a, b, length)) /* ── ADD NEW KERNEL_ENTRY LINES ABOVE THIS LINE ── */ // clang-format on diff --git a/jvector-native/src/main/native/src/jvector_simd_kernels.cpp b/jvector-native/src/main/native/src/jvector_simd_kernels.cpp index 875182af9..f7e540a12 100644 --- a/jvector-native/src/main/native/src/jvector_simd_kernels.cpp +++ b/jvector-native/src/main/native/src/jvector_simd_kernels.cpp @@ -502,6 +502,73 @@ HWY_FLATTEN void sub_scalar_in_place_f32(float *HWY_RESTRICT v1, } } +JV_RENAME_REGS +HWY_FLATTEN void sub_f32(const float *HWY_RESTRICT a, size_t aoffset, + const float *HWY_RESTRICT b, size_t boffset, + float *HWY_RESTRICT res, size_t length) +{ + a += aoffset; + b += boffset; + hn::ScalableTag d; + const size_t lanes = hn::Lanes(d); + size_t i = 0; +#pragma GCC unroll 4 + for (; i + lanes <= length; i += lanes) { + auto va = hn::LoadU(d, a + i); + auto vb = hn::LoadU(d, b + i); + hn::StoreU(hn::Sub(va, vb), d, res + i); + } + if (i < length) { + const size_t rem = length - i; + auto va = hn::LoadN(d, a + i, rem); + auto vb = hn::LoadN(d, b + i, rem); + hn::StoreN(hn::Sub(va, vb), d, res + i, rem); + } +} + +JV_RENAME_REGS +HWY_FLATTEN void sub_scalar_f32(const float *HWY_RESTRICT a, size_t aoffset, + float value, + float *HWY_RESTRICT res, size_t length) +{ + a += aoffset; + hn::ScalableTag d; + const size_t lanes = hn::Lanes(d); + const auto vval = hn::Set(d, value); + size_t i = 0; +#pragma GCC unroll 4 + for (; i + lanes <= length; i += lanes) { + auto va = hn::LoadU(d, a + i); + hn::StoreU(hn::Sub(va, vval), d, res + i); + } + if (i < length) { + const size_t rem = length - i; + auto va = hn::LoadN(d, a + i, rem); + hn::StoreN(hn::Sub(va, vval), d, res + i, rem); + } +} + +JV_RENAME_REGS +HWY_FLATTEN void scale_f32(float *HWY_RESTRICT v, + float multiplier, + size_t length) +{ + hn::ScalableTag d; + const size_t lanes = hn::Lanes(d); + const auto vmult = hn::Set(d, multiplier); + size_t i = 0; +#pragma GCC unroll 4 + for (; i + lanes <= length; i += lanes) { + auto a = hn::LoadU(d, v + i); + hn::StoreU(hn::Mul(a, vmult), d, v + i); + } + if (i < length) { + const size_t rem = length - i; + auto a = hn::LoadN(d, v + i, rem); + hn::StoreN(hn::Mul(a, vmult), d, v + i, rem); + } +} + JV_RENAME_REGS HWY_FLATTEN float max_f32(const float *HWY_RESTRICT v, size_t length) { @@ -520,6 +587,24 @@ HWY_FLATTEN float max_f32(const float *HWY_RESTRICT v, size_t length) return result; } +JV_RENAME_REGS +HWY_FLATTEN float min_f32(const float *HWY_RESTRICT v, size_t length) +{ + hn::ScalableTag d; + const size_t lanes = hn::Lanes(d); + auto accum = hn::Set(d, FLT_MAX); + size_t i = 0; +#pragma GCC unroll 4 + for (; i + lanes <= length; i += lanes) { + accum = hn::Min(accum, hn::LoadU(d, v + i)); + } + float result = hn::ReduceMin(d, accum); + for (; i < length; i++) { + if (v[i] < result) result = v[i]; + } + return result; +} + JV_RENAME_REGS HWY_FLATTEN void min_in_place_f32(float *HWY_RESTRICT v1, const float *HWY_RESTRICT v2, @@ -542,6 +627,24 @@ HWY_FLATTEN void min_in_place_f32(float *HWY_RESTRICT v1, } } +JV_RENAME_REGS +HWY_FLATTEN float sum_f32(const float *HWY_RESTRICT v, size_t length) +{ + hn::ScalableTag d; + const size_t lanes = hn::Lanes(d); + auto accum = hn::Zero(d); + size_t i = 0; +#pragma GCC unroll 4 + for (; i + lanes <= length; i += lanes) { + accum = hn::Add(accum, hn::LoadU(d, v + i)); + } + float result = hn::ReduceSum(d, accum); + for (; i < length; i++) { + result += v[i]; + } + return result; +} + // ============================================================================= // PQ kernels // ============================================================================= @@ -1697,4 +1800,31 @@ HWY_FLATTEN int64_t nvq_cosine_8bit_packed(const float *HWY_RESTRICT vector, return ((int64_t)bmag_bits << 32) | (int64_t)(uint32_t)sum_bits; } +// ============================================================================= +// Hamming distance kernel +// ============================================================================= + +JV_RENAME_REGS +HWY_FLATTEN int32_t hamming_distance(const uint64_t *HWY_RESTRICT a, + const uint64_t *HWY_RESTRICT b, + size_t length) +{ + hn::ScalableTag d; + const size_t lanes = hn::Lanes(d); + auto acc = hn::Zero(d); + size_t i = 0; +#pragma GCC unroll 4 + for (; i + lanes <= length; i += lanes) { + auto va = hn::LoadU(d, a + i); + auto vb = hn::LoadU(d, b + i); + auto xord = hn::Xor(va, vb); + acc = hn::Add(acc, hn::PopulationCount(xord)); + } + uint64_t result = hn::ReduceSum(d, acc); + for (; i < length; i++) { + result += hwy::PopCount(a[i] ^ b[i]); + } + return static_cast(result); +} + } // namespace JV_ISA diff --git a/jvector-native/src/main/native/tests/test_elementwise.cpp b/jvector-native/src/main/native/tests/test_elementwise.cpp index 82dab16f4..0446a3b93 100644 --- a/jvector-native/src/main/native/tests/test_elementwise.cpp +++ b/jvector-native/src/main/native/tests/test_elementwise.cpp @@ -118,6 +118,66 @@ TEST_P(ElementWiseTest, SubScalarInPlace) EXPECT_NEAR(got[i], want[i], 1e-5f) << "sub_scalar_in_place_f32[" << i << "]"; } +// --------------------------------------------------------------------------- +// sub_f32: res[i] = a[aoffset + i] - b[boffset + i] +// --------------------------------------------------------------------------- + +TEST_P(ElementWiseTest, SubF32) +{ + const size_t n = GetParam().length; + auto a = make_vec(n + 7, 1.1f); + auto b = make_vec(n + 7, 0.7f); + std::vector res(n, 0.0f); + const size_t aoffset = 3; + const size_t boffset = 2; + + sub_f32(a.data(), aoffset, b.data(), boffset, res.data(), n); + + for (size_t i = 0; i < n; ++i) { + float want = a[aoffset + i] - b[boffset + i]; + EXPECT_NEAR(res[i], want, 1e-5f) << "sub_f32[" << i << "]"; + } +} + +// --------------------------------------------------------------------------- +// sub_scalar_f32: res[i] = a[aoffset + i] - value +// --------------------------------------------------------------------------- + +TEST_P(ElementWiseTest, SubScalarF32) +{ + const size_t n = GetParam().length; + auto a = make_vec(n + 5, 1.1f); + std::vector res(n, 0.0f); + const size_t aoffset = 2; + const float scalar = 1.618f; + + sub_scalar_f32(a.data(), aoffset, scalar, res.data(), n); + + for (size_t i = 0; i < n; ++i) { + float want = a[aoffset + i] - scalar; + EXPECT_NEAR(res[i], want, 1e-5f) << "sub_scalar_f32[" << i << "]"; + } +} + +// --------------------------------------------------------------------------- +// scale_f32: v[i] *= multiplier +// --------------------------------------------------------------------------- + +TEST_P(ElementWiseTest, ScaleF32) +{ + const size_t n = GetParam().length; + const float multiplier = 2.5f; + auto v = make_vec(n, 1.1f); + auto want = v; + for (size_t i = 0; i < n; ++i) want[i] *= multiplier; + + auto got = v; + scale_f32(got.data(), multiplier, n); + + for (size_t i = 0; i < n; ++i) + EXPECT_NEAR(got[i], want[i], 1e-5f) << "scale_f32[" << i << "]"; +} + // --------------------------------------------------------------------------- // max_f32: returns the maximum element // --------------------------------------------------------------------------- @@ -146,6 +206,71 @@ TEST_P(ElementWiseTest, MaxF32TailElement) EXPECT_FLOAT_EQ(got, 1e6f); } +// --------------------------------------------------------------------------- +// min_f32: returns the minimum element +// --------------------------------------------------------------------------- + +TEST_P(ElementWiseTest, MinF32) +{ + const size_t n = GetParam().length; + auto v = make_vec(n, 0.9f); + + float want = *std::min_element(v.begin(), v.end()); + float got = min_f32(v.data(), n); + + EXPECT_FLOAT_EQ(got, want); +} + +// min_f32 on a vector with a known minimum at the last position (tail element) +TEST_P(ElementWiseTest, MinF32TailElement) +{ + const size_t n = GetParam().length; + auto v = make_vec(n, 0.5f); + // Place the global minimum in the very last element — exercises tail path. + v.back() = -1e6f; + + float got = min_f32(v.data(), n); + + EXPECT_FLOAT_EQ(got, -1e6f); +} + +// --------------------------------------------------------------------------- +// sum_f32: returns the sum of elements +// --------------------------------------------------------------------------- + +TEST_P(ElementWiseTest, SumF32) +{ + const size_t n = GetParam().length; + auto v = make_vec(n, 0.8f); + + float want = 0.0f; + for (size_t i = 0; i < n; ++i) want += v[i]; + float got = sum_f32(v.data(), n); + + EXPECT_NEAR(got, want, 1e-4f * std::max(1.0f, std::abs(want))); +} + +// --------------------------------------------------------------------------- +// hamming_distance: returns population count of XOR of uint64_t arrays +// --------------------------------------------------------------------------- + +TEST_P(ElementWiseTest, HammingDistance) +{ + const size_t n = GetParam().length; + std::vector a(n), b(n); + int32_t want = 0; + for (size_t i = 0; i < n; ++i) { + a[i] = (static_cast(i + 1) * 0x9E3779B97F4A7C15ULL) ^ 0x0123456789ABCDEFULL; + b[i] = (static_cast(i + 1) * 0xC6A4A7935BD1E995ULL) ^ 0xFEDCBA9876543210ULL; + uint64_t xord = a[i] ^ b[i]; + want += static_cast(__builtin_popcountll(xord)); + } + + int32_t got = hamming_distance(a.data(), b.data(), n); + + EXPECT_EQ(got, want) << "HammingDistance failed for length " << n; +} + // --------------------------------------------------------------------------- // min_in_place_f32: v1[i] = min(v1[i], v2[i]) // --------------------------------------------------------------------------- diff --git a/jvector-native/verify_native_lib.sh b/jvector-native/verify_native_lib.sh index ecf98369b..774902fe9 100755 --- a/jvector-native/verify_native_lib.sh +++ b/jvector-native/verify_native_lib.sh @@ -47,6 +47,31 @@ REQUIRED_SYMBOLS=( dot_product_f32 cosine_f32 euclidean_f32 + add_in_place_f32 + add_scalar_in_place_f32 + sub_in_place_f32 + sub_scalar_in_place_f32 + sub_f32 + sub_scalar_f32 + scale_f32 + max_f32 + min_f32 + min_in_place_f32 + sum_f32 + assemble_and_sum_f32 + assemble_and_sum_pq_f32 + pq_decoded_cosine_similarity_f32 + calculate_partial_sums_dot_f32 + calculate_partial_sums_euclidean_f32 + calculate_partial_sums_self_magnitude_f32 + nvq_quantize_8bit + nvq_loss + nvq_uniform_loss + nvq_square_l2_distance_8bit + nvq_dot_product_8bit + nvq_cosine_8bit_packed + nvq_shuffle_query_in_place_8bit + hamming_distance ) WORKDIR=$(mktemp -d)