Skip to content

Commit 16204d5

Browse files
Pierre-vhPierre-vh
committed
[AMDGPU] Canonicalize G_ZEXT of the shift amount in RegBankCombiner
Canonicalize it to a G_AND instead so that ISel patterns can pick it up and ignore it, as the shift instructions only read low bits. G_ZEXT would be lowered to a v/s_and anyway in most cases.
1 parent 0878dd1 commit 16204d5

File tree

3 files changed

+193
-1
lines changed

3 files changed

+193
-1
lines changed

‎llvm/lib/Target/AMDGPU/AMDGPUCombine.td

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,22 @@ def combine_fmul_with_select_to_fldexp : GICombineRule<
134134
[{ return Helper.matchCombineFmulWithSelectToFldexp(*${root}, *${sel}, ${matchinfo}); }]),
135135
(apply [{ Helper.applyBuildFn(*${root}, ${matchinfo}); }])>;
136136

137+
// (shift x, (zext amt)) -> (shift x, (and (anyext amt), mask)
138+
//
139+
// The pattern is longer, but is better for matching during ISel.
140+
class canonicalize_zext_shift_amt<Instruction opc> : GICombineRule<
141+
(defs root:$dst),
142+
(match (G_ZEXT $amt, $amtsrc):$zext,
143+
(opc $dst, $src, $amt):$shift),
144+
(apply [{ applyCanonicalizeZextShiftAmt(*${shift}, *${zext}); }])>;
145+
146+
def canonicalize_zext_lshr : canonicalize_zext_shift_amt<G_LSHR>;
147+
def canonicalize_zext_ashr : canonicalize_zext_shift_amt<G_ASHR>;
148+
def canonicalize_zext_shl : canonicalize_zext_shift_amt<G_SHL>;
149+
150+
def zext_of_shift_amount_combines : GICombineGroup<[
151+
canonicalize_zext_lshr, canonicalize_zext_ashr, canonicalize_zext_shl
152+
]>;
137153

138154
let Predicates = [Has16BitInsts, NotHasMed3_16] in {
139155
// For gfx8, expand f16-fmed3-as-f32 into a min/max f16 sequence. This
@@ -181,5 +197,5 @@ def AMDGPURegBankCombiner : GICombiner<
181197
zext_trunc_fold, int_minmax_to_med3, ptr_add_immed_chain,
182198
fp_minmax_to_clamp, fp_minmax_to_med3, fmed3_intrinsic_to_clamp,
183199
identity_combines, redundant_and, constant_fold_cast_op,
184-
cast_of_cast_combines]> {
200+
cast_of_cast_combines, zext_of_shift_amount_combines]> {
185201
}

‎llvm/lib/Target/AMDGPU/AMDGPURegBankCombiner.cpp

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,8 @@ class AMDGPURegBankCombinerImpl : public Combiner {
8787
void applyMed3(MachineInstr &MI, Med3MatchInfo &MatchInfo) const;
8888
void applyClamp(MachineInstr &MI, Register &Reg) const;
8989

90+
void applyCanonicalizeZextShiftAmt(MachineInstr &MI, MachineInstr &Ext) const;
91+
9092
private:
9193
SIModeRegisterDefaults getMode() const;
9294
bool getIEEE() const;
@@ -362,6 +364,34 @@ void AMDGPURegBankCombinerImpl::applyMed3(MachineInstr &MI,
362364
MI.eraseFromParent();
363365
}
364366

367+
void AMDGPURegBankCombinerImpl::applyCanonicalizeZextShiftAmt(
368+
MachineInstr &MI, MachineInstr &Ext) const {
369+
unsigned ShOpc = MI.getOpcode();
370+
assert(ShOpc == AMDGPU::G_SHL || ShOpc == AMDGPU::G_LSHR ||
371+
ShOpc == AMDGPU::G_ASHR);
372+
assert(Ext.getOpcode() == AMDGPU::G_ZEXT);
373+
374+
Register AmtReg = Ext.getOperand(1).getReg();
375+
Register ShDst = MI.getOperand(0).getReg();
376+
Register ShSrc = MI.getOperand(1).getReg();
377+
378+
LLT ExtAmtTy = MRI.getType(Ext.getOperand(0).getReg());
379+
LLT AmtTy = MRI.getType(AmtReg);
380+
381+
auto &RB = *MRI.getRegBank(AmtReg);
382+
383+
auto NewExt = B.buildAnyExt(ExtAmtTy, AmtReg);
384+
auto Mask = B.buildConstant(
385+
ExtAmtTy, maskTrailingOnes<uint64_t>(AmtTy.getScalarSizeInBits()));
386+
auto And = B.buildAnd(ExtAmtTy, NewExt, Mask);
387+
B.buildInstr(ShOpc, {ShDst}, {ShSrc, And});
388+
389+
MRI.setRegBank(NewExt.getReg(0), RB);
390+
MRI.setRegBank(Mask.getReg(0), RB);
391+
MRI.setRegBank(And.getReg(0), RB);
392+
MI.eraseFromParent();
393+
}
394+
365395
SIModeRegisterDefaults AMDGPURegBankCombinerImpl::getMode() const {
366396
return MF.getInfo<SIMachineFunctionInfo>()->getMode();
367397
}

‎llvm/test/CodeGen/AMDGPU/GlobalISel/combine-shift-amount-zext.mir

Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
1+
# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
2+
# RUN: llc -mtriple=amdgcn -run-pass=amdgpu-regbank-combiner %s -o - | FileCheck %s
3+
4+
---
5+
name: lshr_zext_i16
6+
tracksRegLiveness: true
7+
body: |
8+
bb.0:
9+
liveins: $sgpr0, $sgpr1
10+
11+
; CHECK-LABEL: name: lshr_zext_i16
12+
; CHECK: liveins: $sgpr0, $sgpr1
13+
; CHECK-NEXT: {{ $}}
14+
; CHECK-NEXT: %src:sgpr(s32) = COPY $sgpr0
15+
; CHECK-NEXT: %regamt:sgpr(s32) = COPY $sgpr1
16+
; CHECK-NEXT: [[C:%[0-9]+]]:sgpr(s32) = G_CONSTANT i32 65535
17+
; CHECK-NEXT: [[AND:%[0-9]+]]:sgpr(s32) = G_AND %regamt, [[C]]
18+
; CHECK-NEXT: %res:sgpr(s32) = G_LSHR %src, [[AND]](s32)
19+
; CHECK-NEXT: $sgpr0 = COPY %res(s32)
20+
%src:sgpr(s32) = COPY $sgpr0
21+
%regamt:sgpr(s32) = COPY $sgpr1
22+
%amt:sgpr(s16) = G_TRUNC %regamt
23+
%zextamt:sgpr(s32) = G_ZEXT %amt
24+
%res:sgpr(s32) = G_LSHR %src, %zextamt
25+
$sgpr0 = COPY %res
26+
...
27+
28+
---
29+
name: ashr_zext_i16
30+
tracksRegLiveness: true
31+
body: |
32+
bb.0:
33+
liveins: $sgpr0, $sgpr1
34+
35+
; CHECK-LABEL: name: ashr_zext_i16
36+
; CHECK: liveins: $sgpr0, $sgpr1
37+
; CHECK-NEXT: {{ $}}
38+
; CHECK-NEXT: %src:sgpr(s32) = COPY $sgpr0
39+
; CHECK-NEXT: %regamt:sgpr(s32) = COPY $sgpr1
40+
; CHECK-NEXT: [[C:%[0-9]+]]:sgpr(s32) = G_CONSTANT i32 65535
41+
; CHECK-NEXT: [[AND:%[0-9]+]]:sgpr(s32) = G_AND %regamt, [[C]]
42+
; CHECK-NEXT: %res:sgpr(s32) = G_ASHR %src, [[AND]](s32)
43+
; CHECK-NEXT: $sgpr0 = COPY %res(s32)
44+
%src:sgpr(s32) = COPY $sgpr0
45+
%regamt:sgpr(s32) = COPY $sgpr1
46+
%amt:sgpr(s16) = G_TRUNC %regamt
47+
%zextamt:sgpr(s32) = G_ZEXT %amt
48+
%res:sgpr(s32) = G_ASHR %src, %zextamt
49+
$sgpr0 = COPY %res
50+
...
51+
52+
---
53+
name: shl_zext_i16
54+
tracksRegLiveness: true
55+
body: |
56+
bb.0:
57+
liveins: $sgpr0, $sgpr1
58+
59+
; CHECK-LABEL: name: shl_zext_i16
60+
; CHECK: liveins: $sgpr0, $sgpr1
61+
; CHECK-NEXT: {{ $}}
62+
; CHECK-NEXT: %src:sgpr(s32) = COPY $sgpr0
63+
; CHECK-NEXT: %regamt:sgpr(s32) = COPY $sgpr1
64+
; CHECK-NEXT: [[C:%[0-9]+]]:sgpr(s32) = G_CONSTANT i32 65535
65+
; CHECK-NEXT: [[AND:%[0-9]+]]:sgpr(s32) = G_AND %regamt, [[C]]
66+
; CHECK-NEXT: %res:sgpr(s32) = G_SHL %src, [[AND]](s32)
67+
; CHECK-NEXT: $sgpr0 = COPY %res(s32)
68+
%src:sgpr(s32) = COPY $sgpr0
69+
%regamt:sgpr(s32) = COPY $sgpr1
70+
%amt:sgpr(s16) = G_TRUNC %regamt
71+
%zextamt:sgpr(s32) = G_ZEXT %amt
72+
%res:sgpr(s32) = G_SHL %src, %zextamt
73+
$sgpr0 = COPY %res
74+
...
75+
76+
---
77+
name: lshr_zext_i8
78+
tracksRegLiveness: true
79+
body: |
80+
bb.0:
81+
liveins: $sgpr0, $sgpr1
82+
83+
; CHECK-LABEL: name: lshr_zext_i8
84+
; CHECK: liveins: $sgpr0, $sgpr1
85+
; CHECK-NEXT: {{ $}}
86+
; CHECK-NEXT: %src:sgpr(s32) = COPY $sgpr0
87+
; CHECK-NEXT: %regamt:sgpr(s32) = COPY $sgpr1
88+
; CHECK-NEXT: [[C:%[0-9]+]]:sgpr(s32) = G_CONSTANT i32 255
89+
; CHECK-NEXT: [[AND:%[0-9]+]]:sgpr(s32) = G_AND %regamt, [[C]]
90+
; CHECK-NEXT: %res:sgpr(s32) = G_LSHR %src, [[AND]](s32)
91+
; CHECK-NEXT: $sgpr0 = COPY %res(s32)
92+
%src:sgpr(s32) = COPY $sgpr0
93+
%regamt:sgpr(s32) = COPY $sgpr1
94+
%amt:sgpr(s8) = G_TRUNC %regamt
95+
%zextamt:sgpr(s32) = G_ZEXT %amt
96+
%res:sgpr(s32) = G_LSHR %src, %zextamt
97+
$sgpr0 = COPY %res
98+
...
99+
100+
---
101+
name: ashr_zext_i8
102+
tracksRegLiveness: true
103+
body: |
104+
bb.0:
105+
liveins: $sgpr0, $sgpr1
106+
107+
; CHECK-LABEL: name: ashr_zext_i8
108+
; CHECK: liveins: $sgpr0, $sgpr1
109+
; CHECK-NEXT: {{ $}}
110+
; CHECK-NEXT: %src:sgpr(s32) = COPY $sgpr0
111+
; CHECK-NEXT: %regamt:sgpr(s32) = COPY $sgpr1
112+
; CHECK-NEXT: [[C:%[0-9]+]]:sgpr(s32) = G_CONSTANT i32 255
113+
; CHECK-NEXT: [[AND:%[0-9]+]]:sgpr(s32) = G_AND %regamt, [[C]]
114+
; CHECK-NEXT: %res:sgpr(s32) = G_ASHR %src, [[AND]](s32)
115+
; CHECK-NEXT: $sgpr0 = COPY %res(s32)
116+
%src:sgpr(s32) = COPY $sgpr0
117+
%regamt:sgpr(s32) = COPY $sgpr1
118+
%amt:sgpr(s8) = G_TRUNC %regamt
119+
%zextamt:sgpr(s32) = G_ZEXT %amt
120+
%res:sgpr(s32) = G_ASHR %src, %zextamt
121+
$sgpr0 = COPY %res
122+
...
123+
124+
---
125+
name: shl_zext_i8
126+
tracksRegLiveness: true
127+
body: |
128+
bb.0:
129+
liveins: $sgpr0, $sgpr1
130+
131+
; CHECK-LABEL: name: shl_zext_i8
132+
; CHECK: liveins: $sgpr0, $sgpr1
133+
; CHECK-NEXT: {{ $}}
134+
; CHECK-NEXT: %src:sgpr(s32) = COPY $sgpr0
135+
; CHECK-NEXT: %regamt:sgpr(s32) = COPY $sgpr1
136+
; CHECK-NEXT: [[C:%[0-9]+]]:sgpr(s32) = G_CONSTANT i32 255
137+
; CHECK-NEXT: [[AND:%[0-9]+]]:sgpr(s32) = G_AND %regamt, [[C]]
138+
; CHECK-NEXT: %res:sgpr(s32) = G_SHL %src, [[AND]](s32)
139+
; CHECK-NEXT: $sgpr0 = COPY %res(s32)
140+
%src:sgpr(s32) = COPY $sgpr0
141+
%regamt:sgpr(s32) = COPY $sgpr1
142+
%amt:sgpr(s8) = G_TRUNC %regamt
143+
%zextamt:sgpr(s32) = G_ZEXT %amt
144+
%res:sgpr(s32) = G_SHL %src, %zextamt
145+
$sgpr0 = COPY %res
146+
...

0 commit comments

Comments
 (0)