nir/schedule: Add a callback for backend-specific dependencies

Adds a callback function to nir_schedule_options to give the backend a
chance to add custom dependencies between certain intrinsics. The
callback can assign a class number to the intrinsic and then set a read
or write dependency on that class.

v2: Use a linked-list of schedule nodes for the dependency classes
    instead of a fixed-sized array.

Reviewed-by: Iago Toral Quiroga <itoral@igalia.com>
Reviewed-by: Alejandro Piñeiro <apinheiro@igalia.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/5953>

src/compiler/nir/nir_schedule.c

@@ -119,6 +119,12 @@ typedef struct {
  */
 enum direction { F, R };
 
+struct nir_schedule_class_dep {
+   int klass;
+   nir_schedule_node *node;
+   struct nir_schedule_class_dep *next;
+};
+
 typedef struct {
    nir_schedule_scoreboard *scoreboard;
 
@@ -133,6 +139,8 @@ typedef struct {
    nir_schedule_node *discard;
    nir_schedule_node *jump;
 
+   struct nir_schedule_class_dep *class_deps;
+
    enum direction dir;
 } nir_deps_state;
 
@@ -292,12 +300,52 @@ nir_schedule_ssa_deps(nir_ssa_def *def, void *in_state)
    return true;
 }
 
+static struct nir_schedule_class_dep *
+nir_schedule_get_class_dep(nir_deps_state *state,
+                           int klass)
+{
+   for (struct nir_schedule_class_dep *class_dep = state->class_deps;
+        class_dep != NULL;
+        class_dep = class_dep->next) {
+      if (class_dep->klass == klass)
+         return class_dep;
+   }
+
+   struct nir_schedule_class_dep *class_dep =
+      ralloc(state->reg_map, struct nir_schedule_class_dep);
+
+   class_dep->klass = klass;
+   class_dep->node = NULL;
+   class_dep->next = state->class_deps;
+
+   state->class_deps = class_dep;
+
+   return class_dep;
+}
+
 static void
 nir_schedule_intrinsic_deps(nir_deps_state *state,
                             nir_intrinsic_instr *instr)
 {
    nir_schedule_node *n = nir_schedule_get_node(state->scoreboard->instr_map,
                                                 &instr->instr);
+   const nir_schedule_options *options = state->scoreboard->options;
+   nir_schedule_dependency dep;
+
+   if (options->intrinsic_cb &&
+       options->intrinsic_cb(instr, &dep, options->intrinsic_cb_data)) {
+      struct nir_schedule_class_dep *class_dep =
+         nir_schedule_get_class_dep(state, dep.klass);
+
+      switch (dep.type) {
+      case NIR_SCHEDULE_READ_DEPENDENCY:
+         add_read_dep(state, class_dep->node, n);
+         break;
+      case NIR_SCHEDULE_WRITE_DEPENDENCY:
+         add_write_dep(state, &class_dep->node, n);
+         break;
+      }
+   }
 
    switch (instr->intrinsic) {
    case nir_intrinsic_load_uniform:

src/compiler/nir/nir_schedule.h

@@ -30,6 +30,23 @@
 extern "C" {
 #endif
 
+/**
+ * Struct filled in by the intrinsic_cb callback of nir_schedule_options to
+ * specify a backend-specific dependency on an intrinsic.
+ */
+typedef struct nir_schedule_dependency {
+   /* Which class of dependency this is. The meanings of the classes are
+    * specific to the backend. This must be less than
+    * NIR_SCHEDULE_N_DEPENDENCY_CLASSES.
+    */
+   int klass;
+   /* The type of dependency */
+   enum {
+      NIR_SCHEDULE_READ_DEPENDENCY,
+      NIR_SCHEDULE_WRITE_DEPENDENCY,
+   } type;
+} nir_schedule_dependency;
+
 typedef struct nir_schedule_options {
    /* On some hardware with some stages the inputs and outputs to the shader
     * share the same memory. In that case the scheduler needs to ensure that
@@ -41,6 +58,15 @@ typedef struct nir_schedule_options {
     * will try to reduce register usage.
     */
    int threshold;
+   /* Callback used to add custom dependencies on intrinsics. If it returns
+    * true then a dependency should be added and dep is filled in to describe
+    * it.
+    */
+   bool (* intrinsic_cb)(nir_intrinsic_instr *intr,
+                         nir_schedule_dependency *dep,
+                         void *user_data);
+   /* Data to pass to the callback */
+   void *intrinsic_cb_data;
 } nir_schedule_options;
 
 void nir_schedule(nir_shader *shader, const nir_schedule_options *options);