1/*
2 * Copyright © 2021 Valve 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 * Authors:
24 *    Timur Kristóf
25 *
26 */
27
28#include "nir.h"
29#include "nir_builder.h"
30
31typedef struct
32{
33   struct hash_table *range_ht;
34} opt_offsets_state;
35
36static nir_ssa_def *
37try_extract_const_addition(nir_builder *b, nir_instr *instr, opt_offsets_state *state, unsigned *out_const)
38{
39   if (instr->type != nir_instr_type_alu)
40      return NULL;
41
42   nir_alu_instr *alu = nir_instr_as_alu(instr);
43   if (alu->op != nir_op_iadd ||
44       !nir_alu_src_is_trivial_ssa(alu, 0) ||
45       !nir_alu_src_is_trivial_ssa(alu, 1))
46      return NULL;
47
48   if (!alu->no_unsigned_wrap) {
49      if (!state->range_ht) {
50         /* Cache for nir_unsigned_upper_bound */
51         state->range_ht = _mesa_pointer_hash_table_create(NULL);
52      }
53
54      /* Check if there can really be an unsigned wrap. */
55      nir_ssa_scalar src0 = {alu->src[0].src.ssa, 0};
56      nir_ssa_scalar src1 = {alu->src[1].src.ssa, 0};
57      uint32_t ub0 = nir_unsigned_upper_bound(b->shader, state->range_ht, src0, NULL);
58      uint32_t ub1 = nir_unsigned_upper_bound(b->shader, state->range_ht, src1, NULL);
59
60      if ((UINT32_MAX - ub0) < ub1)
61         return NULL;
62
63      /* We proved that unsigned wrap won't be possible, so we can set the flag too. */
64      alu->no_unsigned_wrap = true;
65   }
66
67   for (unsigned i = 0; i < 2; ++i) {
68      if (nir_src_is_const(alu->src[i].src)) {
69         *out_const += nir_src_as_uint(alu->src[i].src);
70         return alu->src[1 - i].src.ssa;
71      }
72
73      nir_ssa_def *replace_src = try_extract_const_addition(b, alu->src[0].src.ssa->parent_instr, state, out_const);
74      if (replace_src) {
75         b->cursor = nir_before_instr(&alu->instr);
76         return nir_iadd(b, replace_src, alu->src[1 - i].src.ssa);
77      }
78   }
79
80   return NULL;
81}
82
83static bool
84try_fold_load_store(nir_builder *b,
85                    nir_intrinsic_instr *intrin,
86                    opt_offsets_state *state,
87                    unsigned offset_src_idx)
88{
89   /* Assume that BASE is the constant offset of a load/store.
90    * Try to constant-fold additions to the offset source
91    * into the actual const offset of the instruction.
92    */
93
94   unsigned off_const = nir_intrinsic_base(intrin);
95   nir_src *off_src = &intrin->src[offset_src_idx];
96   nir_ssa_def *replace_src = NULL;
97
98   if (!off_src->is_ssa || off_src->ssa->bit_size != 32)
99      return false;
100
101   if (!nir_src_is_const(*off_src)) {
102      nir_ssa_def *r = off_src->ssa;
103      while ((r = try_extract_const_addition(b, r->parent_instr, state, &off_const)))
104         replace_src = r;
105   } else if (nir_src_as_uint(*off_src)) {
106      off_const += nir_src_as_uint(*off_src);
107      b->cursor = nir_before_instr(&intrin->instr);
108      replace_src = nir_imm_zero(b, off_src->ssa->num_components, off_src->ssa->bit_size);
109   }
110
111   if (!replace_src)
112      return false;
113
114   nir_instr_rewrite_src(&intrin->instr, &intrin->src[offset_src_idx], nir_src_for_ssa(replace_src));
115   nir_intrinsic_set_base(intrin, off_const);
116   return true;
117}
118
119static bool
120process_instr(nir_builder *b, nir_instr *instr, void *s)
121{
122   if (instr->type != nir_instr_type_intrinsic)
123      return false;
124
125   opt_offsets_state *state = (opt_offsets_state *) s;
126   nir_intrinsic_instr *intrin = nir_instr_as_intrinsic(instr);
127
128   switch (intrin->intrinsic) {
129   case nir_intrinsic_load_shared:
130      return try_fold_load_store(b, intrin, state, 0);
131   case nir_intrinsic_store_shared:
132      return try_fold_load_store(b, intrin, state, 1);
133   case nir_intrinsic_load_buffer_amd:
134      return try_fold_load_store(b, intrin, state, 1);
135   case nir_intrinsic_store_buffer_amd:
136      return try_fold_load_store(b, intrin, state, 2);
137   default:
138      return false;
139   }
140
141   unreachable("Can't reach here.");
142}
143
144bool
145nir_opt_offsets(nir_shader *shader)
146{
147   opt_offsets_state state;
148   state.range_ht = NULL;
149
150   bool p = nir_shader_instructions_pass(shader, process_instr,
151                                         nir_metadata_block_index |
152                                         nir_metadata_dominance,
153                                         &state);
154
155   if (state.range_ht)
156      _mesa_hash_table_destroy(state.range_ht, NULL);
157
158
159   return p;
160}
161