Skip to content

Commit 781273a

Browse files
authored
[PROTON] Refactor scope id allocation to allow flexible annotations (#8613)
1 parent d0ef9d3 commit 781273a

File tree

5 files changed

+637
-51
lines changed

5 files changed

+637
-51
lines changed

test/Proton/scope_id.mlir

Lines changed: 295 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,52 +1,64 @@
1-
// RUN: triton-opt --split-input-file --test-print-scope-id-allocation -verify-diagnostics -o /dev/null %s
1+
// RUN: triton-opt --split-input-file --test-print-scope-id-allocation -verify-diagnostics=only-expected -o /dev/null %s
22

33
module {
44
// expected-remark @below {{one_scope}}
55
tt.func @one_scope() {
66
// expected-remark @below {{scope id = 0}}
7+
// expected-remark @below {{scope parent id = -1}}
78
proton.record start "name0"
89
// expected-remark @below {{scope id = 0}}
10+
// expected-remark @below {{scope parent id = -1}}
911
proton.record end "name0"
1012
tt.return
1113
}
1214

1315
// expected-remark @below {{two_scopes}}
1416
tt.func @two_scopes() {
1517
// expected-remark @below {{scope id = 1}}
18+
// expected-remark @below {{scope parent id = -1}}
1619
proton.record start "name0"
1720
// expected-remark @below {{scope id = 1}}
21+
// expected-remark @below {{scope parent id = -1}}
1822
proton.record end "name0"
1923
// expected-remark @below {{scope id = 2}}
24+
// expected-remark @below {{scope parent id = -1}}
2025
proton.record start "name1"
2126
// expected-remark @below {{scope id = 2}}
27+
// expected-remark @below {{scope parent id = -1}}
2228
proton.record end "name1"
2329
tt.return
2430
}
2531

2632
// expected-remark @below {{two_scopes_overlap}}
2733
tt.func @two_scopes_overlap() {
2834
// expected-remark @below {{scope id = 3}}
35+
// expected-remark @below {{scope parent id = -1}}
2936
proton.record start "name0"
3037
// expected-remark @below {{scope id = 4}}
38+
// expected-remark @below {{scope parent id = -1}}
3139
proton.record start "name1"
3240
// expected-remark @below {{scope id = 3}}
41+
// expected-remark @below {{scope parent id = -1}}
3342
proton.record end "name0"
3443
// expected-remark @below {{scope id = 4}}
44+
// expected-remark @below {{scope parent id = -1}}
3545
proton.record end "name1"
3646
tt.return
3747
}
3848

39-
// expected-remark @below {{control_flow}}
40-
tt.func @control_flow(%cond: i1) {
49+
// expected-remark @below {{nested_scopes}}
50+
tt.func @nested_scopes() {
4151
// expected-remark @below {{scope id = 5}}
52+
// expected-remark @below {{scope parent id = -1}}
4253
proton.record start "name0"
43-
scf.if %cond {
44-
// expected-remark @below {{scope id = 6}}
45-
proton.record start "name1"
46-
// expected-remark @below {{scope id = 6}}
47-
proton.record end "name1"
48-
}
54+
// expected-remark @below {{scope id = 6}}
55+
// expected-remark @below {{scope parent id = 5}}
56+
proton.record start "name1"
57+
// expected-remark @below {{scope id = 6}}
58+
// expected-remark @below {{scope parent id = 5}}
59+
proton.record end "name1"
4960
// expected-remark @below {{scope id = 5}}
61+
// expected-remark @below {{scope parent id = -1}}
5062
proton.record end "name0"
5163
tt.return
5264
}
@@ -58,18 +70,22 @@ module {
5870
// expected-remark @below {{inner}}
5971
tt.func @inner() {
6072
// expected-remark @below {{scope id = 0}}
73+
// expected-remark @below {{scope parent id = -1}}
6174
proton.record start "name0"
6275
// expected-remark @below {{scope id = 0}}
76+
// expected-remark @below {{scope parent id = -1}}
6377
proton.record end "name0"
6478
tt.return
6579
}
6680

6781
// expected-remark @below {{outer}}
6882
tt.func @outer() {
6983
// expected-remark @below {{scope id = 1}}
84+
// expected-remark @below {{scope parent id = -1}}
7085
proton.record start "name0"
7186
tt.call @inner() : () -> ()
7287
// expected-remark @below {{scope id = 1}}
88+
// expected-remark @below {{scope parent id = -1}}
7389
proton.record end "name0"
7490
tt.return
7591
}
@@ -81,12 +97,16 @@ module {
8197
// expected-remark @below {{duplicate}}
8298
tt.func @duplicate() {
8399
// expected-remark @below {{scope id = 0}}
100+
// expected-remark @below {{scope parent id = -1}}
84101
proton.record start "name0"
85102
// expected-remark @below {{scope id = 0}}
103+
// expected-remark @below {{scope parent id = -1}}
86104
proton.record end "name0"
87105
// expected-remark @below {{scope id = 1}}
106+
// expected-remark @below {{scope parent id = -1}}
88107
proton.record start "name0"
89108
// expected-remark @below {{scope id = 1}}
109+
// expected-remark @below {{scope parent id = -1}}
90110
proton.record end "name0"
91111
tt.return
92112
}
@@ -95,18 +115,279 @@ module {
95115
// -----
96116

97117
module {
98-
// expected-remark @below {{condition}}
99-
tt.func @condition(%cond: i1) {
118+
// expected-remark @below {{cf_reordered}}
119+
tt.func @cf_reordered() {
120+
^entry:
121+
cf.br ^start
122+
^exit:
100123
// expected-remark @below {{scope id = 0}}
124+
// expected-remark @below {{scope parent id = -1}}
125+
proton.record end "name0"
126+
tt.return
127+
^start:
128+
// expected-remark @below {{scope id = 0}}
129+
// expected-remark @below {{scope parent id = -1}}
101130
proton.record start "name0"
131+
cf.br ^exit
132+
}
133+
}
134+
135+
// -----
136+
137+
module {
138+
// expected-remark @below {{scf_cond}}
139+
tt.func @scf_cond(%cond: i1) {
140+
scf.if %cond {
141+
// expected-remark @below {{scope id = 0}}
142+
// expected-remark @below {{scope parent id = -1}}
143+
proton.record start "if_only"
144+
}
145+
// expected-remark @below {{scope id = 0}
146+
// expected-remark @below {{scope parent id = -1}}
147+
proton.record end "if_only"
148+
tt.return
149+
}
150+
}
151+
152+
// -----
153+
154+
module {
155+
tt.func @scf_loop_if(%cond: i1) {
156+
%c0 = arith.constant 0 : index
157+
scf.for %i = %c0 to %c0 step %c0 {
158+
scf.if %cond {
159+
// expected-remark @below {{scope id = 0}}
160+
// expected-remark @below {{scope parent id = -1}}
161+
proton.record start "loop_if"
162+
}
163+
scf.if %cond {
164+
// expected-remark @below {{scope id = 0}}
165+
// expected-remark @below {{scope parent id = -1}}
166+
proton.record end "loop_if"
167+
}
168+
}
169+
tt.return
170+
}
171+
}
172+
173+
// -----
174+
175+
module {
176+
// expected-remark @below {{cf_single_branch}}
177+
tt.func @cf_single_branch(%cond: i1) {
178+
// expected-remark @below {{scope id = 0}}
179+
// expected-remark @below {{scope parent id = -1}}
180+
proton.record start "name0"
181+
cf.cond_br %cond, ^then, ^else
182+
^then: // pred: ^entry
102183
// expected-remark @below {{scope id = 0}}
184+
// expected-remark @below {{scope parent id = -1}}
103185
proton.record end "name0"
104-
scf.if %cond {
186+
cf.br ^merge
187+
^else: // pred: ^entry
188+
cf.br ^merge
189+
^merge: // preds: ^then, ^else
190+
tt.return
191+
}
192+
}
193+
194+
195+
// -----
196+
197+
module {
198+
// expected-remark @below {{warp_specialize_balanced}}
199+
tt.func @warp_specialize_balanced() {
200+
// expected-remark @below {{scope id = 0}}
201+
// expected-remark @below {{scope parent id = -1}}
202+
proton.record start "outer"
203+
ttg.warp_specialize()
204+
default {
105205
// expected-remark @below {{scope id = 1}}
106-
proton.record start "name0"
206+
// expected-remark @below {{scope parent id = -1}}
207+
proton.record start "default"
107208
// expected-remark @below {{scope id = 1}}
108-
proton.record end "name0"
209+
// expected-remark @below {{scope parent id = -1}}
210+
proton.record end "default"
211+
ttg.warp_yield
109212
}
213+
partition0() num_warps(1) {
214+
// expected-remark @below {{scope id = 2}}
215+
// expected-remark @below {{scope parent id = -1}}
216+
proton.record start "partition"
217+
// expected-remark @below {{scope id = 2}}
218+
// expected-remark @below {{scope parent id = -1}}
219+
proton.record end "partition"
220+
ttg.warp_return
221+
} : () -> ()
222+
// expected-remark @below {{scope id = 0}}
223+
// expected-remark @below {{scope parent id = -1}}
224+
proton.record end "outer"
225+
tt.return
226+
}
227+
}
228+
229+
// -----
230+
231+
module {
232+
// expected-remark @below {{cf_loop_closed}}
233+
tt.func @cf_loop_closed() {
234+
^entry:
235+
%c0 = arith.constant 0 : index
236+
cf.br ^loop(%c0 : index)
237+
^exit:
238+
tt.return
239+
^loop(%iv: index):
240+
// expected-remark @below {{scope id = 0}}
241+
// expected-remark @below {{scope parent id = -1}}
242+
proton.record start "loop_body"
243+
%c1 = arith.constant 1 : index
244+
%next = arith.addi %iv, %c1 : index
245+
%c2 = arith.constant 2 : index
246+
%cond = arith.cmpi ult, %next, %c2: index
247+
// expected-remark @below {{scope id = 0}}
248+
// expected-remark @below {{scope parent id = -1}}
249+
proton.record end "loop_body"
250+
cf.cond_br %cond, ^loop(%next : index), ^exit
251+
}
252+
}
253+
254+
// -----
255+
256+
module {
257+
// expected-remark @below {{cf_loop_closed_two_blocks}}
258+
tt.func @cf_loop_closed_two_blocks() {
259+
^entry:
260+
%c0 = arith.constant 0 : index
261+
cf.br ^loop(%c0 : index)
262+
^exit:
263+
tt.return
264+
^loop(%iv: index):
265+
// expected-remark @below {{scope id = 0}}
266+
// expected-remark @below {{scope parent id = -1}}
267+
proton.record start "loop_body"
268+
%c1 = arith.constant 1 : index
269+
%next = arith.addi %iv, %c1 : index
270+
cf.br ^loop_body(%next : index)
271+
^loop_body(%iv_next: index):
272+
%c2 = arith.constant 2 : index
273+
%cond = arith.cmpi ult, %iv_next, %c2: index
274+
// expected-remark @below {{scope id = 0}}
275+
// expected-remark @below {{scope parent id = -1}}
276+
proton.record end "loop_body"
277+
cf.cond_br %cond, ^loop(%iv_next : index), ^exit
278+
}
279+
}
280+
281+
// -----
282+
283+
module {
284+
tt.func @cf_unclosed() {
285+
// expected-error @below {{The scope name 'unclosed' is not properly closed (missing end record)}}
286+
proton.record start "unclosed"
287+
tt.return
288+
}
289+
}
290+
291+
// -----
292+
293+
module {
294+
tt.func @cf_dangling_end() {
295+
// expected-error @below {{The scope name 'dangling' is closed without being opened}}
296+
proton.record end "dangling"
297+
tt.return
298+
}
299+
}
300+
301+
// -----
302+
303+
module {
304+
tt.func @cf_liveness_error(%cond: i1) {
305+
proton.record start "name0"
306+
cf.cond_br %cond, ^then, ^else
307+
^then: // pred: ^entry
308+
proton.record end "name0"
309+
cf.br ^merge
310+
^else: // pred: ^entry
311+
// expected-error @below {{The scope name 'name0' is not properly closed (missing start record)}}
312+
proton.record end "name0"
313+
cf.br ^merge
314+
^merge: // preds: ^then, ^else
315+
tt.return
316+
}
317+
}
318+
319+
// -----
320+
321+
module {
322+
tt.func @cf_branch_unclosed_dangling(%cond: i1) {
323+
cf.cond_br %cond, ^then, ^else
324+
^then: // pred: ^entry
325+
proton.record start "ghost"
326+
cf.br ^merge
327+
^else: // pred: ^entry
328+
// expected-error @below {{The scope name 'ghost' is closed without being opened}}
329+
proton.record end "ghost"
330+
cf.br ^merge
331+
^merge: // preds: ^then, ^else
332+
tt.return
333+
}
334+
}
335+
336+
// -----
337+
338+
module {
339+
tt.func @cf_merge_unclosed(%cond: i1) {
340+
cf.br ^start(%cond : i1)
341+
^start(%cond_arg: i1):
342+
proton.record start "ghost"
343+
cf.cond_br %cond_arg, ^then, ^else
344+
^then: // pred: ^start
345+
proton.record end "ghost"
346+
cf.br ^merge
347+
^else: // pred: ^start
348+
proton.record start "ghost"
349+
cf.br ^merge
350+
^merge: // preds: ^then, ^else
351+
proton.record end "ghost"
352+
tt.return
353+
}
354+
}
355+
356+
// -----
357+
358+
module {
359+
tt.func @cf_loop_unclosed() {
360+
%c0 = arith.constant 0 : index
361+
cf.br ^loop(%c0 : index)
362+
^exit:
363+
tt.return
364+
^loop(%iv: index):
365+
// expected-error @below {{The scope name 'loop' is started without being closed}}
366+
proton.record start "loop"
367+
%c1 = arith.constant 1 : index
368+
%next = arith.addi %iv, %c1 : index
369+
%c2 = arith.constant 2 : index
370+
%cond = arith.cmpi ult, %next, %c2: index
371+
cf.cond_br %cond, ^loop(%next : index), ^exit
372+
}
373+
}
374+
375+
// -----
376+
377+
module {
378+
tt.func @cf_loop_end_before_start() {
379+
%c0 = arith.constant 0 : index
380+
cf.br ^loop(%c0 : index)
381+
^exit:
110382
tt.return
383+
^loop(%iv: index):
384+
// expected-error @below {{The scope name 'loop' has end record that dominates its start record}}
385+
proton.record end "loop"
386+
%c1 = arith.constant 1 : index
387+
%next = arith.addi %iv, %c1 : index
388+
%c2 = arith.constant 2 : index
389+
%cond = arith.cmpi ult, %next, %c2: index
390+
proton.record start "loop"
391+
cf.cond_br %cond, ^loop(%next : index), ^exit
111392
}
112393
}

0 commit comments

Comments
 (0)