Skip to content

Commit 8ac4e41

Browse files
r-barnesfacebook-github-bot
authored andcommitted
Switch //faiss/gpu to use templates instead of macros (facebookresearch#2914)
Summary: Pull Request resolved: facebookresearch#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 Reviewed By: mdouze Differential Revision: D46549341 fbshipit-source-id: 5c3457876fd649e03ebeac89e4d1713f091ee9f5
1 parent e0741ca commit 8ac4e41

11 files changed

+392
-288
lines changed

faiss/gpu/impl/IVFInterleaved.cu

+27-8
Original file line numberDiff line numberDiff line change
@@ -192,24 +192,43 @@ void runIVFInterleavedScan(
192192
// caught for exceptions at a higher level
193193
FAISS_ASSERT(k <= GPU_MAX_SELECTION_K);
194194

195+
const auto ivf_interleaved_call = [&](const auto func) {
196+
func(queries,
197+
listIds,
198+
listData,
199+
listIndices,
200+
indicesOptions,
201+
listLengths,
202+
k,
203+
metric,
204+
useResidual,
205+
residualBase,
206+
scalarQ,
207+
outDistances,
208+
outIndices,
209+
res);
210+
};
211+
195212
if (k == 1) {
196-
IVF_INTERLEAVED_CALL(1);
213+
ivf_interleaved_call(ivfInterleavedScanImpl<IVFINTERLEAVED_1_PARAMS>);
197214
} else if (k <= 32) {
198-
IVF_INTERLEAVED_CALL(32);
215+
ivf_interleaved_call(ivfInterleavedScanImpl<IVFINTERLEAVED_32_PARAMS>);
199216
} else if (k <= 64) {
200-
IVF_INTERLEAVED_CALL(64);
217+
ivf_interleaved_call(ivfInterleavedScanImpl<IVFINTERLEAVED_64_PARAMS>);
201218
} else if (k <= 128) {
202-
IVF_INTERLEAVED_CALL(128);
219+
ivf_interleaved_call(ivfInterleavedScanImpl<IVFINTERLEAVED_128_PARAMS>);
203220
} else if (k <= 256) {
204-
IVF_INTERLEAVED_CALL(256);
221+
ivf_interleaved_call(ivfInterleavedScanImpl<IVFINTERLEAVED_256_PARAMS>);
205222
} else if (k <= 512) {
206-
IVF_INTERLEAVED_CALL(512);
223+
ivf_interleaved_call(ivfInterleavedScanImpl<IVFINTERLEAVED_512_PARAMS>);
207224
} else if (k <= 1024) {
208-
IVF_INTERLEAVED_CALL(1024);
225+
ivf_interleaved_call(
226+
ivfInterleavedScanImpl<IVFINTERLEAVED_1024_PARAMS>);
209227
}
210228
#if GPU_MAX_SELECTION_K >= 2048
211229
else if (k <= 2048) {
212-
IVF_INTERLEAVED_CALL(2048);
230+
ivf_interleaved_call(
231+
ivfInterleavedScanImpl<IVFINTERLEAVED_2048_PARAMS>);
213232
}
214233
#endif
215234
}

faiss/gpu/impl/IVFInterleaved.cuh

+12-203
Large diffs are not rendered by default.

faiss/gpu/impl/scan/IVFInterleaved1.cu

+1-1
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
namespace faiss {
1111
namespace gpu {
1212

13-
IVF_INTERLEAVED_IMPL(128, 1, 1)
13+
IVF_INTERLEAVED_IMPL(IVFINTERLEAVED_1_PARAMS)
1414

1515
}
1616
} // namespace faiss

faiss/gpu/impl/scan/IVFInterleaved1024.cu

+1-1
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
namespace faiss {
1111
namespace gpu {
1212

13-
IVF_INTERLEAVED_IMPL(128, 1024, 8)
13+
IVF_INTERLEAVED_IMPL(IVFINTERLEAVED_1024_PARAMS)
1414

1515
}
1616
} // namespace faiss

faiss/gpu/impl/scan/IVFInterleaved128.cu

+1-1
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
namespace faiss {
1111
namespace gpu {
1212

13-
IVF_INTERLEAVED_IMPL(128, 128, 3)
13+
IVF_INTERLEAVED_IMPL(IVFINTERLEAVED_128_PARAMS)
1414

1515
}
1616
} // namespace faiss

faiss/gpu/impl/scan/IVFInterleaved2048.cu

+1-1
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ namespace faiss {
1111
namespace gpu {
1212

1313
#if GPU_MAX_SELECTION_K >= 2048
14-
IVF_INTERLEAVED_IMPL(64, 2048, 8)
14+
IVF_INTERLEAVED_IMPL(IVFINTERLEAVED_2048_PARAMS)
1515
#endif
1616

1717
} // namespace gpu

faiss/gpu/impl/scan/IVFInterleaved256.cu

+1-1
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
namespace faiss {
1111
namespace gpu {
1212

13-
IVF_INTERLEAVED_IMPL(128, 256, 4)
13+
IVF_INTERLEAVED_IMPL(IVFINTERLEAVED_256_PARAMS)
1414

1515
}
1616
} // namespace faiss

faiss/gpu/impl/scan/IVFInterleaved32.cu

+1-1
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
namespace faiss {
1111
namespace gpu {
1212

13-
IVF_INTERLEAVED_IMPL(128, 32, 2)
13+
IVF_INTERLEAVED_IMPL(IVFINTERLEAVED_32_PARAMS)
1414

1515
}
1616
} // namespace faiss

faiss/gpu/impl/scan/IVFInterleaved512.cu

+1-1
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
namespace faiss {
1111
namespace gpu {
1212

13-
IVF_INTERLEAVED_IMPL(128, 512, 8)
13+
IVF_INTERLEAVED_IMPL(IVFINTERLEAVED_512_PARAMS)
1414

1515
}
1616
} // namespace faiss

faiss/gpu/impl/scan/IVFInterleaved64.cu

+1-1
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
namespace faiss {
1111
namespace gpu {
1212

13-
IVF_INTERLEAVED_IMPL(128, 64, 3)
13+
IVF_INTERLEAVED_IMPL(IVFINTERLEAVED_64_PARAMS)
1414

1515
}
1616
} // namespace faiss

0 commit comments

Comments
 (0)