nir: Add and use nir_foreach_block_unstructured helpers

These are safe to call on either structured or unstructured NIR but
don't provide the nice ordering guarantees of nir_foreach_block and
friends.  While we're here, we use them for a very small selection of
passes which are known to be safe for unstructured control-flow.  The
most important such pass is nir_dominance which is required for
structurizing.

Signed-off-by: Karol Herbst <kherbst@redhat.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/2401>

src/compiler/nir/nir.c

@@ -1590,6 +1590,33 @@ nir_ssa_def_components_read(const nir_ssa_def *def)
    return read_mask;
 }
 
+nir_block *
+nir_block_unstructured_next(nir_block *block)
+{
+   if (block == NULL) {
+      /* nir_foreach_block_unstructured_safe() will call this function on a
+       * NULL block after the last iteration, but it won't use the result so
+       * just return NULL here.
+       */
+      return NULL;
+   }
+
+   nir_cf_node *cf_next = nir_cf_node_next(&block->cf_node);
+   if (cf_next == NULL && block->cf_node.parent->type == nir_cf_node_function)
+      return NULL;
+
+   if (cf_next && cf_next->type == nir_cf_node_block)
+      return nir_cf_node_as_block(cf_next);
+
+   return nir_block_cf_tree_next(block);
+}
+
+nir_block *
+nir_unstructured_start_block(nir_function_impl *impl)
+{
+   return nir_start_block(impl);
+}
+
 nir_block *
 nir_block_cf_tree_next(nir_block *block)
 {
@@ -1775,7 +1802,7 @@ nir_index_blocks(nir_function_impl *impl)
    if (impl->valid_metadata & nir_metadata_block_index)
       return;
 
-   nir_foreach_block(block, impl) {
+   nir_foreach_block_unstructured(block, impl) {
       block->index = index++;
    }
 
@@ -1803,7 +1830,7 @@ nir_index_ssa_defs(nir_function_impl *impl)
 {
    unsigned index = 0;
 
-   nir_foreach_block(block, impl) {
+   nir_foreach_block_unstructured(block, impl) {
       nir_foreach_instr(instr, block)
          nir_foreach_ssa_def(instr, index_ssa_def_cb, &index);
    }

src/compiler/nir/nir.h

@@ -3677,6 +3677,25 @@ void nir_ssa_def_rewrite_uses_after(nir_ssa_def *def, nir_src new_src,
 
 nir_component_mask_t nir_ssa_def_components_read(const nir_ssa_def *def);
 
+
+/** Returns the next block, disregarding structure
+ *
+ * The ordering is deterministic but has no guarantees beyond that.  In
+ * particular, it is not guaranteed to be dominance-preserving.
+ */
+nir_block *nir_block_unstructured_next(nir_block *block);
+nir_block *nir_unstructured_start_block(nir_function_impl *impl);
+
+#define nir_foreach_block_unstructured(block, impl) \
+   for (nir_block *block = nir_unstructured_start_block(impl); block != NULL; \
+        block = nir_block_unstructured_next(block))
+
+#define nir_foreach_block_unstructured_safe(block, impl) \
+   for (nir_block *block = nir_unstructured_start_block(impl), \
+        *next = nir_block_unstructured_next(block); \
+        block != NULL; \
+        block = next, next = nir_block_unstructured_next(block))
+
 /*
  * finds the next basic block in source-code order, returns NULL if there is
  * none

src/compiler/nir/nir_deref.c

@@ -735,7 +735,7 @@ nir_rematerialize_derefs_in_use_blocks_impl(nir_function_impl *impl)
    struct rematerialize_deref_state state = { 0 };
    nir_builder_init(&state.builder, impl);
 
-   nir_foreach_block(block, impl) {
+   nir_foreach_block_unstructured(block, impl) {
       state.block = block;
 
       /* Start each block with a fresh cache */

src/compiler/nir/nir_dominance.c

@@ -131,18 +131,18 @@ calc_dom_children(nir_function_impl* impl)
 {
    void *mem_ctx = ralloc_parent(impl);
 
-   nir_foreach_block(block, impl) {
+   nir_foreach_block_unstructured(block, impl) {
       if (block->imm_dom)
          block->imm_dom->num_dom_children++;
    }
 
-   nir_foreach_block(block, impl) {
+   nir_foreach_block_unstructured(block, impl) {
       block->dom_children = ralloc_array(mem_ctx, nir_block *,
                                          block->num_dom_children);
       block->num_dom_children = 0;
    }
 
-   nir_foreach_block(block, impl) {
+   nir_foreach_block_unstructured(block, impl) {
       if (block->imm_dom) {
          block->imm_dom->dom_children[block->imm_dom->num_dom_children++]
             = block;
@@ -170,20 +170,20 @@ nir_calc_dominance_impl(nir_function_impl *impl)
    nir_metadata_require(impl, nir_metadata_block_index);
 
 
-   nir_foreach_block(block, impl) {
+   nir_foreach_block_unstructured(block, impl) {
       init_block(block, impl);
    }
 
    bool progress = true;
    while (progress) {
       progress = false;
-      nir_foreach_block(block, impl) {
+      nir_foreach_block_unstructured(block, impl) {
          if (block != nir_start_block(impl))
             progress |= calc_dominance(block);
       }
    }
 
-   nir_foreach_block(block, impl) {
+   nir_foreach_block_unstructured(block, impl) {
       calc_dom_frontier(block);
    }
 
@@ -281,7 +281,7 @@ nir_dump_dom_tree_impl(nir_function_impl *impl, FILE *fp)
 {
    fprintf(fp, "digraph doms_%s {\n", impl->function->name);
 
-   nir_foreach_block(block, impl) {
+   nir_foreach_block_unstructured(block, impl) {
       if (block->imm_dom)
          fprintf(fp, "\t%u -> %u\n", block->imm_dom->index, block->index);
    }
@@ -301,7 +301,7 @@ nir_dump_dom_tree(nir_shader *shader, FILE *fp)
 void
 nir_dump_dom_frontier_impl(nir_function_impl *impl, FILE *fp)
 {
-   nir_foreach_block(block, impl) {
+   nir_foreach_block_unstructured(block, impl) {
       fprintf(fp, "DF(%u) = {", block->index);
       set_foreach(block->dom_frontier, entry) {
          nir_block *df = (nir_block *) entry->key;
@@ -325,7 +325,7 @@ nir_dump_cfg_impl(nir_function_impl *impl, FILE *fp)
 {
    fprintf(fp, "digraph cfg_%s {\n", impl->function->name);
 
-   nir_foreach_block(block, impl) {
+   nir_foreach_block_unstructured(block, impl) {
       if (block->successors[0])
          fprintf(fp, "\t%u -> %u\n", block->index, block->successors[0]->index);
       if (block->successors[1])