Add tests for containers when switching active CUDA device

tests/CMakeLists.txt

@@ -186,4 +186,7 @@ ConfigureTest(BINNING_MR_TEST mr/device/binning_mr_tests.cpp)
 # callback memory resource tests
 ConfigureTest(CALLBACK_MR_TEST mr/device/callback_mr_tests.cpp)
 
+# container multidevice tests
+ConfigureTest(CONTAINER_MULTIDEVICE_TEST container_multidevice_tests.cu)
+
 rapids_test_install_relocatable(INSTALL_COMPONENT_SET testing DESTINATION bin/gtests/librmm)

tests/container_multidevice_tests.cu

@@ -0,0 +1,149 @@
+/*
+ * Copyright (c) 2023, 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 "device_check_resource_adaptor.hpp"
+#include "rmm/mr/device/per_device_resource.hpp"
+
+#include <rmm/cuda_stream.hpp>
+#include <rmm/device_buffer.hpp>
+#include <rmm/device_scalar.hpp>
+#include <rmm/device_uvector.hpp>
+
+#include <gtest/gtest.h>
+
+#include <type_traits>
+
+template <typename ContainerType>
+struct ContainerMultiDeviceTest : public ::testing::Test {};
+
+using containers =
+  ::testing::Types<rmm::device_buffer, rmm::device_uvector<int>, rmm::device_scalar<int>>;
+
+TYPED_TEST_CASE(ContainerMultiDeviceTest, containers);
+
+TYPED_TEST(ContainerMultiDeviceTest, CreateDestroyDifferentActiveDevice)
+{
+  // Get the number of cuda devices
+  int num_devices = rmm::get_num_cuda_devices();
+
+  // only run on multidevice systems
+  if (num_devices >= 2) {
+    rmm::cuda_set_device_raii dev{rmm::cuda_device_id{0}};
+    auto* orig_mr = rmm::mr::get_current_device_resource();
+    auto check_mr = device_check_resource_adaptor{orig_mr};
+    rmm::mr::set_current_device_resource(&check_mr);
+
+    {
+      if constexpr (std::is_same_v<TypeParam, rmm::device_scalar<int>>) {
+        auto buf = TypeParam(rmm::cuda_stream_view{});
+        RMM_ASSERT_CUDA_SUCCESS(cudaSetDevice(1));  // force dtor with different active device
+      } else {
+        auto buf = TypeParam(128, rmm::cuda_stream_view{});
+        RMM_ASSERT_CUDA_SUCCESS(cudaSetDevice(1));  // force dtor with different active device
+      }
+    }
+
+    RMM_ASSERT_CUDA_SUCCESS(cudaSetDevice(0));
+    rmm::mr::set_current_device_resource(orig_mr);
+  }
+}
+
+TYPED_TEST(ContainerMultiDeviceTest, CreateMoveDestroyDifferentActiveDevice)
+{
+  // Get the number of cuda devices
+  int num_devices = rmm::get_num_cuda_devices();
+
+  // only run on multidevice systems
+  if (num_devices >= 2) {
+    rmm::cuda_set_device_raii dev{rmm::cuda_device_id{0}};
+    auto* orig_mr = rmm::mr::get_current_device_resource();
+    auto check_mr = device_check_resource_adaptor{orig_mr};
+    rmm::mr::set_current_device_resource(&check_mr);
+
+    {
+      auto buf_1 = []() {
+        if constexpr (std::is_same_v<TypeParam, rmm::device_scalar<int>>) {
+          return TypeParam(rmm::cuda_stream_view{});
+        } else {
+          return TypeParam(128, rmm::cuda_stream_view{});
+        }
+      }();
+
+      {
+        if constexpr (std::is_same_v<TypeParam, rmm::device_scalar<int>>) {
+          // device_vector does not have a constructor that takes a stream
+          auto buf_0 = TypeParam(rmm::cuda_stream_view{});
+          buf_1      = std::move(buf_0);
+        } else {
+          auto buf_0 = TypeParam(128, rmm::cuda_stream_view{});
+          buf_1      = std::move(buf_0);
+        }
+      }
+
+      RMM_ASSERT_CUDA_SUCCESS(cudaSetDevice(1));  // force dtor with different active device
+    }
+
+    RMM_ASSERT_CUDA_SUCCESS(cudaSetDevice(0));
+    rmm::mr::set_current_device_resource(orig_mr);
+  }
+}
+
+TYPED_TEST(ContainerMultiDeviceTest, ResizeDifferentActiveDevice)
+{
+  // Get the number of cuda devices
+  int num_devices = rmm::get_num_cuda_devices();
+
+  // only run on multidevice systems
+  if (num_devices >= 2) {
+    rmm::cuda_set_device_raii dev{rmm::cuda_device_id{0}};
+    auto* orig_mr = rmm::mr::get_current_device_resource();
+    auto check_mr = device_check_resource_adaptor{orig_mr};
+    rmm::mr::set_current_device_resource(&check_mr);
+
+    if constexpr (not std::is_same_v<TypeParam, rmm::device_scalar<int>>) {
+      auto buf = TypeParam(128, rmm::cuda_stream_view{});
+      RMM_ASSERT_CUDA_SUCCESS(cudaSetDevice(1));  // force resize with different active device
+      buf.resize(1024, rmm::cuda_stream_view{});
+    }
+
+    RMM_ASSERT_CUDA_SUCCESS(cudaSetDevice(0));
+    rmm::mr::set_current_device_resource(orig_mr);
+  }
+}
+
+TYPED_TEST(ContainerMultiDeviceTest, ShrinkDifferentActiveDevice)
+{
+  // Get the number of cuda devices
+  int num_devices = rmm::get_num_cuda_devices();
+
+  // only run on multidevice systems
+  if (num_devices >= 2) {
+    rmm::cuda_set_device_raii dev{rmm::cuda_device_id{0}};
+    auto* orig_mr = rmm::mr::get_current_device_resource();
+    auto check_mr = device_check_resource_adaptor{orig_mr};
+    rmm::mr::set_current_device_resource(&check_mr);
+
+    if constexpr (not std::is_same_v<TypeParam, rmm::device_scalar<int>>) {
+      auto buf = TypeParam(128, rmm::cuda_stream_view{});
+      RMM_ASSERT_CUDA_SUCCESS(cudaSetDevice(1));  // force resize with different active device
+      buf.resize(64, rmm::cuda_stream_view{});
+      buf.shrink_to_fit(rmm::cuda_stream_view{});
+    }
+
+    RMM_ASSERT_CUDA_SUCCESS(cudaSetDevice(0));
+    rmm::mr::set_current_device_resource(orig_mr);
+  }
+}

tests/device_buffer_tests.cu

@@ -14,8 +14,6 @@
  * limitations under the License.
  */
 
-#include <gtest/gtest.h>
-
 #include <rmm/cuda_stream.hpp>
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/detail/error.hpp>
@@ -29,9 +27,12 @@
 
 #include <thrust/equal.h>
 #include <thrust/sequence.h>
+
+#include <gtest/gtest.h>
+
 namespace testing {
 namespace thrust = THRUST_NS_QUALIFIER;
-}
+}  // namespace testing
 using namespace testing;
 
 #include <cuda_runtime_api.h>

tests/device_check_resource_adaptor.hpp

@@ -0,0 +1,77 @@
+/*
+ * Copyright (c) 2023, 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 <rmm/cuda_device.hpp>
+#include <rmm/cuda_stream_view.hpp>
+#include <rmm/mr/device/device_memory_resource.hpp>
+
+#include <gtest/gtest.h>
+
+class device_check_resource_adaptor final : public rmm::mr::device_memory_resource {
+ public:
+  device_check_resource_adaptor(rmm::mr::device_memory_resource* upstream)
+    : device_id{rmm::get_current_cuda_device()}, upstream_(upstream)
+  {
+  }
+
+  [[nodiscard]] bool supports_streams() const noexcept override
+  {
+    return upstream_->supports_streams();
+  }
+
+  [[nodiscard]] bool supports_get_mem_info() const noexcept override
+  {
+    return upstream_->supports_get_mem_info();
+  }
+
+  [[nodiscard]] device_memory_resource* get_upstream() const noexcept { return upstream_; }
+
+ private:
+  [[nodiscard]] bool check_device_id() const { return device_id == rmm::get_current_cuda_device(); }
+
+  void* do_allocate(std::size_t bytes, rmm::cuda_stream_view stream) override
+  {
+    bool const is_correct_device = check_device_id();
+    EXPECT_TRUE(is_correct_device);
+    if (is_correct_device) { return upstream_->allocate(bytes, stream); }
+    return nullptr;
+  }
+
+  void do_deallocate(void* ptr, std::size_t bytes, rmm::cuda_stream_view stream) override
+  {
+    bool const is_correct_device = check_device_id();
+    EXPECT_TRUE(is_correct_device);
+    if (is_correct_device) { upstream_->deallocate(ptr, bytes, stream); }
+  }
+
+  [[nodiscard]] bool do_is_equal(
+    rmm::mr::device_memory_resource const& other) const noexcept override
+  {
+    if (this == &other) { return true; }
+    auto const* cast = dynamic_cast<device_check_resource_adaptor const*>(&other);
+    if (cast != nullptr) { return upstream_->is_equal(*cast->get_upstream()); }
+    return upstream_->is_equal(other);
+  }
+
+  [[nodiscard]] std::pair<std::size_t, std::size_t> do_get_mem_info(
+    rmm::cuda_stream_view stream) const override
+  {
+    return upstream_->get_mem_info(stream);
+  }
+
+  rmm::cuda_device_id device_id;
+  rmm::mr::device_memory_resource* upstream_{};
+};