1// RUN: mlir-opt -split-input-file -convert-gpu-to-spirv %s -o - | FileCheck %s
2
3module attributes {gpu.container_module} {
4  func @builtin() {
5    %c0 = constant 1 : index
6    gpu.launch_func @kernels::@builtin_workgroup_id_x
7        blocks in (%c0, %c0, %c0) threads in (%c0, %c0, %c0)
8    return
9  }
10
11  // CHECK-LABEL:  spv.module @{{.*}} Logical GLSL450
12  // CHECK: spv.globalVariable [[WORKGROUPID:@.*]] built_in("WorkgroupId")
13  gpu.module @kernels {
14    gpu.func @builtin_workgroup_id_x() kernel
15      attributes {spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} {
16      // CHECK: [[ADDRESS:%.*]] = spv.mlir.addressof [[WORKGROUPID]]
17      // CHECK-NEXT: [[VEC:%.*]] = spv.Load "Input" [[ADDRESS]]
18      // CHECK-NEXT: {{%.*}} = spv.CompositeExtract [[VEC]]{{\[}}0 : i32{{\]}}
19      %0 = "gpu.block_id"() {dimension = "x"} : () -> index
20      gpu.return
21    }
22  }
23}
24
25// -----
26
27module attributes {gpu.container_module} {
28  func @builtin() {
29    %c0 = constant 1 : index
30    gpu.launch_func @kernels::@builtin_workgroup_id_y
31        blocks in (%c0, %c0, %c0) threads in (%c0, %c0, %c0)
32    return
33  }
34
35  // CHECK-LABEL:  spv.module @{{.*}} Logical GLSL450
36  // CHECK: spv.globalVariable [[WORKGROUPID:@.*]] built_in("WorkgroupId")
37  gpu.module @kernels {
38    gpu.func @builtin_workgroup_id_y() kernel
39      attributes {spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} {
40      // CHECK: [[ADDRESS:%.*]] = spv.mlir.addressof [[WORKGROUPID]]
41      // CHECK-NEXT: [[VEC:%.*]] = spv.Load "Input" [[ADDRESS]]
42      // CHECK-NEXT: {{%.*}} = spv.CompositeExtract [[VEC]]{{\[}}1 : i32{{\]}}
43      %0 = "gpu.block_id"() {dimension = "y"} : () -> index
44      gpu.return
45    }
46  }
47}
48
49// -----
50
51module attributes {gpu.container_module} {
52  func @builtin() {
53    %c0 = constant 1 : index
54    gpu.launch_func @kernels::@builtin_workgroup_id_z
55        blocks in (%c0, %c0, %c0) threads in (%c0, %c0, %c0)
56    return
57  }
58
59  // CHECK-LABEL:  spv.module @{{.*}} Logical GLSL450
60  // CHECK: spv.globalVariable [[WORKGROUPID:@.*]] built_in("WorkgroupId")
61  gpu.module @kernels {
62    gpu.func @builtin_workgroup_id_z() kernel
63      attributes {spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} {
64      // CHECK: [[ADDRESS:%.*]] = spv.mlir.addressof [[WORKGROUPID]]
65      // CHECK-NEXT: [[VEC:%.*]] = spv.Load "Input" [[ADDRESS]]
66      // CHECK-NEXT: {{%.*}} = spv.CompositeExtract [[VEC]]{{\[}}2 : i32{{\]}}
67      %0 = "gpu.block_id"() {dimension = "z"} : () -> index
68      gpu.return
69    }
70  }
71}
72
73// -----
74
75module attributes {gpu.container_module} {
76  func @builtin() {
77    %c0 = constant 1 : index
78    gpu.launch_func @kernels::@builtin_workgroup_size_x
79        blocks in (%c0, %c0, %c0) threads in (%c0, %c0, %c0)
80    return
81  }
82
83  // CHECK-LABEL:  spv.module @{{.*}} Logical GLSL450
84  gpu.module @kernels {
85    gpu.func @builtin_workgroup_size_x() kernel
86      attributes {spv.entry_point_abi = {local_size = dense<[32, 1, 1]>: vector<3xi32>}} {
87      // The constant value is obtained from the spv.entry_point_abi.
88      // Note that this ignores the workgroup size specification in gpu.launch.
89      // We may want to define gpu.workgroup_size and convert it to the entry
90      // point ABI we want here.
91      // CHECK: spv.constant 32 : i32
92      %0 = "gpu.block_dim"() {dimension = "x"} : () -> index
93      gpu.return
94    }
95  }
96}
97
98// -----
99
100module attributes {gpu.container_module} {
101  func @builtin() {
102    %c0 = constant 1 : index
103    gpu.launch_func @kernels::@builtin_workgroup_size_y
104        blocks in (%c0, %c0, %c0) threads in (%c0, %c0, %c0)
105    return
106  }
107
108  // CHECK-LABEL:  spv.module @{{.*}} Logical GLSL450
109  gpu.module @kernels {
110    gpu.func @builtin_workgroup_size_y() kernel
111      attributes {spv.entry_point_abi = {local_size = dense<[32, 4, 1]>: vector<3xi32>}} {
112      // The constant value is obtained from the spv.entry_point_abi.
113      // CHECK: spv.constant 4 : i32
114      %0 = "gpu.block_dim"() {dimension = "y"} : () -> index
115      gpu.return
116    }
117  }
118}
119
120// -----
121
122module attributes {gpu.container_module} {
123  func @builtin() {
124    %c0 = constant 1 : index
125    gpu.launch_func @kernels::@builtin_workgroup_size_z
126        blocks in (%c0, %c0, %c0) threads in (%c0, %c0, %c0)
127    return
128  }
129
130  // CHECK-LABEL:  spv.module @{{.*}} Logical GLSL450
131  gpu.module @kernels {
132    gpu.func @builtin_workgroup_size_z() kernel
133      attributes {spv.entry_point_abi = {local_size = dense<[32, 4, 1]>: vector<3xi32>}} {
134      // The constant value is obtained from the spv.entry_point_abi.
135      // CHECK: spv.constant 1 : i32
136      %0 = "gpu.block_dim"() {dimension = "z"} : () -> index
137      gpu.return
138    }
139  }
140}
141
142// -----
143
144module attributes {gpu.container_module} {
145  func @builtin() {
146    %c0 = constant 1 : index
147    gpu.launch_func @kernels::@builtin_local_id_x
148        blocks in (%c0, %c0, %c0) threads in (%c0, %c0, %c0)
149    return
150  }
151
152  // CHECK-LABEL:  spv.module @{{.*}} Logical GLSL450
153  // CHECK: spv.globalVariable [[LOCALINVOCATIONID:@.*]] built_in("LocalInvocationId")
154  gpu.module @kernels {
155    gpu.func @builtin_local_id_x() kernel
156      attributes {spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} {
157      // CHECK: [[ADDRESS:%.*]] = spv.mlir.addressof [[LOCALINVOCATIONID]]
158      // CHECK-NEXT: [[VEC:%.*]] = spv.Load "Input" [[ADDRESS]]
159      // CHECK-NEXT: {{%.*}} = spv.CompositeExtract [[VEC]]{{\[}}0 : i32{{\]}}
160      %0 = "gpu.thread_id"() {dimension = "x"} : () -> index
161      gpu.return
162    }
163  }
164}
165
166// -----
167
168module attributes {gpu.container_module} {
169  func @builtin() {
170    %c0 = constant 1 : index
171    gpu.launch_func @kernels::@builtin_num_workgroups_x
172        blocks in (%c0, %c0, %c0) threads in (%c0, %c0, %c0)
173    return
174  }
175
176  // CHECK-LABEL:  spv.module @{{.*}} Logical GLSL450
177  // CHECK: spv.globalVariable [[NUMWORKGROUPS:@.*]] built_in("NumWorkgroups")
178  gpu.module @kernels {
179    gpu.func @builtin_num_workgroups_x() kernel
180      attributes {spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} {
181      // CHECK: [[ADDRESS:%.*]] = spv.mlir.addressof [[NUMWORKGROUPS]]
182      // CHECK-NEXT: [[VEC:%.*]] = spv.Load "Input" [[ADDRESS]]
183      // CHECK-NEXT: {{%.*}} = spv.CompositeExtract [[VEC]]{{\[}}0 : i32{{\]}}
184      %0 = "gpu.grid_dim"() {dimension = "x"} : () -> index
185      gpu.return
186    }
187  }
188}
189
190// -----
191
192module attributes {gpu.container_module} {
193  // CHECK-LABEL:  spv.module @{{.*}} Logical GLSL450
194  // CHECK: spv.globalVariable [[SUBGROUPID:@.*]] built_in("SubgroupId")
195  gpu.module @kernels {
196    gpu.func @builtin_subgroup_id() kernel
197      attributes {spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} {
198      // CHECK: [[ADDRESS:%.*]] = spv.mlir.addressof [[SUBGROUPID]]
199      // CHECK-NEXT: {{%.*}} = spv.Load "Input" [[ADDRESS]]
200      %0 = gpu.subgroup_id : index
201      gpu.return
202    }
203  }
204}
205
206// -----
207
208module attributes {gpu.container_module} {
209  // CHECK-LABEL:  spv.module @{{.*}} Logical GLSL450
210  // CHECK: spv.globalVariable [[NUMSUBGROUPS:@.*]] built_in("NumSubgroups")
211  gpu.module @kernels {
212    gpu.func @builtin_num_subgroups() kernel
213      attributes {spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} {
214      // CHECK: [[ADDRESS:%.*]] = spv.mlir.addressof [[NUMSUBGROUPS]]
215      // CHECK-NEXT: {{%.*}} = spv.Load "Input" [[ADDRESS]]
216      %0 = gpu.num_subgroups : index
217      gpu.return
218    }
219  }
220}
221
222// -----
223
224module attributes {gpu.container_module} {
225  // CHECK-LABEL:  spv.module @{{.*}} Logical GLSL450
226  // CHECK: spv.globalVariable [[SUBGROUPSIZE:@.*]] built_in("SubgroupSize")
227  gpu.module @kernels {
228    gpu.func @builtin_subgroup_size() kernel
229      attributes {spv.entry_point_abi = {local_size = dense<[16, 1, 1]>: vector<3xi32>}} {
230      // CHECK: [[ADDRESS:%.*]] = spv.mlir.addressof [[SUBGROUPSIZE]]
231      // CHECK-NEXT: {{%.*}} = spv.Load "Input" [[ADDRESS]]
232      %0 = gpu.subgroup_size : index
233      gpu.return
234    }
235  }
236}
237