[ValueTracking] Support vscale in computeConstantRange()

[nikic]

Summary

Add support for vscale in computeConstantRange(), based on vscale_range attributes. This allows simplifying based on the precise range, rather than a KnownBits approximation (which will be off by a factor of two for the usual case of a power of two upper bound).

[sdesmalen]

llvm/lib/Analysis/ValueTracking.cpp

1160–1164

Is it possible to create test for this?

[nikic]

Rebase over additional tests for the poison case. In this case all comparisons just fold to true (but could also fold to any other value).

[reames]

LGTM

[sdesmalen]

In this case all comparisons just fold to true (but could also fold to any other value).

What is the reason that the IR doesn't explicitly return poison in that case? (i.e. ret i1 poison)

[nikic]

In D146217#4199478, @sdesmalen wrote:

In this case all comparisons just fold to true (but could also fold to any other value).

What is the reason that the IR doesn't explicitly return poison in that case? (i.e. ret i1 poison)

No particular reason, this just hasn't been implemented. (Would need a special case for empty ranges in simplifyICmpWithConstant.)

[sdesmalen]

In D146217#4199479, @nikic wrote:

In D146217#4199478, @sdesmalen wrote:

In this case all comparisons just fold to true (but could also fold to any other value).

What is the reason that the IR doesn't explicitly return poison in that case? (i.e. ret i1 poison)

No particular reason, this just hasn't been implemented. (Would need a special case for empty ranges in simplifyICmpWithConstant.)

Fair enough, I was mostly wondering if there was a fundamental reason for it or if my understanding of poison was lacking here :) Thanks for clarifying!

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) {
+  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)
+    return ConstantRange::getEmpty(BitWidth);
+
+  APInt Min(BitWidth, AttrMin);
+  std::optional<unsigned> AttrMax = Attr.getVScaleRangeMax();
+  if (!AttrMax || (unsigned)llvm::bit_width(*AttrMax) > BitWidth)
+    return ConstantRange(Min, APInt::getZero(BitWidth));
+
+  return ConstantRange(Min, APInt(BitWidth, *AttrMax) + 1);
+}
+
 static void computeKnownBitsFromOperator(const Operator *I,
                                          const APInt &DemandedElts,
                                          KnownBits &Known, unsigned Depth,
                                          const Query &Q) {
   unsigned BitWidth = Known.getBitWidth();
 
   KnownBits Known2(BitWidth);
   switch (I->getOpcode()) {
if (const IntrinsicInst *II = dyn_cast<IntrinsicInst>(I)) {
       case Intrinsic::riscv_vsetvli:
       case Intrinsic::riscv_vsetvlimax:
         // Assume that VL output is >= 65536.
         // TODO: Take SEW and LMUL into account.
         if (BitWidth > 17)
           Known.Zero.setBitsFrom(17);
         break;
       case Intrinsic::vscale: {
-        if (!II->getParent() || !II->getFunction() ||
-            !II->getFunction()->hasFnAttribute(Attribute::VScaleRange))
+        if (!II->getParent() || !II->getFunction())
           break;
 
-        auto Attr = II->getFunction()->getFnAttribute(Attribute::VScaleRange);
-        std::optional<unsigned> VScaleMax = Attr.getVScaleRangeMax();
-
-        if (!VScaleMax)
-          break;
-
-        unsigned VScaleMin = Attr.getVScaleRangeMin();
-
-        // If vscale min = max then we know the exact value at compile time
-        // and hence we know the exact bits.
-        if (VScaleMin == VScaleMax) {
-          Known.One = VScaleMin;
-          Known.Zero = VScaleMin;
-          Known.Zero.flipAllBits();
-          break;
-        }
-
-        unsigned FirstZeroHighBit = llvm::bit_width(*VScaleMax);
-        if (FirstZeroHighBit < BitWidth)
-          Known.Zero.setBitsFrom(FirstZeroHighBit);
-
+        Known = getVScaleRange(II->getFunction(), BitWidth).toKnownBits();
         break;
       }
       }
     }
     break;
   case Instruction::ShuffleVector: {
     auto *Shuf = dyn_cast<ShuffleVectorInst>(I);
     // FIXME: Do we need to handle ConstantExpr involving shufflevectors?
case Intrinsic::abs:
     // If abs of SIGNED_MIN is poison, then the result is [0..SIGNED_MAX],
     // otherwise it is [0..SIGNED_MIN], as -SIGNED_MIN == SIGNED_MIN.
     if (match(II.getOperand(1), m_One()))
       return ConstantRange(APInt::getZero(Width),
                            APInt::getSignedMaxValue(Width) + 1);
 
     return ConstantRange(APInt::getZero(Width),
                          APInt::getSignedMinValue(Width) + 1);
+  case Intrinsic::vscale:
+    if (!II.getParent() || !II.getFunction())
+      break;
+    return getVScaleRange(II.getFunction(), Width);
   default:
     break;
   }
 
   return ConstantRange::getFull(Width);
 }
 
 static void setLimitsForSelectPattern(const SelectInst &SI, APInt &Lower,

llvm/test/Transforms/InstCombine/icmp-vscale.ll

entry:
   %vscale = call i64 @llvm.vscale.i64()
   %num_els = mul i64 %vscale, 32
   %res = icmp ne i64 %num_els, 39488
   ret i1 %res
 }
 
 define i1 @vscale_ule_max() vscale_range(5,10) {
 ; CHECK-LABEL: @vscale_ule_max(
-; CHECK-NEXT:    [[VSCALE:%.*]] = call i16 @llvm.vscale.i16()
-; CHECK-NEXT:    [[RES:%.*]] = icmp ult i16 [[VSCALE]], 11
-; CHECK-NEXT:    ret i1 [[RES]]
+; CHECK-NEXT:    ret i1 true
 ;
   %vscale = call i16 @llvm.vscale.i16()
   %res = icmp ule i16 %vscale, 10
   ret i1 %res
 }
 
 define i1 @vscale_ult_max() vscale_range(5,10) {
 ; CHECK-LABEL: @vscale_ult_max(
 ; CHECK-NEXT:    [[VSCALE:%.*]] = call i16 @llvm.vscale.i16()
 ; CHECK-NEXT:    [[RES:%.*]] = icmp ult i16 [[VSCALE]], 10
 ; CHECK-NEXT:    ret i1 [[RES]]
 ;
   %vscale = call i16 @llvm.vscale.i16()
   %res = icmp ult i16 %vscale, 10
   ret i1 %res
 }
 
 define i1 @vscale_uge_min() vscale_range(5,10) {
 ; CHECK-LABEL: @vscale_uge_min(
-; CHECK-NEXT:    [[VSCALE:%.*]] = call i16 @llvm.vscale.i16()
-; CHECK-NEXT:    [[RES:%.*]] = icmp ugt i16 [[VSCALE]], 4
-; CHECK-NEXT:    ret i1 [[RES]]
+; CHECK-NEXT:    ret i1 true
 ;
   %vscale = call i16 @llvm.vscale.i16()
   %res = icmp uge i16 %vscale, 5
   ret i1 %res
 }
 
 define i1 @vscale_ugt_min() vscale_range(5,10) {
 ; CHECK-LABEL: @vscale_ugt_min(

 ;
   %vscale = call i8 @llvm.vscale.i8()
   %res = icmp ugt i8 %vscale, 5
   ret i1 %res
 }
 
 define i1 @vscale_uge_max_overflow() vscale_range(5,256) {
 ; CHECK-LABEL: @vscale_uge_max_overflow(
-; CHECK-NEXT:    [[VSCALE:%.*]] = call i8 @llvm.vscale.i8()
-; CHECK-NEXT:    [[RES:%.*]] = icmp ugt i8 [[VSCALE]], 4
-; CHECK-NEXT:    ret i1 [[RES]]
+; CHECK-NEXT:    ret i1 true
 ;
   %vscale = call i8 @llvm.vscale.i8()
   %res = icmp uge i8 %vscale, 5
   ret i1 %res
 }
 
 define i1 @vscale_ugt_max_overflow() vscale_range(5,256) {
 ; CHECK-LABEL: @vscale_ugt_max_overflow(
 ; CHECK-NEXT:    [[VSCALE:%.*]] = call i8 @llvm.vscale.i8()
 ; CHECK-NEXT:    [[RES:%.*]] = icmp ugt i8 [[VSCALE]], 5
 ; CHECK-NEXT:    ret i1 [[RES]]
 ;
   %vscale = call i8 @llvm.vscale.i8()
   %res = icmp ugt i8 %vscale, 5
   ret i1 %res
 }
 
 define i1 @vscale_eq_min_overflow() vscale_range(256,300) {
 ; CHECK-LABEL: @vscale_eq_min_overflow(
-; CHECK-NEXT:    [[VSCALE:%.*]] = call i8 @llvm.vscale.i8()
-; CHECK-NEXT:    [[RES:%.*]] = icmp eq i8 [[VSCALE]], 42
-; CHECK-NEXT:    ret i1 [[RES]]
+; CHECK-NEXT:    ret i1 true
 ;
   %vscale = call i8 @llvm.vscale.i8()
   %res = icmp eq i8 %vscale, 42
   ret i1 %res
 }
 
 define i1 @vscale_ult_min_overflow() vscale_range(256,300) {
 ; CHECK-LABEL: @vscale_ult_min_overflow(
-; CHECK-NEXT:    [[VSCALE:%.*]] = call i8 @llvm.vscale.i8()
-; CHECK-NEXT:    [[RES:%.*]] = icmp ult i8 [[VSCALE]], 42
-; CHECK-NEXT:    ret i1 [[RES]]
+; CHECK-NEXT:    ret i1 true
 ;
   %vscale = call i8 @llvm.vscale.i8()
   %res = icmp ult i8 %vscale, 42
   ret i1 %res
 }
 
 define i1 @vscale_ugt_min_overflow() vscale_range(256,300) {
 ; CHECK-LABEL: @vscale_ugt_min_overflow(
-; CHECK-NEXT:    [[VSCALE:%.*]] = call i8 @llvm.vscale.i8()
-; CHECK-NEXT:    [[RES:%.*]] = icmp ult i8 [[VSCALE]], 42
-; CHECK-NEXT:    ret i1 [[RES]]
+; CHECK-NEXT:    ret i1 true
 ;
   %vscale = call i8 @llvm.vscale.i8()
-  %res = icmp ult i8 %vscale, 42
+  %res = icmp ugt i8 %vscale, 42
   ret i1 %res
 }
 
 declare i8 @llvm.vscale.i8()
 declare i16 @llvm.vscale.i16()
 declare i32 @llvm.vscale.i32()
 declare i64 @llvm.vscale.i64()