ra: Add a callback for selecting a register from what's available.

VC4 has had a tension, similar to pre-Sandybridge Intel, where we want to
use low-numbered registers (more parallelism on Intel, fewer delay slots
on vc4), but in order to give instruction scheduling the most freedom to
avoid delays we want to round-robin between registers of the same cost.
Our two heuristics so far have chosen one end or the other of that
tradeoff.

The callback, instead, hands the driver the set of registers that are
available, and the driver gets to make its own choice.  This will be used
in vc4 to round-robin between registers of the same cost, and might be
used in the future for improving bank selection.

Reviewed-by: Nicolai Hähnle <nicolai.haehnle@amd.com>

src/util/register_allocate.c

@@ -174,6 +174,10 @@ struct ra_graph {
     * stack.
     */
    unsigned int stack_optimistic_start;
+
+   unsigned int (*select_reg_callback)(struct ra_graph *g, BITSET_WORD *regs,
+                                       void *data);
+   void *select_reg_callback_data;
 };
 
 /**
@@ -439,6 +443,16 @@ ra_alloc_interference_graph(struct ra_regs *regs, unsigned int count)
    return g;
 }
 
+void ra_set_select_reg_callback(struct ra_graph *g,
+                                unsigned int (*callback)(struct ra_graph *g,
+                                                         BITSET_WORD *regs,
+                                                         void *data),
+                                void *data)
+{
+   g->select_reg_callback = callback;
+   g->select_reg_callback_data = data;
+}
+
 void
 ra_set_node_class(struct ra_graph *g,
                   unsigned int n, unsigned int class)
@@ -555,6 +569,41 @@ ra_any_neighbors_conflict(struct ra_graph *g, unsigned int n, unsigned int r)
    return false;
 }
 
+/* Computes a bitfield of what regs are available for a given register
+ * selection.
+ *
+ * This lets drivers implement a more complicated policy than our simple first
+ * or round robin policies (which don't require knowing the whole bitset)
+ */
+static bool
+ra_compute_available_regs(struct ra_graph *g, unsigned int n, BITSET_WORD *regs)
+{
+   struct ra_class *c = g->regs->classes[g->nodes[n].class];
+
+   /* Populate with the set of regs that are in the node's class. */
+   memcpy(regs, c->regs, BITSET_WORDS(g->regs->count) * sizeof(BITSET_WORD));
+
+   /* Remove any regs that conflict with nodes that we're adjacent to and have
+    * already colored.
+    */
+   for (int i = 0; i < g->nodes[n].adjacency_count; i++) {
+      unsigned int n2 = g->nodes[n].adjacency_list[i];
+      unsigned int r = g->nodes[n2].reg;
+
+      if (!g->nodes[n2].in_stack) {
+         for (int j = 0; j < BITSET_WORDS(g->regs->count); j++)
+            regs[j] &= ~g->regs->regs[r].conflicts[j];
+      }
+   }
+
+   for (int i = 0; i < BITSET_WORDS(g->regs->count); i++) {
+      if (regs[i])
+         return true;
+   }
+
+   return false;
+}
+
 /**
  * Pops nodes from the stack back into the graph, coloring them with
  * registers as they go.
@@ -566,6 +615,10 @@ static bool
 ra_select(struct ra_graph *g)
 {
    int start_search_reg = 0;
+   BITSET_WORD *select_regs = NULL;
+
+   if (g->select_reg_callback)
+      select_regs = malloc(BITSET_WORDS(g->regs->count) * sizeof(BITSET_WORD));
 
    while (g->stack_count != 0) {
       unsigned int ri;
@@ -573,25 +626,34 @@ ra_select(struct ra_graph *g)
       int n = g->stack[g->stack_count - 1];
       struct ra_class *c = g->regs->classes[g->nodes[n].class];
 
-      /* Find the lowest-numbered reg which is not used by a member
-       * of the graph adjacent to us.
-       */
-      for (ri = 0; ri < g->regs->count; ri++) {
-         r = (start_search_reg + ri) % g->regs->count;
-         if (!reg_belongs_to_class(r, c))
-	    continue;
-
-         if (!ra_any_neighbors_conflict(g, n, r))
-            break;
-      }
-
       /* set this to false even if we return here so that
        * ra_get_best_spill_node() considers this node later.
        */
       g->nodes[n].in_stack = false;
 
-      if (ri == g->regs->count)
-	 return false;
+      if (g->select_reg_callback) {
+         if (!ra_compute_available_regs(g, n, select_regs)) {
+            free(select_regs);
+            return false;
+         }
+
+         r = g->select_reg_callback(g, select_regs, g->select_reg_callback_data);
+      } else {
+         /* Find the lowest-numbered reg which is not used by a member
+          * of the graph adjacent to us.
+          */
+         for (ri = 0; ri < g->regs->count; ri++) {
+            r = (start_search_reg + ri) % g->regs->count;
+            if (!reg_belongs_to_class(r, c))
+               continue;
+
+            if (!ra_any_neighbors_conflict(g, n, r))
+               break;
+         }
+
+         if (ri >= g->regs->count)
+            return false;
+      }
 
       g->nodes[n].reg = r;
       g->stack_count--;

src/util/register_allocate.h

@@ -29,6 +29,7 @@
 #define REGISTER_ALLOCATE_H
 
 #include <stdbool.h>
+#include "util/bitset.h"
 
 #ifdef __cplusplus
 extern "C" {
@@ -74,6 +75,11 @@ void ra_set_finalize(struct ra_regs *regs, unsigned int **conflicts);
 struct ra_graph *ra_alloc_interference_graph(struct ra_regs *regs,
 					     unsigned int count);
 void ra_set_node_class(struct ra_graph *g, unsigned int n, unsigned int c);
+void ra_set_select_reg_callback(struct ra_graph *g,
+                                unsigned int (*callback)(struct ra_graph *g,
+                                                         BITSET_WORD *regs,
+                                                         void *data),
+                                void *data);
 void ra_add_node_interference(struct ra_graph *g,
 			      unsigned int n1, unsigned int n2);
 /** @} */