[llvm] [VPlan] Introduce VPlanConstantFolder (PR #125365)

[Florian Hahn via llvm-commits]

================
@@ -937,10 +938,84 @@ static void recursivelyDeleteDeadRecipes(VPValue *V) {
   }
 }
 
+class VPConstantFolder {
+  TargetFolder Folder;
+  VPTypeAnalysis TypeInfo;
+
+public:
+  VPConstantFolder(const DataLayout &DL, const VPTypeAnalysis &TypeInfo)
+      : Folder(DL), TypeInfo(TypeInfo) {}
+
+  Value *tryToConstantFold(VPRecipeBase &R, unsigned Opcode,
+                           ArrayRef<VPValue *> Operands) {
+    SmallVector<Value *, 4> Ops;
+    for (VPValue *Op : Operands) {
+      if (!Op->isLiveIn() || !Op->getLiveInIRValue())
+        return nullptr;
+      Ops.emplace_back(Op->getLiveInIRValue());
+    }
+    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 Folder.FoldBinOp(static_cast<Instruction::BinaryOps>(Opcode),
+                              Ops[0], Ops[1]);
+    case VPInstruction::LogicalAnd:
+      return Folder.FoldSelect(Ops[0], Ops[1],
+                               ConstantInt::getNullValue(Ops[1]->getType()));
+    case VPInstruction::Not:
+      return Folder.FoldBinOp(Instruction::BinaryOps::Xor, Ops[0],
+                              Constant::getAllOnesValue(Ops[0]->getType()));
+    case Instruction::Select:
+      return Folder.FoldSelect(Ops[0], Ops[1], Ops[2]);
+    case Instruction::ICmp:
+    case Instruction::FCmp:
+      return Folder.FoldCmp(cast<VPRecipeWithIRFlags>(R).getPredicate(), Ops[0],
+                            Ops[1]);
+    case Instruction::GetElementPtr:
+    case VPInstruction::PtrAdd:
+      return Folder.FoldGEP(TypeInfo.inferScalarType(R.getVPSingleValue()),
+                            Ops[0], drop_begin(Ops),
+                            cast<VPRecipeWithIRFlags>(R).getGEPNoWrapFlags());
+    case Instruction::InsertElement:
+      return Folder.FoldInsertElement(Ops[0], Ops[1], Ops[2]);
+    case Instruction::ExtractElement:
+      return Folder.FoldExtractElement(Ops[0], Ops[1]);
+    case Instruction::CastOps::SExt:
+    case Instruction::CastOps::ZExt:
+    case Instruction::CastOps::Trunc:
+      return Folder.FoldCast(static_cast<Instruction::CastOps>(Opcode), Ops[0],
+                             TypeInfo.inferScalarType(R.getVPSingleValue()));
+    }
+    return nullptr;
+  }
+};
+
 /// 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);
+  if (TypeSwitch<VPRecipeBase *, bool>(&R)
+          .Case<VPInstruction, VPWidenRecipe, VPWidenCastRecipe,
+                VPReplicateRecipe>([&](auto *I) {
+            VPlan *Plan = R.getParent()->getPlan();
+            Value *V =
+                Folder.tryToConstantFold(R, I->getOpcode(), I->operands());
----------------
fhahn wrote:

```suggestion
                Folder.tryToConstantFold(I, I->getOpcode(), I->operands());
```

https://github.com/llvm/llvm-project/pull/125365