nir: add guess trip count support to loop analysis

This detects an induction variable used as an array index to guess
the trip count of the loop. This enables us to do a partial
unroll of the loop, which can eventually result in the loop being
eliminated.

v2: check if the induction var is used to index more than a single
    array and if so get the size of the smallest array.

Reviewed-by: Ian Romanick <[email protected]>

2 files changed, 86 insertions, 6 deletions

diff --git a/src/compiler/nir/nir.h b/src/compiler/nir/nir.h
index 7be62ba97ae..39b4c2aaf3e 100644
--- a/src/compiler/nir/nir.h
+++ b/src/compiler/nir/nir.h
@@ -1910,6 +1910,7 @@ typedef struct {
 
    /** True when ::break_block is in the else-path of ::nif. */
    bool continue_from_then;
+   bool induction_rhs;
 
    struct list_head loop_terminator_link;
 } nir_loop_terminator;
@@ -1918,6 +1919,9 @@ typedef struct {
    /* Estimated cost (in number of instructions) of the loop */
    unsigned instr_cost;
 
+   /* Guessed trip count based on array indexing */
+   unsigned guessed_trip_count;
+
    /* Maximum number of times the loop is run (if known) */
    unsigned max_trip_count;
 
diff --git a/src/compiler/nir/nir_loop_analyze.c b/src/compiler/nir/nir_loop_analyze.c
index 2ca021e51f1..9cff670051e 100644
--- a/src/compiler/nir/nir_loop_analyze.c
+++ b/src/compiler/nir/nir_loop_analyze.c
@@ -488,6 +488,57 @@ find_array_access_via_induction(loop_info_state *state,
    return 0;
 }
 
+static bool
+guess_loop_limit(loop_info_state *state, nir_const_value *limit_val,
+                 nir_loop_variable *basic_ind)
+{
+   unsigned min_array_size = 0;
+
+   nir_foreach_block_in_cf_node(block, &state->loop->cf_node) {
+      nir_foreach_instr(instr, block) {
+         if (instr->type != nir_instr_type_intrinsic)
+            continue;
+
+         nir_intrinsic_instr *intrin = nir_instr_as_intrinsic(instr);
+
+         /* Check for arrays variably-indexed by a loop induction variable. */
+         if (intrin->intrinsic == nir_intrinsic_load_deref ||
+             intrin->intrinsic == nir_intrinsic_store_deref ||
+             intrin->intrinsic == nir_intrinsic_copy_deref) {
+
+            nir_loop_variable *array_idx = NULL;
+            unsigned array_size =
+               find_array_access_via_induction(state,
+                                               nir_src_as_deref(intrin->src[0]),
+                                               &array_idx);
+            if (basic_ind == array_idx &&
+                (min_array_size == 0 || min_array_size > array_size)) {
+               min_array_size = array_size;
+            }
+
+            if (intrin->intrinsic != nir_intrinsic_copy_deref)
+               continue;
+
+            array_size =
+               find_array_access_via_induction(state,
+                                               nir_src_as_deref(intrin->src[1]),
+                                               &array_idx);
+            if (basic_ind == array_idx &&
+                (min_array_size == 0 || min_array_size > array_size)) {
+               min_array_size = array_size;
+            }
+         }
+      }
+   }
+
+   if (min_array_size) {
+      limit_val->i32[0] = min_array_size;
+      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)
@@ -664,6 +715,7 @@ static void
 find_trip_count(loop_info_state *state)
 {
    bool trip_count_known = true;
+   bool guessed_trip_count = false;
    nir_loop_terminator *limiting_terminator = NULL;
    int max_trip_count = -1;
 
@@ -699,16 +751,33 @@ find_trip_count(loop_info_state *state)
             basic_ind = get_loop_var(alu->src[1].src.ssa, state);
             limit = get_loop_var(alu->src[0].src.ssa, state);
             limit_rhs = false;
+            terminator->induction_rhs = true;
          }
 
-         /* The comparison has to have a basic induction variable
-          * and a constant for us to be able to find trip counts
+         /* The comparison has to have a basic induction variable for us to be
+          * able to find trip counts.
           */
-         if (basic_ind->type != basic_induction || !is_var_constant(limit)) {
+         if (basic_ind->type != basic_induction) {
             trip_count_known = false;
             continue;
          }
 
+         /* Attempt to find a constant limit for the loop */
+         nir_const_value limit_val;
+         if (is_var_constant(limit)) {
+            limit_val =
+               nir_instr_as_load_const(limit->def->parent_instr)->value;
+         } else {
+            trip_count_known = false;
+
+            /* Guess loop limit based on array access */
+            if (!guess_loop_limit(state, &limit_val, basic_ind)) {
+               continue;
+            }
+
+            guessed_trip_count = true;
+         }
+
          /* We have determined that we have the following constants:
           * (With the typical int i = 0; i < x; i++; as an example)
           *    - Upper limit.
@@ -725,9 +794,6 @@ find_trip_count(loop_info_state *state)
             nir_instr_as_load_const(basic_ind->ind->invariant->def->
                                        parent_instr)->value;
 
-         nir_const_value limit_val =
-            nir_instr_as_load_const(limit->def->parent_instr)->value;
-
          int iterations = calculate_iterations(&initial_val, &step_val,
                                                &limit_val,
                                                basic_ind->ind->alu_def, alu,
@@ -737,6 +803,16 @@ find_trip_count(loop_info_state *state)
          /* Where we not able to calculate the iteration count */
          if (iterations == -1) {
             trip_count_known = false;
+            guessed_trip_count = false;
+            continue;
+         }
+
+         if (guessed_trip_count) {
+            guessed_trip_count = false;
+            if (state->loop->info->guessed_trip_count == 0 ||
+                state->loop->info->guessed_trip_count > iterations)
+              state->loop->info->guessed_trip_count = iterations;
+
             continue;
          }