JagdishKolhe opened a new issue, #20056:
URL: https://github.com/apache/tvm/issues/20056
## 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 `String`s 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:
```bash
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:
1. **(a)** Allocate the partitioner's groups in an owning container
(`std::deque<Group>`) so `~Group()` runs; or
2. **(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](https://github.com/user-attachments/files/30409816/repro_fuseopsbypattern_leak.py)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]