[mlir] add lowering from affine.min to std

[ftynse]

Summary

Affine minimum computation will be used in tiling transformation. The
implementation is mostly boilerplate as we already lower the minimum in the
upper bound of an affine loop.

[merge_guards_bot]

Unit tests: pass. 62221 tests passed, 0 failed and 816 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[rriddle]

mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp

274

These should be ///

279

Can we just change expandAffineMap to use ValueRange instead?

[ftynse]

mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp

274

The rest of the file uses //, so I went for consistency here. We can batch-update.

279

It affects more than just that function (I tried and stopped after third jump) so it's better of as a separate NFC change. WIll do.

mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp

Value mlir::lowerAffineLowerBound(AffineForOp op, OpBuilder &builder) {
   auto lbValues = expandAffineMap(builder, op.getLoc(), op.getLowerBoundMap(),
                                   boundOperands);
   if (!lbValues)
     return nullptr;
   return buildMinMaxReductionSeq(op.getLoc(), CmpIPredicate::sgt, *lbValues,
                                  builder);
 }
 
+// Emit instructions that correspond to computing the minimum value amoung the
+// values of a (potentially) multi-output affine map applied to `operands`.
+static Value lowerAffineMapMin(OpBuilder &builder, Location loc, AffineMap map,
+                               ValueRange operands) {
+  if (auto values =
+          expandAffineMap(builder, loc, map, llvm::to_vector<4>(operands)))
+    return buildMinMaxReductionSeq(loc, CmpIPredicate::slt, *values, builder);
+  return nullptr;
+}
+
 // Emit instructions that correspond to the affine map in the upper bound
 // applied to the respective operands, and compute the minimum value across
 // the results.
 Value mlir::lowerAffineUpperBound(AffineForOp op, OpBuilder &builder) {
-  SmallVector<Value, 8> boundOperands(op.getUpperBoundOperands());
-  auto ubValues = expandAffineMap(builder, op.getLoc(), op.getUpperBoundMap(),
-                                  boundOperands);
-  if (!ubValues)
-    return nullptr;
-  return buildMinMaxReductionSeq(op.getLoc(), CmpIPredicate::slt, *ubValues,
-                                 builder);
+  return lowerAffineMapMin(builder, op.getLoc(), op.getUpperBoundMap(),
+                           op.getUpperBoundOperands());
 }
 
 namespace {
+class AffineMinLowering : public OpRewritePattern<AffineMinOp> {
+public:
+  using OpRewritePattern<AffineMinOp>::OpRewritePattern;
+
+  PatternMatchResult matchAndRewrite(AffineMinOp op,
+                                     PatternRewriter &rewriter) const override {
+    Value reduced =
+        lowerAffineMapMin(rewriter, op.getLoc(), op.map(), op.operands());
+    if (!reduced)
+      return matchFailure();
+
+    rewriter.replaceOp(op, reduced);
+    return matchSuccess();
+  }
+};
+
 // Affine terminators are removed.
 class AffineTerminatorLowering : public OpRewritePattern<AffineTerminatorOp> {
 public:
   using OpRewritePattern<AffineTerminatorOp>::OpRewritePattern;
 
   PatternMatchResult matchAndRewrite(AffineTerminatorOp op,
                                      PatternRewriter &rewriter) const override {
     rewriter.replaceOpWithNewOp<loop::TerminatorOp>(op);
PatternMatchResult matchAndRewrite(AffineDmaWaitOp op,
     return matchSuccess();
   }
 };
 
 } // end namespace
 
 void mlir::populateAffineToStdConversionPatterns(
     OwningRewritePatternList &patterns, MLIRContext *ctx) {
+  // clang-format off
   patterns.insert<
-      AffineApplyLowering, AffineDmaStartLowering, AffineDmaWaitLowering,
-      AffineLoadLowering, AffinePrefetchLowering, AffineStoreLowering,
-      AffineForLowering, AffineIfLowering, AffineTerminatorLowering>(ctx);
+      AffineApplyLowering,
+      AffineDmaStartLowering,
+      AffineDmaWaitLowering,
+      AffineLoadLowering,
+      AffineMinLowering,
+      AffinePrefetchLowering,
+      AffineStoreLowering,
+      AffineForLowering,
+      AffineIfLowering,
+      AffineTerminatorLowering>(ctx);
+  // clang-format on
 }
 
 namespace {
 class LowerAffinePass : public FunctionPass<LowerAffinePass> {
   void runOnFunction() override {
     OwningRewritePatternList patterns;
     populateAffineToStdConversionPatterns(patterns, &getContext());
     ConversionTarget target(getContext());

mlir/test/Transforms/lower-affine.mlir

affine.for %i0 = 0 to 10 {
     affine.dma_wait %2[%i0 + %arg0 + 17], %c64 : memref<1xi32>
   }
 // CHECK:       %[[a:.*]] = addi %{{.*}}, %arg0 : index
 // CHECK-NEXT:  %c17 = constant 17 : index
 // CHECK-NEXT:  %[[b:.*]] = addi %[[a]], %c17 : index
 // CHECK-NEXT:  dma_wait %0[%[[b]]], %c64 : memref<1xi32>
   return
 }
+
+// CHECK-LABEL: func @affine_min
+// CHECK-SAME: %[[ARG0:.*]]: index, %[[ARG1:.*]]: index
+func @affine_min(%arg0: index, %arg1: index) -> index{
+  // CHECK: %[[Cm1:.*]] = constant -1
+  // CHECK: %[[neg1:.*]] = muli %[[ARG1]], %[[Cm1:.*]]
+  // CHECK: %[[first:.*]] = addi %[[ARG0]], %[[neg1]]
+  // CHECK: %[[Cm2:.*]] = constant -1
+  // CHECK: %[[neg2:.*]] = muli %[[ARG0]], %[[Cm2:.*]]
+  // CHECK: %[[second:.*]] = addi %[[ARG1]], %[[neg2]]
+  // CHECK: %[[cmp:.*]] = cmpi "slt", %[[first]], %[[second]]
+  // CHECK: select %[[cmp]], %[[first]], %[[second]]
+  %0 = affine.min affine_map<(d0,d1) -> (d0 - d1, d1 - d0)>(%arg0, %arg1)
+  return %0 : index
+}