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

[llvmlistbot at llvm.org]

https://github.com/Baxi-codes created https://github.com/llvm/llvm-project/pull/152812

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()`.

>From f6e92e9d09112b1a86131fe89244c97aa39bc366 Mon Sep 17 00:00:00 2001
From: Baxi-codes <baxidhairya2312 at gmail.com>
Date: Sat, 9 Aug 2025 04:40:28 +0530
Subject: [PATCH] fix: return failure in getAccessRelation if rel ids are not
 in domain

Signed-off-by: Baxi-codes <baxidhairya2312 at gmail.com>
---
 mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

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