Switch //faiss/gpu to use templates instead of macros (#2914)

Summary: Pull Request resolved: #2914 The macros are part of a system to reduce compilation time via separate compilation units. Unfortunately, the parallelization is across C++ template functions instead of NVCC invocations on kernel compilation, which would be much more effective. This diff removes the preprocessor macros and expands them into templates. Compilation time after this diff is given by [this buck2 output](https://www.internalfb.com/buck2/ae9e6b28-a1bd-4d46-8af8-2895e6f182c8) with 1,043s through impl/scan/IVFInterleaved2048.cu Differential Revision: D46549341 fbshipit-source-id: fc423bbf9b53a228ef6fbe2c28d395aef92289d7
facebookresearch · Jun 13, 2023 · 4ea7b36 · 4ea7b36
1 parent 6951466
commit 4ea7b36
Show file tree

Hide file tree

Showing 12 changed files with 393 additions and 288 deletions.
diff --git a/faiss/gpu/CMakeLists.txt b/faiss/gpu/CMakeLists.txt
@@ -135,6 +135,7 @@ set(FAISS_GPU_HEADERS
   impl/RemapIndices.h
   impl/VectorResidual.cuh
   impl/scan/IVFInterleavedImpl.cuh
+  impl/scan/IVFInterleavedKernel.cuh
   impl/IcmEncoder.cuh
   utils/BlockSelectKernel.cuh
   utils/Comparators.cuh

diff --git a/faiss/gpu/impl/IVFInterleaved.cu b/faiss/gpu/impl/IVFInterleaved.cu
@@ -192,24 +192,43 @@ void runIVFInterleavedScan(
     // caught for exceptions at a higher level
     FAISS_ASSERT(k <= GPU_MAX_SELECTION_K);
 
+    const auto ivf_interleaved_call = [&](const auto func) {
+        func(queries,
+             listIds,
+             listData,
+             listIndices,
+             indicesOptions,
+             listLengths,
+             k,
+             metric,
+             useResidual,
+             residualBase,
+             scalarQ,
+             outDistances,
+             outIndices,
+             res);
+    };
+
     if (k == 1) {
-        IVF_INTERLEAVED_CALL(1);
+        ivf_interleaved_call(ivfInterleavedScanImpl<IVFINTERLEAVED_1_PARAMS>);
     } else if (k <= 32) {
-        IVF_INTERLEAVED_CALL(32);
+        ivf_interleaved_call(ivfInterleavedScanImpl<IVFINTERLEAVED_32_PARAMS>);
     } else if (k <= 64) {
-        IVF_INTERLEAVED_CALL(64);
+        ivf_interleaved_call(ivfInterleavedScanImpl<IVFINTERLEAVED_64_PARAMS>);
     } else if (k <= 128) {
-        IVF_INTERLEAVED_CALL(128);
+        ivf_interleaved_call(ivfInterleavedScanImpl<IVFINTERLEAVED_128_PARAMS>);
     } else if (k <= 256) {
-        IVF_INTERLEAVED_CALL(256);
+        ivf_interleaved_call(ivfInterleavedScanImpl<IVFINTERLEAVED_256_PARAMS>);
     } else if (k <= 512) {
-        IVF_INTERLEAVED_CALL(512);
+        ivf_interleaved_call(ivfInterleavedScanImpl<IVFINTERLEAVED_512_PARAMS>);
     } else if (k <= 1024) {
-        IVF_INTERLEAVED_CALL(1024);
+        ivf_interleaved_call(
+                ivfInterleavedScanImpl<IVFINTERLEAVED_1024_PARAMS>);
     }
 #if GPU_MAX_SELECTION_K >= 2048
     else if (k <= 2048) {
-        IVF_INTERLEAVED_CALL(2048);
+        ivf_interleaved_call(
+                ivfInterleavedScanImpl<IVFINTERLEAVED_2048_PARAMS>);
     }
 #endif
 }

diff --git a/faiss/gpu/impl/IVFInterleaved.cuh b/faiss/gpu/impl/IVFInterleaved.cuh
diff --git a/faiss/gpu/impl/scan/IVFInterleaved1.cu b/faiss/gpu/impl/scan/IVFInterleaved1.cu
@@ -10,7 +10,7 @@
 namespace faiss {
 namespace gpu {
 
-IVF_INTERLEAVED_IMPL(128, 1, 1)
+IVF_INTERLEAVED_IMPL(IVFINTERLEAVED_1_PARAMS)
 
 }
 } // namespace faiss
diff --git a/faiss/gpu/impl/scan/IVFInterleaved1024.cu b/faiss/gpu/impl/scan/IVFInterleaved1024.cu
@@ -10,7 +10,7 @@
 namespace faiss {
 namespace gpu {
 
-IVF_INTERLEAVED_IMPL(128, 1024, 8)
+IVF_INTERLEAVED_IMPL(IVFINTERLEAVED_1024_PARAMS)
 
 }
 } // namespace faiss
diff --git a/faiss/gpu/impl/scan/IVFInterleaved128.cu b/faiss/gpu/impl/scan/IVFInterleaved128.cu
@@ -10,7 +10,7 @@
 namespace faiss {
 namespace gpu {
 
-IVF_INTERLEAVED_IMPL(128, 128, 3)
+IVF_INTERLEAVED_IMPL(IVFINTERLEAVED_128_PARAMS)
 
 }
 } // namespace faiss
diff --git a/faiss/gpu/impl/scan/IVFInterleaved2048.cu b/faiss/gpu/impl/scan/IVFInterleaved2048.cu
@@ -11,7 +11,7 @@ namespace faiss {
 namespace gpu {
 
 #if GPU_MAX_SELECTION_K >= 2048
-IVF_INTERLEAVED_IMPL(64, 2048, 8)
+IVF_INTERLEAVED_IMPL(IVFINTERLEAVED_2048_PARAMS)
 #endif
 
 } // namespace gpu

diff --git a/faiss/gpu/impl/scan/IVFInterleaved256.cu b/faiss/gpu/impl/scan/IVFInterleaved256.cu
@@ -10,7 +10,7 @@
 namespace faiss {
 namespace gpu {
 
-IVF_INTERLEAVED_IMPL(128, 256, 4)
+IVF_INTERLEAVED_IMPL(IVFINTERLEAVED_256_PARAMS)
 
 }
 } // namespace faiss
diff --git a/faiss/gpu/impl/scan/IVFInterleaved32.cu b/faiss/gpu/impl/scan/IVFInterleaved32.cu
@@ -10,7 +10,7 @@
 namespace faiss {
 namespace gpu {
 
-IVF_INTERLEAVED_IMPL(128, 32, 2)
+IVF_INTERLEAVED_IMPL(IVFINTERLEAVED_32_PARAMS)
 
 }
 } // namespace faiss
diff --git a/faiss/gpu/impl/scan/IVFInterleaved512.cu b/faiss/gpu/impl/scan/IVFInterleaved512.cu
@@ -10,7 +10,7 @@
 namespace faiss {
 namespace gpu {
 
-IVF_INTERLEAVED_IMPL(128, 512, 8)
+IVF_INTERLEAVED_IMPL(IVFINTERLEAVED_512_PARAMS)
 
 }
 } // namespace faiss
diff --git a/faiss/gpu/impl/scan/IVFInterleaved64.cu b/faiss/gpu/impl/scan/IVFInterleaved64.cu
@@ -10,7 +10,7 @@
 namespace faiss {
 namespace gpu {
 
-IVF_INTERLEAVED_IMPL(128, 64, 3)
+IVF_INTERLEAVED_IMPL(IVFINTERLEAVED_64_PARAMS)
 
 }
 } // namespace faiss