[VPlan] Add VPResultType mixing to manage recipe result types.

Author: fhahn

llvm/lib/Transforms/Vectorize/VPlan.h

@@ -1030,17 +1030,25 @@ class VPInstruction : public VPRecipeWithIRFlags,
   }
 };
 
-/// A specialization of VPInstruction augmenting it with a dedicated result
-/// type, to be used when the opcode and operands of the VPInstruction don't
-/// directly determine the result type.
-class VPInstructionWithType : public VPInstruction {
+class VPResultType {
   /// Scalar result type produced by the recipe.
   Type *ResultTy;
 
+public:
+  VPResultType(Type *ResultTy) : ResultTy(ResultTy) {}
+  virtual ~VPResultType() = default;
+
+  Type *getResultType() const { return ResultTy; }
+};
+
+/// A specialization of VPInstruction augmenting it with a dedicated result
+/// type, to be used when the opcode and operands of the VPInstruction don't
+/// directly determine the result type.
+class VPInstructionWithType : public VPInstruction, public VPResultType {
 public:
   VPInstructionWithType(unsigned Opcode, ArrayRef<VPValue *> Operands,
                         Type *ResultTy, DebugLoc DL, const Twine &Name = "")
-      : VPInstruction(Opcode, Operands, DL, Name), ResultTy(ResultTy) {}
+      : VPInstruction(Opcode, Operands, DL, Name), VPResultType(ResultTy) {}
 
   static inline bool classof(const VPRecipeBase *R) { return isCast(R); }
 
@@ -1065,15 +1073,45 @@ class VPInstructionWithType : public VPInstruction {
     return 0;
   }
 
-  Type *getResultType() const { return ResultTy; }
-
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
   /// Print the recipe.
   void print(raw_ostream &O, const Twine &Indent,
              VPSlotTracker &SlotTracker) const override;
 #endif
 };
 
+template <> struct CastIsPossible<VPResultType, const VPRecipeBase *> {
+  static inline bool isPossible(const VPRecipeBase *f) {
+    return isa<VPInstructionWithType>(f);
+  }
+};
+
+template <>
+struct CastInfo<VPResultType, const VPRecipeBase *>
+    : public CastIsPossible<VPResultType, const VPRecipeBase *> {
+  using Self = CastInfo<VPResultType, const VPRecipeBase *>;
+
+  using CastReturnType =
+      typename cast_retty<VPResultType, VPRecipeBase *>::ret_type;
+
+  static inline VPResultType *doCast(const VPRecipeBase *f) {
+    auto *VPI =
+        const_cast<VPInstructionWithType *>(dyn_cast<VPInstructionWithType>(f));
+    return VPI;
+  }
+
+  // This assumes that you can construct the cast return type from `nullptr`.
+  // This is largely to support legacy use cases - if you don't want this
+  // behavior you should specialize CastInfo for your use case.
+  static inline VPResultType *castFailed() { return nullptr; }
+
+  static inline VPResultType *doCastIfPossible(const VPRecipeBase *f) {
+    if (!Self::isPossible(f))
+      return castFailed();
+    return doCast(f);
+  }
+};
+
 /// A recipe to wrap on original IR instruction not to be modified during
 /// execution, except for PHIs. PHIs are modeled via the VPIRPhi subclass.
 /// Expect PHIs, VPIRInstructions cannot have any operands.

llvm/lib/Transforms/Vectorize/VPlanAnalysis.cpp

@@ -261,9 +261,10 @@ Type *VPTypeAnalysis::inferScalarType(const VPValue *V) {
                 VPPartialReductionRecipe>([this](const VPRecipeBase *R) {
             return inferScalarType(R->getOperand(0));
           })
-          // VPInstructionWithType must be handled before VPInstruction.
-          .Case<VPInstructionWithType, VPWidenIntrinsicRecipe>(
+          // VPResultType must be handled before VPInstruction.
+          .Case<VPResultType, VPWidenIntrinsicRecipe>(
               [](const auto *R) { return R->getResultType(); })
+
           .Case<VPBlendRecipe, VPInstruction, VPWidenRecipe, VPReplicateRecipe,
                 VPWidenCallRecipe, VPWidenMemoryRecipe, VPWidenSelectRecipe>(
               [this](const auto *R) { return inferScalarTypeForRecipe(R); })

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

@@ -1078,7 +1078,7 @@ void VPInstructionWithType::execute(VPTransformState &State) {
   case Instruction::Trunc: {
     Value *Op = State.get(getOperand(0), VPLane(0));
     Value *Cast = State.Builder.CreateCast(Instruction::CastOps(getOpcode()),
-                                           Op, ResultTy);
+                                           Op, getResultType());
     State.set(this, Cast, VPLane(0));
     break;
   }
@@ -1094,7 +1094,7 @@ void VPInstructionWithType::print(raw_ostream &O, const Twine &Indent,
   printAsOperand(O, SlotTracker);
   O << " = " << Instruction::getOpcodeName(getOpcode()) << " ";
   printOperands(O, SlotTracker);
-  O << " to " << *ResultTy;
+  O << " to " << *getResultType();
 }
 #endif