Provide CPU/GPU interoperable device_id type

cpp/CMakeLists.txt

@@ -55,6 +55,7 @@ option(CUDA_STATIC_RUNTIME "Statically link the CUDA toolkit runtime and librari
 option(DETECT_CONDA_ENV "Enable detection of conda environment for dependencies" ON)
 option(DISABLE_DEPRECATION_WARNINGS "Disable deprecaction warnings " ON)
 option(DISABLE_OPENMP "Disable OpenMP" OFF)
+option(DISABLE_CUDA "Disable CUDA in supported RAFT code" OFF)
 option(RAFT_NVTX "Enable nvtx markers" OFF)
 
 set(RAFT_COMPILE_LIBRARIES_DEFAULT OFF)
@@ -276,6 +277,13 @@ target_compile_definitions(raft::raft INTERFACE $<$<BOOL:${RAFT_NVTX}>:NVTX_ENAB
   )
 endif()
 
+##############################################################################
+# - CUDA-free build support --------------------------------------------------
+
+if (DISABLE_CUDA)
+  target_compile_definitions(raft INTERFACE RAFT_DISABLE_CUDA)
+endif()
+
 # ##################################################################################################
 # * raft_distance ------------------------------------------------------------ TODO: Currently, this
 #   package also contains the 'random' namespace (for rmat logic) We couldn't get this to work

cpp/include/raft/core/detail/device_id_base.hpp

@@ -0,0 +1,37 @@
+/*
+ * Copyright (c) 2022, 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 <raft/core/device_support.hpp>
+#include <raft/core/device_type.hpp>
+
+namespace raft {
+namespace detail {
+template <device_type D>
+struct device_id {
+  using value_type = int;
+
+  device_id(value_type device_index = value_type{}) noexcept {}
+  auto value() const noexcept(false)
+  {
+    throw cuda_unsupported{"Attempting to use a GPU device in a non-CUDA build"};
+  }
+  auto rmm_id() const noexcept(false)
+  {
+    throw cuda_unsupported{"Attempting to use a GPU device in a non-CUDA build"};
+  }
+};
+}  // namespace detail
+}  // namespace raft

cpp/include/raft/core/detail/device_id_cpu.hpp

@@ -0,0 +1,34 @@
+/*
+ * Copyright (c) 2022, 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 <raft/core/detail/device_id_base.hpp>
+#include <raft/core/device_type.hpp>
+
+namespace raft {
+namespace detail {
+template <>
+struct device_id<device_type::cpu> {
+  using value_type = int;
+  device_id(value_type dev_id = value_type{}) noexcept : id_{dev_id} {}
+
+  auto value() const noexcept { return id_; }
+  auto rmm_id() const noexcept(false) { throw bad_device_type{"CPU devices have no RMM ID"}; }
+
+ private:
+  value_type id_;
+};
+}  // namespace detail
+}  // namespace raft

cpp/include/raft/core/detail/device_id_gpu.hpp

@@ -0,0 +1,50 @@
+/*
+ * Copyright (c) 2022, 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 <raft/core/detail/device_id_base.hpp>
+#include <raft/core/device_type.hpp>
+#include <raft/util/cudart_utils.hpp>
+#include <rmm/cuda_device.hpp>
+
+namespace raft {
+namespace detail {
+template <>
+struct device_id<device_type::gpu> {
+  using value_type = typename rmm::cuda_device_id::value_type;
+  device_id() noexcept(false)
+    : id_{[]() {
+        auto raw_id = value_type{};
+        RAFT_CUDA_TRY(cudaGetDevice(&raw_id));
+        return raw_id;
+      }()} {};
+  /* We do not mark this constructor as explicit to allow public API
+   * functions to accept `device_id` arguments without requiring
+   * downstream consumers to explicitly construct a device_id. Thus,
+   * consumers can use the type they expect to use when specifying a device
+   * (int), but once we are inside the public API, the device type remains
+   * attached to this value and we can easily convert to the strongly-typed
+   * rmm::cuda_device_id if desired.
+   */
+  device_id(value_type dev_id) noexcept : id_{dev_id} {};
+
+  auto value() const noexcept { return id_.value(); }
+  auto rmm_id() const noexcept { return id_; }
+
+ private:
+  rmm::cuda_device_id id_;
+};
+}  // namespace detail
+}  // namespace raft

cpp/include/raft/core/device_id.hpp

@@ -0,0 +1,31 @@
+/*
+ * Copyright (c) 2022, 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 <raft/core/detail/device_id_base.hpp>
+#include <raft/core/detail/device_id_cpu.hpp>
+#ifndef RAFT_DISABLE_CUDA
+#include <raft/core/detail/device_id_gpu.hpp>
+#endif
+#include <raft/core/device_type.hpp>
+#include <variant>
+
+namespace raft {
+template <device_type D>
+using device_id = detail::device_id<D>;
+
+using device_id_variant = std::variant<device_id<device_type::cpu>, device_id<device_type::gpu>>;
+}  // namespace raft

cpp/include/raft/core/device_support.hpp

@@ -0,0 +1,32 @@
+/*
+ * Copyright (c) 2022, 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 <raft/core/error.hpp>
+#include <string>
+
+namespace raft {
+#ifdef RAFT_DISABLE_CUDA
+auto constexpr static const CUDA_ENABLED = false;
+#else
+auto constexpr static const CUDA_ENABLED = true;
+#endif
+
+struct cuda_unsupported : raft::exception {
+  explicit cuda_unsupported(std::string const& msg) : raft::exception{msg} {}
+  cuda_unsupported() : cuda_unsupported{"CUDA functionality invoked in non-CUDA build"} {}
+};
+
+}  // namespace raft

cpp/include/raft/core/device_type.hpp

@@ -0,0 +1,35 @@
+/*
+ * Copyright (c) 2022, 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 <raft/core/error.hpp>
+#include <raft/core/memory_type.hpp>
+#include <string>
+
+namespace raft {
+enum class device_type { cpu, gpu };
+
+auto constexpr is_compatible(device_type dev_type, memory_type mem_type)
+{
+  return (dev_type == device_type::gpu && is_device_accessible(mem_type)) ||
+         (dev_type == device_type::cpu && is_host_accessible(mem_type));
+}
+
+struct bad_device_type : raft::exception {
+  bad_device_type(std::string const& msg) : raft::exception{msg} {}
+  bad_device_type() : bad_device_type{"Incorrect device type for this operation"} {}
+};
+
+}  // end namespace raft

cpp/test/CMakeLists.txt

@@ -96,6 +96,8 @@ if(BUILD_TESTS)
     CORE_TEST
     PATH
     test/core/logger.cpp
+    test/core/device_id.cpp
+    test/core/device_type.cpp
     test/core/math_device.cu
     test/core/math_host.cpp
     test/core/operators_device.cu

cpp/test/core/device_id.cpp

@@ -0,0 +1,39 @@
+/*
+ * Copyright (c) 2022, 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.
+ */
+#include <gtest/gtest.h>
+#include <raft/core/device_id.hpp>
+#include <raft/core/device_support.hpp>
+#include <raft/core/device_type.hpp>
+
+namespace raft {
+TEST(DeviceID, CPU)
+{
+  auto dev_id = device_id<device_type::cpu>{};
+  ASSERT_EQ(dev_id.value(), 0);
+  ASSERT_THROW(dev_id.rmm_id(), bad_device_type);
+}
+
+TEST(DeviceID, GPU)
+{
+  auto dev_id = device_id<device_type::gpu>{};
+#ifdef RAFT_DISABLE_CUDA
+  ASSERT_THROW(dev_id.rmm_id(), cuda_unsupported);
+  ASSERT_THROW(dev_id.value(), cuda_unsupported);
+#else
+  ASSERT_EQ(dev_id.value(), dev_id.rmm_id().value());
+#endif
+}
+}  // namespace raft

cpp/test/core/device_type.cpp

@@ -0,0 +1,35 @@
+/*
+ * Copyright (c) 2022, 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.
+ */
+#include <gtest/gtest.h>
+#include <raft/core/device_type.hpp>
+
+namespace raft {
+TEST(DeviceType, CPU)
+{
+  static_assert(is_compatible(device_type::cpu, memory_type::host));
+  static_assert(!is_compatible(device_type::cpu, memory_type::device));
+  static_assert(is_compatible(device_type::cpu, memory_type::managed));
+  static_assert(is_compatible(device_type::cpu, memory_type::pinned));
+}
+
+TEST(DeviceType, GPU)
+{
+  static_assert(!is_compatible(device_type::gpu, memory_type::host));
+  static_assert(is_compatible(device_type::gpu, memory_type::device));
+  static_assert(is_compatible(device_type::gpu, memory_type::managed));
+  static_assert(!is_compatible(device_type::gpu, memory_type::pinned));
+}
+}  // namespace raft