[IR Container] Phase 2.3 Basic shared ptr - #5960

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.

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;

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!}

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_.

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_;
 }