[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

10996

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

11027

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

10965

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.

11010

As above.

[dmgreen]

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

10965

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) {
   unsigned NumBits = V.getScalarValueSizeInBits();
   ConstantSDNode *C =
       isConstOrConstSplat(V, AllowUndefs, /*AllowTruncation*/ true);
   return C && (C->getAPIntValue().countTrailingOnes() >= NumBits);
 }
 
 ConstantSDNode *llvm::isConstOrConstSplat(SDValue N, bool AllowUndefs,
                                           bool AllowTruncation) {
+  EVT VT = N.getValueType();
+  APInt DemandedElts = VT.isFixedLengthVector()
+                           ? APInt::getAllOnes(VT.getVectorMinNumElements())
+                           : APInt(1, 1);
+  return isConstOrConstSplat(N, DemandedElts, AllowUndefs, AllowTruncation);
+}
+
+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(&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;
-    }
-  }
-
-  return nullptr;
-}
-
-ConstantSDNode *llvm::isConstOrConstSplat(SDValue N, const APInt &DemandedElts,
-                                          bool AllowUndefs,
-                                          bool AllowTruncation) {
-  if (ConstantSDNode *CN = dyn_cast<ConstantSDNode>(N))
-    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.
+    // TODO: Look into whether we should allow UndefElements in non-DemandedElts
     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;
     }
   }
 
   return nullptr;
 }
 
 ConstantFPSDNode *llvm::isConstOrConstSplatFP(SDValue N, bool AllowUndefs) {
-  if (ConstantFPSDNode *CN = dyn_cast<ConstantFPSDNode>(N))
-    return CN;
-
-  if (BuildVectorSDNode *BV = dyn_cast<BuildVectorSDNode>(N)) {
-    BitVector UndefElements;
-    ConstantFPSDNode *CN = BV->getConstantFPSplatNode(&UndefElements);
-    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;
+  EVT VT = N.getValueType();
+  APInt DemandedElts = VT.isFixedLengthVector()
+                           ? APInt::getAllOnes(VT.getVectorMinNumElements())
+                           : APInt(1, 1);
+  return isConstOrConstSplatFP(N, DemandedElts, AllowUndefs);
 }
 
 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);
+    // TODO: Look into whether we should allow UndefElements in non-DemandedElts
     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;
 }
 
 bool llvm::isNullOrNullSplat(SDValue N, bool AllowUndefs) {
   // TODO: may want to use peekThroughBitcast() here.
   ConstantSDNode *C =
       isConstOrConstSplat(N, AllowUndefs, /*AllowTruncation=*/true);
   return C && C->isZero();

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:    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, ma
-; 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, ma
-; 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, ma
-; 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, ma
-; 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, ma
-; 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, ma
-; 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, ma
-; 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, ma
-; 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
 }