[Mlir-commits] [mlir] [mlir][arith] Canonicalize cmpi select to min/max (PR #201737)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Jun 4 19:54:31 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-arith
Author: Ming Yan (NexMing)
<details>
<summary>Changes</summary>
This PR adds canonicalization patterns that fold `arith.select` + `arith.cmpi` into `arith.minsi`/`arith.maxsi`/`arith.minui`/`arith.maxui` operations.
---
Full diff: https://github.com/llvm/llvm-project/pull/201737.diff
3 Files Affected:
- (modified) mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td (+36)
- (modified) mlir/lib/Dialect/Arith/IR/ArithOps.cpp (+7-1)
- (modified) mlir/test/Dialect/Arith/canonicalize.mlir (+67)
``````````diff
diff --git a/mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td b/mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td
index 01ad0beb7ee0d..86c9b6d3551d8 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td
+++ b/mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td
@@ -319,6 +319,42 @@ def SelectI1ToNot :
(ConstantLikeMatcher ConstantAttr<I1Attr, "1">)),
(Arith_XOrIOp $pred, (Arith_ConstantOp ConstantAttr<I1Attr, "1">))>;
+class SelectCmpIToMinMaxBase<string pred, Op minMaxOp> :
+ Pat<(SelectOp
+ (Arith_CmpIOp
+ ConstantEnumCase<Arith_CmpIPredicateAttr, pred>,
+ $a, $b),
+ $a, $b),
+ (minMaxOp $a, $b)>;
+
+class SelectCmpIInverseToMinMaxBase<string pred, Op minMaxOp> :
+ Pat<(SelectOp
+ (Arith_CmpIOp
+ ConstantEnumCase<Arith_CmpIPredicateAttr, pred>,
+ $a, $b),
+ $b, $a),
+ (minMaxOp $a, $b)>;
+
+// select(cmpi pred, a, b), a, b => min/max(a, b)
+def SelectCmpISltToMinSI : SelectCmpIToMinMaxBase<"slt", Arith_MinSIOp>;
+def SelectCmpISleToMinSI : SelectCmpIToMinMaxBase<"sle", Arith_MinSIOp>;
+def SelectCmpISgtToMaxSI : SelectCmpIToMinMaxBase<"sgt", Arith_MaxSIOp>;
+def SelectCmpISgeToMaxSI : SelectCmpIToMinMaxBase<"sge", Arith_MaxSIOp>;
+def SelectCmpIUltToMinUI : SelectCmpIToMinMaxBase<"ult", Arith_MinUIOp>;
+def SelectCmpIUleToMinUI : SelectCmpIToMinMaxBase<"ule", Arith_MinUIOp>;
+def SelectCmpIUgtToMaxUI : SelectCmpIToMinMaxBase<"ugt", Arith_MaxUIOp>;
+def SelectCmpIUgeToMaxUI : SelectCmpIToMinMaxBase<"uge", Arith_MaxUIOp>;
+
+// select(cmpi pred, a, b), b, a => inverse min/max(a, b)
+def SelectCmpISltToMaxSI : SelectCmpIInverseToMinMaxBase<"slt", Arith_MaxSIOp>;
+def SelectCmpISleToMaxSI : SelectCmpIInverseToMinMaxBase<"sle", Arith_MaxSIOp>;
+def SelectCmpISgtToMinSI : SelectCmpIInverseToMinMaxBase<"sgt", Arith_MinSIOp>;
+def SelectCmpISgeToMinSI : SelectCmpIInverseToMinMaxBase<"sge", Arith_MinSIOp>;
+def SelectCmpIUltToMaxUI : SelectCmpIInverseToMinMaxBase<"ult", Arith_MaxUIOp>;
+def SelectCmpIUleToMaxUI : SelectCmpIInverseToMinMaxBase<"ule", Arith_MaxUIOp>;
+def SelectCmpIUgtToMinUI : SelectCmpIInverseToMinMaxBase<"ugt", Arith_MinUIOp>;
+def SelectCmpIUgeToMinUI : SelectCmpIInverseToMinMaxBase<"uge", Arith_MinUIOp>;
+
//===----------------------------------------------------------------------===//
// IndexCastOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
index 008242daeaf3c..dc3887c3e0b0e 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
@@ -2724,7 +2724,13 @@ struct SelectToExtUI : public OpRewritePattern<arith::SelectOp> {
void arith::SelectOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {
results.add<RedundantSelectFalse, RedundantSelectTrue, SelectNotCond,
- SelectI1ToNot, SelectToExtUI>(context);
+ SelectI1ToNot, SelectCmpISgeToMaxSI, SelectCmpISgeToMinSI,
+ SelectCmpISgtToMaxSI, SelectCmpISgtToMinSI, SelectCmpISleToMaxSI,
+ SelectCmpISleToMinSI, SelectCmpISltToMaxSI, SelectCmpISltToMinSI,
+ SelectCmpIUgeToMaxUI, SelectCmpIUgeToMinUI, SelectCmpIUgtToMaxUI,
+ SelectCmpIUgtToMinUI, SelectCmpIUleToMaxUI, SelectCmpIUleToMinUI,
+ SelectCmpIUltToMaxUI, SelectCmpIUltToMinUI, SelectToExtUI>(
+ context);
}
OpFoldResult arith::SelectOp::fold(FoldAdaptor adaptor) {
diff --git a/mlir/test/Dialect/Arith/canonicalize.mlir b/mlir/test/Dialect/Arith/canonicalize.mlir
index 88bf90af5be27..c12befc76675a 100644
--- a/mlir/test/Dialect/Arith/canonicalize.mlir
+++ b/mlir/test/Dialect/Arith/canonicalize.mlir
@@ -23,6 +23,73 @@ func.func @select_cmp_ne_select(%arg0: i64, %arg1: i64) -> i64 {
return %1 : i64
}
+// CHECK-LABEL: @select_cmp_signed_min_max
+// CHECK-SAME: (%[[ARG0:.+]]: i64, %[[ARG1:.+]]: i64)
+// CHECK-DAG: %[[MIN0:.+]] = arith.minsi %[[ARG0]], %[[ARG1]] : i64
+// CHECK-DAG: %[[MIN1:.+]] = arith.minsi %[[ARG0]], %[[ARG1]] : i64
+// CHECK-DAG: %[[MIN2:.+]] = arith.minsi %[[ARG0]], %[[ARG1]] : i64
+// CHECK-DAG: %[[MIN3:.+]] = arith.minsi %[[ARG0]], %[[ARG1]] : i64
+// CHECK-DAG: %[[MAX0:.+]] = arith.maxsi %[[ARG0]], %[[ARG1]] : i64
+// CHECK-DAG: %[[MAX1:.+]] = arith.maxsi %[[ARG0]], %[[ARG1]] : i64
+// CHECK-DAG: %[[MAX2:.+]] = arith.maxsi %[[ARG0]], %[[ARG1]] : i64
+// CHECK-DAG: %[[MAX3:.+]] = arith.maxsi %[[ARG0]], %[[ARG1]] : i64
+// CHECK: return %[[MIN0]], %[[MIN1]], %[[MIN2]], %[[MIN3]], %[[MAX0]], %[[MAX1]], %[[MAX2]], %[[MAX3]]
+func.func @select_cmp_signed_min_max(%arg0: i64, %arg1: i64) -> (i64, i64, i64, i64, i64, i64, i64, i64) {
+ %slt = arith.cmpi slt, %arg0, %arg1 : i64
+ %sle = arith.cmpi sle, %arg0, %arg1 : i64
+ %sgt = arith.cmpi sgt, %arg0, %arg1 : i64
+ %sge = arith.cmpi sge, %arg0, %arg1 : i64
+ %min0 = arith.select %slt, %arg0, %arg1 : i64
+ %min1 = arith.select %sle, %arg0, %arg1 : i64
+ %min2 = arith.select %sgt, %arg1, %arg0 : i64
+ %min3 = arith.select %sge, %arg1, %arg0 : i64
+ %max0 = arith.select %slt, %arg1, %arg0 : i64
+ %max1 = arith.select %sle, %arg1, %arg0 : i64
+ %max2 = arith.select %sgt, %arg0, %arg1 : i64
+ %max3 = arith.select %sge, %arg0, %arg1 : i64
+ return %min0, %min1, %min2, %min3, %max0, %max1, %max2, %max3 : i64, i64, i64, i64, i64, i64, i64, i64
+}
+
+// CHECK-LABEL: @select_cmp_unsigned_min_max
+// CHECK-SAME: (%[[ARG0:.+]]: i64, %[[ARG1:.+]]: i64)
+// CHECK-DAG: %[[MIN0:.+]] = arith.minui %[[ARG0]], %[[ARG1]] : i64
+// CHECK-DAG: %[[MIN1:.+]] = arith.minui %[[ARG0]], %[[ARG1]] : i64
+// CHECK-DAG: %[[MIN2:.+]] = arith.minui %[[ARG0]], %[[ARG1]] : i64
+// CHECK-DAG: %[[MIN3:.+]] = arith.minui %[[ARG0]], %[[ARG1]] : i64
+// CHECK-DAG: %[[MAX0:.+]] = arith.maxui %[[ARG0]], %[[ARG1]] : i64
+// CHECK-DAG: %[[MAX1:.+]] = arith.maxui %[[ARG0]], %[[ARG1]] : i64
+// CHECK-DAG: %[[MAX2:.+]] = arith.maxui %[[ARG0]], %[[ARG1]] : i64
+// CHECK-DAG: %[[MAX3:.+]] = arith.maxui %[[ARG0]], %[[ARG1]] : i64
+// CHECK: return %[[MIN0]], %[[MIN1]], %[[MIN2]], %[[MIN3]], %[[MAX0]], %[[MAX1]], %[[MAX2]], %[[MAX3]]
+func.func @select_cmp_unsigned_min_max(%arg0: i64, %arg1: i64) -> (i64, i64, i64, i64, i64, i64, i64, i64) {
+ %ult = arith.cmpi ult, %arg0, %arg1 : i64
+ %ule = arith.cmpi ule, %arg0, %arg1 : i64
+ %ugt = arith.cmpi ugt, %arg0, %arg1 : i64
+ %uge = arith.cmpi uge, %arg0, %arg1 : i64
+ %min0 = arith.select %ult, %arg0, %arg1 : i64
+ %min1 = arith.select %ule, %arg0, %arg1 : i64
+ %min2 = arith.select %ugt, %arg1, %arg0 : i64
+ %min3 = arith.select %uge, %arg1, %arg0 : i64
+ %max0 = arith.select %ult, %arg1, %arg0 : i64
+ %max1 = arith.select %ule, %arg1, %arg0 : i64
+ %max2 = arith.select %ugt, %arg0, %arg1 : i64
+ %max3 = arith.select %uge, %arg0, %arg1 : i64
+ return %min0, %min1, %min2, %min3, %max0, %max1, %max2, %max3 : i64, i64, i64, i64, i64, i64, i64, i64
+}
+
+// CHECK-LABEL: @select_cmp_min_max_index_vector
+// CHECK-SAME: (%[[IDX0:.+]]: index, %[[IDX1:.+]]: index, %[[VEC0:.+]]: vector<4xi32>, %[[VEC1:.+]]: vector<4xi32>)
+// CHECK-DAG: %[[IDX:.+]] = arith.minsi %[[IDX0]], %[[IDX1]] : index
+// CHECK-DAG: %[[VEC:.+]] = arith.maxui %[[VEC0]], %[[VEC1]] : vector<4xi32>
+// CHECK: return %[[IDX]], %[[VEC]]
+func.func @select_cmp_min_max_index_vector(%arg0: index, %arg1: index, %arg2: vector<4xi32>, %arg3: vector<4xi32>) -> (index, vector<4xi32>) {
+ %cmp0 = arith.cmpi sle, %arg0, %arg1 : index
+ %res0 = arith.select %cmp0, %arg0, %arg1 : index
+ %cmp1 = arith.cmpi ult, %arg2, %arg3 : vector<4xi32>
+ %res1 = arith.select %cmp1, %arg3, %arg2 : vector<4xi1>, vector<4xi32>
+ return %res0, %res1 : index, vector<4xi32>
+}
+
// CHECK-LABEL: @select_extui
// CHECK: %[[res:.+]] = arith.extui %arg0 : i1 to i64
// CHECK: return %[[res]]
``````````
</details>
https://github.com/llvm/llvm-project/pull/201737
More information about the Mlir-commits
mailing list