nir: add support for user defined loop control

This will allow us to make use of the loop control support in
spirv and the GL support provided by EXT_control_flow_attributes.

Reviewed-by: Bas Nieuwenhuizen <[email protected]>
Bugzilla: https://bugs.freedesktop.org/show_bug.cgi?id=108841

3 files changed, 22 insertions, 5 deletions

diff --git a/src/compiler/nir/nir.h b/src/compiler/nir/nir.h
index e53b501d959..6e9b30a0e17 100644
--- a/src/compiler/nir/nir.h
+++ b/src/compiler/nir/nir.h
@@ -1968,12 +1968,19 @@ typedef struct {
    struct list_head loop_terminator_list;
 } nir_loop_info;
 
+typedef enum {
+   nir_loop_control_none = 0x0,
+   nir_loop_control_unroll = 0x1,
+   nir_loop_control_dont_unroll = 0x2,
+} nir_loop_control;
+
 typedef struct {
    nir_cf_node cf_node;
 
    struct exec_list body; /** < list of nir_cf_node */
 
    nir_loop_info *info;
+   nir_loop_control control;
    bool partially_unrolled;
 } nir_loop;
 
diff --git a/src/compiler/nir/nir_clone.c b/src/compiler/nir/nir_clone.c
index 9594eb8044d..2dbab60d28b 100644
--- a/src/compiler/nir/nir_clone.c
+++ b/src/compiler/nir/nir_clone.c
@@ -552,6 +552,7 @@ static nir_loop *
 clone_loop(clone_state *state, struct exec_list *cf_list, const nir_loop *loop)
 {
    nir_loop *nloop = nir_loop_create(state->ns);
+   nloop->control = loop->control;
    nloop->partially_unrolled = loop->partially_unrolled;
 
    nir_cf_node_insert_end(cf_list, &nloop->cf_node);
diff --git a/src/compiler/nir/nir_opt_loop_unroll.c b/src/compiler/nir/nir_opt_loop_unroll.c
index b2696d4aadb..c0198390749 100644
--- a/src/compiler/nir/nir_opt_loop_unroll.c
+++ b/src/compiler/nir/nir_opt_loop_unroll.c
@@ -770,11 +770,20 @@ partial_unroll(nir_shader *shader, nir_loop *loop, unsigned trip_count)
    _mesa_hash_table_destroy(remap_table, NULL);
 }
 
+/*
+ * Returns true if we should unroll the loop, otherwise false.
+ */
 static bool
-is_loop_small_enough_to_unroll(nir_shader *shader, nir_loop_info *li)
+check_unrolling_restrictions(nir_shader *shader, nir_loop *loop)
 {
-   unsigned max_iter = shader->options->max_unroll_iterations;
+   if (loop->control == nir_loop_control_unroll)
+      return true;
 
+   if (loop->control == nir_loop_control_dont_unroll)
+      return false;
+
+   nir_loop_info *li = loop->info;
+   unsigned max_iter = shader->options->max_unroll_iterations;
    unsigned trip_count =
       li->max_trip_count ? li->max_trip_count : li->guessed_trip_count;
 
@@ -822,7 +831,7 @@ process_loops(nir_shader *sh, nir_cf_node *cf_node, bool *has_nested_loop_out)
    /* Don't attempt to unroll a second inner loop in this pass, wait until the
     * next pass as we have altered the cf.
     */
-   if (!progress) {
+   if (!progress && loop->control != nir_loop_control_dont_unroll) {
 
       /* Check for the classic
        *
@@ -848,7 +857,7 @@ process_loops(nir_shader *sh, nir_cf_node *cf_node, bool *has_nested_loop_out)
          unsigned num_lt = list_length(&loop->info->loop_terminator_list);
          if (!has_nested_loop && num_lt == 1 && !loop->partially_unrolled &&
              loop->info->guessed_trip_count &&
-             is_loop_small_enough_to_unroll(sh, loop->info)) {
+             check_unrolling_restrictions(sh, loop)) {
             partial_unroll(sh, loop, loop->info->guessed_trip_count);
             progress = true;
          }
@@ -857,7 +866,7 @@ process_loops(nir_shader *sh, nir_cf_node *cf_node, bool *has_nested_loop_out)
       if (has_nested_loop || !loop->info->limiting_terminator)
          goto exit;
 
-      if (!is_loop_small_enough_to_unroll(sh, loop->info))
+      if (!check_unrolling_restrictions(sh, loop))
          goto exit;
 
       if (loop->info->exact_trip_count_known) {