anv: Use nir_src_is_const and friends in lowering code
[mesa.git] / src / intel / vulkan / anv_nir_apply_pipeline_layout.c
1 /*
2 * Copyright © 2015 Intel Corporation
3 *
4 * Permission is hereby granted, free of charge, to any person obtaining a
5 * copy of this software and associated documentation files (the "Software"),
6 * to deal in the Software without restriction, including without limitation
7 * the rights to use, copy, modify, merge, publish, distribute, sublicense,
8 * and/or sell copies of the Software, and to permit persons to whom the
9 * Software is furnished to do so, subject to the following conditions:
10 *
11 * The above copyright notice and this permission notice (including the next
12 * paragraph) shall be included in all copies or substantial portions of the
13 * Software.
14 *
15 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
18 * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
20 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
21 * IN THE SOFTWARE.
22 */
23
24 #include "anv_nir.h"
25 #include "program/prog_parameter.h"
26 #include "nir/nir_builder.h"
27 #include "compiler/brw_nir.h"
28
29 struct apply_pipeline_layout_state {
30 nir_shader *shader;
31 nir_builder builder;
32
33 struct anv_pipeline_layout *layout;
34 bool add_bounds_checks;
35
36 unsigned first_image_uniform;
37
38 bool uses_constants;
39 uint8_t constants_offset;
40 struct {
41 BITSET_WORD *used;
42 uint8_t *surface_offsets;
43 uint8_t *sampler_offsets;
44 uint8_t *image_offsets;
45 } set[MAX_SETS];
46 };
47
48 static void
49 add_binding(struct apply_pipeline_layout_state *state,
50 uint32_t set, uint32_t binding)
51 {
52 BITSET_SET(state->set[set].used, binding);
53 }
54
55 static void
56 add_var_binding(struct apply_pipeline_layout_state *state, nir_variable *var)
57 {
58 add_binding(state, var->data.descriptor_set, var->data.binding);
59 }
60
61 static void
62 add_deref_src_binding(struct apply_pipeline_layout_state *state, nir_src src)
63 {
64 nir_deref_instr *deref = nir_src_as_deref(src);
65 add_var_binding(state, nir_deref_instr_get_variable(deref));
66 }
67
68 static void
69 add_tex_src_binding(struct apply_pipeline_layout_state *state,
70 nir_tex_instr *tex, nir_tex_src_type deref_src_type)
71 {
72 int deref_src_idx = nir_tex_instr_src_index(tex, deref_src_type);
73 if (deref_src_idx < 0)
74 return;
75
76 add_deref_src_binding(state, tex->src[deref_src_idx].src);
77 }
78
79 static void
80 get_used_bindings_block(nir_block *block,
81 struct apply_pipeline_layout_state *state)
82 {
83 nir_foreach_instr_safe(instr, block) {
84 switch (instr->type) {
85 case nir_instr_type_intrinsic: {
86 nir_intrinsic_instr *intrin = nir_instr_as_intrinsic(instr);
87 switch (intrin->intrinsic) {
88 case nir_intrinsic_vulkan_resource_index:
89 add_binding(state, nir_intrinsic_desc_set(intrin),
90 nir_intrinsic_binding(intrin));
91 break;
92
93 case nir_intrinsic_image_deref_load:
94 case nir_intrinsic_image_deref_store:
95 case nir_intrinsic_image_deref_atomic_add:
96 case nir_intrinsic_image_deref_atomic_min:
97 case nir_intrinsic_image_deref_atomic_max:
98 case nir_intrinsic_image_deref_atomic_and:
99 case nir_intrinsic_image_deref_atomic_or:
100 case nir_intrinsic_image_deref_atomic_xor:
101 case nir_intrinsic_image_deref_atomic_exchange:
102 case nir_intrinsic_image_deref_atomic_comp_swap:
103 case nir_intrinsic_image_deref_size:
104 case nir_intrinsic_image_deref_samples:
105 case nir_intrinsic_image_deref_load_param_intel:
106 case nir_intrinsic_image_deref_load_raw_intel:
107 case nir_intrinsic_image_deref_store_raw_intel:
108 add_deref_src_binding(state, intrin->src[0]);
109 break;
110
111 case nir_intrinsic_load_constant:
112 state->uses_constants = true;
113 break;
114
115 default:
116 break;
117 }
118 break;
119 }
120 case nir_instr_type_tex: {
121 nir_tex_instr *tex = nir_instr_as_tex(instr);
122 add_tex_src_binding(state, tex, nir_tex_src_texture_deref);
123 add_tex_src_binding(state, tex, nir_tex_src_sampler_deref);
124 break;
125 }
126 default:
127 continue;
128 }
129 }
130 }
131
132 static void
133 lower_res_index_intrinsic(nir_intrinsic_instr *intrin,
134 struct apply_pipeline_layout_state *state)
135 {
136 nir_builder *b = &state->builder;
137
138 b->cursor = nir_before_instr(&intrin->instr);
139
140 uint32_t set = nir_intrinsic_desc_set(intrin);
141 uint32_t binding = nir_intrinsic_binding(intrin);
142
143 uint32_t surface_index = state->set[set].surface_offsets[binding];
144 uint32_t array_size =
145 state->layout->set[set].layout->binding[binding].array_size;
146
147 nir_ssa_def *block_index;
148 if (nir_src_is_const(intrin->src[0])) {
149 unsigned array_index = nir_src_as_uint(intrin->src[0]);
150 array_index = MIN2(array_index, array_size - 1);
151 block_index = nir_imm_int(b, surface_index + array_index);
152 } else {
153 block_index = nir_ssa_for_src(b, intrin->src[0], 1);
154
155 if (state->add_bounds_checks)
156 block_index = nir_umin(b, block_index, nir_imm_int(b, array_size - 1));
157
158 block_index = nir_iadd(b, nir_imm_int(b, surface_index), block_index);
159 }
160
161 assert(intrin->dest.is_ssa);
162 nir_ssa_def_rewrite_uses(&intrin->dest.ssa, nir_src_for_ssa(block_index));
163 nir_instr_remove(&intrin->instr);
164 }
165
166 static void
167 lower_res_reindex_intrinsic(nir_intrinsic_instr *intrin,
168 struct apply_pipeline_layout_state *state)
169 {
170 nir_builder *b = &state->builder;
171
172 /* For us, the resource indices are just indices into the binding table and
173 * array elements are sequential. A resource_reindex just turns into an
174 * add of the two indices.
175 */
176 assert(intrin->src[0].is_ssa && intrin->src[1].is_ssa);
177 nir_ssa_def *new_index = nir_iadd(b, intrin->src[0].ssa,
178 intrin->src[1].ssa);
179
180 assert(intrin->dest.is_ssa);
181 nir_ssa_def_rewrite_uses(&intrin->dest.ssa, nir_src_for_ssa(new_index));
182 nir_instr_remove(&intrin->instr);
183 }
184
185 static void
186 lower_image_intrinsic(nir_intrinsic_instr *intrin,
187 struct apply_pipeline_layout_state *state)
188 {
189 nir_deref_instr *deref = nir_src_as_deref(intrin->src[0]);
190 nir_variable *var = nir_deref_instr_get_variable(deref);
191
192 unsigned set = var->data.descriptor_set;
193 unsigned binding = var->data.binding;
194 unsigned array_size =
195 state->layout->set[set].layout->binding[binding].array_size;
196
197 nir_builder *b = &state->builder;
198 b->cursor = nir_before_instr(&intrin->instr);
199
200 nir_ssa_def *index = NULL;
201 if (deref->deref_type != nir_deref_type_var) {
202 assert(deref->deref_type == nir_deref_type_array);
203 index = nir_ssa_for_src(b, deref->arr.index, 1);
204 if (state->add_bounds_checks)
205 index = nir_umin(b, index, nir_imm_int(b, array_size - 1));
206 } else {
207 index = nir_imm_int(b, 0);
208 }
209
210 if (intrin->intrinsic == nir_intrinsic_image_deref_load_param_intel) {
211 b->cursor = nir_instr_remove(&intrin->instr);
212
213 nir_intrinsic_instr *load =
214 nir_intrinsic_instr_create(b->shader, nir_intrinsic_load_uniform);
215
216 nir_intrinsic_set_base(load, state->first_image_uniform +
217 state->set[set].image_offsets[binding] *
218 BRW_IMAGE_PARAM_SIZE * 4);
219 nir_intrinsic_set_range(load, array_size * BRW_IMAGE_PARAM_SIZE * 4);
220
221 const unsigned param = nir_intrinsic_base(intrin);
222 nir_ssa_def *offset =
223 nir_imul(b, index, nir_imm_int(b, BRW_IMAGE_PARAM_SIZE * 4));
224 offset = nir_iadd(b, offset, nir_imm_int(b, param * 16));
225 load->src[0] = nir_src_for_ssa(offset);
226
227 load->num_components = intrin->dest.ssa.num_components;
228 nir_ssa_dest_init(&load->instr, &load->dest,
229 intrin->dest.ssa.num_components,
230 intrin->dest.ssa.bit_size, NULL);
231 nir_builder_instr_insert(b, &load->instr);
232
233 nir_ssa_def_rewrite_uses(&intrin->dest.ssa,
234 nir_src_for_ssa(&load->dest.ssa));
235 } else {
236 unsigned binding_offset = state->set[set].surface_offsets[binding];
237 index = nir_iadd(b, index, nir_imm_int(b, binding_offset));
238 brw_nir_rewrite_image_intrinsic(intrin, index);
239 }
240 }
241
242 static void
243 lower_load_constant(nir_intrinsic_instr *intrin,
244 struct apply_pipeline_layout_state *state)
245 {
246 nir_builder *b = &state->builder;
247
248 b->cursor = nir_before_instr(&intrin->instr);
249
250 nir_ssa_def *index = nir_imm_int(b, state->constants_offset);
251 nir_ssa_def *offset = nir_iadd(b, nir_ssa_for_src(b, intrin->src[0], 1),
252 nir_imm_int(b, nir_intrinsic_base(intrin)));
253
254 nir_intrinsic_instr *load_ubo =
255 nir_intrinsic_instr_create(b->shader, nir_intrinsic_load_ubo);
256 load_ubo->num_components = intrin->num_components;
257 load_ubo->src[0] = nir_src_for_ssa(index);
258 load_ubo->src[1] = nir_src_for_ssa(offset);
259 nir_ssa_dest_init(&load_ubo->instr, &load_ubo->dest,
260 intrin->dest.ssa.num_components,
261 intrin->dest.ssa.bit_size, NULL);
262 nir_builder_instr_insert(b, &load_ubo->instr);
263
264 nir_ssa_def_rewrite_uses(&intrin->dest.ssa,
265 nir_src_for_ssa(&load_ubo->dest.ssa));
266 nir_instr_remove(&intrin->instr);
267 }
268
269 static void
270 lower_tex_deref(nir_tex_instr *tex, nir_tex_src_type deref_src_type,
271 unsigned *base_index,
272 struct apply_pipeline_layout_state *state)
273 {
274 int deref_src_idx = nir_tex_instr_src_index(tex, deref_src_type);
275 if (deref_src_idx < 0)
276 return;
277
278 nir_deref_instr *deref = nir_src_as_deref(tex->src[deref_src_idx].src);
279 nir_variable *var = nir_deref_instr_get_variable(deref);
280
281 unsigned set = var->data.descriptor_set;
282 unsigned binding = var->data.binding;
283 unsigned array_size =
284 state->layout->set[set].layout->binding[binding].array_size;
285
286 nir_tex_src_type offset_src_type;
287 if (deref_src_type == nir_tex_src_texture_deref) {
288 offset_src_type = nir_tex_src_texture_offset;
289 *base_index = state->set[set].surface_offsets[binding];
290 } else {
291 assert(deref_src_type == nir_tex_src_sampler_deref);
292 offset_src_type = nir_tex_src_sampler_offset;
293 *base_index = state->set[set].sampler_offsets[binding];
294 }
295
296 nir_ssa_def *index = NULL;
297 if (deref->deref_type != nir_deref_type_var) {
298 assert(deref->deref_type == nir_deref_type_array);
299
300 if (nir_src_is_const(deref->arr.index)) {
301 unsigned arr_index = nir_src_as_uint(deref->arr.index);
302 *base_index += MIN2(arr_index, array_size - 1);
303 } else {
304 nir_builder *b = &state->builder;
305
306 /* From VK_KHR_sampler_ycbcr_conversion:
307 *
308 * If sampler Y’CBCR conversion is enabled, the combined image
309 * sampler must be indexed only by constant integral expressions when
310 * aggregated into arrays in shader code, irrespective of the
311 * shaderSampledImageArrayDynamicIndexing feature.
312 */
313 assert(nir_tex_instr_src_index(tex, nir_tex_src_plane) == -1);
314
315 index = nir_ssa_for_src(b, deref->arr.index, 1);
316
317 if (state->add_bounds_checks)
318 index = nir_umin(b, index, nir_imm_int(b, array_size - 1));
319 }
320 }
321
322 if (index) {
323 nir_instr_rewrite_src(&tex->instr, &tex->src[deref_src_idx].src,
324 nir_src_for_ssa(index));
325 tex->src[deref_src_idx].src_type = offset_src_type;
326 } else {
327 nir_tex_instr_remove_src(tex, deref_src_idx);
328 }
329 }
330
331 static uint32_t
332 tex_instr_get_and_remove_plane_src(nir_tex_instr *tex)
333 {
334 int plane_src_idx = nir_tex_instr_src_index(tex, nir_tex_src_plane);
335 if (plane_src_idx < 0)
336 return 0;
337
338 unsigned plane = nir_src_as_uint(tex->src[plane_src_idx].src);
339
340 nir_tex_instr_remove_src(tex, plane_src_idx);
341
342 return plane;
343 }
344
345 static void
346 lower_tex(nir_tex_instr *tex, struct apply_pipeline_layout_state *state)
347 {
348 state->builder.cursor = nir_before_instr(&tex->instr);
349
350 unsigned plane = tex_instr_get_and_remove_plane_src(tex);
351
352 lower_tex_deref(tex, nir_tex_src_texture_deref,
353 &tex->texture_index, state);
354 tex->texture_index += plane;
355
356 lower_tex_deref(tex, nir_tex_src_sampler_deref,
357 &tex->sampler_index, state);
358 tex->sampler_index += plane;
359
360 /* The backend only ever uses this to mark used surfaces. We don't care
361 * about that little optimization so it just needs to be non-zero.
362 */
363 tex->texture_array_size = 1;
364 }
365
366 static void
367 apply_pipeline_layout_block(nir_block *block,
368 struct apply_pipeline_layout_state *state)
369 {
370 nir_foreach_instr_safe(instr, block) {
371 switch (instr->type) {
372 case nir_instr_type_intrinsic: {
373 nir_intrinsic_instr *intrin = nir_instr_as_intrinsic(instr);
374 switch (intrin->intrinsic) {
375 case nir_intrinsic_vulkan_resource_index:
376 lower_res_index_intrinsic(intrin, state);
377 break;
378 case nir_intrinsic_vulkan_resource_reindex:
379 lower_res_reindex_intrinsic(intrin, state);
380 break;
381 case nir_intrinsic_image_deref_load:
382 case nir_intrinsic_image_deref_store:
383 case nir_intrinsic_image_deref_atomic_add:
384 case nir_intrinsic_image_deref_atomic_min:
385 case nir_intrinsic_image_deref_atomic_max:
386 case nir_intrinsic_image_deref_atomic_and:
387 case nir_intrinsic_image_deref_atomic_or:
388 case nir_intrinsic_image_deref_atomic_xor:
389 case nir_intrinsic_image_deref_atomic_exchange:
390 case nir_intrinsic_image_deref_atomic_comp_swap:
391 case nir_intrinsic_image_deref_size:
392 case nir_intrinsic_image_deref_samples:
393 case nir_intrinsic_image_deref_load_param_intel:
394 case nir_intrinsic_image_deref_load_raw_intel:
395 case nir_intrinsic_image_deref_store_raw_intel:
396 lower_image_intrinsic(intrin, state);
397 break;
398 case nir_intrinsic_load_constant:
399 lower_load_constant(intrin, state);
400 break;
401 default:
402 break;
403 }
404 break;
405 }
406 case nir_instr_type_tex:
407 lower_tex(nir_instr_as_tex(instr), state);
408 break;
409 default:
410 continue;
411 }
412 }
413 }
414
415 static void
416 setup_vec4_uniform_value(uint32_t *params, uint32_t offset, unsigned n)
417 {
418 for (unsigned i = 0; i < n; ++i)
419 params[i] = ANV_PARAM_PUSH(offset + i * sizeof(uint32_t));
420
421 for (unsigned i = n; i < 4; ++i)
422 params[i] = BRW_PARAM_BUILTIN_ZERO;
423 }
424
425 void
426 anv_nir_apply_pipeline_layout(struct anv_pipeline *pipeline,
427 struct anv_pipeline_layout *layout,
428 nir_shader *shader,
429 struct brw_stage_prog_data *prog_data,
430 struct anv_pipeline_bind_map *map)
431 {
432 gl_shader_stage stage = shader->info.stage;
433
434 struct apply_pipeline_layout_state state = {
435 .shader = shader,
436 .layout = layout,
437 .add_bounds_checks = pipeline->device->robust_buffer_access,
438 };
439
440 void *mem_ctx = ralloc_context(NULL);
441
442 for (unsigned s = 0; s < layout->num_sets; s++) {
443 const unsigned count = layout->set[s].layout->binding_count;
444 const unsigned words = BITSET_WORDS(count);
445 state.set[s].used = rzalloc_array(mem_ctx, BITSET_WORD, words);
446 state.set[s].surface_offsets = rzalloc_array(mem_ctx, uint8_t, count);
447 state.set[s].sampler_offsets = rzalloc_array(mem_ctx, uint8_t, count);
448 state.set[s].image_offsets = rzalloc_array(mem_ctx, uint8_t, count);
449 }
450
451 nir_foreach_function(function, shader) {
452 if (!function->impl)
453 continue;
454
455 nir_foreach_block(block, function->impl)
456 get_used_bindings_block(block, &state);
457 }
458
459 if (state.uses_constants) {
460 state.constants_offset = map->surface_count;
461 map->surface_to_descriptor[map->surface_count].set =
462 ANV_DESCRIPTOR_SET_SHADER_CONSTANTS;
463 map->surface_count++;
464 }
465
466 for (uint32_t set = 0; set < layout->num_sets; set++) {
467 struct anv_descriptor_set_layout *set_layout = layout->set[set].layout;
468
469 BITSET_WORD b, _tmp;
470 BITSET_FOREACH_SET(b, _tmp, state.set[set].used,
471 set_layout->binding_count) {
472 struct anv_descriptor_set_binding_layout *binding =
473 &set_layout->binding[b];
474
475 if (binding->stage[stage].surface_index >= 0) {
476 state.set[set].surface_offsets[b] = map->surface_count;
477 struct anv_sampler **samplers = binding->immutable_samplers;
478 for (unsigned i = 0; i < binding->array_size; i++) {
479 uint8_t planes = samplers ? samplers[i]->n_planes : 1;
480 for (uint8_t p = 0; p < planes; p++) {
481 map->surface_to_descriptor[map->surface_count++] =
482 (struct anv_pipeline_binding) {
483 .set = set,
484 .binding = b,
485 .index = i,
486 .plane = p,
487 };
488 }
489 }
490 }
491
492 if (binding->stage[stage].sampler_index >= 0) {
493 state.set[set].sampler_offsets[b] = map->sampler_count;
494 struct anv_sampler **samplers = binding->immutable_samplers;
495 for (unsigned i = 0; i < binding->array_size; i++) {
496 uint8_t planes = samplers ? samplers[i]->n_planes : 1;
497 for (uint8_t p = 0; p < planes; p++) {
498 map->sampler_to_descriptor[map->sampler_count++] =
499 (struct anv_pipeline_binding) {
500 .set = set,
501 .binding = b,
502 .index = i,
503 .plane = p,
504 };
505 }
506 }
507 }
508
509 if (binding->stage[stage].image_index >= 0) {
510 state.set[set].image_offsets[b] = map->image_count;
511 map->image_count += binding->array_size;
512 }
513 }
514 }
515
516 if (map->image_count > 0) {
517 assert(map->image_count <= MAX_IMAGES);
518 assert(shader->num_uniforms == prog_data->nr_params * 4);
519 state.first_image_uniform = shader->num_uniforms;
520 uint32_t *param = brw_stage_prog_data_add_params(prog_data,
521 map->image_count *
522 BRW_IMAGE_PARAM_SIZE);
523 struct anv_push_constants *null_data = NULL;
524 const struct brw_image_param *image_param = null_data->images;
525 for (uint32_t i = 0; i < map->image_count; i++) {
526 setup_vec4_uniform_value(param + BRW_IMAGE_PARAM_OFFSET_OFFSET,
527 (uintptr_t)image_param->offset, 2);
528 setup_vec4_uniform_value(param + BRW_IMAGE_PARAM_SIZE_OFFSET,
529 (uintptr_t)image_param->size, 3);
530 setup_vec4_uniform_value(param + BRW_IMAGE_PARAM_STRIDE_OFFSET,
531 (uintptr_t)image_param->stride, 4);
532 setup_vec4_uniform_value(param + BRW_IMAGE_PARAM_TILING_OFFSET,
533 (uintptr_t)image_param->tiling, 3);
534 setup_vec4_uniform_value(param + BRW_IMAGE_PARAM_SWIZZLING_OFFSET,
535 (uintptr_t)image_param->swizzling, 2);
536
537 param += BRW_IMAGE_PARAM_SIZE;
538 image_param ++;
539 }
540 assert(param == prog_data->param + prog_data->nr_params);
541
542 shader->num_uniforms += map->image_count * BRW_IMAGE_PARAM_SIZE * 4;
543 assert(shader->num_uniforms == prog_data->nr_params * 4);
544 }
545
546 nir_foreach_variable(var, &shader->uniforms) {
547 const struct glsl_type *glsl_type = glsl_without_array(var->type);
548
549 if (!glsl_type_is_image(glsl_type))
550 continue;
551
552 enum glsl_sampler_dim dim = glsl_get_sampler_dim(glsl_type);
553
554 const uint32_t set = var->data.descriptor_set;
555 const uint32_t binding = var->data.binding;
556 const uint32_t array_size =
557 layout->set[set].layout->binding[binding].array_size;
558
559 if (!BITSET_TEST(state.set[set].used, binding))
560 continue;
561
562 struct anv_pipeline_binding *pipe_binding =
563 &map->surface_to_descriptor[state.set[set].surface_offsets[binding]];
564 for (unsigned i = 0; i < array_size; i++) {
565 assert(pipe_binding[i].set == set);
566 assert(pipe_binding[i].binding == binding);
567 assert(pipe_binding[i].index == i);
568
569 if (dim == GLSL_SAMPLER_DIM_SUBPASS ||
570 dim == GLSL_SAMPLER_DIM_SUBPASS_MS)
571 pipe_binding[i].input_attachment_index = var->data.index + i;
572
573 pipe_binding[i].write_only =
574 (var->data.image.access & ACCESS_NON_READABLE) != 0;
575 }
576 }
577
578 nir_foreach_function(function, shader) {
579 if (!function->impl)
580 continue;
581
582 nir_builder_init(&state.builder, function->impl);
583 nir_foreach_block(block, function->impl)
584 apply_pipeline_layout_block(block, &state);
585 nir_metadata_preserve(function->impl, nir_metadata_block_index |
586 nir_metadata_dominance);
587 }
588
589 ralloc_free(mem_ctx);
590 }