nir: Add goto_if jump instruction

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

src/broadcom/compiler/nir_to_vir.c

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

src/compiler/nir/nir.c

@@ -477,7 +477,10 @@ nir_jump_instr_create(nir_shader *shader, nir_jump_type type)
 {
    nir_jump_instr *instr = ralloc(shader, nir_jump_instr);
    instr_init(&instr->instr, nir_instr_type_jump);
+   src_init(&instr->condition);
    instr->type = type;
+   instr->target = NULL;
+   instr->else_target = NULL;
    return instr;
 }
 
@@ -1195,6 +1198,15 @@ visit_parallel_copy_src(nir_parallel_copy_instr *instr,
    return true;
 }
 
+static bool
+visit_jump_src(nir_jump_instr *instr, nir_foreach_src_cb cb, void *state)
+{
+   if (instr->type != nir_jump_goto_if)
+      return true;
+
+   return visit_src(&instr->condition, cb, state);
+}
+
 typedef struct {
    void *state;
    nir_foreach_src_cb cb;
@@ -1248,6 +1260,7 @@ nir_foreach_src(nir_instr *instr, nir_foreach_src_cb cb, void *state)
          return false;
       break;
    case nir_instr_type_jump:
+      return visit_jump_src(nir_instr_as_jump(instr), cb, state);
    case nir_instr_type_ssa_undef:
       return true;
 

src/compiler/nir/nir.h

@@ -2293,11 +2293,22 @@ typedef enum {
     * NIR loop is implemented as "while (1) { body }".
     */
    nir_jump_continue,
+
+   /** Jumps for unstructured CFG.
+    *
+    * As within an unstructured CFG we can't rely on block ordering we need to
+    * place explicit jumps at the end of every block.
+    */
+   nir_jump_goto,
+   nir_jump_goto_if,
 } nir_jump_type;
 
 typedef struct {
    nir_instr instr;
    nir_jump_type type;
+   nir_src condition;
+   struct nir_block *target;
+   struct nir_block *else_target;
 } nir_jump_instr;
 
 /* creates a new SSA variable in an undefined state */

src/compiler/nir/nir_builder.h

@@ -1342,10 +1342,32 @@ nir_load_barycentric(nir_builder *build, nir_intrinsic_op op,
 static inline void
 nir_jump(nir_builder *build, nir_jump_type jump_type)
 {
+   assert(jump_type != nir_jump_goto && jump_type != nir_jump_goto_if);
    nir_jump_instr *jump = nir_jump_instr_create(build->shader, jump_type);
    nir_builder_instr_insert(build, &jump->instr);
 }
 
+static inline void
+nir_goto(nir_builder *build, struct nir_block *target)
+{
+   assert(!build->impl->structured);
+   nir_jump_instr *jump = nir_jump_instr_create(build->shader, nir_jump_goto);
+   jump->target = target;
+   nir_builder_instr_insert(build, &jump->instr);
+}
+
+static inline void
+nir_goto_if(nir_builder *build, struct nir_block *target, nir_src cond,
+            struct nir_block *else_target)
+{
+   assert(!build->impl->structured);
+   nir_jump_instr *jump = nir_jump_instr_create(build->shader, nir_jump_goto_if);
+   jump->condition = cond;
+   jump->target = target;
+   jump->else_target = else_target;
+   nir_builder_instr_insert(build, &jump->instr);
+}
+
 static inline nir_ssa_def *
 nir_compare_func(nir_builder *b, enum compare_func func,
                  nir_ssa_def *src0, nir_ssa_def *src1)

src/compiler/nir/nir_control_flow.c

@@ -140,7 +140,7 @@ link_block_to_non_block(nir_block *block, nir_cf_node *node)
 
       unlink_block_successors(block);
       link_blocks(block, first_then_block, first_else_block);
-   } else {
+   } else if (node->type == nir_cf_node_loop) {
       /*
        * For similar reasons as the corresponding case in
        * link_non_block_to_block(), don't worry about if the loop header has
@@ -312,7 +312,7 @@ block_add_normal_succs(nir_block *block)
          nir_block *first_else_block = nir_if_first_else_block(next_if);
 
          link_blocks(block, first_then_block, first_else_block);
-      } else {
+      } else if (next->type == nir_cf_node_loop) {
          nir_loop *next_loop = nir_cf_node_as_loop(next);
 
          nir_block *first_block = nir_loop_first_block(next_loop);
@@ -491,6 +491,14 @@ nir_handle_add_jump(nir_block *block)
       break;
    }
 
+   case nir_jump_goto:
+      link_blocks(block, jump_instr->target, NULL);
+      break;
+
+   case nir_jump_goto_if:
+      link_blocks(block, jump_instr->else_target, jump_instr->target);
+      break;
+
    default:
       unreachable("Invalid jump type");
    }
@@ -624,8 +632,10 @@ cleanup_cf_node(nir_cf_node *node, nir_function_impl *impl)
       /* We need to walk the instructions and clean up defs/uses */
       nir_foreach_instr_safe(instr, block) {
          if (instr->type == nir_instr_type_jump) {
-            nir_jump_type jump_type = nir_instr_as_jump(instr)->type;
-            unlink_jump(block, jump_type, false);
+            nir_jump_instr *jump = nir_instr_as_jump(instr);
+            unlink_jump(block, jump->type, false);
+            if (jump->type == nir_jump_goto_if)
+               nir_instr_rewrite_src(instr, &jump->condition, NIR_SRC_INIT);
          } else {
             nir_foreach_ssa_def(instr, replace_ssa_def_uses, impl);
             nir_instr_remove(instr);

src/compiler/nir/nir_divergence_analysis.c

@@ -593,6 +593,11 @@ visit_jump(nir_jump_instr *jump, struct divergence_state *state)
       return state->divergent_loop_break;
    case nir_jump_return:
       unreachable("NIR divergence analysis: Unsupported return instruction.");
+      break;
+   case nir_jump_goto:
+   case nir_jump_goto_if:
+      unreachable("NIR divergence analysis: Unsupported goto_if instruction.");
+      break;
    }
    return false;
 }

src/compiler/nir/nir_print.c

@@ -1215,6 +1215,19 @@ print_jump_instr(nir_jump_instr *instr, print_state *state)
    case nir_jump_return:
       fprintf(fp, "return");
       break;
+
+   case nir_jump_goto:
+      fprintf(fp, "goto block_%u",
+              instr->target ? instr->target->index : -1);
+      break;
+
+   case nir_jump_goto_if:
+      fprintf(fp, "goto block_%u if ",
+              instr->target ? instr->target->index : -1);
+      print_src(&instr->condition, state);
+      fprintf(fp, " else block_%u",
+              instr->else_target ? instr->else_target->index : -1);
+      break;
    }
 }
 

src/compiler/nir/nir_validate.c

@@ -777,6 +777,8 @@ validate_jump_instr(nir_jump_instr *instr, validate_state *state)
    case nir_jump_return:
       validate_assert(state, block->successors[0] == state->impl->end_block);
       validate_assert(state, block->successors[1] == NULL);
+      validate_assert(state, instr->target == NULL);
+      validate_assert(state, instr->else_target == NULL);
       break;
 
    case nir_jump_break:
@@ -788,6 +790,8 @@ validate_jump_instr(nir_jump_instr *instr, validate_state *state)
          validate_assert(state, block->successors[0] == after);
       }
       validate_assert(state, block->successors[1] == NULL);
+      validate_assert(state, instr->target == NULL);
+      validate_assert(state, instr->else_target == NULL);
       break;
 
    case nir_jump_continue:
@@ -798,6 +802,24 @@ validate_jump_instr(nir_jump_instr *instr, validate_state *state)
          validate_assert(state, block->successors[0] == first);
       }
       validate_assert(state, block->successors[1] == NULL);
+      validate_assert(state, instr->target == NULL);
+      validate_assert(state, instr->else_target == NULL);
+      break;
+
+   case nir_jump_goto:
+      validate_assert(state, !state->impl->structured);
+      validate_assert(state, instr->target == block->successors[0]);
+      validate_assert(state, instr->target != NULL);
+      validate_assert(state, instr->else_target == NULL);
+      break;
+
+   case nir_jump_goto_if:
+      validate_assert(state, !state->impl->structured);
+      validate_assert(state, instr->target == block->successors[1]);
+      validate_assert(state, instr->else_target == block->successors[0]);
+      validate_src(&instr->condition, state, 0, 1);
+      validate_assert(state, instr->target != NULL);
+      validate_assert(state, instr->else_target != NULL);
       break;
 
    default: