[Mesa-dev] [PATCH v2 04/10] nir: calculate trip count for more loops

[Timothy Arceri]

This adds support to loop analysis for loops where the induction
variable is compared to the result of min(variable, constant).

For example:

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

We add a new bool to the loop terminator struct in order to
differentiate terminators with this exit condition.
---
 src/compiler/nir/nir.h                 | 11 +++++++
 src/compiler/nir/nir_loop_analyze.c    | 41 ++++++++++++++++++++++----
 src/compiler/nir/nir_opt_loop_unroll.c |  3 +-
 3 files changed, 49 insertions(+), 6 deletions(-)

diff --git a/src/compiler/nir/nir.h b/src/compiler/nir/nir.h
index 055d4d30d3..4bbc750bd6 100644
--- a/src/compiler/nir/nir.h
+++ b/src/compiler/nir/nir.h
@@ -1881,6 +1881,17 @@ typedef struct {
    bool continue_from_then;
    bool induction_rhs;
 
+   /* This is true if the terminators exact trip count is unknown. For
+    * example:
+    *
+    *    for (int i = 0; i < imin(x, 4); i++)
+    *       ...
+    *
+    * Here loop analysis would have set a max_trip_count of 4 however we dont
+    * know for sure that this is the exact trip count.
+    */
+   bool exact_trip_count_unknown;
+
    struct list_head loop_terminator_link;
 } nir_loop_terminator;
 
diff --git a/src/compiler/nir/nir_loop_analyze.c b/src/compiler/nir/nir_loop_analyze.c
index c46c491963..0b37da6569 100644
--- a/src/compiler/nir/nir_loop_analyze.c
+++ b/src/compiler/nir/nir_loop_analyze.c
@@ -477,6 +477,35 @@ guess_loop_limit(loop_info_state *state, nir_const_value *limit_val,
    return false;
 }
 
+static bool
+try_find_limit_of_alu(nir_loop_variable *limit, nir_const_value *limit_val,
+                      nir_loop_terminator *terminator, loop_info_state *state)
+{
+   if(!is_var_alu(limit))
+      return false;
+
+   nir_alu_instr *limit_alu = nir_instr_as_alu(limit->def->parent_instr);
+
+   if (limit_alu->op == nir_op_imin ||
+       limit_alu->op == nir_op_fmin) {
+      limit = get_loop_var(limit_alu->src[0].src.ssa, state);
+
+      if (!is_var_constant(limit))
+         limit = get_loop_var(limit_alu->src[1].src.ssa, state);
+
+      if (!is_var_constant(limit))
+         return false;
+
+      *limit_val = nir_instr_as_load_const(limit->def->parent_instr)->value;
+
+      terminator->exact_trip_count_unknown = true;
+
+      return true;
+   }
+
+   return false;
+}
+
 static int32_t
 get_iteration(nir_op cond_op, nir_const_value *initial, nir_const_value *step,
               nir_const_value *limit)
@@ -708,12 +737,14 @@ find_trip_count(loop_info_state *state)
          } else {
             trip_count_known = false;
 
-            /* Guess loop limit based on array access */
-            if (!guess_loop_limit(state, &limit_val, basic_ind)) {
-               continue;
-            }
+            if (!try_find_limit_of_alu(limit, &limit_val, terminator, state)) {
+               /* Guess loop limit based on array access */
+               if (!guess_loop_limit(state, &limit_val, basic_ind)) {
+                  continue;
+               }
 
-            guessed_trip_count = true;
+               guessed_trip_count = true;
+            }
          }
 
          /* We have determined that we have the following constants:
diff --git a/src/compiler/nir/nir_opt_loop_unroll.c b/src/compiler/nir/nir_opt_loop_unroll.c
index d8df619b32..9630e0738a 100644
--- a/src/compiler/nir/nir_opt_loop_unroll.c
+++ b/src/compiler/nir/nir_opt_loop_unroll.c
@@ -830,7 +830,8 @@ process_loops(nir_shader *sh, nir_cf_node *cf_node, bool *has_nested_loop_out)
       } else {
          /* Attempt to unroll loops with two terminators. */
          unsigned num_lt = list_length(&loop->info->loop_terminator_list);
-         if (num_lt == 2) {
+         if (num_lt == 2 &&
+             !loop->info->limiting_terminator->exact_trip_count_unknown) {
             bool limiting_term_second = true;
             nir_loop_terminator *terminator =
                list_first_entry(&loop->info->loop_terminator_list,
-- 
2.19.2