1 #pragma once 2 3 #include <c10/core/Allocator.h> 4 #include <c10/core/DeviceType.h> 5 6 // Use of c10::hip namespace here makes hipification easier, because 7 // I don't have to also fix namespaces. Sorry! 8 namespace c10 { namespace hip { 9 10 // Takes a valid HIPAllocator (of any sort) and turns it into 11 // an allocator pretending to be a CUDA allocator. See 12 // Note [Masquerading as CUDA] 13 class HIPAllocatorMasqueradingAsCUDA final : public Allocator { 14 Allocator* allocator_; 15 public: HIPAllocatorMasqueradingAsCUDA(Allocator * allocator)16 explicit HIPAllocatorMasqueradingAsCUDA(Allocator* allocator) 17 : allocator_(allocator) {} allocate(size_t size)18 DataPtr allocate(size_t size) override { 19 DataPtr r = allocator_->allocate(size); 20 r.unsafe_set_device(Device(c10::DeviceType::CUDA, r.device().index())); 21 return r; 22 } raw_deleter()23 DeleterFnPtr raw_deleter() const override { 24 return allocator_->raw_deleter(); 25 } copy_data(void * dest,const void * src,std::size_t count)26 void copy_data(void* dest, const void* src, std::size_t count) const final { 27 allocator_->copy_data(dest, src, count); 28 } 29 }; 30 31 }} // namespace c10::hip 32