Skip to content

Commit 5be4fc2

Browse files
authored
[mlir][Arith] arith.select doesn't need to be emulated for small floats (#161707)
arith.select isn't an arithmetic operation in the sense of things like addf or mulf, which the emulate-unsupported-floats rewrites using extf and truncf. This patch adds select as a legal operation to prevent a pointless conversion aronud conditional moves. Fixes iree-org/iree#22181
1 parent 68b143d commit 5be4fc2

File tree

2 files changed

+13
-1
lines changed

2 files changed

+13
-1
lines changed

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -123,7 +123,8 @@ void mlir::arith::populateEmulateUnsupportedFloatsLegality(
123123
vector::OuterProductOp, vector::ScanOp>(
124124
[&](Operation *op) { return converter.isLegal(op); });
125125
target.addLegalOp<arith::BitcastOp, arith::ExtFOp, arith::TruncFOp,
126-
arith::ConstantOp, vector::SplatOp, vector::BroadcastOp>();
126+
arith::ConstantOp, arith::SelectOp, vector::SplatOp,
127+
vector::BroadcastOp>();
127128
}
128129

129130
void EmulateUnsupportedFloatsPass::runOnOperation() {

mlir/test/Dialect/Arith/emulate-unsupported-floats.mlir

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,3 +85,14 @@ func.func @no_expansion(%x: f32) -> f32 {
8585
%y = arith.addf %x, %c : f32
8686
func.return %y : f32
8787
}
88+
89+
// -----
90+
91+
func.func @no_promote_select(%c: i1, %x: bf16, %y: bf16) -> bf16 {
92+
// CHECK-LABEL: @no_promote_select
93+
// CHECK-SAME: (%[[C:.+]]: i1, %[[X:.+]]: bf16, %[[Y:.+]]: bf16)
94+
// CHECK: %[[Z:.+]] = arith.select %[[C]], %[[X]], %[[Y]] : bf16
95+
// CHECK: return %[[Z]]
96+
%z = arith.select %c, %x, %y : bf16
97+
func.return %z : bf16
98+
}

0 commit comments

Comments
 (0)