src/compiler/nir/nir_opt_large_constants.c

   1 /*
   2  * Copyright © 2018 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 #include "nir.h"
  25 #include "nir_builder.h"
  26 #include "nir_deref.h"
  27
  28 struct var_info {
  29    bool is_constant;
  30    bool found_read;
  31 };
  32
  33 static nir_ssa_def *
  34 build_constant_load(nir_builder *b, nir_deref_instr *deref,
  35                     glsl_type_size_align_func size_align)
  36 {
  37    nir_variable *var = nir_deref_instr_get_variable(deref);
  38
  39    const unsigned bit_size = glsl_get_bit_size(deref->type);
  40    const unsigned num_components = glsl_get_vector_elements(deref->type);
  41
  42    UNUSED unsigned var_size, var_align;
  43    size_align(var->type, &var_size, &var_align);
  44    assert(var->data.location % var_align == 0);
  45
  46    nir_intrinsic_instr *load =
  47       nir_intrinsic_instr_create(b->shader, nir_intrinsic_load_constant);
  48    load->num_components = num_components;
  49    nir_intrinsic_set_base(load, var->data.location);
  50    nir_intrinsic_set_range(load, var_size);
  51    load->src[0] = nir_src_for_ssa(nir_build_deref_offset(b, deref, size_align));
  52    nir_ssa_dest_init(&load->instr, &load->dest,
  53                      num_components, bit_size, NULL);
  54    nir_builder_instr_insert(b, &load->instr);
  55
  56    return &load->dest.ssa;
  57 }
  58
  59 static void
  60 handle_constant_store(nir_builder *b, nir_intrinsic_instr *store,
  61                       glsl_type_size_align_func size_align)
  62 {
  63    nir_deref_instr *deref = nir_src_as_deref(store->src[0]);
  64    assert(!nir_deref_instr_has_indirect(deref));
  65
  66    nir_variable *var = nir_deref_instr_get_variable(deref);
  67
  68    const unsigned bit_size = glsl_get_bit_size(deref->type);
  69    const unsigned num_components = glsl_get_vector_elements(deref->type);
  70
  71    char *dst = (char *)b->shader->constant_data +
  72                var->data.location +
  73                nir_deref_instr_get_const_offset(deref, size_align);
  74
  75    nir_const_value *val = nir_src_as_const_value(store->src[1]);
  76    switch (bit_size) {
  77    case 8:
  78       for (unsigned i = 0; i < num_components; i++)
  79          ((uint8_t *)dst)[i] = val->u8[i];
  80       break;
  81
  82    case 16:
  83       for (unsigned i = 0; i < num_components; i++)
  84          ((uint16_t *)dst)[i] = val->u16[i];
  85       break;
  86
  87    case 32:
  88       for (unsigned i = 0; i < num_components; i++)
  89          ((uint32_t *)dst)[i] = val->u32[i];
  90       break;
  91
  92    case 64:
  93       for (unsigned i = 0; i < num_components; i++)
  94          ((uint64_t *)dst)[i] = val->u64[i];
  95       break;
  96
  97    default:
  98       unreachable("Invalid bit size");
  99    }
 100 }
 101
 102 /** Lower large constant variables to shader constant data
 103  *
 104  * This pass looks for large (type_size(var->type) > threshold) variables
 105  * which are statically constant and moves them into shader constant data.
 106  * This is especially useful when large tables are baked into the shader
 107  * source code because they can be moved into a UBO by the driver to reduce
 108  * register pressure and make indirect access cheaper.
 109  */
 110 bool
 111 nir_opt_large_constants(nir_shader *shader,
 112                         glsl_type_size_align_func size_align,
 113                         unsigned threshold)
 114 {
 115    /* Default to a natural alignment if none is provided */
 116    if (size_align == NULL)
 117       size_align = glsl_get_natural_size_align_bytes;
 118
 119    /* This only works with a single entrypoint */
 120    nir_function_impl *impl = nir_shader_get_entrypoint(shader);
 121
 122    /* This pass can only be run once */
 123    assert(shader->constant_data == NULL && shader->constant_data_size == 0);
 124
 125    /* The index parameter is unused for local variables so we'll use it for
 126     * indexing into our array of variable metadata.
 127     */
 128    unsigned num_locals = 0;
 129    nir_foreach_variable(var, &impl->locals)
 130       var->data.index = num_locals++;
 131
 132    struct var_info *var_infos = malloc(num_locals * sizeof(struct var_info));
 133    for (unsigned i = 0; i < num_locals; i++) {
 134       var_infos[i] = (struct var_info) {
 135          .is_constant = true,
 136          .found_read = false,
 137       };
 138    }
 139
 140    /* First, walk through the shader and figure out what variables we can
 141     * lower to the constant blob.
 142     */
 143    bool first_block = true;
 144    nir_foreach_block(block, impl) {
 145       nir_foreach_instr(instr, block) {
 146          if (instr->type != nir_instr_type_intrinsic)
 147             continue;
 148
 149          nir_intrinsic_instr *intrin = nir_instr_as_intrinsic(instr);
 150
 151          bool src_is_const = false;
 152          nir_deref_instr *src_deref = NULL, *dst_deref = NULL;
 153          switch (intrin->intrinsic) {
 154          case nir_intrinsic_store_deref:
 155             dst_deref = nir_src_as_deref(intrin->src[0]);
 156             src_is_const = nir_src_is_const(intrin->src[1]);
 157             break;
 158
 159          case nir_intrinsic_load_deref:
 160             src_deref = nir_src_as_deref(intrin->src[0]);
 161             break;
 162
 163          case nir_intrinsic_copy_deref:
 164             /* We always assume the src and therefore the dst are not
 165              * constants here. Copy and constant propagation passes should
 166              * have taken care of this in most cases anyway.
 167              */
 168             dst_deref = nir_src_as_deref(intrin->src[0]);
 169             src_deref = nir_src_as_deref(intrin->src[1]);
 170             src_is_const = false;
 171             break;
 172
 173          default:
 174             continue;
 175          }
 176
 177          if (dst_deref && dst_deref->mode == nir_var_local) {
 178             nir_variable *var = nir_deref_instr_get_variable(dst_deref);
 179             assert(var->data.mode == nir_var_local);
 180
 181             /* We only consider variables constant if they only have constant
 182              * stores, all the stores come before any reads, and all stores
 183              * come in the first block.  We also can't handle indirect stores.
 184              */
 185             struct var_info *info = &var_infos[var->data.index];
 186             if (!src_is_const || info->found_read || !first_block ||
 187                 nir_deref_instr_has_indirect(dst_deref))
 188                info->is_constant = false;
 189          }
 190
 191          if (src_deref && src_deref->mode == nir_var_local) {
 192             nir_variable *var = nir_deref_instr_get_variable(src_deref);
 193             assert(var->data.mode == nir_var_local);
 194
 195             var_infos[var->data.index].found_read = true;
 196          }
 197       }
 198       first_block = false;
 199    }
 200
 201    shader->constant_data_size = 0;
 202    nir_foreach_variable(var, &impl->locals) {
 203       struct var_info *info = &var_infos[var->data.index];
 204       if (!info->is_constant)
 205          continue;
 206
 207       unsigned var_size, var_align;
 208       size_align(var->type, &var_size, &var_align);
 209       if (var_size <= threshold || !info->found_read) {
 210          /* Don't bother lowering small stuff or data that's never read */
 211          info->is_constant = false;
 212          continue;
 213       }
 214
 215       var->data.location = ALIGN_POT(shader->constant_data_size, var_align);
 216       shader->constant_data_size = var->data.location + var_size;
 217    }
 218
 219    if (shader->constant_data_size == 0) {
 220       free(var_infos);
 221       return false;
 222    }
 223
 224    shader->constant_data = rzalloc_size(shader, shader->constant_data_size);
 225
 226    nir_builder b;
 227    nir_builder_init(&b, impl);
 228
 229    nir_foreach_block(block, impl) {
 230       nir_foreach_instr_safe(instr, block) {
 231          if (instr->type != nir_instr_type_intrinsic)
 232             continue;
 233
 234          nir_intrinsic_instr *intrin = nir_instr_as_intrinsic(instr);
 235
 236          switch (intrin->intrinsic) {
 237          case nir_intrinsic_load_deref: {
 238             nir_deref_instr *deref = nir_src_as_deref(intrin->src[0]);
 239             if (deref->mode != nir_var_local)
 240                continue;
 241
 242             nir_variable *var = nir_deref_instr_get_variable(deref);
 243             struct var_info *info = &var_infos[var->data.index];
 244             if (info->is_constant) {
 245                b.cursor = nir_after_instr(&intrin->instr);
 246                nir_ssa_def *val = build_constant_load(&b, deref, size_align);
 247                nir_ssa_def_rewrite_uses(&intrin->dest.ssa,
 248                                         nir_src_for_ssa(val));
 249                nir_instr_remove(&intrin->instr);
 250                nir_deref_instr_remove_if_unused(deref);
 251             }
 252             break;
 253          }
 254
 255          case nir_intrinsic_store_deref: {
 256             nir_deref_instr *deref = nir_src_as_deref(intrin->src[0]);
 257             if (deref->mode != nir_var_local)
 258                continue;
 259
 260             nir_variable *var = nir_deref_instr_get_variable(deref);
 261             struct var_info *info = &var_infos[var->data.index];
 262             if (info->is_constant) {
 263                b.cursor = nir_after_instr(&intrin->instr);
 264                handle_constant_store(&b, intrin, size_align);
 265                nir_instr_remove(&intrin->instr);
 266                nir_deref_instr_remove_if_unused(deref);
 267             }
 268             break;
 269          }
 270
 271          case nir_intrinsic_copy_deref: {
 272             nir_deref_instr *deref = nir_src_as_deref(intrin->src[1]);
 273             if (deref->mode != nir_var_local)
 274                continue;
 275
 276             nir_variable *var = nir_deref_instr_get_variable(deref);
 277             struct var_info *info = &var_infos[var->data.index];
 278             if (info->is_constant) {
 279                b.cursor = nir_after_instr(&intrin->instr);
 280                nir_ssa_def *val = build_constant_load(&b, deref, size_align);
 281                nir_store_deref(&b, nir_src_as_deref(intrin->src[0]), val, ~0);
 282                nir_instr_remove(&intrin->instr);
 283                nir_deref_instr_remove_if_unused(deref);
 284             }
 285             break;
 286          }
 287
 288          default:
 289             continue;
 290          }
 291       }
 292    }
 293
 294    /* Clean up the now unused variables */
 295    nir_foreach_variable_safe(var, &impl->locals) {
 296       if (var_infos[var->data.index].is_constant)
 297          exec_node_remove(&var->node);
 298    }
 299
 300    free(var_infos);
 301
 302    nir_metadata_preserve(impl, nir_metadata_block_index |
 303                                nir_metadata_dominance);
 304    return true;
 305 }