[Mlir-commits] [mlir] [mlir][affine] Fix crash in affine-loop-fusion by validating access relation IDs (PR #152812)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: None (Baxi-codes)

<details>
<summary>Changes</summary>

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
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));

``````````

</details>


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