From ec9eae99b133f8bc34714eda182d0d3df19b5ada Mon Sep 17 00:00:00 2001
From: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Date: Thu, 20 Oct 2022 22:15:54 -0400
Subject: [PATCH] agx: Report GPRs to the driver

This needs to be passed onto the hardware.

Signed-off-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/18813>
---
 src/asahi/compiler/agx_compile.c           | 24 ++++++++++------------
 src/asahi/compiler/agx_compile.h           |  5 +++++
 src/asahi/compiler/agx_compiler.h          |  1 +
 src/asahi/compiler/agx_register_allocate.c |  5 +++++
 4 files changed, 22 insertions(+), 13 deletions(-)

diff --git a/src/asahi/compiler/agx_compile.c b/src/asahi/compiler/agx_compile.c
index 2589099fd64..bb920bc293a 100644
--- a/src/asahi/compiler/agx_compile.c
+++ b/src/asahi/compiler/agx_compile.c
@@ -1535,21 +1535,12 @@ agx_set_st_vary_final(agx_context *ctx)
 static int
 agx_dump_stats(agx_context *ctx, unsigned size, char **out)
 {
-   unsigned nr_ins = 0, max_reg = 0;
+   unsigned nr_ins = 0;
 
-   agx_foreach_instr_global(ctx, I) {
-      /* Count instructions */
+   /* Count instructions */
+   agx_foreach_instr_global(ctx, I)
       nr_ins++;
 
-      /* Count registers */
-      agx_foreach_dest(I, d) {
-         if (I->dest[d].type == AGX_INDEX_REGISTER) {
-            max_reg = MAX2(max_reg,
-                           I->dest[d].value + agx_write_registers(I, d) - 1);
-         }
-      }
-   }
-
    /* TODO: Pipe through occupancy */
    unsigned nr_threads = 1;
 
@@ -1557,7 +1548,7 @@ agx_dump_stats(agx_context *ctx, unsigned size, char **out)
            "%s shader: %u inst, %u bytes, %u halfregs, %u threads, "
            "%u loops, %u:%u spills:fills",
            gl_shader_stage_name(ctx->stage),
-           nr_ins, size, max_reg, nr_threads, ctx->loop_count,
+           nr_ins, size, ctx->max_reg, nr_threads, ctx->loop_count,
            ctx->spills, ctx->fills);
 }
 
@@ -1845,6 +1836,13 @@ agx_compile_function_nir(nir_shader *nir, nir_function_impl *impl,
 
    agx_pack_binary(ctx, binary);
 
+   unsigned nr_gprs = ctx->max_reg + 1;
+
+   if (impl->function->is_preamble)
+      out->nr_preamble_gprs = nr_gprs;
+   else
+      out->nr_gprs = nr_gprs;
+
    /* Don't dump statistics for preambles, since they're not worth optimizing */
    if (!impl->function->is_preamble) {
       char *stats;
diff --git a/src/asahi/compiler/agx_compile.h b/src/asahi/compiler/agx_compile.h
index aa89699dc3a..df770bd90e6 100644
--- a/src/asahi/compiler/agx_compile.h
+++ b/src/asahi/compiler/agx_compile.h
@@ -175,6 +175,11 @@ struct agx_shader_info {
 
    /* Is colour output omitted? */
    bool no_colour_output;
+
+   /* Number of 16-bit registers used by the main shader and preamble
+    * respectively.
+    */
+   unsigned nr_gprs, nr_preamble_gprs;
 };
 
 #define AGX_MAX_RTS (8)
diff --git a/src/asahi/compiler/agx_compiler.h b/src/asahi/compiler/agx_compiler.h
index 882d4d5fea2..e26c4977049 100644
--- a/src/asahi/compiler/agx_compiler.h
+++ b/src/asahi/compiler/agx_compiler.h
@@ -411,6 +411,7 @@ typedef struct {
    unsigned loop_count;
    unsigned spills;
    unsigned fills;
+   unsigned max_reg;
 } agx_context;
 
 static inline void
diff --git a/src/asahi/compiler/agx_register_allocate.c b/src/asahi/compiler/agx_register_allocate.c
index b18b05a8122..1d4da15f26e 100644
--- a/src/asahi/compiler/agx_register_allocate.c
+++ b/src/asahi/compiler/agx_register_allocate.c
@@ -350,6 +350,11 @@ agx_ra(agx_context *ctx)
       });
    }
 
+   for (unsigned i = 0; i < ctx->alloc; ++i) {
+      if (ncomps[i])
+         ctx->max_reg = MAX2(ctx->max_reg, ssa_to_reg[i] + ncomps[i] - 1);
+   }
+
    agx_foreach_instr_global(ctx, ins) {
       agx_foreach_src(ins, s) {
          if (ins->src[s].type == AGX_INDEX_NORMAL) {