block size, number of dimensions of each vector
| 726 | |
| 727 | template<int B, int D> // block size, number of dimensions of each vector |
| 728 | inline void float_dists_vertical(const float* X, const float* q, |
| 729 | float* dists_out, int64_t N) |
| 730 | { |
| 731 | static constexpr int packet_width = 8; // how many vecs we operate on at once |
| 732 | static constexpr int nstripes = B / packet_width; // # of rows of 32B per block |
| 733 | static_assert(B % packet_width == 0, "B must be a multiple of packet_width"); |
| 734 | static_assert(B > 0, "B must be > 0"); |
| 735 | const int64_t nblocks = N / B; |
| 736 | assert(N % B == 0); |
| 737 | |
| 738 | __m256 accumulators[nstripes]; |
| 739 | |
| 740 | for (int64_t b = 0; b < nblocks; b++) { // for each block |
| 741 | for (int i = 0; i < nstripes; i++) { |
| 742 | accumulators[i] = _mm256_setzero_ps(); |
| 743 | } |
| 744 | |
| 745 | for (int j = 0; j < D; j++) { // for each dimension |
| 746 | auto q_broadcast = _mm256_set1_ps(q[j]); |
| 747 | for (int i = 0; i < nstripes; i++) { // for each stripe |
| 748 | auto x_col = _mm256_load_ps(X); |
| 749 | X += packet_width; |
| 750 | |
| 751 | auto diff = _mm256_sub_ps(q_broadcast, x_col); |
| 752 | accumulators[i] = fma(diff, diff, accumulators[i]); |
| 753 | // auto prods = fma(diff, diff, accumulators[i]); |
| 754 | // accumulators[i] = _mm256_add_ps(accumulators[i], prods); |
| 755 | } |
| 756 | } |
| 757 | for (uint8_t i = 0; i < nstripes; i++) { // for each stripe |
| 758 | _mm256_store_ps(dists_out, accumulators[i]); |
| 759 | dists_out += packet_width; |
| 760 | } |
| 761 | } |
| 762 | } |
| 763 | |
| 764 | template<int B, int D> // block size, number of dimensions of each vector |
| 765 | inline void byte_dists_vertical(const uint8_t* X, const uint8_t* q, |