[AIE2p] Use tablegen pattern for trunc where dst is 1/4 bitsize

Author: newling

llvm/lib/Target/AIE/aie2p/AIE2PInstrPatterns.td

@@ -15,6 +15,7 @@
 
 include "AIEBaseInstrPatterns.td"
 
+
 // Placeholder for a bare frameindex. This pseudo represents the
 // pointer register to be allocated, initialized with the address
 // represented by the frameindex in its only operand.
@@ -1257,26 +1258,52 @@ defm : Extract_512<i64, v8i64, (i32 c6u:$idx), VEXTRACT_64_vec_extract_imm_vaddS
 // instruction. Modes 0, 2, and 4 correspond to the deinterleaved operation of
 // 1, 2, and 4 bytes respectively on the concatenated src0 and src1 input
 // vectors of the VSHUFFLE instruction.
-class Trunc1024Pat<ValueType DstTy, ValueType SrcTy, Instruction ShuffleInstOpc, int Mode> :
+
+// | 1024-bits -> 512-bits | 512-bits -> 256-bits  | 1024-bits -> 256-bits |
+// +=======================+=======================+=======================|
+// |   v16i64 -> v16i32    |   v8i32  -> v8i16     |   v16i64 -> v16i16    |
+// |   v32i32 -> v32i16    |   v16i16 -> v16i8     |   v32i32 -> v32i8     |
+// |   v64i16 -> v64i8     |   v32i8  -> v32i4     |                       |
+// +=======================+=======================+=======================|
+class Trunc1024To512Pat<ValueType DstTy, ValueType SrcTy, Instruction ShuffleInstOpc, int Mode> :
   Pat<(DstTy (trunc SrcTy:$s1)),
     (ShuffleInstOpc
       (EXTRACT_SUBREG VEC1024:$s1, sub_512_lo ),
       (EXTRACT_SUBREG VEC1024:$s1, sub_512_hi ),
       (MOV_RLC_imm11_pseudo (i32 Mode)))>;
 
-def : Trunc1024Pat<v16i32, v16i64, VSHUFFLE_vec_shuffle_x, 4>;
-def : Trunc1024Pat<v16i32, v16i64, VSHUFFLE_vec_shuffle_bm, 4>;
-def : Trunc1024Pat<v32i16, v32i32, VSHUFFLE_vec_shuffle_x, 2>;
-def : Trunc1024Pat<v64i8, v64i16, VSHUFFLE_vec_shuffle_x, 0>;
+def : Trunc1024To512Pat<v16i32, v16i64, VSHUFFLE_vec_shuffle_x, 4>;
+def : Trunc1024To512Pat<v16i32, v16i64, VSHUFFLE_vec_shuffle_bm, 4>;
+def : Trunc1024To512Pat<v32i16, v32i32, VSHUFFLE_vec_shuffle_x, 2>;
+def : Trunc1024To512Pat<v64i8, v64i16, VSHUFFLE_vec_shuffle_x, 0>;
 
-class Trunc512Pat<ValueType DstTy, ValueType SrcTy, int Mode> :
+class Trunc512To256Pat<ValueType DstTy, ValueType SrcTy, int Mode> :
   Pat<(DstTy (trunc SrcTy:$s1)),
     (EXTRACT_SUBREG
       (VSHUFFLE_vec_shuffle_x VEC512:$s1, VEC512:$s1, (MOV_RLC_imm11_pseudo (i32 Mode))), 
       sub_256_lo)>;
-def : Trunc512Pat<v8i32, v8i64, 4>;
-def : Trunc512Pat<v16i16, v16i32, 2>;
-def : Trunc512Pat<v32i8, v32i16, 0>;
+def : Trunc512To256Pat<v8i32, v8i64, 4>;
+def : Trunc512To256Pat<v16i16, v16i32, 2>;
+def : Trunc512To256Pat<v32i8, v32i16, 0>;
+
+// This is effectively Trunc1024To512 followed by Trunc512To256.
+class Trunc1024To256Pat<ValueType DstTy, ValueType SrcTy, int LargeMode, int SmallMode> :
+  Pat<(DstTy (trunc SrcTy:$s1)),
+  (EXTRACT_SUBREG
+    (VSHUFFLE_vec_shuffle_x
+        (VSHUFFLE_vec_shuffle_x
+          (EXTRACT_SUBREG VEC1024:$s1, sub_512_lo),
+          (EXTRACT_SUBREG VEC1024:$s1, sub_512_hi),
+          (MOV_RLC_imm11_pseudo (i32 LargeMode))),
+        (VSHUFFLE_vec_shuffle_x
+          (EXTRACT_SUBREG VEC1024:$s1, sub_512_lo),
+          (EXTRACT_SUBREG VEC1024:$s1, sub_512_hi),
+          (MOV_RLC_imm11_pseudo (i32 LargeMode))),
+      (MOV_RLC_imm11_pseudo (i32 SmallMode))),
+    sub_256_lo)>;
+
+def : Trunc1024To256Pat<v32i8, v32i32, 2, 0>;
+def : Trunc1024To256Pat<v16i16, v16i64, 4, 2>;
 
 class EventPat<AIE2PInst Inst, dag Imm> :
   Pat<(int_aie2p_event Imm), (Inst)>;

llvm/lib/Target/AIE/aie2p/AIE2PLegalizerInfo.cpp

@@ -281,13 +281,24 @@ AIE2PLegalizerInfo::AIE2PLegalizerInfo(const AIE2PSubtarget &ST)
 
   getActionDefinitionsBuilder(G_TRUNC)
       .legalIf([=](const LegalityQuery &Query) {
+        // Return true if there is a tablegen pattern to lower truncs on vectors
+        // of specific element types and lengths to shuffles.
         const LLT &SrcTy = Query.Types[1];
         const LLT &DstTy = Query.Types[0];
-        return SrcTy.isVector() && DstTy.isVector() &&
-               (SrcTy.getSizeInBits() == 512 ||
-                SrcTy.getSizeInBits() == 1024) &&
-               DstTy.getElementType().getSizeInBits() * 2 ==
-                   SrcTy.getElementType().getSizeInBits();
+
+        if (!SrcTy.isVector() || !DstTy.isVector())
+          return false;
+
+        const auto SrcElmBits = SrcTy.getElementType().getSizeInBits();
+        if (SrcElmBits != 64 && SrcElmBits != 32 && SrcElmBits != 16)
+          return false;
+
+        const TypeSize SrcBits = SrcTy.getSizeInBits();
+        const TypeSize DstBits = DstTy.getSizeInBits();
+
+        return ((SrcBits == 1024 && DstBits == 256) ||
+                (SrcBits == 1024 && DstBits == 512) ||
+                (SrcBits == 512 && DstBits == 256));
       })
       .legalIf([=](const LegalityQuery &Query) {
         const LLT &SrcTy = Query.Types[1];

llvm/test/CodeGen/AIE/aie2p/GlobalIsel/inst-select-trunc.mir

@@ -184,4 +184,52 @@ body:             |
     %0:vregbank(<32 x s8>) = G_TRUNC %1(<32 x s16>)
     PseudoRET implicit $lr, implicit %0
 ...
+---
+name:            v32s8_trunc_v32s32_vec1024
+legalized:       true
+regBankSelected: true
+body:             |
+  bb.1.entry:
+    ; CHECK-LABEL: name: v32s8_trunc_v32s32_vec1024
+    ; CHECK: [[DEF:%[0-9]+]]:vec1024 = IMPLICIT_DEF
+    ; CHECK-NEXT: [[MOV_RLC_imm11_pseudo:%[0-9]+]]:er = MOV_RLC_imm11_pseudo 0
+    ; CHECK-NEXT: [[MOV_RLC_imm11_pseudo1:%[0-9]+]]:er = MOV_RLC_imm11_pseudo 2
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:exo = COPY [[DEF]].sub_512_hi
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:exe = COPY [[DEF]].sub_512_lo
+    ; CHECK-NEXT: [[VSHUFFLE_vec_shuffle_x:%[0-9]+]]:vec512 = VSHUFFLE_vec_shuffle_x [[COPY1]], [[COPY]], [[MOV_RLC_imm11_pseudo1]]
+    ; CHECK-NEXT: [[MOV_RLC_imm11_pseudo2:%[0-9]+]]:er = MOV_RLC_imm11_pseudo 2
+    ; CHECK-NEXT: [[COPY2:%[0-9]+]]:exo = COPY [[DEF]].sub_512_hi
+    ; CHECK-NEXT: [[COPY3:%[0-9]+]]:exe = COPY [[DEF]].sub_512_lo
+    ; CHECK-NEXT: [[VSHUFFLE_vec_shuffle_x1:%[0-9]+]]:vec512 = VSHUFFLE_vec_shuffle_x [[COPY3]], [[COPY2]], [[MOV_RLC_imm11_pseudo2]]
+    ; CHECK-NEXT: [[VSHUFFLE_vec_shuffle_x2:%[0-9]+]]:mshflxdst = VSHUFFLE_vec_shuffle_x [[VSHUFFLE_vec_shuffle_x1]], [[VSHUFFLE_vec_shuffle_x]], [[MOV_RLC_imm11_pseudo]]
+    ; CHECK-NEXT: [[COPY4:%[0-9]+]]:ewl = COPY [[VSHUFFLE_vec_shuffle_x2]].sub_256_lo
+    ; CHECK-NEXT: PseudoRET implicit $lr, implicit [[COPY4]]
+    %1:vregbank(<32 x s32>) = G_IMPLICIT_DEF
+    %0:vregbank(<32 x s8>) = G_TRUNC %1(<32 x s32>)
+    PseudoRET implicit $lr, implicit %0
+...
 
+---
+name:            v16s16_trunc_v16s64_vec1024
+legalized:       true
+regBankSelected: true
+body:             |
+  bb.1.entry:
+    ; CHECK-LABEL: name: v16s16_trunc_v16s64_vec1024
+    ; CHECK: [[DEF:%[0-9]+]]:vec1024 = IMPLICIT_DEF
+    ; CHECK-NEXT: [[MOV_RLC_imm11_pseudo:%[0-9]+]]:er = MOV_RLC_imm11_pseudo 2
+    ; CHECK-NEXT: [[MOV_RLC_imm11_pseudo1:%[0-9]+]]:er = MOV_RLC_imm11_pseudo 4
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:exo = COPY [[DEF]].sub_512_hi
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:exe = COPY [[DEF]].sub_512_lo
+    ; CHECK-NEXT: [[VSHUFFLE_vec_shuffle_x:%[0-9]+]]:vec512 = VSHUFFLE_vec_shuffle_x [[COPY1]], [[COPY]], [[MOV_RLC_imm11_pseudo1]]
+    ; CHECK-NEXT: [[MOV_RLC_imm11_pseudo2:%[0-9]+]]:er = MOV_RLC_imm11_pseudo 4
+    ; CHECK-NEXT: [[COPY2:%[0-9]+]]:exo = COPY [[DEF]].sub_512_hi
+    ; CHECK-NEXT: [[COPY3:%[0-9]+]]:exe = COPY [[DEF]].sub_512_lo
+    ; CHECK-NEXT: [[VSHUFFLE_vec_shuffle_x1:%[0-9]+]]:vec512 = VSHUFFLE_vec_shuffle_x [[COPY3]], [[COPY2]], [[MOV_RLC_imm11_pseudo2]]
+    ; CHECK-NEXT: [[VSHUFFLE_vec_shuffle_x2:%[0-9]+]]:mshflxdst = VSHUFFLE_vec_shuffle_x [[VSHUFFLE_vec_shuffle_x1]], [[VSHUFFLE_vec_shuffle_x]], [[MOV_RLC_imm11_pseudo]]
+    ; CHECK-NEXT: [[COPY4:%[0-9]+]]:ewl = COPY [[VSHUFFLE_vec_shuffle_x2]].sub_256_lo
+    ; CHECK-NEXT: PseudoRET implicit $lr, implicit [[COPY4]]
+    %1:vregbank(<16 x s64>) = G_IMPLICIT_DEF
+    %0:vregbank(<16 x s16>) = G_TRUNC %1(<16 x s64>)
+    PseudoRET implicit $lr, implicit %0
+...