Skip to content

Commit 5654047

Browse files
ValentijnvdBeekValentijnvdBeek
committed
[AIE2] Implement vshuffle selection
1 parent 0e79ca0 commit 5654047

File tree

4 files changed

+186
-32
lines changed

4 files changed

+186
-32
lines changed

‎llvm/lib/Target/AIE/AIE2InstrPatterns.td‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -597,6 +597,18 @@ def : Pat<(int_aie2_vshuffle VEC512:$s1, VEC512:$s2, eR:$mod),
597597
def : Pat<(int_aie2_vshuffle_bf16 VEC512:$s1, VEC512:$s2, eR:$mod),
598598
(VSHUFFLE VEC512:$s1, VEC512:$s2, eR:$mod)>;
599599

600+
// VSHUFFLE generic opcodes translation
601+
def vshuffle_node : SDNode<"AIE2::G_AIE_VSHUFFLE",
602+
SDTypeProfile<1, 3, [SDTCisVec<1>, SDTCisVec<2>, SDTCisInt<3>]>>;
603+
def : GINodeEquiv<G_AIE_VSHUFFLE, vshuffle_node>;
604+
605+
def : Pat<(v16i32 (vshuffle_node (v16i32 VEC512:$v0), (v16i32 VEC512:$v1), (i32 eR:$mode))),
606+
(VSHUFFLE VEC512:$v0, VEC512:$v1, i32:$mode)>;
607+
def : Pat<(v32i16 (vshuffle_node (v32i16 VEC512:$v0), (v32i16 VEC512:$v1), (i32 eR:$mode))),
608+
(VSHUFFLE VEC512:$v0, VEC512:$v1, i32:$mode)>;
609+
def : Pat<(v64i8 (vshuffle_node (v64i8 VEC512:$v0), (v64i8 VEC512:$v1), (i32 eR:$mode))),
610+
(VSHUFFLE VEC512:$v0, VEC512:$v1, i32:$mode)>;
611+
600612
// VSHIFT Intrinsic (shift/shiftx/shift_bytes)
601613
def : Pat<(int_aie2_vshift_I512_I512 VEC512:$s1, VEC512:$s2, 0x0, eR:$shift),
602614
(VSHIFT VEC512:$s1, VEC512:$s2, eR:$shift)>;

‎llvm/lib/Target/AIE/AIE2PreLegalizerCombiner.cpp‎

Lines changed: 42 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,39 @@ AIE2PreLegalizerCombinerImpl::AIE2PreLegalizerCombinerImpl(
8686
{
8787
}
8888

89+
bool createVShuffle(MachineInstr &MI, const LLT TargetTy, const uint8_t Mode) {
90+
MachineIRBuilder MIB(MI);
91+
MachineRegisterInfo &MRI = *MIB.getMRI();
92+
const Register DstReg = MI.getOperand(0).getReg();
93+
const LLT DstTy = MRI.getType(DstReg);
94+
95+
if (DstTy != TargetTy)
96+
return false;
97+
98+
const Register Src1 = MI.getOperand(1).getReg();
99+
const Register Src2 = MI.getOperand(2).getReg();
100+
const Register ShuffleModeReg =
101+
MRI.createGenericVirtualRegister(LLT::scalar(32));
102+
103+
// This combiner only cares about the lower bits, so we can pad the
104+
// vector to cover the case where two separate vectors are shuffled.
105+
// together
106+
MIB.buildConstant(ShuffleModeReg, Mode);
107+
if (MRI.getType(Src1) == TargetTy) {
108+
MIB.buildInstr(AIE2::G_AIE_VSHUFFLE, {DstReg},
109+
{Src1, Src2, ShuffleModeReg});
110+
} else {
111+
// We reuse the same register since we ignore the high part of the vector
112+
const Register TmpRegister = MRI.createGenericVirtualRegister(TargetTy);
113+
MIB.buildConcatVectors(TmpRegister, {Src1, Src2});
114+
MIB.buildInstr(AIE2::G_AIE_VSHUFFLE, {DstReg},
115+
{TmpRegister, TmpRegister, ShuffleModeReg});
116+
}
117+
118+
MI.eraseFromParent();
119+
return true;
120+
}
121+
89122
CombinerHelper::GeneratorType sectionGenerator(const int32_t From,
90123
const int32_t To,
91124
const int32_t Partitions,
@@ -125,34 +158,15 @@ bool AIE2PreLegalizerCombinerImpl::tryCombineShuffleVector(
125158
const LLT V64S8 = LLT::fixed_vector(64, 8);
126159
CombinerHelper::GeneratorType FourPartitions =
127160
sectionGenerator(0, DstNumElts, 4, 1);
128-
if (Helper.matchCombineShuffleVector(MI, FourPartitions, DstNumElts)) {
129-
if (DstTy != V64S8)
130-
return false;
131-
132-
const Register Src1 = MI.getOperand(1).getReg();
133-
const Register Src2 = MI.getOperand(2).getReg();
134-
const Register ShuffleModeReg =
135-
MRI.createGenericVirtualRegister(LLT::scalar(32));
136-
137-
// This combiner only cares about the lower bits, so we can pad the
138-
// vector to cover the case where two separate vectors are shuffled.
139-
// together
140-
MIB.buildConstant(ShuffleModeReg, 35);
141-
142-
if (SrcTy == V64S8) {
143-
MIB.buildInstr(AIE2::G_AIE_VSHUFFLE, {DstReg},
144-
{Src1, Src2, ShuffleModeReg});
145-
} else {
146-
// We reuse the same register since we ignore the high part of the vector
147-
const Register TmpRegister = MRI.createGenericVirtualRegister(V64S8);
148-
MIB.buildConcatVectors(TmpRegister, {Src1, Src2});
149-
MIB.buildInstr(AIE2::G_AIE_VSHUFFLE, {DstReg},
150-
{TmpRegister, TmpRegister, ShuffleModeReg});
151-
}
152-
153-
MI.eraseFromParent();
154-
return true;
155-
}
161+
if (Helper.matchCombineShuffleVector(MI, FourPartitions, DstNumElts))
162+
return createVShuffle(MI, V64S8, 35);
163+
164+
const LLT V32S16 = LLT::fixed_vector(32, 16);
165+
CombinerHelper::GeneratorType FourPartitionByTwo =
166+
sectionGenerator(0, DstNumElts, 4, 2);
167+
if (Helper.matchCombineShuffleVector(MI, FourPartitionByTwo, DstNumElts))
168+
return createVShuffle(MI, V32S16, 29);
169+
156170
return false;
157171
}
158172

‎llvm/test/CodeGen/AIE/aie2/GlobalISel/inst-select-aie-vshuffle.mir‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
2+
#
3+
# This file is licensed under the Apache License v2.0 with LLVM Exceptions.
4+
# See https://llvm.org/LICENSE.txt for license information.
5+
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
#
7+
# (c) Copyright 2023-2024 Advanced Micro Devices, Inc. or its affiliates
8+
#
9+
# RUN: llc -mtriple aie2 -run-pass=instruction-select %s -verify-machineinstrs -o - | FileCheck %s
10+
11+
---
12+
name: vshuffle_32_m35
13+
legalized: true
14+
regBankSelected: true
15+
tracksRegLiveness: true
16+
stack:
17+
- { id: 0, name: "", size: 128, alignment: 32 }
18+
body: |
19+
bb.0.entry:
20+
liveins: $x2
21+
; CHECK-LABEL: name: vshuffle_32_m35
22+
; CHECK: liveins: $x2
23+
; CHECK-NEXT: {{ $}}
24+
; CHECK-NEXT: [[COPY:%[0-9]+]]:vec512 = COPY $x2
25+
; CHECK-NEXT: [[MOV_RLC_imm10_pseudo:%[0-9]+]]:er = MOV_RLC_imm10_pseudo 29
26+
; CHECK-NEXT: [[VSHUFFLE:%[0-9]+]]:vec512 = VSHUFFLE [[COPY]], [[COPY]], [[MOV_RLC_imm10_pseudo]]
27+
; CHECK-NEXT: $x0 = COPY [[VSHUFFLE]]
28+
; CHECK-NEXT: PseudoRET implicit $lr, implicit $x0
29+
%1:vregbank(<16 x s32>) = COPY $x2
30+
%2:gprregbank(s32) = G_CONSTANT i32 29
31+
%0:vregbank(<16 x s32>) = G_AIE_VSHUFFLE %1:vregbank, %1:vregbank, %2:gprregbank(s32)
32+
$x0 = COPY %0:vregbank(<16 x s32>)
33+
PseudoRET implicit $lr, implicit $x0
34+
...
35+
36+
---
37+
name: vshuffle_16_m35
38+
legalized: true
39+
regBankSelected: true
40+
tracksRegLiveness: true
41+
stack:
42+
- { id: 0, name: "", size: 128, alignment: 32 }
43+
body: |
44+
bb.0.entry:
45+
liveins: $x2
46+
; CHECK-LABEL: name: vshuffle_16_m35
47+
; CHECK: liveins: $x2
48+
; CHECK-NEXT: {{ $}}
49+
; CHECK-NEXT: [[COPY:%[0-9]+]]:vec512 = COPY $x2
50+
; CHECK-NEXT: [[MOV_RLC_imm10_pseudo:%[0-9]+]]:er = MOV_RLC_imm10_pseudo 29
51+
; CHECK-NEXT: [[VSHUFFLE:%[0-9]+]]:vec512 = VSHUFFLE [[COPY]], [[COPY]], [[MOV_RLC_imm10_pseudo]]
52+
; CHECK-NEXT: $x0 = COPY [[VSHUFFLE]]
53+
; CHECK-NEXT: PseudoRET implicit $lr, implicit $x0
54+
%1:vregbank(<32 x s16>) = COPY $x2
55+
%2:gprregbank(s32) = G_CONSTANT i32 29
56+
%0:vregbank(<32 x s16>) = G_AIE_VSHUFFLE %1:vregbank, %1:vregbank, %2:gprregbank(s32)
57+
$x0 = COPY %0:vregbank(<32 x s16>)
58+
PseudoRET implicit $lr, implicit $x0
59+
...
60+
61+
---
62+
name: vshuffle_8_m35
63+
legalized: true
64+
regBankSelected: true
65+
tracksRegLiveness: true
66+
stack:
67+
- { id: 0, name: "", size: 128, alignment: 32 }
68+
body: |
69+
bb.0.entry:
70+
liveins: $x2
71+
; CHECK-LABEL: name: vshuffle_8_m35
72+
; CHECK: liveins: $x2
73+
; CHECK-NEXT: {{ $}}
74+
; CHECK-NEXT: [[COPY:%[0-9]+]]:vec512 = COPY $x2
75+
; CHECK-NEXT: [[MOV_RLC_imm10_pseudo:%[0-9]+]]:er = MOV_RLC_imm10_pseudo 29
76+
; CHECK-NEXT: [[VSHUFFLE:%[0-9]+]]:vec512 = VSHUFFLE [[COPY]], [[COPY]], [[MOV_RLC_imm10_pseudo]]
77+
; CHECK-NEXT: $x0 = COPY [[VSHUFFLE]]
78+
; CHECK-NEXT: PseudoRET implicit $lr, implicit $x0
79+
%1:vregbank(<64 x s8>) = COPY $x2
80+
%2:gprregbank(s32) = G_CONSTANT i32 29
81+
%0:vregbank(<64 x s8>) = G_AIE_VSHUFFLE %1:vregbank, %1:vregbank, %2:gprregbank(s32)
82+
$x0 = COPY %0:vregbank(<64 x s8>)
83+
PseudoRET implicit $lr, implicit $x0

‎llvm/test/CodeGen/AIE/aie2/GlobalISel/prelegalizercombiner-shufflevector.mir‎

Lines changed: 49 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -645,12 +645,12 @@ body: |
645645
...
646646

647647
---
648-
name: shuffle_vector_32
648+
name: shuffle_vector_32_4x4
649649
legalized: false
650650
body: |
651651
bb.1.entry:
652652
liveins: $x0, $x1
653-
; CHECK-LABEL: name: shuffle_vector_32
653+
; CHECK-LABEL: name: shuffle_vector_32_4x4
654654
; CHECK: liveins: $x0, $x1
655655
; CHECK-NEXT: {{ $}}
656656
; CHECK-NEXT: [[COPY:%[0-9]+]]:_(<16 x s32>) = COPY $x0
@@ -666,12 +666,12 @@ body: |
666666
...
667667

668668
---
669-
name: shuffle_vector_16
669+
name: shuffle_vector_16_4x4
670670
legalized: false
671671
body: |
672672
bb.1.entry:
673673
liveins: $x0, $x1
674-
; CHECK-LABEL: name: shuffle_vector_16
674+
; CHECK-LABEL: name: shuffle_vector_16_4x4
675675
; CHECK: liveins: $x0, $x1
676676
; CHECK-NEXT: {{ $}}
677677
; CHECK-NEXT: [[COPY:%[0-9]+]]:_(<32 x s16>) = COPY $x0
@@ -730,3 +730,48 @@ body: |
730730
%0:_(<64 x s8>) = G_SHUFFLE_VECTOR %1:_(<64 x s8>), %2:_, shufflemask(0, 16, 32, 48, 1, 17, 33, 49, 2, 18, 34, 50, 3, 19, 35, 51, 4, 20, 36, 52, 5, 21, 37, 53, 6, 22, 38, 54, 7, 23, 39, 55, 8, 24, 40, 56, 9, 25, 41, 57, 10, 26, 42, 58, 11, 27, 43, 59, 12, 28, 44, 60, 13, 29, 45, 61, 14, 30, 46, 62, 15, 31, 47, 63)
731731
$x2 = COPY %0:_(<64 x s8>)
732732
PseudoRET implicit $lr, implicit $x2
733+
...
734+
735+
---
736+
name: shuffle_vector_1024_4x8
737+
legalized: false
738+
body: |
739+
bb.1.entry:
740+
liveins: $x0, $x1
741+
; CHECK-LABEL: name: shuffle_vector_1024_4x8
742+
; CHECK: liveins: $x0, $x1
743+
; CHECK-NEXT: {{ $}}
744+
; CHECK-NEXT: [[COPY:%[0-9]+]]:_(<32 x s16>) = COPY $x0
745+
; CHECK-NEXT: [[COPY1:%[0-9]+]]:_(<32 x s16>) = COPY $x1
746+
; CHECK-NEXT: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 29
747+
; CHECK-NEXT: [[AIE_VSHUFFLE:%[0-9]+]]:_(<32 x s16>) = G_AIE_VSHUFFLE [[COPY]], [[COPY1]], [[C]](s32)
748+
; CHECK-NEXT: $x2 = COPY [[AIE_VSHUFFLE]](<32 x s16>)
749+
; CHECK-NEXT: PseudoRET implicit $lr, implicit $x2
750+
%1:_(<32 x s16>) = COPY $x0
751+
%2:_(<32 x s16>) = COPY $x1
752+
%0:_(<32 x s16>) = G_SHUFFLE_VECTOR %1:_(<32 x s16>), %2:_, shufflemask(0, 1, 8, 9, 16, 17, 24, 25, 2, 3, 10, 11, 18, 19, 26, 27, 4, 5, 12, 13, 20, 21, 28, 29, 6, 7, 14, 15, 22, 23, 30, 31)
753+
$x2 = COPY %0:_(<32 x s16>)
754+
PseudoRET implicit $lr, implicit $x2
755+
...
756+
757+
---
758+
name: shuffle_vector_512_4x8
759+
legalized: false
760+
body: |
761+
bb.1.entry:
762+
liveins: $wl0, $wl1
763+
; CHECK-LABEL: name: shuffle_vector_512_4x8
764+
; CHECK: liveins: $wl0, $wl1
765+
; CHECK-NEXT: {{ $}}
766+
; CHECK-NEXT: [[COPY:%[0-9]+]]:_(<16 x s16>) = COPY $wl0
767+
; CHECK-NEXT: [[COPY1:%[0-9]+]]:_(<16 x s16>) = COPY $wl1
768+
; CHECK-NEXT: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 29
769+
; CHECK-NEXT: [[CONCAT_VECTORS:%[0-9]+]]:_(<32 x s16>) = G_CONCAT_VECTORS [[COPY]](<16 x s16>), [[COPY1]](<16 x s16>)
770+
; CHECK-NEXT: [[AIE_VSHUFFLE:%[0-9]+]]:_(<32 x s16>) = G_AIE_VSHUFFLE [[CONCAT_VECTORS]], [[CONCAT_VECTORS]], [[C]](s32)
771+
; CHECK-NEXT: $x2 = COPY [[AIE_VSHUFFLE]](<32 x s16>)
772+
; CHECK-NEXT: PseudoRET implicit $lr, implicit $x2
773+
%1:_(<16 x s16>) = COPY $wl0
774+
%2:_(<16 x s16>) = COPY $wl1
775+
%0:_(<32 x s16>) = G_SHUFFLE_VECTOR %1:_(<16 x s16>), %2:_, shufflemask(0, 1, 8, 9, 16, 17, 24, 25, 2, 3, 10, 11, 18, 19, 26, 27, 4, 5, 12, 13, 20, 21, 28, 29, 6, 7, 14, 15, 22, 23, 30, 31)
776+
$x2 = COPY %0:_(<32 x s16>)
777+
PseudoRET implicit $lr, implicit $x2

0 commit comments

Comments
 (0)