[SelectionDAG] Teach isConstOrConstSplat about ISD::SPLAT_VECTOR

This improves llvm::isConstOrConstSplat by allowing it to analyze
ISD::SPLAT_VECTOR nodes, in order to allow more constant-folding of
operations using scalable vector types.

Reviewed By: craig.topper

Differential Revision: https://reviews.llvm.org/D94168

Author: frasercrmck

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

@@ -9129,6 +9129,18 @@ ConstantSDNode *llvm::isConstOrConstSplat(SDValue N, bool AllowUndefs,
   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);

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

@@ -3463,8 +3463,11 @@ SDValue TargetLowering::SimplifySetCC(EVT VT, SDValue N0, SDValue N1,
 
   // Ensure that the constant occurs on the RHS and fold constant comparisons.
   // TODO: Handle non-splat vector constants. All undef causes trouble.
+  // FIXME: We can't yet fold constant scalable vector splats, so avoid an
+  // infinite loop here when we encounter one.
   ISD::CondCode SwappedCC = ISD::getSetCCSwappedOperands(Cond);
   if (isConstOrConstSplat(N0) &&
+      (!OpVT.isScalableVector() || !isConstOrConstSplat(N1)) &&
       (DCI.isBeforeLegalizeOps() ||
        isCondCodeLegal(SwappedCC, N0.getSimpleValueType())))
     return DAG.getSetCC(dl, VT, N1, N0, SwappedCC);

llvm/test/CodeGen/AArch64/sve-intrinsics-int-compares-with-imm.ll

@@ -820,7 +820,7 @@ define <vscale x 16 x i1> @wide_cmphi_b(<vscale x 16 x i1> %pg, <vscale x 16 x i
 
 define <vscale x 8 x i1> @ir_cmphi_h(<vscale x 8 x i16> %a) {
 ; CHECK-LABEL: ir_cmphi_h
-; CHECK: cmphi p0.h, p0/z, z0.h, #0
+; CHECK: cmpne p0.h, p0/z, z0.h, #0
 ; CHECK-NEXT: ret
   %elt   = insertelement <vscale x 8 x i16> undef, i16 0, i32 0
   %splat = shufflevector <vscale x 8 x i16> %elt, <vscale x 8 x i16> undef, <vscale x 8 x i32> zeroinitializer
@@ -948,7 +948,7 @@ define <vscale x 16 x i1> @wide_cmphs_b(<vscale x 16 x i1> %pg, <vscale x 16 x i
 
 define <vscale x 8 x i1> @ir_cmphs_h(<vscale x 8 x i16> %a) {
 ; CHECK-LABEL: ir_cmphs_h
-; CHECK: cmphs p0.h, p0/z, z0.h, #0
+; CHECK: ptrue p0.h
 ; CHECK-NEXT: ret
   %elt   = insertelement <vscale x 8 x i16> undef, i16 0, i32 0
   %splat = shufflevector <vscale x 8 x i16> %elt, <vscale x 8 x i16> undef, <vscale x 8 x i32> zeroinitializer
@@ -1076,7 +1076,7 @@ define <vscale x 16 x i1> @wide_cmplo_b(<vscale x 16 x i1> %pg, <vscale x 16 x i
 
 define <vscale x 8 x i1> @ir_cmplo_h(<vscale x 8 x i16> %a) {
 ; CHECK-LABEL: ir_cmplo_h
-; CHECK: cmplo p0.h, p0/z, z0.h, #0
+; CHECK: whilelo p0.h, xzr, xzr
 ; CHECK-NEXT: ret
   %elt   = insertelement <vscale x 8 x i16> undef, i16 0, i32 0
   %splat = shufflevector <vscale x 8 x i16> %elt, <vscale x 8 x i16> undef, <vscale x 8 x i32> zeroinitializer

llvm/test/CodeGen/RISCV/rvv/vdiv-sdnode-rv32.ll

@@ -37,13 +37,9 @@ define <vscale x 1 x i8> @vdiv_vi_nxv1i8_0(<vscale x 1 x i8> %va) {
 }
 
 ; Test V/1 to see if we can optimize it away for scalable vectors.
-; FIXME: We can't.
 define <vscale x 1 x i8> @vdiv_vi_nxv1i8_1(<vscale x 1 x i8> %va) {
 ; CHECK-LABEL: vdiv_vi_nxv1i8_1:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    addi a0, zero, 1
-; CHECK-NEXT:    vsetvli a1, zero, e8,mf8,ta,mu
-; CHECK-NEXT:    vdivu.vx v16, v16, a0
 ; CHECK-NEXT:    ret
   %head = insertelement <vscale x 1 x i8> undef, i8 1, i32 0
   %splat = shufflevector <vscale x 1 x i8> %head, <vscale x 1 x i8> undef, <vscale x 1 x i32> zeroinitializer
@@ -52,13 +48,11 @@ define <vscale x 1 x i8> @vdiv_vi_nxv1i8_1(<vscale x 1 x i8> %va) {
 }
 
 ; Test 0/V to see if we can optimize it away for scalable vectors.
-; FIXME: We can't.
 define <vscale x 1 x i8> @vdiv_iv_nxv1i8_0(<vscale x 1 x i8> %va) {
 ; CHECK-LABEL: vdiv_iv_nxv1i8_0:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetvli a0, zero, e8,mf8,ta,mu
-; CHECK-NEXT:    vmv.v.i v25, 0
-; CHECK-NEXT:    vdivu.vv v16, v25, v16
+; CHECK-NEXT:    vmv.v.i v16, 0
 ; CHECK-NEXT:    ret
   %head = insertelement <vscale x 1 x i8> undef, i8 0, i32 0
   %splat = shufflevector <vscale x 1 x i8> %head, <vscale x 1 x i8> undef, <vscale x 1 x i32> zeroinitializer

llvm/test/CodeGen/RISCV/rvv/vdivu-sdnode-rv32.ll

@@ -37,13 +37,9 @@ define <vscale x 1 x i8> @vdivu_vi_nxv1i8_0(<vscale x 1 x i8> %va) {
 }
 
 ; Test V/1 to see if we can optimize it away for scalable vectors.
-; FIXME: We can't.
 define <vscale x 1 x i8> @vdivu_vi_nxv1i8_1(<vscale x 1 x i8> %va) {
 ; CHECK-LABEL: vdivu_vi_nxv1i8_1:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    addi a0, zero, 1
-; CHECK-NEXT:    vsetvli a1, zero, e8,mf8,ta,mu
-; CHECK-NEXT:    vdiv.vx v16, v16, a0
 ; CHECK-NEXT:    ret
   %head = insertelement <vscale x 1 x i8> undef, i8 1, i32 0
   %splat = shufflevector <vscale x 1 x i8> %head, <vscale x 1 x i8> undef, <vscale x 1 x i32> zeroinitializer
@@ -52,13 +48,11 @@ define <vscale x 1 x i8> @vdivu_vi_nxv1i8_1(<vscale x 1 x i8> %va) {
 }
 
 ; Test 0/V to see if we can optimize it away for scalable vectors.
-; FIXME: We can't.
 define <vscale x 1 x i8> @vdivu_iv_nxv1i8_0(<vscale x 1 x i8> %va) {
 ; CHECK-LABEL: vdivu_iv_nxv1i8_0:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetvli a0, zero, e8,mf8,ta,mu
-; CHECK-NEXT:    vmv.v.i v25, 0
-; CHECK-NEXT:    vdiv.vv v16, v25, v16
+; CHECK-NEXT:    vmv.v.i v16, 0
 ; CHECK-NEXT:    ret
   %head = insertelement <vscale x 1 x i8> undef, i8 0, i32 0
   %splat = shufflevector <vscale x 1 x i8> %head, <vscale x 1 x i8> undef, <vscale x 1 x i32> zeroinitializer