xref: /llvm-project/mlir/test/Target/SPIRV/loop.mlir (revision bdf00e2216280edef1ec91ccc07987db92197b59)
1// RUN: mlir-translate -no-implicit-module -split-input-file -test-spirv-roundtrip %s | FileCheck %s
2
3// Single loop
4
5spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
6  // for (int i = 0; i < count; ++i) {}
7// CHECK-LABEL: @loop
8  spirv.func @loop(%count : i32) -> () "None" {
9    %zero = spirv.Constant 0: i32
10    %one = spirv.Constant 1: i32
11    %var = spirv.Variable init(%zero) : !spirv.ptr<i32, Function>
12
13// CHECK:        spirv.Branch ^bb1
14// CHECK-NEXT: ^bb1:
15// CHECK-NEXT:   spirv.mlir.loop
16    spirv.mlir.loop {
17// CHECK-NEXT:     spirv.Branch ^bb1
18      spirv.Branch ^header
19
20// CHECK-NEXT:   ^bb1:
21    ^header:
22// CHECK-NEXT:     spirv.Load
23      %val0 = spirv.Load "Function" %var : i32
24// CHECK-NEXT:     spirv.SLessThan
25      %cmp = spirv.SLessThan %val0, %count : i32
26// CHECK-NEXT:     spirv.BranchConditional %{{.*}} [1, 1], ^bb2, ^bb4
27      spirv.BranchConditional %cmp [1, 1], ^body, ^merge
28
29// CHECK-NEXT:   ^bb2:
30    ^body:
31      // Do nothing
32// CHECK-NEXT:     spirv.Branch ^bb3
33      spirv.Branch ^continue
34
35// CHECK-NEXT:   ^bb3:
36    ^continue:
37// CHECK-NEXT:     spirv.Load
38      %val1 = spirv.Load "Function" %var : i32
39// CHECK-NEXT:     spirv.Constant 1
40// CHECK-NEXT:     spirv.IAdd
41      %add = spirv.IAdd %val1, %one : i32
42// CHECK-NEXT:     spirv.Store
43      spirv.Store "Function" %var, %add : i32
44// CHECK-NEXT:     spirv.Branch ^bb1
45      spirv.Branch ^header
46
47// CHECK-NEXT:   ^bb4:
48// CHECK-NEXT:     spirv.mlir.merge
49    ^merge:
50      spirv.mlir.merge
51    }
52    spirv.Return
53  }
54
55  spirv.func @main() -> () "None" {
56    spirv.Return
57  }
58  spirv.EntryPoint "GLCompute" @main
59}
60
61// -----
62
63// Single loop with block arguments
64
65spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
66  spirv.GlobalVariable @GV1 bind(0, 0) : !spirv.ptr<!spirv.struct<(!spirv.array<10 x f32, stride=4> [0])>, StorageBuffer>
67  spirv.GlobalVariable @GV2 bind(0, 1) : !spirv.ptr<!spirv.struct<(!spirv.array<10 x f32, stride=4> [0])>, StorageBuffer>
68// CHECK-LABEL: @loop_kernel
69  spirv.func @loop_kernel() "None" {
70    %0 = spirv.mlir.addressof @GV1 : !spirv.ptr<!spirv.struct<(!spirv.array<10 x f32, stride=4> [0])>, StorageBuffer>
71    %1 = spirv.Constant 0 : i32
72    %2 = spirv.AccessChain %0[%1] : !spirv.ptr<!spirv.struct<(!spirv.array<10 x f32, stride=4> [0])>, StorageBuffer>, i32 -> !spirv.ptr<!spirv.array<10 x f32, stride=4>, StorageBuffer>
73    %3 = spirv.mlir.addressof @GV2 : !spirv.ptr<!spirv.struct<(!spirv.array<10 x f32, stride=4> [0])>, StorageBuffer>
74    %5 = spirv.AccessChain %3[%1] : !spirv.ptr<!spirv.struct<(!spirv.array<10 x f32, stride=4> [0])>, StorageBuffer>, i32 -> !spirv.ptr<!spirv.array<10 x f32, stride=4>, StorageBuffer>
75    %6 = spirv.Constant 4 : i32
76    %7 = spirv.Constant 42 : i32
77    %8 = spirv.Constant 2 : i32
78// CHECK:        spirv.Branch ^bb1(%{{.*}} : i32)
79// CHECK-NEXT: ^bb1(%[[OUTARG:.*]]: i32):
80// CHECK-NEXT:   spirv.mlir.loop {
81    spirv.mlir.loop {
82// CHECK-NEXT:     spirv.Branch ^bb1(%[[OUTARG]] : i32)
83      spirv.Branch ^header(%6 : i32)
84// CHECK-NEXT:   ^bb1(%[[HEADARG:.*]]: i32):
85    ^header(%9: i32):
86      %10 = spirv.SLessThan %9, %7 : i32
87// CHECK:          spirv.BranchConditional %{{.*}}, ^bb2, ^bb3
88      spirv.BranchConditional %10, ^body, ^merge
89// CHECK-NEXT:   ^bb2:     // pred: ^bb1
90    ^body:
91      %11 = spirv.AccessChain %2[%9] : !spirv.ptr<!spirv.array<10 x f32, stride=4>, StorageBuffer>, i32 -> !spirv.ptr<f32, StorageBuffer>
92      %12 = spirv.Load "StorageBuffer" %11 : f32
93      %13 = spirv.AccessChain %5[%9] : !spirv.ptr<!spirv.array<10 x f32, stride=4>, StorageBuffer>, i32 -> !spirv.ptr<f32, StorageBuffer>
94      spirv.Store "StorageBuffer" %13, %12 : f32
95// CHECK:          %[[ADD:.*]] = spirv.IAdd
96      %14 = spirv.IAdd %9, %8 : i32
97// CHECK-NEXT:     spirv.Branch ^bb1(%[[ADD]] : i32)
98      spirv.Branch ^header(%14 : i32)
99// CHECK-NEXT:   ^bb3:
100    ^merge:
101// CHECK-NEXT:     spirv.mlir.merge
102      spirv.mlir.merge
103    }
104    spirv.Return
105  }
106  spirv.EntryPoint "GLCompute" @loop_kernel
107  spirv.ExecutionMode @loop_kernel "LocalSize", 1, 1, 1
108}
109
110// -----
111
112// Nested loop
113
114spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
115  // for (int i = 0; i < count; ++i) {
116  //   for (int j = 0; j < count; ++j) { }
117  // }
118// CHECK-LABEL: @loop
119  spirv.func @loop(%count : i32) -> () "None" {
120    %zero = spirv.Constant 0: i32
121    %one = spirv.Constant 1: i32
122    %ivar = spirv.Variable init(%zero) : !spirv.ptr<i32, Function>
123    %jvar = spirv.Variable init(%zero) : !spirv.ptr<i32, Function>
124
125// CHECK:        spirv.Branch ^bb1
126// CHECK-NEXT: ^bb1:
127// CHECK-NEXT:   spirv.mlir.loop control(Unroll)
128    spirv.mlir.loop control(Unroll) {
129// CHECK-NEXT:     spirv.Branch ^bb1
130      spirv.Branch ^header
131
132// CHECK-NEXT:   ^bb1:
133    ^header:
134// CHECK-NEXT:     spirv.Load
135      %ival0 = spirv.Load "Function" %ivar : i32
136// CHECK-NEXT:     spirv.SLessThan
137      %icmp = spirv.SLessThan %ival0, %count : i32
138// CHECK-NEXT:     spirv.BranchConditional %{{.*}}, ^bb2, ^bb5
139      spirv.BranchConditional %icmp, ^body, ^merge
140
141// CHECK-NEXT:   ^bb2:
142    ^body:
143// CHECK-NEXT:     spirv.Constant 0
144// CHECK-NEXT: 		 spirv.Store
145      spirv.Store "Function" %jvar, %zero : i32
146// CHECK-NEXT:     spirv.Branch ^bb3
147// CHECK-NEXT:   ^bb3:
148// CHECK-NEXT:     spirv.mlir.loop control(DontUnroll)
149      spirv.mlir.loop control(DontUnroll) {
150// CHECK-NEXT:       spirv.Branch ^bb1
151        spirv.Branch ^header
152
153// CHECK-NEXT:     ^bb1:
154      ^header:
155// CHECK-NEXT:       spirv.Load
156        %jval0 = spirv.Load "Function" %jvar : i32
157// CHECK-NEXT:       spirv.SLessThan
158        %jcmp = spirv.SLessThan %jval0, %count : i32
159// CHECK-NEXT:       spirv.BranchConditional %{{.*}}, ^bb2, ^bb4
160        spirv.BranchConditional %jcmp, ^body, ^merge
161
162// CHECK-NEXT:     ^bb2:
163      ^body:
164        // Do nothing
165// CHECK-NEXT:       spirv.Branch ^bb3
166        spirv.Branch ^continue
167
168// CHECK-NEXT:     ^bb3:
169      ^continue:
170// CHECK-NEXT:       spirv.Load
171        %jval1 = spirv.Load "Function" %jvar : i32
172// CHECK-NEXT:       spirv.Constant 1
173// CHECK-NEXT:       spirv.IAdd
174        %add = spirv.IAdd %jval1, %one : i32
175// CHECK-NEXT:       spirv.Store
176        spirv.Store "Function" %jvar, %add : i32
177// CHECK-NEXT:       spirv.Branch ^bb1
178        spirv.Branch ^header
179
180// CHECK-NEXT:     ^bb4:
181      ^merge:
182// CHECK-NEXT:       spirv.mlir.merge
183        spirv.mlir.merge
184      } // end inner loop
185
186// CHECK:          spirv.Branch ^bb4
187      spirv.Branch ^continue
188
189// CHECK-NEXT:   ^bb4:
190    ^continue:
191// CHECK-NEXT:     spirv.Load
192      %ival1 = spirv.Load "Function" %ivar : i32
193// CHECK-NEXT:     spirv.Constant 1
194// CHECK-NEXT:     spirv.IAdd
195      %add = spirv.IAdd %ival1, %one : i32
196// CHECK-NEXT:     spirv.Store
197      spirv.Store "Function" %ivar, %add : i32
198// CHECK-NEXT:     spirv.Branch ^bb1
199      spirv.Branch ^header
200
201// CHECK-NEXT:   ^bb5:
202// CHECK-NEXT:     spirv.mlir.merge
203    ^merge:
204      spirv.mlir.merge
205    } // end outer loop
206    spirv.Return
207  }
208
209  spirv.func @main() -> () "None" {
210    spirv.Return
211  }
212  spirv.EntryPoint "GLCompute" @main
213}
214
215
216// -----
217
218// Loop with selection in its header
219
220spirv.module Physical64 OpenCL requires #spirv.vce<v1.0, [Kernel, Linkage, Addresses, Int64], []> {
221// CHECK-LABEL:   @kernel
222// CHECK-SAME:    (%[[INPUT0:.+]]: i64)
223  spirv.func @kernel(%input: i64) "None" {
224// CHECK-NEXT:     %[[VAR:.+]] = spirv.Variable : !spirv.ptr<i1, Function>
225// CHECK-NEXT:     spirv.Branch ^[[BB0:.+]](%[[INPUT0]] : i64)
226// CHECK-NEXT:   ^[[BB0]](%[[INPUT1:.+]]: i64):
227    %cst0_i64 = spirv.Constant 0 : i64
228    %true = spirv.Constant true
229    %false = spirv.Constant false
230// CHECK-NEXT:     spirv.mlir.loop {
231    spirv.mlir.loop {
232// CHECK-NEXT:       spirv.Branch ^[[LOOP_HEADER:.+]](%[[INPUT1]] : i64)
233      spirv.Branch ^loop_header(%input : i64)
234// CHECK-NEXT:     ^[[LOOP_HEADER]](%[[ARG1:.+]]: i64):
235    ^loop_header(%arg1: i64):
236// CHECK-NEXT:       spirv.Branch ^[[LOOP_BODY:.+]]
237// CHECK-NEXT:     ^[[LOOP_BODY]]:
238// CHECK-NEXT:         %[[C0:.+]] = spirv.Constant 0 : i64
239      %gt = spirv.SGreaterThan %arg1, %cst0_i64 : i64
240// CHECK-NEXT:         %[[GT:.+]] = spirv.SGreaterThan %[[ARG1]], %[[C0]] : i64
241// CHECK-NEXT:         spirv.Branch ^[[BB1:.+]]
242// CHECK-NEXT:     ^[[BB1]]:
243      %var = spirv.Variable : !spirv.ptr<i1, Function>
244// CHECK-NEXT:       spirv.mlir.selection {
245      spirv.mlir.selection {
246// CHECK-NEXT:         spirv.BranchConditional %[[GT]], ^[[THEN:.+]], ^[[ELSE:.+]]
247        spirv.BranchConditional %gt, ^then, ^else
248// CHECK-NEXT:       ^[[THEN]]:
249      ^then:
250// CHECK-NEXT:         %true = spirv.Constant true
251// CHECK-NEXT:         spirv.Store "Function" %[[VAR]], %true : i1
252        spirv.Store "Function" %var, %true : i1
253// CHECK-NEXT:         spirv.Branch ^[[SELECTION_MERGE:.+]]
254        spirv.Branch ^selection_merge
255// CHECK-NEXT:       ^[[ELSE]]:
256      ^else:
257// CHECK-NEXT:         %false = spirv.Constant false
258// CHECK-NEXT:         spirv.Store "Function" %[[VAR]], %false : i1
259        spirv.Store "Function" %var, %false : i1
260// CHECK-NEXT:         spirv.Branch ^[[SELECTION_MERGE]]
261        spirv.Branch ^selection_merge
262// CHECK-NEXT:       ^[[SELECTION_MERGE]]:
263      ^selection_merge:
264// CHECK-NEXT:         spirv.mlir.merge
265        spirv.mlir.merge
266// CHECK-NEXT:       }
267      }
268// CHECK-NEXT:       %[[LOAD:.+]] = spirv.Load "Function" %[[VAR]] : i1
269      %load = spirv.Load "Function" %var : i1
270// CHECK-NEXT:       spirv.BranchConditional %[[LOAD]], ^[[CONTINUE:.+]](%[[ARG1]] : i64), ^[[LOOP_MERGE:.+]]
271      spirv.BranchConditional %load, ^continue(%arg1 : i64), ^loop_merge
272// CHECK-NEXT:     ^[[CONTINUE]](%[[ARG2:.+]]: i64):
273    ^continue(%arg2: i64):
274// CHECK-NEXT:       %[[C0:.+]] = spirv.Constant 0 : i64
275// CHECK-NEXT:       %[[LT:.+]] = spirv.SLessThan %[[ARG2]], %[[C0]] : i64
276      %lt = spirv.SLessThan %arg2, %cst0_i64 : i64
277// CHECK-NEXT:       spirv.Store "Function" %[[VAR]], %[[LT]] : i1
278      spirv.Store "Function" %var, %lt : i1
279// CHECK-NEXT:       spirv.Branch ^[[LOOP_HEADER]](%[[ARG2]] : i64)
280      spirv.Branch ^loop_header(%arg2 : i64)
281// CHECK-NEXT:     ^[[LOOP_MERGE]]:
282    ^loop_merge:
283// CHECK-NEXT:       spirv.mlir.merge
284      spirv.mlir.merge
285// CHECK-NEXT:     }
286    }
287// CHECK-NEXT:     spirv.Return
288    spirv.Return
289  }
290}
291