agx/lower_address: Handle large shifts

If we manage to fold in a left shift that's bigger than the hardware can do, we
should at least avoid generating a useless right shift to feed the hardware
rather bailing completely.

For motivation, this form of address arithmetic is encountered when indexing
into arrays with large power-of-two element sizes (array-of-structs).

Signed-off-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/21643>

diff --git a/src/asahi/compiler/agx_nir_lower_address.c b/src/asahi/compiler/agx_nir_lower_address.c
index 1791164..7622c50 100644 (file)
--- a/src/asahi/compiler/agx_nir_lower_address.c
+++ b/src/asahi/compiler/agx_nir_lower_address.c
@@ -143,6 +143,17 @@ match_address(nir_builder *b, nir_ssa_scalar base, int8_t format_shift)
          if (new_shift <= 2) {
             match.offset = shifted;
             match.shift = new_shift;
+         } else if (new_shift > 0) {
+            /* For large shifts, we do need an ishl instruction but we can
+             * shrink the shift to avoid generating an ishr.
+             */
+            assert(new_shift >= 3);
+
+            nir_ssa_def *rewrite =
+               nir_ishl_imm(b, nir_vec_scalars(b, &shifted, 1), new_shift);
+
+            match.offset = nir_get_ssa_scalar(rewrite, 0);
+            match.shift = 0;
          }
       }
    } else {