thrust_allocator_adaptor.hpp

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 */
 
#pragma once
 
#include <rmm/aligned.hpp>
#include <rmm/cuda_device.hpp>
#include <rmm/detail/exec_check_disable.hpp>
#include <rmm/detail/export.hpp>
#include <rmm/detail/thrust_namespace.h>
#include <rmm/mr/per_device_resource.hpp>
#include <rmm/resource_ref.hpp>
 
#include <cuda/memory_resource>
#include <cuda/stream>
#include <thrust/device_malloc_allocator.h>
#include <thrust/device_ptr.h>
#include <thrust/memory.h>
 
RMM_NAMESPACE_BEGIN
namespace mr {
template <typename T>
class thrust_allocator : public thrust::device_malloc_allocator<T> {
 public:
  using Base      = thrust::device_malloc_allocator<T>;  
  using pointer   = typename Base::pointer;              
  using size_type = typename Base::size_type;            

  template <typename U>
  struct rebind {
    using other = thrust_allocator<U>;  
  };

  RMM_EXEC_CHECK_DISABLE
  thrust_allocator() : _stream{cuda::stream_ref{cudaStream_t{cudaStreamDefault}}} {}

  RMM_EXEC_CHECK_DISABLE
  explicit thrust_allocator(cuda::stream_ref stream) : _stream{stream} {}

  RMM_EXEC_CHECK_DISABLE
  thrust_allocator(cuda::stream_ref stream, cuda::mr::any_resource<cuda::mr::device_accessible> mr)
    : _stream{stream}, _mr(std::move(mr))
  {
  }

  RMM_EXEC_CHECK_DISABLE
  thrust_allocator(thrust_allocator const& other)
    : Base(other), _stream{other._stream}, _mr(other._mr), _device{other._device}
  {
  }

  RMM_EXEC_CHECK_DISABLE
  thrust_allocator(thrust_allocator&& other) noexcept
    : Base(std::move(other)),
      _stream{other._stream},
      _mr(std::move(other._mr)),
      _device{other._device}
  {
  }

  thrust_allocator& operator=(thrust_allocator const&) = default;
  thrust_allocator& operator=(thrust_allocator&&) noexcept = default;

  RMM_EXEC_CHECK_DISABLE
  template <typename U>
  thrust_allocator(thrust_allocator<U> const& other)
    : _mr(other.resource()), _stream{other.stream()}, _device{other._device}
  {
  }

  pointer allocate(size_type num)
  {
    cuda_set_device_raii dev{_device};
    return thrust::device_pointer_cast(
      static_cast<T*>(_mr.allocate(_stream, num * sizeof(T), rmm::CUDA_ALLOCATION_ALIGNMENT)));
  }

  void deallocate(pointer ptr, size_type num) noexcept
  {
    cuda_set_device_raii dev{_device};
    return _mr.deallocate(
      _stream, thrust::raw_pointer_cast(ptr), num * sizeof(T), rmm::CUDA_ALLOCATION_ALIGNMENT);
  }

  [[nodiscard]] rmm::device_async_resource_ref get_upstream_resource() const noexcept
  {
    return rmm::device_async_resource_ref{_mr};
  }

  [[nodiscard]] cuda::stream_ref stream() const noexcept { return _stream; }

  RMM_CONSTEXPR_FRIEND void get_property(thrust_allocator const&,
                                         cuda::mr::device_accessible) noexcept
  {
  }
 
 private:
  cuda::stream_ref _stream{cuda::invalid_stream_t{}};
  mutable cuda::mr::any_resource<cuda::mr::device_accessible> _mr{
    rmm::mr::get_current_device_resource_ref()};
  cuda_device_id _device{get_current_cuda_device()};
};
  // end of group
}  // namespace mr
RMM_NAMESPACE_END