RaBitQ: broadcast one sign word per 32 dims in the AVX2 RBQ9 kernel (#5696)

Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5696

`compute_inner_product<SIMDLevel::AVX2>` scores 9-bit RaBitQ codes (`ex_bits == 8`) in groups of 8 dimensions. Before this change, each group loaded one sign byte and broadcast it. The broadcast uses the shuffle port, which limits the loop.

This change makes these changes to the loop:

- The loop loads one 32-bit sign word and broadcasts it once for each 32 dimensions.
- Each group of 8 dimensions ANDs the word with `positions << 8k` and compares the result with `cmpeq`. A `cmpgt` against zero is not correct here, because bit 31 of the shifted positions is a sign bit.
- The loop uses two accumulators.
- An 8-dimension loop handles the remaining full groups. `ip_scalar` handles the tail, as before.

The two accumulators change the float summation order. The benchmark checksums agree to 6 significant digits.

AI explanation of why:
```
  What limits a loop like this

  A loop runs no faster than its tightest limit. This loop has three possible limits:
  1. Micro-op count. The core issues about 4 micro-ops per cycle.
  2. Port 5, the shuffle port. Broadcasts and byte-to-int expansions can run only there.
  3. The FMA chain. With one accumulator, each FMA must wait for the previous one. Each FMA takes 4 cycles, so one group of 8 dims takes at least 4 cycles.

  Measured values, per group of 8 dims (d=768, codes in L2)

  ┌────────────────────────────────────────┬────────┬──────────────┬─────────────┬─────────────────┐
  │                  loop                  │ cycles │ instructions │ port 5 µops │ FMA chain floor │
  ├────────────────────────────────────────┼────────┼──────────────┼─────────────┼─────────────────┤
  │ V1 (contributor), 1 accumulator        │ 5.53   │ 18.4         │ 4.12        │ 4.0             │
  ├────────────────────────────────────────┼────────┼──────────────┼─────────────┼─────────────────┤
  │ V1, 2 accumulators                     │ 5.63   │ 15.4         │ 4.78        │ 2.0             │
  ├────────────────────────────────────────┼────────┼──────────────┼─────────────┼─────────────────┤
  │ sign word, 1 accumulator               │ 4.43   │ 10.4         │ 2.95        │ 4.0             │
  ├────────────────────────────────────────┼────────┼──────────────┼─────────────┼─────────────────┤
  │ sign word, 2 accumulators (D122893433) │ 3.77   │ 10.5         │ 2.99        │ 2.0             │
  └────────────────────────────────────────┴────────┴──────────────┴─────────────┴─────────────────┘

  With 4,096 codes, which do not fit in L2, the values are 6.20, 5.96, 5.15 and 4.66 cycles. The new loop is 1.33x faster than V1 there, which agrees with the
  faiss benchmark (1.33x–1.36x).

  Step 1: the sign-word broadcast removes micro-ops

  In V1, each group of 8 dims gets its sign byte in four instructions:
  movzx  eax, BYTE PTR [sign + i/8]   ; load one byte
  vmovd  xmm1, eax                    ; copy it into a vector register (port 5)
  vpbroadcastd ymm1, xmm1             ; copy it into all 8 lanes (port 5)
  The loop also does its index and loop-control work (shr, mov, add, cmp, jae) once per 8 dims. In total, each group needs 18 instructions.

  The new loop loads a 32-bit sign word once per 32 dims. The compiler turns the memcpy plus _mm256_set1_epi32 into one instruction:
  vpbroadcastd ymm0, DWORD PTR [sign + i/8]
  A broadcast that reads directly from memory runs on the load port only. It does not use port 5. Each of the four groups then needs only vpand and vpcmpeqd
  against its shifted position constant. The compiler loads these constants before the loop, so they cost nothing per pass. The loop-control work also runs once
  per 32 dims, not once per 8.

  Result: instructions per group go from 18.4 to 10.4, and port 5 micro-ops from 4.12 to 2.95. Cycles go from 5.53 to 4.43, which is 1.25x.

  Step 2: the second accumulator shortens the FMA chain

  After step 1, the micro-op limit is about 2.6 cycles and the port 5 limit is about 3.0 cycles. But one accumulator still forces 4 cycles per group. The measured
  4.43 cycles is close to that floor, so the FMA chain is now the limit.

  With two accumulators, two chains run at the same time, and the floor falls to 2.0 cycles. The next limit is port 5 at about 3 cycles. Cycles go from 4.43 to
  3.77, which is 1.18x.

  Why the second accumulator alone did not help V1

  V1 needs 5.53 cycles because of its micro-op count and port 5, not because of the FMA chain: 5.53 is well above the 4-cycle floor. A second accumulator removes a
  limit that V1 does not reach. My earlier standalone test found the same result. Thus each change helps only together with the other:
  - The broadcast change brings the loop down to the FMA floor.
  - The second accumulator then lowers the floor.

  The cmpeq change

  cmpeq is not faster than cmpgt; both cost the same. It is necessary for correctness: the group that tests bits 24–31 uses 128 << 24 = 0x80000000. As a signed
  integer this value is negative, so cmpgt(x, 0) would give the wrong result for that lane.
```

Reviewed By: junjieqi

Differential Revision: D122893433

fbshipit-source-id: e6c42d33e138614589e8f80c059df3652a7ddfcd

faiss/utils/simd_impl/rabitq_avx2.cpp

@@ -590,29 +590,42 @@ float compute_inner_product<SIMDLevel::AVX2>(
         float cb) {
     if (ex_bits == 8) {
         // Eight byte-aligned extra codes, plus their independent sign bits.
-        __m256 acc = _mm256_setzero_ps();
+        // One sign-word broadcast serves 32 dims, to save shuffle-port uops.
+        __m256 acc0 = _mm256_setzero_ps();
+        __m256 acc1 = _mm256_setzero_ps();
         const __m256 weight = _mm256_set1_ps(256.f);
         const __m256 offset = _mm256_set1_ps(cb);
         const __m256i positions =
                 _mm256_setr_epi32(1, 2, 4, 8, 16, 32, 64, 128);
-        size_t i = 0;
-        for (; i + 8 <= d; i += 8) {
+        // cmpeq, not cmpgt: bit 31 of the shifted positions is a sign bit.
+        auto step = [&](__m256& acc, size_t i, __m256i signs, __m256i pos) {
             const __m128i bytes = _mm_loadl_epi64(
                     reinterpret_cast<const __m128i*>(ex_code + i));
             const __m256 extra =
                     _mm256_cvtepi32_ps(_mm256_cvtepu8_epi32(bytes));
-            const __m256i mask = _mm256_cmpgt_epi32(
-                    _mm256_and_si256(
-                            _mm256_set1_epi32(sign_bits[i / 8]), positions),
-                    _mm256_setzero_si256());
+            const __m256i mask =
+                    _mm256_cmpeq_epi32(_mm256_and_si256(signs, pos), pos);
             const __m256 recon = _mm256_add_ps(
                     extra, _mm256_and_ps(_mm256_castsi256_ps(mask), weight));
             acc = _mm256_fmadd_ps(
                     _mm256_loadu_ps(rotated_q + i),
                     _mm256_add_ps(recon, offset),
                     acc);
+        };
+        size_t i = 0;
+        for (; i + 32 <= d; i += 32) {
+            uint32_t word;
+            memcpy(&word, sign_bits + i / 8, sizeof(word));
+            const __m256i signs = _mm256_set1_epi32(static_cast<int>(word));
+            step(acc0, i, signs, positions);
+            step(acc1, i + 8, signs, _mm256_slli_epi32(positions, 8));
+            step(acc0, i + 16, signs, _mm256_slli_epi32(positions, 16));
+            step(acc1, i + 24, signs, _mm256_slli_epi32(positions, 24));
         }
-        return hsum_avx2(acc) +
+        for (; i + 8 <= d; i += 8) {
+            step(acc0, i, _mm256_set1_epi32(sign_bits[i / 8]), positions);
+        }
+        return hsum_avx2(_mm256_add_ps(acc0, acc1)) +
                 ip_scalar(sign_bits, ex_code, rotated_q, i, d, ex_bits, cb);
     }
     if (ex_bits == 1) {