[ARM] MVE i1 splat

[dmgreen]

Summary

We needn't BFI each lane individually into a predicate register when each lane in the same. A simple sign extend and a vmsr will do it.

[samparker]

Nice, LGTM.

llvm/lib/Target/ARM/ARMISelLowering.cpp

if (NumElts == 4) {
     BitsPerBool = 2;
     BoolMask = 0x3;
   } else if (NumElts == 16) {
     BitsPerBool = 1;
     BoolMask = 0x1;
   } else
     return SDValue();
 
+  // If this is a single value copied into all lanes (a splat), we can just sign
+  // extend that single value
+  SDValue FirstOp = Op.getOperand(0);
+  if (!isa<ConstantSDNode>(FirstOp) &&
+      std::all_of(std::next(Op->op_begin()), Op->op_end(),
+                  [&FirstOp](SDUse &U) {
+                    return U.get().isUndef() || U.get() == FirstOp;
+                  })) {
+    SDValue Ext = DAG.getNode(ISD::SIGN_EXTEND_INREG, dl, MVT::i32, FirstOp,
+                              DAG.getValueType(MVT::i1));
+    return DAG.getNode(ARMISD::PREDICATE_CAST, dl, Op.getValueType(), Ext);
+  }
+
   // First create base with bits set where known
   unsigned Bits32 = 0;
   for (unsigned i = 0; i < NumElts; ++i) {
     SDValue V = Op.getOperand(i);
     if (!isa<ConstantSDNode>(V) && !V.isUndef())
       continue;
     bool BitSet = V.isUndef() ? false : cast<ConstantSDNode>(V)->getZExtValue();
     if (BitSet)
       Bits32 |= BoolMask << (i * BitsPerBool);
   }
 
   // Add in unknown nodes
-  // FIXME: Handle splats of the same value better.
   SDValue Base = DAG.getNode(ARMISD::PREDICATE_CAST, dl, VT,
                              DAG.getConstant(Bits32, dl, MVT::i32));
   for (unsigned i = 0; i < NumElts; ++i) {
     SDValue V = Op.getOperand(i);
     if (isa<ConstantSDNode>(V) || V.isUndef())
       continue;
     Base = DAG.getNode(ISD::INSERT_VECTOR_ELT, dl, VT, Base, V,
                        DAG.getConstant(i, dl, MVT::i32));

llvm/test/CodeGen/Thumb2/mve-pred-build-var.ll

entry:
   %r = select <4 x i1> %vc, <4 x i32> %a, <4 x i32> %b
   ret <4 x i32> %r
 }
 
 define arm_aapcs_vfpcc <4 x i32> @build_varN_v4i1(i32 %s, i32 %t, <4 x i32> %a, <4 x i32> %b) {
 ; CHECK-LABEL: build_varN_v4i1:
 ; CHECK:       @ %bb.0: @ %entry
 ; CHECK-NEXT:    cmp r0, r1
-; CHECK-NEXT:    mov.w r1, #0
 ; CHECK-NEXT:    cset r0, lo
 ; CHECK-NEXT:    and r0, r0, #1
 ; CHECK-NEXT:    rsbs r0, r0, #0
-; CHECK-NEXT:    bfi r1, r0, #0, #4
-; CHECK-NEXT:    bfi r1, r0, #4, #4
-; CHECK-NEXT:    bfi r1, r0, #8, #4
-; CHECK-NEXT:    bfi r1, r0, #12, #4
-; CHECK-NEXT:    vmsr p0, r1
+; CHECK-NEXT:    vmsr p0, r0
 ; CHECK-NEXT:    vpsel q0, q0, q1
 ; CHECK-NEXT:    bx lr
 entry:
   %c = icmp ult i32 %s, %t
   %vc1 = insertelement <4 x i1> undef, i1 %c, i64 0
   %vc4 = shufflevector <4 x i1> %vc1, <4 x i1> undef, <4 x i32> zeroinitializer
   %r = select <4 x i1> %vc4, <4 x i32> %a, <4 x i32> %b
   ret <4 x i32> %r
entry:
   %r = select <8 x i1> %vc, <8 x i16> %a, <8 x i16> %b
   ret <8 x i16> %r
 }
 
 define arm_aapcs_vfpcc <8 x i16> @build_varN_v8i1(i32 %s, i32 %t, <8 x i16> %a, <8 x i16> %b) {
 ; CHECK-LABEL: build_varN_v8i1:
 ; CHECK:       @ %bb.0: @ %entry
 ; CHECK-NEXT:    cmp r0, r1
-; CHECK-NEXT:    mov.w r1, #0
 ; CHECK-NEXT:    cset r0, lo
 ; CHECK-NEXT:    and r0, r0, #1
 ; CHECK-NEXT:    rsbs r0, r0, #0
-; CHECK-NEXT:    bfi r1, r0, #0, #2
-; CHECK-NEXT:    bfi r1, r0, #2, #2
-; CHECK-NEXT:    bfi r1, r0, #4, #2
-; CHECK-NEXT:    bfi r1, r0, #6, #2
-; CHECK-NEXT:    bfi r1, r0, #8, #2
-; CHECK-NEXT:    bfi r1, r0, #10, #2
-; CHECK-NEXT:    bfi r1, r0, #12, #2
-; CHECK-NEXT:    bfi r1, r0, #14, #2
-; CHECK-NEXT:    vmsr p0, r1
+; CHECK-NEXT:    vmsr p0, r0
 ; CHECK-NEXT:    vpsel q0, q0, q1
 ; CHECK-NEXT:    bx lr
 entry:
   %c = icmp ult i32 %s, %t
   %vc1 = insertelement <8 x i1> undef, i1 %c, i64 0
   %vc4 = shufflevector <8 x i1> %vc1, <8 x i1> undef, <8 x i32> zeroinitializer
   %r = select <8 x i1> %vc4, <8 x i16> %a, <8 x i16> %b
   ret <8 x i16> %r
entry:
   %r = select <16 x i1> %vc, <16 x i8> %a, <16 x i8> %b
   ret <16 x i8> %r
 }
 
 define arm_aapcs_vfpcc <16 x i8> @build_varN_v16i1(i32 %s, i32 %t, <16 x i8> %a, <16 x i8> %b) {
 ; CHECK-LABEL: build_varN_v16i1:
 ; CHECK:       @ %bb.0: @ %entry
 ; CHECK-NEXT:    cmp r0, r1
-; CHECK-NEXT:    mov.w r1, #0
 ; CHECK-NEXT:    cset r0, lo
 ; CHECK-NEXT:    and r0, r0, #1
 ; CHECK-NEXT:    rsbs r0, r0, #0
-; CHECK-NEXT:    bfi r1, r0, #0, #1
-; CHECK-NEXT:    bfi r1, r0, #1, #1
-; CHECK-NEXT:    bfi r1, r0, #2, #1
-; CHECK-NEXT:    bfi r1, r0, #3, #1
-; CHECK-NEXT:    bfi r1, r0, #4, #1
-; CHECK-NEXT:    bfi r1, r0, #5, #1
-; CHECK-NEXT:    bfi r1, r0, #6, #1
-; CHECK-NEXT:    bfi r1, r0, #7, #1
-; CHECK-NEXT:    bfi r1, r0, #8, #1
-; CHECK-NEXT:    bfi r1, r0, #9, #1
-; CHECK-NEXT:    bfi r1, r0, #10, #1
-; CHECK-NEXT:    bfi r1, r0, #11, #1
-; CHECK-NEXT:    bfi r1, r0, #12, #1
-; CHECK-NEXT:    bfi r1, r0, #13, #1
-; CHECK-NEXT:    bfi r1, r0, #14, #1
-; CHECK-NEXT:    bfi r1, r0, #15, #1
-; CHECK-NEXT:    vmsr p0, r1
+; CHECK-NEXT:    vmsr p0, r0
 ; CHECK-NEXT:    vpsel q0, q0, q1
 ; CHECK-NEXT:    bx lr
 entry:
   %c = icmp ult i32 %s, %t
   %vc1 = insertelement <16 x i1> undef, i1 %c, i64 0
   %vc4 = shufflevector <16 x i1> %vc1, <16 x i1> undef, <16 x i32> zeroinitializer
   %r = select <16 x i1> %vc4, <16 x i8> %a, <16 x i8> %b
   ret <16 x i8> %r