[InstCombine] InstCombine should fold frexp of select to select of frexp

Author: vortex73

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

@@ -4043,6 +4043,52 @@ InstCombinerImpl::foldExtractOfOverflowIntrinsic(ExtractValueInst &EV) {
   return nullptr;
 }
 
+static Value *foldFrexpOfSelect(ExtractValueInst &EV, CallInst *FrexpCall,
+                                SelectInst *SelectInst,
+                                InstCombiner::BuilderTy &Builder) {
+  // Helper to fold frexp of select to select of frexp.
+  Value *Cond = SelectInst->getCondition();
+  Value *TrueVal = SelectInst->getTrueValue();
+  Value *FalseVal = SelectInst->getFalseValue();
+  ConstantFP *ConstOp = nullptr;
+  Value *VarOp = nullptr;
+  bool ConstIsTrue = false;
+
+  if (auto *TrueConst = dyn_cast<ConstantFP>(TrueVal)) {
+    ConstOp = TrueConst;
+    VarOp = FalseVal;
+    ConstIsTrue = true;
+  } else if (auto *FalseConst = dyn_cast<ConstantFP>(FalseVal)) {
+    ConstOp = FalseConst;
+    VarOp = TrueVal;
+    ConstIsTrue = false;
+  }
+
+  if (!ConstOp || !VarOp)
+    return nullptr;
+
+  CallInst *NewFrexp =
+      Builder.CreateCall(FrexpCall->getCalledFunction(), {VarOp}, "frexp");
+
+  Value *NewEV = Builder.CreateExtractValue(NewFrexp, 0, "mantissa");
+
+  APFloat ConstVal = ConstOp->getValueAPF();
+  int Exp = 0;
+  APFloat Mantissa = ConstVal;
+
+  if (ConstVal.isFiniteNonZero()) {
+    Mantissa = frexp(ConstVal, Exp, APFloat::rmNearestTiesToEven);
+  }
+
+  Constant *ConstantMantissa = ConstantFP::get(ConstOp->getType(), Mantissa);
+
+  Value *NewSel = Builder.CreateSelect(
+      Cond, ConstIsTrue ? ConstantMantissa : NewEV,
+      ConstIsTrue ? NewEV : ConstantMantissa, "select.frexp");
+
+  return NewSel;
+}
+
 Instruction *InstCombinerImpl::visitExtractValueInst(ExtractValueInst &EV) {
   Value *Agg = EV.getAggregateOperand();
 
@@ -4052,7 +4098,28 @@ Instruction *InstCombinerImpl::visitExtractValueInst(ExtractValueInst &EV) {
   if (Value *V = simplifyExtractValueInst(Agg, EV.getIndices(),
                                           SQ.getWithInstruction(&EV)))
     return replaceInstUsesWith(EV, V);
-
+  if (EV.getNumIndices() == 1 && EV.getIndices()[0] == 0) {
+    if (auto *FrexpCall = dyn_cast<CallInst>(Agg)) {
+      if (Function *F = FrexpCall->getCalledFunction()) {
+        if (F->getIntrinsicID() == Intrinsic::frexp) {
+          if (auto *SelInst =
+                  dyn_cast<SelectInst>(FrexpCall->getArgOperand(0))) {
+            if (!SelInst->hasOneUse())
+              return nullptr;
+            if (isa<ConstantFP>(SelInst->getTrueValue()) ||
+                isa<ConstantFP>(SelInst->getFalseValue())) {
+              Builder.SetInsertPoint(&EV);
+
+              if (Value *Result =
+                      foldFrexpOfSelect(EV, FrexpCall, SelInst, Builder)) {
+                return replaceInstUsesWith(EV, Result);
+              }
+            }
+          }
+        }
+      }
+    }
+  }
   if (InsertValueInst *IV = dyn_cast<InsertValueInst>(Agg)) {
     // We're extracting from an insertvalue instruction, compare the indices
     const unsigned *exti, *exte, *insi, *inse;

llvm/test/Transforms/InstCombine/select_frexp.ll

@@ -8,10 +8,10 @@ declare void @use(float)
 define float @test_select_frexp_basic(float %x, i1 %cond) {
 ; CHECK-LABEL: define float @test_select_frexp_basic(
 ; CHECK-SAME: float [[X:%.*]], i1 [[COND:%.*]]) {
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND]], float 1.000000e+00, float [[X]]
-; CHECK-NEXT:    [[FREXP:%.*]] = call { float, i32 } @llvm.frexp.f32.i32(float [[SEL]])
+; CHECK-NEXT:    [[FREXP:%.*]] = call { float, i32 } @llvm.frexp.f32.i32(float [[X]])
 ; CHECK-NEXT:    [[FREXP_0:%.*]] = extractvalue { float, i32 } [[FREXP]], 0
-; CHECK-NEXT:    ret float [[FREXP_0]]
+; CHECK-NEXT:    [[SELECT_FREXP:%.*]] = select i1 [[COND]], float 5.000000e-01, float [[FREXP_0]]
+; CHECK-NEXT:    ret float [[SELECT_FREXP]]
 ;
   %sel = select i1 %cond, float 1.000000e+00, float %x
   %frexp = call { float, i32 } @llvm.frexp.f32.i32(float %sel)
@@ -23,10 +23,10 @@ define float @test_select_frexp_basic(float %x, i1 %cond) {
 define float @test_select_frexp_const_false(float %x, i1 %cond) {
 ; CHECK-LABEL: define float @test_select_frexp_const_false(
 ; CHECK-SAME: float [[X:%.*]], i1 [[COND:%.*]]) {
-; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND]], float [[X]], float 1.000000e+00
-; CHECK-NEXT:    [[FREXP:%.*]] = call { float, i32 } @llvm.frexp.f32.i32(float [[SEL]])
+; CHECK-NEXT:    [[FREXP:%.*]] = call { float, i32 } @llvm.frexp.f32.i32(float [[X]])
 ; CHECK-NEXT:    [[FREXP_0:%.*]] = extractvalue { float, i32 } [[FREXP]], 0
-; CHECK-NEXT:    ret float [[FREXP_0]]
+; CHECK-NEXT:    [[SELECT_FREXP:%.*]] = select i1 [[COND]], float [[FREXP_0]], float 5.000000e-01
+; CHECK-NEXT:    ret float [[SELECT_FREXP]]
 ;
   %sel = select i1 %cond, float %x, float 1.000000e+00
   %frexp = call { float, i32 } @llvm.frexp.f32.i32(float %sel)