pan/mdg: Optimize UBO offset calculations

LD_UNIFORM supports constant shifts and biases, just like LD, so take
advantage of that. Will avoid a regression in code quality from lowering
uniforms to UBOs.

No shader-db changes.

Signed-off-by: Alyssa Rosenzweig <alyssa.rosenzweig@collabora.com>
Reviewed-by: Boris Brezillon <boris.brezillon@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/8973>

3 files changed, 32 insertions, 6 deletions

diff --git a/src/panfrost/midgard/compiler.h b/src/panfrost/midgard/compiler.h
index 458d50de6a2..3532c7d7e96 100644
--- a/src/panfrost/midgard/compiler.h
+++ b/src/panfrost/midgard/compiler.h
@@ -538,6 +538,7 @@ void mir_compute_temp_count(compiler_context *ctx);
 #define LDST_SCRATCH 0x2A
 
 void mir_set_offset(compiler_context *ctx, midgard_instruction *ins, nir_src *offset, unsigned seg);
+void mir_set_ubo_offset(midgard_instruction *ins, nir_src *src, unsigned bias);
 
 /* 'Intrinsic' move for aliasing */
 
diff --git a/src/panfrost/midgard/midgard_address.c b/src/panfrost/midgard/midgard_address.c
index 3dbbd72eaca..54e942a3457 100644
--- a/src/panfrost/midgard/midgard_address.c
+++ b/src/panfrost/midgard/midgard_address.c
@@ -199,11 +199,11 @@ mir_match_mov(struct mir_address *address)
 /* Tries to pattern match into mir_address */
 
 static struct mir_address
-mir_match_offset(nir_ssa_def *offset, bool first_free)
+mir_match_offset(nir_ssa_def *offset, bool first_free, bool extend)
 {
         struct mir_address address = {
                 .B = { .def = offset },
-                .type = ITYPE_U64,
+                .type = extend ? ITYPE_U64 : ITYPE_U32,
         };
 
         mir_match_mov(&address);
@@ -211,9 +211,13 @@ mir_match_offset(nir_ssa_def *offset, bool first_free)
         mir_match_mov(&address);
         mir_match_iadd(&address, first_free);
         mir_match_mov(&address);
-        mir_match_u2u64(&address);
-        mir_match_i2i64(&address);
-        mir_match_mov(&address);
+
+        if (extend) {
+                mir_match_u2u64(&address);
+                mir_match_i2i64(&address);
+                mir_match_mov(&address);
+        }
+
         mir_match_ishl(&address);
 
         return address;
@@ -246,7 +250,7 @@ mir_set_offset(compiler_context *ctx, midgard_instruction *ins, nir_src *offset,
 
         bool first_free = (seg == LDST_GLOBAL);
 
-        struct mir_address match = mir_match_offset(offset->ssa, first_free);
+        struct mir_address match = mir_match_offset(offset->ssa, first_free, true);
 
         if (match.A.def) {
                 ins->src[1] = nir_ssa_index(match.A.def);
@@ -272,3 +276,21 @@ mir_set_offset(compiler_context *ctx, midgard_instruction *ins, nir_src *offset,
 
         ins->constants.u32[0] = match.bias;
 }
+
+
+void
+mir_set_ubo_offset(midgard_instruction *ins, nir_src *src, unsigned bias)
+{
+        assert(src->is_ssa);
+        struct mir_address match = mir_match_offset(src->ssa, false, false);
+
+        if (match.B.def) {
+                ins->src[2] = nir_ssa_index(match.B.def);
+
+                for (unsigned i = 0; i < ARRAY_SIZE(ins->swizzle[2]); ++i)
+                        ins->swizzle[2][i] = match.B.comp;
+        }
+
+        ins->load_store.arg_2 |= (match.shift) << 5;
+        ins->constants.u32[0] = match.bias + bias;
+}
diff --git a/src/panfrost/midgard/midgard_compile.c b/src/panfrost/midgard/midgard_compile.c
index c42b1b784f0..f23baf99200 100644
--- a/src/panfrost/midgard/midgard_compile.c
+++ b/src/panfrost/midgard/midgard_compile.c
@@ -1153,6 +1153,9 @@ emit_ubo_read(
                 ins.load_store.arg_2 = 0x1E;
         }
 
+        if (indirect_offset && indirect_offset->is_ssa && !indirect_shift)
+                mir_set_ubo_offset(&ins, indirect_offset, offset);
+
         ins.load_store.arg_1 = index;
 
         return emit_mir_instruction(ctx, ins);