Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

Appearance settings

[Bug][Relax] FuseOpsByPattern leaks native memory on every call (arena-allocated Group::attrs never destructed) #20056

Copy link
Copy link

Description

@JagdishKolhe
Issue body actions

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:

  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

Reactions are currently unavailable

Metadata

Metadata

Assignees

No one assigned

    Labels

    needs-triagePRs or issues that need to be investigated by maintainers to find the right assignees to address itPRs or issues that need to be investigated by maintainers to find the right assignees to address ittype: bug

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions

    Morty Proxy This is a proxified and sanitized view of the page, visit original site.