[InstCombine] Add folds for certain icmp + vscale combinations

[david-arm]

Summary

Consider some simple cases like this:

icmp ugt i64 %vscale, 1024
icmp ugt i64 (mul i64 %vscale, 32), 1024

where

%vscale = call i64 @llvm.vscale.i64()

We may know at compile time what the maximum value of vscale is
and therefore we may know the result of the comparison. In such cases
we can fold examples like those above away completely to return a
boolean value.

Tests added here:

Transforms/InstCombine/icmp-vscale.ll

[dmgreen]

Could we add vscale to ValueTracking::computeKnownBitsFromOperator? I'm not sure it would capture all the non-power-2 cases, but it would be more general and might optimize more cases.

[david-arm]

In D109696#2997656, @dmgreen wrote:

Could we add vscale to ValueTracking::computeKnownBitsFromOperator? I'm not sure it would capture all the non-power-2 cases, but it would be more general and might optimize more cases.

That's a nice suggestion @dmgreen I'll look into it!

[david-arm]

• Rebased off D109883

[david-arm]

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

1433

Hi @dmgreen, I just realised that I can probably remove the check for m_Shl now as this should be covered by D109883

[david-arm]

• Removed folds involving shl as these now work already.

[david-arm]

This patch is probably not worth it anymore, since the cases it deals with are unlikely to occur in practice.

llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp

if (match(Cmp.getOperand(0), m_URem(m_Value(X), m_Value(Y))) &&
     KnownBits YKnown = computeKnownBits(Y, 0, &Cmp);
     if (XKnown.countMaxPopulation() == 1 && YKnown.countMinPopulation() >= 2)
       return new ICmpInst(Pred, X, Cmp.getOperand(1));
   }
 
   return nullptr;
 }
 
+/// Fold icmp involving trivial vscale calculations, e.g.:
+///   icmp Pred vscale, C
+///   icmp Pred (mul vscale, X), C
+///   icmp Pred (shl vscale, X), C
+/// We can use the maximum value of vscale if it exists in order to determine
+/// the result of certain cases at compile time.
+Instruction *InstCombinerImpl::foldICmpWithConstantVScale(ICmpInst &Cmp) {
+  CmpInst::Predicate Pred = Cmp.getPredicate();
+  Value *Op0 = Cmp.getOperand(0), *Op1 = Cmp.getOperand(1);
+
+  ConstantInt *CI = dyn_cast<ConstantInt>(Op1);
+  if (!CI)
+    return nullptr;
+
+  ConstantInt *VScaleMul = nullptr;
+  if (Pred == ICmpInst::ICMP_UGT &&
+      (match(Op0, m_VScale(DL)) ||
+       match(Op0, m_Mul(m_VScale(DL), m_ConstantInt(VScaleMul))))) {
+    // If we don't know the max value of vscale we can't fold anything.
+    if (!Cmp.getFunction()->hasFnAttribute(Attribute::VScaleRange))
+      return nullptr;
+
+    uint64_t MaxVScale = Cmp.getFunction()
+                             ->getFnAttribute(Attribute::VScaleRange)
+                             .getVScaleRangeArgs()
+                             .second;
+    if (MaxVScale == 0)
+      return nullptr;
+
+    // We know the maximum value of vscale * multipler and so we may
+    // be able to completely eliminate the icmp.
+    uint64_t MulVal = VScaleMul ? VScaleMul->getZExtValue() : 1;
+    uint64_t MaxLHSVal = MaxVScale * MulVal;
+
+    if (MaxLHSVal <= CI->getZExtValue())
+      return replaceInstUsesWith(Cmp, ConstantInt::get(Cmp.getType(), 0));
+  }
+
+  return nullptr;
+}
+
 /// Fold icmp Pred X, C.
 /// TODO: This code structure does not make sense. The saturating add fold
 /// should be moved to some other helper and extended as noted below (it is also
 /// possible that code has been made unnecessary - do we canonicalize IR to
 /// overflow/saturating intrinsics or not?).
 Instruction *InstCombinerImpl::foldICmpWithConstant(ICmpInst &Cmp) {
   // Match the following pattern, which is a common idiom when writing
   // overflow-safe integer arithmetic functions. The source performs an addition
if (Instruction *Res = canonicalizeCmpWithConstant(I))
     return Res;
 
   if (Instruction *Res = canonicalizeICmpPredicate(I))
     return Res;
 
   if (Instruction *Res = foldICmpWithConstant(I))
     return Res;
 
+  if (Instruction *Res = foldICmpWithConstantVScale(I))
+    return Res;
+
   if (Instruction *Res = foldICmpWithDominatingICmp(I))
     return Res;
 
   if (Instruction *Res = foldICmpUsingKnownBits(I))
     return Res;
 
   // Test if the ICmpInst instruction is used exclusively by a select as
   // part of a minimum or maximum operation. If so, refrain from doing

llvm/lib/Transforms/InstCombine/InstCombineInternal.h

Instruction *foldFCmpIntToFPConst(FCmpInst &I, Instruction *LHSI,
                                     Constant *RHSC);
   Instruction *foldICmpAddOpConst(Value *X, const APInt &C,
                                   ICmpInst::Predicate Pred);
   Instruction *foldICmpWithCastOp(ICmpInst &ICI);
 
   Instruction *foldICmpUsingKnownBits(ICmpInst &Cmp);
   Instruction *foldICmpWithDominatingICmp(ICmpInst &Cmp);
   Instruction *foldICmpWithConstant(ICmpInst &Cmp);
+  Instruction *foldICmpWithConstantVScale(ICmpInst &Cmp);
   Instruction *foldICmpInstWithConstant(ICmpInst &Cmp);
   Instruction *foldICmpInstWithConstantNotInt(ICmpInst &Cmp);
   Instruction *foldICmpBinOp(ICmpInst &Cmp, const SimplifyQuery &SQ);
   Instruction *foldICmpEquality(ICmpInst &Cmp);
   Instruction *foldIRemByPowerOfTwoToBitTest(ICmpInst &I);
   Instruction *foldSignBitTest(ICmpInst &I);
   Instruction *foldICmpWithZero(ICmpInst &Cmp);
 

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

 ;
 entry:
   %vscale = call i16 @llvm.vscale.i16()
   %num_els = shl i16 %vscale, 5
   %res = icmp ugt i16 %num_els, 1024
   ret i1 %res
 }
 
+define i1 @uge_vscale8_x_4() vscale_range(0,16) {
+; CHECK-LABEL: @uge_vscale8_x_4(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    ret i1 false
+;
+entry:
+  %vscale = call i8 @llvm.vscale.i8()
+  %num_els = mul i8 %vscale, 5
+  %res = icmp uge i8 %num_els, 128
+  ret i1 %res
+}
+
 define i1 @ult_vscale16() vscale_range(0,16) {
 ; CHECK-LABEL: @ult_vscale16(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    ret i1 false
 ;
 entry:
   %vscale = call i16 @llvm.vscale.i16()
   %res = icmp ult i16 1024, %vscale

 ;
 entry:
   %vscale = call i64 @llvm.vscale.i64()
   %num_els = mul i64 %vscale, 32
   %res = icmp ne i64 %num_els, 39488
   ret i1 %res
 }
 
+; Negative tests
+
+define i1 @fail_uge_vscale64_x_32_max32() vscale_range(0,32) {
+; CHECK-LABEL: @fail_uge_vscale64_x_32_max32(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[VSCALE:%.*]] = call i64 @llvm.vscale.i64()
+; CHECK-NEXT:    [[RES:%.*]] = icmp ugt i64 [[VSCALE]], 31
+; CHECK-NEXT:    ret i1 [[RES]]
+;
+entry:
+  %vscale = call i64 @llvm.vscale.i64()
+  %num_els = shl i64 %vscale, 5
+  %res = icmp uge i64 %num_els, 1024
+  ret i1 %res
+}
+
+define i1 @fail_uge_vscale64_x_32_max0() vscale_range(0,0) {
+; CHECK-LABEL: @fail_uge_vscale64_x_32_max0(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[VSCALE:%.*]] = call i64 @llvm.vscale.i64()
+; CHECK-NEXT:    [[NUM_ELS:%.*]] = mul i64 [[VSCALE]], 31
+; CHECK-NEXT:    [[RES:%.*]] = icmp ugt i64 [[NUM_ELS]], 12345677
+; CHECK-NEXT:    ret i1 [[RES]]
+;
+entry:
+  %vscale = call i64 @llvm.vscale.i64()
+  %num_els = mul i64 %vscale, 31
+  %res = icmp uge i64 %num_els, 12345678
+  ret i1 %res
+}
+
+define i1 @fail_uge_vscale64_x_val(i64 %val) vscale_range(0,32) {
+; CHECK-LABEL: @fail_uge_vscale64_x_val(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[VSCALE:%.*]] = call i64 @llvm.vscale.i64()
+; CHECK-NEXT:    [[NUM_ELS:%.*]] = mul i64 [[VSCALE]], [[VAL:%.*]]
+; CHECK-NEXT:    [[RES:%.*]] = icmp ugt i64 [[NUM_ELS]], 1023
+; CHECK-NEXT:    ret i1 [[RES]]
+;
+entry:
+  %vscale = call i64 @llvm.vscale.i64()
+  %num_els = mul i64 %vscale, %val
+  %res = icmp uge i64 %num_els, 1024
+  ret i1 %res
+}
+
 declare i8 @llvm.vscale.i8()
 declare i16 @llvm.vscale.i16()
 declare i32 @llvm.vscale.i32()
 declare i64 @llvm.vscale.i64()