mesa: handle atomic counter lowering for drivers with big ssbo offset aligns

according to the spec, atomic counters can be bound at any offset divisible by 4,
which means that any driver that uses the ssbo lowering pass and doesn't have
a min offset align of 4 is potentially broken

to handle this, use a statevar to inject the misaligned remainder of the offset
into the shader as a uniform. for well-aligned counter binds, the uniform offset
will be 0

Reviewed-by: Marek Olšák <marek.olsak@amd.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/16749>

6 files changed, 58 insertions, 14 deletions

diff --git a/src/compiler/nir/nir.h b/src/compiler/nir/nir.h
index ac62fcadafc..0ffe4390b08 100644
--- a/src/compiler/nir/nir.h
+++ b/src/compiler/nir/nir.h
@@ -5251,7 +5251,7 @@ typedef struct nir_lower_bitmap_options {
 
 void nir_lower_bitmap(nir_shader *shader, const nir_lower_bitmap_options *options);
 
-bool nir_lower_atomics_to_ssbo(nir_shader *shader);
+bool nir_lower_atomics_to_ssbo(nir_shader *shader, unsigned offset_align_state);
 
 typedef enum  {
    nir_lower_int_source_mods = 1 << 0,
diff --git a/src/compiler/nir/nir_lower_atomics_to_ssbo.c b/src/compiler/nir/nir_lower_atomics_to_ssbo.c
index 448f63bdc7c..a7e10c6a063 100644
--- a/src/compiler/nir/nir_lower_atomics_to_ssbo.c
+++ b/src/compiler/nir/nir_lower_atomics_to_ssbo.c
@@ -32,8 +32,29 @@
  * (info.num_ssbos).
  */
 
+static nir_deref_instr *
+deref_offset_var(nir_builder *b, unsigned binding, unsigned offset_align_state)
+{
+   nir_foreach_uniform_variable(var, b->shader) {
+      if (var->num_state_slots != 1)
+         continue;
+      if (var->state_slots[0].tokens[0] == offset_align_state &&
+          var->state_slots[0].tokens[1] == binding)
+         return nir_build_deref_var(b, var);
+   }
+
+   nir_variable *var = nir_variable_create(b->shader, nir_var_uniform, glsl_uint_type(), "offset");
+   var->state_slots = ralloc_array(var, nir_state_slot, 1);
+   var->state_slots[0].tokens[0] = offset_align_state;
+   var->state_slots[0].tokens[1] = binding;
+   var->num_state_slots = 1;
+   var->data.how_declared = nir_var_hidden;
+   b->shader->num_uniforms++;
+   return nir_build_deref_var(b, var);
+}
+
 static bool
-lower_instr(nir_intrinsic_instr *instr, unsigned ssbo_offset, nir_builder *b)
+lower_instr(nir_intrinsic_instr *instr, unsigned ssbo_offset, nir_builder *b, unsigned offset_align_state)
 {
    nir_intrinsic_op op;
 
@@ -84,6 +105,12 @@ lower_instr(nir_intrinsic_instr *instr, unsigned ssbo_offset, nir_builder *b)
 
    nir_ssa_def *buffer = nir_imm_int(b, ssbo_offset + nir_intrinsic_base(instr));
    nir_ssa_def *temp = NULL;
+
+   nir_ssa_def *offset_load = NULL;
+   if (offset_align_state) {
+      nir_deref_instr *deref_offset = deref_offset_var(b, nir_intrinsic_base(instr), offset_align_state);
+      offset_load = nir_load_deref(b, deref_offset);
+   }
    nir_intrinsic_instr *new_instr =
          nir_intrinsic_instr_create(b->shader, op);
 
@@ -123,6 +150,9 @@ lower_instr(nir_intrinsic_instr *instr, unsigned ssbo_offset, nir_builder *b)
       break;
    }
 
+   if (offset_load)
+      new_instr->src[1].ssa = nir_iadd(b, new_instr->src[1].ssa, offset_load);
+
    if (new_instr->intrinsic == nir_intrinsic_load_ssbo) {
       nir_intrinsic_set_align(new_instr, 4, 0);
 
@@ -159,7 +189,7 @@ is_atomic_uint(const struct glsl_type *type)
 }
 
 bool
-nir_lower_atomics_to_ssbo(nir_shader *shader)
+nir_lower_atomics_to_ssbo(nir_shader *shader, unsigned offset_align_state)
 {
    unsigned ssbo_offset = shader->info.num_ssbos;
    bool progress = false;
@@ -172,7 +202,7 @@ nir_lower_atomics_to_ssbo(nir_shader *shader)
             nir_foreach_instr_safe(instr, block) {
                if (instr->type == nir_instr_type_intrinsic)
                   progress |= lower_instr(nir_instr_as_intrinsic(instr),
-                                          ssbo_offset, &builder);
+                                          ssbo_offset, &builder, offset_align_state);
             }
          }
 
diff --git a/src/gallium/drivers/freedreno/ir3/ir3_cmdline.c b/src/gallium/drivers/freedreno/ir3/ir3_cmdline.c
index 8cc3abacf98..263762efdda 100644
--- a/src/gallium/drivers/freedreno/ir3/ir3_cmdline.c
+++ b/src/gallium/drivers/freedreno/ir3/ir3_cmdline.c
@@ -143,7 +143,7 @@ load_glsl(unsigned num_files, char *const *files, gl_shader_stage stage)
    nir_print_shader(nir, stdout);
    NIR_PASS_V(nir, gl_nir_lower_atomics, prog, true);
    NIR_PASS_V(nir, gl_nir_lower_buffers, prog);
-   NIR_PASS_V(nir, nir_lower_atomics_to_ssbo);
+   NIR_PASS_V(nir, nir_lower_atomics_to_ssbo, 0);
    nir_print_shader(nir, stdout);
 
    switch (stage) {
diff --git a/src/gallium/drivers/lima/standalone/lima_compiler_cmdline.c b/src/gallium/drivers/lima/standalone/lima_compiler_cmdline.c
index f7ee352c377..8a37789f817 100644
--- a/src/gallium/drivers/lima/standalone/lima_compiler_cmdline.c
+++ b/src/gallium/drivers/lima/standalone/lima_compiler_cmdline.c
@@ -135,7 +135,7 @@ load_glsl(unsigned num_files, char* const* files, gl_shader_stage stage)
    NIR_PASS_V(nir, nir_lower_var_copies);
    nir_print_shader(nir, stdout);
    NIR_PASS_V(nir, gl_nir_lower_atomics, prog, true);
-   NIR_PASS_V(nir, nir_lower_atomics_to_ssbo);
+   NIR_PASS_V(nir, nir_lower_atomics_to_ssbo, 0);
    nir_print_shader(nir, stdout);
 
    switch (stage) {
diff --git a/src/mesa/state_tracker/st_atom_atomicbuf.c b/src/mesa/state_tracker/st_atom_atomicbuf.c
index eec2f871afe..4bdf897ed4e 100644
--- a/src/mesa/state_tracker/st_atom_atomicbuf.c
+++ b/src/mesa/state_tracker/st_atom_atomicbuf.c
@@ -41,20 +41,23 @@
 
 static void
 st_binding_to_sb(struct gl_buffer_binding *binding,
-                 struct pipe_shader_buffer *sb)
+                 struct pipe_shader_buffer *sb,
+                 unsigned alignment)
 {
    struct gl_buffer_object *st_obj = binding->BufferObject;
 
    if (st_obj && st_obj->buffer) {
+     unsigned offset = 0;
      sb->buffer = st_obj->buffer;
-     sb->buffer_offset = binding->Offset;
-     sb->buffer_size = st_obj->buffer->width0 - binding->Offset;
+     offset = binding->Offset % alignment;
+     sb->buffer_offset = binding->Offset - offset;
+     sb->buffer_size = st_obj->buffer->width0 - sb->buffer_offset;
 
      /* AutomaticSize is FALSE if the buffer was set with BindBufferRange.
       * Take the minimum just to be sure.
       */
      if (!binding->AutomaticSize)
-       sb->buffer_size = MIN2(sb->buffer_size, (unsigned) binding->Size);
+       sb->buffer_size = MIN2(sb->buffer_size, (unsigned) binding->Size + offset);
    } else {
      sb->buffer = NULL;
      sb->buffer_offset = 0;
@@ -82,7 +85,8 @@ st_bind_atomics(struct st_context *st, struct gl_program *prog,
          &prog->sh.data->AtomicBuffers[i];
       struct pipe_shader_buffer sb;
 
-      st_binding_to_sb(&st->ctx->AtomicBufferBindings[atomic->Binding], &sb);
+      st_binding_to_sb(&st->ctx->AtomicBufferBindings[atomic->Binding], &sb,
+                       st->ctx->Const.ShaderStorageBufferOffsetAlignment);
 
       st->pipe->set_shader_buffers(st->pipe, shader_type,
                                    buffer_base + atomic->Binding, 1, &sb, 0x1);
@@ -159,7 +163,7 @@ st_bind_hw_atomic_buffers(struct st_context *st)
       return;
 
    for (i = 0; i < st->ctx->Const.MaxAtomicBufferBindings; i++)
-      st_binding_to_sb(&st->ctx->AtomicBufferBindings[i], &buffers[i]);
+      st_binding_to_sb(&st->ctx->AtomicBufferBindings[i], &buffers[i], 1);
 
    st->pipe->set_hw_atomic_buffers(st->pipe, 0, st->ctx->Const.MaxAtomicBufferBindings, buffers);
 }
diff --git a/src/mesa/state_tracker/st_glsl_to_nir.cpp b/src/mesa/state_tracker/st_glsl_to_nir.cpp
index b12fc3f6c25..77663849023 100644
--- a/src/mesa/state_tracker/st_glsl_to_nir.cpp
+++ b/src/mesa/state_tracker/st_glsl_to_nir.cpp
@@ -526,8 +526,18 @@ st_glsl_to_nir_post_opts(struct st_context *st, struct gl_program *prog,
       nir_var_shader_in | nir_var_shader_out | nir_var_function_temp;
    nir_remove_dead_variables(nir, mask, NULL);
 
-   if (!st->has_hw_atomics && !screen->get_param(screen, PIPE_CAP_NIR_ATOMICS_AS_DEREF))
-      NIR_PASS_V(nir, nir_lower_atomics_to_ssbo);
+   if (!st->has_hw_atomics && !screen->get_param(screen, PIPE_CAP_NIR_ATOMICS_AS_DEREF)) {
+      unsigned align_offset_state = 0;
+      if (st->ctx->Const.ShaderStorageBufferOffsetAlignment > 4) {
+         struct gl_program_parameter_list *params = prog->Parameters;
+         for (unsigned i = 0; i < shader_program->data->NumAtomicBuffers; i++) {
+            gl_state_index16 state[STATE_LENGTH] = { STATE_ATOMIC_COUNTER_OFFSET, (short)shader_program->data->AtomicBuffers[i].Binding };
+            _mesa_add_state_reference(params, state);
+         }
+         align_offset_state = STATE_ATOMIC_COUNTER_OFFSET;
+      }
+      NIR_PASS_V(nir, nir_lower_atomics_to_ssbo, align_offset_state);
+   }
 
    st_set_prog_affected_state_flags(prog);