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

Prior to this PR, each allocator is closely tied with a device.
To enable using a same allocator across different devices of the
same kind when needed, we lift the device to the allocator `Alloc`
interface.

Author: MasterJH5574

include/tvm/runtime/memory/memory_manager.h

@@ -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.
@@ -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_;
 };

src/runtime/memory/memory_manager.cc

@@ -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:
@@ -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;
@@ -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;

src/runtime/memory/naive_allocator.h

@@ -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";
   }
@@ -83,7 +84,6 @@ class NaiveAllocator final : public Allocator {
 
  private:
   std::atomic<size_t> used_memory_;
-  Device device_;
 };
 
 }  // namespace memory

src/runtime/memory/pooled_allocator.h

@@ -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 {
+  virtual 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);
@@ -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 {
+  virtual 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 {};
@@ -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

src/runtime/relax_vm/builtin.cc

@@ -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;
 }

src/runtime/vm/vm.cc

@@ -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) {
@@ -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);