[MLIR][Affine] Fix getAccessRelation for 0-d memrefs

[bondhugula]

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

Fixes: #132163

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-affine

Author: Uday Bondhugula (bondhugula)

Changes

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

Fixes: #132163

---

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

2 Files Affected:

• (modified) mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp (+9-1)
• (modified) mlir/test/Dialect/Affine/scalrep.mlir (+14)

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
+}

[arnab-polymage]

LGTM. Thanks.