zink: use ralloc for spirv_builder as well

Reviewed-By: Mike Blumenkrantz <[email protected]>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/5954>

3 files changed, 64 insertions, 56 deletions

diff --git a/src/gallium/drivers/zink/nir_to_spirv/nir_to_spirv.c b/src/gallium/drivers/zink/nir_to_spirv/nir_to_spirv.c
index 2f55585b415..3d1e56d5f0e 100644
--- a/src/gallium/drivers/zink/nir_to_spirv/nir_to_spirv.c
+++ b/src/gallium/drivers/zink/nir_to_spirv/nir_to_spirv.c
@@ -2171,6 +2171,7 @@ nir_to_spirv(struct nir_shader *s, const struct pipe_stream_output_info *so_info
 
    struct ntv_context ctx = {};
    ctx.mem_ctx = ralloc_context(NULL);
+   ctx.builder.mem_ctx = ctx.mem_ctx;
 
    switch (s->info.stage) {
    case MESA_SHADER_VERTEX:
diff --git a/src/gallium/drivers/zink/nir_to_spirv/spirv_builder.c b/src/gallium/drivers/zink/nir_to_spirv/spirv_builder.c
index a03d08e06eb..8bcad066601 100644
--- a/src/gallium/drivers/zink/nir_to_spirv/spirv_builder.c
+++ b/src/gallium/drivers/zink/nir_to_spirv/spirv_builder.c
@@ -24,6 +24,7 @@
 #include "spirv_builder.h"
 
 #include "util/macros.h"
+#include "util/ralloc.h"
 #include "util/u_bitcast.h"
 #include "util/u_memory.h"
 #include "util/hash_table.h"
@@ -35,11 +36,12 @@
 #include <string.h>
 
 static bool
-spirv_buffer_grow(struct spirv_buffer *b, size_t needed)
+spirv_buffer_grow(struct spirv_buffer *b, void *mem_ctx, size_t needed)
 {
    size_t new_room = MAX3(64, (b->room * 3) / 2, needed);
 
-   uint32_t *new_words = realloc(b->words, new_room * sizeof(uint32_t));
+   uint32_t *new_words = reralloc_size(mem_ctx, b->words,
+                                       new_room * sizeof(uint32_t));
    if (!new_words)
       return false;
 
@@ -49,13 +51,13 @@ spirv_buffer_grow(struct spirv_buffer *b, size_t needed)
 }
 
 static inline bool
-spirv_buffer_prepare(struct spirv_buffer *b, size_t needed)
+spirv_buffer_prepare(struct spirv_buffer *b, void *mem_ctx, size_t needed)
 {
    needed += b->num_words;
    if (b->room >= b->num_words + needed)
       return true;
 
-   return spirv_buffer_grow(b, needed);
+   return spirv_buffer_grow(b, mem_ctx, needed);
 }
 
 static inline void
@@ -66,20 +68,21 @@ spirv_buffer_emit_word(struct spirv_buffer *b, uint32_t word)
 }
 
 static int
-spirv_buffer_emit_string(struct spirv_buffer *b, const char *str)
+spirv_buffer_emit_string(struct spirv_buffer *b, void *mem_ctx,
+                         const char *str)
 {
    int pos = 0;
    uint32_t word = 0;
    while (str[pos] != '\0') {
       word |= str[pos] << (8 * (pos % 4));
       if (++pos % 4 == 0) {
-         spirv_buffer_prepare(b, 1);
+         spirv_buffer_prepare(b, mem_ctx, 1);
          spirv_buffer_emit_word(b, word);
          word = 0;
       }
    }
 
-   spirv_buffer_prepare(b, 1);
+   spirv_buffer_prepare(b, mem_ctx, 1);
    spirv_buffer_emit_word(b, word);
 
    return 1 + pos / 4;
@@ -88,7 +91,7 @@ spirv_buffer_emit_string(struct spirv_buffer *b, const char *str)
 void
 spirv_builder_emit_cap(struct spirv_builder *b, SpvCapability cap)
 {
-   spirv_buffer_prepare(&b->capabilities, 2);
+   spirv_buffer_prepare(&b->capabilities, b->mem_ctx, 2);
    spirv_buffer_emit_word(&b->capabilities, SpvOpCapability | (2 << 16));
    spirv_buffer_emit_word(&b->capabilities, cap);
 }
@@ -97,7 +100,7 @@ void
 spirv_builder_emit_source(struct spirv_builder *b, SpvSourceLanguage lang,
                           uint32_t version)
 {
-   spirv_buffer_prepare(&b->debug_names, 3);
+   spirv_buffer_prepare(&b->debug_names, b->mem_ctx, 3);
    spirv_buffer_emit_word(&b->debug_names, SpvOpSource | (3 << 16));
    spirv_buffer_emit_word(&b->debug_names, lang);
    spirv_buffer_emit_word(&b->debug_names, version);
@@ -108,7 +111,7 @@ spirv_builder_emit_mem_model(struct spirv_builder *b,
                              SpvAddressingModel addr_model,
                              SpvMemoryModel mem_model)
 {
-   spirv_buffer_prepare(&b->memory_model, 3);
+   spirv_buffer_prepare(&b->memory_model, b->mem_ctx, 3);
    spirv_buffer_emit_word(&b->memory_model, SpvOpMemoryModel | (3 << 16));
    spirv_buffer_emit_word(&b->memory_model, addr_model);
    spirv_buffer_emit_word(&b->memory_model, mem_model);
@@ -121,13 +124,13 @@ spirv_builder_emit_entry_point(struct spirv_builder *b,
                                size_t num_interfaces)
 {
    size_t pos = b->entry_points.num_words;
-   spirv_buffer_prepare(&b->entry_points, 3);
+   spirv_buffer_prepare(&b->entry_points, b->mem_ctx, 3);
    spirv_buffer_emit_word(&b->entry_points, SpvOpEntryPoint);
    spirv_buffer_emit_word(&b->entry_points, exec_model);
    spirv_buffer_emit_word(&b->entry_points, entry_point);
-   int len = spirv_buffer_emit_string(&b->entry_points, name);
+   int len = spirv_buffer_emit_string(&b->entry_points, b->mem_ctx, name);
    b->entry_points.words[pos] |= (3 + len + num_interfaces) << 16;
-   spirv_buffer_prepare(&b->entry_points, num_interfaces);
+   spirv_buffer_prepare(&b->entry_points, b->mem_ctx, num_interfaces);
    for (int i = 0; i < num_interfaces; ++i)
         spirv_buffer_emit_word(&b->entry_points, interfaces[i]);
 }
@@ -136,7 +139,7 @@ void
 spirv_builder_emit_exec_mode(struct spirv_builder *b, SpvId entry_point,
                              SpvExecutionMode exec_mode)
 {
-   spirv_buffer_prepare(&b->exec_modes, 3);
+   spirv_buffer_prepare(&b->exec_modes, b->mem_ctx, 3);
    spirv_buffer_emit_word(&b->exec_modes, SpvOpExecutionMode | (3 << 16));
    spirv_buffer_emit_word(&b->exec_modes, entry_point);
    spirv_buffer_emit_word(&b->exec_modes, exec_mode);
@@ -147,10 +150,10 @@ spirv_builder_emit_name(struct spirv_builder *b, SpvId target,
                         const char *name)
 {
    size_t pos = b->debug_names.num_words;
-   spirv_buffer_prepare(&b->debug_names, 2);
+   spirv_buffer_prepare(&b->debug_names, b->mem_ctx, 2);
    spirv_buffer_emit_word(&b->debug_names, SpvOpName);
    spirv_buffer_emit_word(&b->debug_names, target);
-   int len = spirv_buffer_emit_string(&b->debug_names, name);
+   int len = spirv_buffer_emit_string(&b->debug_names, b->mem_ctx, name);
    b->debug_names.words[pos] |= (2 + len) << 16;
 }
 
@@ -160,7 +163,7 @@ emit_decoration(struct spirv_builder *b, SpvId target,
                 size_t num_extra_operands)
 {
    int words = 3 + num_extra_operands;
-   spirv_buffer_prepare(&b->decorations, words);
+   spirv_buffer_prepare(&b->decorations, b->mem_ctx, words);
    spirv_buffer_emit_word(&b->decorations, SpvOpDecorate | (words << 16));
    spirv_buffer_emit_word(&b->decorations, target);
    spirv_buffer_emit_word(&b->decorations, decoration);
@@ -261,7 +264,7 @@ emit_member_decoration(struct spirv_builder *b, SpvId target, uint32_t member,
                        size_t num_extra_operands)
 {
    int words = 4 + num_extra_operands;
-   spirv_buffer_prepare(&b->decorations, words);
+   spirv_buffer_prepare(&b->decorations, b->mem_ctx, words);
    spirv_buffer_emit_word(&b->decorations,
                           SpvOpMemberDecorate | (words << 16));
    spirv_buffer_emit_word(&b->decorations, target);
@@ -284,7 +287,7 @@ SpvId
 spirv_builder_emit_undef(struct spirv_builder *b, SpvId result_type)
 {
    SpvId result = spirv_builder_new_id(b);
-   spirv_buffer_prepare(&b->instructions, 3);
+   spirv_buffer_prepare(&b->instructions, b->mem_ctx, 3);
    spirv_buffer_emit_word(&b->instructions, SpvOpUndef | (3 << 16));
    spirv_buffer_emit_word(&b->instructions, result_type);
    spirv_buffer_emit_word(&b->instructions, result);
@@ -297,7 +300,7 @@ spirv_builder_function(struct spirv_builder *b, SpvId result,
                        SpvFunctionControlMask function_control,
                        SpvId function_type)
 {
-   spirv_buffer_prepare(&b->instructions, 5);
+   spirv_buffer_prepare(&b->instructions, b->mem_ctx, 5);
    spirv_buffer_emit_word(&b->instructions, SpvOpFunction | (5 << 16));
    spirv_buffer_emit_word(&b->instructions, return_type);
    spirv_buffer_emit_word(&b->instructions, result);
@@ -308,14 +311,14 @@ spirv_builder_function(struct spirv_builder *b, SpvId result,
 void
 spirv_builder_function_end(struct spirv_builder *b)
 {
-   spirv_buffer_prepare(&b->instructions, 1);
+   spirv_buffer_prepare(&b->instructions, b->mem_ctx, 1);
    spirv_buffer_emit_word(&b->instructions, SpvOpFunctionEnd | (1 << 16));
 }
 
 void
 spirv_builder_label(struct spirv_builder *b, SpvId label)
 {
-   spirv_buffer_prepare(&b->instructions, 2);
+   spirv_buffer_prepare(&b->instructions, b->mem_ctx, 2);
    spirv_buffer_emit_word(&b->instructions, SpvOpLabel | (2 << 16));
    spirv_buffer_emit_word(&b->instructions, label);
 }
@@ -323,7 +326,7 @@ spirv_builder_label(struct spirv_builder *b, SpvId label)
 void
 spirv_builder_return(struct spirv_builder *b)
 {
-   spirv_buffer_prepare(&b->instructions, 1);
+   spirv_buffer_prepare(&b->instructions, b->mem_ctx, 1);
    spirv_buffer_emit_word(&b->instructions, SpvOpReturn | (1 << 16));
 }
 
@@ -337,7 +340,7 @@ spirv_builder_emit_load(struct spirv_builder *b, SpvId result_type,
 void
 spirv_builder_emit_store(struct spirv_builder *b, SpvId pointer, SpvId object)
 {
-   spirv_buffer_prepare(&b->instructions, 3);
+   spirv_buffer_prepare(&b->instructions, b->mem_ctx, 3);
    spirv_buffer_emit_word(&b->instructions, SpvOpStore | (3 << 16));
    spirv_buffer_emit_word(&b->instructions, pointer);
    spirv_buffer_emit_word(&b->instructions, object);
@@ -351,7 +354,7 @@ spirv_builder_emit_access_chain(struct spirv_builder *b, SpvId result_type,
    SpvId result = spirv_builder_new_id(b);
 
    int words = 4 + num_indexes;
-   spirv_buffer_prepare(&b->instructions, words);
+   spirv_buffer_prepare(&b->instructions, b->mem_ctx, words);
    spirv_buffer_emit_word(&b->instructions, SpvOpAccessChain | (words << 16));
    spirv_buffer_emit_word(&b->instructions, result_type);
    spirv_buffer_emit_word(&b->instructions, result);
@@ -367,7 +370,7 @@ spirv_builder_emit_unop(struct spirv_builder *b, SpvOp op, SpvId result_type,
                         SpvId operand)
 {
    SpvId result = spirv_builder_new_id(b);
-   spirv_buffer_prepare(&b->instructions, 4);
+   spirv_buffer_prepare(&b->instructions, b->mem_ctx, 4);
    spirv_buffer_emit_word(&b->instructions, op | (4 << 16));
    spirv_buffer_emit_word(&b->instructions, result_type);
    spirv_buffer_emit_word(&b->instructions, result);
@@ -380,7 +383,7 @@ spirv_builder_emit_binop(struct spirv_builder *b, SpvOp op, SpvId result_type,
                          SpvId operand0, SpvId operand1)
 {
    SpvId result = spirv_builder_new_id(b);
-   spirv_buffer_prepare(&b->instructions, 5);
+   spirv_buffer_prepare(&b->instructions, b->mem_ctx, 5);
    spirv_buffer_emit_word(&b->instructions, op | (5 << 16));
    spirv_buffer_emit_word(&b->instructions, result_type);
    spirv_buffer_emit_word(&b->instructions, result);
@@ -394,7 +397,7 @@ spirv_builder_emit_triop(struct spirv_builder *b, SpvOp op, SpvId result_type,
                          SpvId operand0, SpvId operand1, SpvId operand2)
 {
    SpvId result = spirv_builder_new_id(b);
-   spirv_buffer_prepare(&b->instructions, 6);
+   spirv_buffer_prepare(&b->instructions, b->mem_ctx, 6);
    spirv_buffer_emit_word(&b->instructions, op | (6 << 16));
    spirv_buffer_emit_word(&b->instructions, result_type);
    spirv_buffer_emit_word(&b->instructions, result);
@@ -413,7 +416,7 @@ spirv_builder_emit_composite_extract(struct spirv_builder *b, SpvId result_type,
 
    assert(num_indexes > 0);
    int words = 4 + num_indexes;
-   spirv_buffer_prepare(&b->instructions, words);
+   spirv_buffer_prepare(&b->instructions, b->mem_ctx, words);
    spirv_buffer_emit_word(&b->instructions,
                           SpvOpCompositeExtract | (words << 16));
    spirv_buffer_emit_word(&b->instructions, result_type);
@@ -434,7 +437,7 @@ spirv_builder_emit_composite_construct(struct spirv_builder *b,
 
    assert(num_constituents > 0);
    int words = 3 + num_constituents;
-   spirv_buffer_prepare(&b->instructions, words);
+   spirv_buffer_prepare(&b->instructions, b->mem_ctx, words);
    spirv_buffer_emit_word(&b->instructions,
                           SpvOpCompositeConstruct | (words << 16));
    spirv_buffer_emit_word(&b->instructions, result_type);
@@ -454,7 +457,7 @@ spirv_builder_emit_vector_shuffle(struct spirv_builder *b, SpvId result_type,
 
    assert(num_components > 0);
    int words = 5 + num_components;
-   spirv_buffer_prepare(&b->instructions, words);
+   spirv_buffer_prepare(&b->instructions, b->mem_ctx, words);
    spirv_buffer_emit_word(&b->instructions, SpvOpVectorShuffle | (words << 16));
    spirv_buffer_emit_word(&b->instructions, result_type);
    spirv_buffer_emit_word(&b->instructions, result);
@@ -473,7 +476,7 @@ spirv_builder_emit_vector_extract(struct spirv_builder *b, SpvId result_type,
    SpvId result = spirv_builder_new_id(b);
 
    int words = 5;
-   spirv_buffer_prepare(&b->instructions, words);
+   spirv_buffer_prepare(&b->instructions, b->mem_ctx, words);
    spirv_buffer_emit_word(&b->instructions, SpvOpVectorExtractDynamic | (words << 16));
    spirv_buffer_emit_word(&b->instructions, result_type);
    spirv_buffer_emit_word(&b->instructions, result);
@@ -491,7 +494,7 @@ spirv_builder_emit_vector_insert(struct spirv_builder *b, SpvId result_type,
    SpvId result = spirv_builder_new_id(b);
 
    int words = 6;
-   spirv_buffer_prepare(&b->instructions, words);
+   spirv_buffer_prepare(&b->instructions, b->mem_ctx, words);
    spirv_buffer_emit_word(&b->instructions, SpvOpVectorInsertDynamic | (words << 16));
    spirv_buffer_emit_word(&b->instructions, result_type);
    spirv_buffer_emit_word(&b->instructions, result);
@@ -504,7 +507,7 @@ spirv_builder_emit_vector_insert(struct spirv_builder *b, SpvId result_type,
 void
 spirv_builder_emit_branch(struct spirv_builder *b, SpvId label)
 {
-   spirv_buffer_prepare(&b->instructions, 2);
+   spirv_buffer_prepare(&b->instructions, b->mem_ctx, 2);
    spirv_buffer_emit_word(&b->instructions, SpvOpBranch | (2 << 16));
    spirv_buffer_emit_word(&b->instructions, label);
 }
@@ -513,7 +516,7 @@ void
 spirv_builder_emit_selection_merge(struct spirv_builder *b, SpvId merge_block,
                                    SpvSelectionControlMask selection_control)
 {
-   spirv_buffer_prepare(&b->instructions, 3);
+   spirv_buffer_prepare(&b->instructions, b->mem_ctx, 3);
    spirv_buffer_emit_word(&b->instructions, SpvOpSelectionMerge | (3 << 16));
    spirv_buffer_emit_word(&b->instructions, merge_block);
    spirv_buffer_emit_word(&b->instructions, selection_control);
@@ -523,7 +526,7 @@ void
 spirv_builder_loop_merge(struct spirv_builder *b, SpvId merge_block,
                          SpvId cont_target, SpvLoopControlMask loop_control)
 {
-   spirv_buffer_prepare(&b->instructions, 4);
+   spirv_buffer_prepare(&b->instructions, b->mem_ctx, 4);
    spirv_buffer_emit_word(&b->instructions, SpvOpLoopMerge | (4 << 16));
    spirv_buffer_emit_word(&b->instructions, merge_block);
    spirv_buffer_emit_word(&b->instructions, cont_target);
@@ -534,7 +537,7 @@ void
 spirv_builder_emit_branch_conditional(struct spirv_builder *b, SpvId condition,
                                       SpvId true_label, SpvId false_label)
 {
-   spirv_buffer_prepare(&b->instructions, 4);
+   spirv_buffer_prepare(&b->instructions, b->mem_ctx, 4);
    spirv_buffer_emit_word(&b->instructions, SpvOpBranchConditional | (4 << 16));
    spirv_buffer_emit_word(&b->instructions, condition);
    spirv_buffer_emit_word(&b->instructions, true_label);
@@ -549,7 +552,7 @@ spirv_builder_emit_phi(struct spirv_builder *b, SpvId result_type,
 
    assert(num_vars > 0);
    int words = 3 + 2 * num_vars;
-   spirv_buffer_prepare(&b->instructions, words);
+   spirv_buffer_prepare(&b->instructions, b->mem_ctx, words);
    spirv_buffer_emit_word(&b->instructions, SpvOpPhi | (words << 16));
    spirv_buffer_emit_word(&b->instructions, result_type);
    spirv_buffer_emit_word(&b->instructions, result);
@@ -570,7 +573,7 @@ spirv_builder_set_phi_operand(struct spirv_builder *b, size_t position,
 void
 spirv_builder_emit_kill(struct spirv_builder *b)
 {
-   spirv_buffer_prepare(&b->instructions, 1);
+   spirv_buffer_prepare(&b->instructions, b->mem_ctx, 1);
    spirv_buffer_emit_word(&b->instructions, SpvOpKill | (1 << 16));
 }
 
@@ -626,7 +629,7 @@ spirv_builder_emit_image_sample(struct spirv_builder *b,
       num_extra_operands++;
    }
 
-   spirv_buffer_prepare(&b->instructions, operands + num_extra_operands);
+   spirv_buffer_prepare(&b->instructions, b->mem_ctx, operands + num_extra_operands);
    spirv_buffer_emit_word(&b->instructions, opcode | ((operands + num_extra_operands) << 16));
    spirv_buffer_emit_word(&b->instructions, result_type);
    spirv_buffer_emit_word(&b->instructions, result);
@@ -644,7 +647,7 @@ spirv_builder_emit_image(struct spirv_builder *b, SpvId result_type,
                          SpvId sampled_image)
 {
    SpvId result = spirv_builder_new_id(b);
-   spirv_buffer_prepare(&b->instructions, 4);
+   spirv_buffer_prepare(&b->instructions, b->mem_ctx, 4);
    spirv_buffer_emit_word(&b->instructions, SpvOpImage | (4 << 16));
    spirv_buffer_emit_word(&b->instructions, result_type);
    spirv_buffer_emit_word(&b->instructions, result);
@@ -680,7 +683,7 @@ spirv_builder_emit_image_fetch(struct spirv_builder *b,
       num_extra_operands++;
    }
 
-   spirv_buffer_prepare(&b->instructions, 5 + num_extra_operands);
+   spirv_buffer_prepare(&b->instructions, b->mem_ctx, 5 + num_extra_operands);
    spirv_buffer_emit_word(&b->instructions, SpvOpImageFetch |
                           ((5 + num_extra_operands) << 16));
    spirv_buffer_emit_word(&b->instructions, result_type);
@@ -706,7 +709,7 @@ spirv_builder_emit_image_query_size(struct spirv_builder *b,
    }
 
    SpvId result = spirv_builder_new_id(b);
-   spirv_buffer_prepare(&b->instructions, words);
+   spirv_buffer_prepare(&b->instructions, b->mem_ctx, words);
    spirv_buffer_emit_word(&b->instructions, opcode | (words << 16));
    spirv_buffer_emit_word(&b->instructions, result_type);
    spirv_buffer_emit_word(&b->instructions, result);
@@ -726,7 +729,7 @@ spirv_builder_emit_ext_inst(struct spirv_builder *b, SpvId result_type,
    SpvId result = spirv_builder_new_id(b);
 
    int words = 5 + num_args;
-   spirv_buffer_prepare(&b->instructions, words);
+   spirv_buffer_prepare(&b->instructions, b->mem_ctx, words);
    spirv_buffer_emit_word(&b->instructions, SpvOpExtInst | (words << 16));
    spirv_buffer_emit_word(&b->instructions, result_type);
    spirv_buffer_emit_word(&b->instructions, result);
@@ -803,12 +806,13 @@ get_type_def(struct spirv_builder *b, SpvOp op, const uint32_t args[],
       if (entry)
          return ((struct spirv_type *)entry->data)->type;
    } else {
-      b->types = _mesa_hash_table_create(NULL, non_aggregate_type_hash,
+      b->types = _mesa_hash_table_create(b->mem_ctx,
+                                         non_aggregate_type_hash,
                                          non_aggregate_type_equals);
       assert(b->types);
    }
 
-   struct spirv_type *type = CALLOC_STRUCT(spirv_type);
+   struct spirv_type *type = rzalloc(b->mem_ctx, struct spirv_type);
    if (!type)
       return 0;
 
@@ -817,7 +821,7 @@ get_type_def(struct spirv_builder *b, SpvOp op, const uint32_t args[],
    type->num_args = num_args;
 
    type->type = spirv_builder_new_id(b);
-   spirv_buffer_prepare(&b->types_const_defs, 2 + num_args);
+   spirv_buffer_prepare(&b->types_const_defs, b->mem_ctx, 2 + num_args);
    spirv_buffer_emit_word(&b->types_const_defs, op | ((2 + num_args) << 16));
    spirv_buffer_emit_word(&b->types_const_defs, type->type);
    for (int i = 0; i < num_args; ++i)
@@ -904,7 +908,7 @@ spirv_builder_type_array(struct spirv_builder *b, SpvId component_type,
                          SpvId length)
 {
    SpvId type = spirv_builder_new_id(b);
-   spirv_buffer_prepare(&b->types_const_defs, 4);
+   spirv_buffer_prepare(&b->types_const_defs, b->mem_ctx, 4);
    spirv_buffer_emit_word(&b->types_const_defs, SpvOpTypeArray | (4 << 16));
    spirv_buffer_emit_word(&b->types_const_defs, type);
    spirv_buffer_emit_word(&b->types_const_defs, component_type);
@@ -918,7 +922,7 @@ spirv_builder_type_struct(struct spirv_builder *b, const SpvId member_types[],
 {
    int words = 2 + num_member_types;
    SpvId type = spirv_builder_new_id(b);
-   spirv_buffer_prepare(&b->types_const_defs, words);
+   spirv_buffer_prepare(&b->types_const_defs, b->mem_ctx, words);
    spirv_buffer_emit_word(&b->types_const_defs, SpvOpTypeStruct | (words << 16));
    spirv_buffer_emit_word(&b->types_const_defs, type);
    for (int i = 0; i < num_member_types; ++i)
@@ -933,7 +937,7 @@ spirv_builder_type_function(struct spirv_builder *b, SpvId return_type,
 {
    int words = 3 + num_parameter_types;
    SpvId type = spirv_builder_new_id(b);
-   spirv_buffer_prepare(&b->types_const_defs, words);
+   spirv_buffer_prepare(&b->types_const_defs, b->mem_ctx, words);
    spirv_buffer_emit_word(&b->types_const_defs, SpvOpTypeFunction | (words << 16));
    spirv_buffer_emit_word(&b->types_const_defs, type);
    spirv_buffer_emit_word(&b->types_const_defs, return_type);
@@ -992,11 +996,12 @@ get_const_def(struct spirv_builder *b, SpvOp op, SpvId type,
       if (entry)
          return ((struct spirv_const *)entry->data)->result;
    } else {
-      b->consts = _mesa_hash_table_create(NULL, const_hash, const_equals);
+      b->consts = _mesa_hash_table_create(b->mem_ctx, const_hash,
+                                          const_equals);
       assert(b->consts);
    }
 
-   struct spirv_const *cnst = CALLOC_STRUCT(spirv_const);
+   struct spirv_const *cnst = rzalloc(b->mem_ctx, struct spirv_const);
    if (!cnst)
       return 0;
 
@@ -1006,7 +1011,7 @@ get_const_def(struct spirv_builder *b, SpvOp op, SpvId type,
    cnst->num_args = num_args;
 
    cnst->result = spirv_builder_new_id(b);
-   spirv_buffer_prepare(&b->types_const_defs, 3 + num_args);
+   spirv_buffer_prepare(&b->types_const_defs, b->mem_ctx, 3 + num_args);
    spirv_buffer_emit_word(&b->types_const_defs, op | ((3 + num_args) << 16));
    spirv_buffer_emit_word(&b->types_const_defs, type);
    spirv_buffer_emit_word(&b->types_const_defs, cnst->result);
@@ -1072,7 +1077,7 @@ spirv_builder_emit_var(struct spirv_builder *b, SpvId type,
                               &b->types_const_defs : &b->instructions;
 
    SpvId ret = spirv_builder_new_id(b);
-   spirv_buffer_prepare(buf, 4);
+   spirv_buffer_prepare(buf, b->mem_ctx, 4);
    spirv_buffer_emit_word(buf, SpvOpVariable | (4 << 16));
    spirv_buffer_emit_word(buf, type);
    spirv_buffer_emit_word(buf, ret);
@@ -1085,10 +1090,10 @@ spirv_builder_import(struct spirv_builder *b, const char *name)
 {
    SpvId result = spirv_builder_new_id(b);
    size_t pos = b->imports.num_words;
-   spirv_buffer_prepare(&b->imports, 2);
+   spirv_buffer_prepare(&b->imports, b->mem_ctx, 2);
    spirv_buffer_emit_word(&b->imports, SpvOpExtInstImport);
    spirv_buffer_emit_word(&b->imports, result);
-   int len = spirv_buffer_emit_string(&b->imports, name);
+   int len = spirv_buffer_emit_string(&b->imports, b->mem_ctx, name);
    b->imports.words[pos] |= (2 + len) << 16;
    return result;
 }
diff --git a/src/gallium/drivers/zink/nir_to_spirv/spirv_builder.h b/src/gallium/drivers/zink/nir_to_spirv/spirv_builder.h
index 00000cf4e99..7b870e94808 100644
--- a/src/gallium/drivers/zink/nir_to_spirv/spirv_builder.h
+++ b/src/gallium/drivers/zink/nir_to_spirv/spirv_builder.h
@@ -39,6 +39,8 @@ struct spirv_buffer {
 };
 
 struct spirv_builder {
+   void *mem_ctx;
+
    struct spirv_buffer capabilities;
    struct spirv_buffer imports;
    struct spirv_buffer memory_model;