pqCodeDistances read out of bounds of pqCentroids in loadingThreads subset

[XuanYang-cn]

I'm recently very interested in faiss, and run some tests on faiss gpu IVFPQ for 128-dimensional vectors, with M=32, and 256 centroids. Here are some questions of mine.

template <typename OutCodeT,
          typename CentroidT,
          int DimsPerSubQuantizer,
          bool L2Distance>
__global__ void
__launch_bounds__(288, 3)
pqCodeDistances(Tensor<float, 2, true> queries,
                int queriesPerBlock,
                Tensor<CentroidT, 2, true> coarseCentroids,
...
  bool isLoadingThread = threadIdx.x >= codesPerSubQuantizer;
  int loadingThreadId = threadIdx.x - codesPerSubQuantizer;
...
  float subQuantizerData[DimsPerSubQuantizer];
  auto code = threadIdx.x;
...
#pragma unroll
  for (int i = 0; i < DimsPerSubQuantizer; ++i) {
    subQuantizerData[i] = pqCentroids[subQuantizer][i][code].ldg();
  }

faiss/faiss/gpu/impl/PQCodeDistances-inl.cuh

Line 51 in d8af513

auto code = threadIdx.x;

faiss/faiss/gpu/impl/PQCodeDistances-inl.cuh

Line 58 in d8af513

subQuantizerData[i] = pqCentroids[subQuantizer][i][code].ldg();

In the test, there are 288 threads and 256 centroids, which means code=288 at line 51. What confuses me most is why line 58 doesn't go wrong when pqCentroids(32x4x256) trying to obtain data out of the range (when code >= 256), and wouldn't it go wrong at some point?

The 2nd question is, why do Loading Threads need subQuantizerData since they don't process subQuantizerData at all.

The 3rd question is, would it be better had line 56-59 been moved to line 155? why or why not?

[wickedfoo]

1. re: line 51, yes, this seems to be a bug, there should be a bounds check here. This would typically not crash because the allocation is within the temporary memory allocation buffer managed by StandardGpuResources, the memory of which is temporarily allocated here in PQScanMultiPassNoPrecomputed-inl.cuh line 593, so any read out of bounds would still hit allocated memory but it will be garbage. In the case where we allocate from overflow memory via cudaMalloc this could crash. I will fix this, thank you very much for finding this issue!

2. The loading threads are a subset of the block threads to allow for double buffering of data. All threads in the block have the same register set, there is no way to segment this within a block. subQuantizerData is stored in registers.

3. It would not be better because it would be reloaded upon each iteration of the loop at line 82 and 125. It only needs to be loaded once.

[wickedfoo]

Fixed in 9b007c7