nir: Add nir_lower_halt_to_return

This is a lowering pass that was implemented by multiple drivers.

Reviewed-by: Mary Guillemard <mary@mary.zone>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/33003>

src/asahi/vulkan/hk_shader.c

@@ -106,29 +106,13 @@ hk_get_spirv_options(struct vk_physical_device *vk_pdev,
    };
 }
 
-static bool
-lower_halt_to_return(nir_builder *b, nir_instr *instr, UNUSED void *_data)
-{
-   if (instr->type != nir_instr_type_jump)
-      return false;
-
-   nir_jump_instr *jump = nir_instr_as_jump(instr);
-   if (jump->type != nir_jump_halt)
-      return false;
-
-   assert(b->impl == nir_shader_get_entrypoint(b->shader));
-   jump->type = nir_jump_return;
-   return true;
-}
-
 void
 hk_preprocess_nir_internal(struct vk_physical_device *vk_pdev, nir_shader *nir)
 {
    /* Must lower before io to temps */
    if (nir->info.stage == MESA_SHADER_FRAGMENT) {
       NIR_PASS(_, nir, nir_lower_terminate_to_demote);
-      NIR_PASS(_, nir, nir_shader_instructions_pass, lower_halt_to_return,
-               nir_metadata_all, NULL);
+      NIR_PASS(_, nir, nir_lower_halt_to_return);
       NIR_PASS(_, nir, nir_lower_returns);
    }
 

src/compiler/nir/meson.build

@@ -175,6 +175,7 @@ else
   'nir_lower_global_vars_to_local.c',
   'nir_lower_goto_ifs.c',
   'nir_lower_gs_intrinsics.c',
+  'nir_lower_halt_to_return.c',
   'nir_lower_helper_writes.c',
   'nir_lower_load_const_to_scalar.c',
   'nir_lower_locals_to_regs.c',

src/compiler/nir/nir.h

@@ -5848,6 +5848,8 @@ typedef enum {
 
 bool nir_lower_gs_intrinsics(nir_shader *shader, nir_lower_gs_intrinsics_flags options);
 
+bool nir_lower_halt_to_return(nir_shader *nir);
+
 bool nir_lower_tess_coord_z(nir_shader *shader, bool triangles);
 
 typedef struct nir_lower_task_shader_options {

src/compiler/nir/nir_lower_halt_to_return.c

@@ -0,0 +1,28 @@
+/*
+ * Copyright © 2025 Valve Corporation
+ * SPDX-License-Identifier: MIT
+ */
+
+#include "nir.h"
+#include "nir_builder.h"
+
+static bool
+pass(nir_builder *b, nir_instr *instr, void *data)
+{
+   if (instr->type != nir_instr_type_jump)
+      return false;
+
+   nir_jump_instr *jump = nir_instr_as_jump(instr);
+   if (jump->type == nir_jump_halt) {
+      jump->type = nir_jump_return;
+      return true;
+   }
+
+   return false;
+}
+
+bool
+nir_lower_halt_to_return(nir_shader *nir)
+{
+   return nir_shader_instructions_pass(nir, pass, nir_metadata_all, NULL);
+}

src/panfrost/compiler/bifrost_compile.c

@@ -5755,21 +5755,6 @@ bifrost_nir_lower_load_output(nir_shader *nir)
       nir_metadata_control_flow, NULL);
 }
 
-static bool
-bi_lower_halt_to_return(nir_builder *b, nir_instr *instr, UNUSED void *_data)
-{
-   if (instr->type != nir_instr_type_jump)
-      return false;
-
-   nir_jump_instr *jump = nir_instr_as_jump(instr);
-   if (jump->type != nir_jump_halt)
-      return false;
-
-   assert(b->impl == nir_shader_get_entrypoint(b->shader));
-   jump->type = nir_jump_return;
-   return true;
-}
-
 /* Bifrost LDEXP.v2f16 takes i16 exponent, while nir_op_ldexp takes i32. Lower
  * to nir_op_ldexp16_pan. */
 static bool
@@ -5831,8 +5816,7 @@ bifrost_preprocess_nir(nir_shader *nir, unsigned gpu_id)
       NIR_PASS(_, nir, nir_lower_terminate_to_demote);
 
    /* Ensure that halt are translated to returns and get ride of them */
-   NIR_PASS(_, nir, nir_shader_instructions_pass, bi_lower_halt_to_return,
-            nir_metadata_all, NULL);
+   NIR_PASS(_, nir, nir_lower_halt_to_return);
    NIR_PASS(_, nir, nir_lower_returns);
 
    /* Lower gl_Position pre-optimisation, but after lowering vars to ssa

src/panfrost/midgard/midgard_compile.c

@@ -370,29 +370,13 @@ lower_vec816_alu(const nir_instr *instr, const void *cb_data)
    return 4;
 }
 
-static bool
-lower_halt_to_return(nir_builder *b, nir_instr *instr, UNUSED void *_data)
-{
-   if (instr->type != nir_instr_type_jump)
-      return false;
-
-   nir_jump_instr *jump = nir_instr_as_jump(instr);
-   if (jump->type != nir_jump_halt)
-      return false;
-
-   assert(b->impl == nir_shader_get_entrypoint(b->shader));
-   jump->type = nir_jump_return;
-   return true;
-}
-
 void
 midgard_preprocess_nir(nir_shader *nir, unsigned gpu_id)
 {
    unsigned quirks = midgard_get_quirks(gpu_id);
 
    /* Ensure that halt are translated to returns and get ride of them */
-   NIR_PASS(_, nir, nir_shader_instructions_pass, lower_halt_to_return,
-            nir_metadata_all, NULL);
+   NIR_PASS(_, nir, nir_lower_halt_to_return);
    NIR_PASS(_, nir, nir_lower_returns);
 
    /* Lower gl_Position pre-optimisation, but after lowering vars to ssa