[Mlir-commits] [mlir] [MLIR][Affine] Remove restriction in slice validity check on symbols (PR #180709)

[Uday Bondhugula]

https://github.com/bondhugula created https://github.com/llvm/llvm-project/pull/180709

Remove restriction in affine analysis utility for checking slice validity. This was unnecessarily bailing out still after the underlying methods were extended. This update enables fusion of affine nests with symbolic bounds.

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

Based on and revived from https://reviews.llvm.org/D148559 from anoopjs.

>From 623d61f335b47d964a53b8dcaf90025e51358725 Mon Sep 17 00:00:00 2001
From: Uday Bondhugula <uday at polymagelabs.com>
Date: Mon, 9 Feb 2026 08:32:56 +0530
Subject: [PATCH] [MLIR][Affine] Remove restriction in slice validity check on
 symbols

Remove restriction in affine analysis utility for checking slice
validity. This was unnecessarily bailing out still after the underlying
methods were extended. This update enables fusion of affine nests with
symbolic bounds.

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

Based on https://reviews.llvm.org/D148559 from anoopjs
---
 mlir/lib/Dialect/Affine/Analysis/Utils.cpp  |  8 +--
 mlir/test/Dialect/Affine/loop-fusion-4.mlir | 59 +++++++++++++++++++++
 2 files changed, 61 insertions(+), 6 deletions(-)

diff --git a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
index 92d8846a4640d..ee893c14c9c6b 100644
--- a/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/Utils.cpp
@@ -1056,12 +1056,6 @@ std::optional<bool> ComputationSliceState::isSliceValid() const {
     LDBG() << "Unable to compute source's domain";
     return std::nullopt;
   }
-  // As the set difference utility currently cannot handle symbols in its
-  // operands, validity of the slice cannot be determined.
-  if (srcConstraints.getNumSymbolVars() > 0) {
-    LDBG() << "Cannot handle symbols in source domain";
-    return std::nullopt;
-  }
   // TODO: Handle local vars in the source domains while using the 'projectOut'
   // utility below. Currently, aligning is not done assuming that there will be
   // no local vars in the source domain.
@@ -1082,6 +1076,8 @@ std::optional<bool> ComputationSliceState::isSliceValid() const {
   // domain completely in terms of source's IVs.
   sliceConstraints.projectOut(ivs.size(),
                               sliceConstraints.getNumVars() - ivs.size());
+  srcConstraints.projectOut(ivs.size(),
+                            srcConstraints.getNumVars() - ivs.size());
 
   LDBG() << "Domain of the source of the slice:\n"
          << "Source constraints:" << srcConstraints
diff --git a/mlir/test/Dialect/Affine/loop-fusion-4.mlir b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
index 04c8c3ee809a1..d6884fb921ad2 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-4.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
@@ -771,3 +771,62 @@ func.func @memref_cast_reused(%arg: memref<*xf32>) {
   // SIBLING-MAXIMAL-NEXT:   affine.store
   return
 }
+
+// -----
+
+// Test with symbolic loop bounds.
+
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @fusion_non_constant_bounds_0
+func.func @fusion_non_constant_bounds_0(%arg0: memref<?xf32>) {
+  %cst = arith.constant 1.000000e+00 : f32
+  %cst_0 = arith.constant 2.000000e+00 : f32
+  %c0 = arith.constant 0 : index
+  %dim = memref.dim %arg0, %c0 : memref<?xf32>
+  affine.for %arg1 = 0 to %dim {
+    %0 = affine.load %arg0[%arg1] : memref<?xf32>
+    %1 = arith.addf %0, %cst : f32
+    affine.store %1, %arg0[%arg1] : memref<?xf32>
+  }
+  affine.for %arg1 = 0 to %dim {
+    %0 = affine.load %arg0[%arg1] : memref<?xf32>
+    %1 = arith.addf %0, %cst_0 : f32
+    affine.store %1, %arg0[%arg1] : memref<?xf32>
+  }
+  // PRODUCER-CONSUMER-MAXIMAL:      affine.for %[[idx:.*]] = 0 to %{{.*}} {
+  // PRODUCER-CONSUMER-MAXIMAL-NEXT:   affine.load %[[arr:.*]][%[[idx]]] : memref<?xf32>
+  // PRODUCER-CONSUMER-MAXIMAL-NEXT:   arith.addf
+  // PRODUCER-CONSUMER-MAXIMAL-NEXT:   affine.store %{{.*}}, %[[arr]][%[[idx]]] : memref<?xf32>
+  // PRODUCER-CONSUMER-MAXIMAL-NEXT:   affine.load %[[arr:.*]][%[[idx]]] : memref<?xf32>
+  // PRODUCER-CONSUMER-MAXIMAL-NEXT:   arith.addf
+  // PRODUCER-CONSUMER-MAXIMAL-NEXT:   affine.store %{{.*}}, %[[arr]][%[[idx]]] : memref<?xf32>
+  // PRODUCER-CONSUMER-MAXIMAL-NEXT: }
+  return
+}
+
+// PRODUCER-CONSUMER-MAXIMAL-LABEL: func @fusion_non_constant_bounds_1
+func.func @fusion_non_constant_bounds_1(%N: index, %M: memref<?xf32>, %cst: f32) {
+  affine.for %i = 0 to %N {
+    affine.store %cst, %M[%i] : memref<?xf32>
+  }
+  affine.for %i = 0 to %N {
+    affine.load %M[%i] : memref<?xf32>
+  }
+  // Should be fused.
+  // PRODUCER-CONSUMER-MAXIMAL: affine.for
+  // PRODUCER-CONSUMER-MAXIMAL-NOT: affine.for
+
+  // Bounds not matching. Still fused; source remains.
+  // PRODUCER-CONSUMER-MAXIMAL:      affine.for
+  // PRODUCER-CONSUMER-MAXIMAL-NEXT:   affine.store
+  affine.for %i = 0 to %N {
+    affine.store %cst, %M[%i] : memref<?xf32>
+  }
+  // PRODUCER-CONSUMER-MAXIMAL:      affine.for
+  // PRODUCER-CONSUMER-MAXIMAL-NEXT:   affine.store
+  // PRODUCER-CONSUMER-MAXIMAL-NEXT:   affine.load
+  affine.for %i = 1 to %N {
+    affine.load %M[%i] : memref<?xf32>
+  }
+
+  return
+}