agx: implement exports

translate export/load_exported instructions into moves to/from the requested
GPRs at shader part boundaries, with coalescing in RA for perf.

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

src/asahi/compiler/agx_compile.c

@@ -1065,6 +1065,40 @@ agx_emit_image_load(agx_builder *b, agx_index dst, nir_intrinsic_instr *intr)
    return NULL;
 }
 
+static agx_instr *
+agx_emit_export(agx_builder *b, unsigned base, nir_src src)
+{
+   agx_builder b_ = *b;
+   agx_cursor after_cursor = agx_after_block(agx_exit_block(b->shader));
+   b_.cursor = after_cursor;
+
+   for (unsigned c = 0; c < nir_src_num_components(src); ++c) {
+      agx_index chan = agx_extract_nir_src(b, src, c);
+      unsigned stride = agx_size_align_16(chan.size);
+
+      agx_export(&b_, chan, base + (c * stride));
+   }
+
+   if (memcmp(&b->cursor, &after_cursor, sizeof(agx_cursor)) == 0) {
+      b->cursor = agx_after_block_logical(b->cursor.block);
+   }
+
+   return NULL;
+}
+
+static agx_instr *
+agx_load_exported_to(agx_builder *b, agx_index dst, unsigned base, unsigned nr)
+{
+   agx_index chans[4] = {0};
+   unsigned stride = agx_size_align_16(dst.size);
+
+   for (unsigned c = 0; c < nr; ++c) {
+      chans[c] = agx_cached_preload(b->shader, base + c * stride, dst.size);
+   }
+
+   return agx_emit_collect_to(b, dst, nr, chans);
+}
+
 static agx_instr *
 agx_emit_image_store(agx_builder *b, nir_intrinsic_instr *instr)
 {
@@ -1457,6 +1491,13 @@ agx_emit_intrinsic(agx_builder *b, nir_intrinsic_instr *instr)
       assert(b->shader->key->is_helper);
       return agx_get_sr_barrier_to(b, dst, AGX_SR_HELPER_ARG_H);
 
+   case nir_intrinsic_load_exported_agx:
+      return agx_load_exported_to(b, dst, nir_intrinsic_base(instr),
+                                  instr->def.num_components);
+
+   case nir_intrinsic_export_agx:
+      return agx_emit_export(b, nir_intrinsic_base(instr), instr->src[0]);
+
    case nir_intrinsic_load_barycentric_sample:
    case nir_intrinsic_load_sample_id:
    case nir_intrinsic_load_sample_pos:

src/asahi/compiler/agx_compiler.h

@@ -844,6 +844,7 @@ instr_after_logical_end(const agx_instr *I)
    case AGX_OPCODE_IF_FCMP:
    case AGX_OPCODE_WHILE_FCMP:
    case AGX_OPCODE_STOP:
+   case AGX_OPCODE_EXPORT:
       return true;
    default:
       return false;

src/asahi/compiler/agx_lower_pseudo.c

@@ -91,6 +91,12 @@ lower(agx_builder *b, agx_instr *I)
          return cmpsel_for_break_if(b, I);
    }
 
+   case AGX_OPCODE_EXPORT:
+      /* We already lowered exports during RA, we just need to remove them late
+       * after inserting waits.
+       */
+      return (void *)true;
+
    default:
       return NULL;
    }

src/asahi/compiler/agx_opcodes.py

@@ -475,6 +475,13 @@ op("unit_test", _, dests = 0, srcs = 1, can_eliminate = False)
 # to be coalesced during RA, rather than lowered to a real move. 
 op("preload", _, srcs = 1, schedule_class = "preload")
 
+# Opposite of preload. Exports a scalar value to a particular register at the
+# end of the shader part. Must only appear after the logical end of the exit
+# block, this avoids special casing the source's liveness. Logically all exports
+# happen in parallel at the end of the shader part.
+op("export", _, dests = 0, srcs = 1, imms = [IMM], can_eliminate = False,
+   schedule_class = "invalid")
+
 # Pseudo-instructions to set the nesting counter. Lowers to r0l writes after RA.
 op("begin_cf", _, dests = 0, can_eliminate = False)
 op("break", _, dests = 0, imms = [NEST, TARGET], can_eliminate = False,

src/asahi/compiler/agx_optimizer.c

@@ -449,7 +449,7 @@ agx_optimizer_forward(agx_context *ctx)
       if (I->op != AGX_OPCODE_COLLECT && I->op != AGX_OPCODE_IMAGE_LOAD &&
           I->op != AGX_OPCODE_TEXTURE_LOAD &&
           I->op != AGX_OPCODE_UNIFORM_STORE &&
-          I->op != AGX_OPCODE_BLOCK_IMAGE_STORE)
+          I->op != AGX_OPCODE_BLOCK_IMAGE_STORE && I->op != AGX_OPCODE_EXPORT)
          agx_optimizer_inline_imm(defs, I);
 
       if (I->op == AGX_OPCODE_IF_ICMP)

src/asahi/compiler/agx_register_allocate.c

@@ -201,6 +201,14 @@ agx_calc_register_demand(agx_context *ctx)
          if (I->op == AGX_OPCODE_PHI)
             continue;
 
+         if (I->op == AGX_OPCODE_PRELOAD) {
+            unsigned size = agx_size_align_16(I->src[0].size);
+            max_demand = MAX2(max_demand, I->src[0].value + size);
+         } else if (I->op == AGX_OPCODE_EXPORT) {
+            unsigned size = agx_size_align_16(I->src[0].size);
+            max_demand = MAX2(max_demand, I->imm + size);
+         }
+
          /* Handle late-kill registers from last instruction */
          demand -= late_kill_count;
          late_kill_count = 0;
@@ -921,8 +929,32 @@ pick_regs(struct ra_ctx *rctx, agx_instr *I, unsigned d)
       }
    }
 
-   /* Try to allocate sources of collects contiguously */
+   /* Try to coalesce scalar exports */
    agx_instr *collect_phi = rctx->src_to_collect_phi[idx.value];
+   if (collect_phi && collect_phi->op == AGX_OPCODE_EXPORT) {
+      unsigned reg = collect_phi->imm;
+
+      if (!BITSET_TEST_RANGE(rctx->used_regs[cls], reg, reg + align - 1) &&
+          (reg % align) == 0)
+         return reg;
+   }
+
+   /* Try to coalesce vector exports */
+   if (collect_phi && collect_phi->op == AGX_OPCODE_SPLIT) {
+      if (collect_phi->dest[0].type == AGX_INDEX_NORMAL) {
+         agx_instr *exp = rctx->src_to_collect_phi[collect_phi->dest[0].value];
+         if (exp && exp->op == AGX_OPCODE_EXPORT) {
+            unsigned reg = exp->imm;
+
+            if (!BITSET_TEST_RANGE(rctx->used_regs[cls], reg,
+                                   reg + align - 1) &&
+                (reg % align) == 0)
+               return reg;
+         }
+      }
+   }
+
+   /* Try to allocate sources of collects contiguously */
    if (collect_phi && collect_phi->op == AGX_OPCODE_COLLECT) {
       agx_instr *collect = collect_phi;
 
@@ -1224,6 +1256,34 @@ agx_insert_parallel_copies(agx_context *ctx, agx_block *block)
    }
 }
 
+static void
+lower_exports(agx_context *ctx)
+{
+   struct agx_copy copies[AGX_NUM_REGS];
+   unsigned nr = 0;
+   agx_block *block = agx_exit_block(ctx);
+
+   agx_foreach_instr_in_block_safe(block, I) {
+      if (I->op != AGX_OPCODE_EXPORT)
+         continue;
+
+      assert(agx_channels(I->src[0]) == 1 && "scalarized in frontend");
+      assert(nr < ARRAY_SIZE(copies));
+
+      copies[nr++] = (struct agx_copy){
+         .dest = I->imm,
+         .src = I->src[0],
+      };
+
+      /* We cannot use fewer registers than we export */
+      ctx->max_reg =
+         MAX2(ctx->max_reg, I->imm + agx_size_align_16(I->src[0].size));
+   }
+
+   agx_builder b = agx_init_builder(ctx, agx_after_block_logical(block));
+   agx_emit_parallel_copies(&b, copies, nr);
+}
+
 void
 agx_ra(agx_context *ctx)
 {
@@ -1305,7 +1365,8 @@ agx_ra(agx_context *ctx)
 
    agx_foreach_instr_global(ctx, I) {
       /* Record collects/phis so we can coalesce when assigning */
-      if (I->op == AGX_OPCODE_COLLECT || I->op == AGX_OPCODE_PHI) {
+      if (I->op == AGX_OPCODE_COLLECT || I->op == AGX_OPCODE_PHI ||
+          I->op == AGX_OPCODE_EXPORT || I->op == AGX_OPCODE_SPLIT) {
          agx_foreach_ssa_src(I, s) {
             src_to_collect_phi[I->src[s].value] = I;
          }
@@ -1456,11 +1517,13 @@ agx_ra(agx_context *ctx)
       }
    }
 
-   /* Insert parallel copies lowering phi nodes */
+   /* Insert parallel copies lowering phi nodes and exports */
    agx_foreach_block(ctx, block) {
       agx_insert_parallel_copies(ctx, block);
    }
 
+   lower_exports(ctx);
+
    agx_foreach_instr_global_safe(ctx, I) {
       switch (I->op) {
       /* Pseudoinstructions for RA must be removed now */

src/asahi/compiler/agx_validate.c

@@ -4,6 +4,7 @@
  * SPDX-License-Identifier: MIT
  */
 
+#include "util/compiler.h"
 #include "agx_compiler.h"
 #include "agx_debug.h"
 #include "agx_opcodes.h"
@@ -42,6 +43,7 @@ agx_validate_block_form(agx_block *block)
 
    agx_foreach_instr_in_block(block, I) {
       switch (I->op) {
+      case AGX_OPCODE_PRELOAD:
       case AGX_OPCODE_ELSE_ICMP:
       case AGX_OPCODE_ELSE_FCMP:
          agx_validate_assert(state == AGX_BLOCK_STATE_CF_ELSE);
@@ -54,6 +56,11 @@ agx_validate_block_form(agx_block *block)
          state = AGX_BLOCK_STATE_PHI;
          break;
 
+      case AGX_OPCODE_EXPORT:
+         agx_validate_assert(agx_num_successors(block) == 0);
+         state = AGX_BLOCK_STATE_CF;
+         break;
+
       default:
          if (instr_after_logical_end(I)) {
             state = AGX_BLOCK_STATE_CF;
@@ -219,6 +226,7 @@ agx_read_registers(const agx_instr *I, unsigned s)
 
    switch (I->op) {
    case AGX_OPCODE_MOV:
+   case AGX_OPCODE_EXPORT:
       /* Tautological */
       return agx_index_size_16(I->src[0]);