VPlan: implement VPlan-level constant-folding

Introduce VPlanConstantFolder, a variation of ConstantFolder for VPlan,
and use it in VPBuilder to constant-fold when all the underlying IR
values passed into the API are constants.

Author: artagnon

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -7941,7 +7941,8 @@ DenseMap<const SCEV *, Value *> LoopVectorizationPlanner::executePlan(
                            OrigLoop->getHeader()->getContext());
   VPlanTransforms::materializeBroadcasts(BestVPlan);
   VPlanTransforms::optimizeForVFAndUF(BestVPlan, BestVF, BestUF, PSE);
-  VPlanTransforms::simplifyRecipes(BestVPlan, *Legal->getWidestInductionType());
+  VPlanTransforms::simplifyRecipes(BestVPlan, *Legal->getWidestInductionType(),
+                                   OrigLoop->getHeader()->getDataLayout());
   VPlanTransforms::narrowInterleaveGroups(
       BestVPlan, BestVF,
       TTI.getRegisterBitWidth(TargetTransformInfo::RGK_FixedWidthVector));
@@ -9232,7 +9233,7 @@ void LoopVectorizationPlanner::buildVPlansWithVPRecipes(ElementCount MinVF,
       if (!HasScalarVF)
         VPlanTransforms::runPass(VPlanTransforms::truncateToMinimalBitwidths,
                                  *Plan, CM.getMinimalBitwidths());
-      VPlanTransforms::optimize(*Plan);
+      VPlanTransforms::optimize(*Plan, OrigLoop->getHeader()->getDataLayout());
       // TODO: try to put it close to addActiveLaneMask().
       // Discard the plan if it is not EVL-compatible
       if (CM.foldTailWithEVL() && !HasScalarVF &&

llvm/lib/Transforms/Vectorize/VPlan.cpp

@@ -1548,7 +1548,7 @@ void LoopVectorizationPlanner::buildVPlans(ElementCount MinVF,
   for (ElementCount VF = MinVF; ElementCount::isKnownLT(VF, MaxVFTimes2);) {
     VFRange SubRange = {VF, MaxVFTimes2};
     if (auto Plan = tryToBuildVPlan(SubRange)) {
-      VPlanTransforms::optimize(*Plan);
+      VPlanTransforms::optimize(*Plan, OrigLoop->getHeader()->getDataLayout());
       // Update the name of the latch of the top-level vector loop region region
       // after optimizations which includes block folding.
       Plan->getVectorLoopRegion()->getExiting()->setName("vector.latch");

llvm/lib/Transforms/Vectorize/VPlanConstantFolder.h

@@ -0,0 +1,157 @@
+//===- VPlanConstantFolder.h - ConstantFolder for VPlan -------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TRANSFORMS_VECTORIZE_VPLANCONSTANTFOLDER_H
+#define LLVM_TRANSFORMS_VECTORIZE_VPLANCONSTANTFOLDER_H
+
+#include "VPlan.h"
+#include "VPlanValue.h"
+#include "llvm/Analysis/TargetFolder.h"
+
+namespace llvm {
+class VPConstantFolder {
+  TargetFolder Folder;
+  VPTypeAnalysis TypeInfo;
+
+  Constant *getIRConstant(VPValue *V) const {
+    if (!V->isLiveIn())
+      return nullptr;
+    return dyn_cast_if_present<Constant>(V->getLiveInIRValue());
+  }
+
+  Value *foldBinOp(Instruction::BinaryOps Opcode, VPValue *LHS,
+                   VPValue *RHS) const {
+    auto *LC = getIRConstant(LHS);
+    auto *RC = getIRConstant(RHS);
+    if (LC && RC)
+      return Folder.FoldBinOp(Opcode, LC, RC);
+    return nullptr;
+  }
+
+  Value *foldNot(VPValue *Op) const {
+    auto *C = getIRConstant(Op);
+    if (C)
+      return Folder.FoldBinOp(Instruction::BinaryOps::Xor, C,
+                              Constant::getAllOnesValue(C->getType()));
+    return nullptr;
+  }
+
+  Value *foldLogicalAnd(VPValue *LHS, VPValue *RHS) const {
+    auto *LC = getIRConstant(LHS);
+    auto *RC = getIRConstant(RHS);
+    if (LC && RC)
+      return Folder.FoldSelect(LC, RC,
+                               ConstantInt::getNullValue(RC->getType()));
+    return nullptr;
+  }
+
+  Value *foldSelect(VPValue *Cond, VPValue *TrueVal, VPValue *FalseVal) const {
+    auto *CC = getIRConstant(Cond);
+    auto *TV = getIRConstant(TrueVal);
+    auto *FV = getIRConstant(FalseVal);
+    if (CC && TV && FV)
+      return Folder.FoldSelect(CC, TV, FV);
+    return nullptr;
+  }
+
+  Value *foldCmp(CmpInst::Predicate Pred, VPValue *LHS, VPValue *RHS) const {
+    auto *LC = getIRConstant(LHS);
+    auto *RC = getIRConstant(RHS);
+    if (LC && RC)
+      return Folder.FoldCmp(Pred, LC, RC);
+    return nullptr;
+  }
+
+  Value *foldGEP(Type *Ty, VPValue *Base, ArrayRef<VPValue *> Offsets,
+                 GEPNoWrapFlags NW) const {
+    auto *BC = getIRConstant(Base);
+    if (!BC)
+      return nullptr;
+    SmallVector<Value *> IdxList;
+    for (auto *O : Offsets) {
+      if (auto *OffsetV = getIRConstant(O))
+        IdxList.emplace_back(OffsetV);
+      else
+        return nullptr;
+    }
+    return Folder.FoldGEP(Ty, BC, IdxList, NW);
+  }
+
+  Value *foldInsertElement(VPValue *Vec, VPValue *NewElt, VPValue *Idx) const {
+    auto *VC = getIRConstant(Vec);
+    auto *EC = getIRConstant(NewElt);
+    auto *IC = getIRConstant(Idx);
+    if (VC && EC && IC)
+      Folder.FoldInsertElement(VC, EC, IC);
+    return nullptr;
+  }
+
+  Value *foldExtractElement(VPValue *Vec, VPValue *Idx) const {
+    auto *VC = getIRConstant(Vec);
+    auto *IC = getIRConstant(Idx);
+    if (VC && IC)
+      Folder.FoldExtractElement(VC, IC);
+    return nullptr;
+  }
+
+  Value *foldCast(Instruction::CastOps Opcode, VPValue *Op,
+                  Type *DestTy) const {
+    auto *C = getIRConstant(Op);
+    if (C)
+      return Folder.FoldCast(Opcode, C, DestTy);
+    return nullptr;
+  }
+
+public:
+  VPConstantFolder(const DataLayout &DL, const VPTypeAnalysis &TypeInfo)
+      : Folder(DL), TypeInfo(TypeInfo) {}
+
+  Value *tryToConstantFold(VPRecipeBase &R, unsigned Opcode,
+                           ArrayRef<VPValue *> Ops) {
+    switch (Opcode) {
+    case Instruction::BinaryOps::Add:
+    case Instruction::BinaryOps::Sub:
+    case Instruction::BinaryOps::Mul:
+    case Instruction::BinaryOps::AShr:
+    case Instruction::BinaryOps::LShr:
+    case Instruction::BinaryOps::And:
+    case Instruction::BinaryOps::Or:
+    case Instruction::BinaryOps::Xor:
+      return foldBinOp(static_cast<Instruction::BinaryOps>(Opcode), Ops[0],
+                       Ops[1]);
+    case VPInstruction::LogicalAnd:
+      return foldLogicalAnd(Ops[0], Ops[1]);
+    case VPInstruction::Not:
+      return foldNot(Ops[0]);
+    case Instruction::Select:
+      return foldSelect(Ops[0], Ops[1], Ops[2]);
+    case Instruction::ICmp:
+    case Instruction::FCmp:
+      return foldCmp(cast<VPRecipeWithIRFlags>(R).getPredicate(), Ops[0],
+                     Ops[1]);
+    case Instruction::GetElementPtr:
+    case VPInstruction::PtrAdd:
+      return foldGEP(TypeInfo.inferScalarType(R.getVPSingleValue()), Ops[0],
+                     Ops.drop_front(),
+                     cast<VPRecipeWithIRFlags>(R).getGEPNoWrapFlags());
+    case Instruction::InsertElement:
+      return foldInsertElement(Ops[0], Ops[1], Ops[2]);
+    case Instruction::ExtractElement:
+      return foldExtractElement(Ops[0], Ops[1]);
+    case Instruction::CastOps::SExt:
+    case Instruction::CastOps::ZExt:
+    case Instruction::CastOps::Trunc:
+      return foldCast(static_cast<Instruction::CastOps>(Opcode), Ops[0],
+                      TypeInfo.inferScalarType(R.getVPSingleValue()));
+    }
+    return nullptr;
+  }
+};
+} // namespace llvm
+
+#endif

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

@@ -16,6 +16,7 @@
 #include "VPlan.h"
 #include "VPlanAnalysis.h"
 #include "VPlanCFG.h"
+#include "VPlanConstantFolder.h"
 #include "VPlanDominatorTree.h"
 #include "VPlanHelpers.h"
 #include "VPlanPatternMatch.h"
@@ -931,9 +932,21 @@ static void recursivelyDeleteDeadRecipes(VPValue *V) {
 }
 
 /// Try to simplify recipe \p R.
-static void simplifyRecipe(VPRecipeBase &R, VPTypeAnalysis &TypeInfo) {
+static void simplifyRecipe(VPRecipeBase &R, VPTypeAnalysis &TypeInfo,
+                           const DataLayout &DL) {
   using namespace llvm::VPlanPatternMatch;
 
+  // Constant folding.
+  VPConstantFolder Folder(DL, TypeInfo);
+  TypeSwitch<VPRecipeBase *, void>(&R)
+      .Case<VPInstruction, VPWidenRecipe, VPWidenCastRecipe, VPReplicateRecipe>(
+          [&](auto *I) {
+            VPlan *Plan = R.getParent()->getPlan();
+            ArrayRef<VPValue *> Ops(I->op_begin(), I->op_end());
+            if (Value *V = Folder.tryToConstantFold(R, I->getOpcode(), Ops))
+              R.getVPSingleValue()->replaceAllUsesWith(Plan->getOrAddLiveIn(V));
+          });
+
   // VPScalarIVSteps can only be simplified after unrolling. VPScalarIVSteps for
   // part 0 can be replaced by their start value, if only the first lane is
   // demanded.
@@ -1049,13 +1062,14 @@ static void simplifyRecipe(VPRecipeBase &R, VPTypeAnalysis &TypeInfo) {
   }
 }
 
-void VPlanTransforms::simplifyRecipes(VPlan &Plan, Type &CanonicalIVTy) {
+void VPlanTransforms::simplifyRecipes(VPlan &Plan, Type &CanonicalIVTy,
+                                      const DataLayout &DL) {
   ReversePostOrderTraversal<VPBlockDeepTraversalWrapper<VPBlockBase *>> RPOT(
       Plan.getEntry());
   VPTypeAnalysis TypeInfo(&CanonicalIVTy);
   for (VPBasicBlock *VPBB : VPBlockUtils::blocksOnly<VPBasicBlock>(RPOT)) {
     for (VPRecipeBase &R : make_early_inc_range(*VPBB)) {
-      simplifyRecipe(R, TypeInfo);
+      simplifyRecipe(R, TypeInfo, DL);
     }
   }
 }
@@ -1306,7 +1320,8 @@ static bool simplifyBranchConditionForVFAndUF(VPlan &Plan, ElementCount BestVF,
 
     VPBlockUtils::connectBlocks(Preheader, Header);
     VPBlockUtils::connectBlocks(ExitingVPBB, Exit);
-    VPlanTransforms::simplifyRecipes(Plan, *CanIVTy);
+    VPlanTransforms::simplifyRecipes(Plan, *CanIVTy,
+                                     PSE.getSE()->getDataLayout());
   } else {
     // The vector region contains header phis for which we cannot remove the
     // loop region yet.
@@ -1774,17 +1789,16 @@ static void removeBranchOnCondTrue(VPlan &Plan) {
     VPBB->back().eraseFromParent();
   }
 }
-
-void VPlanTransforms::optimize(VPlan &Plan) {
+void VPlanTransforms::optimize(VPlan &Plan, const DataLayout &DL) {
   runPass(removeRedundantCanonicalIVs, Plan);
   runPass(removeRedundantInductionCasts, Plan);
 
-  runPass(simplifyRecipes, Plan, *Plan.getCanonicalIV()->getScalarType());
+  runPass(simplifyRecipes, Plan, *Plan.getCanonicalIV()->getScalarType(), DL);
   runPass(simplifyBlends, Plan);
   runPass(removeDeadRecipes, Plan);
   runPass(legalizeAndOptimizeInductions, Plan);
   runPass(removeRedundantExpandSCEVRecipes, Plan);
-  runPass(simplifyRecipes, Plan, *Plan.getCanonicalIV()->getScalarType());
+  runPass(simplifyRecipes, Plan, *Plan.getCanonicalIV()->getScalarType(), DL);
   runPass(removeBranchOnCondTrue, Plan);
   runPass(removeDeadRecipes, Plan);
 

llvm/lib/Transforms/Vectorize/VPlanTransforms.h

@@ -105,7 +105,7 @@ struct VPlanTransforms {
   /// Apply VPlan-to-VPlan optimizations to \p Plan, including induction recipe
   /// optimizations, dead recipe removal, replicate region optimizations and
   /// block merging.
-  static void optimize(VPlan &Plan);
+  static void optimize(VPlan &Plan, const DataLayout &DL);
 
   /// Wrap predicated VPReplicateRecipes with a mask operand in an if-then
   /// region block and remove the mask operand. Optimize the created regions by
@@ -187,7 +187,8 @@ struct VPlanTransforms {
 
   /// Perform instcombine-like simplifications on recipes in \p Plan. Use \p
   /// CanonicalIVTy as type for all un-typed live-ins in VPTypeAnalysis.
-  static void simplifyRecipes(VPlan &Plan, Type &CanonicalIVTy);
+  static void simplifyRecipes(VPlan &Plan, Type &CanonicalIVTy,
+                              const DataLayout &DL);
 
   /// If there's a single exit block, optimize its phi recipes that use exiting
   /// IV values by feeding them precomputed end values instead, possibly taken