SelectionDAG: Teach ComputeKnownBits about VSCALE

[craig.topper]

Summary

This reverts commit 9b92f70d4758f75903ce93feaba5098130820d40. The issue
with the re-applied change was an implicit truncation due to the
multiplication. Although the operations were converted to APInt, the
values were implicitly converted to long due to the typing rules.

Fixes: #59594

[jrtc27]

Please clean up your RISCV test to match the style (use riscv64 as the triple, use UTC, don't put the triple and datalayout in the IR, clean up all the verbose comments, clean up metadata and attributes)

[MaskRay]

Revert "Revert "Reland

Avoid double negation and perhaps just use the original subject with the original commit message.

You can attach a paragraph describing what is fixed in the new patch.

[craig.topper]

llvm/test/CodeGen/RISCV/vscale-demanded-bits.ll

48

Don't you need a vscale_range attribute to hit the bug?

[craig.topper]

Where was there an implicit conversion to long? There is no implicit conversion operator on APInt.

It looks to me like this multiply overflows with vscale max and the APInt got a value of 0 causing the required bits to be 0.

%4 = tail call i8 @llvm.vscale.i8()                                            
%5 = shl i8 %4, 3

[jrtc27]

llvm/test/CodeGen/RISCV/vscale-demanded-bits.ll

2

4

Use UTC, really

6–8

9

I tend to prefer tests that don't have mangled C++ names in them

[craig.topper]

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

1136 ↗

(On Diff #484071)

I'm not sure that hardcoding 64 here is the right fix. I think what we need to check is that the multiply doesn't overflow.

[compnerd]

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

1136 ↗

(On Diff #484071)

@craig.topper I absolutely agree with you. I call this out in the commit message, that I'm just assuming that this is wide enough. This can potentially overflow still. Is there a good way to ensure that we have the proper width?

llvm/test/CodeGen/RISCV/vscale-demanded-bits.ll

4

What do you mean by UTC?

9

Sure, I can rename it, don't really think that it matters though.

[compnerd]

llvm/test/CodeGen/RISCV/vscale-demanded-bits.ll

48

I thought so, except, removing the attributes somehow still reproduces the difference? I don't think that the reduction is the best, and I'm happy to restore the attribute if you like.

[jrtc27]

llvm/test/CodeGen/RISCV/vscale-demanded-bits.ll

4

UpdateTestChecks, which in this case is update_llc_test_checks.py

[craig.topper]

llvm/test/CodeGen/RISCV/vscale-demanded-bits.ll

48

I thought in my testing that it won't make it to the APInt multiply without the attribute?

[MacDue]

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

1136 ↗

(On Diff #484071)

My thought was to do:

APInt Multiplier = Op.getConstantOperandAPInt(0);
unsigned MultiplyBits = Log2_32(*MaxVScale) + 1 + Multiplier.getActiveBits();
APInt VScaleResultUpperbound = APInt(MultiplyBits, *MaxVScale) * Multiplier.sextOrTrunc(MultiplyBits);

Which should always have enough bits to store the full result.

[craig.topper]

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

1136 ↗

(On Diff #484071)

Here's what I was playing with

APInt VScaleResultUpperbound(BitWidth, *MaxVScale);
// TODO: Vscale max with no leading zeros requires special handling.
if (VScaleResultUpperbound.isNegative())
  return false;
bool Overflow;
VScaleResultUpperbound =
    VScaleResultUpperbound.smul_ov(Op.getConstantOperandAPInt(0), Overflow);
if (Overflow)
  return false;
if (VScaleResultUpperbound.isNegative())
  Known.One.setHighBits(VScaleResultUpperbound.countLeadingOnes());
else
  Known.Zero.setHighBits(VScaleResultUpperbound.countLeadingZeros());

[craig.topper]

BTW, why is this in SimplifyDemandedBits instead of computeKnownBits? It doesn't use the DemandedBits.

[compnerd]

In D140347#4006583, @craig.topper wrote:

BTW, why is this in SimplifyDemandedBits instead of computeKnownBits? It doesn't use the DemandedBits.

I couldn't answer that tbh, perhaps @benmxwl-arm can answer that. I am only trying to help restore the patch that I reverted due to it regressing RVV.

llvm/test/CodeGen/RISCV/vscale-demanded-bits.ll

48

Ugh, I think I mixed up the command line outputs and thought it wasn't needed. It is definitely needed on the test function. Thanks for flagging that!

[compnerd]

Partially address feedback from Craig, address feedback from Jessica.

[compnerd]

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

1136 ↗

(On Diff #484071)

I didn't know about the smul_ov function, that does seem better. The use of BitWidth is still a problem. The BitWidth is computed from the width of the DemandedBits which is computed from the type of the result of the operation. In this particular test (luckily identifies the issue!) the result type is i8, which results in the BitWidth being 8, but the vrange is 2,1024, which exceeds the range of i8, and thus when we do scale * multiplier and overflow, we end up with the wrong result. I think that we should either go with the wider bit width here.

[craig.topper]

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

1136 ↗

(On Diff #484071)

You're right. I would need to at least check that the vscale max fits in the bitwidth.

[compnerd]

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

1136 ↗

(On Diff #484071)

Can we somehow be assured that the multiplication between the vscale max and the scalar will not overflow? It seems that we need a wider width than the vscale max unless I'm not counting bits correctly.

llvm/test/CodeGen/RISCV/vscale-demanded-bits.ll

4

Done, but this really does seem unnecessarily stringent testing. What we are looking to verify here is that the index counter is updated, the rest of the code around it really doesn't matter.

[compnerd]

Further reduce test case by creating a hand written synthetic test case. We can get away with a 2-iteration unrolled loop avoiding any phi branches. The previous overflow would improperly remove the increment in between the two iterations.

[compnerd]

Further reduce the test case. We create the splat and shuffle, offset it by vlenb and return the generated splat. This reduces down to the minimal assembly sequence that would exhibit the issue. If we were to accidentally truncate the scale, we would drop the read of vlenb and the adjustment for the shuffle.

[tschuett]

There is as trySextValue in APInt: https://reviews.llvm.org/D139683. Maybe you need a tryMul.

[compnerd]

In D140347#4009008, @tschuett wrote:

There is as trySextValue in APInt: https://reviews.llvm.org/D139683. Maybe you need a tryMul.

From a look, that seems to do what this was doing - extending to 64-bits. It feels that both Craig and I want a tighter bounds on the bit width to compute the scale. But, I can certainly see the value in having a tryMul to mirror the trySextValue.

[benmxwl-arm]

In D140347#4008253, @compnerd wrote:

In D140347#4006583, @craig.topper wrote:

BTW, why is this in SimplifyDemandedBits instead of computeKnownBits? It doesn't use the DemandedBits.

I couldn't answer that tbh, perhaps @benmxwl-arm can answer that. I am only trying to help restore the patch that I reverted due to it regressing RVV.

Sorry, that's my bad. It's perfectly fine to move this over to computeKnownBits().
Just replace return falses with breaks and TLO.DAG.getMachineFunction() with getMachineFunction() and it works fine moving this case over there.

[foad]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

3040–3051

This seems pretty complex. Is the Multiplier operand guaranteed to have the same width as the result, i.e. BitWidth? If so you should be able to do all calculation in that width and the signedness of the multiplier should be irrelevant.

[compnerd]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

3040–3051

No AFAIK, the multiplier has no guarantees of the width. It may be 32-bit, or 64-bit, though LLVM did truncate it further.

[foad]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

3040–3051

Then the documentation for ISD::VSCALE should be clear about whether the multiplier is treated as signed or unsigned. I don't see it mentioned in ISDOpcodes.h.

Anyway, if it is signed, can't you first sextOrTrunc it to BitWdith, and then do all of the rest of the calculations in that width?

[compnerd]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

3040–3051

My interpretation is that it is signed. The LangRef calls it out as being a positive number, thus, it must be signed. As to truncating it to BitWidth, again, that doesn't work. We need it to be the computed width as we need to ensure that we have sufficient space for the multiplication to not overflow (which was the original bug) and we need it to match to perform the operation (requirements from APInt). Is there a way to write this in a simpler way that I don't know about?

[foad]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

3040–3051

The LangRef calls it out as being a positive number

No, that's talking about vscale itself. The multiplier is not mentioned in LangRef because it is specific to the ISD::VSCALE sdag node.

Anyway does this work:

Known.Zero.setBitsFrom(Log2_32(*MaxVScale) + 1);
Known = KnownBits::mul(Known, KnownBits::makeConstant(Multiplier.sextOrTrunc(BitWidth)));

?

[compnerd]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

3040–3051

void llvm::APInt::setBits(unsigned int, unsigned int): Assertion `loBit <= BitWidth && "loBit out of range"' failed.

This is why I had kept the "complex" code as is. I had tried a few different things, I'd really rather prefer the smul_ov rather than the current implementation, but it seemed more complicated than this long-winded way.

[craig.topper]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

3040–3051

The issue here is that we don't that vscale is *MaxVscale. We're trying to calculate an upper bound on the result of the multiply to determine the value of sign bits only.

[foad]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

3040–3051

Fixing the assertion is easy:

Known.Zero.setBitsFrom(std::min(Log2_32(*MaxVScale) + 1, BitWidth));
Known = KnownBits::mul(Known, KnownBits::makeConstant(Multiplier.sextOrTrunc(BitWidth)));

[benmxwl-arm]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

3040–3051

KnownBits::mul() does not perform the same operation (and will under-report the know bits) and does not work for negative multipliers.

(I did try using KnownBits::mul() in an earlier version of this patch)

[compnerd]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

3040–3051

Fixing the assertion is easy:

Truncating the value to BitWidth seems incorrect. But also, as @benmxwl-arm mentions, that isn't exactly the same result.

[foad]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

3040–3051

KnownBits::mul() does not perform the same operation (and will under-report the know bits)

Ack. I understand now that you want a result based on the actual maximum value of vscale, not just the known bits of vscale. Maybe value range propagation would be a better tool for this, but I guess we don't have that in SelectionDAG.

[benmxwl-arm]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

3040

I think this actually should be Multiplier.getSignificantBits() (which takes into account the sign).
Other than that this LGTM, though I'd wait for approval from a more experienced reviewer.

[compnerd]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

3040

Hmm, I think that it _should_ be the same given that Multiplier should never be negative. But I don't see any harm in changing this to Multiplier.getSignificantBits() as it may be less confusing, so seems like a good idea.

[craig.topper]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

3040

Why can't multiplier be negative?

[benmxwl-arm]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

3040

I think the multiplier can be negative (see the _with_negative_multiplier tests). Maybe you're thinking about the vscale intrinsic? The intrinsic has no multiplier and always returns a positive value, the DAG node can be the result of merging a vscale intrinsic and a multiply, and I don't think there's any restriction on the multiplier value or sign.

[compnerd]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

3040

Yes, I was thinking of the vscale intrinsic. Okay, in that case, yes, this should be getSignificantBits.

[compnerd]

Use getSignificantBits

[foad]

I think the multiplier can be negative (see the _with_negative_multiplier tests).

This needs documenting. Speciflcally, for the ISD::VSCALE node, if the output is wider than the input, does it zero- or sign-extend the input? (Or to put it another way, does it do an unsigned extending multiply or a signed extending multiply?)

[craig.topper]

In D140347#4030800, @foad wrote:

I think the multiplier can be negative (see the _with_negative_multiplier tests).

This needs documenting. Speciflcally, for the ISD::VSCALE node, if the output is wider than the input, does it zero- or sign-extend the input? (Or to put it another way, does it do an unsigned extending multiply or a signed extending multiply?)

I'm not sure it's allowed to have a different type. In tablegen it is using SDTIntUnaryOp which requires the input and output types to match.

[craig.topper]

I'm going to take this over and finish it.

[craig.topper]

Use smul_ov instead of trying to size for the worst case.

[nikic]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

3130

Can this whole block be reduced to something like this?

const Function &F = getMachineFunction().getFunction();
const APInt &Multiplier = Op.getConstantOperandAPInt(0);
Known = getVScaleConstantRange(&F, BitWidth).mul(Multiplier).toKnownBits();

[craig.topper]

Use ConstantRange.

[nikic]

LGTM

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

KnownBits SelectionDAG::computeKnownBits(SDValue Op, const APInt &DemandedElts,
   case ISD::MERGE_VALUES:
     return computeKnownBits(Op.getOperand(Op.getResNo()), DemandedElts,
                             Depth + 1);
   case ISD::SPLAT_VECTOR: {
     SDValue SrcOp = Op.getOperand(0);
     assert(SrcOp.getValueSizeInBits() >= BitWidth &&
            "Expected SPLAT_VECTOR implicit truncation");
     // Implicitly truncate the bits to match the official semantics of
     // SPLAT_VECTOR.
     Known = computeKnownBits(SrcOp, Depth + 1).trunc(BitWidth);
     break;
   }
   case ISD::BUILD_VECTOR:
     assert(!Op.getValueType().isScalableVector());
     // Collect the known bits that are shared by every demanded vector element.
     Known.Zero.setAllBits(); Known.One.setAllBits();
     for (unsigned i = 0, e = Op.getNumOperands(); i != e; ++i) {
       if (!DemandedElts[i])
         continue;
 
       SDValue SrcOp = Op.getOperand(i);
       Known2 = computeKnownBits(SrcOp, Depth + 1);
 
       // BUILD_VECTOR can implicitly truncate sources, we must handle this.
       if (SrcOp.getValueSizeInBits() != BitWidth) {
         assert(SrcOp.getValueSizeInBits() > BitWidth &&
                "Expected BUILD_VECTOR implicit truncation");
         Known2 = Known2.trunc(BitWidth);
if (Known.isUnknown())
       break;
     if (!!DemandedRHS) {
       SDValue RHS = Op.getOperand(1);
       Known2 = computeKnownBits(RHS, DemandedRHS, Depth + 1);
       Known = Known.intersectWith(Known2);
     }
     break;
   }
+  case ISD::VSCALE: {
+    const Function &F = getMachineFunction().getFunction();
+    const APInt &Multiplier = Op.getConstantOperandAPInt(0);
+    Known = getVScaleRange(&F, BitWidth).multiply(Multiplier).toKnownBits();
+    break;
+  }
   case ISD::CONCAT_VECTORS: {
     if (Op.getValueType().isScalableVector())
       break;
     // Split DemandedElts and test each of the demanded subvectors.
     Known.Zero.setAllBits(); Known.One.setAllBits();
     EVT SubVectorVT = Op.getOperand(0).getValueType();
     unsigned NumSubVectorElts = SubVectorVT.getVectorNumElements();
     unsigned NumSubVectors = Op.getNumOperands();
for (unsigned i = 0; i != NumSubVectors; ++i) {
         break;
     }
     break;
   }
   case ISD::INSERT_SUBVECTOR: {
     if (Op.getValueType().isScalableVector())
       break;
     // Demand any elements from the subvector and the remainder from the src its
     // inserted into.
     SDValue Src = Op.getOperand(0);
     SDValue Sub = Op.getOperand(1);
     uint64_t Idx = Op.getConstantOperandVal(2);
     unsigned NumSubElts = Sub.getValueType().getVectorNumElements();
     APInt DemandedSubElts = DemandedElts.extractBits(NumSubElts, Idx);
     APInt DemandedSrcElts = DemandedElts;
     DemandedSrcElts.insertBits(APInt::getZero(NumSubElts), Idx);
 

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
 }
 
+define i32 @pow2_vscale_with_negative_multiplier() vscale_range(1,16) {
+; CHECK-LABEL: pow2_vscale_with_negative_multiplier:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    cntd x8
+; CHECK-NEXT:    neg x8, x8
+; CHECK-NEXT:    orr w9, w8, #0xfffffff0
+; CHECK-NEXT:    add w0, w8, w9
+; CHECK-NEXT:    ret
+  %vscale = call i32 @llvm.vscale.i32()
+  %mul = mul i32 %vscale, -2
+  %or_redundant = or i32 %mul, 4294967264
+  %or_required = or i32 %mul, 4294967280
+  %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)

llvm/test/CodeGen/RISCV/vscale-demanded-bits.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple riscv64 -mattr +v -filetype asm -o - %s | FileCheck %s
+
+declare i8 @llvm.vscale.i8()
+declare <vscale x 8 x i8> @llvm.experimental.stepvector.nxv8i8()
+
+define <vscale x 8 x i8> @f() #0 {
+; CHECK-LABEL: f:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    csrr a0, vlenb
+; CHECK-NEXT:    vsetvli a1, zero, e8, m1, ta, ma
+; CHECK-NEXT:    vid.v v8
+; CHECK-NEXT:    vadd.vx v8, v8, a0
+; CHECK-NEXT:    ret
+entry:
+  %0 = tail call i8 @llvm.vscale.i8()
+  %1 = shl i8 %0, 3
+  %.splat.insert = insertelement <vscale x 8 x i8> poison, i8 %1, i64 0
+  %.splat = shufflevector <vscale x 8 x i8> %.splat.insert, <vscale x 8 x i8> poison, <vscale x 8 x i32> zeroinitializer
+  %2 = tail call <vscale x 8 x i8> @llvm.experimental.stepvector.nxv8i8()
+  %3 = add <vscale x 8 x i8> %2, %.splat
+  ret <vscale x 8 x i8> %3
+}
+
+attributes #0 = { vscale_range(2,1024) }