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

llvmlistbot at llvm.org llvmlistbot at llvm.org
Tue Mar 31 15:49:50 PDT 2026


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

>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 1/4] 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));

>From 26b9d9bd23c1dfbf74b0a8686aafb40d3ccc7033 Mon Sep 17 00:00:00 2001
From: Baxi-codes <baxidhairya2312 at gmail.com>
Date: Wed, 13 Aug 2025 15:52:55 +0530
Subject: [PATCH 2/4] Replace std::find with llvm::is_contained to check
 membership

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

diff --git a/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp b/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
index d5d8143a55ec3..74295e8b90e79 100644
--- a/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
@@ -21,6 +21,7 @@
 #include "mlir/Dialect/Affine/IR/AffineValueMap.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
@@ -475,7 +476,7 @@ LogicalResult MemRefAccess::getAccessRelation(IntegerRelation &rel) const {
       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()) {
+    if (llvm::is_contained(domainIds, idi)) {
       return failure();
     }
   }

>From 31494e5b20ffca98dc096e5084ed6a51f5587769 Mon Sep 17 00:00:00 2001
From: Baxi-codes <baxidhairya2312 at gmail.com>
Date: Wed, 25 Feb 2026 05:43:35 +0530
Subject: [PATCH 3/4] Add test and minor fix

Signed-off-by: Baxi-codes <baxidhairya2312 at gmail.com>
---
 .../Dialect/Affine/Analysis/AffineAnalysis.cpp  |  7 +++++--
 mlir/test/Dialect/Affine/loop-fusion-4.mlir     | 17 +++++++++++++++++
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp b/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
index 74295e8b90e79..64d77b985faf7 100644
--- a/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
@@ -471,12 +471,15 @@ LogicalResult MemRefAccess::getAccessRelation(IntegerRelation &rel) const {
   // of the access map is a subset of the domain of access, the domain ids of
   // `rel` should be a subset of ids of `domain`. If not, return failure.
 
-  for (unsigned i = 0, e = rel.getNumDimVars(); i < e; ++i) {
+  for (unsigned i = 0; i <rel.getNumDimVars(); ++i) {
     if (rel.getVarKindAt(i) != VarKind::SetDim)
       continue;
     Identifier idi = rel.getIds(VarKind::SetDim)[i];
+    if (!idi.hasValue()) {
+      continue;
+    }
     ArrayRef<Identifier> domainIds = domain.getIds(VarKind::SetDim);
-    if (llvm::is_contained(domainIds, idi)) {
+    if (!llvm::is_contained(domainIds, idi)) {
       return failure();
     }
   }
diff --git a/mlir/test/Dialect/Affine/loop-fusion-4.mlir b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
index 04c8c3ee809a1..842bfc5b38fb1 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-4.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
@@ -771,3 +771,20 @@ func.func @memref_cast_reused(%arg: memref<*xf32>) {
   // SIBLING-MAXIMAL-NEXT:   affine.store
   return
 }
+
+
+// Exercises fix for crash reported at https://github.com/llvm/llvm-project/issues/54541
+
+func.func @load_linearized_index() {
+  %0 = memref.alloc() : memref<128xi8>
+  %2 = arith.constant 1 : i8
+  affine.for %i0 = 0 to 128 {
+    %3 = affine.linearize_index [%i0] by (1) : index
+    affine.store %2, %0[%3] : memref<128xi8>
+  }
+  affine.for %i0 = 0 to 128 {
+    %3 = affine.linearize_index [%i0] by (1) : index
+    %4 = affine.load %0[%3] : memref<128xi8>
+  }
+  return
+}
\ No newline at end of file

>From 063cdf994161b6d1e422c03d11352626b92d7708 Mon Sep 17 00:00:00 2001
From: Baxi-codes <46259792+Baxi-codes at users.noreply.github.com>
Date: Wed, 1 Apr 2026 04:03:17 +0530
Subject: [PATCH 4/4] Remove merge markers

---
 mlir/test/Dialect/Affine/loop-fusion-4.mlir | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/mlir/test/Dialect/Affine/loop-fusion-4.mlir b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
index 52abdf9414223..dc6ec8c9e0310 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-4.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
@@ -772,8 +772,6 @@ func.func @memref_cast_reused(%arg: memref<*xf32>) {
   return
 }
 
-<<<<<<< affine-access-relation-failure-fix
-
 // Exercises fix for crash reported at https://github.com/llvm/llvm-project/issues/54541
 
 func.func @load_linearized_index() {
@@ -789,8 +787,6 @@ func.func @load_linearized_index() {
   }
   return
 }
-=======
-// -----
 
 // Test with symbolic loop bounds.
 
@@ -902,4 +898,3 @@ func.func @high_trip_count(%arg0: memref<1024x4096xf32>, %arg1: memref<8192x4096
   }
   return %alloc : memref<1024x8192xf32>
 }
->>>>>>> main



More information about the Mlir-commits mailing list