[mlir][affine] Fix unfolded bounding maps for affine.for

Loop bounds of affine.for didn't perform foldings like affine.load, affine.store.
Bound maps shall be more composed, leaving most affine.apply become dead.

This resolves the bug listed on https://bugs.llvm.org/show_bug.cgi?id=45203

Differential Revision: https://reviews.llvm.org/D99323

diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 2b12b3a..3bd51bc 100644 (file)
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -1578,9 +1578,11 @@ static LogicalResult canonicalizeLoopBounds(AffineForOp forOp) {
   auto prevLbMap = lbMap;
   auto prevUbMap = ubMap;
 
+  composeAffineMapAndOperands(&lbMap, &lbOperands);
   canonicalizeMapAndOperands(&lbMap, &lbOperands);
   lbMap = removeDuplicateExprs(lbMap);
 
+  composeAffineMapAndOperands(&ubMap, &ubOperands);
   canonicalizeMapAndOperands(&ubMap, &ubOperands);
   ubMap = removeDuplicateExprs(ubMap);
 

diff --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir
index 8ede55a..d766dd9 100644 (file)
--- a/mlir/test/Dialect/Affine/canonicalize.mlir
+++ b/mlir/test/Dialect/Affine/canonicalize.mlir
@@ -884,3 +884,23 @@ func @dont_merge_affine_max_if_not_single_sym(%i0: index, %i1: index, %i2: index
   %1 = affine.max affine_map<()[s0, s1] -> (s0 + 4, 7 + s1)> ()[%0, %i2]
   return %1: index
 }
+
+// -----
+
+// Ensure bounding maps of affine.for are composed.
+
+// CHECK-DAG: #[[$MAP0]] = affine_map<()[s0] -> (s0 - 2)>
+// CHECK-DAG: #[[$MAP1]] = affine_map<()[s0] -> (s0 + 2)>
+
+// CHECK-LABEL: func @compose_affine_for_bounds
+// CHECK-SAME:   %[[N:.*]]: index)
+// CHECK: affine.for %{{.*}} = #[[$MAP0]]()[%[[N]]] to #[[$MAP1]]()[%[[N]]] {
+
+func @compose_affine_for_bounds(%N: index) {
+  %u = affine.apply affine_map<(d0) -> (d0 + 2)>(%N)
+  %l = affine.apply affine_map<(d0) -> (d0 - 2)>(%N)
+  affine.for %i = %l to %u {
+    "foo"() : () -> ()
+  }
+  return
+}