Migrate standalone Fusion/HostIrContainer to shared IrContainer

Four code paths created Fusions with fresh IrContainers (default
constructor) instead of sharing the source Fusion's container. This
broke the getCurFusion()-based traversal in iter_visitor.cpp which
assumes all Vals are in the same shared container.

Changes:
- fusion_segmenter.cpp: Welford translation test copy now shares
  the source Fusion's IrContainer
- host_ir/container.h: Add shared-container constructor to
  HostIrContainer (forwarding to Fusion's protected constructor)
- communication_executor.cpp, host_ir/lowering.cpp, host_ir/lower.cpp:
  Use shared-container constructor for HostIrContainer creation

Author: mdavis36

csrc/fusion_segmenter.cpp

@@ -2803,8 +2803,10 @@ bool TranslateApplicableWelford::wouldTranslateToPersistent(
           [fusion](WelfordOp* welford) { return welford->fusion() == fusion; }),
       "Welfords in given vector not in the same fusion");
 
-  // Make initial `in-progress copy`
-  auto test_copy = std::make_unique<Fusion>();
+  // Make initial `in-progress copy` — share the source IrContainer so that
+  // traversal via getCurFusion() sees all Vals in the same container.
+  auto test_copy =
+      std::unique_ptr<Fusion>(new Fusion(fusion->ir_container_ptr()));
   auto original_to_test_map = Fusion::copy(fusion, test_copy.get());
 
   std::vector<WelfordOp*> copied_welfords;

csrc/host_ir/container.h

@@ -23,6 +23,8 @@ namespace nvfuser::hir {
 class HostIrContainer final : public Fusion {
  public:
   HostIrContainer() = default;
+  explicit HostIrContainer(std::shared_ptr<IrContainer> container)
+      : Fusion(std::move(container)) {}
   HostIrContainer(const HostIrContainer&) = delete;
   HostIrContainer& operator=(const HostIrContainer&) = delete;
 

csrc/host_ir/lower.cpp

@@ -118,7 +118,8 @@ std::unique_ptr<hir::HostIrContainer> HostIrLower::lower(
   RuntimeWorkSpace workspace = prepareRuntimeOrder(*staged_fusion);
   // Create the HostIrContainer representing the host program. Each segment of
   // the segmented fusion will be translated to a HostIR
-  auto hic = std::make_unique<hir::HostIrContainer>();
+  auto hic = std::make_unique<hir::HostIrContainer>(
+      staged_fusion->completeFusion()->ir_container_ptr());
   FusionGuard fg(hic.get());
   IrCloner ir_cloner(hic.get());
   auto clone =

csrc/host_ir/lowering.cpp

@@ -396,7 +396,8 @@ std::unique_ptr<hir::HostIrContainer> lowerSegmentedFusionToHostIr(
     const SegmentedFusion& segmented_fusion,
     const std::vector<LaunchParams>& launch_params_per_segment,
     std::vector<std::unique_ptr<ExecutorAbstract>>& executors) {
-  auto hic = std::make_unique<hir::HostIrContainer>();
+  auto hic = std::make_unique<hir::HostIrContainer>(
+      segmented_fusion.completeFusion()->ir_container_ptr());
   IrCloner ir_cloner =
       Fusion::copy(segmented_fusion.completeFusion(), hic.get());
 

csrc/runtime/communication_executor.cpp

@@ -54,7 +54,8 @@ void CommunicationExecutor::compile(Fusion* fusion) {
     FusionProfiler::segment(group_id_).startCompile();
   }
 
-  host_ir_container_ = std::make_unique<hir::HostIrContainer>();
+  host_ir_container_ =
+      std::make_unique<hir::HostIrContainer>(fusion->ir_container_ptr());
   IrCloner cloner = Fusion::copy(fusion, host_ir_container_.get());
   if (fusion->isA<hir::HostIrContainer>()) {
     for (Expr* e : fusion->as<hir::HostIrContainer>()->topLevelExprs()) {