[IR Container] Phase 2 Per-fusion statement tracking

[mdavis36]

No description provided.

[github-actions]

Description

• Transition Fusion::ir_container_ from unique_ptr to shared_ptr to enable container sharing

• Add per-fusion statement tracking with per_fusion_vals_/per_fusion_exprs_ maps in IrContainer

• Update Fusion accessors (vals(), deterministic_exprs(), etc.) to return per-Fusion filtered results

• Add Fusion tracking API (addFusion/removeFusion/transferFusion/sharingCount) to IrContainer

• Temporarily disable parallel compilation during shared_ptr transition

Changes walkthrough

	Relevant files
Enhancement	fusion.cpp Update Fusion methods for shared_ptr and per-fusion tracking csrc/fusion.cpp Update Fusion constructor/destructor to use shared_ptr and track fusion ownership Modify clear() to use removeStatementsOwnedBy() for per-Fusion cleanup Update registerVal/registerExpr to populate per-fusion tracking maps Update removeVal/removeExpr to clean up per-fusion maps Fix swap() to transfer statement ownership between Fusions Update copy() to pass destination fusion for proper ownership tracking +22/-23  container.cpp Implement per-fusion statement tracking and ownership management csrc/ir/container.cpp Add per-fusion tracking maps (per_fusion_vals_, per_fusion_exprs_) and sharing_fusions_ set Implement Fusion tracking methods (addFusion, removeFusion, transferFusion, etc.) Add per-Fusion filtered accessors (valsOwnedBy, exprsOwnedBy, deterministicValsOwnedBy, etc.) Add transferStatementOwnership() and removeStatementsOwnedBy() for ownership management Update copy() to accept destination fusion parameter Update swap() to handle per-fusion tracking maps +161/-4  fusion.h Update Fusion interface for shared_ptr and per-fusion accessors csrc/fusion.h Change ir_container_ from unique_ptr to shared_ptr Add ir_container_ptr() accessor method Update deterministic_vals() and deterministic_exprs() to return per-Fusion filtered results Update vals() and unordered_exprs() to return per-Fusion filtered results Update deterministic_vals_map() and deterministic_exprs_map() to return per-Fusion filtered results +15/-14  container.h Update IrContainer interface for shared ownership and per-fusion tracking csrc/ir/container.h Remove parent_ pointer and add sharing_fusions_ set for multiple fusion tracking Add per-fusion tracking maps (per_fusion_vals_, per_fusion_exprs_) Add Fusion tracking API declarations (addFusion, removeFusion, transferFusion, etc.) Add per-Fusion filtered accessor declarations (valsOwnedBy, exprsOwnedBy, etc.) Update copy() signature to accept destination fusion parameter +29/-9
Configuration changes	fusion_kernel_runtime.cpp Temporarily disable parallel compilation during transition csrc/runtime/fusion_kernel_runtime.cpp Add kPhase2DisableParallelCompile flag to temporarily disable parallel compilation Update compileFusionParallel() to check the temporary disable flag +7/-2

PR Reviewer Guide

Here are some key observations to aid the review process:

🧪 No relevant tests

⚡ Recommended focus areas for review

Memory Management The change from unique_ptr to shared_ptr for ir_container_ suggests multiple Fusions can now share the same IrContainer. The transferStatementOwnership and removeStatementsOwnedBy methods need careful review to ensure proper memory management and avoid dangling pointers when Fusions are destroyed. Fusion::Fusion() : ir_container_(std::make_shared<IrContainer>()) { ir_container_->addFusion(this); } Performance Regression The addition of kPhase2DisableParallelCompile flag suggests potential concurrency issues with the new per-fusion tracking. This may cause performance regressions in multi-segment fusion compilation. Need to verify this is temporary and understand the underlying thread safety concerns. // TODO: Remove when std::shared_mutex is added to IrContainer. constexpr bool kPhase2DisableParallelCompile = true; API Compatibility Multiple method signatures have changed (e.g., deterministic_vals() now returns deque<Val*> instead of const deque<Val*>, copy() now takes additional Fusion* parameter). This breaks the const-correctness and may impact existing code that depends on these APIs. std::deque<Val*> deterministicValsOwnedBy( const Fusion* fusion) const noexcept; std::deque<Expr*> deterministicExprsOwnedBy( const Fusion* fusion) const noexcept; std::unordered_map<Val*, int64_t> deterministicValsMapOwnedBy( const Fusion* fusion) const noexcept; std::unordered_map<Expr*, int64_t> deterministicExprsMapOwnedBy( const Fusion* fusion) const noexcept;

[mdavis36]

!test