[mlir][Arith] Remove expansions of integer min and max ops

[krzysz00]

Summary

As of several months ago, both ArithToLLVM and ArithToSPIRV have
native support for integer min and max operations. Since these are all
the targets available in MLIR core, the need to "expand" arith.minui,
arith.minsi, arith,maxsi, and arith.manxui to more primitive
operations is to longer present.

Therefore, the expanding of integer min and max operations in Arith,
while correct, is likely to lead to performance loss by way of
misoptimization further down the line, and is no longer needed for
anyone's correctness.

This change may break downstream tests, but will not affect the
semantics of MLIR programs.

arith.minf and arith.maxf have a lot of underlying complexity due to
the many different possible NaN and signed zero semantics available on
various platforms, and so removing their expansion is left to a future
commit.

[ThomasRaoux]

Not directly related to this patch but the lowering of MaxFOp/MinFOp to LLVM doesn't look correct for NaN. MaxFOp defines the semantic for NaN as: If one of the arguments is NaN, then the result is also NaN. (https://github.com/llvm/llvm-project/blob/7df1da08ff9229c12884b408b0f614b8f20d6d5a/mlir/include/mlir/Dialect/Arith/IR/ArithOps.td#L851) while MaxFOp is lowered to llvm intrinsic maxnum (https://github.com/llvm/llvm-project/blob/7df1da08ff9229c12884b408b0f614b8f20d6d5a/mlir/lib/Conversion/ArithToLLVM/ArithToLLVM.cpp#L56) which has the semantic: If either operand is a NaN, returns the other non-NaN operand. Returns NaN only if both operands are NaN. (https://llvm.org/docs/LangRef.html#llvm-maxnum-intrinsic)

We should probably fix the lowering to LLVM dialect first to avoid regressing functionality. What do you think?

[kuhar]

In D140856#4022395, @ThomasRaoux wrote:

Not directly related to this patch but the lowering of MaxFOp/MinFOp to LLVM doesn't look correct for NaN. MaxFOp defines the semantic for NaN as: If one of the arguments is NaN, then the result is also NaN. (https://github.com/llvm/llvm-project/blob/7df1da08ff9229c12884b408b0f614b8f20d6d5a/mlir/include/mlir/Dialect/Arith/IR/ArithOps.td#L851) while MaxFOp is lowered to llvm intrinsic maxnum (https://github.com/llvm/llvm-project/blob/7df1da08ff9229c12884b408b0f614b8f20d6d5a/mlir/lib/Conversion/ArithToLLVM/ArithToLLVM.cpp#L56) which has the semantic: If either operand is a NaN, returns the other non-NaN operand. Returns NaN only if both operands are NaN. (https://llvm.org/docs/LangRef.html#llvm-maxnum-intrinsic)

We should probably fix the lowering to LLVM dialect first to avoid regressing functionality. What do you think?

+1. I think this is what D137786 by @gflegar was trying to address.

[krzysz00]

I agree that it'd be a good idea to fix up the LLVM lowering first.

Do y'all think that'l take long enough that I should amend this patch to just remove the integer max/min cases and come back for the float ones later?

[Mogball]

+1 from me, but you might need to add that other patch as a dependency

[krzysz00]

Make this only for integers to dodge the minf/maxf question

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

Value compareRes =
         rewriter.create<arith::OrIOp>(loc, firstTerm, secondTerm);
     // Perform substitution and return success.
     rewriter.replaceOpWithNewOp<arith::SelectOp>(op, compareRes, negRes,
                                                  posRes);
     return success();
   }
 };
 
-template <typename OpTy, arith::CmpFPredicate pred>
-struct MaxMinFOpConverter : public OpRewritePattern<OpTy> {
-public:
-  using OpRewritePattern<OpTy>::OpRewritePattern;
-
-  LogicalResult matchAndRewrite(OpTy op,
-                                PatternRewriter &rewriter) const final {
-    Value lhs = op.getLhs();
-    Value rhs = op.getRhs();
-
-    Location loc = op.getLoc();
-    // If any operand is NaN, 'cmp' will be true (and 'select' returns 'lhs').
-    static_assert(pred == arith::CmpFPredicate::UGT ||
-                      pred == arith::CmpFPredicate::ULT,
-                  "pred must be either UGT or ULT");
-    Value cmp = rewriter.create<arith::CmpFOp>(loc, pred, lhs, rhs);
-    Value select = rewriter.create<arith::SelectOp>(loc, cmp, lhs, rhs);
-
-    // Handle the case where rhs is NaN: 'isNaN(rhs) ? rhs : select'.
-    Value isNaN = rewriter.create<arith::CmpFOp>(loc, arith::CmpFPredicate::UNO,
-                                                 rhs, rhs);
-    rewriter.replaceOpWithNewOp<arith::SelectOp>(op, isNaN, rhs, select);
-    return success();
-  }
-};
-
-template <typename OpTy, arith::CmpIPredicate pred>
-struct MaxMinIOpConverter : public OpRewritePattern<OpTy> {
-public:
-  using OpRewritePattern<OpTy>::OpRewritePattern;
-  LogicalResult matchAndRewrite(OpTy op,
-                                PatternRewriter &rewriter) const final {
-    Value lhs = op.getLhs();
-    Value rhs = op.getRhs();
-
-    Location loc = op.getLoc();
-    Value cmp = rewriter.create<arith::CmpIOp>(loc, pred, lhs, rhs);
-    rewriter.replaceOpWithNewOp<arith::SelectOp>(op, cmp, lhs, rhs);
-    return success();
-  }
-};
-
 struct ArithExpandOpsPass
     : public arith::impl::ArithExpandOpsBase<ArithExpandOpsPass> {
   void runOnOperation() override {
     RewritePatternSet patterns(&getContext());
     ConversionTarget target(getContext());
 
     arith::populateArithExpandOpsPatterns(patterns);
 
     target.addLegalDialect<arith::ArithDialect>();
     // clang-format off
     target.addIllegalOp<
       arith::CeilDivSIOp,
       arith::CeilDivUIOp,
-      arith::FloorDivSIOp,
-      arith::MaxFOp,
-      arith::MaxSIOp,
-      arith::MaxUIOp,
-      arith::MinFOp,
-      arith::MinSIOp,
-      arith::MinUIOp
+      arith::FloorDivSIOp
     >();
     // clang-format on
     if (failed(applyPartialConversion(getOperation(), target,
                                       std::move(patterns))))
       signalPassFailure();
   }
 };
 
 } // namespace
 
 void mlir::arith::populateCeilFloorDivExpandOpsPatterns(
     RewritePatternSet &patterns) {
   patterns
       .add<CeilDivSIOpConverter, CeilDivUIOpConverter, FloorDivSIOpConverter>(
           patterns.getContext());
 }
 
 void mlir::arith::populateArithExpandOpsPatterns(RewritePatternSet &patterns) {
   populateCeilFloorDivExpandOpsPatterns(patterns);
-  // clang-format off
-  patterns.add<
-    MaxMinFOpConverter<MaxFOp, arith::CmpFPredicate::UGT>,
-    MaxMinFOpConverter<MinFOp, arith::CmpFPredicate::ULT>,
-    MaxMinIOpConverter<MaxSIOp, arith::CmpIPredicate::sgt>,
-    MaxMinIOpConverter<MaxUIOp, arith::CmpIPredicate::ugt>,
-    MaxMinIOpConverter<MinSIOp, arith::CmpIPredicate::slt>,
-    MaxMinIOpConverter<MinUIOp, arith::CmpIPredicate::ult>
-   >(patterns.getContext());
-  // clang-format on
 }
 
 std::unique_ptr<Pass> mlir::arith::createArithExpandOpsPass() {
   return std::make_unique<ArithExpandOpsPass>();
 }

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

 // CHECK:           [[ZERO:%.+]] = arith.constant 0 : index
 // CHECK:           [[ISZERO:%.+]] = arith.cmpi eq, %arg0, [[ZERO]] : index
 // CHECK:           [[ONE:%.+]] = arith.constant 1 : index
 // CHECK:           [[SUB:%.+]] = arith.subi %arg0, [[ONE]] : index
 // CHECK:           [[DIV:%.+]] = arith.divui [[SUB]], %arg1 : index
 // CHECK:           [[REM:%.+]] = arith.addi [[DIV]], [[ONE]] : index
 // CHECK:           [[RES:%.+]] = arith.select [[ISZERO]], [[ZERO]], [[REM]] : index
 }
-
-// -----
-
-// CHECK-LABEL: func @maxf
-func.func @maxf(%a: f32, %b: f32) -> f32 {
-  %result = arith.maxf %a, %b : f32
-  return %result : f32
-}
-// CHECK-SAME: %[[LHS:.*]]: f32, %[[RHS:.*]]: f32)
-// CHECK-NEXT: %[[CMP:.*]] = arith.cmpf ugt, %[[LHS]], %[[RHS]] : f32
-// CHECK-NEXT: %[[SELECT:.*]] = arith.select %[[CMP]], %[[LHS]], %[[RHS]] : f32
-// CHECK-NEXT: %[[IS_NAN:.*]] = arith.cmpf uno, %[[RHS]], %[[RHS]] : f32
-// CHECK-NEXT: %[[RESULT:.*]] = arith.select %[[IS_NAN]], %[[RHS]], %[[SELECT]] : f32
-// CHECK-NEXT: return %[[RESULT]] : f32
-
-// -----
-
-// CHECK-LABEL: func @maxf_vector
-func.func @maxf_vector(%a: vector<4xf16>, %b: vector<4xf16>) -> vector<4xf16> {
-  %result = arith.maxf %a, %b : vector<4xf16>
-  return %result : vector<4xf16>
-}
-// CHECK-SAME: %[[LHS:.*]]: vector<4xf16>, %[[RHS:.*]]: vector<4xf16>)
-// CHECK-NEXT: %[[CMP:.*]] = arith.cmpf ugt, %[[LHS]], %[[RHS]] : vector<4xf16>
-// CHECK-NEXT: %[[SELECT:.*]] = arith.select %[[CMP]], %[[LHS]], %[[RHS]]
-// CHECK-NEXT: %[[IS_NAN:.*]] = arith.cmpf uno, %[[RHS]], %[[RHS]] : vector<4xf16>
-// CHECK-NEXT: %[[RESULT:.*]] = arith.select %[[IS_NAN]], %[[RHS]], %[[SELECT]]
-// CHECK-NEXT: return %[[RESULT]] : vector<4xf16>
-
-// -----
-
-// CHECK-LABEL: func @minf
-func.func @minf(%a: f32, %b: f32) -> f32 {
-  %result = arith.minf %a, %b : f32
-  return %result : f32
-}
-// CHECK-SAME: %[[LHS:.*]]: f32, %[[RHS:.*]]: f32)
-// CHECK-NEXT: %[[CMP:.*]] = arith.cmpf ult, %[[LHS]], %[[RHS]] : f32
-// CHECK-NEXT: %[[SELECT:.*]] = arith.select %[[CMP]], %[[LHS]], %[[RHS]] : f32
-// CHECK-NEXT: %[[IS_NAN:.*]] = arith.cmpf uno, %[[RHS]], %[[RHS]] : f32
-// CHECK-NEXT: %[[RESULT:.*]] = arith.select %[[IS_NAN]], %[[RHS]], %[[SELECT]] : f32
-// CHECK-NEXT: return %[[RESULT]] : f32
-
-
-// -----
-
-// CHECK-LABEL: func @maxsi
-func.func @maxsi(%a: i32, %b: i32) -> i32 {
-  %result = arith.maxsi %a, %b : i32
-  return %result : i32
-}
-// CHECK-SAME: %[[LHS:.*]]: i32, %[[RHS:.*]]: i32)
-// CHECK-NEXT: %[[CMP:.*]] = arith.cmpi sgt, %[[LHS]], %[[RHS]] : i32
-
-// -----
-
-// CHECK-LABEL: func @minsi
-func.func @minsi(%a: i32, %b: i32) -> i32 {
-  %result = arith.minsi %a, %b : i32
-  return %result : i32
-}
-// CHECK-SAME: %[[LHS:.*]]: i32, %[[RHS:.*]]: i32)
-// CHECK-NEXT: %[[CMP:.*]] = arith.cmpi slt, %[[LHS]], %[[RHS]] : i32
-
-
-// -----
-
-// CHECK-LABEL: func @maxui
-func.func @maxui(%a: i32, %b: i32) -> i32 {
-  %result = arith.maxui %a, %b : i32
-  return %result : i32
-}
-// CHECK-SAME: %[[LHS:.*]]: i32, %[[RHS:.*]]: i32)
-// CHECK-NEXT: %[[CMP:.*]] = arith.cmpi ugt, %[[LHS]], %[[RHS]] : i32
-
-
-// -----
-
-// CHECK-LABEL: func @minui
-func.func @minui(%a: i32, %b: i32) -> i32 {
-  %result = arith.minui %a, %b : i32
-  return %result : i32
-}
-// CHECK-SAME: %[[LHS:.*]]: i32, %[[RHS:.*]]: i32)
-// CHECK-NEXT: %[[CMP:.*]] = arith.cmpi ult, %[[LHS]], %[[RHS]] : i32