nir: implement loop invariant code motion (LICM) pass

This simple LICM pass hoists all loop-invariant instructions
from the loops' top-level control flow, skipping any nested CF.
The hoisted instructions are placed right before the loop.

Reviewed-by: Georg Lehmann <dadschoorse@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/28783>

src/compiler/nir/meson.build

@@ -234,6 +234,7 @@ files_libnir = files(
   'nir_opt_if.c',
   'nir_opt_intrinsics.c',
   'nir_opt_large_constants.c',
+  'nir_opt_licm.c',
   'nir_opt_load_store_vectorize.c',
   'nir_opt_loop.c',
   'nir_opt_loop_unroll.c',

src/compiler/nir/nir.h

@@ -6606,6 +6606,7 @@ bool nir_opt_large_constants(nir_shader *shader,
                              glsl_type_size_align_func size_align,
                              unsigned threshold);
 
+bool nir_opt_licm(nir_shader *shader);
 bool nir_opt_loop(nir_shader *shader);
 
 bool nir_opt_loop_unroll(nir_shader *shader);

src/compiler/nir/nir_opt_licm.c

@@ -0,0 +1,144 @@
+/*
+ * Copyright 2024 Valve Corporation
+ * SPDX-License-Identifier: MIT
+ */
+
+#include "nir.h"
+
+static bool
+defined_before_loop(nir_src *src, void *state)
+{
+   unsigned *loop_preheader_idx = state;
+   return src->ssa->parent_instr->block->index <= *loop_preheader_idx;
+}
+
+static bool
+is_instr_loop_invariant(nir_instr *instr, unsigned loop_preheader_idx)
+{
+   switch (instr->type) {
+   case nir_instr_type_load_const:
+   case nir_instr_type_undef:
+      return true;
+
+   case nir_instr_type_intrinsic:
+      if (!nir_intrinsic_can_reorder(nir_instr_as_intrinsic(instr)))
+         return false;
+      FALLTHROUGH;
+
+   case nir_instr_type_alu:
+   case nir_instr_type_tex:
+   case nir_instr_type_deref:
+      return nir_foreach_src(instr, defined_before_loop, &loop_preheader_idx);
+
+   case nir_instr_type_phi:
+   case nir_instr_type_call:
+   case nir_instr_type_jump:
+   default:
+      return false;
+   }
+}
+
+static bool
+visit_block(nir_block *block, nir_block *preheader)
+{
+   bool progress = false;
+   nir_foreach_instr_safe(instr, block) {
+      if (is_instr_loop_invariant(instr, preheader->index)) {
+         nir_instr_remove(instr);
+         nir_instr_insert_after_block(preheader, instr);
+         progress = true;
+      }
+   }
+
+   return progress;
+}
+
+static bool
+should_optimize_loop(nir_loop *loop)
+{
+   /* Ignore loops without back-edge */
+   if (nir_loop_first_block(loop)->predecessors->entries == 1)
+      return false;
+
+   nir_foreach_block_in_cf_node(block, &loop->cf_node) {
+      /* Check for an early exit inside the loop. */
+      nir_foreach_instr(instr, block) {
+         if (instr->type == nir_instr_type_intrinsic) {
+            nir_intrinsic_instr *intrin = nir_instr_as_intrinsic(instr);
+            if (intrin->intrinsic == nir_intrinsic_terminate ||
+                intrin->intrinsic == nir_intrinsic_terminate_if)
+               return false;
+         }
+      }
+
+      /* The loop must not contains any return statement. */
+      if (nir_block_ends_in_return_or_halt(block))
+         return false;
+   }
+
+   return true;
+}
+
+static bool
+visit_cf_list(struct exec_list *list, nir_block *preheader, nir_block *exit)
+{
+   bool progress = false;
+
+   foreach_list_typed(nir_cf_node, node, node, list) {
+      switch (node->type) {
+      case nir_cf_node_block: {
+         /* By only visiting blocks which dominate the loop exit, 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 loop preheader.
+          */
+         nir_block *block = nir_cf_node_as_block(node);
+         if (exit && nir_block_dominates(block, exit))
+            progress |= visit_block(block, preheader);
+         break;
+      }
+      case nir_cf_node_if: {
+         nir_if *nif = nir_cf_node_as_if(node);
+         progress |= visit_cf_list(&nif->then_list, preheader, exit);
+         progress |= visit_cf_list(&nif->else_list, preheader, exit);
+         break;
+      }
+      case nir_cf_node_loop: {
+         nir_loop *loop = nir_cf_node_as_loop(node);
+         bool opt = should_optimize_loop(loop);
+         nir_block *inner_preheader = opt ? nir_cf_node_cf_tree_prev(node) : preheader;
+         nir_block *inner_exit = opt ? nir_cf_node_cf_tree_next(node) : exit;
+         progress |= visit_cf_list(&loop->body, inner_preheader, inner_exit);
+         progress |= visit_cf_list(&loop->continue_list, inner_preheader, inner_exit);
+         break;
+      }
+      case nir_cf_node_function:
+         unreachable("NIR LICM: Unsupported cf_node type.");
+      }
+   }
+
+   return progress;
+}
+
+bool
+nir_opt_licm(nir_shader *shader)
+{
+   bool progress = false;
+
+   nir_foreach_function_impl(impl, shader) {
+      nir_metadata_require(impl, nir_metadata_block_index |
+                                    nir_metadata_dominance);
+
+      if (visit_cf_list(&impl->body, NULL, NULL)) {
+         progress = true;
+         nir_metadata_preserve(impl, nir_metadata_block_index |
+                                        nir_metadata_dominance);
+      } else {
+         nir_metadata_preserve(impl, nir_metadata_all);
+      }
+   }
+
+   return progress;
+}