[Mesa-dev] [PATCH v2 05/10] nir: unroll some loops with a variable limit

[Timothy Arceri]

For some loops can have a single terminator but the exact trip
count is still unknown. For example:

   for (int i = 0; i < imin(x, 4); i++)
      ...

Shader-db results radeonsi (all affected are from Tropico 5):

Totals from affected shaders:
SGPRS: 200 -> 208 (4.00 %)
VGPRS: 164 -> 148 (-9.76 %)
Spilled SGPRs: 0 -> 0 (0.00 %)
Spilled VGPRs: 0 -> 0 (0.00 %)
Private memory VGPRs: 0 -> 0 (0.00 %)
Scratch size: 0 -> 0 (0.00 %) dwords per thread
Code Size: 7208 -> 8672 (20.31 %) bytes
LDS: 0 -> 0 (0.00 %) blocks
Max Waves: 23 -> 27 (17.39 %)
Wait states: 0 -> 0 (0.00 %)

vkpipeline-db results RADV (Unrolls some Skyrim VR shaders):

Totals from affected shaders:
SGPRS: 304 -> 304 (0.00 %)
VGPRS: 300 -> 292 (-2.67 %)
Spilled SGPRs: 0 -> 0 (0.00 %)
Spilled VGPRs: 0 -> 0 (0.00 %)
Private memory VGPRs: 0 -> 0 (0.00 %)
Scratch size: 0 -> 0 (0.00 %) dwords per thread
Code Size: 15516 -> 26388 (70.07 %) bytes
LDS: 0 -> 0 (0.00 %) blocks
Max Waves: 29 -> 29 (0.00 %)
Wait states: 0 -> 0 (0.00 %)

v2: fix bug where last iteration would get optimised away by
    mistake.
---
 src/compiler/nir/nir_opt_loop_unroll.c | 55 ++++++++++++++++++++++++++
 1 file changed, 55 insertions(+)

diff --git a/src/compiler/nir/nir_opt_loop_unroll.c b/src/compiler/nir/nir_opt_loop_unroll.c
index 9630e0738a..70e6c67bde 100644
--- a/src/compiler/nir/nir_opt_loop_unroll.c
+++ b/src/compiler/nir/nir_opt_loop_unroll.c
@@ -460,6 +460,55 @@ complex_unroll(nir_loop *loop, nir_loop_terminator *unlimit_term,
    _mesa_hash_table_destroy(remap_table, NULL);
 }
 
+/**
+ * Unroll loops where we only have a single terminator but the exact trip
+ * count is unknown. For example:
+ *
+ *    for (int i = 0; i < imin(x, 4); i++)
+ *       ...
+ */
+static void
+complex_unroll_single_terminator(nir_loop *loop)
+{
+   assert(list_length(&loop->info->loop_terminator_list) == 1);
+   assert(loop->info->limiting_terminator);
+   assert(nir_is_trivial_loop_if(loop->info->limiting_terminator->nif,
+                                 loop->info->limiting_terminator->break_block));
+
+   nir_loop_terminator *terminator = loop->info->limiting_terminator;
+
+   loop_prepare_for_unroll(loop);
+
+   /* Pluck out the loop header */
+   nir_cf_list lp_header;
+   nir_cf_extract(&lp_header, nir_before_block(nir_loop_first_block(loop)),
+                  nir_before_cf_node(&terminator->nif->cf_node));
+
+   struct hash_table *remap_table =
+      _mesa_hash_table_create(NULL, _mesa_hash_pointer,
+                              _mesa_key_pointer_equal);
+
+   /* We need to clone the loop one extra time in order to clone the lcssa
+    * vars for the last iteration (they are inside the following ifs break
+    * branch). We leave other passes to clean up this redundant if.
+    */
+   unsigned num_times_to_clone = loop->info->max_trip_count + 1;
+
+   nir_cf_list lp_body;
+   nir_cf_node *unroll_loc =
+      complex_unroll_loop_body(loop, terminator, &lp_header, &lp_body,
+                               remap_table, num_times_to_clone);
+
+   /* Delete the original loop header and body */
+   nir_cf_delete(&lp_header);
+   nir_cf_delete(&lp_body);
+
+   /* The original loop has been replaced so remove it. */
+   nir_cf_node_remove(&loop->cf_node);
+
+   _mesa_hash_table_destroy(remap_table, NULL);
+}
+
 /* Unrolls the classic wrapper loops e.g
  *
  *    do {
@@ -856,6 +905,12 @@ process_loops(nir_shader *sh, nir_cf_node *cf_node, bool *has_nested_loop_out)
             }
             progress = true;
          }
+
+         if (num_lt == 1) {
+            assert(loop->info->limiting_terminator->exact_trip_count_unknown);
+            complex_unroll_single_terminator(loop);
+            progress = true;
+         }
       }
    }
 
-- 
2.19.2