[mlir][affine] Fix crash in affine-loop-fusion by validating access relation IDs

[Baxi-codes]

Fixes #149507

MemRefAccess::getAccessRelation assumed that access relation’s IDs are always a subset of domain relation IDs.
When this invariant was violated (e.g., in the affine-loop-fusion pass), it led to crashes.

This change adds an explicit check: if any ID is missing from the domain, the function returns failure().

[github-actions]

Thank you for submitting a Pull Request (PR) to the LLVM Project!

This PR will be automatically labeled and the relevant teams will be notified.

If you wish to, you can add reviewers by using the "Reviewers" section on this page.

If this is not working for you, it is probably because you do not have write permissions for the repository. In which case you can instead tag reviewers by name in a comment by using @ followed by their GitHub username.

If you have received no comments on your PR for a week, you can request a review by "ping"ing the PR by adding a comment “Ping”. The common courtesy "ping" rate is once a week. Please remember that you are asking for valuable time from other developers.

If you have further questions, they may be answered by the LLVM GitHub User Guide.

You can also ask questions in a comment on this PR, on the LLVM Discord or on the forums.

[llvmbot]

@llvm/pr-subscribers-mlir-affine

@llvm/pr-subscribers-mlir

Author: None (Baxi-codes)

Changes

Fixes #149507

MemRefAccess::getAccessRelation assumed that access relation’s IDs are always a subset of domain relation IDs.
When this invariant was violated (e.g., in the affine-loop-fusion pass), it led to crashes.

This change adds an explicit check: if any ID is missing from the domain, the function returns failure().

---

Full diff: https://github.com/llvm/llvm-project/pull/152812.diff

1 Files Affected:

• (modified) mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp (+12-1)

diff --git a/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp b/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
index 4d2d8738aa4ad..d5d8143a55ec3 100644
--- a/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
@@ -468,7 +468,18 @@ LogicalResult MemRefAccess::getAccessRelation(IntegerRelation &rel) const {
 
   // Merge and align domain ids of `rel` with ids of `domain`. Since the domain
   // of the access map is a subset of the domain of access, the domain ids of
-  // `rel` are guranteed to be a subset of ids of `domain`.
+  // `rel` should be a subset of ids of `domain`. If not, return failure.
+
+  for (unsigned i = 0, e = rel.getNumDimVars(); i < e; ++i) {
+    if (rel.getVarKindAt(i) != VarKind::SetDim)
+      continue;
+    Identifier idi = rel.getIds(VarKind::SetDim)[i];
+    ArrayRef<Identifier> domainIds = domain.getIds(VarKind::SetDim);
+    if (std::find(domainIds.begin(), domainIds.end(), idi) == domainIds.end()) {
+      return failure();
+    }
+  }
+
   unsigned inserts = 0;
   for (unsigned i = 0, e = domain.getNumDimVars(); i < e; ++i) {
     const Identifier domainIdi = Identifier(domain.getValue(i));

[Baxi-codes]

Ping

[bondhugula]

Can you add the test case (from the linked issue)? You can trim it to minimal. Thanks.