[Mesa-dev] [PATCH 10/12] nir: add helper for cloning loops

[Timothy Arceri]

---
 src/compiler/nir/nir.h       |  3 +++
 src/compiler/nir/nir_clone.c | 64 +++++++++++++++++++++++++++++++++++++++-----
 2 files changed, 60 insertions(+), 7 deletions(-)

diff --git a/src/compiler/nir/nir.h b/src/compiler/nir/nir.h
index c20f46b..0f35a9e 100644
--- a/src/compiler/nir/nir.h
+++ b/src/compiler/nir/nir.h
@@ -2360,6 +2360,9 @@ void nir_print_instr(const nir_instr *instr, FILE *fp);
 
 nir_shader *nir_shader_clone(void *mem_ctx, const nir_shader *s);
 nir_function_impl *nir_function_impl_clone(const nir_function_impl *fi);
+void nir_clone_loop_list(struct exec_list *dst, const struct exec_list *list,
+                         struct hash_table *remap_table,
+                         struct hash_table *phi_remap, nir_shader *ns);
 nir_constant *nir_constant_clone(const nir_constant *c, nir_variable *var);
 nir_variable *nir_variable_clone(const nir_variable *c, nir_shader *shader);
 
diff --git a/src/compiler/nir/nir_clone.c b/src/compiler/nir/nir_clone.c
index 8808333..071afc9 100644
--- a/src/compiler/nir/nir_clone.c
+++ b/src/compiler/nir/nir_clone.c
@@ -35,9 +35,17 @@ typedef struct {
    /* True if we are cloning an entire shader. */
    bool global_clone;
 
+   /* This allows us to clone a loop body without having to add srcs from
+    * outside the loop to the remap table. This is useful for loop unrolling.
+    */
+   bool allow_remap_fallback;
+
    /* maps orig ptr -> cloned ptr: */
    struct hash_table *remap_table;
 
+   /* used for remaping when cloning loop body for loop unrolling */
+   struct hash_table *phi_remap_table;
+
    /* List of phi sources. */
    struct list_head phi_srcs;
 
@@ -46,11 +54,20 @@ typedef struct {
 } clone_state;
 
 static void
-init_clone_state(clone_state *state, bool global)
+init_clone_state(clone_state *state, struct hash_table *remap_table,
+                 bool global, bool allow_remap_fallback)
 {
    state->global_clone = global;
-   state->remap_table = _mesa_hash_table_create(NULL, _mesa_hash_pointer,
-                                                _mesa_key_pointer_equal);
+   state->allow_remap_fallback = allow_remap_fallback;
+
+   state->phi_remap_table = NULL;
+   if (remap_table) {
+      state->remap_table = remap_table;
+   } else {
+      state->remap_table = _mesa_hash_table_create(NULL, _mesa_hash_pointer,
+                                                   _mesa_key_pointer_equal);
+   }
+
    list_inithead(&state->phi_srcs);
 }
 
@@ -72,16 +89,32 @@ _lookup_ptr(clone_state *state, const void *ptr, bool global)
       return (void *)ptr;
 
    entry = _mesa_hash_table_search(state->remap_table, ptr);
-   assert(entry && "Failed to find pointer!");
    if (!entry)
-      return NULL;
+      return state->allow_remap_fallback ? (void *)ptr : NULL;
 
    return entry->data;
 }
 
+/**
+ * Updates a phi remap table used for unrolling loops.
+ */
+static void
+update_phi_remap_table(clone_state *state, const void *ptr, void *nptr)
+{
+   if (state->phi_remap_table == NULL)
+      return;
+
+   struct hash_entry *hte;
+   hash_table_foreach(state->phi_remap_table, hte) {
+      if (hte->data == ptr)
+         hte->data = nptr;
+   }
+}
+
 static void
 add_remap(clone_state *state, void *nptr, const void *ptr)
 {
+   update_phi_remap_table(state, ptr, nptr);
    _mesa_hash_table_insert(state->remap_table, ptr, nptr);
 }
 
@@ -613,6 +646,23 @@ fixup_phi_srcs(clone_state *state)
    assert(list_empty(&state->phi_srcs));
 }
 
+void
+nir_clone_loop_list(struct exec_list *dst, const struct exec_list *list,
+                    struct hash_table *remap_table,
+                    struct hash_table *phi_remap, nir_shader *ns)
+{
+   clone_state state;
+   init_clone_state(&state, remap_table, false, true);
+
+   /* We use the same shader */
+   state.ns = ns;
+   state.phi_remap_table = phi_remap;
+
+   clone_cf_list(&state, dst, list);
+
+   fixup_phi_srcs(&state);
+}
+
 static nir_function_impl *
 clone_function_impl(clone_state *state, const nir_function_impl *fi)
 {
@@ -646,7 +696,7 @@ nir_function_impl *
 nir_function_impl_clone(const nir_function_impl *fi)
 {
    clone_state state;
-   init_clone_state(&state, false);
+   init_clone_state(&state, NULL, false, false);
 
    /* We use the same shader */
    state.ns = fi->function->shader;
@@ -686,7 +736,7 @@ nir_shader *
 nir_shader_clone(void *mem_ctx, const nir_shader *s)
 {
    clone_state state;
-   init_clone_state(&state, true);
+   init_clone_state(&state, NULL, true, false);
 
    nir_shader *ns = nir_shader_create(mem_ctx, s->stage, s->options);
    state.ns = ns;
-- 
2.7.4