diff --git a/ffi/include/tvm/ffi/container/array.h b/ffi/include/tvm/ffi/container/array.h index 30402d9ae68c..97fe5916822d 100644 --- a/ffi/include/tvm/ffi/container/array.h +++ b/ffi/include/tvm/ffi/container/array.h @@ -42,8 +42,18 @@ namespace tvm { namespace ffi { /*! \brief array node content in array */ -class ArrayObj : public Object, public details::InplaceArrayBase { +class ArrayObj : public Object, public details::InplaceArrayBase { public: + ~ArrayObj() { + Any* begin = MutableBegin(); + for (int64_t i = 0; i < size_; ++i) { + (begin + i)->Any::~Any(); + } + if (data_deleter_ != nullptr) { + data_deleter_(data_); + } + } + /*! \return The size of the array */ size_t size() const { return this->size_; } @@ -52,10 +62,22 @@ class ArrayObj : public Object, public details::InplaceArrayBase * \param i The index * \return the i-th element. */ - const Any at(int64_t i) const { return this->operator[](i); } + const Any& at(int64_t i) const { return this->operator[](i); } + + /*! + * \brief Read i-th element from array. + * \param i The index + * \return the i-th element. + */ + const Any& operator[](int64_t i) const { + if (i >= size_) { + TVM_FFI_THROW(IndexError) << "Index " << i << " out of bounds " << size_; + } + return static_cast(data_)[i]; + } /*! \return begin constant iterator */ - const Any* begin() const { return static_cast(InplaceArrayBase::AddressOf(0)); } + const Any* begin() const { return static_cast(data_); } /*! \return end constant iterator */ const Any* end() const { return begin() + size_; } @@ -68,7 +90,12 @@ class ArrayObj : public Object, public details::InplaceArrayBase * \param i The index * \param item The value to be set */ - void SetItem(int64_t i, Any item) { this->operator[](i) = std::move(item); } + void SetItem(int64_t i, Any item) { + if (i >= size_) { + TVM_FFI_THROW(IndexError) << "Index " << i << " out of bounds " << size_; + } + static_cast(data_)[i] = std::move(item); + } /*! * \brief Constructs a container and copy from another @@ -138,21 +165,31 @@ class ArrayObj : public Object, public details::InplaceArrayBase size_t GetSize() const { return this->size_; } /*! \return begin mutable iterator */ - Any* MutableBegin() const { return static_cast(InplaceArrayBase::AddressOf(0)); } + Any* MutableBegin() const { return static_cast(this->data_); } /*! \return end mutable iterator */ Any* MutableEnd() const { return MutableBegin() + size_; } + /*! + * \brief Emplace a new element at the back of the array + * \param args The arguments to construct the new element + */ + template + void EmplaceInit(size_t idx, Args&&... args) { + Any* itr = MutableBegin() + idx; + new (itr) Any(std::forward(args)...); + } + /*! * \brief Create an ArrayObj with the given capacity. * \param n Required capacity * \return Ref-counted ArrayObj requested */ static ObjectPtr Empty(int64_t n = kInitSize) { - TVM_FFI_ICHECK_GE(n, 0); ObjectPtr p = make_inplace_array_object(n); p->capacity_ = n; p->size_ = 0; + p->data_ = p->AddressOf(0); return p; } @@ -235,11 +272,17 @@ class ArrayObj : public Object, public details::InplaceArrayBase return this; } + /*! \brief Data pointer to the first element of the array */ + void* data_; /*! \brief Number of elements used */ int64_t size_; - /*! \brief Number of elements allocated */ int64_t capacity_; + /*! + * \brief Optional data deleter when data is allocated separately + * and its deletion is not managed by ArrayObj::deleter_. + */ + void (*data_deleter_)(void*) = nullptr; /*! \brief Initial size of ArrayObj */ static constexpr int64_t kInitSize = 4; @@ -472,6 +515,12 @@ class Array : public ObjectRef { p->EmplaceInit(p->size_++, item); } + template + void emplace_back(Args&&... args) { + ArrayObj* p = CopyOnWrite(1); + p->EmplaceInit(p->size_++, std::forward(args)...); + } + /*! * \brief Insert an element into the given position * \param position An iterator pointing to the insertion point diff --git a/ffi/include/tvm/ffi/container/map.h b/ffi/include/tvm/ffi/container/map.h index a738c1e229f6..2cc6bf069a40 100644 --- a/ffi/include/tvm/ffi/container/map.h +++ b/ffi/include/tvm/ffi/container/map.h @@ -56,6 +56,11 @@ class MapObj : public Object { using mapped_type = Any; /*! \brief Type of value stored in the hash map */ using KVType = std::pair; + /*! \brief Type of raw storage of the key-value pair in the hash map */ + struct KVRawStorageType { + TVMFFIAny first; + TVMFFIAny second; + }; /*! \brief Iterator class */ class iterator; @@ -206,17 +211,29 @@ class MapObj : public Object { * \return The object created */ static inline ObjectPtr CopyFrom(MapObj* from); - /*! \brief number of slots minus 1 */ - uint64_t slots_; + /*! + * \brief data pointer to the data region of the map. + * \note For immutable inplace small map we do not need data_, + * but we keep it here for future compact with mutable container. + */ + void* data_; /*! \brief number of entries in the container */ uint64_t size_; + /*! \brief number of slots */ + uint64_t slots_; + /*! + * \brief Optional data deleter when data is allocated separately + * and its deletion is not managed by MapObj::deleter_. + */ + void (*data_deleter_)(void*) = nullptr; // Reference class template friend class Map; }; /*! \brief A specialization of small-sized hash map */ -class SmallMapObj : public MapObj, public details::InplaceArrayBase { +class SmallMapObj : public MapObj, + public details::InplaceArrayBase { private: static constexpr uint64_t kInitSize = 2; static constexpr uint64_t kMaxSize = 4; @@ -225,8 +242,19 @@ class SmallMapObj : public MapObj, public details::InplaceArrayBase(data_); + for (uint64_t index = 0; index < size_; ++index) { + // call destructor to destroy the item in `begin + index` + // Explicit call Any::~Any() to destroy the Any object + // Favor this over ~KVType as MSVC may not support ~KVType (need the original name) + (begin + index)->first.Any::~Any(); + (begin + index)->second.Any::~Any(); + } + if (data_deleter_ != nullptr) { + data_deleter_(data_); + } + } /*! * \brief Count the number of times a key exists in the SmallMapObj * \param key The indexing key @@ -267,7 +295,7 @@ class SmallMapObj : public MapObj, public details::InplaceArrayBase(AddressOf(0)); + KVType* ptr = static_cast(data_); for (uint64_t i = 0; i < size_; ++i, ++ptr) { if (AnyEqual()(ptr->first, key)) { return iterator(i, this); @@ -290,7 +318,7 @@ class SmallMapObj : public MapObj, public details::InplaceArrayBase= size_) { return; } - KVType* begin = static_cast(AddressOf(0)); + KVType* begin = static_cast(data_); // call destructor to destroy the item in `begin + index` // Explicit call Any::~Any() to destroy the Any object // Favor this over ~KVType as MSVC may not support ~KVType (need the original name) @@ -314,6 +342,7 @@ class SmallMapObj : public MapObj, public details::InplaceArrayBase Empty(uint64_t n = kInitSize) { using ::tvm::ffi::make_inplace_array_object; ObjectPtr p = make_inplace_array_object(n); + p->data_ = p->AddressOf(0); p->size_ = 0; p->slots_ = n; return p; @@ -329,7 +358,7 @@ class SmallMapObj : public MapObj, public details::InplaceArrayBase static ObjectPtr CreateFromRange(uint64_t n, IterType first, IterType last) { ObjectPtr p = Empty(n); - KVType* ptr = static_cast(p->AddressOf(0)); + KVType* ptr = static_cast(p->data_); for (; first != last; ++first, ++p->size_) { new (ptr++) KVType(*first); } @@ -341,7 +370,7 @@ class SmallMapObj : public MapObj, public details::InplaceArrayBase CopyFrom(SmallMapObj* from) { - KVType* first = static_cast(from->AddressOf(0)); + KVType* first = static_cast(from->data_); KVType* last = first + from->size_; return CreateFromRange(from->size_, first, last); } @@ -358,7 +387,7 @@ class SmallMapObj : public MapObj, public details::InplaceArrayBasesize_ < map_node->slots_) { - KVType* ptr = static_cast(map_node->AddressOf(map_node->size_)); + KVType* ptr = static_cast(map_node->data_) + map_node->size_; new (ptr) KVType(std::move(kv)); ++map_node->size_; return; @@ -387,7 +416,7 @@ class SmallMapObj : public MapObj, public details::InplaceArrayBase(AddressOf(index)); } + KVType* DeRefItr(uint64_t index) const { return static_cast(data_) + index; } /*! \brief A size function used by InplaceArrayBase */ uint64_t GetSize() const { return size_; } @@ -487,6 +516,12 @@ class DenseMapObj : public MapObj { static_assert(sizeof(Block) == kBlockCap * (sizeof(ItemType) + 1), "sizeof(Block) incorrect"); static_assert(std::is_standard_layout::value, "Block is not standard layout"); + /*! + * \brief Deleter for the Block + * \param data The pointer to the Block + */ + static void BlockDeleter(void* data) { delete[] static_cast(data); } + public: using MapObj::iterator; @@ -533,6 +568,7 @@ class DenseMapObj : public MapObj { iterator end() const { return iterator(kInvalidIndex, this); } private: + Block* GetBlock(size_t index) const { return static_cast(data_) + index; } /*! * \brief Unlink the entry from iterator list * \param node The node to be unlinked @@ -783,8 +819,8 @@ class DenseMapObj : public MapObj { void Reset() { uint64_t n_blocks = CalcNumBlocks(this->slots_); for (uint64_t bi = 0; bi < n_blocks; ++bi) { - uint8_t* meta_ptr = data_[bi].bytes; - ItemType* data_ptr = reinterpret_cast(data_[bi].bytes + kBlockCap); + uint8_t* meta_ptr = GetBlock(bi)->bytes; + ItemType* data_ptr = reinterpret_cast(GetBlock(bi)->bytes + kBlockCap); for (int j = 0; j < kBlockCap; ++j, ++meta_ptr, ++data_ptr) { uint8_t& meta = *meta_ptr; if (meta != uint8_t(kProtectedSlot) && meta != uint8_t(kEmptySlot)) { @@ -798,8 +834,12 @@ class DenseMapObj : public MapObj { /*! \brief Release the memory acquired by the container without deleting its entries stored inside */ void ReleaseMemory() { - delete[] data_; + if (data_ != nullptr) { + TVM_FFI_ICHECK(data_deleter_ != nullptr); + data_deleter_(data_); + } data_ = nullptr; + data_deleter_ = nullptr; slots_ = 0; size_ = 0; fib_shift_ = 63; @@ -813,9 +853,14 @@ class DenseMapObj : public MapObj { static ObjectPtr Empty(uint32_t fib_shift, uint64_t n_slots) { TVM_FFI_ICHECK_GT(n_slots, uint64_t(SmallMapObj::kMaxSize)); ObjectPtr p = make_object(); - uint64_t n_blocks = CalcNumBlocks(n_slots - 1); - Block* block = p->data_ = new Block[n_blocks]; - p->slots_ = n_slots - 1; + uint64_t n_blocks = CalcNumBlocks(n_slots); + Block* block = new Block[n_blocks]; + p->data_ = block; + // assign block deleter so even if we take re-alloc data + // in another shared-lib that may have different malloc/free behavior + // it will still be safe. + p->data_deleter_ = BlockDeleter; + p->slots_ = n_slots; p->size_ = 0; p->fib_shift_ = fib_shift; p->iter_list_head_ = kInvalidIndex; @@ -834,16 +879,20 @@ class DenseMapObj : public MapObj { ObjectPtr p = make_object(); uint64_t n_blocks = CalcNumBlocks(from->slots_); p->data_ = new Block[n_blocks]; + // assign block deleter so even if we take re-alloc data + // in another shared-lib that may have different malloc/free behavior + // it will still be safe. + p->data_deleter_ = BlockDeleter; p->slots_ = from->slots_; p->size_ = from->size_; p->fib_shift_ = from->fib_shift_; p->iter_list_head_ = from->iter_list_head_; p->iter_list_tail_ = from->iter_list_tail_; for (uint64_t bi = 0; bi < n_blocks; ++bi) { - uint8_t* meta_ptr_from = from->data_[bi].bytes; - ItemType* data_ptr_from = reinterpret_cast(from->data_[bi].bytes + kBlockCap); - uint8_t* meta_ptr_to = p->data_[bi].bytes; - ItemType* data_ptr_to = reinterpret_cast(p->data_[bi].bytes + kBlockCap); + uint8_t* meta_ptr_from = from->GetBlock(bi)->bytes; + ItemType* data_ptr_from = reinterpret_cast(from->GetBlock(bi)->bytes + kBlockCap); + uint8_t* meta_ptr_to = p->GetBlock(bi)->bytes; + ItemType* data_ptr_to = reinterpret_cast(p->GetBlock(bi)->bytes + kBlockCap); for (int j = 0; j < kBlockCap; ++j, ++meta_ptr_from, ++data_ptr_from, ++meta_ptr_to, ++data_ptr_to) { uint8_t& meta = *meta_ptr_to = *meta_ptr_from; @@ -872,7 +921,7 @@ class DenseMapObj : public MapObj { } TVM_FFI_ICHECK_GT(map_node->slots_, uint64_t(SmallMapObj::kMaxSize)); // Otherwise, start rehash - ObjectPtr p = Empty(map_node->fib_shift_ - 1, map_node->slots_ * 2 + 2); + ObjectPtr p = Empty(map_node->fib_shift_ - 1, map_node->slots_ * 2); // need to insert in the same order as the original map for (uint64_t index = map_node->iter_list_head_; index != kInvalidIndex;) { @@ -898,7 +947,7 @@ class DenseMapObj : public MapObj { * \brief Check whether the hash table is full * \return A boolean indicating whether hash table is full */ - bool IsFull() const { return size_ + 1 > (slots_ + 1) * kMaxLoadFactor; } + bool IsFull() const { return size_ + 1 > slots_ * kMaxLoadFactor; } /*! * \brief Increment the pointer * \param index The pointer to be incremented @@ -942,10 +991,7 @@ class DenseMapObj : public MapObj { return node.IsHead() ? node : ListNode(); } /*! \brief Construct the number of blocks in the hash table */ - static uint64_t CalcNumBlocks(uint64_t n_slots_m1) { - uint64_t n_slots = n_slots_m1 > 0 ? n_slots_m1 + 1 : 0; - return (n_slots + kBlockCap - 1) / kBlockCap; - } + static uint64_t CalcNumBlocks(uint64_t n_slots) { return (n_slots + kBlockCap - 1) / kBlockCap; } /*! * \brief Calculate the power-of-2 table size given the lower-bound of required capacity. * \param cap The lower-bound of the required capacity @@ -985,7 +1031,7 @@ class DenseMapObj : public MapObj { ListNode() : index(0), block(nullptr) {} /*! \brief Construct from position */ ListNode(uint64_t index, const DenseMapObj* self) - : index(index), block(self->data_ + (index / kBlockCap)) {} + : index(index), block(self->GetBlock(index / kBlockCap)) {} /*! \brief Metadata on the entry */ uint8_t& Meta() const { return *(block->bytes + index % kBlockCap); } /*! \brief Data on the entry */ @@ -1030,6 +1076,7 @@ class DenseMapObj : public MapObj { Meta() = 0b10000000; new (&Item()) ItemType(std::move(v)); } + /*! \brief If the entry has next entry on the linked list */ bool HasNext() const { return NextProbeLocation(Meta() & 0b01111111) != 0; } /*! \brief Move the entry to the next entry on the linked list */ @@ -1037,12 +1084,13 @@ class DenseMapObj : public MapObj { uint64_t offset = NextProbeLocation(meta & 0b01111111); if (offset == 0) { index = 0; - block = nullptr; return false; } - index = (index + offset) & (self->slots_); - block = self->data_ + (index / kBlockCap); + // the probing will go to next position and round back to stay within the + // correct range of the slots + index = (index + offset) % self->slots_; + block = self->GetBlock(index / kBlockCap); return true; } /*! \brief Move the entry to the next entry on the linked list */ @@ -1060,7 +1108,9 @@ class DenseMapObj : public MapObj { /*! \brief Get the next empty jump */ bool GetNextEmpty(const DenseMapObj* self, uint8_t* jump, ListNode* result) const { for (uint8_t idx = 1; idx < kNumJumpDists; ++idx) { - ListNode candidate((index + NextProbeLocation(idx)) & (self->slots_), self); + // the probing will go to next position and round back to stay within the + // correct range of the slots + ListNode candidate((index + NextProbeLocation(idx)) % self->slots_, self); if (candidate.IsEmpty()) { *jump = idx; *result = candidate; @@ -1078,8 +1128,6 @@ class DenseMapObj : public MapObj { protected: /*! \brief fib shift in Fibonacci Hashing */ uint32_t fib_shift_; - /*! \brief array of data blocks */ - Block* data_; /*! \brief the head of iterator list */ uint64_t iter_list_head_ = kInvalidIndex; /*! \brief the tail of iterator list */ @@ -1118,7 +1166,7 @@ class DenseMapObj : public MapObj { friend class MapObj; }; -#define TVM_DISPATCH_MAP(base, var, body) \ +#define TVM_FFI_DISPATCH_MAP(base, var, body) \ { \ using TSmall = SmallMapObj*; \ using TDense = DenseMapObj*; \ @@ -1132,28 +1180,28 @@ class DenseMapObj : public MapObj { } \ } -#define TVM_DISPATCH_MAP_CONST(base, var, body) \ - { \ - using TSmall = const SmallMapObj*; \ - using TDense = const DenseMapObj*; \ - uint64_t slots = base->slots_; \ - if (slots <= SmallMapObj::kMaxSize) { \ - TSmall var = static_cast(base); \ - body; \ - } else { \ - TDense var = static_cast(base); \ - body; \ - } \ +#define TVM_FFI_DISPATCH_MAP_CONST(base, var, body) \ + { \ + using TSmall = const SmallMapObj*; \ + using TDense = const DenseMapObj*; \ + uint64_t slots = base->slots_; \ + if (slots <= SmallMapObj::kMaxSize) { \ + TSmall var = static_cast(base); \ + body; \ + } else { \ + TDense var = static_cast(base); \ + body; \ + } \ } inline MapObj::iterator::pointer MapObj::iterator::operator->() const { TVM_FFI_MAP_FAIL_IF_CHANGED() - TVM_DISPATCH_MAP_CONST(self, p, { return p->DeRefItr(index); }); + TVM_FFI_DISPATCH_MAP_CONST(self, p, { return p->DeRefItr(index); }); } inline MapObj::iterator& MapObj::iterator::operator++() { TVM_FFI_MAP_FAIL_IF_CHANGED() - TVM_DISPATCH_MAP_CONST(self, p, { + TVM_FFI_DISPATCH_MAP_CONST(self, p, { index = p->IncItr(index); return *this; }); @@ -1161,42 +1209,42 @@ inline MapObj::iterator& MapObj::iterator::operator++() { inline MapObj::iterator& MapObj::iterator::operator--() { TVM_FFI_MAP_FAIL_IF_CHANGED() - TVM_DISPATCH_MAP_CONST(self, p, { + TVM_FFI_DISPATCH_MAP_CONST(self, p, { index = p->DecItr(index); return *this; }); } inline size_t MapObj::count(const key_type& key) const { - TVM_DISPATCH_MAP_CONST(this, p, { return p->count(key); }); + TVM_FFI_DISPATCH_MAP_CONST(this, p, { return p->count(key); }); } inline const MapObj::mapped_type& MapObj::at(const MapObj::key_type& key) const { - TVM_DISPATCH_MAP_CONST(this, p, { return p->at(key); }); + TVM_FFI_DISPATCH_MAP_CONST(this, p, { return p->at(key); }); } inline MapObj::mapped_type& MapObj::at(const MapObj::key_type& key) { - TVM_DISPATCH_MAP(this, p, { return p->at(key); }); + TVM_FFI_DISPATCH_MAP(this, p, { return p->at(key); }); } inline MapObj::iterator MapObj::begin() const { - TVM_DISPATCH_MAP_CONST(this, p, { return p->begin(); }); + TVM_FFI_DISPATCH_MAP_CONST(this, p, { return p->begin(); }); } inline MapObj::iterator MapObj::end() const { - TVM_DISPATCH_MAP_CONST(this, p, { return p->end(); }); + TVM_FFI_DISPATCH_MAP_CONST(this, p, { return p->end(); }); } inline MapObj::iterator MapObj::find(const MapObj::key_type& key) const { - TVM_DISPATCH_MAP_CONST(this, p, { return p->find(key); }); + TVM_FFI_DISPATCH_MAP_CONST(this, p, { return p->find(key); }); } inline void MapObj::erase(const MapObj::iterator& position) { - TVM_DISPATCH_MAP(this, p, { return p->erase(position); }); + TVM_FFI_DISPATCH_MAP(this, p, { return p->erase(position); }); } -#undef TVM_DISPATCH_MAP -#undef TVM_DISPATCH_MAP_CONST +#undef TVM_FFI_DISPATCH_MAP +#undef TVM_FFI_DISPATCH_MAP_CONST inline ObjectPtr MapObj::Empty() { return SmallMapObj::Empty(); } diff --git a/ffi/include/tvm/ffi/container/tuple.h b/ffi/include/tvm/ffi/container/tuple.h index 10303f0aecab..27f08e7fc9f4 100644 --- a/ffi/include/tvm/ffi/container/tuple.h +++ b/ffi/include/tvm/ffi/container/tuple.h @@ -130,10 +130,7 @@ class Tuple : public ObjectRef { private: static ObjectPtr MakeDefaultTupleNode() { - ObjectPtr p = make_inplace_array_object(sizeof...(Types)); - p->capacity_ = sizeof...(Types); - // immeidate set size to 0, to ensure exception safety - p->size_ = 0; + ObjectPtr p = ArrayObj::Empty(sizeof...(Types)); Any* itr = p->MutableBegin(); // increase size after each new to ensure exception safety ((new (itr++) Any(Types()), p->size_++), ...); @@ -142,10 +139,7 @@ class Tuple : public ObjectRef { template static ObjectPtr MakeTupleNode(UTypes&&... args) { - ObjectPtr p = make_inplace_array_object(sizeof...(Types)); - p->capacity_ = sizeof...(Types); - // immeidate set size to 0, to ensure exception safety - p->size_ = 0; + ObjectPtr p = ArrayObj::Empty(sizeof...(Types)); Any* itr = p->MutableBegin(); // increase size after each new to ensure exception safety ((new (itr++) Any(Types(std::forward(args))), p->size_++), ...); @@ -155,10 +149,7 @@ class Tuple : public ObjectRef { /*! \brief Copy on write */ void CopyIfNotUnique() { if (!data_.unique()) { - ObjectPtr p = make_inplace_array_object(sizeof...(Types)); - p->capacity_ = sizeof...(Types); - // immeidate set size to 0, to ensure exception safety - p->size_ = 0; + ObjectPtr p = ArrayObj::Empty(sizeof...(Types)); Any* itr = p->MutableBegin(); const Any* read = GetArrayObj()->begin(); // increase size after each new to ensure exception safety