nir/lower_shader_calls: put inserted instructions into a dummy block

When moving code into the main block or loop blocks, put the code into
its own :

    if(true) { ... }

block so that we avoid break/continue/return issues.

v2: Also take care of the main block with return instructions

v3: Make deletion more obvious with dummy if blocks (Jason)

v4: Fixup assert for loops (Lionel)

Signed-off-by: Lionel Landwerlin <lionel.g.landwerlin@intel.com>
Fixes: 8dfb240b1f ("nir: Add raytracing shader call lowering pass.")
Reviewed-by: Jason Ekstrand <jason.ekstrand@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/16036>

src/compiler/nir/nir_lower_shader_calls.c

@@ -866,8 +866,24 @@ flatten_resume_if_ladder(nir_builder *b,
       }
 
       case nir_cf_node_if: {
-         assert(!before_cursor);
          nir_if *_if = nir_cf_node_as_if(child);
+
+         /* Because of the dummy blocks inserted in the first if block of the
+          * loops, it's possible we find an empty if block that contains our
+          * cursor. At this point, the block should still be empty and we can
+          * just skip it and consider we're after the cursor.
+          */
+         if (cf_node_contains_block(&_if->cf_node,
+                                    nir_cursor_current_block(b->cursor))) {
+            /* Some sanity checks to verify this is actually a dummy block */
+            assert(nir_src_as_bool(_if->condition) == true);
+            assert(nir_cf_list_is_empty_block(&_if->then_list));
+            assert(nir_cf_list_is_empty_block(&_if->else_list));
+            before_cursor = false;
+            break;
+         }
+         assert(!before_cursor);
+
          if (flatten_resume_if_ladder(b, &_if->cf_node, &_if->then_list,
                                       false, resume_instr, remat)) {
             resume_node = child;
@@ -900,17 +916,26 @@ flatten_resume_if_ladder(nir_builder *b,
             nir_block *header = nir_loop_first_block(loop);
             nir_if *_if = nir_cf_node_as_if(nir_cf_node_next(&header->cf_node));
 
-            /* We want to place anything re-materialized from inside the loop
-             * at the top of the resume half of the loop.
-             */
             nir_builder bl;
             nir_builder_init(&bl, b->impl);
             bl.cursor = nir_before_cf_list(&_if->then_list);
-
+            /* We want to place anything re-materialized from inside the loop
+             * at the top of the resume half of the loop.
+             *
+             * Because we're inside a loop, we might run into a break/continue
+             * instructions. We can't place those within a block of
+             * instructions, they need to be at the end of a block. So we
+             * build our own dummy block to place them.
+             */
+            nir_push_if(&bl, nir_imm_true(&bl));
+            {
                ASSERTED bool found =
                   flatten_resume_if_ladder(&bl, &_if->cf_node, &_if->then_list,
                                            true, resume_instr, remat);
                assert(found);
+            }
+            nir_pop_if(&bl, NULL);
+
             resume_node = child;
             goto found_resume;
          } else {
@@ -986,7 +1011,23 @@ found_resume:
     * cursor.  Delete everything else.
     */
    if (child_list_contains_cursor) {
-      nir_cf_extract(&cf_list, b->cursor, nir_after_cf_list(child_list));
+      /* If the cursor is in child_list, then we're either a loop or function
+       * that contains the cursor. Cursors are always placed in a wrapper if
+       * (true) to deal with break/continue and early returns. We've already
+       * moved everything interesting inside the wrapper if and we want to
+       * remove whatever is left after it.
+       */
+      nir_block *cursor_block = nir_cursor_current_block(b->cursor);
+      nir_if *wrapper_if = nir_cf_node_as_if(cursor_block->cf_node.parent);
+      assert(wrapper_if->cf_node.parent == parent_node);
+      /* The wrapper if blocks are either put into the body of the main
+       * function, or within the resume if block of the loops.
+       */
+      assert(parent_node->type == nir_cf_node_function ||
+             (parent_node->type == nir_cf_node_if &&
+              parent_node->parent->type == nir_cf_node_loop));
+      nir_cf_extract(&cf_list, nir_after_cf_node(&wrapper_if->cf_node),
+                     nir_after_cf_list(child_list));
    } else {
       nir_cf_list_extract(&cf_list, child_list);
    }
@@ -1032,10 +1073,15 @@ lower_resume(nir_shader *shader, int call_idx)
    nir_builder b;
    nir_builder_init(&b, impl);
    b.cursor = nir_before_cf_list(&impl->body);
+
+   nir_push_if(&b, nir_imm_true(&b));
+   {
       ASSERTED bool found =
          flatten_resume_if_ladder(&b, &impl->cf_node, &impl->body,
                                   true, resume_instr, &remat);
       assert(found);
+   }
+   nir_pop_if(&b, NULL);
 
    ralloc_free(mem_ctx);
 
@@ -1167,6 +1213,8 @@ nir_lower_shader_calls(nir_shader *shader,
       nir_instr *resume_instr = lower_resume(resume_shaders[i], i);
       replace_resume_with_halt(resume_shaders[i], resume_instr);
       nir_opt_remove_phis(resume_shaders[i]);
+      /* Remove the dummy blocks added by flatten_resume_if_ladder() */
+      nir_opt_if(resume_shaders[i], false);
    }
 
    *resume_shaders_out = resume_shaders;