nir/opt_intrinsic: Optimize bcsel(b, shuffle(x, i), shuffle(x, j))

The shuffles provided by the SPV_INTEL_subgroups extension generate

    bcsel(b, shuffle(x, i), shuffle(y, j))

In the case where x and y are the same, we can turn this into a shuffle
with the bcsel on the index which lets us drop a whole shuffle.

Reviewed-by: Caio Marcelo de Oliveira Filho <caio.oliveira@intel.com>
Reviewed-by: Ian Romanick <ian.d.romanick@intel.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/7366>

src/compiler/nir/nir_opt_intrinsics.c

@@ -28,6 +28,89 @@
  * \file nir_opt_intrinsics.c
  */
 
+static bool
+src_is_single_use_shuffle(nir_src src, nir_ssa_def **data, nir_ssa_def **index)
+{
+   nir_intrinsic_instr *shuffle = nir_src_as_intrinsic(src);
+   if (shuffle == NULL || shuffle->intrinsic != nir_intrinsic_shuffle)
+      return false;
+
+   /* This is only called when src is part of an ALU op so requiring no if
+    * uses is reasonable.  If we ever want to use this from an if statement,
+    * we can change it then.
+    */
+   if (!list_is_empty(&shuffle->dest.ssa.if_uses) ||
+       !list_is_singular(&shuffle->dest.ssa.uses))
+      return false;
+
+   assert(shuffle->src[0].is_ssa);
+   assert(shuffle->src[1].is_ssa);
+
+   *data = shuffle->src[0].ssa;
+   *index = shuffle->src[1].ssa;
+
+   return true;
+}
+
+static nir_ssa_def *
+try_opt_bcsel_of_shuffle(nir_builder *b, nir_alu_instr *alu)
+{
+   assert(alu->op == nir_op_bcsel);
+
+   if (!nir_alu_src_is_trivial_ssa(alu, 0))
+      return NULL;
+
+   nir_ssa_def *data1, *index1;
+   if (!nir_alu_src_is_trivial_ssa(alu, 1) ||
+       !src_is_single_use_shuffle(alu->src[1].src, &data1, &index1))
+      return NULL;
+
+   nir_ssa_def *data2, *index2;
+   if (!nir_alu_src_is_trivial_ssa(alu, 2) ||
+       !src_is_single_use_shuffle(alu->src[2].src, &data2, &index2))
+      return NULL;
+
+   if (data1 != data2)
+      return NULL;
+
+   nir_ssa_def *index = nir_bcsel(b, alu->src[0].src.ssa, index1, index2);
+   nir_intrinsic_instr *shuffle =
+      nir_intrinsic_instr_create(b->shader, nir_intrinsic_shuffle);
+   shuffle->src[0] = nir_src_for_ssa(index);
+   shuffle->src[1] = nir_src_for_ssa(data1);
+   shuffle->num_components = alu->dest.dest.ssa.num_components;
+   nir_ssa_dest_init(&shuffle->instr, &shuffle->dest,
+                     alu->dest.dest.ssa.num_components,
+                     alu->dest.dest.ssa.bit_size, NULL);
+   nir_builder_instr_insert(b, &shuffle->instr);
+
+   return &shuffle->dest.ssa;
+}
+
+static bool
+opt_intrinsics_alu(nir_builder *b, nir_alu_instr *alu)
+{
+   nir_ssa_def *replacement = NULL;
+
+   switch (alu->op) {
+   case nir_op_bcsel:
+      replacement = try_opt_bcsel_of_shuffle(b, alu);
+      break;
+
+   default:
+      break;
+   }
+
+   if (replacement) {
+      nir_ssa_def_rewrite_uses(&alu->dest.dest.ssa,
+                               nir_src_for_ssa(replacement));
+      nir_instr_remove(&alu->instr);
+      return true;
+   } else {
+      return false;
+   }
+}
+
 static bool
 opt_intrinsics_intrin(nir_builder *b, nir_intrinsic_instr *intrin,
                       const struct nir_shader_compiler_options *options)
@@ -89,6 +172,11 @@ opt_intrinsics_impl(nir_function_impl *impl,
          b.cursor = nir_before_instr(instr);
 
          switch (instr->type) {
+         case nir_instr_type_alu:
+            if (opt_intrinsics_alu(&b, nir_instr_as_alu(instr)))
+               progress = true;
+            break;
+
          case nir_instr_type_intrinsic:
             if (opt_intrinsics_intrin(&b, nir_instr_as_intrinsic(instr),
                                       options))