Skip to content

[Refactor][Runtime] Always specify device in allocator interface #16738

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Mar 19, 2024
Merged

[Refactor][Runtime] Always specify device in allocator interface #16738

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 5 additions & 6 deletions include/tvm/runtime/memory/memory_manager.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,19 +71,22 @@ class Allocator {
/*! \brief Return the allocator type. */
inline AllocatorType type() const { return type_; }
/*! \brief Allocate a buffer given a size, alignment and type.
* \param dev The device where the array is allocated.
* \param nbytes The size of the buffer.
* \param alignment The alignment of the buffer.
* \param type_hint A type hint to the allocator.
* \return A sized allocation in the form of a buffer.
*/
TVM_DLL virtual Buffer Alloc(size_t nbytes, size_t alignment, DLDataType type_hint) = 0;
TVM_DLL virtual Buffer Alloc(Device dev, size_t nbytes, size_t alignment,
DLDataType type_hint) = 0;
/*! \brief Allocate a buffer given a shape and type.
* \param dev The device where the array is allocated.
* \param shape The shape of the tensor.
* \param type_hint A type hint to the allocator.
* \param mem_scope A memory scope of the buffer.
* \return A sized allocation in the form of a buffer.
*/
TVM_DLL virtual Buffer Alloc(ShapeTuple shape, DLDataType type_hint,
TVM_DLL virtual Buffer Alloc(Device dev, ShapeTuple shape, DLDataType type_hint,
const std::string& mem_scope = "") = 0;
/*! \brief Free a buffer allocated by the allocator.
* \param buffer The buffer to free.
Expand All @@ -96,10 +99,6 @@ class Allocator {
*/
TVM_DLL virtual size_t UsedMemory() const = 0;

protected:
TVM_DLL virtual Buffer Alloc(Device dev, ShapeTuple shape, DLDataType type_hint,
const std::string& mem_scope);

private:
AllocatorType type_;
};
Expand Down
10 changes: 5 additions & 5 deletions src/runtime/memory/memory_manager.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -138,12 +138,12 @@ Allocator* MemoryManager::GetOrCreateAllocator(Device dev, AllocatorType type) {
switch (type) {
case kNaive: {
VLOG(1) << "New naive allocator for " << dev;
alloc.reset(new NaiveAllocator(dev));
alloc.reset(new NaiveAllocator());
break;
}
case kPooled: {
VLOG(1) << "New pooled allocator for " << dev;
alloc.reset(new PooledAllocator(dev));
alloc.reset(new PooledAllocator());
break;
}
default:
Expand Down Expand Up @@ -194,9 +194,9 @@ NDArray Allocator::Empty(ShapeTuple shape, DLDataType dtype, DLDevice dev,
size_t alignment = GetDataAlignment(container->dl_tensor);
Buffer* buffer = new Buffer;
if (!mem_scope.defined() || mem_scope.value().empty() || mem_scope.value() == "global") {
*buffer = this->Alloc(size, alignment, dtype);
*buffer = this->Alloc(dev, size, alignment, dtype);
} else {
*buffer = this->Alloc(shape, dtype, mem_scope.value());
*buffer = this->Alloc(dev, shape, dtype, mem_scope.value());
}
container->manager_ctx = reinterpret_cast<void*>(buffer);
container->dl_tensor.data = buffer->data;
Expand All @@ -210,7 +210,7 @@ Buffer Allocator::Alloc(Device dev, ShapeTuple shape, DLDataType type_hint,
NDArray::Container container(nullptr, shape, type_hint, dev);
size_t size = DeviceAPI::Get(dev)->GetDataSize(container.dl_tensor);
size_t alignment = GetDataAlignment(container.dl_tensor);
return Alloc(size, alignment, type_hint);
return Alloc(dev, size, alignment, type_hint);
}
LOG(FATAL) << "Allocator cannot allocate data space with "
<< "specified memory scope: " << mem_scope;
Expand Down
22 changes: 11 additions & 11 deletions src/runtime/memory/naive_allocator.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,46 +35,47 @@ namespace memory {

class NaiveAllocator final : public Allocator {
public:
explicit NaiveAllocator(Device dev) : Allocator(kNaive), used_memory_(0), device_(dev) {}
explicit NaiveAllocator() : Allocator(kNaive), used_memory_(0) {}

Buffer Alloc(size_t nbytes, size_t alignment, DLDataType type_hint) override {
Buffer Alloc(Device dev, size_t nbytes, size_t alignment, DLDataType type_hint) final {
Buffer buf;
buf.device = device_;
buf.device = dev;
buf.size = nbytes;
buf.alloc_type = kNaive;
buf.data = DeviceAPI::Get(device_)->AllocDataSpace(device_, nbytes, alignment, type_hint);
buf.data = DeviceAPI::Get(dev)->AllocDataSpace(dev, nbytes, alignment, type_hint);
used_memory_.fetch_add(nbytes, std::memory_order_relaxed);
DLOG(INFO) << "allocate " << nbytes << " B, used memory " << used_memory_ << " B";
return buf;
}

Buffer Alloc(ShapeTuple shape, DLDataType type_hint, const std::string& mem_scope) override {
Buffer Alloc(Device dev, ShapeTuple shape, DLDataType type_hint,
const std::string& mem_scope) final {
Buffer buf;
size_t nbytes = 1;
for (int i = 0; i < static_cast<int>(shape.size()); ++i) {
nbytes *= static_cast<size_t>(shape[i]);
}
nbytes *= (type_hint.bits * type_hint.lanes + 7) / 8;
buf.device = device_;
buf.device = dev;
if (mem_scope.empty() || mem_scope == "global") {
auto tmp_buf = Allocator::Alloc(device_, shape, type_hint, mem_scope);
auto tmp_buf = Allocator::Alloc(dev, shape, type_hint, mem_scope);
buf.size = tmp_buf.size;
buf.data = tmp_buf.data;
buf.alloc_type = kNaive;
return buf;
}

buf.size = nbytes;
buf.data = DeviceAPI::Get(device_)->AllocDataSpace(device_, shape.size(), shape.data(),
type_hint, String(mem_scope));
buf.data = DeviceAPI::Get(dev)->AllocDataSpace(dev, shape.size(), shape.data(), type_hint,
String(mem_scope));
used_memory_.fetch_add(nbytes, std::memory_order_relaxed);
DLOG(INFO) << "allocate " << nbytes << " B, used memory " << used_memory_ << " B";
buf.alloc_type = kNaive;
return buf;
}

void Free(const Buffer& buffer) override {
DeviceAPI::Get(device_)->FreeDataSpace(buffer.device, buffer.data);
DeviceAPI::Get(buffer.device)->FreeDataSpace(buffer.device, buffer.data);
used_memory_.fetch_sub(buffer.size, std::memory_order_relaxed);
DLOG(INFO) << "free " << buffer.size << " B, used memory " << used_memory_ << " B";
}
Expand All @@ -83,7 +84,6 @@ class NaiveAllocator final : public Allocator {

private:
std::atomic<size_t> used_memory_;
Device device_;
};

} // namespace memory
Expand Down
18 changes: 9 additions & 9 deletions src/runtime/memory/pooled_allocator.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,12 +40,12 @@ class PooledAllocator final : public Allocator {
public:
static constexpr size_t kDefaultPageSize = 4096;

explicit PooledAllocator(Device dev, size_t page_size = kDefaultPageSize)
: Allocator(kPooled), page_size_(page_size), used_memory_(0), device_(dev) {}
explicit PooledAllocator(size_t page_size = kDefaultPageSize)
: Allocator(kPooled), page_size_(page_size), used_memory_(0) {}

~PooledAllocator() { ReleaseAll(); }

Buffer Alloc(size_t nbytes, size_t alignment, DLDataType type_hint) override {
Buffer Alloc(Device dev, size_t nbytes, size_t alignment, DLDataType type_hint) override {
std::lock_guard<std::recursive_mutex> lock(mu_);
size_t size = ((nbytes + page_size_ - 1) / page_size_) * page_size_;
auto&& it = memory_pool_.find(size);
Expand All @@ -56,26 +56,27 @@ class PooledAllocator final : public Allocator {
return ret;
}
Buffer buf;
buf.device = device_;
buf.device = dev;
buf.size = size;
buf.alloc_type = kPooled;
try {
buf.data = DeviceAPI::Get(device_)->AllocDataSpace(device_, size, alignment, type_hint);
buf.data = DeviceAPI::Get(dev)->AllocDataSpace(dev, size, alignment, type_hint);
} catch (InternalError& err) {
LOG(WARNING) << "PooledAllocator got InternalError during allocation: " << err.message();
LOG(WARNING) << "Trying to release all unused memory and reallocate...";
ReleaseAll();
buf.data = DeviceAPI::Get(device_)->AllocDataSpace(device_, size, alignment, type_hint);
buf.data = DeviceAPI::Get(dev)->AllocDataSpace(dev, size, alignment, type_hint);
}

used_memory_.fetch_add(size, std::memory_order_relaxed);
VLOG(1) << "allocate " << size << " B, used memory " << used_memory_ << " B";
return buf;
}

Buffer Alloc(ShapeTuple shape, DLDataType type_hint, const std::string& mem_scope) override {
Buffer Alloc(Device dev, ShapeTuple shape, DLDataType type_hint,
const std::string& mem_scope) override {
if (mem_scope.empty() || mem_scope == "global") {
return Allocator::Alloc(device_, shape, type_hint, mem_scope);
return Allocator::Alloc(dev, shape, type_hint, mem_scope);
}
LOG(FATAL) << "This alloc should be implemented";
return {};
Expand Down Expand Up @@ -113,7 +114,6 @@ class PooledAllocator final : public Allocator {
std::atomic<size_t> used_memory_;
std::unordered_map<size_t, std::vector<Buffer>> memory_pool_;
std::recursive_mutex mu_;
Device device_;
};

} // namespace memory
Expand Down
3 changes: 2 additions & 1 deletion src/runtime/relax_vm/builtin.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -347,7 +347,8 @@ Storage VMAllocStorage(void* ctx_ptr, ShapeTuple buffer_shape, Index device_inde
auto* alloc = vm->allocators[device_index];
ICHECK(alloc) << "Did you forget to init the VirtualMachine with devices?";

storage_obj->buffer = alloc->Alloc(buffer_shape, dtype_hint, mem_scope);
storage_obj->buffer =
alloc->Alloc(vm->devices[device_index], buffer_shape, dtype_hint, mem_scope);
Storage storage(storage_obj);
return storage;
}
Expand Down
8 changes: 5 additions & 3 deletions src/runtime/vm/vm.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -823,6 +823,7 @@ void VirtualMachine::RunLoop(const std::vector<Index>& output_tensor_reg_indices

auto storage_obj = SimpleObjAllocator().make_object<StorageObj>();
Allocator* allocator = GetAllocator(instr.alloc_storage.device_index);
Device device = devices_[instr.alloc_storage.device_index];
ICHECK(allocator) << "Did you forget to init the VirtualMachine with devices?";

if (instr.alloc_storage.ndim > 0) {
Expand All @@ -844,15 +845,16 @@ void VirtualMachine::RunLoop(const std::vector<Index>& output_tensor_reg_indices
shape_.resize(instr.alloc_storage.ndim);
shape_.assign(instr.alloc_storage.shape,
instr.alloc_storage.shape + instr.alloc_storage.ndim);
storage_obj->buffer =
allocator->Alloc(ShapeTuple(shape_), instr.alloc_storage.dtype_hint, mem_scope);
storage_obj->buffer = allocator->Alloc(device, ShapeTuple(shape_),
instr.alloc_storage.dtype_hint, mem_scope);
} else {
auto size = LoadScalarInt(instr.alloc_storage.allocation_size);
auto alignment = instr.alloc_storage.alignment;
VLOG(2) << "allocating with allocation_size=" << size << ", alignment=" << alignment
<< ", dtype_hint=" << DLDataType2String(instr.alloc_storage.dtype_hint)
<< ", device_index=" << instr.alloc_storage.device_index;
storage_obj->buffer = allocator->Alloc(size, alignment, instr.alloc_storage.dtype_hint);
storage_obj->buffer =
allocator->Alloc(device, size, alignment, instr.alloc_storage.dtype_hint);
}
Storage storage(storage_obj);
WriteRegister(instr.dst, storage);
Expand Down
20 changes: 10 additions & 10 deletions tests/cpp/runtime/memory/memory_manager_tests.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@ TEST_F(TvmVMMemoryManagerTest, NaiveAllocBasic) {
Device dev = {kDLCPU, 0};
Allocator* allocator = MemoryManagerWrapper::GetOrCreateAllocator(dev, kNaive);
EXPECT_EQ(allocator->UsedMemory(), 0);
auto buff = allocator->Alloc(64, 32, DataType::Float(32));
auto buff = allocator->Alloc(dev, 64, 32, DataType::Float(32));
EXPECT_EQ(allocator->UsedMemory(), 64);
allocator->Free(buff);
EXPECT_EQ(allocator->UsedMemory(), 0);
Expand All @@ -65,7 +65,7 @@ TEST_F(TvmVMMemoryManagerTest, PooledAllocBasic) {
size_t size = ((nbytes + page_size - 1) / page_size) * page_size;
Allocator* allocator = MemoryManagerWrapper::GetOrCreateAllocator(dev, kPooled);
EXPECT_EQ(allocator->UsedMemory(), 0);
auto buff = allocator->Alloc(nbytes, 32, DataType::Float(32));
auto buff = allocator->Alloc(dev, nbytes, 32, DataType::Float(32));
EXPECT_EQ(allocator->UsedMemory(), size);
allocator->Free(buff);
EXPECT_EQ(allocator->UsedMemory(), size);
Expand Down Expand Up @@ -108,13 +108,13 @@ TEST_F(TvmVMMemoryManagerTest, NaiveAllocWithShape) {
auto dt = DataType::Float(32);
size_t nbytes = 1 * 3 * 6 * 6 * dt.bytes();
ShapeTuple shape = {1, 3, 6, 6};
auto buff = allocator->Alloc(shape, dt);
auto buff = allocator->Alloc(dev, shape, dt);
EXPECT_EQ(allocator->UsedMemory(), nbytes);
allocator->Free(buff);
EXPECT_EQ(allocator->UsedMemory(), 0);

try {
auto texture = allocator->Alloc(shape, dt, "global.texture");
auto texture = allocator->Alloc(dev, shape, dt, "global.texture");
(void)texture;
FAIL();
} catch (std::exception& e) {
Expand All @@ -134,13 +134,13 @@ TEST_F(TvmVMMemoryManagerTest, PooledAllocWithShape) {
size_t page_size = PooledAllocator::kDefaultPageSize;
size_t size = ((nbytes + page_size - 1) / page_size) * page_size;
ShapeTuple shape = {1, 3, 6, 6};
auto buff = allocator->Alloc(shape, dt);
auto buff = allocator->Alloc(dev, shape, dt);
EXPECT_EQ(allocator->UsedMemory(), size);
allocator->Free(buff);
EXPECT_EQ(allocator->UsedMemory(), size);

try {
auto texture = allocator->Alloc(shape, dt, "global.texture");
auto texture = allocator->Alloc(dev, shape, dt, "global.texture");
(void)texture;
FAIL();
} catch (std::exception& e) {
Expand All @@ -162,12 +162,12 @@ TEST_F(TvmVMMemoryManagerTest, NaiveAllocOpenCLTexture) {
auto dt = DataType::Float(32);
size_t nbytes = 1 * 3 * 6 * 6 * dt.bytes();
ShapeTuple shape = {1, 3, 6, 6};
auto buff = allocator->Alloc(shape, dt);
auto buff = allocator->Alloc(dev, shape, dt);
EXPECT_EQ(allocator->UsedMemory(), nbytes);
allocator->Free(buff);
EXPECT_EQ(allocator->UsedMemory(), 0);

auto texture = allocator->Alloc(shape, dt, "global.texture");
auto texture = allocator->Alloc(dev, shape, dt, "global.texture");
EXPECT_EQ(allocator->UsedMemory(), nbytes);
allocator->Free(texture);
EXPECT_EQ(allocator->UsedMemory(), 0);
Expand All @@ -187,13 +187,13 @@ TEST_F(TvmVMMemoryManagerTest, PooledAllocOpenCLTexture) {
size_t page_size = PooledAllocator::kDefaultPageSize;
size_t size = ((nbytes + page_size - 1) / page_size) * page_size;
ShapeTuple shape = {1, 3, 6, 6};
auto buff = allocator->Alloc(shape, dt);
auto buff = allocator->Alloc(dev, shape, dt);
EXPECT_EQ(allocator->UsedMemory(), size);
allocator->Free(buff);
EXPECT_EQ(allocator->UsedMemory(), size);

try {
auto texture = allocator->Alloc(shape, dt, "global.texture");
auto texture = allocator->Alloc(dev, shape, dt, "global.texture");
(void)texture;
FAIL();
} catch (std::exception& e) {
Expand Down