xref: /aosp_15_r20/external/mesa3d/src/compiler/nir/nir_lower_scratch.c (revision 6104692788411f58d303aa86923a9ff6ecaded22)
1 /*
2  * Copyright © 2016 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 /*
25  * This lowering pass converts references to variables with loads/stores to
26  * scratch space based on a few configurable parameters.
27  */
28 
29 #include "nir.h"
30 #include "nir_builder.h"
31 #include "nir_deref.h"
32 
33 static void
lower_load_store(nir_builder * b,nir_intrinsic_instr * intrin,glsl_type_size_align_func size_align)34 lower_load_store(nir_builder *b,
35                  nir_intrinsic_instr *intrin,
36                  glsl_type_size_align_func size_align)
37 {
38    b->cursor = nir_before_instr(&intrin->instr);
39 
40    nir_deref_instr *deref = nir_src_as_deref(intrin->src[0]);
41    nir_variable *var = nir_deref_instr_get_variable(deref);
42 
43    nir_def *offset =
44       nir_iadd_imm(b, nir_build_deref_offset(b, deref, size_align),
45                    var->data.location);
46 
47    unsigned align, UNUSED size;
48    size_align(deref->type, &size, &align);
49 
50    if (intrin->intrinsic == nir_intrinsic_load_deref) {
51       unsigned bit_size = intrin->def.bit_size;
52       nir_def *value = nir_load_scratch(
53          b, intrin->num_components, bit_size == 1 ? 32 : bit_size, offset, .align_mul = align);
54       if (bit_size == 1)
55          value = nir_b2b1(b, value);
56 
57       nir_def_rewrite_uses(&intrin->def, value);
58    } else {
59       assert(intrin->intrinsic == nir_intrinsic_store_deref);
60 
61       nir_def *value = intrin->src[1].ssa;
62       if (value->bit_size == 1)
63          value = nir_b2b32(b, value);
64 
65       nir_store_scratch(b, value, offset, .align_mul = align,
66                         .write_mask = nir_intrinsic_write_mask(intrin));
67    }
68 
69    nir_instr_remove(&intrin->instr);
70    nir_deref_instr_remove_if_unused(deref);
71 }
72 
73 static bool
only_used_for_load_store(nir_deref_instr * deref)74 only_used_for_load_store(nir_deref_instr *deref)
75 {
76    nir_foreach_use(src, &deref->def) {
77       if (!nir_src_parent_instr(src))
78          return false;
79       if (nir_src_parent_instr(src)->type == nir_instr_type_deref) {
80          if (!only_used_for_load_store(nir_instr_as_deref(nir_src_parent_instr(src))))
81             return false;
82       } else if (nir_src_parent_instr(src)->type != nir_instr_type_intrinsic) {
83          return false;
84       } else {
85          nir_intrinsic_instr *intrin = nir_instr_as_intrinsic(nir_src_parent_instr(src));
86          if (intrin->intrinsic != nir_intrinsic_load_deref &&
87              intrin->intrinsic != nir_intrinsic_store_deref)
88             return false;
89       }
90    }
91    return true;
92 }
93 
94 bool
nir_lower_vars_to_scratch(nir_shader * shader,nir_variable_mode modes,int size_threshold,glsl_type_size_align_func variable_size_align,glsl_type_size_align_func scratch_layout_size_align)95 nir_lower_vars_to_scratch(nir_shader *shader,
96                           nir_variable_mode modes,
97                           int size_threshold,
98                           glsl_type_size_align_func variable_size_align,
99                           glsl_type_size_align_func scratch_layout_size_align)
100 {
101    struct set *set = _mesa_pointer_set_create(NULL);
102 
103    /* First, we walk the instructions and flag any variables we want to lower
104     * by removing them from their respective list and setting the mode to 0.
105     */
106    nir_foreach_function_impl(impl, shader) {
107       nir_foreach_block(block, impl) {
108          nir_foreach_instr(instr, block) {
109             if (instr->type != nir_instr_type_intrinsic)
110                continue;
111 
112             nir_intrinsic_instr *intrin = nir_instr_as_intrinsic(instr);
113             if (intrin->intrinsic != nir_intrinsic_load_deref &&
114                 intrin->intrinsic != nir_intrinsic_store_deref)
115                continue;
116 
117             nir_deref_instr *deref = nir_src_as_deref(intrin->src[0]);
118             if (!nir_deref_mode_is_one_of(deref, modes))
119                continue;
120 
121             if (!nir_deref_instr_has_indirect(nir_src_as_deref(intrin->src[0])))
122                continue;
123 
124             nir_variable *var = nir_deref_instr_get_variable(deref);
125             if (!var)
126                continue;
127 
128             /* We set var->mode to 0 to indicate that a variable will be moved
129              * to scratch.  Don't assign a scratch location twice.
130              */
131             if (var->data.mode == 0)
132                continue;
133 
134             unsigned var_size, var_align;
135             variable_size_align(var->type, &var_size, &var_align);
136             if (var_size <= size_threshold)
137                continue;
138 
139             _mesa_set_add(set, var);
140          }
141       }
142    }
143 
144    if (set->entries == 0) {
145       _mesa_set_destroy(set, NULL);
146       return false;
147    }
148 
149    bool progress = false;
150 
151    nir_foreach_function_impl(impl, shader) {
152       nir_foreach_block(block, impl) {
153          nir_foreach_instr_safe(instr, block) {
154             if (instr->type != nir_instr_type_deref)
155                continue;
156 
157             nir_deref_instr *deref = nir_instr_as_deref(instr);
158 
159             if (nir_deref_instr_remove_if_unused(deref)) {
160                progress = true;
161                continue;
162             }
163 
164             if (deref->deref_type != nir_deref_type_var)
165                continue;
166 
167             struct set_entry *entry = _mesa_set_search(set, deref->var);
168             if (!entry)
169                continue;
170 
171             if (!only_used_for_load_store(deref))
172                _mesa_set_remove(set, entry);
173          }
174       }
175    }
176 
177    set_foreach(set, entry) {
178       nir_variable *var = (void *)entry->key;
179 
180       /* Remove it from its list */
181       exec_node_remove(&var->node);
182       /* Invalid mode used to flag "moving to scratch" */
183       var->data.mode = 0;
184 
185       /* We don't allocate space here as iteration in this loop is
186        * non-deterministic due to the nir_variable pointers. */
187       var->data.location = INT_MAX;
188    }
189 
190    nir_foreach_function_impl(impl, shader) {
191       nir_builder build = nir_builder_create(impl);
192 
193       bool impl_progress = false;
194       nir_foreach_block(block, impl) {
195          nir_foreach_instr_safe(instr, block) {
196             if (instr->type != nir_instr_type_intrinsic)
197                continue;
198 
199             nir_intrinsic_instr *intrin = nir_instr_as_intrinsic(instr);
200             if (intrin->intrinsic != nir_intrinsic_load_deref &&
201                 intrin->intrinsic != nir_intrinsic_store_deref)
202                continue;
203 
204             nir_variable *var = nir_intrinsic_get_var(intrin, 0);
205             /* Variables flagged for lowering above have mode == 0 */
206             if (!var || var->data.mode)
207                continue;
208 
209             if (var->data.location == INT_MAX) {
210                unsigned var_size, var_align;
211                scratch_layout_size_align(var->type, &var_size, &var_align);
212 
213                var->data.location = ALIGN_POT(shader->scratch_size, var_align);
214                shader->scratch_size = var->data.location + var_size;
215             }
216 
217             lower_load_store(&build, intrin, scratch_layout_size_align);
218             impl_progress = true;
219          }
220       }
221 
222       if (impl_progress) {
223          progress = true;
224          nir_metadata_preserve(impl, nir_metadata_control_flow);
225       } else {
226          nir_metadata_preserve(impl, nir_metadata_all);
227       }
228    }
229 
230    _mesa_set_destroy(set, NULL);
231 
232    return progress;
233 }
234