* allow multiple allocators to coexist for the same device.

Using available allocator instead of requested is leading to an unpexpected crash

Author: srkreddy1238

include/tvm/runtime/memory/memory_manager.h

@@ -39,18 +39,20 @@ namespace tvm {
 namespace runtime {
 namespace memory {
 
+enum AllocatorType {
+  kNaive = 1,
+  kPooled,
+};
+
 struct Buffer {
   /*! \brief The pointer to the allocated block of memory. */
   void* data{nullptr};
   /*! \brief The size of the block. */
   size_t size{0};
   /*! \brief The context of the allocated buffers. */
   Device device;
-};
-
-enum AllocatorType {
-  kNaive = 1,
-  kPooled,
+  /*! \brief The allocator that created this buffer. */
+  AllocatorType alloc_type;
 };
 
 class Allocator {
@@ -113,16 +115,18 @@ class MemoryManager {
   /*!
    * \brief Get an allocator given the context.
    * \param dev The TVM device
+   * \param type The allocator type
    * \return The memory allocator.
    */
-  static Allocator* GetAllocator(Device dev);
+  static Allocator* GetAllocator(Device dev, AllocatorType type);
 
  private:
   MemoryManager() {}
 
  protected:
   std::mutex mu_;
-  std::unordered_map<Device, std::unique_ptr<Allocator>> allocators_;
+  std::unordered_map<Device, std::unordered_map<AllocatorType, std::unique_ptr<Allocator>>>
+      allocators_;
 };
 
 /*! \brief An object representing a storage allocation. */
@@ -138,7 +142,7 @@ class StorageObj : public Object {
   static void Deleter(Object* ptr);
 
   ~StorageObj() {
-    auto alloc = MemoryManager::Global()->GetAllocator(buffer.device);
+    auto alloc = MemoryManager::Global()->GetAllocator(buffer.device, buffer.alloc_type);
     alloc->Free(buffer);
   }
 

src/runtime/memory/memory_manager.cc

@@ -37,7 +37,7 @@ static void BufferDeleter(Object* obj) {
   auto* ptr = static_cast<NDArray::Container*>(obj);
   ICHECK(ptr->manager_ctx != nullptr);
   Buffer* buffer = reinterpret_cast<Buffer*>(ptr->manager_ctx);
-  MemoryManager::GetAllocator(buffer->device)->Free(*(buffer));
+  MemoryManager::GetAllocator(buffer->device, buffer->alloc_type)->Free(*(buffer));
   delete buffer;
   delete ptr;
 }
@@ -122,6 +122,9 @@ Allocator* MemoryManager::GetOrCreateAllocator(Device dev, AllocatorType type) {
   MemoryManager* m = MemoryManager::Global();
   std::lock_guard<std::mutex> lock(m->mu_);
   if (m->allocators_.find(dev) == m->allocators_.end()) {
+    m->allocators_.emplace(dev, std::unordered_map<AllocatorType, std::unique_ptr<Allocator>>());
+  }
+  if (m->allocators_.at(dev).find(type) == m->allocators_.at(dev).end()) {
     std::unique_ptr<Allocator> alloc;
     switch (type) {
       case kNaive: {
@@ -138,26 +141,29 @@ Allocator* MemoryManager::GetOrCreateAllocator(Device dev, AllocatorType type) {
         LOG(FATAL) << "Unknown allocator type: " << type;
     }
     auto ret = alloc.get();
-    m->allocators_.emplace(dev, std::move(alloc));
+    m->allocators_.at(dev).emplace(type, std::move(alloc));
     return ret;
   }
-  auto alloc = m->allocators_.at(dev).get();
-  if (alloc->type() != type) {
+  auto alloc = m->allocators_.at(dev).at(type).get();
+  /*if (alloc->type() != type) {
     LOG(WARNING) << "The type of existing allocator for " << dev
                  << " is different from the request type (" << alloc->type() << " vs " << type
                  << ")";
-  }
+  }*/
   return alloc;
 }
 
-Allocator* MemoryManager::GetAllocator(Device dev) {
+Allocator* MemoryManager::GetAllocator(Device dev, AllocatorType type) {
   MemoryManager* m = MemoryManager::Global();
   std::lock_guard<std::mutex> lock(m->mu_);
   auto it = m->allocators_.find(dev);
   if (it == m->allocators_.end()) {
     LOG(FATAL) << "Allocator for " << dev << " has not been created yet.";
   }
-  return it->second.get();
+  if (it->second.find(type) == it->second.end()) {
+    LOG(FATAL) << "Allocator for " << dev << " of type " << type << " has not been created yet.";
+  }
+  return it->second.at(type).get();
 }
 
 NDArray Allocator::Empty(ShapeTuple shape, DLDataType dtype, DLDevice dev,

src/runtime/memory/naive_allocator.h

@@ -41,6 +41,7 @@ class NaiveAllocator final : public Allocator {
     Buffer buf;
     buf.device = device_;
     buf.size = nbytes;
+    buf.alloc_type = kNaive;
     buf.data = DeviceAPI::Get(device_)->AllocDataSpace(device_, nbytes, alignment, type_hint);
     used_memory_.fetch_add(nbytes, std::memory_order_relaxed);
     DLOG(INFO) << "allocate " << nbytes << " B, used memory " << used_memory_ << " B";
@@ -59,6 +60,7 @@ class NaiveAllocator final : public Allocator {
       auto tmp_buf = Allocator::Alloc(device_, shape, type_hint, mem_scope);
       buf.size = tmp_buf.size;
       buf.data = tmp_buf.data;
+      buf.alloc_type = kNaive;
       return buf;
     }
 
@@ -67,6 +69,7 @@ class NaiveAllocator final : public Allocator {
                                                        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;
   }
 

src/runtime/memory/pooled_allocator.h

@@ -58,6 +58,7 @@ class PooledAllocator final : public Allocator {
     Buffer buf;
     buf.device = device_;
     buf.size = size;
+    buf.alloc_type = kPooled;
     try {
       buf.data = DeviceAPI::Get(device_)->AllocDataSpace(device_, size, alignment, type_hint);
     } catch (InternalError& err) {