[VPlan] Simplify VPDerivedIVRecipe.

Author: fhahn

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

@@ -384,7 +384,10 @@ Value *VPInstruction::generate(VPTransformState &State) {
   IRBuilderBase &Builder = State.Builder;
 
   if (Instruction::isBinaryOp(getOpcode())) {
-    bool OnlyFirstLaneUsed = vputils::onlyFirstLaneUsed(this);
+    bool OnlyFirstLaneUsed =
+        vputils::onlyFirstLaneUsed(this) || all_of(operands(), [](VPValue *Op) {
+          return vputils::isUniformAfterVectorization(Op);
+        });
     Value *A = State.get(getOperand(0), OnlyFirstLaneUsed);
     Value *B = State.get(getOperand(1), OnlyFirstLaneUsed);
     auto *Res =
@@ -692,12 +695,8 @@ void VPInstruction::execute(VPTransformState &State) {
   if (!hasResult())
     return;
   assert(GeneratedValue && "generate must produce a value");
-  assert(
-      (GeneratedValue->getType()->isVectorTy() == !GeneratesPerFirstLaneOnly ||
-       State.VF.isScalar()) &&
-      "scalar value but not only first lane defined");
   State.set(this, GeneratedValue,
-            /*IsScalar*/ GeneratesPerFirstLaneOnly);
+            /*IsScalar*/ !GeneratedValue->getType()->isVectorTy());
 }
 
 bool VPInstruction::onlyFirstLaneUsed(const VPValue *Op) const {

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

@@ -969,6 +969,32 @@ static void simplifyRecipe(VPRecipeBase &R, VPTypeAnalysis &TypeInfo) {
 
   if (match(&R, m_c_Mul(m_VPValue(A), m_SpecificInt(1))))
     return R.getVPSingleValue()->replaceAllUsesWith(A);
+
+  if (auto *DerivedIV = dyn_cast<VPDerivedIVRecipe>(&R)) {
+    VPValue *Start = DerivedIV->getStartValue();
+    VPValue *Step = DerivedIV->getStepValue();
+    VPValue *Op = DerivedIV->getOperand(1);
+    if (Step->isLiveIn() && Start->isLiveIn()) {
+      auto *StartC = dyn_cast<ConstantInt>(Start->getLiveInIRValue());
+      auto *StepC = dyn_cast<ConstantInt>(Step->getLiveInIRValue());
+      if (Op->isLiveIn()) {
+        auto *OpC = dyn_cast_if_present<ConstantInt>(Op->getUnderlyingValue());
+        if (OpC && OpC->getType() == StepC->getType() && StartC->getType()) {
+          APInt Folded =
+              StartC->getValue() + OpC->getValue() * StepC->getValue();
+          DerivedIV->replaceAllUsesWith(
+              R.getParent()->getPlan()->getOrAddLiveIn(
+                  ConstantInt::get(TypeInfo.getContext(), Folded)));
+        }
+      }
+      if (StartC && StartC->isZero() && StepC && StepC->isOne() &&
+          TypeInfo.inferScalarType(DerivedIV) ==
+              TypeInfo.inferScalarType(DerivedIV->getOperand(1))) {
+        DerivedIV->replaceAllUsesWith(DerivedIV->getOperand(1));
+        DerivedIV->eraseFromParent();
+      }
+    }
+  }
 }
 
 /// Move loop-invariant recipes out of the vector loop region in \p Plan.