Skip to content

Commit ceb2b9c

Browse files
paulwalker-armpaulwalker-arm
authored
[LLVM][DAGCombiner] fold (shl (X * vscale(C0)), C1) -> (X * vscale(C0 << C1)). (#150651)
1 parent 77363fb commit ceb2b9c

File tree

3 files changed

+85
-41
lines changed

3 files changed

+85
-41
lines changed

‎llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10622,6 +10622,19 @@ SDValue DAGCombiner::visitSHL(SDNode *N) {
1062210622
return DAG.getVScale(DL, VT, C0 << C1);
1062310623
}
1062410624

10625+
SDValue X;
10626+
APInt VS0;
10627+
10628+
// fold (shl (X * vscale(VS0)), C1) -> (X * vscale(VS0 << C1))
10629+
if (N1C && sd_match(N0, m_Mul(m_Value(X), m_VScale(m_ConstInt(VS0))))) {
10630+
SDNodeFlags Flags;
10631+
Flags.setNoUnsignedWrap(N->getFlags().hasNoUnsignedWrap() &&
10632+
N0->getFlags().hasNoUnsignedWrap());
10633+
10634+
SDValue VScale = DAG.getVScale(DL, VT, VS0 << N1C->getAPIntValue());
10635+
return DAG.getNode(ISD::MUL, DL, VT, X, VScale, Flags);
10636+
}
10637+
1062510638
// Fold (shl step_vector(C0), C1) to (step_vector(C0 << C1)).
1062610639
APInt ShlVal;
1062710640
if (N0.getOpcode() == ISD::STEP_VECTOR &&

‎llvm/test/CodeGen/AArch64/sve-vscale-combine.ll

Lines changed: 68 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -1,24 +1,25 @@
1-
; RUN: llc -mtriple=aarch64--linux-gnu -mattr=+sve --asm-verbose=false < %s |FileCheck %s
1+
; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
2+
; RUN: llc -mattr=+sve < %s | FileCheck %s
23

3-
declare i32 @llvm.vscale.i32()
4-
declare i64 @llvm.vscale.i64()
4+
target triple = "aarch64-unknown-linux-gnu"
55

66
; Fold (add (vscale * C0), (vscale * C1)) to (vscale * (C0 + C1)).
77
define i64 @combine_add_vscale_i64() nounwind {
88
; CHECK-LABEL: combine_add_vscale_i64:
9-
; CHECK-NOT: add
10-
; CHECK-NEXT: cntd x0
11-
; CHECK-NEXT: ret
9+
; CHECK: // %bb.0:
10+
; CHECK-NEXT: cntd x0
11+
; CHECK-NEXT: ret
1212
%vscale = call i64 @llvm.vscale.i64()
1313
%add = add i64 %vscale, %vscale
1414
ret i64 %add
1515
}
1616

1717
define i32 @combine_add_vscale_i32() nounwind {
1818
; CHECK-LABEL: combine_add_vscale_i32:
19-
; CHECK-NOT: add
20-
; CHECK-NEXT: cntd x0
21-
; CHECK-NEXT: ret
19+
; CHECK: // %bb.0:
20+
; CHECK-NEXT: cntd x0
21+
; CHECK-NEXT: // kill: def $w0 killed $w0 killed $x0
22+
; CHECK-NEXT: ret
2223
%vscale = call i32 @llvm.vscale.i32()
2324
%add = add i32 %vscale, %vscale
2425
ret i32 %add
@@ -28,19 +29,20 @@ define i32 @combine_add_vscale_i32() nounwind {
2829
; In this test, C0 = 1, C1 = 32.
2930
define i64 @combine_mul_vscale_i64() nounwind {
3031
; CHECK-LABEL: combine_mul_vscale_i64:
31-
; CHECK-NOT: mul
32-
; CHECK-NEXT: rdvl x0, #2
33-
; CHECK-NEXT: ret
32+
; CHECK: // %bb.0:
33+
; CHECK-NEXT: rdvl x0, #2
34+
; CHECK-NEXT: ret
3435
%vscale = call i64 @llvm.vscale.i64()
3536
%mul = mul i64 %vscale, 32
3637
ret i64 %mul
3738
}
3839

3940
define i32 @combine_mul_vscale_i32() nounwind {
4041
; CHECK-LABEL: combine_mul_vscale_i32:
41-
; CHECK-NOT: mul
42-
; CHECK-NEXT: rdvl x0, #3
43-
; CHECK-NEXT: ret
42+
; CHECK: // %bb.0:
43+
; CHECK-NEXT: rdvl x0, #3
44+
; CHECK-NEXT: // kill: def $w0 killed $w0 killed $x0
45+
; CHECK-NEXT: ret
4446
%vscale = call i32 @llvm.vscale.i32()
4547
%mul = mul i32 %vscale, 48
4648
ret i32 %mul
@@ -49,23 +51,23 @@ define i32 @combine_mul_vscale_i32() nounwind {
4951
; Canonicalize (sub X, (vscale * C)) to (add X, (vscale * -C))
5052
define i64 @combine_sub_vscale_i64(i64 %in) nounwind {
5153
; CHECK-LABEL: combine_sub_vscale_i64:
52-
; CHECK-NOT: sub
53-
; CHECK-NEXT: rdvl x8, #-1
54-
; CHECK-NEXT: asr x8, x8, #4
55-
; CHECK-NEXT: add x0, x0, x8
56-
; CHECK-NEXT: ret
54+
; CHECK: // %bb.0:
55+
; CHECK-NEXT: rdvl x8, #-1
56+
; CHECK-NEXT: asr x8, x8, #4
57+
; CHECK-NEXT: add x0, x0, x8
58+
; CHECK-NEXT: ret
5759
%vscale = call i64 @llvm.vscale.i64()
5860
%sub = sub i64 %in, %vscale
5961
ret i64 %sub
6062
}
6163

6264
define i32 @combine_sub_vscale_i32(i32 %in) nounwind {
6365
; CHECK-LABEL: combine_sub_vscale_i32:
64-
; CHECK-NOT: sub
65-
; CHECK-NEXT: rdvl x8, #-1
66-
; CHECK-NEXT: asr x8, x8, #4
67-
; CHECK-NEXT: add w0, w0, w8
68-
; CHECK-NEXT: ret
66+
; CHECK: // %bb.0:
67+
; CHECK-NEXT: rdvl x8, #-1
68+
; CHECK-NEXT: asr x8, x8, #4
69+
; CHECK-NEXT: add w0, w0, w8
70+
; CHECK-NEXT: ret
6971
%vscale = call i32 @llvm.vscale.i32()
7072
%sub = sub i32 %in, %vscale
7173
ret i32 %sub
@@ -75,12 +77,13 @@ define i32 @combine_sub_vscale_i32(i32 %in) nounwind {
7577
; (sub X, (vscale * C)) to (add X, (vscale * -C))
7678
define i64 @multiple_uses_sub_vscale_i64(i64 %x, i64 %y) nounwind {
7779
; CHECK-LABEL: multiple_uses_sub_vscale_i64:
78-
; CHECK-NEXT: rdvl x8, #1
79-
; CHECK-NEXT: lsr x8, x8, #4
80-
; CHECK-NEXT: sub x9, x0, x8
81-
; CHECK-NEXT: add x8, x1, x8
82-
; CHECK-NEXT: mul x0, x9, x8
83-
; CHECK-NEXT: ret
80+
; CHECK: // %bb.0:
81+
; CHECK-NEXT: rdvl x8, #1
82+
; CHECK-NEXT: lsr x8, x8, #4
83+
; CHECK-NEXT: sub x9, x0, x8
84+
; CHECK-NEXT: add x8, x1, x8
85+
; CHECK-NEXT: mul x0, x9, x8
86+
; CHECK-NEXT: ret
8487
%vscale = call i64 @llvm.vscale.i64()
8588
%sub = sub i64 %x, %vscale
8689
%add = add i64 %y, %vscale
@@ -95,20 +98,48 @@ define i64 @multiple_uses_sub_vscale_i64(i64 %x, i64 %y) nounwind {
9598
; Hence, the immediate for RDVL is #1.
9699
define i64 @combine_shl_vscale_i64() nounwind {
97100
; CHECK-LABEL: combine_shl_vscale_i64:
98-
; CHECK-NOT: shl
99-
; CHECK-NEXT: rdvl x0, #1
100-
; CHECK-NEXT: ret
101+
; CHECK: // %bb.0:
102+
; CHECK-NEXT: rdvl x0, #1
103+
; CHECK-NEXT: ret
101104
%vscale = call i64 @llvm.vscale.i64()
102105
%shl = shl i64 %vscale, 4
103106
ret i64 %shl
104107
}
105108

106109
define i32 @combine_shl_vscale_i32() nounwind {
107110
; CHECK-LABEL: combine_shl_vscale_i32:
108-
; CHECK-NOT: shl
109-
; CHECK-NEXT: rdvl x0, #1
110-
; CHECK-NEXT: ret
111+
; CHECK: // %bb.0:
112+
; CHECK-NEXT: rdvl x0, #1
113+
; CHECK-NEXT: // kill: def $w0 killed $w0 killed $x0
114+
; CHECK-NEXT: ret
111115
%vscale = call i32 @llvm.vscale.i32()
112116
%shl = shl i32 %vscale, 4
113117
ret i32 %shl
114118
}
119+
120+
define i64 @combine_shl_mul_vscale(i64 %a) nounwind {
121+
; CHECK-LABEL: combine_shl_mul_vscale:
122+
; CHECK: // %bb.0:
123+
; CHECK-NEXT: cnth x8
124+
; CHECK-NEXT: mul x0, x0, x8
125+
; CHECK-NEXT: ret
126+
%vscale = tail call i64 @llvm.vscale.i64()
127+
%mul = mul i64 %a, %vscale
128+
%shl = shl i64 %mul, 3
129+
ret i64 %shl
130+
}
131+
132+
define i64 @combine_shl_mul_vscale_commuted(i64 %a) nounwind {
133+
; CHECK-LABEL: combine_shl_mul_vscale_commuted:
134+
; CHECK: // %bb.0:
135+
; CHECK-NEXT: cnth x8
136+
; CHECK-NEXT: mul x0, x0, x8
137+
; CHECK-NEXT: ret
138+
%vscale = tail call i64 @llvm.vscale.i64()
139+
%mul = mul i64 %vscale, %a
140+
%shl = shl i64 %mul, 3
141+
ret i64 %shl
142+
}
143+
144+
declare i32 @llvm.vscale.i32()
145+
declare i64 @llvm.vscale.i64()

‎llvm/test/Transforms/LoopStrengthReduce/AArch64/vscale-fixups.ll

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -383,14 +383,14 @@ define void @vscale_squared_offset(ptr %alloc) #0 {
383383
; COMMON-LABEL: vscale_squared_offset:
384384
; COMMON: // %bb.0: // %entry
385385
; COMMON-NEXT: rdvl x9, #1
386+
; COMMON-NEXT: rdvl x10, #4
386387
; COMMON-NEXT: fmov z0.s, #4.00000000
387-
; COMMON-NEXT: mov x8, xzr
388388
; COMMON-NEXT: lsr x9, x9, #4
389389
; COMMON-NEXT: fmov z1.s, #8.00000000
390-
; COMMON-NEXT: cntw x10
390+
; COMMON-NEXT: mov x8, xzr
391391
; COMMON-NEXT: ptrue p0.s, vl1
392-
; COMMON-NEXT: umull x9, w9, w9
393-
; COMMON-NEXT: lsl x9, x9, #6
392+
; COMMON-NEXT: umull x9, w9, w10
393+
; COMMON-NEXT: cntw x10
394394
; COMMON-NEXT: cmp x8, x10
395395
; COMMON-NEXT: b.ge .LBB6_2
396396
; COMMON-NEXT: .LBB6_1: // %for.body

0 commit comments

Comments
 (0)