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/LoopVectorizationPlanner.h

@@ -25,7 +25,7 @@
 #define LLVM_TRANSFORMS_VECTORIZE_LOOPVECTORIZATIONPLANNER_H
 
 #include "VPlan.h"
-#include "llvm/ADT/SmallSet.h"
+#include "VPlanConstantFolder.h"
 #include "llvm/Support/InstructionCost.h"
 
 namespace llvm {
@@ -46,6 +46,7 @@ struct VFRange;
 class VPBuilder {
   VPBasicBlock *BB = nullptr;
   VPBasicBlock::iterator InsertPt = VPBasicBlock::iterator();
+  VPConstantFolder Folder;
 
   /// Insert \p VPI in BB at InsertPt if BB is set.
   template <typename T> T *tryInsertInstruction(T *R) {
@@ -66,6 +67,11 @@ class VPBuilder {
     return createInstruction(Opcode, ArrayRef<VPValue *>(Operands), DL, Name);
   }
 
+  VPValue *getOrAddLiveIn(Value *V) {
+    assert(BB && "Expected insertion point to be set");
+    return BB->getPlan()->getOrAddLiveIn(V);
+  }
+
 public:
   VPBuilder() = default;
   VPBuilder(VPBasicBlock *InsertBB) { setInsertPoint(InsertBB); }
@@ -180,31 +186,45 @@ class VPBuilder {
 
   VPValue *createNot(VPValue *Operand, DebugLoc DL = {},
                      const Twine &Name = "") {
+    if (auto *V = Folder.foldNot(Operand))
+      if (BB)
+        return getOrAddLiveIn(V);
     return createInstruction(VPInstruction::Not, {Operand}, DL, Name);
   }
 
   VPValue *createAnd(VPValue *LHS, VPValue *RHS, DebugLoc DL = {},
                      const Twine &Name = "") {
+    if (auto *V = Folder.foldAnd(LHS, RHS))
+      if (BB)
+        return getOrAddLiveIn(V);
     return createInstruction(Instruction::BinaryOps::And, {LHS, RHS}, DL, Name);
   }
 
   VPValue *createOr(VPValue *LHS, VPValue *RHS, DebugLoc DL = {},
                     const Twine &Name = "") {
-
+    if (auto *V = Folder.foldOr(LHS, RHS))
+      if (BB)
+        return getOrAddLiveIn(V);
     return tryInsertInstruction(new VPInstruction(
         Instruction::BinaryOps::Or, {LHS, RHS},
         VPRecipeWithIRFlags::DisjointFlagsTy(false), DL, Name));
   }
 
   VPValue *createLogicalAnd(VPValue *LHS, VPValue *RHS, DebugLoc DL = {},
                             const Twine &Name = "") {
+    if (auto *V = Folder.foldLogicalAnd(LHS, RHS))
+      if (BB)
+        return getOrAddLiveIn(V);
     return tryInsertInstruction(
         new VPInstruction(VPInstruction::LogicalAnd, {LHS, RHS}, DL, Name));
   }
 
   VPValue *createSelect(VPValue *Cond, VPValue *TrueVal, VPValue *FalseVal,
                         DebugLoc DL = {}, const Twine &Name = "",
                         std::optional<FastMathFlags> FMFs = std::nullopt) {
+    if (auto *V = Folder.foldSelect(Cond, TrueVal, FalseVal))
+      if (BB)
+        return getOrAddLiveIn(V);
     auto *Select =
         FMFs ? new VPInstruction(Instruction::Select, {Cond, TrueVal, FalseVal},
                                  *FMFs, DL, Name)
@@ -220,17 +240,26 @@ class VPBuilder {
                       DebugLoc DL = {}, const Twine &Name = "") {
     assert(Pred >= CmpInst::FIRST_ICMP_PREDICATE &&
            Pred <= CmpInst::LAST_ICMP_PREDICATE && "invalid predicate");
+    if (auto *V = Folder.foldCmp(Pred, A, B))
+      if (BB)
+        return getOrAddLiveIn(V);
     return tryInsertInstruction(
         new VPInstruction(Instruction::ICmp, Pred, A, B, DL, Name));
   }
 
-  VPInstruction *createPtrAdd(VPValue *Ptr, VPValue *Offset, DebugLoc DL = {},
-                              const Twine &Name = "") {
+  VPValue *createPtrAdd(VPValue *Ptr, VPValue *Offset, DebugLoc DL = {},
+                        const Twine &Name = "") {
+    if (auto *V = Folder.foldPtrAdd(Ptr, Offset, GEPNoWrapFlags::none()))
+      if (BB)
+        return getOrAddLiveIn(V);
     return tryInsertInstruction(
         new VPInstruction(Ptr, Offset, GEPNoWrapFlags::none(), DL, Name));
   }
   VPValue *createInBoundsPtrAdd(VPValue *Ptr, VPValue *Offset, DebugLoc DL = {},
                                 const Twine &Name = "") {
+    if (auto *V = Folder.foldPtrAdd(Ptr, Offset, GEPNoWrapFlags::inBounds()))
+      if (BB)
+        return getOrAddLiveIn(V);
     return tryInsertInstruction(
         new VPInstruction(Ptr, Offset, GEPNoWrapFlags::inBounds(), DL, Name));
   }
@@ -246,14 +275,20 @@ class VPBuilder {
         new VPDerivedIVRecipe(Kind, FPBinOp, Start, Current, Step, Name));
   }
 
-  VPScalarCastRecipe *createScalarCast(Instruction::CastOps Opcode, VPValue *Op,
-                                       Type *ResultTy, DebugLoc DL) {
+  VPValue *createScalarCast(Instruction::CastOps Opcode, VPValue *Op,
+                            Type *ResultTy, DebugLoc DL) {
+    if (auto *V = Folder.foldCast(Opcode, Op, ResultTy))
+      if (BB)
+        return getOrAddLiveIn(V);
     return tryInsertInstruction(
         new VPScalarCastRecipe(Opcode, Op, ResultTy, DL));
   }
 
-  VPWidenCastRecipe *createWidenCast(Instruction::CastOps Opcode, VPValue *Op,
-                                     Type *ResultTy) {
+  VPValue *createWidenCast(Instruction::CastOps Opcode, VPValue *Op,
+                           Type *ResultTy) {
+    if (auto *V = Folder.foldCast(Opcode, Op, ResultTy))
+      if (BB)
+        return getOrAddLiveIn(V);
     return tryInsertInstruction(new VPWidenCastRecipe(Opcode, Op, ResultTy));
   }
 

llvm/lib/Transforms/Vectorize/VPlanConstantFolder.h

@@ -0,0 +1,98 @@
+//===- 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "VPlanValue.h"
+#include "llvm/IR/ConstantFolder.h"
+
+#ifndef LLVM_TRANSFORMS_VECTORIZE_VPLANCONSTANTFOLDER_H
+#define LLVM_TRANSFORMS_VECTORIZE_VPLANCONSTANTFOLDER_H
+
+namespace llvm {
+class VPConstantFolder {
+private:
+  ConstantFolder Folder;
+
+  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;
+  }
+
+public:
+  Value *foldAnd(VPValue *LHS, VPValue *RHS) const {
+    return foldBinOp(Instruction::BinaryOps::And, LHS, RHS);
+  }
+
+  Value *foldOr(VPValue *LHS, VPValue *RHS) const {
+    return foldBinOp(Instruction::BinaryOps::Or, LHS, RHS);
+  }
+
+  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 *foldPtrAdd(VPValue *Base, VPValue *Offset, GEPNoWrapFlags NW) const {
+    auto *BC = getIRConstant(Base);
+    auto *OC = getIRConstant(Offset);
+    if (BC && OC) {
+      auto &Ctx = BC->getType()->getContext();
+      return Folder.FoldGEP(Type::getInt8Ty(Ctx), BC, OC, NW);
+    }
+    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;
+  }
+};
+} // namespace llvm
+
+#endif

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

@@ -533,8 +533,8 @@ createScalarIVSteps(VPlan &Plan, InductionDescriptor::InductionKind Kind,
                     VPBuilder &Builder) {
   VPBasicBlock *HeaderVPBB = Plan.getVectorLoopRegion()->getEntryBasicBlock();
   VPCanonicalIVPHIRecipe *CanonicalIV = Plan.getCanonicalIV();
-  VPSingleDefRecipe *BaseIV = Builder.createDerivedIV(
-      Kind, FPBinOp, StartV, CanonicalIV, Step, "offset.idx");
+  VPValue *BaseIV = Builder.createDerivedIV(Kind, FPBinOp, StartV, CanonicalIV,
+                                            Step, "offset.idx");
 
   // Truncate base induction if needed.
   Type *CanonicalIVType = CanonicalIV->getScalarType();