[DAG] visitVECTOR_SHUFFLE - MergeInnerShuffle - improve shuffle(shuffle(x,y),shuffle(x,y)) merging

[RKSimon]

Summary

MergeInnerShuffle currently attempts to merge shuffle(shuffle(x,y),z) patterns into a single shuffle, using 1 or 2 of the x,y,z ops.

However if we already match 2 ops we might be able to handle the third op if its also a shuffle that references one of the previous ops, allowing us to handle some cases like:

shuffle(shuffle(x,y),shuffle(x,y))
shuffle(shuffle(shuffle(x,z),y),z)
shuffle(shuffle(x,shuffle(x,y)),z)
etc.

This isn't an exhaustive match and is dependent on the order the candidate ops are encountered - if one of the matched ops was a shuffle that was peek-able we don't go back and try to split that, I haven't found much need for that amount of analysis yet.

This is a preliminary patch that will allow us to later improve x86 HADD/HSUB matching - but needs to be reviewed separately as its in generic code and affects existing Thumb2 tests.

[dmgreen]

LGTM. We sometimes generate a lot of shuffles in an attempt to do lane interleaving and I know the simplification of them isn't always what it could be once all the lowering has happened. I thought more happened through simplifying buildvectors but apparently not. This looks like a good continuation to the existing code.

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

for (unsigned i = 0; i != NumElts; ++i) {
       if (!SV1.getNode() || SV1 == CurrentVec) {
         // Ok. CurrentVec is the right hand side.
         // Update the mask accordingly.
         SV1 = CurrentVec;
         Mask.push_back(Idx + NumElts);
         continue;
       }
 
+      // Last chance - see if the vector is another shuffle and if it
+      // uses one of the existing candidate shuffle ops.
+      if (auto *CurrentSVN = dyn_cast<ShuffleVectorSDNode>(CurrentVec)) {
+        int InnerIdx = CurrentSVN->getMaskElt(Idx);
+        if (InnerIdx < 0) {
+          Mask.push_back(-1);
+          continue;
+        }
+        SDValue InnerVec = (InnerIdx < (int)NumElts)
+                               ? CurrentSVN->getOperand(0)
+                               : CurrentSVN->getOperand(1);
+        if (InnerVec.isUndef()) {
+          Mask.push_back(-1);
+          continue;
+        }
+        InnerIdx %= NumElts;
+        if (InnerVec == SV0) {
+          Mask.push_back(InnerIdx);
+          continue;
+        }
+        if (InnerVec == SV1) {
+          Mask.push_back(InnerIdx + NumElts);
+          continue;
+        }
+      }
+
       // Bail out if we cannot convert the shuffle pair into a single shuffle.
       return false;
     }
     return true;
   };
 
   // Try to fold according to rules:
   //   shuffle(shuffle(A, B, M0), C, M1) -> shuffle(A, B, M2)

llvm/test/CodeGen/Thumb2/mve-vmull-loop.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=thumbv8.1m.main-none-none-eabi -mattr=+mve -verify-machineinstrs %s -o - | FileCheck %s --check-prefix=CHECK
 
 define arm_aapcs_vfpcc void @test32(i32* noalias nocapture readonly %x, i32* noalias nocapture readonly %y, i32* nocapture %z, i32 %n) {
 ; CHECK-LABEL: test32:
 ; CHECK:       @ %bb.0: @ %entry
 ; CHECK-NEXT:    .save {r4, r5, r7, lr}
 ; CHECK-NEXT:    push {r4, r5, r7, lr}
 ; CHECK-NEXT:    cmp r3, #1
 ; CHECK-NEXT:    blt .LBB0_2
 ; CHECK-NEXT:  .LBB0_1: @ %vector.body
 ; CHECK-NEXT:    @ =>This Inner Loop Header: Depth=1
-; CHECK-NEXT:    vldrw.u32 q1, [r0], #16
-; CHECK-NEXT:    vldrw.u32 q2, [r1], #16
+; CHECK-NEXT:    vldrw.u32 q0, [r0], #16
+; CHECK-NEXT:    vldrw.u32 q1, [r1], #16
 ; CHECK-NEXT:    subs r3, #4
-; CHECK-NEXT:    vmullt.s32 q0, q2, q1
-; CHECK-NEXT:    vmullb.s32 q3, q2, q1
-; CHECK-NEXT:    vmov r5, s3
-; CHECK-NEXT:    vmov r12, s2
-; CHECK-NEXT:    vmov r7, s1
+; CHECK-NEXT:    vmullb.s32 q2, q1, q0
+; CHECK-NEXT:    vmullt.s32 q3, q1, q0
+; CHECK-NEXT:    vmov r5, s11
+; CHECK-NEXT:    vmov r12, s10
 ; CHECK-NEXT:    lsrl r12, r5, #31
-; CHECK-NEXT:    vmov r4, s0
-; CHECK-NEXT:    lsrl r4, r7, #31
-; CHECK-NEXT:    vmov q0[2], q0[0], r4, r12
-; CHECK-NEXT:    vmov r12, s14
-; CHECK-NEXT:    vmov q0[3], q0[1], r7, r5
+; CHECK-NEXT:    vmov r4, s8
+; CHECK-NEXT:    vmov r5, s9
+; CHECK-NEXT:    lsrl r4, r5, #31
+; CHECK-NEXT:    vmov q2[2], q2[0], r4, r12
 ; CHECK-NEXT:    vmov r5, s15
-; CHECK-NEXT:    vmov r7, s13
+; CHECK-NEXT:    vmov r12, s14
 ; CHECK-NEXT:    lsrl r12, r5, #31
 ; CHECK-NEXT:    vmov r4, s12
-; CHECK-NEXT:    lsrl r4, r7, #31
-; CHECK-NEXT:    vmov q1[2], q1[0], r4, r12
-; CHECK-NEXT:    vmov q1[3], q1[1], r7, r5
-; CHECK-NEXT:    vmov.f32 s8, s6
-; CHECK-NEXT:    vmov.f32 s9, s7
-; CHECK-NEXT:    vmov.f32 s6, s0
-; CHECK-NEXT:    vmov.f32 s7, s1
-; CHECK-NEXT:    vmov.f32 s10, s2
-; CHECK-NEXT:    vmov.f32 s5, s6
-; CHECK-NEXT:    vmov.f32 s11, s3
-; CHECK-NEXT:    vmov.f32 s6, s8
-; CHECK-NEXT:    vmov.f32 s7, s10
-; CHECK-NEXT:    vstrb.8 q1, [r2], #16
+; CHECK-NEXT:    vmov r5, s13
+; CHECK-NEXT:    lsrl r4, r5, #31
+; CHECK-NEXT:    vmov q2[3], q2[1], r4, r12
+; CHECK-NEXT:    vstrb.8 q2, [r2], #16
 ; CHECK-NEXT:    bne .LBB0_1
 ; CHECK-NEXT:  .LBB0_2: @ %for.cond.cleanup
 ; CHECK-NEXT:    pop {r4, r5, r7, pc}
 entry:
   %0 = and i32 %n, 3
   %cmp = icmp eq i32 %0, 0
   %cmp113 = icmp sgt i32 %n, 0
   br i1 %cmp113, label %vector.body, label %for.cond.cleanup

llvm/test/CodeGen/Thumb2/mve-vst2.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=thumbv8.1m.main-none-none-eabi -mattr=+mve.fp,+fp64 -verify-machineinstrs %s -o - | FileCheck %s
 
 ; i32
 
 define void @vst2_v2i32(<2 x i32> *%src, <4 x i32> *%dst) {
 ; CHECK-LABEL: vst2_v2i32:
 ; CHECK:       @ %bb.0: @ %entry
-; CHECK-NEXT:    ldrd r12, r3, [r0]
-; CHECK-NEXT:    ldrd r2, r0, [r0, #8]
-; CHECK-NEXT:    vmov q0[2], q0[0], r12, r3
-; CHECK-NEXT:    vmov.f64 d2, d1
-; CHECK-NEXT:    vmov q2[2], q2[0], r2, r0
-; CHECK-NEXT:    vmov.f32 s5, s3
-; CHECK-NEXT:    vmov.f32 s2, s8
-; CHECK-NEXT:    vmov.f32 s3, s9
-; CHECK-NEXT:    vmov.f32 s6, s10
-; CHECK-NEXT:    vmov.f32 s1, s2
-; CHECK-NEXT:    vmov.f32 s7, s11
-; CHECK-NEXT:    vmov.f32 s2, s4
-; CHECK-NEXT:    vmov.f32 s3, s6
+; CHECK-NEXT:    ldm.w r0, {r2, r3, r12}
+; CHECK-NEXT:    ldr r0, [r0, #12]
+; CHECK-NEXT:    vmov q0[2], q0[0], r2, r3
+; CHECK-NEXT:    vmov q0[3], q0[1], r12, r0
 ; CHECK-NEXT:    vstrw.32 q0, [r1]
 ; CHECK-NEXT:    bx lr
 entry:
   %s1 = getelementptr <2 x i32>, <2 x i32>* %src, i32 0
   %l1 = load <2 x i32>, <2 x i32>* %s1, align 4
   %s2 = getelementptr <2 x i32>, <2 x i32>* %src, i32 1
   %l2 = load <2 x i32>, <2 x i32>* %s2, align 4
   %s = shufflevector <2 x i32> %l1, <2 x i32> %l2, <4 x i32> <i32 0, i32 2, i32 1, i32 3>

 }
 
 ; f32
 
 define void @vst2_v2f32(<2 x float> *%src, <4 x float> *%dst) {
 ; CHECK-LABEL: vst2_v2f32:
 ; CHECK:       @ %bb.0: @ %entry
 ; CHECK-NEXT:    vldr s0, [r0]
-; CHECK-NEXT:    vldr s4, [r0, #4]
+; CHECK-NEXT:    vldr s2, [r0, #4]
 ; CHECK-NEXT:    vldr s1, [r0, #8]
-; CHECK-NEXT:    vldr s5, [r0, #12]
-; CHECK-NEXT:    vmov.f32 s2, s4
-; CHECK-NEXT:    vmov.f32 s3, s5
+; CHECK-NEXT:    vldr s3, [r0, #12]
 ; CHECK-NEXT:    vstrw.32 q0, [r1]
 ; CHECK-NEXT:    bx lr
 entry:
   %s1 = getelementptr <2 x float>, <2 x float>* %src, i32 0
   %l1 = load <2 x float>, <2 x float>* %s1, align 4
   %s2 = getelementptr <2 x float>, <2 x float>* %src, i32 1
   %l2 = load <2 x float>, <2 x float>* %s2, align 4
   %s = shufflevector <2 x float> %l1, <2 x float> %l2, <4 x i32> <i32 0, i32 2, i32 1, i32 3>