[Mesa-dev] [PATCH 1/4] nir: Add a small pass to rematerialize derefs per-block

[Iago Toral]

Hi Jason,

I left a few comments in patches 1 and 4, feel free to ignore them if
you think they are not relevant. Either way the series is:

Reviewed-by: Iago Toral Quiroga <itoral at igalia.com>

On Mon, 2018-09-17 at 09:43 -0500, Jason Ekstrand wrote:
> This pass re-materializes deref instructions on a per-block basis to
> ensure that every use of a deref occurs in the same block as the
> instruction which uses it.
> ---
>  src/compiler/nir/nir.h       |   1 +
>  src/compiler/nir/nir_deref.c | 132
> +++++++++++++++++++++++++++++++++++
>  2 files changed, 133 insertions(+)
> 
> diff --git a/src/compiler/nir/nir.h b/src/compiler/nir/nir.h
> index 599f469a714..e0df95c391c 100644
> --- a/src/compiler/nir/nir.h
> +++ b/src/compiler/nir/nir.h
> @@ -3012,6 +3012,7 @@ bool nir_convert_from_ssa(nir_shader *shader,
> bool phi_webs_only);
>  
>  bool nir_lower_phis_to_regs_block(nir_block *block);
>  bool nir_lower_ssa_defs_to_regs_block(nir_block *block);
> +bool nir_rematerialize_derefs_in_use_blocks_impl(nir_function_impl
> *impl);
>  
>  bool nir_opt_algebraic(nir_shader *shader);
>  bool nir_opt_algebraic_before_ffma(nir_shader *shader);
> diff --git a/src/compiler/nir/nir_deref.c
> b/src/compiler/nir/nir_deref.c
> index 097ea8f1046..87a54790c95 100644
> --- a/src/compiler/nir/nir_deref.c
> +++ b/src/compiler/nir/nir_deref.c
> @@ -24,6 +24,7 @@
>  #include "nir.h"
>  #include "nir_builder.h"
>  #include "nir_deref.h"
> +#include "util/hash_table.h"
>  
>  void
>  nir_deref_path_init(nir_deref_path *path,
> @@ -379,3 +380,134 @@ nir_compare_derefs(nir_deref_instr *a,
> nir_deref_instr *b)
>  
>     return result;
>  }
> +
> +struct rematerialize_deref_state {
> +   bool progress;
> +   nir_builder builder;
> +   nir_block *block;
> +   struct hash_table *cache;
> +};
> +
> +static nir_deref_instr *
> +rematerialize_deref_in_block(nir_deref_instr *deref,
> +                             struct rematerialize_deref_state
> *state)
> +{
> +   if (deref->instr.block == state->block)
> +      return deref;
> +
> +   if (!state->cache) {
> +      state->cache= _mesa_hash_table_create(NULL,
> _mesa_hash_pointer,
> +                                            _mesa_key_pointer_equal)
> ;
> +   }
> +
> +   struct hash_entry *cached = _mesa_hash_table_search(state->cache, 
> deref);
> +   if (cached)
> +      return cached->data;
> +
> +   nir_builder *b = &state->builder;
> +   nir_deref_instr *new_deref =
> +      nir_deref_instr_create(b->shader, deref->deref_type);
> +   new_deref->mode = deref->mode;
> +   new_deref->type = deref->type;
> +
> +   if (deref->deref_type == nir_deref_type_var) {
> +      new_deref->var = deref->var;
> +   } else {
> +      nir_deref_instr *parent = nir_src_as_deref(deref->parent);
> +      if (parent) {
> +         parent = rematerialize_deref_in_block(parent, state);
> +         new_deref->parent = nir_src_for_ssa(&parent->dest.ssa);
> +      } else {
> +         nir_src_copy(&new_deref->parent, &deref->parent,
> new_deref);
> +      }
> +   }
> +
> +   switch (deref->deref_type) {
> +   case nir_deref_type_var:
> +   case nir_deref_type_array_wildcard:
> +   case nir_deref_type_cast:
> +      /* Nothing more to do */
> +      break;
> +
> +   case nir_deref_type_array:
> +      assert(!nir_src_as_deref(deref->arr.index));
> +      nir_src_copy(&new_deref->arr.index, &deref->arr.index,
> new_deref);
> +      break;
> +
> +   case nir_deref_type_struct:
> +      new_deref->strct.index = deref->strct.index;
> +      break;
> +
> +   default:
> +      unreachable("Invalid deref instruction type");
> +   }
> +
> +   nir_ssa_dest_init(&new_deref->instr, &new_deref->dest,
> +                     deref->dest.ssa.num_components,
> +                     deref->dest.ssa.bit_size,
> +                     deref->dest.ssa.name);
> +   nir_builder_instr_insert(b, &new_deref->instr);
> +
> +   return new_deref;
> +}
> +
> +static bool
> +rematerialize_deref_src(nir_src *src, void *_state)
> +{
> +   struct rematerialize_deref_state *state = _state;
> +
> +   nir_deref_instr *deref = nir_src_as_deref(*src);
> +   if (!deref)
> +      return true;
> +
> +   nir_deref_instr *block_deref =
> rematerialize_deref_in_block(deref, state);
> +
> +   nir_instr_rewrite_src(src->parent_instr, src,
> +                         nir_src_for_ssa(&block_deref->dest.ssa));
> +   nir_deref_instr_remove_if_unused(deref);
> +   state->progress = true;
> +
> +   return true;
> +}
> +
> +/** Re-materialize derefs in every block
> + *
> + * This pass re-materializes deref instructions in every block in
> which it is
> + * used.  After this pass has been run, every use of a deref will be
> of a
> + * deref in the same block as the use.  Also, all unused derefs will
> be
> + * deleted as a side-effect.
> + */
> +bool
> +nir_rematerialize_derefs_in_use_blocks_impl(nir_function_impl *impl)
> +{
> +   struct rematerialize_deref_state state = { };
> +   nir_builder_init(&state.builder, impl);
> +
> +   nir_foreach_block(block, impl) {
> +      state.block = block;
> +
> +      /* Start each block with a fresh cache */
> +      if (state.cache)
> +         _mesa_hash_table_clear(state.cache, NULL);
> +
> +      nir_foreach_instr_safe(instr, block) {
> +         if (instr->type == nir_instr_type_deref) {
> +            nir_deref_instr_remove_if_unused(nir_instr_as_deref(inst
> r));
> +            continue;
> +         }
> +
> +         state.builder.cursor = nir_before_instr(instr);
> +         nir_foreach_src(instr, rematerialize_deref_src, &state);
> +      }
> +
> +#ifndef NDEBUG
> +      nir_if *following_if = nir_block_get_following_if(block);
> +      if (following_if)
> +         assert(!nir_src_as_deref(following_if->condition));
> +#endif
> +   }
> +
> +   _mesa_hash_table_destroy(state.cache, NULL);
> +
> +   return state.progress;
> +}