From 0f3ed5ed4b99d34324f3a74e64cbe5e2b4215a1e Mon Sep 17 00:00:00 2001
From: Narayan Sreekumar <nsreekumar6@gmail.com>
Date: Sat, 28 Dec 2024 00:31:40 +0530
Subject: [PATCH] [InstCombine] InstCombine should fold frexp of select to
 select of frexp

---
 .../InstCombine/InstCombineSelect.cpp         | 57 +++++++++++++++++++
 1 file changed, 57 insertions(+)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
index 3d251d662bd53d..9bcf70c4122bac 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -3813,6 +3813,42 @@ static Value *foldSelectIntoAddConstant(SelectInst &SI,
   return nullptr;
 }
 
+static Instruction *foldFrexpOfSelect(ExtractValueInst *EV, CallInst *FrexpCall,
+                                      SelectInst *SelectInst) {
+  // A helper that folds frexp of select to select of frexp.
+  IRBuilder<> Builder(EV);
+
+  Value *Cond = SelectInst->getCondition();
+  Value *TrueVal = SelectInst->getTrueValue();
+  Value *FalseVal = SelectInst->getFalseValue();
+
+  ConstantFP *ConstOp = nullptr;
+  Value *VarOp = nullptr;
+  bool ConstIsTrue = false;
+
+  if ((ConstOp = dyn_cast<ConstantFP>(TrueVal))) {
+    VarOp = FalseVal;
+    ConstIsTrue = true;
+  } else {
+    ConstOp = dyn_cast<ConstantFP>(FalseVal);
+    VarOp = TrueVal;
+    ConstIsTrue = false;
+  }
+
+  CallInst *NewFrexp =
+      Builder.CreateCall(FrexpCall->getCalledFunction(), {VarOp}, "frexp");
+  Value *NewEV = Builder.CreateExtractValue(NewFrexp, 0);
+
+  APFloat ConstVal = ConstOp->getValueAPF();
+  int Exp;
+  APFloat Mantissa = frexp(ConstVal, Exp, APFloat::rmNearestTiesToEven);
+  Constant *ConstantMantissa = ConstantFP::get(ConstOp->getType(), Mantissa);
+  Value *NewSel =
+      Builder.CreateSelect(Cond, ConstIsTrue ? ConstantMantissa : NewEV,
+                           ConstIsTrue ? NewEV : ConstantMantissa);
+  return cast<Instruction>(NewSel);
+}
+
 Instruction *InstCombinerImpl::visitSelectInst(SelectInst &SI) {
   Value *CondVal = SI.getCondition();
   Value *TrueVal = SI.getTrueValue();
@@ -3829,6 +3865,27 @@ Instruction *InstCombinerImpl::visitSelectInst(SelectInst &SI) {
   if (Instruction *I = canonicalizeScalarSelectOfVecs(SI, *this))
     return I;
 
+  for (User *U : SI.users()) {
+    if (auto *EV = dyn_cast<ExtractValueInst>(U)) {
+      if (!EV->hasOneUse())
+        continue;
+      if (auto *FrexpCall = dyn_cast<CallInst>(EV->getAggregateOperand())) {
+        if (Function *F = FrexpCall->getCalledFunction()) {
+          if (F->getIntrinsicID() == Intrinsic::frexp &&
+              EV->getNumIndices() == 1 && EV->getIndices()[0] == 0) {
+            Value *SelectOp = FrexpCall->getArgOperand(0);
+            if (auto *SelInst = dyn_cast<SelectInst>(SelectOp)) {
+              if (isa<ConstantFP>(SelInst->getTrueValue()) ||
+                  isa<ConstantFP>(SelInst->getFalseValue())) {
+                return foldFrexpOfSelect(EV, FrexpCall, SelInst);
+              }
+            }
+          }
+        }
+      }
+    }
+  }
+
   // If the type of select is not an integer type or if the condition and
   // the selection type are not both scalar nor both vector types, there is no
   // point in attempting to match these patterns.