nir/loop_analyze: record induction variables for each loop

For being used by uniform inline lowering pass.

Reviewed-by: Marek Olšák <marek.olsak@amd.com>
Signed-off-by: Qiang Yu <yuq825@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/11950>

src/compiler/nir/nir.h

@@ -3016,6 +3016,17 @@ typedef struct {
    struct list_head loop_terminator_link;
 } nir_loop_terminator;
 
+typedef struct {
+   /* Induction variable. */
+   nir_ssa_def *def;
+
+   /* Init statement with only uniform. */
+   nir_src *init_src;
+
+   /* Update statement with only uniform. */
+   nir_alu_src *update_src;
+} nir_loop_induction_variable;
+
 typedef struct {
    /* Estimated cost (in number of instructions) of the loop */
    unsigned instr_cost;
@@ -3042,6 +3053,10 @@ typedef struct {
 
    /* A list of loop_terminators terminating this loop. */
    struct list_head loop_terminator_list;
+
+   /* array of induction variables for this loop */
+   nir_loop_induction_variable *induction_vars;
+   unsigned num_induction_vars;
 } nir_loop_info;
 
 typedef enum {

src/compiler/nir/nir_loop_analyze.c

@@ -59,6 +59,9 @@ typedef struct {
    /* True if variable is in a nested loop */
    bool in_nested_loop;
 
+   /* Could be a basic_induction if following uniforms are inlined */
+   nir_src *init_src;
+   nir_alu_src *update_src;
 } nir_loop_variable;
 
 typedef struct {
@@ -86,6 +89,8 @@ get_loop_var(nir_ssa_def *value, loop_info_state *state)
       var->def = value;
       var->in_if_branch = false;
       var->in_nested_loop = false;
+      var->init_src = NULL;
+      var->update_src = NULL;
       if (value->parent_instr->type == nir_instr_type_load_const)
          var->type = invariant;
       else
@@ -319,10 +324,46 @@ alu_src_has_identity_swizzle(nir_alu_instr *alu, unsigned src_idx)
    return true;
 }
 
+static bool
+is_only_uniform_src(nir_src *src)
+{
+   if (!src->is_ssa)
+      return false;
+
+   nir_instr *instr = src->ssa->parent_instr;
+
+   switch (instr->type) {
+   case nir_instr_type_alu: {
+      /* Return true if all sources return true. */
+      nir_alu_instr *alu = nir_instr_as_alu(instr);
+      for (unsigned i = 0; i < nir_op_infos[alu->op].num_inputs; i++) {
+         if (!is_only_uniform_src(&alu->src[i].src))
+             return false;
+      }
+      return true;
+   }
+
+   case nir_instr_type_intrinsic: {
+      nir_intrinsic_instr *inst = nir_instr_as_intrinsic(instr);
+      /* current uniform inline only support load ubo */
+      return inst->intrinsic == nir_intrinsic_load_ubo;
+   }
+
+   case nir_instr_type_load_const:
+      /* Always return true for constants. */
+      return true;
+
+   default:
+      return false;
+   }
+}
+
 static bool
 compute_induction_information(loop_info_state *state)
 {
    bool found_induction_var = false;
+   unsigned num_induction_vars = 0;
+
    list_for_each_entry_safe(nir_loop_variable, var, &state->process_list,
                             process_link) {
 
@@ -343,6 +384,7 @@ compute_induction_information(loop_info_state *state)
       nir_phi_instr *phi = nir_instr_as_phi(var->def->parent_instr);
       nir_basic_induction_var *biv = rzalloc(state, nir_basic_induction_var);
 
+      nir_src *init_src = NULL;
       nir_loop_variable *alu_src_var = NULL;
       nir_foreach_phi_src(src, phi) {
          nir_loop_variable *src_var = get_loop_var(src->src.ssa, state);
@@ -369,19 +411,29 @@ compute_induction_information(loop_info_state *state)
 
          if (!src_var->in_loop && !biv->def_outside_loop) {
             biv->def_outside_loop = src_var->def;
+            init_src = &src->src;
          } else if (is_var_alu(src_var) && !biv->alu) {
             alu_src_var = src_var;
             nir_alu_instr *alu = nir_instr_as_alu(src_var->def->parent_instr);
 
             if (nir_op_infos[alu->op].num_inputs == 2) {
                for (unsigned i = 0; i < 2; i++) {
-                  /* Is one of the operands const, and the other the phi.  The
+                  /* Is one of the operands const or uniform, and the other the phi.
-                   * phi source can't be swizzled in any way.
+                   * The phi source can't be swizzled in any way.
                    */
-                  if (nir_src_is_const(alu->src[i].src) &&
+                  if (alu->src[1-i].src.ssa == &phi->dest.ssa &&
-                      alu->src[1-i].src.ssa == &phi->dest.ssa &&
+                      alu_src_has_identity_swizzle(alu, 1 - i)) {
-                      alu_src_has_identity_swizzle(alu, 1 - i))
+                     nir_src *src = &alu->src[i].src;
-                     biv->alu = alu;
+                     if (nir_src_is_const(*src))
+                        biv->alu = alu;
+                     else if (is_only_uniform_src(src)) {
+                        /* Update value of induction variable is a statement
+                         * contains only uniform and constant
+                         */
+                        var->update_src = alu->src + i;
+                        biv->alu = alu;
+                     }
+                  }
                }
             }
 
@@ -393,18 +445,64 @@ compute_induction_information(loop_info_state *state)
          }
       }
 
-      if (biv->alu && biv->def_outside_loop &&
+      if (biv->alu && biv->def_outside_loop) {
-          biv->def_outside_loop->parent_instr->type == nir_instr_type_load_const) {
+         nir_instr *inst = biv->def_outside_loop->parent_instr;
-         alu_src_var->type = basic_induction;
+         if (inst->type == nir_instr_type_load_const)  {
-         alu_src_var->ind = biv;
+            /* Initial value of induction variable is a constant */
-         var->type = basic_induction;
+            if (var->update_src) {
-         var->ind = biv;
+               alu_src_var->update_src = var->update_src;
+               ralloc_free(biv);
+            } else {
+               alu_src_var->type = basic_induction;
+               alu_src_var->ind = biv;
+               var->type = basic_induction;
+               var->ind = biv;
 
-         found_induction_var = true;
+               found_induction_var = true;
+            }
+            num_induction_vars += 2;
+         } else if (is_only_uniform_src(init_src)) {
+            /* Initial value of induction variable is a uniform */
+            var->init_src = init_src;
+
+            alu_src_var->init_src = var->init_src;
+            alu_src_var->update_src = var->update_src;
+
+            num_induction_vars += 2;
+            ralloc_free(biv);
+         } else {
+            var->update_src = NULL;
+            ralloc_free(biv);
+         }
       } else {
+         var->update_src = NULL;
          ralloc_free(biv);
       }
    }
+
+   nir_loop_info *info = state->loop->info;
+   ralloc_free(info->induction_vars);
+   info->num_induction_vars = 0;
+
+   /* record induction variables into nir_loop_info */
+   if (num_induction_vars) {
+      info->induction_vars = ralloc_array(info, nir_loop_induction_variable,
+                                          num_induction_vars);
+
+      list_for_each_entry(nir_loop_variable, var, &state->process_list,
+                          process_link) {
+         if (var->type == basic_induction || var->init_src || var->update_src) {
+            nir_loop_induction_variable *ivar =
+               &info->induction_vars[info->num_induction_vars++];
+             ivar->def = var->def;
+             ivar->init_src = var->init_src;
+             ivar->update_src = var->update_src;
+         }
+      }
+      /* don't overflow */
+      assert(info->num_induction_vars <= num_induction_vars);
+   }
+
    return found_induction_var;
 }