[Mlir-commits] [mlir] [mlir][arith] Gate min/max expansion in arith-expand behind include-min-max (PR #211884)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Jul 24 11:58:41 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-gpu
Author: Md Abdullah Shahneous Bari (mshahneo)
<details>
<summary>Changes</summary>
`arith-expand` expanded `arith.maximumf`/`minimumf`/`maxnumf`/`minnumf` and the signed/unsigned integer max/min ops into `cmpf`/`cmpi` + `select` sequences. These ops also have a direct arith-to-llvm lowering to the `llvm.intr.maximum`/`minimum`/... intrinsics, which are a single hardware instruction on many targets. Pipelines that run arith-to-llvm after arith-expand (e.g. the GPU-to-XeVM pipeline) therefore paid a large, avoidable overhead.
Add an `include-min-max` option (default `true`, preserving existing behavior) that controls whether these min/max ops are expanded. The min/max converters are factored into a new `populateExpandMinMaxPatterns`; the ceil/floor-div and scaling ext/trunc expansions (which have no LLVM lowering) always run. Set `include-min-max=false` in the GPU-to-XeVM pipeline so the intrinsic lowering is used.
---
Full diff: https://github.com/llvm/llvm-project/pull/211884.diff
5 Files Affected:
- (modified) mlir/include/mlir/Dialect/Arith/Transforms/Passes.h (+7)
- (modified) mlir/include/mlir/Dialect/Arith/Transforms/Passes.td (+8)
- (modified) mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp (+31-12)
- (modified) mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp (+7)
- (added) mlir/test/Dialect/Arith/expand-ops-min-max.mlir (+64)
``````````diff
diff --git a/mlir/include/mlir/Dialect/Arith/Transforms/Passes.h b/mlir/include/mlir/Dialect/Arith/Transforms/Passes.h
index 5a07a01d0928a..4895443a7ba01 100644
--- a/mlir/include/mlir/Dialect/Arith/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/Arith/Transforms/Passes.h
@@ -68,6 +68,13 @@ void populateExpandF8E8M0Patterns(RewritePatternSet &patterns);
/// Add patterns to expand scaling ExtF/TruncF ops to equivalent arith ops
void populateExpandScalingExtTruncPatterns(RewritePatternSet &patterns);
+/// Add patterns to expand the min/max ops (`arith.maximumf`/`minimumf`/
+/// `maxnumf`/`minnumf` and the signed/unsigned integer `max`/`min`) into
+/// `cmpf`/`cmpi` + `select` sequences. These ops also have a direct
+/// arith-to-llvm lowering, so pipelines that run arith-to-llvm may prefer to
+/// skip this expansion.
+void populateExpandMinMaxPatterns(RewritePatternSet &patterns);
+
/// Add patterns to expand `arith.flush_denormals` into integer arithmetic
/// (bitcast + bit masks + compare + select). Only matches IEEE-like
/// floating-point types.
diff --git a/mlir/include/mlir/Dialect/Arith/Transforms/Passes.td b/mlir/include/mlir/Dialect/Arith/Transforms/Passes.td
index 27e9146ec3606..8fb7439854406 100644
--- a/mlir/include/mlir/Dialect/Arith/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Arith/Transforms/Passes.td
@@ -25,6 +25,14 @@ def ArithExpandOpsPass : Pass<"arith-expand"> {
/*default=*/"false",
"Enable expansion of `arith.flush_denormals` on IEEE-like "
"floating-point types">,
+ Option<"includeMinMax", "include-min-max", "bool", /*default=*/"true",
+ "Enable expansion of the min/max ops (maximumf/minimumf/maxnumf/"
+ "minnumf and the signed/unsigned integer max/min). These ops "
+ "also have a direct arith-to-llvm lowering to the "
+ "`llvm.intr.maximum/minimum/maxnum/minnum` (and smax/umax/...) "
+ "intrinsics; disable this so pipelines that run arith-to-llvm "
+ "can use those single-instruction lowerings instead of the "
+ "cmpf/select expansion">,
];
}
diff --git a/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp b/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp
index c9217c57a5f25..6e607311300a5 100644
--- a/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp
@@ -823,7 +823,8 @@ struct ArithExpandOpsPass
RewritePatternSet patterns(&getContext());
ConversionTarget target(getContext());
- arith::populateArithExpandOpsPatterns(patterns);
+ arith::populateCeilFloorDivExpandOpsPatterns(patterns);
+ arith::populateExpandScalingExtTruncPatterns(patterns);
target.addLegalDialect<arith::ArithDialect>();
target.addLegalDialect<vector::VectorDialect>();
@@ -833,17 +834,31 @@ struct ArithExpandOpsPass
arith::CeilDivSIOp,
arith::CeilDivUIOp,
arith::FloorDivSIOp,
- arith::MaxSIOp,
- arith::MaxUIOp,
- arith::MinSIOp,
- arith::MinUIOp,
- arith::MaximumFOp,
- arith::MinimumFOp,
- arith::MaxNumFOp,
- arith::MinNumFOp,
arith::ScalingExtFOp,
arith::ScalingTruncFOp
>();
+ // clang-format on
+
+ // The min/max ops also have a direct arith-to-llvm lowering to the
+ // `llvm.intr.maximum`/`minimum`/... intrinsics, which are a single hardware
+ // instruction on many targets. Only expand them into cmpf/cmpi + select
+ // when requested, so pipelines that run arith-to-llvm can keep the
+ // intrinsic lowering.
+ if (includeMinMax) {
+ arith::populateExpandMinMaxPatterns(patterns);
+ // clang-format off
+ target.addIllegalOp<
+ arith::MaxSIOp,
+ arith::MaxUIOp,
+ arith::MinSIOp,
+ arith::MinUIOp,
+ arith::MaximumFOp,
+ arith::MinimumFOp,
+ arith::MaxNumFOp,
+ arith::MinNumFOp
+ >();
+ // clang-format on
+ }
if (includeBf16)
arith::populateExpandBFloat16Patterns(patterns);
@@ -936,9 +951,7 @@ void mlir::arith::populateExpandFlushDenormalsPatterns(
patterns.add<FlushDenormalsOpConverter>(patterns.getContext());
}
-void mlir::arith::populateArithExpandOpsPatterns(RewritePatternSet &patterns) {
- populateCeilFloorDivExpandOpsPatterns(patterns);
- populateExpandScalingExtTruncPatterns(patterns);
+void mlir::arith::populateExpandMinMaxPatterns(RewritePatternSet &patterns) {
// clang-format off
patterns.add<
MaxMinIOpConverter<MaxSIOp, arith::CmpIPredicate::sgt>,
@@ -952,3 +965,9 @@ void mlir::arith::populateArithExpandOpsPatterns(RewritePatternSet &patterns) {
>(patterns.getContext());
// clang-format on
}
+
+void mlir::arith::populateArithExpandOpsPatterns(RewritePatternSet &patterns) {
+ populateCeilFloorDivExpandOpsPatterns(patterns);
+ populateExpandScalingExtTruncPatterns(patterns);
+ populateExpandMinMaxPatterns(patterns);
+}
diff --git a/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp b/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp
index 4dc9e2acfe235..7841e34793bc5 100644
--- a/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp
+++ b/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp
@@ -111,6 +111,13 @@ void buildGPUPassPipeline(OpPassManager &pm,
{
arith::ArithExpandOpsPassOptions arithExpandOptions;
arithExpandOptions.includeF8E8M0 = true;
+ // Do not expand the min/max ops (arith.maximumf/minimumf/...) into
+ // cmpf + select here: they are lowered directly to the
+ // `llvm.intr.maximum`/`minimum`/... intrinsics (single hardware
+ // instructions) by the later arith-to-llvm conversion. Expanding them
+ // regresses e.g. flash-attention softmax, whose running-max is
+ // arith.maximumf-heavy.
+ arithExpandOptions.includeMinMax = false;
pm.addNestedPass<gpu::GPUModuleOp>(
arith::createArithExpandOpsPass(arithExpandOptions));
}
diff --git a/mlir/test/Dialect/Arith/expand-ops-min-max.mlir b/mlir/test/Dialect/Arith/expand-ops-min-max.mlir
new file mode 100644
index 0000000000000..92d6c8ad773d5
--- /dev/null
+++ b/mlir/test/Dialect/Arith/expand-ops-min-max.mlir
@@ -0,0 +1,64 @@
+// Default (include-min-max=true): the min/max ops are expanded into
+// cmpf/cmpi + select sequences.
+// RUN: mlir-opt %s -arith-expand -split-input-file | FileCheck %s --check-prefix=EXPAND
+
+// include-min-max=false: the min/max ops are left untouched so that a later
+// arith-to-llvm conversion can lower them to the single-instruction
+// llvm.intr.maximum/minimum/... intrinsics.
+// RUN: mlir-opt %s -arith-expand="include-min-max=false" -split-input-file | FileCheck %s --check-prefix=KEEP
+
+// EXPAND-LABEL: func @maximumf
+// KEEP-LABEL: func @maximumf
+func.func @maximumf(%a: f32, %b: f32) -> f32 {
+ // EXPAND: arith.cmpf ugt
+ // EXPAND: arith.select
+ // EXPAND: arith.cmpf uno
+ // EXPAND: arith.select
+ // EXPAND-NOT: arith.maximumf
+ // KEEP: arith.maximumf
+ // KEEP-NOT: arith.select
+ %result = arith.maximumf %a, %b : f32
+ return %result : f32
+}
+
+// -----
+
+// EXPAND-LABEL: func @minnumf
+// KEEP-LABEL: func @minnumf
+func.func @minnumf(%a: f32, %b: f32) -> f32 {
+ // EXPAND: arith.cmpf ult
+ // EXPAND: arith.select
+ // EXPAND-NOT: arith.minnumf
+ // KEEP: arith.minnumf
+ // KEEP-NOT: arith.select
+ %result = arith.minnumf %a, %b : f32
+ return %result : f32
+}
+
+// -----
+
+// EXPAND-LABEL: func @maxsi
+// KEEP-LABEL: func @maxsi
+func.func @maxsi(%a: i32, %b: i32) -> i32 {
+ // EXPAND: arith.cmpi sgt
+ // EXPAND: arith.select
+ // EXPAND-NOT: arith.maxsi
+ // KEEP: arith.maxsi
+ // KEEP-NOT: arith.select
+ %result = arith.maxsi %a, %b : i32
+ return %result : i32
+}
+
+// -----
+
+// Even with min/max expansion disabled, the other expansions (here
+// ceildivsi) still run.
+
+// EXPAND-LABEL: func @ceildivi_still_expands
+// KEEP-LABEL: func @ceildivi_still_expands
+func.func @ceildivi_still_expands(%a: i32, %b: i32) -> i32 {
+ // EXPAND-NOT: arith.ceildivsi
+ // KEEP-NOT: arith.ceildivsi
+ %result = arith.ceildivsi %a, %b : i32
+ return %result : i32
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/211884
More information about the Mlir-commits
mailing list