Module: Mesa Branch: main Commit: 9808ef0349d30baee22548d830ff9bc285cdfb81 URL: http://cgit.freedesktop.org/mesa/mesa/commit/?id=9808ef0349d30baee22548d830ff9bc285cdfb81
Author: Daniel Schürmann <dan...@schuermann.dev> Date: Mon Jul 31 22:26:50 2023 +0200 nir/opt_loop: move loop control-flow optimizations into separate pass This new pass aims to simplify loop control-flow by reducing the number of break and continue statements. It also supersedes nir_opt_trivial_continues(). For this purpose, it implements 3 optimizations: - opt_loop_terminator(), as previously - opt_loop_merge_break_continue(), similar to opt_merge_breaks() incl. continues - opt_loop_last_block(), a generalization of opt_if_loop_last_continue() Reviewed-by: Konstantin Seurer <konstantin.seu...@gmail.com> Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/24940> --- src/compiler/nir/meson.build | 1 + src/compiler/nir/nir.h | 2 + src/compiler/nir/nir_opt_loop.c | 344 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 347 insertions(+) diff --git a/src/compiler/nir/meson.build b/src/compiler/nir/meson.build index 66011a67126..7b6e0cd197a 100644 --- a/src/compiler/nir/meson.build +++ b/src/compiler/nir/meson.build @@ -252,6 +252,7 @@ files_libnir = files( 'nir_opt_intrinsics.c', 'nir_opt_large_constants.c', 'nir_opt_load_store_vectorize.c', + 'nir_opt_loop.c', 'nir_opt_loop_unroll.c', 'nir_opt_memcpy.c', 'nir_opt_move.c', diff --git a/src/compiler/nir/nir.h b/src/compiler/nir/nir.h index 7853f0c02bc..730a320d812 100644 --- a/src/compiler/nir/nir.h +++ b/src/compiler/nir/nir.h @@ -6314,6 +6314,8 @@ bool nir_opt_large_constants(nir_shader *shader, glsl_type_size_align_func size_align, unsigned threshold); +bool nir_opt_loop(nir_shader *shader); + bool nir_opt_loop_unroll(nir_shader *shader); typedef enum { diff --git a/src/compiler/nir/nir_opt_loop.c b/src/compiler/nir/nir_opt_loop.c new file mode 100644 index 00000000000..34a60fc3446 --- /dev/null +++ b/src/compiler/nir/nir_opt_loop.c @@ -0,0 +1,344 @@ +/* + * Copyright 2023 Valve Corporation + * SPDX-License-Identifier: MIT + */ + +#include "nir.h" +#include "nir_control_flow.h" + +static bool +is_block_empty(nir_block *block) +{ + return nir_cf_node_is_last(&block->cf_node) && + exec_list_is_empty(&block->instr_list); +} + +static bool +nir_block_ends_in_continue(nir_block *block) +{ + if (exec_list_is_empty(&block->instr_list)) + return false; + + nir_instr *instr = nir_block_last_instr(block); + return instr->type == nir_instr_type_jump && + nir_instr_as_jump(instr)->type == nir_jump_continue; +} + +/** + * This optimization tries to merge two equal jump instructions (break or + * continue) into a single one. + * + * This optimization turns + * + * loop { + * ... + * if (cond) { + * do_work_1(); + * break; + * } else { + * do_work_2(); + * break; + * } + * } + * + * into: + * + * loop { + * ... + * if (cond) { + * do_work_1(); + * } else { + * do_work_2(); + * } + * break; + * } + * + * It does the same with continue statements, respectively. + * + */ +static bool +opt_loop_merge_break_continue(nir_if *nif) +{ + nir_block *after_if = nir_cf_node_cf_tree_next(&nif->cf_node); + + /* The block after the IF must have no predecessors and be empty. */ + if (after_if->predecessors->entries > 0 || !is_block_empty(after_if)) + return false; + + nir_block *last_then = nir_if_last_then_block(nif); + nir_block *last_else = nir_if_last_else_block(nif); + const bool then_break = nir_block_ends_in_break(last_then); + const bool else_break = nir_block_ends_in_break(last_else); + const bool then_cont = nir_block_ends_in_continue(last_then); + const bool else_cont = nir_block_ends_in_continue(last_else); + + /* If both branch legs end with the same jump instruction, + * merge the statement after the branch + */ + if ((then_break && else_break) || (then_cont && else_cont)) { + nir_lower_phis_to_regs_block(last_then->successors[0]); + nir_instr_remove_v(nir_block_last_instr(last_then)); + nir_instr *jump = nir_block_last_instr(last_else); + nir_instr_remove_v(jump); + nir_instr_insert(nir_after_block(after_if), jump); + return true; + } + + return false; +} + +/** + * This optimization simplifies potential loop terminators which then allows + * other passes such as opt_if_simplification() and loop unrolling to progress + * further: + * + * if (cond) { + * ... then block instructions ... + * } else { + * ... + * break; + * } + * + * into: + * + * if (cond) { + * } else { + * ... + * break; + * } + * ... then block instructions ... + */ +static bool +opt_loop_terminator(nir_if *nif) +{ + nir_block *break_blk = NULL; + nir_block *continue_from_blk = NULL; + nir_block *first_continue_from_blk = NULL; + + nir_block *last_then = nir_if_last_then_block(nif); + nir_block *last_else = nir_if_last_else_block(nif); + + if (nir_block_ends_in_break(last_then)) { + break_blk = last_then; + continue_from_blk = last_else; + first_continue_from_blk = nir_if_first_else_block(nif); + } else if (nir_block_ends_in_break(last_else)) { + break_blk = last_else; + continue_from_blk = last_then; + first_continue_from_blk = nir_if_first_then_block(nif); + } + + /* Continue if the if-statement contained no jumps at all */ + if (!break_blk) + return false; + + /* If the continue from block is empty then return as there is nothing to + * move. + */ + if (is_block_empty(first_continue_from_blk)) + return false; + + if (nir_block_ends_in_jump(continue_from_blk)) { + /* Let nir_opt_dead_cf() clean up any dead code. */ + if (!is_block_empty(nir_cf_node_cf_tree_next(&nif->cf_node))) + return false; + + /* We are about to move the predecessor. */ + nir_lower_phis_to_regs_block(continue_from_blk->successors[0]); + } + + /* Even though this if statement has a jump on one side, we may still have + * phis afterwards. Single-source phis can be produced by loop unrolling + * or dead control-flow passes and are perfectly legal. Run a quick phi + * removal on the block after the if to clean up any such phis. + */ + nir_opt_remove_phis_block(nir_cf_node_as_block(nir_cf_node_next(&nif->cf_node))); + + /* Finally, move the continue from branch after the if-statement. */ + nir_cf_list tmp; + nir_cf_extract(&tmp, nir_before_block(first_continue_from_blk), + nir_after_block(continue_from_blk)); + nir_cf_reinsert(&tmp, nir_after_cf_node(&nif->cf_node)); + + return true; +} + +/** + * This optimization tries to merge the jump instruction (break or continue) + * of a block with an equal one from a previous IF. + * + * This optimization turns: + * + * loop { + * ... + * if (cond) { + * do_work_1(); + * break; + * } else { + * } + * do_work_2(); + * break; + * } + * + * into: + * + * loop { + * ... + * if (cond) { + * do_work_1(); + * } else { + * do_work_2(); + * } + * break; + * } + * + * It does the same with continue statements, respectively. + * + */ +static bool +opt_loop_last_block(nir_block *block, bool is_loop_last_block) +{ + /* If this block has no predecessors, let nir_opt_dead_cf() do the cleanup */ + if (block->predecessors->entries == 0) + return false; + + const bool has_break = nir_block_ends_in_break(block); + const bool has_continue = nir_block_ends_in_continue(block) || is_loop_last_block; + + /* This is already handled by opt_loop_last_block(is_loop_last_block=false) */ + if (is_loop_last_block && has_break) + return false; + + if (!has_continue && !has_break) + return false; + + bool progress = false; + + /* First remove any "trivial" continue, i.e. those that are at the tail + * of a loop where we can just delete the continue instruction and + * control-flow will naturally take us back to the top of the loop. + */ + if (is_loop_last_block && nir_block_ends_in_continue(block)) { + nir_lower_phis_to_regs_block(block->successors[0]); + nir_instr_remove_v(nir_block_last_instr(block)); + progress = true; + } + + /* Walk backwards and check for previous IF statements whether one of the + * branch legs ends with an equal jump instruction as this block. + */ + for (nir_cf_node *prev = nir_cf_node_prev(&block->cf_node); prev != NULL; prev = nir_cf_node_prev(prev)) { + /* Skip blocks and nested loops */ + if (prev->type != nir_cf_node_if) + continue; + + nir_if *nif = nir_cf_node_as_if(prev); + nir_block *then_block = nir_if_last_then_block(nif); + nir_block *else_block = nir_if_last_else_block(nif); + if (!nir_block_ends_in_jump(then_block) && !nir_block_ends_in_jump(else_block)) + continue; + + bool merge_into_then = (has_continue && nir_block_ends_in_continue(else_block)) || + (has_break && nir_block_ends_in_break(else_block)); + bool merge_into_else = (has_continue && nir_block_ends_in_continue(then_block)) || + (has_break && nir_block_ends_in_break(then_block)); + + if (!merge_into_then && !merge_into_else) + continue; + + /* If there are single-source phis after the IF, get rid of them first */ + nir_opt_remove_phis_block(nir_cf_node_cf_tree_next(prev)); + + /* We are about to remove one predecessor. */ + nir_lower_phis_to_regs_block(block->successors[0]); + + nir_cf_list tmp; + nir_cf_extract(&tmp, nir_after_cf_node(prev), nir_after_block_before_jump(block)); + + if (merge_into_then) { + nir_cf_reinsert(&tmp, nir_after_block(then_block)); + nir_instr_remove_v(nir_block_last_instr(else_block)); + } else { + nir_cf_reinsert(&tmp, nir_after_block(else_block)); + nir_instr_remove_v(nir_block_last_instr(then_block)); + } + + /* Because we split the current block, the pointer is not valid anymore. */ + block = nir_cf_node_cf_tree_next(prev); + progress = true; + } + + /* Revisit these blocks with is_loop_last_block=true to optimize with the implicit continue. */ + if (is_loop_last_block && is_block_empty(block)) { + nir_cf_node *prev = nir_cf_node_prev(&block->cf_node); + if (prev && prev->type == nir_cf_node_if) { + nir_if *nif = nir_cf_node_as_if(prev); + progress |= opt_loop_last_block(nir_if_last_then_block(nif), true); + progress |= opt_loop_last_block(nir_if_last_else_block(nif), true); + } + } + + return progress; +} + +static bool +opt_loop_cf_list(struct exec_list *cf_list) +{ + bool progress = false; + foreach_list_typed_safe(nir_cf_node, cf_node, node, cf_list) { + switch (cf_node->type) { + case nir_cf_node_block: { + nir_block *block = nir_cf_node_as_block(cf_node); + progress |= opt_loop_last_block(block, false); + break; + } + + case nir_cf_node_if: { + nir_if *nif = nir_cf_node_as_if(cf_node); + progress |= opt_loop_cf_list(&nif->then_list); + progress |= opt_loop_cf_list(&nif->else_list); + progress |= opt_loop_merge_break_continue(nif); + progress |= opt_loop_terminator(nif); + break; + } + + case nir_cf_node_loop: { + nir_loop *loop = nir_cf_node_as_loop(cf_node); + assert(!nir_loop_has_continue_construct(loop)); + progress |= opt_loop_cf_list(&loop->body); + progress |= opt_loop_last_block(nir_loop_last_block(loop), true); + break; + } + + case nir_cf_node_function: + unreachable("Invalid cf type"); + } + } + + return progress; +} + +/** + * This pass aims to simplify loop control-flow by reducing the number + * of break and continue statements. + */ +bool +nir_opt_loop(nir_shader *shader) +{ + bool progress = false; + + nir_foreach_function_impl(impl, shader) { + /* First we run the simple pass to get rid of pesky continues */ + if (opt_loop_cf_list(&impl->body)) { + nir_metadata_preserve(impl, nir_metadata_none); + + /* If that made progress, we're no longer really in SSA form. */ + nir_lower_reg_intrinsics_to_ssa_impl(impl); + progress = true; + } else { + nir_metadata_preserve(impl, nir_metadata_all); + } + } + + return progress; +}
