nir: add nir_loop_has_back_edge helper

Signed-off-by: Rhys Perry <pendingchaos02@gmail.com>
Acked-by: Alyssa Rosenzweig <alyssa.rosenzweig@intel.com>
Reviewed-by: Georg Lehmann <dadschoorse@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/40242>

src/compiler/nir/nir.c

@@ -3906,6 +3906,18 @@ nir_block_contains_work(nir_block *block)
    return false;
 }
 
+bool
+nir_loop_has_back_edge(nir_loop *loop)
+{
+   nir_block *header = nir_loop_first_block(loop);
+   nir_block *preheader = nir_cf_node_as_block(nir_cf_node_prev(&loop->cf_node));
+   /* We also check whether the preheader is a predecessor of the header in
+    * case there is a jump right before the loop. Usually nir_opt_dead_cf will
+    * remove this, but we might call this function before that pass. */
+   bool has_preheader = nir_block_has_pred(header, preheader);
+   return nir_block_num_preds(header) - has_preheader > 0;
+}
+
 nir_op
 nir_atomic_op_to_alu(nir_atomic_op op)
 {

src/compiler/nir/nir.h

@@ -3916,6 +3916,8 @@ nir_loop_last_continue_block(nir_loop *loop)
    return nir_cf_node_as_block(exec_node_data(nir_cf_node, tail, node));
 }
 
+bool nir_loop_has_back_edge(nir_loop *loop);
+
 /**
  * Return true if this list of cf_nodes contains a single empty block.
  */

src/compiler/nir/nir_opt_dce.c

@@ -181,8 +181,7 @@ dce_cf_list(struct exec_list *cf_list, BITSET_WORD *defs_live,
          /* Fast path if the loop has no continues: we can remove instructions
           * as we mark the others live.
           */
-         nir_block *header = nir_loop_first_block(loop);
-         if (nir_block_num_preds(header) == 1 && nir_block_has_pred(header, inner_state.preheader)) {
+         if (!nir_loop_has_back_edge(loop)) {
             progress |= dce_cf_list(&loop->body, defs_live, parent_loop, dead_instrs);
             break;
          }

src/compiler/nir/nir_opt_licm.c

@@ -58,7 +58,7 @@ static bool
 should_optimize_loop(nir_loop *loop)
 {
    /* Ignore loops without back-edge */
-   if (nir_block_num_preds(nir_loop_first_block(loop)) == 1)
+   if (!nir_loop_has_back_edge(loop))
       return false;
 
    nir_foreach_block_in_cf_node(block, &loop->cf_node) {

src/compiler/nir/nir_opt_load_store_vectorize.c

@@ -1848,25 +1848,24 @@ static void
 add_entries_from_predecessor(struct vectorize_ctx *ctx, nir_block *block)
 {
    nir_cf_node *parent = block->cf_node.parent;
-   bool is_loop_header = false;
+   nir_loop *loop = NULL;
    if (parent->type == nir_cf_node_loop) {
-      nir_loop *loop = nir_cf_node_as_loop(parent);
-      if (block == nir_loop_first_block(loop))
-         is_loop_header = true;
+      loop = nir_cf_node_as_loop(parent);
+      if (block != nir_loop_first_block(loop))
+         loop = NULL;
    }
 
    for (unsigned i = 0; i < nir_num_variable_modes; i++) {
       struct entry *entry = NULL;
 
-      if (is_loop_header) {
+      if (loop) {
          /* If this is a loop header, just take the last entries of the preheader. */
          nir_block *preheader = nir_block_cf_tree_prev(block);
          entry = ctx->per_block_ctx[preheader->index].last_entry[i];
 
          /* If this isn't reorderable, we would have to consider the loop back-edges to safely use
           * it, in case there is an interfering store in the loop. */
-         bool has_continue = nir_block_num_preds(block) > 1 || !nir_block_has_pred(block, preheader);
-         if (entry && !(entry->access & ACCESS_CAN_REORDER) && has_continue)
+         if (entry && !(entry->access & ACCESS_CAN_REORDER) && nir_loop_has_back_edge(loop))
             entry = NULL;
       } else {
          /* If all predecessor entries are the same, the entry dominates the block. */

src/compiler/nir/nir_opt_sink.c

@@ -301,7 +301,7 @@ get_innermost_loop(nir_cf_node *node)
    for (; node != NULL; node = node->parent) {
       if (node->type == nir_cf_node_loop) {
          nir_loop *loop = nir_cf_node_as_loop(node);
-         if (nir_block_num_preds(nir_loop_first_block(loop)) > 1)
+         if (nir_loop_has_back_edge(loop))
             return loop;
       }
    }
@@ -341,10 +341,10 @@ adjust_block_for_loops(nir_block *use_block, nir_block *def_block,
       }
 
       nir_cf_node *next = nir_cf_node_next(&cur_block->cf_node);
-      if (next && next->type == nir_cf_node_loop &&
-          nir_block_num_preds(nir_block_cf_tree_next(cur_block)) > 1) {
+      if (next && next->type == nir_cf_node_loop) {
          nir_loop *following_loop = nir_cf_node_as_loop(next);
-         if (loop_contains_block(following_loop, use_block)) {
+         if (nir_loop_has_back_edge(following_loop) &&
+             loop_contains_block(following_loop, use_block)) {
             use_block = cur_block;
             continue;
          }

src/compiler/nir/nir_to_lcssa.c

@@ -342,7 +342,7 @@ convert_to_lcssa(nir_cf_node *cf_node, lcssa_state *state)
           * The variance then depends on all (nested) break conditions.
           * We don't consider this, but assume all not_invariant.
           */
-         if (nir_block_num_preds(nir_loop_first_block(loop)) == 1)
+         if (!nir_loop_has_back_edge(loop))
             goto end;
 
          nir_foreach_block_in_cf_node(block, cf_node) {