diff --git a/faiss/cppcontrib/SaDecodeKernels.h b/faiss/cppcontrib/SaDecodeKernels.h index b0e7efe1f..ca6e3cc9d 100644 --- a/faiss/cppcontrib/SaDecodeKernels.h +++ b/faiss/cppcontrib/SaDecodeKernels.h @@ -313,7 +313,7 @@ #ifdef __AVX2__ #include #include -#elif defined(__ARM_NEON) +#elif defined(__ARM_NEON) || defined(__aarch64__) #include #include #else diff --git a/faiss/cppcontrib/sa_decode/Level2-neon-inl.h b/faiss/cppcontrib/sa_decode/Level2-neon-inl.h index 51daec426..497bf9804 100644 --- a/faiss/cppcontrib/sa_decode/Level2-neon-inl.h +++ b/faiss/cppcontrib/sa_decode/Level2-neon-inl.h @@ -142,6 +142,57 @@ inline float32x4x2_t elementaryBlock8x1bAccum( return {result0, result1}; } +// Processes 8 float values. +// Returns { +// val[0] = {*coarse[0..1] + *fine0[0..1], *coarse[2..3] + *fine1[0..1]}; +// val[1] = {*coarse[4..5] + *fine2[0..1], *coarse[6..7] + *fine3[0..1]}; +// } +inline float32x4x2_t elementaryBlock2x4b( + const float* const __restrict coarse, + const float* const __restrict fine0, + const float* const __restrict fine1, + const float* const __restrict fine2, + const float* const __restrict fine3) { + const auto fine0Value = vld1_f32(fine0); + const auto fine1Value = vld1_f32(fine1); + const auto fine2Value = vld1_f32(fine2); + const auto fine3Value = vld1_f32(fine3); + + const auto coarseValue0 = vld1q_f32(coarse); + const auto coarseValue1 = vld1q_f32(coarse + 4); + + const auto fineResult0 = vcombine_f32(fine0Value, fine1Value); + const auto fineResult1 = vcombine_f32(fine2Value, fine3Value); + + return {vaddq_f32(fineResult0, coarseValue0), + vaddq_f32(fineResult1, coarseValue1)}; +} + +// Processes 8 float values. +// Returns { +// val[0] = existingValue.val[0] + weight * (*coarse[0..1] + *fine0[0..1], +// *coarse[2..3] + *fine1[0..1]); val[1] = existingValue.val[1] + weight * +// (*coarse[4..5] + *fine2[0..1], *coarse[6..7] + *fine3[0..1]); +// } +inline float32x4x2_t elementaryBlock2x4bAccum( + const float* const __restrict coarse, + const float* const __restrict fine0, + const float* const __restrict fine1, + const float* const __restrict fine2, + const float* const __restrict fine3, + const float weight, + const float32x4x2_t existingValue) { + const auto combinedValue = + elementaryBlock2x4b(coarse, fine0, fine1, fine2, fine3); + + const auto weightNeon = vdupq_n_f32(weight); + const auto result0 = + vfmaq_f32(existingValue.val[0], weightNeon, combinedValue.val[0]); + const auto result1 = + vfmaq_f32(existingValue.val[1], weightNeon, combinedValue.val[1]); + return {result0, result1}; +} + // The following code uses template-based for-loop unrolling, // because the compiler does not do that on its own as needed. // The idea is the following: @@ -168,12 +219,846 @@ template < intptr_t COARSE_BITS, intptr_t FINE_BITS, intptr_t CPOS, + bool FINE_SIZE_EQ_2 = FINE_SIZE == 2, bool FINE_SIZE_EQ_4 = FINE_SIZE == 4, bool QPOS_LEFT_GE_8 = (FINE_SIZE - CPOS % FINE_SIZE >= 8), bool QPOS_LEFT_GE_4 = (FINE_SIZE - CPOS % FINE_SIZE >= 4), bool DIM_EQ_CPOS = DIM == CPOS> struct Index2LevelDecoderImpl; +template < + intptr_t DIM, + intptr_t COARSE_SIZE, + intptr_t COARSE_BITS, + intptr_t FINE_BITS, + intptr_t CPOS, + bool QPOS_LEFT_GE_8, + bool QPOS_LEFT_GE_4> +struct Index2LevelDecoderImpl< + DIM, + COARSE_SIZE, + 2, + COARSE_BITS, + FINE_BITS, + CPOS, + true, + false, + QPOS_LEFT_GE_8, + QPOS_LEFT_GE_4, + false> { + static constexpr intptr_t FINE_SIZE = 2; + + static constexpr intptr_t coarseCentroidIdx = CPOS / COARSE_SIZE; + static constexpr intptr_t coarseCentroidOffset = CPOS % COARSE_SIZE; + static constexpr intptr_t fineCentroidIdx = CPOS / FINE_SIZE; + static constexpr intptr_t fineCentroidOffset = CPOS % FINE_SIZE; + + static constexpr intptr_t QPOS_LEFT = FINE_SIZE - fineCentroidOffset; + + // coarse quantizer storage + static constexpr intptr_t COARSE_TABLE_BYTES = (1 << COARSE_BITS); + + // coarse quantizer bytes start from 0 + // fine quantizer bytes start from N_COARSE_ELEMENTS_BYTES + static constexpr intptr_t N_COARSE_ELEMENTS = DIM / COARSE_SIZE; + static constexpr intptr_t N_COARSE_ELEMENTS_BITS = + N_COARSE_ELEMENTS * COARSE_BITS; + static constexpr intptr_t N_COARSE_ELEMENTS_BYTES = + (N_COARSE_ELEMENTS_BITS + 7) / 8; + + static constexpr intptr_t FINE_TABLE_BYTES = (1 << FINE_BITS); + + // process 1 sample + static void store( + const float* const __restrict pqCoarseCentroids0, + const float* const __restrict pqFineCentroids0, + const uint8_t* const __restrict code0, + float* const __restrict outputStore) { + // coarse quantizer + const uint8_t* const __restrict coarse0 = code0; + + // fine quantizer + const uint8_t* const __restrict fine0 = code0 + N_COARSE_ELEMENTS_BYTES; + + // process chunks, 2 float + // but 8 floats per loop + + const intptr_t coarseCode0 = detail:: + UintReader:: + get(coarse0); + const intptr_t fineCode0a = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0b = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0c = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0d = detail:: + UintReader::get( + fine0); + + const auto storeValue = elementaryBlock2x4b( + pqCoarseCentroids0 + + (coarseCentroidIdx * COARSE_TABLE_BYTES + coarseCode0) * + COARSE_SIZE + + coarseCentroidOffset, + pqFineCentroids0 + + ((fineCentroidIdx + 0) * FINE_TABLE_BYTES + + fineCode0a) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids0 + + ((fineCentroidIdx + 1) * FINE_TABLE_BYTES + + fineCode0b) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids0 + + ((fineCentroidIdx + 2) * FINE_TABLE_BYTES + + fineCode0c) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids0 + + ((fineCentroidIdx + 3) * FINE_TABLE_BYTES + + fineCode0d) * + FINE_SIZE + + fineCentroidOffset); + + vst1q_f32(outputStore + CPOS, storeValue.val[0]); + vst1q_f32(outputStore + CPOS + 4, storeValue.val[1]); + + // next + Index2LevelDecoderImpl< + DIM, + COARSE_SIZE, + FINE_SIZE, + COARSE_BITS, + FINE_BITS, + CPOS + 8>:: + store(pqCoarseCentroids0, pqFineCentroids0, code0, outputStore); + } + + // process 1 sample + static void accum( + const float* const __restrict pqCoarseCentroids0, + const float* const __restrict pqFineCentroids0, + const uint8_t* const __restrict code0, + const float weight0, + float* const __restrict outputAccum) { + // coarse quantizer + const uint8_t* const __restrict coarse0 = code0; + + // fine quantizer + const uint8_t* const __restrict fine0 = code0 + N_COARSE_ELEMENTS_BYTES; + + // process chunks, 2 float + // but 8 floats per loop + + const intptr_t coarseCode0 = detail:: + UintReader:: + get(coarse0); + const intptr_t fineCode0a = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0b = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0c = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0d = detail:: + UintReader::get( + fine0); + + auto existingValue0 = vld1q_f32(outputAccum + CPOS); + auto existingValue1 = vld1q_f32(outputAccum + CPOS + 4); + + auto existingValue = elementaryBlock2x4bAccum( + pqCoarseCentroids0 + + (coarseCentroidIdx * COARSE_TABLE_BYTES + coarseCode0) * + COARSE_SIZE + + coarseCentroidOffset, + pqFineCentroids0 + + ((fineCentroidIdx + 0) * FINE_TABLE_BYTES + + fineCode0a) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids0 + + ((fineCentroidIdx + 1) * FINE_TABLE_BYTES + + fineCode0b) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids0 + + ((fineCentroidIdx + 2) * FINE_TABLE_BYTES + + fineCode0c) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids0 + + ((fineCentroidIdx + 3) * FINE_TABLE_BYTES + + fineCode0d) * + FINE_SIZE + + fineCentroidOffset, + weight0, + {existingValue0, existingValue1}); + + vst1q_f32(outputAccum + CPOS, existingValue.val[0]); + vst1q_f32(outputAccum + CPOS + 4, existingValue.val[1]); + + // next + Index2LevelDecoderImpl< + DIM, + COARSE_SIZE, + FINE_SIZE, + COARSE_BITS, + FINE_BITS, + CPOS + 8>:: + accum(pqCoarseCentroids0, + pqFineCentroids0, + code0, + weight0, + outputAccum); + } + + // Process 2 samples. + // Each code uses its own coarse pq centroids table and fine pq centroids + // table. + static void accum( + const float* const __restrict pqCoarseCentroids0, + const float* const __restrict pqFineCentroids0, + const uint8_t* const __restrict code0, + const float weight0, + const float* const __restrict pqCoarseCentroids1, + const float* const __restrict pqFineCentroids1, + const uint8_t* const __restrict code1, + const float weight1, + float* const __restrict outputAccum) { + // coarse quantizer + const uint8_t* const __restrict coarse0 = code0; + const uint8_t* const __restrict coarse1 = code1; + + // fine quantizer + const uint8_t* const __restrict fine0 = code0 + N_COARSE_ELEMENTS_BYTES; + const uint8_t* const __restrict fine1 = code1 + N_COARSE_ELEMENTS_BYTES; + + // process chunks, 2 float + // but 8 floats per loop + + const intptr_t coarseCode0 = detail:: + UintReader:: + get(coarse0); + const intptr_t fineCode0a = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0b = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0c = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0d = detail:: + UintReader::get( + fine0); + const intptr_t coarseCode1 = detail:: + UintReader:: + get(coarse1); + const intptr_t fineCode1a = detail:: + UintReader::get( + fine1); + const intptr_t fineCode1b = detail:: + UintReader::get( + fine1); + const intptr_t fineCode1c = detail:: + UintReader::get( + fine1); + const intptr_t fineCode1d = detail:: + UintReader::get( + fine1); + + auto existingValue0 = vld1q_f32(outputAccum + CPOS); + auto existingValue1 = vld1q_f32(outputAccum + CPOS + 4); + + auto existingValue = elementaryBlock2x4bAccum( + pqCoarseCentroids0 + + (coarseCentroidIdx * COARSE_TABLE_BYTES + coarseCode0) * + COARSE_SIZE + + coarseCentroidOffset, + pqFineCentroids0 + + ((fineCentroidIdx + 0) * FINE_TABLE_BYTES + + fineCode0a) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids0 + + ((fineCentroidIdx + 1) * FINE_TABLE_BYTES + + fineCode0b) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids0 + + ((fineCentroidIdx + 2) * FINE_TABLE_BYTES + + fineCode0c) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids0 + + ((fineCentroidIdx + 3) * FINE_TABLE_BYTES + + fineCode0d) * + FINE_SIZE + + fineCentroidOffset, + weight0, + {existingValue0, existingValue1}); + + existingValue = elementaryBlock2x4bAccum( + pqCoarseCentroids1 + + (coarseCentroidIdx * COARSE_TABLE_BYTES + coarseCode1) * + COARSE_SIZE + + coarseCentroidOffset, + pqFineCentroids1 + + ((fineCentroidIdx + 0) * FINE_TABLE_BYTES + + fineCode1a) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids1 + + ((fineCentroidIdx + 1) * FINE_TABLE_BYTES + + fineCode1b) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids1 + + ((fineCentroidIdx + 2) * FINE_TABLE_BYTES + + fineCode1c) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids1 + + ((fineCentroidIdx + 3) * FINE_TABLE_BYTES + + fineCode1d) * + FINE_SIZE + + fineCentroidOffset, + weight1, + existingValue); + + vst1q_f32(outputAccum + CPOS, existingValue.val[0]); + vst1q_f32(outputAccum + CPOS + 4, existingValue.val[1]); + + // next + Index2LevelDecoderImpl< + DIM, + COARSE_SIZE, + FINE_SIZE, + COARSE_BITS, + FINE_BITS, + CPOS + 8>:: + accum(pqCoarseCentroids0, + pqFineCentroids0, + code0, + weight0, + pqCoarseCentroids1, + pqFineCentroids1, + code1, + weight1, + outputAccum); + } + + // Process 2 samples. + // Coarse pq centroids table and fine pq centroids table are shared among + // codes. + static void accum( + const float* const __restrict pqCoarseCentroids, + const float* const __restrict pqFineCentroids, + const uint8_t* const __restrict code0, + const float weight0, + const uint8_t* const __restrict code1, + const float weight1, + float* const __restrict outputAccum) { + // coarse quantizer + const uint8_t* const __restrict coarse0 = code0; + const uint8_t* const __restrict coarse1 = code1; + + // fine quantizer + const uint8_t* const __restrict fine0 = code0 + N_COARSE_ELEMENTS_BYTES; + const uint8_t* const __restrict fine1 = code1 + N_COARSE_ELEMENTS_BYTES; + + // process chunks, 2 float + // but 8 floats per loop + + const intptr_t coarseCode0 = detail:: + UintReader:: + get(coarse0); + const intptr_t fineCode0a = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0b = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0c = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0d = detail:: + UintReader::get( + fine0); + const intptr_t coarseCode1 = detail:: + UintReader:: + get(coarse1); + const intptr_t fineCode1a = detail:: + UintReader::get( + fine1); + const intptr_t fineCode1b = detail:: + UintReader::get( + fine1); + const intptr_t fineCode1c = detail:: + UintReader::get( + fine1); + const intptr_t fineCode1d = detail:: + UintReader::get( + fine1); + + auto existingValue0 = vld1q_f32(outputAccum + CPOS); + auto existingValue1 = vld1q_f32(outputAccum + CPOS + 4); + + auto existingValue = elementaryBlock2x4bAccum( + pqCoarseCentroids + + (coarseCentroidIdx * COARSE_TABLE_BYTES + coarseCode0) * + COARSE_SIZE + + coarseCentroidOffset, + pqFineCentroids + + ((fineCentroidIdx + 0) * FINE_TABLE_BYTES + + fineCode0a) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids + + ((fineCentroidIdx + 1) * FINE_TABLE_BYTES + + fineCode0b) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids + + ((fineCentroidIdx + 2) * FINE_TABLE_BYTES + + fineCode0c) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids + + ((fineCentroidIdx + 3) * FINE_TABLE_BYTES + + fineCode0d) * + FINE_SIZE + + fineCentroidOffset, + weight0, + {existingValue0, existingValue1}); + + existingValue = elementaryBlock2x4bAccum( + pqCoarseCentroids + + (coarseCentroidIdx * COARSE_TABLE_BYTES + coarseCode1) * + COARSE_SIZE + + coarseCentroidOffset, + pqFineCentroids + + ((fineCentroidIdx + 0) * FINE_TABLE_BYTES + + fineCode1a) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids + + ((fineCentroidIdx + 1) * FINE_TABLE_BYTES + + fineCode1b) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids + + ((fineCentroidIdx + 2) * FINE_TABLE_BYTES + + fineCode1c) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids + + ((fineCentroidIdx + 3) * FINE_TABLE_BYTES + + fineCode1d) * + FINE_SIZE + + fineCentroidOffset, + weight1, + existingValue); + + vst1q_f32(outputAccum + CPOS, existingValue.val[0]); + vst1q_f32(outputAccum + CPOS + 4, existingValue.val[1]); + + // next + Index2LevelDecoderImpl< + DIM, + COARSE_SIZE, + FINE_SIZE, + COARSE_BITS, + FINE_BITS, + CPOS + 8>:: + accum(pqCoarseCentroids, + pqFineCentroids, + code0, + weight0, + code1, + weight1, + outputAccum); + } + + // Process 3 samples. + // Each code uses its own coarse pq centroids table and fine pq centroids + // table. + static void accum( + const float* const __restrict pqCoarseCentroids0, + const float* const __restrict pqFineCentroids0, + const uint8_t* const __restrict code0, + const float weight0, + const float* const __restrict pqCoarseCentroids1, + const float* const __restrict pqFineCentroids1, + const uint8_t* const __restrict code1, + const float weight1, + const float* const __restrict pqCoarseCentroids2, + const float* const __restrict pqFineCentroids2, + const uint8_t* const __restrict code2, + const float weight2, + float* const __restrict outputAccum) { + // coarse quantizer + const uint8_t* const __restrict coarse0 = code0; + const uint8_t* const __restrict coarse1 = code1; + const uint8_t* const __restrict coarse2 = code2; + + // fine quantizer + const uint8_t* const __restrict fine0 = code0 + N_COARSE_ELEMENTS_BYTES; + const uint8_t* const __restrict fine1 = code1 + N_COARSE_ELEMENTS_BYTES; + const uint8_t* const __restrict fine2 = code2 + N_COARSE_ELEMENTS_BYTES; + + // process chunks, 2 float + // but 8 floats per loop + + const intptr_t coarseCode0 = detail:: + UintReader:: + get(coarse0); + const intptr_t fineCode0a = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0b = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0c = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0d = detail:: + UintReader::get( + fine0); + const intptr_t coarseCode1 = detail:: + UintReader:: + get(coarse1); + const intptr_t fineCode1a = detail:: + UintReader::get( + fine1); + const intptr_t fineCode1b = detail:: + UintReader::get( + fine1); + const intptr_t fineCode1c = detail:: + UintReader::get( + fine1); + const intptr_t fineCode1d = detail:: + UintReader::get( + fine1); + const intptr_t coarseCode2 = detail:: + UintReader:: + get(coarse2); + const intptr_t fineCode2a = detail:: + UintReader::get( + fine2); + const intptr_t fineCode2b = detail:: + UintReader::get( + fine2); + const intptr_t fineCode2c = detail:: + UintReader::get( + fine2); + const intptr_t fineCode2d = detail:: + UintReader::get( + fine2); + + auto existingValue0 = vld1q_f32(outputAccum + CPOS); + auto existingValue1 = vld1q_f32(outputAccum + CPOS + 4); + + auto existingValue = elementaryBlock2x4bAccum( + pqCoarseCentroids0 + + (coarseCentroidIdx * COARSE_TABLE_BYTES + coarseCode0) * + COARSE_SIZE + + coarseCentroidOffset, + pqFineCentroids0 + + ((fineCentroidIdx + 0) * FINE_TABLE_BYTES + + fineCode0a) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids0 + + ((fineCentroidIdx + 1) * FINE_TABLE_BYTES + + fineCode0b) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids0 + + ((fineCentroidIdx + 2) * FINE_TABLE_BYTES + + fineCode0c) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids0 + + ((fineCentroidIdx + 3) * FINE_TABLE_BYTES + + fineCode0d) * + FINE_SIZE + + fineCentroidOffset, + weight0, + {existingValue0, existingValue1}); + + existingValue = elementaryBlock2x4bAccum( + pqCoarseCentroids1 + + (coarseCentroidIdx * COARSE_TABLE_BYTES + coarseCode1) * + COARSE_SIZE + + coarseCentroidOffset, + pqFineCentroids1 + + ((fineCentroidIdx + 0) * FINE_TABLE_BYTES + + fineCode1a) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids1 + + ((fineCentroidIdx + 1) * FINE_TABLE_BYTES + + fineCode1b) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids1 + + ((fineCentroidIdx + 2) * FINE_TABLE_BYTES + + fineCode1c) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids1 + + ((fineCentroidIdx + 3) * FINE_TABLE_BYTES + + fineCode1d) * + FINE_SIZE + + fineCentroidOffset, + weight1, + existingValue); + + existingValue = elementaryBlock2x4bAccum( + pqCoarseCentroids2 + + (coarseCentroidIdx * COARSE_TABLE_BYTES + coarseCode2) * + COARSE_SIZE + + coarseCentroidOffset, + pqFineCentroids2 + + ((fineCentroidIdx + 0) * FINE_TABLE_BYTES + + fineCode2a) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids2 + + ((fineCentroidIdx + 1) * FINE_TABLE_BYTES + + fineCode2b) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids2 + + ((fineCentroidIdx + 2) * FINE_TABLE_BYTES + + fineCode2c) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids2 + + ((fineCentroidIdx + 3) * FINE_TABLE_BYTES + + fineCode2d) * + FINE_SIZE + + fineCentroidOffset, + weight2, + existingValue); + + vst1q_f32(outputAccum + CPOS, existingValue.val[0]); + vst1q_f32(outputAccum + CPOS + 4, existingValue.val[1]); + + // next + Index2LevelDecoderImpl< + DIM, + COARSE_SIZE, + FINE_SIZE, + COARSE_BITS, + FINE_BITS, + CPOS + 8>:: + accum(pqCoarseCentroids0, + pqFineCentroids0, + code0, + weight0, + pqCoarseCentroids1, + pqFineCentroids1, + code1, + weight1, + pqCoarseCentroids2, + pqFineCentroids2, + code2, + weight2, + outputAccum); + } + + // Process 3 samples. + // Coarse pq centroids table and fine pq centroids table are shared among + // codes. + static void accum( + const float* const __restrict pqCoarseCentroids, + const float* const __restrict pqFineCentroids, + const uint8_t* const __restrict code0, + const float weight0, + const uint8_t* const __restrict code1, + const float weight1, + const uint8_t* const __restrict code2, + const float weight2, + float* const __restrict outputAccum) { + // coarse quantizer + const uint8_t* const __restrict coarse0 = code0; + const uint8_t* const __restrict coarse1 = code1; + const uint8_t* const __restrict coarse2 = code2; + + // fine quantizer + const uint8_t* const __restrict fine0 = code0 + N_COARSE_ELEMENTS_BYTES; + const uint8_t* const __restrict fine1 = code1 + N_COARSE_ELEMENTS_BYTES; + const uint8_t* const __restrict fine2 = code2 + N_COARSE_ELEMENTS_BYTES; + + // process chunks, 2 float + // but 8 floats per loop + + const intptr_t coarseCode0 = detail:: + UintReader:: + get(coarse0); + const intptr_t fineCode0a = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0b = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0c = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0d = detail:: + UintReader::get( + fine0); + const intptr_t coarseCode1 = detail:: + UintReader:: + get(coarse1); + const intptr_t fineCode1a = detail:: + UintReader::get( + fine1); + const intptr_t fineCode1b = detail:: + UintReader::get( + fine1); + const intptr_t fineCode1c = detail:: + UintReader::get( + fine1); + const intptr_t fineCode1d = detail:: + UintReader::get( + fine1); + const intptr_t coarseCode2 = detail:: + UintReader:: + get(coarse2); + const intptr_t fineCode2a = detail:: + UintReader::get( + fine2); + const intptr_t fineCode2b = detail:: + UintReader::get( + fine2); + const intptr_t fineCode2c = detail:: + UintReader::get( + fine2); + const intptr_t fineCode2d = detail:: + UintReader::get( + fine2); + + auto existingValue0 = vld1q_f32(outputAccum + CPOS); + auto existingValue1 = vld1q_f32(outputAccum + CPOS + 4); + + auto existingValue = elementaryBlock2x4bAccum( + pqCoarseCentroids + + (coarseCentroidIdx * COARSE_TABLE_BYTES + coarseCode0) * + COARSE_SIZE + + coarseCentroidOffset, + pqFineCentroids + + ((fineCentroidIdx + 0) * FINE_TABLE_BYTES + + fineCode0a) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids + + ((fineCentroidIdx + 1) * FINE_TABLE_BYTES + + fineCode0b) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids + + ((fineCentroidIdx + 2) * FINE_TABLE_BYTES + + fineCode0c) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids + + ((fineCentroidIdx + 3) * FINE_TABLE_BYTES + + fineCode0d) * + FINE_SIZE + + fineCentroidOffset, + weight0, + {existingValue0, existingValue1}); + + existingValue = elementaryBlock2x4bAccum( + pqCoarseCentroids + + (coarseCentroidIdx * COARSE_TABLE_BYTES + coarseCode1) * + COARSE_SIZE + + coarseCentroidOffset, + pqFineCentroids + + ((fineCentroidIdx + 0) * FINE_TABLE_BYTES + + fineCode1a) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids + + ((fineCentroidIdx + 1) * FINE_TABLE_BYTES + + fineCode1b) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids + + ((fineCentroidIdx + 2) * FINE_TABLE_BYTES + + fineCode1c) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids + + ((fineCentroidIdx + 3) * FINE_TABLE_BYTES + + fineCode1d) * + FINE_SIZE + + fineCentroidOffset, + weight1, + existingValue); + + existingValue = elementaryBlock2x4bAccum( + pqCoarseCentroids + + (coarseCentroidIdx * COARSE_TABLE_BYTES + coarseCode2) * + COARSE_SIZE + + coarseCentroidOffset, + pqFineCentroids + + ((fineCentroidIdx + 0) * FINE_TABLE_BYTES + + fineCode2a) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids + + ((fineCentroidIdx + 1) * FINE_TABLE_BYTES + + fineCode2b) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids + + ((fineCentroidIdx + 2) * FINE_TABLE_BYTES + + fineCode2c) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids + + ((fineCentroidIdx + 3) * FINE_TABLE_BYTES + + fineCode2d) * + FINE_SIZE + + fineCentroidOffset, + weight2, + existingValue); + + vst1q_f32(outputAccum + CPOS, existingValue.val[0]); + vst1q_f32(outputAccum + CPOS + 4, existingValue.val[1]); + + // next + Index2LevelDecoderImpl< + DIM, + COARSE_SIZE, + FINE_SIZE, + COARSE_BITS, + FINE_BITS, + CPOS + 8>:: + accum(pqCoarseCentroids, + pqFineCentroids, + code0, + weight0, + code1, + weight1, + code2, + weight2, + outputAccum); + } +}; + template < intptr_t DIM, intptr_t COARSE_SIZE, @@ -189,6 +1074,7 @@ struct Index2LevelDecoderImpl< COARSE_BITS, FINE_BITS, CPOS, + false, true, QPOS_LEFT_GE_8, QPOS_LEFT_GE_4, @@ -829,6 +1715,7 @@ struct Index2LevelDecoderImpl< FINE_BITS, CPOS, false, + false, true, true, false> { @@ -1353,6 +2240,7 @@ struct Index2LevelDecoderImpl< CPOS, false, false, + false, true, false> { static constexpr intptr_t coarseCentroidIdx = CPOS / COARSE_SIZE; @@ -1856,6 +2744,7 @@ template < intptr_t FINE_SIZE, intptr_t COARSE_BITS, intptr_t FINE_BITS, + bool FINE_SIZE_EQ_2, bool FINE_SIZE_EQ_4, bool QPOS_LEFT_GE_8, bool QPOS_LEFT_GE_4> @@ -1866,6 +2755,7 @@ struct Index2LevelDecoderImpl< COARSE_BITS, FINE_BITS, DIM, + FINE_SIZE_EQ_2, FINE_SIZE_EQ_4, QPOS_LEFT_GE_8, QPOS_LEFT_GE_4, diff --git a/faiss/cppcontrib/sa_decode/PQ-neon-inl.h b/faiss/cppcontrib/sa_decode/PQ-neon-inl.h index 1ee278e25..4a85c04c2 100644 --- a/faiss/cppcontrib/sa_decode/PQ-neon-inl.h +++ b/faiss/cppcontrib/sa_decode/PQ-neon-inl.h @@ -122,6 +122,49 @@ inline float32x4x2_t elementaryBlock8x1bAccum( return {result0, result1}; } +// Processes 8 float values. +// Returns { +// val[0] = {*fine0[0..1], *fine1[0..1]}; +// val[1] = {*fine2[0..1], *fine3[0..1]}; +// } +inline float32x4x2_t elementaryBlock2x4b( + const float* const __restrict fine0, + const float* const __restrict fine1, + const float* const __restrict fine2, + const float* const __restrict fine3) { + const auto fine0Value = vld1_f32(fine0); + const auto fine1Value = vld1_f32(fine1); + const auto fine2Value = vld1_f32(fine2); + const auto fine3Value = vld1_f32(fine3); + + const auto result0 = vcombine_f32(fine0Value, fine1Value); + const auto result1 = vcombine_f32(fine2Value, fine3Value); + + return {result0, result1}; +} + +// Processes 8 float values. +// Returns { +// val[0] = existingValue.val[0] + weight * {*fine0[0..1], *fine1[0..1]}; +// val[1] = existingValue.val[1] + weight * {*fine2[0..1], *fine3[0..1]}; +// } +inline float32x4x2_t elementaryBlock2x4bAccum( + const float* const __restrict fine0, + const float* const __restrict fine1, + const float* const __restrict fine2, + const float* const __restrict fine3, + const float weight, + const float32x4x2_t existingValue) { + const auto fineValue = elementaryBlock2x4b(fine0, fine1, fine2, fine3); + + const auto weightNeon = vdupq_n_f32(weight); + const auto result0 = + vfmaq_f32(existingValue.val[0], weightNeon, fineValue.val[0]); + const auto result1 = + vfmaq_f32(existingValue.val[1], weightNeon, fineValue.val[1]); + return {result0, result1}; +} + // The following code uses template-based for-loop unrolling, // because the compiler does not do that on its own as needed. // The idea is the following: @@ -146,12 +189,656 @@ template < intptr_t FINE_SIZE, intptr_t FINE_BITS, intptr_t CPOS, + bool FINE_SIZE_EQ_2 = FINE_SIZE == 2, bool FINE_SIZE_EQ_4 = FINE_SIZE == 4, bool QPOS_LEFT_GE_8 = (FINE_SIZE - CPOS % FINE_SIZE >= 8), bool QPOS_LEFT_GE_4 = (FINE_SIZE - CPOS % FINE_SIZE >= 4), bool DIM_EQ_CPOS = DIM == CPOS> struct IndexPQDecoderImpl; +template < + intptr_t DIM, + intptr_t CPOS, + intptr_t FINE_BITS, + bool QPOS_LEFT_GE_8, + bool QPOS_LEFT_GE_4> +struct IndexPQDecoderImpl< + DIM, + 2, + FINE_BITS, + CPOS, + true, + false, + QPOS_LEFT_GE_8, + QPOS_LEFT_GE_4, + false> { + static constexpr intptr_t FINE_SIZE = 2; + + static constexpr intptr_t fineCentroidIdx = CPOS / FINE_SIZE; + static constexpr intptr_t fineCentroidOffset = CPOS % FINE_SIZE; + + static constexpr intptr_t QPOS_LEFT = FINE_SIZE - fineCentroidOffset; + + static constexpr intptr_t FINE_TABLE_BYTES = (1 << FINE_BITS); + + // process 1 sample + static void store( + const float* const __restrict pqFineCentroids0, + const uint8_t* const __restrict code0, + float* const __restrict outputStore) { + // fine quantizer + const uint8_t* const __restrict fine0 = code0; + + // process chunks, 2 float + // but 8 floats per loop + + const intptr_t fineCode0a = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0b = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0c = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0d = detail:: + UintReader::get( + fine0); + + const auto storeValue = elementaryBlock2x4b( + pqFineCentroids0 + + ((fineCentroidIdx + 0) * FINE_TABLE_BYTES + + fineCode0a) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids0 + + ((fineCentroidIdx + 1) * FINE_TABLE_BYTES + + fineCode0b) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids0 + + ((fineCentroidIdx + 2) * FINE_TABLE_BYTES + + fineCode0c) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids0 + + ((fineCentroidIdx + 3) * FINE_TABLE_BYTES + + fineCode0d) * + FINE_SIZE + + fineCentroidOffset); + + vst1q_f32(outputStore + CPOS, storeValue.val[0]); + vst1q_f32(outputStore + CPOS + 4, storeValue.val[1]); + + // next + IndexPQDecoderImpl::store( + pqFineCentroids0, code0, outputStore); + } + + // process 1 sample + static void accum( + const float* const __restrict pqFineCentroids0, + const uint8_t* const __restrict code0, + const float weight0, + float* const __restrict outputAccum) { + // fine quantizer + const uint8_t* const __restrict fine0 = code0; + + // process chunks, 2 float + // but 8 floats per loop + + const intptr_t fineCode0a = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0b = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0c = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0d = detail:: + UintReader::get( + fine0); + + auto existingValue0 = vld1q_f32(outputAccum + CPOS); + auto existingValue1 = vld1q_f32(outputAccum + CPOS + 4); + + auto existingValue = elementaryBlock2x4bAccum( + pqFineCentroids0 + + ((fineCentroidIdx + 0) * FINE_TABLE_BYTES + + fineCode0a) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids0 + + ((fineCentroidIdx + 1) * FINE_TABLE_BYTES + + fineCode0b) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids0 + + ((fineCentroidIdx + 2) * FINE_TABLE_BYTES + + fineCode0c) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids0 + + ((fineCentroidIdx + 3) * FINE_TABLE_BYTES + + fineCode0d) * + FINE_SIZE + + fineCentroidOffset, + weight0, + {existingValue0, existingValue1}); + + vst1q_f32(outputAccum + CPOS, existingValue.val[0]); + vst1q_f32(outputAccum + CPOS + 4, existingValue.val[1]); + + // next + IndexPQDecoderImpl::accum( + pqFineCentroids0, code0, weight0, outputAccum); + } + + // Process 2 samples. + // Each code uses its own fine pq centroids table. + static void accum( + const float* const __restrict pqFineCentroids0, + const uint8_t* const __restrict code0, + const float weight0, + const float* const __restrict pqFineCentroids1, + const uint8_t* const __restrict code1, + const float weight1, + float* const __restrict outputAccum) { + // fine quantizer + const uint8_t* const __restrict fine0 = code0; + const uint8_t* const __restrict fine1 = code1; + + // process chunks, 2 float + // but 8 floats per loop + + const intptr_t fineCode0a = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0b = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0c = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0d = detail:: + UintReader::get( + fine0); + const intptr_t fineCode1a = detail:: + UintReader::get( + fine1); + const intptr_t fineCode1b = detail:: + UintReader::get( + fine1); + const intptr_t fineCode1c = detail:: + UintReader::get( + fine1); + const intptr_t fineCode1d = detail:: + UintReader::get( + fine1); + + auto existingValue0 = vld1q_f32(outputAccum + CPOS); + auto existingValue1 = vld1q_f32(outputAccum + CPOS + 4); + + auto existingValue = elementaryBlock2x4bAccum( + pqFineCentroids0 + + ((fineCentroidIdx + 0) * FINE_TABLE_BYTES + + fineCode0a) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids0 + + ((fineCentroidIdx + 1) * FINE_TABLE_BYTES + + fineCode0b) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids0 + + ((fineCentroidIdx + 2) * FINE_TABLE_BYTES + + fineCode0c) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids0 + + ((fineCentroidIdx + 3) * FINE_TABLE_BYTES + + fineCode0d) * + FINE_SIZE + + fineCentroidOffset, + weight0, + {existingValue0, existingValue1}); + + existingValue = elementaryBlock2x4bAccum( + pqFineCentroids1 + + ((fineCentroidIdx + 0) * FINE_TABLE_BYTES + + fineCode1a) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids1 + + ((fineCentroidIdx + 1) * FINE_TABLE_BYTES + + fineCode1b) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids1 + + ((fineCentroidIdx + 2) * FINE_TABLE_BYTES + + fineCode1c) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids1 + + ((fineCentroidIdx + 3) * FINE_TABLE_BYTES + + fineCode1d) * + FINE_SIZE + + fineCentroidOffset, + weight1, + existingValue); + + vst1q_f32(outputAccum + CPOS, existingValue.val[0]); + vst1q_f32(outputAccum + CPOS + 4, existingValue.val[1]); + + // next + IndexPQDecoderImpl::accum( + pqFineCentroids0, + code0, + weight0, + pqFineCentroids1, + code1, + weight1, + outputAccum); + } + + // Process 2 samples. + // Fine pq centroids table is shared among codes. + static void accum( + const float* const __restrict pqFineCentroids, + const uint8_t* const __restrict code0, + const float weight0, + const uint8_t* const __restrict code1, + const float weight1, + float* const __restrict outputAccum) { + // fine quantizer + const uint8_t* const __restrict fine0 = code0; + const uint8_t* const __restrict fine1 = code1; + + // process chunks, 2 float + // but 8 floats per loop + + const intptr_t fineCode0a = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0b = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0c = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0d = detail:: + UintReader::get( + fine0); + const intptr_t fineCode1a = detail:: + UintReader::get( + fine1); + const intptr_t fineCode1b = detail:: + UintReader::get( + fine1); + const intptr_t fineCode1c = detail:: + UintReader::get( + fine1); + const intptr_t fineCode1d = detail:: + UintReader::get( + fine1); + + auto existingValue0 = vld1q_f32(outputAccum + CPOS); + auto existingValue1 = vld1q_f32(outputAccum + CPOS + 4); + + auto existingValue = elementaryBlock2x4bAccum( + pqFineCentroids + + ((fineCentroidIdx + 0) * FINE_TABLE_BYTES + + fineCode0a) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids + + ((fineCentroidIdx + 1) * FINE_TABLE_BYTES + + fineCode0b) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids + + ((fineCentroidIdx + 2) * FINE_TABLE_BYTES + + fineCode0c) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids + + ((fineCentroidIdx + 3) * FINE_TABLE_BYTES + + fineCode0d) * + FINE_SIZE + + fineCentroidOffset, + weight0, + {existingValue0, existingValue1}); + + existingValue = elementaryBlock2x4bAccum( + pqFineCentroids + + ((fineCentroidIdx + 0) * FINE_TABLE_BYTES + + fineCode1a) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids + + ((fineCentroidIdx + 1) * FINE_TABLE_BYTES + + fineCode1b) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids + + ((fineCentroidIdx + 2) * FINE_TABLE_BYTES + + fineCode1c) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids + + ((fineCentroidIdx + 3) * FINE_TABLE_BYTES + + fineCode1d) * + FINE_SIZE + + fineCentroidOffset, + weight1, + existingValue); + + vst1q_f32(outputAccum + CPOS, existingValue.val[0]); + vst1q_f32(outputAccum + CPOS + 4, existingValue.val[1]); + + // next + IndexPQDecoderImpl::accum( + pqFineCentroids, code0, weight0, code1, weight1, outputAccum); + } + + // Process 3 samples. + // Each code uses its own fine pq centroids table. + static void accum( + const float* const __restrict pqFineCentroids0, + const uint8_t* const __restrict code0, + const float weight0, + const float* const __restrict pqFineCentroids1, + const uint8_t* const __restrict code1, + const float weight1, + const float* const __restrict pqFineCentroids2, + const uint8_t* const __restrict code2, + const float weight2, + float* const __restrict outputAccum) { + // fine quantizer + const uint8_t* const __restrict fine0 = code0; + const uint8_t* const __restrict fine1 = code1; + const uint8_t* const __restrict fine2 = code2; + + // process chunks, 2 float + // but 8 floats per loop + + const intptr_t fineCode0a = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0b = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0c = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0d = detail:: + UintReader::get( + fine0); + const intptr_t fineCode1a = detail:: + UintReader::get( + fine1); + const intptr_t fineCode1b = detail:: + UintReader::get( + fine1); + const intptr_t fineCode1c = detail:: + UintReader::get( + fine1); + const intptr_t fineCode1d = detail:: + UintReader::get( + fine1); + const intptr_t fineCode2a = detail:: + UintReader::get( + fine2); + const intptr_t fineCode2b = detail:: + UintReader::get( + fine2); + const intptr_t fineCode2c = detail:: + UintReader::get( + fine2); + const intptr_t fineCode2d = detail:: + UintReader::get( + fine2); + + auto existingValue0 = vld1q_f32(outputAccum + CPOS); + auto existingValue1 = vld1q_f32(outputAccum + CPOS + 4); + + auto existingValue = elementaryBlock2x4bAccum( + pqFineCentroids0 + + ((fineCentroidIdx + 0) * FINE_TABLE_BYTES + + fineCode0a) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids0 + + ((fineCentroidIdx + 1) * FINE_TABLE_BYTES + + fineCode0b) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids0 + + ((fineCentroidIdx + 2) * FINE_TABLE_BYTES + + fineCode0c) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids0 + + ((fineCentroidIdx + 3) * FINE_TABLE_BYTES + + fineCode0d) * + FINE_SIZE + + fineCentroidOffset, + weight0, + {existingValue0, existingValue1}); + + existingValue = elementaryBlock2x4bAccum( + pqFineCentroids1 + + ((fineCentroidIdx + 0) * FINE_TABLE_BYTES + + fineCode1a) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids1 + + ((fineCentroidIdx + 1) * FINE_TABLE_BYTES + + fineCode1b) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids1 + + ((fineCentroidIdx + 2) * FINE_TABLE_BYTES + + fineCode1c) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids1 + + ((fineCentroidIdx + 3) * FINE_TABLE_BYTES + + fineCode1d) * + FINE_SIZE + + fineCentroidOffset, + weight1, + existingValue); + + existingValue = elementaryBlock2x4bAccum( + pqFineCentroids2 + + ((fineCentroidIdx + 0) * FINE_TABLE_BYTES + + fineCode2a) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids2 + + ((fineCentroidIdx + 1) * FINE_TABLE_BYTES + + fineCode2b) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids2 + + ((fineCentroidIdx + 2) * FINE_TABLE_BYTES + + fineCode2c) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids2 + + ((fineCentroidIdx + 3) * FINE_TABLE_BYTES + + fineCode2d) * + FINE_SIZE + + fineCentroidOffset, + weight2, + existingValue); + + vst1q_f32(outputAccum + CPOS, existingValue.val[0]); + vst1q_f32(outputAccum + CPOS + 4, existingValue.val[1]); + + // next + IndexPQDecoderImpl::accum( + pqFineCentroids0, + code0, + weight0, + pqFineCentroids1, + code1, + weight1, + pqFineCentroids2, + code2, + weight2, + outputAccum); + } + + // Process 3 samples. + // Fine pq centroids table is shared among codes. + static void accum( + const float* const __restrict pqFineCentroids, + const uint8_t* const __restrict code0, + const float weight0, + const uint8_t* const __restrict code1, + const float weight1, + const uint8_t* const __restrict code2, + const float weight2, + float* const __restrict outputAccum) { + // fine quantizer + const uint8_t* const __restrict fine0 = code0; + const uint8_t* const __restrict fine1 = code1; + const uint8_t* const __restrict fine2 = code2; + + // process chunks, 2 float + // but 8 floats per loop + + const intptr_t fineCode0a = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0b = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0c = detail:: + UintReader::get( + fine0); + const intptr_t fineCode0d = detail:: + UintReader::get( + fine0); + const intptr_t fineCode1a = detail:: + UintReader::get( + fine1); + const intptr_t fineCode1b = detail:: + UintReader::get( + fine1); + const intptr_t fineCode1c = detail:: + UintReader::get( + fine1); + const intptr_t fineCode1d = detail:: + UintReader::get( + fine1); + const intptr_t fineCode2a = detail:: + UintReader::get( + fine2); + const intptr_t fineCode2b = detail:: + UintReader::get( + fine2); + const intptr_t fineCode2c = detail:: + UintReader::get( + fine2); + const intptr_t fineCode2d = detail:: + UintReader::get( + fine2); + + auto existingValue0 = vld1q_f32(outputAccum + CPOS); + auto existingValue1 = vld1q_f32(outputAccum + CPOS + 4); + + auto existingValue = elementaryBlock2x4bAccum( + pqFineCentroids + + ((fineCentroidIdx + 0) * FINE_TABLE_BYTES + + fineCode0a) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids + + ((fineCentroidIdx + 1) * FINE_TABLE_BYTES + + fineCode0b) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids + + ((fineCentroidIdx + 2) * FINE_TABLE_BYTES + + fineCode0c) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids + + ((fineCentroidIdx + 3) * FINE_TABLE_BYTES + + fineCode0d) * + FINE_SIZE + + fineCentroidOffset, + weight0, + {existingValue0, existingValue1}); + + existingValue = elementaryBlock2x4bAccum( + pqFineCentroids + + ((fineCentroidIdx + 0) * FINE_TABLE_BYTES + + fineCode1a) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids + + ((fineCentroidIdx + 1) * FINE_TABLE_BYTES + + fineCode1b) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids + + ((fineCentroidIdx + 2) * FINE_TABLE_BYTES + + fineCode1c) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids + + ((fineCentroidIdx + 3) * FINE_TABLE_BYTES + + fineCode1d) * + FINE_SIZE + + fineCentroidOffset, + weight1, + existingValue); + + existingValue = elementaryBlock2x4bAccum( + pqFineCentroids + + ((fineCentroidIdx + 0) * FINE_TABLE_BYTES + + fineCode2a) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids + + ((fineCentroidIdx + 1) * FINE_TABLE_BYTES + + fineCode2b) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids + + ((fineCentroidIdx + 2) * FINE_TABLE_BYTES + + fineCode2c) * + FINE_SIZE + + fineCentroidOffset, + pqFineCentroids + + ((fineCentroidIdx + 3) * FINE_TABLE_BYTES + + fineCode2d) * + FINE_SIZE + + fineCentroidOffset, + weight2, + existingValue); + + vst1q_f32(outputAccum + CPOS, existingValue.val[0]); + vst1q_f32(outputAccum + CPOS + 4, existingValue.val[1]); + + // next + IndexPQDecoderImpl::accum( + pqFineCentroids, + code0, + weight0, + code1, + weight1, + code2, + weight2, + outputAccum); + } +}; + template < intptr_t DIM, intptr_t CPOS, @@ -163,6 +850,7 @@ struct IndexPQDecoderImpl< 4, FINE_BITS, CPOS, + false, true, QPOS_LEFT_GE_8, QPOS_LEFT_GE_4, @@ -609,6 +1297,7 @@ struct IndexPQDecoderImpl< FINE_BITS, CPOS, false, + false, true, true, false> { @@ -938,6 +1627,7 @@ struct IndexPQDecoderImpl< CPOS, false, false, + false, true, false> { static constexpr intptr_t fineCentroidIdx = CPOS / FINE_SIZE; @@ -1252,6 +1942,7 @@ template < intptr_t DIM, intptr_t FINE_SIZE, intptr_t FINE_BITS, + bool FINE_SIZE_EQ_2, bool FINE_SIZE_EQ_4, bool QPOS_LEFT_GE_8, bool QPOS_LEFT_GE_4> @@ -1260,6 +1951,7 @@ struct IndexPQDecoderImpl< FINE_SIZE, FINE_BITS, DIM, + FINE_SIZE_EQ_2, FINE_SIZE_EQ_4, QPOS_LEFT_GE_8, QPOS_LEFT_GE_4,