nir,agx: Move AGX's loop (generalized) to shared NIR code.

When I went to use opt_reassociate for tu, I was advised that you want to
do this loop to get the best results.  If everyone needs it, let's make it
common code and explain what's going on.

In the process, also make it skip work appropriately when there's no
progress.

Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/36342>

src/asahi/compiler/agx_compile.c

@@ -3168,23 +3168,8 @@ agx_optimize_nir(nir_shader *nir, bool soft_fault, uint16_t *preamble_size,
    /* Reassociate before forming preambles because it makes preambles more
     * effective. Clean up after.
     */
-   for (unsigned i = 0; i < 4; ++i) {
-      nir_reassociate_options opts = nir_reassociate_scalar_math;
-      if (i < 2)
-         opts |= nir_reassociate_cse_heuristic;
-
-      NIR_PASS(_, nir, nir_opt_reassociate, opts);
-
-      do {
-         progress = false;
-
-         NIR_PASS(progress, nir, nir_opt_algebraic);
-         NIR_PASS(progress, nir, nir_opt_constant_folding);
-         NIR_PASS(progress, nir, nir_copy_prop);
-         NIR_PASS(progress, nir, nir_opt_cse);
-         NIR_PASS(progress, nir, nir_opt_dce);
-      } while (progress);
-   }
+   nir_opt_reassociate_loop(
+      nir, nir_reassociate_scalar_math | nir_reassociate_cse_heuristic);
 
    /* Lower fmin/fmax before optimizing preambles so we can see across uniform
     * expressions. Do it after nir_opt_reassociate because nir_opt_reassociate

src/compiler/nir/nir.h

@@ -6321,6 +6321,8 @@ typedef enum {
 
 bool nir_opt_reassociate(nir_shader *shader, nir_reassociate_options opts);
 
+bool nir_opt_reassociate_loop(nir_shader *shader, nir_reassociate_options opts);
+
 bool nir_opt_reassociate_bfi(nir_shader *shader);
 
 bool nir_opt_rematerialize_compares(nir_shader *shader);

src/compiler/nir/nir_opt_reassociate.c

@@ -550,3 +550,52 @@ nir_opt_reassociate(nir_shader *nir, nir_reassociate_options opts)
    util_dynarray_fini(&chains);
    return progress;
 }
+
+/* Helper loop for doing reassociation.
+ *
+ * You need to do two passes per set of flags you run with, because in the first
+ * pass we may reassociate constants together, which can then be folded,
+ * resulting in a cleanup, then you need to rebalance again. If you do both
+ * heuristics, you want to do the full CSE path first, then finalize with just
+ * the scalar math once the CSE work has been done.
+ */
+bool
+nir_opt_reassociate_loop(nir_shader *nir, nir_reassociate_options in_opts)
+{
+   bool any_progress = false;
+
+   for (unsigned i = 0; i < 2; ++i) {
+      nir_reassociate_options opts = in_opts;
+      if (i >= 1) {
+         if (in_opts == (nir_reassociate_cse_heuristic | nir_reassociate_scalar_math))
+            opts = nir_reassociate_scalar_math;
+         else
+            break;
+      }
+
+      /* For this set of opt flags, reassociate, then clean up constant folding,
+       * then re-reassociate to rebalance.
+       */
+      for (unsigned j = 0; j < 2; j++) {
+         bool progress = false;
+         NIR_PASS(progress, nir, nir_opt_reassociate, opts);
+         if (!progress)
+            break;
+
+         any_progress = true;
+
+         do {
+            progress = false;
+
+            NIR_PASS(progress, nir, nir_opt_algebraic);
+            NIR_PASS(progress, nir, nir_opt_constant_folding);
+            NIR_PASS(progress, nir, nir_copy_prop);
+            NIR_PASS(progress, nir, nir_opt_cse);
+            NIR_PASS(progress, nir, nir_opt_dce);
+            any_progress |= progress;
+         } while (progress);
+      }
+   }
+
+   return any_progress;
+}