[AArch64] Handle vector with two different values with efficient vector mask

[jaykang10]

Summary

Following @dmgreen's comment on https://reviews.llvm.org/D148347, when we lower BUILD_VECTOR to VECTOR_SHUFFL, we could generate efficient vector mask.

[dmgreen]

LGTM Thanks.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

if (canUseVECTOR_CONCAT) {
         SDValue DUP2 =
             LowerBUILD_VECTOR(DAG.getBuildVector(SubVT, dl, Ops2), DAG);
         SDValue CONCAT_VECTORS =
             DAG.getNode(ISD::CONCAT_VECTORS, dl, VT, DUP1, DUP2);
         return CONCAT_VECTORS;
       }
     }
 
-    // Let's try to generate two DUPs and VECTOR_SHUFFLE. For example,
+    // Let's try to generate VECTOR_SHUFFLE. For example,
     //
     //  t24: v8i8 = BUILD_VECTOR t25, t25, t25, t25, t26, t26, t26, t26
     //  ==>
-    //    t28: v8i8 = AArch64ISD::DUP t25
-    //    t30: v8i8 = AArch64ISD::DUP t26
-    //  t31: v8i8 = vector_shuffle<0,0,0,0,8,8,8,8> t28, t30
+    //    t27: v8i8 = BUILD_VECTOR t26, t26, t26, t26, t26, t26, t26, t26
+    //    t28: v8i8 = BUILD_VECTOR t25, t25, t25, t25, t25, t25, t25, t25
+    //  t29: v8i8 = vector_shuffle<0,1,2,3,12,13,14,15> t27, t28
     if (NumElts >= 8) {
       SmallVector<int, 16> MaskVec;
       // Build mask for VECTOR_SHUFLLE.
       SDValue FirstLaneVal = Op.getOperand(0);
       for (unsigned i = 0; i < NumElts; ++i) {
         SDValue Val = Op.getOperand(i);
         if (FirstLaneVal == Val)
-          MaskVec.push_back(0);
+          MaskVec.push_back(i);
         else
-          MaskVec.push_back(NumElts);
+          MaskVec.push_back(i + NumElts);
       }
 
       SmallVector<SDValue, 8> Ops1(NumElts, Vals[0]);
       SmallVector<SDValue, 8> Ops2(NumElts, Vals[1]);
-      SDValue DUP1 = LowerBUILD_VECTOR(DAG.getBuildVector(VT, dl, Ops1), DAG);
-      SDValue DUP2 = LowerBUILD_VECTOR(DAG.getBuildVector(VT, dl, Ops2), DAG);
+      SDValue VEC1 = DAG.getBuildVector(VT, dl, Ops1);
+      SDValue VEC2 = DAG.getBuildVector(VT, dl, Ops2);
       SDValue VECTOR_SHUFFLE =
-          DAG.getVectorShuffle(VT, dl, DUP1, DUP2, MaskVec);
+          DAG.getVectorShuffle(VT, dl, VEC1, VEC2, MaskVec);
       return VECTOR_SHUFFLE;
     }
   }
 
   // If all else fails, just use a sequence of INSERT_VECTOR_ELT when we
   // know the default expansion would otherwise fall back on something even
   // worse. For a vector with one or two non-undef values, that's
   // scalar_to_vector for the elements followed by a shuffle (provided the

llvm/test/CodeGen/AArch64/build-vector-two-dup.ll

entry:
   %lane2 = shufflevector <8 x i8> %3, <8 x i8> poison, <8 x i32> zeroinitializer
   %shuffle.i = shufflevector <8 x i8> %lane, <8 x i8> %lane2, <16 x i32> <i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 15>
   ret <16 x i8> %shuffle.i
 }
 
 define <8 x i8> @test6(ptr nocapture noundef readonly %a, ptr nocapture noundef readonly %b) {
 ; CHECK-LABEL: test6:
 ; CHECK:       // %bb.0: // %entry
-; CHECK-NEXT:    ld1r { v0.8b }, [x1]
-; CHECK-NEXT:    adrp x8, .LCPI5_0
-; CHECK-NEXT:    ld1r { v1.8b }, [x0]
-; CHECK-NEXT:    mov v1.d[1], v0.d[0]
-; CHECK-NEXT:    ldr d0, [x8, :lo12:.LCPI5_0]
-; CHECK-NEXT:    tbl v0.8b, { v1.16b }, v0.8b
+; CHECK-NEXT:    ld1r { v1.8b }, [x1]
+; CHECK-NEXT:    ld1r { v0.8b }, [x0]
+; CHECK-NEXT:    mov v0.s[1], v1.s[1]
+; CHECK-NEXT:    // kill: def $d0 killed $d0 killed $q0
 ; CHECK-NEXT:    ret
 entry:
   %0 = load i8, ptr %a, align 1
   %1 = insertelement <4 x i8> poison, i8 %0, i64 0
   %lane = shufflevector <4 x i8> %1, <4 x i8> poison, <4 x i32> zeroinitializer
   %2 = load i8, ptr %b, align 1
   %3 = insertelement <4 x i8> poison, i8 %2, i64 0
   %lane2 = shufflevector <4 x i8> %3, <4 x i8> poison, <4 x i32> zeroinitializer
   %shuffle.i = shufflevector <4 x i8> %lane, <4 x i8> %lane2, <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
   ret <8 x i8> %shuffle.i
 }
 
 define <8 x i8> @test7(ptr nocapture noundef readonly %a, ptr nocapture noundef readonly %b) {
 ; CHECK-LABEL: test7:
 ; CHECK:       // %bb.0: // %entry
-; CHECK-NEXT:    ld1r { v0.8b }, [x0]
-; CHECK-NEXT:    adrp x8, .LCPI6_0
-; CHECK-NEXT:    ld1r { v1.8b }, [x1]
-; CHECK-NEXT:    mov v1.d[1], v0.d[0]
-; CHECK-NEXT:    ldr d0, [x8, :lo12:.LCPI6_0]
-; CHECK-NEXT:    tbl v0.8b, { v1.16b }, v0.8b
+; CHECK-NEXT:    ld1r { v1.8b }, [x0]
+; CHECK-NEXT:    ld1r { v0.8b }, [x1]
+; CHECK-NEXT:    mov v0.s[1], v1.s[1]
+; CHECK-NEXT:    // kill: def $d0 killed $d0 killed $q0
 ; CHECK-NEXT:    ret
 entry:
   %0 = load i8, ptr %a, align 1
   %1 = insertelement <4 x i8> poison, i8 %0, i64 0
   %lane = shufflevector <4 x i8> %1, <4 x i8> poison, <4 x i32> zeroinitializer
   %2 = load i8, ptr %b, align 1
   %3 = insertelement <4 x i8> poison, i8 %2, i64 0
   %lane2 = shufflevector <4 x i8> %3, <4 x i8> poison, <4 x i32> zeroinitializer