nir: Add a halt instruction type

Halt is like a return for the entire shader or exit() if you prefer to
think of it that way.  Once an invocation hits a halt, it's 100% dead.
Any writes to output variables which happened before the halt do,
however, still apply.

Reviewed-by: Caio Marcelo de Oliveira Filho <caio.oliveira@intel.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/7356>

src/broadcom/compiler/nir_to_vir.c

@@ -2955,6 +2955,7 @@ ntq_emit_jump(struct v3d_compile *c, nir_jump_instr *jump)
                 unreachable("All returns shouold be lowered\n");
                 break;
 
+        case nir_jump_halt:
         case nir_jump_goto:
         case nir_jump_goto_if:
                 unreachable("not supported\n");

src/compiler/nir/nir.h

@@ -2560,6 +2560,23 @@ typedef enum {
     */
    nir_jump_return,
 
+   /** Immediately exit the current shader
+    *
+    * This instruction is roughly the equivalent of C's "exit()" in that it
+    * immediately terminates the current shader invocation.  From a CFG
+    * perspective, it looks like a jump to nir_function_impl::end_block but
+    * it actually jumps to the end block of the shader entrypoint.  A halt
+    * instruction in the shader entrypoint itself is semantically identical
+    * to a return.
+    *
+    * For shaders with built-in I/O, any outputs written prior to a halt
+    * instruction remain written and any outputs not written prior to the
+    * halt have undefined values.  It does NOT cause an implicit discard of
+    * written results.  If one wants discard results in a fragment shader,
+    * for instance, a discard or demote intrinsic is required.
+    */
+   nir_jump_halt,
+
    /** Break out of the inner-most loop
     *
     * This has the same semantics as C's "break" statement.

src/compiler/nir/nir_control_flow.c

@@ -473,6 +473,7 @@ nir_handle_add_jump(nir_block *block)
 
    switch (jump_instr->type) {
    case nir_jump_return:
+   case nir_jump_halt:
       link_blocks(block, impl->end_block, NULL);
       break;
 
@@ -734,6 +735,53 @@ nir_cf_extract(nir_cf_list *extracted, nir_cursor begin, nir_cursor end)
    stitch_blocks(block_before, block_after);
 }
 
+static void
+relink_jump_halt_cf_node(nir_cf_node *node, nir_block *end_block)
+{
+   switch (node->type) {
+   case nir_cf_node_block: {
+      nir_block *block = nir_cf_node_as_block(node);
+      nir_instr *last_instr = nir_block_last_instr(block);
+      if (last_instr == NULL || last_instr->type != nir_instr_type_jump)
+         break;
+
+      nir_jump_instr *jump = nir_instr_as_jump(last_instr);
+      /* We can't move a CF list from one function to another while we still
+       * have returns.
+       */
+      assert(jump->type != nir_jump_return);
+
+      if (jump->type == nir_jump_halt) {
+         unlink_block_successors(block);
+         link_blocks(block, end_block, NULL);
+      }
+      break;
+   }
+
+   case nir_cf_node_if: {
+      nir_if *if_stmt = nir_cf_node_as_if(node);
+      foreach_list_typed(nir_cf_node, child, node, &if_stmt->then_list)
+         relink_jump_halt_cf_node(child, end_block);
+      foreach_list_typed(nir_cf_node, child, node, &if_stmt->else_list)
+         relink_jump_halt_cf_node(child, end_block);
+      break;
+   }
+
+   case nir_cf_node_loop: {
+      nir_loop *loop = nir_cf_node_as_loop(node);
+      foreach_list_typed(nir_cf_node, child, node, &loop->body)
+         relink_jump_halt_cf_node(child, end_block);
+      break;
+   }
+
+   case nir_cf_node_function:
+      unreachable("Cannot insert a function in a function");
+
+   default:
+      unreachable("Invalid CF node type");
+   }
+}
+
 void
 nir_cf_reinsert(nir_cf_list *cf_list, nir_cursor cursor)
 {
@@ -742,6 +790,13 @@ nir_cf_reinsert(nir_cf_list *cf_list, nir_cursor cursor)
    if (exec_list_is_empty(&cf_list->list))
       return;
 
+   nir_function_impl *cursor_impl =
+      nir_cf_node_get_function(&nir_cursor_current_block(cursor)->cf_node);
+   if (cf_list->impl != cursor_impl) {
+      foreach_list_typed(nir_cf_node, node, node, &cf_list->list)
+         relink_jump_halt_cf_node(node, cursor_impl->end_block);
+   }
+
    split_block_cursor(cursor, &before, &after);
 
    foreach_list_typed_safe(nir_cf_node, node, node, &cf_list->list) {

src/compiler/nir/nir_divergence_analysis.c

@@ -595,6 +595,9 @@ visit_jump(nir_jump_instr *jump, struct divergence_state *state)
       if (state->divergent_loop_cf)
          state->divergent_loop_break = true;
       return state->divergent_loop_break;
+   case nir_jump_halt:
+      /* This totally kills invocations so it doesn't add divergence */
+      break;
    case nir_jump_return:
       unreachable("NIR divergence analysis: Unsupported return instruction.");
       break;

src/compiler/nir/nir_opt_dead_cf.c

@@ -218,15 +218,17 @@ node_is_dead(nir_cf_node *node)
          if (instr->type == nir_instr_type_call)
             return false;
 
-         /* Return instructions can cause us to skip over other side-effecting
-          * instructions after the loop, so consider them to have side effects
-          * here.
+         /* Return and halt instructions can cause us to skip over other
+          * side-effecting instructions after the loop, so consider them to
+          * have side effects here.
           *
           * When the block is not inside a loop, break and continue might also
           * cause a skip.
           */
          if (instr->type == nir_instr_type_jump &&
-             (!inside_loop || nir_instr_as_jump(instr)->type == nir_jump_return))
+             (!inside_loop ||
+              nir_instr_as_jump(instr)->type == nir_jump_return ||
+              nir_instr_as_jump(instr)->type == nir_jump_halt))
             return false;
 
          if (instr->type == nir_instr_type_intrinsic) {

src/compiler/nir/nir_print.c

@@ -1330,6 +1330,10 @@ print_jump_instr(nir_jump_instr *instr, print_state *state)
       fprintf(fp, "return");
       break;
 
+   case nir_jump_halt:
+      fprintf(fp, "halt");
+      break;
+
    case nir_jump_goto:
       fprintf(fp, "goto block_%u",
               instr->target ? instr->target->index : -1);

src/compiler/nir/nir_validate.c

@@ -906,6 +906,7 @@ validate_jump_instr(nir_jump_instr *instr, validate_state *state)
 
    switch (instr->type) {
    case nir_jump_return:
+   case nir_jump_halt:
       validate_assert(state, block->successors[0] == state->impl->end_block);
       validate_assert(state, block->successors[1] == NULL);
       validate_assert(state, instr->target == NULL);