[LV] Fold tail by masking - handle reductions

[Ayal]

Summary

This addresses the "TODO: handle reductions when tail is folded by masking" left by D50480.

When folding the tail (scalar leftover iterations) by introducing an additional (masked) vector iteration at the end, first note that
the header phi of a reduction can remain intact - it takes care of "accumulating the partial sums" of a reduction across all iterations except for the last one. It is only this last iteration which now needs to accumulate these partial sums under a mask, effecting the live-out values. This can be accomplished by introducing a Select instruction, choosing between the last and penultimate (i.e., header phi) values of the partial sums.

[fhahn]

LGTM, thanks!

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

3690

I think the assertions here could be a bit stronger, by additionally ensuring we only have a single SelectInst as user and all other users are PHINodes, instead of breaking once we find a SelectInst.

[dcaballe]

Thanks, Ayal! LGTM

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

3682

chosing -> choosing?

[Kryukova_l_p]

@Ayal Hi.
Found a problem caused by these changes.
Take a look, please. I posted about it here https://bugs.llvm.org/show_bug.cgi?id=47390

I attach the file here too, after the bug point:

reduced-simplified.ll18 KBDownload

Run like this:
opt -passes=loop-vectorize 'reduced-simplified.ll' -S

Could you take a look?
Thanks in advance.

[Ayal]

In D66720#2413499, @Kryukova_l_p wrote:

@Ayal Hi.
Found a problem caused by these changes.
Take a look, please. I posted about it here https://bugs.llvm.org/show_bug.cgi?id=47390

I attach the file here too, after the bug point:

reduced-simplified.ll18 KBDownload

Run like this:
opt -passes=loop-vectorize 'reduced-simplified.ll' -S

Could you take a look?
Thanks in advance.

Thanks @Kryukova_l_p for reporting; the fix in D92017 should be pushed shortly.

llvm/lib/Transforms/Vectorize/LoopVectorizationLegality.cpp

if (!PrimaryInduction) {
     reportVectorizationFailure(
         "No primary induction, cannot fold tail by masking",
         "Missing a primary induction variable in the loop, which is "
         "needed in order to fold tail by masking as required.",
         "NoPrimaryInduction", ORE, TheLoop);
     return false;
   }
 
-  // TODO: handle reductions when tail is folded by masking.
-  if (!Reductions.empty()) {
-    reportVectorizationFailure(
-        "Loop has reductions, cannot fold tail by masking",
-        "Cannot fold tail by masking in the presence of reductions.",
-        "ReductionFoldingTailByMasking", ORE, TheLoop);
-    return false;
-  }
+  SmallPtrSet<const Value *, 8> ReductionLiveOuts;
 
-  // TODO: handle outside users when tail is folded by masking.
+  for (auto &Reduction : *getReductionVars())
+    ReductionLiveOuts.insert(Reduction.second.getLoopExitInstr());
+
+  // TODO: handle non-reduction outside users when tail is folded by masking.
   for (auto *AE : AllowedExit) {
-    // Check that all users of allowed exit values are inside the loop.
+    // Check that all users of allowed exit values are inside the loop or
+    // are the live-out of a reduction.
+    if (ReductionLiveOuts.count(AE))
+      continue;
     for (User *U : AE->users()) {
       Instruction *UI = cast<Instruction>(U);
       if (TheLoop->contains(UI))
         continue;
       reportVectorizationFailure(
           "Cannot fold tail by masking, loop has an outside user for",
           "Cannot fold tail by masking in the presence of live outs.",
           "LiveOutFoldingTailByMasking", ORE, TheLoop, UI);

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

void InnerLoopVectorizer::fixReduction(PHINode *Phi) {
   // Before each round, move the insertion point right between
   // the PHIs and the values we are going to write.
   // This allows us to write both PHINodes and the extractelement
   // instructions.
   Builder.SetInsertPoint(&*LoopMiddleBlock->getFirstInsertionPt());
 
   setDebugLocFromInst(Builder, LoopExitInst);
 
+  // If tail is folded by masking, the vector value to leave the loop should be
+  // a Select chosing between the vectorized LoopExitInst and vectorized Phi,
+  // instead of the former.
+  if (Cost->foldTailByMasking()) {
+    for (unsigned Part = 0; Part < UF; ++Part) {
+      Value *VecLoopExitInst =
+          VectorLoopValueMap.getVectorValue(LoopExitInst, Part);
+      SelectInst *Sel = nullptr;
+      for (User *U : VecLoopExitInst->users()) {
+        if ((Sel = dyn_cast<SelectInst>(U)))
+          break;
+        assert(isa<PHINode>(U) && "Non-Phi user of reduction exit instruction");
+      }
+      assert(Sel && "No Select of reduction exit instruction under fold tail");
+      VectorLoopValueMap.resetVectorValue(LoopExitInst, Part, Sel);
+    }
+  }
+
   // If the vector reduction can be performed in a smaller type, we truncate
   // then extend the loop exit value to enable InstCombine to evaluate the
   // entire expression in the smaller type.
   if (VF > 1 && Phi->getType() != RdxDesc.getRecurrenceType()) {
     Type *RdxVecTy = VectorType::get(RdxDesc.getRecurrenceType(), VF);
     Builder.SetInsertPoint(
         LI->getLoopFor(LoopVectorBody)->getLoopLatch()->getTerminator());
     VectorParts RdxParts(UF);
if (BB == Latch)
       continue;
     BranchInst *Branch = dyn_cast<BranchInst>(BB->getTerminator());
     if (Branch && Branch->isConditional())
       NeedDef.insert(Branch->getCondition());
   }
 
   // If the tail is to be folded by masking, the primary induction variable
   // needs to be represented in VPlan for it to model early-exit masking.
-  if (CM.foldTailByMasking())
+  // Also, both the Phi and the live-out instruction of each reduction are
+  // required in order to introduce a select between them in VPlan.
+  if (CM.foldTailByMasking()) {
     NeedDef.insert(Legal->getPrimaryInduction());
+    for (auto &Reduction : *Legal->getReductionVars()) {
+      NeedDef.insert(Reduction.first);
+      NeedDef.insert(Reduction.second.getLoopExitInstr());
+    }
+  }
 
   // Collect instructions from the original loop that will become trivially dead
   // in the vectorized loop. We don't need to vectorize these instructions. For
   // example, original induction update instructions can become dead because we
   // separately emit induction "steps" when generating code for the new loop.
   // Similarly, we create a new latch condition when setting up the structure
   // of the new loop, so the old one can become dead.
   SmallPtrSet<Instruction *, 4> DeadInstructions;
VPlanPtr LoopVectorizationPlanner::buildVPlanWithVPRecipes(
   // may also be empty, such as the last one VPBB, reflecting original
   // basic-blocks with no recipes.
   VPBasicBlock *PreEntry = cast<VPBasicBlock>(Plan->getEntry());
   assert(PreEntry->empty() && "Expecting empty pre-entry block.");
   VPBlockBase *Entry = Plan->setEntry(PreEntry->getSingleSuccessor());
   VPBlockUtils::disconnectBlocks(PreEntry, Entry);
   delete PreEntry;
 
+  // Finally, if tail is folded by masking, introduce selects between the phi
+  // and the live-out instruction of each reduction, at the end of the latch.
+  if (CM.foldTailByMasking()) {
+    Builder.setInsertPoint(VPBB);
+    auto *Cond = RecipeBuilder.createBlockInMask(OrigLoop->getHeader(), Plan);
+    for (auto &Reduction : *Legal->getReductionVars()) {
+      VPValue *Phi = Plan->getVPValue(Reduction.first);
+      VPValue *Red = Plan->getVPValue(Reduction.second.getLoopExitInstr());
+      Builder.createNaryOp(Instruction::Select, {Cond, Red, Phi});
+    }
+  }
+
   std::string PlanName;
   raw_string_ostream RSO(PlanName);
   unsigned VF = Range.Start;
   Plan->addVF(VF);
   RSO << "Initial VPlan for VF={" << VF;
   for (VF *= 2; VF < Range.End; VF *= 2) {
     Plan->addVF(VF);
     RSO << "," << VF;

llvm/lib/Transforms/Vectorize/VPlan.cpp

void VPInstruction::generateInstruction(VPTransformState &State,
   }
   case VPInstruction::ICmpULE: {
     Value *IV = State.get(getOperand(0), Part);
     Value *TC = State.get(getOperand(1), Part);
     Value *V = Builder.CreateICmpULE(IV, TC);
     State.set(this, V, Part);
     break;
   }
+  case Instruction::Select: {
+    Value *Cond = State.get(getOperand(0), Part);
+    Value *Op1 = State.get(getOperand(1), Part);
+    Value *Op2 = State.get(getOperand(2), Part);
+    Value *V = Builder.CreateSelect(Cond, Op1, Op2);
+    State.set(this, V, Part);
+    break;
+  }
   default:
     llvm_unreachable("Unsupported opcode for instruction");
   }
 }
 
 void VPInstruction::execute(VPTransformState &State) {
   assert(!State.Instance && "VPInstruction executing an Instance");
   for (unsigned Part = 0; Part < State.UF; ++Part)

llvm/test/Transforms/LoopVectorize/X86/tail_loop_folding.ll

for.body:
   %add = add nsw i32 %1, %0
   %arrayidx4 = getelementptr inbounds i32, i32* %A, i64 %indvars.iv
   store i32 %add, i32* %arrayidx4, align 4
   %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
   %exitcond = icmp eq i64 %indvars.iv.next, 430
   br i1 %exitcond, label %for.cond.cleanup, label %for.body, !llvm.loop !10
 }
 
+; Check that fold tail under optsize passes the reduction live-out value
+; through a select.
+; int reduction_i32(int *A, int *B, int N) {
+;   int sum = 0;
+;   for (int i = 0; i < N; ++i)
+;     sum += (A[i] + B[i]);
+;   return sum;
+; }
+;
+define i32 @reduction_i32(i32* nocapture readonly %A, i32* nocapture readonly %B, i32 %N) #0 {
+; CHECK-LABEL: @reduction_i32(
+; CHECK:       vector.body:
+; CHECK-NEXT:    [[INDEX:%.*]] = phi i64 [ 0, %vector.ph ], [ [[INDEX_NEXT:%.*]], %vector.body ]
+; CHECK-NEXT:    [[ACCUM_PHI:%.*]] = phi <8 x i32> [ zeroinitializer, %vector.ph ], [ [[ACCUM:%.*]], %vector.body ]
+; CHECK:         [[ICMPULE:%.*]] = icmp ule <8 x i64>
+; CHECK:         [[LOAD1:%.*]] = call <8 x i32> @llvm.masked.load.v8i32.p0v8i32(<8 x i32>* {{.*}}, i32 4, <8 x i1> [[ICMPULE]], <8 x i32> undef)
+; CHECK:         [[LOAD2:%.*]] = call <8 x i32> @llvm.masked.load.v8i32.p0v8i32(<8 x i32>* {{.*}}, i32 4, <8 x i1> [[ICMPULE]], <8 x i32> undef)
+; CHECK-NEXT:    [[ADD:%.*]] = add nsw <8 x i32> [[LOAD2]], [[LOAD1]]
+; CHECK-NEXT:    [[ACCUM]] = add nuw nsw <8 x i32> [[ADD]], [[ACCUM_PHI]]
+; CHECK:         [[LIVEOUT:%.*]] = select <8 x i1> [[ICMPULE]], <8 x i32> [[ACCUM]], <8 x i32> [[ACCUM_PHI]]
+; CHECK-NEXT:    [[INDEX_NEXT]] = add i64 [[INDEX]], 8
+; CHECK:       middle.block:
+; CHECK-NEXT:    [[RDX_SHUF:%.*]] = shufflevector <8 x i32> [[LIVEOUT]], <8 x i32> undef, <8 x i32> <i32 4, i32 5, i32 6, i32 7, i32 undef, i32 undef, i32 undef, i32 undef>
+; CHECK-NEXT:    [[BIN_RDX:%.*]] = add <8 x i32> [[LIVEOUT]], [[RDX_SHUF]]
+; CHECK-NEXT:    [[RDX_SHUF4:%.*]] = shufflevector <8 x i32> [[BIN_RDX]], <8 x i32> undef, <8 x i32> <i32 2, i32 3, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef>
+; CHECK-NEXT:    [[BIN_RDX5:%.*]] = add <8 x i32> [[BIN_RDX]], [[RDX_SHUF4]]
+; CHECK-NEXT:    [[RDX_SHUF6:%.*]] = shufflevector <8 x i32> [[BIN_RDX5]], <8 x i32> undef, <8 x i32> <i32 1, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef, i32 undef>
+; CHECK-NEXT:    [[BIN_RDX7:%.*]] = add <8 x i32> [[BIN_RDX5]], [[RDX_SHUF6]]
+; CHECK-NEXT:    [[TMP17:%.*]] = extractelement <8 x i32> [[BIN_RDX7]], i32 0
+; CHECK-NEXT:    br i1 true, label %for.cond.cleanup, label %scalar.ph
+; CHECK:       scalar.ph:
+; CHECK:       for.cond.cleanup:
+; CHECK-NEXT:    [[SUM_1_LCSSA:%.*]] = phi i32 [ {{.*}}, %for.body ], [ [[TMP17]], %middle.block ]
+; CHECK-NEXT:    ret i32 [[SUM_1_LCSSA]]
+;
+entry:
+  br label %for.body
+
+for.body:
+  %indvars.iv = phi i64 [ %indvars.iv.next, %for.body ], [ 0, %entry ]
+  %sum.0 = phi i32 [ %sum.1, %for.body ], [ 0, %entry ]
+  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+  %arrayidxA = getelementptr inbounds i32, i32* %A, i64 %indvars.iv
+  %0 = load i32, i32* %arrayidxA, align 4
+  %arrayidxB = getelementptr inbounds i32, i32* %B, i64 %indvars.iv
+  %1 = load i32, i32* %arrayidxB, align 4
+  %add = add nsw i32 %1, %0
+  %sum.1 = add nuw nsw i32 %add, %sum.0
+  %lftr.wideiv = trunc i64 %indvars.iv.next to i32
+  %exitcond = icmp eq i32 %lftr.wideiv, %N
+  br i1 %exitcond, label %for.cond.cleanup, label %for.body
+
+for.cond.cleanup:
+  ret i32 %sum.1
+}
+
 ; CHECK:      !0 = distinct !{!0, !1}
 ; CHECK-NEXT: !1 = !{!"llvm.loop.isvectorized", i32 1}
 ; CHECK-NEXT: !2 = distinct !{!2, !3, !1}
 ; CHECK-NEXT: !3 = !{!"llvm.loop.unroll.runtime.disable"}
 ; CHECK-NEXT: !4 = distinct !{!4, !1}
 ; CHECK-NEXT: !5 = distinct !{!5, !3, !1}
 
 attributes #0 = { nounwind optsize uwtable "target-cpu"="core-avx2" "target-features"="+avx,+avx2" }
 
 !6 = distinct !{!6, !7, !8}
 !7 = !{!"llvm.loop.vectorize.predicate.enable", i1 true}
 !8 = !{!"llvm.loop.vectorize.enable", i1 true}
 
 !10 = distinct !{!10, !11, !12}
 !11 = !{!"llvm.loop.vectorize.predicate.enable", i1 false}
 !12 = !{!"llvm.loop.vectorize.enable", i1 true}