Skip to content

Commit

Permalink
Roll back the changes to device_u{buffer, vector}
Browse files Browse the repository at this point in the history 
Those are potentially breaking and we want to do that in a followup
  • Loading branch information
@miscco
miscco committed Nov 2, 2023
1 parent 849c880 commit 2cf03a7
Show file tree
Hide file tree
Showing 5 changed files with 64 additions and 77 deletions.
30 changes: 10 additions & 20 deletions include/rmm/device_buffer.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,8 +27,6 @@
#include <stdexcept>
#include <utility>

#include <cuda/memory_resource>

namespace rmm {
/**
* @addtogroup data_containers
Expand Down Expand Up @@ -81,8 +79,6 @@ namespace rmm {
*```
*/
class device_buffer {
using async_resource_ref = cuda::mr::async_resource_ref<cuda::mr::device_accessible>;

public:
// The copy constructor and copy assignment operator without a stream are deleted because they
// provide no way to specify an explicit stream
Expand Down Expand Up @@ -110,7 +106,7 @@ class device_buffer {
*/
explicit device_buffer(std::size_t size,
cuda_stream_view stream,
async_resource_ref mr = mr::get_current_device_resource())
mr::device_memory_resource* mr = mr::get_current_device_resource())
: _stream{stream}, _mr{mr}
{
allocate_async(size);
Expand Down Expand Up @@ -138,7 +134,7 @@ class device_buffer {
device_buffer(void const* source_data,
std::size_t size,
cuda_stream_view stream,
async_resource_ref mr = rmm::mr::get_current_device_resource())
mr::device_memory_resource* mr = mr::get_current_device_resource())
: _stream{stream}, _mr{mr}
{
allocate_async(size);
Expand Down Expand Up @@ -168,7 +164,7 @@ class device_buffer {
*/
device_buffer(device_buffer const& other,
cuda_stream_view stream,
async_resource_ref mr = rmm::mr::get_current_device_resource())
rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource())
: device_buffer{other.data(), other.size(), stream, mr}
{
}
Expand Down Expand Up @@ -240,6 +236,7 @@ class device_buffer {
~device_buffer() noexcept
{
deallocate_async();
_mr = nullptr;
_stream = cuda_stream_view{};
}

Expand Down Expand Up @@ -400,23 +397,16 @@ class device_buffer {
/**
* @briefreturn{The async_resource_ref used to allocate and deallocate}
*/
[[nodiscard]] async_resource_ref memory_resource() const noexcept { return _mr; }

/**
* @brief Enables the `cuda::mr::device_accessible` property
*
* This property declares that a `device_buffer` provides device accessible memory
*/
friend void get_property(device_buffer const&, cuda::mr::device_accessible) noexcept {}
[[nodiscard]] mr::device_memory_resource* memory_resource() const noexcept { return _mr; }

private:
void* _data{nullptr}; ///< Pointer to device memory allocation
std::size_t _size{}; ///< Requested size of the device memory allocation
std::size_t _capacity{}; ///< The actual size of the device memory allocation
cuda_stream_view _stream{}; ///< Stream to use for device memory deallocation
async_resource_ref _mr{
rmm::mr::get_current_device_resource()}; ///< The memory resource used to
///< allocate/deallocate device memory
mr::device_memory_resource* _mr{
mr::get_current_device_resource()}; ///< The memory resource used to
///< allocate/deallocate device memory

/**
* @brief Allocates the specified amount of memory and updates the size/capacity accordingly.
Expand All @@ -431,7 +421,7 @@ class device_buffer {
{
_size = bytes;
_capacity = bytes;
_data = (bytes > 0) ? _mr.allocate_async(bytes, stream()) : nullptr;
_data = (bytes > 0) ? memory_resource()->allocate(bytes, stream()) : nullptr;
}

/**
Expand All @@ -445,7 +435,7 @@ class device_buffer {
*/
void deallocate_async() noexcept
{
if (capacity() > 0) { _mr.deallocate_async(data(), capacity(), stream()); }
if (capacity() > 0) { memory_resource()->deallocate(data(), capacity(), stream()); }
_size = 0;
_capacity = 0;
_data = nullptr;
Expand Down
27 changes: 9 additions & 18 deletions include/rmm/device_uvector.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,8 +26,6 @@
#include <cstddef>
#include <vector>

#include <cuda/memory_resource>

namespace rmm {
/**
* @addtogroup data_containers
Expand Down Expand Up @@ -74,7 +72,6 @@ namespace rmm {
*/
template <typename T>
class device_uvector {
using async_resource_ref = cuda::mr::async_resource_ref<cuda::mr::device_accessible>;
static_assert(std::is_trivially_copyable<T>::value,
"device_uvector only supports types that are trivially copyable.");

Expand Down Expand Up @@ -124,9 +121,10 @@ class device_uvector {
* @param stream The stream on which to perform the allocation
* @param mr The resource used to allocate the device storage
*/
explicit device_uvector(std::size_t size,
cuda_stream_view stream,
async_resource_ref mr = rmm::mr::get_current_device_resource())
explicit device_uvector(
std::size_t size,
cuda_stream_view stream,
rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource())
: _storage{elements_to_bytes(size), stream, mr}
{
}
Expand All @@ -140,9 +138,10 @@ class device_uvector {
* @param stream The stream on which to perform the copy
* @param mr The resource used to allocate device memory for the new vector
*/
explicit device_uvector(device_uvector const& other,
cuda_stream_view stream,
async_resource_ref mr = rmm::mr::get_current_device_resource())
explicit device_uvector(
device_uvector const& other,
cuda_stream_view stream,
rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource())
: _storage{other._storage, stream, mr}
{
}
Expand Down Expand Up @@ -527,19 +526,11 @@ class device_uvector {
/**
* @briefreturn{The async_resource_ref used to allocate and deallocate the device storage}
*/
[[nodiscard]] cuda::mr::async_resource_ref<cuda::mr::device_accessible> memory_resource()
const noexcept
[[nodiscard]] mr::device_memory_resource* memory_resource() const noexcept
{
return _storage.memory_resource();
}

/**
* @brief Enables the `cuda::mr::device_accessible` property
*
* This property declares that a `device_uvector` provides device accessible memory
*/
friend void get_property(device_uvector const&, cuda::mr::device_accessible) noexcept {}

/**
* @briefreturn{Stream most recently specified for allocation/deallocation}
*/
Expand Down
32 changes: 16 additions & 16 deletions python/rmm/tests/test_rmm.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -630,11 +630,11 @@ def test_statistics_resource_adaptor(stats_mr):
del buffers[i]

assert stats_mr.allocation_counts == {
"current_bytes": 5040,
"current_bytes": 5000,
"current_count": 5,
"peak_bytes": 10080,
"peak_bytes": 10000,
"peak_count": 10,
"total_bytes": 10080,
"total_bytes": 10000,
"total_count": 10,
}

Expand All @@ -646,19 +646,19 @@ def test_statistics_resource_adaptor(stats_mr):
buffers.append(rmm.DeviceBuffer(size=1000))

assert mr2.allocation_counts == {
"current_bytes": 2016,
"current_bytes": 2000,
"current_count": 2,
"peak_bytes": 2016,
"peak_bytes": 2000,
"peak_count": 2,
"total_bytes": 2016,
"total_bytes": 2000,
"total_count": 2,
}
assert stats_mr.allocation_counts == {
"current_bytes": 7056,
"current_bytes": 7000,
"current_count": 7,
"peak_bytes": 10080,
"peak_bytes": 10000,
"peak_count": 10,
"total_bytes": 12096,
"total_bytes": 12000,
"total_count": 12,
}

Expand All @@ -668,17 +668,17 @@ def test_statistics_resource_adaptor(stats_mr):
assert mr2.allocation_counts == {
"current_bytes": 0,
"current_count": 0,
"peak_bytes": 2016,
"peak_bytes": 2000,
"peak_count": 2,
"total_bytes": 2016,
"total_bytes": 2000,
"total_count": 2,
}
assert stats_mr.allocation_counts == {
"current_bytes": 0,
"current_count": 0,
"peak_bytes": 10080,
"peak_bytes": 10000,
"peak_count": 10,
"total_bytes": 12096,
"total_bytes": 12000,
"total_count": 12,
}
gc.collect()
Expand All @@ -696,7 +696,7 @@ def test_tracking_resource_adaptor():
for i in range(9, 0, -2):
del buffers[i]

assert mr.get_allocated_bytes() == 5040
assert mr.get_allocated_bytes() == 5000

# Push a new Tracking adaptor
mr2 = rmm.mr.TrackingResourceAdaptor(mr, capture_stacks=True)
Expand All @@ -705,8 +705,8 @@ def test_tracking_resource_adaptor():
for _ in range(2):
buffers.append(rmm.DeviceBuffer(size=1000))

assert mr2.get_allocated_bytes() == 2016
assert mr.get_allocated_bytes() == 7056
assert mr2.get_allocated_bytes() == 2000
assert mr.get_allocated_bytes() == 7000

# Ensure we get back a non-empty string for the allocations
assert len(mr.get_outstanding_allocations_str()) > 0
Expand Down
Loading

0 comments on commit 2cf03a7

Please sign in to comment.