Expected behavior
Running relax.transform.FuseOpsByPattern repeatedly (e.g. once per layer in a
custom quantization/sensitivity loop) should return process memory to a steady state
after each call. Resident memory should not grow monotonically across calls.
Actual behavior
Each FuseOpsByPattern invocation leaks native (C++) memory that is never freed
for the lifetime of the process. It is invisible to Python (tracemalloc stays
flat) and unreclaimable by gc.collect() or malloc_trim(), so long-running
pipelines that fuse many times grow without bound and eventually OOM.
The leak is per pattern processed per call, and independent of
bind_constants / annotate_codegen. With the reproducer below on TVM 0.24.0:
| config |
slope |
PyTraced |
reclaimed by gc+malloc_trim |
| 400 ops, 100 patterns |
+4.96 MiB/iter |
flat (0.02 MiB) |
no |
| 400 ops, 1 pattern |
+0.07 MiB/iter |
flat |
no |
Root cause: PatternBasedPartitioner (src/relax/transform/fuse_ops.cc)
allocates its union-find Group nodes from a support::Arena
(arena_->make<Group>()). support::Arena frees its pages with delete[] and
never runs the objects' destructors — this is documented in
src/support/arena.h (make<T>: "The type T must be simple type ... Otherwise
the destructor needs to be called explicitly"). But
GraphPartitioner::Group (src/relax/analysis/graph_partitioner.h) is not
trivially destructible: it holds ffi::Map<String, Any> attrs. So every group's
attrs map (set via parent_group->attrs.Set(attr::kComposite, ...)) and its
interned Strings leak. Confirmed with valgrind — all leak records are
"definitely/indirectly lost" (orphaned, not a global cache) rooted in
PatternBasedPartitioner::VisitBinding_ / VisitVarDef.
The same arena-allocated Group (with attrs) pattern also appears in the plain
FuseOps path (GraphPartitioner) and in MergeCompositeFunctions, which have
the same latent leak.
Environment
- TVM: 0.24.0 (
v0.24.0-13535-g20a91a4e8e); leaking code path is unmodified from upstream main
- Python: 3.10.12
- OS: Linux (Ubuntu), kernel 6.8, glibc 2.35, x86_64
- Build: LLVM backend, Release
Steps to reproduce
Minimal script (public APIs, no external data): see attached
repro_fuseopsbypattern_leak.py. Run:
python repro_fuseopsbypattern_leak.py # 400 ops, 100 patterns -> ~+5 MiB/iter, monotonic
python repro_fuseopsbypattern_leak.py 400 1 30 # 1 pattern -> near-flat (leak scales with #patterns)
It builds a small Relax module, calls FuseOpsByPattern in a loop (dropping the
result each iteration), and prints RSS, RSS-after-gc, and tracemalloc. RSS climbs
monotonically while tracemalloc stays flat and gc/malloc_trim reclaim nothing →
native leak.
Possible fix
Group (non-trivial) shouldn't live in a trivial-only arena. Two options:
- (a) Allocate the partitioner's groups in an owning container
(std::deque<Group>) so ~Group() runs; or
- (b) Teach
support::Arena to register destructors for
non-trivially-destructible types (à la protobuf Arena), guarded by
if constexpr (!std::is_trivially_destructible_v<T>) so trivial types are
unaffected.
I have a working, verified fix for (a) — fused output is byte-identical and
the leak drops to ~0 — and can open a PR.
repro_fuseopsbypattern_leak.py
Expected behavior
Running
relax.transform.FuseOpsByPatternrepeatedly (e.g. once per layer in acustom quantization/sensitivity loop) should return process memory to a steady state
after each call. Resident memory should not grow monotonically across calls.
Actual behavior
Each
FuseOpsByPatterninvocation leaks native (C++) memory that is never freedfor the lifetime of the process. It is invisible to Python (
tracemallocstaysflat) and unreclaimable by
gc.collect()ormalloc_trim(), so long-runningpipelines that fuse many times grow without bound and eventually OOM.
The leak is per pattern processed per call, and independent of
bind_constants/annotate_codegen. With the reproducer below on TVM 0.24.0:Root cause:
PatternBasedPartitioner(src/relax/transform/fuse_ops.cc)allocates its union-find
Groupnodes from asupport::Arena(
arena_->make<Group>()).support::Arenafrees its pages withdelete[]andnever runs the objects' destructors — this is documented in
src/support/arena.h(make<T>: "The type T must be simple type ... Otherwisethe destructor needs to be called explicitly"). But
GraphPartitioner::Group(src/relax/analysis/graph_partitioner.h) is nottrivially destructible: it holds
ffi::Map<String, Any> attrs. So every group'sattrsmap (set viaparent_group->attrs.Set(attr::kComposite, ...)) and itsinterned
Strings leak. Confirmed with valgrind — all leak records are"definitely/indirectly lost" (orphaned, not a global cache) rooted in
PatternBasedPartitioner::VisitBinding_/VisitVarDef.The same arena-allocated
Group(withattrs) pattern also appears in the plainFuseOpspath (GraphPartitioner) and inMergeCompositeFunctions, which havethe same latent leak.
Environment
v0.24.0-13535-g20a91a4e8e); leaking code path is unmodified from upstreammainSteps to reproduce
Minimal script (public APIs, no external data): see attached
repro_fuseopsbypattern_leak.py. Run:It builds a small Relax module, calls
FuseOpsByPatternin a loop (dropping theresult each iteration), and prints RSS, RSS-after-gc, and tracemalloc. RSS climbs
monotonically while tracemalloc stays flat and gc/malloc_trim reclaim nothing →
native leak.
Possible fix
Group(non-trivial) shouldn't live in a trivial-only arena. Two options:(
std::deque<Group>) so~Group()runs; orsupport::Arenato register destructors fornon-trivially-destructible types (à la protobuf
Arena), guarded byif constexpr (!std::is_trivially_destructible_v<T>)so trivial types areunaffected.
I have a working, verified fix for (a) — fused output is byte-identical and
the leak drops to ~0 — and can open a PR.
repro_fuseopsbypattern_leak.py