nir/lower_subgroups: add filter

this will be useful for AGX, which has many reductions (but not all) in
hardware with the logic too backend-specific to encode with bitflags.

Signed-off-by: Alyssa Rosenzweig <alyssa@rosenzweig.io>
Reviewed-by: Daniel Schürmann <daniel@schuermann.dev>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/28993>

src/compiler/nir/nir.h

@@ -5796,6 +5796,12 @@ bool nir_lower_is_helper_invocation(nir_shader *shader);
 bool nir_lower_single_sampled(nir_shader *shader);
 
 typedef struct nir_lower_subgroups_options {
+   /* In addition to the boolean lowering options below, this optional callback
+    * will filter instructions for lowering if non-NULL. The data passed will be
+    * this options struct itself.
+    */
+   nir_instr_filter_cb filter;
+
    uint8_t subgroup_size;
    uint8_t ballot_bit_size;
    uint8_t ballot_components;

src/compiler/nir/nir_lower_subgroups.c

@@ -1090,8 +1090,8 @@ bool
 nir_lower_subgroups(nir_shader *shader,
                     const nir_lower_subgroups_options *options)
 {
-   return nir_shader_lower_instructions(shader,
-                                        lower_subgroups_filter,
+   void *filter = options->filter ? options->filter : lower_subgroups_filter;
+   return nir_shader_lower_instructions(shader, filter,
                                         lower_subgroups_instr,
                                         (void *)options);
 }