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

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

Fixes: llvm#132163

Author: bondhugula

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

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