1/*
2 * Copyright © 2019 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
21 * DEALINGS IN THE SOFTWARE.
22 */
23
24#include "compiler/nir/nir.h"
25#include "compiler/nir/nir_builder.h"
26#include "gl_nir.h"
27#include "ir_uniform.h"
28
29#include "util/compiler.h"
30#include "main/shader_types.h"
31
32static nir_ssa_def *
33get_block_array_index(nir_builder *b, nir_deref_instr *deref,
34                      const struct gl_shader_program *shader_program)
35{
36   unsigned array_elements = 1;
37
38   /* Build a block name such as "block[2][0]" for finding in the list of
39    * blocks later on as well as an optional dynamic index which gets added
40    * to the block index later.
41    */
42   int const_array_offset = 0;
43   const char *block_name = "";
44   nir_ssa_def *nonconst_index = NULL;
45   while (deref->deref_type == nir_deref_type_array) {
46      nir_deref_instr *parent = nir_deref_instr_parent(deref);
47      assert(parent && glsl_type_is_array(parent->type));
48      unsigned arr_size = glsl_get_length(parent->type);
49
50      if (nir_src_is_const(deref->arr.index)) {
51         unsigned arr_index = nir_src_as_uint(deref->arr.index);
52
53         /* We're walking the deref from the tail so prepend the array index */
54         block_name = ralloc_asprintf(b->shader, "[%u]%s", arr_index,
55                                      block_name);
56
57         const_array_offset += arr_index * array_elements;
58      } else {
59         nir_ssa_def *arr_index = nir_ssa_for_src(b, deref->arr.index, 1);
60         arr_index = nir_umin(b, arr_index, nir_imm_int(b, arr_size - 1));
61         nir_ssa_def *arr_offset = nir_amul_imm(b, arr_index, array_elements);
62         if (nonconst_index)
63            nonconst_index = nir_iadd(b, nonconst_index, arr_offset);
64         else
65            nonconst_index = arr_offset;
66
67         /* We're walking the deref from the tail so prepend the array index */
68         block_name = ralloc_asprintf(b->shader, "[0]%s", block_name);
69      }
70
71      array_elements *= arr_size;
72      deref = parent;
73   }
74
75   assert(deref->deref_type == nir_deref_type_var);
76   int binding = const_array_offset + deref->var->data.binding;
77   block_name = ralloc_asprintf(b->shader, "%s%s",
78                                glsl_get_type_name(deref->var->interface_type),
79                                block_name);
80
81   struct gl_linked_shader *linked_shader =
82      shader_program->_LinkedShaders[b->shader->info.stage];
83
84   unsigned num_blocks;
85   struct gl_uniform_block **blocks;
86   if (nir_deref_mode_is(deref, nir_var_mem_ubo)) {
87      num_blocks = linked_shader->Program->info.num_ubos;
88      blocks = linked_shader->Program->sh.UniformBlocks;
89   } else {
90      assert(nir_deref_mode_is(deref, nir_var_mem_ssbo));
91      num_blocks = linked_shader->Program->info.num_ssbos;
92      blocks = linked_shader->Program->sh.ShaderStorageBlocks;
93   }
94
95   /* Block names are optional with ARB_gl_spirv so use the binding instead. */
96   bool use_bindings = shader_program->data->spirv;
97
98   for (unsigned i = 0; i < num_blocks; i++) {
99      if (( use_bindings && binding == blocks[i]->Binding) ||
100          (!use_bindings && strcmp(block_name, blocks[i]->name.string) == 0)) {
101         deref->var->data.driver_location = i - const_array_offset;
102         if (nonconst_index)
103            return nir_iadd_imm(b, nonconst_index, i);
104         else
105            return nir_imm_int(b, i);
106      }
107   }
108
109   /* TODO: Investigate if we could change the code to assign Bindings to the
110    * blocks that were not explicitly assigned, so we can always compare
111    * bindings.
112    */
113
114   if (use_bindings)
115      unreachable("Failed to find the block by binding");
116   else
117      unreachable("Failed to find the block by name");
118}
119
120static void
121get_block_index_offset(nir_variable *var,
122                       const struct gl_shader_program *shader_program,
123                       gl_shader_stage stage,
124                       unsigned *index, unsigned *offset)
125{
126
127   struct gl_linked_shader *linked_shader =
128      shader_program->_LinkedShaders[stage];
129
130   unsigned num_blocks;
131   struct gl_uniform_block **blocks;
132   if (var->data.mode == nir_var_mem_ubo) {
133      num_blocks = linked_shader->Program->info.num_ubos;
134      blocks = linked_shader->Program->sh.UniformBlocks;
135   } else {
136      assert(var->data.mode == nir_var_mem_ssbo);
137      num_blocks = linked_shader->Program->info.num_ssbos;
138      blocks = linked_shader->Program->sh.ShaderStorageBlocks;
139   }
140
141   /* Block names are optional with ARB_gl_spirv so use the binding instead. */
142   bool use_bindings = shader_program->data->spirv;
143
144   for (unsigned i = 0; i < num_blocks; i++) {
145      const char *block_name = glsl_get_type_name(var->interface_type);
146      if (( use_bindings && blocks[i]->Binding == var->data.binding) ||
147          (!use_bindings && strcmp(block_name, blocks[i]->name.string) == 0)) {
148         var->data.driver_location = i;
149         *index = i;
150         *offset = blocks[i]->Uniforms[var->data.location].Offset;
151         return;
152      }
153   }
154
155   if (use_bindings)
156      unreachable("Failed to find the block by binding");
157   else
158      unreachable("Failed to find the block by name");
159}
160
161static bool
162lower_buffer_interface_derefs_impl(nir_function_impl *impl,
163                                   const struct gl_shader_program *shader_program)
164{
165   bool progress = false;
166
167   nir_builder b;
168   nir_builder_init(&b, impl);
169
170   /* this must be a separate loop before the main pass in order to ensure that
171    * access info is fully propagated prior to the info being lost during rewrites
172    */
173   nir_foreach_block(block, impl) {
174      nir_foreach_instr(instr, block) {
175         if (instr->type != nir_instr_type_intrinsic)
176            continue;
177
178         nir_intrinsic_instr *intrin = nir_instr_as_intrinsic(instr);
179         if (intrin->intrinsic == nir_intrinsic_load_deref ||
180             intrin->intrinsic == nir_intrinsic_store_deref) {
181            nir_variable *var = nir_intrinsic_get_var(intrin, 0);
182            assert(var);
183            nir_intrinsic_set_access(intrin, nir_intrinsic_access(intrin) | var->data.access);
184         }
185      }
186   }
187
188   nir_foreach_block(block, impl) {
189      nir_foreach_instr_safe(instr, block) {
190         switch (instr->type) {
191         case nir_instr_type_deref: {
192            nir_deref_instr *deref = nir_instr_as_deref(instr);
193            if (!nir_deref_mode_is_one_of(deref, nir_var_mem_ubo |
194                                                 nir_var_mem_ssbo))
195               break;
196
197            /* We use nir_address_format_32bit_index_offset */
198            assert(deref->dest.is_ssa);
199            assert(deref->dest.ssa.bit_size == 32);
200            deref->dest.ssa.num_components = 2;
201
202            progress = true;
203
204            b.cursor = nir_before_instr(&deref->instr);
205
206            unsigned offset = 0;
207            nir_ssa_def *ptr;
208            if (deref->deref_type == nir_deref_type_var &&
209                !glsl_type_is_interface(glsl_without_array(deref->var->type))) {
210               /* This variable is contained in an interface block rather than
211                * containing one.  We need the block index and its offset
212                * inside that block
213                */
214               unsigned index;
215               get_block_index_offset(deref->var, shader_program,
216                                      b.shader->info.stage,
217                                      &index, &offset);
218               ptr = nir_imm_ivec2(&b, index, offset);
219            } else if (glsl_type_is_interface(deref->type)) {
220               /* This is the last deref before the block boundary.
221                * Everything after this point is a byte offset and will be
222                * handled by nir_lower_explicit_io().
223                */
224               nir_ssa_def *index = get_block_array_index(&b, deref,
225                                                          shader_program);
226               ptr = nir_vec2(&b, index, nir_imm_int(&b, offset));
227            } else {
228               /* This will get handled by nir_lower_explicit_io(). */
229               break;
230            }
231
232            nir_deref_instr *cast = nir_build_deref_cast(&b, ptr, deref->modes,
233                                                         deref->type, 0);
234            /* Set the alignment on the cast so that we get good alignment out
235             * of nir_lower_explicit_io.  Our offset to the start of the UBO
236             * variable is always a constant, so we can use the maximum
237             * align_mul.
238             */
239            cast->cast.align_mul = NIR_ALIGN_MUL_MAX;
240            cast->cast.align_offset = offset % NIR_ALIGN_MUL_MAX;
241
242            nir_ssa_def_rewrite_uses(&deref->dest.ssa,
243                                     &cast->dest.ssa);
244            nir_deref_instr_remove_if_unused(deref);
245            break;
246         }
247
248         case nir_instr_type_intrinsic: {
249            nir_intrinsic_instr *intrin = nir_instr_as_intrinsic(instr);
250            switch (intrin->intrinsic) {
251            case nir_intrinsic_load_deref: {
252               nir_deref_instr *deref = nir_src_as_deref(intrin->src[0]);
253               if (!nir_deref_mode_is_one_of(deref, nir_var_mem_ubo |
254                                                    nir_var_mem_ssbo))
255                  break;
256
257               /* UBO and SSBO Booleans are 32-bit integers where any non-zero
258                * value is considered true.  NIR Booleans, on the other hand
259                * are 1-bit values until you get to a very late stage of the
260                * compilation process.  We need to turn those 1-bit loads into
261                * a 32-bit load wrapped in an i2b to get a proper NIR boolean
262                * from the SSBO.
263                */
264               if (glsl_type_is_boolean(deref->type)) {
265                  assert(intrin->dest.is_ssa);
266                  b.cursor = nir_after_instr(&intrin->instr);
267                  intrin->dest.ssa.bit_size = 32;
268                  nir_ssa_def *bval = nir_i2b(&b, &intrin->dest.ssa);
269                  nir_ssa_def_rewrite_uses_after(&intrin->dest.ssa,
270                                                 bval,
271                                                 bval->parent_instr);
272                  progress = true;
273               }
274               break;
275            }
276
277            case nir_intrinsic_store_deref: {
278               nir_deref_instr *deref = nir_src_as_deref(intrin->src[0]);
279               if (!nir_deref_mode_is_one_of(deref, nir_var_mem_ubo |
280                                                    nir_var_mem_ssbo))
281                  break;
282
283               /* SSBO Booleans are 32-bit integers where any non-zero value
284                * is considered true.  NIR Booleans, on the other hand are
285                * 1-bit values until you get to a very late stage of the
286                * compilation process.  We need to turn those 1-bit stores
287                * into a b2i32 followed by a 32-bit store.  Technically the
288                * value we write doesn't have to be 0/1 so once Booleans are
289                * lowered to 32-bit values, we have an unneeded sanitation
290                * step but in practice it doesn't cost much.
291                */
292               if (glsl_type_is_boolean(deref->type)) {
293                  assert(intrin->src[1].is_ssa);
294                  b.cursor = nir_before_instr(&intrin->instr);
295                  nir_ssa_def *ival = nir_b2i32(&b, intrin->src[1].ssa);
296                  nir_instr_rewrite_src(&intrin->instr, &intrin->src[1],
297                                        nir_src_for_ssa(ival));
298                  progress = true;
299               }
300               break;
301            }
302
303            case nir_intrinsic_copy_deref:
304               unreachable("copy_deref should be lowered by now");
305               break;
306
307            default:
308               /* Nothing to do */
309               break;
310            }
311            break;
312         }
313
314         default:
315            break; /* Nothing to do */
316         }
317      }
318   }
319
320   if (progress) {
321      nir_metadata_preserve(impl, nir_metadata_block_index |
322                                  nir_metadata_dominance);
323   } else {
324      nir_metadata_preserve(impl, nir_metadata_all);
325   }
326
327   return progress;
328}
329
330bool
331gl_nir_lower_buffers(nir_shader *shader,
332                     const struct gl_shader_program *shader_program)
333{
334   bool progress = false;
335
336   nir_foreach_variable_with_modes(var, shader, nir_var_mem_ubo | nir_var_mem_ssbo) {
337      var->data.driver_location = -1;
338      progress = true;
339   }
340
341   /* First, we lower the derefs to turn block variable and array derefs into
342    * a nir_address_format_32bit_index_offset pointer.  From there forward,
343    * we leave the derefs in place and let nir_lower_explicit_io handle them.
344    */
345   nir_foreach_function(function, shader) {
346      if (function->impl &&
347          lower_buffer_interface_derefs_impl(function->impl, shader_program))
348         progress = true;
349   }
350
351   /* If that did something, we validate and then call nir_lower_explicit_io
352    * to finish the process.
353    */
354   if (progress) {
355      nir_validate_shader(shader, "Lowering buffer interface derefs");
356      nir_lower_explicit_io(shader, nir_var_mem_ubo | nir_var_mem_ssbo,
357                            nir_address_format_32bit_index_offset);
358   }
359
360   return progress;
361}
362