1 /*
2 * Copyright 2023 Valve Corpoation
3 * Copyright 2020 Raspberry Pi Ltd
4 * SPDX-License-Identifier: MIT
5 */
6
7 #include "nir.h"
8 #include "nir_builder.h"
9 #include "nir_intrinsics_indices.h"
10
11 static void
rewrite_offset(nir_builder * b,nir_intrinsic_instr * instr,uint32_t type_sz,uint32_t offset_src,nir_def * size)12 rewrite_offset(nir_builder *b, nir_intrinsic_instr *instr,
13 uint32_t type_sz, uint32_t offset_src, nir_def *size)
14 {
15 /* Compute the maximum offset being accessed and if it is out of bounds
16 * rewrite it to 0 to ensure the access is within bounds.
17 */
18 const uint32_t access_size = instr->num_components * type_sz;
19 nir_def *max_access_offset =
20 nir_iadd_imm(b, instr->src[offset_src].ssa, access_size - 1);
21 nir_def *offset =
22 nir_bcsel(b, nir_uge(b, max_access_offset, size), nir_imm_int(b, 0),
23 instr->src[offset_src].ssa);
24
25 /* Rewrite offset */
26 nir_src_rewrite(&instr->src[offset_src], offset);
27 }
28
29 /*
30 * Wrap a intrinsic in an if, predicated on a "valid" condition. If the
31 * intrinsic produces a destination, it will be zero in the invalid case.
32 */
33 static void
wrap_in_if(nir_builder * b,nir_intrinsic_instr * instr,nir_def * valid)34 wrap_in_if(nir_builder *b, nir_intrinsic_instr *instr, nir_def *valid)
35 {
36 bool has_dest = nir_intrinsic_infos[instr->intrinsic].has_dest;
37 nir_def *res, *zero;
38
39 if (has_dest) {
40 zero = nir_imm_zero(b, instr->def.num_components,
41 instr->def.bit_size);
42 }
43
44 nir_push_if(b, valid);
45 {
46 nir_instr *orig = nir_instr_clone(b->shader, &instr->instr);
47 nir_builder_instr_insert(b, orig);
48
49 if (has_dest)
50 res = &nir_instr_as_intrinsic(orig)->def;
51 }
52 nir_pop_if(b, NULL);
53
54 if (has_dest)
55 nir_def_rewrite_uses(&instr->def, nir_if_phi(b, res, zero));
56
57 /* We've cloned and wrapped, so drop original instruction */
58 nir_instr_remove(&instr->instr);
59 }
60
61 static void
lower_buffer_load(nir_builder * b,nir_intrinsic_instr * instr,const nir_lower_robust_access_options * opts)62 lower_buffer_load(nir_builder *b,
63 nir_intrinsic_instr *instr,
64 const nir_lower_robust_access_options *opts)
65 {
66 uint32_t type_sz = instr->def.bit_size / 8;
67 nir_def *size;
68 nir_def *index = instr->src[0].ssa;
69
70 if (instr->intrinsic == nir_intrinsic_load_ubo) {
71 size = nir_get_ubo_size(b, 32, index);
72 } else {
73 size = nir_get_ssbo_size(b, index);
74 }
75
76 rewrite_offset(b, instr, type_sz, 1, size);
77 }
78
79 static void
lower_buffer_store(nir_builder * b,nir_intrinsic_instr * instr)80 lower_buffer_store(nir_builder *b, nir_intrinsic_instr *instr)
81 {
82 uint32_t type_sz = nir_src_bit_size(instr->src[0]) / 8;
83 rewrite_offset(b, instr, type_sz, 2,
84 nir_get_ssbo_size(b, instr->src[1].ssa));
85 }
86
87 static void
lower_buffer_atomic(nir_builder * b,nir_intrinsic_instr * instr)88 lower_buffer_atomic(nir_builder *b, nir_intrinsic_instr *instr)
89 {
90 rewrite_offset(b, instr, 4, 1, nir_get_ssbo_size(b, instr->src[0].ssa));
91 }
92
93 static void
lower_buffer_shared(nir_builder * b,nir_intrinsic_instr * instr)94 lower_buffer_shared(nir_builder *b, nir_intrinsic_instr *instr)
95 {
96 uint32_t type_sz, offset_src;
97 if (instr->intrinsic == nir_intrinsic_load_shared) {
98 offset_src = 0;
99 type_sz = instr->def.bit_size / 8;
100 } else if (instr->intrinsic == nir_intrinsic_store_shared) {
101 offset_src = 1;
102 type_sz = nir_src_bit_size(instr->src[0]) / 8;
103 } else {
104 /* atomic */
105 offset_src = 0;
106 type_sz = 4;
107 }
108
109 rewrite_offset(b, instr, type_sz, offset_src,
110 nir_imm_int(b, b->shader->info.shared_size));
111 }
112
113 static bool
lower_image(nir_builder * b,nir_intrinsic_instr * instr,const nir_lower_robust_access_options * opts)114 lower_image(nir_builder *b,
115 nir_intrinsic_instr *instr,
116 const nir_lower_robust_access_options *opts)
117 {
118 enum glsl_sampler_dim dim = nir_intrinsic_image_dim(instr);
119 bool atomic = (instr->intrinsic == nir_intrinsic_image_atomic ||
120 instr->intrinsic == nir_intrinsic_image_atomic_swap);
121 if (!opts->lower_image &&
122 !(opts->lower_buffer_image && dim == GLSL_SAMPLER_DIM_BUF) &&
123 !(opts->lower_image_atomic && atomic))
124 return false;
125
126 uint32_t num_coords = nir_image_intrinsic_coord_components(instr);
127 bool is_array = nir_intrinsic_image_array(instr);
128 nir_def *coord = instr->src[1].ssa;
129
130 /* Get image size. imageSize for cubes returns the size of a single face. */
131 unsigned size_components = num_coords;
132 if (dim == GLSL_SAMPLER_DIM_CUBE && !is_array)
133 size_components -= 1;
134
135 nir_def *size =
136 nir_image_size(b, size_components, 32,
137 instr->src[0].ssa, nir_imm_int(b, 0),
138 .image_array = is_array, .image_dim = dim);
139
140 if (dim == GLSL_SAMPLER_DIM_CUBE) {
141 nir_def *z = is_array ? nir_imul_imm(b, nir_channel(b, size, 2), 6)
142 : nir_imm_int(b, 6);
143
144 size = nir_vec3(b, nir_channel(b, size, 0), nir_channel(b, size, 1), z);
145 }
146
147 /* Only execute if coordinates are in-bounds. Otherwise, return zero. */
148 wrap_in_if(b, instr, nir_ball(b, nir_ult(b, coord, size)));
149 return true;
150 }
151
152 static bool
lower(nir_builder * b,nir_instr * instr,void * _opts)153 lower(nir_builder *b, nir_instr *instr, void *_opts)
154 {
155 const nir_lower_robust_access_options *opts = _opts;
156 if (instr->type != nir_instr_type_intrinsic)
157 return false;
158
159 nir_intrinsic_instr *intr = nir_instr_as_intrinsic(instr);
160 b->cursor = nir_before_instr(instr);
161
162 switch (intr->intrinsic) {
163 case nir_intrinsic_image_load:
164 case nir_intrinsic_image_store:
165 case nir_intrinsic_image_atomic:
166 case nir_intrinsic_image_atomic_swap:
167 return lower_image(b, intr, opts);
168
169 case nir_intrinsic_load_ubo:
170 if (opts->lower_ubo) {
171 lower_buffer_load(b, intr, opts);
172 return true;
173 }
174 return false;
175
176 case nir_intrinsic_load_ssbo:
177 if (opts->lower_ssbo) {
178 lower_buffer_load(b, intr, opts);
179 return true;
180 }
181 return false;
182 case nir_intrinsic_store_ssbo:
183 if (opts->lower_ssbo) {
184 lower_buffer_store(b, intr);
185 return true;
186 }
187 return false;
188 case nir_intrinsic_ssbo_atomic:
189 if (opts->lower_ssbo) {
190 lower_buffer_atomic(b, intr);
191 return true;
192 }
193 return false;
194
195 case nir_intrinsic_store_shared:
196 case nir_intrinsic_load_shared:
197 case nir_intrinsic_shared_atomic:
198 case nir_intrinsic_shared_atomic_swap:
199 if (opts->lower_shared) {
200 lower_buffer_shared(b, intr);
201 return true;
202 }
203 return false;
204
205 default:
206 return false;
207 }
208 }
209
210 bool
nir_lower_robust_access(nir_shader * s,const nir_lower_robust_access_options * opts)211 nir_lower_robust_access(nir_shader *s,
212 const nir_lower_robust_access_options *opts)
213 {
214 return nir_shader_instructions_pass(s, lower, nir_metadata_block_index | nir_metadata_dominance,
215 (void *)opts);
216 }
217