spirv,treewide: rework specialization constant

With SPV_KHR_constant_data, it's allowed to specialize array of
constants.

RustiCL changes are from Karol Herbst <kherbst@redhat.com>.

Signed-off-by: Samuel Pitoiset <samuel.pitoiset@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/41046>

src/amd/vulkan/radv_shader.c

@@ -496,9 +496,7 @@ radv_shader_spirv_to_nir(struct radv_device *device, struct radv_shader_stage *s
          if (instance->debug_flags & dump_flags)
             spirv_print_asm(stderr, (const uint32_t *)stage->spirv.data, stage->spirv.size / 4);
       }
-
-      uint32_t num_spec_entries = 0;
-      struct nir_spirv_specialization *spec_entries = vk_spec_info_to_nir_spirv(stage->spec_info, &num_spec_entries);
+      struct nir_spirv_specialization *spec = vk_spec_info_to_nir_spirv(stage->spec_info);
       struct radv_shader_debug_data spirv_debug_data = {
          .debug_report = &instance->vk.debug_report,
          .object = stage->spirv.object,
@@ -537,13 +535,13 @@ radv_shader_spirv_to_nir(struct radv_device *device, struct radv_shader_stage *s
          .buffer_descriptor_size = pdev->vk.properties.bufferDescriptorSize,
          .buffer_descriptor_alignment = pdev->vk.properties.bufferDescriptorAlignment,
       };
-      nir = spirv_to_nir(spirv, stage->spirv.size / 4, spec_entries, num_spec_entries, stage->stage, stage->entrypoint,
-                         &spirv_options, &pdev->nir_options[stage->stage]);
+      nir = spirv_to_nir(spirv, stage->spirv.size / 4, spec, stage->stage, stage->entrypoint, &spirv_options,
+                         &pdev->nir_options[stage->stage]);
       nir->info.internal |= is_internal;
       assert(nir->info.stage == stage->stage);
       nir_validate_shader(nir, "after spirv_to_nir");
 
-      free(spec_entries);
+      vtn_free_specialization(spec);
 
       const struct nir_lower_sysvals_to_varyings_options sysvals_to_varyings = {
          .point_coord = true,

src/asahi/clc/asahi_clc.c

@@ -88,8 +88,8 @@ compile(void *memctx, const uint32_t *spirv, size_t spirv_size)
 
    assert(spirv_size % 4 == 0);
    nir_shader *nir =
-      spirv_to_nir(spirv, spirv_size / 4, NULL, 0, MESA_SHADER_KERNEL,
-                   "library", &spirv_options, nir_options);
+      spirv_to_nir(spirv, spirv_size / 4, NULL, MESA_SHADER_KERNEL, "library",
+                   &spirv_options, nir_options);
    nir_validate_shader(nir, "after spirv_to_nir");
    ralloc_steal(memctx, nir);
 

src/compiler/clc/nir_load_libclc.c

@@ -368,7 +368,7 @@ nir_load_libclc_shader(unsigned ptr_bit_size,
 
    assert(clc.size % SPIRV_WORD_SIZE == 0);
    nir_shader *nir = spirv_to_nir(clc.data, clc.size / SPIRV_WORD_SIZE,
-                                  NULL, 0, MESA_SHADER_KERNEL, NULL,
+                                  NULL, MESA_SHADER_KERNEL, NULL,
                                   &spirv_lib_options, nir_options);
    nir_validate_shader(nir, "after nir_load_clc_shader");
 

src/compiler/spirv/gl_spirv.c

@@ -62,9 +62,15 @@ spec_constant_decoration_cb(struct vtn_builder *b, struct vtn_value *v,
    if (dec->decoration != SpvDecorationSpecId)
       return;
 
-   for (unsigned i = 0; i < b->num_specializations; i++) {
-      if (b->specializations[i].id == dec->operands[0]) {
-         b->specializations[i].defined_on_module = true;
+   if (!b->specialization)
+      return;
+
+   for (unsigned i = 0; i < b->specialization->num_entries; i++) {
+      struct nir_spirv_specialization_entry *entry =
+         &b->specialization->entries[i];
+
+      if (entry->id == dec->operands[0]) {
+         entry->defined_on_module = true;
          return;
       }
    }
@@ -209,7 +215,7 @@ vtn_validate_handle_constant_instruction(struct vtn_builder *b, SpvOp opcode,
 enum spirv_verify_result
 spirv_verify_gl_specialization_constants(
    const uint32_t *words, size_t word_count,
-   struct nir_spirv_specialization *spec, unsigned num_spec,
+   struct nir_spirv_specialization *spec,
    mesa_shader_stage stage, const char *entry_point_name)
 {
    /* vtn_warn/vtn_log uses debug.func. Setting a null to prevent crash. Not
@@ -248,8 +254,7 @@ spirv_verify_gl_specialization_constants(
       return SPIRV_VERIFY_ENTRY_POINT_NOT_FOUND;
    }
 
-   b->specializations = spec;
-   b->num_specializations = num_spec;
+   b->specialization = spec;
 
    /* Handle constant instructions (we don't need to handle
     * variables or types for gl_spirv)
@@ -259,9 +264,12 @@ spirv_verify_gl_specialization_constants(
 
    ralloc_free(b);
 
-   for (unsigned i = 0; i < num_spec; i++) {
-      if (!spec[i].defined_on_module)
-         return SPIRV_VERIFY_UNKNOWN_SPEC_INDEX;
+   if (spec) {
+      for (unsigned i = 0; i < spec->num_entries; i++) {
+         const struct nir_spirv_specialization_entry *entry = &spec->entries[i];
+         if (!entry->defined_on_module)
+            return SPIRV_VERIFY_UNKNOWN_SPEC_INDEX;
+      }
    }
 
    return SPIRV_VERIFY_OK;

src/compiler/spirv/nir_spirv.h

@@ -16,12 +16,18 @@ extern "C" {
 
 struct spirv_capabilities;
 
-struct nir_spirv_specialization {
+struct nir_spirv_specialization_entry {
    uint32_t id;
-   nir_const_value value;
+   uint32_t size;
+   uint8_t *data;
    bool defined_on_module;
 };
 
+struct nir_spirv_specialization {
+   uint32_t num_entries;
+   struct nir_spirv_specialization_entry *entries;
+};
+
 enum nir_spirv_debug_level {
    NIR_SPIRV_DEBUG_LEVEL_INVALID = -1,
    NIR_SPIRV_DEBUG_LEVEL_INFO,
@@ -150,18 +156,27 @@ enum spirv_verify_result {
 
 enum spirv_verify_result spirv_verify_gl_specialization_constants(
    const uint32_t *words, size_t word_count,
-   struct nir_spirv_specialization *spec, unsigned num_spec,
+   struct nir_spirv_specialization *spec,
    mesa_shader_stage stage, const char *entry_point_name);
 
 nir_shader *spirv_to_nir(const uint32_t *words, size_t word_count,
-                         struct nir_spirv_specialization *specializations,
-                         unsigned num_specializations,
+                         struct nir_spirv_specialization *specialization,
                          mesa_shader_stage stage, const char *entry_point_name,
                          const struct spirv_to_nir_options *options,
                          const nir_shader_compiler_options *nir_options);
 
 void spirv_print_asm(FILE *fp, const uint32_t *words, size_t word_count);
 
+struct nir_spirv_specialization *
+vtn_alloc_specialization(uint32_t num_entries);
+
+bool
+vtn_add_specialization_entry(struct nir_spirv_specialization *spec, uint32_t slot,
+                             uint32_t entry_id, uint32_t entry_size,
+                             const void *entry_data, bool defined_on_module);
+
+void vtn_free_specialization(struct nir_spirv_specialization *spec);
+
 #ifdef __cplusplus
 }
 #endif

src/compiler/spirv/spirv2nir.c

@@ -270,7 +270,7 @@ int main(int argc, char **argv)
    if (entry_point.stage == MESA_SHADER_KERNEL)
       spirv_opts.environment = NIR_SPIRV_OPENCL;
 
-   nir_shader *nir = spirv_to_nir(map, word_count, NULL, 0,
+   nir_shader *nir = spirv_to_nir(map, word_count, NULL,
                                   entry_point.stage, entry_point.name,
                                   &spirv_opts, &nir_opts);
 

src/compiler/spirv/spirv_to_nir.c

@@ -2546,10 +2546,15 @@ spec_constant_decoration_cb(struct vtn_builder *b, UNUSED struct vtn_value *val,
    if (dec->decoration != SpvDecorationSpecId)
       return;
 
-   nir_const_value *value = data;
-   for (unsigned i = 0; i < b->num_specializations; i++) {
-      if (b->specializations[i].id == dec->operands[0]) {
-         *value = b->specializations[i].value;
+   if (!b->specialization)
+      return;
+
+   for (unsigned i = 0; i < b->specialization->num_entries; i++) {
+      const struct nir_spirv_specialization_entry *entry =
+         &b->specialization->entries[i];
+
+      if (entry->id == dec->operands[0] && entry->size > 0) {
+         memcpy(data, entry->data, entry->size);
          return;
       }
    }
@@ -7481,7 +7486,7 @@ can_remove(nir_variable *var, void *data)
 
 nir_shader *
 spirv_to_nir(const uint32_t *words, size_t word_count,
-             struct nir_spirv_specialization *spec, unsigned num_spec,
+             struct nir_spirv_specialization *spec,
              mesa_shader_stage stage, const char *entry_point_name,
              const struct spirv_to_nir_options *options,
              const nir_shader_compiler_options *nir_options)
@@ -7574,8 +7579,7 @@ spirv_to_nir(const uint32_t *words, size_t word_count,
       ralloc_free(b->shader);
       ralloc_free(b);
       nir_shader* result = spirv_to_nir(replacement_words, replacement_size / sizeof(uint32_t),
-                                        spec, num_spec,
-                                        stage, entry_point_name, options,
+                                        spec, stage, entry_point_name, options,
                                         nir_options);
 
       free((void *)replacement_words);
@@ -7636,8 +7640,7 @@ spirv_to_nir(const uint32_t *words, size_t word_count,
       vtn_foreach_execution_mode(b, b->entry_point,
                                  vtn_handle_execution_mode, NULL);
 
-   b->specializations = spec;
-   b->num_specializations = num_spec;
+   b->specialization = spec;
 
    /* Handle all variable, type, and constant instructions */
    words = vtn_foreach_instruction(b, words, word_end,
@@ -7909,3 +7912,61 @@ vtn_dump_values(struct vtn_builder *b, FILE *f)
    }
    fprintf(f, "===\n");
 }
+
+struct nir_spirv_specialization *
+vtn_alloc_specialization(uint32_t num_entries)
+{
+   struct nir_spirv_specialization *spec;
+
+   spec = calloc(1, sizeof(*spec));
+   if (!spec)
+      return NULL;
+
+   spec->entries = calloc(num_entries, sizeof(*spec->entries));
+   if (!spec->entries) {
+      free(spec);
+      return NULL;
+   }
+
+   spec->num_entries = num_entries;
+   return spec;
+}
+
+bool
+vtn_add_specialization_entry(struct nir_spirv_specialization *spec, uint32_t slot,
+                             uint32_t entry_id, uint32_t entry_size,
+                             const void *entry_data, bool defined_on_module)
+{
+   struct nir_spirv_specialization_entry *entry = &spec->entries[slot];
+
+   assert(!entry->data);
+
+   entry->id = entry_id;
+   entry->size = entry_size;
+   entry->defined_on_module = defined_on_module;
+
+   if (entry_size > 0) {
+      entry->data = malloc(entry_size);
+      if (!entry->data)
+         return false;
+
+      memcpy(entry->data, entry_data, entry_size);
+   }
+
+   return true;
+}
+
+void
+vtn_free_specialization(struct nir_spirv_specialization *spec)
+{
+   if (!spec)
+      return;
+
+   for (uint32_t i = 0; i < spec->num_entries; i++) {
+      struct nir_spirv_specialization_entry *entry = &spec->entries[i];
+      free(entry->data);
+   }
+
+   free(spec->entries);
+   free(spec);
+}

src/compiler/spirv/tests/helpers.h

@@ -44,7 +44,7 @@ protected:
 
    void get_nir(size_t num_words, const uint32_t *words, mesa_shader_stage stage = MESA_SHADER_COMPUTE)
    {
-      shader = spirv_to_nir(words, num_words, NULL, 0,
+      shader = spirv_to_nir(words, num_words, NULL,
                             stage, "main", &spirv_options, &nir_options);
    }
 

src/compiler/spirv/vtn_bindgen2.c

@@ -123,7 +123,7 @@ compile(void *memctx, const uint32_t *spirv, size_t spirv_size)
 
    assert(spirv_size % 4 == 0);
    nir_shader *nir =
-      spirv_to_nir(spirv, spirv_size / 4, NULL, 0, MESA_SHADER_KERNEL,
+      spirv_to_nir(spirv, spirv_size / 4, NULL, MESA_SHADER_KERNEL,
                    "library", &spirv_options, nir_options);
    nir_validate_shader(nir, "after spirv_to_nir");
    ralloc_steal(memctx, nir);

src/compiler/spirv/vtn_private.h

@@ -651,8 +651,7 @@ struct vtn_builder {
     */
    struct set *vars_used_indirectly;
 
-   unsigned num_specializations;
-   struct nir_spirv_specialization *specializations;
+   struct nir_spirv_specialization *specialization;
 
    unsigned value_id_bound;
    struct vtn_value *values;

src/freedreno/vulkan/tu_shader.cc

@@ -69,9 +69,7 @@ tu_spirv_to_nir_library(struct tu_device *dev,
    spirv_to_nir_options spirv_options = tu_spirv_options;
    spirv_options.create_library = true;
 
-   nir_shader *nir =
-      spirv_to_nir(words, word_count, NULL, 0, MESA_SHADER_COMPUTE,
-                   "main", &spirv_options, nir_options);
+   nir_shader *nir = spirv_to_nir(words, word_count, NULL, MESA_SHADER_COMPUTE, "main", &spirv_options, nir_options);
 
    NIR_PASS(_, nir, nir_lower_system_values);
 

src/gallium/drivers/zink/zink_compiler.c

@@ -3522,8 +3522,7 @@ zink_shader_spirv_compile(struct zink_screen *screen, struct zink_shader *zs, st
          .push_const_addr_format = nir_address_format_logical,
          .shared_addr_format = nir_address_format_32bit_offset,
       };
-      uint32_t num_spec_entries = 0;
-      struct nir_spirv_specialization *spec_entries = NULL;
+      struct nir_spirv_specialization *spec = NULL;
       VkSpecializationInfo sinfo = {0};
       VkSpecializationMapEntry me[3];
       uint32_t size[3] = {1,1,1};
@@ -3538,14 +3537,14 @@ zink_shader_spirv_compile(struct zink_screen *screen, struct zink_shader *zs, st
             me[i].constantID = ids[i];
             me[i].offset = i * sizeof(uint32_t);
          }
-         spec_entries = vk_spec_info_to_nir_spirv(&sinfo, &num_spec_entries);
+         spec = vk_spec_info_to_nir_spirv(&sinfo);
       }
       nir_shader *nir = spirv_to_nir(spirv->words, spirv->num_words,
-                         spec_entries, num_spec_entries,
-                         clamp_stage(&zs->info), "main", &spirv_options, &screen->nir_options);
+                         spec, clamp_stage(&zs->info), "main", &spirv_options,
+                         &screen->nir_options);
       assert(nir);
       ralloc_free(nir);
-      free(spec_entries);
+      vtn_free_specialization(spec);
    }
 #endif
 

src/gallium/frontends/rusticl/core/kernel.rs

@@ -1225,7 +1225,7 @@ pub(super) fn convert_spirv_to_nir(
     build: &DeviceProgramBuild,
     name: &str,
     args: &[spirv::SPIRVKernelArg],
-    spec_constants: &HashMap<u32, nir_const_value>,
+    spec_constants: &mut HashMap<u32, Vec<u8>>,
     dev: &'static Device,
 ) -> SPIRVToNirResult {
     let cache = dev.screen().shader_cache();

src/gallium/frontends/rusticl/core/program.rs

@@ -96,7 +96,7 @@ impl_cl_type_trait!(cl_program, Program, CL_INVALID_PROGRAM);
 pub struct ProgramBuild {
     pub builds_by_device: HashMap<&'static Device, DeviceProgramBuild>,
     pub kernel_info: HashMap<String, Arc<KernelInfo>>,
-    spec_constants: HashMap<u32, nir_const_value>,
+    spec_constants: HashMap<u32, Vec<u8>>,
     kernels: Vec<String>,
 }
 
@@ -138,7 +138,7 @@ impl ProgramBuild {
                 }
 
                 let build_result =
-                    convert_spirv_to_nir(build, kernel_name, &args, &self.spec_constants, dev);
+                    convert_spirv_to_nir(build, kernel_name, &args, &mut self.spec_constants, dev);
                 kernel_info_set.insert(build_result.kernel_info);
 
                 self.builds_by_device.get_mut(dev).unwrap().kernels.insert(
@@ -194,7 +194,7 @@ impl DeviceProgramBuild {
         &self,
         cache: Option<&DiskCacheBorrowed>,
         name: &str,
-        spec_constants: &HashMap<u32, nir_const_value>,
+        spec_constants: &HashMap<u32, Vec<u8>>,
     ) -> Option<cache_key> {
         if let Some(cache) = cache {
             assert_eq!(self.status, CL_BUILD_SUCCESS as cl_build_status);
@@ -205,10 +205,7 @@ impl DeviceProgramBuild {
 
             for (k, v) in spec_constants {
                 bin.extend_from_slice(&k.to_ne_bytes());
-                unsafe {
-                    // SAFETY: we fully initialize this union
-                    bin.extend_from_slice(&v.u64_.to_ne_bytes());
-                }
+                bin.extend_from_slice(v);
             }
 
             Some(cache.gen_key(&bin))
@@ -225,19 +222,25 @@ impl DeviceProgramBuild {
         &self,
         kernel: &str,
         device: &Device,
-        spec_constants: &HashMap<u32, nir_const_value>,
+        spec_constants: &mut HashMap<u32, Vec<u8>>,
     ) -> NirShader {
         assert_eq!(self.status, CL_BUILD_SUCCESS as cl_build_status);
 
         let mut spec_constants: Vec<_> = spec_constants
-            .iter()
-            .map(|(&id, &value)| nir_spirv_specialization {
+            .iter_mut()
+            .map(|(&id, value)| nir_spirv_specialization_entry {
                 id: id,
-                value: value,
+                size: value.len() as u32,
+                data: value.as_mut_ptr(),
                 defined_on_module: true,
             })
             .collect();
 
+        let mut spec_constants = nir_spirv_specialization {
+            num_entries: spec_constants.len() as u32,
+            entries: spec_constants.as_mut_ptr(),
+        };
+
         let mut log = Platform::dbg().program.then(Vec::new);
         let nir = self.spirv.as_ref().unwrap().to_nir(
             kernel,
@@ -912,17 +915,7 @@ impl Program {
 
     pub fn set_spec_constant(&self, spec_id: u32, data: &[u8]) {
         let mut lock = self.build_info();
-        let mut val = nir_const_value::default();
-
-        match data.len() {
-            1 => val.u8_ = u8::from_ne_bytes(data.try_into().unwrap()),
-            2 => val.u16_ = u16::from_ne_bytes(data.try_into().unwrap()),
-            4 => val.u32_ = u32::from_ne_bytes(data.try_into().unwrap()),
-            8 => val.u64_ = u64::from_ne_bytes(data.try_into().unwrap()),
-            _ => unreachable!("Spec constant with invalid size!"),
-        };
-
-        lock.spec_constants.insert(spec_id, val);
+        lock.spec_constants.insert(spec_id, data.to_owned());
     }
 }
 

src/gallium/frontends/rusticl/mesa/compiler/clc/spirv.rs

@@ -325,7 +325,7 @@ impl SPIRVBin {
         nir_options: *const nir_shader_compiler_options,
         spirv_caps: &spirv_capabilities,
         libclc: &NirShader,
-        spec_constants: &mut [nir_spirv_specialization],
+        spec_constants: &mut nir_spirv_specialization,
         address_bits: u32,
         log: Option<&mut Vec<String>>,
     ) -> Option<NirShader> {
@@ -337,8 +337,7 @@ impl SPIRVBin {
             spirv_to_nir(
                 self.spirv.data.cast(),
                 self.spirv.size / 4,
-                spec_constants.as_mut_ptr(),
-                spec_constants.len() as u32,
+                spec_constants,
                 mesa_shader_stage::MESA_SHADER_KERNEL,
                 c_entry.as_ptr(),
                 &spirv_options,

src/imagination/pco/uscgen/pco_clc.c

@@ -166,7 +166,6 @@ spv_to_nir(void *mem_ctx, uint32_t *spirv_map, unsigned spirv_len)
    nir_shader *nir = spirv_to_nir(spirv_map,
                                   spirv_len / 4,
                                   NULL,
-                                  0,
                                   MESA_SHADER_KERNEL,
                                   "library",
                                   &precomp_spirv_options,

src/intel/vulkan/anv_pipeline_cache.c

@@ -357,7 +357,7 @@ anv_load_fp64_shader(struct anv_device *device)
 
    nir_shader* nir =
       spirv_to_nir(float64_spv_source, sizeof(float64_spv_source) / 4,
-                   NULL, 0, MESA_SHADER_VERTEX, "main",
+                   NULL, MESA_SHADER_VERTEX, "main",
                    &spirv_options, nir_options);
 
    assert(nir != NULL);

src/kosmickrisp/clc/kk_clc.c

@@ -83,8 +83,8 @@ compile(void *memctx, const uint32_t *spirv, size_t spirv_size)
 
    assert(spirv_size % 4 == 0);
    nir_shader *nir =
-      spirv_to_nir(spirv, spirv_size / 4, NULL, 0, MESA_SHADER_KERNEL,
-                   "library", &spirv_options, nir_options);
+      spirv_to_nir(spirv, spirv_size / 4, NULL, MESA_SHADER_KERNEL, "library",
+                   &spirv_options, nir_options);
    nir_validate_shader(nir, "after spirv_to_nir");
    ralloc_steal(memctx, nir);
 

src/kosmickrisp/kosmicomp.c

@@ -172,8 +172,8 @@ main(int argc, char **argv)
       fprintf(stderr, "Couldn't guess shader stage from %s\n", argv[1]);
       return 4;
    }
-   nir_shader *shader = spirv_to_nir(words, nwords, NULL, 0, stage, "main",
-                                     &options, &nir_options);
+   nir_shader *shader =
+      spirv_to_nir(words, nwords, NULL, stage, "main", &options, &nir_options);
    if (!shader) {
       fprintf(stderr, "Compilation failed!\n");
       return 3;

src/mesa/main/glspirv.c

@@ -249,14 +249,14 @@ _mesa_spirv_to_nir(struct gl_context *ctx,
    const char *entry_point_name = spirv_data->SpirVEntryPoint;
    assert(entry_point_name);
 
-   struct nir_spirv_specialization *spec_entries =
-      calloc(sizeof(*spec_entries),
-             spirv_data->NumSpecializationConstants);
+   struct nir_spirv_specialization *spec =
+      vtn_alloc_specialization(spirv_data->NumSpecializationConstants);
 
    for (unsigned i = 0; i < spirv_data->NumSpecializationConstants; ++i) {
-      spec_entries[i].id = spirv_data->SpecializationConstantsIndex[i];
-      spec_entries[i].value.u32 = spirv_data->SpecializationConstantsValue[i];
-      spec_entries[i].defined_on_module = false;
+      vtn_add_specialization_entry(spec, i, spirv_data->SpecializationConstantsIndex[i],
+                                   sizeof(uint32_t),
+                                   &spirv_data->SpecializationConstantsValue[i],
+                                   false);
    }
 
    struct spirv_capabilities spirv_caps;
@@ -281,11 +281,11 @@ _mesa_spirv_to_nir(struct gl_context *ctx,
    nir_shader *nir =
       spirv_to_nir((const uint32_t *) &spirv_module->Binary[0],
                    spirv_module->Length / 4,
-                   spec_entries, spirv_data->NumSpecializationConstants,
+                   spec,
                    stage, entry_point_name,
                    &spirv_options,
                    options);
-   free(spec_entries);
+   vtn_free_specialization(spec);
 
    assert(nir);
    assert(nir->info.stage == stage);
@@ -350,7 +350,6 @@ _mesa_SpecializeShaderARB(GLuint shader,
 {
    GET_CURRENT_CONTEXT(ctx);
    struct gl_shader *sh;
-   struct nir_spirv_specialization *spec_entries = NULL;
 
    if (!ctx->Extensions.ARB_gl_spirv) {
       _mesa_error(ctx, GL_INVALID_OPERATION, "glSpecializeShaderARB");
@@ -396,19 +395,18 @@ _mesa_SpecializeShaderARB(GLuint shader,
     * parsing the module. However, flagging them during specialization is okay,
     * since it makes no difference in terms of application-visible state.
     */
-   spec_entries = calloc(sizeof(*spec_entries), numSpecializationConstants);
+   struct nir_spirv_specialization *spec =
+      vtn_alloc_specialization(numSpecializationConstants);
 
    for (unsigned i = 0; i < numSpecializationConstants; ++i) {
-      spec_entries[i].id = pConstantIndex[i];
-      spec_entries[i].value.u32 = pConstantValue[i];
-      spec_entries[i].defined_on_module = false;
+      vtn_add_specialization_entry(spec, i, pConstantIndex[i], sizeof(uint32_t),
+                                   &pConstantValue[i], false);
    }
 
    enum spirv_verify_result r = spirv_verify_gl_specialization_constants(
       (uint32_t *)&spirv_data->SpirVModule->Binary[0],
       spirv_data->SpirVModule->Length / 4,
-      spec_entries, numSpecializationConstants,
-      sh->Stage, pEntryPoint);
+      spec, sh->Stage, pEntryPoint);
 
    switch (r) {
    case SPIRV_VERIFY_OK:
@@ -425,10 +423,11 @@ _mesa_SpecializeShaderARB(GLuint shader,
       goto end;
    case SPIRV_VERIFY_UNKNOWN_SPEC_INDEX:
       for (unsigned i = 0; i < numSpecializationConstants; ++i) {
-         if (spec_entries[i].defined_on_module == false) {
+         const struct nir_spirv_specialization_entry *entry = &spec->entries[i];
+         if (entry->defined_on_module == false) {
             _mesa_error(ctx, GL_INVALID_VALUE,
                         "glSpecializeShaderARB(constant \"%i\" does not exist "
-                        "in shader)", spec_entries[i].id);
+                        "in shader)", entry->id);
             break;
          }
       }
@@ -456,5 +455,5 @@ _mesa_SpecializeShaderARB(GLuint shader,
    }
 
  end:
-   free(spec_entries);
+   vtn_free_specialization(spec);
 }

src/microsoft/clc/clc_compiler.c

@@ -793,12 +793,20 @@ clc_spirv_to_dxil(struct clc_libclc *lib,
 
    glsl_type_singleton_init_or_ref();
 
+   struct nir_spirv_specialization *spec = NULL;
+   if (consts && consts->num_specializations > 0) {
+      spec = vtn_alloc_specialization(consts->num_specializations);
+
+      for (unsigned i = 0; i < consts->num_specializations; i++) {
+         vtn_add_specialization_entry(spec, i, consts->specializations[i].id,
+                                      sizeof(consts->specializations[i].value),
+                                      &consts->specializations[i].value, false);
+      }
+   }
    nir = spirv_to_nir(linked_spirv->data, linked_spirv->size / 4,
-                      consts ? (struct nir_spirv_specialization *)consts->specializations : NULL,
-                      consts ? consts->num_specializations : 0,
-                      MESA_SHADER_KERNEL, entrypoint,
-                      &spirv_options,
-                      &nir_options);
+                      spec, MESA_SHADER_KERNEL, entrypoint,
+                      &spirv_options, &nir_options);
+   vtn_free_specialization(spec);
    if (!nir) {
       clc_error(logger, "spirv_to_nir() failed");
       goto err_free_dxil;

src/microsoft/spirv_to_dxil/spirv2dxil.c

@@ -102,7 +102,7 @@ compile_shader(const char *filename, mesa_shader_stage shader_stage, struct shad
 
    shader->nir = spirv_to_nir(
       (const uint32_t *)file_contents, word_count, NULL,
-      0, (mesa_shader_stage)shader_stage, shader->entry_point,
+      (mesa_shader_stage)shader_stage, shader->entry_point,
       spirv_opts, &nir_options);
    free(file_contents);
    if (!shader->nir) {

src/microsoft/spirv_to_dxil/spirv_to_dxil.c

@@ -71,10 +71,20 @@ spirv_to_dxil(const uint32_t *words, size_t word_count,
    // have been already converted to zero-base.
    nir_options.lower_base_vertex = conf->first_vertex_and_base_instance_mode != DXIL_SPIRV_SYSVAL_TYPE_ZERO;
 
-   nir_shader *nir = spirv_to_nir(
-      words, word_count, (struct nir_spirv_specialization *)specializations,
-      num_specializations, (mesa_shader_stage)stage, entry_point_name,
-      spirv_opts, &nir_options);
+   struct nir_spirv_specialization *spec = NULL;
+   if (specializations && num_specializations > 0) {
+      spec = vtn_alloc_specialization(num_specializations);
+
+      for (unsigned i = 0; i < num_specializations; i++) {
+         vtn_add_specialization_entry(spec, i, specializations[i].id,
+                                      sizeof(specializations[i].value),
+                                      &specializations[i].value, false);
+      }
+   }
+   nir_shader *nir = spirv_to_nir(words, word_count, spec,
+                                  (mesa_shader_stage)stage, entry_point_name,
+                                  spirv_opts, &nir_options);
+   vtn_free_specialization(spec);
    if (!nir) {
       glsl_type_singleton_decref();
       return false;

src/microsoft/spirv_to_dxil/spirv_to_dxil.h

@@ -64,7 +64,6 @@ typedef union {
    uint64_t u64;
 } dxil_spirv_const_value;
 
-// Copy of nir_spirv_specialization
 struct dxil_spirv_specialization {
    uint32_t id;
    dxil_spirv_const_value value;

src/panfrost/clc/pan_compile.c

@@ -99,8 +99,8 @@ compile(void *memctx, const uint32_t *spirv, size_t spirv_size, unsigned arch)
    const nir_shader_compiler_options *nir_options = get_compiler_options(arch);
 
    nir_shader *nir =
-      spirv_to_nir(spirv, spirv_size / 4, NULL, 0, MESA_SHADER_KERNEL,
-                   "library", &spirv_options, nir_options);
+      spirv_to_nir(spirv, spirv_size / 4, NULL, MESA_SHADER_KERNEL, "library",
+                   &spirv_options, nir_options);
    nir_validate_shader(nir, "after spirv_to_nir");
    nir_validate_ssa_dominance(nir, "after spirv_to_nir");
    ralloc_steal(memctx, nir);

src/vulkan/runtime/vk_nir.c

@@ -149,15 +149,13 @@ vk_spirv_to_nir(struct vk_device *device,
    spirv_options_local.buffer_descriptor_alignment =
       device->physical->properties.bufferDescriptorAlignment;
 
-   uint32_t num_spec_entries = 0;
-   struct nir_spirv_specialization *spec_entries =
-      vk_spec_info_to_nir_spirv(spec_info, &num_spec_entries);
+   struct nir_spirv_specialization *spec =
+      vk_spec_info_to_nir_spirv(spec_info);
 
    nir_shader *nir = spirv_to_nir(spirv_data, spirv_size_B / 4,
-                                  spec_entries, num_spec_entries,
-                                  stage, entrypoint_name,
+                                  spec, stage, entrypoint_name,
                                   &spirv_options_local, nir_options);
-   free(spec_entries);
+   vtn_free_specialization(spec);
 
    if (nir == NULL)
       return NULL;

src/vulkan/util/vk_util.c

@@ -84,63 +84,34 @@ vk_warn_non_conformant_implementation(const char *driver_name)
 }
 
 struct nir_spirv_specialization*
-vk_spec_info_to_nir_spirv(const VkSpecializationInfo *spec_info,
-                          uint32_t *out_num_spec_entries)
+vk_spec_info_to_nir_spirv(const VkSpecializationInfo *vk_spec_info)
 {
-   if (spec_info == NULL || spec_info->mapEntryCount == 0)
+   if (vk_spec_info == NULL || vk_spec_info->mapEntryCount == 0)
       return NULL;
 
-   uint32_t num_spec_entries = spec_info->mapEntryCount;
-   struct nir_spirv_specialization *spec_entries =
-      calloc(num_spec_entries, sizeof(*spec_entries));
+   struct nir_spirv_specialization *spec =
+      vtn_alloc_specialization(vk_spec_info->mapEntryCount);
+   if (!spec)
+      return NULL;
 
-   for (uint32_t i = 0; i < num_spec_entries; i++) {
-      VkSpecializationMapEntry entry = spec_info->pMapEntries[i];
-      const void *data = (uint8_t *)spec_info->pData + entry.offset;
-      assert((uint8_t *)data + entry.size <=
-             (uint8_t *)spec_info->pData + spec_info->dataSize);
+   for (uint32_t i = 0; i < vk_spec_info->mapEntryCount; i++) {
+      VkSpecializationMapEntry vk_entry = vk_spec_info->pMapEntries[i];
+      const void *vk_data = (uint8_t *)vk_spec_info->pData + vk_entry.offset;
 
-      spec_entries[i].id = spec_info->pMapEntries[i].constantID;
-      switch (entry.size) {
-      case 8:
-         memcpy(&spec_entries[i].value.u64, data, entry.size);
-         break;
-      case 4:
-         memcpy(&spec_entries[i].value.u32, data, entry.size);
-         break;
-      case 2:
-         memcpy(&spec_entries[i].value.u16, data, entry.size);
-         break;
-      case 1:
-         memcpy(&spec_entries[i].value.u8, data, entry.size);
-         break;
-      case 0:
-      default:
-         /* The Vulkan spec says:
-          *
-          *    "For a constantID specialization constant declared in a
-          *    shader, size must match the byte size of the constantID. If
-          *    the specialization constant is of type boolean, size must be
-          *    the byte size of VkBool32."
-          *
-          * Therefore, since only scalars can be decorated as
-          * specialization constants, we can assume that if it doesn't have
-          * a size of 1, 2, 4, or 8, any use in a shader would be invalid
-          * usage.  The spec further says:
-          *
-          *    "If a constantID value is not a specialization constant ID
-          *    used in the shader, that map entry does not affect the
-          *    behavior of the pipeline."
-          *
-          * so we should ignore any invalid specialization constants rather
-          * than crash or error out when we see one.
-          */
-         break;
-      }
+      assert((uint8_t *)vk_data + vk_entry.size <=
+             (uint8_t *)vk_spec_info->pData + vk_spec_info->dataSize);
+
+      if (!vtn_add_specialization_entry(spec, i,
+                                        vk_spec_info->pMapEntries[i].constantID,
+                                        vk_entry.size, vk_data, false))
+         goto fail;
    }
 
-   *out_num_spec_entries = num_spec_entries;
-   return spec_entries;
+   return spec;
+
+fail:
+   vtn_free_specialization(spec);
+   return NULL;
 }
 
 enum mesa_prim

src/vulkan/util/vk_util.h

@@ -338,8 +338,7 @@ mesa_to_vk_shader_stage(mesa_shader_stage mesa_stage)
 struct nir_spirv_specialization;
 
 struct nir_spirv_specialization*
-vk_spec_info_to_nir_spirv(const VkSpecializationInfo *spec_info,
-                          uint32_t *out_num_spec_entries);
+vk_spec_info_to_nir_spirv(const VkSpecializationInfo *vk_spec_info);
 
 static inline uint8_t
 vk_index_type_to_bytes(VkIndexType type)