nir/opt_peephole_select: collapse nested IFs if applicable

Single-sided nested IFs can sometimes be collapsed
even if they cannot be flattened.
This optimization re-uses block_check_for_allowed_instrs()
to determine if it is beneficial to collapse the IFs.
Additionally, it is required that the phis of the outer IF
become trivial after this optimization, so that no additional
bcsel instructions are added.
This optimization turns

   if (cond1) {
      <allowed instruction>
      if (cond2) {
         <any code>
      } else {
      }
   } else {
   }

into

   <allowed instruction>
   if (cond1 && cond2) {
      <any code>
   } else {
   }

Totals from 17044 (12.35% of 138013) affected shaders (RAVEN):
SGPRs: 1246416 -> 1246256 (-0.01%); split: -0.01%, +0.00%
VGPRs: 802752 -> 802736 (-0.00%); split: -0.01%, +0.01%
SpillSGPRs: 45857 -> 45850 (-0.02%); split: -0.07%, +0.05%
CodeSize: 85318240 -> 85208592 (-0.13%); split: -0.15%, +0.02%
Instrs: 16769049 -> 16738195 (-0.18%); split: -0.20%, +0.02%
Cycles: 947328732 -> 947145796 (-0.02%); split: -0.03%, +0.01%
VMEM: 7271539 -> 7274090 (+0.04%); split: +0.05%, -0.01%
SMEM: 925983 -> 927374 (+0.15%); split: +0.19%, -0.04%
VClause: 294334 -> 294340 (+0.00%); split: -0.00%, +0.00%
SClause: 633600 -> 634048 (+0.07%); split: -0.01%, +0.08%
Copies: 1589650 -> 1580573 (-0.57%); split: -0.66%, +0.09%
Branches: 540830 -> 525767 (-2.79%); split: -2.79%, +0.00%
PreSGPRs: 902500 -> 902415 (-0.01%); split: -0.02%, +0.01%
PreVGPRs: 759992 -> 760019 (+0.00%); split: -0.00%, +0.01%

Reviewed-by: Rhys Perry <pendingchaos02@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/7478>

diff --git a/src/compiler/nir/nir_opt_peephole_select.c b/src/compiler/nir/nir_opt_peephole_select.c
index 2d7bc04..840d0fe 100644 (file)
--- a/src/compiler/nir/nir_opt_peephole_select.c
+++ b/src/compiler/nir/nir_opt_peephole_select.c
@@ -26,6 +26,7 @@
  */
 
 #include "nir.h"
+#include "nir/nir_builder.h"
 #include "nir_control_flow.h"
 #include "nir_search_helpers.h"
 
@@ -211,6 +212,116 @@ block_check_for_allowed_instrs(nir_block *block, unsigned *count,
    return true;
 }
 
+/**
+ * Try to collapse nested ifs:
+ * This optimization turns
+ *
+ * if (cond1) {
+ *   <allowed instruction>
+ *   if (cond2) {
+ *     <any code>
+ *   } else {
+ *   }
+ * } else {
+ * }
+ *
+ * into
+ *
+ * <allowed instruction>
+ * if (cond1 && cond2) {
+ *   <any code>
+ * } else {
+ * }
+ *
+ */
+static bool
+nir_opt_collapse_if(nir_if *if_stmt, nir_shader *shader, unsigned limit,
+                    bool indirect_load_ok, bool expensive_alu_ok)
+{
+   /* the if has to be nested */
+   if (if_stmt->cf_node.parent->type != nir_cf_node_if)
+      return false;
+
+   /* check if the else block is empty */
+   if (!nir_cf_list_is_empty_block(&if_stmt->else_list))
+      return false;
+
+   /* this opt doesn't make much sense if the branch is empty */
+   if (nir_cf_list_is_empty_block(&if_stmt->then_list))
+      return false;
+
+   /* the nested if has to be the only cf_node:
+    * i.e. <block> <if_stmt> <block> */
+   nir_if *parent_if = nir_cf_node_as_if(if_stmt->cf_node.parent);
+   if (exec_list_length(&parent_if->then_list) != 3)
+      return false;
+
+   /* check if the else block of the parent if is empty */
+   if (!nir_cf_list_is_empty_block(&parent_if->else_list))
+      return false;
+
+   /* check if the block after the nested if is empty except for phis */
+   nir_block *last = nir_if_last_then_block(parent_if);
+   nir_instr *last_instr = nir_block_last_instr(last);
+   if (last_instr && last_instr->type != nir_instr_type_phi)
+      return false;
+
+   /* check if all outer phis become trivial after merging the ifs */
+   nir_foreach_instr(instr, last) {
+      nir_phi_instr *phi = nir_instr_as_phi(instr);
+      nir_phi_src *else_src =
+         nir_phi_get_src_from_block(phi, nir_if_first_else_block(if_stmt));
+
+      nir_foreach_use (src, &phi->dest.ssa) {
+         assert(src->parent_instr->type == nir_instr_type_phi);
+         nir_phi_src *phi_src =
+            nir_phi_get_src_from_block(nir_instr_as_phi(src->parent_instr),
+                                       nir_if_first_else_block(parent_if));
+         if (phi_src->src.ssa != else_src->src.ssa)
+            return false;
+      }
+   }
+
+   /* check if the block before the nested if matches the requirements */
+   nir_block *first = nir_if_first_then_block(parent_if);
+   unsigned count = 0;
+   if (!block_check_for_allowed_instrs(first, &count, limit != 0,
+                                       indirect_load_ok, expensive_alu_ok))
+      return false;
+
+   if (count > limit)
+      return false;
+
+   /* trivialize succeeding phis */
+   nir_foreach_instr(instr, last) {
+      nir_phi_instr *phi = nir_instr_as_phi(instr);
+      nir_foreach_use_safe(src, &phi->dest.ssa) {
+         nir_phi_src *phi_src =
+            nir_phi_get_src_from_block(nir_instr_as_phi(src->parent_instr),
+                                       nir_if_first_else_block(parent_if));
+         nir_instr_rewrite_src(src->parent_instr, &phi_src->src,
+                               nir_src_for_ssa(&phi->dest.ssa));
+      }
+   }
+
+   /* combine the conditions */
+   struct nir_builder b;
+   nir_builder_init(&b, nir_cf_node_get_function(&if_stmt->cf_node)->function->impl);
+   b.cursor = nir_before_cf_node(&if_stmt->cf_node);
+   nir_ssa_def *cond = nir_iand(&b, if_stmt->condition.ssa,
+                                parent_if->condition.ssa);
+   nir_if_rewrite_condition(if_stmt, nir_src_for_ssa(cond));
+
+   /* move the whole inner if before the parent if */
+   nir_cf_list tmp;
+   nir_cf_extract(&tmp, nir_before_block(first),
+                        nir_after_block(last));
+   nir_cf_reinsert(&tmp, nir_before_cf_node(&parent_if->cf_node));
+
+   /* The now empty parent if will be cleaned up by other passes */
+   return true;
+}
+
 static bool
 nir_opt_peephole_select_block(nir_block *block, nir_shader *shader,
                               unsigned limit, bool indirect_load_ok,
@@ -225,6 +336,11 @@ nir_opt_peephole_select_block(nir_block *block, nir_shader *shader,
 
    nir_if *if_stmt = nir_cf_node_as_if(prev_node);
 
+   /* first, try to collapse the if */
+   if (nir_opt_collapse_if(if_stmt, shader, limit,
+                           indirect_load_ok, expensive_alu_ok))
+      return true;
+
    if (if_stmt->control == nir_selection_control_dont_flatten)
       return false;