Skip to content

Commit 55be370

Browse files
authored
[InstCombine] Fold frexp of select to select of frexp (#121227)
This patch implements an optimization to push select operations through frexp when one of the select operands is a constant. When we encounter: ``` define float @src(float %x, i1 %bool) { %select = select i1 %bool, float 1.000000e+00, float %x %frexp = tail call { float, i32 } @llvm.frexp.f32.i32(float %select) %frexp.0 = extractvalue { float, i32 } %frexp, 0 ret float %frexp.0 } ``` We transform it to: ``` define float @tgt(float %x, i1 %bool) { %frexp = tail call { float, i32 } @llvm.frexp.f32.i32(float %x) %frexp.0 = extractvalue { float, i32 } %frexp, 0 %select = select i1 %bool, float 5.000000e-01, float %frexp.0 ret float %select } ``` Fixes #92542
1 parent a325622 commit 55be370

File tree

2 files changed

+244
-0
lines changed

2 files changed

+244
-0
lines changed

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
//===----------------------------------------------------------------------===//
3434

3535
#include "InstCombineInternal.h"
36+
#include "llvm/ADT/APFloat.h"
3637
#include "llvm/ADT/APInt.h"
3738
#include "llvm/ADT/ArrayRef.h"
3839
#include "llvm/ADT/DenseMap.h"
@@ -4069,6 +4070,49 @@ InstCombinerImpl::foldExtractOfOverflowIntrinsic(ExtractValueInst &EV) {
40694070
return nullptr;
40704071
}
40714072

4073+
static Value *foldFrexpOfSelect(ExtractValueInst &EV, IntrinsicInst *FrexpCall,
4074+
SelectInst *SelectInst,
4075+
InstCombiner::BuilderTy &Builder) {
4076+
// Helper to fold frexp of select to select of frexp.
4077+
4078+
if (!SelectInst->hasOneUse() || !FrexpCall->hasOneUse())
4079+
return nullptr;
4080+
Value *Cond = SelectInst->getCondition();
4081+
Value *TrueVal = SelectInst->getTrueValue();
4082+
Value *FalseVal = SelectInst->getFalseValue();
4083+
4084+
const APFloat *ConstVal = nullptr;
4085+
Value *VarOp = nullptr;
4086+
bool ConstIsTrue = false;
4087+
4088+
if (match(TrueVal, m_APFloat(ConstVal))) {
4089+
VarOp = FalseVal;
4090+
ConstIsTrue = true;
4091+
} else if (match(FalseVal, m_APFloat(ConstVal))) {
4092+
VarOp = TrueVal;
4093+
ConstIsTrue = false;
4094+
} else {
4095+
return nullptr;
4096+
}
4097+
4098+
Builder.SetInsertPoint(&EV);
4099+
4100+
CallInst *NewFrexp =
4101+
Builder.CreateCall(FrexpCall->getCalledFunction(), {VarOp}, "frexp");
4102+
NewFrexp->copyIRFlags(FrexpCall);
4103+
4104+
Value *NewEV = Builder.CreateExtractValue(NewFrexp, 0, "mantissa");
4105+
4106+
int Exp;
4107+
APFloat Mantissa = frexp(*ConstVal, Exp, APFloat::rmNearestTiesToEven);
4108+
4109+
Constant *ConstantMantissa = ConstantFP::get(TrueVal->getType(), Mantissa);
4110+
4111+
Value *NewSel = Builder.CreateSelectFMF(
4112+
Cond, ConstIsTrue ? ConstantMantissa : NewEV,
4113+
ConstIsTrue ? NewEV : ConstantMantissa, SelectInst, "select.frexp");
4114+
return NewSel;
4115+
}
40724116
Instruction *InstCombinerImpl::visitExtractValueInst(ExtractValueInst &EV) {
40734117
Value *Agg = EV.getAggregateOperand();
40744118

@@ -4079,6 +4123,15 @@ Instruction *InstCombinerImpl::visitExtractValueInst(ExtractValueInst &EV) {
40794123
SQ.getWithInstruction(&EV)))
40804124
return replaceInstUsesWith(EV, V);
40814125

4126+
Value *Cond, *TrueVal, *FalseVal;
4127+
if (match(&EV, m_ExtractValue<0>(m_Intrinsic<Intrinsic::frexp>(m_Select(
4128+
m_Value(Cond), m_Value(TrueVal), m_Value(FalseVal)))))) {
4129+
auto *SelInst =
4130+
cast<SelectInst>(cast<IntrinsicInst>(Agg)->getArgOperand(0));
4131+
if (Value *Result =
4132+
foldFrexpOfSelect(EV, cast<IntrinsicInst>(Agg), SelInst, Builder))
4133+
return replaceInstUsesWith(EV, Result);
4134+
}
40824135
if (InsertValueInst *IV = dyn_cast<InsertValueInst>(Agg)) {
40834136
// We're extracting from an insertvalue instruction, compare the indices
40844137
const unsigned *exti, *exte, *insi, *inse;
Lines changed: 191 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,191 @@
1+
; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
2+
; RUN: opt -passes=instcombine -S < %s | FileCheck %s
3+
4+
declare { float, i32 } @llvm.frexp.f32.i32(float)
5+
declare void @use(float)
6+
7+
; Basic test case - constant in true position
8+
define float @test_select_frexp_basic(float %x, i1 %cond) {
9+
; CHECK-LABEL: define float @test_select_frexp_basic(
10+
; CHECK-SAME: float [[X:%.*]], i1 [[COND:%.*]]) {
11+
; CHECK-NEXT: [[FREXP:%.*]] = call { float, i32 } @llvm.frexp.f32.i32(float [[X]])
12+
; CHECK-NEXT: [[FREXP_0:%.*]] = extractvalue { float, i32 } [[FREXP]], 0
13+
; CHECK-NEXT: [[SELECT_FREXP:%.*]] = select i1 [[COND]], float 5.000000e-01, float [[FREXP_0]]
14+
; CHECK-NEXT: ret float [[SELECT_FREXP]]
15+
;
16+
%sel = select i1 %cond, float 1.000000e+00, float %x
17+
%frexp = call { float, i32 } @llvm.frexp.f32.i32(float %sel)
18+
%frexp.0 = extractvalue { float, i32 } %frexp, 0
19+
ret float %frexp.0
20+
}
21+
22+
; Test with constant in false position
23+
define float @test_select_frexp_const_false(float %x, i1 %cond) {
24+
; CHECK-LABEL: define float @test_select_frexp_const_false(
25+
; CHECK-SAME: float [[X:%.*]], i1 [[COND:%.*]]) {
26+
; CHECK-NEXT: [[FREXP:%.*]] = call { float, i32 } @llvm.frexp.f32.i32(float [[X]])
27+
; CHECK-NEXT: [[FREXP_0:%.*]] = extractvalue { float, i32 } [[FREXP]], 0
28+
; CHECK-NEXT: [[SELECT_FREXP:%.*]] = select i1 [[COND]], float [[FREXP_0]], float 5.000000e-01
29+
; CHECK-NEXT: ret float [[SELECT_FREXP]]
30+
;
31+
%sel = select i1 %cond, float %x, float 1.000000e+00
32+
%frexp = call { float, i32 } @llvm.frexp.f32.i32(float %sel)
33+
%frexp.0 = extractvalue { float, i32 } %frexp, 0
34+
ret float %frexp.0
35+
}
36+
37+
; Multi-use test
38+
define float @test_select_frexp_multi_use(float %x, i1 %cond) {
39+
; CHECK-LABEL: define float @test_select_frexp_multi_use(
40+
; CHECK-SAME: float [[X:%.*]], i1 [[COND:%.*]]) {
41+
; CHECK-NEXT: [[SEL:%.*]] = select i1 [[COND]], float 1.000000e+00, float [[X]]
42+
; CHECK-NEXT: call void @use(float [[SEL]])
43+
; CHECK-NEXT: [[FREXP:%.*]] = call { float, i32 } @llvm.frexp.f32.i32(float [[SEL]])
44+
; CHECK-NEXT: [[FREXP_0:%.*]] = extractvalue { float, i32 } [[FREXP]], 0
45+
; CHECK-NEXT: ret float [[FREXP_0]]
46+
;
47+
%sel = select i1 %cond, float 1.000000e+00, float %x
48+
call void @use(float %sel)
49+
%frexp = call { float, i32 } @llvm.frexp.f32.i32(float %sel)
50+
%frexp.0 = extractvalue { float, i32 } %frexp, 0
51+
ret float %frexp.0
52+
}
53+
54+
; Vector test - splat constant
55+
define <2 x float> @test_select_frexp_vec_splat(<2 x float> %x, <2 x i1> %cond) {
56+
; CHECK-LABEL: define <2 x float> @test_select_frexp_vec_splat(
57+
; CHECK-SAME: <2 x float> [[X:%.*]], <2 x i1> [[COND:%.*]]) {
58+
; CHECK-NEXT: [[FREXP:%.*]] = call { <2 x float>, <2 x i32> } @llvm.frexp.v2f32.v2i32(<2 x float> [[X]])
59+
; CHECK-NEXT: [[FREXP_0:%.*]] = extractvalue { <2 x float>, <2 x i32> } [[FREXP]], 0
60+
; CHECK-NEXT: [[SELECT_FREXP:%.*]] = select <2 x i1> [[COND]], <2 x float> splat (float 5.000000e-01), <2 x float> [[FREXP_0]]
61+
; CHECK-NEXT: ret <2 x float> [[SELECT_FREXP]]
62+
;
63+
%sel = select <2 x i1> %cond, <2 x float> <float 1.000000e+00, float 1.000000e+00>, <2 x float> %x
64+
%frexp = call { <2 x float>, <2 x i32> } @llvm.frexp.v2f32.v2i32(<2 x float> %sel)
65+
%frexp.0 = extractvalue { <2 x float>, <2 x i32> } %frexp, 0
66+
ret <2 x float> %frexp.0
67+
}
68+
69+
; Vector test with poison
70+
define <2 x float> @test_select_frexp_vec_poison(<2 x float> %x, <2 x i1> %cond) {
71+
; CHECK-LABEL: define <2 x float> @test_select_frexp_vec_poison(
72+
; CHECK-SAME: <2 x float> [[X:%.*]], <2 x i1> [[COND:%.*]]) {
73+
; CHECK-NEXT: [[SEL:%.*]] = select <2 x i1> [[COND]], <2 x float> <float 1.000000e+00, float poison>, <2 x float> [[X]]
74+
; CHECK-NEXT: [[FREXP:%.*]] = call { <2 x float>, <2 x i32> } @llvm.frexp.v2f32.v2i32(<2 x float> [[SEL]])
75+
; CHECK-NEXT: [[FREXP_0:%.*]] = extractvalue { <2 x float>, <2 x i32> } [[FREXP]], 0
76+
; CHECK-NEXT: ret <2 x float> [[FREXP_0]]
77+
;
78+
%sel = select <2 x i1> %cond, <2 x float> <float 1.000000e+00, float poison>, <2 x float> %x
79+
%frexp = call { <2 x float>, <2 x i32> } @llvm.frexp.v2f32.v2i32(<2 x float> %sel)
80+
%frexp.0 = extractvalue { <2 x float>, <2 x i32> } %frexp, 0
81+
ret <2 x float> %frexp.0
82+
}
83+
84+
; Vector test - non-splat (should not fold)
85+
define <2 x float> @test_select_frexp_vec_nonsplat(<2 x float> %x, <2 x i1> %cond) {
86+
; CHECK-LABEL: define <2 x float> @test_select_frexp_vec_nonsplat(
87+
; CHECK-SAME: <2 x float> [[X:%.*]], <2 x i1> [[COND:%.*]]) {
88+
; CHECK-NEXT: [[SEL:%.*]] = select <2 x i1> [[COND]], <2 x float> <float 1.000000e+00, float 2.000000e+00>, <2 x float> [[X]]
89+
; CHECK-NEXT: [[FREXP:%.*]] = call { <2 x float>, <2 x i32> } @llvm.frexp.v2f32.v2i32(<2 x float> [[SEL]])
90+
; CHECK-NEXT: [[FREXP_0:%.*]] = extractvalue { <2 x float>, <2 x i32> } [[FREXP]], 0
91+
; CHECK-NEXT: ret <2 x float> [[FREXP_0]]
92+
;
93+
%sel = select <2 x i1> %cond, <2 x float> <float 1.000000e+00, float 2.000000e+00>, <2 x float> %x
94+
%frexp = call { <2 x float>, <2 x i32> } @llvm.frexp.v2f32.v2i32(<2 x float> %sel)
95+
%frexp.0 = extractvalue { <2 x float>, <2 x i32> } %frexp, 0
96+
ret <2 x float> %frexp.0
97+
}
98+
99+
; Negative test - both operands non-constant
100+
define float @test_select_frexp_no_const(float %x, float %y, i1 %cond) {
101+
; CHECK-LABEL: define float @test_select_frexp_no_const(
102+
; CHECK-SAME: float [[X:%.*]], float [[Y:%.*]], i1 [[COND:%.*]]) {
103+
; CHECK-NEXT: [[SEL:%.*]] = select i1 [[COND]], float [[X]], float [[Y]]
104+
; CHECK-NEXT: [[FREXP:%.*]] = call { float, i32 } @llvm.frexp.f32.i32(float [[SEL]])
105+
; CHECK-NEXT: [[FREXP_0:%.*]] = extractvalue { float, i32 } [[FREXP]], 0
106+
; CHECK-NEXT: ret float [[FREXP_0]]
107+
;
108+
%sel = select i1 %cond, float %x, float %y
109+
%frexp = call { float, i32 } @llvm.frexp.f32.i32(float %sel)
110+
%frexp.0 = extractvalue { float, i32 } %frexp, 0
111+
ret float %frexp.0
112+
}
113+
114+
; Negative test - extracting exp instead of mantissa
115+
define i32 @test_select_frexp_extract_exp(float %x, i1 %cond) {
116+
; CHECK-LABEL: define i32 @test_select_frexp_extract_exp(
117+
; CHECK-SAME: float [[X:%.*]], i1 [[COND:%.*]]) {
118+
; CHECK-NEXT: [[SEL:%.*]] = select i1 [[COND]], float 1.000000e+00, float [[X]]
119+
; CHECK-NEXT: [[FREXP:%.*]] = call { float, i32 } @llvm.frexp.f32.i32(float [[SEL]])
120+
; CHECK-NEXT: [[FREXP_1:%.*]] = extractvalue { float, i32 } [[FREXP]], 1
121+
; CHECK-NEXT: ret i32 [[FREXP_1]]
122+
;
123+
%sel = select i1 %cond, float 1.000000e+00, float %x
124+
%frexp = call { float, i32 } @llvm.frexp.f32.i32(float %sel)
125+
%frexp.1 = extractvalue { float, i32 } %frexp, 1
126+
ret i32 %frexp.1
127+
}
128+
129+
; Test with fast math flags
130+
define float @test_select_frexp_fast_math_select(float %x, i1 %cond) {
131+
; CHECK-LABEL: define float @test_select_frexp_fast_math_select(
132+
; CHECK-SAME: float [[X:%.*]], i1 [[COND:%.*]]) {
133+
; CHECK-NEXT: [[FREXP1:%.*]] = call { float, i32 } @llvm.frexp.f32.i32(float [[X]])
134+
; CHECK-NEXT: [[MANTISSA:%.*]] = extractvalue { float, i32 } [[FREXP1]], 0
135+
; CHECK-NEXT: [[SELECT_FREXP:%.*]] = select nnan ninf nsz i1 [[COND]], float 5.000000e-01, float [[MANTISSA]]
136+
; CHECK-NEXT: ret float [[SELECT_FREXP]]
137+
;
138+
%sel = select nnan ninf nsz i1 %cond, float 1.000000e+00, float %x
139+
%frexp = call { float, i32 } @llvm.frexp.f32.i32(float %sel)
140+
%frexp.0 = extractvalue { float, i32 } %frexp, 0
141+
ret float %frexp.0
142+
}
143+
144+
145+
; Test vector case with fast math flags
146+
define <2 x float> @test_select_frexp_vec_fast_math(<2 x float> %x, <2 x i1> %cond) {
147+
; CHECK-LABEL: define <2 x float> @test_select_frexp_vec_fast_math(
148+
; CHECK-SAME: <2 x float> [[X:%.*]], <2 x i1> [[COND:%.*]]) {
149+
; CHECK-NEXT: [[FREXP1:%.*]] = call { <2 x float>, <2 x i32> } @llvm.frexp.v2f32.v2i32(<2 x float> [[X]])
150+
; CHECK-NEXT: [[MANTISSA:%.*]] = extractvalue { <2 x float>, <2 x i32> } [[FREXP1]], 0
151+
; CHECK-NEXT: [[SELECT_FREXP:%.*]] = select nnan ninf nsz <2 x i1> [[COND]], <2 x float> splat (float 5.000000e-01), <2 x float> [[MANTISSA]]
152+
; CHECK-NEXT: ret <2 x float> [[SELECT_FREXP]]
153+
;
154+
%sel = select nnan ninf nsz <2 x i1> %cond, <2 x float> <float 1.000000e+00, float 1.000000e+00>, <2 x float> %x
155+
%frexp = call { <2 x float>, <2 x i32> } @llvm.frexp.v2f32.v2i32(<2 x float> %sel)
156+
%frexp.0 = extractvalue { <2 x float>, <2 x i32> } %frexp, 0
157+
ret <2 x float> %frexp.0
158+
}
159+
160+
; Test with scalable vectors with constant at True Position
161+
define <vscale x 2 x float> @test_select_frexp_scalable_vec0(<vscale x 2 x float> %x, <vscale x 2 x i1> %cond) {
162+
; CHECK-LABEL: define <vscale x 2 x float> @test_select_frexp_scalable_vec0(
163+
; CHECK-SAME: <vscale x 2 x float> [[X:%.*]], <vscale x 2 x i1> [[COND:%.*]]) {
164+
; CHECK-NEXT: [[FREXP1:%.*]] = call { <vscale x 2 x float>, <vscale x 2 x i32> } @llvm.frexp.nxv2f32.nxv2i32(<vscale x 2 x float> [[X]])
165+
; CHECK-NEXT: [[MANTISSA:%.*]] = extractvalue { <vscale x 2 x float>, <vscale x 2 x i32> } [[FREXP1]], 0
166+
; CHECK-NEXT: [[SELECT_FREXP:%.*]] = select <vscale x 2 x i1> [[COND]], <vscale x 2 x float> splat (float 5.000000e-01), <vscale x 2 x float> [[MANTISSA]]
167+
; CHECK-NEXT: ret <vscale x 2 x float> [[SELECT_FREXP]]
168+
;
169+
%sel = select <vscale x 2 x i1> %cond, <vscale x 2 x float> splat (float 1.000000e+00), <vscale x 2 x float> %x
170+
%frexp = call { <vscale x 2 x float>, <vscale x 2 x i32> } @llvm.frexp.nxv2f32.nxv2i32(<vscale x 2 x float> %sel)
171+
%frexp.0 = extractvalue { <vscale x 2 x float>, <vscale x 2 x i32> } %frexp, 0
172+
ret <vscale x 2 x float> %frexp.0
173+
}
174+
175+
; Test with scalable vectors with constant at False Position
176+
define <vscale x 2 x float> @test_select_frexp_scalable_vec1(<vscale x 2 x float> %x, <vscale x 2 x i1> %cond) {
177+
; CHECK-LABEL: define <vscale x 2 x float> @test_select_frexp_scalable_vec1(
178+
; CHECK-SAME: <vscale x 2 x float> [[X:%.*]], <vscale x 2 x i1> [[COND:%.*]]) {
179+
; CHECK-NEXT: [[FREXP1:%.*]] = call { <vscale x 2 x float>, <vscale x 2 x i32> } @llvm.frexp.nxv2f32.nxv2i32(<vscale x 2 x float> [[X]])
180+
; CHECK-NEXT: [[MANTISSA:%.*]] = extractvalue { <vscale x 2 x float>, <vscale x 2 x i32> } [[FREXP1]], 0
181+
; CHECK-NEXT: [[SELECT_FREXP:%.*]] = select <vscale x 2 x i1> [[COND]], <vscale x 2 x float> [[MANTISSA]], <vscale x 2 x float> splat (float 5.000000e-01)
182+
; CHECK-NEXT: ret <vscale x 2 x float> [[SELECT_FREXP]]
183+
;
184+
%sel = select <vscale x 2 x i1> %cond, <vscale x 2 x float> %x, <vscale x 2 x float> splat (float 1.000000e+00)
185+
%frexp = call { <vscale x 2 x float>, <vscale x 2 x i32> } @llvm.frexp.nxv2f32.nxv2i32(<vscale x 2 x float> %sel)
186+
%frexp.0 = extractvalue { <vscale x 2 x float>, <vscale x 2 x i32> } %frexp, 0
187+
ret <vscale x 2 x float> %frexp.0
188+
}
189+
190+
declare { <2 x float>, <2 x i32> } @llvm.frexp.v2f32.v2i32(<2 x float>)
191+
declare { <vscale x 2 x float>, <vscale x 2 x i32> } @llvm.frexp.nxv2f32.nxv2i32(<vscale x 2 x float>)

0 commit comments

Comments
 (0)