[Mesa-dev] [PATCH 1/2] nir: update nir_lower_returns to only predicate instructions when needed

[Timothy Arceri]

Unless an if statement contains nested returns we can simply add
any following instructions to the branch without the return.

V2: fix handling if_nested_return value when there is a sibling if/loop
that doesn't contain a return. (Spotted by Ken)
---
 src/compiler/nir/nir_lower_returns.c | 37 ++++++++++++++++++++++++++++++------
 1 file changed, 31 insertions(+), 6 deletions(-)

diff --git a/src/compiler/nir/nir_lower_returns.c b/src/compiler/nir/nir_lower_returns.c
index cf49d5b..5eec984 100644
--- a/src/compiler/nir/nir_lower_returns.c
+++ b/src/compiler/nir/nir_lower_returns.c
@@ -30,6 +30,8 @@ struct lower_returns_state {
    struct exec_list *cf_list;
    nir_loop *loop;
    nir_variable *return_flag;
+   /* Are there other return statments nested in the current if */
+   bool if_nested_return;
 };
 
 static bool lower_returns_in_cf_list(struct exec_list *cf_list,
@@ -82,8 +84,10 @@ lower_returns_in_loop(nir_loop *loop, struct lower_returns_state *state)
     * flag set to true.  We need to predicate everything following the loop
     * on the return flag.
     */
-   if (progress)
+   if (progress) {
       predicate_following(&loop->cf_node, state);
+      state->if_nested_return = true;
+   }
 
    return progress;
 }
@@ -91,10 +95,13 @@ lower_returns_in_loop(nir_loop *loop, struct lower_returns_state *state)
 static bool
 lower_returns_in_if(nir_if *if_stmt, struct lower_returns_state *state)
 {
-   bool progress;
+   bool progress, then_progress;
 
-   progress = lower_returns_in_cf_list(&if_stmt->then_list, state);
-   progress = lower_returns_in_cf_list(&if_stmt->else_list, state) || progress;
+   bool if_nested_return = state->if_nested_return;
+   state->if_nested_return = false;
+
+   then_progress = lower_returns_in_cf_list(&if_stmt->then_list, state);
+   progress = lower_returns_in_cf_list(&if_stmt->else_list, state) || then_progress;
 
    /* If either of the recursive calls made progress, then there were
     * returns inside of the body of the if.  If we're in a loop, then these
@@ -106,8 +113,25 @@ lower_returns_in_if(nir_if *if_stmt, struct lower_returns_state *state)
     * after a return, we need to predicate everything following on the
     * return flag.
     */
-   if (progress && !state->loop)
-      predicate_following(&if_stmt->cf_node, state);
+   if (progress && !state->loop) {
+      if (state->if_nested_return) {
+         predicate_following(&if_stmt->cf_node, state);
+      } else {
+         /* If there are no nested returns we can just add the instructions to
+          * the end of the branch that doesn't have the return.
+          */
+         nir_cf_list list;
+         nir_cf_extract(&list, nir_after_cf_node(&if_stmt->cf_node),
+                        nir_after_cf_list(state->cf_list));
+
+         if (then_progress)
+            nir_cf_reinsert(&list, nir_after_cf_list(&if_stmt->else_list));
+         else
+            nir_cf_reinsert(&list, nir_after_cf_list(&if_stmt->then_list));
+      }
+   }
+
+   state->if_nested_return = progress || if_nested_return;
 
    return progress;
 }
@@ -221,6 +245,7 @@ nir_lower_returns_impl(nir_function_impl *impl)
    state.cf_list = &impl->body;
    state.loop = NULL;
    state.return_flag = NULL;
+   state.if_nested_return = false;
    nir_builder_init(&state.builder, impl);
 
    bool progress = lower_returns_in_cf_list(&impl->body, &state);
-- 
2.9.3