[InstCombine] Refactor and Preserve fast math flags

Author: vortex73

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

@@ -33,6 +33,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "InstCombineInternal.h"
+#include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
@@ -4043,52 +4044,57 @@ InstCombinerImpl::foldExtractOfOverflowIntrinsic(ExtractValueInst &EV) {
   return nullptr;
 }
 
-static Value *foldFrexpOfSelect(ExtractValueInst &EV, CallInst *FrexpCall,
+static Value *foldFrexpOfSelect(ExtractValueInst &EV, IntrinsicInst *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;
+
+  const APFloat *ConstVal = nullptr;
   Value *VarOp = nullptr;
   bool ConstIsTrue = false;
 
-  if (auto *TrueConst = dyn_cast<ConstantFP>(TrueVal)) {
-    ConstOp = TrueConst;
+  if (match(TrueVal, m_APFloat(ConstVal))) {
     VarOp = FalseVal;
     ConstIsTrue = true;
-  } else if (auto *FalseConst = dyn_cast<ConstantFP>(FalseVal)) {
-    ConstOp = FalseConst;
+  } else if (match(FalseVal, m_APFloat(ConstVal))) {
     VarOp = TrueVal;
     ConstIsTrue = false;
+  } else {
+    return nullptr;
   }
 
-  if (!ConstOp || !VarOp)
-    return nullptr;
+  Builder.SetInsertPoint(&EV);
 
   CallInst *NewFrexp =
       Builder.CreateCall(FrexpCall->getCalledFunction(), {VarOp}, "frexp");
+  NewFrexp->copyIRFlags(FrexpCall);
 
   Value *NewEV = Builder.CreateExtractValue(NewFrexp, 0, "mantissa");
 
-  APFloat ConstVal = ConstOp->getValueAPF();
-  int Exp = 0;
-  APFloat Mantissa = ConstVal;
+  int Exp;
+  APFloat Mantissa = frexp(*ConstVal, Exp, APFloat::rmNearestTiesToEven);
 
-  if (ConstVal.isFiniteNonZero()) {
-    Mantissa = frexp(ConstVal, Exp, APFloat::rmNearestTiesToEven);
+  Constant *ConstantMantissa;
+  if (auto *VecTy = dyn_cast<VectorType>(TrueVal->getType())) {
+    SmallVector<Constant *, 4> Elems(
+        VecTy->getElementCount().getFixedValue(),
+        ConstantFP::get(VecTy->getElementType(), Mantissa));
+    ConstantMantissa = ConstantVector::get(Elems);
+  } else {
+    ConstantMantissa = ConstantFP::get(TrueVal->getType(), Mantissa);
   }
 
-  Constant *ConstantMantissa = ConstantFP::get(ConstOp->getType(), Mantissa);
-
   Value *NewSel = Builder.CreateSelect(
       Cond, ConstIsTrue ? ConstantMantissa : NewEV,
       ConstIsTrue ? NewEV : ConstantMantissa, "select.frexp");
+  if (auto *NewSelInst = dyn_cast<Instruction>(NewSel))
+    NewSelInst->copyFastMathFlags(SelectInst);
 
   return NewSel;
 }
-
 Instruction *InstCombinerImpl::visitExtractValueInst(ExtractValueInst &EV) {
   Value *Agg = EV.getAggregateOperand();
 
@@ -4099,20 +4105,12 @@ Instruction *InstCombinerImpl::visitExtractValueInst(ExtractValueInst &EV) {
                                           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 (isa<ConstantFP>(SelInst->getTrueValue()) ||
-                isa<ConstantFP>(SelInst->getFalseValue())) {
-              Builder.SetInsertPoint(&EV);
-
-              if (Value *Result =
-                      foldFrexpOfSelect(EV, FrexpCall, SelInst, Builder)) {
-                return replaceInstUsesWith(EV, Result);
-              }
-            }
+    if (auto *FrexpCall = dyn_cast<IntrinsicInst>(Agg)) {
+      if (FrexpCall->getIntrinsicID() == Intrinsic::frexp) {
+        if (auto *SelInst = dyn_cast<SelectInst>(FrexpCall->getArgOperand(0))) {
+          if (Value *Result =
+                  foldFrexpOfSelect(EV, FrexpCall, SelInst, Builder)) {
+            return replaceInstUsesWith(EV, Result);
           }
         }
       }

llvm/test/Transforms/InstCombine/select_frexp.ll

@@ -56,10 +56,10 @@ define float @test_select_frexp_multi_use(float %x, i1 %cond) {
 define <2 x float> @test_select_frexp_vec_splat(<2 x float> %x, <2 x i1> %cond) {
 ; CHECK-LABEL: define <2 x float> @test_select_frexp_vec_splat(
 ; CHECK-SAME: <2 x float> [[X:%.*]], <2 x i1> [[COND:%.*]]) {
-; CHECK-NEXT:    [[SEL:%.*]] = select <2 x i1> [[COND]], <2 x float> splat (float 1.000000e+00), <2 x float> [[X]]
-; CHECK-NEXT:    [[FREXP:%.*]] = call { <2 x float>, <2 x i32> } @llvm.frexp.v2f32.v2i32(<2 x float> [[SEL]])
+; CHECK-NEXT:    [[FREXP:%.*]] = call { <2 x float>, <2 x i32> } @llvm.frexp.v2f32.v2i32(<2 x float> [[X]])
 ; CHECK-NEXT:    [[FREXP_0:%.*]] = extractvalue { <2 x float>, <2 x i32> } [[FREXP]], 0
-; CHECK-NEXT:    ret <2 x float> [[FREXP_0]]
+; CHECK-NEXT:    [[SELECT_FREXP:%.*]] = select <2 x i1> [[COND]], <2 x float> splat (float 5.000000e-01), <2 x float> [[FREXP_0]]
+; CHECK-NEXT:    ret <2 x float> [[SELECT_FREXP]]
 ;
   %sel = select <2 x i1> %cond, <2 x float> <float 1.000000e+00, float 1.000000e+00>, <2 x float> %x
   %frexp = call { <2 x float>, <2 x i32> } @llvm.frexp.v2f32.v2i32(<2 x float> %sel)
@@ -127,4 +127,35 @@ define i32 @test_select_frexp_extract_exp(float %x, i1 %cond) {
   ret i32 %frexp.1
 }
 
+; Test with fast math flags
+define float @test_select_frexp_fast_math_select(float %x, i1 %cond) {
+; CHECK-LABEL: define float @test_select_frexp_fast_math_select(
+; CHECK-SAME: float [[X:%.*]], i1 [[COND:%.*]]) {
+; CHECK-NEXT:    [[FREXP1:%.*]] = call { float, i32 } @llvm.frexp.f32.i32(float [[X]])
+; CHECK-NEXT:    [[MANTISSA:%.*]] = extractvalue { float, i32 } [[FREXP1]], 0
+; CHECK-NEXT:    [[SELECT_FREXP:%.*]] = select nnan ninf nsz i1 [[COND]], float 5.000000e-01, float [[MANTISSA]]
+; CHECK-NEXT:    ret float [[SELECT_FREXP]]
+;
+  %sel = select nnan ninf nsz i1 %cond, float 1.000000e+00, float %x
+  %frexp = call { float, i32 } @llvm.frexp.f32.i32(float %sel)
+  %frexp.0 = extractvalue { float, i32 } %frexp, 0
+  ret float %frexp.0
+}
+
+
+; Test vector case with fast math flags
+define <2 x float> @test_select_frexp_vec_fast_math(<2 x float> %x, <2 x i1> %cond) {
+; CHECK-LABEL: define <2 x float> @test_select_frexp_vec_fast_math(
+; CHECK-SAME: <2 x float> [[X:%.*]], <2 x i1> [[COND:%.*]]) {
+; CHECK-NEXT:    [[FREXP1:%.*]] = call { <2 x float>, <2 x i32> } @llvm.frexp.v2f32.v2i32(<2 x float> [[X]])
+; CHECK-NEXT:    [[MANTISSA:%.*]] = extractvalue { <2 x float>, <2 x i32> } [[FREXP1]], 0
+; CHECK-NEXT:    [[SELECT_FREXP:%.*]] = select nnan ninf nsz <2 x i1> [[COND]], <2 x float> splat (float 5.000000e-01), <2 x float> [[MANTISSA]]
+; CHECK-NEXT:    ret <2 x float> [[SELECT_FREXP]]
+;
+  %sel = select nnan ninf nsz <2 x i1> %cond, <2 x float> <float 1.000000e+00, float 1.000000e+00>, <2 x float> %x
+  %frexp = call { <2 x float>, <2 x i32> } @llvm.frexp.v2f32.v2i32(<2 x float> %sel)
+  %frexp.0 = extractvalue { <2 x float>, <2 x i32> } %frexp, 0
+  ret <2 x float> %frexp.0
+}
+
 declare { <2 x float>, <2 x i32> } @llvm.frexp.v2f32.v2i32(<2 x float>)