[Mlir-commits] [mlir] 9cc417c - [mlir][affine] Fix unfolded bounding maps for affine.for

[Uday Bondhugula]

Author: eopXD
Date: 2021-04-13T00:12:43+05:30
New Revision: 9cc417cbca1cece0d55fa3d1e15682943a06139e

URL: https://github.com/llvm/llvm-project/commit/9cc417cbca1cece0d55fa3d1e15682943a06139e
DIFF: https://github.com/llvm/llvm-project/commit/9cc417cbca1cece0d55fa3d1e15682943a06139e.diff

LOG: [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

Added: 
    

Modified: 
    mlir/lib/Dialect/Affine/IR/AffineOps.cpp
    mlir/test/Dialect/Affine/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 2b12b3ad16167..3bd51bc5673b7 100644
--- 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 8ede55a9d284c..d766dd9956f8a 100644
--- 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
+}