agx: Pipe in nir_register

This is kind of lazy...

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

src/asahi/compiler/agx_compile.c

@@ -1163,6 +1163,7 @@ agx_optimize_nir(nir_shader *nir)
 
    NIR_PASS_V(nir, nir_opt_sink, move_all);
    NIR_PASS_V(nir, nir_opt_move, move_all);
+   NIR_PASS_V(nir, nir_convert_from_ssa, true);
 }
 
 /* ABI: position first, then user, then psiz */
@@ -1346,11 +1347,47 @@ agx_compile_shader_nir(nir_shader *nir,
       if (!func->impl)
          continue;
 
+      /* TODO: Handle phi nodes instead of just convert_from_ssa and yolo'ing
+       * the mapping of nir_register to hardware registers and guaranteeing bad
+       * performance and breaking spilling... */
+      ctx->nir_regalloc = rzalloc_array(ctx, unsigned, func->impl->reg_alloc);
+
+      /* Leave the last 4 registers for hacky p-copy lowering */
+      unsigned nir_regalloc = AGX_NUM_REGS - (4 * 2);
+
+      /* Assign backwards so we don't need to guess a size */
+      nir_foreach_register(reg, &func->impl->registers) {
+         /* Ensure alignment */
+         if (reg->bit_size >= 32 && (nir_regalloc & 1))
+            nir_regalloc--;
+
+         unsigned size = DIV_ROUND_UP(reg->bit_size * reg->num_components, 16);
+         nir_regalloc -= size;
+         ctx->nir_regalloc[reg->index] = nir_regalloc;
+      }
+
       ctx->alloc += func->impl->ssa_alloc;
       emit_cf_list(ctx, &func->impl->body);
       break; /* TODO: Multi-function shaders */
    }
 
+   /* TODO: Actual RA... this way passes don't need to deal nir_register */
+   agx_foreach_instr_global(ctx, I) {
+      agx_foreach_dest(I, d) {
+         if (I->dest[d].type == AGX_INDEX_NIR_REGISTER) {
+            I->dest[d].type = AGX_INDEX_REGISTER;
+            I->dest[d].value = ctx->nir_regalloc[I->dest[d].value];
+         }
+      }
+
+      agx_foreach_src(I, s) {
+         if (I->src[s].type == AGX_INDEX_NIR_REGISTER) {
+            I->src[s].type = AGX_INDEX_REGISTER;
+            I->src[s].value = ctx->nir_regalloc[I->src[s].value];
+         }
+      }
+   }
+
    /* Terminate the shader after the exit block */
    agx_block *last_block = list_last_entry(&ctx->blocks, agx_block, link);
    agx_builder _b = agx_init_builder(ctx, agx_after_block(last_block));

src/asahi/compiler/agx_compiler.h

@@ -52,6 +52,7 @@ enum agx_index_type {
    AGX_INDEX_IMMEDIATE = 2,
    AGX_INDEX_UNIFORM = 3,
    AGX_INDEX_REGISTER = 4,
+   AGX_INDEX_NIR_REGISTER = 5,
 };
 
 enum agx_size {
@@ -119,6 +120,16 @@ agx_register(uint8_t imm, enum agx_size size)
    };
 }
 
+static inline agx_index
+agx_nir_register(unsigned imm, enum agx_size size)
+{
+   return (agx_index) {
+      .type = AGX_INDEX_NIR_REGISTER,
+      .value = imm,
+      .size = size
+   };
+}
+
 /* Also in half-words */
 static inline agx_index
 agx_uniform(uint8_t imm, enum agx_size size)
@@ -348,6 +359,10 @@ typedef struct {
    /* Remapping table for varyings indexed by driver_location */
    unsigned varyings[AGX_MAX_VARYINGS];
 
+   /* Handling phi nodes is still TODO while we bring up other parts of the
+    * driver. YOLO the mapping of nir_register to fixed hardware registers */
+   unsigned *nir_regalloc;
+
    /* Place to start pushing new values */
    unsigned push_base;
 
@@ -407,7 +422,11 @@ agx_size_for_bits(unsigned bits)
 static inline agx_index
 agx_src_index(nir_src *src)
 {
-   assert(src->is_ssa);
+   if (!src->is_ssa) {
+      return agx_nir_register(src->reg.reg->index,
+            agx_size_for_bits(nir_src_bit_size(*src)));
+   }
+
    return agx_get_index(src->ssa->index,
          agx_size_for_bits(nir_src_bit_size(*src)));
 }
@@ -415,7 +434,11 @@ agx_src_index(nir_src *src)
 static inline agx_index
 agx_dest_index(nir_dest *dst)
 {
-   assert(dst->is_ssa);
+   if (!dst->is_ssa) {
+      return agx_nir_register(dst->reg.reg->index,
+            agx_size_for_bits(nir_dest_bit_size(*dst)));
+   }
+
    return agx_get_index(dst->ssa.index,
          agx_size_for_bits(nir_dest_bit_size(*dst)));
 }