nir/opt_licm: add filter callback

Speculative hoisting is only possible with the filter callback.

Reviewed-by: Daniel Schürmann <daniel@schuermann.dev>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/41453>

src/amd/vulkan/radv_pipeline.c

@@ -371,7 +371,7 @@ radv_postprocess_nir(const struct radv_compiler_info *compiler_info, const struc
    nir_move_options sink_opts = nir_move_const_undef | nir_move_copies | nir_dont_move_byte_word_vecs;
 
    if (!stage->key.optimisations_disabled) {
-      NIR_PASS(_, stage->nir, nir_opt_licm);
+      NIR_PASS(_, stage->nir, nir_opt_licm, NULL);
 
       if (stage->stage == MESA_SHADER_VERTEX) {
          /* Always load all VS inputs at the top to eliminate needless VMEM->s_wait->VMEM sequences.

src/compiler/nir/nir.h

@@ -6755,7 +6755,11 @@ bool nir_opt_large_constants(nir_shader *shader,
                              glsl_type_size_align_func size_align,
                              unsigned threshold);
 
-bool nir_opt_licm(nir_shader *shader);
+typedef bool (*nir_opt_licm_filter_cb)(nir_instr *instr, nir_loop *loop,
+                                       bool instr_block_dominates_exit);
+
+bool nir_opt_licm(nir_shader *shader,
+                  nir_opt_licm_filter_cb filter);
 bool nir_opt_loop(nir_shader *shader);
 
 bool nir_opt_loop_unroll(nir_shader *shader);

src/compiler/nir/nir_opt_licm.c

@@ -6,7 +6,10 @@
 #include "nir.h"
 
 typedef struct {
+   nir_opt_licm_filter_cb filter;
+
    nir_loop *loop;
+   bool current_block_dominates_exit;
 } licm_state;
 
 static bool
@@ -37,7 +40,10 @@ is_instr_loop_invariant(nir_instr *instr, licm_state *state)
    case nir_instr_type_alu:
    case nir_instr_type_tex:
    case nir_instr_type_deref:
-      return nir_foreach_src(instr, defined_before_loop, state);
+      return nir_foreach_src(instr, defined_before_loop, state) &&
+             (!state->filter ||
+              state->filter(instr, state->loop,
+                            state->current_block_dominates_exit));
 
    case nir_instr_type_phi:
    case nir_instr_type_call:
@@ -51,8 +57,6 @@ is_instr_loop_invariant(nir_instr *instr, licm_state *state)
 static bool
 visit_block(nir_block *block, licm_state *state)
 {
-   assert(state->loop);
-
    bool progress = false;
    nir_foreach_instr_safe(instr, block) {
       if (is_instr_loop_invariant(instr, state)) {
@@ -123,17 +127,24 @@ visit_cf_list(struct exec_list *list, licm_state *state)
             }
          }
 
-         /* By only visiting blocks which dominate the block after the loop,
-          * we ensure that we don't speculatively hoist any instructions
-          * which otherwise might not be executed.
-          *
-          * Note, that the proper check would be whether this block
-          * postdominates the block before the loop.
-          */
          nir_block *block = nir_cf_node_as_block(node);
-         if (state->loop &&
-             nir_block_dominates(block, nir_loop_successor_block(state->loop)))
-            progress |= visit_block(block, state);
+         if (state->loop) {
+            state->current_block_dominates_exit =
+               nir_block_dominates(block, nir_loop_successor_block(state->loop));
+
+            /* By only visiting blocks which dominate the block after the loop,
+             * we ensure that we don't speculatively hoist any instructions
+             * which otherwise might not be executed.
+             *
+             * Note, that the proper check would be whether this block
+             * postdominates the block before the loop.
+             *
+             * If filter != NULL, speculative hoisting is controlled
+             * by the callback.
+             */
+            if (state->current_block_dominates_exit || state->filter)
+               progress |= visit_block(block, state);
+         }
 
          if (next && next->type == nir_cf_node_loop && !optimize_loop) {
             nir_loop *loop = nir_cf_node_as_loop(next);
@@ -164,10 +175,15 @@ visit_cf_list(struct exec_list *list, licm_state *state)
    return progress;
 }
 
+/* Loop Invariant Code Motion.
+ *
+ * Speculative hoisting is only possible with filter != NULL, and the filter
+ * callback is expected to determine which instructions are speculatable.
+ */
 bool
-nir_opt_licm(nir_shader *shader)
+nir_opt_licm(nir_shader *shader, nir_opt_licm_filter_cb filter)
 {
-   licm_state state = {0};
+   licm_state state = {filter};
    bool progress = false;
 
    nir_foreach_function_impl(impl, shader) {

src/imagination/pco/pco_nir.c

@@ -802,7 +802,7 @@ void pco_lower_nir(pco_ctx *ctx, nir_shader *nir, pco_data *data)
 
    NIR_PASS(_, nir, nir_lower_memory_model);
 
-   NIR_PASS(_, nir, nir_opt_licm);
+   NIR_PASS(_, nir, nir_opt_licm, NULL);
 
    NIR_PASS(_, nir, nir_lower_memcpy);