Add pattern to canonicalize for loop bounds

- add pattern to canonicalize affine.for loop bounds (using
  canonicalizeMapAndOperands)
- rename AffineForLoopBoundFolder -> AffineForLoopBoundFolder for
  consistency

Signed-off-by: Uday Bondhugula <uday@polymagelabs.com>

Closes #111

COPYBARA_INTEGRATE_REVIEW=tensorflow/mlir#111 from bondhugula:bound-canonicalize ee8fb7f43a7ffd45f6df3f53c95098d8b7e494c7
PiperOrigin-RevId: 269041220

third_party/mlir/lib/Dialect/AffineOps/AffineOps.cpp

@@ -1358,7 +1358,7 @@ struct AffineForEmptyLoopFolder : public OpRewritePattern<AffineForOp> {
 };
 
 /// This is a pattern to fold constant loop bounds.
-struct AffineForLoopBoundFolder : public OpRewritePattern<AffineForOp> {
+struct AffineForOpBoundFolder : public OpRewritePattern<AffineForOp> {
   using OpRewritePattern<AffineForOp>::OpRewritePattern;
 
   PatternMatchResult matchAndRewrite(AffineForOp forOp,
@@ -1413,11 +1413,44 @@ struct AffineForLoopBoundFolder : public OpRewritePattern<AffineForOp> {
     return matchSuccess();
   }
 };
+
+// This is a pattern to canonicalize affine for op loop bounds.
+struct AffineForOpBoundCanonicalizer : public OpRewritePattern<AffineForOp> {
+  using OpRewritePattern<AffineForOp>::OpRewritePattern;
+
+  PatternMatchResult matchAndRewrite(AffineForOp forOp,
+                                     PatternRewriter &rewriter) const override {
+    SmallVector<Value *, 4> lbOperands(forOp.getLowerBoundOperands());
+    SmallVector<Value *, 4> ubOperands(forOp.getUpperBoundOperands());
+
+    auto lbMap = forOp.getLowerBoundMap();
+    auto ubMap = forOp.getUpperBoundMap();
+    auto prevLbMap = lbMap;
+    auto prevUbMap = ubMap;
+
+    canonicalizeMapAndOperands(&lbMap, &lbOperands);
+    canonicalizeMapAndOperands(&ubMap, &ubOperands);
+
+    // Any canonicalization change always leads to updated map(s).
+    if (lbMap == prevLbMap && ubMap == prevUbMap)
+      return matchFailure();
+
+    if (lbMap != prevLbMap)
+      forOp.setLowerBound(lbOperands, lbMap);
+    if (ubMap != prevUbMap)
+      forOp.setUpperBound(ubOperands, ubMap);
+
+    rewriter.updatedRootInPlace(forOp);
+    return matchSuccess();
+  }
+};
+
 } // end anonymous namespace
 
 void AffineForOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
                                               MLIRContext *context) {
-  results.insert<AffineForEmptyLoopFolder, AffineForLoopBoundFolder>(context);
+  results.insert<AffineForEmptyLoopFolder, AffineForOpBoundFolder,
+                 AffineForOpBoundCanonicalizer>(context);
 }
 
 AffineBound AffineForOp::getLowerBound() {