From f7c382f31763dd8db2163d14f5c65e4feb25827f Mon Sep 17 00:00:00 2001 From: Honghua Cao <49267856+Beya2019@users.noreply.github.com> Date: Wed, 14 Apr 2021 07:22:24 +0800 Subject: [PATCH] [RUNTIME] Add clear() function in tvm::Map class (#7826) Co-authored-by: honghua.cao --- include/tvm/runtime/container.h | 7 +++++++ tests/cpp/container_test.cc | 10 ++++++++++ 2 files changed, 17 insertions(+) diff --git a/include/tvm/runtime/container.h b/include/tvm/runtime/container.h index 362582f4dab9a..90a29380f9e0e 100644 --- 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 index 41632ff8d5614..63819308a666e 100644 --- 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 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;