From d1e0c7b88e48f5b9eb4bbb39b80eeaa20d0a7e90 Mon Sep 17 00:00:00 2001 From: rayx Date: Fri, 13 Mar 2026 13:34:28 +0800 Subject: [PATCH] fix: use static array (#222) * fix: use static array * fix: clang-format * fix: clang-format * fix: clang-format --- ...ner_product_distance_batch_impl_fp16_avx2.cc | 9 ++++----- ...r_product_distance_batch_impl_fp16_avx512.cc | 16 +++++++--------- ...ner_product_distance_batch_impl_fp32_avx2.cc | 17 +++++++++++------ ...ner_product_distance_batch_impl_int8_avx2.cc | 10 ++++++---- ...r_product_distance_batch_impl_int8_avx512.cc | 9 ++++++--- 5 files changed, 34 insertions(+), 27 deletions(-) diff --git a/src/ailego/math_batch/inner_product_distance_batch_impl_fp16_avx2.cc b/src/ailego/math_batch/inner_product_distance_batch_impl_fp16_avx2.cc index 59320de..d6fe475 100644 --- a/src/ailego/math_batch/inner_product_distance_batch_impl_fp16_avx2.cc +++ b/src/ailego/math_batch/inner_product_distance_batch_impl_fp16_avx2.cc @@ -28,8 +28,7 @@ compute_one_to_many_inner_product_avx2_fp16( const ailego::Float16 *query, const ailego::Float16 **ptrs, std::array &prefetch_ptrs, size_t dimensionality, float *results) { - std::array<__m256, dp_batch> accs; - + __m256 accs[dp_batch]; for (size_t i = 0; i < dp_batch; ++i) { accs[i] = _mm256_setzero_ps(); } @@ -42,8 +41,8 @@ compute_one_to_many_inner_product_avx2_fp16( __m256 q1 = _mm256_cvtph_ps(_mm256_castsi256_si128(q)); __m256 q2 = _mm256_cvtph_ps(_mm256_extractf128_si256(q, 1)); - std::array<__m256, dp_batch> data_regs_1; - std::array<__m256, dp_batch> data_regs_2; + __m256 data_regs_1[dp_batch]; + __m256 data_regs_2[dp_batch]; for (size_t i = 0; i < dp_batch; ++i) { __m256i m = _mm256_loadu_si256(reinterpret_cast(ptrs[i] + dim)); @@ -68,7 +67,7 @@ compute_one_to_many_inner_product_avx2_fp16( __m256 q = _mm256_cvtph_ps( _mm_loadu_si128(reinterpret_cast(query + dim))); - std::array<__m256, dp_batch> data_regs; + __m256 data_regs[dp_batch]; for (size_t i = 0; i < dp_batch; ++i) { data_regs[i] = _mm256_cvtph_ps( _mm_loadu_si128(reinterpret_cast(ptrs[i] + dim))); diff --git a/src/ailego/math_batch/inner_product_distance_batch_impl_fp16_avx512.cc b/src/ailego/math_batch/inner_product_distance_batch_impl_fp16_avx512.cc index 1fbe5b2..e06820e 100644 --- a/src/ailego/math_batch/inner_product_distance_batch_impl_fp16_avx512.cc +++ b/src/ailego/math_batch/inner_product_distance_batch_impl_fp16_avx512.cc @@ -27,8 +27,7 @@ compute_one_to_many_inner_product_avx512fp16_fp16( const ailego::Float16 *query, const ailego::Float16 **ptrs, std::array &prefetch_ptrs, size_t dimensionality, float *results) { - std::array<__m512h, dp_batch> accs; - + __m512h accs[dp_batch]; for (size_t i = 0; i < dp_batch; ++i) { accs[i] = _mm512_setzero_ph(); } @@ -37,7 +36,7 @@ compute_one_to_many_inner_product_avx512fp16_fp16( for (; dim + 32 <= dimensionality; dim += 32) { __m512h q = _mm512_loadu_ph(query + dim); - std::array<__m512h, dp_batch> data_regs; + __m512h data_regs[dp_batch]; for (size_t i = 0; i < dp_batch; ++i) { data_regs[i] = _mm512_loadu_ph(ptrs[i] + dim); } @@ -83,8 +82,7 @@ compute_one_to_many_inner_product_avx512f_fp16( const ailego::Float16 *query, const ailego::Float16 **ptrs, std::array &prefetch_ptrs, size_t dimensionality, float *results) { - std::array<__m512, dp_batch> accs; - + __m512 accs[dp_batch]; for (size_t i = 0; i < dp_batch; ++i) { accs[i] = _mm512_setzero_ps(); } @@ -97,8 +95,8 @@ compute_one_to_many_inner_product_avx512f_fp16( __m512 q1 = _mm512_cvtph_ps(_mm512_castsi512_si256(q)); __m512 q2 = _mm512_cvtph_ps(_mm512_extracti64x4_epi64(q, 1)); - std::array<__m512, dp_batch> data_regs_1; - std::array<__m512, dp_batch> data_regs_2; + __m512 data_regs_1[dp_batch]; + __m512 data_regs_2[dp_batch]; for (size_t i = 0; i < dp_batch; ++i) { __m512i m = _mm512_loadu_si512(reinterpret_cast(ptrs[i] + dim)); @@ -123,7 +121,7 @@ compute_one_to_many_inner_product_avx512f_fp16( __m512 q = _mm512_cvtph_ps( _mm256_loadu_si256(reinterpret_cast(query + dim))); - std::array<__m512, dp_batch> data_regs; + __m512 data_regs[dp_batch]; for (size_t i = 0; i < dp_batch; ++i) { data_regs[i] = _mm512_cvtph_ps( _mm256_loadu_si256(reinterpret_cast(ptrs[i] + dim))); @@ -133,7 +131,7 @@ compute_one_to_many_inner_product_avx512f_fp16( dim += 16; } - std::array<__m256, dp_batch> acc_new; + __m256 acc_new[dp_batch]; for (size_t i = 0; i < dp_batch; ++i) { acc_new[i] = _mm256_add_ps( _mm512_castps512_ps256(accs[i]), diff --git a/src/ailego/math_batch/inner_product_distance_batch_impl_fp32_avx2.cc b/src/ailego/math_batch/inner_product_distance_batch_impl_fp32_avx2.cc index 0e54064..ffda66e 100644 --- a/src/ailego/math_batch/inner_product_distance_batch_impl_fp32_avx2.cc +++ b/src/ailego/math_batch/inner_product_distance_batch_impl_fp32_avx2.cc @@ -39,14 +39,15 @@ compute_one_to_many_inner_product_avx2_fp32( const ValueType *query, const ValueType **ptrs, std::array &prefetch_ptrs, size_t dimensionality, float *results) { - std::array<__m256, dp_batch> accs; + __m256 accs[dp_batch]; for (size_t i = 0; i < dp_batch; ++i) { accs[i] = _mm256_setzero_ps(); } size_t dim = 0; for (; dim + 8 <= dimensionality; dim += 8) { __m256 q = _mm256_loadu_ps(query + dim); - std::array<__m256, dp_batch> data_regs; + + __m256 data_regs[dp_batch]; for (size_t i = 0; i < dp_batch; ++i) { data_regs[i] = _mm256_loadu_ps(ptrs[i] + dim); } @@ -59,13 +60,15 @@ compute_one_to_many_inner_product_avx2_fp32( accs[i] = _mm256_fnmadd_ps(q, data_regs[i], accs[i]); } } - std::array<__m128, dp_batch> sum128_regs; + + __m128 sum128_regs[dp_batch]; for (size_t i = 0; i < dp_batch; ++i) { sum128_regs[i] = sum_top_bottom_avx(accs[i]); } if (dim + 4 <= dimensionality) { __m128 q = _mm_loadu_ps(query + dim); - std::array<__m128, dp_batch> data_regs; + + __m128 data_regs[dp_batch]; for (size_t i = 0; i < dp_batch; ++i) { data_regs[i] = _mm_loadu_ps(ptrs[i] + dim); } @@ -81,7 +84,8 @@ compute_one_to_many_inner_product_avx2_fp32( } if (dim + 2 <= dimensionality) { __m128 q = _mm_setzero_ps(); - std::array<__m128, dp_batch> data_regs; + + __m128 data_regs[dp_batch]; for (size_t i = 0; i < dp_batch; ++i) { data_regs[i] = _mm_setzero_ps(); } @@ -95,7 +99,8 @@ compute_one_to_many_inner_product_avx2_fp32( } dim += 2; } - std::array res; + + float res[dp_batch]; for (size_t i = 0; i < dp_batch; ++i) { res[i] = sum4(sum128_regs[i]); } diff --git a/src/ailego/math_batch/inner_product_distance_batch_impl_int8_avx2.cc b/src/ailego/math_batch/inner_product_distance_batch_impl_int8_avx2.cc index 23d3566..66d7e15 100644 --- a/src/ailego/math_batch/inner_product_distance_batch_impl_int8_avx2.cc +++ b/src/ailego/math_batch/inner_product_distance_batch_impl_int8_avx2.cc @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -#include +#include #include #include #include @@ -27,14 +27,15 @@ compute_one_to_many_inner_product_avx2_int8( const int8_t *query, const int8_t **ptrs, std::array &prefetch_ptrs, size_t dimensionality, float *results) { - std::vector<__m256i> accs(dp_batch); + __m256i accs[dp_batch]; for (size_t i = 0; i < dp_batch; ++i) { accs[i] = _mm256_setzero_si256(); } size_t dim = 0; for (; dim + 32 <= dimensionality; dim += 32) { __m256i q = _mm256_loadu_si256((const __m256i *)(query + dim)); - std::vector<__m256i> data_regs(dp_batch); + + __m256i data_regs[dp_batch]; for (size_t i = 0; i < dp_batch; ++i) { data_regs[i] = _mm256_loadu_si256((const __m256i *)(ptrs[i] + dim)); } @@ -63,7 +64,8 @@ compute_one_to_many_inner_product_avx2_int8( _mm256_add_epi32(accs[i], _mm256_add_epi32(prod_lo[i], prod_hi[i])); } } - std::array temp_results; + + int temp_results[dp_batch]; for (size_t i = 0; i < dp_batch; ++i) { __m128i lo = _mm256_castsi256_si128(accs[i]); __m128i hi = _mm256_extracti128_si256(accs[i], 1); diff --git a/src/ailego/math_batch/inner_product_distance_batch_impl_int8_avx512.cc b/src/ailego/math_batch/inner_product_distance_batch_impl_int8_avx512.cc index 1e10583..2caf83c 100644 --- a/src/ailego/math_batch/inner_product_distance_batch_impl_int8_avx512.cc +++ b/src/ailego/math_batch/inner_product_distance_batch_impl_int8_avx512.cc @@ -52,7 +52,7 @@ static void compute_one_to_many_inner_product_avx512_vnni_int8( const int8_t *query, const int8_t **ptrs, std::array &prefetch_ptrs, size_t dimensionality, float *results) { - std::array<__m512i, dp_batch> accs; + __m512i accs[dp_batch]; for (size_t i = 0; i < dp_batch; ++i) { accs[i] = _mm512_setzero_si512(); } @@ -60,7 +60,8 @@ static void compute_one_to_many_inner_product_avx512_vnni_int8( for (; dim + 64 <= dimensionality; dim += 64) { __m512i q = _mm512_loadu_si512(reinterpret_cast(query + dim)); - std::array<__m512i, dp_batch> data_regs; + + __m512i data_regs[dp_batch]; for (size_t i = 0; i < dp_batch; ++i) { data_regs[i] = _mm512_loadu_si512(reinterpret_cast(ptrs[i] + dim)); @@ -74,7 +75,8 @@ static void compute_one_to_many_inner_product_avx512_vnni_int8( accs[i] = _mm512_dpbusd_epi32(accs[i], q, data_regs[i]); } } - std::array temp_results{}; + + int temp_results[dp_batch]{}; for (size_t i = 0; i < dp_batch; ++i) { temp_results[i] = _mm512_reduce_add_epi32(accs[i]); } @@ -88,6 +90,7 @@ static void compute_one_to_many_inner_product_avx512_vnni_int8( results[i] = static_cast(temp_results[i]); } } + // // #elif defined(__AVX512BW__) // // TODO: this version is problematic