[Runtime] Support clear global memory allocators

This PR supports clearing up all the allocated memory.

Prior to this PR, all the allocated memory are managed in the pool
of memory manager. The allocated memory in the pool is on hold and
never freed. Consequently, the pool size always goes up monotonically
within a single run in a process.

While good to save time of memory allocation, in some cases (e.g.,
on mobile phones which may have running memory limit) we need to
clear the pool and free all the memory in order to prevent the pool
from endlessly growing up and some of allocated memory not being
effectively utilized (fragmentation).

Therefore, this PR introduces a PackedFunc that helps clean up
the memory manager, releasing all the allocated memory. Runtime
apps can decide when to invoke this PackedFunc and clean up the
pool. Usually, this will happen at some app "reset" or "reload"
stage.

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: