nir/lower_explicit_io: add support for offset_shift

The goal here is to generate addresses that are a right-shifted version
of the actual byte address and record the shift amount in the
offset_shift index. While we could just insert a ushr at the end of
deref chains, this will prevent the shift to be optimized away in many
cases. Instead, we try to extract the shift from the array strides and
struct offsets that make up the deref chain, and only insert a ushr when
absolutely necessary (i.e., for casts). This means we have to walk the
entire deref chain at once for accesses that support offset_shift and we
don't use the standard algorithm of replacing each deref one at a time.

To be able to legally right-shift casts, we use the alignment
information and never shift more than what the alignment could support.
It should also be noted that casts generally have two sources: something
provided by the driver (e.g., a Vulkan resource index) or a variable
pointer coming from a phi/bcsel. For the latter, the entire access chain
consists of multiple parts that are ended by either a phi/bcsel or an
access. Only the part the ends in an access is handled by this new
algorithm; the other parts are handled as usual. This is necessary
because we have no way to encode the offset shift or to even know how
much we would be able to shift without knowing how it is accessed.

This commit adds the general implementation for lowering accesses using
offset_shift and adds a compiler option for drivers to enable it for
SSBO accesses.

Signed-off-by: Job Noorman <jnoorman@igalia.com>
Reviewed-by: Emma Anholt <emma@anholt.net>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/35092>

src/compiler/nir/nir_lower_explicit_io.c

@@ -118,6 +118,26 @@ nir_build_addr_iadd_imm(nir_builder *b, nir_def *addr,
                      addr_get_offset_bit_size(addr, addr_format)));
 }
 
+static nir_def *
+build_addr_ushr_imm(nir_builder *b, nir_def *addr,
+                    nir_address_format addr_format, unsigned shift)
+{
+   switch (addr_format) {
+   case nir_address_format_vec2_index_32bit_offset:
+      assert(addr->num_components == 3);
+      return nir_vector_insert_imm(
+         b, addr, nir_ushr_imm(b, nir_channel(b, addr, 2), shift), 2);
+
+   case nir_address_format_32bit_index_offset:
+      assert(addr->num_components == 2);
+      return nir_vector_insert_imm(
+         b, addr, nir_ushr_imm(b, nir_channel(b, addr, 1), shift), 1);
+
+   default:
+      UNREACHABLE("Unsupported address format");
+   }
+}
+
 static nir_def *
 build_addr_for_var(nir_builder *b, nir_variable *var,
                    nir_address_format addr_format)
@@ -768,6 +788,9 @@ build_explicit_io_load(nir_builder *b, nir_intrinsic_instr *intrin,
       nir_intrinsic_set_range(load, range);
    }
 
+   if (addr_shift.shift)
+      nir_intrinsic_set_offset_shift(load, addr_shift.shift);
+
    load->num_components = num_components;
    nir_def_init(&load->instr, &load->def, num_components, bit_size);
 
@@ -966,6 +989,9 @@ build_explicit_io_store(nir_builder *b, nir_intrinsic_instr *intrin,
 
    nir_intrinsic_set_align(store, align_mul, align_offset);
 
+   if (addr_shift.shift)
+      nir_intrinsic_set_offset_shift(store, addr_shift.shift);
+
    assert(value->num_components == 1 ||
           value->num_components == intrin->num_components);
    store->num_components = value->num_components;
@@ -1081,6 +1107,9 @@ build_explicit_io_atomic(nir_builder *b, nir_intrinsic_instr *intrin,
    if (nir_intrinsic_has_access(atomic))
       nir_intrinsic_set_access(atomic, nir_intrinsic_access(intrin));
 
+   if (addr_shift.shift)
+      nir_intrinsic_set_offset_shift(atomic, addr_shift.shift);
+
    assert(intrin->def.num_components == 1);
    nir_def_init(&atomic->instr, &atomic->def, 1,
                 intrin->def.bit_size);
@@ -1102,9 +1131,14 @@ build_explicit_io_atomic(nir_builder *b, nir_intrinsic_instr *intrin,
    }
 }
 
+/* The resulting address will be right-shifted by `shift` and any bits that
+ * got shifted-out by that will be put in `leftover`. So the final byte
+ * address will be: (ret << shift) + *leftover.
+ */
 static nir_def *
 explicit_io_offset_from_deref(nir_builder *b, nir_deref_instr *deref,
-                              unsigned offset_bit_size)
+                              unsigned offset_bit_size, unsigned shift,
+                              unsigned *leftover)
 {
    switch (deref->deref_type) {
    case nir_deref_type_var:
@@ -1118,6 +1152,35 @@ explicit_io_offset_from_deref(nir_builder *b, nir_deref_instr *deref,
       nir_def *index = deref->arr.index.ssa;
       nir_def *offset;
 
+      unsigned max_stride_shift = ffs(stride) - 1;
+      unsigned stride_shift = MIN2(shift, max_stride_shift);
+
+      if (stride_shift < shift) {
+         /* The stride isn't aligned enough to fully shift right. Try to apply
+          * the leftover shift to the index. We can only do this (without
+          * losing precision) if the index is constant.
+          */
+         assert(nir_src_is_const(deref->arr.index));
+
+         unsigned index_shift = shift - stride_shift;
+         int64_t const_index = nir_src_as_int(deref->arr.index);
+
+         if (!util_is_aligned(const_index, (uintmax_t)1 << index_shift)) {
+            assert(leftover);
+
+            /* The index isn't aligned enough either. Just put the full offset
+             * in `leftover` and return zero.
+             */
+            *leftover = stride * const_index;
+            return nir_imm_intN_t(b, 0, deref->arr.index.ssa->bit_size);
+         }
+
+         index = nir_imm_intN_t(b, const_index >> index_shift,
+                                deref->arr.index.ssa->bit_size);
+      }
+
+      stride >>= stride_shift;
+
       /* If the access chain has been declared in-bounds, then we know it doesn't
        * overflow the type.  For nir_deref_type_array, this implies it cannot be
        * negative. Also, since types in NIR have a maximum 32-bit size, we know the
@@ -1144,6 +1207,18 @@ explicit_io_offset_from_deref(nir_builder *b, nir_deref_instr *deref,
       int offset = glsl_get_struct_field_offset(parent->type,
                                                 deref->strct.index);
       assert(offset >= 0);
+
+      if (!util_is_aligned(offset, (uintmax_t)1 << shift)) {
+         assert(leftover);
+
+         /* The offset isn't aligned enough to fully shift right. Just put it
+          * in leftover and return zero.
+          */
+         *leftover = offset;
+         return nir_imm_intN_t(b, 0, offset_bit_size);
+      }
+
+      offset >>= shift;
       return nir_imm_intN_t(b, offset, offset_bit_size);
    }
 
@@ -1165,12 +1240,114 @@ nir_explicit_io_address_from_deref(nir_builder *b, nir_deref_instr *deref,
    }
 
    nir_def *offset = explicit_io_offset_from_deref(
-      b, deref, addr_get_offset_bit_size(base_addr, addr_format));
+      b, deref, addr_get_offset_bit_size(base_addr, addr_format), 0, NULL);
    return offset ? nir_build_addr_iadd(b, base_addr, addr_format, deref->modes,
                                        offset)
                  : base_addr;
 }
 
+/* Walk the full deref chain and return the resulting address. The resulting
+ * address will be right-shifted by `shift` and any bits that got shifted-out
+ * by that will be put in `leftover`. So the final byte address will be:
+ *    (ret << shift) + *leftover.
+ *
+ * Note on `shift` and `leftover`: we try to shift-right the resulting address
+ * by shifting-right the array strides and struct offsets. As long as `shift`
+ * isn't larger than the alignment, this should generally work without
+ * requiring the `leftover` value (i.e., all intermediary strides and offsets
+ * should be properly aligned). However, `leftover` is necessary in cases like
+ * this:
+ *
+ * struct {
+ *    uint16_t a;
+ *    uint8_t b;
+ *    uint8_t c[3];
+ * } s;
+ * use s.c[1];
+ *
+ * s.c[1] is 2-byte aligned but neither s.c[]'s base offset (3) nor its stride
+ * (1) are. In this case, the returned address will be zero while `leftover`
+ * is set to 4.
+ */
+static nir_def *
+explicit_io_address_from_deref_chain_aux(nir_builder *b, nir_deref_instr *deref,
+                                         nir_address_format addr_format,
+                                         nir_variable_mode modes,
+                                         unsigned shift, unsigned *leftover)
+{
+   unsigned offset_bit_size =
+      addr_get_offset_bit_size(&deref->def, addr_format);
+
+   switch (deref->deref_type) {
+   case nir_deref_type_var:
+      UNREACHABLE("Unsupported deref type");
+      return NULL;
+
+   case nir_deref_type_cast:
+      if (nir_src_is_deref(deref->parent)) {
+         return explicit_io_address_from_deref_chain_aux(
+            b, nir_deref_instr_parent(deref), addr_format, modes, shift,
+            leftover);
+      } else {
+         /* For casts of non-deref instructions, there's nothing we can do
+          * besides simply right-shifting the result. As long as the cast's
+          * alignment is large enough, this should be fine. This is ensured in
+          * build_deref_addr by taking the alignment information into account.
+          * Note that it's correct to use a logical shift as the cast has to
+          * be in-bounds which means its result cannot be negative.
+          */
+         return build_addr_ushr_imm(b, deref->parent.ssa, addr_format, shift);
+      }
+
+   default: {
+      unsigned offset_leftover = 0;
+      nir_def *offset = explicit_io_offset_from_deref(b, deref, offset_bit_size,
+                                                      shift, &offset_leftover);
+      assert(offset);
+
+      unsigned parent_leftover = 0;
+      nir_def *parent_addr = explicit_io_address_from_deref_chain_aux(
+         b, nir_deref_instr_parent(deref), addr_format, modes, shift,
+         &parent_leftover);
+
+      if (offset_leftover || parent_leftover) {
+         assert(leftover && *leftover == 0);
+         *leftover = offset_leftover + parent_leftover;
+      }
+
+      return nir_build_addr_iadd(b, parent_addr, addr_format, modes, offset);
+   }
+   }
+}
+
+static nir_def *
+explicit_io_address_from_deref_chain(nir_builder *b, nir_deref_instr *deref,
+                                     nir_address_format addr_format,
+                                     nir_variable_mode modes, unsigned shift,
+                                     unsigned comp_offset)
+{
+   unsigned leftover = 0;
+   nir_def *addr = explicit_io_address_from_deref_chain_aux(
+      b, deref, addr_format, modes, shift, &leftover);
+
+   unsigned extra_offset = comp_offset + leftover;
+   assert(util_is_aligned(extra_offset, (uintmax_t)1 << shift));
+
+   return nir_build_addr_iadd_imm(b, addr, addr_format, modes,
+                                  extra_offset >> shift);
+}
+
+static unsigned
+get_max_shift(nir_intrinsic_instr *intrin,
+              const nir_shader_compiler_options *options)
+{
+   if (options->max_offset_shift) {
+      return options->max_offset_shift(intrin, options->cb_data);
+   }
+
+   return 0;
+}
+
 static nir_io_offset
 build_addr(nir_builder *b, nir_intrinsic_instr *intrin, nir_def *base_addr,
            nir_address_format addr_format, unsigned comp_offset,
@@ -1180,9 +1357,28 @@ build_addr(nir_builder *b, nir_intrinsic_instr *intrin, nir_def *base_addr,
    assert(deref);
 
    nir_io_offset addr;
-   addr.def = nir_build_addr_iadd_imm(b, base_addr, addr_format, deref->modes,
-                                      comp_offset);
-   addr.shift = 0;
+   unsigned max_shift = get_max_shift(intrin, b->shader->options);
+
+   if (max_shift) {
+      /* In order to calculate a shifted address, we have to walk the full
+       * deref chain. In that case, providing a base_addr that's not the
+       * actual deref probably won't produce the desired result.
+       */
+      assert(base_addr == &deref->def);
+
+      /* Don't try to shift more than the alignment would allow. This ensures
+       * we can just right-shift casts.
+       */
+      unsigned align = nir_combined_align(align_mul, align_offset);
+      addr.shift = MIN2(max_shift, util_logbase2(align));
+      addr.def = explicit_io_address_from_deref_chain(
+         b, deref, addr_format, deref->modes, addr.shift, comp_offset);
+   } else {
+      addr.def = nir_build_addr_iadd_imm(b, base_addr, addr_format,
+                                         deref->modes, comp_offset);
+      addr.shift = 0;
+   }
+
    return addr;
 }
 
@@ -1645,10 +1841,21 @@ nir_lower_explicit_io_impl(nir_function_impl *impl, nir_variable_mode modes,
  * consistent format, pointers can safely be conjured from thin air by the
  * driver, stored to variables, passed through phis, etc.
  *
- * The one exception to the simple algorithm described above is for handling
+ * One exception to the simple algorithm described above is for handling
  * row-major matrices in which case we may look down one additional level of
  * the deref chain.
  *
+ * Another exception is when lowering accesses using offset_shift. The goal
+ * here is to generate addresses that are a right-shifted version of the
+ * actual byte address and record the shift amount in the offset_shift index.
+ * While we could just insert a ushr at the end of deref chains, this will
+ * prevent the shift to be optimized away in many cases. Instead, we try to
+ * extract the shift from the array strides and struct offsets that make up
+ * the deref chain, and only insert a ushr when absolutely necessary (i.e.,
+ * for casts). This means we have to walk the entire deref chain at once for
+ * accesses that support offset_shift and we don't use the algorithm described
+ * above.
+ *
  * This pass is also capable of handling OpenCL generic pointers.  If the
  * address mode is global, it will lower any ambiguous (more than one mode)
  * access to global and pass through the deref_mode_is run-time checks as

src/compiler/nir/nir_shader_compiler_options.h

@@ -840,6 +840,17 @@ typedef struct nir_shader_compiler_options {
     * the next shader according to default_varying_estimate_instr_cost.
     */
    unsigned max_varying_expression_cost;
+
+   /**
+    * Used by nir_lower_explicit_io to determine the maximum offset_shift to
+    * use when lowering the deref address of the given intrinsic.
+    */
+   unsigned (*max_offset_shift)(nir_intrinsic_instr *, const void *);
+
+   /**
+    * Passed to the callbacks that accept a data pointer.
+    */
+   const void *cb_data;
 } nir_shader_compiler_options;
 
 #ifdef __cplusplus