From 987879ca4bdcae0d959266fd39196123007fa45e Mon Sep 17 00:00:00 2001 From: Yunsong Wang Date: Thu, 13 Jun 2024 19:27:11 -0700 Subject: [PATCH] Fix the pool size alignment issue (#16024) 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: https://github.com/rapidsai/cudf/pull/16024 --- cpp/src/utilities/pinned_memory.cpp | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/cpp/src/utilities/pinned_memory.cpp b/cpp/src/utilities/pinned_memory.cpp index 5d2e3ac332a..e90b7969b4d 100644 --- a/cpp/src/utilities/pinned_memory.cpp +++ b/cpp/src/utilities/pinned_memory.cpp @@ -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_); @@ -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};