SNMG ANN

cpp/CMakeLists.txt

@@ -470,6 +470,10 @@ add_library(
   ${CUVS_MG_ALGOS}
 )
 
+if(BUILD_MG_ALGOS)
+  target_compile_definitions(cuvs PUBLIC CUVS_BUILD_MG_ALGOS=1)
+endif()
+
 target_compile_options(
   cuvs INTERFACE $<$<COMPILE_LANG_AND_ID:CUDA,NVIDIA>:--expt-extended-lambda
                  --expt-relaxed-constexpr>

cpp/include/cuvs/neighbors/ann_mg.hpp

@@ -16,34 +16,16 @@
 
 #pragma once
 
-#include <nccl.h>
 #include <raft/core/device_resources.hpp>
 #include <rmm/mr/device/per_device_resource.hpp>
 
 #include <cuvs/neighbors/cagra.hpp>
 #include <cuvs/neighbors/ivf_flat.hpp>
 #include <cuvs/neighbors/ivf_pq.hpp>
 
-/**
- * @brief Error checking macro for NCCL runtime API functions.
- *
- * Invokes a NCCL runtime API function call, if the call does not return ncclSuccess, throws an
- * exception detailing the NCCL error that occurred
- */
-#define RAFT_NCCL_TRY(call)                        \
-  do {                                             \
-    ncclResult_t const status = (call);            \
-    if (ncclSuccess != status) {                   \
-      std::string msg{};                           \
-      SET_ERROR_MSG(msg,                           \
-                    "NCCL error encountered at: ", \
-                    "call='%s', Reason=%d:%s",     \
-                    #call,                         \
-                    status,                        \
-                    ncclGetErrorString(status));   \
-      throw raft::logic_error(msg);                \
-    }                                              \
-  } while (0);
+#ifndef NO_NCCL_FORWARD_DECLARATION
+class ncclComm_t {};
+#endif
 
 namespace cuvs::neighbors::mg {
 enum parallel_mode { REPLICATED, SHARDED };

cpp/src/neighbors/ann_mg/ann_mg.cuh

@@ -16,12 +16,18 @@
 
 #pragma once
 
+#ifdef CUVS_BUILD_MG_ALGOS
+
 #include "../detail/knn_merge_parts.cuh"
-#include <cuvs/neighbors/ann_mg.hpp>
-#include <cuvs/neighbors/common.hpp>
 #include <raft/core/serialize.hpp>
 #include <raft/util/cuda_dev_essentials.cuh>
 
+#include "nccl_helpers.cuh"
+#define NO_NCCL_FORWARD_DECLARATION
+#include <cuvs/neighbors/ann_mg.hpp>
+#undef NO_NCCL_FORWARD_DECLARATION
+#include <cuvs/neighbors/common.hpp>
+
 namespace cuvs::neighbors::mg {
 using namespace cuvs::neighbors;
 using namespace raft;
@@ -741,3 +747,11 @@ ann_mg_index<cagra::index<T, IdxT>, T, IdxT> distribute_cagra(
 }
 
 }  // namespace cuvs::neighbors::mg::detail
+
+#else
+
+static_assert(false,
+              "FORBIDEN_MG_ALGORITHM_IMPORT\n\n"
+              "Please recompile the cuVS library with MG algorithms.\n");
+
+#endif

cpp/src/neighbors/ann_mg/ann_mg_cagra_float_uint32_t.cu

@@ -23,8 +23,6 @@
  *
  */
 
-#include <cuvs/neighbors/ann_mg.hpp>
-
 #include "ann_mg.cuh"
 
 namespace cuvs::neighbors::mg {

cpp/src/neighbors/ann_mg/ann_mg_cagra_int8_t_uint32_t.cu

@@ -23,8 +23,6 @@
  *
  */
 
-#include <cuvs/neighbors/ann_mg.hpp>
-
 #include "ann_mg.cuh"
 
 namespace cuvs::neighbors::mg {

cpp/src/neighbors/ann_mg/ann_mg_cagra_uint8_t_uint32_t.cu

@@ -23,8 +23,6 @@
  *
  */
 
-#include <cuvs/neighbors/ann_mg.hpp>
-
 #include "ann_mg.cuh"
 
 namespace cuvs::neighbors::mg {

cpp/src/neighbors/ann_mg/ann_mg_flat_float_int64_t.cu

@@ -23,8 +23,6 @@
  *
  */
 
-#include <cuvs/neighbors/ann_mg.hpp>
-
 #include "ann_mg.cuh"
 
 namespace cuvs::neighbors::mg {

cpp/src/neighbors/ann_mg/ann_mg_flat_int8_t_int64_t.cu

@@ -23,8 +23,6 @@
  *
  */
 
-#include <cuvs/neighbors/ann_mg.hpp>
-
 #include "ann_mg.cuh"
 
 namespace cuvs::neighbors::mg {

cpp/src/neighbors/ann_mg/ann_mg_flat_uint8_t_int64_t.cu

@@ -23,8 +23,6 @@
  *
  */
 
-#include <cuvs/neighbors/ann_mg.hpp>
-
 #include "ann_mg.cuh"
 
 namespace cuvs::neighbors::mg {

cpp/src/neighbors/ann_mg/ann_mg_pq_float_int64_t.cu

@@ -23,8 +23,6 @@
  *
  */
 
-#include <cuvs/neighbors/ann_mg.hpp>
-
 #include "ann_mg.cuh"
 
 namespace cuvs::neighbors::mg {

cpp/src/neighbors/ann_mg/ann_mg_pq_int8_t_int64_t.cu

@@ -23,8 +23,6 @@
  *
  */
 
-#include <cuvs/neighbors/ann_mg.hpp>
-
 #include "ann_mg.cuh"
 
 namespace cuvs::neighbors::mg {

cpp/src/neighbors/ann_mg/ann_mg_pq_uint8_t_int64_t.cu

@@ -23,8 +23,6 @@
  *
  */
 
-#include <cuvs/neighbors/ann_mg.hpp>
-
 #include "ann_mg.cuh"
 
 namespace cuvs::neighbors::mg {

cpp/src/neighbors/ann_mg/generate_ann_mg.py

@@ -37,7 +37,6 @@
  *
  */
 
-#include <cuvs/neighbors/ann_mg.hpp>
 """
 
 include_macro = """

cpp/src/neighbors/ann_mg/nccl_clique.cu

@@ -14,9 +14,13 @@
  * limitations under the License.
  */
 
-#include <cuvs/neighbors/ann_mg.hpp>
 #include <raft/comms/std_comms.hpp>
 
+#include "nccl_helpers.cuh"
+#define NO_NCCL_FORWARD_DECLARATION
+#include <cuvs/neighbors/ann_mg.hpp>
+#undef NO_NCCL_FORWARD_DECLARATION
+
 namespace cuvs::neighbors::mg {
 
 nccl_clique::nccl_clique(const std::vector<int>& device_ids)

cpp/src/neighbors/ann_mg/nccl_helpers.cuh

@@ -0,0 +1,40 @@
+/*
+ * Copyright (c) 2024, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <nccl.h>
+
+/**
+ * @brief Error checking macro for NCCL runtime API functions.
+ *
+ * Invokes a NCCL runtime API function call, if the call does not return ncclSuccess, throws an
+ * exception detailing the NCCL error that occurred
+ */
+#define RAFT_NCCL_TRY(call)                        \
+  do {                                             \
+    ncclResult_t const status = (call);            \
+    if (ncclSuccess != status) {                   \
+      std::string msg{};                           \
+      SET_ERROR_MSG(msg,                           \
+                    "NCCL error encountered at: ", \
+                    "call='%s', Reason=%d:%s",     \
+                    #call,                         \
+                    status,                        \
+                    ncclGetErrorString(status));   \
+      throw raft::logic_error(msg);                \
+    }                                              \
+  } while (0);