[TargetLowering] Teach DemandedBits about VSCALE

[benmxwl-arm]

Summary

This allows DemandedBits to see the result of VSCALE will be at most
VScaleMax * some compile-time constant. This relies on the vscale_range()
attribute being present on the function, with a max set. (This is done by
default when clang is targeting AArch64+SVE).

Using this various redundant operations (zexts, sexts, ands, ors, etc)
can be eliminated.

[benmxwl-arm]

@dmgreen Thanks a lot for your review of D138424. Just adding you here as it's the natural continuation.

[paulwalker-arm]

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

1141–1145

ISD::VSCALE's operand is guaranteed to be a constant so you can use Op.getConstantOperandAPInt(0) * (*MaxVScale) without need to test it first.

[benmxwl-arm]

Remove unnecessary dyn_cast<ConstantSDNode>

[benmxwl-arm]

Done (and now rebased on main fixing up some Optional -> std::optional changes)

[paulwalker-arm]

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

1144

Would this be more accurate as ~? Also means we don't need to worry about the negation overflowing, which I doubt is a real concern anyway.

[benmxwl-arm]

Replaced negation with ~ and added a new test case where this makes a difference.
(Previosuly the redundant or in the new test case would not be removed)

[uabelho]

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

1141

With an UBSan build we get this

../lib/CodeGen/SelectionDAG/TargetLowering.cpp:1137:20: runtime error: signed integer overflow: 16 * -9223372036854775808 cannot be represented in type 'long'

when running lit tests

Failed Tests (1):
  LLVM :: CodeGen/AArch64/sve-gather-scatter-addr-opts.ll

I get it also on latest main (d2b070d3c95c5).

[compnerd]

I'm seeing a potential miscompilation on RISCV64 with this particular change: https://github.com/llvm/llvm-project/issues/59594

[compnerd]

Given the UB and potential miscompile, reverted in 9b92f70d4758f75903ce93feaba5098130820d40.

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

if (!AssumeSingleUse && !Op.getNode()->hasOneUse()) {
     return TLO.CombineTo(Op, TLO.DAG.getUNDEF(VT));
   } else if (Depth >= SelectionDAG::MaxRecursionDepth) {
     // Limit search depth.
     return false;
   }
 
   KnownBits Known2;
   switch (Op.getOpcode()) {
+  case ISD::VSCALE: {
+    Function const &F = TLO.DAG.getMachineFunction().getFunction();
+    Attribute const &Attr = F.getFnAttribute(Attribute::VScaleRange);
+    if (!Attr.isValid())
+      return false;
+    Optional<unsigned> MaxVScale = Attr.getVScaleRangeMax();
+    if (!MaxVScale.has_value())
+      return false;
+    int64_t VScaleResultUpperbound = *MaxVScale;
+    if (auto *MulImm = dyn_cast<ConstantSDNode>(Op.getOperand(0))) {
+      VScaleResultUpperbound *= MulImm->getSExtValue();
+    } else {
+      return false;
+    }
+    bool Negative = VScaleResultUpperbound < 0;
+    if (Negative)
+      VScaleResultUpperbound = -VScaleResultUpperbound;
+    unsigned RequiredBits = Log2_64(VScaleResultUpperbound) + 1;
+    if (RequiredBits < BitWidth)
+      (Negative ? Known.One : Known.Zero).setHighBits(BitWidth - RequiredBits);
+    return false;
+  }
   case ISD::SCALAR_TO_VECTOR: {
     if (!DemandedElts[0])
       return TLO.CombineTo(Op, TLO.DAG.getUNDEF(VT));
 
     KnownBits SrcKnown;
     SDValue Src = Op.getOperand(0);
     unsigned SrcBitWidth = Src.getScalarValueSizeInBits();
     APInt SrcDemandedBits = DemandedBits.zext(SrcBitWidth);

llvm/test/CodeGen/AArch64/vscale-and-sve-cnt-demandedbits.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=aarch64 -mattr=+sve < %s | FileCheck %s
 
 ; This tests that various ands, sexts, and zexts (and other operations)
 ; operating on vscale or the SVE count instructions can be eliminated
 ; (via demanded bits) due to their known limited range.
 
 ; On AArch64 vscale can be at most 16 (for a 2048-bit vector).
 ; The counting instructions (sans multiplier) have a value of at most 256
 ; (for a 2048-bit vector of i8s).
 
 define i32 @vscale_and_elimination() vscale_range(1,16) {
 ; CHECK-LABEL: vscale_and_elimination:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    rdvl x8, #1
 ; CHECK-NEXT:    lsr x8, x8, #4
-; CHECK-NEXT:    and w9, w8, #0x1f
-; CHECK-NEXT:    and w8, w8, #0xfffffffc
-; CHECK-NEXT:    add w0, w9, w8
+; CHECK-NEXT:    and w9, w8, #0x1c
+; CHECK-NEXT:    add w0, w8, w9
 ; CHECK-NEXT:    ret
   %vscale = call i32 @llvm.vscale.i32()
   %and_redundant = and i32 %vscale, 31
   %and_required = and i32 %vscale, 17179869180
   %result = add i32 %and_redundant, %and_required
   ret i32 %result
 }
 
; CHECK-NEXT:    ret
   %result = add i64 %and_redundant, %and_required
   ret i64 %result
 }
 
 define i64 @vscale_trunc_zext() vscale_range(1,16) {
 ; CHECK-LABEL: vscale_trunc_zext:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    rdvl x8, #1
-; CHECK-NEXT:    lsr x8, x8, #4
-; CHECK-NEXT:    and x0, x8, #0xffffffff
+; CHECK-NEXT:    lsr x0, x8, #4
 ; CHECK-NEXT:    ret
   %vscale = call i32 @llvm.vscale.i32()
   %zext = zext i32 %vscale to i64
   ret i64 %zext
 }
 
 define i64 @vscale_trunc_sext() vscale_range(1,16) {
 ; CHECK-LABEL: vscale_trunc_sext:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    rdvl x8, #1
-; CHECK-NEXT:    lsr x8, x8, #4
-; CHECK-NEXT:    sxtw x0, w8
+; CHECK-NEXT:    lsr x0, x8, #4
 ; CHECK-NEXT:    ret
   %vscale = call i32 @llvm.vscale.i32()
   %sext = sext i32 %vscale to i64
   ret i64 %sext
 }
 
 define i64 @count_bytes_trunc_zext() {
 ; CHECK-LABEL: count_bytes_trunc_zext:

 
 define i32 @vscale_with_multiplier() vscale_range(1,16) {
 ; CHECK-LABEL: vscale_with_multiplier:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    rdvl x8, #1
 ; CHECK-NEXT:    mov w9, #5
 ; CHECK-NEXT:    lsr x8, x8, #4
 ; CHECK-NEXT:    mul x8, x8, x9
-; CHECK-NEXT:    and w9, w8, #0x7f
-; CHECK-NEXT:    and w8, w8, #0x3f
-; CHECK-NEXT:    add w0, w9, w8
+; CHECK-NEXT:    and w9, w8, #0x3f
+; CHECK-NEXT:    add w0, w8, w9
 ; CHECK-NEXT:    ret
   %vscale = call i32 @llvm.vscale.i32()
   %mul = mul i32 %vscale, 5
   %and_redundant = and i32 %mul, 127
   %and_required = and i32 %mul, 63
   %result = add i32 %and_redundant, %and_required
   ret i32 %result
 }
 
 define i32 @vscale_with_negative_multiplier() vscale_range(1,16) {
 ; CHECK-LABEL: vscale_with_negative_multiplier:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    rdvl x8, #1
 ; CHECK-NEXT:    mov x9, #-5
 ; CHECK-NEXT:    lsr x8, x8, #4
 ; CHECK-NEXT:    mul x8, x8, x9
-; CHECK-NEXT:    orr w9, w8, #0xffffff80
-; CHECK-NEXT:    and w8, w8, #0xffffffc0
-; CHECK-NEXT:    add w0, w9, w8
+; CHECK-NEXT:    and w9, w8, #0xffffffc0
+; CHECK-NEXT:    add w0, w8, w9
 ; CHECK-NEXT:    ret
   %vscale = call i32 @llvm.vscale.i32()
   %mul = mul i32 %vscale, -5
   %or_redundant = or i32 %mul, 4294967168
   %or_required = and i32 %mul, 4294967232
   %result = add i32 %or_redundant, %or_required
   ret i32 %result
 }
 
 declare i32 @llvm.vscale.i32()
 declare i64 @llvm.aarch64.sve.cntb(i32 %pattern)
 declare i64 @llvm.aarch64.sve.cnth(i32 %pattern)
 declare i64 @llvm.aarch64.sve.cntw(i32 %pattern)
 declare i64 @llvm.aarch64.sve.cntd(i32 %pattern)