[mlir][arith] Add narrowing patterns for `max*i` and `min*i`

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

struct DivUIPattern final : BinaryOpNarrowingPattern<arith::DivUIOp> {
 
   // Unlike multiplication, division requires only one more result bit.
   unsigned getResultBitsProduced(unsigned operandBits) const override {
     return operandBits + 1;
   }
 };
 
 //===----------------------------------------------------------------------===//
+// Min/Max Patterns
+//===----------------------------------------------------------------------===//
+
+template <typename MinMaxOp, ExtensionKind Kind>
+struct MinMaxPattern final : BinaryOpNarrowingPattern<MinMaxOp> {
+  using BinaryOpNarrowingPattern<MinMaxOp>::BinaryOpNarrowingPattern;
+
+  // This optimization only applies to signed arguments.
+  bool isSupported(ExtensionOp ext) const override {
+    return ext.getKind() == Kind;
+  }
+
+  unsigned getResultBitsProduced(unsigned operandBits) const override {
+    return operandBits;
+  }
+};
+using MaxSIPattern = MinMaxPattern<arith::MaxSIOp, ExtensionKind::Sign>;
+using MaxUIPattern = MinMaxPattern<arith::MaxUIOp, ExtensionKind::Zero>;
+using MinSIPattern = MinMaxPattern<arith::MinSIOp, ExtensionKind::Sign>;
+using MinUIPattern = MinMaxPattern<arith::MinUIOp, ExtensionKind::Zero>;
+
+//===----------------------------------------------------------------------===//
 // *IToFPOp Patterns
 //===----------------------------------------------------------------------===//
 
 template <typename IToFPOp, ExtensionKind Extension>
 struct IToFPPattern final : NarrowingPattern<IToFPOp> {
   using NarrowingPattern<IToFPOp>::NarrowingPattern;
 
   LogicalResult matchAndRewrite(IToFPOp op,
void populateArithIntNarrowingPatterns(
     RewritePatternSet &patterns, const ArithIntNarrowingOptions &options) {
   // Add commute patterns with a higher benefit. This is to expose more
   // optimization opportunities to narrowing patterns.
   patterns.add<ExtensionOverExtract, ExtensionOverExtractElement,
                ExtensionOverExtractStridedSlice, ExtensionOverInsert>(
       patterns.getContext(), options, PatternBenefit(2));
 
   patterns.add<AddIPattern, SubIPattern, MulIPattern, DivSIPattern,
-               DivUIPattern, SIToFPPattern, UIToFPPattern>(
+               DivUIPattern, MaxSIPattern, MaxUIPattern, MinSIPattern,
+               MinUIPattern, SIToFPPattern, UIToFPPattern>(
       patterns.getContext(), options);
 }
 
 } // namespace mlir::arith

mlir/test/Dialect/Arith/int-narrowing.mlir

 // CHECK-NEXT:    return %[[RET]] : f16
 func.func @uitofp_extsi_i16(%a: i16) -> f16 {
   %b = arith.extsi %a : i16 to i32
   %f = arith.uitofp %b : i32 to f16
   return %f : f16
 }
 
 //===----------------------------------------------------------------------===//
+// arith.maxsi
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: func.func @maxsi_extsi_i8
+// CHECK-SAME:    (%[[LHS:.+]]: i8, %[[RHS:.+]]: i8)
+// CHECK-NEXT:    %[[MAX:.+]]  = arith.maxsi %[[LHS]], %[[RHS]] : i8
+// CHECK-NEXT:    %[[RET:.+]]  = arith.extsi %[[MAX]] : i8 to i32
+// CHECK-NEXT:    return %[[RET]] : i32
+func.func @maxsi_extsi_i8(%lhs: i8, %rhs: i8) -> i32 {
+  %a = arith.extsi %lhs : i8 to i32
+  %b = arith.extsi %rhs : i8 to i32
+  %r = arith.maxsi %a, %b : i32
+  return %r : i32
+}
+
+// This patterns should only apply to `arith.maxsi` ops with sign-extended
+// arguments.
+//
+// CHECK-LABEL: func.func @maxsi_extui_i8
+// CHECK-SAME:    (%[[ARG0:.+]]: i8, %[[ARG1:.+]]: i8)
+// CHECK-NEXT:    %[[EXT0:.+]] = arith.extui %[[ARG0]] : i8 to i32
+// CHECK-NEXT:    %[[EXT1:.+]] = arith.extui %[[ARG1]] : i8 to i32
+// CHECK-NEXT:    %[[MAX:.+]]  = arith.maxsi %[[EXT0]], %[[EXT1]] : i32
+// CHECK-NEXT:    return %[[MAX]] : i32
+func.func @maxsi_extui_i8(%lhs: i8, %rhs: i8) -> i32 {
+  %a = arith.extui %lhs : i8 to i32
+  %b = arith.extui %rhs : i8 to i32
+  %r = arith.maxsi %a, %b : i32
+  return %r : i32
+}
+
+//===----------------------------------------------------------------------===//
+// arith.maxui
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: func.func @maxui_extui_i8
+// CHECK-SAME:    (%[[LHS:.+]]: i8, %[[RHS:.+]]: i8)
+// CHECK-NEXT:    %[[MAX:.+]]  = arith.maxui %[[LHS]], %[[RHS]] : i8
+// CHECK-NEXT:    %[[RET:.+]]  = arith.extui %[[MAX]] : i8 to i32
+// CHECK-NEXT:    return %[[RET]] : i32
+func.func @maxui_extui_i8(%lhs: i8, %rhs: i8) -> i32 {
+  %a = arith.extui %lhs : i8 to i32
+  %b = arith.extui %rhs : i8 to i32
+  %r = arith.maxui %a, %b : i32
+  return %r : i32
+}
+
+// This patterns should only apply to `arith.maxsi` ops with zero-extended
+// arguments.
+//
+// CHECK-LABEL: func.func @maxui_extsi_i8
+// CHECK-SAME:    (%[[ARG0:.+]]: i8, %[[ARG1:.+]]: i8)
+// CHECK-NEXT:    %[[EXT0:.+]] = arith.extsi %[[ARG0]] : i8 to i32
+// CHECK-NEXT:    %[[EXT1:.+]] = arith.extsi %[[ARG1]] : i8 to i32
+// CHECK-NEXT:    %[[MAX:.+]]  = arith.maxui %[[EXT0]], %[[EXT1]] : i32
+// CHECK-NEXT:    return %[[MAX]] : i32
+func.func @maxui_extsi_i8(%lhs: i8, %rhs: i8) -> i32 {
+  %a = arith.extsi %lhs : i8 to i32
+  %b = arith.extsi %rhs : i8 to i32
+  %r = arith.maxui %a, %b : i32
+  return %r : i32
+}
+
+//===----------------------------------------------------------------------===//
+// arith.minsi
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: func.func @minsi_extsi_i8
+// CHECK-SAME:    (%[[LHS:.+]]: i8, %[[RHS:.+]]: i8)
+// CHECK-NEXT:    %[[min:.+]]  = arith.minsi %[[LHS]], %[[RHS]] : i8
+// CHECK-NEXT:    %[[RET:.+]]  = arith.extsi %[[min]] : i8 to i32
+// CHECK-NEXT:    return %[[RET]] : i32
+func.func @minsi_extsi_i8(%lhs: i8, %rhs: i8) -> i32 {
+  %a = arith.extsi %lhs : i8 to i32
+  %b = arith.extsi %rhs : i8 to i32
+  %r = arith.minsi %a, %b : i32
+  return %r : i32
+}
+
+// This patterns should only apply to `arith.minsi` ops with sign-extended
+// arguments.
+//
+// CHECK-LABEL: func.func @minsi_extui_i8
+// CHECK-SAME:    (%[[ARG0:.+]]: i8, %[[ARG1:.+]]: i8)
+// CHECK-NEXT:    %[[EXT0:.+]] = arith.extui %[[ARG0]] : i8 to i32
+// CHECK-NEXT:    %[[EXT1:.+]] = arith.extui %[[ARG1]] : i8 to i32
+// CHECK-NEXT:    %[[min:.+]]  = arith.minsi %[[EXT0]], %[[EXT1]] : i32
+// CHECK-NEXT:    return %[[min]] : i32
+func.func @minsi_extui_i8(%lhs: i8, %rhs: i8) -> i32 {
+  %a = arith.extui %lhs : i8 to i32
+  %b = arith.extui %rhs : i8 to i32
+  %r = arith.minsi %a, %b : i32
+  return %r : i32
+}
+
+//===----------------------------------------------------------------------===//
+// arith.minui
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: func.func @minui_extui_i8
+// CHECK-SAME:    (%[[LHS:.+]]: i8, %[[RHS:.+]]: i8)
+// CHECK-NEXT:    %[[min:.+]]  = arith.minui %[[LHS]], %[[RHS]] : i8
+// CHECK-NEXT:    %[[RET:.+]]  = arith.extui %[[min]] : i8 to i32
+// CHECK-NEXT:    return %[[RET]] : i32
+func.func @minui_extui_i8(%lhs: i8, %rhs: i8) -> i32 {
+  %a = arith.extui %lhs : i8 to i32
+  %b = arith.extui %rhs : i8 to i32
+  %r = arith.minui %a, %b : i32
+  return %r : i32
+}
+
+// This patterns should only apply to `arith.minsi` ops with zero-extended
+// arguments.
+//
+// CHECK-LABEL: func.func @minui_extsi_i8
+// CHECK-SAME:    (%[[ARG0:.+]]: i8, %[[ARG1:.+]]: i8)
+// CHECK-NEXT:    %[[EXT0:.+]] = arith.extsi %[[ARG0]] : i8 to i32
+// CHECK-NEXT:    %[[EXT1:.+]] = arith.extsi %[[ARG1]] : i8 to i32
+// CHECK-NEXT:    %[[min:.+]]  = arith.minui %[[EXT0]], %[[EXT1]] : i32
+// CHECK-NEXT:    return %[[min]] : i32
+func.func @minui_extsi_i8(%lhs: i8, %rhs: i8) -> i32 {
+  %a = arith.extsi %lhs : i8 to i32
+  %b = arith.extsi %rhs : i8 to i32
+  %r = arith.minui %a, %b : i32
+  return %r : i32
+}
+
+//===----------------------------------------------------------------------===//
 // Commute Extension over Vector Ops
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: func.func @extsi_over_extract_3xi16
 // CHECK-SAME:    (%[[ARG:.+]]: vector<3xi16>)
 // CHECK-NEXT:    %[[EXTR:.+]] = vector.extract %[[ARG]][1] : vector<3xi16>
 // CHECK-NEXT:    %[[RET:.+]]  = arith.sitofp %[[EXTR]] : i16 to f16
 // CHECK-NEXT:    return %[[RET]] : f16