This is an automated email from the ASF dual-hosted git repository.
ruihangl pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 5a54d969e8 [FFI] Fix SmallMapInit with duplicated keys (#18178)
5a54d969e8 is described below
commit 5a54d969e82230b617904688f3e0f154d287b2ca
Author: Tianqi Chen <[email protected]>
AuthorDate: Thu Jul 31 14:48:10 2025 -0400
[FFI] Fix SmallMapInit with duplicated keys (#18178)
This PR fixes Small map init when there are duplicated keys
---
ffi/include/tvm/ffi/container/map.h | 30 ++++++++++++++++++++----------
ffi/tests/cpp/test_map.cc | 7 +++++++
2 files changed, 27 insertions(+), 10 deletions(-)
diff --git a/ffi/include/tvm/ffi/container/map.h
b/ffi/include/tvm/ffi/container/map.h
index fc94ab15a8..7892e20085 100644
--- a/ffi/include/tvm/ffi/container/map.h
+++ b/ffi/include/tvm/ffi/container/map.h
@@ -1264,17 +1264,27 @@ inline ObjectPtr<Object>
MapObj::CreateFromRange(IterType first, IterType last)
}
uint64_t cap = static_cast<uint64_t>(_cap);
if (cap < SmallMapObj::kMaxSize) {
- return SmallMapObj::CreateFromRange(cap, first, last);
- }
- uint32_t fib_shift;
- uint64_t n_slots;
- DenseMapObj::CalcTableSize(cap, &fib_shift, &n_slots);
- ObjectPtr<Object> obj = DenseMapObj::Empty(fib_shift, n_slots);
- for (; first != last; ++first) {
- KVType kv(*first);
- DenseMapObj::InsertMaybeReHash(std::move(kv), &obj);
+ if (cap < 2) {
+ return SmallMapObj::CreateFromRange(cap, first, last);
+ }
+ // need to insert to avoid duplicate keys
+ ObjectPtr<Object> obj = SmallMapObj::Empty(cap);
+ for (; first != last; ++first) {
+ KVType kv(*first);
+ SmallMapObj::InsertMaybeReHash(std::move(kv), &obj);
+ }
+ return obj;
+ } else {
+ uint32_t fib_shift;
+ uint64_t n_slots;
+ DenseMapObj::CalcTableSize(cap, &fib_shift, &n_slots);
+ ObjectPtr<Object> obj = DenseMapObj::Empty(fib_shift, n_slots);
+ for (; first != last; ++first) {
+ KVType kv(*first);
+ DenseMapObj::InsertMaybeReHash(std::move(kv), &obj);
+ }
+ return obj;
}
- return obj;
}
inline void MapObj::InsertMaybeReHash(KVType&& kv, ObjectPtr<Object>* map) {
diff --git a/ffi/tests/cpp/test_map.cc b/ffi/tests/cpp/test_map.cc
index b7c977fd34..98d8427c23 100644
--- a/ffi/tests/cpp/test_map.cc
+++ b/ffi/tests/cpp/test_map.cc
@@ -356,4 +356,11 @@ TEST(Map, EmptyIter) {
// now m0 is dense map with all empty slots
EXPECT_EQ(m0.begin(), m0.end());
}
+
+TEST(Map, DuplicatedKeysInit) {
+ std::vector<std::pair<String, int>> data = {{"a", 1}, {"a", 2}, {"a", 3}};
+ Map<String, int> map(data.begin(), data.end());
+ EXPECT_EQ(map.size(), 1);
+ EXPECT_EQ(map["a"], 3);
+}
} // namespace
