nir: Fix num_ssbos when lowering atomic counters

Otherwise it's impossible to know the maximum SSBO index for both
internal TGSI shaders from TTN (which don't have any notion of atomic
counters and no offset) as well as shaders from GLSL.

I fixed everything I could find while grepping for num_ssbos and
num_abos, which hopefully is everything (iris was the only user I could
find that uses it in a meaningful way).

Reviewed-by: Marek Olšák <marek.olsak@amd.com>

src/broadcom/compiler/vir.c

@@ -812,8 +812,7 @@ v3d_nir_lower_fs_early(struct v3d_compile *c)
          * enabling early_fragment_tests even if the user didn't.
          */
         if (!(c->s->info.num_images ||
-              c->s->info.num_ssbos ||
-              c->s->info.num_abos)) {
+              c->s->info.num_ssbos)) {
                 c->s->info.fs.early_fragment_tests = true;
         }
 }

src/compiler/nir/nir_lower_atomics_to_ssbo.c

@@ -239,6 +239,27 @@ nir_lower_atomics_to_ssbo(nir_shader *shader, unsigned ssbo_offset)
             replaced |= (1 << var->data.binding);
          }
       }
+
+      /* Make sure that shader->info.num_ssbos still reflects the maximum SSBO
+       * index that can be used in the shader.
+       */
+      if (shader->info.num_ssbos > 0) {
+         shader->info.num_ssbos += ssbo_offset;
+      } else {
+         /* We can't use num_abos, because it only represents the number of
+          * active atomic counters, and currently unlike SSBO's they aren't
+          * compacted so num_abos actually isn't a bound on the index passed
+          * to nir_intrinsic_atomic_counter_*. e.g. if we have a single atomic
+          * counter declared like:
+          *
+          * layout(binding=1) atomic_uint counter0;
+          *
+          * then when we lower accesses to it the atomic_counter_* intrinsics
+          * will have 1 as the index but num_abos will still be 1.
+          * */
+         shader->info.num_ssbos = util_last_bit(replaced);
+      }
+      shader->info.num_abos = 0;
    }
 
    return progress;

src/gallium/drivers/iris/iris_program.c

@@ -691,10 +691,7 @@ iris_setup_binding_table(const struct gen_device_info *devinfo,
     */
    bt->sizes[IRIS_SURFACE_GROUP_UBO] = num_cbufs + 1;
 
-   /* The first IRIS_MAX_ABOs indices in the SSBO group are for atomics, real
-    * SSBOs start after that.  Compaction will remove unused ABOs.
-    */
-   bt->sizes[IRIS_SURFACE_GROUP_SSBO] = IRIS_MAX_ABOS + info->num_ssbos;
+   bt->sizes[IRIS_SURFACE_GROUP_SSBO] = info->num_ssbos;
 
    for (int i = 0; i < IRIS_SURFACE_GROUP_COUNT; i++)
       assert(bt->sizes[i] <= SURFACE_GROUP_MAX_ELEMENTS);