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

Co-authored-by: honghua.cao <honghua.cao@streamcomputing.com>
neo-ai · May 11, 2021 · b79e144 · b79e144
1 parent 5db21ea
commit b79e144
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;