[VP] IR expansion for fabs/fsqrt/fma/fmadd

Add basic handling for VP ops that can expand to FP intrinsics

Fixes #60464

Differential Revision: https://reviews.llvm.org/D149052

1 files changed, 44 insertions, 0 deletions

diff --git a/llvm/lib/CodeGen/ExpandVectorPredication.cpp b/llvm/lib/CodeGen/ExpandVectorPredication.cpp
index 5ee76ff..fdf0be0 100644
--- a/llvm/lib/CodeGen/ExpandVectorPredication.cpp
+++ b/llvm/lib/CodeGen/ExpandVectorPredication.cpp
@@ -171,6 +171,10 @@ struct CachingVPExpander {
   Value *expandPredicationInBinaryOperator(IRBuilder<> &Builder,
                                            VPIntrinsic &PI);
 
+  /// Lower this VP fp call to a unpredicated fp call.
+  Value *expandPredicationToFPCall(IRBuilder<> &Builder, VPIntrinsic &PI,
+                                   unsigned UnpredicatedIntrinsicID);
+
   /// Lower this VP reduction to a call to an unpredicated reduction intrinsic.
   Value *expandPredicationInReduction(IRBuilder<> &Builder,
                                       VPReductionIntrinsic &PI);
@@ -271,6 +275,38 @@ CachingVPExpander::expandPredicationInBinaryOperator(IRBuilder<> &Builder,
   return NewBinOp;
 }
 
+Value *CachingVPExpander::expandPredicationToFPCall(
+    IRBuilder<> &Builder, VPIntrinsic &VPI, unsigned UnpredicatedIntrinsicID) {
+  assert((maySpeculateLanes(VPI) || VPI.canIgnoreVectorLengthParam()) &&
+         "Implicitly dropping %evl in non-speculatable operator!");
+
+  switch (UnpredicatedIntrinsicID) {
+  case Intrinsic::fabs:
+  case Intrinsic::sqrt: {
+    Value *Op0 = VPI.getOperand(0);
+    Function *Fn = Intrinsic::getDeclaration(
+        VPI.getModule(), UnpredicatedIntrinsicID, {VPI.getType()});
+    Value *NewOp = Builder.CreateCall(Fn, {Op0}, VPI.getName());
+    replaceOperation(*NewOp, VPI);
+    return NewOp;
+  }
+  case Intrinsic::experimental_constrained_fma:
+  case Intrinsic::experimental_constrained_fmuladd: {
+    Value *Op0 = VPI.getOperand(0);
+    Value *Op1 = VPI.getOperand(1);
+    Value *Op2 = VPI.getOperand(2);
+    Function *Fn = Intrinsic::getDeclaration(
+        VPI.getModule(), UnpredicatedIntrinsicID, {VPI.getType()});
+    Value *NewOp =
+        Builder.CreateConstrainedFPCall(Fn, {Op0, Op1, Op2}, VPI.getName());
+    replaceOperation(*NewOp, VPI);
+    return NewOp;
+  }
+  }
+
+  return nullptr;
+}
+
 static Value *getNeutralReductionElement(const VPReductionIntrinsic &VPI,
                                          Type *EltTy) {
   bool Negative = false;
@@ -565,6 +601,10 @@ Value *CachingVPExpander::expandPredication(VPIntrinsic &VPI) {
   switch (VPI.getIntrinsicID()) {
   default:
     break;
+  case Intrinsic::vp_fabs:
+    return expandPredicationToFPCall(Builder, VPI, Intrinsic::fabs);
+  case Intrinsic::vp_sqrt:
+    return expandPredicationToFPCall(Builder, VPI, Intrinsic::sqrt);
   case Intrinsic::vp_load:
   case Intrinsic::vp_store:
   case Intrinsic::vp_gather:
@@ -572,6 +612,10 @@ Value *CachingVPExpander::expandPredication(VPIntrinsic &VPI) {
     return expandPredicationInMemoryIntrinsic(Builder, VPI);
   }
 
+  if (auto CID = VPI.getConstrainedIntrinsicID())
+    if (Value *Call = expandPredicationToFPCall(Builder, VPI, *CID))
+      return Call;
+
   return &VPI;
 }