nir: Tag entrypoint for easy recognition by nir_shader_get_entrypoint()

We're going to have multiple functions, so nir_shader_get_entrypoint()
needs to do something a little smarter.

Reviewed-by: Kenneth Graunke <[email protected]>
Reviewed-by: Jason Ekstrand <[email protected]>

7 files changed, 30 insertions, 6 deletions

diff --git a/src/compiler/glsl/glsl_to_nir.cpp b/src/compiler/glsl/glsl_to_nir.cpp
index 3e198e0cbef..fac9a6a0f50 100644
--- a/src/compiler/glsl/glsl_to_nir.cpp
+++ b/src/compiler/glsl/glsl_to_nir.cpp
@@ -478,6 +478,8 @@ nir_visitor::create_function(ir_function_signature *ir)
       return;
 
    nir_function *func = nir_function_create(shader, ir->function_name());
+   if (strcmp(ir->function_name(), "main") == 0)
+      func->is_entrypoint = true;
 
    func->num_params = ir->parameters.length() +
                       (ir->return_type != glsl_type::void_type);
diff --git a/src/compiler/nir/nir.c b/src/compiler/nir/nir.c
index 99e5925cd0b..3345266fa69 100644
--- a/src/compiler/nir/nir.c
+++ b/src/compiler/nir/nir.c
@@ -208,6 +208,7 @@ nir_function_create(nir_shader *shader, const char *name)
    func->num_params = 0;
    func->params = NULL;
    func->impl = NULL;
+   func->is_entrypoint = false;
 
    return func;
 }
diff --git a/src/compiler/nir/nir.h b/src/compiler/nir/nir.h
index 318ffb33cef..0ed13a3f931 100644
--- a/src/compiler/nir/nir.h
+++ b/src/compiler/nir/nir.h
@@ -2098,6 +2098,8 @@ typedef struct nir_function {
     * If the function is only declared and not implemented, this is NULL.
     */
    nir_function_impl *impl;
+
+   bool is_entrypoint;
 } nir_function;
 
 typedef struct nir_shader_compiler_options {
@@ -2278,20 +2280,32 @@ typedef struct nir_shader {
    unsigned constant_data_size;
 } nir_shader;
 
+#define nir_foreach_function(func, shader) \
+   foreach_list_typed(nir_function, func, node, &(shader)->functions)
+
 static inline nir_function_impl *
 nir_shader_get_entrypoint(nir_shader *shader)
 {
-   assert(exec_list_length(&shader->functions) == 1);
-   struct exec_node *func_node = exec_list_get_head(&shader->functions);
-   nir_function *func = exec_node_data(nir_function, func_node, node);
+   nir_function *func = NULL;
+
+   nir_foreach_function(function, shader) {
+      assert(func == NULL);
+      if (function->is_entrypoint) {
+         func = function;
+#ifndef NDEBUG
+         break;
+#endif
+      }
+   }
+
+   if (!func)
+      return NULL;
+
    assert(func->num_params == 0);
    assert(func->impl);
    return func->impl;
 }
 
-#define nir_foreach_function(func, shader) \
-   foreach_list_typed(nir_function, func, node, &(shader)->functions)
-
 nir_shader *nir_shader_create(void *mem_ctx,
                               gl_shader_stage stage,
                               const nir_shader_compiler_options *options,
diff --git a/src/compiler/nir/nir_builder.h b/src/compiler/nir/nir_builder.h
index 7dbdd5724c2..2a36eb3c91b 100644
--- a/src/compiler/nir/nir_builder.h
+++ b/src/compiler/nir/nir_builder.h
@@ -55,6 +55,7 @@ nir_builder_init_simple_shader(nir_builder *build, void *mem_ctx,
 {
    build->shader = nir_shader_create(mem_ctx, stage, options, NULL);
    nir_function *func = nir_function_create(build->shader, "main");
+   func->is_entrypoint = true;
    build->exact = false;
    build->impl = nir_function_impl_create(func);
    build->cursor = nir_after_cf_list(&build->impl->body);
diff --git a/src/compiler/nir/nir_clone.c b/src/compiler/nir/nir_clone.c
index e690b9165c9..c7c41ef5c6b 100644
--- a/src/compiler/nir/nir_clone.c
+++ b/src/compiler/nir/nir_clone.c
@@ -684,6 +684,7 @@ clone_function(clone_state *state, const nir_function *fxn, nir_shader *ns)
    nfxn->num_params = fxn->num_params;
    nfxn->params = ralloc_array(state->ns, nir_parameter, fxn->num_params);
    memcpy(nfxn->params, fxn->params, sizeof(nir_parameter) * fxn->num_params);
+   nfxn->is_entrypoint = fxn->is_entrypoint;
 
    /* At first glance, it looks like we should clone the function_impl here.
     * However, call instructions need to be able to reference at least the
diff --git a/src/compiler/nir/nir_serialize.c b/src/compiler/nir/nir_serialize.c
index b4187213716..808df193754 100644
--- a/src/compiler/nir/nir_serialize.c
+++ b/src/compiler/nir/nir_serialize.c
@@ -1048,6 +1048,8 @@ write_function(write_ctx *ctx, const nir_function *fxn)
       blob_write_uint32(ctx->blob, val);
    }
 
+   blob_write_uint32(ctx->blob, fxn->is_entrypoint);
+
    /* At first glance, it looks like we should write the function_impl here.
     * However, call instructions need to be able to reference at least the
     * function and those will get processed as we write the function_impls.
@@ -1072,6 +1074,8 @@ read_function(read_ctx *ctx)
       fxn->params[i].num_components = val & 0xff;
       fxn->params[i].bit_size = (val >> 8) & 0xff;
    }
+
+   fxn->is_entrypoint = blob_read_uint32(ctx->blob);
 }
 
 void
diff --git a/src/compiler/spirv/spirv_to_nir.c b/src/compiler/spirv/spirv_to_nir.c
index eaad406ecd4..e3dc619c6f7 100644
--- a/src/compiler/spirv/spirv_to_nir.c
+++ b/src/compiler/spirv/spirv_to_nir.c
@@ -4420,5 +4420,6 @@ spirv_to_nir(const uint32_t *words, size_t word_count,
 
    ralloc_free(b);
 
+   entry_point->is_entrypoint = true;
    return entry_point;
 }