[Mesa-dev] [PATCH v2 05/14] nir: fix up phi nodes when removing cf nodes

[Connor Abbott]

When we deleted jumps like 'break' and 'continue,' we weren't removing
the phi sources that corresponded to them. Also, if we had a loop like

loop {
   ...
   break;
}

and we deleted the break at the end because it was unreachable, we
wouldn't add an undefined source to the phi nodes at the beginning of
the loop. Finally, we weren't updating the phi sources when stitching
two basic blocks together after deleting a control-flow node. This patch
fixes these issues by adding codepaths to the helpers for
adding/removing control flow nodes that deal with updating phi sources.

There's one somewhat unrelated fix to stitch_blocks() snuck in, where we
weren't updating the successors correctly if the earlier block had a
jump, but a lot of the logic for it overlaps with the phi node changes
and it seemed hard to split it out so I've left it in for now.

Signed-off-by: Connor Abbott <cwabbott0 at gmail.com>
---
 src/glsl/nir/nir.c | 99 +++++++++++++++++++++++++++++++++++++++++++++++++++++-
 src/glsl/nir/nir.h |  3 ++
 2 files changed, 101 insertions(+), 1 deletion(-)

diff --git a/src/glsl/nir/nir.c b/src/glsl/nir/nir.c
index dc6d63f..bd4f6cc 100644
--- a/src/glsl/nir/nir.c
+++ b/src/glsl/nir/nir.c
@@ -945,8 +945,62 @@ handle_jump(nir_block *block)
 }
 
 static void
+add_undef_phi_src(nir_block *pred, nir_block *block)
+{
+   nir_function_impl *impl = nir_cf_node_get_function(&block->cf_node);
+   void *mem_ctx = ralloc_parent(block);
+
+   nir_foreach_instr(block, instr) {
+      if (instr->type != nir_instr_type_phi)
+         break;
+
+      nir_phi_instr *phi = nir_instr_as_phi(instr);
+
+      nir_ssa_undef_instr *undef =
+         nir_ssa_undef_instr_create(mem_ctx, phi->dest.ssa.num_components);
+
+      nir_instr_insert_before_cf_list(&impl->body, &undef->instr);
+
+      nir_phi_src *src = ralloc(phi, nir_phi_src);
+      src->pred = pred;
+      src->src.parent_instr = &phi->instr;
+      src->src.is_ssa = true;
+      src->src.ssa = &undef->def;
+
+      list_addtail(&src->src.use_link, &src->src.ssa->uses);
+
+      exec_list_push_tail(&phi->srcs, &src->node);
+   }
+}
+
+static void
 handle_remove_jump(nir_block *block, nir_jump_type type)
 {
+   /*
+    * If the jump we're removing is a break or continue, then we're removing
+    * predecessors from a basic block that may have phi nodes at the beginning
+    * so we need to update them by removing the source corresponding to us.
+    */
+
+   if (type == nir_jump_break || type == nir_jump_continue) {
+      nir_block *succ = block->successors[0];
+
+      nir_foreach_instr(succ, instr) {
+         if (instr->type != nir_instr_type_phi)
+            break;
+
+         nir_phi_instr *phi = nir_instr_as_phi(instr);
+
+         nir_foreach_phi_src_safe(phi, src) {
+            if (src->pred == block) {
+               list_del(&src->src.use_link);
+               exec_node_remove(&src->node);
+               break;
+            }
+         }
+      }
+   }
+
    unlink_block_successors(block);
 
    if (exec_node_is_tail_sentinel(block->cf_node.node.next)) {
@@ -957,6 +1011,7 @@ handle_remove_jump(nir_block *block, nir_jump_type type)
          nir_block *next_block = nir_cf_node_as_block(next);
 
          link_blocks(block, next_block, NULL);
+         add_undef_phi_src(block, next_block);
       } else {
          assert(parent->type == nir_cf_node_loop);
          nir_loop *loop = nir_cf_node_as_loop(parent);
@@ -966,6 +1021,7 @@ handle_remove_jump(nir_block *block, nir_jump_type type)
          nir_block *head_block = nir_cf_node_as_block(head);
 
          link_blocks(block, head_block, NULL);
+         add_undef_phi_src(block, head_block);
       }
    } else {
       nir_cf_node *next = nir_cf_node_next(&block->cf_node);
@@ -990,6 +1046,7 @@ handle_remove_jump(nir_block *block, nir_jump_type type)
          nir_block *first_block = nir_cf_node_as_block(first);
 
          link_blocks(block, first_block, NULL);
+         add_undef_phi_src(block, first_block);
       }
    }
 
@@ -1008,6 +1065,7 @@ handle_remove_jump(nir_block *block, nir_jump_type type)
 
          last_block->successors[1] = next_block;
          block_add_pred(next_block, last_block);
+         add_undef_phi_src(last_block, next_block);
       }
    }
 
@@ -1195,7 +1253,46 @@ stitch_blocks(nir_block *before, nir_block *after)
     * TODO: special case when before is empty and after isn't?
     */
 
-   move_successors(after, before);
+   /*
+    * First, we have to consider the case where 'after' may have some phi
+    * nodes that refer to it, so we have to update those to refer to 'before'
+    * instead, unless 'before' ends in a jump in which case those phi sources
+    * are unreachable and we should delete them.
+    */
+
+   bool before_ends_in_jump = !exec_list_is_empty(&before->instr_list) &&
+                              nir_block_last_instr(before)->type == nir_instr_type_jump;
+
+   assert(!before_ends_in_jump || exec_list_is_empty(&after->instr_list));
+
+   for (unsigned i = 0; i < 2; i++) {
+      nir_block *succ = after->successors[i];
+      if (!succ)
+         break;
+
+      nir_foreach_instr(succ, instr) {
+         if (instr->type != nir_instr_type_phi)
+            break;
+
+         nir_phi_instr *phi = nir_instr_as_phi(instr);
+         nir_foreach_phi_src(phi, src) {
+            if (src->pred == after) {
+               if (before_ends_in_jump) {
+                  list_del(&src->src.use_link);
+                  exec_node_remove(&src->node);
+               } else {
+                  src->pred = before;
+               }
+               break;
+            }
+         }
+      }
+   }
+
+   if (before_ends_in_jump)
+      unlink_block_successors(after);
+   else
+      move_successors(after, before);
 
    foreach_list_typed(nir_instr, instr, node, &after->instr_list) {
       instr->block = before;
diff --git a/src/glsl/nir/nir.h b/src/glsl/nir/nir.h
index 697d37e..5ae261a 100644
--- a/src/glsl/nir/nir.h
+++ b/src/glsl/nir/nir.h
@@ -1107,6 +1107,9 @@ typedef struct {
 #define nir_foreach_phi_src(phi, entry) \
    foreach_list_typed(nir_phi_src, entry, node, &(phi)->srcs)
 
+#define nir_foreach_phi_src_safe(phi, entry) \
+   foreach_list_typed_safe(nir_phi_src, entry, node, &(phi)->srcs)
+
 typedef struct {
    nir_instr instr;
 
-- 
2.1.0