-
Notifications
You must be signed in to change notification settings - Fork 54
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Prevent self mapping in the AlmostExact graph #3926
base: main
Are you sure you want to change the base?
Conversation
!test --diff |
Review updated until commit 37ac5d4 Description
Changes walkthrough 📝
PR Reviewer Guide 🔍Here are some key observations to aid the review process:
|
!test --diff |
1 similar comment
!test --diff |
fc99e68
to
e0e55e9
Compare
!test --diff |
// Make sure there's no self mapping in the Exact graph as that | ||
// would invalidate lowering assumptions. | ||
if (!allow_self_mapping_) { | ||
assertNoSelfMapping(graph); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Moved this check to buildExactGraph
// had self mapping (see issue #3919). ALMOSTEXACT is used in | ||
// indexing, which assumes the graph has no self mapping. | ||
if (!allow_self_mapping_) { | ||
assertNoSelfMapping(almost_exact_graph); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Added this check for AlmostExact
@@ -116,7 +116,7 @@ class IdModel : public PolymorphicBase { | |||
const std::vector<Expr*>& exprs, | |||
const std::vector<TensorView*>& additional_tvs = {}, | |||
bool build_graphs = false, | |||
bool allow_self_mapping = false, | |||
bool allow_self_mapping = true, |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Changed to not check self mapping by default since that restriction only matters for indexing.
!test --diff |
!test --diff |
!test --dev |
Fixes #3919
The root cause of #3919 is there's a tensor that has a self mapping loop domain. Specifically:
In this case, because the tv0 merge just uses broadcast IDs, all of the two input IDs and the output ID are mapped in the AlmostExact graph (see here). That is also the case with ComputeAtMap. That means
b1
andb2
are mapped, which in turn means the two logical IDs of tv1 are also mapped since they use the same resize op.While the mapping of
b1
andb2
may not have any actual effect as they are just broadcast IDs, the mapping of the logical IDs of tv1 is problematic since they are concrete IDs. In the case of #3919, that manifests as the error in the contiguity analysis since it assumes no self mapping.The fix in this PR is simply avoiding mapping
b1
andb2
. More specifically, I extendedValGraph
to exclude certain Vals from mapping. In the case of the AlmostExact graph, even when two IDs are trivially mapped, we no longer map them if they can result in self mapping. It may not sound ideal as they do have the same extent, but until we address the self mapping limitation, I think this is a reasonable workaround.Note: I am not 100% confident if this might have any negative side effect. That is, the AlmostExact graph now lacks some mappings that do have the same extent. However, self mapping would certainly break TensorIndexer, so I think this is still reasonable.