[DAG] Teach isConstOrConstSplat about SPLAT_VECTORs

[dmgreen]

Summary

This teaches the DemandedElts version of isConstOrConstSplat about SPLAT_VECTORS, in the same way as the non-DemandedElts version by calling the demanded-bits version from the non-demanded-bits version.

[dmgreen]

Update to just call the demanded-bits version of the function from the demanded-bits version, and do the same for isConstOrConstSplatFP.

[RKSimon]

I've no objections, but the experts on SVE targets need to have final say.

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

10672

Please can you add a TODO - about whether we should allow UndefElements in non-DemandedElts?

10712

Please can you add a TODO - about whether we should allow UndefElements in non-DemandedElts?

[craig.topper]

RISCV changes LGTM

[RKSimon]

LGTM (with the addition of the TODO comments I mentioned before)

[dmgreen]

Thanks for the review. I think there is enough in now that this can go ahead. Hopefully it won't hit the same issues seen on Hexagon in the other patches. Let me know if anything does look off.

[paulwalker-arm]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

10616

Based on @reames's demanded elts patches I think this should be isFixedLengthVector()? because we now use a 1bit APInt to represent demanding all elements of a scalable vector.

10700

As above.

[dmgreen]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

10616

Yes. Those are already in the version I am about to commit. I was just running ninja check-all first and it was taking a long time for some reason.

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

bool llvm::isBitwiseNot(SDValue V, bool AllowUndefs) {
   ConstantSDNode *C =
       isConstOrConstSplat(V, AllowUndefs, /*AllowTruncation*/ true);
   return C && (C->getAPIntValue().countTrailingOnes() >= NumBits);
 }
 
 ConstantSDNode *llvm::isConstOrConstSplat(SDValue N, bool AllowUndefs,
                                           bool AllowTruncation) {
   if (ConstantSDNode *CN = dyn_cast<ConstantSDNode>(N))
     return CN;
 
   // SplatVectors can truncate their operands. Ignore that case here unless
   // AllowTruncation is set.
   if (N->getOpcode() == ISD::SPLAT_VECTOR) {
     EVT VecEltVT = N->getValueType(0).getVectorElementType();
     if (auto *CN = dyn_cast<ConstantSDNode>(N->getOperand(0))) {
       EVT CVT = CN->getValueType(0);
       assert(CVT.bitsGE(VecEltVT) && "Illegal splat_vector element extension");

 }
 
 ConstantSDNode *llvm::isConstOrConstSplat(SDValue N, const APInt &DemandedElts,
                                           bool AllowUndefs,
                                           bool AllowTruncation) {
   if (ConstantSDNode *CN = dyn_cast<ConstantSDNode>(N))
     return CN;
 
+  // SplatVectors can truncate their operands. Ignore that case here unless
+  // AllowTruncation is set.
+  if (N->getOpcode() == ISD::SPLAT_VECTOR) {
+    EVT VecEltVT = N->getValueType(0).getVectorElementType();
+    if (auto *CN = dyn_cast<ConstantSDNode>(N->getOperand(0))) {
+      EVT CVT = CN->getValueType(0);
+      assert(CVT.bitsGE(VecEltVT) && "Illegal splat_vector element extension");
+      if (AllowTruncation || CVT == VecEltVT)
+        return CN;
+    }
+  }
+
   if (BuildVectorSDNode *BV = dyn_cast<BuildVectorSDNode>(N)) {
     BitVector UndefElements;
     ConstantSDNode *CN = BV->getConstantSplatNode(DemandedElts, &UndefElements);
 
     // BuildVectors can truncate their operands. Ignore that case here unless
     // AllowTruncation is set.
     if (CN && (UndefElements.none() || AllowUndefs)) {
       EVT CVT = CN->getValueType(0);
       EVT NSVT = N.getValueType().getScalarType();
       assert(CVT.bitsGE(NSVT) && "Illegal build vector element extension");
       if (AllowTruncation || (CVT == NSVT))
         return CN;
     }
   }
 
if (CN && (UndefElements.none() || AllowUndefs))
       return CN;
   }
 
   if (N.getOpcode() == ISD::SPLAT_VECTOR)
     if (ConstantFPSDNode *CN = dyn_cast<ConstantFPSDNode>(N.getOperand(0)))
       return CN;
 
   return nullptr;
 }
 
 ConstantFPSDNode *llvm::isConstOrConstSplatFP(SDValue N,
                                               const APInt &DemandedElts,
                                               bool AllowUndefs) {
   if (ConstantFPSDNode *CN = dyn_cast<ConstantFPSDNode>(N))
     return CN;
 
   if (BuildVectorSDNode *BV = dyn_cast<BuildVectorSDNode>(N)) {
     BitVector UndefElements;
     ConstantFPSDNode *CN =
         BV->getConstantFPSplatNode(DemandedElts, &UndefElements);
     if (CN && (UndefElements.none() || AllowUndefs))
       return CN;
   }
 
   return nullptr;
 }
 
 bool llvm::isNullOrNullSplat(SDValue N, bool AllowUndefs) {
   // TODO: may want to use peekThroughBitcast() here.

llvm/test/CodeGen/AArch64/sve-knownbits.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sve -verify-machineinstrs < %s | FileCheck %s
 
 define <vscale x 8 x i16> @test_knownzero(<vscale x 8 x i16> %x) {
 ; CHECK-LABEL: test_knownzero:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    mov z0.h, #0 // =0x0
 ; CHECK-NEXT:    ret
   %a1 = shl <vscale x 8 x i16> %x, shufflevector (<vscale x 8 x i16> insertelement (<vscale x 8 x i16> poison, i16 8, i32 0), <vscale x 8 x i16> poison, <vscale x 8 x i32> zeroinitializer)
   %a2 = and <vscale x 8 x i16> %a1, shufflevector (<vscale x 8 x i16> insertelement (<vscale x 8 x i16> poison, i16 8, i32 0), <vscale x 8 x i16> poison, <vscale x 8 x i32> zeroinitializer)
   ret <vscale x 8 x i16> %a2
 }
 
 define <vscale x 4 x i32> @asrlsr(<vscale x 4 x i64> %va) {
 ; CHECK-LABEL: asrlsr:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    asr z1.d, z1.d, #15
-; CHECK-NEXT:    asr z0.d, z0.d, #15
+; CHECK-NEXT:    lsr z1.d, z1.d, #15
+; CHECK-NEXT:    lsr z0.d, z0.d, #15
 ; CHECK-NEXT:    uzp1 z0.s, z0.s, z1.s
 ; CHECK-NEXT:    ret
   %head = insertelement <vscale x 4 x i32> poison, i32 15, i32 0
   %splat = shufflevector <vscale x 4 x i32> %head, <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer
   %vb = zext <vscale x 4 x i32> %splat to <vscale x 4 x i64>
   %x = ashr <vscale x 4 x i64> %va, %vb
   %y = trunc <vscale x 4 x i64> %x to <vscale x 4 x i32>
   ret <vscale x 4 x i32> %y
 }

llvm/test/CodeGen/RISCV/rvv/vnsra-sdnode.ll

; CHECK-NEXT:    ret
   %y = trunc <vscale x 1 x i64> %x to <vscale x 1 x i32>
   ret <vscale x 1 x i32> %y
 }
 
 define <vscale x 1 x i32> @vnsra_wi_i32_nxv1i32_sext(<vscale x 1 x i64> %va) {
 ; CHECK-LABEL: vnsra_wi_i32_nxv1i32_sext:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetvli a0, zero, e32, mf2, ta, mu
-; CHECK-NEXT:    vnsra.wi v8, v8, 15
+; CHECK-NEXT:    vnsrl.wi v8, v8, 15
 ; CHECK-NEXT:    ret
   %head = insertelement <vscale x 1 x i32> poison, i32 15, i32 0
   %splat = shufflevector <vscale x 1 x i32> %head, <vscale x 1 x i32> poison, <vscale x 1 x i32> zeroinitializer
   %vb = sext <vscale x 1 x i32> %splat to <vscale x 1 x i64>
   %x = ashr <vscale x 1 x i64> %va, %vb
   %y = trunc <vscale x 1 x i64> %x to <vscale x 1 x i32>
   ret <vscale x 1 x i32> %y
 }
; CHECK-NEXT:    ret
   %y = trunc <vscale x 2 x i64> %x to <vscale x 2 x i32>
   ret <vscale x 2 x i32> %y
 }
 
 define <vscale x 2 x i32> @vnsra_wi_i32_nxv2i32_sext(<vscale x 2 x i64> %va) {
 ; CHECK-LABEL: vnsra_wi_i32_nxv2i32_sext:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetvli a0, zero, e32, m1, ta, mu
-; CHECK-NEXT:    vnsra.wi v10, v8, 15
+; CHECK-NEXT:    vnsrl.wi v10, v8, 15
 ; CHECK-NEXT:    vmv.v.v v8, v10
 ; CHECK-NEXT:    ret
   %head = insertelement <vscale x 2 x i32> poison, i32 15, i32 0
   %splat = shufflevector <vscale x 2 x i32> %head, <vscale x 2 x i32> poison, <vscale x 2 x i32> zeroinitializer
   %vb = sext <vscale x 2 x i32> %splat to <vscale x 2 x i64>
   %x = ashr <vscale x 2 x i64> %va, %vb
   %y = trunc <vscale x 2 x i64> %x to <vscale x 2 x i32>
   ret <vscale x 2 x i32> %y
; CHECK-NEXT:    ret
   %y = trunc <vscale x 4 x i64> %x to <vscale x 4 x i32>
   ret <vscale x 4 x i32> %y
 }
 
 define <vscale x 4 x i32> @vnsra_wi_i32_nxv4i32_sext(<vscale x 4 x i64> %va) {
 ; CHECK-LABEL: vnsra_wi_i32_nxv4i32_sext:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetvli a0, zero, e32, m2, ta, mu
-; CHECK-NEXT:    vnsra.wi v12, v8, 15
+; CHECK-NEXT:    vnsrl.wi v12, v8, 15
 ; CHECK-NEXT:    vmv.v.v v8, v12
 ; CHECK-NEXT:    ret
   %head = insertelement <vscale x 4 x i32> poison, i32 15, i32 0
   %splat = shufflevector <vscale x 4 x i32> %head, <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer
   %vb = sext <vscale x 4 x i32> %splat to <vscale x 4 x i64>
   %x = ashr <vscale x 4 x i64> %va, %vb
   %y = trunc <vscale x 4 x i64> %x to <vscale x 4 x i32>
   ret <vscale x 4 x i32> %y
; CHECK-NEXT:    ret
   %y = trunc <vscale x 8 x i64> %x to <vscale x 8 x i32>
   ret <vscale x 8 x i32> %y
 }
 
 define <vscale x 8 x i32> @vnsra_wi_i32_nxv8i32_sext(<vscale x 8 x i64> %va) {
 ; CHECK-LABEL: vnsra_wi_i32_nxv8i32_sext:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetvli a0, zero, e32, m4, ta, mu
-; CHECK-NEXT:    vnsra.wi v16, v8, 15
+; CHECK-NEXT:    vnsrl.wi v16, v8, 15
 ; CHECK-NEXT:    vmv.v.v v8, v16
 ; CHECK-NEXT:    ret
   %head = insertelement <vscale x 8 x i32> poison, i32 15, i32 0
   %splat = shufflevector <vscale x 8 x i32> %head, <vscale x 8 x i32> poison, <vscale x 8 x i32> zeroinitializer
   %vb = sext <vscale x 8 x i32> %splat to <vscale x 8 x i64>
   %x = ashr <vscale x 8 x i64> %va, %vb
   %y = trunc <vscale x 8 x i64> %x to <vscale x 8 x i32>
   ret <vscale x 8 x i32> %y
; CHECK-NEXT:    ret
   %y = trunc <vscale x 1 x i64> %x to <vscale x 1 x i32>
   ret <vscale x 1 x i32> %y
 }
 
 define <vscale x 1 x i32> @vnsra_wi_i32_nxv1i32_zext(<vscale x 1 x i64> %va) {
 ; CHECK-LABEL: vnsra_wi_i32_nxv1i32_zext:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetvli a0, zero, e32, mf2, ta, mu
-; CHECK-NEXT:    vnsra.wi v8, v8, 15
+; CHECK-NEXT:    vnsrl.wi v8, v8, 15
 ; CHECK-NEXT:    ret
   %head = insertelement <vscale x 1 x i32> poison, i32 15, i32 0
   %splat = shufflevector <vscale x 1 x i32> %head, <vscale x 1 x i32> poison, <vscale x 1 x i32> zeroinitializer
   %vb = zext <vscale x 1 x i32> %splat to <vscale x 1 x i64>
   %x = ashr <vscale x 1 x i64> %va, %vb
   %y = trunc <vscale x 1 x i64> %x to <vscale x 1 x i32>
   ret <vscale x 1 x i32> %y
 }
; CHECK-NEXT:    ret
   %y = trunc <vscale x 2 x i64> %x to <vscale x 2 x i32>
   ret <vscale x 2 x i32> %y
 }
 
 define <vscale x 2 x i32> @vnsra_wi_i32_nxv2i32_zext(<vscale x 2 x i64> %va) {
 ; CHECK-LABEL: vnsra_wi_i32_nxv2i32_zext:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetvli a0, zero, e32, m1, ta, mu
-; CHECK-NEXT:    vnsra.wi v10, v8, 15
+; CHECK-NEXT:    vnsrl.wi v10, v8, 15
 ; CHECK-NEXT:    vmv.v.v v8, v10
 ; CHECK-NEXT:    ret
   %head = insertelement <vscale x 2 x i32> poison, i32 15, i32 0
   %splat = shufflevector <vscale x 2 x i32> %head, <vscale x 2 x i32> poison, <vscale x 2 x i32> zeroinitializer
   %vb = zext <vscale x 2 x i32> %splat to <vscale x 2 x i64>
   %x = ashr <vscale x 2 x i64> %va, %vb
   %y = trunc <vscale x 2 x i64> %x to <vscale x 2 x i32>
   ret <vscale x 2 x i32> %y
; CHECK-NEXT:    ret
   %y = trunc <vscale x 4 x i64> %x to <vscale x 4 x i32>
   ret <vscale x 4 x i32> %y
 }
 
 define <vscale x 4 x i32> @vnsra_wi_i32_nxv4i32_zext(<vscale x 4 x i64> %va) {
 ; CHECK-LABEL: vnsra_wi_i32_nxv4i32_zext:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetvli a0, zero, e32, m2, ta, mu
-; CHECK-NEXT:    vnsra.wi v12, v8, 15
+; CHECK-NEXT:    vnsrl.wi v12, v8, 15
 ; CHECK-NEXT:    vmv.v.v v8, v12
 ; CHECK-NEXT:    ret
   %head = insertelement <vscale x 4 x i32> poison, i32 15, i32 0
   %splat = shufflevector <vscale x 4 x i32> %head, <vscale x 4 x i32> poison, <vscale x 4 x i32> zeroinitializer
   %vb = zext <vscale x 4 x i32> %splat to <vscale x 4 x i64>
   %x = ashr <vscale x 4 x i64> %va, %vb
   %y = trunc <vscale x 4 x i64> %x to <vscale x 4 x i32>
   ret <vscale x 4 x i32> %y
; CHECK-NEXT:    ret
   %y = trunc <vscale x 8 x i64> %x to <vscale x 8 x i32>
   ret <vscale x 8 x i32> %y
 }
 
 define <vscale x 8 x i32> @vnsra_wi_i32_nxv8i32_zext(<vscale x 8 x i64> %va) {
 ; CHECK-LABEL: vnsra_wi_i32_nxv8i32_zext:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetvli a0, zero, e32, m4, ta, mu
-; CHECK-NEXT:    vnsra.wi v16, v8, 15
+; CHECK-NEXT:    vnsrl.wi v16, v8, 15
 ; CHECK-NEXT:    vmv.v.v v8, v16
 ; CHECK-NEXT:    ret
   %head = insertelement <vscale x 8 x i32> poison, i32 15, i32 0
   %splat = shufflevector <vscale x 8 x i32> %head, <vscale x 8 x i32> poison, <vscale x 8 x i32> zeroinitializer
   %vb = zext <vscale x 8 x i32> %splat to <vscale x 8 x i64>
   %x = ashr <vscale x 8 x i64> %va, %vb
   %y = trunc <vscale x 8 x i64> %x to <vscale x 8 x i32>
   ret <vscale x 8 x i32> %y
 }