[MLIR][Affine] Remove restriction in slice validity check on symbols

[bondhugula]

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: #61784

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

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-affine

Author: Uday Bondhugula (bondhugula)

Changes

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: #61784

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

---

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

2 Files Affected:

• (modified) mlir/lib/Dialect/Affine/Analysis/Utils.cpp (+2-6)
• (modified) mlir/test/Dialect/Affine/loop-fusion-4.mlir (+59)

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

[bondhugula]

This was posted and reviewed earlier here: https://reviews.llvm.org/D148559, but was somehow lost. Affine fusion with symbolic/non-constant bounds has always been possible, but for these now outdated guards.