system_memory_resource.hpp

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
 * SPDX-License-Identifier: Apache-2.0
 */
#pragma once
 
#include <rmm/aligned.hpp>
#include <rmm/cuda_device.hpp>
#include <rmm/detail/aligned.hpp>
#include <rmm/detail/error.hpp>
#include <rmm/detail/export.hpp>
#include <rmm/detail/format.hpp>
 
#include <cuda/memory_resource>
#include <cuda/stream_ref>
#include <cuda_runtime_api.h>
 
#include <cstddef>
#include <string>
 
namespace RMM_NAMESPACE {
namespace mr {
 
namespace detail {
static bool is_system_memory_supported(cuda_device_id device_id)
{
  // Check if pageable memory access is supported
  int pageableMemoryAccess;
  RMM_CUDA_TRY(cudaDeviceGetAttribute(
    &pageableMemoryAccess, cudaDevAttrPageableMemoryAccess, device_id.value()));
  return pageableMemoryAccess == 1;
}
}  // namespace detail
 
class system_memory_resource final {
 public:
  system_memory_resource()
  {
    RMM_EXPECTS(rmm::mr::detail::is_system_memory_supported(rmm::get_current_cuda_device()),
                "System memory allocator is not supported with this hardware/software version.");
  }
  ~system_memory_resource()                             = default;
  system_memory_resource(system_memory_resource const&) = default;  
  system_memory_resource(system_memory_resource&&)      = default;  
  system_memory_resource& operator=(system_memory_resource const&) =
    default;  
  system_memory_resource& operator=(system_memory_resource&&) =
    default;  
 
  void* allocate([[maybe_unused]] cuda::stream_ref stream,
                 std::size_t bytes,
                 [[maybe_unused]] std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT)
  {
    try {
      return rmm::detail::aligned_host_allocate(
        bytes, rmm::CUDA_ALLOCATION_ALIGNMENT, [](std::size_t size) {
          return ::operator new(size);
        });
    } catch (std::bad_alloc const& e) {
      auto const msg = std::string("Failed to allocate ") + rmm::detail::format_bytes(bytes) +
                       std::string("of memory: ") + e.what();
      RMM_FAIL(msg.c_str(), rmm::out_of_memory);
    }
  }
 
  void deallocate(cuda::stream_ref stream,
                  void* ptr,
                  std::size_t bytes,
                  [[maybe_unused]] std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) noexcept
  {
    // With `cudaFree`, the CUDA runtime keeps track of dependent operations and does implicit
    // synchronization. However, with SAM, since `free` is immediate, we need to wait for in-flight
    // CUDA operations to finish before freeing the memory, to avoid potential use-after-free errors
    // or race conditions.
    RMM_ASSERT_CUDA_SUCCESS_SAFE_SHUTDOWN(cudaStreamSynchronize(stream.get()));
 
    rmm::detail::aligned_host_deallocate(
      ptr, bytes, rmm::CUDA_ALLOCATION_ALIGNMENT, [](void* ptr) { ::operator delete(ptr); });
  }
 
  void* allocate_sync(std::size_t bytes, std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT)
  {
    auto* ptr = allocate(cuda::stream_ref{cudaStream_t{nullptr}}, bytes, alignment);
    RMM_CUDA_TRY(cudaStreamSynchronize(cudaStream_t{nullptr}));
    return ptr;
  }
 
  void deallocate_sync(void* ptr,
                       std::size_t bytes,
                       std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) noexcept
  {
    deallocate(cuda::stream_ref{cudaStream_t{nullptr}}, ptr, bytes, alignment);
  }
 
  RMM_CONSTEXPR_FRIEND void get_property(system_memory_resource const&,
                                         cuda::mr::device_accessible) noexcept
  {
  }
 
  RMM_CONSTEXPR_FRIEND void get_property(system_memory_resource const&,
                                         cuda::mr::host_accessible) noexcept
  {
  }
 
  [[nodiscard]] bool operator==(system_memory_resource const&) const noexcept { return true; }
 
  [[nodiscard]] bool operator!=(system_memory_resource const&) const noexcept { return false; }
};
 
// static property checks
static_assert(cuda::mr::synchronous_resource<system_memory_resource>);
static_assert(cuda::mr::resource<system_memory_resource>);
static_assert(
  cuda::mr::synchronous_resource_with<system_memory_resource, cuda::mr::device_accessible>);
static_assert(
  cuda::mr::synchronous_resource_with<system_memory_resource, cuda::mr::host_accessible>);
static_assert(cuda::mr::resource_with<system_memory_resource, cuda::mr::device_accessible>);
static_assert(cuda::mr::resource_with<system_memory_resource, cuda::mr::host_accessible>);  // end of group
}  // namespace mr
}  // namespace RMM_NAMESPACE