[Runtime] Support clear global memory allocators

include/tvm/runtime/memory/memory_manager.h

@@ -89,6 +89,8 @@ class Allocator {
    *  \param buffer The buffer to free.
    */
   virtual void Free(const Buffer& buffer) = 0;
+  /*! \brief Clear the allocated memory. */
+  virtual void Clear();
   /*! \brief The amount of memory currently allocated.
    *  \return The amount of memory currently allocated.
    */
@@ -119,6 +121,8 @@ class MemoryManager {
    * \return The memory allocator.
    */
   static Allocator* GetAllocator(Device dev, AllocatorType type);
+  /*! \brief Clear the allocators. */
+  static void Clear();
 
  private:
   MemoryManager() {}

src/runtime/memory/memory_manager.cc

@@ -22,6 +22,7 @@
  * \brief Allocate and manage memory for the runtime.
  */
 #include <tvm/runtime/memory/memory_manager.h>
+#include <tvm/runtime/registry.h>
 
 #include <memory>
 #include <utility>
@@ -166,6 +167,16 @@ Allocator* MemoryManager::GetAllocator(Device dev, AllocatorType type) {
   return it->second.at(type).get();
 }
 
+void MemoryManager::Clear() {
+  MemoryManager* m = MemoryManager::Global();
+  std::lock_guard<std::mutex> lock(m->mu_);
+  for (const auto& [device, allocators] : m->allocators_) {
+    for (const auto& [allocator_type, allocator] : allocators) {
+      allocator->Clear();
+    }
+  }
+}
+
 NDArray Allocator::Empty(ShapeTuple shape, DLDataType dtype, DLDevice dev,
                          Optional<String> mem_scope) {
   VerifyDataType(dtype);
@@ -198,6 +209,14 @@ Buffer Allocator::Alloc(Device dev, ShapeTuple shape, DLDataType type_hint,
   return {};
 }
 
+void Allocator::Clear() {
+  // This function by default does nothing.
+  // For naive allocator, no explicit manual clear is needed.
+  // Pooled allocator will override this method.
+}
+
+TVM_REGISTER_GLOBAL("vm.builtin.memory_manager.clear").set_body_typed(MemoryManager::Clear);
+
 }  // namespace memory
 }  // namespace runtime
 }  // namespace tvm

src/runtime/memory/pooled_allocator.h

@@ -90,6 +90,8 @@ class PooledAllocator final : public Allocator {
     VLOG(1) << "reclaim buffer " << buffer.size;
   }
 
+  void Clear() override { ReleaseAll(); }
+
   size_t UsedMemory() const override { return used_memory_.load(std::memory_order_relaxed); }
 
  private: