From cd13ed165d31421c043cc89ab561cdbb7e1c56ea Mon Sep 17 00:00:00 2001
From: Kateryna Muts <kateryna.muts@amd.com>
Date: Mon, 27 Jan 2025 08:34:22 +0000
Subject: [PATCH] [AIE2P] Instruction selection for G_AIE_VECTOR_SUBVECTOR

---
 .../Target/AIE/aie2p/AIE2PInstrPatterns.td    |  30 +++
 .../lib/Target/AIE/aie2p/AIE2PRegisterInfo.td |   1 +
 .../inst-select-extract-subvector.mir         | 254 ++++++++++++++++++
 3 files changed, 285 insertions(+)
 create mode 100644 llvm/test/CodeGen/AIE/aie2p/GlobalIsel/inst-select-extract-subvector.mir

diff --git a/llvm/lib/Target/AIE/aie2p/AIE2PInstrPatterns.td b/llvm/lib/Target/AIE/aie2p/AIE2PInstrPatterns.td
index 8b35540e418c..a3aecfc37657 100644
--- a/llvm/lib/Target/AIE/aie2p/AIE2PInstrPatterns.td
+++ b/llvm/lib/Target/AIE/aie2p/AIE2PInstrPatterns.td
@@ -801,6 +801,36 @@ def : Pat<(int_aie2p_vinsert64_bf512 VEC512:$s1, eR29:$idx, eL:$s0),
 def : Pat<(int_aie2p_vinsert32_accfloat ACC512:$s1, eR29:$idx, eR:$s0),
            (COPY_TO_REGCLASS (VINSERT_32_mR29_insert (COPY_TO_REGCLASS ACC512:$s1, VEC512), eR29:$idx, eR:$s0), ACC512)>;
 
+// VEXTRACT
+def extract_subvec_node : SDNode<"AIE2P::G_AIE_EXTRACT_SUBVECTOR", SDTypeProfile<1, 2, []>>;
+def : GINodeEquiv<G_AIE_EXTRACT_SUBVECTOR , extract_subvec_node>;
+// G_AIE_EXTRACT_SUBVECTOR extracts only 32-bit or 64-bit subvector.
+def : Pat<(v4i8(extract_subvec_node(v64i8 VEC512:$s0), (i32 c6u:$idx))), 
+          (VEXTRACT_32_vec_extract_imm_vaddSign1 VEC512:$s0, c6u:$idx)>;
+def : Pat<(v8i8(extract_subvec_node(v64i8 VEC512:$s0), (i32 c6u:$idx))), 
+          (VEXTRACT_64_vec_extract_imm_vaddSign1 VEC512:$s0, c6u:$idx)>;
+def : Pat<(v2i16(extract_subvec_node(v32i16 VEC512:$s0), (i32 c6u:$idx))), 
+          (VEXTRACT_32_vec_extract_imm_vaddSign1 VEC512:$s0, c6u:$idx)>;
+def : Pat<(v4i16(extract_subvec_node(v32i16 VEC512:$s0), (i32 c6u:$idx))), 
+          (VEXTRACT_64_vec_extract_imm_vaddSign1 VEC512:$s0, c6u:$idx)>;
+def : Pat<(v2i32(extract_subvec_node(v16i32 VEC512:$s0), (i32 c6u:$idx))), 
+          (VEXTRACT_64_vec_extract_imm_vaddSign1 VEC512:$s0, c6u:$idx)>;
+def : Pat<(v2i32(extract_subvec_node(v8i64 VEC512:$s0), (i32 c6u:$idx))), 
+          (VEXTRACT_64_vec_extract_imm_vaddSign1 VEC512:$s0, c6u:$idx)>;
+
+def : Pat<(v4i8(extract_subvec_node(v64i8 VEC512:$s0), (i32 eRS4:$idx))), 
+          (VEXTRACT_32_vec_extract_r_vaddSign1 VEC512:$s0, eRS4:$idx)>;
+def : Pat<(v8i8(extract_subvec_node(v64i8 VEC512:$s0), (i32 eRS4:$idx))), 
+          (VEXTRACT_64_vec_extract_r_vaddSign1 VEC512:$s0, eRS4:$idx)>;
+def : Pat<(v2i16(extract_subvec_node(v32i16 VEC512:$s0), (i32 eRS4:$idx))), 
+          (VEXTRACT_32_vec_extract_r_vaddSign1 VEC512:$s0, eRS4:$idx)>;
+def : Pat<(v4i16(extract_subvec_node(v32i16 VEC512:$s0), (i32 eRS4:$idx))), 
+          (VEXTRACT_64_vec_extract_r_vaddSign1 VEC512:$s0, eRS4:$idx)>;
+def : Pat<(v2i32(extract_subvec_node(v16i32 VEC512:$s0), (i32 eRS4:$idx))), 
+          (VEXTRACT_64_vec_extract_r_vaddSign1 VEC512:$s0, eRS4:$idx)>;
+def : Pat<(v2i32(extract_subvec_node(v8i64 VEC512:$s0), (i32 eRS4:$idx))), 
+          (VEXTRACT_64_vec_extract_r_vaddSign1 VEC512:$s0, eRS4:$idx)>;
+
 // VBCST
 def : Pat<(int_aie2p_vbroadcast32_bf512 eR:$s0), (VBCST_32 eR:$s0)>;
 def : Pat<(int_aie2p_vbroadcast64_bf512 eL:$s0), (VBCST_64 eL:$s0)>;
diff --git a/llvm/lib/Target/AIE/aie2p/AIE2PRegisterInfo.td b/llvm/lib/Target/AIE/aie2p/AIE2PRegisterInfo.td
index 2c776a09bb0a..48b9cee1c8a9 100644
--- a/llvm/lib/Target/AIE/aie2p/AIE2PRegisterInfo.td
+++ b/llvm/lib/Target/AIE/aie2p/AIE2PRegisterInfo.td
@@ -158,6 +158,7 @@ class AIE2PRegisterClass <int size, int align, list<ValueType> regTypes, dag reg
   def eR16 : AIE2PScalarRegisterClass<(add r16)>;
   def mR16_vcompare : AIE2PScalarRegisterClass<(add eR16)>;
 
+  def eRS4 : AIE2PScalarRegisterClass<(add r16, r17, r18, r19)>;
   def eRS16
       : AIE2PScalarRegisterClass<(add r16, r17, r18, r19, r20, r21, r22,
                                  r23, r24, r25, r26, r27, r28, r29, r30,
diff --git a/llvm/test/CodeGen/AIE/aie2p/GlobalIsel/inst-select-extract-subvector.mir b/llvm/test/CodeGen/AIE/aie2p/GlobalIsel/inst-select-extract-subvector.mir
new file mode 100644
index 000000000000..11b0a38d59a9
--- /dev/null
+++ b/llvm/test/CodeGen/AIE/aie2p/GlobalIsel/inst-select-extract-subvector.mir
@@ -0,0 +1,254 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+#
+# This file is licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+#
+# (c) Copyright 2025 Advanced Micro Devices, Inc. or its affiliates
+# RUN: llc -mtriple aie2p -run-pass=instruction-select %s -verify-machineinstrs -o - | FileCheck %s
+
+---
+name:            extract_subvector_imm_4xs8Dst
+alignment:       16
+legalized:       true
+regBankSelected: true
+body:             |
+  bb.1.entry:
+    liveins: $x2
+    ; CHECK-LABEL: name: extract_subvector_imm_4xs8Dst
+    ; CHECK: liveins: $x2
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:vec512 = COPY $x2
+    ; CHECK-NEXT: [[VEXTRACT_32_vec_extract_imm_vaddSign1_:%[0-9]+]]:er = VEXTRACT_32_vec_extract_imm_vaddSign1 [[COPY]], 1, implicit $vaddsign1
+    ; CHECK-NEXT: PseudoRET implicit $lr, implicit [[VEXTRACT_32_vec_extract_imm_vaddSign1_]]
+    %1:vregbank(<64 x s8>) = COPY $x2
+    %2:gprregbank(s32) = G_CONSTANT i32 1
+    %0:gprregbank(<4 x s8>) = G_AIE_EXTRACT_SUBVECTOR %1:vregbank(<64 x s8>), %2:gprregbank(s32)
+    PseudoRET implicit $lr, implicit %0
+...
+
+---
+name:            extract_subvector_imm_8xs8Dst
+alignment:       16
+legalized:       true
+regBankSelected: true
+body:             |
+  bb.1.entry:
+    liveins: $x2
+    ; CHECK-LABEL: name: extract_subvector_imm_8xs8Dst
+    ; CHECK: liveins: $x2
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:vec512 = COPY $x2
+    ; CHECK-NEXT: [[VEXTRACT_64_vec_extract_imm_vaddSign1_:%[0-9]+]]:el = VEXTRACT_64_vec_extract_imm_vaddSign1 [[COPY]], 1, implicit $vaddsign1
+    ; CHECK-NEXT: PseudoRET implicit $lr, implicit [[VEXTRACT_64_vec_extract_imm_vaddSign1_]]
+    %1:vregbank(<64 x s8>) = COPY $x2
+    %2:gprregbank(s32) = G_CONSTANT i32 1
+    %0:gprregbank(<8 x s8>) = G_AIE_EXTRACT_SUBVECTOR %1:vregbank(<64 x s8>), %2:gprregbank(s32)
+    PseudoRET implicit $lr, implicit %0
+...
+
+---
+name:            extract_subvector_imm_2xs16Dst
+alignment:       16
+legalized:       true
+regBankSelected: true
+body:             |
+  bb.1.entry:
+    liveins: $x2
+    ; CHECK-LABEL: name: extract_subvector_imm_2xs16Dst
+    ; CHECK: liveins: $x2
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:vec512 = COPY $x2
+    ; CHECK-NEXT: [[VEXTRACT_32_vec_extract_imm_vaddSign1_:%[0-9]+]]:er = VEXTRACT_32_vec_extract_imm_vaddSign1 [[COPY]], 1, implicit $vaddsign1
+    ; CHECK-NEXT: PseudoRET implicit $lr, implicit [[VEXTRACT_32_vec_extract_imm_vaddSign1_]]
+    %1:vregbank(<32 x s16>) = COPY $x2
+    %2:gprregbank(s32) = G_CONSTANT i32 1
+    %0:gprregbank(<2 x s16>) = G_AIE_EXTRACT_SUBVECTOR %1:vregbank(<32 x s16>), %2:gprregbank(s32)
+    PseudoRET implicit $lr, implicit %0
+...
+
+---
+name:            extract_subvector_imm_4xs16Dst
+alignment:       16
+legalized:       true
+regBankSelected: true
+body:             |
+  bb.1.entry:
+    liveins: $x2
+    ; CHECK-LABEL: name: extract_subvector_imm_4xs16Dst
+    ; CHECK: liveins: $x2
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:vec512 = COPY $x2
+    ; CHECK-NEXT: [[VEXTRACT_64_vec_extract_imm_vaddSign1_:%[0-9]+]]:el = VEXTRACT_64_vec_extract_imm_vaddSign1 [[COPY]], 1, implicit $vaddsign1
+    ; CHECK-NEXT: PseudoRET implicit $lr, implicit [[VEXTRACT_64_vec_extract_imm_vaddSign1_]]
+    %1:vregbank(<32 x s16>) = COPY $x2
+    %2:gprregbank(s32) = G_CONSTANT i32 1
+    %0:gprregbank(<4 x s16>) = G_AIE_EXTRACT_SUBVECTOR %1:vregbank(<32 x s16>), %2:gprregbank(s32)
+    PseudoRET implicit $lr, implicit %0
+...
+
+---
+name:            extract_subvector_imm_2xs32Dst_16x32Src
+alignment:       16
+legalized:       true
+regBankSelected: true
+body:             |
+  bb.1.entry:
+    liveins: $x2
+    ; CHECK-LABEL: name: extract_subvector_imm_2xs32Dst_16x32Src
+    ; CHECK: liveins: $x2
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:vec512 = COPY $x2
+    ; CHECK-NEXT: [[VEXTRACT_64_vec_extract_imm_vaddSign1_:%[0-9]+]]:el = VEXTRACT_64_vec_extract_imm_vaddSign1 [[COPY]], 1, implicit $vaddsign1
+    ; CHECK-NEXT: PseudoRET implicit $lr, implicit [[VEXTRACT_64_vec_extract_imm_vaddSign1_]]
+    %1:vregbank(<16 x s32>) = COPY $x2
+    %2:gprregbank(s32) = G_CONSTANT i32 1
+    %0:gprregbank(<2 x s32>) = G_AIE_EXTRACT_SUBVECTOR %1:vregbank(<16 x s32>), %2:gprregbank(s32)
+    PseudoRET implicit $lr, implicit %0
+...
+
+---
+name:            extract_subvector_imm_2xs32Dst_8x64Src
+alignment:       16
+legalized:       true
+regBankSelected: true
+body:             |
+  bb.1.entry:
+    liveins: $x2
+    ; CHECK-LABEL: name: extract_subvector_imm_2xs32Dst_8x64Src
+    ; CHECK: liveins: $x2
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:vec512 = COPY $x2
+    ; CHECK-NEXT: [[VEXTRACT_64_vec_extract_imm_vaddSign1_:%[0-9]+]]:el = VEXTRACT_64_vec_extract_imm_vaddSign1 [[COPY]], 1, implicit $vaddsign1
+    ; CHECK-NEXT: PseudoRET implicit $lr, implicit [[VEXTRACT_64_vec_extract_imm_vaddSign1_]]
+    %1:vregbank(<8 x s64>) = COPY $x2
+    %2:gprregbank(s32) = G_CONSTANT i32 1
+    %0:gprregbank(<2 x s32>) = G_AIE_EXTRACT_SUBVECTOR %1:vregbank(<8 x s64>), %2:gprregbank(s32)
+    PseudoRET implicit $lr, implicit %0
+...
+
+---
+name:            extract_subvector_reg_4xs8Dst
+alignment:       16
+legalized:       true
+regBankSelected: true
+body:             |
+  bb.1.entry:
+    liveins: $r0, $r1, $x2, $x4
+    ; CHECK-LABEL: name: extract_subvector_reg_4xs8Dst
+    ; CHECK: liveins: $r0, $r1, $x2, $x4
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:vec512 = COPY $x2
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:er = COPY $r0
+    ; CHECK-NEXT: [[VEXTRACT_32_vec_extract_r_vaddSign1_:%[0-9]+]]:er = VEXTRACT_32_vec_extract_r_vaddSign1 [[COPY]], [[COPY1]], implicit $vaddsign1
+    ; CHECK-NEXT: PseudoRET implicit $lr, implicit [[VEXTRACT_32_vec_extract_r_vaddSign1_]]
+    %1:vregbank(<64 x s8>) = COPY $x2
+    %2:gprregbank(s32) = COPY $r0
+    %0:gprregbank(<4 x s8>) = G_AIE_EXTRACT_SUBVECTOR %1:vregbank(<64 x s8>), %2:gprregbank(s32)
+    PseudoRET implicit $lr, implicit %0
+...
+
+---
+name:            extract_subvector_reg_8xs8Dst
+alignment:       16
+legalized:       true
+regBankSelected: true
+body:             |
+  bb.1.entry:
+    liveins: $r0, $r1, $x2, $x4
+    ; CHECK-LABEL: name: extract_subvector_reg_8xs8Dst
+    ; CHECK: liveins: $r0, $r1, $x2, $x4
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:vec512 = COPY $x2
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:er = COPY $r0
+    ; CHECK-NEXT: [[VEXTRACT_64_vec_extract_r_vaddSign1_:%[0-9]+]]:el = VEXTRACT_64_vec_extract_r_vaddSign1 [[COPY]], [[COPY1]], implicit $vaddsign1
+    ; CHECK-NEXT: PseudoRET implicit $lr, implicit [[VEXTRACT_64_vec_extract_r_vaddSign1_]]
+    %1:vregbank(<64 x s8>) = COPY $x2
+    %2:gprregbank(s32) = COPY $r0
+    %0:gprregbank(<8 x s8>) = G_AIE_EXTRACT_SUBVECTOR %1:vregbank(<64 x s8>), %2:gprregbank(s32)
+    PseudoRET implicit $lr, implicit %0
+...
+
+---
+name:            extract_subvector_reg_2xs16Dst
+alignment:       16
+legalized:       true
+regBankSelected: true
+body:             |
+  bb.1.entry:
+    liveins: $r0, $r1, $x2, $x4
+    ; CHECK-LABEL: name: extract_subvector_reg_2xs16Dst
+    ; CHECK: liveins: $r0, $r1, $x2, $x4
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:vec512 = COPY $x2
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:er = COPY $r0
+    ; CHECK-NEXT: [[VEXTRACT_32_vec_extract_r_vaddSign1_:%[0-9]+]]:er = VEXTRACT_32_vec_extract_r_vaddSign1 [[COPY]], [[COPY1]], implicit $vaddsign1
+    ; CHECK-NEXT: PseudoRET implicit $lr, implicit [[VEXTRACT_32_vec_extract_r_vaddSign1_]]
+    %1:vregbank(<32 x s16>) = COPY $x2
+    %2:gprregbank(s32) = COPY $r0
+    %0:gprregbank(<2 x s16>) = G_AIE_EXTRACT_SUBVECTOR %1:vregbank(<32 x s16>), %2:gprregbank(s32)
+    PseudoRET implicit $lr, implicit %0
+...
+
+---
+name:            extract_subvector_reg_4xs16Dst
+alignment:       16
+legalized:       true
+regBankSelected: true
+body:             |
+  bb.1.entry:
+    liveins: $r0, $r1, $x2, $x4
+    ; CHECK-LABEL: name: extract_subvector_reg_4xs16Dst
+    ; CHECK: liveins: $r0, $r1, $x2, $x4
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:vec512 = COPY $x2
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:er = COPY $r0
+    ; CHECK-NEXT: [[VEXTRACT_64_vec_extract_r_vaddSign1_:%[0-9]+]]:el = VEXTRACT_64_vec_extract_r_vaddSign1 [[COPY]], [[COPY1]], implicit $vaddsign1
+    ; CHECK-NEXT: PseudoRET implicit $lr, implicit [[VEXTRACT_64_vec_extract_r_vaddSign1_]]
+    %1:vregbank(<32 x s16>) = COPY $x2
+    %2:gprregbank(s32) = COPY $r0
+    %0:gprregbank(<4 x s16>) = G_AIE_EXTRACT_SUBVECTOR %1:vregbank(<32 x s16>), %2:gprregbank(s32)
+    PseudoRET implicit $lr, implicit %0
+...
+
+---
+name:            extract_subvector_reg_2xs32Dst_16x32Src
+alignment:       16
+legalized:       true
+regBankSelected: true
+body:             |
+  bb.1.entry:
+    liveins: $r0, $r1, $x2, $x4
+    ; CHECK-LABEL: name: extract_subvector_reg_2xs32Dst_16x32Src
+    ; CHECK: liveins: $r0, $r1, $x2, $x4
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:vec512 = COPY $x2
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:er = COPY $r0
+    ; CHECK-NEXT: [[VEXTRACT_64_vec_extract_r_vaddSign1_:%[0-9]+]]:el = VEXTRACT_64_vec_extract_r_vaddSign1 [[COPY]], [[COPY1]], implicit $vaddsign1
+    ; CHECK-NEXT: PseudoRET implicit $lr, implicit [[VEXTRACT_64_vec_extract_r_vaddSign1_]]
+    %1:vregbank(<16 x s32>) = COPY $x2
+    %2:gprregbank(s32) = COPY $r0
+    %0:gprregbank(<2 x s32>) = G_AIE_EXTRACT_SUBVECTOR %1:vregbank(<16 x s32>), %2:gprregbank(s32)
+    PseudoRET implicit $lr, implicit %0
+...
+
+---
+name:            extract_subvector_reg_2xs32Dst_8x64Src
+alignment:       16
+legalized:       true
+regBankSelected: true
+body:             |
+  bb.1.entry:
+    liveins: $r0, $r1, $x2, $x4
+    ; CHECK-LABEL: name: extract_subvector_reg_2xs32Dst_8x64Src
+    ; CHECK: liveins: $r0, $r1, $x2, $x4
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:vec512 = COPY $x2
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:er = COPY $r0
+    ; CHECK-NEXT: [[VEXTRACT_64_vec_extract_r_vaddSign1_:%[0-9]+]]:el = VEXTRACT_64_vec_extract_r_vaddSign1 [[COPY]], [[COPY1]], implicit $vaddsign1
+    ; CHECK-NEXT: PseudoRET implicit $lr, implicit [[VEXTRACT_64_vec_extract_r_vaddSign1_]]
+    %1:vregbank(<8 x s64>) = COPY $x2
+    %2:gprregbank(s32) = COPY $r0
+    %0:gprregbank(<2 x s32>) = G_AIE_EXTRACT_SUBVECTOR %1:vregbank(<8 x s64>), %2:gprregbank(s32)
+    PseudoRET implicit $lr, implicit %0
+...