[RUNTIME] Add clear() function in tvm::Map class (apache#7826)

Co-authored-by: honghua.cao <honghua.cao@streamcomputing.com>
trevor-m · May 6, 2021 · f7c382f · f7c382f
1 parent 3a49066
commit f7c382f
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 0 deletions.
diff --git a/include/tvm/runtime/container.h b/include/tvm/runtime/container.h
@@ -2980,6 +2980,13 @@ class Map : public ObjectRef {
   }
   /*! \return whether array is empty */
   bool empty() const { return size() == 0; }
+  /*! \brief Release reference to all the elements */
+  void clear() {
+    MapNode* n = GetMapNode();
+    if (n != nullptr) {
+      data_ = MapNode::Empty();
+    }
+  }
   /*!
    * \brief set the Map.
    * \param key The index key.

diff --git a/tests/cpp/container_test.cc b/tests/cpp/container_test.cc
@@ -313,6 +313,16 @@ TEST(Map, Mutate) {
   ICHECK(it == dict2.end());
 }
 
+TEST(Map, Clear) {
+  using namespace tvm;
+  Var x("x");
+  auto z = max(x + 1 + 2, 100);
+  Map<PrimExpr, PrimExpr> dict{{x, z}, {z, 2}};
+  ICHECK(dict.size() == 2);
+  dict.clear();
+  ICHECK(dict.size() == 0);
+}
+
 TEST(Map, Iterator) {
   using namespace tvm;
   PrimExpr a = 1, b = 2;