WIP: nir: Copy propagate from bcsel source to phi node

TODO: Deal with source modifiers in the bcsel.

Bugzilla: https://bugs.freedesktop.org/show_bug.cgi?id=109216
Fixes: 8fb8ebfbb05 ("intel/compiler: More peephole select")

1 files changed, 148 insertions, 0 deletions

diff --git a/src/compiler/nir/nir_opt_if.c b/src/compiler/nir/nir_opt_if.c
index acd5218f6fa..434ad7ae08a 100644
--- a/src/compiler/nir/nir_opt_if.c
+++ b/src/compiler/nir/nir_opt_if.c
@@ -379,6 +379,153 @@ opt_simpify_bcsel_of_phi(nir_loop *loop)
    return progress;
 }
 
+/**
+ * Attempt to replace a bcsel with a phi node.
+ *
+ * In order for this pass to make progress, several conditions must be met.
+ *
+ * 1. The condition of the bcsel must be a phi node that selects between a
+ *    constant value from the block preceeding the loop and a constant value
+ *    from inside the loop.
+ *
+ * 2. The bcsel source selected by the constant from the loop continue must
+ *    not be a phi node.  Phi nodes can often be removed by
+ *    \c opt_simpify_bcsel_of_phi.
+ *
+ * 3. The result of the bcsel must be used as a source of a phi node in the
+ *    same block.
+ *
+ * 4. The phi node source for the entry block must be the same as the bcsel
+ *    source that would be selected from the entry block.
+ *
+ * If these conditions are met, the use of the bcsel in the phi node source is
+ * replaced with the bcsel source selected by the constant from the loop
+ * continue.
+ *
+ * This allows code such as
+ *
+ *        vec1 32 ssa_51 = load_const (0xffffffff)
+ *        // succs: block_1
+ *        loop {
+ *                block block_1:
+ *                // preds: block_0 block_9
+ *                vec1 32 ssa_56 = phi block_0: ssa_3, block_9: ssa_59
+ *                vec1 32 ssa_57 = phi block_0: ssa_14, block_9: ssa_51
+ *                vec1 32 ssa_58 = iadd ssa_56, ssa_2
+ *                vec1 32 ssa_59 = b32csel ssa_57, ssa_58, ssa_3
+ *
+ * to be converted to
+ *
+ *        vec1 32 ssa_51 = load_const (0xffffffff)
+ *        // succs: block_1
+ *        loop {
+ *                block block_1:
+ *                // preds: block_0 block_9
+ *                vec1 32 ssa_56 = phi block_0: ssa_3, block_9: ssa_58
+ *                vec1 32 ssa_57 = phi block_0: ssa_14, block_9: ssa_51
+ *                vec1 32 ssa_58 = iadd ssa_56, ssa_2
+ */
+static bool
+opt_simpify_phi_of_bcsel(nir_loop *loop)
+{
+   bool progress = false;
+   nir_block *header_block = nir_loop_first_block(loop);
+   nir_block *const prev_block =
+      nir_cf_node_as_block(nir_cf_node_prev(&loop->cf_node));
+
+   /* It would be insane if this were not true */
+   assert(_mesa_set_search(header_block->predecessors, prev_block));
+
+   /* The loop must have exactly one continue block which could be a block
+    * ending in a continue instruction or the "natural" continue from the
+    * last block in the loop back to the top.
+    */
+   if (header_block->predecessors->entries != 2)
+      return false;
+
+   nir_foreach_instr_safe(instr, header_block) {
+      nir_alu_instr *const bcsel =
+         get_bcsel_with_phi_as_condition(instr, header_block);
+      if (bcsel == NULL)
+         continue;
+
+      nir_phi_instr *const cond_phi =
+         nir_instr_as_phi(bcsel->src[0].src.ssa->parent_instr);
+
+      /* We already know we have exactly one continue */
+      assert(exec_list_length(&cond_phi->srcs) == 2);
+
+      uint32_t entry_val = 0, continue_val = 0;
+      if (!phi_has_constant_from_outside_and_one_from_inside_loop(cond_phi,
+                                                                  prev_block,
+                                                                  &entry_val,
+                                                                  &continue_val))
+         continue;
+
+      /* If they both execute or both don't execute, this is a job for
+       * nir_dead_cf, not this pass.
+       */
+      if ((entry_val && continue_val) || (!entry_val && !continue_val))
+         continue;
+
+      const unsigned entry_src = entry_val ? 1 : 2;
+      const unsigned continue_src = entry_val ? 2 : 1;
+
+      nir_foreach_use_safe(bcsel_use_src, &bcsel->dest.dest.ssa) {
+         assert(bcsel_use_src->is_ssa);
+
+         if (bcsel_use_src->parent_instr->type != nir_instr_type_phi)
+            continue;
+
+         if (bcsel_use_src->parent_instr->block != header_block)
+            continue;
+
+         nir_phi_instr *use_phi = nir_instr_as_phi(bcsel_use_src->parent_instr);
+
+         /* We already know we have exactly one continue */
+         assert(exec_list_length(&use_phi->srcs) == 2);
+
+         nir_phi_src *entry_phi_src = NULL;
+         nir_phi_src *continue_phi_src = NULL;
+
+         nir_foreach_phi_src(src, use_phi) {
+            assert(src->src.is_ssa);
+
+            if (src->pred != prev_block) {
+               continue_phi_src = src;
+            } else {
+               assert(src->pred == prev_block);
+
+               entry_phi_src =
+                  src->src.ssa == bcsel->src[entry_src].src.ssa ? src : NULL;
+            }
+         }
+
+         if (entry_phi_src != NULL && continue_phi_src != NULL) {
+            /* Rewrite the phi to use the bcsel source. */
+            nir_instr_rewrite_src(&use_phi->instr,
+                                  &continue_phi_src->src,
+                                  bcsel->src[continue_src].src);
+
+            /* Rewrite the non-phi node users of the bcsel to use the phi. */
+            nir_foreach_use_safe(use_src, &bcsel->dest.dest.ssa) {
+               if (use_src->parent_instr->type == nir_instr_type_phi)
+                  continue;
+
+               nir_instr_rewrite_src(use_src->parent_instr,
+                                     use_src,
+                                     nir_src_for_ssa(&use_phi->dest.ssa));
+            }
+
+            progress = true;
+            break;
+         }
+      }
+   }
+
+   return progress;
+}
+
 static bool
 is_block_empty(nir_block *block)
 {
@@ -947,6 +1094,7 @@ opt_if_cf_list(nir_builder *b, struct exec_list *cf_list)
          nir_loop *loop = nir_cf_node_as_loop(cf_node);
          progress |= opt_if_cf_list(b, &loop->body);
          progress |= opt_simpify_bcsel_of_phi(loop);
+         progress |= opt_simpify_phi_of_bcsel(loop);
          progress |= opt_peel_loop_initial_if(loop);
          progress |= opt_if_loop_last_continue(loop);
          break;