[AIE2P] Fix combine G_SHUFFLE_VEC to G_AIE_VSEL for <64xs8>

Xilinx · Feb 6, 2025 · ca266a2 · ca266a2
1 parent 9af5b07
commit ca266a2
Show file tree

Hide file tree

Showing 4 changed files with 53 additions and 6 deletions.
diff --git a/llvm/lib/Target/AIE/AIECombinerHelper.cpp b/llvm/lib/Target/AIE/AIECombinerHelper.cpp
@@ -1828,7 +1828,9 @@ bool llvm::matchShuffleToVSel(MachineInstr &MI, MachineRegisterInfo &MRI,
   }
 
   MatchInfo = [=, &TII](MachineIRBuilder &B) {
-    MachineInstrBuilder MaskReg = B.buildConstant(LLT::scalar(32), DstMask);
+    const unsigned ScalarSize = NumMaskElems == 64 ? 64 : 32;
+    MachineInstrBuilder MaskReg =
+        B.buildConstant(LLT::scalar(ScalarSize), DstMask);
     const unsigned VSelOpc = TII.getGenericVSelOpcode();
     B.buildInstr(VSelOpc, {DstReg}, {Src1Reg, Src2Reg, MaskReg});
   };

diff --git a/llvm/lib/Target/AIE/aie2p/AIE2PInstrPatterns.td b/llvm/lib/Target/AIE/aie2p/AIE2PInstrPatterns.td
@@ -1062,7 +1062,7 @@ def vsel_node : SDNode<"AIE2P::G_AIE_VSEL",
                        SDTypeProfile<1, 3, []>>;
 def : GINodeEquiv<G_AIE_VSEL, vsel_node>;
 
-def : Pat<(v64i8(vsel_node (v64i8 VEC512:$src1), (v64i8 VEC512:$src2), (v2i32 eL:$lsel))), 
+def : Pat<(v64i8(vsel_node (v64i8 VEC512:$src1), (v64i8 VEC512:$src2), (i64 eL:$lsel))), 
   (VSEL_8 VEC512:$src1, VEC512:$src2, eL:$lsel)>;
 def : Pat<(v32i16(vsel_node (v32i16 VEC512:$src1), (v32i16 VEC512:$src2), (i32 eRS16:$rsel))), 
   (VSEL_16 VEC512:$src1, VEC512:$src2, eRS16:$rsel)>;

diff --git a/llvm/test/CodeGen/AIE/GlobalISel/prelegalizercombiner-shuffle-vector.mir b/llvm/test/CodeGen/AIE/GlobalISel/prelegalizercombiner-shuffle-vector.mir
@@ -140,7 +140,7 @@ body:             |
     PseudoRET implicit $lr, implicit %0
 ...
 ---
-name: shuffle_vector_vsel_valid_mask_ordered_indices
+name: shuffle_vector_vsel_valid_mask_16_s32
 alignment:       16
 exposesReturnsTwice: false
 legalized:       false
@@ -149,7 +149,7 @@ body:             |
   bb.1.entry:
     liveins: $x2
 
-    ; CHECK-LABEL: name: shuffle_vector_vsel_valid_mask_ordered_indices
+    ; CHECK-LABEL: name: shuffle_vector_vsel_valid_mask_16_s32
     ; CHECK: liveins: $x2
     ; CHECK-NEXT: {{  $}}
     ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(<16 x s32>) = COPY $x2
@@ -163,6 +163,52 @@ body:             |
     PseudoRET implicit $lr, implicit %0
 ...
 ---
+name: shuffle_vector_vsel_valid_mask_32_s16
+alignment:       16
+exposesReturnsTwice: false
+legalized:       false
+tracksRegLiveness: true
+body:             |
+  bb.1.entry:
+    liveins: $x2
+
+    ; CHECK-LABEL: name: shuffle_vector_vsel_valid_mask_32_s16
+    ; CHECK: liveins: $x2
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(<32 x s16>) = COPY $x2
+    ; CHECK-NEXT: [[DEF:%[0-9]+]]:_(<32 x s16>) = G_IMPLICIT_DEF
+    ; CHECK-NEXT: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 10
+    ; CHECK-NEXT: [[AIE_VSEL:%[0-9]+]]:_(<32 x s16>) = G_AIE_VSEL [[DEF]], [[COPY]], [[C]](s32)
+    ; CHECK-NEXT: PseudoRET implicit $lr, implicit [[AIE_VSEL]](<32 x s16>)
+    %1:_(<32 x s16>) = COPY $x2
+    %8:_(<32 x s16>) = G_IMPLICIT_DEF
+    %0:_(<32 x s16>) = G_SHUFFLE_VECTOR %8(<32 x s16>), %1, shufflemask(0, 33, 2, 35, 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)
+    PseudoRET implicit $lr, implicit %0
+...
+---
+name: shuffle_vector_vsel_valid_mask_64_s8
+alignment:       16
+exposesReturnsTwice: false
+legalized:       false
+tracksRegLiveness: true
+body:             |
+  bb.1.entry:
+    liveins: $x2
+
+    ; CHECK-LABEL: name: shuffle_vector_vsel_valid_mask_64_s8
+    ; CHECK: liveins: $x2
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(<64 x s8>) = COPY $x2
+    ; CHECK-NEXT: [[DEF:%[0-9]+]]:_(<64 x s8>) = G_IMPLICIT_DEF
+    ; CHECK-NEXT: [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 5
+    ; CHECK-NEXT: [[AIE_VSEL:%[0-9]+]]:_(<64 x s8>) = G_AIE_VSEL [[DEF]], [[COPY]], [[C]](s64)
+    ; CHECK-NEXT: PseudoRET implicit $lr, implicit [[AIE_VSEL]](<64 x s8>)
+    %1:_(<64 x s8>) = COPY $x2
+    %8:_(<64 x s8>) = G_IMPLICIT_DEF
+    %0:_(<64 x s8>) = G_SHUFFLE_VECTOR %8(<64 x s8>), %1, shufflemask(64, 1, 66, 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, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63)
+    PseudoRET implicit $lr, implicit %0
+...
+---
 name: shuffle_vector_vsel_valid_mask_ordered_indices_with_undef
 alignment:       16
 exposesReturnsTwice: false

diff --git a/llvm/test/CodeGen/AIE/aie2p/GlobalIsel/inst-select-vsel.mir b/llvm/test/CodeGen/AIE/aie2p/GlobalIsel/inst-select-vsel.mir
@@ -111,8 +111,7 @@ body:             |
     %4:gprregbank(s32) = COPY $r0
     %5:gprregbank(s32) = COPY $r1
     %3:gprregbank(s64) = G_MERGE_VALUES %4:gprregbank(s32), %5:gprregbank(s32)
-    %6:gprregbank(<2 x s32>) = G_BITCAST %3:gprregbank(s64)
-    %0:vregbank(<64 x s8>) = G_AIE_VSEL %1:vregbank(<64 x s8>), %2:vregbank(<64 x s8>), %6:gprregbank(<2 x s32>)
+    %0:vregbank(<64 x s8>) = G_AIE_VSEL %1:vregbank(<64 x s8>), %2:vregbank(<64 x s8>), %3:gprregbank(s64)
     $x0 = COPY %0:vregbank(<64 x s8>)
     PseudoRET implicit $lr, implicit $x0
 ...