agx: Mark the logical ends of blocks

We need to insert parallel copies at the logical end of blocks, before branches.
Add a pseudo instruction signaling that. Cribbed from ACO.

Signed-off-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/16268>

src/asahi/compiler/agx_compile.c

@@ -977,6 +977,20 @@ agx_emit_tex(agx_builder *b, nir_tex_instr *instr)
    agx_emit_cached_split(b, dst, 4);
 }
 
+/*
+ * Mark the logical end of the current block by emitting a p_logical_end marker.
+ * Note if an unconditional jump is emitted (for instance, to break out of a
+ * loop from inside an if), the block has already reached its logical end so we
+ * don't re-emit p_logical_end. The validator checks this, and correct register
+ * allocation depends on it.
+ */
+static void
+agx_emit_logical_end(agx_builder *b)
+{
+   if (!b->shader->current_block->unconditional_jumps)
+      agx_p_logical_end(b);
+}
+
 /* NIR loops are treated as a pair of AGX loops:
  *
  *    do {
@@ -1014,6 +1028,9 @@ agx_emit_jump(agx_builder *b, nir_jump_instr *instr)
    /* Update the counter and flush */
    agx_index r0l = agx_register(0, false);
    agx_mov_to(b, r0l, agx_immediate(nestings));
+
+   /* Jumps must come at the end of a block */
+   agx_emit_logical_end(b);
    agx_pop_exec(b, 0);
 
    ctx->current_block->unconditional_jumps = true;
@@ -1104,6 +1121,7 @@ emit_if(agx_context *ctx, nir_if *nif)
    agx_builder _b = agx_init_builder(ctx, agx_after_block(first_block));
    agx_index cond = agx_src_index(&nif->condition);
 
+   agx_emit_logical_end(&_b);
    agx_if_icmp(&_b, cond, agx_zero(), 1, AGX_ICOND_UEQ, true);
    ctx->loop_nesting++;
 
@@ -1112,6 +1130,7 @@ emit_if(agx_context *ctx, nir_if *nif)
    agx_block *end_then = ctx->current_block;
 
    _b.cursor = agx_after_block(ctx->current_block);
+   agx_emit_logical_end(&_b);
    agx_else_icmp(&_b, cond, agx_zero(), 1, AGX_ICOND_UEQ, false);
 
    agx_block *else_block = emit_cf_list(ctx, &nif->else_list);
@@ -1125,6 +1144,7 @@ emit_if(agx_context *ctx, nir_if *nif)
    agx_block_add_successor(end_else, ctx->after_block);
 
    _b.cursor = agx_after_block(ctx->current_block);
+   agx_emit_logical_end(&_b);
    agx_pop_exec(&_b, 1);
    ctx->loop_nesting--;
 }
@@ -1143,6 +1163,7 @@ emit_loop(agx_context *ctx, nir_loop *nloop)
 
    /* Make room for break/continue nesting (TODO: skip if no divergent CF) */
    agx_builder _b = agx_init_builder(ctx, agx_after_block(ctx->current_block));
+   agx_emit_logical_end(&_b);
    agx_push_exec(&_b, 2);
 
    /* Fallthrough to body */
@@ -1155,6 +1176,7 @@ emit_loop(agx_context *ctx, nir_loop *nloop)
    /* Fix up the nesting counter via an always true while_icmp, and branch back
     * to start of loop if any lanes are active */
    _b.cursor = agx_after_block(ctx->current_block);
+   agx_emit_logical_end(&_b);
    agx_while_icmp(&_b, agx_zero(), agx_zero(), 2, AGX_ICOND_UEQ, false);
    agx_jmp_exec_any(&_b, start_block);
    agx_pop_exec(&_b, 2);

src/asahi/compiler/agx_opcodes.py

@@ -252,6 +252,9 @@ op("xor", _, srcs = 2)
 op("and", _, srcs = 2)
 op("or", _, srcs = 2)
 
+# Indicates the logical end of the block, before final branches/control flow
+op("p_logical_end", _, dests = 0, srcs = 0, can_eliminate = False)
+
 op("p_combine", _, srcs = 4)
 op("p_split", _, srcs = 1, dests = 4)
 op("p_extract", _, srcs = 1, imms = [COMPONENT])

src/asahi/compiler/agx_register_allocate.c

@@ -302,6 +302,11 @@ agx_ra(agx_context *ctx)
       }
    }
 
+   agx_foreach_instr_global_safe(ctx, I) {
+      if (I->op == AGX_OPCODE_P_LOGICAL_END)
+         agx_remove_instruction(I);
+   }
+
    free(ssa_to_reg);
    free(ncomps);
    free(alloc);