This is an automated email from the ASF dual-hosted git repository.
tlopex 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 33ffa5a3b4 [Fix][Relax] Run destructors for non-trivially-destructible
types in Arena (#20163)
33ffa5a3b4 is described below
commit 33ffa5a3b415802aaed197622e14ce95c8b05819
Author: Kryptonite <[email protected]>
AuthorDate: Tue Aug 25 00:15:32 2026 +0300
[Fix][Relax] Run destructors for non-trivially-destructible types in Arena
(#20163)
## Summary
`support::Arena` never calls destructors but frees raw memory pages.
`GraphPartitioner::Group` violates this by holding a ref counted `attrs`
map, which leaks native memory on every pass invocation since `~Group()`
is never called.
Fixes #20056.
## Fix
Implements option (b) from the issue: `Arena::make<T>()` now registers
invokes a destructor for non-trivially-destructible types, guarded by
`if constexpr` so trivially destructible types pay no extra cost. Fixes
the leak at the source instead of per call site.
Went with option (b) instead of (a) as the same pattern also exists in
[`LiftedFunctionRewritePlan`] in
([`rewrite_cuda_graph.cc`](https://github.com/apache/tvm/blob/4db33674decd09d8e093d4e260776f9b949b43ca/src/relax/transform/rewrite_cuda_graph.cc#L76))
which would fix `Group` but leave that and any future case unfixed.
## Testing
Reproduced the leak using the script attached to the issue (updated for
the `TensorStructInfo` -> `Type` unification in #19853, otherwise
unchanged): ~4.96 MiB/iter before this fix, ~0 MiB/iter after. Existing
Relax fusion/partitioning tests pass.
---
src/support/arena.h | 55 ++++++++++++++++++++++++++++++--------
tests/cpp/support_test.cc | 67 +++++++++++++++++++++++++++++++++++++++++++++++
2 files changed, 111 insertions(+), 11 deletions(-)
diff --git a/src/support/arena.h b/src/support/arena.h
index 79036326f6..6cf4a28180 100644
--- a/src/support/arena.h
+++ b/src/support/arena.h
@@ -43,13 +43,6 @@
namespace tvm {
namespace support {
-namespace {
-template <typename T> // For lvalues (T is T&),
-T&& forward(T&& param) { // take/return lvalue refs.
- return static_cast<T&&>(param); // For rvalues (T is T),
-} // take/return rvalue refs.
-} // namespace
-
/*!
* \brief An arena page header.
*/
@@ -64,6 +57,21 @@ struct ArenaPageHeader {
size_t offset;
};
+/*!
+ * \brief A pending destructor call for one arena-allocated, non-trivially
+ * destructible object. Forms a singly-linked list of all such objects so
+ * their destructors can be invoked before the arena's pages are freed or
+ * recycled.
+ */
+struct ArenaDeleter {
+ /*! \brief The next pending destructor call. */
+ ArenaDeleter* next;
+ /*! \brief The object to destroy. */
+ void* obj;
+ /*! \brief Function that destroys obj (a type-erased call to ~T()). */
+ void (*dtor)(void*);
+};
+
/*!
* \brief Arena allocator that allocates memory from continuous
* chunk and frees them all only during destruction.
@@ -83,11 +91,13 @@ class GenericArena {
/*! \brief Free all pages. */
void FreeAll() {
+ RunDeleters();
FreePageList(&head_);
FreePageList(&free_list_);
}
/*! \brief Recycle all the pages in the arena */
void RecycleAll() {
+ RunDeleters();
// put all the current list to the free list.
tail_->next = free_list_;
// allocate the first in the free list to head
@@ -115,18 +125,39 @@ class GenericArena {
* \tparam Args Arguments to the constructor.
*
* \return The allocated object.
- * \note The type T must be simple type, or only contain
- * memory allocated from the same arena.
- * Otherwise the destructor needs to be called explicitly.
+ * \note If T is not trivially destructible, its destructor is recorded and
+ * invoked automatically when the arena's pages are freed or recycled.
*/
template <typename T, typename... Args>
T* make(Args&&... args) {
T* ptr = allocate_<T>();
- new (ptr) T(forward<Args>(args)...);
+ new (ptr) T(std::forward<Args>(args)...);
+ if constexpr (!std::is_trivially_destructible<T>::value) {
+ RegisterDeleter(ptr);
+ }
return ptr;
}
private:
+ /*!
+ * \brief Record ptr's destructor to be called by RunDeleters, before the
+ * arena's pages are freed or recycled.
+ */
+ template <typename T>
+ void RegisterDeleter(T* ptr) {
+ ArenaDeleter* node = allocate_<ArenaDeleter>();
+ node->obj = ptr;
+ node->dtor = [](void* p) { static_cast<T*>(p)->~T(); };
+ node->next = deleters_;
+ deleters_ = node;
+ }
+ /*! \brief Invoke and clear all pending destructor calls registered so far.
*/
+ void RunDeleters() {
+ for (ArenaDeleter* d = deleters_; d != nullptr; d = d->next) {
+ d->dtor(d->obj);
+ }
+ deleters_ = nullptr;
+ }
/*! \brief internal page allocator. */
PageAllocator alloc_;
/* \brief The head of the allocated list. */
@@ -135,6 +166,8 @@ class GenericArena {
ArenaPageHeader* tail_{nullptr};
/* \brief List of free pages. */
ArenaPageHeader* free_list_{nullptr};
+ /*! \brief Pending destructor calls for non-trivially-destructible objects.
*/
+ ArenaDeleter* deleters_{nullptr};
/*!
* \brief Align ptr by upper bound.
* \param offset The offset value.
diff --git a/tests/cpp/support_test.cc b/tests/cpp/support_test.cc
index 87f14dce02..243b3f848c 100644
--- a/tests/cpp/support_test.cc
+++ b/tests/cpp/support_test.cc
@@ -20,6 +20,9 @@
#include <gtest/gtest.h>
#include <tvm/runtime/logging.h>
+#include <memory>
+
+#include "../../src/support/arena.h"
#include "../../src/support/utils.h"
namespace tvm {
@@ -43,5 +46,69 @@ TEST(StartsWithTests, Basic) {
EXPECT_FALSE(::tvm::support::StartsWith("abc", "abcd"));
}
+namespace {
+// A non-trivially-destructible type: destructing it has an observable
+// side effect (incrementing a counter), unlike a plain-data struct.
+struct DtorCounter {
+ explicit DtorCounter(int* counter) : counter(counter) {}
+ ~DtorCounter() { (*counter)++; }
+ int* counter;
+};
+} // namespace
+
+TEST(ArenaTests, MakeRunsDestructorOnFreeAll) {
+ int destroyed = 0;
+ {
+ support::Arena arena;
+ for (int i = 0; i < 8; ++i) {
+ arena.make<DtorCounter>(&destroyed);
+ }
+ EXPECT_EQ(destroyed, 0);
+ // FreeAll() may be called directly (not only via ~Arena()), e.g. by
+ // MinRPCServer. It must run pending destructors itself rather than
+ // relying on the caller to do so, or on ~Arena() running afterwards.
+ arena.FreeAll();
+ EXPECT_EQ(destroyed, 8);
+ }
+ // ~Arena() must not re-run the same destructors after an explicit
+ // FreeAll(), and must not touch the now-freed ArenaDeleter bookkeeping.
+ EXPECT_EQ(destroyed, 8);
+}
+
+TEST(ArenaTests, MakeRunsDestructorOnRecycleAll) {
+ int destroyed = 0;
+ support::Arena arena;
+ arena.make<DtorCounter>(&destroyed);
+ arena.RecycleAll();
+ EXPECT_EQ(destroyed, 1);
+
+ arena.make<DtorCounter>(&destroyed);
+ EXPECT_EQ(destroyed, 1);
+}
+
+TEST(ArenaTests, TrivialTypeUnaffected) {
+ support::Arena arena;
+ int* x = arena.make<int>(42);
+ EXPECT_EQ(*x, 42);
+ // No crash/UB freeing an arena that only ever held trivially
+ // destructible objects.
+ arena.FreeAll();
+}
+
+namespace {
+struct MoveOnlyHolder {
+ explicit MoveOnlyHolder(std::unique_ptr<int> value) :
value(std::move(value)) {}
+ std::unique_ptr<int> value;
+};
+} // namespace
+
+TEST(ArenaTests, MakeAcceptsMoveOnlyArgument) {
+ support::Arena arena;
+ // Regression test: make<T>(std::make_unique<...>(...)) must not be
ambiguous with std::forward
+ // via ADL
+ auto* holder = arena.make<MoveOnlyHolder>(std::make_unique<int>(7));
+ EXPECT_EQ(*holder->value, 7);
+}
+
} // namespace test
} // namespace tvm
