Skip to content

Commit

Permalink
more
Browse files Browse the repository at this point in the history
  • Loading branch information
vuule committed Mar 1, 2025
1 parent d91b2f9 commit 4678c91
Showing 1 changed file with 3 additions and 5 deletions.
8 changes: 3 additions & 5 deletions cpp/src/utilities/host_memory.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
#include <rmm/cuda_device.hpp>
#include <rmm/mr/device/pool_memory_resource.hpp>
#include <rmm/mr/pinned_host_memory_resource.hpp>

namespace cudf {

namespace {
Expand All @@ -48,17 +49,14 @@ class fixed_pinned_pool_memory_resource {
pool_size_{rmm::align_up(size, rmm::CUDA_ALLOCATION_ALIGNMENT)},
pool_{new host_pooled_mr(upstream_mr_, pool_size_, pool_size_)}
{
CUDF_LOG_INFO("Pinned pool size = %zu", 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_);
pool_end_ = static_cast<void*>(static_cast<uint8_t*>(pool_begin_) + pool_size_);
pool_->deallocate_async(pool_begin_, pool_size_, stream_);
}

fixed_pinned_pool_memory_resource(fixed_pinned_pool_memory_resource const&) = delete;
fixed_pinned_pool_memory_resource(fixed_pinned_pool_memory_resource&&) = delete;
fixed_pinned_pool_memory_resource& operator=(fixed_pinned_pool_memory_resource const&) = delete;
fixed_pinned_pool_memory_resource& operator=(fixed_pinned_pool_memory_resource&&) = delete;

void* allocate_async(std::size_t bytes, std::size_t alignment, cuda::stream_ref stream)
{
if (bytes <= pool_size_) {
Expand Down

0 comments on commit 4678c91

Please sign in to comment.