[IR Container] Phase 2.3 Basic shared ptr

[mdavis36]

Summary

Replace unique_ptr<IrContainer> with shared_ptr<IrContainer> in Fusion and add the container-side registration API that tracks which Fusions share a given container.

This is the foundational change of Phase 2. The pointer type change alone has no behavioral impact — single-Fusion containers behave identically under shared_ptr. The tracking infrastructure (sharing_fusions_, addFusion/removeFusion) lays the groundwork for all subsequent tasks.

Parallel compilation is disabled in-code via kPhase2DisableParallelCompile = true as a precaution during the transition. This ensures CI runs serial compilation for later PRs without requiring environment variables. Parallel compilation is restored in #5971 .

Relationship to Phase 2

This is the core type change that enables the shared container model:

Phase 1 (unique_ptr — exclusive ownership):
  Fusion A ──→ unique_ptr<IrContainer_A> ──→ {vals_A, exprs_A}
  Fusion B ──→ unique_ptr<IrContainer_B> ──→ {vals_B, exprs_B}

This PR (Basic shared_ptr — exclusive ownership):
  Fusion A ──→ shared_ptr<IrContainer_A> ──→ {vals_A, exprs_A}
  Fusion B ──→ shared_ptr<IrContainer_B> ──→ {vals_B, exprs_B}

Phase 2 (shared_ptr — shared storage):
  Fusion A ─┐
             ├──→ shared_ptr<IrContainer> ──→ {vals_A, vals_B, exprs_A, exprs_B}
  Fusion B ─┘

  IrContainer tracks: sharing_fusions_ = {A, B}

Without this change, no subsequent Phase 2 work (per-Fusion tracking, shared-container copy/move, thread safety) is possible. The tracking API is the mechanism that makes container sharing safe — it allows the container to know its consumers, enabling correct cleanup when Fusions are destroyed.

CI Risk

Low. For single-Fusion containers (all existing usage), shared_ptr is behaviorally identical to unique_ptr. No accessor paths change. Parallel compilation is serialized by the in-code constant.

[mdavis36]

!test

[github-actions]

Review updated until commit fab1748

Description

• Replace unique_ptr<IrContainer> with shared_ptr<IrContainer> in Fusion for future container sharing

• Add Fusion tracking API to IrContainer (addFusion, removeFusion, transferFusion, sharingCount, hasMultipleFusions, sharingFusions)

• Remove IrContainer::parent_ backpointer as 1:1 relationship no longer holds

• Disable parallel compilation via kPhase2DisableParallelCompile = true during transition

Changes walkthrough

	Relevant files
Enhancement	fusion.h Convert ir_container_ to shared_ptr                                            csrc/fusion.h Change ir_container_ from std::unique_ptr to std::shared_ptr Add ir_container_ptr() method to expose the shared_ptr Update protected accessor comments for direct container access +5/-2      fusion.cpp Update Fusion to use shared_ptr and tracking API                  csrc/fusion.cpp Use std::make_shared in constructor and call addFusion(this) for registration Call removeFusion(this) in destructor for cleanup Update swap to remove parent pointer updates (now updates Statement pointers only) Pass destination Fusion to IrContainer::copy for proper tracking +10/-11  container.h Add Fusion tracking API to IrContainer                                      csrc/ir/container.h Remove parent_ member variable and parent() method Add tracking API: addFusion, removeFusion, transferFusion, sharingCount, hasMultipleFusions, sharingFusions Add sharing_fusions_ unordered_set to track Fusions using this container Update copy signature to accept Fusion* dest_fusion parameter +11/-9    container.cpp Implement IrContainer tracking and remove parent backpointer csrc/ir/container.cpp Remove std::swap(a.parent_, b.parent_) from swap function Update copy to use dest_fusion for IrCloner instead of parent Update inContainer to check sharing_fusions_ instead of parent_ Implement all new tracking methods (addFusion, removeFusion, transferFusion, etc.) +31/-5
Configuration changes	fusion_kernel_runtime.cpp Disable parallel compilation during shared_ptr transition csrc/runtime/fusion_kernel_runtime.cpp Add kPhase2DisableParallelCompile = true constant to disable parallel compile during transition Use the constant in compileFusionParallel to conditionally disable parallel compilation +7/-2

PR Reviewer Guide

Here are some key observations to aid the review process:

🧪 PR contains tests

⚡ Recommended focus areas for review

Thread Safety The sharing_fusions_ set is modified by addFusion, removeFusion, and transferFusion without any synchronization. When parallel compilation is re-enabled in #5971, concurrent access to these methods from multiple threads could cause race conditions. Consider if std::shared_mutex or other synchronization is needed. void IrContainer::addFusion(Fusion* fusion) { sharing_fusions_.insert(fusion); } void IrContainer::removeFusion(Fusion* fusion) { sharing_fusions_.erase(fusion); } void IrContainer::transferFusion(Fusion* from, Fusion* to) { sharing_fusions_.erase(from); sharing_fusions_.insert(to); } Incomplete Move Semantics The move assignment operator (Fusion::operator=(Fusion&& other)) calls swap() which transfers container contents but does not update the sharing_fusions_ tracking. After a move, the destination Fusion should call addFusion(this) on the container, and the source should call removeFusion(this). This needs verification to ensure proper tracking after moves. Fusion& Fusion::operator=(Fusion&& other) noexcept { FUSER_PERF_SCOPE("Fusion move assign"); clear(); swap(*this, other); return *this; } Hardcoded Parallel Compile Disable The kPhase2DisableParallelCompile = true is hardcoded. While this is intentional for the transition, it should be clearly documented when and how this will be reverted in #5971 to ensure this temporary workaround doesn't become permanent technical debt. // TODO: Remove when std::shared_mutex is added to IrContainer. constexpr bool kPhase2DisableParallelCompile = true;

[mdavis36]

!test

[greptile-apps]

Greptile Summary

This PR is the core Phase 2.3 step that converts Fusion::ir_container_ from unique_ptr<IrContainer> to shared_ptr<IrContainer> and replaces the single parent_ back-pointer in IrContainer with a sharing_fusions_ set, along with the registration API (addFusion/removeFusion/transferFusion) that will let multiple Fusions safely share one container in later Phase 2 PRs. For all existing single-Fusion containers the behaviour is identical to before. Parallel compilation is serialised in-code via kPhase2DisableParallelCompile = true as a safety measure during the transition.

Key observations:

• Destructor ordering risk for shared containers: Fusion::~Fusion() calls clear() (which destroys all vals/exprs in the container) before calling removeFusion(this). In Phase 2.3 this is fine, but once real container sharing is enabled a destroying Fusion will wipe the backing store out from under every other co-owning Fusion before deregistering. The destructor logic should gate the clear() call on sharingCount() == 0.
• Dangling Fusion* in inContainer: const_stmt->container() returns the raw Fusion* stored in Statement::ir_container_. The new ir_container_ptr() accessor is the first mechanism that lets the container outlive its Fusions; once used, surviving statements will hold a dangling raw pointer. This needs to be resolved (e.g. switching ir_container_ in Statement to a weak reference) before actual sharing is activated.
• ir_container_ptr() missing lifecycle contract docs: There is no documentation requiring callers to invoke addFusion()/removeFusion() around any new Fusion that takes ownership of the returned pointer.
• Dead code from kPhase2DisableParallelCompile = true: The else thread-launch branch in compileFusionParallel is permanently unreachable while the constant is true, which will produce warnings from static-analysis tools.

Confidence Score: 3/5

• Safe to merge for Phase 2.3 (single-Fusion containers, no actual sharing), but contains latent correctness issues that must be resolved before real container sharing is enabled in later PRs.
• The changes are behaviourally equivalent to the previous unique_ptr model for all existing call sites, and the kPhase2DisableParallelCompile guard removes the one area of concurrent risk. However, two forward-looking correctness hazards are introduced in the same changeset that makes them possible: (1) the destructor's clear()-before-removeFusion() ordering will silently corrupt shared containers the moment real sharing lands, and (2) inContainer's use of a raw Fusion* becomes UB once ir_container_ptr() allows the container to outlive its Fusions. These are not regressions today but they are structural traps for the very next Phase 2 PR.
• csrc/fusion.cpp (destructor ordering) and csrc/ir/container.cpp (inContainer raw-pointer dereference) need the most attention before Phase 2 sharing is activated.

Important Files Changed

Filename	Overview
csrc/ir/container.h	Replaces single parent_ back-pointer with sharing_fusions_ unordered_set; adds addFusion/removeFusion/transferFusion tracking API; copy() now requires explicit Fusion* dest_fusion argument. Logic is correct for Phase 2.3 (1-to-1 containers), but the public-API methods have no lifecycle documentation.
csrc/ir/container.cpp	Implements the new tracking API and updates inContainer to check sharing_fusions_; IrContainer::swap no longer swaps the removed parent_ field. The inContainer consistency check calls const_stmt->container() (a raw Fusion*) which becomes a dangling pointer once a container outlives its Fusions — safe now but a forward risk once ir_container_ptr() is actually used.
csrc/fusion.h	Upgrades ir_container_ from unique_ptr to shared_ptr and exposes ir_container_ptr() for future sharing scenarios. The new accessor lacks documentation about the mandatory addFusion()/removeFusion() lifecycle contract that callers must satisfy.
csrc/fusion.cpp	Default constructor switched to make_shared + addFusion; destructor now calls removeFusion after clear(). The ordering (clear() before removeFusion) is safe in Phase 2.3 but will destroy shared content for co-owning Fusions once real sharing lands. Fusion::copy correctly threads dest_fusion through to IrCloner.
csrc/runtime/fusion_kernel_runtime.cpp	Introduces kPhase2DisableParallelCompile = true to serialise compilation during the transition. Correct and well-commented, but the constexpr true makes the parallel else branch unreachable dead code that static analysers will flag.

Sequence Diagram

sequenceDiagram
    participant FA as Fusion A
    participant CA as IrContainer A
    participant FB as Fusion B
    participant CB as IrContainer B

    note over FA,CA: Phase 2.3 construction
    FA->>CA: make_shared IrContainer
    FA->>CA: addFusion(A) → sharing_fusions={A}
    FB->>CB: make_shared IrContainer
    FB->>CB: addFusion(B) → sharing_fusions={B}

    note over FA,CB: Fusion::swap(A, B)
    FA->>CA: IrContainer::swap swaps vals and exprs contents
    FA->>CA: update stmt->ir_container_ = A
    FB->>CB: update stmt->ir_container_ = B
    note over CA,CB: sharing_fusions unchanged — each container still tracks its original Fusion

    note over FA,CA: Destruction of Fusion A
    FA->>CA: clear() destroys vals_up_ and exprs_up_
    FA->>CA: removeFusion(A) → sharing_fusions={}
    CA-->>FA: shared_ptr refcount drops to 0 IrContainer A destroyed

    note over CA,CB: Future risk — ir_container_ptr() can let container outlive Fusion leaving stmt->ir_container_ dangling

Loading

_{Last reviewed commit: dcc99d7}

[greptile-apps]

_{5 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Additional Comments (1)

csrc/fusion.cpp
Stale comment references removed code

These comments reference "parent backpointers" and say "each Fusion owns a different IrContainer", both of which are artifacts of the old parent_ pointer model. Since parent_ was removed and the comment no longer accurately describes the code, it should be updated. The actual purpose of the loop below is to update Statement::ir_container_ pointers so that statements point to the correct owning Fusion after the content swap.

  // After swapping container contents, update Statement::ir_container_
  // pointers so each Statement points to the Fusion whose container now
  // holds it.

[greptile-apps]

Additional Comments (4)

csrc/fusion.h, line 168
const method returns mutable shared_ptr

ir_container_ptr() is declared const but returns a std::shared_ptr<IrContainer> (non-const), which allows callers to mutate the IrContainer — and even call addFusion/removeFusion — through a const Fusion reference. This breaks const-correctness. Consider returning std::shared_ptr<const IrContainer> instead, or providing two overloads:

std::shared_ptr<IrContainer> ir_container_ptr();
std::shared_ptr<const IrContainer> ir_container_ptr() const;

---

csrc/ir/container.h, line 138
Tracking API is fully public without access restriction

addFusion, removeFusion, transferFusion, sharingCount, hasMultipleFusions, and sharingFusions are all placed in the public section, meaning any code — not just Fusion — can manipulate sharing_fusions_. Since the invariant "a Fusion and its container stay in sync" is safety-critical for Phase 2, it would be safer to use the existing passkey pattern (similar to IrBuilderPasskey / IrContainerPasskey) to restrict who can call addFusion/removeFusion/transferFusion:

class IrContainerFusionPasskey {
  friend class Fusion;
 private:
  explicit IrContainerFusionPasskey() = default;
};
// Then:
void addFusion(IrContainerFusionPasskey, Fusion* fusion);
void removeFusion(IrContainerFusionPasskey, Fusion* fusion);
void transferFusion(IrContainerFusionPasskey, Fusion* from, Fusion* to);

At minimum, Fusion is already a friend via the existing friend class Fusion; declaration, so these mutation methods could be moved to private and relied upon through friendship alone.

---

csrc/runtime/fusion_kernel_runtime.cpp, line 29
No runtime escape hatch for the parallel-compile disable

kPhase2DisableParallelCompile is a constexpr bool, so it cannot be overridden at runtime. Unlike isOptionDisabled(DisableOption::ParallelCompile), there is no way to re-enable parallel compilation for testing without recompiling. If any CI job or developer needs to exercise the parallel code path during Phase 2, they would have to change the constant and rebuild. Consider documenting this explicitly in the comment (or using a weak/env-based runtime override) so that the limitation is obvious to future readers before #5971 lands.

---

csrc/ir/container.h, line 138
sharingFusions() returns a reference to an internal mutable container

sharingFusions() returns const std::unordered_set<Fusion*>&. Callers that hold this reference while addFusion / removeFusion is concurrently (or even serially) called risk iterator invalidation, since unordered_set rehashes on insertion. Given that thread-safety is planned for a later Phase 2 PR, it would be safer to return by value here, or at least document the caveat clearly:

  const std::unordered_set<Fusion*> sharingFusions() const;

[greptile-apps]

Additional Comments (4)

csrc/ir/container.cpp, line 162
removeFusion silently ignores unregistered fusions

removeFusion calls erase on sharing_fusions_, which silently does nothing if fusion was never registered via addFusion. This could mask lifecycle bugs — for example, if a Fusion is double-removed or removed without ever being added.

Consider adding a debug/NVF_ERROR assertion that fusion is present before erasing:

void IrContainer::removeFusion(Fusion* fusion) {
  NVF_ERROR(
      sharing_fusions_.count(fusion) > 0,
      "Attempted to remove a Fusion that is not registered with this container.");
  sharing_fusions_.erase(fusion);
}

---

csrc/ir/container.cpp, line 122
IrContainer::clear() leaves sharing_fusions_ intact

IrContainer::clear() resets all vals/exprs but does not touch sharing_fusions_. This is presumably intentional (the container is kept alive and the fusion-tracking must survive a clear() call), but the invariant is subtle: after clear(), the container is empty yet still claims to be shared by its registered fusions.

A comment explaining why sharing_fusions_ is deliberately not cleared here would help future maintainers understand the design contract, especially when Phase 2 introduces actual container sharing.

---

csrc/runtime/fusion_kernel_runtime.cpp, line 28
Compile-time constant shadows the runtime ParallelCompile disable option

With kPhase2DisableParallelCompile = true as a compile-time constant, the isOptionDisabled(DisableOption::ParallelCompile) path in the condition below is only reachable for disabling parallel compilation — the runtime option can no longer be used to enable it while the constant is set. This means any tests or workflows that rely on DisableOption::ParallelCompile being togglable at runtime will see changed behavior until #5971 lands.

The intent is documented in the PR and the TODO comment, but it's worth ensuring CI tests that exercise parallel compilation explicitly account for this temporary change in behavior.

---

csrc/ir/container.h, line 135
transferFusion is defined but never called

transferFusion is added as infrastructure for future phases but has no callers in this PR. While this is intentional (laying the groundwork for Phase 2 container sharing), leaving it unused may invite confusion about when it should be called vs. removeFusion + addFusion. A brief doc comment explaining the intended use case (e.g., "called when moving a Fusion object to update the container's tracking") would clarify its purpose.

_{Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!}

[mdavis36]

This additional argument is only temporary while IrContainer::copy is a valid call from Fusion.

[mdavis36]

!test

[greptile-apps]

Additional Comments (3)

csrc/ir/container.h, line 138
Mutation tracking methods should be private

addFusion, removeFusion, and transferFusion are declared public, but they directly manipulate the internal sharing_fusions_ tracking set. Any external code that holds a pointer to an IrContainer can call these to add/remove arbitrary Fusion* entries, corrupting the ownership invariant that inContainer relies on.

Since Fusion is already declared a friend class IrContainer (line 97), all three mutation methods can be moved to private:/protected: and still be accessible from Fusion without any other changes. The read-only query methods (sharingCount, hasMultipleFusions, sharingFusions) are a different story and can remain public if needed.

 private:
  void addFusion(Fusion* fusion);
  void removeFusion(Fusion* fusion);
  void transferFusion(Fusion* from, Fusion* to);

 public:
  size_t sharingCount() const;
  bool hasMultipleFusions() const;
  const std::unordered_set<Fusion*>& sharingFusions() const;

---

csrc/ir/container.cpp, line 167
transferFusion is defined but never called

transferFusion is defined here and declared in container.h, but a search of the entire csrc/ tree shows it is not called anywhere in this PR or the existing codebase. This is dead code — if it is infrastructure for a later Phase 2 task, a comment to that effect (similar to how kPhase2DisableParallelCompile is annotated with a TODO) would make the intent clear. If it is truly unused for now, consider removing it until it is actually needed to keep the API surface minimal.

---

csrc/fusion.cpp, line 268
Move constructor leaves other's container tracking inconsistent

The move constructor delegates to the default constructor first (which calls addFusion(this) on a freshly created IrContainer), then calls swap(*this, other). IrContainer::swap swaps the contents (vals, exprs, name counters) but intentionally does not swap sharing_fusions_. This means after construction:

• this->ir_container_ → new_container, whose sharing_fusions_ = {this} ✓
• other.ir_container_ → other_original_container, whose sharing_fusions_ = {&other} ✓ (and now holds empty/default-constructed data)

When other is later destroyed its destructor calls removeFusion(&other) on other_original_container, which is correct.

However, the other_original_container that other ends up owning was created by this's default constructor and had addFusion(this) called on it, making sharing_fusions_ = {this} — then after IrContainer::swap, that container is now owned by other, yet sharing_fusions_ still says {this}. The destructor of other then calls removeFusion(&other), erasing &other — but {this} is already there and is never erased from other_original_container.

Concretely:

new_container created → new_container.sharing_fusions_ = {this}
IrContainer::swap(new_container, other_original_container):
  - new_container now has other's data;   new_container.sharing_fusions_ stays {this}  ✓
  - other_original holds empty data;  other_original.sharing_fusions_ stays {&other} — but was briefly {this}

The new_container's sharing_fusions_ ends up correct. The other_original_container path is: starts as {this} (from default ctor), then after IrContainer::swap its contents become empty, and ultimately other's destructor calls removeFusion(&other) erasing &other — but &other was never added to it. This means removeFusion(&other) on other_original_container is effectively a no-op erase on the wrong element, and other_original_container.sharing_fusions_ is left with {this} as a dangling pointer when the container is destroyed.

While this dangling entry lives only until other_original_container is destroyed (when other's ir_container_ shared_ptr refcount drops to 0), any code that calls inContainer on the half-dead other_original_container between the move construction and other's destruction would see a stale/dangling pointer in sharing_fusions_.

[greptile-apps]

Additional Comments (4)

csrc/fusion.cpp, line 290
clear() before removeFusion() is unsafe for shared containers

clear() calls ir_container()->clear(), which destroys every Val and Expr in the container before removeFusion(this) runs. In Phase 2.3 each container is still exclusively owned by one Fusion, so this ordering is harmless. However, once ir_container_ptr() is used in a later PR to share a container between two Fusions, a Fusion being destroyed will nuke all vals/exprs belonging to the other still-live Fusion before it unregisters itself.

The safer ordering would be to call removeFusion(this) first, and only then clear the container content (or, better, decide based on sharingCount() whether to clear at all):

Fusion::~Fusion() {
  if (ir_container_) {
    ir_container_->removeFusion(this);
    // Only clear the backing store when this was the last owner.
    if (ir_container_->sharingCount() == 0) {
      clear();
    }
  }
}

As written, the invariant "a container is cleared only when no Fusion still needs it" is violated the moment actual sharing is switched on.

---

csrc/ir/container.cpp, line 138
container() returns a raw Fusion* that can dangle once sharing is live

const_stmt->container() returns the raw Fusion* stored in Statement::ir_container_. As long as containers are 1-to-1 with Fusions this is safe — the Fusion's destructor calls clear(), which destroys all statements before any pointer goes stale.

The new ir_container_ptr() accessor (added in fusion.h) is the first step toward letting the IrContainer outlive its Fusions. The moment a container is shared and one owning Fusion is destroyed:

1. ~Fusion() calls clear(), destroying the container's vals_up_ / exprs_up_ deques — which in turn destroys the Val/Expr objects themselves.
2. Any remaining Fusion still holds a live shared_ptr<IrContainer>, but every statement it contained has been deleted.
3. On a later inContainer() call, const_stmt->container() dereferences a dangling pointer.

The root cause is that Statement::ir_container_ stores a raw Fusion* rather than something whose lifetime is managed alongside the container. This should be tracked as a prerequisite before actual container sharing is enabled, to avoid silent UB.

---

csrc/runtime/fusion_kernel_runtime.cpp, line 443
Constant true makes the else branch permanently dead code

Because kPhase2DisableParallelCompile is a constexpr bool set to true, the short-circuit evaluation of || means the else block (thread-pool launch) can never execute. Most static analyzers (and the compiler with -Wunreachable-code) will flag the dead else path.

Consider wrapping the parallel path in an #if directive, or disabling the warning locally, to silence diagnostics and make the intent more explicit:

if (num_groups == 1 ||
    isOptionDisabled(DisableOption::ParallelCompile)) {
  compileKernel(group_runtime_inputs, group_to_run);
#ifndef PHASE2_DISABLE_PARALLEL_COMPILE
} else {
  // launch compileKernel thread here
  ...
}
#endif

Alternatively, a single NOLINT comment beside kPhase2DisableParallelCompile to acknowledge the deliberate dead code would be sufficient for now.

---

csrc/fusion.h, line 168
ir_container_ptr() needs a doc-comment warning callers to call addFusion()

This accessor returns the raw shared_ptr<IrContainer>, making it trivial for a caller to store the pointer and use it to construct (or replace) another Fusion's container without ever calling addFusion() — leaving sharing_fusions_ out of sync.

A short doc-comment spelling out the contract would prevent misuse:

  /// Returns the shared_ptr to the underlying IrContainer.
  /// IMPORTANT: Any Fusion that takes ownership of this container
  /// must call `ir_container_->addFusion(new_fusion)` and, on
  /// destruction, `ir_container_->removeFusion(new_fusion)` to keep
  /// the sharing_fusions_ tracking accurate.
  std::shared_ptr<IrContainer> ir_container_ptr() const {
    return ir_container_;
  }

[mdavis36]

!test