Skip to content

Commit 9226688

Browse files
MacDueMacDue
authored
[mlir][ArmSME] Add a tests showing liveness issues in the tile allocator (#90447)
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 0606747 commit 9226688

File tree

1 file changed

+183
-0
lines changed

1 file changed

+183
-0
lines changed

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

Lines changed: 183 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,183 @@
1+
// RUN: mlir-opt %s -allocate-arm-sme-tiles -split-input-file -verify-diagnostics | FileCheck %s --check-prefix=CHECK-BAD
2+
3+
// This file tests some aspects of liveness issues in the SME tile allocator.
4+
// These tests were designed with a new liveness-based tile allocator in mind
5+
// (where the names of test cases make more sense), with the current tile
6+
// allocator these tests all give incorrect results (which is documented by
7+
// `CHECK-BAD`).
8+
9+
// Incorrect result! The second `move_vector_to_tile_slice` overwrites the first (which is still live).
10+
//
11+
// CHECK-BAD-LABEL: @constant_with_multiple_users
12+
// CHECK-BAD: %[[ZERO_TILE:.*]] = arm_sme.zero {tile_id = 0 : i32} : vector<[4]x[4]xf32>
13+
// CHECK-BAD: %[[INSERT_TILE_1:.*]] = arm_sme.move_vector_to_tile_slice %{{.*}} {tile_id = 0 : i32} : vector<[4]xf32> into vector<[4]x[4]xf32>
14+
// CHECK-BAD: %[[INSERT_TILE_0:.*]] = arm_sme.move_vector_to_tile_slice %{{.*}} {tile_id = 0 : i32} : vector<[4]xf32> into vector<[4]x[4]xf32>
15+
func.func @constant_with_multiple_users(%a: vector<[4]xf32>, %b: vector<[4]xf32>, %index: index) {
16+
%zero = arm_sme.zero : vector<[4]x[4]xf32>
17+
%tile_a = arm_sme.move_vector_to_tile_slice %a, %zero, %index : vector<[4]xf32> into vector<[4]x[4]xf32>
18+
%tile_b = arm_sme.move_vector_to_tile_slice %b, %zero, %index : vector<[4]xf32> into vector<[4]x[4]xf32>
19+
"test.some_use"(%tile_a) : (vector<[4]x[4]xf32>) -> ()
20+
"test.some_use"(%tile_b) : (vector<[4]x[4]xf32>) -> ()
21+
return
22+
}
23+
24+
// -----
25+
26+
// (No tile IDs -- the current tile allocator ignores this case)
27+
28+
// CHECK-BAD-LABEL: @value_with_multiple_users
29+
// CHECK-BAD-NOT: tile_id
30+
func.func @value_with_multiple_users(%tile: vector<[4]x[4]xf32>, %a: vector<[4]xf32>, %b: vector<[4]xf32>, %index: index) {
31+
// A future allocator should error here (as `%tile` would need to be copied).
32+
%tile_a = arm_sme.move_vector_to_tile_slice %a, %tile, %index : vector<[4]xf32> into vector<[4]x[4]xf32>
33+
%tile_b = arm_sme.move_vector_to_tile_slice %b, %tile, %index : vector<[4]xf32> into vector<[4]x[4]xf32>
34+
"test.some_use"(%tile_a) : (vector<[4]x[4]xf32>) -> ()
35+
"test.some_use"(%tile_b) : (vector<[4]x[4]xf32>) -> ()
36+
return
37+
}
38+
39+
// -----
40+
41+
// CHECK-BAD-LABEL: @reuse_tiles_after_initial_use
42+
func.func @reuse_tiles_after_initial_use() {
43+
// CHECK-BAD: arm_sme.get_tile {tile_id = 0 : i32}
44+
// CHECK-BAD: arm_sme.get_tile {tile_id = 1 : i32}
45+
// CHECK-BAD: arm_sme.get_tile {tile_id = 2 : i32}
46+
// CHECK-BAD: arm_sme.get_tile {tile_id = 3 : i32}
47+
%tile_a = arm_sme.get_tile : vector<[4]x[4]xf32>
48+
%tile_b = arm_sme.get_tile : vector<[4]x[4]xf32>
49+
%tile_c = arm_sme.get_tile : vector<[4]x[4]xf32>
50+
%tile_d = arm_sme.get_tile : vector<[4]x[4]xf32>
51+
"test.dummy"(): () -> ()
52+
"test.dummy"(): () -> ()
53+
"test.dummy"(): () -> ()
54+
"test.some_use"(%tile_a) : (vector<[4]x[4]xf32>) -> ()
55+
"test.some_use"(%tile_b) : (vector<[4]x[4]xf32>) -> ()
56+
"test.some_use"(%tile_c) : (vector<[4]x[4]xf32>) -> ()
57+
"test.some_use"(%tile_d) : (vector<[4]x[4]xf32>) -> ()
58+
// -> Spills after the fourth tile (unnecessary):
59+
// CHECK-BAD: arm_sme.zero {tile_id = 16 : i32}
60+
// CHECK-BAD: arm_sme.zero {tile_id = 17 : i32}
61+
// CHECK-BAD: arm_sme.zero {tile_id = 18 : i32}
62+
// CHECK-BAD: arm_sme.zero {tile_id = 19 : i32}
63+
// Unnecessary spills:
64+
// expected-warning @below {{failed to allocate SME virtual tile to operation, all tile operations will go through memory, expect degraded performance}}
65+
%tile_1 = arm_sme.zero : vector<[4]x[4]xf32>
66+
// expected-warning @below {{failed to allocate SME virtual tile to operation, all tile operations will go through memory, expect degraded performance}}
67+
%tile_2 = arm_sme.zero : vector<[4]x[4]xf32>
68+
// expected-warning @below {{failed to allocate SME virtual tile to operation, all tile operations will go through memory, expect degraded performance}}
69+
%tile_3 = arm_sme.zero : vector<[4]x[4]xf32>
70+
// expected-warning @below {{failed to allocate SME virtual tile to operation, all tile operations will go through memory, expect degraded performance}}
71+
%tile_4 = arm_sme.zero : vector<[4]x[4]xf32>
72+
"test.dummy"(): () -> ()
73+
"test.dummy"(): () -> ()
74+
"test.dummy"(): () -> ()
75+
"test.some_use"(%tile_1) : (vector<[4]x[4]xf32>) -> ()
76+
"test.some_use"(%tile_2) : (vector<[4]x[4]xf32>) -> ()
77+
"test.some_use"(%tile_3) : (vector<[4]x[4]xf32>) -> ()
78+
"test.some_use"(%tile_4) : (vector<[4]x[4]xf32>) -> ()
79+
return
80+
}
81+
82+
// -----
83+
84+
// Incorrect result! Both branches should yield the result via the same tile.
85+
//
86+
// CHECK-BAD-LABEL: @non_overlapping_branches
87+
// CHECK-BAD: arm_sme.zero {tile_id = 0 : i32} : vector<[4]x[4]xf32>
88+
// CHECK-BAD: arm_sme.get_tile {tile_id = 1 : i32} : vector<[4]x[4]xf32>
89+
func.func @non_overlapping_branches(%cond: i1) {
90+
%tile = scf.if %cond -> vector<[4]x[4]xf32> {
91+
%zero = arm_sme.zero : vector<[4]x[4]xf32>
92+
scf.yield %zero : vector<[4]x[4]xf32>
93+
} else {
94+
%undef = arm_sme.get_tile : vector<[4]x[4]xf32>
95+
scf.yield %undef : vector<[4]x[4]xf32>
96+
}
97+
"test.some_use"(%tile) : (vector<[4]x[4]xf32>) -> ()
98+
return
99+
}
100+
101+
// -----
102+
103+
// Incorrect result! Everything assigned to tile 0 (which means values that are still live are overwritten).
104+
//
105+
// CHECK-BAD-LABEL: @constant_loop_init_with_multiple_users
106+
// CHECK-BAD: arm_sme.zero {tile_id = 0 : i32} : vector<[4]x[4]xf32>
107+
// CHECK-BAD: arm_sme.move_vector_to_tile_slice {{.*}} {tile_id = 0 : i32} : vector<[4]xf32> into vector<[4]x[4]xf32>
108+
// CHECK-BAD: arm_sme.move_vector_to_tile_slice {{.*}} {tile_id = 0 : i32} : vector<[4]xf32> into vector<[4]x[4]xf32>
109+
func.func @constant_loop_init_with_multiple_users(%a: vector<[4]xf32>, %b: vector<[4]xf32>) {
110+
%c0 = arith.constant 0 : index
111+
%c1 = arith.constant 1 : index
112+
%c10 = arith.constant 10 : index
113+
%init = arm_sme.zero : vector<[4]x[4]xf32>
114+
%tile_a = scf.for %i = %c0 to %c10 step %c1 iter_args(%iter = %init) -> vector<[4]x[4]xf32> {
115+
%new_tile = arm_sme.move_vector_to_tile_slice %a, %iter, %i : vector<[4]xf32> into vector<[4]x[4]xf32>
116+
scf.yield %new_tile : vector<[4]x[4]xf32>
117+
}
118+
%tile_b = scf.for %i = %c0 to %c10 step %c1 iter_args(%iter = %init) -> vector<[4]x[4]xf32> {
119+
%new_tile = arm_sme.move_vector_to_tile_slice %a, %iter, %i : vector<[4]xf32> into vector<[4]x[4]xf32>
120+
scf.yield %new_tile : vector<[4]x[4]xf32>
121+
}
122+
"test.some_use"(%tile_a) : (vector<[4]x[4]xf32>) -> ()
123+
"test.some_use"(%tile_b) : (vector<[4]x[4]xf32>) -> ()
124+
return
125+
}
126+
127+
// -----
128+
129+
// Incorrect result! Everything assigned to tile 0 (which means values that are still live are overwritten).
130+
//
131+
// CHECK-BAD-LABEL: @run_out_of_tiles_but_avoid_spill
132+
// CHECK-BAD: arm_sme.zero {tile_id = 0 : i32}
133+
// CHECK-BAD-COUNT-4: arm_sme.move_vector_to_tile_slice {{.*}} {tile_id = 0 : i32} : vector<[4]xf32> into vector<[4]x[4]xf32>
134+
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>) {
135+
%init = arm_sme.zero : vector<[4]x[4]xf32>
136+
%c0 = arith.constant 0 : index
137+
%c1 = arith.constant 1 : index
138+
%c10 = arith.constant 10 : index
139+
scf.for %i = %c0 to %c10 step %c1 {
140+
%tile_a, %tile_b, %tile_c, %tile_d = scf.for %j = %c0 to %c10 step %c1
141+
iter_args(%iter_a = %init, %iter_b = %init, %iter_c = %init, %iter_d = %init)
142+
-> (vector<[4]x[4]xf32>, vector<[4]x[4]xf32> , vector<[4]x[4]xf32> , vector<[4]x[4]xf32>) {
143+
%new_a = arm_sme.move_vector_to_tile_slice %a, %iter_a, %i : vector<[4]xf32> into vector<[4]x[4]xf32>
144+
%new_b = arm_sme.move_vector_to_tile_slice %b, %iter_b, %i : vector<[4]xf32> into vector<[4]x[4]xf32>
145+
%new_c = arm_sme.move_vector_to_tile_slice %c, %iter_c, %i : vector<[4]xf32> into vector<[4]x[4]xf32>
146+
%new_d = arm_sme.move_vector_to_tile_slice %d, %iter_d, %i : vector<[4]xf32> into vector<[4]x[4]xf32>
147+
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>
148+
}
149+
"test.some_use"(%tile_a) : (vector<[4]x[4]xf32>) -> ()
150+
"test.some_use"(%tile_b) : (vector<[4]x[4]xf32>) -> ()
151+
"test.some_use"(%tile_c) : (vector<[4]x[4]xf32>) -> ()
152+
"test.some_use"(%tile_d) : (vector<[4]x[4]xf32>) -> ()
153+
}
154+
return
155+
}
156+
157+
// -----
158+
159+
// Incorrect result! Everything other than zero assigned to tile 1 (which means values that are still live are overwritten).
160+
//
161+
// CHECK-BAD-LABEL: @avoidable_spill
162+
// CHECK-BAD: arm_sme.zero {tile_id = 0 : i32}
163+
// CHECK-BAD: arm_sme.get_tile {tile_id = 1 : i32}
164+
// CHECK-BAD-COUNT-4: arm_sme.move_vector_to_tile_slice {{.*}} {tile_id = 1 : i32}
165+
func.func @avoidable_spill(%a: vector<[4]xf32>, %b: vector<[4]xf32>, %c: vector<[4]xf32>, %d: vector<[4]xf32>) {
166+
%zero = arm_sme.zero : vector<[4]x[4]xf32>
167+
%tile = arm_sme.get_tile : vector<[4]x[4]xf32>
168+
%c0 = arith.constant 0 : index
169+
%c1 = arith.constant 1 : index
170+
%c10 = arith.constant 10 : index
171+
scf.for %i = %c0 to %c10 step %c1 {
172+
"test.some_use"(%zero) : (vector<[4]x[4]xf32>) -> ()
173+
%tile_a = arm_sme.move_vector_to_tile_slice %a, %tile, %c0 : vector<[4]xf32> into vector<[4]x[4]xf32>
174+
%tile_b = arm_sme.move_vector_to_tile_slice %b, %tile, %c0 : vector<[4]xf32> into vector<[4]x[4]xf32>
175+
%tile_c = arm_sme.move_vector_to_tile_slice %c, %tile, %c0 : vector<[4]xf32> into vector<[4]x[4]xf32>
176+
%tile_d = arm_sme.move_vector_to_tile_slice %d, %tile, %c0 : vector<[4]xf32> into vector<[4]x[4]xf32>
177+
"test.some_use"(%tile_a) : (vector<[4]x[4]xf32>) -> ()
178+
"test.some_use"(%tile_b) : (vector<[4]x[4]xf32>) -> ()
179+
"test.some_use"(%tile_c) : (vector<[4]x[4]xf32>) -> ()
180+
"test.some_use"(%tile_d) : (vector<[4]x[4]xf32>) -> ()
181+
}
182+
return
183+
}

0 commit comments

Comments
 (0)