[VPlan] Support struct return types for widen intrinsics (NFC).

[arcbbb]

Splitting from #151300, vp_load_ff returns a struct type that cannot be widened by toVectorizedTy.
This patch adds isVectorIntrinsicWithStructReturnScalarAtField and widen each struct element type independently.

[llvmbot]

@llvm/pr-subscribers-llvm-analysis
@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-vectorizers

Author: Shih-Po Hung (arcbbb)

Changes

Splitting from #151300, vp_load_ff returns a struct type that cannot be widened by toVectorizedTy.
This patch adds isVectorIntrinsicWithStructReturnScalarAtField and widen each struct element type independently.

---

Full diff: https://github.com/llvm/llvm-project/pull/165218.diff

4 Files Affected:

• (modified) llvm/include/llvm/Analysis/VectorUtils.h (+5)
• (modified) llvm/lib/Analysis/VectorUtils.cpp (+14)
• (modified) llvm/lib/Transforms/Vectorize/LoopVectorize.cpp (+16-1)
• (modified) llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp (+24-3)

diff --git a/llvm/include/llvm/Analysis/VectorUtils.h b/llvm/include/llvm/Analysis/VectorUtils.h
index ce94906ee7c00..f0c064f5490a7 100644
--- a/llvm/include/llvm/Analysis/VectorUtils.h
+++ b/llvm/include/llvm/Analysis/VectorUtils.h
@@ -164,6 +164,11 @@ LLVM_ABI bool
 isVectorIntrinsicWithOverloadTypeAtArg(Intrinsic::ID ID, int OpdIdx,
                                        const TargetTransformInfo *TTI);
 
+/// Identifies if the vector form of the intrinsic that returns a struct has
+/// a scalar element at the struct element index \p RetIdx.
+LLVM_ABI bool isVectorIntrinsicWithStructReturnScalarAtField(Intrinsic::ID ID,
+                                                             int RetIdx);
+
 /// Identifies if the vector form of the intrinsic that returns a struct is
 /// overloaded at the struct element index \p RetIdx. /// \p TTI is used to
 /// consider target specific intrinsics, if no target specific intrinsics
diff --git a/llvm/lib/Analysis/VectorUtils.cpp b/llvm/lib/Analysis/VectorUtils.cpp
index 091d94843698c..1f2fc80501cba 100644
--- a/llvm/lib/Analysis/VectorUtils.cpp
+++ b/llvm/lib/Analysis/VectorUtils.cpp
@@ -217,6 +217,16 @@ bool llvm::isVectorIntrinsicWithOverloadTypeAtArg(
   }
 }
 
+bool llvm::isVectorIntrinsicWithStructReturnScalarAtField(Intrinsic::ID ID,
+                                                          int RetIdx) {
+  switch (ID) {
+  case Intrinsic::vp_load_ff:
+    return RetIdx == 1;
+  default:
+    return false;
+  }
+}
+
 bool llvm::isVectorIntrinsicWithStructReturnOverloadAtField(
     Intrinsic::ID ID, int RetIdx, const TargetTransformInfo *TTI) {
 
@@ -224,6 +234,10 @@ bool llvm::isVectorIntrinsicWithStructReturnOverloadAtField(
     return TTI->isTargetIntrinsicWithStructReturnOverloadAtField(ID, RetIdx);
 
   switch (ID) {
+  case Intrinsic::modf:
+  case Intrinsic::sincos:
+  case Intrinsic::sincospi:
+    return false;
   case Intrinsic::frexp:
     return RetIdx == 0 || RetIdx == 1;
   default:
diff --git a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
index facb0fabdf57e..93419100cea35 100644
--- a/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ b/llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -4100,7 +4100,22 @@ static bool willGenerateVectors(VPlan &Plan, ElementCount VF,
       Type *ScalarTy = TypeInfo.inferScalarType(ToCheck);
       if (!Visited.insert({ScalarTy}).second)
         continue;
-      Type *WideTy = toVectorizedTy(ScalarTy, VF);
+      Type *WideTy;
+      if (auto *WI = dyn_cast<VPWidenIntrinsicRecipe>(&R);
+          WI && ScalarTy->isStructTy()) {
+        auto *StructTy = cast<StructType>(ScalarTy);
+        SmallVector<Type *, 2> Tys;
+        for (unsigned I = 0, E = StructTy->getNumElements(); I != E; ++I) {
+          Type *ElementTy = StructTy->getStructElementType(I);
+          if (!isVectorIntrinsicWithStructReturnScalarAtField(
+                  WI->getVectorIntrinsicID(), I))
+            ElementTy = toVectorizedTy(ElementTy, VF);
+          Tys.push_back(ElementTy);
+        }
+        WideTy = StructType::create(Tys);
+      } else
+        WideTy = toVectorizedTy(ScalarTy, VF);
+
       if (any_of(getContainedTypes(WideTy), WillGenerateTargetVectors))
         return true;
     }
diff --git a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
index 931a5b7582c4e..7cb15914fb2ce 100644
--- a/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
+++ b/llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp
@@ -1738,7 +1738,16 @@ void VPWidenIntrinsicRecipe::execute(VPTransformState &State) {
 
   SmallVector<Type *, 2> TysForDecl;
   // Add return type if intrinsic is overloaded on it.
-  if (isVectorIntrinsicWithOverloadTypeAtArg(VectorIntrinsicID, -1, State.TTI))
+  if (ResultTy->isStructTy()) {
+    auto *StructTy = cast<StructType>(ResultTy);
+    for (unsigned I = 0, E = StructTy->getNumElements(); I != E; ++I) {
+      if (isVectorIntrinsicWithStructReturnOverloadAtField(VectorIntrinsicID, I,
+                                                           State.TTI))
+        TysForDecl.push_back(
+            toVectorizedTy(StructTy->getStructElementType(I), State.VF));
+    }
+  } else if (isVectorIntrinsicWithOverloadTypeAtArg(VectorIntrinsicID, -1,
+                                                    State.TTI))
     TysForDecl.push_back(VectorType::get(getResultType(), State.VF));
   SmallVector<Value *, 4> Args;
   for (const auto &I : enumerate(operands())) {
@@ -1802,8 +1811,20 @@ static InstructionCost getCostForIntrinsics(Intrinsic::ID ID,
     Arguments.push_back(V);
   }
 
-  Type *ScalarRetTy = Ctx.Types.inferScalarType(&R);
-  Type *RetTy = VF.isVector() ? toVectorizedTy(ScalarRetTy, VF) : ScalarRetTy;
+  Type *RetTy = Ctx.Types.inferScalarType(&R);
+
+  if (VF.isVector() && RetTy->isStructTy()) {
+    auto *StructTy = cast<StructType>(RetTy);
+    SmallVector<Type *> Tys;
+    for (unsigned I = 0, E = StructTy->getNumElements(); I != E; ++I) {
+      Type *ElementTy = StructTy->getStructElementType(I);
+      if (!isVectorIntrinsicWithStructReturnScalarAtField(ID, I))
+        ElementTy = toVectorizedTy(ElementTy, VF);
+      Tys.push_back(ElementTy);
+    }
+    RetTy = StructType::get(StructTy->getContext(), Tys);
+  } else if (VF.isVector())
+    RetTy = toVectorizedTy(RetTy, VF);
   SmallVector<Type *> ParamTys;
   for (const VPValue *Op : Operands) {
     ParamTys.push_back(VF.isVector()