[Mlir-commits] [mlir] 2170d77 - [MLIR][Affine] Fix getAccessRelation for 0-d memrefs (#132347)

[llvmlistbot at llvm.org]

Author: Uday Bondhugula
Date: 2025-03-21T14:35:20+05:30
New Revision: 2170d77e5dfd0a096ab16ecb3c3d82b86dc83a45

URL: https://github.com/llvm/llvm-project/commit/2170d77e5dfd0a096ab16ecb3c3d82b86dc83a45
DIFF: https://github.com/llvm/llvm-project/commit/2170d77e5dfd0a096ab16ecb3c3d82b86dc83a45.diff

LOG: [MLIR][Affine] Fix getAccessRelation for 0-d memrefs (#132347)

Fix getAccessRelation for 0-d memref accesses in certain cases. Fixes
corner case crashes when using scalrep, affine dep analysis, etc.

Fixes: https://github.com/llvm/llvm-project/issues/132163

Added: 
    

Modified: 
    mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
    mlir/test/Dialect/Affine/scalrep.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp b/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
index 84b76d33c3e67..6f79665c2bb60 100644
--- a/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
@@ -489,8 +489,11 @@ LogicalResult MemRefAccess::getAccessRelation(IntegerRelation &rel) const {
 
   // Append domain constraints to `rel`.
   IntegerRelation domainRel = domain;
-  if (rel.getSpace().isUsingIds() && !domainRel.getSpace().isUsingIds())
+  // For 0-d spaces, there will be no IDs. Enable if that's the case.
+  if (!domainRel.getSpace().isUsingIds())
     domainRel.resetIds();
+  if (!rel.getSpace().isUsingIds())
+    rel.resetIds();
   domainRel.appendVar(VarKind::Range, accessValueMap.getNumResults());
   domainRel.mergeAndAlignSymbols(rel);
   domainRel.mergeLocalVars(rel);
@@ -660,6 +663,11 @@ DependenceResult mlir::affine::checkMemrefAccessDependence(
   // `srcAccess` to the iteration domain of `dstAccess` which access the same
   // memory locations.
   dstRel.inverse();
+  // For 0-d spaces, there will be no IDs. Enable if that's the case.
+  if (!dstRel.getSpace().isUsingIds())
+    dstRel.resetIds();
+  if (!srcRel.getSpace().isUsingIds())
+    srcRel.resetIds();
   dstRel.mergeAndCompose(srcRel);
   dstRel.convertVarKind(VarKind::Domain, 0, dstRel.getNumDomainVars(),
                         VarKind::Range, 0);

diff  --git a/mlir/test/Dialect/Affine/scalrep.mlir b/mlir/test/Dialect/Affine/scalrep.mlir
index 092597860c8d9..901e15911f4ea 100644
--- a/mlir/test/Dialect/Affine/scalrep.mlir
+++ b/mlir/test/Dialect/Affine/scalrep.mlir
@@ -983,3 +983,17 @@ func.func @scf_for_if(%arg0: memref<?xi32>, %arg1: i32) -> i32 attributes {llvm.
   %3 = affine.load %0[0] : memref<1xi32>
   return %3 : i32
 }
+
+// CHECK-LABEL: func @zero_d_memrefs
+func.func @zero_d_memrefs() {
+  // CHECK: %[[C0:.*]] = arith.constant 0
+  %c0_i32 = arith.constant 0 : i32
+  %alloc_0 = memref.alloc() {alignment = 64 : i64} : memref<i32>
+  affine.store %c0_i32, %alloc_0[] : memref<i32>
+  affine.for %arg0 = 0 to 9 {
+    %2 = affine.load %alloc_0[] : memref<i32>
+    arith.addi %2, %2 : i32
+    // CHECK: arith.addi %[[C0]], %[[C0]]
+  }
+  return
+}