nir/loop_analyze: Track induction variables with uniform increments

As an intentional side effect, induction variables with constant
increments will now have their nir_loop_induction_variable::update_src
field point to the load_const source. Previously this pointer would be
NULL.

v2: Update unit tests and commit message.

Reviewed-by: Caio Oliveira <caio.oliveira@intel.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/21289>

src/compiler/nir/nir_loop_analyze.c

@@ -466,13 +466,7 @@ compute_induction_information(loop_info_state *state)
                    */
                   if (alu->src[1-i].src.ssa == &phi->dest.ssa &&
                       alu_src_has_identity_swizzle(alu, 1 - i)) {
-                     nir_src *src = &alu->src[i].src;
+                     if (is_only_uniform_src(&alu->src[i].src)) {
-                     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;
                      }
@@ -492,17 +486,14 @@ compute_induction_information(loop_info_state *state)
          nir_instr *inst = biv->def_outside_loop->parent_instr;
          if (inst->type == nir_instr_type_load_const)  {
             /* Initial value of induction variable is a constant */
-            if (var->update_src) {
+            alu_src_var->init_src = var->init_src;
             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;
-            }
             num_induction_vars += 2;
          } else if (is_only_uniform_src(init_src)) {
             /* Initial value of induction variable is a uniform */
@@ -1217,12 +1208,12 @@ find_trip_count(loop_info_state *state, unsigned execution_mode)
 
       nir_const_value initial_val = nir_ssa_scalar_as_const_value(initial_s);
 
-      /* We are guaranteed by earlier code that at least one of these sources
+      /* We are not guaranteed by that at one of these sources is a constant.
-       * is a constant but we don't know which.
+       * Try to find one.
        */
       nir_const_value step_val;
       memset(&step_val, 0, sizeof(step_val));
-      UNUSED bool found_step_value = false;
+      bool found_step_value = false;
       assert(nir_op_infos[ind_var->alu->op].num_inputs == 2);
       for (unsigned i = 0; i < 2; i++) {
          nir_ssa_scalar alu_src = nir_ssa_scalar_chase_alu_src(alu_s, i);
@@ -1232,7 +1223,9 @@ find_trip_count(loop_info_state *state, unsigned execution_mode)
             break;
          }
       }
-      assert(found_step_value);
+
+      if (!found_step_value)
+         continue;
 
       int iterations = calculate_iterations(initial_val, step_val, limit_val,
                                             ind_var->alu, cond,

src/compiler/nir/tests/loop_analyze_tests.cpp

@@ -248,15 +248,17 @@ TEST_F(nir_loop_analyze_test, infinite_loop_feq)
    EXPECT_EQ(2, loop->info->num_induction_vars);
    ASSERT_NE((void *)0, loop->info->induction_vars);
 
-   /* Since the initializer and increments are constants, the init_src and
+   /* Since the initializer is a constant, the init_src field will be
-    * update_src fields will be NULL. The def field should not be NULL.
+    * NULL. The def field should not be NULL. The update_src field should
+    * point to a load_const.
     */
    const nir_loop_induction_variable *const ivars = loop->info->induction_vars;
 
    for (unsigned i = 0; i < loop->info->num_induction_vars; i++) {
       EXPECT_NE((void *)0, ivars[i].def);
       EXPECT_EQ((void *)0, ivars[i].init_src);
-      EXPECT_EQ((void *)0, ivars[i].update_src);
+      ASSERT_NE((void *)0, ivars[i].update_src);
+      EXPECT_TRUE(nir_src_is_const(ivars[i].update_src->src));
    }
 }
 
@@ -328,15 +330,17 @@ TEST_F(nir_loop_analyze_test, zero_iterations_ine)
    EXPECT_EQ(2, loop->info->num_induction_vars);
    ASSERT_NE((void *)0, loop->info->induction_vars);
 
-   /* Since the initializer and increments are constants, the init_src and
+   /* Since the initializer is a constant, the init_src field will be
-    * update_src fields will be NULL. The def field should not be NULL.
+    * NULL. The def field should not be NULL. The update_src field should
+    * point to a load_const.
     */
    const nir_loop_induction_variable *const ivars = loop->info->induction_vars;
 
    for (unsigned i = 0; i < loop->info->num_induction_vars; i++) {
       EXPECT_NE((void *)0, ivars[i].def);
       EXPECT_EQ((void *)0, ivars[i].init_src);
-      EXPECT_EQ((void *)0, ivars[i].update_src);
+      ASSERT_NE((void *)0, ivars[i].update_src);
+      EXPECT_TRUE(nir_src_is_const(ivars[i].update_src->src));
    }
 }
 
@@ -405,15 +409,17 @@ TEST_F(nir_loop_analyze_test, one_iteration_uge)
    EXPECT_EQ(2, loop->info->num_induction_vars);
    ASSERT_NE((void *)0, loop->info->induction_vars);
 
-   /* Since the initializer and increments are constants, the init_src and
+   /* Since the initializer is a constant, the init_src field will be
-    * update_src fields will be NULL. The def field should not be NULL.
+    * NULL. The def field should not be NULL. The update_src field should
+    * point to a load_const.
     */
    const nir_loop_induction_variable *const ivars = loop->info->induction_vars;
 
    for (unsigned i = 0; i < loop->info->num_induction_vars; i++) {
       EXPECT_NE((void *)0, ivars[i].def);
       EXPECT_EQ((void *)0, ivars[i].init_src);
-      EXPECT_EQ((void *)0, ivars[i].update_src);
+      ASSERT_NE((void *)0, ivars[i].update_src);
+      EXPECT_TRUE(nir_src_is_const(ivars[i].update_src->src));
    }
 }
 
@@ -482,15 +488,17 @@ TEST_F(nir_loop_analyze_test, one_iteration_ine)
    EXPECT_EQ(2, loop->info->num_induction_vars);
    ASSERT_NE((void *)0, loop->info->induction_vars);
 
-   /* Since the initializer and increments are constants, the init_src and
+   /* Since the initializer is a constant, the init_src field will be
-    * update_src fields will be NULL. The def field should not be NULL.
+    * NULL. The def field should not be NULL. The update_src field should
+    * point to a load_const.
     */
    const nir_loop_induction_variable *const ivars = loop->info->induction_vars;
 
    for (unsigned i = 0; i < loop->info->num_induction_vars; i++) {
       EXPECT_NE((void *)0, ivars[i].def);
       EXPECT_EQ((void *)0, ivars[i].init_src);
-      EXPECT_EQ((void *)0, ivars[i].update_src);
+      ASSERT_NE((void *)0, ivars[i].update_src);
+      EXPECT_TRUE(nir_src_is_const(ivars[i].update_src->src));
    }
 }
 
@@ -559,15 +567,17 @@ TEST_F(nir_loop_analyze_test, one_iteration_ieq)
    EXPECT_EQ(2, loop->info->num_induction_vars);
    ASSERT_NE((void *)0, loop->info->induction_vars);
 
-   /* Since the initializer and increments are constants, the init_src and
+   /* Since the initializer is a constant, the init_src field will be
-    * update_src fields will be NULL. The def field should not be NULL.
+    * NULL. The def field should not be NULL. The update_src field should
+    * point to a load_const.
     */
    const nir_loop_induction_variable *const ivars = loop->info->induction_vars;
 
    for (unsigned i = 0; i < loop->info->num_induction_vars; i++) {
       EXPECT_NE((void *)0, ivars[i].def);
       EXPECT_EQ((void *)0, ivars[i].init_src);
-      EXPECT_EQ((void *)0, ivars[i].update_src);
+      ASSERT_NE((void *)0, ivars[i].update_src);
+      EXPECT_TRUE(nir_src_is_const(ivars[i].update_src->src));
    }
 }
 
@@ -636,15 +646,17 @@ TEST_F(nir_loop_analyze_test, one_iteration_easy_fneu)
    EXPECT_EQ(2, loop->info->num_induction_vars);
    ASSERT_NE((void *)0, loop->info->induction_vars);
 
-   /* Since the initializer and increments are constants, the init_src and
+   /* Since the initializer is a constant, the init_src field will be
-    * update_src fields will be NULL. The def field should not be NULL.
+    * NULL. The def field should not be NULL. The update_src field should
+    * point to a load_const.
     */
    const nir_loop_induction_variable *const ivars = loop->info->induction_vars;
 
    for (unsigned i = 0; i < loop->info->num_induction_vars; i++) {
       EXPECT_NE((void *)0, ivars[i].def);
       EXPECT_EQ((void *)0, ivars[i].init_src);
-      EXPECT_EQ((void *)0, ivars[i].update_src);
+      ASSERT_NE((void *)0, ivars[i].update_src);
+      EXPECT_TRUE(nir_src_is_const(ivars[i].update_src->src));
    }
 }
 
@@ -718,15 +730,17 @@ TEST_F(nir_loop_analyze_test, one_iteration_fneu)
    EXPECT_EQ(2, loop->info->num_induction_vars);
    ASSERT_NE((void *)0, loop->info->induction_vars);
 
-   /* Since the initializer and increments are constants, the init_src and
+   /* Since the initializer is a constant, the init_src field will be
-    * update_src fields will be NULL. The def field should not be NULL.
+    * NULL. The def field should not be NULL. The update_src field should
+    * point to a load_const.
     */
    const nir_loop_induction_variable *const ivars = loop->info->induction_vars;
 
    for (unsigned i = 0; i < loop->info->num_induction_vars; i++) {
       EXPECT_NE((void *)0, ivars[i].def);
       EXPECT_EQ((void *)0, ivars[i].init_src);
-      EXPECT_EQ((void *)0, ivars[i].update_src);
+      ASSERT_NE((void *)0, ivars[i].update_src);
+      EXPECT_TRUE(nir_src_is_const(ivars[i].update_src->src));
    }
 }
 
@@ -797,15 +811,17 @@ TEST_F(nir_loop_analyze_test, zero_iterations_ine_inverted)
    EXPECT_EQ(2, loop->info->num_induction_vars);
    ASSERT_NE((void *)0, loop->info->induction_vars);
 
-   /* Since the initializer and increments are constants, the init_src and
+   /* Since the initializer is a constant, the init_src field will be
-    * update_src fields will be NULL. The def field should not be NULL.
+    * NULL. The def field should not be NULL. The update_src field should
+    * point to a load_const.
     */
    const nir_loop_induction_variable *const ivars = loop->info->induction_vars;
 
    for (unsigned i = 0; i < loop->info->num_induction_vars; i++) {
       EXPECT_NE((void *)0, ivars[i].def);
       EXPECT_EQ((void *)0, ivars[i].init_src);
-      EXPECT_EQ((void *)0, ivars[i].update_src);
+      ASSERT_NE((void *)0, ivars[i].update_src);
+      EXPECT_TRUE(nir_src_is_const(ivars[i].update_src->src));
    }
 }
 
@@ -876,14 +892,16 @@ TEST_F(nir_loop_analyze_test, five_iterations_ige_inverted)
    EXPECT_EQ(2, loop->info->num_induction_vars);
    ASSERT_NE((void *)0, loop->info->induction_vars);
 
-   /* Since the initializer and increments are constants, the init_src and
+   /* Since the initializer is a constant, the init_src field will be
-    * update_src fields will be NULL. The def field should not be NULL.
+    * NULL. The def field should not be NULL. The update_src field should
+    * point to a load_const.
     */
    const nir_loop_induction_variable *const ivars = loop->info->induction_vars;
 
    for (unsigned i = 0; i < loop->info->num_induction_vars; i++) {
       EXPECT_NE((void *)0, ivars[i].def);
       EXPECT_EQ((void *)0, ivars[i].init_src);
-      EXPECT_EQ((void *)0, ivars[i].update_src);
+      ASSERT_NE((void *)0, ivars[i].update_src);
+      EXPECT_TRUE(nir_src_is_const(ivars[i].update_src->src));
    }
 }