[TargetLowering][AArch64] Teach DemandedBits about SVE count intrinsics

[benmxwl-arm]

Summary

This allows DemandedBits to see that the SVE count intrinsics (CNTB,
CNTH, CNTW, CNTD) sans multiplier will only ever produce small
positive integers. The maximum value you could get here is 256, which
is CNTB on a machine with a 2048bit vector size (the maximum for SVE).

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

[dmgreen]

This may be better as two separate patches, one for the vscale and one for the cnt's

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

1140

Could this just use construct known bits of both sides and use KnownBits::mul? It might be able to get value out of the low bits then too.

1147

The Negative = needn't be part of the if.

[benmxwl-arm]

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

1140

Please correct me if I've made a silly mistake here, but it does not seem to be exactly KnownBits::mul:

Optional<unsigned> MaxVScale = Attr.getVScaleRangeMax();
if (!MaxVScale.has_value())
  return false;
if (auto *MulImm = dyn_cast<ConstantSDNode>(Op.getOperand(0))) {
  unsigned RequiredBits = Log2_64(*MaxVScale) + 1;
  if (RequiredBits >= BitWidth)
    return false;
  Known.Zero.setHighBits(BitWidth - RequiredBits);
  Known = KnownBits::mul(Known, KnownBits::makeConstant(MulImm->getAPIntValue()));
}
return false;

If the MaxVScale is 16 (5 bits)

The known zero bits are everything above Log2(16 * Mul) + 1.
The above snippet seems to end up with Log2((2^5 - 1) * Mul) + 1 (which is off by 1 bit)

Also it seems that KnowBits::mul can't handle negative multipliers and always reports no known bits in that case.

[benmxwl-arm]

This patch now only handles teaching DemandedBits about the SVE count intrinsics.
The VSCALE changes have been split out to: https://reviews.llvm.org/D138508

[dmgreen]

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

1140

Ah - I was worried about that - because it works on bits and not ranges, the results are not as exact as they could otherwise be. I was also considering the vscale turning into a shift, but that may be SVE specific.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

23323

Can you move this into the switch under INTRINSIC_WO_CHAIN. Just so it doesn't need to be checked for every instruction.

23329

Is there a reason why this is this based on SVEMaxBitsPerVector and not the maximum value in VScaleRange? Or a combo of both if vscale_range is unbounded for some reason.

23330

Nit: Log2_32?

[benmxwl-arm]

• Moved CNTx intrinsic check into switch
• Now base the MaxElements off VScaleRange if present, otherwise fallback to SVEMaxBitsPerVector
• Switched to Log2_32

[benmxwl-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

23329

No real reason, it just was a little simpler. But it does not add much extra complexity doing the combo check.

[benmxwl-arm]

Simplified to use Subtarget->getMaxSVEVectorSizeInBits() (which is already based off the vscale_range() attribute).

This now depends on https://reviews.llvm.org/D138575 (currently getMaxSVEVectorSizeInBits() only checks for hasSVE(),
which means it breaks for SME). Until that lands there will be a few test breakages.

[dmgreen]

Sounds good. This LGTM if no one else has any other comments.

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

bool TargetLowering::SimplifyDemandedBits(
   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);
     if (SimplifyDemandedBits(Src, SrcDemandedBits, SrcKnown, TLO, Depth + 1))
       return true;
 
     // Upper elements are undef, so only get the knownbits if we just demand
     // the bottom element.
     if (DemandedElts == 1)
       Known = SrcKnown.anyextOrTrunc(BitWidth);
     break;
   }
   case ISD::BUILD_VECTOR:
     // Collect the known bits that are shared by every demanded element.
     // TODO: Call SimplifyDemandedBits for non-constant demanded elements.
     Known = TLO.DAG.computeKnownBits(Op, DemandedElts, Depth);
     return false; // Don't fall through, will infinitely loop.
   case ISD::LOAD: {
     auto *LD = cast<LoadSDNode>(Op);
if (isOperationLegalOrCustom(ISD::SHL, VT)) {
         if (unsigned ShAmt = getShiftLeftAmt(Op1))
           return foldMul(ISD::ADD, Op1.getOperand(0), Op0, ShAmt);
       }
     }
 
     [[fallthrough]];
   }
   default:
-    if (Op.getOpcode() >= ISD::BUILTIN_OP_END) {
+    // We also ask the target about intrinsics (which could be specific to it).
+    if (Op.getOpcode() >= ISD::BUILTIN_OP_END ||
+        Op.getOpcode() == ISD::INTRINSIC_WO_CHAIN) {
       if (SimplifyDemandedBitsForTargetNode(Op, DemandedBits, DemandedElts,
                                             Known, TLO, Depth))
         return true;
       break;
     }
 
     // Just use computeKnownBits to compute output bits.
     Known = TLO.DAG.computeKnownBits(Op, DemandedElts, Depth);

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

 #include <algorithm>
 #include <bitset>
 #include <cassert>
 #include <cctype>
 #include <cstdint>
 #include <cstdlib>
 #include <iterator>
 #include <limits>
+#include <optional>
 #include <tuple>
 #include <utility>
 #include <vector>
 
 using namespace llvm;
 using namespace llvm::PatternMatch;
 
 #define DEBUG_TYPE "aarch64-lower"
if (Lg2 == 1) {
     Created.push_back(Negs.getNode());
     Created.push_back(AndPos.getNode());
     Created.push_back(AndNeg.getNode());
   }
 
   return CSNeg;
 }
 
-static bool IsSVECntIntrinsic(SDValue S) {
+static std::optional<unsigned> IsSVECntIntrinsic(SDValue S) {
   switch(getIntrinsicID(S.getNode())) {
   default:
     break;
   case Intrinsic::aarch64_sve_cntb:
+    return 8;
   case Intrinsic::aarch64_sve_cnth:
+    return 16;
   case Intrinsic::aarch64_sve_cntw:
+    return 32;
   case Intrinsic::aarch64_sve_cntd:
-    return true;
+    return 64;
   }
-  return false;
+  return {};
 }
 
 /// Calculates what the pre-extend type is, based on the extension
 /// operation node provided by \p Extend.
 ///
 /// In the case that \p Extend is a SIGN_EXTEND or a ZERO_EXTEND, the
 /// pre-extend type is pulled directly from the operand, while other extend
 /// operations need a bit more inspection to get this information.
case AArch64ISD::VSHL: {
 
     if ((ZeroBits & UnusedBits) != ZeroBits)
       return false;
 
     // All bits that are zeroed by (VSHL (VLSHR Val X) X) are not
     // used - simplify to just Val.
     return TLO.CombineTo(Op, ShiftR->getOperand(0));
   }
+  case ISD::INTRINSIC_WO_CHAIN: {
+    if (auto ElementSize = IsSVECntIntrinsic(Op)) {
+      unsigned MaxSVEVectorSizeInBits = Subtarget->getMaxSVEVectorSizeInBits();
+      if (!MaxSVEVectorSizeInBits)
+        MaxSVEVectorSizeInBits = AArch64::SVEMaxBitsPerVector;
+      unsigned MaxElements = MaxSVEVectorSizeInBits / *ElementSize;
+      // The SVE count intrinsics don't support the multiplier immediate so we
+      // don't have to account for that here. The value returned may be slightly
+      // over the true required bits, as this is based on the "ALL" pattern. The
+      // other patterns are also exposed by these intrinsics, but they all
+      // return a value that's strictly less than "ALL".
+      unsigned RequiredBits = Log2_32(MaxElements) + 1;
+      unsigned BitWidth = Known.Zero.getBitWidth();
+      if (RequiredBits < BitWidth)
+        Known.Zero.setHighBits(BitWidth - RequiredBits);
+      return false;
+    }
+  }
   }
 
   return TargetLowering::SimplifyDemandedBitsForTargetNode(
       Op, OriginalDemandedBits, OriginalDemandedElts, Known, TLO, Depth);
 }
 
 bool AArch64TargetLowering::isTargetCanonicalConstantNode(SDValue Op) const {
   return Op.getOpcode() == AArch64ISD::DUP ||
          Op.getOpcode() == AArch64ISD::MOVI ||
          (Op.getOpcode() == ISD::EXTRACT_SUBVECTOR &&
           Op.getOperand(0).getOpcode() == AArch64ISD::DUP) ||
          TargetLowering::isTargetCanonicalConstantNode(Op);
 }
 
 bool AArch64TargetLowering::isConstantUnsignedBitfieldExtractLegal(
     unsigned Opc, LLT Ty1, LLT Ty2) const {
   return Ty1 == Ty2 && (Ty1 == LLT::scalar(32) || Ty1 == LLT::scalar(64));
 }

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

; CHECK-NEXT:    ret
   %result = add i32 %and_redundant, %and_required
   ret i32 %result
 }
 
 define i64 @cntb_and_elimination() {
 ; CHECK-LABEL: cntb_and_elimination:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    cntb x8
-; CHECK-NEXT:    and x9, x8, #0x1ff
-; CHECK-NEXT:    and x8, x8, #0x3fffffffc
-; CHECK-NEXT:    add x0, x9, x8
+; CHECK-NEXT:    and x9, x8, #0x1fc
+; CHECK-NEXT:    add x0, x8, x9
 ; CHECK-NEXT:    ret
   %cntb = call i64 @llvm.aarch64.sve.cntb(i32 31)
   %and_redundant = and i64 %cntb, 511
   %and_required = and i64 %cntb, 17179869180
   %result = add i64 %and_redundant, %and_required
   ret i64 %result
 }
 
 define i64 @cnth_and_elimination() {
 ; CHECK-LABEL: cnth_and_elimination:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    cnth x8
-; CHECK-NEXT:    and x9, x8, #0x3ff
-; CHECK-NEXT:    and x8, x8, #0x3fffffffc
-; CHECK-NEXT:    add x0, x9, x8
+; CHECK-NEXT:    and x9, x8, #0xfc
+; CHECK-NEXT:    add x0, x8, x9
 ; CHECK-NEXT:    ret
   %cnth = call i64 @llvm.aarch64.sve.cnth(i32 31)
   %and_redundant = and i64 %cnth, 1023
   %and_required = and i64 %cnth, 17179869180
   %result = add i64 %and_redundant, %and_required
   ret i64 %result
 }
 
 define i64 @cntw_and_elimination() {
 ; CHECK-LABEL: cntw_and_elimination:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    cntw x8
-; CHECK-NEXT:    and x9, x8, #0x7f
-; CHECK-NEXT:    and x8, x8, #0x3fffffffc
-; CHECK-NEXT:    add x0, x9, x8
+; CHECK-NEXT:    and x9, x8, #0x7c
+; CHECK-NEXT:    add x0, x8, x9
 ; CHECK-NEXT:    ret
   %cntw = call i64 @llvm.aarch64.sve.cntw(i32 31)
   %and_redundant = and i64 %cntw, 127
   %and_required = and i64 %cntw, 17179869180
   %result = add i64 %and_redundant, %and_required
   ret i64 %result
 }
 
 define i64 @cntd_and_elimination() {
 ; CHECK-LABEL: cntd_and_elimination:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    cntd x8
-; CHECK-NEXT:    and x9, x8, #0x3f
-; CHECK-NEXT:    and x8, x8, #0x3fffffffc
-; CHECK-NEXT:    add x0, x9, x8
+; CHECK-NEXT:    and x9, x8, #0x3c
+; CHECK-NEXT:    add x0, x8, x9
 ; CHECK-NEXT:    ret
   %cntd = call i64 @llvm.aarch64.sve.cntd(i32 31)
   %and_redundant = and i64 %cntd, 63
   %and_required = and i64 %cntd, 17179869180
   %result = add i64 %and_redundant, %and_required
   ret i64 %result
 }
 
; 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:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    cntb x8
-; CHECK-NEXT:    and x0, x8, #0xffffffff
+; CHECK-NEXT:    cntb x0
 ; CHECK-NEXT:    ret
   %cnt = call i64 @llvm.aarch64.sve.cntb(i32 31)
   %trunc = trunc i64 %cnt to i32
   %zext = zext i32 %trunc to i64
   ret i64 %zext
 }
 
 define i64 @count_halfs_trunc_zext() {
 ; CHECK-LABEL: count_halfs_trunc_zext:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    cnth x8
-; CHECK-NEXT:    and x0, x8, #0xffffffff
+; CHECK-NEXT:    cnth x0
 ; CHECK-NEXT:    ret
   %cnt = call i64 @llvm.aarch64.sve.cnth(i32 31)
   %trunc = trunc i64 %cnt to i32
   %zext = zext i32 %trunc to i64
   ret i64 %zext
 }
 
 define i64 @count_words_trunc_zext() {
 ; CHECK-LABEL: count_words_trunc_zext:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    cntw x8
-; CHECK-NEXT:    and x0, x8, #0xffffffff
+; CHECK-NEXT:    cntw x0
 ; CHECK-NEXT:    ret
   %cnt = call i64 @llvm.aarch64.sve.cntw(i32 31)
   %trunc = trunc i64 %cnt to i32
   %zext = zext i32 %trunc to i64
   ret i64 %zext
 }
 
 define i64 @count_doubles_trunc_zext() {
 ; CHECK-LABEL: count_doubles_trunc_zext:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    cntd x8
-; CHECK-NEXT:    and x0, x8, #0xffffffff
+; CHECK-NEXT:    cntd x0
 ; CHECK-NEXT:    ret
   %cnt = call i64 @llvm.aarch64.sve.cntd(i32 31)
   %trunc = trunc i64 %cnt to i32
   %zext = zext i32 %trunc to i64
   ret i64 %zext
 }
 
 define i64 @count_bytes_trunc_sext() {
 ; CHECK-LABEL: count_bytes_trunc_sext:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    cntb x8
-; CHECK-NEXT:    sxtw x0, w8
+; CHECK-NEXT:    cntb x0
 ; CHECK-NEXT:    ret
   %cnt = call i64 @llvm.aarch64.sve.cntb(i32 31)
   %trunc = trunc i64 %cnt to i32
   %sext = sext i32 %trunc to i64
   ret i64 %sext
 }
 
 define i64 @count_halfs_trunc_sext() {
 ; CHECK-LABEL: count_halfs_trunc_sext:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    cnth x8
-; CHECK-NEXT:    sxtw x0, w8
+; CHECK-NEXT:    cnth x0
 ; CHECK-NEXT:    ret
   %cnt = call i64 @llvm.aarch64.sve.cnth(i32 31)
   %trunc = trunc i64 %cnt to i32
   %sext = sext i32 %trunc to i64
   ret i64 %sext
 }
 
 define i64 @count_words_trunc_sext() {
 ; CHECK-LABEL: count_words_trunc_sext:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    cntw x8
-; CHECK-NEXT:    sxtw x0, w8
+; CHECK-NEXT:    cntw x0
 ; CHECK-NEXT:    ret
   %cnt = call i64 @llvm.aarch64.sve.cntw(i32 31)
   %trunc = trunc i64 %cnt to i32
   %sext = sext i32 %trunc to i64
   ret i64 %sext
 }
 
 define i64 @count_doubles_trunc_sext() {
 ; CHECK-LABEL: count_doubles_trunc_sext:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    cntd x8
-; CHECK-NEXT:    sxtw x0, w8
+; CHECK-NEXT:    cntd x0
 ; CHECK-NEXT:    ret
   %cnt = call i64 @llvm.aarch64.sve.cntd(i32 31)
   %trunc = trunc i64 %cnt to i32
   %sext = sext i32 %trunc to i64
   ret i64 %sext
 }
 
 define i32 @vscale_with_multiplier() vscale_range(1,16) {