[SCEV] Recognize vscale intrinsics

[nikic]

Summary

Now that SCEV has a dedicated vscale node type, we should also map vscale intrinsics to it. To make sure this does not regress ranges (which were KnownBits based previously), add support for vscale to getRangeRef() as well.

Depends on D146217.

[reames]

LGTM w/minor comment

llvm/include/llvm/Analysis/ValueTracking.h

657

Can you rephrase this comment? This doesn't just bound the vscale intrinsic, it also bounds the vscale type param.

llvm/include/llvm/Analysis/ValueTracking.h

OverflowResult computeOverflowForSignedSub(const Value *LHS, const Value *RHS,
                                            const DominatorTree *DT);
 
 /// Returns true if the arithmetic part of the \p WO 's result is
 /// used only along the paths control dependent on the computation
 /// not overflowing, \p WO being an <op>.with.overflow intrinsic.
 bool isOverflowIntrinsicNoWrap(const WithOverflowInst *WO,
                                const DominatorTree &DT);
 
+/// Determine the possible constant range of vscale with the given bit width,
+/// based on the vscale_range function attribute.
+ConstantRange getVScaleRange(const Function *F, unsigned BitWidth);
+
 /// Determine the possible constant range of an integer or vector of integer
 /// value. This is intended as a cheap, non-recursive check.
 ConstantRange computeConstantRange(const Value *V, bool ForSigned,
                                    bool UseInstrInfo = true,
                                    AssumptionCache *AC = nullptr,
                                    const Instruction *CtxI = nullptr,
                                    const DominatorTree *DT = nullptr,
                                    unsigned Depth = 0);

llvm/lib/Analysis/ScalarEvolution.cpp

else
           APInt::getSignedMinValue(BitWidth),
           APInt::getSignedMaxValue(BitWidth).ashr(TZ).shl(TZ) + 1);
   }
 
   switch (S->getSCEVType()) {
   case scConstant:
     llvm_unreachable("Already handled above.");
   case scVScale:
-    return setRange(S, SignHint, std::move(ConservativeResult));
+    return setRange(S, SignHint, getVScaleRange(&F, BitWidth));
   case scTruncate: {
     const SCEVTruncateExpr *Trunc = cast<SCEVTruncateExpr>(S);
     ConstantRange X = getRangeRef(Trunc->getOperand(), SignHint, Depth + 1);
     return setRange(
         Trunc, SignHint,
         ConservativeResult.intersectWith(X.truncate(BitWidth), RangeType));
   }
   case scZeroExtend: {
if (auto *II = dyn_cast<IntrinsicInst>(U)) {
         return getAddExpr(ClampedX, Y, SCEV::FlagNUW);
       }
       case Intrinsic::start_loop_iterations:
       case Intrinsic::annotation:
       case Intrinsic::ptr_annotation:
         // A start_loop_iterations or llvm.annotation or llvm.prt.annotation is
         // just eqivalent to the first operand for SCEV purposes.
         return getSCEV(II->getArgOperand(0));
+      case Intrinsic::vscale:
+        return getVScale(II->getType());
       default:
         break;
       }
     }
     break;
   }
 
   return getUnknown(V);

llvm/lib/Analysis/ValueTracking.cpp

KnownBits llvm::analyzeKnownBitsFromAndXorOr(
   APInt DemandedElts =
       FVTy ? APInt::getAllOnes(FVTy->getNumElements()) : APInt(1, 1);
 
   return getKnownBitsFromAndXorOr(
       I, DemandedElts, KnownLHS, KnownRHS, Depth,
       Query(DL, AC, safeCxtI(I, CxtI), DT, UseInstrInfo, ORE));
 }
 
-static ConstantRange getVScaleRange(const Function *F, unsigned BitWidth) {
+ConstantRange llvm::getVScaleRange(const Function *F, unsigned BitWidth) {
   Attribute Attr = F->getFnAttribute(Attribute::VScaleRange);
   // Without vscale_range, we only know that vscale is non-zero.
   if (!Attr.isValid())
     return ConstantRange(APInt(BitWidth, 1), APInt::getZero(BitWidth));
 
   unsigned AttrMin = Attr.getVScaleRangeMin();
   // Minimum is larger than vscale width, result is always poison.
   if ((unsigned)llvm::bit_width(AttrMin) > BitWidth)

llvm/test/Analysis/ScalarEvolution/scalable-vector.ll

;
   getelementptr <vscale x 1 x i64>, ptr %p, i32 1
   ret void
 }
 
 define void @vscale_gep_range(ptr %p) vscale_range(2, 16) {
 ; CHECK-LABEL: 'vscale_gep_range'
 ; CHECK-NEXT:  Classifying expressions for: @vscale_gep_range
 ; CHECK-NEXT:    %1 = getelementptr <vscale x 4 x i32>, ptr null, i32 3
-; CHECK-NEXT:    --> ((48 * vscale) + null) U: [0,-15) S: [-9223372036854775808,9223372036854775793)
+; CHECK-NEXT:    --> ((48 * vscale)<nuw><nsw> + null) U: [96,769) S: [96,769)
 ; CHECK-NEXT:    %2 = getelementptr <vscale x 1 x i64>, ptr %p, i32 1
-; CHECK-NEXT:    --> ((8 * vscale) + %p) U: full-set S: full-set
+; CHECK-NEXT:    --> ((8 * vscale)<nuw><nsw> + %p) U: full-set S: full-set
 ; CHECK-NEXT:  Determining loop execution counts for: @vscale_gep_range
 ;
   getelementptr <vscale x 4 x i32>, ptr null, i32 3
   getelementptr <vscale x 1 x i64>, ptr %p, i32 1
   ret void
 }
 
 define i64 @vscale_no_range() {
 ; CHECK-LABEL: 'vscale_no_range'
 ; CHECK-NEXT:  Classifying expressions for: @vscale_no_range
 ; CHECK-NEXT:    %vscale = call i64 @llvm.vscale.i64()
-; CHECK-NEXT:    --> %vscale U: [1,0) S: [1,0)
+; CHECK-NEXT:    --> vscale U: [1,0) S: [1,0)
 ; CHECK-NEXT:  Determining loop execution counts for: @vscale_no_range
 ;
   %vscale = call i64 @llvm.vscale.i64()
   ret i64 %vscale
 }
 
 define i64 @vscale_min_max_range() vscale_range(2, 16) {
 ; CHECK-LABEL: 'vscale_min_max_range'
 ; CHECK-NEXT:  Classifying expressions for: @vscale_min_max_range
 ; CHECK-NEXT:    %vscale = call i64 @llvm.vscale.i64()
-; CHECK-NEXT:    --> %vscale U: [1,32) S: [1,32)
+; CHECK-NEXT:    --> vscale U: [2,17) S: [2,17)
 ; CHECK-NEXT:  Determining loop execution counts for: @vscale_min_max_range
 ;
   %vscale = call i64 @llvm.vscale.i64()
   ret i64 %vscale
 }
 
 define i64 @vscale_min_range() vscale_range(2, 0) {
 ; CHECK-LABEL: 'vscale_min_range'
 ; CHECK-NEXT:  Classifying expressions for: @vscale_min_range
 ; CHECK-NEXT:    %vscale = call i64 @llvm.vscale.i64()
-; CHECK-NEXT:    --> %vscale U: [1,0) S: [1,0)
+; CHECK-NEXT:    --> vscale U: [2,0) S: [2,0)
 ; CHECK-NEXT:  Determining loop execution counts for: @vscale_min_range
 ;
   %vscale = call i64 @llvm.vscale.i64()
   ret i64 %vscale
 }
 
 define i64 @vscale_exact_range() vscale_range(2) {
 ; CHECK-LABEL: 'vscale_exact_range'
 ; CHECK-NEXT:  Classifying expressions for: @vscale_exact_range
 ; CHECK-NEXT:    %vscale = call i64 @llvm.vscale.i64()
-; CHECK-NEXT:    --> %vscale U: [2,3) S: [2,3)
+; CHECK-NEXT:    --> vscale U: [2,3) S: [2,3)
 ; CHECK-NEXT:  Determining loop execution counts for: @vscale_exact_range
 ;
   %vscale = call i64 @llvm.vscale.i64()
   ret i64 %vscale
 }
 
 declare i64 @llvm.vscale.i64()