nir/lower_tex: Add lower_index_to_offset

Some backends can handle a constant texture index or a dynamic texture index but
not a constant texture index plus a dynamic texture offset. Add a nir_lower_tex
option to lower to one of these options.

v2: Use more straightforward code proposed by Faith.

Signed-off-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Reviewed-by: Emma Anholt <emma@anholt.net> [v1]
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/21546>

src/compiler/nir/nir.h

@@ -5365,6 +5365,12 @@ typedef struct nir_lower_tex_options {
     */
    bool lower_array_layer_round_even;
 
+   /* If true, texture_index (sampler_index) will be zero if a texture_offset
+    * (sampler_offset) source is present. This is convenient for backends that
+    * support indirect indexing of textures (samplers) but not offsetting it.
+    */
+   bool lower_index_to_offset;
+
    /**
     * Payload data to be sent to callback / filter functions.
     */

src/compiler/nir/nir_lower_tex.c

@@ -1369,6 +1369,40 @@ nir_lower_lod_zero_width(nir_builder *b, nir_tex_instr *tex)
    nir_ssa_def_rewrite_uses_after(&tex->dest.ssa, def, def->parent_instr);
 }
 
+static bool
+lower_index_to_offset(nir_builder *b, nir_tex_instr *tex)
+{
+   bool progress = false;
+   b->cursor = nir_before_instr(&tex->instr);
+
+   for (unsigned i = 0; i < tex->num_srcs; i++) {
+      unsigned *index;
+      switch (tex->src[i].src_type) {
+      case nir_tex_src_texture_offset:
+         index = &tex->texture_index;
+         break;
+      case nir_tex_src_sampler_offset:
+         index = &tex->sampler_index;
+         break;
+      default:
+         continue;
+      }
+
+      /* If there's no base index, there's nothing to lower */
+      if ((*index) == 0)
+         continue;
+
+      assert(tex->src[i].src.is_ssa);
+      nir_ssa_def *sum = nir_iadd_imm(b, tex->src[i].src.ssa, *index);
+      nir_instr_rewrite_src(&tex->instr, &tex->src[i].src,
+                            nir_src_for_ssa(sum));
+      *index = 0;
+      progress = true;
+   }
+
+   return progress;
+}
+
 static bool
 nir_lower_tex_block(nir_block *block, nir_builder *b,
                     const nir_lower_tex_options *options,
@@ -1393,6 +1427,9 @@ nir_lower_tex_block(nir_block *block, nir_builder *b,
       if ((1 << tex->sampler_index) & options->saturate_s)
          sat_mask |= (1 << 0);    /* .x */
 
+      if (options->lower_index_to_offset)
+         progress |= lower_index_to_offset(b, tex);
+
       /* If we are clamping any coords, we must lower projector first
        * as clamping happens *after* projection:
        */