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>

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,

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);
             }
          }
 

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) {

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) {

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);
 }

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);