[mlir] simplify affine maps and operands in affine.min/max

[ftynse]

Summary

Affine dialect already has a map+operand simplification infrastructure in
place. Plug the recently added affine.min/max operations into this
infrastructure and add a simple test. More complex behavior of the simplifier
is already tested by other ops.

Addresses https://bugs.llvm.org/show_bug.cgi?id=45008.

[andydavis1]

Thanks Alex!

[rriddle]

mlir/include/mlir/Dialect/AffineOps/AffineOps.td

246

Weird spacing? builder-> getIndexType()

mlir/include/mlir/Dialect/AffineOps/AffineOps.td

def AffineIfOp : Affine_Op<"if", [ImplicitAffineTerminator]> {
 
   let hasFolder = 1;
 }
 
 class AffineMinMaxOpBase<string mnemonic, list<OpTrait> traits = []> :
     Op<Affine_Dialect, mnemonic, traits> {
   let arguments = (ins AffineMapAttr:$map, Variadic<Index>:$operands);
   let results = (outs Index);
+
+  let builders = [
+    OpBuilder<"Builder *builder, OperationState &result, AffineMap affineMap, "
+              "ValueRange mapOperands",
+    [{
+      build(builder, result, builder->getIndexType(), affineMap, mapOperands);
+    }]>
+  ];
+
   let extraClassDeclaration = [{
     static StringRef getMapAttrName() { return "map"; }
+    AffineMap getAffineMap() { return map(); }
+    ValueRange getMapOperands() { return operands(); }
   }];
   let verifier = [{ return ::verifyAffineMinMaxOp(*this); }];
   let printer = [{ return ::printAffineMinMaxOp(p, *this); }];
   let parser = [{ return ::parseAffineMinMaxOp<$cppClass>(parser, result); }];
   let hasFolder = 1;
+  let hasCanonicalizer = 1;
 }
 
 def AffineMinOp : AffineMinMaxOpBase<"min", [NoSideEffect]> {
   let summary = "min operation";
   let description = [{
     The "min" operation computes the minimum value result from a multi-result
     affine map.
 

mlir/lib/Dialect/AffineOps/AffineOps.cpp

struct SimplifyAffineOp : public OpRewritePattern<AffineOpTy> {
   void replaceAffineOp(PatternRewriter &rewriter, AffineOpTy affineOp,
                        AffineMap map, ArrayRef<Value> mapOperands) const;
 
   PatternMatchResult matchAndRewrite(AffineOpTy affineOp,
                                      PatternRewriter &rewriter) const override {
     static_assert(std::is_same<AffineOpTy, AffineLoadOp>::value ||
                       std::is_same<AffineOpTy, AffinePrefetchOp>::value ||
                       std::is_same<AffineOpTy, AffineStoreOp>::value ||
-                      std::is_same<AffineOpTy, AffineApplyOp>::value,
+                      std::is_same<AffineOpTy, AffineApplyOp>::value ||
+                      std::is_same<AffineOpTy, AffineMinOp>::value ||
+                      std::is_same<AffineOpTy, AffineMaxOp>::value,
                   "affine load/store/apply op expected");
     auto map = affineOp.getAffineMap();
     AffineMap oldMap = map;
     auto oldOperands = affineOp.getMapOperands();
     SmallVector<Value, 8> resultOperands(oldOperands);
     composeAffineMapAndOperands(&map, &resultOperands);
     if (map == oldMap && std::equal(oldOperands.begin(), oldOperands.end(),
                                     resultOperands.begin()))

 }
 template <>
 void SimplifyAffineOp<AffineStoreOp>::replaceAffineOp(
     PatternRewriter &rewriter, AffineStoreOp store, AffineMap map,
     ArrayRef<Value> mapOperands) const {
   rewriter.replaceOpWithNewOp<AffineStoreOp>(
       store, store.getValueToStore(), store.getMemRef(), map, mapOperands);
 }
-template <>
-void SimplifyAffineOp<AffineApplyOp>::replaceAffineOp(
-    PatternRewriter &rewriter, AffineApplyOp apply, AffineMap map,
+
+// Generic version for ops that don't have extra operands.
+template <typename AffineOpTy>
+void SimplifyAffineOp<AffineOpTy>::replaceAffineOp(
+    PatternRewriter &rewriter, AffineOpTy op, AffineMap map,
     ArrayRef<Value> mapOperands) const {
-  rewriter.replaceOpWithNewOp<AffineApplyOp>(apply, map, mapOperands);
+  rewriter.replaceOpWithNewOp<AffineOpTy>(op, map, mapOperands);
 }
 } // end anonymous namespace.
 
 void AffineApplyOp::getCanonicalizationPatterns(
     OwningRewritePatternList &results, MLIRContext *context) {
   results.insert<SimplifyAffineOp<AffineApplyOp>>(context);
 }
 
if (intAttr.getInt() < min) {
       minIndex = i;
     }
   }
   if (minIndex < 0)
     return {};
   return results[minIndex];
 }
 
+void AffineMinOp::getCanonicalizationPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *context) {
+  patterns.insert<SimplifyAffineOp<AffineMinOp>>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // AffineMaxOp
 //===----------------------------------------------------------------------===//
 //
 //   %0 = affine.max (d0) -> (1000, d0 + 512) (%i0)
 //
 
 OpFoldResult AffineMaxOp::fold(ArrayRef<Attribute> operands) {
if (intAttr.getInt() > max) {
       maxIndex = i;
     }
   }
   if (maxIndex < 0)
     return {};
   return results[maxIndex];
 }
 
+void AffineMaxOp::getCanonicalizationPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *context) {
+  patterns.insert<SimplifyAffineOp<AffineMaxOp>>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // AffinePrefetchOp
 //===----------------------------------------------------------------------===//
 
 //
 // affine.prefetch %0[%i, %j + 5], read, locality<3>, data : memref<400x400xi32>
 //
 static ParseResult parseAffinePrefetchOp(OpAsmParser &parser,

mlir/test/AffineOps/canonicalize.mlir

func @affine_max(%arg0 : index, %arg1 : index, %arg2 : index) {
   %c3 = constant 3 : index
   %c20 = constant 20 : index
   %0 = affine.max affine_map<(d0)[s0] -> (1000, d0 floordiv 4, (s0 mod 5) + 1)> (%c20)[%c3]
   "op0"(%0) : (index) -> ()
   // CHECK:       %[[CST:.*]] = constant 1000 : index
   // CHECK-NEXT:  "op0"(%[[CST]]) : (index) -> ()
   // CHECK-NEXT:  return
   return
-}
+}
+
+// -----
+
+// CHECK: #[[map:.*]] = affine_map<(d0, d1) -> (d0, d1 - 2)>
+
+func @affine_min(%arg0: index) {
+  affine.for %i = 0 to %arg0 {
+    affine.for %j = 0 to %arg0 {
+      %c2 = constant 2 : index
+      // CHECK: affine.min #[[map]]
+      %0 = affine.min affine_map<(d0,d1,d2)->(d0, d1 - d2)>(%i, %j, %c2)
+      "consumer"(%0) : (index) -> ()
+    }
+  }
+  return
+}