nir: Add preamble functions

These are functions that run before the entrypoint at least once per
draw and write their results via store_preamble, and then are loaded in
the rest of the shader via load_preamble.

We will add users in the following commits.

Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/13148>

src/compiler/nir/nir.c

@@ -426,6 +426,7 @@ nir_function_create(nir_shader *shader, const char *name)
    func->params = NULL;
    func->impl = NULL;
    func->is_entrypoint = false;
+   func->is_preamble = false;
 
    return func;
 }
@@ -541,6 +542,7 @@ nir_function_impl_create_bare(nir_shader *shader)
    nir_function_impl *impl = ralloc(shader, nir_function_impl);
 
    impl->function = NULL;
+   impl->preamble = NULL;
 
    cf_init(&impl->cf_node, nir_cf_node_function);
 

src/compiler/nir/nir.h

@@ -2942,6 +2942,13 @@ typedef struct {
    /** pointer to the function of which this is an implementation */
    struct nir_function *function;
 
+   /**
+    * For entrypoints, a pointer to a nir_function_impl which runs before
+    * it, once per draw or dispatch, communicating via store_preamble and
+    * load_preamble intrinsics. If NULL then there is no preamble.
+    */
+   struct nir_function *preamble;
+
    struct exec_list body; /** < list of nir_cf_node */
 
    nir_block *end_block;
@@ -3114,6 +3121,7 @@ typedef struct nir_function {
    nir_function_impl *impl;
 
    bool is_entrypoint;
+   bool is_preamble;
 } nir_function;
 
 typedef enum {

src/compiler/nir/nir_clone.c

@@ -671,6 +671,9 @@ clone_function_impl(clone_state *state, const nir_function_impl *fi)
 {
    nir_function_impl *nfi = nir_function_impl_create_bare(state->ns);
 
+   if (fi->preamble)
+      nfi->preamble = remap_global(state, fi->preamble);
+
    clone_var_list(state, &nfi->locals, &fi->locals);
    clone_reg_list(state, &nfi->registers, &fi->registers);
    nfi->reg_alloc = fi->reg_alloc;
@@ -717,6 +720,7 @@ clone_function(clone_state *state, const nir_function *fxn, nir_shader *ns)
            memcpy(nfxn->params, fxn->params, sizeof(nir_parameter) * fxn->num_params);
    }
    nfxn->is_entrypoint = fxn->is_entrypoint;
+   nfxn->is_preamble = fxn->is_preamble;
 
    /* 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
@@ -743,9 +747,9 @@ nir_shader_clone(void *mem_ctx, const nir_shader *s)
       clone_function(&state, fxn, ns);
 
    /* Only after all functions are cloned can we clone the actual function
-    * implementations.  This is because nir_call_instrs need to reference the
-    * functions of other functions and we don't know what order the functions
-    * will have in the list.
+    * implementations.  This is because nir_call_instrs and preambles need to
+    * reference the functions of other functions and we don't know what order
+    * the functions will have in the list.
     */
    nir_foreach_function(fxn, s) {
       nir_function *nfxn = remap_global(&state, fxn);

src/compiler/nir/nir_intrinsics.py

@@ -1047,6 +1047,13 @@ system_value("printf_buffer_address", 1, bit_sizes=[32,64])
 system_value("mesh_view_count", 1)
 load("mesh_view_indices", [1], [BASE, RANGE], [CAN_ELIMINATE, CAN_REORDER])
 
+# Used to pass values from the preamble to the main shader.
+# This should use something similar to Vulkan push constants and load_preamble
+# should be relatively cheap.
+# For now we only support accesses with a constant offset.
+load("preamble", [], indices=[BASE], flags=[CAN_ELIMINATE, CAN_REORDER])
+store("preamble", [], indices=[BASE])
+
 # IR3-specific version of most SSBO intrinsics. The only different
 # compare to the originals is that they add an extra source to hold
 # the dword-offset, which is needed by the backend code apart from

src/compiler/nir/nir_print.c

@@ -1600,6 +1600,10 @@ print_function_impl(nir_function_impl *impl, print_state *state)
 
    fprintf(fp, "{\n");
 
+   if (impl->preamble) {
+      fprintf(fp, "\tpreamble %s\n", impl->preamble->name);
+   }
+
    nir_foreach_function_temp_variable(var, impl) {
       fprintf(fp, "\t");
       print_var_decl(var, state);

src/compiler/nir/nir_serialize.c

@@ -1955,6 +1955,10 @@ static void
 write_function_impl(write_ctx *ctx, const nir_function_impl *fi)
 {
    blob_write_uint8(ctx->blob, fi->structured);
+   blob_write_uint8(ctx->blob, !!fi->preamble);
+
+   if (fi->preamble)
+      blob_write_uint32(ctx->blob, write_lookup_object(ctx, fi->preamble));
 
    write_var_list(ctx, &fi->locals);
    write_reg_list(ctx, &fi->registers);
@@ -1971,6 +1975,10 @@ read_function_impl(read_ctx *ctx, nir_function *fxn)
    fi->function = fxn;
 
    fi->structured = blob_read_uint8(ctx->blob);
+   bool preamble = blob_read_uint8(ctx->blob);
+
+   if (preamble)
+      fi->preamble = read_object(ctx);
 
    read_var_list(ctx, &fi->locals);
    read_reg_list(ctx, &fi->registers);
@@ -1987,11 +1995,15 @@ read_function_impl(read_ctx *ctx, nir_function *fxn)
 static void
 write_function(write_ctx *ctx, const nir_function *fxn)
 {
-   uint32_t flags = fxn->is_entrypoint;
-   if (fxn->name)
+   uint32_t flags = 0;
+   if (fxn->is_entrypoint)
+      flags |= 0x1;
+   if (fxn->is_preamble)
       flags |= 0x2;
-   if (fxn->impl)
+   if (fxn->name)
       flags |= 0x4;
+   if (fxn->impl)
+      flags |= 0x8;
    blob_write_uint32(ctx->blob, flags);
    if (fxn->name)
       blob_write_string(ctx->blob, fxn->name);
@@ -2017,7 +2029,7 @@ static void
 read_function(read_ctx *ctx)
 {
    uint32_t flags = blob_read_uint32(ctx->blob);
-   bool has_name = flags & 0x2;
+   bool has_name = flags & 0x4;
    char *name = has_name ? blob_read_string(ctx->blob) : NULL;
 
    nir_function *fxn = nir_function_create(ctx->nir, name);
@@ -2033,7 +2045,8 @@ read_function(read_ctx *ctx)
    }
 
    fxn->is_entrypoint = flags & 0x1;
-   if (flags & 0x4)
+   fxn->is_preamble = flags & 0x2;
+   if (flags & 0x8)
       fxn->impl = NIR_SERIALIZE_FUNC_HAS_IMPL;
 }
 

src/compiler/nir/nir_validate.c

@@ -1619,6 +1619,11 @@ validate_function_impl(nir_function_impl *impl, validate_state *state)
    validate_assert(state, impl->function->impl == impl);
    validate_assert(state, impl->cf_node.parent == NULL);
 
+   if (impl->preamble) {
+      validate_assert(state, impl->function->is_entrypoint);
+      validate_assert(state, impl->preamble->is_preamble);
+   }
+
    validate_assert(state, exec_list_is_empty(&impl->end_block->instr_list));
    validate_assert(state, impl->end_block->successors[0] == NULL);
    validate_assert(state, impl->end_block->successors[1] == NULL);