nir/find_array_copies: Handle cast derefs

Reviewed-by: Caio Marcelo de Oliveira Filho <caio.oliveira@intel.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/6210>

1 files changed, 59 insertions, 10 deletions

diff --git a/src/compiler/nir/nir_opt_find_array_copies.c b/src/compiler/nir/nir_opt_find_array_copies.c
index 0739a1f3719..0e50ca18af9 100644
--- a/src/compiler/nir/nir_opt_find_array_copies.c
+++ b/src/compiler/nir/nir_opt_find_array_copies.c
@@ -53,7 +53,9 @@ struct match_node {
 
 struct match_state {
    /* Map from nir_variable * -> match_node */
-   struct hash_table *table;
+   struct hash_table *var_nodes;
+   /* Map from cast nir_deref_instr * -> match_node */
+   struct hash_table *cast_nodes;
 
    unsigned cur_instr;
 
@@ -89,12 +91,25 @@ node_for_deref(nir_deref_instr *instr, struct match_node *parent,
    unsigned idx;
    switch (instr->deref_type) {
    case nir_deref_type_var: {
-      struct hash_entry *entry = _mesa_hash_table_search(state->table, instr->var);
+      struct hash_entry *entry =
+         _mesa_hash_table_search(state->var_nodes, instr->var);
       if (entry) {
          return entry->data;
       } else {
          struct match_node *node = create_match_node(instr->type, state);
-         _mesa_hash_table_insert(state->table, instr->var, node);
+         _mesa_hash_table_insert(state->var_nodes, instr->var, node);
+         return node;
+      }
+   }
+
+   case nir_deref_type_cast: {
+      struct hash_entry *entry =
+         _mesa_hash_table_search(state->cast_nodes, instr);
+      if (entry) {
+         return entry->data;
+      } else {
+         struct match_node *node = create_match_node(instr->type, state);
+         _mesa_hash_table_insert(state->cast_nodes, instr, node);
          return node;
       }
    }
@@ -223,6 +238,17 @@ _foreach_aliasing(nir_deref_instr **deref, match_cb cb,
    }
 }
 
+static void
+_foreach_child(match_cb cb, struct match_node *node, struct match_state *state)
+{
+   if (node->num_children == 0) {
+      cb(node, state);
+   } else {
+      for (unsigned i = 0; i < node->num_children; i++)
+         _foreach_child(cb, node->children[i], state);
+   }
+}
+
 /* Given a deref path, find all the leaf deref nodes that alias it. */
 
 static void
@@ -230,11 +256,32 @@ foreach_aliasing_node(nir_deref_path *path,
                       match_cb cb,
                       struct match_state *state)
 {
-   assert(path->path[0]->deref_type == nir_deref_type_var);
-   struct hash_entry *entry = _mesa_hash_table_search(state->table,
-                                                      path->path[0]->var);
-   if (entry)
-      _foreach_aliasing(&path->path[1], cb, entry->data, state);
+   if (path->path[0]->deref_type == nir_deref_type_var) {
+      struct hash_entry *entry = _mesa_hash_table_search(state->var_nodes,
+                                                         path->path[0]->var);
+      if (entry)
+         _foreach_aliasing(&path->path[1], cb, entry->data, state);
+
+      hash_table_foreach(state->cast_nodes, entry)
+         _foreach_child(cb, entry->data, state);
+   } else {
+      /* Casts automatically alias anything that isn't a cast */
+      assert(path->path[0]->deref_type == nir_deref_type_cast);
+      hash_table_foreach(state->var_nodes, entry)
+         _foreach_child(cb, entry->data, state);
+
+      /* Casts alias other casts if the casts are different or if they're the
+       * same and the path from the cast may alias as per the usual rules.
+       */
+      hash_table_foreach(state->cast_nodes, entry) {
+         const nir_deref_instr *cast = entry->key;
+         assert(cast->deref_type == nir_deref_type_cast);
+         if (cast == path->path[0])
+            _foreach_aliasing(&path->path[1], cb, entry->data, state);
+         else
+            _foreach_child(cb, entry->data, state);
+      }
+   }
 }
 
 static nir_deref_instr *
@@ -472,7 +519,8 @@ opt_find_array_copies_block(nir_builder *b, nir_block *block,
 
    unsigned next_index = 0;
 
-   _mesa_hash_table_clear(state->table, NULL);
+   _mesa_hash_table_clear(state->var_nodes, NULL);
+   _mesa_hash_table_clear(state->cast_nodes, NULL);
 
    nir_foreach_instr(instr, block) {
       if (instr->type != nir_instr_type_intrinsic)
@@ -576,7 +624,8 @@ opt_find_array_copies_impl(nir_function_impl *impl)
 
    struct match_state s;
    s.dead_ctx = ralloc_context(NULL);
-   s.table = _mesa_pointer_hash_table_create(s.dead_ctx);
+   s.var_nodes = _mesa_pointer_hash_table_create(s.dead_ctx);
+   s.cast_nodes = _mesa_pointer_hash_table_create(s.dead_ctx);
    nir_builder_init(&s.builder, impl);
 
    nir_foreach_block(block, impl) {