1 /*
2  * Copyright © 2020 Mike Blumenkrantz
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  * Authors:
24  *    Mike Blumenkrantz <michael.blumenkrantz@gmail.com>
25  */
26 
27 #include "nir.h"
28 #include "nir_builder.h"
29 
30 bool nir_lower_dynamic_bo_access(nir_shader *shader);
31 /**
32  * This pass converts dynamic UBO/SSBO block indices to constant indices by generating
33  * conditional chains which reduce to single values.
34  *
35  * This is needed by anything which intends to convert GLSL-like shaders to SPIRV,
36  * as SPIRV requires explicit load points for UBO/SSBO variables and has no instruction for
37  * loading based on an offset in the underlying driver's binding table
38  */
39 
40 
41 /* generate a single ssa value which conditionally selects the right value that
42  * was previously loaded by the load_ubo conditional chain
43  */
44 static nir_ssa_def *
recursive_generate_bo_ssa_def(nir_builder * b,nir_intrinsic_instr * instr,nir_ssa_def * index,unsigned start,unsigned end)45 recursive_generate_bo_ssa_def(nir_builder *b, nir_intrinsic_instr *instr, nir_ssa_def *index, unsigned start, unsigned end)
46 {
47    if (start == end - 1) {
48       nir_intrinsic_instr *new_instr = nir_intrinsic_instr_create(b->shader, instr->intrinsic);
49       new_instr->src[0] = nir_src_for_ssa(nir_imm_int(b, start));
50       for (unsigned i = 0; i < nir_intrinsic_infos[instr->intrinsic].num_srcs; i++) {
51          if (i)
52             nir_src_copy(&new_instr->src[i], &instr->src[i]);
53       }
54       if (instr->intrinsic != nir_intrinsic_load_ubo_vec4) {
55          nir_intrinsic_set_align(new_instr, nir_intrinsic_align_mul(instr), nir_intrinsic_align_offset(instr));
56          if (instr->intrinsic != nir_intrinsic_load_ssbo)
57             nir_intrinsic_set_range(new_instr, nir_intrinsic_range(instr));
58       }
59       new_instr->num_components = instr->num_components;
60       nir_ssa_dest_init(&new_instr->instr, &new_instr->dest,
61                         nir_dest_num_components(instr->dest),
62                         nir_dest_bit_size(instr->dest), NULL);
63       nir_builder_instr_insert(b, &new_instr->instr);
64       return &new_instr->dest.ssa;
65    }
66 
67    unsigned mid = start + (end - start) / 2;
68    return nir_build_alu(b, nir_op_bcsel, nir_build_alu(b, nir_op_ilt, index, nir_imm_int(b, mid), NULL, NULL),
69       recursive_generate_bo_ssa_def(b, instr, index, start, mid),
70       recursive_generate_bo_ssa_def(b, instr, index, mid, end),
71       NULL
72    );
73 }
74 
75 static void
generate_store_ssbo_ssa_def(nir_builder * b,nir_intrinsic_instr * instr,nir_ssa_def * index,unsigned start,unsigned end)76 generate_store_ssbo_ssa_def(nir_builder *b, nir_intrinsic_instr *instr, nir_ssa_def *index, unsigned start, unsigned end)
77 {
78    if (start == end - 1) {
79       nir_intrinsic_instr *new_instr = nir_instr_as_intrinsic(nir_instr_clone(b->shader, &instr->instr));
80       new_instr->src[1] = nir_src_for_ssa(nir_imm_int(b, start));
81       nir_builder_instr_insert(b, &new_instr->instr);
82    } else {
83       int mid = start + (end - start) / 2;
84       nir_ssa_def *mid_idx = nir_imm_int(b, mid);
85       nir_push_if(b, nir_ilt(b, index, mid_idx));
86       generate_store_ssbo_ssa_def(b, instr, index, start, mid);
87       nir_push_else(b, NULL);
88       generate_store_ssbo_ssa_def(b, instr, index, mid, end);
89       nir_pop_if(b, NULL);
90    }
91 }
92 
93 static bool
lower_dynamic_bo_access_instr(nir_builder * b,nir_instr * instr_,UNUSED void * cb_data)94 lower_dynamic_bo_access_instr(nir_builder *b,
95                               nir_instr *instr_,
96                               UNUSED void *cb_data)
97 {
98    if (instr_->type != nir_instr_type_intrinsic)
99       return false;
100 
101    nir_intrinsic_instr *instr = nir_instr_as_intrinsic(instr_);
102 
103    if (instr->intrinsic != nir_intrinsic_load_ubo &&
104        instr->intrinsic != nir_intrinsic_load_ubo_vec4 &&
105        instr->intrinsic != nir_intrinsic_get_ssbo_size &&
106        instr->intrinsic != nir_intrinsic_load_ssbo &&
107        instr->intrinsic != nir_intrinsic_store_ssbo)
108       return false;
109    /* block index src is 1 for this op */
110    unsigned block_idx = instr->intrinsic == nir_intrinsic_store_ssbo;
111    if (nir_src_is_const(instr->src[block_idx]))
112       return false;
113    b->cursor = nir_after_instr(&instr->instr);
114    bool ssbo_mode = instr->intrinsic != nir_intrinsic_load_ubo && instr->intrinsic != nir_intrinsic_load_ubo_vec4;
115    unsigned first_idx = UINT_MAX, last_idx;
116    if (ssbo_mode) {
117       nir_foreach_variable_with_modes(var, b->shader, nir_var_mem_ssbo)
118          first_idx = MIN2(first_idx, var->data.driver_location);
119       last_idx = first_idx + b->shader->info.num_ssbos;
120    } else {
121       /* skip 0 index if uniform_0 is one we created previously */
122       first_idx = !b->shader->info.first_ubo_is_default_ubo;
123       last_idx = first_idx + b->shader->info.num_ubos;
124    }
125 
126    if (instr->intrinsic != nir_intrinsic_store_ssbo) {
127       /* now create the composite dest with a bcsel chain based on the original value */
128       nir_ssa_def *new_dest = recursive_generate_bo_ssa_def(b, instr,
129                                                           instr->src[block_idx].ssa,
130                                                           first_idx, last_idx);
131 
132       /* now use the composite dest in all cases where the original dest (from the dynamic index)
133        * was used and remove the dynamically-indexed load_*bo instruction
134        */
135       nir_ssa_def_rewrite_uses_after(&instr->dest.ssa, new_dest,
136                                      &instr->instr);
137    } else
138       generate_store_ssbo_ssa_def(b, instr, instr->src[block_idx].ssa, first_idx, last_idx);
139    nir_instr_remove(&instr->instr);
140 
141    return true;
142 }
143 
144 bool
nir_lower_dynamic_bo_access(nir_shader * shader)145 nir_lower_dynamic_bo_access(nir_shader *shader)
146 {
147    return nir_shader_instructions_pass(shader,
148                                        lower_dynamic_bo_access_instr,
149                                        nir_metadata_dominance,
150                                        NULL);
151 }
152