nir/lower_outputs_to_temporaries: Take a nir_function entrypoint

4 files changed, 15 insertions, 7 deletions

diff --git a/src/glsl/nir/glsl_to_nir.cpp b/src/glsl/nir/glsl_to_nir.cpp
index cadcea50346..3cc42d229a0 100644
--- a/src/glsl/nir/glsl_to_nir.cpp
+++ b/src/glsl/nir/glsl_to_nir.cpp
@@ -145,7 +145,16 @@ glsl_to_nir(const struct gl_shader_program *shader_prog,
    v2.run(sh->ir);
    visit_exec_list(sh->ir, &v1);
 
-   nir_lower_outputs_to_temporaries(shader);
+   nir_function *main = NULL;
+   nir_foreach_function(shader, func) {
+      if (strcmp(func->name, "main") == 0) {
+         main = func;
+         break;
+      }
+   }
+   assert(main);
+
+   nir_lower_outputs_to_temporaries(shader, main);
 
    shader->info.name = ralloc_asprintf(shader, "GLSL%d", shader_prog->Name);
    if (shader_prog->Label)
diff --git a/src/glsl/nir/nir.h b/src/glsl/nir/nir.h
index b413b38220e..a050be4d53b 100644
--- a/src/glsl/nir/nir.h
+++ b/src/glsl/nir/nir.h
@@ -2036,9 +2036,8 @@ bool nir_lower_global_vars_to_local(nir_shader *shader);
 
 bool nir_lower_locals_to_regs(nir_shader *shader);
 
-void nir_lower_outputs_to_temporaries(nir_shader *shader);
-
-void nir_lower_outputs_to_temporaries(nir_shader *shader);
+void nir_lower_outputs_to_temporaries(nir_shader *shader,
+                                      nir_function *entrypoint);
 
 void nir_shader_gather_info(nir_shader *shader, nir_function_impl *entrypoint);
 
diff --git a/src/glsl/nir/nir_lower_outputs_to_temporaries.c b/src/glsl/nir/nir_lower_outputs_to_temporaries.c
index 71b06b81fcc..70d85138552 100644
--- a/src/glsl/nir/nir_lower_outputs_to_temporaries.c
+++ b/src/glsl/nir/nir_lower_outputs_to_temporaries.c
@@ -74,7 +74,7 @@ emit_output_copies_block(nir_block *block, void *state)
 }
 
 void
-nir_lower_outputs_to_temporaries(nir_shader *shader)
+nir_lower_outputs_to_temporaries(nir_shader *shader, nir_function *entrypoint)
 {
    struct lower_outputs_state state;
 
@@ -114,7 +114,7 @@ nir_lower_outputs_to_temporaries(nir_shader *shader)
           * before each EmitVertex call.
           */
          nir_foreach_block(function->impl, emit_output_copies_block, &state);
-      } else if (strcmp(function->name, "main") == 0) {
+      } else if (function == entrypoint) {
          /* For all other shader types, we need to do the copies right before
           * the jumps to the end block.
           */
diff --git a/src/glsl/nir/spirv/spirv_to_nir.c b/src/glsl/nir/spirv/spirv_to_nir.c
index 5b31f7e7e2a..c4b0c50c52e 100644
--- a/src/glsl/nir/spirv/spirv_to_nir.c
+++ b/src/glsl/nir/spirv/spirv_to_nir.c
@@ -3738,7 +3738,7 @@ spirv_to_nir(const uint32_t *words, size_t word_count,
    /* Because we can still have output reads in NIR, we need to lower
     * outputs to temporaries before we are truely finished.
     */
-   nir_lower_outputs_to_temporaries(entry_point->shader);
+   nir_lower_outputs_to_temporaries(entry_point->shader, entry_point);
 
    return entry_point;
 }