Skip to content

Commit

Permalink
Fix the pool size alignment issue (#16024)
Browse files Browse the repository at this point in the history
This PR fixes a pool size alignment bug.

Authors:
  - Yunsong Wang (https://github.com/PointKernel)

Approvers:
  - Mark Harris (https://github.com/harrism)
  - Vukasin Milovanovic (https://github.com/vuule)
  - David Wendt (https://github.com/davidwendt)

URL: #16024
  • Loading branch information
PointKernel authored Jun 14, 2024
1 parent 31d909b commit 987879c
Showing 1 changed file with 5 additions and 7 deletions.
12 changes: 5 additions & 7 deletions cpp/src/utilities/pinned_memory.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -43,9 +43,11 @@ class fixed_pinned_pool_memory_resource {

public:
fixed_pinned_pool_memory_resource(size_t size)
: pool_size_{size}, pool_{new host_pooled_mr(upstream_mr_, size, size)}
: // rmm requires the pool size to be a multiple of 256 bytes
pool_size_{rmm::align_up(size, rmm::CUDA_ALLOCATION_ALIGNMENT)},
pool_{new host_pooled_mr(upstream_mr_, pool_size_, pool_size_)}
{
if (pool_size_ == 0) { return; }
CUDF_LOG_INFO("Pinned pool size = {}", pool_size_);

// Allocate full size from the pinned pool to figure out the beginning and end address
pool_begin_ = pool_->allocate_async(pool_size_, stream_);
Expand Down Expand Up @@ -145,12 +147,8 @@ CUDF_EXPORT rmm::host_device_async_resource_ref& make_default_pinned_mr(
return std::min(total / 200, size_t{100} * 1024 * 1024);
}();

// rmm requires the pool size to be a multiple of 256 bytes
auto const aligned_size = rmm::align_up(size, rmm::RMM_DEFAULT_HOST_ALIGNMENT);
CUDF_LOG_INFO("Pinned pool size = {}", aligned_size);

// make the pool with max size equal to the initial size
return fixed_pinned_pool_memory_resource{aligned_size};
return fixed_pinned_pool_memory_resource{size};
}();

static rmm::host_device_async_resource_ref mr_ref{mr};
Expand Down

0 comments on commit 987879c

Please sign in to comment.