xref: /aosp_15_r20/external/mesa3d/src/compiler/nir/nir_lower_alu_width.c (revision 6104692788411f58d303aa86923a9ff6ecaded22)
1 /*
2  * Copyright © 2014-2015 Broadcom
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 "nir.h"
25 #include "nir_builder.h"
26 
27 struct alu_width_data {
28    nir_vectorize_cb cb;
29    const void *data;
30 };
31 
32 /** @file nir_lower_alu_width.c
33  *
34  * Replaces nir_alu_instr operations with more than one channel used in the
35  * arguments with individual per-channel operations.
36  *
37  * Optionally, a callback function which returns the max vectorization width
38  * per instruction can be provided.
39  *
40  * The max vectorization width must be a power of 2.
41  */
42 
43 static bool
inst_is_vector_alu(const nir_instr * instr,const void * _state)44 inst_is_vector_alu(const nir_instr *instr, const void *_state)
45 {
46    if (instr->type != nir_instr_type_alu)
47       return false;
48 
49    nir_alu_instr *alu = nir_instr_as_alu(instr);
50 
51    /* There is no ALU instruction which has a scalar destination, scalar
52     * src[0], and some other vector source.
53     */
54    return alu->def.num_components > 1 ||
55           nir_op_infos[alu->op].input_sizes[0] > 1;
56 }
57 
58 /* Checks whether all operands of an ALU instruction are swizzled
59  * within the targeted vectorization width.
60  *
61  * The assumption here is that a vecN instruction can only swizzle
62  * within the first N channels of the values it consumes, irrespective
63  * of the capabilities of the instruction which produced those values.
64  * If we assume values are packed consistently (i.e., they always start
65  * at the beginning of a hardware register), we can actually access any
66  * aligned group of N channels so long as we stay within the group.
67  * This means for a vectorization width of 4 that only swizzles from
68  * either [xyzw] or [abcd] etc are allowed.  For a width of 2 these are
69  * swizzles from either [xy] or [zw] etc.
70  */
71 static bool
alu_is_swizzled_in_bounds(const nir_alu_instr * alu,unsigned width)72 alu_is_swizzled_in_bounds(const nir_alu_instr *alu, unsigned width)
73 {
74    for (unsigned i = 0; i < nir_op_infos[alu->op].num_inputs; i++) {
75       if (nir_op_infos[alu->op].input_sizes[i] == 1)
76          continue;
77 
78       unsigned mask = ~(width - 1);
79       for (unsigned j = 1; j < alu->def.num_components; j++) {
80          if ((alu->src[i].swizzle[0] & mask) != (alu->src[i].swizzle[j] & mask))
81             return false;
82       }
83    }
84 
85    return true;
86 }
87 
88 static void
nir_alu_ssa_dest_init(nir_alu_instr * alu,unsigned num_components,unsigned bit_size)89 nir_alu_ssa_dest_init(nir_alu_instr *alu, unsigned num_components,
90                       unsigned bit_size)
91 {
92    nir_def_init(&alu->instr, &alu->def, num_components, bit_size);
93 }
94 
95 static nir_def *
lower_reduction(nir_alu_instr * alu,nir_op chan_op,nir_op merge_op,nir_builder * builder,bool reverse_order)96 lower_reduction(nir_alu_instr *alu, nir_op chan_op, nir_op merge_op,
97                 nir_builder *builder, bool reverse_order)
98 {
99    unsigned num_components = nir_op_infos[alu->op].input_sizes[0];
100 
101    nir_def *last = NULL;
102    for (int i = 0; i < num_components; i++) {
103       int channel = reverse_order ? num_components - 1 - i : i;
104       nir_alu_instr *chan = nir_alu_instr_create(builder->shader, chan_op);
105       nir_alu_ssa_dest_init(chan, 1, alu->def.bit_size);
106       nir_alu_src_copy(&chan->src[0], &alu->src[0]);
107       chan->src[0].swizzle[0] = chan->src[0].swizzle[channel];
108       if (nir_op_infos[chan_op].num_inputs > 1) {
109          assert(nir_op_infos[chan_op].num_inputs == 2);
110          nir_alu_src_copy(&chan->src[1], &alu->src[1]);
111          chan->src[1].swizzle[0] = chan->src[1].swizzle[channel];
112       }
113       chan->exact = alu->exact;
114       chan->fp_fast_math = alu->fp_fast_math;
115 
116       nir_builder_instr_insert(builder, &chan->instr);
117 
118       if (i == 0) {
119          last = &chan->def;
120       } else {
121          last = nir_build_alu(builder, merge_op,
122                               last, &chan->def, NULL, NULL);
123       }
124    }
125 
126    return last;
127 }
128 
129 static inline bool
will_lower_ffma(nir_shader * shader,unsigned bit_size)130 will_lower_ffma(nir_shader *shader, unsigned bit_size)
131 {
132    switch (bit_size) {
133    case 16:
134       return shader->options->lower_ffma16;
135    case 32:
136       return shader->options->lower_ffma32;
137    case 64:
138       return shader->options->lower_ffma64;
139    }
140    unreachable("bad bit size");
141 }
142 
143 static nir_def *
lower_fdot(nir_alu_instr * alu,nir_builder * builder)144 lower_fdot(nir_alu_instr *alu, nir_builder *builder)
145 {
146    /* Reversed order can result in lower instruction count because it
147     * creates more MAD/FMA in the case of fdot(a, vec4(b, 1.0)).
148     * Some games expect xyzw order, so only reverse the order for imprecise fdot.
149     */
150    bool reverse_order = !builder->exact;
151 
152    /* If we don't want to lower ffma, create several ffma instead of fmul+fadd
153     * and fusing later because fusing is not possible for exact fdot instructions.
154     */
155    if (will_lower_ffma(builder->shader, alu->def.bit_size))
156       return lower_reduction(alu, nir_op_fmul, nir_op_fadd, builder, reverse_order);
157 
158    unsigned num_components = nir_op_infos[alu->op].input_sizes[0];
159 
160    nir_def *prev = NULL;
161    for (int i = 0; i < num_components; i++) {
162       int channel = reverse_order ? num_components - 1 - i : i;
163       nir_alu_instr *instr = nir_alu_instr_create(
164          builder->shader, prev ? nir_op_ffma : nir_op_fmul);
165       nir_alu_ssa_dest_init(instr, 1, alu->def.bit_size);
166       for (unsigned j = 0; j < 2; j++) {
167          nir_alu_src_copy(&instr->src[j], &alu->src[j]);
168          instr->src[j].swizzle[0] = alu->src[j].swizzle[channel];
169       }
170       if (i != 0)
171          instr->src[2].src = nir_src_for_ssa(prev);
172       instr->exact = builder->exact;
173       instr->fp_fast_math = builder->fp_fast_math;
174 
175       nir_builder_instr_insert(builder, &instr->instr);
176 
177       prev = &instr->def;
178    }
179 
180    return prev;
181 }
182 
183 static nir_def *
lower_alu_instr_width(nir_builder * b,nir_instr * instr,void * _data)184 lower_alu_instr_width(nir_builder *b, nir_instr *instr, void *_data)
185 {
186    struct alu_width_data *data = _data;
187    nir_alu_instr *alu = nir_instr_as_alu(instr);
188    unsigned num_src = nir_op_infos[alu->op].num_inputs;
189    unsigned i, chan;
190 
191    b->exact = alu->exact;
192    b->fp_fast_math = alu->fp_fast_math;
193 
194    unsigned num_components = alu->def.num_components;
195    unsigned target_width = 1;
196 
197    if (data->cb) {
198       target_width = data->cb(instr, data->data);
199       assert(util_is_power_of_two_or_zero(target_width));
200       if (target_width == 0)
201          return NULL;
202    }
203 
204 #define LOWER_REDUCTION(name, chan, merge) \
205    case name##2:                           \
206    case name##3:                           \
207    case name##4:                           \
208    case name##8:                           \
209    case name##16:                          \
210       return lower_reduction(alu, chan, merge, b, true);
211 
212    switch (alu->op) {
213    case nir_op_vec16:
214    case nir_op_vec8:
215    case nir_op_vec5:
216    case nir_op_vec4:
217    case nir_op_vec3:
218    case nir_op_vec2:
219    case nir_op_cube_amd:
220       /* We don't need to scalarize these ops, they're the ones generated to
221        * group up outputs into a value that can be SSAed.
222        */
223       return NULL;
224 
225    case nir_op_pack_half_2x16: {
226       if (!b->shader->options->lower_pack_half_2x16)
227          return NULL;
228 
229       nir_def *src_vec2 = nir_ssa_for_alu_src(b, alu, 0);
230       return nir_pack_half_2x16_split(b, nir_channel(b, src_vec2, 0),
231                                       nir_channel(b, src_vec2, 1));
232    }
233 
234    case nir_op_unpack_unorm_4x8:
235    case nir_op_unpack_snorm_4x8:
236    case nir_op_unpack_unorm_2x16:
237    case nir_op_unpack_snorm_2x16:
238    case nir_op_mqsad_4x8:
239       /* There is no scalar version of these ops, unless we were to break it
240        * down to bitshifts and math (which is definitely not intended).
241        */
242       return NULL;
243 
244    case nir_op_unpack_half_2x16: {
245       if (!b->shader->options->lower_unpack_half_2x16)
246          return NULL;
247 
248       nir_def *packed = nir_ssa_for_alu_src(b, alu, 0);
249       return nir_vec2(b,
250                       nir_unpack_half_2x16_split_x(b, packed),
251                       nir_unpack_half_2x16_split_y(b, packed));
252    }
253 
254    case nir_op_pack_uvec2_to_uint: {
255       assert(b->shader->options->lower_pack_snorm_2x16 ||
256              b->shader->options->lower_pack_unorm_2x16);
257 
258       nir_def *word = nir_extract_u16(b, nir_ssa_for_alu_src(b, alu, 0),
259                                       nir_imm_int(b, 0));
260       return nir_ior(b, nir_ishl(b, nir_channel(b, word, 1), nir_imm_int(b, 16)),
261                      nir_channel(b, word, 0));
262    }
263 
264    case nir_op_pack_uvec4_to_uint: {
265       assert(b->shader->options->lower_pack_snorm_4x8 ||
266              b->shader->options->lower_pack_unorm_4x8);
267 
268       if (b->shader->options->has_pack_32_4x8)
269          return nir_pack_32_4x8(b, nir_u2u8(b, nir_ssa_for_alu_src(b, alu, 0)));
270 
271       nir_def *byte = nir_extract_u8(b, nir_ssa_for_alu_src(b, alu, 0),
272                                      nir_imm_int(b, 0));
273       return nir_ior(b, nir_ior(b, nir_ishl(b, nir_channel(b, byte, 3), nir_imm_int(b, 24)), nir_ishl(b, nir_channel(b, byte, 2), nir_imm_int(b, 16))),
274                      nir_ior(b, nir_ishl(b, nir_channel(b, byte, 1), nir_imm_int(b, 8)),
275                              nir_channel(b, byte, 0)));
276    }
277 
278    case nir_op_fdph: {
279       nir_def *src0_vec = nir_ssa_for_alu_src(b, alu, 0);
280       nir_def *src1_vec = nir_ssa_for_alu_src(b, alu, 1);
281 
282       /* Only use reverse order for imprecise fdph, see explanation in lower_fdot. */
283       bool reverse_order = !b->exact;
284       if (will_lower_ffma(b->shader, alu->def.bit_size)) {
285          nir_def *sum[4];
286          for (unsigned i = 0; i < 3; i++) {
287             int dest = reverse_order ? 3 - i : i;
288             sum[dest] = nir_fmul(b, nir_channel(b, src0_vec, i),
289                                  nir_channel(b, src1_vec, i));
290          }
291          sum[reverse_order ? 0 : 3] = nir_channel(b, src1_vec, 3);
292 
293          return nir_fadd(b, nir_fadd(b, nir_fadd(b, sum[0], sum[1]), sum[2]), sum[3]);
294       } else if (reverse_order) {
295          nir_def *sum = nir_channel(b, src1_vec, 3);
296          for (int i = 2; i >= 0; i--)
297             sum = nir_ffma(b, nir_channel(b, src0_vec, i), nir_channel(b, src1_vec, i), sum);
298          return sum;
299       } else {
300          nir_def *sum = nir_fmul(b, nir_channel(b, src0_vec, 0), nir_channel(b, src1_vec, 0));
301          sum = nir_ffma(b, nir_channel(b, src0_vec, 1), nir_channel(b, src1_vec, 1), sum);
302          sum = nir_ffma(b, nir_channel(b, src0_vec, 2), nir_channel(b, src1_vec, 2), sum);
303          return nir_fadd(b, sum, nir_channel(b, src1_vec, 3));
304       }
305    }
306 
307    case nir_op_pack_64_2x32: {
308       if (!b->shader->options->lower_pack_64_2x32)
309          return NULL;
310 
311       nir_def *src_vec2 = nir_ssa_for_alu_src(b, alu, 0);
312       return nir_pack_64_2x32_split(b, nir_channel(b, src_vec2, 0),
313                                     nir_channel(b, src_vec2, 1));
314    }
315    case nir_op_pack_64_4x16: {
316       if (!b->shader->options->lower_pack_64_4x16)
317          return NULL;
318 
319       nir_def *src_vec4 = nir_ssa_for_alu_src(b, alu, 0);
320       nir_def *xy = nir_pack_32_2x16_split(b, nir_channel(b, src_vec4, 0),
321                                            nir_channel(b, src_vec4, 1));
322       nir_def *zw = nir_pack_32_2x16_split(b, nir_channel(b, src_vec4, 2),
323                                            nir_channel(b, src_vec4, 3));
324 
325       return nir_pack_64_2x32_split(b, xy, zw);
326    }
327    case nir_op_pack_32_2x16: {
328       if (!b->shader->options->lower_pack_32_2x16)
329          return NULL;
330 
331       nir_def *src_vec2 = nir_ssa_for_alu_src(b, alu, 0);
332       return nir_pack_32_2x16_split(b, nir_channel(b, src_vec2, 0),
333                                     nir_channel(b, src_vec2, 1));
334    }
335    case nir_op_unpack_64_2x32:
336    case nir_op_unpack_64_4x16:
337    case nir_op_unpack_32_2x16:
338    case nir_op_unpack_32_4x8:
339    case nir_op_unpack_double_2x32_dxil:
340       return NULL;
341 
342    case nir_op_fdot2:
343    case nir_op_fdot3:
344    case nir_op_fdot4:
345    case nir_op_fdot8:
346    case nir_op_fdot16:
347       return lower_fdot(alu, b);
348 
349       LOWER_REDUCTION(nir_op_ball_fequal, nir_op_feq, nir_op_iand);
350       LOWER_REDUCTION(nir_op_ball_iequal, nir_op_ieq, nir_op_iand);
351       LOWER_REDUCTION(nir_op_bany_fnequal, nir_op_fneu, nir_op_ior);
352       LOWER_REDUCTION(nir_op_bany_inequal, nir_op_ine, nir_op_ior);
353       LOWER_REDUCTION(nir_op_b8all_fequal, nir_op_feq8, nir_op_iand);
354       LOWER_REDUCTION(nir_op_b8all_iequal, nir_op_ieq8, nir_op_iand);
355       LOWER_REDUCTION(nir_op_b8any_fnequal, nir_op_fneu8, nir_op_ior);
356       LOWER_REDUCTION(nir_op_b8any_inequal, nir_op_ine8, nir_op_ior);
357       LOWER_REDUCTION(nir_op_b16all_fequal, nir_op_feq16, nir_op_iand);
358       LOWER_REDUCTION(nir_op_b16all_iequal, nir_op_ieq16, nir_op_iand);
359       LOWER_REDUCTION(nir_op_b16any_fnequal, nir_op_fneu16, nir_op_ior);
360       LOWER_REDUCTION(nir_op_b16any_inequal, nir_op_ine16, nir_op_ior);
361       LOWER_REDUCTION(nir_op_b32all_fequal, nir_op_feq32, nir_op_iand);
362       LOWER_REDUCTION(nir_op_b32all_iequal, nir_op_ieq32, nir_op_iand);
363       LOWER_REDUCTION(nir_op_b32any_fnequal, nir_op_fneu32, nir_op_ior);
364       LOWER_REDUCTION(nir_op_b32any_inequal, nir_op_ine32, nir_op_ior);
365       LOWER_REDUCTION(nir_op_fall_equal, nir_op_seq, nir_op_fmin);
366       LOWER_REDUCTION(nir_op_fany_nequal, nir_op_sne, nir_op_fmax);
367 
368    default:
369       break;
370    }
371 
372    if (num_components == 1)
373       return NULL;
374 
375    if (num_components <= target_width) {
376       /* If the ALU instr is swizzled outside the target width,
377        * reduce the target width.
378        */
379       if (alu_is_swizzled_in_bounds(alu, target_width))
380          return NULL;
381       else
382          target_width = DIV_ROUND_UP(num_components, 2);
383    }
384 
385    nir_alu_instr *vec = nir_alu_instr_create(b->shader, nir_op_vec(num_components));
386 
387    for (chan = 0; chan < num_components; chan += target_width) {
388       unsigned components = MIN2(target_width, num_components - chan);
389       nir_alu_instr *lower = nir_alu_instr_create(b->shader, alu->op);
390 
391       for (i = 0; i < num_src; i++) {
392          nir_alu_src_copy(&lower->src[i], &alu->src[i]);
393 
394          /* We only handle same-size-as-dest (input_sizes[] == 0) or scalar
395           * args (input_sizes[] == 1).
396           */
397          assert(nir_op_infos[alu->op].input_sizes[i] < 2);
398          for (int j = 0; j < components; j++) {
399             unsigned src_chan = nir_op_infos[alu->op].input_sizes[i] == 1 ? 0 : chan + j;
400             lower->src[i].swizzle[j] = alu->src[i].swizzle[src_chan];
401          }
402       }
403 
404       nir_alu_ssa_dest_init(lower, components, alu->def.bit_size);
405       lower->exact = alu->exact;
406       lower->fp_fast_math = alu->fp_fast_math;
407 
408       for (i = 0; i < components; i++) {
409          vec->src[chan + i].src = nir_src_for_ssa(&lower->def);
410          vec->src[chan + i].swizzle[0] = i;
411       }
412 
413       nir_builder_instr_insert(b, &lower->instr);
414    }
415 
416    return nir_builder_alu_instr_finish_and_insert(b, vec);
417 }
418 
419 bool
nir_lower_alu_width(nir_shader * shader,nir_vectorize_cb cb,const void * _data)420 nir_lower_alu_width(nir_shader *shader, nir_vectorize_cb cb, const void *_data)
421 {
422    struct alu_width_data data = {
423       .cb = cb,
424       .data = _data,
425    };
426 
427    return nir_shader_lower_instructions(shader,
428                                         inst_is_vector_alu,
429                                         lower_alu_instr_width,
430                                         &data);
431 }
432 
433 struct alu_to_scalar_data {
434    nir_instr_filter_cb cb;
435    const void *data;
436 };
437 
438 static uint8_t
scalar_cb(const nir_instr * instr,const void * data)439 scalar_cb(const nir_instr *instr, const void *data)
440 {
441    /* return vectorization-width = 1 for filtered instructions */
442    const struct alu_to_scalar_data *filter = data;
443    return filter->cb(instr, filter->data) ? 1 : 0;
444 }
445 
446 bool
nir_lower_alu_to_scalar(nir_shader * shader,nir_instr_filter_cb cb,const void * _data)447 nir_lower_alu_to_scalar(nir_shader *shader, nir_instr_filter_cb cb, const void *_data)
448 {
449    struct alu_to_scalar_data data = {
450       .cb = cb,
451       .data = _data,
452    };
453 
454    return nir_lower_alu_width(shader, cb ? scalar_cb : NULL, &data);
455 }
456 
457 static bool
lower_alu_vec8_16_src(nir_builder * b,nir_alu_instr * alu,void * _data)458 lower_alu_vec8_16_src(nir_builder *b, nir_alu_instr *alu, void *_data)
459 {
460    const nir_op_info *info = &nir_op_infos[alu->op];
461 
462    bool changed = false;
463    b->cursor = nir_before_instr(&alu->instr);
464    for (int i = 0; i < info->num_inputs; i++) {
465       if (alu->src[i].src.ssa->num_components < 8 || info->input_sizes[i])
466          continue;
467 
468       changed = true;
469       nir_def *comps[4];
470       for (int c = 0; c < alu->def.num_components; c++) {
471          unsigned swizzle = alu->src[i].swizzle[c];
472          alu->src[i].swizzle[c] = c;
473 
474          nir_const_value *const_val = nir_src_as_const_value(alu->src[i].src);
475          if (const_val) {
476             comps[c] = nir_build_imm(b, 1, alu->src[i].src.ssa->bit_size, &const_val[swizzle]);
477          } else {
478             comps[c] = nir_swizzle(b, alu->src[i].src.ssa, &swizzle, 1);
479          }
480       }
481       nir_def *src = nir_vec(b, comps, alu->def.num_components);
482       nir_src_rewrite(&alu->src[i].src, src);
483    }
484 
485    return changed;
486 }
487 
488 bool
nir_lower_alu_vec8_16_srcs(nir_shader * shader)489 nir_lower_alu_vec8_16_srcs(nir_shader *shader)
490 {
491    return nir_shader_alu_pass(shader, lower_alu_vec8_16_src,
492                               nir_metadata_control_flow,
493                               NULL);
494 }
495