faiss 0.6.4 → 0.6.5
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- checksums.yaml +4 -4
- data/CHANGELOG.md +4 -0
- data/ext/faiss/extconf.rb +1 -1
- data/lib/faiss/version.rb +1 -1
- data/vendor/faiss/faiss/AutoTune.cpp +175 -0
- data/vendor/faiss/faiss/Clustering.cpp +150 -60
- data/vendor/faiss/faiss/Clustering.h +29 -0
- data/vendor/faiss/faiss/IVFlib.cpp +99 -0
- data/vendor/faiss/faiss/IVFlib.h +15 -0
- data/vendor/faiss/faiss/Index.cpp +44 -0
- data/vendor/faiss/faiss/Index.h +19 -14
- data/vendor/faiss/faiss/IndexBinary.cpp +8 -0
- data/vendor/faiss/faiss/IndexBinary.h +8 -0
- data/vendor/faiss/faiss/IndexBinaryHNSW.cpp +82 -4
- data/vendor/faiss/faiss/IndexBinaryHNSW.h +19 -0
- data/vendor/faiss/faiss/IndexBinaryIVF.cpp +34 -0
- data/vendor/faiss/faiss/IndexFastScan.cpp +33 -17
- data/vendor/faiss/faiss/IndexFastScan.h +6 -3
- data/vendor/faiss/faiss/IndexFlat.cpp +56 -0
- data/vendor/faiss/faiss/IndexFlat.h +31 -0
- data/vendor/faiss/faiss/IndexFlatCodes.cpp +6 -1
- data/vendor/faiss/faiss/IndexHNSW.cpp +456 -10
- data/vendor/faiss/faiss/IndexHNSW.h +38 -0
- data/vendor/faiss/faiss/IndexIDMap.cpp +32 -0
- data/vendor/faiss/faiss/IndexIDMap.h +7 -0
- data/vendor/faiss/faiss/IndexIVF.cpp +281 -31
- data/vendor/faiss/faiss/IndexIVF.h +53 -1
- data/vendor/faiss/faiss/IndexIVFFastScan.cpp +46 -14
- data/vendor/faiss/faiss/IndexIVFFastScan.h +4 -2
- data/vendor/faiss/faiss/IndexIVFFlat.cpp +33 -0
- data/vendor/faiss/faiss/IndexIVFFlat.h +16 -0
- data/vendor/faiss/faiss/IndexIVFPQFastScan.cpp +1 -0
- data/vendor/faiss/faiss/IndexIVFRaBitQ.cpp +3 -20
- data/vendor/faiss/faiss/IndexIVFRaBitQFastScan.cpp +50 -14
- data/vendor/faiss/faiss/IndexIVFRaBitQFastScan.h +5 -2
- data/vendor/faiss/faiss/IndexIVFSQFastScan.cpp +855 -0
- data/vendor/faiss/faiss/IndexIVFSQFastScan.h +145 -0
- data/vendor/faiss/faiss/IndexNNDescent.cpp +0 -3
- data/vendor/faiss/faiss/IndexPiPNN.cpp +201 -0
- data/vendor/faiss/faiss/IndexPiPNN.h +83 -0
- data/vendor/faiss/faiss/IndexRaBitQFastScan.cpp +6 -3
- data/vendor/faiss/faiss/IndexRaBitQFastScan.h +10 -4
- data/vendor/faiss/faiss/IndexReplicas.cpp +20 -0
- data/vendor/faiss/faiss/IndexReplicas.h +9 -0
- data/vendor/faiss/faiss/IndexSQFastScan.cpp +648 -0
- data/vendor/faiss/faiss/IndexSQFastScan.h +129 -0
- data/vendor/faiss/faiss/IndexShards.cpp +10 -0
- data/vendor/faiss/faiss/IndexShards.h +5 -0
- data/vendor/faiss/faiss/SuperKMeans.cpp +57 -15
- data/vendor/faiss/faiss/SuperKMeans.h +9 -4
- data/vendor/faiss/faiss/VectorTransform.cpp +3 -2
- data/vendor/faiss/faiss/clone_index.cpp +24 -0
- data/vendor/faiss/faiss/gpu/Dataset.h +112 -0
- data/vendor/faiss/faiss/gpu/GpuCloner.cpp +16 -13
- data/vendor/faiss/faiss/gpu/GpuCloner.h +4 -0
- data/vendor/faiss/faiss/gpu/GpuClonerOptions.h +3 -0
- data/vendor/faiss/faiss/gpu/GpuIndex.h +4 -13
- data/vendor/faiss/faiss/gpu/GpuIndexBinaryCagra.h +26 -13
- data/vendor/faiss/faiss/gpu/GpuIndexCagra.h +45 -22
- data/vendor/faiss/faiss/gpu/GpuIndexFlat.h +11 -0
- data/vendor/faiss/faiss/gpu/GpuIndexIVF.h +6 -0
- data/vendor/faiss/faiss/gpu/test/TestGpuDatasetHeader.cpp +18 -0
- data/vendor/faiss/faiss/gpu/test/TestGpuIndexFlat.cpp +215 -0
- data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFFlat.cpp +98 -0
- data/vendor/faiss/faiss/gpu/utils/DeviceUtils.h +11 -0
- data/vendor/faiss/faiss/gpu/utils/StackDeviceMemory.cpp +1 -1
- data/vendor/faiss/faiss/gpu_metal/MetalCloner.h +11 -0
- data/vendor/faiss/faiss/gpu_metal/MetalDistance.h +92 -0
- data/vendor/faiss/faiss/gpu_metal/MetalIndex.h +4 -0
- data/vendor/faiss/faiss/gpu_metal/MetalIndexBinaryFlat.h +69 -0
- data/vendor/faiss/faiss/gpu_metal/MetalIndexIVFPQ.h +38 -0
- data/vendor/faiss/faiss/gpu_metal/MetalKernels.h +106 -0
- data/vendor/faiss/faiss/gpu_metal/MetalPythonBridge.h +10 -0
- data/vendor/faiss/faiss/impl/AuxIndexStructures.cpp +2 -0
- data/vendor/faiss/faiss/impl/AuxIndexStructures.h +1 -6
- data/vendor/faiss/faiss/impl/ClusteringHelpers.cpp +134 -45
- data/vendor/faiss/faiss/impl/ClusteringHelpers.h +36 -0
- data/vendor/faiss/faiss/impl/ClusteringInitialization.cpp +286 -136
- data/vendor/faiss/faiss/impl/DistanceComputer.cpp +63 -0
- data/vendor/faiss/faiss/impl/DistanceComputer.h +3 -29
- data/vendor/faiss/faiss/impl/HNSW.cpp +46 -10
- data/vendor/faiss/faiss/impl/HNSW.h +72 -0
- data/vendor/faiss/faiss/impl/LocalSearchQuantizer.cpp +6 -19
- data/vendor/faiss/faiss/impl/NNDescent.cpp +2 -1
- data/vendor/faiss/faiss/impl/NSG.cpp +3 -2
- data/vendor/faiss/faiss/impl/Panorama.h +3 -2
- data/vendor/faiss/faiss/impl/PiPNN.cpp +822 -0
- data/vendor/faiss/faiss/impl/PiPNN.h +122 -0
- data/vendor/faiss/faiss/impl/ProductQuantizer.cpp +54 -11
- data/vendor/faiss/faiss/impl/RaBitQUtils.h +1 -1
- data/vendor/faiss/faiss/impl/RaBitQuantizer.cpp +81 -57
- data/vendor/faiss/faiss/impl/RaBitQuantizer.h +15 -0
- data/vendor/faiss/faiss/impl/ScalarQuantizer.cpp +29 -0
- data/vendor/faiss/faiss/impl/ScalarQuantizer.h +1 -6
- data/vendor/faiss/faiss/impl/approx_topk/avx2.cpp +0 -4
- data/vendor/faiss/faiss/impl/approx_topk/neon.cpp +0 -4
- data/vendor/faiss/faiss/impl/binary_hamming/IndexBinaryHash_impl.h +5 -2
- data/vendor/faiss/faiss/impl/binary_hamming/avx2.cpp +0 -4
- data/vendor/faiss/faiss/impl/binary_hamming/avx512.cpp +0 -4
- data/vendor/faiss/faiss/impl/binary_hamming/neon.cpp +0 -4
- data/vendor/faiss/faiss/impl/binary_hamming/rvv.cpp +0 -4
- data/vendor/faiss/faiss/impl/expanded_scanners.h +88 -0
- data/vendor/faiss/faiss/impl/fast_scan/FastScanDistancePostProcessing.h +4 -0
- data/vendor/faiss/faiss/impl/fast_scan/accumulate_loops.h +61 -64
- data/vendor/faiss/faiss/impl/fast_scan/accumulate_loops_512.h +13 -14
- data/vendor/faiss/faiss/impl/fast_scan/decompose_qbs.h +17 -129
- data/vendor/faiss/faiss/impl/fast_scan/dispatching.h +25 -19
- data/vendor/faiss/faiss/impl/fast_scan/fast_scan.cpp +5 -0
- data/vendor/faiss/faiss/impl/fast_scan/impl-avx2.cpp +0 -4
- data/vendor/faiss/faiss/impl/fast_scan/impl-avx512.cpp +0 -4
- data/vendor/faiss/faiss/impl/fast_scan/impl-neon.cpp +0 -88
- data/vendor/faiss/faiss/impl/fast_scan/impl-riscv.cpp +0 -4
- data/vendor/faiss/faiss/impl/fast_scan/kernels_neon.h +155 -0
- data/vendor/faiss/faiss/impl/fast_scan/kernels_simd256.h +22 -1
- data/vendor/faiss/faiss/impl/fast_scan/kernels_simd512.h +3 -2
- data/vendor/faiss/faiss/impl/fast_scan/simd_result_handlers.h +41 -23
- data/vendor/faiss/faiss/impl/hnsw/avx2.cpp +0 -4
- data/vendor/faiss/faiss/impl/hnsw/avx512.cpp +0 -4
- data/vendor/faiss/faiss/impl/index_read.cpp +172 -0
- data/vendor/faiss/faiss/impl/index_write.cpp +95 -0
- data/vendor/faiss/faiss/impl/pipnn/HashPrune.cpp +159 -0
- data/vendor/faiss/faiss/impl/pipnn/HashPrune.h +69 -0
- data/vendor/faiss/faiss/impl/pipnn/Partition.cpp +587 -0
- data/vendor/faiss/faiss/impl/pipnn/Partition.h +93 -0
- data/vendor/faiss/faiss/impl/pipnn/kernels.cpp +326 -0
- data/vendor/faiss/faiss/impl/pipnn/kernels.h +92 -0
- data/vendor/faiss/faiss/impl/platform_macros.h +45 -0
- data/vendor/faiss/faiss/impl/polysemous_training/avx512.cpp +0 -4
- data/vendor/faiss/faiss/impl/pq_code_distance/IVFPQScanner_impl.h +79 -4
- data/vendor/faiss/faiss/impl/pq_code_distance/PQDistanceComputer_impl.h +6 -0
- data/vendor/faiss/faiss/impl/pq_code_distance/avx2.cpp +0 -4
- data/vendor/faiss/faiss/impl/pq_code_distance/avx512.cpp +0 -4
- data/vendor/faiss/faiss/impl/pq_code_distance/neon.cpp +0 -4
- data/vendor/faiss/faiss/impl/pq_code_distance/pq_code_distance-avx2.h +4 -0
- data/vendor/faiss/faiss/impl/pq_code_distance/pq_code_distance-avx512.h +4 -1
- data/vendor/faiss/faiss/impl/pq_code_distance/pq_code_distance-generic.cpp +19 -0
- data/vendor/faiss/faiss/impl/pq_code_distance/pq_code_distance-generic.h +6 -6
- data/vendor/faiss/faiss/impl/pq_code_distance/pq_code_distance-inl.h +131 -6
- data/vendor/faiss/faiss/impl/pq_code_distance/pq_code_distance-sve.cpp +0 -4
- data/vendor/faiss/faiss/impl/pq_code_distance/pq_rvv_kernel.h +89 -0
- data/vendor/faiss/faiss/impl/pq_code_distance/pq_scan_impl.h +21 -10
- data/vendor/faiss/faiss/impl/pq_code_distance/rvv.cpp +45 -18
- data/vendor/faiss/faiss/impl/residual_quantizer_encode_steps.cpp +2 -1
- data/vendor/faiss/faiss/impl/result_handler/avx2.cpp +0 -4
- data/vendor/faiss/faiss/impl/result_handler/avx512.cpp +0 -4
- data/vendor/faiss/faiss/impl/scalar_quantizer/EDENQuantizer-avx2.cpp +0 -4
- data/vendor/faiss/faiss/impl/scalar_quantizer/EDENQuantizer-avx512.cpp +0 -4
- data/vendor/faiss/faiss/impl/scalar_quantizer/distance_computers.h +5 -1
- data/vendor/faiss/faiss/impl/scalar_quantizer/quantizers.h +28 -52
- data/vendor/faiss/faiss/impl/scalar_quantizer/scanners.h +59 -11
- data/vendor/faiss/faiss/impl/scalar_quantizer/sq-avx2.cpp +226 -77
- data/vendor/faiss/faiss/impl/scalar_quantizer/sq-avx512-impl.h +482 -103
- data/vendor/faiss/faiss/impl/scalar_quantizer/sq-avx512-spr.cpp +2 -149
- data/vendor/faiss/faiss/impl/scalar_quantizer/sq-avx512.cpp +3 -698
- data/vendor/faiss/faiss/impl/scalar_quantizer/sq-dispatch.h +66 -57
- data/vendor/faiss/faiss/impl/scalar_quantizer/sq-neon.cpp +101 -96
- data/vendor/faiss/faiss/impl/scalar_quantizer/sq-rvv.cpp +3185 -144
- data/vendor/faiss/faiss/impl/scalar_quantizer/sq8_coefficients.h +48 -0
- data/vendor/faiss/faiss/impl/scalar_quantizer/sq8_rvv_kernel.h +546 -0
- data/vendor/faiss/faiss/impl/scalar_quantizer/training.cpp +2 -2
- data/vendor/faiss/faiss/impl/simd_dispatch.h +7 -0
- data/vendor/faiss/faiss/impl/simdlib/simdlib_avx2.h +14 -0
- data/vendor/faiss/faiss/impl/simdlib/simdlib_dispatch.h +7 -8
- data/vendor/faiss/faiss/impl/sq_fastscan_utils.h +232 -0
- data/vendor/faiss/faiss/index_factory.cpp +76 -0
- data/vendor/faiss/faiss/invlists/InvertedLists.cpp +3 -1
- data/vendor/faiss/faiss/svs/IndexSVSVamana.h +3 -2
- data/vendor/faiss/faiss/utils/AlignedTable.h +6 -1
- data/vendor/faiss/faiss/utils/Heap.cpp +4 -3
- data/vendor/faiss/faiss/utils/approx_topk_hamming/approx_topk_hamming.h +16 -5
- data/vendor/faiss/faiss/utils/bf16.h +36 -0
- data/vendor/faiss/faiss/utils/distances.cpp +226 -15
- data/vendor/faiss/faiss/utils/distances.h +34 -0
- data/vendor/faiss/faiss/utils/distances_dispatch.h +6 -4
- data/vendor/faiss/faiss/utils/distances_fused/avx512.cpp +2 -5
- data/vendor/faiss/faiss/utils/distances_fused/distances_fused.cpp +0 -14
- data/vendor/faiss/faiss/utils/distances_fused/rvv.cpp +25 -0
- data/vendor/faiss/faiss/utils/distances_fused/simdlib_based.cpp +0 -4
- data/vendor/faiss/faiss/utils/distances_fused/simdlib_based_neon.cpp +0 -4
- data/vendor/faiss/faiss/utils/distances_fused/simdlib_kernel-inl.h +2 -1
- data/vendor/faiss/faiss/utils/distances_simd.cpp +9 -0
- data/vendor/faiss/faiss/utils/fp16-fp16c.h +4 -0
- data/vendor/faiss/faiss/utils/fp16-inl.h +26 -52
- data/vendor/faiss/faiss/utils/hamming_distance/hamming_avx2.cpp +0 -4
- data/vendor/faiss/faiss/utils/hamming_distance/hamming_avx512.cpp +0 -4
- data/vendor/faiss/faiss/utils/hamming_distance/hamming_avx512_vpopcnt.cpp +0 -4
- data/vendor/faiss/faiss/utils/hamming_distance/hamming_computer-avx512.h +33 -5
- data/vendor/faiss/faiss/utils/hamming_distance/hamming_computer-avx512_vpopcnt.h +15 -30
- data/vendor/faiss/faiss/utils/hamming_distance/hamming_computer-neon.h +0 -4
- data/vendor/faiss/faiss/utils/hamming_distance/hamming_impl.h +2 -1
- data/vendor/faiss/faiss/utils/hamming_distance/hamming_neon.cpp +0 -4
- data/vendor/faiss/faiss/utils/hamming_distance/hamming_rvv.cpp +0 -4
- data/vendor/faiss/faiss/utils/pq_code_distance.h +8 -4
- data/vendor/faiss/faiss/utils/prefetch.h +1 -26
- data/vendor/faiss/faiss/utils/quantize_lut.cpp +26 -27
- data/vendor/faiss/faiss/utils/quantize_lut.h +11 -0
- data/vendor/faiss/faiss/utils/rabitq_simd.h +51 -16
- data/vendor/faiss/faiss/utils/simd_impl/distances_aarch64.cpp +9 -4
- data/vendor/faiss/faiss/utils/simd_impl/distances_arm_sve.cpp +37 -5
- data/vendor/faiss/faiss/utils/simd_impl/distances_avx2.cpp +64 -5
- data/vendor/faiss/faiss/utils/simd_impl/distances_avx512.cpp +64 -5
- data/vendor/faiss/faiss/utils/simd_impl/distances_rvv.cpp +353 -71
- data/vendor/faiss/faiss/utils/simd_impl/distances_rvv_kernel.h +68 -0
- data/vendor/faiss/faiss/utils/simd_impl/exhaustive_L2sqr_blas_cmax.h +93 -0
- data/vendor/faiss/faiss/utils/simd_impl/fp16_kernels.h +76 -0
- data/vendor/faiss/faiss/utils/simd_impl/fp16_kernels_avx2.cpp +52 -0
- data/vendor/faiss/faiss/utils/simd_impl/fp16_kernels_avx512.cpp +52 -0
- data/vendor/faiss/faiss/utils/simd_impl/fp16_kernels_neon.cpp +48 -0
- data/vendor/faiss/faiss/utils/simd_impl/masked_sum-avx2.h +44 -0
- data/vendor/faiss/faiss/utils/simd_impl/masked_sum-avx512.h +43 -0
- data/vendor/faiss/faiss/utils/simd_impl/masked_sum.h +38 -0
- data/vendor/faiss/faiss/utils/simd_impl/partitioning_avx2.cpp +0 -4
- data/vendor/faiss/faiss/utils/simd_impl/partitioning_neon.cpp +0 -4
- data/vendor/faiss/faiss/utils/simd_impl/rabitq_avx2.cpp +54 -15
- data/vendor/faiss/faiss/utils/simd_impl/rabitq_avx512.cpp +38 -21
- data/vendor/faiss/faiss/utils/simd_impl/rabitq_avx512_vpopcnt.cpp +51 -14
- data/vendor/faiss/faiss/utils/simd_impl/rabitq_neon.cpp +297 -10
- data/vendor/faiss/faiss/utils/simd_impl/rabitq_rvv.cpp +0 -4
- data/vendor/faiss/faiss/utils/simd_impl/super_kmeans_kernels_avx2.cpp +0 -4
- data/vendor/faiss/faiss/utils/simd_impl/super_kmeans_kernels_avx512.cpp +0 -4
- data/vendor/faiss/faiss/utils/simd_impl/super_kmeans_kernels_sve.cpp +0 -4
- data/vendor/faiss/faiss/utils/simd_levels.cpp +19 -3
- data/vendor/faiss/faiss/utils/simd_levels.h +74 -1
- data/vendor/faiss/faiss/utils/transpose/transpose-avx2-inl.h +0 -4
- data/vendor/faiss/faiss/utils/transpose/transpose-avx512-inl.h +0 -4
- data/vendor/faiss/faiss/utils/utils.cpp +80 -0
- data/vendor/faiss/faiss/utils/utils.h +24 -0
- metadata +33 -1
|
@@ -5,16 +5,21 @@
|
|
|
5
5
|
* LICENSE file in the root directory of this source tree.
|
|
6
6
|
*/
|
|
7
7
|
|
|
8
|
-
#ifdef COMPILE_SIMD_RISCV_RVV
|
|
9
|
-
|
|
10
8
|
#include <faiss/impl/scalar_quantizer/codecs.h>
|
|
11
9
|
#include <faiss/impl/scalar_quantizer/distance_computers.h>
|
|
12
10
|
#include <faiss/impl/scalar_quantizer/quantizers.h>
|
|
13
11
|
#include <faiss/impl/scalar_quantizer/scanners.h>
|
|
14
12
|
#include <faiss/impl/scalar_quantizer/similarities.h>
|
|
13
|
+
#include <faiss/impl/scalar_quantizer/sq8_coefficients.h>
|
|
14
|
+
#include <faiss/impl/scalar_quantizer/sq8_rvv_kernel.h>
|
|
15
15
|
|
|
16
16
|
#include <riscv_vector.h>
|
|
17
17
|
#include <cmath>
|
|
18
|
+
#include <cstring>
|
|
19
|
+
|
|
20
|
+
#if !defined(__riscv_zvfhmin) && !defined(__riscv_zvfh)
|
|
21
|
+
#error "RISCV_RVV scalar quantizers require Zvfhmin; compile this file with -march including _zvfhmin (e.g. rv64gcv_zvfhmin)"
|
|
22
|
+
#endif
|
|
18
23
|
|
|
19
24
|
namespace faiss {
|
|
20
25
|
|
|
@@ -82,39 +87,379 @@ struct QuantizerTemplate<
|
|
|
82
87
|
};
|
|
83
88
|
|
|
84
89
|
template <>
|
|
85
|
-
struct QuantizerFP16<SIMDLevel::RISCV_RVV>
|
|
90
|
+
struct QuantizerFP16<SIMDLevel::RISCV_RVV>
|
|
91
|
+
: QuantizerFP16Base<SIMDLevel::RISCV_RVV> {
|
|
86
92
|
QuantizerFP16(size_t d, const std::vector<float>& trained)
|
|
87
|
-
:
|
|
93
|
+
: QuantizerFP16Base<SIMDLevel::RISCV_RVV>(d, trained) {}
|
|
94
|
+
|
|
95
|
+
void decode_vector(const uint8_t* code, float* x) const final {
|
|
96
|
+
const auto* code16 = reinterpret_cast<const uint16_t*>(code);
|
|
97
|
+
size_t i = 0;
|
|
98
|
+
const size_t d = this->d;
|
|
99
|
+
while (i < d) {
|
|
100
|
+
const size_t vl = __riscv_vsetvl_e16m4(d - i);
|
|
101
|
+
vuint16m4_t vh = __riscv_vle16_v_u16m4(code16 + i, vl);
|
|
102
|
+
vbool4_t is_nan = __riscv_vmsgtu_vx_u16m4_b4(
|
|
103
|
+
__riscv_vand_vx_u16m4(vh, 0x7FFF, vl), 0x7C00, vl);
|
|
104
|
+
vfloat32m8_t vf = __riscv_vfwcvt_f_f_v_f32m8_m(
|
|
105
|
+
__riscv_vmnot_m_b4(is_nan, vl),
|
|
106
|
+
__riscv_vreinterpret_v_u16m4_f16m4(vh),
|
|
107
|
+
vl);
|
|
108
|
+
vf = __riscv_vfsub_vf_f32m8_m(
|
|
109
|
+
__riscv_vmnot_m_b4(is_nan, vl), vf, 0.0f, vl);
|
|
110
|
+
if (__riscv_vfirst_m_b4(is_nan, vl) >= 0) {
|
|
111
|
+
vint32m8_t bits = __riscv_vsext_vf2_i32m8(
|
|
112
|
+
__riscv_vreinterpret_v_u16m4_i16m4(vh), vl);
|
|
113
|
+
bits = __riscv_vsll_vx_i32m8(bits, 13, vl);
|
|
114
|
+
bits = __riscv_vand_vx_i32m8(bits, -0x70002000, vl);
|
|
115
|
+
bits = __riscv_vor_vx_i32m8(bits, 0x70000000, vl);
|
|
116
|
+
vf = __riscv_vmerge_vvm_f32m8(
|
|
117
|
+
vf,
|
|
118
|
+
__riscv_vreinterpret_v_i32m8_f32m8(bits),
|
|
119
|
+
is_nan,
|
|
120
|
+
vl);
|
|
121
|
+
}
|
|
122
|
+
__riscv_vse32_v_f32m8(x + i, vf, vl);
|
|
123
|
+
i += vl;
|
|
124
|
+
}
|
|
125
|
+
}
|
|
88
126
|
};
|
|
89
127
|
|
|
90
128
|
template <>
|
|
91
|
-
struct QuantizerBF16<SIMDLevel::RISCV_RVV>
|
|
129
|
+
struct QuantizerBF16<SIMDLevel::RISCV_RVV>
|
|
130
|
+
: QuantizerBF16Base<SIMDLevel::RISCV_RVV> {
|
|
92
131
|
QuantizerBF16(size_t d, const std::vector<float>& trained)
|
|
93
|
-
:
|
|
132
|
+
: QuantizerBF16Base<SIMDLevel::RISCV_RVV>(d, trained) {}
|
|
133
|
+
|
|
134
|
+
void decode_vector(const uint8_t* code, float* x) const final {
|
|
135
|
+
const auto* code16 = reinterpret_cast<const uint16_t*>(code);
|
|
136
|
+
size_t i = 0;
|
|
137
|
+
const size_t d = this->d;
|
|
138
|
+
while (i < d) {
|
|
139
|
+
const size_t vl = __riscv_vsetvl_e16m4(d - i);
|
|
140
|
+
vuint16m4_t v = __riscv_vle16_v_u16m4(code16 + i, vl);
|
|
141
|
+
vuint32m8_t w = __riscv_vzext_vf2_u32m8(v, vl);
|
|
142
|
+
w = __riscv_vsll_vx_u32m8(w, 16, vl);
|
|
143
|
+
__riscv_vse32_v_f32m8(
|
|
144
|
+
x + i, __riscv_vreinterpret_v_u32m8_f32m8(w), vl);
|
|
145
|
+
i += vl;
|
|
146
|
+
}
|
|
147
|
+
}
|
|
94
148
|
};
|
|
95
149
|
|
|
96
150
|
template <>
|
|
97
151
|
struct Quantizer8bitDirect<SIMDLevel::RISCV_RVV>
|
|
98
|
-
:
|
|
152
|
+
: Quantizer8bitDirectBase<SIMDLevel::RISCV_RVV> {
|
|
99
153
|
Quantizer8bitDirect(size_t d, const std::vector<float>& trained)
|
|
100
|
-
:
|
|
154
|
+
: Quantizer8bitDirectBase<SIMDLevel::RISCV_RVV>(d, trained) {}
|
|
155
|
+
|
|
156
|
+
void decode_vector(const uint8_t* code, float* x) const final {
|
|
157
|
+
size_t i = 0;
|
|
158
|
+
const size_t d = this->d;
|
|
159
|
+
while (i < d) {
|
|
160
|
+
const size_t vl = __riscv_vsetvl_e8m2(d - i);
|
|
161
|
+
vuint8m2_t b = __riscv_vle8_v_u8m2(code + i, vl);
|
|
162
|
+
vuint16m4_t u = __riscv_vzext_vf2_u16m4(b, vl);
|
|
163
|
+
vfloat32m8_t f = __riscv_vfwcvt_f_xu_v_f32m8(u, vl);
|
|
164
|
+
__riscv_vse32_v_f32m8(x + i, f, vl);
|
|
165
|
+
i += vl;
|
|
166
|
+
}
|
|
167
|
+
}
|
|
168
|
+
|
|
169
|
+
void encode_vector(const float* x, uint8_t* code) const final {
|
|
170
|
+
std::size_t i = 0;
|
|
171
|
+
while (i < this->d) {
|
|
172
|
+
const std::size_t vl = __riscv_vsetvl_e32m4(this->d - i);
|
|
173
|
+
const vfloat32m4_t input = __riscv_vle32_v_f32m4(x + i, vl);
|
|
174
|
+
const vuint32m4_t converted =
|
|
175
|
+
__riscv_vfcvt_rtz_xu_f_v_u32m4(input, vl);
|
|
176
|
+
const vuint16m2_t half = __riscv_vnsrl_wx_u16m2(converted, 0, vl);
|
|
177
|
+
const vuint8m1_t bytes = __riscv_vnsrl_wx_u8m1(half, 0, vl);
|
|
178
|
+
__riscv_vse8_v_u8m1(code + i, bytes, vl);
|
|
179
|
+
i += vl;
|
|
180
|
+
}
|
|
181
|
+
}
|
|
101
182
|
};
|
|
102
183
|
|
|
103
184
|
template <>
|
|
104
185
|
struct Quantizer8bitDirectSigned<SIMDLevel::RISCV_RVV>
|
|
105
|
-
:
|
|
186
|
+
: Quantizer8bitDirectSignedBase<SIMDLevel::RISCV_RVV> {
|
|
106
187
|
Quantizer8bitDirectSigned(size_t d, const std::vector<float>& trained)
|
|
107
|
-
:
|
|
188
|
+
: Quantizer8bitDirectSignedBase<SIMDLevel::RISCV_RVV>(d, trained) {}
|
|
189
|
+
|
|
190
|
+
void decode_vector(const uint8_t* code, float* x) const final {
|
|
191
|
+
size_t i = 0;
|
|
192
|
+
const size_t d = this->d;
|
|
193
|
+
while (i < d) {
|
|
194
|
+
const size_t vl = __riscv_vsetvl_e8m2(d - i);
|
|
195
|
+
vint8m2_t v = __riscv_vle8_v_i8m2(
|
|
196
|
+
reinterpret_cast<const int8_t*>(code) + i, vl);
|
|
197
|
+
v = __riscv_vxor_vx_i8m2(v, -128, vl);
|
|
198
|
+
vint16m4_t v16 = __riscv_vsext_vf2_i16m4(v, vl);
|
|
199
|
+
vfloat32m8_t f = __riscv_vfwcvt_f_x_v_f32m8(v16, vl);
|
|
200
|
+
__riscv_vse32_v_f32m8(x + i, f, vl);
|
|
201
|
+
i += vl;
|
|
202
|
+
}
|
|
203
|
+
}
|
|
204
|
+
};
|
|
205
|
+
|
|
206
|
+
template <>
|
|
207
|
+
struct QuantizerLloydMax<1, SIMDLevel::RISCV_RVV>
|
|
208
|
+
: QuantizerLloydMaxBase<1, SIMDLevel::RISCV_RVV> {
|
|
209
|
+
using Base = QuantizerLloydMaxBase<1, SIMDLevel::RISCV_RVV>;
|
|
210
|
+
|
|
211
|
+
QuantizerLloydMax(size_t d, const std::vector<float>& trained)
|
|
212
|
+
: Base(d, trained) {}
|
|
213
|
+
|
|
214
|
+
FAISS_ALWAYS_INLINE vfloat32m8_t
|
|
215
|
+
reconstruct_m8_components(const uint8_t* code, size_t i, size_t vl) const {
|
|
216
|
+
const size_t byte_base = i >> 3;
|
|
217
|
+
const size_t byte_vl = ((i + vl - 1) >> 3) - byte_base + 1;
|
|
218
|
+
vuint8m2_t packed = __riscv_vle8_v_u8m2(code + byte_base, byte_vl);
|
|
219
|
+
vuint32m8_t packed32 = __riscv_vzext_vf4_u32m8(packed, byte_vl);
|
|
220
|
+
vuint32m8_t comp =
|
|
221
|
+
__riscv_vadd_vx_u32m8(__riscv_vid_v_u32m8(vl), i, vl);
|
|
222
|
+
vuint32m8_t rel = __riscv_vsub_vx_u32m8(
|
|
223
|
+
__riscv_vsrl_vx_u32m8(comp, 3, vl), byte_base, vl);
|
|
224
|
+
vuint32m8_t bytes = __riscv_vrgather_vv_u32m8(packed32, rel, vl);
|
|
225
|
+
vuint32m8_t shift = __riscv_vand_vx_u32m8(comp, 7, vl);
|
|
226
|
+
vuint32m8_t idx = __riscv_vand_vx_u32m8(
|
|
227
|
+
__riscv_vsrl_vv_u32m8(bytes, shift, vl), 1, vl);
|
|
228
|
+
vuint32m8_t off = __riscv_vsll_vx_u32m8(idx, 2, vl);
|
|
229
|
+
return __riscv_vluxei32_v_f32m8(this->centroids, off, vl);
|
|
230
|
+
}
|
|
231
|
+
|
|
232
|
+
void decode_vector(const uint8_t* code, float* x) const final {
|
|
233
|
+
size_t i = 0;
|
|
234
|
+
while (i < this->d) {
|
|
235
|
+
size_t vl = __riscv_vsetvl_e32m8(this->d - i);
|
|
236
|
+
vfloat32m8_t v = reconstruct_m8_components(code, i, vl);
|
|
237
|
+
__riscv_vse32_v_f32m8(x + i, v, vl);
|
|
238
|
+
i += vl;
|
|
239
|
+
}
|
|
240
|
+
}
|
|
241
|
+
};
|
|
242
|
+
|
|
243
|
+
template <>
|
|
244
|
+
struct QuantizerLloydMax<2, SIMDLevel::RISCV_RVV>
|
|
245
|
+
: QuantizerLloydMaxBase<2, SIMDLevel::RISCV_RVV> {
|
|
246
|
+
using Base = QuantizerLloydMaxBase<2, SIMDLevel::RISCV_RVV>;
|
|
247
|
+
|
|
248
|
+
QuantizerLloydMax(size_t d, const std::vector<float>& trained)
|
|
249
|
+
: Base(d, trained) {}
|
|
250
|
+
|
|
251
|
+
FAISS_ALWAYS_INLINE vfloat32m8_t
|
|
252
|
+
reconstruct_m8_components(const uint8_t* code, size_t i, size_t vl) const {
|
|
253
|
+
const size_t byte_base = i >> 2;
|
|
254
|
+
const size_t byte_vl = ((i + vl - 1) >> 2) - byte_base + 1;
|
|
255
|
+
vuint8m2_t packed = __riscv_vle8_v_u8m2(code + byte_base, byte_vl);
|
|
256
|
+
vuint32m8_t packed32 = __riscv_vzext_vf4_u32m8(packed, byte_vl);
|
|
257
|
+
vuint32m8_t comp =
|
|
258
|
+
__riscv_vadd_vx_u32m8(__riscv_vid_v_u32m8(vl), i, vl);
|
|
259
|
+
vuint32m8_t rel = __riscv_vsub_vx_u32m8(
|
|
260
|
+
__riscv_vsrl_vx_u32m8(comp, 2, vl), byte_base, vl);
|
|
261
|
+
vuint32m8_t bytes = __riscv_vrgather_vv_u32m8(packed32, rel, vl);
|
|
262
|
+
vuint32m8_t shift = __riscv_vsll_vx_u32m8(
|
|
263
|
+
__riscv_vand_vx_u32m8(comp, 3, vl), 1, vl);
|
|
264
|
+
vuint32m8_t idx = __riscv_vand_vx_u32m8(
|
|
265
|
+
__riscv_vsrl_vv_u32m8(bytes, shift, vl), 3, vl);
|
|
266
|
+
vuint32m8_t off = __riscv_vsll_vx_u32m8(idx, 2, vl);
|
|
267
|
+
return __riscv_vluxei32_v_f32m8(this->centroids, off, vl);
|
|
268
|
+
}
|
|
269
|
+
|
|
270
|
+
void decode_vector(const uint8_t* code, float* x) const final {
|
|
271
|
+
size_t i = 0;
|
|
272
|
+
while (i < this->d) {
|
|
273
|
+
size_t vl = __riscv_vsetvl_e32m8(this->d - i);
|
|
274
|
+
vfloat32m8_t v = reconstruct_m8_components(code, i, vl);
|
|
275
|
+
__riscv_vse32_v_f32m8(x + i, v, vl);
|
|
276
|
+
i += vl;
|
|
277
|
+
}
|
|
278
|
+
}
|
|
279
|
+
};
|
|
280
|
+
|
|
281
|
+
template <>
|
|
282
|
+
struct QuantizerLloydMax<3, SIMDLevel::RISCV_RVV>
|
|
283
|
+
: QuantizerLloydMaxBase<3, SIMDLevel::RISCV_RVV> {
|
|
284
|
+
using Base = QuantizerLloydMaxBase<3, SIMDLevel::RISCV_RVV>;
|
|
285
|
+
|
|
286
|
+
QuantizerLloydMax(size_t d, const std::vector<float>& trained)
|
|
287
|
+
: Base(d, trained) {}
|
|
288
|
+
|
|
289
|
+
FAISS_ALWAYS_INLINE vfloat32m8_t
|
|
290
|
+
reconstruct_m8_components(const uint8_t* code, size_t i, size_t vl) const {
|
|
291
|
+
vuint32m8_t idx0 =
|
|
292
|
+
__riscv_vadd_vx_u32m8(__riscv_vid_v_u32m8(vl), i, vl);
|
|
293
|
+
vuint32m8_t bitpos = __riscv_vmul_vx_u32m8(idx0, 3, vl);
|
|
294
|
+
vuint32m8_t byteoff = __riscv_vsrl_vx_u32m8(bitpos, 3, vl);
|
|
295
|
+
vuint32m8_t shift = __riscv_vand_vx_u32m8(bitpos, 7, vl);
|
|
296
|
+
size_t last = i + vl - 1;
|
|
297
|
+
size_t last_used = (3 * last + 2) >> 3;
|
|
298
|
+
vuint8m2_t lo = __riscv_vluxei32_v_u8m2(code, byteoff, vl);
|
|
299
|
+
vuint8m2_t hi = __riscv_vluxei32_v_u8m2(
|
|
300
|
+
code,
|
|
301
|
+
__riscv_vminu_vx_u32m8(
|
|
302
|
+
__riscv_vadd_vx_u32m8(byteoff, 1, vl), last_used, vl),
|
|
303
|
+
vl);
|
|
304
|
+
vuint32m8_t w = __riscv_vor_vv_u32m8(
|
|
305
|
+
__riscv_vzext_vf4_u32m8(lo, vl),
|
|
306
|
+
__riscv_vsll_vx_u32m8(__riscv_vzext_vf4_u32m8(hi, vl), 8, vl),
|
|
307
|
+
vl);
|
|
308
|
+
vuint32m8_t idx = __riscv_vand_vx_u32m8(
|
|
309
|
+
__riscv_vsrl_vv_u32m8(w, shift, vl), 7, vl);
|
|
310
|
+
vuint32m8_t off = __riscv_vsll_vx_u32m8(idx, 2, vl);
|
|
311
|
+
return __riscv_vluxei32_v_f32m8(this->centroids, off, vl);
|
|
312
|
+
}
|
|
313
|
+
|
|
314
|
+
void decode_vector(const uint8_t* code, float* x) const final {
|
|
315
|
+
size_t i = 0;
|
|
316
|
+
while (i < this->d) {
|
|
317
|
+
size_t vl = __riscv_vsetvl_e32m8(this->d - i);
|
|
318
|
+
vfloat32m8_t v = reconstruct_m8_components(code, i, vl);
|
|
319
|
+
__riscv_vse32_v_f32m8(x + i, v, vl);
|
|
320
|
+
i += vl;
|
|
321
|
+
}
|
|
322
|
+
}
|
|
323
|
+
};
|
|
324
|
+
|
|
325
|
+
template <>
|
|
326
|
+
struct QuantizerLloydMax<4, SIMDLevel::RISCV_RVV>
|
|
327
|
+
: QuantizerLloydMaxBase<4, SIMDLevel::RISCV_RVV> {
|
|
328
|
+
using Base = QuantizerLloydMaxBase<4, SIMDLevel::RISCV_RVV>;
|
|
329
|
+
|
|
330
|
+
QuantizerLloydMax(size_t d, const std::vector<float>& trained)
|
|
331
|
+
: Base(d, trained) {}
|
|
332
|
+
|
|
333
|
+
FAISS_ALWAYS_INLINE vfloat32m8_t
|
|
334
|
+
reconstruct_m8_components(const uint8_t* code, size_t i, size_t vl) const {
|
|
335
|
+
const size_t byte_base = i >> 1;
|
|
336
|
+
const size_t byte_vl = ((i + vl - 1) >> 1) - byte_base + 1;
|
|
337
|
+
vuint8m2_t packed = __riscv_vle8_v_u8m2(code + byte_base, byte_vl);
|
|
338
|
+
vuint32m8_t packed32 = __riscv_vzext_vf4_u32m8(packed, byte_vl);
|
|
339
|
+
vuint32m8_t comp =
|
|
340
|
+
__riscv_vadd_vx_u32m8(__riscv_vid_v_u32m8(vl), i, vl);
|
|
341
|
+
vuint32m8_t rel = __riscv_vsub_vx_u32m8(
|
|
342
|
+
__riscv_vsrl_vx_u32m8(comp, 1, vl), byte_base, vl);
|
|
343
|
+
vuint32m8_t bytes = __riscv_vrgather_vv_u32m8(packed32, rel, vl);
|
|
344
|
+
vuint32m8_t lo = __riscv_vand_vx_u32m8(bytes, 0xf, vl);
|
|
345
|
+
vuint32m8_t hi = __riscv_vsrl_vx_u32m8(bytes, 4, vl);
|
|
346
|
+
vbool4_t odd = __riscv_vmsne_vx_u32m8_b4(
|
|
347
|
+
__riscv_vand_vx_u32m8(comp, 1, vl), 0, vl);
|
|
348
|
+
vuint32m8_t idx = __riscv_vmerge_vvm_u32m8(lo, hi, odd, vl);
|
|
349
|
+
vuint32m8_t off = __riscv_vsll_vx_u32m8(idx, 2, vl);
|
|
350
|
+
return __riscv_vluxei32_v_f32m8(this->centroids, off, vl);
|
|
351
|
+
}
|
|
352
|
+
|
|
353
|
+
void decode_vector(const uint8_t* code, float* x) const final {
|
|
354
|
+
size_t i = 0;
|
|
355
|
+
while (i < this->d) {
|
|
356
|
+
size_t vl = __riscv_vsetvl_e32m8(this->d - i);
|
|
357
|
+
vfloat32m8_t v = reconstruct_m8_components(code, i, vl);
|
|
358
|
+
__riscv_vse32_v_f32m8(x + i, v, vl);
|
|
359
|
+
i += vl;
|
|
360
|
+
}
|
|
361
|
+
}
|
|
362
|
+
};
|
|
363
|
+
|
|
364
|
+
template <>
|
|
365
|
+
struct QuantizerLloydMax<8, SIMDLevel::RISCV_RVV>
|
|
366
|
+
: QuantizerLloydMaxBase<8, SIMDLevel::RISCV_RVV> {
|
|
367
|
+
using Base = QuantizerLloydMaxBase<8, SIMDLevel::RISCV_RVV>;
|
|
368
|
+
|
|
369
|
+
QuantizerLloydMax(size_t d, const std::vector<float>& trained)
|
|
370
|
+
: Base(d, trained) {}
|
|
371
|
+
|
|
372
|
+
FAISS_ALWAYS_INLINE vfloat32m8_t
|
|
373
|
+
reconstruct_m8_components(const uint8_t* code, size_t i, size_t vl) const {
|
|
374
|
+
vuint8m2_t vb = __riscv_vle8_v_u8m2(code + i, vl);
|
|
375
|
+
vuint32m8_t off =
|
|
376
|
+
__riscv_vsll_vx_u32m8(__riscv_vzext_vf4_u32m8(vb, vl), 2, vl);
|
|
377
|
+
return __riscv_vluxei32_v_f32m8(this->centroids, off, vl);
|
|
378
|
+
}
|
|
379
|
+
|
|
380
|
+
void decode_vector(const uint8_t* code, float* x) const final {
|
|
381
|
+
size_t i = 0;
|
|
382
|
+
while (i < this->d) {
|
|
383
|
+
size_t vl = __riscv_vsetvl_e32m8(this->d - i);
|
|
384
|
+
vfloat32m8_t v = reconstruct_m8_components(code, i, vl);
|
|
385
|
+
__riscv_vse32_v_f32m8(x + i, v, vl);
|
|
386
|
+
i += vl;
|
|
387
|
+
}
|
|
388
|
+
}
|
|
108
389
|
};
|
|
109
390
|
|
|
110
391
|
template <>
|
|
111
392
|
struct SimilarityL2<SIMDLevel::RISCV_RVV> : SimilarityL2<SIMDLevel::NONE> {
|
|
112
393
|
using SimilarityL2<SIMDLevel::NONE>::SimilarityL2;
|
|
394
|
+
|
|
395
|
+
static constexpr SIMDLevel simd_level = SIMDLevel::RISCV_RVV;
|
|
396
|
+
|
|
397
|
+
FAISS_ALWAYS_INLINE void begin_m8() {
|
|
398
|
+
yi = y;
|
|
399
|
+
}
|
|
400
|
+
|
|
401
|
+
static FAISS_ALWAYS_INLINE vfloat32m8_t zero_m8(size_t vl) {
|
|
402
|
+
return __riscv_vfmv_v_f_f32m8(0.0f, vl);
|
|
403
|
+
}
|
|
404
|
+
|
|
405
|
+
FAISS_ALWAYS_INLINE vfloat32m8_t
|
|
406
|
+
add_m8_components(vfloat32m8_t accu, vfloat32m8_t x, size_t vl) {
|
|
407
|
+
vfloat32m8_t yiv = __riscv_vle32_v_f32m8(yi, vl);
|
|
408
|
+
yi += vl;
|
|
409
|
+
vfloat32m8_t tmp = __riscv_vfsub_vv_f32m8(yiv, x, vl);
|
|
410
|
+
return __riscv_vfmacc_vv_f32m8_tu(accu, tmp, tmp, vl);
|
|
411
|
+
}
|
|
412
|
+
|
|
413
|
+
static FAISS_ALWAYS_INLINE vfloat32m8_t add_m8_components_2(
|
|
414
|
+
vfloat32m8_t accu,
|
|
415
|
+
vfloat32m8_t x,
|
|
416
|
+
vfloat32m8_t y_2,
|
|
417
|
+
size_t vl) {
|
|
418
|
+
vfloat32m8_t tmp = __riscv_vfsub_vv_f32m8(y_2, x, vl);
|
|
419
|
+
return __riscv_vfmacc_vv_f32m8_tu(accu, tmp, tmp, vl);
|
|
420
|
+
}
|
|
421
|
+
|
|
422
|
+
static FAISS_ALWAYS_INLINE float result_m8(vfloat32m8_t accu, size_t vl) {
|
|
423
|
+
vfloat32m1_t zero = __riscv_vfmv_v_f_f32m1(0.0f, 1);
|
|
424
|
+
vfloat32m1_t sum = __riscv_vfredusum_vs_f32m8_f32m1(accu, zero, vl);
|
|
425
|
+
return __riscv_vfmv_f_s_f32m1_f32(sum);
|
|
426
|
+
}
|
|
113
427
|
};
|
|
114
428
|
|
|
115
429
|
template <>
|
|
116
430
|
struct SimilarityIP<SIMDLevel::RISCV_RVV> : SimilarityIP<SIMDLevel::NONE> {
|
|
117
431
|
using SimilarityIP<SIMDLevel::NONE>::SimilarityIP;
|
|
432
|
+
|
|
433
|
+
static constexpr SIMDLevel simd_level = SIMDLevel::RISCV_RVV;
|
|
434
|
+
|
|
435
|
+
FAISS_ALWAYS_INLINE void begin_m8() {
|
|
436
|
+
yi = y;
|
|
437
|
+
}
|
|
438
|
+
|
|
439
|
+
static FAISS_ALWAYS_INLINE vfloat32m8_t zero_m8(size_t vl) {
|
|
440
|
+
return __riscv_vfmv_v_f_f32m8(0.0f, vl);
|
|
441
|
+
}
|
|
442
|
+
|
|
443
|
+
FAISS_ALWAYS_INLINE vfloat32m8_t
|
|
444
|
+
add_m8_components(vfloat32m8_t accu, vfloat32m8_t x, size_t vl) {
|
|
445
|
+
vfloat32m8_t yiv = __riscv_vle32_v_f32m8(yi, vl);
|
|
446
|
+
yi += vl;
|
|
447
|
+
return __riscv_vfmacc_vv_f32m8_tu(accu, yiv, x, vl);
|
|
448
|
+
}
|
|
449
|
+
|
|
450
|
+
static FAISS_ALWAYS_INLINE vfloat32m8_t add_m8_components_2(
|
|
451
|
+
vfloat32m8_t accu,
|
|
452
|
+
vfloat32m8_t x1,
|
|
453
|
+
vfloat32m8_t x2,
|
|
454
|
+
size_t vl) {
|
|
455
|
+
return __riscv_vfmacc_vv_f32m8_tu(accu, x1, x2, vl);
|
|
456
|
+
}
|
|
457
|
+
|
|
458
|
+
static FAISS_ALWAYS_INLINE float result_m8(vfloat32m8_t accu, size_t vl) {
|
|
459
|
+
vfloat32m1_t zero = __riscv_vfmv_v_f_f32m1(0.0f, 1);
|
|
460
|
+
vfloat32m1_t sum = __riscv_vfredusum_vs_f32m8_f32m1(accu, zero, vl);
|
|
461
|
+
return __riscv_vfmv_f_s_f32m1_f32(sum);
|
|
462
|
+
}
|
|
118
463
|
};
|
|
119
464
|
|
|
120
465
|
/*************************************************************************
|
|
@@ -125,7 +470,14 @@ struct SimilarityIP<SIMDLevel::RISCV_RVV> : SimilarityIP<SIMDLevel::NONE> {
|
|
|
125
470
|
* falls through to scalar code. Callers and the dispatcher don't know or care.
|
|
126
471
|
************************************************************************/
|
|
127
472
|
|
|
473
|
+
template <class Quantizer>
|
|
474
|
+
inline constexpr bool has_reconstruct_m8_v =
|
|
475
|
+
requires(const Quantizer& q, const uint8_t* code, size_t i, size_t vl) {
|
|
476
|
+
q.reconstruct_m8_components(code, i, vl);
|
|
477
|
+
};
|
|
478
|
+
|
|
128
479
|
template <class Quantizer, class Similarity>
|
|
480
|
+
requires(!has_reconstruct_m8_v<Quantizer>)
|
|
129
481
|
struct DCTemplate<Quantizer, Similarity, SIMDLevel::RISCV_RVV>
|
|
130
482
|
: DCTemplate<Quantizer, Similarity, SIMDLevel::NONE> {
|
|
131
483
|
using Base = DCTemplate<Quantizer, Similarity, SIMDLevel::NONE>;
|
|
@@ -139,20 +491,102 @@ struct DistanceComputerByte<Similarity, SIMDLevel::RISCV_RVV>
|
|
|
139
491
|
using Base::Base;
|
|
140
492
|
};
|
|
141
493
|
|
|
142
|
-
|
|
143
|
-
|
|
144
|
-
|
|
145
|
-
|
|
146
|
-
|
|
147
|
-
|
|
148
|
-
|
|
149
|
-
|
|
150
|
-
|
|
151
|
-
|
|
152
|
-
|
|
153
|
-
*
|
|
154
|
-
|
|
155
|
-
|
|
494
|
+
template <class Quantizer, class Similarity>
|
|
495
|
+
requires(has_reconstruct_m8_v<Quantizer>)
|
|
496
|
+
struct DCTemplate<Quantizer, Similarity, SIMDLevel::RISCV_RVV>
|
|
497
|
+
: SQDistanceComputer {
|
|
498
|
+
using Sim = Similarity;
|
|
499
|
+
|
|
500
|
+
Quantizer quant;
|
|
501
|
+
|
|
502
|
+
DCTemplate(size_t d, const std::vector<float>& trained)
|
|
503
|
+
: quant(d, trained) {}
|
|
504
|
+
|
|
505
|
+
float compute_distance(const float* x, const uint8_t* code) const {
|
|
506
|
+
if (quant.d == 0) {
|
|
507
|
+
return 0.0f;
|
|
508
|
+
}
|
|
509
|
+
Similarity sim(x);
|
|
510
|
+
sim.begin_m8();
|
|
511
|
+
const size_t first_vl = __riscv_vsetvl_e32m8(quant.d);
|
|
512
|
+
vfloat32m8_t accu = Sim::zero_m8(first_vl);
|
|
513
|
+
size_t i = 0;
|
|
514
|
+
while (i < quant.d) {
|
|
515
|
+
size_t vl = __riscv_vsetvl_e32m8(quant.d - i);
|
|
516
|
+
vfloat32m8_t xi = quant.reconstruct_m8_components(code, i, vl);
|
|
517
|
+
accu = sim.add_m8_components(accu, xi, vl);
|
|
518
|
+
i += vl;
|
|
519
|
+
}
|
|
520
|
+
return Sim::result_m8(accu, first_vl);
|
|
521
|
+
}
|
|
522
|
+
|
|
523
|
+
float compute_code_distance(const uint8_t* code1, const uint8_t* code2)
|
|
524
|
+
const {
|
|
525
|
+
if (quant.d == 0) {
|
|
526
|
+
return 0.0f;
|
|
527
|
+
}
|
|
528
|
+
Similarity sim(nullptr);
|
|
529
|
+
sim.begin_m8();
|
|
530
|
+
const size_t first_vl = __riscv_vsetvl_e32m8(quant.d);
|
|
531
|
+
vfloat32m8_t accu = Sim::zero_m8(first_vl);
|
|
532
|
+
size_t i = 0;
|
|
533
|
+
while (i < quant.d) {
|
|
534
|
+
size_t vl = __riscv_vsetvl_e32m8(quant.d - i);
|
|
535
|
+
vfloat32m8_t x1 = quant.reconstruct_m8_components(code1, i, vl);
|
|
536
|
+
vfloat32m8_t x2 = quant.reconstruct_m8_components(code2, i, vl);
|
|
537
|
+
accu = Sim::add_m8_components_2(accu, x1, x2, vl);
|
|
538
|
+
i += vl;
|
|
539
|
+
}
|
|
540
|
+
return Sim::result_m8(accu, first_vl);
|
|
541
|
+
}
|
|
542
|
+
|
|
543
|
+
void set_query(const float* x) final {
|
|
544
|
+
q = x;
|
|
545
|
+
}
|
|
546
|
+
|
|
547
|
+
float symmetric_dis(idx_t i, idx_t j) override {
|
|
548
|
+
return compute_code_distance(
|
|
549
|
+
codes + i * code_size, codes + j * code_size);
|
|
550
|
+
}
|
|
551
|
+
|
|
552
|
+
float query_to_code(const uint8_t* code) const final {
|
|
553
|
+
return compute_distance(q, code);
|
|
554
|
+
}
|
|
555
|
+
|
|
556
|
+
void query_to_codes_batch_4(
|
|
557
|
+
const uint8_t* code_0,
|
|
558
|
+
const uint8_t* code_1,
|
|
559
|
+
const uint8_t* code_2,
|
|
560
|
+
const uint8_t* code_3,
|
|
561
|
+
float& dis0,
|
|
562
|
+
float& dis1,
|
|
563
|
+
float& dis2,
|
|
564
|
+
float& dis3) const final {
|
|
565
|
+
dis0 = compute_distance(q, code_0);
|
|
566
|
+
dis1 = compute_distance(q, code_1);
|
|
567
|
+
dis2 = compute_distance(q, code_2);
|
|
568
|
+
dis3 = compute_distance(q, code_3);
|
|
569
|
+
}
|
|
570
|
+
};
|
|
571
|
+
|
|
572
|
+
// * Fast path — QT_4bit_uniform + L2 (float domain)
|
|
573
|
+
// *
|
|
574
|
+
// * 4-bit UNIFORM scaling: every component reconstructs as
|
|
575
|
+
// * recon(c) = vmin + vdiff * (c + 0.5) / 15 = c0 + a * c
|
|
576
|
+
// * with a = vdiff / 15 and c0 = vmin + 0.5 * a.
|
|
577
|
+
// *
|
|
578
|
+
// * L2 distance is evaluated in the exact affine float domain,
|
|
579
|
+
// * L2(q, code) = sum_i (q_i - recon(c_i))^2 = sum_i (e_i - a*c_i)^2
|
|
580
|
+
// * with e_i = q_i - c0, so the ORIGINAL float query is retained (no
|
|
581
|
+
// * integer-grid pre-quantization) and the result matches the scalar
|
|
582
|
+
// * NONE reference to float-reassociation precision. vdiff == 0
|
|
583
|
+
// * degenerates naturally (a == 0 -> L2 = sum_i (q_i - vmin)^2).
|
|
584
|
+
// *
|
|
585
|
+
// * The packed 4-bit code interleaves dimensions (even -> low nibble,
|
|
586
|
+
// * odd -> high nibble), so set_query deinterleaves e_i into e_lo/e_hi
|
|
587
|
+
// * once per query. The hot loop unpacks nibbles, converts to f32 and
|
|
588
|
+
// * accumulates (e - a*c)^2 via the scalar-operand vfnmsac.vf + vfmacc.
|
|
589
|
+
// ************************************************************************/
|
|
156
590
|
|
|
157
591
|
template <>
|
|
158
592
|
struct DCTemplate<
|
|
@@ -165,123 +599,119 @@ struct DCTemplate<
|
|
|
165
599
|
using Sim = SimilarityL2<SIMDLevel::RISCV_RVV>;
|
|
166
600
|
|
|
167
601
|
size_t d;
|
|
168
|
-
|
|
169
|
-
float vdiff
|
|
170
|
-
float
|
|
171
|
-
std::vector<
|
|
172
|
-
std::vector<uint8_t> q_hi;
|
|
602
|
+
size_t half; // bytes per code = ceil(d/2)
|
|
603
|
+
float a; // vdiff / 15
|
|
604
|
+
float c0; // vmin + 0.5 * a
|
|
605
|
+
std::vector<float> e_lo, e_hi; // q_i - c0, deinterleaved per query
|
|
173
606
|
|
|
174
607
|
DCTemplate(size_t d_in, const std::vector<float>& trained)
|
|
175
608
|
: d(d_in),
|
|
176
|
-
|
|
177
|
-
|
|
178
|
-
|
|
179
|
-
|
|
180
|
-
|
|
181
|
-
final_scale_sq = final_scale * final_scale;
|
|
609
|
+
half((d_in + 1) / 2),
|
|
610
|
+
e_lo(half, 0.0f),
|
|
611
|
+
e_hi(half, 0.0f) {
|
|
612
|
+
a = trained[1] / 15.0f;
|
|
613
|
+
c0 = trained[0] + 0.5f * a;
|
|
182
614
|
}
|
|
183
615
|
|
|
184
616
|
void set_query(const float* x) final {
|
|
185
|
-
|
|
186
|
-
|
|
617
|
+
q = x;
|
|
618
|
+
// e_i = q_i - c0, deinterleaved. Padding slot (odd d) keeps 0.
|
|
187
619
|
for (size_t i = 0; i < d; i++) {
|
|
188
|
-
float val = (x[i] - vmin) * inv_scale;
|
|
189
|
-
int code = static_cast<int>(val);
|
|
190
|
-
if (code < 0) {
|
|
191
|
-
code = 0;
|
|
192
|
-
}
|
|
193
|
-
if (code > 15) {
|
|
194
|
-
code = 15;
|
|
195
|
-
}
|
|
196
620
|
if (i % 2 == 0) {
|
|
197
|
-
|
|
621
|
+
e_lo[i / 2] = x[i] - c0;
|
|
198
622
|
} else {
|
|
199
|
-
|
|
623
|
+
e_hi[i / 2] = x[i] - c0;
|
|
200
624
|
}
|
|
201
625
|
}
|
|
202
626
|
}
|
|
203
627
|
|
|
204
|
-
///
|
|
205
|
-
///
|
|
206
|
-
///
|
|
207
|
-
|
|
208
|
-
|
|
209
|
-
|
|
210
|
-
|
|
211
|
-
|
|
212
|
-
|
|
213
|
-
|
|
214
|
-
|
|
215
|
-
|
|
216
|
-
size_t vl = __riscv_vsetvl_e8m1(remaining);
|
|
217
|
-
if (vl & 1) {
|
|
218
|
-
vl -= 1; // keep even; tail handled on next iter or scalar
|
|
219
|
-
}
|
|
220
|
-
if (vl == 0) {
|
|
221
|
-
break;
|
|
222
|
-
}
|
|
223
|
-
const size_t byte_vl = vl / 2;
|
|
224
|
-
|
|
225
|
-
vuint8m1_t packed = __riscv_vle8_v_u8m1(code + i / 2, byte_vl);
|
|
226
|
-
vuint8m1_t ql = __riscv_vle8_v_u8m1(q_lo.data() + i / 2, byte_vl);
|
|
227
|
-
vuint8m1_t qh = __riscv_vle8_v_u8m1(q_hi.data() + i / 2, byte_vl);
|
|
228
|
-
|
|
229
|
-
vuint8m1_t lo_nib = __riscv_vand_vx_u8m1(packed, 0x0F, byte_vl);
|
|
230
|
-
vuint8m1_t hi_nib = __riscv_vsrl_vx_u8m1(packed, 4, byte_vl);
|
|
231
|
-
|
|
232
|
-
// |ql - lo| and |qh - hi| fit in u8 (values are in [0, 15]).
|
|
233
|
-
vuint8m1_t d_lo = __riscv_vsub_vv_u8m1(
|
|
234
|
-
__riscv_vmaxu_vv_u8m1(ql, lo_nib, byte_vl),
|
|
235
|
-
__riscv_vminu_vv_u8m1(ql, lo_nib, byte_vl),
|
|
236
|
-
byte_vl);
|
|
237
|
-
vuint8m1_t d_hi = __riscv_vsub_vv_u8m1(
|
|
238
|
-
__riscv_vmaxu_vv_u8m1(qh, hi_nib, byte_vl),
|
|
239
|
-
__riscv_vminu_vv_u8m1(qh, hi_nib, byte_vl),
|
|
240
|
-
byte_vl);
|
|
241
|
-
|
|
242
|
-
// Square via widening multiply (each byte squared fits in u16,
|
|
243
|
-
// since max byte value is 15 -> 225).
|
|
244
|
-
vuint16m2_t sq_lo = __riscv_vwmulu_vv_u16m2(d_lo, d_lo, byte_vl);
|
|
245
|
-
vuint16m2_t sq_hi = __riscv_vwmulu_vv_u16m2(d_hi, d_hi, byte_vl);
|
|
246
|
-
vuint16m2_t sq_sum = __riscv_vadd_vv_u16m2(sq_lo, sq_hi, byte_vl);
|
|
247
|
-
|
|
248
|
-
// Reduce to a scalar u32 (safe: byte_vl * 450 fits in u32 for
|
|
249
|
-
// any realistic d).
|
|
250
|
-
vuint32m1_t zero = __riscv_vmv_v_x_u32m1(0, 1);
|
|
251
|
-
vuint32m1_t red =
|
|
252
|
-
__riscv_vwredsumu_vs_u16m2_u32m1(sq_sum, zero, byte_vl);
|
|
253
|
-
acc += __riscv_vmv_x_s_u32m1_u32(red);
|
|
628
|
+
/// Exact affine float-domain L2 over the packed 4-bit code.
|
|
629
|
+
/// Per byte chunk: unpack nibbles (e8), widen+convert to f32, then
|
|
630
|
+
/// t = e - a*c (vfnmsac.vf, scalar a) and acc += t*t (vfmacc).
|
|
631
|
+
/// vsetvl hoisted; _tu on the tail keeps lanes past vt intact.
|
|
632
|
+
/// Odd d: the last byte's high nibble is padding, so it is excluded
|
|
633
|
+
/// from the vector domain (nbv = nb - 1) and the low nibble is
|
|
634
|
+
/// processed in scalar to stay exact for any producer.
|
|
635
|
+
float compute_l2(const uint8_t* code) const {
|
|
636
|
+
const size_t nb = half; // bytes per code
|
|
637
|
+
const size_t nbv = (d & 1) ? (nb - 1) : nb; // vector-domain bytes
|
|
638
|
+
float total = 0.0f;
|
|
639
|
+
size_t b = 0;
|
|
254
640
|
|
|
255
|
-
|
|
641
|
+
// Hoist vsetvl: VLMAX for e8m1 (== f32m4 lanes). d==0 guard.
|
|
642
|
+
const size_t vlb = __riscv_vsetvl_e8m1(nbv > 0 ? nbv : 1);
|
|
643
|
+
vfloat32m4_t acc = __riscv_vfmv_v_f_f32m4(0.0f, vlb);
|
|
644
|
+
|
|
645
|
+
for (; b + vlb <= nbv; b += vlb) {
|
|
646
|
+
vuint8m1_t packed = __riscv_vle8_v_u8m1(code + b, vlb);
|
|
647
|
+
vuint8m1_t lo = __riscv_vand_vx_u8m1(packed, 0x0F, vlb);
|
|
648
|
+
vuint8m1_t hi = __riscv_vsrl_vx_u8m1(packed, 4, vlb);
|
|
649
|
+
|
|
650
|
+
vfloat32m4_t clo = __riscv_vfcvt_f_xu_v_f32m4(
|
|
651
|
+
__riscv_vzext_vf4_u32m4(lo, vlb), vlb);
|
|
652
|
+
vfloat32m4_t t_lo = __riscv_vle32_v_f32m4(e_lo.data() + b, vlb);
|
|
653
|
+
t_lo = __riscv_vfnmsac_vf_f32m4(t_lo, a, clo, vlb);
|
|
654
|
+
acc = __riscv_vfmacc_vv_f32m4(acc, t_lo, t_lo, vlb);
|
|
655
|
+
|
|
656
|
+
vfloat32m4_t chi = __riscv_vfcvt_f_xu_v_f32m4(
|
|
657
|
+
__riscv_vzext_vf4_u32m4(hi, vlb), vlb);
|
|
658
|
+
vfloat32m4_t t_hi = __riscv_vle32_v_f32m4(e_hi.data() + b, vlb);
|
|
659
|
+
t_hi = __riscv_vfnmsac_vf_f32m4(t_hi, a, chi, vlb);
|
|
660
|
+
acc = __riscv_vfmacc_vv_f32m4(acc, t_hi, t_hi, vlb);
|
|
256
661
|
}
|
|
257
|
-
|
|
258
|
-
|
|
259
|
-
|
|
260
|
-
|
|
261
|
-
|
|
262
|
-
|
|
263
|
-
|
|
662
|
+
|
|
663
|
+
// Tail: fewer than vlb bytes left — one shorter-vl pass into the
|
|
664
|
+
// same accumulator; _tu keeps lanes past vt intact.
|
|
665
|
+
if (b < nbv) {
|
|
666
|
+
const size_t vt = __riscv_vsetvl_e8m1(nbv - b);
|
|
667
|
+
vuint8m1_t packed = __riscv_vle8_v_u8m1(code + b, vt);
|
|
668
|
+
vuint8m1_t lo = __riscv_vand_vx_u8m1(packed, 0x0F, vt);
|
|
669
|
+
vuint8m1_t hi = __riscv_vsrl_vx_u8m1(packed, 4, vt);
|
|
670
|
+
|
|
671
|
+
vfloat32m4_t clo = __riscv_vfcvt_f_xu_v_f32m4(
|
|
672
|
+
__riscv_vzext_vf4_u32m4(lo, vt), vt);
|
|
673
|
+
vfloat32m4_t t_lo = __riscv_vle32_v_f32m4(e_lo.data() + b, vt);
|
|
674
|
+
t_lo = __riscv_vfnmsac_vf_f32m4(t_lo, a, clo, vt);
|
|
675
|
+
acc = __riscv_vfmacc_vv_f32m4_tu(acc, t_lo, t_lo, vt);
|
|
676
|
+
|
|
677
|
+
vfloat32m4_t chi = __riscv_vfcvt_f_xu_v_f32m4(
|
|
678
|
+
__riscv_vzext_vf4_u32m4(hi, vt), vt);
|
|
679
|
+
vfloat32m4_t t_hi = __riscv_vle32_v_f32m4(e_hi.data() + b, vt);
|
|
680
|
+
t_hi = __riscv_vfnmsac_vf_f32m4(t_hi, a, chi, vt);
|
|
681
|
+
acc = __riscv_vfmacc_vv_f32m4_tu(acc, t_hi, t_hi, vt);
|
|
264
682
|
}
|
|
265
|
-
|
|
683
|
+
|
|
684
|
+
// Horizontal reduce over all vlb lanes.
|
|
685
|
+
vfloat32m1_t red = __riscv_vfredusum_vs_f32m4_f32m1(
|
|
686
|
+
acc, __riscv_vfmv_v_f_f32m1(0.0f, 1), vlb);
|
|
687
|
+
total = __riscv_vfmv_f_s_f32m1_f32(red);
|
|
688
|
+
|
|
689
|
+
// Odd d: last dimension lives in the low nibble of the last byte.
|
|
690
|
+
if (d & 1) {
|
|
691
|
+
float diff = e_lo[nb - 1] - a * float(code[nb - 1] & 0x0F);
|
|
692
|
+
total += diff * diff;
|
|
693
|
+
}
|
|
694
|
+
return total;
|
|
266
695
|
}
|
|
267
696
|
|
|
268
697
|
float query_to_code(const uint8_t* code) const final {
|
|
269
|
-
return
|
|
698
|
+
return compute_l2(code);
|
|
270
699
|
}
|
|
271
700
|
|
|
272
701
|
float symmetric_dis(idx_t i, idx_t j) override {
|
|
273
702
|
// Not on the critical path for most workloads; reconstruct both
|
|
274
703
|
// codes into nibbles scalar-style and compute squared distance.
|
|
704
|
+
// recon1 - recon2 = a * (c1 - c2).
|
|
275
705
|
const uint8_t* c1 = codes + i * code_size;
|
|
276
706
|
const uint8_t* c2 = codes + j * code_size;
|
|
277
|
-
|
|
707
|
+
float acc = 0.0f;
|
|
278
708
|
for (size_t k = 0; k < d; k++) {
|
|
279
|
-
uint8_t
|
|
280
|
-
uint8_t
|
|
281
|
-
|
|
709
|
+
uint8_t n1 = (k % 2 == 0) ? (c1[k / 2] & 0x0F) : (c1[k / 2] >> 4);
|
|
710
|
+
uint8_t n2 = (k % 2 == 0) ? (c2[k / 2] & 0x0F) : (c2[k / 2] >> 4);
|
|
711
|
+
float diff = a * float(int(n1) - int(n2));
|
|
282
712
|
acc += diff * diff;
|
|
283
713
|
}
|
|
284
|
-
return
|
|
714
|
+
return acc;
|
|
285
715
|
}
|
|
286
716
|
|
|
287
717
|
void query_to_codes_batch_4(
|
|
@@ -293,45 +723,2656 @@ struct DCTemplate<
|
|
|
293
723
|
float& dis1,
|
|
294
724
|
float& dis2,
|
|
295
725
|
float& dis3) const final {
|
|
296
|
-
|
|
297
|
-
|
|
298
|
-
|
|
299
|
-
|
|
300
|
-
dis2 = static_cast<float>(accumulate_int_l2(code_2)) * final_scale_sq;
|
|
301
|
-
dis3 = static_cast<float>(accumulate_int_l2(code_3)) * final_scale_sq;
|
|
726
|
+
dis0 = compute_l2(code_0);
|
|
727
|
+
dis1 = compute_l2(code_1);
|
|
728
|
+
dis2 = compute_l2(code_2);
|
|
729
|
+
dis3 = compute_l2(code_3);
|
|
302
730
|
}
|
|
303
731
|
};
|
|
304
732
|
|
|
305
|
-
|
|
306
|
-
|
|
307
|
-
|
|
733
|
+
// * Fast path — QT_4bit (NON_UNIFORM) + L2
|
|
734
|
+
// *
|
|
735
|
+
// * Per-dimension affine reconstruction:
|
|
736
|
+
// * recon_i(c) = vmin[i] + vdiff[i] * (c + 0.5) / 15
|
|
737
|
+
// * = rmin_i + a_i * c, with a_i = vdiff[i] / 15,
|
|
738
|
+
// * rmin_i = vmin[i] + 0.5 * a_i
|
|
739
|
+
// * L2 contribution per dim: (q_i - recon_i)^2 = (e_i - a_i * c_i)^2
|
|
740
|
+
// * where e_i = q_i - rmin_i is precomputed once per query in set_query.
|
|
741
|
+
// *
|
|
742
|
+
// * Unlike the UNIFORM variant (above) the scale a_i differs per dimension,
|
|
743
|
+
// * so the sum cannot be pulled out of the integer domain — the kernel stays
|
|
744
|
+
// * in f32 and evaluates one fused multiply-sub + one fused multiply-add per
|
|
745
|
+
// * nibble stream.
|
|
746
|
+
// *
|
|
747
|
+
// * The packed 4-bit code interleaves dimensions (even dim -> low nibble,
|
|
748
|
+
// * odd dim -> high nibble), so the constructor deinterleaves a_i / rmin_i
|
|
749
|
+
// * into _lo/_hi streams once, and set_query does the same for e_i. The hot
|
|
750
|
+
// * loop then works on contiguous streams: per 16-byte code chunk (32 dims,
|
|
751
|
+
// * e8m1 -> f32m4) it unpacks nibbles, converts to f32 and accumulates
|
|
752
|
+
// * squared differences. vsetvl is hoisted (0 explicit vsetvl in the loop).
|
|
753
|
+
// *
|
|
754
|
+
// * Padding: for odd d the last byte's high nibble is not a real dimension;
|
|
755
|
+
// * a_hi/e_hi keep 0 in that slot so its contribution is exactly 0.
|
|
756
|
+
// ************************************************************************/
|
|
757
|
+
|
|
758
|
+
template <>
|
|
759
|
+
struct DCTemplate<
|
|
760
|
+
QuantizerTemplate<
|
|
761
|
+
Codec4bit<SIMDLevel::RISCV_RVV>,
|
|
762
|
+
QuantizerTemplateScaling::NON_UNIFORM,
|
|
763
|
+
SIMDLevel::RISCV_RVV>,
|
|
764
|
+
SimilarityL2<SIMDLevel::RISCV_RVV>,
|
|
765
|
+
SIMDLevel::RISCV_RVV> : SQDistanceComputer {
|
|
766
|
+
using Sim = SimilarityL2<SIMDLevel::RISCV_RVV>;
|
|
767
|
+
|
|
768
|
+
size_t d;
|
|
769
|
+
size_t half; // bytes per code = ceil(d/2)
|
|
770
|
+
std::vector<float> a_lo, a_hi; // a_i = vdiff[i]/15, deinterleaved
|
|
771
|
+
std::vector<float> rmin_lo, rmin_hi; // vmin[i] + 0.5*a_i, deinterleaved
|
|
772
|
+
std::vector<float> e_lo, e_hi; // q_i - rmin_i, per query
|
|
773
|
+
|
|
774
|
+
DCTemplate(size_t d_in, const std::vector<float>& trained)
|
|
775
|
+
: d(d_in),
|
|
776
|
+
half((d_in + 1) / 2),
|
|
777
|
+
a_lo(half, 0.0f),
|
|
778
|
+
a_hi(half, 0.0f),
|
|
779
|
+
rmin_lo(half, 0.0f),
|
|
780
|
+
rmin_hi(half, 0.0f),
|
|
781
|
+
e_lo(half, 0.0f),
|
|
782
|
+
e_hi(half, 0.0f) {
|
|
783
|
+
const float* vmin = trained.data();
|
|
784
|
+
const float* vdiff = trained.data() + d_in;
|
|
785
|
+
for (size_t i = 0; i < d_in; i++) {
|
|
786
|
+
float a = vdiff[i] / 15.0f;
|
|
787
|
+
float rmin = vmin[i] + 0.5f * a;
|
|
788
|
+
if (i % 2 == 0) {
|
|
789
|
+
a_lo[i / 2] = a;
|
|
790
|
+
rmin_lo[i / 2] = rmin;
|
|
791
|
+
} else {
|
|
792
|
+
a_hi[i / 2] = a;
|
|
793
|
+
rmin_hi[i / 2] = rmin;
|
|
794
|
+
}
|
|
795
|
+
}
|
|
796
|
+
}
|
|
308
797
|
|
|
309
|
-
|
|
310
|
-
|
|
311
|
-
|
|
312
|
-
|
|
313
|
-
|
|
314
|
-
|
|
315
|
-
|
|
316
|
-
|
|
317
|
-
float result = 0;
|
|
318
|
-
for (size_t byte_idx = 0; byte_idx < (d + 7) / 8; byte_idx++) {
|
|
319
|
-
uint8_t b = bits[byte_idx];
|
|
320
|
-
size_t base = byte_idx * 8;
|
|
321
|
-
size_t end = std::min(base + 8, d);
|
|
322
|
-
for (size_t j = base; j < end; j++) {
|
|
323
|
-
if (b & (1 << (j - base))) {
|
|
324
|
-
result += arr[j];
|
|
798
|
+
void set_query(const float* x) final {
|
|
799
|
+
q = x;
|
|
800
|
+
// e_i = q_i - rmin_i, deinterleaved. Padding slot (odd d) keeps 0.
|
|
801
|
+
for (size_t i = 0; i < d; i++) {
|
|
802
|
+
if (i % 2 == 0) {
|
|
803
|
+
e_lo[i / 2] = x[i] - rmin_lo[i / 2];
|
|
804
|
+
} else {
|
|
805
|
+
e_hi[i / 2] = x[i] - rmin_hi[i / 2];
|
|
325
806
|
}
|
|
326
807
|
}
|
|
327
808
|
}
|
|
328
|
-
return result;
|
|
329
|
-
}
|
|
330
809
|
|
|
331
|
-
|
|
332
|
-
|
|
810
|
+
/// Full-precision vector L2 over the packed code.
|
|
811
|
+
/// Hot loop body per byte chunk: unpack nibbles (e8), widen+convert to
|
|
812
|
+
/// f32, then t = e - a*c (vfnmsac) and acc += t*t (vfmacc).
|
|
813
|
+
float compute_l2(const uint8_t* code) const {
|
|
814
|
+
const size_t nb = half; // total bytes to process
|
|
815
|
+
size_t b = 0;
|
|
333
816
|
|
|
334
|
-
|
|
335
|
-
|
|
817
|
+
// Hoist vsetvl: VLMAX for e8m1 (== f32 lanes per m4 group chunk).
|
|
818
|
+
// d==0 => nb==0: guard so vsetvl(0) doesn't return 0 and the
|
|
819
|
+
// chunk loop below (b + vlb <= nb with vlb==0) never stalls.
|
|
820
|
+
size_t vlb = __riscv_vsetvl_e8m1(nb > 0 ? nb : 1);
|
|
821
|
+
vfloat32m4_t acc = __riscv_vfmv_v_f_f32m4(0.0f, vlb);
|
|
822
|
+
|
|
823
|
+
for (; b + vlb <= nb; b += vlb) {
|
|
824
|
+
vuint8m1_t packed = __riscv_vle8_v_u8m1(code + b, vlb);
|
|
825
|
+
vuint8m1_t lo = __riscv_vand_vx_u8m1(packed, 0x0F, vlb);
|
|
826
|
+
vuint8m1_t hi = __riscv_vsrl_vx_u8m1(packed, 4, vlb);
|
|
827
|
+
|
|
828
|
+
// Low-nibble stream (even dims)
|
|
829
|
+
vfloat32m4_t clo = __riscv_vfcvt_f_xu_v_f32m4(
|
|
830
|
+
__riscv_vzext_vf4_u32m4(lo, vlb), vlb);
|
|
831
|
+
vfloat32m4_t t_lo = __riscv_vle32_v_f32m4(e_lo.data() + b, vlb);
|
|
832
|
+
t_lo = __riscv_vfnmsac_vv_f32m4(
|
|
833
|
+
t_lo,
|
|
834
|
+
__riscv_vle32_v_f32m4(a_lo.data() + b, vlb),
|
|
835
|
+
clo,
|
|
836
|
+
vlb);
|
|
837
|
+
acc = __riscv_vfmacc_vv_f32m4(acc, t_lo, t_lo, vlb);
|
|
838
|
+
|
|
839
|
+
// High-nibble stream (odd dims)
|
|
840
|
+
vfloat32m4_t chi = __riscv_vfcvt_f_xu_v_f32m4(
|
|
841
|
+
__riscv_vzext_vf4_u32m4(hi, vlb), vlb);
|
|
842
|
+
vfloat32m4_t t_hi = __riscv_vle32_v_f32m4(e_hi.data() + b, vlb);
|
|
843
|
+
t_hi = __riscv_vfnmsac_vv_f32m4(
|
|
844
|
+
t_hi,
|
|
845
|
+
__riscv_vle32_v_f32m4(a_hi.data() + b, vlb),
|
|
846
|
+
chi,
|
|
847
|
+
vlb);
|
|
848
|
+
acc = __riscv_vfmacc_vv_f32m4(acc, t_hi, t_hi, vlb);
|
|
849
|
+
}
|
|
850
|
+
|
|
851
|
+
// Tail: fewer than vlb bytes left; accumulates into the first
|
|
852
|
+
// lanes of acc (safe: reduction below covers all vlb lanes).
|
|
853
|
+
if (b < nb) {
|
|
854
|
+
size_t vt = __riscv_vsetvl_e8m1(nb - b);
|
|
855
|
+
|
|
856
|
+
vuint8m1_t packed = __riscv_vle8_v_u8m1(code + b, vt);
|
|
857
|
+
vuint8m1_t lo = __riscv_vand_vx_u8m1(packed, 0x0F, vt);
|
|
858
|
+
vuint8m1_t hi = __riscv_vsrl_vx_u8m1(packed, 4, vt);
|
|
859
|
+
|
|
860
|
+
vfloat32m4_t clo = __riscv_vfcvt_f_xu_v_f32m4(
|
|
861
|
+
__riscv_vzext_vf4_u32m4(lo, vt), vt);
|
|
862
|
+
vfloat32m4_t t_lo = __riscv_vle32_v_f32m4(e_lo.data() + b, vt);
|
|
863
|
+
t_lo = __riscv_vfnmsac_vv_f32m4(
|
|
864
|
+
t_lo, __riscv_vle32_v_f32m4(a_lo.data() + b, vt), clo, vt);
|
|
865
|
+
acc = __riscv_vfmacc_vv_f32m4_tu(acc, t_lo, t_lo, vt);
|
|
866
|
+
|
|
867
|
+
vfloat32m4_t chi = __riscv_vfcvt_f_xu_v_f32m4(
|
|
868
|
+
__riscv_vzext_vf4_u32m4(hi, vt), vt);
|
|
869
|
+
vfloat32m4_t t_hi = __riscv_vle32_v_f32m4(e_hi.data() + b, vt);
|
|
870
|
+
t_hi = __riscv_vfnmsac_vv_f32m4(
|
|
871
|
+
t_hi, __riscv_vle32_v_f32m4(a_hi.data() + b, vt), chi, vt);
|
|
872
|
+
acc = __riscv_vfmacc_vv_f32m4_tu(acc, t_hi, t_hi, vt);
|
|
873
|
+
}
|
|
874
|
+
|
|
875
|
+
// Horizontal reduce over all vlb lanes
|
|
876
|
+
vfloat32m1_t red = __riscv_vfredusum_vs_f32m4_f32m1(
|
|
877
|
+
acc, __riscv_vfmv_v_f_f32m1(0.0f, 1), vlb);
|
|
878
|
+
return __riscv_vfmv_f_s_f32m1_f32(red);
|
|
879
|
+
}
|
|
880
|
+
|
|
881
|
+
float query_to_code(const uint8_t* code) const final {
|
|
882
|
+
return compute_l2(code);
|
|
883
|
+
}
|
|
884
|
+
|
|
885
|
+
float symmetric_dis(idx_t i, idx_t j) override {
|
|
886
|
+
// Not on the benchmark-critical path; scalar per-dim evaluation.
|
|
887
|
+
// recon1 - recon2 = a_k * (c1_k - c2_k)
|
|
888
|
+
const uint8_t* c1 = codes + i * code_size;
|
|
889
|
+
const uint8_t* c2 = codes + j * code_size;
|
|
890
|
+
float acc = 0;
|
|
891
|
+
for (size_t k = 0; k < d; k++) {
|
|
892
|
+
uint8_t n1 = (k % 2 == 0) ? (c1[k / 2] & 0x0F) : (c1[k / 2] >> 4);
|
|
893
|
+
uint8_t n2 = (k % 2 == 0) ? (c2[k / 2] & 0x0F) : (c2[k / 2] >> 4);
|
|
894
|
+
float a = (k % 2 == 0) ? a_lo[k / 2] : a_hi[k / 2];
|
|
895
|
+
float diff = a * float(int(n1) - int(n2));
|
|
896
|
+
acc += diff * diff;
|
|
897
|
+
}
|
|
898
|
+
return acc;
|
|
899
|
+
}
|
|
900
|
+
|
|
901
|
+
void query_to_codes_batch_4(
|
|
902
|
+
const uint8_t* code_0,
|
|
903
|
+
const uint8_t* code_1,
|
|
904
|
+
const uint8_t* code_2,
|
|
905
|
+
const uint8_t* code_3,
|
|
906
|
+
float& dis0,
|
|
907
|
+
float& dis1,
|
|
908
|
+
float& dis2,
|
|
909
|
+
float& dis3) const final {
|
|
910
|
+
dis0 = compute_l2(code_0);
|
|
911
|
+
dis1 = compute_l2(code_1);
|
|
912
|
+
dis2 = compute_l2(code_2);
|
|
913
|
+
dis3 = compute_l2(code_3);
|
|
914
|
+
}
|
|
915
|
+
};
|
|
916
|
+
|
|
917
|
+
// * Fast path — QT_4bit (NON_UNIFORM) + IP
|
|
918
|
+
// *
|
|
919
|
+
// * Per-dimension affine reconstruction:
|
|
920
|
+
// * recon_i(c) = vmin[i] + vdiff[i] * (c + 0.5) / 15
|
|
921
|
+
// * = rmin_i + a_i * c, with a_i = vdiff[i] / 15,
|
|
922
|
+
// * rmin_i = vmin[i] + 0.5 * a_i
|
|
923
|
+
// * The inner product against the query decomposes into a query-only
|
|
924
|
+
// * constant plus a coefficient dot product over the integer codes:
|
|
925
|
+
// * IP(q, code) = sum_i q_i * recon_i(c_i)
|
|
926
|
+
// * = sum_i q_i * rmin_i + sum_i (q_i * a_i) * c_i
|
|
927
|
+
// * = K_q + sum_i b_i * c_i
|
|
928
|
+
// * K_q and b_i = q_i * a_i depend only on the query, so set_query
|
|
929
|
+
// * precomputes them once per query (amortized over all codes).
|
|
930
|
+
// *
|
|
931
|
+
// * The packed 4-bit code interleaves dimensions (even dim -> low nibble,
|
|
932
|
+
// * odd dim -> high nibble), so the constructor deinterleaves a_i / rmin_i
|
|
933
|
+
// * into _lo/_hi streams once and set_query does the same for b_i. The hot
|
|
934
|
+
// * loop works on contiguous streams: per 16-byte code chunk (32 dims,
|
|
935
|
+
// * e8m1 -> f32m4) it unpacks nibbles, converts to f32 and accumulates
|
|
936
|
+
// * b * c into a single accumulator (1 load + 1 vfmacc per nibble stream —
|
|
937
|
+
// * lighter than the L2 kernel). vsetvl is hoisted (0 explicit vsetvl in
|
|
938
|
+
// * the loop).
|
|
939
|
+
// *
|
|
940
|
+
// * Padding: for odd d the last byte's high nibble is not a real dimension;
|
|
941
|
+
// * b_hi keeps 0 in that slot so its contribution is exactly 0.
|
|
942
|
+
// ************************************************************************/
|
|
943
|
+
|
|
944
|
+
template <>
|
|
945
|
+
struct DCTemplate<
|
|
946
|
+
QuantizerTemplate<
|
|
947
|
+
Codec4bit<SIMDLevel::RISCV_RVV>,
|
|
948
|
+
QuantizerTemplateScaling::NON_UNIFORM,
|
|
949
|
+
SIMDLevel::RISCV_RVV>,
|
|
950
|
+
SimilarityIP<SIMDLevel::RISCV_RVV>,
|
|
951
|
+
SIMDLevel::RISCV_RVV> : SQDistanceComputer {
|
|
952
|
+
using Sim = SimilarityIP<SIMDLevel::RISCV_RVV>;
|
|
953
|
+
|
|
954
|
+
size_t d;
|
|
955
|
+
size_t half; // bytes per code = ceil(d/2)
|
|
956
|
+
std::vector<float> a_lo, a_hi; // a_i = vdiff[i]/15, deinterleaved
|
|
957
|
+
std::vector<float> rmin_lo, rmin_hi; // vmin[i] + 0.5*a_i, deinterleaved
|
|
958
|
+
std::vector<float> b_lo, b_hi; // q_i * a_i, per query
|
|
959
|
+
float k_q; // sum_i q_i * rmin_i, per query
|
|
336
960
|
|
|
337
|
-
|
|
961
|
+
DCTemplate(size_t d_in, const std::vector<float>& trained)
|
|
962
|
+
: d(d_in),
|
|
963
|
+
half((d_in + 1) / 2),
|
|
964
|
+
a_lo(half, 0.0f),
|
|
965
|
+
a_hi(half, 0.0f),
|
|
966
|
+
rmin_lo(half, 0.0f),
|
|
967
|
+
rmin_hi(half, 0.0f),
|
|
968
|
+
b_lo(half, 0.0f),
|
|
969
|
+
b_hi(half, 0.0f),
|
|
970
|
+
k_q(0.0f) {
|
|
971
|
+
const float* vmin = trained.data();
|
|
972
|
+
const float* vdiff = trained.data() + d_in;
|
|
973
|
+
for (size_t i = 0; i < d_in; i++) {
|
|
974
|
+
float a = vdiff[i] / 15.0f;
|
|
975
|
+
float rmin = vmin[i] + 0.5f * a;
|
|
976
|
+
if (i % 2 == 0) {
|
|
977
|
+
a_lo[i / 2] = a;
|
|
978
|
+
rmin_lo[i / 2] = rmin;
|
|
979
|
+
} else {
|
|
980
|
+
a_hi[i / 2] = a;
|
|
981
|
+
rmin_hi[i / 2] = rmin;
|
|
982
|
+
}
|
|
983
|
+
}
|
|
984
|
+
}
|
|
985
|
+
|
|
986
|
+
void set_query(const float* x) final {
|
|
987
|
+
q = x;
|
|
988
|
+
// b_i = q_i * a_i, deinterleaved; K_q = sum_i q_i * rmin_i.
|
|
989
|
+
// Padding slot (odd d) keeps b_hi = 0 so its contribution is 0.
|
|
990
|
+
float acc = 0;
|
|
991
|
+
for (size_t i = 0; i < d; i++) {
|
|
992
|
+
if (i % 2 == 0) {
|
|
993
|
+
b_lo[i / 2] = x[i] * a_lo[i / 2];
|
|
994
|
+
acc += x[i] * rmin_lo[i / 2];
|
|
995
|
+
} else {
|
|
996
|
+
b_hi[i / 2] = x[i] * a_hi[i / 2];
|
|
997
|
+
acc += x[i] * rmin_hi[i / 2];
|
|
998
|
+
}
|
|
999
|
+
}
|
|
1000
|
+
k_q = acc;
|
|
1001
|
+
}
|
|
1002
|
+
|
|
1003
|
+
/// IP = K_q + sum_i b_i * c_i over the packed code.
|
|
1004
|
+
/// Hot loop per byte chunk: unpack nibbles, widen to f32, acc += b*c.
|
|
1005
|
+
/// Software-pipelined: the next chunk's load is issued at the top of
|
|
1006
|
+
/// this iteration (prologue + rotate) to stretch the load-use distance.
|
|
1007
|
+
float compute_ip(const uint8_t* code) const {
|
|
1008
|
+
const size_t nb = half; // total bytes to process
|
|
1009
|
+
size_t bpos = 0;
|
|
1010
|
+
|
|
1011
|
+
// Hoist vsetvl: VLMAX for e8m1 (== f32 lanes per m4 group chunk).
|
|
1012
|
+
// d==0 => nb==0: guard so the `if (nb >= vlb)` prologue + chunk
|
|
1013
|
+
// loop (bpos + 2*vlb <= nb) don't stall on vlb==0.
|
|
1014
|
+
size_t vlb = __riscv_vsetvl_e8m1(nb > 0 ? nb : 1);
|
|
1015
|
+
vfloat32m4_t acc = __riscv_vfmv_v_f_f32m4(0.0f, vlb);
|
|
1016
|
+
|
|
1017
|
+
if (nb >= vlb) {
|
|
1018
|
+
// Prologue: preload chunk 0.
|
|
1019
|
+
vuint8m1_t packed = __riscv_vle8_v_u8m1(code, vlb);
|
|
1020
|
+
|
|
1021
|
+
// Main loop: while a full NEXT chunk exists, issue its load
|
|
1022
|
+
// first, then process the current one.
|
|
1023
|
+
for (; bpos + 2 * vlb <= nb; bpos += vlb) {
|
|
1024
|
+
__builtin_prefetch(code + bpos + 256, 0, 0);
|
|
1025
|
+
vuint8m1_t packed_next =
|
|
1026
|
+
__riscv_vle8_v_u8m1(code + bpos + vlb, vlb);
|
|
1027
|
+
|
|
1028
|
+
vuint8m1_t lo = __riscv_vand_vx_u8m1(packed, 0x0F, vlb);
|
|
1029
|
+
vuint8m1_t hi = __riscv_vsrl_vx_u8m1(packed, 4, vlb);
|
|
1030
|
+
|
|
1031
|
+
// Low-nibble stream (even dims): acc += b_lo * c_lo
|
|
1032
|
+
vfloat32m4_t clo = __riscv_vfcvt_f_xu_v_f32m4(
|
|
1033
|
+
__riscv_vzext_vf4_u32m4(lo, vlb), vlb);
|
|
1034
|
+
acc = __riscv_vfmacc_vv_f32m4(
|
|
1035
|
+
acc,
|
|
1036
|
+
__riscv_vle32_v_f32m4(b_lo.data() + bpos, vlb),
|
|
1037
|
+
clo,
|
|
1038
|
+
vlb);
|
|
1039
|
+
|
|
1040
|
+
// High-nibble stream (odd dims): acc += b_hi * c_hi
|
|
1041
|
+
vfloat32m4_t chi = __riscv_vfcvt_f_xu_v_f32m4(
|
|
1042
|
+
__riscv_vzext_vf4_u32m4(hi, vlb), vlb);
|
|
1043
|
+
acc = __riscv_vfmacc_vv_f32m4(
|
|
1044
|
+
acc,
|
|
1045
|
+
__riscv_vle32_v_f32m4(b_hi.data() + bpos, vlb),
|
|
1046
|
+
chi,
|
|
1047
|
+
vlb);
|
|
1048
|
+
|
|
1049
|
+
packed = packed_next; // rotate
|
|
1050
|
+
}
|
|
1051
|
+
|
|
1052
|
+
// Epilogue: process the last full chunk (already loaded).
|
|
1053
|
+
{
|
|
1054
|
+
vuint8m1_t lo = __riscv_vand_vx_u8m1(packed, 0x0F, vlb);
|
|
1055
|
+
vuint8m1_t hi = __riscv_vsrl_vx_u8m1(packed, 4, vlb);
|
|
1056
|
+
|
|
1057
|
+
vfloat32m4_t clo = __riscv_vfcvt_f_xu_v_f32m4(
|
|
1058
|
+
__riscv_vzext_vf4_u32m4(lo, vlb), vlb);
|
|
1059
|
+
acc = __riscv_vfmacc_vv_f32m4(
|
|
1060
|
+
acc,
|
|
1061
|
+
__riscv_vle32_v_f32m4(b_lo.data() + bpos, vlb),
|
|
1062
|
+
clo,
|
|
1063
|
+
vlb);
|
|
1064
|
+
|
|
1065
|
+
vfloat32m4_t chi = __riscv_vfcvt_f_xu_v_f32m4(
|
|
1066
|
+
__riscv_vzext_vf4_u32m4(hi, vlb), vlb);
|
|
1067
|
+
acc = __riscv_vfmacc_vv_f32m4(
|
|
1068
|
+
acc,
|
|
1069
|
+
__riscv_vle32_v_f32m4(b_hi.data() + bpos, vlb),
|
|
1070
|
+
chi,
|
|
1071
|
+
vlb);
|
|
1072
|
+
|
|
1073
|
+
bpos += vlb;
|
|
1074
|
+
}
|
|
1075
|
+
}
|
|
1076
|
+
|
|
1077
|
+
// Tail: fewer than vlb bytes left; accumulates into the first
|
|
1078
|
+
// lanes of acc (safe: reduction below covers all vlb lanes).
|
|
1079
|
+
if (bpos < nb) {
|
|
1080
|
+
size_t vt = __riscv_vsetvl_e8m1(nb - bpos);
|
|
1081
|
+
|
|
1082
|
+
vuint8m1_t packed = __riscv_vle8_v_u8m1(code + bpos, vt);
|
|
1083
|
+
vuint8m1_t lo = __riscv_vand_vx_u8m1(packed, 0x0F, vt);
|
|
1084
|
+
vuint8m1_t hi = __riscv_vsrl_vx_u8m1(packed, 4, vt);
|
|
1085
|
+
|
|
1086
|
+
vfloat32m4_t clo = __riscv_vfcvt_f_xu_v_f32m4(
|
|
1087
|
+
__riscv_vzext_vf4_u32m4(lo, vt), vt);
|
|
1088
|
+
acc = __riscv_vfmacc_vv_f32m4_tu(
|
|
1089
|
+
acc,
|
|
1090
|
+
__riscv_vle32_v_f32m4(b_lo.data() + bpos, vt),
|
|
1091
|
+
clo,
|
|
1092
|
+
vt);
|
|
1093
|
+
|
|
1094
|
+
vfloat32m4_t chi = __riscv_vfcvt_f_xu_v_f32m4(
|
|
1095
|
+
__riscv_vzext_vf4_u32m4(hi, vt), vt);
|
|
1096
|
+
acc = __riscv_vfmacc_vv_f32m4_tu(
|
|
1097
|
+
acc,
|
|
1098
|
+
__riscv_vle32_v_f32m4(b_hi.data() + bpos, vt),
|
|
1099
|
+
chi,
|
|
1100
|
+
vt);
|
|
1101
|
+
}
|
|
1102
|
+
|
|
1103
|
+
// Horizontal reduce over all vlb lanes
|
|
1104
|
+
vfloat32m1_t red = __riscv_vfredusum_vs_f32m4_f32m1(
|
|
1105
|
+
acc, __riscv_vfmv_v_f_f32m1(0.0f, 1), vlb);
|
|
1106
|
+
return k_q + __riscv_vfmv_f_s_f32m1_f32(red);
|
|
1107
|
+
}
|
|
1108
|
+
|
|
1109
|
+
float query_to_code(const uint8_t* code) const final {
|
|
1110
|
+
return compute_ip(code);
|
|
1111
|
+
}
|
|
1112
|
+
|
|
1113
|
+
float symmetric_dis(idx_t i, idx_t j) override {
|
|
1114
|
+
// Not on the benchmark-critical path; scalar per-dim evaluation.
|
|
1115
|
+
// IP(recon(c1), recon(c2)) = sum_k recon1_k * recon2_k
|
|
1116
|
+
const uint8_t* c1 = codes + i * code_size;
|
|
1117
|
+
const uint8_t* c2 = codes + j * code_size;
|
|
1118
|
+
float acc = 0;
|
|
1119
|
+
for (size_t k = 0; k < d; k++) {
|
|
1120
|
+
uint8_t n1 = (k % 2 == 0) ? (c1[k / 2] & 0x0F) : (c1[k / 2] >> 4);
|
|
1121
|
+
uint8_t n2 = (k % 2 == 0) ? (c2[k / 2] & 0x0F) : (c2[k / 2] >> 4);
|
|
1122
|
+
float a = (k % 2 == 0) ? a_lo[k / 2] : a_hi[k / 2];
|
|
1123
|
+
float rmin = (k % 2 == 0) ? rmin_lo[k / 2] : rmin_hi[k / 2];
|
|
1124
|
+
float r1 = rmin + a * float(n1);
|
|
1125
|
+
float r2 = rmin + a * float(n2);
|
|
1126
|
+
acc += r1 * r2;
|
|
1127
|
+
}
|
|
1128
|
+
return acc;
|
|
1129
|
+
}
|
|
1130
|
+
|
|
1131
|
+
void query_to_codes_batch_4(
|
|
1132
|
+
const uint8_t* code_0,
|
|
1133
|
+
const uint8_t* code_1,
|
|
1134
|
+
const uint8_t* code_2,
|
|
1135
|
+
const uint8_t* code_3,
|
|
1136
|
+
float& dis0,
|
|
1137
|
+
float& dis1,
|
|
1138
|
+
float& dis2,
|
|
1139
|
+
float& dis3) const final {
|
|
1140
|
+
dis0 = compute_ip(code_0);
|
|
1141
|
+
dis1 = compute_ip(code_1);
|
|
1142
|
+
dis2 = compute_ip(code_2);
|
|
1143
|
+
dis3 = compute_ip(code_3);
|
|
1144
|
+
}
|
|
1145
|
+
};
|
|
1146
|
+
|
|
1147
|
+
// * Fast path — QT_4bit_uniform + IP
|
|
1148
|
+
// *
|
|
1149
|
+
// * Uniform scaling: recon(c) = vmin + vdiff*(c + 0.5)/15 = c0 + a*c,
|
|
1150
|
+
// * with a = vdiff/15, c0 = vmin + 0.5*a. The inner product folds into a
|
|
1151
|
+
// * query-only constant plus one integer-code dot product over the
|
|
1152
|
+
// * per-dim pre-scaled query b_i = q_i*a_i:
|
|
1153
|
+
// * IP(q, code) = sum_i q_i*c0 + sum_i (q_i*a_i)*c_i = K_q + S
|
|
1154
|
+
// * K_q is precomputed in set_query; the hot loop only evaluates S.
|
|
1155
|
+
// * Pre-scaling per dim (as the non-uniform IP kernel already does)
|
|
1156
|
+
// * keeps every term inside the scalar per-dim semantics: a == 0 or
|
|
1157
|
+
// * c0 == 0 zeroes its term exactly, so a huge finite query can never
|
|
1158
|
+
// * overflow a bare c0*sum(q) or sum(q_i*c_i) intermediate into
|
|
1159
|
+
// * 0*inf == NaN.
|
|
1160
|
+
// *
|
|
1161
|
+
// * The packed code interleaves dims (even -> low nibble, odd -> high), so
|
|
1162
|
+
// * set_query deinterleaves b into b_lo/b_hi. Odd-d padding nibble keeps
|
|
1163
|
+
// * b_hi = 0. vdiff == 0 degenerates naturally (b = 0 -> IP = K_q).
|
|
1164
|
+
// ************************************************************************/
|
|
1165
|
+
|
|
1166
|
+
template <>
|
|
1167
|
+
struct DCTemplate<
|
|
1168
|
+
QuantizerTemplate<
|
|
1169
|
+
Codec4bit<SIMDLevel::RISCV_RVV>,
|
|
1170
|
+
QuantizerTemplateScaling::UNIFORM,
|
|
1171
|
+
SIMDLevel::RISCV_RVV>,
|
|
1172
|
+
SimilarityIP<SIMDLevel::RISCV_RVV>,
|
|
1173
|
+
SIMDLevel::RISCV_RVV> : SQDistanceComputer {
|
|
1174
|
+
using Sim = SimilarityIP<SIMDLevel::RISCV_RVV>;
|
|
1175
|
+
|
|
1176
|
+
size_t d;
|
|
1177
|
+
size_t half; // bytes per code = ceil(d/2)
|
|
1178
|
+
float vmin;
|
|
1179
|
+
float vdiff;
|
|
1180
|
+
float a; // vdiff / 15
|
|
1181
|
+
float c0; // vmin + 0.5 * a
|
|
1182
|
+
std::vector<float> b_lo, b_hi; // q_i * a_i, per query, deinterleaved
|
|
1183
|
+
float k_q; // sum_i (q_i * c0), per query
|
|
1184
|
+
|
|
1185
|
+
DCTemplate(size_t d_in, const std::vector<float>& trained)
|
|
1186
|
+
: d(d_in),
|
|
1187
|
+
half((d_in + 1) / 2),
|
|
1188
|
+
vmin(trained[0]),
|
|
1189
|
+
vdiff(trained[1]),
|
|
1190
|
+
b_lo(half, 0.0f),
|
|
1191
|
+
b_hi(half, 0.0f),
|
|
1192
|
+
k_q(0.0f) {
|
|
1193
|
+
a = vdiff / 15.0f;
|
|
1194
|
+
c0 = vmin + 0.5f * a;
|
|
1195
|
+
}
|
|
1196
|
+
|
|
1197
|
+
void set_query(const float* x) final {
|
|
1198
|
+
q = x;
|
|
1199
|
+
// b_i = q_i * a_i and K_q = sum_i (q_i * c0), both accumulated
|
|
1200
|
+
// per dim. Per-dim pre-scaling keeps every term inside the
|
|
1201
|
+
// scalar per-dim semantics: a == 0 / c0 == 0 zero the term
|
|
1202
|
+
// exactly, so no bare c0*sum(q) or sum(q_i*c_i) intermediate
|
|
1203
|
+
// can overflow a finite query into 0*inf == NaN. Padding slot
|
|
1204
|
+
// (odd d) keeps b_hi = 0 so the padding nibble contributes 0.
|
|
1205
|
+
float kq = 0;
|
|
1206
|
+
for (size_t i = 0; i < d; i++) {
|
|
1207
|
+
kq += x[i] * c0;
|
|
1208
|
+
if (i % 2 == 0) {
|
|
1209
|
+
b_lo[i / 2] = x[i] * a;
|
|
1210
|
+
} else {
|
|
1211
|
+
b_hi[i / 2] = x[i] * a;
|
|
1212
|
+
}
|
|
1213
|
+
}
|
|
1214
|
+
k_q = kq;
|
|
1215
|
+
}
|
|
1216
|
+
|
|
1217
|
+
/// S = sum_i b_i * c_i over the packed code. Software-pipelined hot
|
|
1218
|
+
/// loop (prologue + rotate) with dual accumulators per nibble stream.
|
|
1219
|
+
float compute_qc_dot(const uint8_t* code) const {
|
|
1220
|
+
const size_t nb = half; // total bytes to process
|
|
1221
|
+
size_t bpos = 0;
|
|
1222
|
+
|
|
1223
|
+
// Hoist vsetvl; d==0 => nb==0, guard so the chunk loop can't
|
|
1224
|
+
// stall on vlb==0.
|
|
1225
|
+
size_t vlb = __riscv_vsetvl_e8m1(nb > 0 ? nb : 1);
|
|
1226
|
+
// Dual accumulators (one per nibble stream) break the serial
|
|
1227
|
+
// chain of two dependent FMAs into two independent chains.
|
|
1228
|
+
vfloat32m4_t acc_lo = __riscv_vfmv_v_f_f32m4(0.0f, vlb);
|
|
1229
|
+
vfloat32m4_t acc_hi = __riscv_vfmv_v_f_f32m4(0.0f, vlb);
|
|
1230
|
+
|
|
1231
|
+
if (nb >= vlb) {
|
|
1232
|
+
// Prologue: preload chunk 0.
|
|
1233
|
+
vuint8m1_t packed = __riscv_vle8_v_u8m1(code, vlb);
|
|
1234
|
+
|
|
1235
|
+
// Main loop: while a full NEXT chunk exists, issue its load
|
|
1236
|
+
// first, then process the current one.
|
|
1237
|
+
for (; bpos + 2 * vlb <= nb; bpos += vlb) {
|
|
1238
|
+
__builtin_prefetch(code + bpos + 512, 0, 0);
|
|
1239
|
+
vuint8m1_t packed_next =
|
|
1240
|
+
__riscv_vle8_v_u8m1(code + bpos + vlb, vlb);
|
|
1241
|
+
|
|
1242
|
+
vuint8m1_t lo = __riscv_vand_vx_u8m1(packed, 0x0F, vlb);
|
|
1243
|
+
vuint8m1_t hi = __riscv_vsrl_vx_u8m1(packed, 4, vlb);
|
|
1244
|
+
|
|
1245
|
+
// Issue both b-stream loads early to hide L1 load-use
|
|
1246
|
+
// latency ahead of the consuming FMAs.
|
|
1247
|
+
vfloat32m4_t vb_lo =
|
|
1248
|
+
__riscv_vle32_v_f32m4(b_lo.data() + bpos, vlb);
|
|
1249
|
+
vfloat32m4_t vb_hi =
|
|
1250
|
+
__riscv_vle32_v_f32m4(b_hi.data() + bpos, vlb);
|
|
1251
|
+
|
|
1252
|
+
// Low-nibble stream (even dims): acc_lo += b_lo * c_lo
|
|
1253
|
+
vfloat32m4_t clo = __riscv_vfwcvt_f_xu_v_f32m4(
|
|
1254
|
+
__riscv_vzext_vf2_u16m2(lo, vlb), vlb);
|
|
1255
|
+
acc_lo = __riscv_vfmacc_vv_f32m4(acc_lo, vb_lo, clo, vlb);
|
|
1256
|
+
|
|
1257
|
+
// High-nibble stream (odd dims): acc_hi += b_hi * c_hi
|
|
1258
|
+
vfloat32m4_t chi = __riscv_vfwcvt_f_xu_v_f32m4(
|
|
1259
|
+
__riscv_vzext_vf2_u16m2(hi, vlb), vlb);
|
|
1260
|
+
acc_hi = __riscv_vfmacc_vv_f32m4(acc_hi, vb_hi, chi, vlb);
|
|
1261
|
+
|
|
1262
|
+
packed = packed_next; // rotate
|
|
1263
|
+
}
|
|
1264
|
+
|
|
1265
|
+
// Epilogue: process the last full chunk (already loaded).
|
|
1266
|
+
{
|
|
1267
|
+
vuint8m1_t lo = __riscv_vand_vx_u8m1(packed, 0x0F, vlb);
|
|
1268
|
+
vuint8m1_t hi = __riscv_vsrl_vx_u8m1(packed, 4, vlb);
|
|
1269
|
+
|
|
1270
|
+
vfloat32m4_t clo = __riscv_vfwcvt_f_xu_v_f32m4(
|
|
1271
|
+
__riscv_vzext_vf2_u16m2(lo, vlb), vlb);
|
|
1272
|
+
acc_lo = __riscv_vfmacc_vv_f32m4(
|
|
1273
|
+
acc_lo,
|
|
1274
|
+
__riscv_vle32_v_f32m4(b_lo.data() + bpos, vlb),
|
|
1275
|
+
clo,
|
|
1276
|
+
vlb);
|
|
1277
|
+
|
|
1278
|
+
vfloat32m4_t chi = __riscv_vfwcvt_f_xu_v_f32m4(
|
|
1279
|
+
__riscv_vzext_vf2_u16m2(hi, vlb), vlb);
|
|
1280
|
+
acc_hi = __riscv_vfmacc_vv_f32m4(
|
|
1281
|
+
acc_hi,
|
|
1282
|
+
__riscv_vle32_v_f32m4(b_hi.data() + bpos, vlb),
|
|
1283
|
+
chi,
|
|
1284
|
+
vlb);
|
|
1285
|
+
|
|
1286
|
+
bpos += vlb;
|
|
1287
|
+
}
|
|
1288
|
+
}
|
|
1289
|
+
|
|
1290
|
+
// Tail: fewer than vlb bytes left; accumulates into the first
|
|
1291
|
+
// lanes of acc (safe: reduction below covers all vlb lanes).
|
|
1292
|
+
if (bpos < nb) {
|
|
1293
|
+
size_t vt = __riscv_vsetvl_e8m1(nb - bpos);
|
|
1294
|
+
|
|
1295
|
+
vuint8m1_t packed = __riscv_vle8_v_u8m1(code + bpos, vt);
|
|
1296
|
+
vuint8m1_t lo = __riscv_vand_vx_u8m1(packed, 0x0F, vt);
|
|
1297
|
+
vuint8m1_t hi = __riscv_vsrl_vx_u8m1(packed, 4, vt);
|
|
1298
|
+
|
|
1299
|
+
vfloat32m4_t clo = __riscv_vfwcvt_f_xu_v_f32m4(
|
|
1300
|
+
__riscv_vzext_vf2_u16m2(lo, vt), vt);
|
|
1301
|
+
acc_lo = __riscv_vfmacc_vv_f32m4_tu(
|
|
1302
|
+
acc_lo,
|
|
1303
|
+
__riscv_vle32_v_f32m4(b_lo.data() + bpos, vt),
|
|
1304
|
+
clo,
|
|
1305
|
+
vt);
|
|
1306
|
+
|
|
1307
|
+
vfloat32m4_t chi = __riscv_vfwcvt_f_xu_v_f32m4(
|
|
1308
|
+
__riscv_vzext_vf2_u16m2(hi, vt), vt);
|
|
1309
|
+
acc_hi = __riscv_vfmacc_vv_f32m4_tu(
|
|
1310
|
+
acc_hi,
|
|
1311
|
+
__riscv_vle32_v_f32m4(b_hi.data() + bpos, vt),
|
|
1312
|
+
chi,
|
|
1313
|
+
vt);
|
|
1314
|
+
}
|
|
1315
|
+
|
|
1316
|
+
// Merge the two chains, then one horizontal reduce over vlb lanes
|
|
1317
|
+
vfloat32m4_t acc = __riscv_vfadd_vv_f32m4(acc_lo, acc_hi, vlb);
|
|
1318
|
+
vfloat32m1_t red = __riscv_vfredusum_vs_f32m4_f32m1(
|
|
1319
|
+
acc, __riscv_vfmv_v_f_f32m1(0.0f, 1), vlb);
|
|
1320
|
+
return __riscv_vfmv_f_s_f32m1_f32(red);
|
|
1321
|
+
}
|
|
1322
|
+
|
|
1323
|
+
float query_to_code(const uint8_t* code) const final {
|
|
1324
|
+
return k_q + compute_qc_dot(code);
|
|
1325
|
+
}
|
|
1326
|
+
|
|
1327
|
+
float symmetric_dis(idx_t i, idx_t j) override {
|
|
1328
|
+
// Not on the benchmark-critical path; scalar per-dim evaluation.
|
|
1329
|
+
// IP(recon(c1), recon(c2)) = sum_k (c0 + a*n1) * (c0 + a*n2)
|
|
1330
|
+
const uint8_t* c1 = codes + i * code_size;
|
|
1331
|
+
const uint8_t* c2 = codes + j * code_size;
|
|
1332
|
+
float acc = 0;
|
|
1333
|
+
for (size_t k = 0; k < d; k++) {
|
|
1334
|
+
uint8_t n1 = (k % 2 == 0) ? (c1[k / 2] & 0x0F) : (c1[k / 2] >> 4);
|
|
1335
|
+
uint8_t n2 = (k % 2 == 0) ? (c2[k / 2] & 0x0F) : (c2[k / 2] >> 4);
|
|
1336
|
+
float r1 = c0 + a * float(n1);
|
|
1337
|
+
float r2 = c0 + a * float(n2);
|
|
1338
|
+
acc += r1 * r2;
|
|
1339
|
+
}
|
|
1340
|
+
return acc;
|
|
1341
|
+
}
|
|
1342
|
+
|
|
1343
|
+
void query_to_codes_batch_4(
|
|
1344
|
+
const uint8_t* code_0,
|
|
1345
|
+
const uint8_t* code_1,
|
|
1346
|
+
const uint8_t* code_2,
|
|
1347
|
+
const uint8_t* code_3,
|
|
1348
|
+
float& dis0,
|
|
1349
|
+
float& dis1,
|
|
1350
|
+
float& dis2,
|
|
1351
|
+
float& dis3) const final {
|
|
1352
|
+
dis0 = k_q + compute_qc_dot(code_0);
|
|
1353
|
+
dis1 = k_q + compute_qc_dot(code_1);
|
|
1354
|
+
dis2 = k_q + compute_qc_dot(code_2);
|
|
1355
|
+
dis3 = k_q + compute_qc_dot(code_3);
|
|
1356
|
+
}
|
|
1357
|
+
};
|
|
1358
|
+
|
|
1359
|
+
// * Fast path — QT_6bit (NON_UNIFORM) + L2
|
|
1360
|
+
// *
|
|
1361
|
+
// * recon_i(c) = vmin[i] + vdiff[i]*(c + 0.5)/63 = rmin_i + a_i*c_i,
|
|
1362
|
+
// * with a_i = vdiff[i]/63, rmin_i = vmin[i] + 0.5*a_i. L2 contribution
|
|
1363
|
+
// * per dim: (q_i - recon_i)^2 = (e_i - a_i*c_i)^2, e_i = q_i - rmin_i.
|
|
1364
|
+
// *
|
|
1365
|
+
// * Codec6bit packs 4 dims into 3 bytes (b0,b1,b2); no lo/hi symmetry, so
|
|
1366
|
+
// * the kernel deinterleaves by (i & 3) into FOUR streams (a0..a3 / e0..e3).
|
|
1367
|
+
// * Hot loop: vlseg3e8 deinterleaves the 3-byte groups, 6-bit fields are
|
|
1368
|
+
// * extracted with u8 ALU ops, each stream widens to f32 and accumulates
|
|
1369
|
+
// * t = e - a*c (vfnmsac), acc += t*t (vfmacc). Software-pipelined
|
|
1370
|
+
// * (prologue + rotate). Tail dims (d % 4) are decoded scalar.
|
|
1371
|
+
// ************************************************************************
|
|
1372
|
+
|
|
1373
|
+
// Raw 6-bit field extract, scalar (exact Codec6bit<NONE> bit semantics).
|
|
1374
|
+
static inline uint8_t sq6_decode_raw(const uint8_t* code, size_t i) {
|
|
1375
|
+
const uint8_t* p = code + (i >> 2) * 3;
|
|
1376
|
+
switch (i & 3) {
|
|
1377
|
+
case 0:
|
|
1378
|
+
return p[0] & 0x3F;
|
|
1379
|
+
case 1:
|
|
1380
|
+
return uint8_t((p[0] >> 6) | ((p[1] & 0x0F) << 2));
|
|
1381
|
+
case 2:
|
|
1382
|
+
return uint8_t((p[1] >> 4) | ((p[2] & 0x03) << 4));
|
|
1383
|
+
default:
|
|
1384
|
+
return uint8_t(p[2] >> 2);
|
|
1385
|
+
}
|
|
1386
|
+
}
|
|
1387
|
+
|
|
1388
|
+
template <>
|
|
1389
|
+
struct DCTemplate<
|
|
1390
|
+
QuantizerTemplate<
|
|
1391
|
+
Codec6bit<SIMDLevel::RISCV_RVV>,
|
|
1392
|
+
QuantizerTemplateScaling::NON_UNIFORM,
|
|
1393
|
+
SIMDLevel::RISCV_RVV>,
|
|
1394
|
+
SimilarityL2<SIMDLevel::RISCV_RVV>,
|
|
1395
|
+
SIMDLevel::RISCV_RVV> : SQDistanceComputer {
|
|
1396
|
+
using Sim = SimilarityL2<SIMDLevel::RISCV_RVV>;
|
|
1397
|
+
|
|
1398
|
+
size_t d;
|
|
1399
|
+
size_t ng; // full 4-dim groups = d / 4
|
|
1400
|
+
// a_i = vdiff[i]/63 and e_i = q_i - rmin_i, deinterleaved by (i & 3)
|
|
1401
|
+
std::vector<float> a0, a1, a2, a3;
|
|
1402
|
+
std::vector<float> e0, e1, e2, e3;
|
|
1403
|
+
// Interleaved copies for the scalar tail / symmetric_dis
|
|
1404
|
+
std::vector<float> a_all, rmin_all;
|
|
1405
|
+
|
|
1406
|
+
DCTemplate(size_t d_in, const std::vector<float>& trained)
|
|
1407
|
+
: d(d_in),
|
|
1408
|
+
ng(d_in / 4),
|
|
1409
|
+
a0(ng, 0.0f),
|
|
1410
|
+
a1(ng, 0.0f),
|
|
1411
|
+
a2(ng, 0.0f),
|
|
1412
|
+
a3(ng, 0.0f),
|
|
1413
|
+
e0(ng, 0.0f),
|
|
1414
|
+
e1(ng, 0.0f),
|
|
1415
|
+
e2(ng, 0.0f),
|
|
1416
|
+
e3(ng, 0.0f),
|
|
1417
|
+
a_all(d_in, 0.0f),
|
|
1418
|
+
rmin_all(d_in, 0.0f) {
|
|
1419
|
+
const float* vmin = trained.data();
|
|
1420
|
+
const float* vdiff = trained.data() + d_in;
|
|
1421
|
+
for (size_t i = 0; i < d_in; i++) {
|
|
1422
|
+
float a = vdiff[i] / 63.0f;
|
|
1423
|
+
float rmin = vmin[i] + 0.5f * a;
|
|
1424
|
+
a_all[i] = a;
|
|
1425
|
+
rmin_all[i] = rmin;
|
|
1426
|
+
if (i < 4 * ng) {
|
|
1427
|
+
size_t k = i >> 2;
|
|
1428
|
+
switch (i & 3) {
|
|
1429
|
+
case 0:
|
|
1430
|
+
a0[k] = a;
|
|
1431
|
+
break;
|
|
1432
|
+
case 1:
|
|
1433
|
+
a1[k] = a;
|
|
1434
|
+
break;
|
|
1435
|
+
case 2:
|
|
1436
|
+
a2[k] = a;
|
|
1437
|
+
break;
|
|
1438
|
+
default:
|
|
1439
|
+
a3[k] = a;
|
|
1440
|
+
break;
|
|
1441
|
+
}
|
|
1442
|
+
}
|
|
1443
|
+
}
|
|
1444
|
+
}
|
|
1445
|
+
|
|
1446
|
+
void set_query(const float* x) final {
|
|
1447
|
+
q = x;
|
|
1448
|
+
// e_i = q_i - rmin_i, deinterleaved into the 4 group-position
|
|
1449
|
+
// streams (amortized: called once per query, 1/n weight).
|
|
1450
|
+
for (size_t k = 0; k < ng; k++) {
|
|
1451
|
+
const float* xi = x + 4 * k;
|
|
1452
|
+
const float* rm = rmin_all.data() + 4 * k;
|
|
1453
|
+
e0[k] = xi[0] - rm[0];
|
|
1454
|
+
e1[k] = xi[1] - rm[1];
|
|
1455
|
+
e2[k] = xi[2] - rm[2];
|
|
1456
|
+
e3[k] = xi[3] - rm[3];
|
|
1457
|
+
}
|
|
1458
|
+
// Tail dims (d % 4) are handled scalar in compute_l2 via q.
|
|
1459
|
+
}
|
|
1460
|
+
|
|
1461
|
+
/// Full-precision vector L2 over the packed 6-bit code.
|
|
1462
|
+
float compute_l2(const uint8_t* code) const {
|
|
1463
|
+
const size_t ngf = ng;
|
|
1464
|
+
float total = 0.0f;
|
|
1465
|
+
size_t g = 0;
|
|
1466
|
+
|
|
1467
|
+
if (ngf > 0) {
|
|
1468
|
+
// Hoist vsetvl: VLMAX for e8m1 (== f32 lanes per m4 group)
|
|
1469
|
+
const size_t vl = __riscv_vsetvl_e8m1(ngf);
|
|
1470
|
+
vfloat32m4_t acc = __riscv_vfmv_v_f_f32m4(0.0f, vl);
|
|
1471
|
+
|
|
1472
|
+
const float* pa0 = a0.data();
|
|
1473
|
+
const float* pa1 = a1.data();
|
|
1474
|
+
const float* pa2 = a2.data();
|
|
1475
|
+
const float* pa3 = a3.data();
|
|
1476
|
+
const float* pe0 = e0.data();
|
|
1477
|
+
const float* pe1 = e1.data();
|
|
1478
|
+
const float* pe2 = e2.data();
|
|
1479
|
+
const float* pe3 = e3.data();
|
|
1480
|
+
|
|
1481
|
+
// Software-pipelined vlseg3e8: preload the next chunk's code
|
|
1482
|
+
// block, then compute on the already-loaded current chunk.
|
|
1483
|
+
if (ngf >= vl) {
|
|
1484
|
+
// Prologue: preload chunk 0.
|
|
1485
|
+
vuint8m1x3_t seg = __riscv_vlseg3e8_v_u8m1x3(code + 3 * g, vl);
|
|
1486
|
+
g += vl;
|
|
1487
|
+
|
|
1488
|
+
// Main loop: while a full next chunk exists, issue its
|
|
1489
|
+
// load first, then process the current one.
|
|
1490
|
+
for (; g + vl <= ngf; g += vl) {
|
|
1491
|
+
vuint8m1x3_t seg_next =
|
|
1492
|
+
__riscv_vlseg3e8_v_u8m1x3(code + 3 * g, vl);
|
|
1493
|
+
vuint8m1_t b0 = __riscv_vget_v_u8m1x3_u8m1(seg, 0);
|
|
1494
|
+
vuint8m1_t b1 = __riscv_vget_v_u8m1x3_u8m1(seg, 1);
|
|
1495
|
+
vuint8m1_t b2 = __riscv_vget_v_u8m1x3_u8m1(seg, 2);
|
|
1496
|
+
|
|
1497
|
+
// Light extracts (c0, c3) first so their FMA chains
|
|
1498
|
+
// issue early; heavy extracts (c1, c2) overlap.
|
|
1499
|
+
// Accumulation order: c0, c3, c1, c2.
|
|
1500
|
+
vuint8m1_t c0 = __riscv_vand_vx_u8m1(b0, 0x3F, vl);
|
|
1501
|
+
vuint8m1_t c3 = __riscv_vsrl_vx_u8m1(b2, 2, vl);
|
|
1502
|
+
|
|
1503
|
+
vfloat32m4_t f0 = __riscv_vfwcvt_f_xu_v_f32m4(
|
|
1504
|
+
__riscv_vzext_vf2_u16m2(c0, vl), vl);
|
|
1505
|
+
vfloat32m4_t t0 = __riscv_vle32_v_f32m4(pe0 + g - vl, vl);
|
|
1506
|
+
t0 = __riscv_vfnmsac_vv_f32m4(
|
|
1507
|
+
t0,
|
|
1508
|
+
__riscv_vle32_v_f32m4(pa0 + g - vl, vl),
|
|
1509
|
+
f0,
|
|
1510
|
+
vl);
|
|
1511
|
+
acc = __riscv_vfmacc_vv_f32m4(acc, t0, t0, vl);
|
|
1512
|
+
|
|
1513
|
+
vfloat32m4_t f3 = __riscv_vfwcvt_f_xu_v_f32m4(
|
|
1514
|
+
__riscv_vzext_vf2_u16m2(c3, vl), vl);
|
|
1515
|
+
vfloat32m4_t t3 = __riscv_vle32_v_f32m4(pe3 + g - vl, vl);
|
|
1516
|
+
t3 = __riscv_vfnmsac_vv_f32m4(
|
|
1517
|
+
t3,
|
|
1518
|
+
__riscv_vle32_v_f32m4(pa3 + g - vl, vl),
|
|
1519
|
+
f3,
|
|
1520
|
+
vl);
|
|
1521
|
+
acc = __riscv_vfmacc_vv_f32m4(acc, t3, t3, vl);
|
|
1522
|
+
|
|
1523
|
+
vuint8m1_t c1 = __riscv_vmacc_vx_u8m1(
|
|
1524
|
+
__riscv_vsrl_vx_u8m1(b0, 6, vl),
|
|
1525
|
+
4,
|
|
1526
|
+
__riscv_vand_vx_u8m1(b1, 0x0F, vl),
|
|
1527
|
+
vl);
|
|
1528
|
+
|
|
1529
|
+
vfloat32m4_t f1 = __riscv_vfwcvt_f_xu_v_f32m4(
|
|
1530
|
+
__riscv_vzext_vf2_u16m2(c1, vl), vl);
|
|
1531
|
+
vfloat32m4_t t1 = __riscv_vle32_v_f32m4(pe1 + g - vl, vl);
|
|
1532
|
+
t1 = __riscv_vfnmsac_vv_f32m4(
|
|
1533
|
+
t1,
|
|
1534
|
+
__riscv_vle32_v_f32m4(pa1 + g - vl, vl),
|
|
1535
|
+
f1,
|
|
1536
|
+
vl);
|
|
1537
|
+
acc = __riscv_vfmacc_vv_f32m4(acc, t1, t1, vl);
|
|
1538
|
+
|
|
1539
|
+
vuint8m1_t c2 = __riscv_vmacc_vx_u8m1(
|
|
1540
|
+
__riscv_vsrl_vx_u8m1(b1, 4, vl),
|
|
1541
|
+
16,
|
|
1542
|
+
__riscv_vand_vx_u8m1(b2, 0x03, vl),
|
|
1543
|
+
vl);
|
|
1544
|
+
|
|
1545
|
+
vfloat32m4_t f2 = __riscv_vfwcvt_f_xu_v_f32m4(
|
|
1546
|
+
__riscv_vzext_vf2_u16m2(c2, vl), vl);
|
|
1547
|
+
vfloat32m4_t t2 = __riscv_vle32_v_f32m4(pe2 + g - vl, vl);
|
|
1548
|
+
t2 = __riscv_vfnmsac_vv_f32m4(
|
|
1549
|
+
t2,
|
|
1550
|
+
__riscv_vle32_v_f32m4(pa2 + g - vl, vl),
|
|
1551
|
+
f2,
|
|
1552
|
+
vl);
|
|
1553
|
+
acc = __riscv_vfmacc_vv_f32m4(acc, t2, t2, vl);
|
|
1554
|
+
|
|
1555
|
+
seg = seg_next; // rotate
|
|
1556
|
+
}
|
|
1557
|
+
|
|
1558
|
+
// Epilogue: process the last full chunk (already loaded).
|
|
1559
|
+
{
|
|
1560
|
+
vuint8m1_t b0 = __riscv_vget_v_u8m1x3_u8m1(seg, 0);
|
|
1561
|
+
vuint8m1_t b1 = __riscv_vget_v_u8m1x3_u8m1(seg, 1);
|
|
1562
|
+
vuint8m1_t b2 = __riscv_vget_v_u8m1x3_u8m1(seg, 2);
|
|
1563
|
+
|
|
1564
|
+
vuint8m1_t c0 = __riscv_vand_vx_u8m1(b0, 0x3F, vl);
|
|
1565
|
+
vuint8m1_t c3 = __riscv_vsrl_vx_u8m1(b2, 2, vl);
|
|
1566
|
+
|
|
1567
|
+
vfloat32m4_t f0 = __riscv_vfwcvt_f_xu_v_f32m4(
|
|
1568
|
+
__riscv_vzext_vf2_u16m2(c0, vl), vl);
|
|
1569
|
+
vfloat32m4_t t0 = __riscv_vle32_v_f32m4(pe0 + g - vl, vl);
|
|
1570
|
+
t0 = __riscv_vfnmsac_vv_f32m4(
|
|
1571
|
+
t0,
|
|
1572
|
+
__riscv_vle32_v_f32m4(pa0 + g - vl, vl),
|
|
1573
|
+
f0,
|
|
1574
|
+
vl);
|
|
1575
|
+
acc = __riscv_vfmacc_vv_f32m4(acc, t0, t0, vl);
|
|
1576
|
+
|
|
1577
|
+
vfloat32m4_t f3 = __riscv_vfwcvt_f_xu_v_f32m4(
|
|
1578
|
+
__riscv_vzext_vf2_u16m2(c3, vl), vl);
|
|
1579
|
+
vfloat32m4_t t3 = __riscv_vle32_v_f32m4(pe3 + g - vl, vl);
|
|
1580
|
+
t3 = __riscv_vfnmsac_vv_f32m4(
|
|
1581
|
+
t3,
|
|
1582
|
+
__riscv_vle32_v_f32m4(pa3 + g - vl, vl),
|
|
1583
|
+
f3,
|
|
1584
|
+
vl);
|
|
1585
|
+
acc = __riscv_vfmacc_vv_f32m4(acc, t3, t3, vl);
|
|
1586
|
+
|
|
1587
|
+
vuint8m1_t c1 = __riscv_vmacc_vx_u8m1(
|
|
1588
|
+
__riscv_vsrl_vx_u8m1(b0, 6, vl),
|
|
1589
|
+
4,
|
|
1590
|
+
__riscv_vand_vx_u8m1(b1, 0x0F, vl),
|
|
1591
|
+
vl);
|
|
1592
|
+
|
|
1593
|
+
vfloat32m4_t f1 = __riscv_vfwcvt_f_xu_v_f32m4(
|
|
1594
|
+
__riscv_vzext_vf2_u16m2(c1, vl), vl);
|
|
1595
|
+
vfloat32m4_t t1 = __riscv_vle32_v_f32m4(pe1 + g - vl, vl);
|
|
1596
|
+
t1 = __riscv_vfnmsac_vv_f32m4(
|
|
1597
|
+
t1,
|
|
1598
|
+
__riscv_vle32_v_f32m4(pa1 + g - vl, vl),
|
|
1599
|
+
f1,
|
|
1600
|
+
vl);
|
|
1601
|
+
acc = __riscv_vfmacc_vv_f32m4(acc, t1, t1, vl);
|
|
1602
|
+
|
|
1603
|
+
vuint8m1_t c2 = __riscv_vmacc_vx_u8m1(
|
|
1604
|
+
__riscv_vsrl_vx_u8m1(b1, 4, vl),
|
|
1605
|
+
16,
|
|
1606
|
+
__riscv_vand_vx_u8m1(b2, 0x03, vl),
|
|
1607
|
+
vl);
|
|
1608
|
+
|
|
1609
|
+
vfloat32m4_t f2 = __riscv_vfwcvt_f_xu_v_f32m4(
|
|
1610
|
+
__riscv_vzext_vf2_u16m2(c2, vl), vl);
|
|
1611
|
+
vfloat32m4_t t2 = __riscv_vle32_v_f32m4(pe2 + g - vl, vl);
|
|
1612
|
+
t2 = __riscv_vfnmsac_vv_f32m4(
|
|
1613
|
+
t2,
|
|
1614
|
+
__riscv_vle32_v_f32m4(pa2 + g - vl, vl),
|
|
1615
|
+
f2,
|
|
1616
|
+
vl);
|
|
1617
|
+
acc = __riscv_vfmacc_vv_f32m4(acc, t2, t2, vl);
|
|
1618
|
+
}
|
|
1619
|
+
}
|
|
1620
|
+
|
|
1621
|
+
// Tail groups: one shorter-vl pass into the same accumulator.
|
|
1622
|
+
if (g < ngf) {
|
|
1623
|
+
const size_t vt = __riscv_vsetvl_e8m1(ngf - g);
|
|
1624
|
+
|
|
1625
|
+
vuint8m1x3_t seg = __riscv_vlseg3e8_v_u8m1x3(code + 3 * g, vt);
|
|
1626
|
+
vuint8m1_t b0 = __riscv_vget_v_u8m1x3_u8m1(seg, 0);
|
|
1627
|
+
vuint8m1_t b1 = __riscv_vget_v_u8m1x3_u8m1(seg, 1);
|
|
1628
|
+
vuint8m1_t b2 = __riscv_vget_v_u8m1x3_u8m1(seg, 2);
|
|
1629
|
+
|
|
1630
|
+
vuint8m1_t c0 = __riscv_vand_vx_u8m1(b0, 0x3F, vt);
|
|
1631
|
+
vuint8m1_t c3 = __riscv_vsrl_vx_u8m1(b2, 2, vt);
|
|
1632
|
+
|
|
1633
|
+
vfloat32m4_t f0 = __riscv_vfwcvt_f_xu_v_f32m4(
|
|
1634
|
+
__riscv_vzext_vf2_u16m2(c0, vt), vt);
|
|
1635
|
+
vfloat32m4_t t0 = __riscv_vle32_v_f32m4(pe0 + g, vt);
|
|
1636
|
+
t0 = __riscv_vfnmsac_vv_f32m4(
|
|
1637
|
+
t0, __riscv_vle32_v_f32m4(pa0 + g, vt), f0, vt);
|
|
1638
|
+
acc = __riscv_vfmacc_vv_f32m4_tu(acc, t0, t0, vt);
|
|
1639
|
+
|
|
1640
|
+
vfloat32m4_t f3 = __riscv_vfwcvt_f_xu_v_f32m4(
|
|
1641
|
+
__riscv_vzext_vf2_u16m2(c3, vt), vt);
|
|
1642
|
+
vfloat32m4_t t3 = __riscv_vle32_v_f32m4(pe3 + g, vt);
|
|
1643
|
+
t3 = __riscv_vfnmsac_vv_f32m4(
|
|
1644
|
+
t3, __riscv_vle32_v_f32m4(pa3 + g, vt), f3, vt);
|
|
1645
|
+
acc = __riscv_vfmacc_vv_f32m4_tu(acc, t3, t3, vt);
|
|
1646
|
+
|
|
1647
|
+
vuint8m1_t c1 = __riscv_vmacc_vx_u8m1(
|
|
1648
|
+
__riscv_vsrl_vx_u8m1(b0, 6, vt),
|
|
1649
|
+
4,
|
|
1650
|
+
__riscv_vand_vx_u8m1(b1, 0x0F, vt),
|
|
1651
|
+
vt);
|
|
1652
|
+
|
|
1653
|
+
vfloat32m4_t f1 = __riscv_vfwcvt_f_xu_v_f32m4(
|
|
1654
|
+
__riscv_vzext_vf2_u16m2(c1, vt), vt);
|
|
1655
|
+
vfloat32m4_t t1 = __riscv_vle32_v_f32m4(pe1 + g, vt);
|
|
1656
|
+
t1 = __riscv_vfnmsac_vv_f32m4(
|
|
1657
|
+
t1, __riscv_vle32_v_f32m4(pa1 + g, vt), f1, vt);
|
|
1658
|
+
acc = __riscv_vfmacc_vv_f32m4_tu(acc, t1, t1, vt);
|
|
1659
|
+
|
|
1660
|
+
vuint8m1_t c2 = __riscv_vmacc_vx_u8m1(
|
|
1661
|
+
__riscv_vsrl_vx_u8m1(b1, 4, vt),
|
|
1662
|
+
16,
|
|
1663
|
+
__riscv_vand_vx_u8m1(b2, 0x03, vt),
|
|
1664
|
+
vt);
|
|
1665
|
+
|
|
1666
|
+
vfloat32m4_t f2 = __riscv_vfwcvt_f_xu_v_f32m4(
|
|
1667
|
+
__riscv_vzext_vf2_u16m2(c2, vt), vt);
|
|
1668
|
+
vfloat32m4_t t2 = __riscv_vle32_v_f32m4(pe2 + g, vt);
|
|
1669
|
+
t2 = __riscv_vfnmsac_vv_f32m4(
|
|
1670
|
+
t2, __riscv_vle32_v_f32m4(pa2 + g, vt), f2, vt);
|
|
1671
|
+
acc = __riscv_vfmacc_vv_f32m4_tu(acc, t2, t2, vt);
|
|
1672
|
+
}
|
|
1673
|
+
|
|
1674
|
+
// Single horizontal reduction over all vl lanes.
|
|
1675
|
+
vfloat32m1_t red = __riscv_vfredusum_vs_f32m4_f32m1(
|
|
1676
|
+
acc, __riscv_vfmv_v_f_f32m1(0.0f, 1), vl);
|
|
1677
|
+
total = __riscv_vfmv_f_s_f32m1_f32(red);
|
|
1678
|
+
}
|
|
1679
|
+
|
|
1680
|
+
// Scalar tail: dims beyond the last full group (d % 4).
|
|
1681
|
+
for (size_t i = 4 * ngf; i < d; i++) {
|
|
1682
|
+
float r = rmin_all[i] + a_all[i] * float(sq6_decode_raw(code, i));
|
|
1683
|
+
float diff = q[i] - r;
|
|
1684
|
+
total += diff * diff;
|
|
1685
|
+
}
|
|
1686
|
+
return total;
|
|
1687
|
+
}
|
|
1688
|
+
|
|
1689
|
+
float query_to_code(const uint8_t* code) const final {
|
|
1690
|
+
return compute_l2(code);
|
|
1691
|
+
}
|
|
1692
|
+
|
|
1693
|
+
float symmetric_dis(idx_t i, idx_t j) override {
|
|
1694
|
+
// Not on the benchmark-critical path; scalar per-dim evaluation.
|
|
1695
|
+
// recon1 - recon2 = a_k * (c1_k - c2_k)
|
|
1696
|
+
const uint8_t* c1 = codes + i * code_size;
|
|
1697
|
+
const uint8_t* c2 = codes + j * code_size;
|
|
1698
|
+
float acc = 0;
|
|
1699
|
+
for (size_t k = 0; k < d; k++) {
|
|
1700
|
+
int n1 = sq6_decode_raw(c1, k);
|
|
1701
|
+
int n2 = sq6_decode_raw(c2, k);
|
|
1702
|
+
float diff = a_all[k] * float(n1 - n2);
|
|
1703
|
+
acc += diff * diff;
|
|
1704
|
+
}
|
|
1705
|
+
return acc;
|
|
1706
|
+
}
|
|
1707
|
+
|
|
1708
|
+
void query_to_codes_batch_4(
|
|
1709
|
+
const uint8_t* code_0,
|
|
1710
|
+
const uint8_t* code_1,
|
|
1711
|
+
const uint8_t* code_2,
|
|
1712
|
+
const uint8_t* code_3,
|
|
1713
|
+
float& dis0,
|
|
1714
|
+
float& dis1,
|
|
1715
|
+
float& dis2,
|
|
1716
|
+
float& dis3) const final {
|
|
1717
|
+
dis0 = compute_l2(code_0);
|
|
1718
|
+
dis1 = compute_l2(code_1);
|
|
1719
|
+
dis2 = compute_l2(code_2);
|
|
1720
|
+
dis3 = compute_l2(code_3);
|
|
1721
|
+
}
|
|
1722
|
+
};
|
|
1723
|
+
|
|
1724
|
+
// * Fast path — QT_6bit (NON_UNIFORM) + IP
|
|
1725
|
+
// *
|
|
1726
|
+
// * recon_i(c) = rmin_i + a_i*c_i, a_i = vdiff[i]/63, rmin_i = vmin[i] +
|
|
1727
|
+
// * 0.5*a_i. IP decomposes into a query-only constant plus a coefficient
|
|
1728
|
+
// * dot product over the integer codes:
|
|
1729
|
+
// * IP(q, code) = sum_i q_i*rmin_i + sum_i (q_i*a_i)*c_i = K_q + sum_i
|
|
1730
|
+
// b_i*c_i
|
|
1731
|
+
// * K_q and b_i = q_i*a_i are precomputed in set_query. Codec6bit packs 4
|
|
1732
|
+
// * dims into 3 bytes, so b_i is deinterleaved by (i & 3) into four streams.
|
|
1733
|
+
// *
|
|
1734
|
+
// * Hot loop: software-pipelined vlseg3e8, 6-bit fields extracted with u8
|
|
1735
|
+
// * ALU ops, each stream widens to f32 and accumulates acc += b*c (vfmacc).
|
|
1736
|
+
// * Tail dims (d % 4) are decoded scalar.
|
|
1737
|
+
// ************************************************************************/
|
|
1738
|
+
|
|
1739
|
+
template <>
|
|
1740
|
+
struct DCTemplate<
|
|
1741
|
+
QuantizerTemplate<
|
|
1742
|
+
Codec6bit<SIMDLevel::RISCV_RVV>,
|
|
1743
|
+
QuantizerTemplateScaling::NON_UNIFORM,
|
|
1744
|
+
SIMDLevel::RISCV_RVV>,
|
|
1745
|
+
SimilarityIP<SIMDLevel::RISCV_RVV>,
|
|
1746
|
+
SIMDLevel::RISCV_RVV> : SQDistanceComputer {
|
|
1747
|
+
using Sim = SimilarityIP<SIMDLevel::RISCV_RVV>;
|
|
1748
|
+
|
|
1749
|
+
size_t d;
|
|
1750
|
+
size_t ng; // full 4-dim groups = d / 4
|
|
1751
|
+
// b_i = q_i * a_i, deinterleaved by (i & 3), rebuilt per query
|
|
1752
|
+
std::vector<float> b0, b1, b2, b3;
|
|
1753
|
+
// Interleaved constants for set_query / scalar tail / symmetric_dis
|
|
1754
|
+
std::vector<float> a_all, rmin_all;
|
|
1755
|
+
float k_q; // sum_i q_i * rmin_i, per query
|
|
1756
|
+
|
|
1757
|
+
DCTemplate(size_t d_in, const std::vector<float>& trained)
|
|
1758
|
+
: d(d_in),
|
|
1759
|
+
ng(d_in / 4),
|
|
1760
|
+
b0(ng, 0.0f),
|
|
1761
|
+
b1(ng, 0.0f),
|
|
1762
|
+
b2(ng, 0.0f),
|
|
1763
|
+
b3(ng, 0.0f),
|
|
1764
|
+
a_all(d_in, 0.0f),
|
|
1765
|
+
rmin_all(d_in, 0.0f),
|
|
1766
|
+
k_q(0.0f) {
|
|
1767
|
+
const float* vmin = trained.data();
|
|
1768
|
+
const float* vdiff = trained.data() + d_in;
|
|
1769
|
+
for (size_t i = 0; i < d_in; i++) {
|
|
1770
|
+
float a = vdiff[i] / 63.0f;
|
|
1771
|
+
a_all[i] = a;
|
|
1772
|
+
rmin_all[i] = vmin[i] + 0.5f * a;
|
|
1773
|
+
}
|
|
1774
|
+
}
|
|
1775
|
+
|
|
1776
|
+
void set_query(const float* x) final {
|
|
1777
|
+
q = x;
|
|
1778
|
+
// b_i = q_i * a_i deinterleaved into the 4 group-position streams;
|
|
1779
|
+
// K_q = sum over ALL dims of q_i * rmin_i (tail dims included, so
|
|
1780
|
+
// the scalar tail in compute_ip only adds the b_i * c_i part).
|
|
1781
|
+
float acc = 0;
|
|
1782
|
+
for (size_t k = 0; k < ng; k++) {
|
|
1783
|
+
const float* xi = x + 4 * k;
|
|
1784
|
+
const float* aa = a_all.data() + 4 * k;
|
|
1785
|
+
const float* rm = rmin_all.data() + 4 * k;
|
|
1786
|
+
b0[k] = xi[0] * aa[0];
|
|
1787
|
+
b1[k] = xi[1] * aa[1];
|
|
1788
|
+
b2[k] = xi[2] * aa[2];
|
|
1789
|
+
b3[k] = xi[3] * aa[3];
|
|
1790
|
+
acc += xi[0] * rm[0] + xi[1] * rm[1] + xi[2] * rm[2] +
|
|
1791
|
+
xi[3] * rm[3];
|
|
1792
|
+
}
|
|
1793
|
+
for (size_t i = 4 * ng; i < d; i++) {
|
|
1794
|
+
acc += x[i] * rmin_all[i];
|
|
1795
|
+
}
|
|
1796
|
+
k_q = acc;
|
|
1797
|
+
}
|
|
1798
|
+
|
|
1799
|
+
/// S = sum_i b_i * c_i over the packed 6-bit code; returns K_q + S.
|
|
1800
|
+
float compute_ip(const uint8_t* code) const {
|
|
1801
|
+
const size_t ngf = ng;
|
|
1802
|
+
float total = 0.0f;
|
|
1803
|
+
size_t g = 0;
|
|
1804
|
+
|
|
1805
|
+
if (ngf > 0) {
|
|
1806
|
+
// Hoist vsetvl: VLMAX for e8m1 (== f32 lanes per m4 group)
|
|
1807
|
+
const size_t vl = __riscv_vsetvl_e8m1(ngf);
|
|
1808
|
+
vfloat32m4_t acc = __riscv_vfmv_v_f_f32m4(0.0f, vl);
|
|
1809
|
+
|
|
1810
|
+
const float* pb0 = b0.data();
|
|
1811
|
+
const float* pb1 = b1.data();
|
|
1812
|
+
const float* pb2 = b2.data();
|
|
1813
|
+
const float* pb3 = b3.data();
|
|
1814
|
+
|
|
1815
|
+
// Software-pipelined vlseg3e8 (depth 1): preload the NEXT
|
|
1816
|
+
// chunk's code block at the top of THIS iteration, then
|
|
1817
|
+
// compute on the already-loaded current chunk.
|
|
1818
|
+
if (ngf >= vl) {
|
|
1819
|
+
// Prologue: preload chunk 0.
|
|
1820
|
+
vuint8m1x3_t seg = __riscv_vlseg3e8_v_u8m1x3(code + 3 * g, vl);
|
|
1821
|
+
g += vl;
|
|
1822
|
+
|
|
1823
|
+
// Main loop: while a full next chunk exists, issue its
|
|
1824
|
+
// load first, then process the current one.
|
|
1825
|
+
for (; g + vl <= ngf; g += vl) {
|
|
1826
|
+
vuint8m1x3_t seg_next =
|
|
1827
|
+
__riscv_vlseg3e8_v_u8m1x3(code + 3 * g, vl);
|
|
1828
|
+
vuint8m1_t r0 = __riscv_vget_v_u8m1x3_u8m1(seg, 0);
|
|
1829
|
+
vuint8m1_t r1 = __riscv_vget_v_u8m1x3_u8m1(seg, 1);
|
|
1830
|
+
vuint8m1_t r2 = __riscv_vget_v_u8m1x3_u8m1(seg, 2);
|
|
1831
|
+
|
|
1832
|
+
// Light extracts (c0, c3) first so their FMA chains
|
|
1833
|
+
// issue early; heavy extracts (c1, c2) overlap.
|
|
1834
|
+
// Accumulation order: 0, 3, 1, 2.
|
|
1835
|
+
vuint8m1_t c0 = __riscv_vand_vx_u8m1(r0, 0x3F, vl);
|
|
1836
|
+
vuint8m1_t c3 = __riscv_vsrl_vx_u8m1(r2, 2, vl);
|
|
1837
|
+
|
|
1838
|
+
vfloat32m4_t vb0 = __riscv_vle32_v_f32m4(pb0 + g - vl, vl);
|
|
1839
|
+
vfloat32m4_t f0 = __riscv_vfwcvt_f_xu_v_f32m4(
|
|
1840
|
+
__riscv_vzext_vf2_u16m2(c0, vl), vl);
|
|
1841
|
+
acc = __riscv_vfmacc_vv_f32m4(acc, vb0, f0, vl);
|
|
1842
|
+
|
|
1843
|
+
vfloat32m4_t vb3 = __riscv_vle32_v_f32m4(pb3 + g - vl, vl);
|
|
1844
|
+
vfloat32m4_t f3 = __riscv_vfwcvt_f_xu_v_f32m4(
|
|
1845
|
+
__riscv_vzext_vf2_u16m2(c3, vl), vl);
|
|
1846
|
+
acc = __riscv_vfmacc_vv_f32m4(acc, vb3, f3, vl);
|
|
1847
|
+
|
|
1848
|
+
vuint8m1_t c1 = __riscv_vmacc_vx_u8m1(
|
|
1849
|
+
__riscv_vsrl_vx_u8m1(r0, 6, vl),
|
|
1850
|
+
4,
|
|
1851
|
+
__riscv_vand_vx_u8m1(r1, 0x0F, vl),
|
|
1852
|
+
vl);
|
|
1853
|
+
vfloat32m4_t vb1 = __riscv_vle32_v_f32m4(pb1 + g - vl, vl);
|
|
1854
|
+
vfloat32m4_t f1 = __riscv_vfwcvt_f_xu_v_f32m4(
|
|
1855
|
+
__riscv_vzext_vf2_u16m2(c1, vl), vl);
|
|
1856
|
+
acc = __riscv_vfmacc_vv_f32m4(acc, vb1, f1, vl);
|
|
1857
|
+
|
|
1858
|
+
vuint8m1_t c2 = __riscv_vmacc_vx_u8m1(
|
|
1859
|
+
__riscv_vsrl_vx_u8m1(r1, 4, vl),
|
|
1860
|
+
16,
|
|
1861
|
+
__riscv_vand_vx_u8m1(r2, 0x03, vl),
|
|
1862
|
+
vl);
|
|
1863
|
+
vfloat32m4_t vb2 = __riscv_vle32_v_f32m4(pb2 + g - vl, vl);
|
|
1864
|
+
vfloat32m4_t f2 = __riscv_vfwcvt_f_xu_v_f32m4(
|
|
1865
|
+
__riscv_vzext_vf2_u16m2(c2, vl), vl);
|
|
1866
|
+
acc = __riscv_vfmacc_vv_f32m4(acc, vb2, f2, vl);
|
|
1867
|
+
|
|
1868
|
+
seg = seg_next; // rotate
|
|
1869
|
+
}
|
|
1870
|
+
|
|
1871
|
+
// Epilogue: process the last full chunk (already loaded).
|
|
1872
|
+
{
|
|
1873
|
+
vuint8m1_t r0 = __riscv_vget_v_u8m1x3_u8m1(seg, 0);
|
|
1874
|
+
vuint8m1_t r1 = __riscv_vget_v_u8m1x3_u8m1(seg, 1);
|
|
1875
|
+
vuint8m1_t r2 = __riscv_vget_v_u8m1x3_u8m1(seg, 2);
|
|
1876
|
+
|
|
1877
|
+
// Light extracts (c0, c3) first so their FMA chains
|
|
1878
|
+
// issue early; heavy extracts (c1, c2) overlap.
|
|
1879
|
+
// Accumulation order: 0, 3, 1, 2.
|
|
1880
|
+
vuint8m1_t c0 = __riscv_vand_vx_u8m1(r0, 0x3F, vl);
|
|
1881
|
+
vuint8m1_t c3 = __riscv_vsrl_vx_u8m1(r2, 2, vl);
|
|
1882
|
+
|
|
1883
|
+
vfloat32m4_t vb0 = __riscv_vle32_v_f32m4(pb0 + g - vl, vl);
|
|
1884
|
+
vfloat32m4_t f0 = __riscv_vfwcvt_f_xu_v_f32m4(
|
|
1885
|
+
__riscv_vzext_vf2_u16m2(c0, vl), vl);
|
|
1886
|
+
acc = __riscv_vfmacc_vv_f32m4(acc, vb0, f0, vl);
|
|
1887
|
+
|
|
1888
|
+
vfloat32m4_t vb3 = __riscv_vle32_v_f32m4(pb3 + g - vl, vl);
|
|
1889
|
+
vfloat32m4_t f3 = __riscv_vfwcvt_f_xu_v_f32m4(
|
|
1890
|
+
__riscv_vzext_vf2_u16m2(c3, vl), vl);
|
|
1891
|
+
acc = __riscv_vfmacc_vv_f32m4(acc, vb3, f3, vl);
|
|
1892
|
+
|
|
1893
|
+
vuint8m1_t c1 = __riscv_vmacc_vx_u8m1(
|
|
1894
|
+
__riscv_vsrl_vx_u8m1(r0, 6, vl),
|
|
1895
|
+
4,
|
|
1896
|
+
__riscv_vand_vx_u8m1(r1, 0x0F, vl),
|
|
1897
|
+
vl);
|
|
1898
|
+
vfloat32m4_t vb1 = __riscv_vle32_v_f32m4(pb1 + g - vl, vl);
|
|
1899
|
+
vfloat32m4_t f1 = __riscv_vfwcvt_f_xu_v_f32m4(
|
|
1900
|
+
__riscv_vzext_vf2_u16m2(c1, vl), vl);
|
|
1901
|
+
acc = __riscv_vfmacc_vv_f32m4(acc, vb1, f1, vl);
|
|
1902
|
+
|
|
1903
|
+
vuint8m1_t c2 = __riscv_vmacc_vx_u8m1(
|
|
1904
|
+
__riscv_vsrl_vx_u8m1(r1, 4, vl),
|
|
1905
|
+
16,
|
|
1906
|
+
__riscv_vand_vx_u8m1(r2, 0x03, vl),
|
|
1907
|
+
vl);
|
|
1908
|
+
vfloat32m4_t vb2 = __riscv_vle32_v_f32m4(pb2 + g - vl, vl);
|
|
1909
|
+
vfloat32m4_t f2 = __riscv_vfwcvt_f_xu_v_f32m4(
|
|
1910
|
+
__riscv_vzext_vf2_u16m2(c2, vl), vl);
|
|
1911
|
+
acc = __riscv_vfmacc_vv_f32m4(acc, vb2, f2, vl);
|
|
1912
|
+
}
|
|
1913
|
+
}
|
|
1914
|
+
|
|
1915
|
+
// Tail groups: one shorter-vl pass into the same accumulator
|
|
1916
|
+
// (only the first vt lanes are touched; reduction below
|
|
1917
|
+
// covers all vl lanes).
|
|
1918
|
+
if (g < ngf) {
|
|
1919
|
+
const size_t vt = __riscv_vsetvl_e8m1(ngf - g);
|
|
1920
|
+
|
|
1921
|
+
vuint8m1x3_t seg = __riscv_vlseg3e8_v_u8m1x3(code + 3 * g, vt);
|
|
1922
|
+
vuint8m1_t r0 = __riscv_vget_v_u8m1x3_u8m1(seg, 0);
|
|
1923
|
+
vuint8m1_t r1 = __riscv_vget_v_u8m1x3_u8m1(seg, 1);
|
|
1924
|
+
vuint8m1_t r2 = __riscv_vget_v_u8m1x3_u8m1(seg, 2);
|
|
1925
|
+
|
|
1926
|
+
vuint8m1_t c0 = __riscv_vand_vx_u8m1(r0, 0x3F, vt);
|
|
1927
|
+
vuint8m1_t c1 = __riscv_vmacc_vx_u8m1(
|
|
1928
|
+
__riscv_vsrl_vx_u8m1(r0, 6, vt),
|
|
1929
|
+
4,
|
|
1930
|
+
__riscv_vand_vx_u8m1(r1, 0x0F, vt),
|
|
1931
|
+
vt);
|
|
1932
|
+
vuint8m1_t c2 = __riscv_vmacc_vx_u8m1(
|
|
1933
|
+
__riscv_vsrl_vx_u8m1(r1, 4, vt),
|
|
1934
|
+
16,
|
|
1935
|
+
__riscv_vand_vx_u8m1(r2, 0x03, vt),
|
|
1936
|
+
vt);
|
|
1937
|
+
vuint8m1_t c3 = __riscv_vsrl_vx_u8m1(r2, 2, vt);
|
|
1938
|
+
|
|
1939
|
+
vfloat32m4_t f0 = __riscv_vfwcvt_f_xu_v_f32m4(
|
|
1940
|
+
__riscv_vzext_vf2_u16m2(c0, vt), vt);
|
|
1941
|
+
acc = __riscv_vfmacc_vv_f32m4_tu(
|
|
1942
|
+
acc, __riscv_vle32_v_f32m4(pb0 + g, vt), f0, vt);
|
|
1943
|
+
|
|
1944
|
+
vfloat32m4_t f1 = __riscv_vfwcvt_f_xu_v_f32m4(
|
|
1945
|
+
__riscv_vzext_vf2_u16m2(c1, vt), vt);
|
|
1946
|
+
acc = __riscv_vfmacc_vv_f32m4_tu(
|
|
1947
|
+
acc, __riscv_vle32_v_f32m4(pb1 + g, vt), f1, vt);
|
|
1948
|
+
|
|
1949
|
+
vfloat32m4_t f2 = __riscv_vfwcvt_f_xu_v_f32m4(
|
|
1950
|
+
__riscv_vzext_vf2_u16m2(c2, vt), vt);
|
|
1951
|
+
acc = __riscv_vfmacc_vv_f32m4_tu(
|
|
1952
|
+
acc, __riscv_vle32_v_f32m4(pb2 + g, vt), f2, vt);
|
|
1953
|
+
|
|
1954
|
+
vfloat32m4_t f3 = __riscv_vfwcvt_f_xu_v_f32m4(
|
|
1955
|
+
__riscv_vzext_vf2_u16m2(c3, vt), vt);
|
|
1956
|
+
acc = __riscv_vfmacc_vv_f32m4_tu(
|
|
1957
|
+
acc, __riscv_vle32_v_f32m4(pb3 + g, vt), f3, vt);
|
|
1958
|
+
}
|
|
1959
|
+
|
|
1960
|
+
// Single horizontal reduction over all vl lanes.
|
|
1961
|
+
vfloat32m1_t red = __riscv_vfredusum_vs_f32m4_f32m1(
|
|
1962
|
+
acc, __riscv_vfmv_v_f_f32m1(0.0f, 1), vl);
|
|
1963
|
+
total = __riscv_vfmv_f_s_f32m1_f32(red);
|
|
1964
|
+
}
|
|
1965
|
+
|
|
1966
|
+
// Scalar tail: dims beyond the last full group (d % 4). K_q
|
|
1967
|
+
// already contains their q_i * rmin_i part; add b_i * c_i.
|
|
1968
|
+
for (size_t i = 4 * ngf; i < d; i++) {
|
|
1969
|
+
total += q[i] * a_all[i] * float(sq6_decode_raw(code, i));
|
|
1970
|
+
}
|
|
1971
|
+
return k_q + total;
|
|
1972
|
+
}
|
|
1973
|
+
|
|
1974
|
+
float query_to_code(const uint8_t* code) const final {
|
|
1975
|
+
return compute_ip(code);
|
|
1976
|
+
}
|
|
1977
|
+
|
|
1978
|
+
float symmetric_dis(idx_t i, idx_t j) override {
|
|
1979
|
+
// Not on the benchmark-critical path; scalar per-dim evaluation.
|
|
1980
|
+
// IP(recon(c1), recon(c2)) = sum_k recon1_k * recon2_k
|
|
1981
|
+
const uint8_t* c1 = codes + i * code_size;
|
|
1982
|
+
const uint8_t* c2 = codes + j * code_size;
|
|
1983
|
+
float acc = 0;
|
|
1984
|
+
for (size_t k = 0; k < d; k++) {
|
|
1985
|
+
float r1 = rmin_all[k] + a_all[k] * float(sq6_decode_raw(c1, k));
|
|
1986
|
+
float r2 = rmin_all[k] + a_all[k] * float(sq6_decode_raw(c2, k));
|
|
1987
|
+
acc += r1 * r2;
|
|
1988
|
+
}
|
|
1989
|
+
return acc;
|
|
1990
|
+
}
|
|
1991
|
+
|
|
1992
|
+
void query_to_codes_batch_4(
|
|
1993
|
+
const uint8_t* code_0,
|
|
1994
|
+
const uint8_t* code_1,
|
|
1995
|
+
const uint8_t* code_2,
|
|
1996
|
+
const uint8_t* code_3,
|
|
1997
|
+
float& dis0,
|
|
1998
|
+
float& dis1,
|
|
1999
|
+
float& dis2,
|
|
2000
|
+
float& dis3) const final {
|
|
2001
|
+
dis0 = compute_ip(code_0);
|
|
2002
|
+
dis1 = compute_ip(code_1);
|
|
2003
|
+
dis2 = compute_ip(code_2);
|
|
2004
|
+
dis3 = compute_ip(code_3);
|
|
2005
|
+
}
|
|
2006
|
+
};
|
|
2007
|
+
|
|
2008
|
+
/**********************************************************
|
|
2009
|
+
* QT_8bit (NON_UNIFORM) + L2 — full RVV specialization
|
|
2010
|
+
*
|
|
2011
|
+
* Per-dimension affine reconstruction:
|
|
2012
|
+
* recon_i(c) = vmin[i] + vdiff[i] * (c + 0.5) / 255
|
|
2013
|
+
* = rmin_i + a_i * c, with a_i = vdiff[i] / 255,
|
|
2014
|
+
* rmin_i = vmin[i] + 0.5 * a_i
|
|
2015
|
+
* L2 distance per dimension contributes (q_i - recon_i(c))²
|
|
2016
|
+
* = (e_i - a_i * c_i)² with e_i = q_i - rmin_i.
|
|
2017
|
+
* The constructor precomputes a_i / rmin_i; set_query precomputes
|
|
2018
|
+
* e_i once per query (amortized over all codes of the scan).
|
|
2019
|
+
*
|
|
2020
|
+
* Unlike the 4/6-bit codecs there is no bit-field unpacking: each
|
|
2021
|
+
* code byte is one dimension, so the kernel is a single contiguous
|
|
2022
|
+
* stream. Hot loop per vl = VLMAX(e8m1) dims (16 at VLEN=128):
|
|
2023
|
+
* vle8 -> vzext_vf2 (u8->u16m2) -> vfwcvt (u16->f32m4), then
|
|
2024
|
+
* t = e - a*c (vfnmsac) and acc += t*t (vfmacc).
|
|
2025
|
+
* vsetvl is hoisted: 0 explicit vsetvl in the hot loop, 1 hoisted
|
|
2026
|
+
* + 1 for the tail (d % vl != 0; absent for the benchmark d=768).
|
|
2027
|
+
**********************************************************/
|
|
2028
|
+
|
|
2029
|
+
template <>
|
|
2030
|
+
struct DCTemplate<
|
|
2031
|
+
QuantizerTemplate<
|
|
2032
|
+
Codec8bit<SIMDLevel::RISCV_RVV>,
|
|
2033
|
+
QuantizerTemplateScaling::NON_UNIFORM,
|
|
2034
|
+
SIMDLevel::RISCV_RVV>,
|
|
2035
|
+
SimilarityL2<SIMDLevel::RISCV_RVV>,
|
|
2036
|
+
SIMDLevel::RISCV_RVV> : SQDistanceComputer {
|
|
2037
|
+
using Sim = SimilarityL2<SIMDLevel::RISCV_RVV>;
|
|
2038
|
+
|
|
2039
|
+
size_t d;
|
|
2040
|
+
std::vector<float> a_v; // a_i = vdiff[i] / 255
|
|
2041
|
+
std::vector<float> rmin_v; // vmin[i] + 0.5 * a_i
|
|
2042
|
+
std::vector<float> e_v; // q_i - rmin_i, per query
|
|
2043
|
+
|
|
2044
|
+
DCTemplate(size_t d_in, const std::vector<float>& trained)
|
|
2045
|
+
: d(d_in), a_v(d_in, 0.0f), rmin_v(d_in, 0.0f), e_v(d_in, 0.0f) {
|
|
2046
|
+
sq8::make_coefficients(trained.data(), d_in, a_v.data(), rmin_v.data());
|
|
2047
|
+
}
|
|
2048
|
+
|
|
2049
|
+
void set_query(const float* x) final {
|
|
2050
|
+
q = x;
|
|
2051
|
+
// e_i = q_i - rmin_i (amortized: once per query, 1/n weight)
|
|
2052
|
+
for (size_t i = 0; i < d; i++) {
|
|
2053
|
+
e_v[i] = x[i] - rmin_v[i];
|
|
2054
|
+
}
|
|
2055
|
+
}
|
|
2056
|
+
|
|
2057
|
+
/// Full-precision vector L2 over the 1-byte-per-dim code.
|
|
2058
|
+
float compute_l2(const uint8_t* code) const {
|
|
2059
|
+
return sq8_rvv::l2_distance(code, a_v.data(), e_v.data(), d);
|
|
2060
|
+
}
|
|
2061
|
+
|
|
2062
|
+
float query_to_code(const uint8_t* code) const final {
|
|
2063
|
+
return compute_l2(code);
|
|
2064
|
+
}
|
|
2065
|
+
|
|
2066
|
+
float symmetric_dis(idx_t i, idx_t j) override {
|
|
2067
|
+
// Not on the benchmark-critical path; scalar per-dim evaluation.
|
|
2068
|
+
// recon1 - recon2 = a_k * (c1_k - c2_k)
|
|
2069
|
+
const uint8_t* c1 = codes + i * code_size;
|
|
2070
|
+
const uint8_t* c2 = codes + j * code_size;
|
|
2071
|
+
float acc = 0;
|
|
2072
|
+
for (size_t k = 0; k < d; k++) {
|
|
2073
|
+
float diff = a_v[k] * float(int(c1[k]) - int(c2[k]));
|
|
2074
|
+
acc += diff * diff;
|
|
2075
|
+
}
|
|
2076
|
+
return acc;
|
|
2077
|
+
}
|
|
2078
|
+
|
|
2079
|
+
void query_to_codes_batch_4(
|
|
2080
|
+
const uint8_t* code_0,
|
|
2081
|
+
const uint8_t* code_1,
|
|
2082
|
+
const uint8_t* code_2,
|
|
2083
|
+
const uint8_t* code_3,
|
|
2084
|
+
float& dis0,
|
|
2085
|
+
float& dis1,
|
|
2086
|
+
float& dis2,
|
|
2087
|
+
float& dis3) const final {
|
|
2088
|
+
dis0 = compute_l2(code_0);
|
|
2089
|
+
dis1 = compute_l2(code_1);
|
|
2090
|
+
dis2 = compute_l2(code_2);
|
|
2091
|
+
dis3 = compute_l2(code_3);
|
|
2092
|
+
}
|
|
2093
|
+
};
|
|
2094
|
+
|
|
2095
|
+
/**********************************************************
|
|
2096
|
+
* QT_8bit (NON_UNIFORM) + IP — full RVV specialization
|
|
2097
|
+
*
|
|
2098
|
+
* recon_i(c) = rmin_i + a_i*c_i, a_i = vdiff[i]/255, rmin_i = vmin[i] +
|
|
2099
|
+
* 0.5*a_i. IP decomposes into a query-only constant plus a coefficient
|
|
2100
|
+
* dot product over the integer codes:
|
|
2101
|
+
* IP(q, code) = sum_i q_i*rmin_i + sum_i (q_i*a_i)*c_i = K_q + sum_i
|
|
2102
|
+
* b_i*c_i K_q and b_i = q_i*a_i are precomputed in set_query. Each code byte is
|
|
2103
|
+
* one dim (no bit unpacking): hot loop vle8 -> vzext -> vfcvt -> vle32 ->
|
|
2104
|
+
* vfmacc, software-pipelined.
|
|
2105
|
+
**********************************************************/
|
|
2106
|
+
|
|
2107
|
+
template <>
|
|
2108
|
+
struct DCTemplate<
|
|
2109
|
+
QuantizerTemplate<
|
|
2110
|
+
Codec8bit<SIMDLevel::RISCV_RVV>,
|
|
2111
|
+
QuantizerTemplateScaling::NON_UNIFORM,
|
|
2112
|
+
SIMDLevel::RISCV_RVV>,
|
|
2113
|
+
SimilarityIP<SIMDLevel::RISCV_RVV>,
|
|
2114
|
+
SIMDLevel::RISCV_RVV> : SQDistanceComputer {
|
|
2115
|
+
using Sim = SimilarityIP<SIMDLevel::RISCV_RVV>;
|
|
2116
|
+
|
|
2117
|
+
size_t d;
|
|
2118
|
+
std::vector<float> a_v; // a_i = vdiff[i] / 255
|
|
2119
|
+
std::vector<float> rmin_v; // vmin[i] + 0.5 * a_i
|
|
2120
|
+
std::vector<float> b_v; // b_i = q_i * a_i, per query
|
|
2121
|
+
float k_q; // sum_i q_i * rmin_i, per query
|
|
2122
|
+
|
|
2123
|
+
DCTemplate(size_t d_in, const std::vector<float>& trained)
|
|
2124
|
+
: d(d_in),
|
|
2125
|
+
a_v(d_in, 0.0f),
|
|
2126
|
+
rmin_v(d_in, 0.0f),
|
|
2127
|
+
// Pad b_v by 2*VLMAX so the pipelined next-iteration
|
|
2128
|
+
// vle32 loads stay in bounds (padding lanes never used).
|
|
2129
|
+
b_v(d_in + 2 * __riscv_vsetvlmax_e8m1(), 0.0f),
|
|
2130
|
+
k_q(0.0f) {
|
|
2131
|
+
sq8::make_coefficients(trained.data(), d_in, a_v.data(), rmin_v.data());
|
|
2132
|
+
}
|
|
2133
|
+
|
|
2134
|
+
void set_query(const float* x) final {
|
|
2135
|
+
q = x;
|
|
2136
|
+
// b_i = q_i * a_i and K_q = sum_i q_i * rmin_i, once per
|
|
2137
|
+
// query (amortized 1/n over the codes of the scan).
|
|
2138
|
+
float acc = 0.0f;
|
|
2139
|
+
for (size_t i = 0; i < d; i++) {
|
|
2140
|
+
b_v[i] = x[i] * a_v[i];
|
|
2141
|
+
acc += x[i] * rmin_v[i];
|
|
2142
|
+
}
|
|
2143
|
+
k_q = acc;
|
|
2144
|
+
}
|
|
2145
|
+
|
|
2146
|
+
/// S = sum_i b_i * c_i over the 1-byte-per-dim code; returns
|
|
2147
|
+
/// K_q + S.
|
|
2148
|
+
float compute_ip(const uint8_t* code) const {
|
|
2149
|
+
const float* pb = b_v.data();
|
|
2150
|
+
|
|
2151
|
+
// Hoist vsetvl: VLMAX for e8m1 (== f32 lanes per m4 group).
|
|
2152
|
+
// d==0: guard so vsetvl(0) doesn't return 0 and the prologue
|
|
2153
|
+
// `if (i + vl <= d)` + chunk loop don't stall.
|
|
2154
|
+
const size_t vl = __riscv_vsetvl_e8m1(d > 0 ? d : 1);
|
|
2155
|
+
vfloat32m4_t acc = __riscv_vfmv_v_f_f32m4(0.0f, vl);
|
|
2156
|
+
// Second accumulator: the 2x-unrolled loop feeds acc/acc1
|
|
2157
|
+
// alternately, halving the serial FMA chain.
|
|
2158
|
+
vfloat32m4_t acc1 = __riscv_vfmv_v_f_f32m4(0.0f, vl);
|
|
2159
|
+
|
|
2160
|
+
size_t i = 0;
|
|
2161
|
+
if (i + vl <= d) {
|
|
2162
|
+
// Software pipeline depth 1: preload the next chunk's code
|
|
2163
|
+
// bytes at the top of this iteration.
|
|
2164
|
+
vuint8m1_t c8 = __riscv_vle8_v_u8m1(code + i, vl); // prologue
|
|
2165
|
+
i += vl;
|
|
2166
|
+
|
|
2167
|
+
// 2x unroll + dual accumulators; the b stream is also
|
|
2168
|
+
// software-pipelined (loads past d land in b_v's zero
|
|
2169
|
+
// padding and are discarded on exit).
|
|
2170
|
+
if (i + 2 * vl <= d) {
|
|
2171
|
+
// b-prologue: preload both halves of iteration 0.
|
|
2172
|
+
vfloat32m4_t vb0 = __riscv_vle32_v_f32m4(pb + i - vl, vl);
|
|
2173
|
+
vfloat32m4_t vb1 = __riscv_vle32_v_f32m4(pb + i, vl);
|
|
2174
|
+
for (; i + 2 * vl <= d; i += 2 * vl) {
|
|
2175
|
+
vuint8m1_t c8_n1 = __riscv_vle8_v_u8m1(code + i, vl);
|
|
2176
|
+
vfloat32m4_t vb0_next =
|
|
2177
|
+
__riscv_vle32_v_f32m4(pb + i + vl, vl);
|
|
2178
|
+
vfloat32m4_t cf0 = __riscv_vfcvt_f_xu_v_f32m4(
|
|
2179
|
+
__riscv_vzext_vf4_u32m4(c8, vl), vl);
|
|
2180
|
+
acc = __riscv_vfmacc_vv_f32m4(acc, vb0, cf0, vl);
|
|
2181
|
+
|
|
2182
|
+
vuint8m1_t c8_n2 = __riscv_vle8_v_u8m1(code + i + vl, vl);
|
|
2183
|
+
vfloat32m4_t vb1_next =
|
|
2184
|
+
__riscv_vle32_v_f32m4(pb + i + 2 * vl, vl);
|
|
2185
|
+
vfloat32m4_t cf1 = __riscv_vfcvt_f_xu_v_f32m4(
|
|
2186
|
+
__riscv_vzext_vf4_u32m4(c8_n1, vl), vl);
|
|
2187
|
+
acc1 = __riscv_vfmacc_vv_f32m4(acc1, vb1, cf1, vl);
|
|
2188
|
+
|
|
2189
|
+
c8 = c8_n2; // rotate code stream
|
|
2190
|
+
vb0 = vb0_next; // rotate b stream
|
|
2191
|
+
vb1 = vb1_next;
|
|
2192
|
+
}
|
|
2193
|
+
}
|
|
2194
|
+
|
|
2195
|
+
// Odd full chunk left beyond the current one: process the
|
|
2196
|
+
// current chunk and advance the pipeline by one.
|
|
2197
|
+
if (i + vl <= d) {
|
|
2198
|
+
vuint8m1_t c8_n1 = __riscv_vle8_v_u8m1(code + i, vl);
|
|
2199
|
+
vfloat32m4_t vb = __riscv_vle32_v_f32m4(pb + i - vl, vl);
|
|
2200
|
+
vfloat32m4_t cf = __riscv_vfcvt_f_xu_v_f32m4(
|
|
2201
|
+
__riscv_vzext_vf4_u32m4(c8, vl), vl);
|
|
2202
|
+
acc = __riscv_vfmacc_vv_f32m4(acc, vb, cf, vl);
|
|
2203
|
+
c8 = c8_n1;
|
|
2204
|
+
i += vl;
|
|
2205
|
+
}
|
|
2206
|
+
|
|
2207
|
+
// Epilogue: process the last full chunk (already loaded).
|
|
2208
|
+
{
|
|
2209
|
+
vfloat32m4_t vb = __riscv_vle32_v_f32m4(pb + i - vl, vl);
|
|
2210
|
+
vfloat32m4_t cf = __riscv_vfcvt_f_xu_v_f32m4(
|
|
2211
|
+
__riscv_vzext_vf4_u32m4(c8, vl), vl);
|
|
2212
|
+
acc1 = __riscv_vfmacc_vv_f32m4(acc1, vb, cf, vl);
|
|
2213
|
+
}
|
|
2214
|
+
}
|
|
2215
|
+
|
|
2216
|
+
// Tail: fewer than vl dims left; accumulates into the first
|
|
2217
|
+
// lanes of acc (safe: reduction below covers all vl lanes).
|
|
2218
|
+
if (i < d) {
|
|
2219
|
+
size_t vt = __riscv_vsetvl_e8m1(d - i);
|
|
2220
|
+
vuint8m1_t c8 = __riscv_vle8_v_u8m1(code + i, vt);
|
|
2221
|
+
vfloat32m4_t cf = __riscv_vfcvt_f_xu_v_f32m4(
|
|
2222
|
+
__riscv_vzext_vf4_u32m4(c8, vt), vt);
|
|
2223
|
+
acc = __riscv_vfmacc_vv_f32m4_tu(
|
|
2224
|
+
acc, __riscv_vle32_v_f32m4(pb + i, vt), cf, vt);
|
|
2225
|
+
}
|
|
2226
|
+
|
|
2227
|
+
// Merge the two accumulators, then horizontal reduce.
|
|
2228
|
+
acc = __riscv_vfadd_vv_f32m4(acc, acc1, vl);
|
|
2229
|
+
// Horizontal reduce over all vl lanes
|
|
2230
|
+
vfloat32m1_t red = __riscv_vfredusum_vs_f32m4_f32m1(
|
|
2231
|
+
acc, __riscv_vfmv_v_f_f32m1(0.0f, 1), vl);
|
|
2232
|
+
return k_q + __riscv_vfmv_f_s_f32m1_f32(red);
|
|
2233
|
+
}
|
|
2234
|
+
|
|
2235
|
+
float query_to_code(const uint8_t* code) const final {
|
|
2236
|
+
return compute_ip(code);
|
|
2237
|
+
}
|
|
2238
|
+
|
|
2239
|
+
float symmetric_dis(idx_t i, idx_t j) override {
|
|
2240
|
+
// Not on the benchmark-critical path; scalar per-dim
|
|
2241
|
+
// evaluation: IP(recon(c1), recon(c2)) = sum_k r1_k * r2_k
|
|
2242
|
+
const uint8_t* c1 = codes + i * code_size;
|
|
2243
|
+
const uint8_t* c2 = codes + j * code_size;
|
|
2244
|
+
float acc = 0;
|
|
2245
|
+
for (size_t k = 0; k < d; k++) {
|
|
2246
|
+
float r1 = rmin_v[k] + a_v[k] * float(c1[k]);
|
|
2247
|
+
float r2 = rmin_v[k] + a_v[k] * float(c2[k]);
|
|
2248
|
+
acc += r1 * r2;
|
|
2249
|
+
}
|
|
2250
|
+
return acc;
|
|
2251
|
+
}
|
|
2252
|
+
|
|
2253
|
+
void query_to_codes_batch_4(
|
|
2254
|
+
const uint8_t* code_0,
|
|
2255
|
+
const uint8_t* code_1,
|
|
2256
|
+
const uint8_t* code_2,
|
|
2257
|
+
const uint8_t* code_3,
|
|
2258
|
+
float& dis0,
|
|
2259
|
+
float& dis1,
|
|
2260
|
+
float& dis2,
|
|
2261
|
+
float& dis3) const final {
|
|
2262
|
+
dis0 = compute_ip(code_0);
|
|
2263
|
+
dis1 = compute_ip(code_1);
|
|
2264
|
+
dis2 = compute_ip(code_2);
|
|
2265
|
+
dis3 = compute_ip(code_3);
|
|
2266
|
+
}
|
|
2267
|
+
};
|
|
2268
|
+
|
|
2269
|
+
/**********************************************************
|
|
2270
|
+
* QT_8bit_uniform + L2 — float-domain RVV specialization
|
|
2271
|
+
*
|
|
2272
|
+
* 8-bit UNIFORM scaling: every component reconstructs as
|
|
2273
|
+
* recon(c) = vmin + vdiff * (c + 0.5) / 255 = c0 + a * c
|
|
2274
|
+
* with a = vdiff / 255 and c0 = vmin + 0.5 * a.
|
|
2275
|
+
*
|
|
2276
|
+
* L2 distance is evaluated in the exact affine float domain,
|
|
2277
|
+
* L2(q, code) = sum_i (q_i - recon(c_i))^2
|
|
2278
|
+
* = sum_i (e_i - a * c_i)^2 with e_i = q_i - c0,
|
|
2279
|
+
* so the ORIGINAL float query is retained — no integer-grid
|
|
2280
|
+
* pre-quantization — and the result matches the scalar NONE
|
|
2281
|
+
* reference to float-reassociation precision. vdiff == 0
|
|
2282
|
+
* degenerates naturally (a == 0 -> L2 = sum_i (q_i - vmin)^2).
|
|
2283
|
+
*
|
|
2284
|
+
* Hot loop per vl = VLMAX(e8m1) dims (16 at VLEN=128): vle8 ->
|
|
2285
|
+
* vzext_vf4 (u8 -> u32m4) -> vfcvt (u32 -> f32m4), then
|
|
2286
|
+
* t = e - a*c via the scalar-operand vfnmsac.vf (a is a SHARED
|
|
2287
|
+
* scalar constant — no per-dimension coefficient stream) and
|
|
2288
|
+
* acc += t*t (vfmacc). vsetvl hoisted: 1 hoisted + 1 for the tail.
|
|
2289
|
+
**********************************************************/
|
|
2290
|
+
|
|
2291
|
+
template <>
|
|
2292
|
+
struct DCTemplate<
|
|
2293
|
+
QuantizerTemplate<
|
|
2294
|
+
Codec8bit<SIMDLevel::RISCV_RVV>,
|
|
2295
|
+
QuantizerTemplateScaling::UNIFORM,
|
|
2296
|
+
SIMDLevel::RISCV_RVV>,
|
|
2297
|
+
SimilarityL2<SIMDLevel::RISCV_RVV>,
|
|
2298
|
+
SIMDLevel::RISCV_RVV> : SQDistanceComputer {
|
|
2299
|
+
using Sim = SimilarityL2<SIMDLevel::RISCV_RVV>;
|
|
2300
|
+
|
|
2301
|
+
size_t d;
|
|
2302
|
+
float a; // vdiff / 255
|
|
2303
|
+
float c0; // vmin + 0.5 * a
|
|
2304
|
+
std::vector<float> e_v; // q_i - c0, per query
|
|
2305
|
+
|
|
2306
|
+
DCTemplate(size_t d_in, const std::vector<float>& trained)
|
|
2307
|
+
: d(d_in), e_v(d_in, 0.0f) {
|
|
2308
|
+
a = trained[1] / 255.0f;
|
|
2309
|
+
c0 = trained[0] + 0.5f * a;
|
|
2310
|
+
}
|
|
2311
|
+
|
|
2312
|
+
void set_query(const float* x) final {
|
|
2313
|
+
q = x;
|
|
2314
|
+
// e_i = q_i - c0, once per query (amortized over all codes).
|
|
2315
|
+
for (size_t i = 0; i < d; i++) {
|
|
2316
|
+
e_v[i] = x[i] - c0;
|
|
2317
|
+
}
|
|
2318
|
+
}
|
|
2319
|
+
|
|
2320
|
+
/// Exact affine float-domain L2: sum_i (e_i - a*c_i)^2.
|
|
2321
|
+
/// Hot loop per chunk: vle8 (code byte) -> vzext_vf4 -> vfcvt to
|
|
2322
|
+
/// f32, then t = e - a*c (vfnmsac.vf, a shared scalar) and
|
|
2323
|
+
/// acc += t*t (vfmacc). 5 vector ops per chunk, single f32m4
|
|
2324
|
+
/// accumulator, single vfredusum at the end. vsetvl hoisted.
|
|
2325
|
+
float compute_l2(const uint8_t* code) const {
|
|
2326
|
+
const float* pe = e_v.data();
|
|
2327
|
+
|
|
2328
|
+
// Hoist vsetvl: VLMAX for e8m1 (== f32m4 lanes).
|
|
2329
|
+
// d==0 guard: vsetvl(0) would return 0 and stall the prologue.
|
|
2330
|
+
const size_t vl = __riscv_vsetvl_e8m1(d > 0 ? d : 1);
|
|
2331
|
+
vfloat32m4_t acc = __riscv_vfmv_v_f_f32m4(0.0f, vl);
|
|
2332
|
+
|
|
2333
|
+
size_t i = 0;
|
|
2334
|
+
if (i + vl <= d) {
|
|
2335
|
+
vuint8m1_t c8 = __riscv_vle8_v_u8m1(code + i, vl); // prologue
|
|
2336
|
+
i += vl;
|
|
2337
|
+
for (; i + vl <= d; i += vl) {
|
|
2338
|
+
vuint8m1_t c8_next = __riscv_vle8_v_u8m1(code + i, vl);
|
|
2339
|
+
vfloat32m4_t cf = __riscv_vfcvt_f_xu_v_f32m4(
|
|
2340
|
+
__riscv_vzext_vf4_u32m4(c8, vl), vl);
|
|
2341
|
+
vfloat32m4_t t = __riscv_vle32_v_f32m4(pe + i - vl, vl);
|
|
2342
|
+
t = __riscv_vfnmsac_vf_f32m4(t, a, cf, vl);
|
|
2343
|
+
acc = __riscv_vfmacc_vv_f32m4(acc, t, t, vl);
|
|
2344
|
+
c8 = c8_next;
|
|
2345
|
+
}
|
|
2346
|
+
// Epilogue: process the last full chunk (already loaded).
|
|
2347
|
+
{
|
|
2348
|
+
vfloat32m4_t cf = __riscv_vfcvt_f_xu_v_f32m4(
|
|
2349
|
+
__riscv_vzext_vf4_u32m4(c8, vl), vl);
|
|
2350
|
+
vfloat32m4_t t = __riscv_vle32_v_f32m4(pe + i - vl, vl);
|
|
2351
|
+
t = __riscv_vfnmsac_vf_f32m4(t, a, cf, vl);
|
|
2352
|
+
acc = __riscv_vfmacc_vv_f32m4(acc, t, t, vl);
|
|
2353
|
+
}
|
|
2354
|
+
}
|
|
2355
|
+
|
|
2356
|
+
// Tail: fewer than vl dims left — one shorter-vl pass into the
|
|
2357
|
+
// same accumulator; _tu keeps lanes past vt (accumulated by
|
|
2358
|
+
// prior chunks) intact for the full-vl reduction below.
|
|
2359
|
+
if (i < d) {
|
|
2360
|
+
const size_t vt = __riscv_vsetvl_e8m1(d - i);
|
|
2361
|
+
vuint8m1_t c8 = __riscv_vle8_v_u8m1(code + i, vt);
|
|
2362
|
+
vfloat32m4_t cf = __riscv_vfcvt_f_xu_v_f32m4(
|
|
2363
|
+
__riscv_vzext_vf4_u32m4(c8, vt), vt);
|
|
2364
|
+
vfloat32m4_t t = __riscv_vle32_v_f32m4(pe + i, vt);
|
|
2365
|
+
t = __riscv_vfnmsac_vf_f32m4(t, a, cf, vt);
|
|
2366
|
+
acc = __riscv_vfmacc_vv_f32m4_tu(acc, t, t, vt);
|
|
2367
|
+
}
|
|
2368
|
+
|
|
2369
|
+
// Horizontal reduce over all vl lanes.
|
|
2370
|
+
vfloat32m1_t red = __riscv_vfredusum_vs_f32m4_f32m1(
|
|
2371
|
+
acc, __riscv_vfmv_v_f_f32m1(0.0f, 1), vl);
|
|
2372
|
+
return __riscv_vfmv_f_s_f32m1_f32(red);
|
|
2373
|
+
}
|
|
2374
|
+
|
|
2375
|
+
float query_to_code(const uint8_t* code) const final {
|
|
2376
|
+
return compute_l2(code);
|
|
2377
|
+
}
|
|
2378
|
+
|
|
2379
|
+
float symmetric_dis(idx_t i, idx_t j) override {
|
|
2380
|
+
// Not on the benchmark-critical path; scalar per-dim evaluation.
|
|
2381
|
+
// recon1 - recon2 = a * (c1_k - c2_k). Scale each difference
|
|
2382
|
+
// BEFORE squaring: factoring a*a out of the integer sum turns
|
|
2383
|
+
// identical codes into 0 * inf == NaN when a*a overflows (huge
|
|
2384
|
+
// finite vdiff), while the scalar formulation returns exact 0.
|
|
2385
|
+
const uint8_t* c1 = codes + i * code_size;
|
|
2386
|
+
const uint8_t* c2 = codes + j * code_size;
|
|
2387
|
+
float acc = 0;
|
|
2388
|
+
for (size_t k = 0; k < d; k++) {
|
|
2389
|
+
float diff = a * float(int(c1[k]) - int(c2[k]));
|
|
2390
|
+
acc += diff * diff;
|
|
2391
|
+
}
|
|
2392
|
+
return acc;
|
|
2393
|
+
}
|
|
2394
|
+
|
|
2395
|
+
void query_to_codes_batch_4(
|
|
2396
|
+
const uint8_t* code_0,
|
|
2397
|
+
const uint8_t* code_1,
|
|
2398
|
+
const uint8_t* code_2,
|
|
2399
|
+
const uint8_t* code_3,
|
|
2400
|
+
float& dis0,
|
|
2401
|
+
float& dis1,
|
|
2402
|
+
float& dis2,
|
|
2403
|
+
float& dis3) const final {
|
|
2404
|
+
dis0 = compute_l2(code_0);
|
|
2405
|
+
dis1 = compute_l2(code_1);
|
|
2406
|
+
dis2 = compute_l2(code_2);
|
|
2407
|
+
dis3 = compute_l2(code_3);
|
|
2408
|
+
}
|
|
2409
|
+
};
|
|
2410
|
+
|
|
2411
|
+
/**********************************************************
|
|
2412
|
+
* QT_8bit_uniform + IP — float-domain RVV specialization
|
|
2413
|
+
*
|
|
2414
|
+
* 8-bit UNIFORM scaling: every component reconstructs as
|
|
2415
|
+
* recon(c) = vmin + vdiff * (c + 0.5) / 255 = c0 + scale * c
|
|
2416
|
+
* with scale = vdiff / 255 and c0 = vmin + 0.5 * scale (SHARED
|
|
2417
|
+
* scalar constants). The inner product folds into a query-only
|
|
2418
|
+
* constant plus one integer-code dot product over the per-dim
|
|
2419
|
+
* pre-scaled query b_i = q_i * scale:
|
|
2420
|
+
* IP(q, code) = sum_i q_i*c0 + sum_i (q_i * scale) * c_i
|
|
2421
|
+
* = K_q + S
|
|
2422
|
+
* K_q = sum_i (q_i * c0) is precomputed once per query in set_query
|
|
2423
|
+
* (amortized 1/n over the codes of the scan), so the hot loop only
|
|
2424
|
+
* evaluates S = sum_i b_i * c_i.
|
|
2425
|
+
*
|
|
2426
|
+
* The ORIGINAL float query is retained — no fixed-point / integer
|
|
2427
|
+
* pre-quantization of the query — so scores and rankings match the
|
|
2428
|
+
* scalar NONE reference to float-reassociation precision (same
|
|
2429
|
+
* semantics class as the QT_8bit_uniform+L2 and QT_8bit_nonuniform
|
|
2430
|
+
* IP float-domain kernels). Per-dim pre-scaling (as the
|
|
2431
|
+
* QT_8bit_nonuniform+IP kernel already does) keeps every term
|
|
2432
|
+
* inside the scalar per-dim semantics: scale == 0 or c0 == 0 zeroes
|
|
2433
|
+
* its term exactly, so a huge finite query can never overflow a
|
|
2434
|
+
* bare c0*sum(q) or sum(q_i*c_i) intermediate into 0*inf == NaN.
|
|
2435
|
+
* No special-casing.
|
|
2436
|
+
*
|
|
2437
|
+
* Hot loop per vl = VLMAX(e8m1) dims (16 at VLEN=128): vle8 ->
|
|
2438
|
+
* vzext_vf4 (u8 -> u32m4) -> vfcvt (u32 -> f32m4), then
|
|
2439
|
+
* acc += b * c (vfmacc). 4 vector ops per chunk, single f32m4
|
|
2440
|
+
* accumulator, single vfredusum at the end. vsetvl hoisted:
|
|
2441
|
+
* 1 hoisted + 1 for the tail.
|
|
2442
|
+
**********************************************************/
|
|
2443
|
+
|
|
2444
|
+
template <>
|
|
2445
|
+
struct DCTemplate<
|
|
2446
|
+
QuantizerTemplate<
|
|
2447
|
+
Codec8bit<SIMDLevel::RISCV_RVV>,
|
|
2448
|
+
QuantizerTemplateScaling::UNIFORM,
|
|
2449
|
+
SIMDLevel::RISCV_RVV>,
|
|
2450
|
+
SimilarityIP<SIMDLevel::RISCV_RVV>,
|
|
2451
|
+
SIMDLevel::RISCV_RVV> : SQDistanceComputer {
|
|
2452
|
+
using Sim = SimilarityIP<SIMDLevel::RISCV_RVV>;
|
|
2453
|
+
|
|
2454
|
+
size_t d;
|
|
2455
|
+
float scale; // vdiff / 255
|
|
2456
|
+
float c0; // vmin + 0.5 * scale
|
|
2457
|
+
float k_q; // sum_i (q_i * c0), per query
|
|
2458
|
+
// b_i = q_i * scale, per query. Padded by 2*VLMAX so the pipelined
|
|
2459
|
+
// next-iteration vle32 loads stay in bounds (padding lanes unused).
|
|
2460
|
+
std::vector<float> b;
|
|
2461
|
+
|
|
2462
|
+
DCTemplate(size_t d_in, const std::vector<float>& trained)
|
|
2463
|
+
: d(d_in), k_q(0.0f), b(d_in + 2 * __riscv_vsetvlmax_e8m1(), 0.0f) {
|
|
2464
|
+
scale = trained[1] / 255.0f;
|
|
2465
|
+
c0 = trained[0] + 0.5f * scale;
|
|
2466
|
+
}
|
|
2467
|
+
|
|
2468
|
+
void set_query(const float* x) final {
|
|
2469
|
+
q = x;
|
|
2470
|
+
// b_i = q_i * scale and K_q = sum_i (q_i * c0), both accumulated
|
|
2471
|
+
// per dim (amortized over codes). Per-dim pre-scaling keeps every
|
|
2472
|
+
// term inside the scalar per-dim semantics: scale == 0 /
|
|
2473
|
+
// c0 == 0 zero the term exactly, so no bare c0*sum(q) or
|
|
2474
|
+
// sum(q_i*c_i) intermediate can overflow a finite query into
|
|
2475
|
+
// 0*inf == NaN.
|
|
2476
|
+
float kq = 0.0f;
|
|
2477
|
+
for (size_t i = 0; i < d; i++) {
|
|
2478
|
+
kq += x[i] * c0;
|
|
2479
|
+
b[i] = x[i] * scale;
|
|
2480
|
+
}
|
|
2481
|
+
k_q = kq;
|
|
2482
|
+
}
|
|
2483
|
+
|
|
2484
|
+
/// S = sum_i b_i * c_i over the 1-byte-per-dim code; returns
|
|
2485
|
+
/// K_q + S.
|
|
2486
|
+
float compute_ip(const uint8_t* code) const {
|
|
2487
|
+
const float* pb = b.data();
|
|
2488
|
+
|
|
2489
|
+
// Hoist vsetvl: VLMAX for e8m1 (== f32 lanes per m4 group).
|
|
2490
|
+
// d==0: guard so vsetvl(0) doesn't return 0 and the prologue
|
|
2491
|
+
// `if (i + vl <= d)` + chunk loop don't stall.
|
|
2492
|
+
const size_t vl = __riscv_vsetvl_e8m1(d > 0 ? d : 1);
|
|
2493
|
+
vfloat32m4_t acc = __riscv_vfmv_v_f_f32m4(0.0f, vl);
|
|
2494
|
+
|
|
2495
|
+
size_t i = 0;
|
|
2496
|
+
if (i + vl <= d) {
|
|
2497
|
+
// Software pipeline depth 1 — preload the NEXT chunk's
|
|
2498
|
+
// code bytes at the top of THIS iteration.
|
|
2499
|
+
vuint8m1_t c8 = __riscv_vle8_v_u8m1(code + i, vl); // prologue
|
|
2500
|
+
i += vl;
|
|
2501
|
+
|
|
2502
|
+
for (; i + vl <= d; i += vl) {
|
|
2503
|
+
vuint8m1_t c8_next = __riscv_vle8_v_u8m1(code + i, vl);
|
|
2504
|
+
vfloat32m4_t cf = __riscv_vfcvt_f_xu_v_f32m4(
|
|
2505
|
+
__riscv_vzext_vf4_u32m4(c8, vl), vl);
|
|
2506
|
+
vfloat32m4_t vb = __riscv_vle32_v_f32m4(pb + i - vl, vl);
|
|
2507
|
+
acc = __riscv_vfmacc_vv_f32m4(acc, vb, cf, vl);
|
|
2508
|
+
c8 = c8_next; // rotate
|
|
2509
|
+
}
|
|
2510
|
+
|
|
2511
|
+
// Epilogue: process the last full chunk (already loaded).
|
|
2512
|
+
{
|
|
2513
|
+
vfloat32m4_t cf = __riscv_vfcvt_f_xu_v_f32m4(
|
|
2514
|
+
__riscv_vzext_vf4_u32m4(c8, vl), vl);
|
|
2515
|
+
vfloat32m4_t vb = __riscv_vle32_v_f32m4(pb + i - vl, vl);
|
|
2516
|
+
acc = __riscv_vfmacc_vv_f32m4(acc, vb, cf, vl);
|
|
2517
|
+
}
|
|
2518
|
+
}
|
|
2519
|
+
|
|
2520
|
+
// Tail: fewer than vl dims left — one shorter-vl pass into the
|
|
2521
|
+
// same accumulator; _tu keeps lanes past vt (accumulated by
|
|
2522
|
+
// prior chunks) intact for the full-vl reduction below.
|
|
2523
|
+
if (i < d) {
|
|
2524
|
+
const size_t vt = __riscv_vsetvl_e8m1(d - i);
|
|
2525
|
+
vuint8m1_t c8 = __riscv_vle8_v_u8m1(code + i, vt);
|
|
2526
|
+
vfloat32m4_t cf = __riscv_vfcvt_f_xu_v_f32m4(
|
|
2527
|
+
__riscv_vzext_vf4_u32m4(c8, vt), vt);
|
|
2528
|
+
vfloat32m4_t vb = __riscv_vle32_v_f32m4(pb + i, vt);
|
|
2529
|
+
acc = __riscv_vfmacc_vv_f32m4_tu(acc, vb, cf, vt);
|
|
2530
|
+
}
|
|
2531
|
+
|
|
2532
|
+
// Horizontal reduce over all vl lanes.
|
|
2533
|
+
vfloat32m1_t red = __riscv_vfredusum_vs_f32m4_f32m1(
|
|
2534
|
+
acc, __riscv_vfmv_v_f_f32m1(0.0f, 1), vl);
|
|
2535
|
+
return k_q + __riscv_vfmv_f_s_f32m1_f32(red);
|
|
2536
|
+
}
|
|
2537
|
+
|
|
2538
|
+
float query_to_code(const uint8_t* code) const final {
|
|
2539
|
+
return compute_ip(code);
|
|
2540
|
+
}
|
|
2541
|
+
|
|
2542
|
+
float symmetric_dis(idx_t i, idx_t j) override {
|
|
2543
|
+
// Not on the benchmark-critical path; scalar per-dim
|
|
2544
|
+
// evaluation: IP(recon(c1), recon(c2)) = sum_k r1_k * r2_k
|
|
2545
|
+
const uint8_t* code1 = codes + i * code_size;
|
|
2546
|
+
const uint8_t* code2 = codes + j * code_size;
|
|
2547
|
+
float acc = 0;
|
|
2548
|
+
for (size_t k = 0; k < d; k++) {
|
|
2549
|
+
float r1 = c0 + scale * float(code1[k]);
|
|
2550
|
+
float r2 = c0 + scale * float(code2[k]);
|
|
2551
|
+
acc += r1 * r2;
|
|
2552
|
+
}
|
|
2553
|
+
return acc;
|
|
2554
|
+
}
|
|
2555
|
+
|
|
2556
|
+
void query_to_codes_batch_4(
|
|
2557
|
+
const uint8_t* code_0,
|
|
2558
|
+
const uint8_t* code_1,
|
|
2559
|
+
const uint8_t* code_2,
|
|
2560
|
+
const uint8_t* code_3,
|
|
2561
|
+
float& dis0,
|
|
2562
|
+
float& dis1,
|
|
2563
|
+
float& dis2,
|
|
2564
|
+
float& dis3) const final {
|
|
2565
|
+
dis0 = query_to_code(code_0);
|
|
2566
|
+
dis1 = query_to_code(code_1);
|
|
2567
|
+
dis2 = query_to_code(code_2);
|
|
2568
|
+
dis3 = query_to_code(code_3);
|
|
2569
|
+
}
|
|
2570
|
+
};
|
|
2571
|
+
|
|
2572
|
+
/**********************************************************
|
|
2573
|
+
* QT_8bit_direct + L2 — full RVV specialization
|
|
2574
|
+
*
|
|
2575
|
+
* Direct 8-bit storage: recon_i(c) = c_i (no affine, no trained
|
|
2576
|
+
* params), so L2(q, code) = sum_i (q_i - c_i)^2.
|
|
2577
|
+
*
|
|
2578
|
+
* Float-domain kernel: the ORIGINAL float query is retained —
|
|
2579
|
+
* no integer-grid pre-quantization — so scores and rankings
|
|
2580
|
+
* match the scalar NONE reference to float-reassociation
|
|
2581
|
+
* precision (the semantics class of every other RVV SQ kernel;
|
|
2582
|
+
* an earlier integer-domain variant truncated the query to
|
|
2583
|
+
* bytes, which changed results for fractional queries).
|
|
2584
|
+
*
|
|
2585
|
+
* Hot loop per vl = VLMAX(e8m1) dims (16 at VLEN=128): vle8 ->
|
|
2586
|
+
* vzext_vf4 (u8 -> u32m4) -> vfcvt (u32 -> f32m4), then
|
|
2587
|
+
* t = q - c (vfsub) and acc += t*t (vfmacc). 5 vector ops per
|
|
2588
|
+
* chunk, single f32m4 accumulator, single vfredusum at the end.
|
|
2589
|
+
* vsetvl hoisted: 1 hoisted + 1 for the tail. f32 accumulators
|
|
2590
|
+
* saturate near 3.4e38, so no integer-flush strip-mining is
|
|
2591
|
+
* needed at any d.
|
|
2592
|
+
**********************************************************/
|
|
2593
|
+
|
|
2594
|
+
template <>
|
|
2595
|
+
struct DCTemplate<
|
|
2596
|
+
Quantizer8bitDirect<SIMDLevel::RISCV_RVV>,
|
|
2597
|
+
SimilarityL2<SIMDLevel::RISCV_RVV>,
|
|
2598
|
+
SIMDLevel::RISCV_RVV> : SQDistanceComputer {
|
|
2599
|
+
using Sim = SimilarityL2<SIMDLevel::RISCV_RVV>;
|
|
2600
|
+
|
|
2601
|
+
size_t d;
|
|
2602
|
+
|
|
2603
|
+
DCTemplate(size_t d_in, const std::vector<float>& /* unused */) : d(d_in) {}
|
|
2604
|
+
|
|
2605
|
+
void set_query(const float* x) final {
|
|
2606
|
+
q = x;
|
|
2607
|
+
}
|
|
2608
|
+
|
|
2609
|
+
/// Float-domain L2: sum_i (q_i - c_i)^2 in an f32m4 vector
|
|
2610
|
+
/// accumulator; single vfredusum at the end.
|
|
2611
|
+
float compute_l2(const uint8_t* code) const {
|
|
2612
|
+
const float* pq = q;
|
|
2613
|
+
|
|
2614
|
+
// Hoist vsetvl: VLMAX for e8m1 (== f32m4 lanes).
|
|
2615
|
+
// d==0 guard: vsetvl(0) would return 0 and stall the loop.
|
|
2616
|
+
const size_t vl = __riscv_vsetvl_e8m1(d > 0 ? d : 1);
|
|
2617
|
+
vfloat32m4_t acc = __riscv_vfmv_v_f_f32m4(0.0f, vl);
|
|
2618
|
+
|
|
2619
|
+
size_t i = 0;
|
|
2620
|
+
for (; i + vl <= d; i += vl) {
|
|
2621
|
+
vuint8m1_t c8 = __riscv_vle8_v_u8m1(code + i, vl);
|
|
2622
|
+
vfloat32m4_t cf = __riscv_vfcvt_f_xu_v_f32m4(
|
|
2623
|
+
__riscv_vzext_vf4_u32m4(c8, vl), vl);
|
|
2624
|
+
vfloat32m4_t vq = __riscv_vle32_v_f32m4(pq + i, vl);
|
|
2625
|
+
vfloat32m4_t t = __riscv_vfsub_vv_f32m4(vq, cf, vl);
|
|
2626
|
+
acc = __riscv_vfmacc_vv_f32m4(acc, t, t, vl);
|
|
2627
|
+
}
|
|
2628
|
+
|
|
2629
|
+
// Tail: fewer than vl dims left — one shorter-vl pass into the
|
|
2630
|
+
// same accumulator; _tu keeps lanes past vt (accumulated by
|
|
2631
|
+
// prior chunks) intact for the full-vl reduction below.
|
|
2632
|
+
if (i < d) {
|
|
2633
|
+
const size_t vt = __riscv_vsetvl_e8m1(d - i);
|
|
2634
|
+
vuint8m1_t c8 = __riscv_vle8_v_u8m1(code + i, vt);
|
|
2635
|
+
vfloat32m4_t cf = __riscv_vfcvt_f_xu_v_f32m4(
|
|
2636
|
+
__riscv_vzext_vf4_u32m4(c8, vt), vt);
|
|
2637
|
+
vfloat32m4_t vq = __riscv_vle32_v_f32m4(pq + i, vt);
|
|
2638
|
+
vfloat32m4_t t = __riscv_vfsub_vv_f32m4(vq, cf, vt);
|
|
2639
|
+
acc = __riscv_vfmacc_vv_f32m4_tu(acc, t, t, vt);
|
|
2640
|
+
}
|
|
2641
|
+
|
|
2642
|
+
// Horizontal reduce over all vl lanes.
|
|
2643
|
+
vfloat32m1_t red = __riscv_vfredusum_vs_f32m4_f32m1(
|
|
2644
|
+
acc, __riscv_vfmv_v_f_f32m1(0.0f, 1), vl);
|
|
2645
|
+
return __riscv_vfmv_f_s_f32m1_f32(red);
|
|
2646
|
+
}
|
|
2647
|
+
|
|
2648
|
+
float query_to_code(const uint8_t* code) const final {
|
|
2649
|
+
return compute_l2(code);
|
|
2650
|
+
}
|
|
2651
|
+
|
|
2652
|
+
float symmetric_dis(idx_t i, idx_t j) override {
|
|
2653
|
+
// Not on the benchmark-critical path; exact integer-domain
|
|
2654
|
+
// scalar evaluation (codes are raw bytes).
|
|
2655
|
+
const uint8_t* c1 = codes + i * code_size;
|
|
2656
|
+
const uint8_t* c2 = codes + j * code_size;
|
|
2657
|
+
int64_t acc = 0;
|
|
2658
|
+
for (size_t k = 0; k < d; k++) {
|
|
2659
|
+
int diff = int(c1[k]) - int(c2[k]);
|
|
2660
|
+
acc += diff * diff;
|
|
2661
|
+
}
|
|
2662
|
+
return static_cast<float>(acc);
|
|
2663
|
+
}
|
|
2664
|
+
|
|
2665
|
+
void query_to_codes_batch_4(
|
|
2666
|
+
const uint8_t* code_0,
|
|
2667
|
+
const uint8_t* code_1,
|
|
2668
|
+
const uint8_t* code_2,
|
|
2669
|
+
const uint8_t* code_3,
|
|
2670
|
+
float& dis0,
|
|
2671
|
+
float& dis1,
|
|
2672
|
+
float& dis2,
|
|
2673
|
+
float& dis3) const final {
|
|
2674
|
+
dis0 = compute_l2(code_0);
|
|
2675
|
+
dis1 = compute_l2(code_1);
|
|
2676
|
+
dis2 = compute_l2(code_2);
|
|
2677
|
+
dis3 = compute_l2(code_3);
|
|
2678
|
+
}
|
|
2679
|
+
};
|
|
2680
|
+
|
|
2681
|
+
/**********************************************************
|
|
2682
|
+
* QT_8bit_direct + IP — full RVV specialization
|
|
2683
|
+
*
|
|
2684
|
+
* Direct 8-bit storage: recon_i(c) = c_i (no affine, no trained
|
|
2685
|
+
* params), so IP(q, code) = sum_i q_i * c_i.
|
|
2686
|
+
*
|
|
2687
|
+
* Float-domain kernel: the ORIGINAL float query is retained —
|
|
2688
|
+
* no integer-grid pre-quantization — so scores and rankings
|
|
2689
|
+
* match the scalar NONE reference to float-reassociation
|
|
2690
|
+
* precision. Hot loop per vl = VLMAX(e8m1) dims: vle8 ->
|
|
2691
|
+
* vzext_vf4 -> vfcvt (u8 code -> f32), then acc += q * c
|
|
2692
|
+
* (vfmacc) with the query loaded as f32. 4 vector ops per
|
|
2693
|
+
* chunk, single f32m4 accumulator, single vfredusum at the
|
|
2694
|
+
* end. vsetvl hoisted: 1 hoisted + 1 for the tail.
|
|
2695
|
+
**********************************************************/
|
|
2696
|
+
|
|
2697
|
+
template <>
|
|
2698
|
+
struct DCTemplate<
|
|
2699
|
+
Quantizer8bitDirect<SIMDLevel::RISCV_RVV>,
|
|
2700
|
+
SimilarityIP<SIMDLevel::RISCV_RVV>,
|
|
2701
|
+
SIMDLevel::RISCV_RVV> : SQDistanceComputer {
|
|
2702
|
+
using Sim = SimilarityIP<SIMDLevel::RISCV_RVV>;
|
|
2703
|
+
|
|
2704
|
+
size_t d;
|
|
2705
|
+
|
|
2706
|
+
DCTemplate(size_t d_in, const std::vector<float>& /* unused */) : d(d_in) {}
|
|
2707
|
+
|
|
2708
|
+
void set_query(const float* x) final {
|
|
2709
|
+
q = x;
|
|
2710
|
+
}
|
|
2711
|
+
|
|
2712
|
+
/// Float-domain IP: sum_i q_i * c_i in an f32m4 vector
|
|
2713
|
+
/// accumulator; single vfredusum at the end.
|
|
2714
|
+
float compute_ip(const uint8_t* code) const {
|
|
2715
|
+
const float* pq = q;
|
|
2716
|
+
|
|
2717
|
+
// Hoist vsetvl: VLMAX for e8m1 (== f32m4 lanes).
|
|
2718
|
+
// d==0 guard: vsetvl(0) would return 0 and stall the loop.
|
|
2719
|
+
const size_t vl = __riscv_vsetvl_e8m1(d > 0 ? d : 1);
|
|
2720
|
+
vfloat32m4_t acc = __riscv_vfmv_v_f_f32m4(0.0f, vl);
|
|
2721
|
+
|
|
2722
|
+
size_t i = 0;
|
|
2723
|
+
for (; i + vl <= d; i += vl) {
|
|
2724
|
+
vuint8m1_t c8 = __riscv_vle8_v_u8m1(code + i, vl);
|
|
2725
|
+
vfloat32m4_t cf = __riscv_vfcvt_f_xu_v_f32m4(
|
|
2726
|
+
__riscv_vzext_vf4_u32m4(c8, vl), vl);
|
|
2727
|
+
vfloat32m4_t vq = __riscv_vle32_v_f32m4(pq + i, vl);
|
|
2728
|
+
acc = __riscv_vfmacc_vv_f32m4(acc, vq, cf, vl);
|
|
2729
|
+
}
|
|
2730
|
+
|
|
2731
|
+
// Tail: fewer than vl dims left — one shorter-vl pass into the
|
|
2732
|
+
// same accumulator; _tu keeps lanes past vt (accumulated by
|
|
2733
|
+
// prior chunks) intact for the full-vl reduction below.
|
|
2734
|
+
if (i < d) {
|
|
2735
|
+
const size_t vt = __riscv_vsetvl_e8m1(d - i);
|
|
2736
|
+
vuint8m1_t c8 = __riscv_vle8_v_u8m1(code + i, vt);
|
|
2737
|
+
vfloat32m4_t cf = __riscv_vfcvt_f_xu_v_f32m4(
|
|
2738
|
+
__riscv_vzext_vf4_u32m4(c8, vt), vt);
|
|
2739
|
+
vfloat32m4_t vq = __riscv_vle32_v_f32m4(pq + i, vt);
|
|
2740
|
+
acc = __riscv_vfmacc_vv_f32m4_tu(acc, vq, cf, vt);
|
|
2741
|
+
}
|
|
2742
|
+
|
|
2743
|
+
// Horizontal reduce over all vl lanes.
|
|
2744
|
+
vfloat32m1_t red = __riscv_vfredusum_vs_f32m4_f32m1(
|
|
2745
|
+
acc, __riscv_vfmv_v_f_f32m1(0.0f, 1), vl);
|
|
2746
|
+
return __riscv_vfmv_f_s_f32m1_f32(red);
|
|
2747
|
+
}
|
|
2748
|
+
|
|
2749
|
+
float query_to_code(const uint8_t* code) const final {
|
|
2750
|
+
return compute_ip(code);
|
|
2751
|
+
}
|
|
2752
|
+
|
|
2753
|
+
float symmetric_dis(idx_t i, idx_t j) override {
|
|
2754
|
+
// Not on the benchmark-critical path; exact integer-domain
|
|
2755
|
+
// scalar evaluation (codes are raw bytes).
|
|
2756
|
+
const uint8_t* c1 = codes + i * code_size;
|
|
2757
|
+
const uint8_t* c2 = codes + j * code_size;
|
|
2758
|
+
int64_t acc = 0;
|
|
2759
|
+
for (size_t k = 0; k < d; k++) {
|
|
2760
|
+
acc += int(c1[k]) * int(c2[k]);
|
|
2761
|
+
}
|
|
2762
|
+
return static_cast<float>(acc);
|
|
2763
|
+
}
|
|
2764
|
+
|
|
2765
|
+
void query_to_codes_batch_4(
|
|
2766
|
+
const uint8_t* code_0,
|
|
2767
|
+
const uint8_t* code_1,
|
|
2768
|
+
const uint8_t* code_2,
|
|
2769
|
+
const uint8_t* code_3,
|
|
2770
|
+
float& dis0,
|
|
2771
|
+
float& dis1,
|
|
2772
|
+
float& dis2,
|
|
2773
|
+
float& dis3) const final {
|
|
2774
|
+
dis0 = compute_ip(code_0);
|
|
2775
|
+
dis1 = compute_ip(code_1);
|
|
2776
|
+
dis2 = compute_ip(code_2);
|
|
2777
|
+
dis3 = compute_ip(code_3);
|
|
2778
|
+
}
|
|
2779
|
+
};
|
|
2780
|
+
|
|
2781
|
+
/**********************************************************
|
|
2782
|
+
* QT_8bit_direct_signed + L2 — full RVV specialization
|
|
2783
|
+
*
|
|
2784
|
+
* Signed direct storage (Quantizer8bitDirectSigned): recon_i(c) =
|
|
2785
|
+
* c_i - 128, so L2(q, code) = sum_i (q_i - (c_i - 128))^2.
|
|
2786
|
+
*
|
|
2787
|
+
* Float-domain kernel: the ORIGINAL float query is retained —
|
|
2788
|
+
* no integer-grid pre-quantization — so scores and rankings
|
|
2789
|
+
* match the scalar NONE reference to float-reassociation
|
|
2790
|
+
* precision. Hot loop per vl = VLMAX(e8m1) dims: vle8 ->
|
|
2791
|
+
* vzext_vf4 -> vfcvt (u8 code -> f32) -> vfsub 128 (decode to
|
|
2792
|
+
* the signed value domain), then t = q - c' and acc += t*t.
|
|
2793
|
+
* vsetvl hoisted: 1 hoisted + 1 for the tail. f32 accumulators
|
|
2794
|
+
* saturate near 3.4e38, so no integer-flush strip-mining is
|
|
2795
|
+
* needed at any d.
|
|
2796
|
+
**********************************************************/
|
|
2797
|
+
|
|
2798
|
+
template <>
|
|
2799
|
+
struct DCTemplate<
|
|
2800
|
+
Quantizer8bitDirectSigned<SIMDLevel::RISCV_RVV>,
|
|
2801
|
+
SimilarityL2<SIMDLevel::RISCV_RVV>,
|
|
2802
|
+
SIMDLevel::RISCV_RVV> : SQDistanceComputer {
|
|
2803
|
+
using Sim = SimilarityL2<SIMDLevel::RISCV_RVV>;
|
|
2804
|
+
|
|
2805
|
+
size_t d;
|
|
2806
|
+
|
|
2807
|
+
DCTemplate(size_t d_in, const std::vector<float>& /* unused */) : d(d_in) {}
|
|
2808
|
+
|
|
2809
|
+
void set_query(const float* x) final {
|
|
2810
|
+
q = x;
|
|
2811
|
+
}
|
|
2812
|
+
|
|
2813
|
+
/// Float-domain L2: sum_i (q_i - (c_i - 128))^2 in an f32m4
|
|
2814
|
+
/// vector accumulator; single vfredusum at the end.
|
|
2815
|
+
float compute_l2(const uint8_t* code) const {
|
|
2816
|
+
const float* pq = q;
|
|
2817
|
+
|
|
2818
|
+
// Hoist vsetvl: VLMAX for e8m1 (== f32m4 lanes).
|
|
2819
|
+
// d==0 guard: vsetvl(0) would return 0 and stall the loop.
|
|
2820
|
+
const size_t vl = __riscv_vsetvl_e8m1(d > 0 ? d : 1);
|
|
2821
|
+
vfloat32m4_t acc = __riscv_vfmv_v_f_f32m4(0.0f, vl);
|
|
2822
|
+
|
|
2823
|
+
size_t i = 0;
|
|
2824
|
+
for (; i + vl <= d; i += vl) {
|
|
2825
|
+
vuint8m1_t c8 = __riscv_vle8_v_u8m1(code + i, vl);
|
|
2826
|
+
vfloat32m4_t cf = __riscv_vfcvt_f_xu_v_f32m4(
|
|
2827
|
+
__riscv_vzext_vf4_u32m4(c8, vl), vl);
|
|
2828
|
+
// Decode the stored byte into the signed value domain:
|
|
2829
|
+
// c' = c - 128.
|
|
2830
|
+
cf = __riscv_vfsub_vf_f32m4(cf, 128.0f, vl);
|
|
2831
|
+
vfloat32m4_t vq = __riscv_vle32_v_f32m4(pq + i, vl);
|
|
2832
|
+
vfloat32m4_t t = __riscv_vfsub_vv_f32m4(vq, cf, vl);
|
|
2833
|
+
acc = __riscv_vfmacc_vv_f32m4(acc, t, t, vl);
|
|
2834
|
+
}
|
|
2835
|
+
|
|
2836
|
+
// Tail: fewer than vl dims left — one shorter-vl pass into the
|
|
2837
|
+
// same accumulator; _tu keeps lanes past vt (accumulated by
|
|
2838
|
+
// prior chunks) intact for the full-vl reduction below.
|
|
2839
|
+
if (i < d) {
|
|
2840
|
+
const size_t vt = __riscv_vsetvl_e8m1(d - i);
|
|
2841
|
+
vuint8m1_t c8 = __riscv_vle8_v_u8m1(code + i, vt);
|
|
2842
|
+
vfloat32m4_t cf = __riscv_vfcvt_f_xu_v_f32m4(
|
|
2843
|
+
__riscv_vzext_vf4_u32m4(c8, vt), vt);
|
|
2844
|
+
cf = __riscv_vfsub_vf_f32m4(cf, 128.0f, vt);
|
|
2845
|
+
vfloat32m4_t vq = __riscv_vle32_v_f32m4(pq + i, vt);
|
|
2846
|
+
vfloat32m4_t t = __riscv_vfsub_vv_f32m4(vq, cf, vt);
|
|
2847
|
+
acc = __riscv_vfmacc_vv_f32m4_tu(acc, t, t, vt);
|
|
2848
|
+
}
|
|
2849
|
+
|
|
2850
|
+
// Horizontal reduce over all vl lanes.
|
|
2851
|
+
vfloat32m1_t red = __riscv_vfredusum_vs_f32m4_f32m1(
|
|
2852
|
+
acc, __riscv_vfmv_v_f_f32m1(0.0f, 1), vl);
|
|
2853
|
+
return __riscv_vfmv_f_s_f32m1_f32(red);
|
|
2854
|
+
}
|
|
2855
|
+
|
|
2856
|
+
float query_to_code(const uint8_t* code) const final {
|
|
2857
|
+
return compute_l2(code);
|
|
2858
|
+
}
|
|
2859
|
+
|
|
2860
|
+
float symmetric_dis(idx_t i, idx_t j) override {
|
|
2861
|
+
// Not on the benchmark-critical path; exact integer-domain
|
|
2862
|
+
// scalar evaluation ((c1-128)-(c2-128) == c1-c2).
|
|
2863
|
+
const uint8_t* c1 = codes + i * code_size;
|
|
2864
|
+
const uint8_t* c2 = codes + j * code_size;
|
|
2865
|
+
int64_t acc = 0;
|
|
2866
|
+
for (size_t k = 0; k < d; k++) {
|
|
2867
|
+
int diff = int(c1[k]) - int(c2[k]);
|
|
2868
|
+
acc += diff * diff;
|
|
2869
|
+
}
|
|
2870
|
+
return static_cast<float>(acc);
|
|
2871
|
+
}
|
|
2872
|
+
|
|
2873
|
+
void query_to_codes_batch_4(
|
|
2874
|
+
const uint8_t* code_0,
|
|
2875
|
+
const uint8_t* code_1,
|
|
2876
|
+
const uint8_t* code_2,
|
|
2877
|
+
const uint8_t* code_3,
|
|
2878
|
+
float& dis0,
|
|
2879
|
+
float& dis1,
|
|
2880
|
+
float& dis2,
|
|
2881
|
+
float& dis3) const final {
|
|
2882
|
+
dis0 = compute_l2(code_0);
|
|
2883
|
+
dis1 = compute_l2(code_1);
|
|
2884
|
+
dis2 = compute_l2(code_2);
|
|
2885
|
+
dis3 = compute_l2(code_3);
|
|
2886
|
+
}
|
|
2887
|
+
};
|
|
2888
|
+
|
|
2889
|
+
/**********************************************************
|
|
2890
|
+
* QT_8bit_direct_signed + IP — full RVV specialization
|
|
2891
|
+
*
|
|
2892
|
+
* Signed direct storage: recon_i(c) = c_i - 128, so IP(q, code) =
|
|
2893
|
+
* sum_i q_i*(c_i - 128).
|
|
2894
|
+
*
|
|
2895
|
+
* Float-domain kernel: the ORIGINAL float query is retained —
|
|
2896
|
+
* no integer-grid pre-quantization — so scores and rankings
|
|
2897
|
+
* match the scalar NONE reference to float-reassociation
|
|
2898
|
+
* precision. Hot loop per vl = VLMAX(e8m1) dims: vle8 ->
|
|
2899
|
+
* vzext_vf4 -> vfcvt (u8 code -> f32) -> vfsub 128 (decode to
|
|
2900
|
+
* the signed value domain), then acc += q * c' (vfmacc) with
|
|
2901
|
+
* the query loaded as f32. vsetvl hoisted: 1 hoisted + 1 for
|
|
2902
|
+
* the tail.
|
|
2903
|
+
**********************************************************/
|
|
2904
|
+
|
|
2905
|
+
template <>
|
|
2906
|
+
struct DCTemplate<
|
|
2907
|
+
Quantizer8bitDirectSigned<SIMDLevel::RISCV_RVV>,
|
|
2908
|
+
SimilarityIP<SIMDLevel::RISCV_RVV>,
|
|
2909
|
+
SIMDLevel::RISCV_RVV> : SQDistanceComputer {
|
|
2910
|
+
using Sim = SimilarityIP<SIMDLevel::RISCV_RVV>;
|
|
2911
|
+
|
|
2912
|
+
size_t d;
|
|
2913
|
+
|
|
2914
|
+
DCTemplate(size_t d_in, const std::vector<float>& /* unused */) : d(d_in) {}
|
|
2915
|
+
|
|
2916
|
+
void set_query(const float* x) final {
|
|
2917
|
+
q = x;
|
|
2918
|
+
}
|
|
2919
|
+
|
|
2920
|
+
/// Float-domain signed IP: sum_i q_i * (c_i - 128) in an f32m4
|
|
2921
|
+
/// vector accumulator; single vfredusum at the end.
|
|
2922
|
+
float compute_ip(const uint8_t* code) const {
|
|
2923
|
+
const float* pq = q;
|
|
2924
|
+
|
|
2925
|
+
// Hoist vsetvl: VLMAX for e8m1 (== f32m4 lanes).
|
|
2926
|
+
// d==0 guard: vsetvl(0) would return 0 and stall the loop.
|
|
2927
|
+
const size_t vl = __riscv_vsetvl_e8m1(d > 0 ? d : 1);
|
|
2928
|
+
vfloat32m4_t acc = __riscv_vfmv_v_f_f32m4(0.0f, vl);
|
|
2929
|
+
|
|
2930
|
+
size_t i = 0;
|
|
2931
|
+
for (; i + vl <= d; i += vl) {
|
|
2932
|
+
vuint8m1_t c8 = __riscv_vle8_v_u8m1(code + i, vl);
|
|
2933
|
+
vfloat32m4_t cf = __riscv_vfcvt_f_xu_v_f32m4(
|
|
2934
|
+
__riscv_vzext_vf4_u32m4(c8, vl), vl);
|
|
2935
|
+
// Decode the stored byte into the signed value domain:
|
|
2936
|
+
// c' = c - 128.
|
|
2937
|
+
cf = __riscv_vfsub_vf_f32m4(cf, 128.0f, vl);
|
|
2938
|
+
vfloat32m4_t vq = __riscv_vle32_v_f32m4(pq + i, vl);
|
|
2939
|
+
acc = __riscv_vfmacc_vv_f32m4(acc, vq, cf, vl);
|
|
2940
|
+
}
|
|
2941
|
+
|
|
2942
|
+
// Tail: fewer than vl dims left — one shorter-vl pass into the
|
|
2943
|
+
// same accumulator; _tu keeps lanes past vt (accumulated by
|
|
2944
|
+
// prior chunks) intact for the full-vl reduction below.
|
|
2945
|
+
if (i < d) {
|
|
2946
|
+
const size_t vt = __riscv_vsetvl_e8m1(d - i);
|
|
2947
|
+
vuint8m1_t c8 = __riscv_vle8_v_u8m1(code + i, vt);
|
|
2948
|
+
vfloat32m4_t cf = __riscv_vfcvt_f_xu_v_f32m4(
|
|
2949
|
+
__riscv_vzext_vf4_u32m4(c8, vt), vt);
|
|
2950
|
+
cf = __riscv_vfsub_vf_f32m4(cf, 128.0f, vt);
|
|
2951
|
+
vfloat32m4_t vq = __riscv_vle32_v_f32m4(pq + i, vt);
|
|
2952
|
+
acc = __riscv_vfmacc_vv_f32m4_tu(acc, vq, cf, vt);
|
|
2953
|
+
}
|
|
2954
|
+
|
|
2955
|
+
// Horizontal reduce over all vl lanes.
|
|
2956
|
+
vfloat32m1_t red = __riscv_vfredusum_vs_f32m4_f32m1(
|
|
2957
|
+
acc, __riscv_vfmv_v_f_f32m1(0.0f, 1), vl);
|
|
2958
|
+
return __riscv_vfmv_f_s_f32m1_f32(red);
|
|
2959
|
+
}
|
|
2960
|
+
|
|
2961
|
+
float query_to_code(const uint8_t* code) const final {
|
|
2962
|
+
return compute_ip(code);
|
|
2963
|
+
}
|
|
2964
|
+
|
|
2965
|
+
float symmetric_dis(idx_t i, idx_t j) override {
|
|
2966
|
+
// Not on the benchmark-critical path; exact integer-domain
|
|
2967
|
+
// scalar evaluation ((c1-128)*(c2-128)).
|
|
2968
|
+
const uint8_t* c1 = codes + i * code_size;
|
|
2969
|
+
const uint8_t* c2 = codes + j * code_size;
|
|
2970
|
+
int64_t acc = 0;
|
|
2971
|
+
for (size_t k = 0; k < d; k++) {
|
|
2972
|
+
acc += (int(c1[k]) - 128) * (int(c2[k]) - 128);
|
|
2973
|
+
}
|
|
2974
|
+
return static_cast<float>(acc);
|
|
2975
|
+
}
|
|
2976
|
+
|
|
2977
|
+
void query_to_codes_batch_4(
|
|
2978
|
+
const uint8_t* code_0,
|
|
2979
|
+
const uint8_t* code_1,
|
|
2980
|
+
const uint8_t* code_2,
|
|
2981
|
+
const uint8_t* code_3,
|
|
2982
|
+
float& dis0,
|
|
2983
|
+
float& dis1,
|
|
2984
|
+
float& dis2,
|
|
2985
|
+
float& dis3) const final {
|
|
2986
|
+
dis0 = compute_ip(code_0);
|
|
2987
|
+
dis1 = compute_ip(code_1);
|
|
2988
|
+
dis2 = compute_ip(code_2);
|
|
2989
|
+
dis3 = compute_ip(code_3);
|
|
2990
|
+
}
|
|
2991
|
+
};
|
|
2992
|
+
|
|
2993
|
+
// * Fast path — QT_bf16 + L2
|
|
2994
|
+
// *
|
|
2995
|
+
// * bf16 code: each dim is the high 16 bits of an f32, stored as uint16;
|
|
2996
|
+
// * decode_bf16(v) = reinterpret_f32(u32(v) << 16). The query stays in full
|
|
2997
|
+
// * f32 precision: L2 = sum_i (q_i - decode_bf16(c_i))^2.
|
|
2998
|
+
// *
|
|
2999
|
+
// * march has no zvfbfmin/zvfbfwma, so bf16→f32 widening is manual: vle16 →
|
|
3000
|
+
// * vzext_vf2 → vsll 16 → reinterpret. Hot loop: code load + 2-op widen + q
|
|
3001
|
+
// * load + vfsub + vfmacc.
|
|
3002
|
+
// ************************************************************************/
|
|
3003
|
+
|
|
3004
|
+
/// Widen a bf16 (u16) vector chunk to f32: (u32(v) << 16) reinterpreted.
|
|
3005
|
+
static inline vfloat32m4_t bf16_widen_f32m4(vuint16m2_t v, size_t vl) {
|
|
3006
|
+
vuint32m4_t w = __riscv_vzext_vf2_u32m4(v, vl);
|
|
3007
|
+
w = __riscv_vsll_vx_u32m4(w, 16, vl);
|
|
3008
|
+
return __riscv_vreinterpret_v_u32m4_f32m4(w);
|
|
3009
|
+
}
|
|
3010
|
+
|
|
3011
|
+
template <>
|
|
3012
|
+
struct DCTemplate<
|
|
3013
|
+
QuantizerBF16<SIMDLevel::RISCV_RVV>,
|
|
3014
|
+
SimilarityL2<SIMDLevel::RISCV_RVV>,
|
|
3015
|
+
SIMDLevel::RISCV_RVV> : SQDistanceComputer {
|
|
3016
|
+
using Sim = SimilarityL2<SIMDLevel::RISCV_RVV>;
|
|
3017
|
+
|
|
3018
|
+
size_t d;
|
|
3019
|
+
|
|
3020
|
+
DCTemplate(size_t d_in, const std::vector<float>&) : d(d_in) {}
|
|
3021
|
+
|
|
3022
|
+
void set_query(const float* x) final {
|
|
3023
|
+
q = x;
|
|
3024
|
+
}
|
|
3025
|
+
|
|
3026
|
+
/// Direct-form L2 between the f32 query and a bf16 code.
|
|
3027
|
+
float compute_l2(const float* qf, const uint8_t* code8) const {
|
|
3028
|
+
const uint16_t* code = (const uint16_t*)code8;
|
|
3029
|
+
size_t i = 0;
|
|
3030
|
+
// Hoist vsetvl: VLMAX for e16m2 (== f32m4 lanes), reused across
|
|
3031
|
+
// the hot loop; 0 vsetvl inside.
|
|
3032
|
+
const size_t vl = __riscv_vsetvl_e16m2(d > 0 ? d : 1);
|
|
3033
|
+
vfloat32m4_t acc = __riscv_vfmv_v_f_f32m4(0.0f, vl);
|
|
3034
|
+
|
|
3035
|
+
for (; i + vl <= d; i += vl) {
|
|
3036
|
+
vuint16m2_t vc = __riscv_vle16_v_u16m2(code + i, vl);
|
|
3037
|
+
vfloat32m4_t fc = bf16_widen_f32m4(vc, vl);
|
|
3038
|
+
vfloat32m4_t fq = __riscv_vle32_v_f32m4(qf + i, vl);
|
|
3039
|
+
vfloat32m4_t t = __riscv_vfsub_vv_f32m4(fq, fc, vl);
|
|
3040
|
+
acc = __riscv_vfmacc_vv_f32m4(acc, t, t, vl);
|
|
3041
|
+
}
|
|
3042
|
+
|
|
3043
|
+
// Tail: fewer than vl dims left — one shorter-vl pass into the
|
|
3044
|
+
// same accumulator (only the first vt lanes are touched).
|
|
3045
|
+
if (i < d) {
|
|
3046
|
+
const size_t vt = __riscv_vsetvl_e16m2(d - i);
|
|
3047
|
+
vuint16m2_t vc = __riscv_vle16_v_u16m2(code + i, vt);
|
|
3048
|
+
vfloat32m4_t fc = bf16_widen_f32m4(vc, vt);
|
|
3049
|
+
vfloat32m4_t fq = __riscv_vle32_v_f32m4(qf + i, vt);
|
|
3050
|
+
vfloat32m4_t t = __riscv_vfsub_vv_f32m4(fq, fc, vt);
|
|
3051
|
+
acc = __riscv_vfmacc_vv_f32m4_tu(acc, t, t, vt);
|
|
3052
|
+
}
|
|
3053
|
+
|
|
3054
|
+
// Single horizontal reduction over all vl lanes.
|
|
3055
|
+
vfloat32m1_t red = __riscv_vfredusum_vs_f32m4_f32m1(
|
|
3056
|
+
acc, __riscv_vfmv_v_f_f32m1(0.0f, 1), vl);
|
|
3057
|
+
return __riscv_vfmv_f_s_f32m1_f32(red);
|
|
3058
|
+
}
|
|
3059
|
+
|
|
3060
|
+
float query_to_code(const uint8_t* code) const final {
|
|
3061
|
+
return compute_l2(q, code);
|
|
3062
|
+
}
|
|
3063
|
+
|
|
3064
|
+
float symmetric_dis(idx_t i, idx_t j) override {
|
|
3065
|
+
// Not on the benchmark-critical path; scalar per-dim evaluation.
|
|
3066
|
+
const uint16_t* a = (const uint16_t*)(codes + i * code_size);
|
|
3067
|
+
const uint16_t* b = (const uint16_t*)(codes + j * code_size);
|
|
3068
|
+
float accu = 0;
|
|
3069
|
+
for (size_t k = 0; k < d; k++) {
|
|
3070
|
+
float diff = decode_bf16(a[k]) - decode_bf16(b[k]);
|
|
3071
|
+
accu += diff * diff;
|
|
3072
|
+
}
|
|
3073
|
+
return accu;
|
|
3074
|
+
}
|
|
3075
|
+
|
|
3076
|
+
void query_to_codes_batch_4(
|
|
3077
|
+
const uint8_t* code_0,
|
|
3078
|
+
const uint8_t* code_1,
|
|
3079
|
+
const uint8_t* code_2,
|
|
3080
|
+
const uint8_t* code_3,
|
|
3081
|
+
float& dis0,
|
|
3082
|
+
float& dis1,
|
|
3083
|
+
float& dis2,
|
|
3084
|
+
float& dis3) const final {
|
|
3085
|
+
dis0 = compute_l2(q, code_0);
|
|
3086
|
+
dis1 = compute_l2(q, code_1);
|
|
3087
|
+
dis2 = compute_l2(q, code_2);
|
|
3088
|
+
dis3 = compute_l2(q, code_3);
|
|
3089
|
+
}
|
|
3090
|
+
};
|
|
3091
|
+
|
|
3092
|
+
// * Fast path — QT_bf16 + IP
|
|
3093
|
+
// *
|
|
3094
|
+
// * Same bf16 layout as L2 above; query stays in full f32 precision:
|
|
3095
|
+
// * IP = sum_i q_i * decode_bf16(c_i)
|
|
3096
|
+
// * Hot loop: code load + 2-op widen + q load + vfmacc.
|
|
3097
|
+
// ************************************************************************/
|
|
3098
|
+
|
|
3099
|
+
template <>
|
|
3100
|
+
struct DCTemplate<
|
|
3101
|
+
QuantizerBF16<SIMDLevel::RISCV_RVV>,
|
|
3102
|
+
SimilarityIP<SIMDLevel::RISCV_RVV>,
|
|
3103
|
+
SIMDLevel::RISCV_RVV> : SQDistanceComputer {
|
|
3104
|
+
using Sim = SimilarityIP<SIMDLevel::RISCV_RVV>;
|
|
3105
|
+
|
|
3106
|
+
size_t d;
|
|
3107
|
+
|
|
3108
|
+
DCTemplate(size_t d_in, const std::vector<float>&) : d(d_in) {}
|
|
3109
|
+
|
|
3110
|
+
void set_query(const float* x) final {
|
|
3111
|
+
q = x;
|
|
3112
|
+
}
|
|
3113
|
+
|
|
3114
|
+
/// Direct-form inner product between the f32 query and a bf16 code.
|
|
3115
|
+
float compute_ip(const float* qf, const uint8_t* code8) const {
|
|
3116
|
+
const uint16_t* code = (const uint16_t*)code8;
|
|
3117
|
+
size_t i = 0;
|
|
3118
|
+
// Hoist vsetvl: VLMAX for e16m2 (== f32m4 lanes), reused across
|
|
3119
|
+
// the hot loop; 0 vsetvl inside.
|
|
3120
|
+
const size_t vl = __riscv_vsetvl_e16m2(d > 0 ? d : 1);
|
|
3121
|
+
vfloat32m4_t acc = __riscv_vfmv_v_f_f32m4(0.0f, vl);
|
|
3122
|
+
|
|
3123
|
+
for (; i + vl <= d; i += vl) {
|
|
3124
|
+
vuint16m2_t vc = __riscv_vle16_v_u16m2(code + i, vl);
|
|
3125
|
+
vfloat32m4_t fc = bf16_widen_f32m4(vc, vl);
|
|
3126
|
+
vfloat32m4_t fq = __riscv_vle32_v_f32m4(qf + i, vl);
|
|
3127
|
+
acc = __riscv_vfmacc_vv_f32m4(acc, fq, fc, vl);
|
|
3128
|
+
}
|
|
3129
|
+
|
|
3130
|
+
// Tail: fewer than vl dims left — one shorter-vl pass into the
|
|
3131
|
+
// same accumulator (only the first vt lanes are touched).
|
|
3132
|
+
if (i < d) {
|
|
3133
|
+
const size_t vt = __riscv_vsetvl_e16m2(d - i);
|
|
3134
|
+
vuint16m2_t vc = __riscv_vle16_v_u16m2(code + i, vt);
|
|
3135
|
+
vfloat32m4_t fc = bf16_widen_f32m4(vc, vt);
|
|
3136
|
+
vfloat32m4_t fq = __riscv_vle32_v_f32m4(qf + i, vt);
|
|
3137
|
+
acc = __riscv_vfmacc_vv_f32m4_tu(acc, fq, fc, vt);
|
|
3138
|
+
}
|
|
3139
|
+
|
|
3140
|
+
// Single horizontal reduction over all vl lanes.
|
|
3141
|
+
vfloat32m1_t red = __riscv_vfredusum_vs_f32m4_f32m1(
|
|
3142
|
+
acc, __riscv_vfmv_v_f_f32m1(0.0f, 1), vl);
|
|
3143
|
+
return __riscv_vfmv_f_s_f32m1_f32(red);
|
|
3144
|
+
}
|
|
3145
|
+
|
|
3146
|
+
float query_to_code(const uint8_t* code) const final {
|
|
3147
|
+
return compute_ip(q, code);
|
|
3148
|
+
}
|
|
3149
|
+
|
|
3150
|
+
float symmetric_dis(idx_t i, idx_t j) override {
|
|
3151
|
+
// Not on the benchmark-critical path; scalar per-dim evaluation.
|
|
3152
|
+
const uint16_t* a = (const uint16_t*)(codes + i * code_size);
|
|
3153
|
+
const uint16_t* b = (const uint16_t*)(codes + j * code_size);
|
|
3154
|
+
float accu = 0;
|
|
3155
|
+
for (size_t k = 0; k < d; k++) {
|
|
3156
|
+
accu += decode_bf16(a[k]) * decode_bf16(b[k]);
|
|
3157
|
+
}
|
|
3158
|
+
return accu;
|
|
3159
|
+
}
|
|
3160
|
+
|
|
3161
|
+
void query_to_codes_batch_4(
|
|
3162
|
+
const uint8_t* code_0,
|
|
3163
|
+
const uint8_t* code_1,
|
|
3164
|
+
const uint8_t* code_2,
|
|
3165
|
+
const uint8_t* code_3,
|
|
3166
|
+
float& dis0,
|
|
3167
|
+
float& dis1,
|
|
3168
|
+
float& dis2,
|
|
3169
|
+
float& dis3) const final {
|
|
3170
|
+
dis0 = compute_ip(q, code_0);
|
|
3171
|
+
dis1 = compute_ip(q, code_1);
|
|
3172
|
+
dis2 = compute_ip(q, code_2);
|
|
3173
|
+
dis3 = compute_ip(q, code_3);
|
|
3174
|
+
}
|
|
3175
|
+
};
|
|
3176
|
+
|
|
3177
|
+
// * Fast path — QT_fp16 + L2
|
|
3178
|
+
// *
|
|
3179
|
+
// * fp16 code: each dim is an IEEE-754 binary16 stored as uint16. The
|
|
3180
|
+
// * query stays in full f32 precision: L2 = sum_i (q_i - decode_fp16(c_i))^2.
|
|
3181
|
+
// *
|
|
3182
|
+
// * Toolchain note: march is rv64gcv_zvfhmin — NO full zvfh, so f16
|
|
3183
|
+
// * vector arithmetic is unavailable; only vle16/vse16, vfwcvt.f.f.v,
|
|
3184
|
+
// * vfncvt.f.f.w are legal on f16.
|
|
3185
|
+
// *
|
|
3186
|
+
// * Hot loop per vl=8 dims (e16m1 -> f32m2): vle16 c + vfwcvt + vle32 q
|
|
3187
|
+
// * + vfsub + vfmacc, single f32m2 accumulator.
|
|
3188
|
+
// ************************************************************************/
|
|
3189
|
+
|
|
3190
|
+
template <>
|
|
3191
|
+
struct DCTemplate<
|
|
3192
|
+
QuantizerFP16<SIMDLevel::RISCV_RVV>,
|
|
3193
|
+
SimilarityL2<SIMDLevel::RISCV_RVV>,
|
|
3194
|
+
SIMDLevel::RISCV_RVV> : SQDistanceComputer {
|
|
3195
|
+
using Sim = SimilarityL2<SIMDLevel::RISCV_RVV>;
|
|
3196
|
+
|
|
3197
|
+
size_t d;
|
|
3198
|
+
|
|
3199
|
+
DCTemplate(size_t d_in, const std::vector<float>&) : d(d_in) {}
|
|
3200
|
+
|
|
3201
|
+
void set_query(const float* x) final {
|
|
3202
|
+
q = x;
|
|
3203
|
+
}
|
|
3204
|
+
|
|
3205
|
+
/// Direct-form L2 between the f32 query and an fp16 code (m1/m2).
|
|
3206
|
+
float compute_l2(const float* qf, const uint8_t* code8) const {
|
|
3207
|
+
const _Float16* code = (const _Float16*)code8;
|
|
3208
|
+
size_t i = 0;
|
|
3209
|
+
// Hoist vsetvl: VLMAX for e16m1 (== f32m2 lanes), reused across
|
|
3210
|
+
// the hot loop; 0 vsetvl inside.
|
|
3211
|
+
const size_t vl = __riscv_vsetvl_e16m1(d > 0 ? d : 1);
|
|
3212
|
+
vfloat32m2_t acc = __riscv_vfmv_v_f_f32m2(0.0f, vl);
|
|
3213
|
+
|
|
3214
|
+
for (; i + vl <= d; i += vl) {
|
|
3215
|
+
vfloat16m1_t vc = __riscv_vle16_v_f16m1(code + i, vl);
|
|
3216
|
+
vfloat32m2_t fc = __riscv_vfwcvt_f_f_v_f32m2(vc, vl);
|
|
3217
|
+
vfloat32m2_t fq = __riscv_vle32_v_f32m2(qf + i, vl);
|
|
3218
|
+
vfloat32m2_t t = __riscv_vfsub_vv_f32m2(fq, fc, vl);
|
|
3219
|
+
acc = __riscv_vfmacc_vv_f32m2(acc, t, t, vl);
|
|
3220
|
+
}
|
|
3221
|
+
|
|
3222
|
+
// Tail: fewer than vl dims left — one shorter-vl pass into the
|
|
3223
|
+
// same accumulator (only the first vt lanes are touched).
|
|
3224
|
+
if (i < d) {
|
|
3225
|
+
const size_t vt = __riscv_vsetvl_e16m1(d - i);
|
|
3226
|
+
vfloat16m1_t vc = __riscv_vle16_v_f16m1(code + i, vt);
|
|
3227
|
+
vfloat32m2_t fc = __riscv_vfwcvt_f_f_v_f32m2(vc, vt);
|
|
3228
|
+
vfloat32m2_t fq = __riscv_vle32_v_f32m2(qf + i, vt);
|
|
3229
|
+
vfloat32m2_t t = __riscv_vfsub_vv_f32m2(fq, fc, vt);
|
|
3230
|
+
acc = __riscv_vfmacc_vv_f32m2_tu(acc, t, t, vt);
|
|
3231
|
+
}
|
|
3232
|
+
|
|
3233
|
+
// Single horizontal reduction over all vl lanes.
|
|
3234
|
+
vfloat32m1_t red = __riscv_vfredusum_vs_f32m2_f32m1(
|
|
3235
|
+
acc, __riscv_vfmv_v_f_f32m1(0.0f, 1), vl);
|
|
3236
|
+
return __riscv_vfmv_f_s_f32m1_f32(red);
|
|
3237
|
+
}
|
|
3238
|
+
|
|
3239
|
+
float query_to_code(const uint8_t* code) const final {
|
|
3240
|
+
return compute_l2(q, code);
|
|
3241
|
+
}
|
|
3242
|
+
|
|
3243
|
+
float symmetric_dis(idx_t i, idx_t j) override {
|
|
3244
|
+
// Not on the benchmark-critical path; scalar per-dim evaluation.
|
|
3245
|
+
const uint16_t* a = (const uint16_t*)(codes + i * code_size);
|
|
3246
|
+
const uint16_t* b = (const uint16_t*)(codes + j * code_size);
|
|
3247
|
+
float accu = 0;
|
|
3248
|
+
for (size_t k = 0; k < d; k++) {
|
|
3249
|
+
float diff = decode_fp16(a[k]) - decode_fp16(b[k]);
|
|
3250
|
+
accu += diff * diff;
|
|
3251
|
+
}
|
|
3252
|
+
return accu;
|
|
3253
|
+
}
|
|
3254
|
+
|
|
3255
|
+
void query_to_codes_batch_4(
|
|
3256
|
+
const uint8_t* code_0,
|
|
3257
|
+
const uint8_t* code_1,
|
|
3258
|
+
const uint8_t* code_2,
|
|
3259
|
+
const uint8_t* code_3,
|
|
3260
|
+
float& dis0,
|
|
3261
|
+
float& dis1,
|
|
3262
|
+
float& dis2,
|
|
3263
|
+
float& dis3) const final {
|
|
3264
|
+
dis0 = compute_l2(q, code_0);
|
|
3265
|
+
dis1 = compute_l2(q, code_1);
|
|
3266
|
+
dis2 = compute_l2(q, code_2);
|
|
3267
|
+
dis3 = compute_l2(q, code_3);
|
|
3268
|
+
}
|
|
3269
|
+
};
|
|
3270
|
+
|
|
3271
|
+
// * Fast path — QT_fp16 + IP
|
|
3272
|
+
// *
|
|
3273
|
+
// * fp16 code: each dim is an IEEE-754 binary16 stored as uint16. The
|
|
3274
|
+
// * query stays in full f32 precision: IP = sum_i q_i * decode_fp16(c_i).
|
|
3275
|
+
// * Software-pipelined (prologue + rotate): preload the next chunk's f16
|
|
3276
|
+
// * code at loop top, then compute on the current chunk.
|
|
3277
|
+
// *
|
|
3278
|
+
// * Toolchain: march=rv64gcv_zvfhmin — NO full zvfh.
|
|
3279
|
+
// ************************************************************************/
|
|
3280
|
+
|
|
3281
|
+
template <>
|
|
3282
|
+
struct DCTemplate<
|
|
3283
|
+
QuantizerFP16<SIMDLevel::RISCV_RVV>,
|
|
3284
|
+
SimilarityIP<SIMDLevel::RISCV_RVV>,
|
|
3285
|
+
SIMDLevel::RISCV_RVV> : SQDistanceComputer {
|
|
3286
|
+
using Sim = SimilarityIP<SIMDLevel::RISCV_RVV>;
|
|
3287
|
+
|
|
3288
|
+
size_t d;
|
|
3289
|
+
|
|
3290
|
+
DCTemplate(size_t d_in, const std::vector<float>&) : d(d_in) {}
|
|
3291
|
+
|
|
3292
|
+
void set_query(const float* x) final {
|
|
3293
|
+
q = x;
|
|
3294
|
+
}
|
|
3295
|
+
|
|
3296
|
+
/// Direct-form IP, software-pipelined (prologue + rotate): preload the
|
|
3297
|
+
/// next chunk's f16 code, then compute on the current chunk.
|
|
3298
|
+
float compute_ip(const float* qf, const uint8_t* code8) const {
|
|
3299
|
+
const _Float16* code = (const _Float16*)code8;
|
|
3300
|
+
size_t i = 0;
|
|
3301
|
+
const size_t vl = __riscv_vsetvl_e16m1(d > 0 ? d : 1);
|
|
3302
|
+
vfloat32m2_t acc = __riscv_vfmv_v_f_f32m2(0.0f, vl);
|
|
3303
|
+
|
|
3304
|
+
if (i + vl <= d) {
|
|
3305
|
+
// Prologue: preload chunk 0.
|
|
3306
|
+
vfloat16m1_t vc = __riscv_vle16_v_f16m1(code, vl);
|
|
3307
|
+
i += vl;
|
|
3308
|
+
|
|
3309
|
+
// Main loop: while a full next chunk exists, issue its load
|
|
3310
|
+
// first, then process the current one.
|
|
3311
|
+
for (; i + vl <= d; i += vl) {
|
|
3312
|
+
vfloat16m1_t vc_next = __riscv_vle16_v_f16m1(code + i, vl);
|
|
3313
|
+
|
|
3314
|
+
vfloat32m2_t fc = __riscv_vfwcvt_f_f_v_f32m2(vc, vl);
|
|
3315
|
+
vfloat32m2_t fq = __riscv_vle32_v_f32m2(qf + i - vl, vl);
|
|
3316
|
+
acc = __riscv_vfmacc_vv_f32m2(acc, fq, fc, vl);
|
|
3317
|
+
|
|
3318
|
+
vc = vc_next; // rotate
|
|
3319
|
+
}
|
|
3320
|
+
|
|
3321
|
+
// Epilogue: process the last chunk (already loaded).
|
|
3322
|
+
{
|
|
3323
|
+
vfloat32m2_t fc = __riscv_vfwcvt_f_f_v_f32m2(vc, vl);
|
|
3324
|
+
vfloat32m2_t fq = __riscv_vle32_v_f32m2(qf + i - vl, vl);
|
|
3325
|
+
acc = __riscv_vfmacc_vv_f32m2(acc, fq, fc, vl);
|
|
3326
|
+
}
|
|
3327
|
+
}
|
|
3328
|
+
|
|
3329
|
+
// Tail: fewer than vl dims left.
|
|
3330
|
+
if (i < d) {
|
|
3331
|
+
const size_t vt = __riscv_vsetvl_e16m1(d - i);
|
|
3332
|
+
vfloat16m1_t vc = __riscv_vle16_v_f16m1(code + i, vt);
|
|
3333
|
+
vfloat32m2_t fc = __riscv_vfwcvt_f_f_v_f32m2(vc, vt);
|
|
3334
|
+
vfloat32m2_t fq = __riscv_vle32_v_f32m2(qf + i, vt);
|
|
3335
|
+
acc = __riscv_vfmacc_vv_f32m2_tu(acc, fq, fc, vt);
|
|
3336
|
+
}
|
|
3337
|
+
|
|
3338
|
+
vfloat32m1_t red = __riscv_vfredusum_vs_f32m2_f32m1(
|
|
3339
|
+
acc, __riscv_vfmv_v_f_f32m1(0.0f, 1), vl);
|
|
3340
|
+
return __riscv_vfmv_f_s_f32m1_f32(red);
|
|
3341
|
+
}
|
|
3342
|
+
|
|
3343
|
+
float query_to_code(const uint8_t* code) const final {
|
|
3344
|
+
return compute_ip(q, code);
|
|
3345
|
+
}
|
|
3346
|
+
|
|
3347
|
+
float symmetric_dis(idx_t i, idx_t j) override {
|
|
3348
|
+
// Not on the benchmark-critical path; scalar per-dim evaluation.
|
|
3349
|
+
const uint16_t* a = (const uint16_t*)(codes + i * code_size);
|
|
3350
|
+
const uint16_t* b = (const uint16_t*)(codes + j * code_size);
|
|
3351
|
+
float accu = 0;
|
|
3352
|
+
for (size_t k = 0; k < d; k++) {
|
|
3353
|
+
accu += decode_fp16(a[k]) * decode_fp16(b[k]);
|
|
3354
|
+
}
|
|
3355
|
+
return accu;
|
|
3356
|
+
}
|
|
3357
|
+
|
|
3358
|
+
void query_to_codes_batch_4(
|
|
3359
|
+
const uint8_t* code_0,
|
|
3360
|
+
const uint8_t* code_1,
|
|
3361
|
+
const uint8_t* code_2,
|
|
3362
|
+
const uint8_t* code_3,
|
|
3363
|
+
float& dis0,
|
|
3364
|
+
float& dis1,
|
|
3365
|
+
float& dis2,
|
|
3366
|
+
float& dis3) const final {
|
|
3367
|
+
dis0 = compute_ip(q, code_0);
|
|
3368
|
+
dis1 = compute_ip(q, code_1);
|
|
3369
|
+
dis2 = compute_ip(q, code_2);
|
|
3370
|
+
dis3 = compute_ip(q, code_3);
|
|
3371
|
+
}
|
|
3372
|
+
};
|
|
3373
|
+
|
|
3374
|
+
} // namespace scalar_quantizer
|
|
3375
|
+
} // namespace faiss
|
|
3376
|
+
|
|
3377
|
+
#define THE_LEVEL_TO_DISPATCH SIMDLevel::RISCV_RVV
|
|
3378
|
+
#include <faiss/impl/scalar_quantizer/sq-dispatch.h>
|