From c24373d907a135d3a5920449cf9bc91211e34f03 Mon Sep 17 00:00:00 2001
From: Job Noorman <jnoorman@igalia.com>
Date: Fri, 28 Feb 2025 09:09:08 +0100
Subject: [PATCH] ir3/sched: unblock a0.x/a1.x after last use

We currently keep a0.x/a1.x unnecessarily blocked until we have to spill
it, even if all its uses have been scheduled already. This causes
unnecessary spills and potentially schedules address writes later than
we'd like.

Fix this by keeping track of the number of uses that are still
unscheduled, unblocking address writes when it drops to zero.

Signed-off-by: Job Noorman <jnoorman@igalia.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/33816>
---
 src/freedreno/ir3/ir3.h       | 12 ++++++++++++
 src/freedreno/ir3/ir3_sched.c | 24 ++++++++++++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/src/freedreno/ir3/ir3.h b/src/freedreno/ir3/ir3.h
index b11ec10b011..e438464c8df 100644
--- a/src/freedreno/ir3/ir3.h
+++ b/src/freedreno/ir3/ir3.h
@@ -1512,6 +1512,18 @@ writes_addr1(struct ir3_instruction *instr)
    return false;
 }
 
+static inline bool
+reads_addr0(struct ir3_instruction *instr)
+{
+   return instr->address && instr->address->num == regid(REG_A0, 0);
+}
+
+static inline bool
+reads_addr1(struct ir3_instruction *instr)
+{
+   return instr->address && instr->address->num == regid(REG_A0, 1);
+}
+
 static inline bool
 writes_pred(struct ir3_instruction *instr)
 {
diff --git a/src/freedreno/ir3/ir3_sched.c b/src/freedreno/ir3/ir3_sched.c
index ddda8bf8b70..e68c9b24d9a 100644
--- a/src/freedreno/ir3/ir3_sched.c
+++ b/src/freedreno/ir3/ir3_sched.c
@@ -80,6 +80,8 @@ struct ir3_sched_ctx {
    struct ir3_instruction *scheduled; /* last scheduled instr */
    struct ir3_instruction *addr0;     /* current a0.x user, if any */
    struct ir3_instruction *addr1;     /* current a1.x user, if any */
+   unsigned addr0_uses;               /* number of unscheduled uses of addr0 */
+   unsigned addr1_uses;               /* number of unscheduled uses of addr1 */
 
    struct ir3_instruction *split; /* most-recently-split a0/a1 producer */
 
@@ -256,11 +258,31 @@ schedule(struct ir3_sched_ctx *ctx, struct ir3_instruction *instr)
    if (writes_addr0(instr)) {
       assert(ctx->addr0 == NULL);
       ctx->addr0 = instr;
+      ctx->addr0_uses = instr->uses->entries;
    }
 
    if (writes_addr1(instr)) {
       assert(ctx->addr1 == NULL);
       ctx->addr1 = instr;
+      ctx->addr1_uses = instr->uses->entries;
+   }
+
+   if (reads_addr0(instr)) {
+      assert(instr->address->def->instr == ctx->addr0);
+      assert(ctx->addr0_uses > 0);
+
+      if (--ctx->addr0_uses == 0) {
+         ctx->addr0 = NULL;
+      }
+   }
+
+   if (reads_addr1(instr)) {
+      assert(instr->address->def->instr == ctx->addr1);
+      assert(ctx->addr1_uses > 0);
+
+      if (--ctx->addr1_uses == 0) {
+         ctx->addr1 = NULL;
+      }
    }
 
    instr->flags |= IR3_INSTR_MARK;
@@ -914,8 +936,10 @@ split_addr(struct ir3_sched_ctx *ctx, struct ir3_instruction **addr,
             new_addr = split_instr(ctx, *addr);
             /* original addr is scheduled, but new one isn't: */
             new_addr->flags &= ~IR3_INSTR_MARK;
+            new_addr->uses = _mesa_pointer_set_create(ctx);
          }
          indirect->address->def = new_addr->dsts[0];
+         _mesa_set_add(new_addr->uses, indirect);
          /* don't need to remove old dag edge since old addr is
           * already scheduled:
           */