From 9cc417cbca1cece0d55fa3d1e15682943a06139e Mon Sep 17 00:00:00 2001
From: eopXD <eopxd@skymizer.com>
Date: Mon, 12 Apr 2021 17:52:58 +0530
Subject: [PATCH] [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
---
 mlir/lib/Dialect/Affine/IR/AffineOps.cpp   |  2 ++
 mlir/test/Dialect/Affine/canonicalize.mlir | 20 ++++++++++++++++++++
 2 files changed, 22 insertions(+)

diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 2b12b3a..3bd51bc 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 8ede55a..d766dd9 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
+}
-- 
2.7.4