From c7288cbe5469cf6023795f5d8bd5c2e0947553d9 Mon Sep 17 00:00:00 2001 From: Michael Schellenberger Costa Date: Tue, 30 Jan 2024 10:59:15 +0000 Subject: [PATCH 1/5] Move `rmm::device_scalar` over to `cuda::mr::async_resource_ref` This rewrites `rmm::device_scalar` to take a `cuda::mr::async_resource_ref` instead of a plain `rmm::device_memory_resource*`. This is completely opaque to the user as the underlying `rmm::device_uvector` already takes a `cuda::mr::async_resource_ref` Addresses #1446 --- include/rmm/device_scalar.hpp | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/include/rmm/device_scalar.hpp b/include/rmm/device_scalar.hpp index 8e99905ce..15cab8f61 100644 --- a/include/rmm/device_scalar.hpp +++ b/include/rmm/device_scalar.hpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019-2021, NVIDIA CORPORATION. + * Copyright (c) 2019-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. @@ -18,9 +18,10 @@ #include #include -#include #include +#include + #include namespace rmm { @@ -39,6 +40,8 @@ namespace rmm { */ template class device_scalar { + using async_resource_ref = cuda::mr::async_resource_ref; + public: static_assert(std::is_trivially_copyable::value, "Scalar type must be trivially copyable"); @@ -92,9 +95,8 @@ class device_scalar { * @param stream Stream on which to perform asynchronous allocation. * @param mr Optional, resource with which to allocate. */ - explicit device_scalar( - cuda_stream_view stream, - rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource()) + explicit device_scalar(cuda_stream_view stream, + async_resource_ref mr = rmm::mr::get_current_device_resource()) : _storage{1, stream, mr} { } @@ -115,10 +117,9 @@ class device_scalar { * @param stream Optional, stream on which to perform allocation and copy. * @param mr Optional, resource with which to allocate. */ - explicit device_scalar( - value_type const& initial_value, - cuda_stream_view stream, - rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource()) + explicit device_scalar(value_type const& initial_value, + cuda_stream_view stream, + async_resource_ref mr = rmm::mr::get_current_device_resource()) : _storage{1, stream, mr} { set_value_async(initial_value, stream); @@ -138,7 +139,7 @@ class device_scalar { */ device_scalar(device_scalar const& other, cuda_stream_view stream, - rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource()) + async_resource_ref mr = rmm::mr::get_current_device_resource()) : _storage{other._storage, stream, mr} { } From 2d9c76bc3c787de498fc0366220765509adaa1f5 Mon Sep 17 00:00:00 2001 From: Michael Schellenberger Costa Date: Tue, 30 Jan 2024 16:30:53 +0000 Subject: [PATCH 2/5] Also replace usage in the `device_scalar` tests --- tests/device_scalar_tests.cpp | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/device_scalar_tests.cpp b/tests/device_scalar_tests.cpp index 7fbdaec29..a05c788b5 100644 --- a/tests/device_scalar_tests.cpp +++ b/tests/device_scalar_tests.cpp @@ -23,6 +23,8 @@ #include +#include + #include #include #include @@ -33,10 +35,12 @@ template class rmm::device_scalar; template struct DeviceScalarTest : public ::testing::Test { + using async_resource_ref = cuda::mr::async_resource_ref; + std::default_random_engine generator{}; T value{}; rmm::cuda_stream stream{}; - rmm::mr::device_memory_resource* mr{rmm::mr::get_current_device_resource()}; + async_resource_ref mr{rmm::mr::get_current_device_resource()}; DeviceScalarTest() : value{random_value()} {} From 58d276e3fbe4b5da3771fe4515502d08af0dd0d9 Mon Sep 17 00:00:00 2001 From: Michael Schellenberger Costa Date: Thu, 1 Feb 2024 08:13:08 +0000 Subject: [PATCH 3/5] Adopt `rmm::device_async_resource_ref` alias --- include/rmm/device_scalar.hpp | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/include/rmm/device_scalar.hpp b/include/rmm/device_scalar.hpp index 15cab8f61..f5df54c97 100644 --- a/include/rmm/device_scalar.hpp +++ b/include/rmm/device_scalar.hpp @@ -19,8 +19,7 @@ #include #include #include - -#include +#include #include @@ -40,8 +39,6 @@ namespace rmm { */ template class device_scalar { - using async_resource_ref = cuda::mr::async_resource_ref; - public: static_assert(std::is_trivially_copyable::value, "Scalar type must be trivially copyable"); @@ -96,7 +93,7 @@ class device_scalar { * @param mr Optional, resource with which to allocate. */ explicit device_scalar(cuda_stream_view stream, - async_resource_ref mr = rmm::mr::get_current_device_resource()) + rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource()) : _storage{1, stream, mr} { } @@ -119,7 +116,7 @@ class device_scalar { */ explicit device_scalar(value_type const& initial_value, cuda_stream_view stream, - async_resource_ref mr = rmm::mr::get_current_device_resource()) + rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource()) : _storage{1, stream, mr} { set_value_async(initial_value, stream); @@ -139,7 +136,7 @@ class device_scalar { */ device_scalar(device_scalar const& other, cuda_stream_view stream, - async_resource_ref mr = rmm::mr::get_current_device_resource()) + rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource()) : _storage{other._storage, stream, mr} { } From ccf253762537c0aa672aced4896f093d9e8a0dd1 Mon Sep 17 00:00:00 2001 From: Michael Schellenberger Costa Date: Thu, 1 Feb 2024 09:41:37 +0000 Subject: [PATCH 4/5] Also adopt the tests --- tests/device_scalar_tests.cpp | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tests/device_scalar_tests.cpp b/tests/device_scalar_tests.cpp index a05c788b5..5a7825533 100644 --- a/tests/device_scalar_tests.cpp +++ b/tests/device_scalar_tests.cpp @@ -20,6 +20,7 @@ #include #include #include +#include #include @@ -35,12 +36,10 @@ template class rmm::device_scalar; template struct DeviceScalarTest : public ::testing::Test { - using async_resource_ref = cuda::mr::async_resource_ref; - std::default_random_engine generator{}; T value{}; rmm::cuda_stream stream{}; - async_resource_ref mr{rmm::mr::get_current_device_resource()}; + rmm::device_async_resource_ref mr{rmm::mr::get_current_device_resource()}; DeviceScalarTest() : value{random_value()} {} From 5b0789a4966b941058ee9ca8d297967b2b941213 Mon Sep 17 00:00:00 2001 From: Michael Schellenberger Costa Date: Thu, 1 Feb 2024 12:18:11 +0000 Subject: [PATCH 5/5] Do not use `rmm::` inside `rmm` namespace itself --- include/rmm/device_scalar.hpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/include/rmm/device_scalar.hpp b/include/rmm/device_scalar.hpp index f5df54c97..762ba1612 100644 --- a/include/rmm/device_scalar.hpp +++ b/include/rmm/device_scalar.hpp @@ -93,7 +93,7 @@ class device_scalar { * @param mr Optional, resource with which to allocate. */ explicit device_scalar(cuda_stream_view stream, - rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource()) + device_async_resource_ref mr = mr::get_current_device_resource()) : _storage{1, stream, mr} { } @@ -116,7 +116,7 @@ class device_scalar { */ explicit device_scalar(value_type const& initial_value, cuda_stream_view stream, - rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource()) + device_async_resource_ref mr = mr::get_current_device_resource()) : _storage{1, stream, mr} { set_value_async(initial_value, stream); @@ -136,7 +136,7 @@ class device_scalar { */ device_scalar(device_scalar const& other, cuda_stream_view stream, - rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource()) + device_async_resource_ref mr = mr::get_current_device_resource()) : _storage{other._storage, stream, mr} { }