From 2170d77e5dfd0a096ab16ecb3c3d82b86dc83a45 Mon Sep 17 00:00:00 2001
From: Uday Bondhugula <uday@polymagelabs.com>
Date: Fri, 21 Mar 2025 14:35:20 +0530
Subject: [PATCH] [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
---
 .../lib/Dialect/Affine/Analysis/AffineAnalysis.cpp | 10 +++++++++-
 mlir/test/Dialect/Affine/scalrep.mlir              | 14 ++++++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp b/mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp
index 84b76d33c3e6..6f79665c2bb6 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 092597860c8d..901e15911f4e 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
+}