PR #5961 Review Fixes — Per-Fusion Statement Tracking

[mdavis36]

Summary

Review fixes for PR #5961 (Per-Fusion statement tracking):

• O(n²) → O(n): Optimize removeStatementsOwnedBy with std::erase_if
• Per-Fusion counts: Convert numExprs()/numVals() to return per-Fusion counts instead of global
• StatementGuard fixes: Snapshot and compare per-Fusion counts for correct LIFO rollback in shared containers
• LIFO assertions: Verify tail elements belong to this Fusion before popping

Tests

All tests pass:

• ✅ StatementGuardTest.ExecuteAfterGuard
• ✅ StatementGuardTest.LazySpecialValsNotDangling
• ✅ FusionCopy_CUDA
• ✅ FusionMove_CUDA

[mdavis36]

!test

[github-actions]

Description

• Optimize statement removal performance from O(n²) to O(n) using std::erase_if

• Convert numExprs/numVals to return per-Fusion counts instead of global container sizes

• Add numValsExcludingShortcuts() helper for correct LIFO rollback in StatementGuard

• Add LIFO assertions to verify tail elements belong to current Fusion before popping

Changes walkthrough

Relevant files

Enhancement

fusion.cpp Fix removeStatementsCreatedAfter for per-Fusion counts and LIFO verification csrc/fusion.cpp Update removeStatementsCreatedAfter to use per-Fusion counts via exprsOwnedBy(this) Replace global deque size checks with LIFO assertions for tail element ownership Change val removal to use numValsExcludingShortcuts() for correct shortcut handling Add NVF_ERROR assertions to verify LIFO invariant before popping expressions and values +11/-17  container.cpp Optimize removeStatementsOwnedBy performance with std::erase_if csrc/ir/container.cpp Replace manual iterate-and-erase loops with std::erase_if for O(n) performance Optimize both val and expr removal in removeStatementsOwnedBy function Fix O(n²) performance issue when destroying Fusions with many statements +14/-14  statement_guard.cpp Update StatementGuard to use per-Fusion counts excluding shortcuts csrc/statement_guard.cpp Change StatementGuard to snapshot numValsExcludingShortcuts() instead of numVals() Ensure snapshot consistency with LIFO rollback logic in removeStatementsCreatedAfter Properly handle shortcut singletons that persist across guard scopes +1/-1      fusion.h Convert count methods to per-Fusion and add shortcut exclusion helper csrc/fusion.h Convert numExprs() to return per-Fusion count using exprsOwnedBy(this) Convert numVals() to return per-Fusion count using valsOwnedBy(this) Add numValsExcludingShortcuts() helper that subtracts shortcut val pointers Update comments to clarify per-Fusion counting behavior +16/-3

PR Reviewer Guide

Here are some key observations to aid the review process:

🧪 PR contains tests

⚡ Recommended focus areas for review

LIFO Invariant Validation The new NVF_ERROR checks verify that tail elements belong to the correct Fusion before popping. This is crucial for maintaining the LIFO invariant in per-fusion statement tracking. The error messages are clear and provide context for debugging. c->per_fusion_exprs_[this].count(e) > 0, "removeStatementsCreatedAfter: tail expr belongs to another Fusion"); Performance Optimization The optimization from manual iterator manipulation to std::erase_if is a good improvement. This changes the complexity from O(n²) to O(n) for the removeStatementsOwnedBy function. The lambda-based approach is cleaner and more maintainable. std::erase_if(vals_up_, [&](const std::unique_ptr<Val>& v) { if (owned.count(v.get()) > 0) { vals_.erase(v.get()); return true; } return false; }); Shortcut Value Handling The new numValsExcludingShortcuts() method correctly handles the edge case where shortcut values (zero_val_, one_val_, etc.) should persist across StatementGuard scopes. This ensures LIFO ordering works correctly by excluding these singleton values from the count used for rollback. int64_t numValsExcludingShortcuts() const noexcept { int64_t count = std::ssize(ir_container()->valsOwnedBy(this)); count -= (zero_val_ != nullptr) + (one_val_ != nullptr) + (true_val_ != nullptr) + (false_val_ != nullptr) + (magic_zero_val_ != nullptr); return count; }

[greptile-apps]

Greptile Summary

Applies review fixes from PR #5961 for per-Fusion statement tracking in shared containers:

• Performance: Optimizes removeStatementsOwnedBy from O(n²) to O(n) using std::erase_if
• Per-Fusion semantics: Converts numExprs()/numVals() to return per-Fusion counts instead of global counts
• LIFO correctness: Updates removeStatementsCreatedAfter to use per-Fusion counts with assertions verifying tail elements belong to this Fusion before popping
• Shortcut vals: Adds numValsExcludingShortcuts() to properly handle singleton shortcut vals (zero_val_, one_val_, etc.) that should persist across StatementGuard scopes

All changes are well-implemented with defensive assertions to catch violations of the LIFO invariant. The optimizations are algorithmically sound and the per-Fusion counting ensures correct behavior when multiple Fusions share an IrContainer.

Confidence Score: 5/5

• This PR is safe to merge with no issues found
• All changes are correct, well-tested (tests pass), include defensive assertions, and implement sound algorithmic optimizations. The per-Fusion counting logic properly handles edge cases like shortcut vals, and LIFO assertions ensure proper StatementGuard nesting
• No files require special attention

Important Files Changed

Filename	Overview
csrc/ir/container.cpp	Optimizes removeStatementsOwnedBy from O(n²) to O(n) using std::erase_if, correctly maintains both deque and set consistency
csrc/fusion.h	Converts numExprs()/numVals() to per-Fusion counts; adds numValsExcludingShortcuts() for correct StatementGuard rollback with singleton vals
csrc/fusion.cpp	Updates removeStatementsCreatedAfter to use per-Fusion counts and adds LIFO assertions to verify tail elements belong to this Fusion
csrc/statement_guard.cpp	Updates snapshot to use numValsExcludingShortcuts() for correct handling of singleton shortcut vals across guard scopes

_{Last reviewed commit: 0abaf11}

[greptile-apps]

_{4 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}