Move cagra template instantiations from raft_lib to NEIGHBORS_TEST

rapidsai · May 14, 2023 · 60184c1 · 60184c1
1 parent 97f7a48
commit 60184c1
Show file tree

Hide file tree

Showing 9 changed files with 50 additions and 34 deletions.
diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt
@@ -306,33 +306,6 @@ if(RAFT_COMPILE_LIBRARY)
     src/neighbors/brute_force_knn_int64_t_float_uint32_t.cu
     src/neighbors/brute_force_knn_int_float_int.cu
     src/neighbors/brute_force_knn_uint32_t_float_uint32_t.cu
-    src/neighbors/cagra_build_float_uint32.cu
-    src/neighbors/cagra_prune_float_uint32.cu
-    src/neighbors/cagra_search_float_uint32.cu
-    src/neighbors/detail/cagra/search_multi_cta_float_uint32_dim128_t8.cu
-    src/neighbors/detail/cagra/search_multi_cta_float_uint32_dim256_t16.cu
-    src/neighbors/detail/cagra/search_multi_cta_float_uint32_dim512_t32.cu
-    src/neighbors/detail/cagra/search_multi_cta_float_uint32_dim1024_t32.cu  
-    src/neighbors/detail/cagra/search_multi_cta_int8_uint32_dim128_t8.cu
-    src/neighbors/detail/cagra/search_multi_cta_int8_uint32_dim256_t16.cu
-    src/neighbors/detail/cagra/search_multi_cta_int8_uint32_dim512_t32.cu
-    src/neighbors/detail/cagra/search_multi_cta_int8_uint32_dim1024_t32.cu
-    src/neighbors/detail/cagra/search_multi_cta_uint8_uint32_dim128_t8.cu
-    src/neighbors/detail/cagra/search_multi_cta_uint8_uint32_dim256_t16.cu
-    src/neighbors/detail/cagra/search_multi_cta_uint8_uint32_dim512_t32.cu
-    src/neighbors/detail/cagra/search_multi_cta_uint8_uint32_dim1024_t32.cu
-    src/neighbors/detail/cagra/search_single_cta_float_uint32_dim128_t8.cu
-    src/neighbors/detail/cagra/search_single_cta_float_uint32_dim256_t16.cu
-    src/neighbors/detail/cagra/search_single_cta_float_uint32_dim512_t32.cu
-    src/neighbors/detail/cagra/search_single_cta_float_uint32_dim1024_t32.cu
-    src/neighbors/detail/cagra/search_single_cta_int8_uint32_dim128_t8.cu
-    src/neighbors/detail/cagra/search_single_cta_int8_uint32_dim256_t16.cu
-    src/neighbors/detail/cagra/search_single_cta_int8_uint32_dim512_t32.cu
-    src/neighbors/detail/cagra/search_single_cta_int8_uint32_dim1024_t32.cu
-    src/neighbors/detail/cagra/search_single_cta_uint8_uint32_dim128_t8.cu
-    src/neighbors/detail/cagra/search_single_cta_uint8_uint32_dim256_t16.cu
-    src/neighbors/detail/cagra/search_single_cta_uint8_uint32_dim512_t32.cu
-    src/neighbors/detail/cagra/search_single_cta_uint8_uint32_dim1024_t32.cu
     src/neighbors/detail/ivf_flat_interleaved_scan_float_float_int64_t.cu
     src/neighbors/detail/ivf_flat_interleaved_scan_int8_t_int32_t_int64_t.cu
     src/neighbors/detail/ivf_flat_interleaved_scan_uint8_t_uint32_t_int64_t.cu

diff --git a/cpp/include/raft/neighbors/detail/cagra/search_multi_cta_kernel-ext.cuh b/cpp/include/raft/neighbors/detail/cagra/search_multi_cta_kernel-ext.cuh
@@ -20,7 +20,7 @@
 namespace raft::neighbors::experimental::cagra::detail {
 namespace multi_cta_search {
 
-#ifdef RAFT_EXPLICIT_INSTANTIATE_ONLY
+#ifdef RAFT_EXPLICIT_INSTANTIATE_ONLY_CAGRA
 
 template <unsigned TEAM_SIZE,
           unsigned BLOCK_SIZE,

diff --git a/cpp/include/raft/neighbors/detail/cagra/search_multi_cta_kernel.cuh b/cpp/include/raft/neighbors/detail/cagra/search_multi_cta_kernel.cuh
@@ -15,10 +15,10 @@
  */
 #pragma once
 
-#ifndef RAFT_EXPLICIT_INSTANTIATE_ONLY
+#ifndef RAFT_EXPLICIT_INSTANTIATE_ONLY_CAGRA
 #include "search_multi_cta_kernel-inl.cuh"
 #endif
 
-#ifdef RAFT_COMPILED
+#ifdef RAFT_COMPILED_CAGRA
 #include "search_multi_cta_kernel-ext.cuh"
 #endif
diff --git a/cpp/include/raft/neighbors/detail/cagra/search_single_cta_kernel-ext.cuh b/cpp/include/raft/neighbors/detail/cagra/search_single_cta_kernel-ext.cuh
@@ -19,7 +19,7 @@
 namespace raft::neighbors::experimental::cagra::detail {
 namespace single_cta_search {
 
-#ifdef RAFT_EXPLICIT_INSTANTIATE_ONLY
+#ifdef RAFT_EXPLICIT_INSTANTIATE_ONLY_CAGRA
 
 template <unsigned TEAM_SIZE,
           unsigned BLOCK_SIZE,

diff --git a/cpp/include/raft/neighbors/detail/cagra/search_single_cta_kernel.cuh b/cpp/include/raft/neighbors/detail/cagra/search_single_cta_kernel.cuh
@@ -15,10 +15,10 @@
  */
 #pragma once
 
-#ifndef RAFT_EXPLICIT_INSTANTIATE_ONLY
+#ifndef RAFT_EXPLICIT_INSTANTIATE_ONLY_CAGRA
 #include "search_single_cta_kernel-inl.cuh"
 #endif
 
-#ifdef RAFT_COMPILED
+#ifdef RAFT_COMPILED_CAGRA
 #include "search_single_cta_kernel-ext.cuh"
 #endif
diff --git a/cpp/test/CMakeLists.txt b/cpp/test/CMakeLists.txt
@@ -332,6 +332,35 @@ if(BUILD_TESTS)
     test/neighbors/epsilon_neighborhood.cu
     test/neighbors/refine.cu
     test/neighbors/selection.cu
+    # TODO(tfeher): move the cagra template instantiation .cu files back to raft_lib
+    # once https://github.com/rapidsai/raft/issues/1459 is resolved.
+    src/neighbors/cagra_build_float_uint32.cu
+    src/neighbors/cagra_prune_float_uint32.cu
+    src/neighbors/cagra_search_float_uint32.cu
+    src/neighbors/detail/cagra/search_multi_cta_float_uint32_dim128_t8.cu
+    src/neighbors/detail/cagra/search_multi_cta_float_uint32_dim256_t16.cu
+    src/neighbors/detail/cagra/search_multi_cta_float_uint32_dim512_t32.cu
+    src/neighbors/detail/cagra/search_multi_cta_float_uint32_dim1024_t32.cu  
+    src/neighbors/detail/cagra/search_multi_cta_int8_uint32_dim128_t8.cu
+    src/neighbors/detail/cagra/search_multi_cta_int8_uint32_dim256_t16.cu
+    src/neighbors/detail/cagra/search_multi_cta_int8_uint32_dim512_t32.cu
+    src/neighbors/detail/cagra/search_multi_cta_int8_uint32_dim1024_t32.cu
+    src/neighbors/detail/cagra/search_multi_cta_uint8_uint32_dim128_t8.cu
+    src/neighbors/detail/cagra/search_multi_cta_uint8_uint32_dim256_t16.cu
+    src/neighbors/detail/cagra/search_multi_cta_uint8_uint32_dim512_t32.cu
+    src/neighbors/detail/cagra/search_multi_cta_uint8_uint32_dim1024_t32.cu
+    src/neighbors/detail/cagra/search_single_cta_float_uint32_dim128_t8.cu
+    src/neighbors/detail/cagra/search_single_cta_float_uint32_dim256_t16.cu
+    src/neighbors/detail/cagra/search_single_cta_float_uint32_dim512_t32.cu
+    src/neighbors/detail/cagra/search_single_cta_float_uint32_dim1024_t32.cu
+    src/neighbors/detail/cagra/search_single_cta_int8_uint32_dim128_t8.cu
+    src/neighbors/detail/cagra/search_single_cta_int8_uint32_dim256_t16.cu
+    src/neighbors/detail/cagra/search_single_cta_int8_uint32_dim512_t32.cu
+    src/neighbors/detail/cagra/search_single_cta_int8_uint32_dim1024_t32.cu
+    src/neighbors/detail/cagra/search_single_cta_uint8_uint32_dim128_t8.cu
+    src/neighbors/detail/cagra/search_single_cta_uint8_uint32_dim256_t16.cu
+    src/neighbors/detail/cagra/search_single_cta_uint8_uint32_dim512_t32.cu
+    src/neighbors/detail/cagra/search_single_cta_uint8_uint32_dim1024_t32.cu
     OPTIONAL
     LIB
     EXPLICIT_INSTANTIATE_ONLY
@@ -370,7 +399,7 @@ if(BUILD_TESTS)
     LIB
     EXPLICIT_INSTANTIATE_ONLY
   )
-
+  
   ConfigureTest(
     NAME UTILS_TEST PATH test/core/seive.cu test/util/bitonic_sort.cu test/util/cudart_utils.cpp
     test/util/device_atomics.cu test/util/integer_utils.cpp test/util/pow2_utils.cu

diff --git a/cpp/test/neighbors/ann_cagra/test_float_uint32_t.cu b/cpp/test/neighbors/ann_cagra/test_float_uint32_t.cu
@@ -16,6 +16,11 @@
 
 #include <gtest/gtest.h>
 
+// TODO(tfeher): remove these definitions once https://github.com/rapidsai/raft/issues/1459
+// is resolved
+#define RAFT_EXPLICIT_INSTANTIATE_ONLY_CAGRA
+#define RAFT_COMPILED_CAGRA
+
 #include "../ann_cagra.cuh"
 
 namespace raft::neighbors::experimental::cagra {

diff --git a/cpp/test/neighbors/ann_cagra/test_int8_t_uint32_t.cu b/cpp/test/neighbors/ann_cagra/test_int8_t_uint32_t.cu
@@ -15,6 +15,10 @@
  */
 
 #include <gtest/gtest.h>
+// TODO(tfeher): remove these definitions once https://github.com/rapidsai/raft/issues/1459
+// is resolved
+#define RAFT_EXPLICIT_INSTANTIATE_ONLY_CAGRA
+#define RAFT_COMPILED_CAGRA
 
 #include "../ann_cagra.cuh"
 

diff --git a/cpp/test/neighbors/ann_cagra/test_uint8_t_uint32_t.cu b/cpp/test/neighbors/ann_cagra/test_uint8_t_uint32_t.cu
@@ -16,6 +16,11 @@
 
 #include <gtest/gtest.h>
 
+// TODO(tfeher): remove these definitions once https://github.com/rapidsai/raft/issues/1459
+// is resolved
+#define RAFT_EXPLICIT_INSTANTIATE_ONLY_CAGRA
+#define RAFT_COMPILED_CAGRA
+
 #include "../ann_cagra.cuh"
 
 namespace raft::neighbors::experimental::cagra {