pinned_host_memory_resource.hpp

Go to the documentation of this file.

/*
 * Copyright (c) 2024-2025, NVIDIA CORPORATION.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#pragma once
 
#include <rmm/aligned.hpp>
#include <rmm/detail/aligned.hpp>
#include <rmm/detail/error.hpp>
#include <rmm/detail/export.hpp>
#include <rmm/detail/nvtx/ranges.hpp>
 
#include <cuda/memory_resource>
#include <cuda/stream_ref>
#include <cuda_runtime_api.h>
 
#include <cstddef>
 
namespace RMM_NAMESPACE {
namespace mr {
 
class pinned_host_memory_resource {
 public:
  // Disable clang-tidy complaining about the easily swappable size and alignment parameters
  // of allocate and deallocate
  // NOLINTBEGIN(bugprone-easily-swappable-parameters)
 
  static void* allocate(std::size_t bytes,
                        [[maybe_unused]] std::size_t alignment = rmm::RMM_DEFAULT_HOST_ALIGNMENT)
  {
    RMM_FUNC_RANGE();
 
    // don't allocate anything if the user requested zero bytes
    if (0 == bytes) { return nullptr; }
 
    return rmm::detail::aligned_host_allocate(bytes, alignment, [](std::size_t size) {
      void* ptr{nullptr};
      RMM_CUDA_TRY_ALLOC(cudaHostAlloc(&ptr, size, cudaHostAllocDefault), size);
      return ptr;
    });
  }
 
  static void deallocate(void* ptr,
                         std::size_t bytes,
                         std::size_t alignment = rmm::RMM_DEFAULT_HOST_ALIGNMENT) noexcept
  {
    RMM_FUNC_RANGE();
 
    rmm::detail::aligned_host_deallocate(
      ptr, bytes, alignment, [](void* ptr) { RMM_ASSERT_CUDA_SUCCESS(cudaFreeHost(ptr)); });
  }
 
  static void* allocate_async(std::size_t bytes, [[maybe_unused]] cuda::stream_ref stream)
  {
    RMM_FUNC_RANGE();
 
    return allocate(bytes);
  }
 
  static void* allocate_async(std::size_t bytes,
                              std::size_t alignment,
                              [[maybe_unused]] cuda::stream_ref stream)
  {
    RMM_FUNC_RANGE();
 
    return allocate(bytes, alignment);
  }
 
  static void deallocate_async(void* ptr,
                               std::size_t bytes,
                               [[maybe_unused]] cuda::stream_ref stream) noexcept
  {
    RMM_FUNC_RANGE();
 
    return deallocate(ptr, bytes);
  }
 
  static void deallocate_async(void* ptr,
                               std::size_t bytes,
                               std::size_t alignment,
                               [[maybe_unused]] cuda::stream_ref stream) noexcept
  {
    RMM_FUNC_RANGE();
 
    return deallocate(ptr, bytes, alignment);
  }
  // NOLINTEND(bugprone-easily-swappable-parameters)
 
  bool operator==(const pinned_host_memory_resource&) const { return true; }
 
  bool operator!=(const pinned_host_memory_resource&) const { return false; }
 
  friend void get_property(pinned_host_memory_resource const&, cuda::mr::device_accessible) noexcept
  {
  }
 
  friend void get_property(pinned_host_memory_resource const&, cuda::mr::host_accessible) noexcept
  {
  }
};
 
static_assert(cuda::mr::async_resource_with<pinned_host_memory_resource,
                                            cuda::mr::device_accessible,
                                            cuda::mr::host_accessible>);
  // end of group
}  // namespace mr
}  // namespace RMM_NAMESPACE