nir/opt_if: also merge break statements with ones after the branch

This optimizations turns

     loop {
        ...
        if (cond1) {
           if (cond2) {
              do_work_1();
              break;
           } else {
              do_work_2();
           }
           do_work_3();
           break;
        } else {
           ...
        }
     }

 into:

     loop {
        ...
        if (cond1) {
           if (cond2) {
              do_work_1();
           } else {
              do_work_2();
              do_work_3();
           }
           break;
        } else {
           ...
        }
     }

As this optimizations moves code into the NIF statement,
it re-iterates on the branch legs in case of success.

Reviewed-by: Emma Anholt <emma@anholt.net>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/7587>

src/compiler/nir/nir_opt_if.c

@@ -157,7 +157,8 @@ opt_peel_loop_initial_if(nir_loop *loop)
       return false;
 
    nir_if *nif = nir_cf_node_as_if(if_node);
-   assert(nif->condition.is_ssa);
+   if (!nif->condition.is_ssa)
+      return false;
 
    nir_ssa_def *cond = nif->condition.ssa;
    if (cond->parent_instr->type != nir_instr_type_phi)
@@ -932,7 +933,9 @@ opt_if_simplification(nir_builder *b, nir_if *nif)
 
 /**
  * This optimization tries to merge two break statements into a single break.
- * For this purpose, it checks if both branch legs end in a break.
+ * For this purpose, it checks if both branch legs end in a break or
+ * if one branch leg ends in a break, and the other one does so after the
+ * branch.
  *
  * This optimization turns
  *
@@ -959,6 +962,40 @@ opt_if_simplification(nir_builder *b, nir_if *nif)
  *        break;
  *     }
  *
+ * but also situations like
+ *
+ *     loop {
+ *        ...
+ *        if (cond1) {
+ *           if (cond2) {
+ *              do_work_1();
+ *              break;
+ *           } else {
+ *              do_work_2();
+ *           }
+ *           do_work_3();
+ *           break;
+ *        } else {
+ *           ...
+ *        }
+ *     }
+ *
+ *  into:
+ *
+ *     loop {
+ *        ...
+ *        if (cond1) {
+ *           if (cond2) {
+ *              do_work_1();
+ *           } else {
+ *              do_work_2();
+ *              do_work_3();
+ *           }
+ *           break;
+ *        } else {
+ *           ...
+ *        }
+ *     }
  */
 static bool
 opt_merge_breaks(nir_if *nif)
@@ -985,6 +1022,46 @@ opt_merge_breaks(nir_if *nif)
       return true;
     }
 
+   /* Single break: If there's a break after the branch and the non-breaking
+    * side of the if falls through to it, then hoist that code after up into
+    * the if and leave just a single break there.
+    */
+   if (then_break || else_break) {
+
+      /* At least one branch leg must fall-through */
+      if (nir_block_ends_in_jump(last_then) && nir_block_ends_in_jump(last_else))
+         return false;
+
+      /* Check if there is a single break after the IF */
+      nir_cf_node *first = nir_cf_node_next(&nif->cf_node);
+      nir_cf_node *last = first;
+      while (!nir_cf_node_is_last(last)) {
+         if (contains_other_jump (last, NULL))
+            return false;
+         last = nir_cf_node_next(last);
+      }
+
+      assert(last->type == nir_cf_node_block);
+      if (!nir_block_ends_in_break(nir_cf_node_as_block(last)))
+         return false;
+
+      /* Hoist the code from after the IF into the falling-through branch leg */
+      nir_opt_remove_phis_block(nir_cf_node_as_block(first));
+      nir_block *break_block = then_break ? last_then : last_else;
+      nir_lower_phis_to_regs_block(break_block->successors[0]);
+
+      nir_cf_list tmp;
+      nir_cf_extract(&tmp, nir_before_cf_node(first),
+                           nir_after_block_before_jump(nir_cf_node_as_block(last)));
+      if (then_break)
+         nir_cf_reinsert(&tmp, nir_after_block(last_else));
+      else
+         nir_cf_reinsert(&tmp, nir_after_block(last_then));
+
+      nir_instr_remove_v(nir_block_last_instr(break_block));
+      return true;
+   }
+
    return false;
 }
 
@@ -1518,7 +1595,13 @@ opt_if_regs_cf_list(struct exec_list *cf_list)
          nir_if *nif = nir_cf_node_as_if(cf_node);
          progress |= opt_if_regs_cf_list(&nif->then_list);
          progress |= opt_if_regs_cf_list(&nif->else_list);
-         progress |= opt_merge_breaks(nif);
+         if (opt_merge_breaks(nif)) {
+            /* This optimization might move blocks
+             * from after the NIF into the NIF */
+            progress = true;
+            opt_if_regs_cf_list(&nif->then_list);
+            opt_if_regs_cf_list(&nif->else_list);
+         }
          break;
       }
 

src/gallium/drivers/d3d12/ci/d3d12-quick_shader.txt

@@ -7658,12 +7658,9 @@ spec/ext_shader_image_load_formatted/execution/image_checkerboard: skip
 spec/glsl-1.10/execution/fs-dfdx-accuracy: warn
 spec/glsl-1.10/execution/fs-dfdy-accuracy: warn
 spec/glsl-1.10/execution/fs-discard-deep-branch: warn
-spec/glsl-1.10/execution/fs-nested-return-in-loop-nested_in_if: fail
 spec/glsl-1.10/execution/glsl-fs-uniform-sampler-array: fail
 spec/glsl-1.10/execution/samplers/glsl-fs-sampler-numbering-3: fail
 spec/glsl-1.10/execution/samplers/in-parameter-array: fail
-spec/glsl-1.10/execution/vs-nested-return-sibling-loop: fail
-spec/glsl-1.10/execution/vs-nested-return-sibling-loop2: fail
 spec/glsl-1.10/preprocessor/extension-defined-test: skip
 spec/glsl-1.10/preprocessor/extension-if-1: skip
 spec/glsl-1.30/execution/interpolation/interpolation-mixed: fail
@@ -12986,8 +12983,8 @@ spec/oes_viewport_array/viewport-gs-writes-out-of-range: skip
 summary:
        name:  results
        ----  --------
-       pass:     7259
+       pass:     7262
-       fail:       58
+       fail:       55
       crash:       23
        skip:    12879
     timeout:        0