[AIE2P] Combine G_SHUFFLE_VECTOR into Unmerge+Concat

Xilinx · Feb 14, 2025 · 97f321a · 97f321a
1 parent a9a9e6f
commit 97f321a
Show file tree

Hide file tree

Showing 3 changed files with 112 additions and 23 deletions.
diff --git a/llvm/lib/Target/AIE/AIECombinerHelper.cpp b/llvm/lib/Target/AIE/AIECombinerHelper.cpp
@@ -1974,6 +1974,20 @@ buildExtractSubvector(MachineIRBuilder &B, MachineRegisterInfo &MRI,
   return B.buildInstr(Opc, {DstVecReg}, {NewSrcReg, Cst});
 }
 
+static void buildUnmergeVector(MachineIRBuilder &B, MachineRegisterInfo &MRI,
+                               Register DstReg, Register SrcReg,
+                               unsigned NumSubVectors, unsigned SubIdx) {
+  const LLT DstTy = MRI.getType(DstReg);
+  SmallVector<Register, 4> SubVecs;
+  for (unsigned I = 0; I < NumSubVectors; I++) {
+    if (I == (unsigned)SubIdx)
+      SubVecs.push_back(DstReg);
+    else
+      SubVecs.push_back(MRI.createGenericVirtualRegister(DstTy));
+  }
+  B.buildUnmerge(SubVecs, SrcReg);
+}
+
 /// Match something like this:
 ///  %1:_(<16 x s32>) = COPY $x0
 ///  %2:_(<16 x s32>) = COPY $x1
@@ -1990,18 +2004,10 @@ static bool matchShuffleToUnmerge(MachineInstr &MI, MachineRegisterInfo &MRI,
                                   unsigned NumSubVectors) {
   const Register DstReg = MI.getOperand(0).getReg();
   const Register Src1Reg = MI.getOperand(1).getReg();
-  const LLT DstTy = MRI.getType(DstReg);
 
   // TODO: Select into G_EXTRACT_SUBVECTOR once it is more widely supported
   MatchInfo = [=, &MRI](MachineIRBuilder &B) {
-    SmallVector<Register, 4> SubVecs;
-    for (unsigned I = 0; I < NumSubVectors; I++) {
-      if (I == (unsigned)SubIdx)
-        SubVecs.push_back(DstReg);
-      else
-        SubVecs.push_back(MRI.createGenericVirtualRegister(DstTy));
-    }
-    B.buildUnmerge(SubVecs, Src1Reg);
+    buildUnmergeVector(B, MRI, DstReg, Src1Reg, NumSubVectors, SubIdx);
   };
   return true;
 }
@@ -2227,20 +2233,36 @@ static bool matchShuffleToSubvecBroadcast(MachineInstr &MI,
   const LLT ElemTy = Src1Ty.getElementType();
   const LLT DstSubvecType =
       LLT::fixed_vector(SplatMaskLen, ElemTy.getSizeInBits());
+  const unsigned SubIdx = SplatMaskStart / SplatMaskLen;
+  Register ExtractSubvecDstReg =
+      MRI.createGenericVirtualRegister(DstSubvecType);
 
   // Check whether we can extract the subvector
-  if (!checkExtractSubvectorPrerequisites(TII, DstSubvecType, Src1Ty))
-    return false;
+  bool canExtractSubvector =
+      checkExtractSubvectorPrerequisites(TII, DstSubvecType, Src1Ty);
+  if (canExtractSubvector) {
+    MatchInfo = [=, &MRI, &TII](MachineIRBuilder &B) {
+      auto Extract = buildExtractSubvector(B, MRI, TII, ExtractSubvecDstReg,
+                                           Src1Reg, SubIdx);
+      buildBroadcastVector(B, MRI, Extract.getReg(0), DstReg);
+    };
+    return true;
+  }
 
-  MatchInfo = [=, &MRI, &TII](MachineIRBuilder &B) {
-    Register ExtractSubvecDstReg =
-        MRI.createGenericVirtualRegister(DstSubvecType);
-    auto Extract =
-        buildExtractSubvector(B, MRI, TII, ExtractSubvecDstReg, Src1Reg,
-                              SplatMaskStart / SplatMaskLen);
-    buildBroadcastVector(B, MRI, Extract.getReg(0), DstReg);
-  };
-  return true;
+  // If we cannot extract the subvector, we try to apply UNMERGE + CONCAT
+  const unsigned NumSubVectors = NumSrcElems / SplatMaskLen;
+  if (NumDstElems == SplatMaskLen * 2) {
+    MatchInfo = [=, &MRI](MachineIRBuilder &B) {
+      buildUnmergeVector(B, MRI, ExtractSubvecDstReg, Src1Reg, NumSubVectors,
+                         SubIdx);
+      const SmallVector<Register, 2> ConcatOps = {ExtractSubvecDstReg,
+                                                  ExtractSubvecDstReg};
+      B.buildConcatVectors({DstReg}, ConcatOps);
+    };
+    return true;
+  }
+
+  return false;
 }
 
 /// Match something like this:

diff --git a/llvm/test/CodeGen/AIE/GlobalISel/prelegalizercombiner-shuffle-vector.mir b/llvm/test/CodeGen/AIE/GlobalISel/prelegalizercombiner-shuffle-vector.mir
@@ -776,9 +776,9 @@ body:             |
     ; CHECK: liveins: $dm0, $dm1
     ; CHECK-NEXT: {{  $}}
     ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(<64 x s32>) = COPY $dm0
-    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:_(<64 x s32>) = COPY $dm1
-    ; CHECK-NEXT: [[SHUF:%[0-9]+]]:_(<64 x s32>) = G_SHUFFLE_VECTOR [[COPY]](<64 x s32>), [[COPY1]], shufflemask(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, undef, undef, undef, undef, undef, undef, undef, undef, undef, undef, undef, undef, undef, undef, undef, undef, undef, undef, undef, undef, undef, undef, undef, undef, undef, undef, undef, undef, undef, undef, undef, undef)
-    ; CHECK-NEXT: PseudoRET implicit $lr, implicit [[SHUF]](<64 x s32>)
+    ; CHECK-NEXT: [[UV:%[0-9]+]]:_(<32 x s32>), [[UV1:%[0-9]+]]:_(<32 x s32>) = G_UNMERGE_VALUES [[COPY]](<64 x s32>)
+    ; CHECK-NEXT: [[CONCAT_VECTORS:%[0-9]+]]:_(<64 x s32>) = G_CONCAT_VECTORS [[UV]](<32 x s32>), [[UV]](<32 x s32>)
+    ; CHECK-NEXT: PseudoRET implicit $lr, implicit [[CONCAT_VECTORS]](<64 x s32>)
     %1:_(<64 x s32>) = COPY $dm0
     %2:_(<64 x s32>) = COPY $dm1
     %0:_(<64 x s32>) = G_SHUFFLE_VECTOR %1(<64 x s32>), %2(<64 x s32>), shufflemask(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1)
@@ -928,3 +928,40 @@ body:             |
     %0:_(<16 x s32>) = G_SHUFFLE_VECTOR %1(<16 x s32>), %2(<16 x s32>), shufflemask(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15)
     PseudoRET implicit $lr, implicit %0
 ...
+# Test G_SHUFFLE_VECTOR to UNMERGE+CONCAT
+---
+name: shuffle_vector_to_unmerge_concat_lo
+tracksRegLiveness: true
+body:             |
+  bb.1:
+    liveins: $x0, $x1
+    ; CHECK-LABEL: name: shuffle_vector_to_unmerge_concat_lo
+    ; CHECK: liveins: $x0, $x1
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(<16 x s32>) = COPY $x0
+    ; CHECK-NEXT: [[AIE_UNPAD_VECTOR:%[0-9]+]]:_(<8 x s32>) = G_AIE_UNPAD_VECTOR [[COPY]](<16 x s32>)
+    ; CHECK-NEXT: [[CONCAT_VECTORS:%[0-9]+]]:_(<16 x s32>) = G_CONCAT_VECTORS [[AIE_UNPAD_VECTOR]](<8 x s32>), [[AIE_UNPAD_VECTOR]](<8 x s32>)
+    ; CHECK-NEXT: PseudoRET implicit $lr, implicit [[CONCAT_VECTORS]](<16 x s32>)
+    %1:_(<16 x s32>) = COPY $x0
+    %2:_(<16 x s32>) = COPY $x1
+    %0:_(<16 x s32>) = G_SHUFFLE_VECTOR %1(<16 x s32>), %2(<16 x s32>), shufflemask(0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7)
+    PseudoRET implicit $lr, implicit %0
+...
+---
+name: shuffle_vector_to_unmerge_concat_hi
+tracksRegLiveness: true
+body:             |
+  bb.1:
+    liveins: $x0, $x1
+    ; CHECK-LABEL: name: shuffle_vector_to_unmerge_concat_hi
+    ; CHECK: liveins: $x0, $x1
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(<16 x s32>) = COPY $x0
+    ; CHECK-NEXT: [[UV:%[0-9]+]]:_(<8 x s32>), [[UV1:%[0-9]+]]:_(<8 x s32>) = G_UNMERGE_VALUES [[COPY]](<16 x s32>)
+    ; CHECK-NEXT: [[CONCAT_VECTORS:%[0-9]+]]:_(<16 x s32>) = G_CONCAT_VECTORS [[UV1]](<8 x s32>), [[UV1]](<8 x s32>)
+    ; CHECK-NEXT: PseudoRET implicit $lr, implicit [[CONCAT_VECTORS]](<16 x s32>)
+    %1:_(<16 x s32>) = COPY $x0
+    %2:_(<16 x s32>) = COPY $x1
+    %0:_(<16 x s32>) = G_SHUFFLE_VECTOR %1(<16 x s32>), %2(<16 x s32>), shufflemask(8, 9, 10, 11, 12, 13, 14, 15, 8, 9, 10, 11, 12, 13, 14, 15)
+    PseudoRET implicit $lr, implicit %0
+...
diff --git a/llvm/test/CodeGen/AIE/aie2p/shufflevec.ll b/llvm/test/CodeGen/AIE/aie2p/shufflevec.ll
@@ -521,3 +521,33 @@ entry:
   %shuffle = shufflevector <64 x i32> %a, <64 x i32> %b, <32 x i32> <i32 0, i32  1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7, i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15, i32 16, i32 17, i32 18, i32 19, i32 20, i32 21, i32 22, i32 23, i32 24, i32 25, i32 26, i32 27, i32 28, i32 29, i32 30, i32 31>
   ret <32 x i32> %shuffle
 }
+
+define <16 x i32> @shuffle_vector_to_unmerge_concat_lo(<16 x i32> noundef %a, <16 x i32> noundef %b) {
+; CHECK-LABEL: shuffle_vector_to_unmerge_concat_lo:
+; CHECK:         .p2align 4
+; CHECK-NEXT:  // %bb.0: // %entry
+; CHECK-NEXT:    nopa ; nopb ; nops ; ret lr; nopm ; nopv
+; CHECK-NEXT:    nopx // Delay Slot 5
+; CHECK-NEXT:    nop // Delay Slot 4
+; CHECK-NEXT:    vmov x0, x2 // Delay Slot 3
+; CHECK-NEXT:    vmov wh0, wl0 // Delay Slot 2
+; CHECK-NEXT:    nop // Delay Slot 1
+entry:
+  %shuffle = shufflevector <16 x i32> %a, <16 x i32> %b, <16 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7, i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
+  ret <16 x i32> %shuffle
+}
+
+define <16 x i32> @shuffle_vector_to_unmerge_concat_hi(<16 x i32> noundef %a, <16 x i32> noundef %b) {
+; CHECK-LABEL: shuffle_vector_to_unmerge_concat_hi:
+; CHECK:         .p2align 4
+; CHECK-NEXT:  // %bb.0: // %entry
+; CHECK-NEXT:    nopa ; nopb ; nops ; ret lr; nopm ; nopv
+; CHECK-NEXT:    nopx // Delay Slot 5
+; CHECK-NEXT:    nop // Delay Slot 4
+; CHECK-NEXT:    vmov x0, x2 // Delay Slot 3
+; CHECK-NEXT:    vmov wl0, wh0 // Delay Slot 2
+; CHECK-NEXT:    nop // Delay Slot 1
+entry:
+  %shuffle = shufflevector <16 x i32> %a, <16 x i32> %b, <16 x i32> <i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15, i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15>
+  ret <16 x i32> %shuffle
+}