system_memory_resource.hpp

Go to the documentation of this file.

/*
 * Copyright (c) 2024, 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/cuda_device.hpp>
#include <rmm/cuda_stream_view.hpp>
#include <rmm/detail/error.hpp>
#include <rmm/detail/export.hpp>
#include <rmm/mr/device/device_memory_resource.hpp>
 
namespace RMM_NAMESPACE {
namespace mr {
 
namespace detail {
static bool is_system_memory_supported(cuda_device_id device_id)
{
  int pageableMemoryAccess;
  RMM_CUDA_TRY(cudaDeviceGetAttribute(
    &pageableMemoryAccess, cudaDevAttrPageableMemoryAccess, device_id.value()));
  return pageableMemoryAccess == 1;
}
}  // namespace detail
 
class system_memory_resource final : public device_memory_resource {
 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() override                    = 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;  
 
 private:
  void* do_allocate(std::size_t bytes, [[maybe_unused]] cuda_stream_view stream) override
  {
    try {
      return rmm::detail::aligned_host_allocate(
        bytes, CUDA_ALLOCATION_ALIGNMENT, [](std::size_t size) { return ::operator new(size); });
    } catch (std::bad_alloc const& e) {
      RMM_FAIL("Failed to allocate memory: " + std::string{e.what()}, rmm::out_of_memory);
    }
  }
 
  void do_deallocate(void* ptr,
                     [[maybe_unused]] std::size_t bytes,
                     cuda_stream_view stream) override
  {
    // 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.
    stream.synchronize();
 
    rmm::detail::aligned_host_deallocate(
      ptr, bytes, CUDA_ALLOCATION_ALIGNMENT, [](void* ptr) { ::operator delete(ptr); });
  }
 
  [[nodiscard]] bool do_is_equal(device_memory_resource const& other) const noexcept override
  {
    return dynamic_cast<system_memory_resource const*>(&other) != nullptr;
  }
  friend void get_property(system_memory_resource const&, cuda::mr::device_accessible) noexcept {}
 
  friend void get_property(system_memory_resource const&, cuda::mr::host_accessible) noexcept {}
};
 
// static property checks
static_assert(cuda::mr::async_resource_with<system_memory_resource, cuda::mr::device_accessible>);
static_assert(cuda::mr::async_resource_with<system_memory_resource, cuda::mr::host_accessible>);  // end of group
}  // namespace mr
}  // namespace RMM_NAMESPACE