From 1bcd4d823c464af5e8430b52c443f4ff7ef1aa9c Mon Sep 17 00:00:00 2001 From: Dmitriy Smirnov Date: Wed, 5 Jan 2022 12:11:15 +0000 Subject: [PATCH 1/3] Attempt to prevent concurrent update in Map Calling Map::Set invalidates exising iterators to protect from using already deleted data due to re-hashing Change-Id: Ib6b580758e74c8b77ed560932d87b643bd6c9402 --- include/tvm/runtime/container/map.h | 22 +++++++++++++++++++--- tests/cpp/container_test.cc | 13 +++++++++++++ 2 files changed, 32 insertions(+), 3 deletions(-) diff --git a/include/tvm/runtime/container/map.h b/include/tvm/runtime/container/map.h index 977dbfbaaaa1..fc2d46ed9ef7 100644 --- a/include/tvm/runtime/container/map.h +++ b/include/tvm/runtime/container/map.h @@ -234,9 +234,10 @@ class MapNode : public Object { using pointer = KVType*; using reference = KVType&; /*! \brief Default constructor */ - iterator() : index(0), self(nullptr) {} + iterator() : state_marker(0), index(0), self(nullptr) {} /*! \brief Compare iterators */ bool operator==(const iterator& other) const { + fail_if_changed(); return index == other.index && self == other.self; } /*! \brief Compare iterators */ @@ -244,27 +245,37 @@ class MapNode : public Object { /*! \brief De-reference iterators */ pointer operator->() const; /*! \brief De-reference iterators */ - reference operator*() const { return *((*this).operator->()); } + reference operator*() const { + fail_if_changed(); + return *((*this).operator->()); + } /*! \brief Prefix self increment, e.g. ++iter */ iterator& operator++(); /*! \brief Prefix self decrement, e.g. --iter */ iterator& operator--(); /*! \brief Suffix self increment */ iterator operator++(int) { + fail_if_changed(); iterator copy = *this; ++(*this); return copy; } /*! \brief Suffix self decrement */ iterator operator--(int) { + fail_if_changed(); iterator copy = *this; --(*this); return copy; } protected: + uint64_t state_marker; + inline void fail_if_changed() const { + ICHECK(state_marker == self->state_marker) << "Concurrent modification of the Map"; + } /*! \brief Construct by value */ - iterator(uint64_t index, const MapNode* self) : index(index), self(self) {} + iterator(uint64_t index, const MapNode* self) + : state_marker(self->state_marker), index(index), self(self) {} /*! \brief The position on the array */ uint64_t index; /*! \brief The container it points to */ @@ -280,6 +291,7 @@ class MapNode : public Object { static inline ObjectPtr Empty(); protected: + uint64_t state_marker; /*! * \brief Create the map using contents from the given iterators. * \param first Begin of iterator @@ -1118,10 +1130,12 @@ class DenseMapNode : public MapNode { } inline MapNode::iterator::pointer MapNode::iterator::operator->() const { + fail_if_changed(); TVM_DISPATCH_MAP_CONST(self, p, { return p->DeRefItr(index); }); } inline MapNode::iterator& MapNode::iterator::operator++() { + fail_if_changed(); TVM_DISPATCH_MAP_CONST(self, p, { index = p->IncItr(index); return *this; @@ -1129,6 +1143,7 @@ inline MapNode::iterator& MapNode::iterator::operator++() { } inline MapNode::iterator& MapNode::iterator::operator--() { + fail_if_changed(); TVM_DISPATCH_MAP_CONST(self, p, { index = p->DecItr(index); return *this; @@ -1200,6 +1215,7 @@ inline ObjectPtr MapNode::CreateFromRange(IterType first, IterType last) inline void MapNode::InsertMaybeReHash(const KVType& kv, ObjectPtr* map) { constexpr uint64_t kSmallMapMaxSize = SmallMapNode::kMaxSize; MapNode* base = static_cast(map->get()); + base->state_marker++; if (base->slots_ < kSmallMapMaxSize) { SmallMapNode::InsertMaybeReHash(kv, map); } else if (base->slots_ == kSmallMapMaxSize) { diff --git a/tests/cpp/container_test.cc b/tests/cpp/container_test.cc index 019fde069878..fa96066bc5b1 100644 --- a/tests/cpp/container_test.cc +++ b/tests/cpp/container_test.cc @@ -380,6 +380,19 @@ TEST(Map, Erase) { } } +TEST(Map, Race) { + using namespace tvm::runtime; + Map m; + + m.Set(1, 1); + Map::iterator it = m.begin(); + EXPECT_NO_THROW({ auto& kv = *it; }); + + m.Set(2, 2); + // changed. iterator should be re-obtained + EXPECT_ANY_THROW({ auto& kv = *it; }); +} + TEST(String, MoveFromStd) { using namespace std; string source = "this is a string"; From b5599b73d5f8ad3664a81ce0b0f23a5ee3f7766f Mon Sep 17 00:00:00 2001 From: Dmitriy Smirnov Date: Fri, 14 Jan 2022 13:45:47 +0000 Subject: [PATCH 2/3] Migrated to using TVM_LOG_DEBUG Now uses TVM_LOG_DEBUG Map state_marker made atomic Change-Id: I090c4b33e6edaa977cccba11f8d1c6ff3fbca430 --- include/tvm/runtime/container/map.h | 42 ++++++++++++++++++++--------- tests/cpp/container_test.cc | 2 ++ 2 files changed, 32 insertions(+), 12 deletions(-) diff --git a/include/tvm/runtime/container/map.h b/include/tvm/runtime/container/map.h index fc2d46ed9ef7..726aad8f0af2 100644 --- a/include/tvm/runtime/container/map.h +++ b/include/tvm/runtime/container/map.h @@ -29,6 +29,7 @@ #endif #include +#include #include #include @@ -38,6 +39,13 @@ namespace tvm { namespace runtime { +#if TVM_LOG_DEBUG +#define TVM_MAP_FAIL_IF_CHANGED() \ + ICHECK(state_marker == self->state_marker) << "Concurrent modification of the Map"; +#else +#define TVM_MAP_FAIL_IF_CHANGED() +#endif // TVM_LOG_DEBUG + #if (USE_FALLBACK_STL_MAP != 0) /*! \brief Shared content of all specializations of hash map */ @@ -233,11 +241,15 @@ class MapNode : public Object { using value_type = KVType; using pointer = KVType*; using reference = KVType&; - /*! \brief Default constructor */ +/*! \brief Default constructor */ +#if TVM_LOG_DEBUG iterator() : state_marker(0), index(0), self(nullptr) {} +#else + iterator() : index(0), self(nullptr) {} +#endif // TVM_LOG_DEBUG /*! \brief Compare iterators */ bool operator==(const iterator& other) const { - fail_if_changed(); + TVM_MAP_FAIL_IF_CHANGED() return index == other.index && self == other.self; } /*! \brief Compare iterators */ @@ -246,7 +258,7 @@ class MapNode : public Object { pointer operator->() const; /*! \brief De-reference iterators */ reference operator*() const { - fail_if_changed(); + TVM_MAP_FAIL_IF_CHANGED() return *((*this).operator->()); } /*! \brief Prefix self increment, e.g. ++iter */ @@ -255,27 +267,29 @@ class MapNode : public Object { iterator& operator--(); /*! \brief Suffix self increment */ iterator operator++(int) { - fail_if_changed(); + TVM_MAP_FAIL_IF_CHANGED() iterator copy = *this; ++(*this); return copy; } /*! \brief Suffix self decrement */ iterator operator--(int) { - fail_if_changed(); + TVM_MAP_FAIL_IF_CHANGED() iterator copy = *this; --(*this); return copy; } protected: +#if TVM_LOG_DEBUG uint64_t state_marker; - inline void fail_if_changed() const { - ICHECK(state_marker == self->state_marker) << "Concurrent modification of the Map"; - } /*! \brief Construct by value */ iterator(uint64_t index, const MapNode* self) : state_marker(self->state_marker), index(index), self(self) {} + +#else + iterator(uint64_t index, const MapNode* self) : index(index), self(self) {} +#endif // TVM_LOG_DEBUG /*! \brief The position on the array */ uint64_t index; /*! \brief The container it points to */ @@ -291,7 +305,9 @@ class MapNode : public Object { static inline ObjectPtr Empty(); protected: - uint64_t state_marker; +#if TVM_LOG_DEBUG + std::atomic state_marker; +#endif // TVM_LOG_DEBUG /*! * \brief Create the map using contents from the given iterators. * \param first Begin of iterator @@ -1130,12 +1146,12 @@ class DenseMapNode : public MapNode { } inline MapNode::iterator::pointer MapNode::iterator::operator->() const { - fail_if_changed(); + TVM_MAP_FAIL_IF_CHANGED() TVM_DISPATCH_MAP_CONST(self, p, { return p->DeRefItr(index); }); } inline MapNode::iterator& MapNode::iterator::operator++() { - fail_if_changed(); + TVM_MAP_FAIL_IF_CHANGED() TVM_DISPATCH_MAP_CONST(self, p, { index = p->IncItr(index); return *this; @@ -1143,7 +1159,7 @@ inline MapNode::iterator& MapNode::iterator::operator++() { } inline MapNode::iterator& MapNode::iterator::operator--() { - fail_if_changed(); + TVM_MAP_FAIL_IF_CHANGED() TVM_DISPATCH_MAP_CONST(self, p, { index = p->DecItr(index); return *this; @@ -1215,7 +1231,9 @@ inline ObjectPtr MapNode::CreateFromRange(IterType first, IterType last) inline void MapNode::InsertMaybeReHash(const KVType& kv, ObjectPtr* map) { constexpr uint64_t kSmallMapMaxSize = SmallMapNode::kMaxSize; MapNode* base = static_cast(map->get()); +#if TVM_LOG_DEBUG base->state_marker++; +#endif // TVM_LOG_DEBUG if (base->slots_ < kSmallMapMaxSize) { SmallMapNode::InsertMaybeReHash(kv, map); } else if (base->slots_ == kSmallMapMaxSize) { diff --git a/tests/cpp/container_test.cc b/tests/cpp/container_test.cc index fa96066bc5b1..32ec346c8796 100644 --- a/tests/cpp/container_test.cc +++ b/tests/cpp/container_test.cc @@ -380,6 +380,7 @@ TEST(Map, Erase) { } } +#if TVM_LOG_DEBUG TEST(Map, Race) { using namespace tvm::runtime; Map m; @@ -392,6 +393,7 @@ TEST(Map, Race) { // changed. iterator should be re-obtained EXPECT_ANY_THROW({ auto& kv = *it; }); } +#endif // TVM_LOG_DEBUG TEST(String, MoveFromStd) { using namespace std; From e4fb820bdc1c9f423dbb22a3adeb55ae3464cf4c Mon Sep 17 00:00:00 2001 From: Dmitriy Smirnov Date: Mon, 11 Apr 2022 12:18:21 +0100 Subject: [PATCH 3/3] removed usage of atomics Change-Id: I7bd930cb52d58ca10fd49a5fe8f5d48b3e955d0a --- include/tvm/runtime/container/map.h | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/include/tvm/runtime/container/map.h b/include/tvm/runtime/container/map.h index 726aad8f0af2..4c76a3b0ad4f 100644 --- a/include/tvm/runtime/container/map.h +++ b/include/tvm/runtime/container/map.h @@ -29,7 +29,6 @@ #endif #include -#include #include #include @@ -306,7 +305,7 @@ class MapNode : public Object { protected: #if TVM_LOG_DEBUG - std::atomic state_marker; + uint64_t state_marker; #endif // TVM_LOG_DEBUG /*! * \brief Create the map using contents from the given iterators.