shared_ptr<IrContainer> transition and Fusion tracking infrastructure

Change Fusion::ir_container_ from unique_ptr to shared_ptr to enable
future container sharing between Fusions. Add Fusion tracking API to
IrContainer (addFusion/removeFusion/transferFusion/sharingCount).
Disable parallel compilation during the shared_ptr transition.

Author: mdavis36

csrc/fusion.cpp

@@ -254,8 +254,9 @@ IrCloner Fusion::copy(const Fusion* from, Fusion* to) {
 }
 
 // Default constructor
-Fusion::Fusion() : ir_container_(std::make_unique<IrContainer>()) {
+Fusion::Fusion() : ir_container_(std::make_shared<IrContainer>()) {
   ir_container_->parent_ = this;
+  ir_container_->addFusion(this);
 }
 
 // Copy constructor
@@ -287,6 +288,9 @@ Fusion& Fusion::operator=(Fusion&& other) noexcept {
 
 Fusion::~Fusion() {
   clear();
+  if (ir_container_) {
+    ir_container_->removeFusion(this);
+  }
 }
 
 void Fusion::clear() noexcept {
@@ -350,9 +354,7 @@ void Fusion::removeExpr(Expr* expr) {
   auto expr_in_deque = std::find_if(
       c->exprs_up_.begin(),
       c->exprs_up_.end(),
-      [expr](std::unique_ptr<Expr>& expr_up) {
-        return expr_up.get() == expr;
-      });
+      [expr](std::unique_ptr<Expr>& expr_up) { return expr_up.get() == expr; });
   NVF_ERROR(
       expr_in_deque != c->exprs_up_.end(),
       "Wanted to remove an expression but its unique ptr is missing.");

csrc/fusion.h

@@ -148,7 +148,6 @@ class NVF_API Fusion : public PolymorphicBase {
   typedef std::unordered_map<int, std::vector<int64_t>> PermutationMap;
 
  protected:
-  // Direct access to underlying container
   IrContainer* ir_container() {
     NVF_ERROR(
         ir_container_.get() != nullptr,
@@ -163,6 +162,10 @@ class NVF_API Fusion : public PolymorphicBase {
     return ir_container_.get();
   }
 
+  std::shared_ptr<IrContainer> ir_container_ptr() const {
+    return ir_container_;
+  }
+
  public:
   // Registration (public API with passkey)
   virtual void registerStmt(IrBuilderPasskey, Statement* stmt) {
@@ -635,7 +638,7 @@ class NVF_API Fusion : public PolymorphicBase {
   std::unique_ptr<std::vector<TensorView*>> all_tvs_ptr_ = nullptr;
 
   inline static const std::string exact_mappings_key = "exact_mappings";
-  std::unique_ptr<IrContainer> ir_container_;
+  std::shared_ptr<IrContainer> ir_container_;
 
   Val* zero_val_ = nullptr;
   Val* one_val_ = nullptr;

csrc/ir/container.cpp

@@ -157,4 +157,29 @@ bool IrContainer::inContainer(const Statement* const_stmt) const {
   return true;
 }
 
+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);
+}
+
+size_t IrContainer::sharingCount() const {
+  return sharing_fusions_.size();
+}
+
+bool IrContainer::hasMultipleFusions() const {
+  return sharing_fusions_.size() > 1;
+}
+
+const std::unordered_set<Fusion*>& IrContainer::sharingFusions() const {
+  return sharing_fusions_;
+}
+
 } // namespace nvfuser

csrc/ir/container.h

@@ -133,10 +133,16 @@ class IrContainer {
     return parent_;
   }
 
+  void addFusion(Fusion* fusion);
+  void removeFusion(Fusion* fusion);
+  void transferFusion(Fusion* from, Fusion* to);
+  size_t sharingCount() const;
+  bool hasMultipleFusions() const;
+  const std::unordered_set<Fusion*>& sharingFusions() const;
+
  private:
-  // Parent Fusion that owns this container (for pure composition pattern)
-  // Used by Statement::fusion() to navigate back to owning Fusion
   Fusion* parent_ = nullptr;
+  std::unordered_set<Fusion*> sharing_fusions_;
 };
 
 } // namespace nvfuser

csrc/runtime/fusion_kernel_runtime.cpp

@@ -28,6 +28,9 @@
 
 namespace nvfuser {
 
+// TODO: Remove when std::shared_mutex is added to IrContainer.
+constexpr bool kPhase2DisableParallelCompile = true;
+
 namespace {
 // Replace CUDA tensor with Meta tensor because storing tensors can cause
 // out-of-memory issues. Other arguments are returned as-is.
@@ -454,7 +457,8 @@ void FusionKernelRuntime::compileFusionParallel(KernelArgumentHolder args) {
   try {
     for (const auto& [group_to_run, group_runtime_inputs] :
          zip(runtime_workspace_.group_run_order, all_runtime_inputs)) {
-      if (num_groups == 1 || isOptionDisabled(DisableOption::ParallelCompile)) {
+      if (num_groups == 1 || kPhase2DisableParallelCompile ||
+          isOptionDisabled(DisableOption::ParallelCompile)) {
         compileKernel(group_runtime_inputs, group_to_run);
       } else {
         // launch compileKernel thread here
@@ -488,7 +492,8 @@ void FusionKernelRuntime::compileFusionParallel(KernelArgumentHolder args) {
     throw;
   }
 
-  if (num_groups != 1 && !isOptionDisabled(DisableOption::ParallelCompile)) {
+  if (num_groups != 1 && !kPhase2DisableParallelCompile &&
+      !isOptionDisabled(DisableOption::ParallelCompile)) {
     // Wait until all segments finish compiling
     getThreadPool()->waitWorkComplete();
     NVF_ERROR(