nir: pass a callback to nir_lower_robust_access

rather than try to enumerate everything a driver might want with an unmanageable
collection of booleans, just do a filter callback + data. this ends up simpler
overall, and will allow Intel to use this pass for just 64-bit images without
needing to add even more booleans.

while we're churning the pass signature, also do a quick port to
nir_shader_intrinsics_pass

Signed-off-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Reviewed-by: Lionel Landwerlin <lionel.g.landwerlin@intel.com>
Reviewed-by: Iago Toral Quiroga <itoral@igalia.com> [NIR and V3D]
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/32907>

src/asahi/vulkan/hk_shader.c

@@ -578,6 +578,17 @@ lower_subpass_dim(nir_builder *b, nir_instr *instr, UNUSED void *_data)
    return true;
 }
 
+static bool
+should_lower_robust(const nir_intrinsic_instr *intr, const void *_)
+{
+   /* The hardware is robust, but our software image atomics are not. Unlike the
+    * GL driver, we don't use the common buffer image lowering, using the
+    * agx_nir_lower_texture lowering for robustImageAccess2 semantics.
+    */
+   return intr->intrinsic == nir_intrinsic_image_deref_atomic ||
+          intr->intrinsic == nir_intrinsic_image_deref_atomic_swap;
+}
+
 void
 hk_lower_nir(struct hk_device *dev, nir_shader *nir,
              const struct vk_pipeline_robustness_state *rs, bool is_multiview,
@@ -633,16 +644,7 @@ hk_lower_nir(struct hk_device *dev, nir_shader *nir,
    /* Turn cache flushes into image coherency bits while we still have derefs */
    NIR_PASS(_, nir, nir_lower_memory_model);
 
-   /* Images accessed through the texture or PBE hardware are robust, so we
-    * don't set lower_image. (There are some sticky details around txf but
-    * they're handled by agx_nir_lower_texture). However, image atomics are
-    * software so require robustness lowering.
-    */
-   nir_lower_robust_access_options robustness = {
-      .lower_image_atomic = true,
-   };
-
-   NIR_PASS(_, nir, nir_lower_robust_access, &robustness);
+   NIR_PASS(_, nir, nir_lower_robust_access, should_lower_robust, NULL);
 
    /* We must do early lowering before hk_nir_lower_descriptors, since this will
     * create lod_bias_agx instructions.

src/broadcom/compiler/vir.c

@@ -1673,6 +1673,32 @@ v3d_nir_lower_subgroup_intrinsics(nir_shader *s, struct v3d_compile *c)
         return progress;
 }
 
+static bool
+should_lower_robustness(const nir_intrinsic_instr *intr, const void *data)
+{
+        const struct v3d_key *key = data;
+
+        switch (intr->intrinsic) {
+        case nir_intrinsic_load_ubo:
+                return key->robust_uniform_access;
+
+        case nir_intrinsic_load_ssbo:
+        case nir_intrinsic_store_ssbo:
+        case nir_intrinsic_ssbo_atomic:
+        case nir_intrinsic_ssbo_atomic_swap:
+                return key->robust_storage_access;
+
+        case nir_intrinsic_image_load:
+        case nir_intrinsic_image_store:
+        case nir_intrinsic_image_atomic:
+        case nir_intrinsic_image_atomic_swap:
+                return key->robust_image_access;
+
+        default:
+                return false;
+        }
+}
+
 static void
 v3d_attempt_compile(struct v3d_compile *c)
 {
@@ -1744,13 +1770,8 @@ v3d_attempt_compile(struct v3d_compile *c)
                 NIR_PASS(_, c->s, nir_copy_prop);
                 NIR_PASS(_, c->s, nir_opt_constant_folding);
 
-                nir_lower_robust_access_options opts = {
-                   .lower_image = c->key->robust_image_access,
-                   .lower_ssbo = c->key->robust_storage_access,
-                   .lower_ubo = c->key->robust_uniform_access,
-                };
-
-                NIR_PASS(_, c->s, nir_lower_robust_access, &opts);
+                NIR_PASS(_, c->s, nir_lower_robust_access,
+                         should_lower_robustness, c->key);
         }
 
         NIR_PASS(_, c->s, nir_lower_vars_to_scratch,

src/compiler/nir/nir.h

@@ -3843,6 +3843,9 @@ typedef enum {
  */
 typedef bool (*nir_instr_filter_cb)(const nir_instr *, const void *);
 
+/* Like nir_instr_filter_cb but specialized to intrinsics */
+typedef bool (*nir_intrin_filter_cb)(const nir_intrinsic_instr *, const void *);
+
 /** A vectorization width callback
  *
  * Returns the maximum vectorization width per instruction.
@@ -5958,42 +5961,8 @@ typedef struct {
 bool nir_lower_mem_access_bit_sizes(nir_shader *shader,
                                     const nir_lower_mem_access_bit_sizes_options *options);
 
-typedef struct {
-   /* Lower load_ubo to be robust. Out-of-bounds loads will return UNDEFINED
-    * values (not necessarily zero).
-    */
-   bool lower_ubo;
-
-   /* Lower load_ssbo/store_ssbo/ssbo_atomic(_swap) to be robust. Out-of-bounds
-    * loads and atomics will return UNDEFINED values (not necessarily zero).
-    * Out-of-bounds stores and atomics CORRUPT the contents of the SSBO.
-    *
-    * This suffices for robustBufferAccess but not robustBufferAccess2.
-    */
-   bool lower_ssbo;
-
-   /* Lower all image_load/image_store/image_atomic(_swap) instructions to be
-    * robust.  Out-of-bounds loads will return ZERO.
-    *
-    * This suffices for robustImageAccess but not robustImageAccess2.
-    */
-   bool lower_image;
-
-   /* Lower all buffer image instructions as above. Implied by lower_image. */
-   bool lower_buffer_image;
-
-   /* Lower image_atomic(_swap) for all dimensions. Implied by lower_image. */
-   bool lower_image_atomic;
-
-   /* Vulkan's robustBufferAccess feature is only concerned with buffers that
-    * are bound through descriptor sets, so shared memory is not included, but
-    * it may be useful to enable this for debugging.
-    */
-   bool lower_shared;
-} nir_lower_robust_access_options;
-
 bool nir_lower_robust_access(nir_shader *s,
-                             const nir_lower_robust_access_options *opts);
+                             nir_intrin_filter_cb filter, const void *data);
 
 /* clang-format off */
 typedef bool (*nir_should_vectorize_mem_func)(unsigned align_mul,

src/compiler/nir/nir_lower_robust_access.c

@@ -59,9 +59,7 @@ wrap_in_if(nir_builder *b, nir_intrinsic_instr *instr, nir_def *valid)
 }
 
 static void
-lower_buffer_load(nir_builder *b,
-                  nir_intrinsic_instr *instr,
-                  const nir_lower_robust_access_options *opts)
+lower_buffer_load(nir_builder *b, nir_intrinsic_instr *instr)
 {
    uint32_t type_sz = instr->def.bit_size / 8;
    nir_def *size;
@@ -110,21 +108,10 @@ lower_buffer_shared(nir_builder *b, nir_intrinsic_instr *instr)
                   nir_imm_int(b, b->shader->info.shared_size));
 }
 
-static bool
-lower_image(nir_builder *b,
-            nir_intrinsic_instr *instr,
-            const nir_lower_robust_access_options *opts, bool deref)
+static void
+lower_image(nir_builder *b, nir_intrinsic_instr *instr, bool deref)
 {
    enum glsl_sampler_dim dim = nir_intrinsic_image_dim(instr);
-   bool atomic = (instr->intrinsic == nir_intrinsic_image_atomic ||
-                  instr->intrinsic == nir_intrinsic_image_atomic_swap ||
-                  instr->intrinsic == nir_intrinsic_image_deref_atomic ||
-                  instr->intrinsic == nir_intrinsic_image_deref_atomic_swap);
-   if (!opts->lower_image &&
-       !(opts->lower_buffer_image && dim == GLSL_SAMPLER_DIM_BUF) &&
-       !(opts->lower_image_atomic && atomic))
-      return false;
-
    uint32_t num_coords = nir_image_intrinsic_coord_components(instr);
    bool is_array = nir_intrinsic_image_array(instr);
    nir_def *coord = instr->src[1].ssa;
@@ -165,77 +152,73 @@ lower_image(nir_builder *b,
 
    /* Only execute if coordinates are in-bounds. Otherwise, return zero. */
    wrap_in_if(b, instr, in_bounds);
-   return true;
 }
 
+struct pass_opts {
+   nir_intrin_filter_cb filter;
+   const void *data;
+};
+
 static bool
-lower(nir_builder *b, nir_instr *instr, void *_opts)
+lower(nir_builder *b, nir_intrinsic_instr *intr, void *_opts)
 {
-   const nir_lower_robust_access_options *opts = _opts;
-   if (instr->type != nir_instr_type_intrinsic)
+   const struct pass_opts *opts = _opts;
+   if (!opts->filter(intr, opts->data))
       return false;
 
-   nir_intrinsic_instr *intr = nir_instr_as_intrinsic(instr);
-   b->cursor = nir_before_instr(instr);
+   b->cursor = nir_before_instr(&intr->instr);
 
    switch (intr->intrinsic) {
    case nir_intrinsic_image_load:
    case nir_intrinsic_image_store:
    case nir_intrinsic_image_atomic:
    case nir_intrinsic_image_atomic_swap:
-      return lower_image(b, intr, opts, false);
+      lower_image(b, intr, false);
+      return true;
 
    case nir_intrinsic_image_deref_load:
    case nir_intrinsic_image_deref_store:
    case nir_intrinsic_image_deref_atomic:
    case nir_intrinsic_image_deref_atomic_swap:
-      return lower_image(b, intr, opts, true);
+      lower_image(b, intr, true);
+      return true;
 
    case nir_intrinsic_load_ubo:
-      if (opts->lower_ubo) {
-         lower_buffer_load(b, intr, opts);
-         return true;
-      }
-      return false;
-
    case nir_intrinsic_load_ssbo:
-      if (opts->lower_ssbo) {
-         lower_buffer_load(b, intr, opts);
-         return true;
-      }
-      return false;
+      lower_buffer_load(b, intr);
+      return true;
    case nir_intrinsic_store_ssbo:
-      if (opts->lower_ssbo) {
-         lower_buffer_store(b, intr);
-         return true;
-      }
-      return false;
+      lower_buffer_store(b, intr);
+      return true;
    case nir_intrinsic_ssbo_atomic:
-      if (opts->lower_ssbo) {
-         lower_buffer_atomic(b, intr);
-         return true;
-      }
-      return false;
+      lower_buffer_atomic(b, intr);
+      return true;
 
    case nir_intrinsic_store_shared:
    case nir_intrinsic_load_shared:
    case nir_intrinsic_shared_atomic:
    case nir_intrinsic_shared_atomic_swap:
-      if (opts->lower_shared) {
-         lower_buffer_shared(b, intr);
-         return true;
-      }
-      return false;
+      /* Vulkan's robustBufferAccess feature is only concerned with buffers that
+       * are bound through descriptor sets, so shared memory is not included,
+       * but this lowering may be useful for debugging.
+       */
+      lower_buffer_shared(b, intr);
+      return true;
 
    default:
-      return false;
+      unreachable("driver requested lowering for unsupported intrinsic");
    }
 }
 
+/*
+ * Buffer/image robustness lowering with robustBufferAccess/robustImageAccess
+ * semantics. This is sufficient for GL, but not for D3D. However, Vulkan
+ * drivers get buffer robustness lowered via nir_lower_explicit_io.
+ */
 bool
-nir_lower_robust_access(nir_shader *s,
-                        const nir_lower_robust_access_options *opts)
+nir_lower_robust_access(nir_shader *s, nir_intrin_filter_cb filter,
+                        const void *data)
 {
-   return nir_shader_instructions_pass(s, lower, nir_metadata_control_flow,
-                                       (void *)opts);
+   struct pass_opts opt = { .filter = filter, .data = data };
+   return nir_shader_intrinsics_pass(s, lower, nir_metadata_control_flow, &opt);
 }

src/gallium/drivers/asahi/agx_state.c

@@ -1788,6 +1788,37 @@ glsl_type_size(const struct glsl_type *type, bool bindless)
    return glsl_count_attribute_slots(type, false);
 }
 
+static bool
+should_lower_robustness(const nir_intrinsic_instr *intr, const void *data)
+{
+   const bool *gl_robust = data;
+
+   switch (intr->intrinsic) {
+   /* The texture/PBE hardware is robust, but our buffer image implementation
+    * is not. Lower robustness only for buffer images.
+    */
+   case nir_intrinsic_image_load:
+   case nir_intrinsic_image_store:
+      return nir_intrinsic_image_dim(intr) == GLSL_SAMPLER_DIM_BUF;
+
+   /* Image atomics are lowered to raw memory access */
+   case nir_intrinsic_image_atomic:
+   case nir_intrinsic_image_atomic_swap:
+      return true;
+
+   /* UBOs/SSBOs are lowered to raw pointers */
+   case nir_intrinsic_load_ubo:
+   case nir_intrinsic_load_ssbo:
+   case nir_intrinsic_store_ssbo:
+   case nir_intrinsic_ssbo_atomic:
+   case nir_intrinsic_ssbo_atomic_swap:
+      return *gl_robust;
+
+   default:
+      return false;
+   }
+}
+
 static void
 agx_shader_initialize(struct agx_device *dev, struct agx_uncompiled_shader *so,
                       nir_shader *nir, bool support_lod_bias, bool robust)
@@ -1798,24 +1829,10 @@ agx_shader_initialize(struct agx_device *dev, struct agx_uncompiled_shader *so,
    blob_init(&so->early_serialized_nir);
    nir_serialize(&so->early_serialized_nir, nir, true);
 
-   nir_lower_robust_access_options robustness = {
-      /* Images accessed through the texture or PBE hardware are robust, so we
-       * don't set lower_image. However, buffer images and image atomics are
-       * lowered so require robustness lowering.
-       */
-      .lower_buffer_image = true,
-      .lower_image_atomic = true,
-
-      /* Buffer access is based on raw pointers and hence needs lowering to be
-         robust */
-      .lower_ubo = robust,
-      .lower_ssbo = robust,
-   };
-
    /* We need to lower robustness before bindings, since robustness lowering
     * affects the bindings used.
     */
-   NIR_PASS(_, nir, nir_lower_robust_access, &robustness);
+   NIR_PASS(_, nir, nir_lower_robust_access, should_lower_robustness, &robust);
 
    /* Similarly, we need to do early texture lowering before bindings */
    NIR_PASS(_, nir, agx_nir_lower_texture_early, support_lod_bias);