Skip to content

Commit 3298946

Browse files
MacDueMacDue
committed
[mlir][ArmSME] Add a tests showing liveness issues in the tile allocator
This test shows a few cases (not at all complete) where the current ArmSME tile allocator produces incorrect results. The plan is to resolve these issues with a future tile allocator that uses liveness information.
1 parent df762a1 commit 3298946

File tree

1 file changed

+272
-0
lines changed

1 file changed

+272
-0
lines changed

‎mlir/test/Dialect/ArmSME/tile-allocation-liveness.mlir

Lines changed: 272 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,272 @@
1+
// RUN: mlir-opt %s -allocate-arm-sme-tiles -split-input-file -verify-diagnostics | FileCheck %s
2+
3+
// This file tests some simple aspects of using liveness in the SME tile allocator.
4+
5+
// Note: This is an XFAIL the new allocator is not yet upstream, and the current
6+
// allocator gives incorrect results for these tests.
7+
// XFAIL: *
8+
9+
// CHECK-LIVE-RANGE: ========== Coalesced Live Ranges:
10+
// CHECK-LIVE-RANGE-NEXT: @constant_with_multiple_users
11+
// CHECK-LIVE-RANGE: ^bb0:
12+
// CHECK-LIVE-RANGE: S arm_sme.zero
13+
// CHECK-LIVE-RANGE-NEXT: |S arm_sme.move_vector_to_tile_slice
14+
// CHECK-LIVE-RANGE-NEXT: || arm_sme.move_vector_to_tile_slice
15+
// CHECK-LIVE-RANGE-NEXT: |E test.some_use
16+
// CHECK-LIVE-RANGE-NEXT: E test.some_use
17+
18+
// CHECK-LABEL: @constant_with_multiple_users(
19+
// CHECK-SAME: %[[VECTOR_A:.*]]: vector<[4]xf32>, %[[VECTOR_B:.*]]: vector<[4]xf32>
20+
func.func @constant_with_multiple_users(%a: vector<[4]xf32>, %b: vector<[4]xf32>, %index: index) {
21+
// CHECK-NEXT: %[[ZERO_TILE_0:.*]] = arm_sme.zero {tile_id = 0 : i32} : vector<[4]x[4]xf32>
22+
// CHECK-NEXT: %[[ZERO_TILE_1:.*]] = arm_sme.zero {tile_id = 1 : i32} : vector<[4]x[4]xf32>
23+
// CHECK-NEXT: %[[INSERT_TILE_1:.*]] = arm_sme.move_vector_to_tile_slice %[[VECTOR_A]], %[[ZERO_TILE_1]], %{{.*}} {tile_id = 1 : i32} : vector<[4]xf32> into vector<[4]x[4]xf32>
24+
// CHECK-NEXT: %[[INSERT_TILE_0:.*]] = arm_sme.move_vector_to_tile_slice %[[VECTOR_B]], %[[ZERO_TILE_0]], %{{.*}} {tile_id = 0 : i32} : vector<[4]xf32> into vector<[4]x[4]xf32>
25+
%zero = arm_sme.zero : vector<[4]x[4]xf32>
26+
%tile_a = arm_sme.move_vector_to_tile_slice %a, %zero, %index : vector<[4]xf32> into vector<[4]x[4]xf32>
27+
%tile_b = arm_sme.move_vector_to_tile_slice %b, %zero, %index : vector<[4]xf32> into vector<[4]x[4]xf32>
28+
"test.some_use"(%tile_a) : (vector<[4]x[4]xf32>) -> ()
29+
"test.some_use"(%tile_b) : (vector<[4]x[4]xf32>) -> ()
30+
return
31+
}
32+
33+
// -----
34+
35+
// CHECK-LIVE-RANGE: ========== Coalesced Live Ranges:
36+
// CHECK-LIVE-RANGE-NEXT: @value_with_multiple_users
37+
// CHECK-LIVE-RANGE: ^bb0:
38+
// CHECK-LIVE-RANGE-NEXT: |S arm_sme.move_vector_to_tile_slice
39+
// CHECK-LIVE-RANGE-NEXT: || arm_sme.move_vector_to_tile_slice
40+
// CHECK-LIVE-RANGE-NEXT: |E test.some_use
41+
// CHECK-LIVE-RANGE-NEXT: E test.some_use
42+
43+
func.func @value_with_multiple_users(%tile: vector<[4]x[4]xf32>, %a: vector<[4]xf32>, %b: vector<[4]xf32>, %index: index) {
44+
// expected-error@below {{op failed to rectify tile operand with tile result (move required)}}
45+
%tile_a = arm_sme.move_vector_to_tile_slice %a, %tile, %index : vector<[4]xf32> into vector<[4]x[4]xf32>
46+
%tile_b = arm_sme.move_vector_to_tile_slice %b, %tile, %index : vector<[4]xf32> into vector<[4]x[4]xf32>
47+
"test.some_use"(%tile_a) : (vector<[4]x[4]xf32>) -> ()
48+
"test.some_use"(%tile_b) : (vector<[4]x[4]xf32>) -> ()
49+
return
50+
}
51+
52+
// -----
53+
54+
// CHECK-LIVE-RANGE: ========== Coalesced Live Ranges:
55+
// CHECK-LIVE-RANGE-NEXT: @reuse_tiles_after_initial_use
56+
// CHECK-LIVE-RANGE: ^bb0:
57+
// CHECK-LIVE-RANGE-NEXT: S arm_sme.get_tile
58+
// CHECK-LIVE-RANGE-NEXT: |S arm_sme.get_tile
59+
// CHECK-LIVE-RANGE-NEXT: ||S arm_sme.get_tile
60+
// CHECK-LIVE-RANGE-NEXT: |||S arm_sme.get_tile
61+
// CHECK-LIVE-RANGE-NEXT: |||| test.dummy
62+
// CHECK-LIVE-RANGE-NEXT: |||| test.dummy
63+
// CHECK-LIVE-RANGE-NEXT: |||| test.dummy
64+
// CHECK-LIVE-RANGE-NEXT: E||| test.some_use
65+
// CHECK-LIVE-RANGE-NEXT: E|| test.some_use
66+
// CHECK-LIVE-RANGE-NEXT: E| test.some_use
67+
// CHECK-LIVE-RANGE-NEXT: E test.some_use
68+
// CHECK-LIVE-RANGE-NEXT: S arm_sme.zero
69+
// CHECK-LIVE-RANGE-NEXT: |S arm_sme.zero
70+
// CHECK-LIVE-RANGE-NEXT: ||S arm_sme.zero
71+
// CHECK-LIVE-RANGE-NEXT: |||S arm_sme.zero
72+
// CHECK-LIVE-RANGE-NEXT: |||| test.dummy
73+
// CHECK-LIVE-RANGE-NEXT: |||| test.dummy
74+
// CHECK-LIVE-RANGE-NEXT: |||| test.dummy
75+
// CHECK-LIVE-RANGE-NEXT: E||| test.some_use
76+
// CHECK-LIVE-RANGE-NEXT: E|| test.some_use
77+
// CHECK-LIVE-RANGE-NEXT: E| test.some_use
78+
// CHECK-LIVE-RANGE-NEXT: E test.some_use
79+
80+
// CHECK-LABEL: @reuse_tiles_after_initial_use
81+
func.func @reuse_tiles_after_initial_use() {
82+
// CHECK: arm_sme.get_tile {tile_id = 0 : i32}
83+
// CHECK: arm_sme.get_tile {tile_id = 1 : i32}
84+
// CHECK: arm_sme.get_tile {tile_id = 2 : i32}
85+
// CHECK: arm_sme.get_tile {tile_id = 3 : i32}
86+
%tile_a = arm_sme.get_tile : vector<[4]x[4]xf32>
87+
%tile_b = arm_sme.get_tile : vector<[4]x[4]xf32>
88+
%tile_c = arm_sme.get_tile : vector<[4]x[4]xf32>
89+
%tile_d = arm_sme.get_tile : vector<[4]x[4]xf32>
90+
"test.dummy"(): () -> ()
91+
"test.dummy"(): () -> ()
92+
"test.dummy"(): () -> ()
93+
"test.some_use"(%tile_a) : (vector<[4]x[4]xf32>) -> ()
94+
"test.some_use"(%tile_b) : (vector<[4]x[4]xf32>) -> ()
95+
"test.some_use"(%tile_c) : (vector<[4]x[4]xf32>) -> ()
96+
"test.some_use"(%tile_d) : (vector<[4]x[4]xf32>) -> ()
97+
// CHECK: arm_sme.zero {tile_id = 0 : i32}
98+
// CHECK: arm_sme.zero {tile_id = 1 : i32}
99+
// CHECK: arm_sme.zero {tile_id = 2 : i32}
100+
// CHECK: arm_sme.zero {tile_id = 3 : i32}
101+
%tile_1 = arm_sme.zero : vector<[4]x[4]xf32>
102+
%tile_2 = arm_sme.zero : vector<[4]x[4]xf32>
103+
%tile_3 = arm_sme.zero : vector<[4]x[4]xf32>
104+
%tile_4 = arm_sme.zero : vector<[4]x[4]xf32>
105+
"test.dummy"(): () -> ()
106+
"test.dummy"(): () -> ()
107+
"test.dummy"(): () -> ()
108+
"test.some_use"(%tile_1) : (vector<[4]x[4]xf32>) -> ()
109+
"test.some_use"(%tile_2) : (vector<[4]x[4]xf32>) -> ()
110+
"test.some_use"(%tile_3) : (vector<[4]x[4]xf32>) -> ()
111+
"test.some_use"(%tile_4) : (vector<[4]x[4]xf32>) -> ()
112+
return
113+
}
114+
115+
// -----
116+
117+
// CHECK-LIVE-RANGE: ========== Coalesced Live Ranges:
118+
// CHECK-LIVE-RANGE-NEXT: @non_overlapping_branches
119+
// CHECK-LIVE-RANGE: ^bb1:
120+
// CHECK-LIVE-RANGE-NEXT: S arm_sme.zero
121+
// CHECK-LIVE-RANGE-NEXT: | arm_sme.copy_tile
122+
// CHECK-LIVE-RANGE-NEXT: E cf.br
123+
// CHECK-LIVE-RANGE-NEXT: ^bb2:
124+
// CHECK-LIVE-RANGE-NEXT: S arm_sme.get_tile
125+
// CHECK-LIVE-RANGE-NEXT: | arm_sme.copy_tile
126+
// CHECK-LIVE-RANGE-NEXT: E cf.br
127+
128+
// CHECK-LABEL: @non_overlapping_branches
129+
func.func @non_overlapping_branches(%cond: i1) {
130+
// CHECK: arm_sme.zero {tile_id = 0 : i32} : vector<[4]x[4]xf32>
131+
// CHECK: arm_sme.get_tile {tile_id = 0 : i32} : vector<[4]x[4]xf32>
132+
%tile = scf.if %cond -> vector<[4]x[4]xf32> {
133+
// ^bb1:
134+
%zero = arm_sme.zero : vector<[4]x[4]xf32>
135+
scf.yield %zero : vector<[4]x[4]xf32>
136+
} else {
137+
// ^bb2:
138+
%undef = arm_sme.get_tile : vector<[4]x[4]xf32>
139+
scf.yield %undef : vector<[4]x[4]xf32>
140+
}
141+
"test.some_use"(%tile) : (vector<[4]x[4]xf32>) -> ()
142+
return
143+
}
144+
145+
// -----
146+
147+
// CHECK-LIVE-RANGE: ========== Coalesced Live Ranges:
148+
// <deliberately omitted>
149+
150+
// CHECK-LABEL: @constant_loop_init_with_multiple_users
151+
func.func @constant_loop_init_with_multiple_users(%a: vector<[4]xf32>, %b: vector<[4]xf32>) {
152+
// CHECK: arm_sme.zero {tile_id = 0 : i32} : vector<[4]x[4]xf32>
153+
// CHECK: arm_sme.zero {tile_id = 1 : i32} : vector<[4]x[4]xf32>
154+
// CHECK: arm_sme.move_vector_to_tile_slice {{.*}} {tile_id = 1 : i32} : vector<[4]xf32> into vector<[4]x[4]xf32>
155+
// CHECK: arm_sme.move_vector_to_tile_slice {{.*}} {tile_id = 0 : i32} : vector<[4]xf32> into vector<[4]x[4]xf32>
156+
%c0 = arith.constant 0 : index
157+
%c1 = arith.constant 1 : index
158+
%c10 = arith.constant 10 : index
159+
%init = arm_sme.zero : vector<[4]x[4]xf32>
160+
%tile_a = scf.for %i = %c0 to %c10 step %c1 iter_args(%iter = %init) -> vector<[4]x[4]xf32> {
161+
%new_tile = arm_sme.move_vector_to_tile_slice %a, %iter, %i : vector<[4]xf32> into vector<[4]x[4]xf32>
162+
scf.yield %new_tile : vector<[4]x[4]xf32>
163+
}
164+
%tile_b = scf.for %i = %c0 to %c10 step %c1 iter_args(%iter = %init) -> vector<[4]x[4]xf32> {
165+
%new_tile = arm_sme.move_vector_to_tile_slice %a, %iter, %i : vector<[4]xf32> into vector<[4]x[4]xf32>
166+
scf.yield %new_tile : vector<[4]x[4]xf32>
167+
}
168+
"test.some_use"(%tile_a) : (vector<[4]x[4]xf32>) -> ()
169+
"test.some_use"(%tile_b) : (vector<[4]x[4]xf32>) -> ()
170+
return
171+
}
172+
173+
// -----
174+
175+
// CHECK-LIVE-RANGE: ========== Coalesced Live Ranges:
176+
// CHECK-LIVE-RANGE-NEXT: @run_out_of_tiles_but_avoid_spill
177+
// CHECK-LIVE-RANGE: ^bb2:
178+
// CHECK-LIVE-RANGE-NEXT: |S arm_sme.copy_tile
179+
// CHECK-LIVE-RANGE-NEXT: ||S arm_sme.copy_tile
180+
// CHECK-LIVE-RANGE-NEXT: |||S arm_sme.copy_tile
181+
// CHECK-LIVE-RANGE-NEXT: ||||S arm_sme.copy_tile
182+
// CHECK-LIVE-RANGE-NEXT: EEEEE cf.br
183+
184+
// Note in the live ranges (above) there is five tile values, but we only have four tiles.
185+
186+
// CHECK-LABEL: @run_out_of_tiles_but_avoid_spill
187+
func.func @run_out_of_tiles_but_avoid_spill(%a: vector<[4]xf32>, %b: vector<[4]xf32>, %c: vector<[4]xf32>, %d: vector<[4]xf32>) {
188+
%init = arm_sme.zero : vector<[4]x[4]xf32>
189+
%c0 = arith.constant 0 : index
190+
%c1 = arith.constant 1 : index
191+
%c10 = arith.constant 10 : index
192+
// Live = %init
193+
scf.for %i = %c0 to %c10 step %c1 {
194+
// CHECK: arm_sme.zero {tile_id = 1 : i32}
195+
// CHECK: arm_sme.zero {tile_id = 2 : i32}
196+
// CHECK: arm_sme.zero {tile_id = 3 : i32}
197+
// CHECK: arm_sme.zero {tile_id = 0 : i32}
198+
%tile_a, %tile_b, %tile_c, %tile_d = scf.for %j = %c0 to %c10 step %c1
199+
iter_args(%iter_a = %init, %iter_b = %init, %iter_c = %init, %iter_d = %init)
200+
-> (vector<[4]x[4]xf32>, vector<[4]x[4]xf32> , vector<[4]x[4]xf32> , vector<[4]x[4]xf32>) {
201+
// ^bb2:
202+
// CHECK: arm_sme.move_vector_to_tile_slice {{.*}} {tile_id = 1 : i32} : vector<[4]xf32> into vector<[4]x[4]xf32>
203+
// CHECK: arm_sme.move_vector_to_tile_slice {{.*}} {tile_id = 2 : i32} : vector<[4]xf32> into vector<[4]x[4]xf32>
204+
// CHECK: arm_sme.move_vector_to_tile_slice {{.*}} {tile_id = 3 : i32} : vector<[4]xf32> into vector<[4]x[4]xf32>
205+
// CHECK: arm_sme.move_vector_to_tile_slice {{.*}} {tile_id = 0 : i32} : vector<[4]xf32> into vector<[4]x[4]xf32>
206+
%new_a = arm_sme.move_vector_to_tile_slice %a, %iter_a, %i : vector<[4]xf32> into vector<[4]x[4]xf32>
207+
%new_b = arm_sme.move_vector_to_tile_slice %b, %iter_b, %i : vector<[4]xf32> into vector<[4]x[4]xf32>
208+
%new_c = arm_sme.move_vector_to_tile_slice %c, %iter_c, %i : vector<[4]xf32> into vector<[4]x[4]xf32>
209+
%new_d = arm_sme.move_vector_to_tile_slice %d, %iter_d, %i : vector<[4]xf32> into vector<[4]x[4]xf32>
210+
scf.yield %new_a, %new_b, %new_c, %new_d : vector<[4]x[4]xf32>, vector<[4]x[4]xf32>, vector<[4]x[4]xf32>, vector<[4]x[4]xf32>
211+
}
212+
// Live = %init, %tile_a, %tile_b, %tile_c, %tile_d (out of tiles!)
213+
// This should be resolved by duplicating the arm_sme.zero (from folding
214+
// arm_sme.copy_tile operations inserted by the tile allocator).
215+
"test.some_use"(%tile_a) : (vector<[4]x[4]xf32>) -> ()
216+
"test.some_use"(%tile_b) : (vector<[4]x[4]xf32>) -> ()
217+
"test.some_use"(%tile_c) : (vector<[4]x[4]xf32>) -> ()
218+
"test.some_use"(%tile_d) : (vector<[4]x[4]xf32>) -> ()
219+
}
220+
return
221+
}
222+
223+
// -----
224+
225+
// We should be able to avoid spills like this, but logic handling this case is
226+
// not implemented yet. Note tile ID >= 16 means a spill/in-memory tile.
227+
228+
// CHECK-LIVE-RANGE: ========== Coalesced Live Ranges:
229+
// CHECK-LIVE-RANGE-NEXT: @avoidable_spill
230+
// CHECK-LIVE-RANGE: ^bb2:
231+
// CHECK-LIVE-RANGE-NEXT: || test.some_use
232+
// CHECK-LIVE-RANGE-NEXT: ||S arm_sme.move_vector_to_tile_slice
233+
// CHECK-LIVE-RANGE-NEXT: |||S arm_sme.move_vector_to_tile_slice
234+
// CHECK-LIVE-RANGE-NEXT: ||||S arm_sme.move_vector_to_tile_slice
235+
// CHECK-LIVE-RANGE-NEXT: |||||S arm_sme.move_vector_to_tile_slice
236+
// CHECK-LIVE-RANGE-NEXT: ||E||| test.some_use
237+
// CHECK-LIVE-RANGE-NEXT: || E|| test.some_use
238+
// CHECK-LIVE-RANGE-NEXT: || E| test.some_use
239+
// CHECK-LIVE-RANGE-NEXT: || E test.some_use
240+
// CHECK-LIVE-RANGE-NEXT: || arith.addi
241+
// CHECK-LIVE-RANGE-NEXT: EE cf.br
242+
243+
// Note in the live ranges (above) there is two constant live-ins (first two ranges),
244+
// which gives six overlapping live ranges. The allocator currently will spill the
245+
// first constant (which results in a real spill at it's use), however, this could
246+
// be avoided by using the knowledge that at the first "test.some_use" there's
247+
// actually only two live ranges (so we can fix this be duplicating the constant).
248+
249+
// CHECK-LABEL: @avoidable_spill
250+
func.func @avoidable_spill(%a: vector<[4]xf32>, %b: vector<[4]xf32>, %c: vector<[4]xf32>, %d: vector<[4]xf32>) {
251+
// CHECK: arm_sme.zero {tile_id = 16 : i32} : vector<[4]x[4]xf32>
252+
%zero = arm_sme.zero : vector<[4]x[4]xf32>
253+
%tile = arm_sme.get_tile : vector<[4]x[4]xf32>
254+
%c0 = arith.constant 0 : index
255+
%c1 = arith.constant 1 : index
256+
%c10 = arith.constant 10 : index
257+
scf.for %i = %c0 to %c10 step %c1 {
258+
// So spilled here (unnecessarily).
259+
// The arm_sme.zero op could be moved into the loop to avoid this.
260+
"test.some_use"(%zero) : (vector<[4]x[4]xf32>) -> ()
261+
%tile_a = arm_sme.move_vector_to_tile_slice %a, %tile, %c0 : vector<[4]xf32> into vector<[4]x[4]xf32>
262+
%tile_b = arm_sme.move_vector_to_tile_slice %b, %tile, %c0 : vector<[4]xf32> into vector<[4]x[4]xf32>
263+
%tile_c = arm_sme.move_vector_to_tile_slice %c, %tile, %c0 : vector<[4]xf32> into vector<[4]x[4]xf32>
264+
%tile_d = arm_sme.move_vector_to_tile_slice %d, %tile, %c0 : vector<[4]xf32> into vector<[4]x[4]xf32>
265+
// %zero is still live here (due the the backedge)
266+
"test.some_use"(%tile_a) : (vector<[4]x[4]xf32>) -> ()
267+
"test.some_use"(%tile_b) : (vector<[4]x[4]xf32>) -> ()
268+
"test.some_use"(%tile_c) : (vector<[4]x[4]xf32>) -> ()
269+
"test.some_use"(%tile_d) : (vector<[4]x[4]xf32>) -> ()
270+
}
271+
return
272+
}

0 commit comments

Comments
 (0)