thrust_allocator_adaptor.hpp

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION.
 * SPDX-License-Identifier: Apache-2.0
 */
 
#pragma once
 
#include <rmm/cuda_device.hpp>
#include <rmm/detail/cuda_memory_resource.hpp>
#include <rmm/detail/export.hpp>
#include <rmm/detail/thrust_namespace.h>
#include <rmm/mr/device/per_device_resource.hpp>
#include <rmm/resource_ref.hpp>
 
#include <thrust/device_malloc_allocator.h>
#include <thrust/device_ptr.h>
#include <thrust/memory.h>
 
namespace RMM_NAMESPACE {
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>;  
  };
 
  thrust_allocator() = default;
 
  explicit thrust_allocator(cuda_stream_view stream) : _stream{stream} {}
 
  thrust_allocator(cuda_stream_view stream, rmm::device_async_resource_ref mr)
    : _stream{stream}, _mr(mr)
  {
  }
 
  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))));
  }
 
  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));
  }
 
  [[nodiscard]] rmm::device_async_resource_ref get_upstream_resource() const noexcept
  {
    return _mr;
  }
 
  [[nodiscard]] cuda_stream_view stream() const noexcept { return _stream; }
 
  friend void get_property(thrust_allocator const&, cuda::mr::device_accessible) noexcept {}
 
 private:
  cuda_stream_view _stream{};
  rmm::device_async_resource_ref _mr{rmm::mr::get_current_device_resource_ref()};
  cuda_device_id _device{get_current_cuda_device()};
};  // end of group
}  // namespace mr
}  // namespace RMM_NAMESPACE