Check consistency of multiplier

[zasdfgbnm]

No description provided.

[zasdfgbnm]

!test

[github-actions]

Review updated until commit 438a860

Description

• Add consistency checking for multipliers in unshardedSizes function

• Introduce extent_to_multiplier_map parameter to track and validate multiplier consistency

• Update PrecomputedValues to own the extent-to-multiplier mapping

• Modify ExpressionEvaluator to provide access to the multiplier map

• Update all unshardedSizes call sites to pass the consistency map

Changes walkthrough

Relevant files

Enhancement

9 files execution_utils.h Update unshardedSizes function signature                                  +2/-1      execution_utils.cpp Implement multiplier consistency validation logic                +23/-1    evaluator_common.h Add extent_to_multiplier_map ownership to PrecomputedValues +8/-0      expr_evaluator.h Add getExtentToMultiplierMap accessor method                          +5/-0      evaluator_common.cpp Update unshardedSizes call with consistency map                    +1/-1      expr_evaluator.cpp Update unshardedSizes call with consistency map                    +1/-1      allocations.cpp Update multiple unshardedSizes calls with consistency map +5/-4      executor.cpp Update unshardedSizes call with consistency map                    +1/-1      tensor_metadata.cpp Update unshardedSizes calls with consistency map                  +2/-2

PR Reviewer Guide

Here are some key observations to aid the review process:

🧪 No relevant tests

⚡ Recommended focus areas for review

Missing Tests This PR adds consistency checking for multipliers but no tests are included. Given that this is a correctness feature that validates consistency across sharded tensors, tests should be added to verify the consistency checking logic works correctly and that it properly detects and reports inconsistencies. // Check consistency: for the same extent, we should always get the same multiplier // Only perform this check if a map is provided if (extent_to_multiplier_map) { Val* extent = sharded_id->extent(); auto it = extent_to_multiplier_map->find(extent); if (it != extent_to_multiplier_map->end()) { NVF_ERROR( it->second == multiplier, "Inconsistent multiplier for extent ", extent->toString(), ": expected ", it->second, " but got ", multiplier); } else { (*extent_to_multiplier_map)[extent] = multiplier; } } else { // NVF_ERROR(false, "Extent to multiplier map not provided"); } Memory Ownership Clarification The extent_to_multiplier_map is owned by PrecomputedValues but accessed through ExpressionEvaluator. While this works with the current design, consider documenting the ownership semantics and lifetime expectations to prevent potential dangling pointer issues if the design evolves. std::unordered_map<Val*, int64_t>* getExtentToMultiplierMap() { return &extent_to_multiplier_map_; }

[zasdfgbnm]

!test

[zasdfgbnm]

!test

[zasdfgbnm]

!test

[zasdfgbnm]

!test

[zasdfgbnm]

!test