[LV] Extend FindLastIV to unsigned case

In an effort to not have two different RecurKinds, one for the signed
case, and another for the unsigned case, introduce
RecurrenceDescriptor::isReduxSigned() to indicate whether the the
RecurKind is of the signed or unsigned variant. Demonstrate its use by
extending FindLastIV to the unsigned case.

Author: artagnon

llvm/include/llvm/Analysis/IVDescriptors.h

@@ -80,12 +80,13 @@ class RecurrenceDescriptor {
 
   RecurrenceDescriptor(Value *Start, Instruction *Exit, StoreInst *Store,
                        RecurKind K, FastMathFlags FMF, Instruction *ExactFP,
-                       Type *RT, bool Signed, bool Ordered,
-                       SmallPtrSetImpl<Instruction *> &CI,
+                       Type *RT, bool IsResultSigned, bool IsReduxSigned,
+                       bool Ordered, SmallPtrSetImpl<Instruction *> &CI,
                        unsigned MinWidthCastToRecurTy)
       : IntermediateStore(Store), StartValue(Start), LoopExitInstr(Exit),
         Kind(K), FMF(FMF), ExactFPMathInst(ExactFP), RecurrenceType(RT),
-        IsSigned(Signed), IsOrdered(Ordered),
+        IsResultSigned(IsResultSigned), IsReduxSigned(IsReduxSigned),
+        IsOrdered(Ordered),
         MinWidthCastToRecurrenceType(MinWidthCastToRecurTy) {
     CastInsts.insert_range(CI);
   }
@@ -97,12 +98,14 @@ class RecurrenceDescriptor {
         : IsRecurrence(IsRecur), PatternLastInst(I),
           RecKind(RecurKind::None), ExactFPMathInst(ExactFP) {}
 
-    InstDesc(Instruction *I, RecurKind K, Instruction *ExactFP = nullptr)
-        : IsRecurrence(true), PatternLastInst(I), RecKind(K),
-          ExactFPMathInst(ExactFP) {}
+    InstDesc(Instruction *I, RecurKind K, bool IsSigned = false)
+        : IsRecurrence(true), IsSigned(IsSigned), PatternLastInst(I),
+          RecKind(K) {}
 
     bool isRecurrence() const { return IsRecurrence; }
 
+    bool isSigned() const { return IsSigned; }
+
     bool needsExactFPMath() const { return ExactFPMathInst != nullptr; }
 
     Instruction *getExactFPMathInst() const { return ExactFPMathInst; }
@@ -114,13 +117,15 @@ class RecurrenceDescriptor {
   private:
     // Is this instruction a recurrence candidate.
     bool IsRecurrence;
+    // Is this recurrence a signed variant.
+    bool IsSigned = false;
     // The last instruction in a min/max pattern (select of the select(icmp())
     // pattern), or the current recurrence instruction otherwise.
     Instruction *PatternLastInst;
     // If this is a min/max pattern.
     RecurKind RecKind;
     // Recurrence does not allow floating-point reassociation.
-    Instruction *ExactFPMathInst;
+    Instruction *ExactFPMathInst = nullptr;
   };
 
   /// Returns a struct describing if the instruction 'I' can be a recurrence
@@ -272,8 +277,9 @@ class RecurrenceDescriptor {
   Value *getSentinelValue() const {
     assert(isFindLastIVRecurrenceKind(Kind) && "Unexpected recurrence kind");
     Type *Ty = StartValue->getType();
-    return ConstantInt::get(Ty,
-                            APInt::getSignedMinValue(Ty->getIntegerBitWidth()));
+    unsigned BW = Ty->getIntegerBitWidth();
+    return ConstantInt::get(Ty, isReduxSigned() ? APInt::getSignedMinValue(BW)
+                                                : APInt::getMinValue(BW));
   }
 
   /// Returns a reference to the instructions used for type-promoting the
@@ -285,8 +291,11 @@ class RecurrenceDescriptor {
     return MinWidthCastToRecurrenceType;
   }
 
-  /// Returns true if all source operands of the recurrence are SExtInsts.
-  bool isSigned() const { return IsSigned; }
+  /// Returns true if the reduction result is signed.
+  bool isResultSigned() const { return IsResultSigned; }
+
+  /// Returns true if the reduction redux is signed.
+  bool isReduxSigned() const { return IsReduxSigned; }
 
   /// Expose an ordered FP reduction to the instance users.
   bool isOrdered() const { return IsOrdered; }
@@ -322,8 +331,10 @@ class RecurrenceDescriptor {
   Instruction *ExactFPMathInst = nullptr;
   // The type of the recurrence.
   Type *RecurrenceType = nullptr;
-  // True if all source operands of the recurrence are SExtInsts.
-  bool IsSigned = false;
+  // True if reduction result is signed.
+  bool IsResultSigned = false;
+  // True if reduction redux is signed.
+  bool IsReduxSigned = false;
   // True if this recurrence can be treated as an in-order reduction.
   // Currently only a non-reassociative FAdd can be considered in-order,
   // if it is also the only FAdd in the PHI's use chain.

llvm/lib/Analysis/IVDescriptors.cpp

@@ -88,7 +88,7 @@ static Instruction *lookThroughAnd(PHINode *Phi, Type *&RT,
 }
 
 /// Compute the minimal bit width needed to represent a reduction whose exit
-/// instruction is given by Exit.
+/// instruction is given by Exit, along with its signedness.
 static std::pair<Type *, bool> computeRecurrenceType(Instruction *Exit,
                                                      DemandedBits *DB,
                                                      AssumptionCache *AC,
@@ -255,7 +255,7 @@ bool RecurrenceDescriptor::AddReductionVar(
   SmallPtrSet<Instruction *, 4> CastInsts;
   unsigned MinWidthCastToRecurrenceType;
   Instruction *Start = Phi;
-  bool IsSigned = false;
+  bool IsResultSigned = false, IsReduxSigned = false;
 
   SmallPtrSet<Instruction *, 8> VisitedInsts;
   SmallVector<Instruction *, 8> Worklist;
@@ -396,6 +396,7 @@ bool RecurrenceDescriptor::AddReductionVar(
       //       state accurate while processing the worklist?
       if (ReduxDesc.getRecKind() != RecurKind::None)
         Kind = ReduxDesc.getRecKind();
+      IsReduxSigned = ReduxDesc.isSigned();
     }
 
     bool IsASelect = isa<SelectInst>(Cur);
@@ -565,7 +566,7 @@ bool RecurrenceDescriptor::AddReductionVar(
     //       smaller type. We should just generate a correctly typed expression
     //       to begin with.
     Type *ComputedType;
-    std::tie(ComputedType, IsSigned) =
+    std::tie(ComputedType, IsResultSigned) =
         computeRecurrenceType(ExitInstruction, DB, AC, DT);
     if (ComputedType != RecurrenceType)
       return false;
@@ -595,8 +596,9 @@ bool RecurrenceDescriptor::AddReductionVar(
 
   // Save the description of this reduction variable.
   RecurrenceDescriptor RD(RdxStart, ExitInstruction, IntermediateStore, Kind,
-                          FMF, ExactFPMathInst, RecurrenceType, IsSigned,
-                          IsOrdered, CastInsts, MinWidthCastToRecurrenceType);
+                          FMF, ExactFPMathInst, RecurrenceType, IsResultSigned,
+                          IsReduxSigned, IsOrdered, CastInsts,
+                          MinWidthCastToRecurrenceType);
   RedDes = RD;
 
   return true;
@@ -700,47 +702,59 @@ RecurrenceDescriptor::isFindLastIVPattern(Loop *TheLoop, PHINode *OrigPhi,
                                      m_Value(NonRdxPhi)))))
     return InstDesc(false, I);
 
-  auto IsIncreasingLoopInduction = [&](Value *V) {
+  // Returns a non-nullopt boolean indicating the signedness of the recurrence
+  // when a valid FindLastIV pattern is found.
+  auto GetInductionSignedness = [&](Value *V) -> std::optional<bool> {
     Type *Ty = V->getType();
     if (!SE.isSCEVable(Ty))
-      return false;
+      return std::nullopt;
 
     auto *AR = dyn_cast<SCEVAddRecExpr>(SE.getSCEV(V));
     if (!AR || AR->getLoop() != TheLoop)
-      return false;
+      return std::nullopt;
 
     const SCEV *Step = AR->getStepRecurrence(SE);
     if (!SE.isKnownPositive(Step))
-      return false;
+      return std::nullopt;
 
-    const ConstantRange IVRange = SE.getSignedRange(AR);
-    unsigned NumBits = Ty->getIntegerBitWidth();
     // Keep the minimum value of the recurrence type as the sentinel value.
     // The maximum acceptable range for the increasing induction variable,
     // called the valid range, will be defined as
     //   [<sentinel value> + 1, <sentinel value>)
-    // where <sentinel value> is SignedMin(<recurrence type>)
+    // where <sentinel value> is [Signed|Unsigned]Min(<recurrence type>)
     // TODO: This range restriction can be lifted by adding an additional
     // virtual OR reduction.
-    const APInt Sentinel = APInt::getSignedMinValue(NumBits);
-    const ConstantRange ValidRange =
-        ConstantRange::getNonEmpty(Sentinel + 1, Sentinel);
-    LLVM_DEBUG(dbgs() << "LV: FindLastIV valid range is " << ValidRange
-                      << ", and the signed range of " << *AR << " is "
-                      << IVRange << "\n");
-    // Ensure the induction variable does not wrap around by verifying that its
-    // range is fully contained within the valid range.
-    return ValidRange.contains(IVRange);
+    auto CheckRange = [&](bool IsSigned) {
+      const ConstantRange IVRange =
+          IsSigned ? SE.getSignedRange(AR) : SE.getUnsignedRange(AR);
+      unsigned NumBits = Ty->getIntegerBitWidth();
+      const APInt Sentinel = IsSigned ? APInt::getSignedMinValue(NumBits)
+                                      : APInt::getMinValue(NumBits);
+      const ConstantRange ValidRange =
+          ConstantRange::getNonEmpty(Sentinel + 1, Sentinel);
+      LLVM_DEBUG(dbgs() << "LV: FindLastIV valid range is " << ValidRange
+                        << ", and the range of " << *AR << " is " << IVRange
+                        << "\n");
+
+      // Ensure the induction variable does not wrap around by verifying that
+      // its range is fully contained within the valid range.
+      return ValidRange.contains(IVRange);
+    };
+    if (CheckRange(true))
+      return true;
+    if (CheckRange(false))
+      return false;
+    return std::nullopt;
   };
 
   // We are looking for selects of the form:
   //   select(cmp(), phi, increasing_loop_induction) or
   //   select(cmp(), increasing_loop_induction, phi)
   // TODO: Support for monotonically decreasing induction variable
-  if (!IsIncreasingLoopInduction(NonRdxPhi))
-    return InstDesc(false, I);
+  if (auto IsSigned = GetInductionSignedness(NonRdxPhi))
+    return InstDesc(I, RecurKind::FindLastIV, *IsSigned);
 
-  return InstDesc(I, RecurKind::FindLastIV);
+  return InstDesc(false, I);
 }
 
 RecurrenceDescriptor::InstDesc

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

@@ -9203,7 +9203,7 @@ void LoopVectorizationPlanner::adjustRecipesForReductions(
       auto *Trunc =
           new VPWidenCastRecipe(Instruction::Trunc, NewExitingVPV, RdxTy);
       auto *Extnd =
-          RdxDesc.isSigned()
+          RdxDesc.isResultSigned()
               ? new VPWidenCastRecipe(Instruction::SExt, Trunc, PhiTy)
               : new VPWidenCastRecipe(Instruction::ZExt, Trunc, PhiTy);
 

llvm/lib/Transforms/Vectorize/VPlanRecipes.cpp

@@ -621,8 +621,9 @@ Value *VPInstruction::generate(VPTransformState &State) {
     unsigned UF = getNumOperands() - 2;
     Value *ReducedPartRdx = State.get(getOperand(2));
     for (unsigned Part = 1; Part < UF; ++Part) {
-      ReducedPartRdx = createMinMaxOp(Builder, RecurKind::SMax, ReducedPartRdx,
-                                      State.get(getOperand(2 + Part)));
+      ReducedPartRdx = createMinMaxOp(
+          Builder, RdxDesc.isReduxSigned() ? RecurKind::SMax : RecurKind::UMax,
+          ReducedPartRdx, State.get(getOperand(2 + Part)));
     }
 
     return createFindLastIVReduction(Builder, ReducedPartRdx,
@@ -697,7 +698,7 @@ Value *VPInstruction::generate(VPTransformState &State) {
       // If the reduction can be performed in a smaller type, we need to extend
       // the reduction to the wider type before we branch to the original loop.
       if (PhiTy != RdxDesc.getRecurrenceType())
-        ReducedPartRdx = RdxDesc.isSigned()
+        ReducedPartRdx = RdxDesc.isResultSigned()
                              ? Builder.CreateSExt(ReducedPartRdx, PhiTy)
                              : Builder.CreateZExt(ReducedPartRdx, PhiTy);
     }