From 1a2d0d3f314f5efb7a7b2957c5226a8ba9f6b3db Mon Sep 17 00:00:00 2001
From: Emma Anholt <emma@anholt.net>
Date: Fri, 21 Nov 2025 16:52:40 -0800
Subject: [PATCH] nir: Optimistically unroll loops using induction var as a
 sample id.

On the assumption that nobody will use a sample id greater than the sample
count, have loop unrolling guess based on the driver's max sample count.
This unrolls a simple resolve shader with a uniform max samples on ir3 to:

  value = vec4(0);
  if (max_samples > 0) {
    value += txf_ms(coord, 0);
    if (max_samples > 1 {
      value += txf_ms(coord, 1);
      if (max_samples > 2){
        value += txf_ms(coord, 2);
        if (max_samples > 3) {
          value += txf_ms(coord, 3);
          for (i = 4; i < max_samples; i++)
            value += txf_ms(coord, i);
          }
        }
     }
   }
   ...

This is only worth a 1% win on our microbenchmark as-is, but if we could
flatten those ifs out and pull the fadds out to the end, avoiding syncs
per load would be a big win.  This seems like a first step.

I've taken a shot at updating drivers to set the value, and tried to leave
notes in places that drivers might update, and want to follow up with
updating the compiler option.

This affects over half the DX11 apps in shader-db-private.

Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/38585>
---
 src/compiler/nir/nir_loop_analyze.c | 83 ++++++++++++++++++-----------
 1 file changed, 53 insertions(+), 30 deletions(-)

diff --git a/src/compiler/nir/nir_loop_analyze.c b/src/compiler/nir/nir_loop_analyze.c
index 3f78a524c8d..6396e7c642e 100644
--- a/src/compiler/nir/nir_loop_analyze.c
+++ b/src/compiler/nir/nir_loop_analyze.c
@@ -418,46 +418,69 @@ find_array_access_via_induction(loop_info_state *state,
    return 0;
 }
 
+static void
+guess_loop_limit_by_intrinsic(loop_info_state *state, nir_intrinsic_instr *intrin, unsigned *min_loop_count)
+{
+   /* 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_induction_variable *array_idx = NULL;
+      unsigned array_size =
+         find_array_access_via_induction(state,
+                                         nir_src_as_deref(intrin->src[0]),
+                                         &array_idx);
+      if (array_idx)
+         *min_loop_count = MIN2(*min_loop_count, array_size);
+
+      if (intrin->intrinsic == nir_intrinsic_copy_deref) {
+         array_size =
+            find_array_access_via_induction(state,
+                                            nir_src_as_deref(intrin->src[1]),
+                                            &array_idx);
+         if (array_idx)
+            *min_loop_count = MIN2(*min_loop_count, array_size);
+      }
+   }
+}
+
+static void
+guess_loop_limit_by_tex(loop_info_state *state, nir_tex_instr *tex, unsigned *min_loop_count)
+{
+   nir_def *sample = nir_get_tex_src(tex, nir_tex_src_ms_index);
+   if (sample) {
+      nir_loop_induction_variable *var = get_loop_var(sample, state);
+      if (var) {
+         const nir_function_impl *impl = nir_cf_node_get_function(&state->loop->cf_node);
+         const nir_shader_compiler_options *options = impl->function->shader->options;
+         *min_loop_count = MIN2(*min_loop_count, options->max_samples);
+      }
+   }
+}
+
 static unsigned
 guess_loop_limit(loop_info_state *state)
 {
-   unsigned min_array_size = UINT_MAX;
+   unsigned min_loop_count = UINT_MAX;
 
    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_induction_variable *array_idx = NULL;
-            unsigned array_size =
-               find_array_access_via_induction(state,
-                                               nir_src_as_deref(intrin->src[0]),
-                                               &array_idx);
-            if (array_idx)
-               min_array_size = MIN2(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 (array_idx)
-               min_array_size = MIN2(min_array_size, array_size);
+         switch (instr->type) {
+         case nir_instr_type_intrinsic:
+            guess_loop_limit_by_intrinsic(state, nir_instr_as_intrinsic(instr), &min_loop_count);
+            break;
+         case nir_instr_type_tex:
+            guess_loop_limit_by_tex(state, nir_instr_as_tex(instr), &min_loop_count);
+            break;
+         default:
+            break;
          }
       }
    }
 
-   if (min_array_size != UINT_MAX)
-      return min_array_size;
+   if (min_loop_count != UINT_MAX)
+      return min_loop_count;
    else
       return 0;
 }