[VPlan] Support to widen call and select intructions in VPlan native path

[Kazhuu]

Add support to widen call and select instructions in VPlan native path by using correct recipes when such instruction is encountered. This is already used by inner loop vectorizer. I modified my previously committed VPlan test to also test for the call instruction widening.

Previously call and select instructions get handled by the wrong recipe and resulted in unreachable instruction errors like this one: https://bugs.llvm.org/show_bug.cgi?id=48139.

[Kazhuu]

Accidentally duplicated this with arcanist. Original diff here: https://reviews.llvm.org/D97136

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

if (EnableVPlanPredication) {
 
     // Avoid running transformation to recipes until masked code generation in
     // VPlan-native path is in place.
     return Plan;
   }
 
   SmallPtrSet<Instruction *, 1> DeadInstructions;
   VPlanTransforms::VPInstructionsToVPRecipes(
-      OrigLoop, Plan, Legal->getInductionVars(), DeadInstructions);
+      PSE, OrigLoop, Plan, Legal->getInductionVars(), DeadInstructions);
   return Plan;
 }
 
 // Adjust the recipes for any inloop reductions. The chain of instructions
 // leading from the loop exit instr to the phi need to be converted to
 // reductions, with one operand being vector and the other being the scalar
 // reduction chain.
 void LoopVectorizationPlanner::adjustRecipesForInLoopReductions(

llvm/lib/Transforms/Vectorize/VPlanTransforms.h

 #include "llvm/Transforms/Vectorize/LoopVectorizationLegality.h"
 
 namespace llvm {
 
 struct VPlanTransforms {
   /// Replaces the VPInstructions in \p Plan with corresponding
   /// widen recipes.
   static void VPInstructionsToVPRecipes(
-      Loop *OrigLoop, VPlanPtr &Plan,
+      PredicatedScalarEvolution &SPE, Loop *OrigLoop, VPlanPtr &Plan,
       LoopVectorizationLegality::InductionList &Inductions,
       SmallPtrSetImpl<Instruction *> &DeadInstructions);
 };
 
 } // namespace llvm
 
 #endif // LLVM_TRANSFORMS_VECTORIZE_VPLANTRANSFORMS_H

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

 //===----------------------------------------------------------------------===//
 
 #include "VPlanTransforms.h"
 #include "llvm/ADT/PostOrderIterator.h"
 
 using namespace llvm;
 
 void VPlanTransforms::VPInstructionsToVPRecipes(
-    Loop *OrigLoop, VPlanPtr &Plan,
+    PredicatedScalarEvolution &PSE, Loop *OrigLoop, VPlanPtr &Plan,
     LoopVectorizationLegality::InductionList &Inductions,
     SmallPtrSetImpl<Instruction *> &DeadInstructions) {
 
   auto *TopRegion = cast<VPRegionBlock>(Plan->getEntry());
   ReversePostOrderTraversal<VPBlockBase *> RPOT(TopRegion->getEntry());
 
   for (VPBlockBase *Base : RPOT) {
     // Do not widen instructions in pre-header and exit blocks.
for (auto I = VPBB->begin(), E = VPBB->end(); I != E;) {
             *Store, Plan->getOrAddVPValue(getLoadStorePointerOperand(Inst)),
             Plan->getOrAddVPValue(Store->getValueOperand()), nullptr /*Mask*/);
       else if (PHINode *Phi = dyn_cast<PHINode>(Inst)) {
         InductionDescriptor II = Inductions.lookup(Phi);
         if (II.getKind() == InductionDescriptor::IK_IntInduction ||
             II.getKind() == InductionDescriptor::IK_FpInduction) {
           VPValue *Start = Plan->getOrAddVPValue(II.getStartValue());
           NewRecipe = new VPWidenIntOrFpInductionRecipe(Phi, Start, nullptr);
+
         } else
           NewRecipe = new VPWidenPHIRecipe(Phi);
       } else if (GetElementPtrInst *GEP = dyn_cast<GetElementPtrInst>(Inst)) {
         NewRecipe = new VPWidenGEPRecipe(
             GEP, Plan->mapToVPValues(GEP->operands()), OrigLoop);
+      } else if (CallInst *CI = dyn_cast<CallInst>(Inst)) {
+        NewRecipe =
+            new VPWidenCallRecipe(*CI, Plan->mapToVPValues(CI->arg_operands()));
+      } else if (SelectInst *SI = dyn_cast<SelectInst>(Inst)) {
+        bool InvariantCond = PSE.getSE()->isLoopInvariant(
+            PSE.getSCEV(SI->getOperand(0)), OrigLoop);
+        NewRecipe = new VPWidenSelectRecipe(
+            *SI, Plan->mapToVPValues(SI->operands()), InvariantCond);
       } else
         NewRecipe =
             new VPWidenRecipe(*Inst, Plan->mapToVPValues(Inst->operands()));
 
       NewRecipe->insertBefore(Ingredient);
       if (NewRecipe->getNumDefinedValues() == 1)
         VPInst->replaceAllUsesWith(NewRecipe->getVPValue());
       else

llvm/test/Transforms/LoopVectorize/vplan-vectorize-inner-loop-reduction.ll

 ; RUN: opt -loop-vectorize -force-vector-width=4 -enable-vplan-native-path -S %s | FileCheck %s
 
-; Vectorize explict marked outer loop using vplan native path. Inner loop
-; contains simple double add reduction. IR is compiled and modified by hand
-; from following C code:
-; void inner_loop_reduction(const double* restrict in_a, const double* restrict in_b, double* restrict out)
-; {
-;     #pragma clang loop vectorize(enable)
-;     for (int i = 0; i < 1000; ++i) {
-;         double a = in_a[i];
-;         double b = in_b[i];
-;         for (int j = 0; j < 10000; ++j) {
-;             a = a + b;
-;         }
-;         out[i] = a;
-;     }
-; }
+; Test that VPlan native path is able to vectorize explictly marked outer loop.
+; Inner loop contains simple double add reduction and outer loop contains
+; intrincis call.
+
+declare double @llvm.sqrt.f64(double %0)
 define void @inner_loop_reduction(double* noalias nocapture readonly %a.in, double* noalias nocapture readonly %b.in, double* noalias nocapture %c.out) {
 ; CHECK-LABEL: @inner_loop_reduction(
 
 ; CHECK: vector.body:
 ; CHECK-NEXT: %[[FOR1_INDEX:.*]] = phi i64 [ 0, %[[LABEL_PR:.*]] ], [ %{{.*}}, %[[LABEL_FOR1_LATCH:.*]] ]
 ; CHECK: %[[VEC_INDEX:.*]] = phi <4 x i64> [ <i64 0, i64 1, i64 2, i64 3>, %[[LABEL_PR]] ], [ %{{.*}}, %[[LABEL_FOR1_LATCH]] ]
 ; CHECK-NEXT: %[[A_PTR:.*]] = getelementptr inbounds double, double* %a.in, <4 x i64> %[[VEC_INDEX]]
 ; CHECK-NEXT: %[[MASKED_GATHER1:.*]] = call <4 x double> @llvm.masked.gather.v4f64.v4p0f64(<4 x double*> %[[A_PTR]], i32 8, <4 x i1> <i1 true, i1 true, i1 true, i1 true>, <4 x double> undef)
 ; CHECK-NEXT: %[[B_PTR:.*]] = getelementptr inbounds double, double* %b.in, <4 x i64> %[[VEC_INDEX]]
 ; CHECK-NEXT: %[[MASKED_GATHER2:.*]] = call <4 x double> @llvm.masked.gather.v4f64.v4p0f64(<4 x double*> %[[B_PTR]], i32 8, <4 x i1> <i1 true, i1 true, i1 true, i1 true>, <4 x double> undef)
+; CHECK-NEXT: %[[B_SQRT:.*]] = call <4 x double> @llvm.sqrt.v4f64(<4 x double> %[[MASKED_GATHER2]])
 ; CHECK-NEXT: br label %[[FOR2_HEADER:.*]]
 
 ; CHECK: [[FOR2_HEADER]]:
 ; CHECK-NEXT: %[[FOR2_INDEX:.*]] = phi <4 x i32> [ %[[FOR2_INDEX_NEXT:.*]], %[[FOR2_HEADER]] ], [ zeroinitializer, %vector.body ]
 ; CHECK-NEXT: %[[REDUCTION:.*]] = phi <4 x double> [ %[[REDUCTION_NEXT:.*]], %[[FOR2_HEADER]] ], [ %[[MASKED_GATHER1]], %vector.body ]
-; CHECK-NEXT: %[[REDUCTION_NEXT]] = fadd <4 x double> %[[MASKED_GATHER2]], %[[REDUCTION]]
+; CHECK-NEXT: %[[REDUCTION_NEXT]] = fadd <4 x double> %[[B_SQRT]], %[[REDUCTION]]
 ; CHECK-NEXT: %[[FOR2_INDEX_NEXT]] = add nuw nsw <4 x i32> %[[FOR2_INDEX]], <i32 1, i32 1, i32 1, i32 1>
 ; CHECK-NEXT: %[[VEC_PTR:.*]] = icmp eq <4 x i32> %[[FOR2_INDEX_NEXT]], <i32 10000, i32 10000, i32 10000, i32 10000>
 ; CHECK-NEXT: %[[EXIT_COND:.*]] = extractelement <4 x i1> %[[VEC_PTR]], i32 0
 ; CHECK-NEXT: br i1 %[[EXIT_COND]], label %[[FOR1_LATCH:.*]], label %{{.*}}
 
 ; CHECK: [[FOR1_LATCH]]:
 ; CHECK-NEXT: %[[REDUCTION:.*]] = phi <4 x double> [ %[[REDUCTION_NEXT]], %[[FOR2_HEADER]] ]
 ; CHECK-NEXT: %[[C_PTR:.*]] = getelementptr inbounds double, double* %c.out, <4 x i64> %[[VEC_INDEX]]
 ; CHECK-NEXT: call void @llvm.masked.scatter.v4f64.v4p0f64(<4 x double> %[[REDUCTION]], <4 x double*> %[[C_PTR]], i32 8, <4 x i1> <i1 true, i1 true, i1 true, i1 true>)
 ; CHECK-NEXT: %[[VEC_INDEX_NEXT:.*]] = add nuw nsw <4 x i64> %[[VEC_INDEX]], <i64 1, i64 1, i64 1, i64 1>
 ; CHECK-NEXT: %[[VEC_PTR:.*]] = icmp eq <4 x i64> %[[VEC_INDEX_NEXT]], <i64 1000, i64 1000, i64 1000, i64 1000>
 ; CHECK-NEXT: %{{.*}} = extractelement <4 x i1> %[[VEC_PTR]], i32 0
 ; CHECK-NEXT: %[[FOR1_INDEX_NEXT:.*]] = add i64 %[[FOR1_INDEX]], 4
 ; CHECK-NEXT: %{{.*}} = add <4 x i64> %[[VEC_INDEX]], <i64 4, i64 4, i64 4, i64 4>
 ; CHECK-NEXT: %[[EXIT_COND:.*]] = icmp eq i64 %[[FOR1_INDEX_NEXT]], 1000
 ; CHECK-NEXT: br i1 %[[EXIT_COND]], label %{{.*}}, label %vector.body
 
 entry:
   br label %for1.header
 
-for1.header:                                              ; preds = %entry
+for1.header:
   %indvar1 = phi i64 [ 0, %entry ], [ %indvar11, %for1.latch ]
   %a.ptr = getelementptr inbounds double, double* %a.in, i64 %indvar1
   %a = load double, double* %a.ptr, align 8
   %b.ptr = getelementptr inbounds double, double* %b.in, i64 %indvar1
   %b = load double, double* %b.ptr, align 8
+  %b.sqrt = call double @llvm.sqrt.f64(double %b)
   br label %for2.header
 
-for2.header:                                              ; preds = %for1.header, %for2.header
+for2.header:
   %indvar2 = phi i32 [ 0, %for1.header ], [ %indvar21, %for2.header ]
   %a.reduction = phi double [ %a, %for1.header ], [ %a.reduction1, %for2.header ]
-  %a.reduction1 = fadd double %b, %a.reduction
+  %a.reduction1 = fadd double %b.sqrt, %a.reduction
   %indvar21 = add nuw nsw i32 %indvar2, 1
   %for2.cond = icmp eq i32 %indvar21, 10000
   br i1 %for2.cond, label %for1.latch, label %for2.header
 
-for1.latch:                                               ; preds = %for2.header
+for1.latch:
   %c.ptr = getelementptr inbounds double, double* %c.out, i64 %indvar1
   store double %a.reduction1, double* %c.ptr, align 8
   %indvar11 = add nuw nsw i64 %indvar1, 1
   %for1.cond = icmp eq i64 %indvar11, 1000
   br i1 %for1.cond, label %exit, label %for1.header, !llvm.loop !0
 
-exit:                                                    ; preds = %for1.latch
+exit:
   ret void
 }
 
 !0 = distinct !{!0, !1}
 !1 = !{!"llvm.loop.vectorize.enable", i1 true}

llvm/unittests/Transforms/Vectorize/VPlanHCFGTest.cpp

subgraph cluster_N0 {
     ]
   }
 }
 )";
   EXPECT_EQ(ExpectedStr, FullDump);
 
   LoopVectorizationLegality::InductionList Inductions;
   SmallPtrSet<Instruction *, 1> DeadInstructions;
-  VPlanTransforms::VPInstructionsToVPRecipes(LI->getLoopFor(LoopHeader), Plan,
-                                             Inductions, DeadInstructions);
+  VPlanTransforms::VPInstructionsToVPRecipes(
+      *PSE, LI->getLoopFor(LoopHeader), Plan, Inductions, DeadInstructions);
 }
 
 TEST_F(VPlanHCFGTest, testVPInstructionToVPRecipesInner) {
   const char *ModuleString =
       "define void @f(i32* %A, i64 %N) {\n"
       "entry:\n"
       "  br label %for.body\n"
       "for.body:\n"
TEST_F(VPlanHCFGTest, testVPInstructionToVPRecipesInner) {
   Module &M = parseModule(ModuleString);
 
   Function *F = M.getFunction("f");
   BasicBlock *LoopHeader = F->getEntryBlock().getSingleSuccessor();
   auto Plan = buildHCFG(LoopHeader);
 
   LoopVectorizationLegality::InductionList Inductions;
   SmallPtrSet<Instruction *, 1> DeadInstructions;
-  VPlanTransforms::VPInstructionsToVPRecipes(LI->getLoopFor(LoopHeader), Plan,
-                                             Inductions, DeadInstructions);
+  VPlanTransforms::VPInstructionsToVPRecipes(
+      *PSE, LI->getLoopFor(LoopHeader), Plan, Inductions, DeadInstructions);
 
   VPBlockBase *Entry = Plan->getEntry()->getEntryBasicBlock();
   EXPECT_NE(nullptr, Entry->getSingleSuccessor());
   EXPECT_EQ(0u, Entry->getNumPredecessors());
   EXPECT_EQ(1u, Entry->getNumSuccessors());
 
   VPBasicBlock *VecBB = Entry->getSingleSuccessor()->getEntryBasicBlock();
   EXPECT_EQ(7u, VecBB->size());

llvm/unittests/Transforms/Vectorize/VPlanTestBase.h

 #ifndef LLVM_UNITTESTS_TRANSFORMS_VECTORIZE_VPLANTESTBASE_H
 #define LLVM_UNITTESTS_TRANSFORMS_VECTORIZE_VPLANTESTBASE_H
 
 #include "../lib/Transforms/Vectorize/VPlan.h"
 #include "../lib/Transforms/Vectorize/VPlanHCFGBuilder.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/BasicAliasAnalysis.h"
 #include "llvm/Analysis/LoopInfo.h"
+#include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/AsmParser/Parser.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/Support/SourceMgr.h"
 #include "gtest/gtest.h"
 
 namespace llvm {
 
 /// Helper class to create a module from an assembly string and VPlans for a
 /// given loop entry block.
 class VPlanTestBase : public testing::Test {
 protected:
+  TargetLibraryInfoImpl TLII;
+  TargetLibraryInfo TLI;
+  DataLayout DL;
+
   std::unique_ptr<LLVMContext> Ctx;
   std::unique_ptr<Module> M;
   std::unique_ptr<LoopInfo> LI;
   std::unique_ptr<DominatorTree> DT;
-
-  VPlanTestBase() : Ctx(new LLVMContext) {}
+  std::unique_ptr<AssumptionCache> AC;
+  std::unique_ptr<ScalarEvolution> SE;
+  std::unique_ptr<PredicatedScalarEvolution> PSE;
+
+  VPlanTestBase()
+      : TLII(), TLI(TLII),
+        DL("e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-"
+           "f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:"
+           "16:32:64-S128"),
+        Ctx(new LLVMContext) {}
 
   Module &parseModule(const char *ModuleString) {
     SMDiagnostic Err;
     M = parseAssemblyString(ModuleString, Err, *Ctx);
     EXPECT_TRUE(M);
     return *M;
   }
 
-  void doAnalysis(Function &F) {
+  void doAnalysis(Function &F, BasicBlock *LoopHeader) {
     DT.reset(new DominatorTree(F));
     LI.reset(new LoopInfo(*DT));
+    AC.reset(new AssumptionCache(F));
+    SE.reset(new ScalarEvolution(F, TLI, *AC, *DT, *LI));
+    PSE.reset(new PredicatedScalarEvolution(*SE, *LI->getLoopFor(LoopHeader)));
   }
 
   VPlanPtr buildHCFG(BasicBlock *LoopHeader) {
-    doAnalysis(*LoopHeader->getParent());
+    doAnalysis(*LoopHeader->getParent(), LoopHeader);
 
     auto Plan = std::make_unique<VPlan>();
     VPlanHCFGBuilder HCFGBuilder(LI->getLoopFor(LoopHeader), LI.get(), *Plan);
     HCFGBuilder.buildHierarchicalCFG();
     return Plan;
   }
 
   /// Build the VPlan plain CFG for the loop starting from \p LoopHeader.
   VPlanPtr buildPlainCFG(BasicBlock *LoopHeader) {
-    doAnalysis(*LoopHeader->getParent());
+    doAnalysis(*LoopHeader->getParent(), LoopHeader);
 
     auto Plan = std::make_unique<VPlan>();
     VPlanHCFGBuilder HCFGBuilder(LI->getLoopFor(LoopHeader), LI.get(), *Plan);
     VPRegionBlock *TopRegion = HCFGBuilder.buildPlainCFG();
     Plan->setEntry(TopRegion);
     return Plan;
   }
 };
 
 } // namespace llvm
 
 #endif // LLVM_UNITTESTS_TRANSFORMS_VECTORIZE_VPLANTESTBASE_H