nir: Add range_base to atomic_counter and an option to use it

Some drivers may encode constant offsets in the instruction, so
make it possible for the drivers to request lowering the atomic
uniform offset into the range_base variable of the intrinsic.

v2: drop patch to use build-in array offset evaluation, it makes
    problems with zink, and update the code accordingly
v3: always initialize range base

Signed-off-by: Gert Wollny <gert.wollny@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/19980>

src/compiler/glsl/gl_nir_lower_atomics.c

@@ -111,7 +111,14 @@ lower_deref_instr(nir_builder *b, nir_intrinsic_instr *instr,
 
    b->cursor = nir_before_instr(&instr->instr);
 
-   nir_ssa_def *offset = nir_imm_int(b, var->data.offset);
+   int offset_value = 0;
+   int range_base = 0;
+   if (!b->shader->options->lower_atomic_offset_to_range_base)
+      offset_value = var->data.offset;
+   else
+      range_base = var->data.offset;
+
+   nir_ssa_def *offset = nir_imm_int(b, offset_value);
    for (nir_deref_instr *d = deref; d->deref_type != nir_deref_type_var;
         d = nir_deref_instr_parent(d)) {
       assert(d->deref_type == nir_deref_type_array);
@@ -130,6 +137,8 @@ lower_deref_instr(nir_builder *b, nir_intrinsic_instr *instr,
     * opcode.
     */
    instr->intrinsic = op;
+   nir_intrinsic_set_range_base(instr, range_base);
+
    nir_instr_rewrite_src(&instr->instr, &instr->src[0],
                          nir_src_for_ssa(offset));
    nir_intrinsic_set_base(instr, idx);

src/compiler/nir/nir.h

@@ -3765,6 +3765,11 @@ typedef struct nir_shader_compiler_options {
     *  of adding it to the image index.
     */
    bool lower_image_offset_to_range_base;
+
+   /** store the variable offset into the instrinsic range_base instead
+    *  of adding it to the atomic source
+    */
+   bool lower_atomic_offset_to_range_base;
 } nir_shader_compiler_options;
 
 typedef struct nir_shader {

src/compiler/nir/nir_intrinsics.py

@@ -587,15 +587,15 @@ intrinsic("rt_trace_ray", src_comp=[-1, 1, 1, 1, 1, 1, 3, 1, 3, 1, -1],
 # undefined.")
 def atomic(name, flags=[]):
     intrinsic(name + "_deref", src_comp=[-1], dest_comp=1, flags=flags)
-    intrinsic(name, src_comp=[1], dest_comp=1, indices=[BASE], flags=flags)
+    intrinsic(name, src_comp=[1], dest_comp=1, indices=[BASE, RANGE_BASE], flags=flags)
 
 def atomic2(name):
     intrinsic(name + "_deref", src_comp=[-1, 1], dest_comp=1)
-    intrinsic(name, src_comp=[1, 1], dest_comp=1, indices=[BASE])
+    intrinsic(name, src_comp=[1, 1], dest_comp=1, indices=[BASE, RANGE_BASE])
 
 def atomic3(name):
     intrinsic(name + "_deref", src_comp=[-1, 1, 1], dest_comp=1)
-    intrinsic(name, src_comp=[1, 1, 1], dest_comp=1, indices=[BASE])
+    intrinsic(name, src_comp=[1, 1, 1], dest_comp=1, indices=[BASE, RANGE_BASE])
 
 atomic("atomic_counter_inc")
 atomic("atomic_counter_pre_dec")

src/compiler/nir/nir_lower_atomics_to_ssbo.c

@@ -153,6 +153,10 @@ lower_instr(nir_intrinsic_instr *instr, unsigned ssbo_offset, nir_builder *b, un
    if (offset_load)
       new_instr->src[1].ssa = nir_iadd(b, new_instr->src[1].ssa, offset_load);
 
+   if (nir_intrinsic_range_base(instr))
+      new_instr->src[1].ssa = nir_iadd(b, new_instr->src[1].ssa,
+                                       nir_imm_int(b, nir_intrinsic_range_base(instr)));
+
    if (new_instr->intrinsic == nir_intrinsic_load_ssbo) {
       nir_intrinsic_set_align(new_instr, 4, 0);