ir3: add support for precolored sources in predicate RA

To support predt/predf which always read from p0.x, we need to support
precolored sources for the predicates RA.

This patch implements this as follows: whenever a precolored source is
encountered whose def isn't live in the correct register, reload it into
the correct one. To make sure we don't reload too often, two precautions
are made. First, we precolor all defs of precolored sources and try do
use that register when allocating one for a def. Second, since currently
only p0.x is used for precoloring, we try not to allocate it whenever
there are outstanding precolored defs.

Signed-off-by: Job Noorman <jnoorman@igalia.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/27982>

src/freedreno/ir3/ir3_ra_predicates.c

@@ -48,6 +48,12 @@ struct ra_predicates_ctx {
    unsigned num_regs;
    struct ir3_liveness *liveness;
    struct block_liveness *blocks_liveness;
+
+   /* Number of precolored defs that have not been processed yet. When this
+    * drops to zero, we can stop trying to avoid allocating p0.x (the only
+    * register currently used for precoloring).
+    */
+   unsigned outstanding_precolored_defs;
 };
 
 static bool
@@ -61,9 +67,33 @@ has_free_regs(struct ra_predicates_ctx *ctx, struct block_liveness *live)
    return false;
 }
 
+static bool
+try_avoid_comp(struct ra_predicates_ctx *ctx, struct block_liveness *live,
+               unsigned comp)
+{
+   /* Currently, only p0.x is ever used for a precolored register so just try to
+    * avoid that one if we have any precolored defs.
+    */
+   return comp == 0 && ctx->outstanding_precolored_defs > 0;
+}
+
+static bool
+reg_is_free(struct ra_predicates_ctx *ctx, struct block_liveness *live,
+            unsigned comp)
+{
+   assert(comp < ctx->num_regs);
+
+   return live->live_defs[comp].def == NULL;
+}
+
 static unsigned
 alloc_reg_comp(struct ra_predicates_ctx *ctx, struct block_liveness *live)
 {
+   for (unsigned i = 0; i < ctx->num_regs; ++i) {
+      if (live->live_defs[i].def == NULL && !try_avoid_comp(ctx, live, i))
+         return i;
+   }
+
    for (unsigned i = 0; i < ctx->num_regs; ++i) {
       if (live->live_defs[i].def == NULL)
          return i;
@@ -96,6 +126,19 @@ static struct live_def *
 alloc_reg(struct ra_predicates_ctx *ctx, struct block_liveness *live,
           struct ir3_register *def, struct ir3_register *reloaded_def)
 {
+   /* Try to assign the precolored register if it's free. If not, use normal
+    * allocation and reload whenever a precolored source needs it.
+    * NOTE: this means we currently only support precolored sources, not dests.
+    */
+   if (def->num != INVALID_REG) {
+      assert(ctx->outstanding_precolored_defs > 0);
+      ctx->outstanding_precolored_defs--;
+      unsigned comp = reg_comp(def);
+
+      if (reg_is_free(ctx, live, comp))
+         return assign_reg(ctx, live, def, reloaded_def, comp);
+   }
+
    unsigned comp = alloc_reg_comp(ctx, live);
    return assign_reg(ctx, live, def, reloaded_def, comp);
 }
@@ -271,12 +314,28 @@ ra_block(struct ra_predicates_ctx *ctx, struct ir3_block *block)
             continue;
 
          struct live_def *live_def = find_live_def(ctx, live, src->def);
-         if (live_def == NULL)
+         if (src->num != INVALID_REG &&
+             (!live_def || get_def(live_def)->num != src->num)) {
+            /* If src is precolored and its def is either not live or is live in
+             * the wrong register, reload it into the correct one.
+             */
+            unsigned comp = reg_comp(src);
+
+            if (!reg_is_free(ctx, live, comp))
+               free_reg(ctx, live, get_def(&live->live_defs[comp]));
+            if (live_def)
+               free_reg(ctx, live, get_def(live_def));
+
+            live_def = reload_into(ctx, live, src->def, instr, comp);
+         } else if (!live_def) {
             live_def = reload(ctx, live, src->def, instr);
+         }
 
          assert(live_def != NULL);
 
          struct ir3_register *def = get_def(live_def);
+
+         assert((src->num == INVALID_REG) || (src->num == def->num));
          src->num = def->num;
          src->def = def;
 
@@ -425,6 +484,37 @@ init_block_liveness(struct ra_predicates_ctx *ctx, struct ir3_block *block)
           sizeof(struct live_def) * ctx->num_regs);
 }
 
+static void
+precolor_def(struct ra_predicates_ctx *ctx, struct ir3_register *def)
+{
+   foreach_ssa_use (use, def->instr) {
+      foreach_src (src, use) {
+         if (src->def != def)
+            continue;
+         if (src->num == INVALID_REG)
+            continue;
+
+         def->num = src->num;
+         ctx->outstanding_precolored_defs++;
+
+         /* We can only precolor a def once. */
+         return;
+      }
+   }
+}
+
+/* Precolor the defs of precolored sources so that we can try to assign the
+ * correct register immediately.
+ */
+static void
+precolor_defs(struct ra_predicates_ctx *ctx)
+{
+   for (unsigned i = 1; i < ctx->liveness->definitions_count; ++i) {
+      struct ir3_register *def = ctx->liveness->definitions[i];
+      precolor_def(ctx, def);
+   }
+}
+
 void
 ir3_ra_predicates(struct ir3_shader_variant *v)
 {
@@ -437,6 +527,7 @@ ir3_ra_predicates(struct ir3_shader_variant *v)
       rzalloc_array(ctx, struct block_liveness, ctx->liveness->block_count);
    ir3_count_instructions_ra(ctx->ir);
    ir3_find_ssa_uses_for(ctx->ir, ctx, is_predicate_use);
+   precolor_defs(ctx);
 
    foreach_block (block, &v->ir->block_list) {
       init_block_liveness(ctx, block);