[mlir] Add min/max operations to Standard.

[pifon2a]

Summary

RFC: Add min/max ops

I was following the naming style for Arith dialect in
https://reviews.llvm.org/D110200,
i.e. similar to DivSIOp and DivUIOp I defined MaxSIOp, MaxUIOp.

When Arith PR is landed, I will migrate these ops as well.

[pifon2a]

Updated documentation

[mehdi_amini]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

1260

Wrong copy/paste here.

1333

ditto

[aartbik]

mlir/docs/Rationale/Rationale.md

346–347

Did you keep this paragraph on purpose. The however, reflects back on a now deleted part?

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

1263

this can be worked out a bit better

1336

same

1361

smaller

1385

samller

[pifon2a]

Addressed the comments. Thank you Mehdi and Aart!

mlir/docs/Rationale/Rationale.md

346–347

Thanks for noticing

[aartbik]

One last typo

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

1336

maximum -> minimum

[pifon2a]

Last typo

[aartbik]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

1264

Just curious, the NaN are actually anything in the range 0x7f800001 and up. Does the comparison deal with all these correctly? Also, perhaps mention Float.NEGATIVE_INFINITY and Float.POSITIVE_INFINITY as other interesting "corner" values.

mlir/test/Dialect/Standard/expand-ops.mlir

150

looks like you forgot the CHECKs here

[aartbik]

mlir/test/Dialect/Standard/expand-ops.mlir

150

also, was the intention to have a maxf_vector as well?

[dcaballe]

Thanks, Alex! LGTM.
Could you please add some tests to Dialect/Standard/ops.mlir and Dialect/Standard/invalid.mlir.

[pifon2a]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

1264

We are not comparing lhs or rhs directly to NaN, we use cmpf 'uno' instead. It should handle these correctly. I don't think that +/- inf are interesting cases for maximum/minimum.

mlir/test/Dialect/Standard/expand-ops.mlir

150

do you mean minf_vector? I don't think it would test anything that was not tested in maxf_vector test.

[pifon2a]

Address the comments.

[pifon2a]

In D110540#3025929, @dcaballe wrote:

Thanks, Alex! LGTM.
Could you please add some tests to Dialect/Standard/ops.mlir and Dialect/Standard/invalid.mlir.

I added the tests to ops.mlir. I am not sure what we would test in invalid.mlir, since the ops are just standard binary ops that were tested multiple times.

[nicolasvasilache]

Thanks @pifon2a !

mlir/docs/Rationale/Rationale.md

 integer equality comparison than `cmpi 0, %lhs, %rhs` where it is unclear what
 gets compared to what else. This syntactic sugar is possible thanks to parser
 logic redefinitions for custom assembly form of non-builtin operations.
 Supporting it in the full notation would have required changing how the main
 parsing algorithm works and may have unexpected repercussions. While it had been
 possible to store the predicate as string attribute, it would have rendered
 impossible to implement switching logic based on the comparison kind and made
 attribute validity checks (one out of ten possible kinds) more complex.
 
-### 'select' operation to implement min/max
-
-Although `min` and `max` operations are likely to occur as a result of
-transforming affine loops in ML functions, we did not make them first-class
-operations. Instead, we provide the `select` operation that can be combined with
-`cmpi` to implement the minimum and maximum computation. Although they now
-require two operations, they are likely to be emitted automatically during the
-transformation inside MLIR. On the other hand, there are multiple benefits of
-introducing `select`: standalone min/max would concern themselves with the
-signedness of the comparison, already taken into account by `cmpi`; `select` can
-support floats transparently if used after a float-comparison operation; the
-lower-level targets provide `select`-like instructions making the translation
-trivial.
-
-This operation could have been implemented with additional control flow: `%r =
-select %cond, %t, %f` is equivalent to
-
-```mlir
-^bb0:
-  cond_br %cond, ^bb1(%t), ^bb1(%f)
-^bb1(%r):
-```
-
-However, this control flow granularity is not available in the ML functions
-where min/max, and thus `select`, are likely to appear. In addition, simpler
-control flow may be beneficial for optimization in general.
-
 ### Regions
 
 #### Attributes of type 'Block'
 
 We considered representing regions through `ArrayAttr`s containing a list of a
 special type `IRBlockAttr`, which in turn would contain a list of operations.
 All attributes in MLIR are unique’d within the context, which would make the IR
 inside the regions immortal for no good reason.
 

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

let description = [{
     integer, the value is truncated.
   }];
 
   let hasFolder = 1;
   let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
+// MaxFOp
+//===----------------------------------------------------------------------===//
+
+def MaxFOp : FloatBinaryOp<"maxf"> {
+  let summary = "floating-point maximum operation";
+  let description = [{
+    Syntax:
+
+    ```
+    operation ::= ssa-id `=` `maxf` ssa-use `,` ssa-use `:` type
+    ```
+
+    Returns the maximum of the two arguments, treating -0.0 as less than +0.0.
+    If one of the arguments is NaN, then the result is also NaN.
+
+    Example:
+
+    ```mlir
+    // Scalar floating-point maximum.
+    %a = maxf %b, %c : f64
+    ```
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// MaxSIOp
+//===----------------------------------------------------------------------===//
+
+def MaxSIOp : IntBinaryOp<"maxsi"> {
+  let summary = "signed integer maximum operation";
+  let description = [{
+    Syntax:
+
+    ```
+    operation ::= ssa-id `=` `maxsi` ssa-use `,` ssa-use `:` type
+    ```
+
+    Returns the larger of %a and %b comparing the values as signed integers.
+
+    Example:
+
+    ```mlir
+    // Scalar signed integer maximum.
+    %a = maxsi %b, %c : i64
+    ```
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// MaxUIOp
+//===----------------------------------------------------------------------===//
+
+def MaxUIOp : IntBinaryOp<"maxui"> {
+  let summary = "unsigned integer maximum operation";
+  let description = [{
+    Syntax:
+
+    ```
+    operation ::= ssa-id `=` `maxui` ssa-use `,` ssa-use `:` type
+    ```
+
+    Returns the larger of %a and %b comparing the values as unsigned integers.
+
+    Example:
+
+    ```mlir
+    // Scalar unsigned integer maximum.
+    %a = maxui %b, %c : i64
+    ```
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// MinFOp
+//===----------------------------------------------------------------------===//
+
+def MinFOp : FloatBinaryOp<"minf"> {
+  let summary = "floating-point minimum operation";
+  let description = [{
+    Syntax:
+
+    ```
+    operation ::= ssa-id `=` `minf` ssa-use `,` ssa-use `:` type
+    ```
+
+    Returns the minimum of the two arguments, treating -0.0 as less than +0.0.
+    If one of the arguments is NaN, then the result is also NaN.
+
+    Example:
+
+    ```mlir
+    // Scalar floating-point minimum.
+    %a = minf %b, %c : f64
+    ```
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// MinSIOp
+//===----------------------------------------------------------------------===//
+
+def MinSIOp : IntBinaryOp<"minsi"> {
+  let summary = "signed integer minimum operation";
+  let description = [{
+    Syntax:
+
+    ```
+    operation ::= ssa-id `=` `minsi` ssa-use `,` ssa-use `:` type
+    ```
+
+    Returns the smaller of %a and %b comparing the values as signed integers.
+
+    Example:
+
+    ```mlir
+    // Scalar signed integer minimum.
+    %a = minsi %b, %c : i64
+    ```
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// MinUIOp
+//===----------------------------------------------------------------------===//
+
+def MinUIOp : IntBinaryOp<"minui"> {
+  let summary = "unsigned integer minimum operation";
+  let description = [{
+    Syntax:
+
+    ```
+    operation ::= ssa-id `=` `minui` ssa-use `,` ssa-use `:` type
+    ```
+
+    Returns the smaller of %a and %b comparing the values as unsigned integers.
+
+    Example:
+
+    ```mlir
+    // Scalar unsigned integer minimum.
+    %a = minui %b, %c : i64
+    ```
+  }];
+}
+
+//===----------------------------------------------------------------------===//
 // MulFOp
 //===----------------------------------------------------------------------===//
 
 def MulFOp : FloatBinaryOp<"mulf"> {
   let summary = "floating point multiplication operation";
   let description = [{
     Syntax:
 

mlir/lib/Dialect/StandardOps/Transforms/ExpandOps.cpp

LogicalResult matchAndRewrite(SignedFloorDivIOp op,
     Value compareRes = rewriter.create<OrOp>(loc, firstTerm, secondTerm);
     Value res = rewriter.create<SelectOp>(loc, compareRes, negRes, posRes);
     // Perform substitution and return success.
     rewriter.replaceOp(op, {res});
     return success();
   }
 };
 
+static Type getElementTypeOrSelf(Type type) {
+  if (auto st = type.dyn_cast<ShapedType>())
+    return st.getElementType();
+  return type;
+}
+
+template <typename OpTy, CmpFPredicate pred>
+struct MaxMinFOpConverter : public OpRewritePattern<OpTy> {
+public:
+  using OpRewritePattern<OpTy>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(OpTy op,
+                                PatternRewriter &rewriter) const final {
+    Value lhs = op.lhs();
+    Value rhs = op.rhs();
+
+    Location loc = op.getLoc();
+    Value cmp = rewriter.create<CmpFOp>(loc, pred, lhs, rhs);
+    Value select = rewriter.create<SelectOp>(loc, cmp, lhs, rhs);
+
+    auto floatType = getElementTypeOrSelf(lhs.getType()).cast<FloatType>();
+    Value isNaN = rewriter.create<CmpFOp>(loc, CmpFPredicate::UNO, lhs, rhs);
+
+    Value nan = rewriter.create<ConstantFloatOp>(
+        loc, APFloat::getQNaN(floatType.getFloatSemantics()), floatType);
+    if (VectorType vectorType = lhs.getType().dyn_cast<VectorType>())
+      nan = rewriter.create<SplatOp>(loc, vectorType, nan);
+
+    rewriter.replaceOpWithNewOp<SelectOp>(op, isNaN, nan, select);
+    return success();
+  }
+};
+
+template <typename OpTy, CmpIPredicate pred>
+struct MaxMinIOpConverter : public OpRewritePattern<OpTy> {
+public:
+  using OpRewritePattern<OpTy>::OpRewritePattern;
+  LogicalResult matchAndRewrite(OpTy op,
+                                PatternRewriter &rewriter) const final {
+    Value lhs = op.lhs();
+    Value rhs = op.rhs();
+
+    Location loc = op.getLoc();
+    Value cmp = rewriter.create<CmpIOp>(loc, pred, lhs, rhs);
+    rewriter.replaceOpWithNewOp<SelectOp>(op, cmp, lhs, rhs);
+    return success();
+  }
+};
+
 struct StdExpandOpsPass : public StdExpandOpsBase<StdExpandOpsPass> {
   void runOnFunction() override {
     MLIRContext &ctx = getContext();
 
     RewritePatternSet patterns(&ctx);
     populateStdExpandOpsPatterns(patterns);
 
     ConversionTarget target(getContext());
 
     target.addLegalDialect<memref::MemRefDialect, StandardOpsDialect>();
     target.addDynamicallyLegalOp<AtomicRMWOp>([](AtomicRMWOp op) {
       return op.kind() != AtomicRMWKind::maxf &&
              op.kind() != AtomicRMWKind::minf;
     });
     target.addDynamicallyLegalOp<memref::ReshapeOp>([](memref::ReshapeOp op) {
       return !op.shape().getType().cast<MemRefType>().hasStaticShape();
     });
-    target.addIllegalOp<SignedCeilDivIOp>();
-    target.addIllegalOp<SignedFloorDivIOp>();
+    // clang-format off
+    target.addIllegalOp<
+      MaxFOp,
+      MaxSIOp,
+      MaxUIOp,
+      MinFOp,
+      MinSIOp,
+      MinUIOp,
+      SignedCeilDivIOp,
+      SignedFloorDivIOp
+    >();
+    // clang-format on
     if (failed(
             applyPartialConversion(getFunction(), target, std::move(patterns))))
       signalPassFailure();
   }
 };
 
 } // namespace
 
 void mlir::populateStdExpandOpsPatterns(RewritePatternSet &patterns) {
-  patterns.add<AtomicRMWOpConverter, MemRefReshapeOpConverter,
-               SignedCeilDivIOpConverter, SignedFloorDivIOpConverter>(
-      patterns.getContext());
+  // clang-format off
+  patterns.add<
+    AtomicRMWOpConverter,
+    MaxMinFOpConverter<MaxFOp, CmpFPredicate::OGT>,
+    MaxMinFOpConverter<MinFOp, CmpFPredicate::OLT>,
+    MaxMinIOpConverter<MaxSIOp, CmpIPredicate::sgt>,
+    MaxMinIOpConverter<MaxUIOp, CmpIPredicate::ugt>,
+    MaxMinIOpConverter<MinSIOp, CmpIPredicate::slt>,
+    MaxMinIOpConverter<MinUIOp, CmpIPredicate::ult>,
+    MemRefReshapeOpConverter,
+    SignedCeilDivIOpConverter,
+    SignedFloorDivIOpConverter
+  >(patterns.getContext());
+  // clang-format on
 }
 
 std::unique_ptr<Pass> mlir::createStdExpandOpsPass() {
   return std::make_unique<StdExpandOpsPass>();
 }

mlir/test/Dialect/Standard/expand-ops.mlir

 // CHECK: [[C0:%.*]] = constant 0 : index
 // CHECK: [[DIM_0:%.*]] = memref.load [[SHAPE]]{{\[}}[[C0]]] : memref<3xi32>
 // CHECK: [[SIZE_0:%.*]] = index_cast [[DIM_0]] : i32 to index
 
 // CHECK: [[RESULT:%.*]] = memref.reinterpret_cast [[SRC]]
 // CHECK-SAME: to offset: [0], sizes: {{\[}}[[SIZE_0]], [[SIZE_1]], 8],
 // CHECK-SAME: strides: {{\[}}[[STRIDE_0]], [[STRIDE_1]], [[C1]]]
 // CHECK-SAME: : memref<*xf32> to memref<?x?x8xf32>
+
+// -----
+
+// CHECK-LABEL: func @maxf
+func @maxf(%a: f32, %b: f32) -> f32 {
+  %result = maxf(%a, %b): (f32, f32) -> f32
+  return %result : f32
+}
+// CHECK-SAME: %[[LHS:.*]]: f32, %[[RHS:.*]]: f32)
+// CHECK-NEXT: %[[CMP:.*]] = cmpf ogt, %[[LHS]], %[[RHS]] : f32
+// CHECK-NEXT: %[[SELECT:.*]] = select %[[CMP]], %[[LHS]], %[[RHS]] : f32
+// CHECK-NEXT: %[[IS_NAN:.*]] = cmpf uno, %[[LHS]], %[[RHS]] : f32
+// CHECK-NEXT: %[[NAN:.*]] = constant 0x7FC00000 : f32
+// CHECK-NEXT: %[[RESULT:.*]] = select %[[IS_NAN]], %[[NAN]], %[[SELECT]] : f32
+// CHECK-NEXT: return %[[RESULT]] : f32
+
+// -----
+
+// CHECK-LABEL: func @maxf_vector
+func @maxf_vector(%a: vector<4xf16>, %b: vector<4xf16>) -> vector<4xf16> {
+  %result = maxf(%a, %b): (vector<4xf16>, vector<4xf16>) -> vector<4xf16>
+  return %result : vector<4xf16>
+}
+// CHECK-SAME: %[[LHS:.*]]: vector<4xf16>, %[[RHS:.*]]: vector<4xf16>)
+// CHECK-NEXT: %[[CMP:.*]] = cmpf ogt, %[[LHS]], %[[RHS]] : vector<4xf16>
+// CHECK-NEXT: %[[SELECT:.*]] = select %[[CMP]], %[[LHS]], %[[RHS]]
+// CHECK-NEXT: %[[IS_NAN:.*]] = cmpf uno, %[[LHS]], %[[RHS]] : vector<4xf16>
+// CHECK-NEXT: %[[NAN:.*]] = constant 0x7E00 : f16
+// CHECK-NEXT: %[[SPLAT_NAN:.*]] = splat %[[NAN]] : vector<4xf16>
+// CHECK-NEXT: %[[RESULT:.*]] = select %[[IS_NAN]], %[[SPLAT_NAN]], %[[SELECT]]
+// CHECK-NEXT: return %[[RESULT]] : vector<4xf16>
+
+// -----
+
+// CHECK-LABEL: func @minf
+func @minf(%a: f32, %b: f32) -> f32 {
+  %result = minf(%a, %b): (f32, f32) -> f32
+  return %result : f32
+}
+// CHECK-SAME: %[[LHS:.*]]: f32, %[[RHS:.*]]: f32)
+// CHECK-NEXT: %[[CMP:.*]] = cmpf olt, %[[LHS]], %[[RHS]] : f32
+// CHECK-NEXT: %[[SELECT:.*]] = select %[[CMP]], %[[LHS]], %[[RHS]] : f32
+// CHECK-NEXT: %[[IS_NAN:.*]] = cmpf uno, %[[LHS]], %[[RHS]] : f32
+// CHECK-NEXT: %[[NAN:.*]] = constant 0x7FC00000 : f32
+// CHECK-NEXT: %[[RESULT:.*]] = select %[[IS_NAN]], %[[NAN]], %[[SELECT]] : f32
+// CHECK-NEXT: return %[[RESULT]] : f32
+
+
+// -----
+
+// CHECK-LABEL: func @maxsi
+func @maxsi(%a: i32, %b: i32) -> i32 {
+  %result = maxsi(%a, %b): (i32, i32) -> i32
+  return %result : i32
+}
+// CHECK-SAME: %[[LHS:.*]]: i32, %[[RHS:.*]]: i32)
+// CHECK-NEXT: %[[CMP:.*]] = cmpi sgt, %[[LHS]], %[[RHS]] : i32
+
+// -----
+
+// CHECK-LABEL: func @minsi
+func @minsi(%a: i32, %b: i32) -> i32 {
+  %result = minsi(%a, %b): (i32, i32) -> i32
+  return %result : i32
+}
+// CHECK-SAME: %[[LHS:.*]]: i32, %[[RHS:.*]]: i32)
+// CHECK-NEXT: %[[CMP:.*]] = cmpi slt, %[[LHS]], %[[RHS]] : i32
+
+
+// -----
+
+// CHECK-LABEL: func @maxui
+func @maxui(%a: i32, %b: i32) -> i32 {
+  %result = maxui(%a, %b): (i32, i32) -> i32
+  return %result : i32
+}
+// CHECK-SAME: %[[LHS:.*]]: i32, %[[RHS:.*]]: i32)
+// CHECK-NEXT: %[[CMP:.*]] = cmpi ugt, %[[LHS]], %[[RHS]] : i32
+
+
+// -----
+
+// CHECK-LABEL: func @minui
+func @minui(%a: i32, %b: i32) -> i32 {
+  %result = minui(%a, %b): (i32, i32) -> i32
+  return %result : i32
+}
+// CHECK-SAME: %[[LHS:.*]]: i32, %[[RHS:.*]]: i32)
+// CHECK-NEXT: %[[CMP:.*]] = cmpi ult, %[[LHS]], %[[RHS]] : i32

mlir/test/Dialect/Standard/ops.mlir

func @constant_complex_f64() -> complex<f64> {
   return %result : complex<f64>
 }
 
 // CHECK-LABEL: func @bitcast(
 func @bitcast(%arg : f32) -> i32 {
   %res = bitcast %arg : f32 to i32
   return %res : i32
 }
+
+// CHECK-LABEL: func @maximum
+func @maximum(%v1: vector<4xf32>, %v2: vector<4xf32>,
+               %f1: f32, %f2: f32,
+               %i1: i32, %i2: i32) {
+  %max_vector = maxf(%v1, %v2)
+    : (vector<4xf32>, vector<4xf32>) -> vector<4xf32>
+  %max_float = maxf(%f1, %f2) : (f32, f32) -> f32
+  %max_signed = maxsi(%i1, %i2) : (i32, i32) -> i32
+  %max_unsigned = maxui(%i1, %i2) : (i32, i32) -> i32
+  return
+}
+
+// CHECK-LABEL: func @minimum
+func @minimum(%v1: vector<4xf32>, %v2: vector<4xf32>,
+               %f1: f32, %f2: f32,
+               %i1: i32, %i2: i32) {
+  %min_vector = minf(%v1, %v2)
+    : (vector<4xf32>, vector<4xf32>) -> vector<4xf32>
+  %min_float = minf(%f1, %f2) : (f32, f32) -> f32
+  %min_signed = minsi(%i1, %i2) : (i32, i32) -> i32
+  %min_unsigned = minui(%i1, %i2) : (i32, i32) -> i32
+  return
+}