spirv: add support for doubles to OpSpecConstant

v2 (Jason):
- Fix indent in radv change
- Add vtn_u64_literal() helper to take 64 bits (Jason)

Signed-off-by: Samuel Iglesias Gonsálvez <[email protected]>
Reviewed-by: Jason Ekstrand <[email protected]>

5 files changed, 55 insertions, 8 deletions

diff --git a/src/amd/vulkan/radv_pipeline.c b/src/amd/vulkan/radv_pipeline.c
index 75785ec921d..efb9cacf924 100644
--- a/src/amd/vulkan/radv_pipeline.c
+++ b/src/amd/vulkan/radv_pipeline.c
@@ -188,7 +188,10 @@ radv_shader_compile_to_nir(struct radv_device *device,
 				assert(data + entry.size <= spec_info->pData + spec_info->dataSize);
 
 				spec_entries[i].id = spec_info->pMapEntries[i].constantID;
-				spec_entries[i].data = *(const uint32_t *)data;
+				if (spec_info->dataSize == 8)
+					spec_entries[i].data64 = *(const uint64_t *)data;
+				else
+					spec_entries[i].data32 = *(const uint32_t *)data;
 			}
 		}
 		const struct nir_spirv_supported_extensions supported_ext = {
diff --git a/src/compiler/spirv/nir_spirv.h b/src/compiler/spirv/nir_spirv.h
index e0112ef4ad4..116b0a374cd 100644
--- a/src/compiler/spirv/nir_spirv.h
+++ b/src/compiler/spirv/nir_spirv.h
@@ -38,7 +38,10 @@ extern "C" {
 
 struct nir_spirv_specialization {
    uint32_t id;
-   uint32_t data;
+   union {
+      uint32_t data32;
+      uint64_t data64;
+   };
 };
 
 struct nir_spirv_supported_extensions {
diff --git a/src/compiler/spirv/spirv_to_nir.c b/src/compiler/spirv/spirv_to_nir.c
index f203ebc7ef7..b67189e07a6 100644
--- a/src/compiler/spirv/spirv_to_nir.c
+++ b/src/compiler/spirv/spirv_to_nir.c
@@ -31,6 +31,14 @@
 #include "nir/nir_constant_expressions.h"
 #include "spirv_info.h"
 
+struct spec_constant_value {
+   bool is_double;
+   union {
+      uint32_t data32;
+      uint64_t data64;
+   };
+};
+
 void
 _vtn_warn(const char *file, int line, const char *msg, ...)
 {
@@ -942,11 +950,14 @@ spec_constant_decoration_cb(struct vtn_builder *b, struct vtn_value *v,
    if (dec->decoration != SpvDecorationSpecId)
       return;
 
-   uint32_t *const_value = data;
+   struct spec_constant_value *const_value = data;
 
    for (unsigned i = 0; i < b->num_specializations; i++) {
       if (b->specializations[i].id == dec->literals[0]) {
-         *const_value = b->specializations[i].data;
+         if (const_value->is_double)
+            const_value->data64 = b->specializations[i].data64;
+         else
+            const_value->data32 = b->specializations[i].data32;
          return;
       }
    }
@@ -956,8 +967,22 @@ static uint32_t
 get_specialization(struct vtn_builder *b, struct vtn_value *val,
                    uint32_t const_value)
 {
-   vtn_foreach_decoration(b, val, spec_constant_decoration_cb, &const_value);
-   return const_value;
+   struct spec_constant_value data;
+   data.is_double = false;
+   data.data32 = const_value;
+   vtn_foreach_decoration(b, val, spec_constant_decoration_cb, &data);
+   return data.data32;
+}
+
+static uint64_t
+get_specialization64(struct vtn_builder *b, struct vtn_value *val,
+                   uint64_t const_value)
+{
+   struct spec_constant_value data;
+   data.is_double = true;
+   data.data64 = const_value;
+   vtn_foreach_decoration(b, val, spec_constant_decoration_cb, &data);
+   return data.data64;
 }
 
 static void
@@ -1017,10 +1042,17 @@ vtn_handle_constant(struct vtn_builder *b, SpvOp opcode,
       }
       break;
    }
-   case SpvOpSpecConstant:
+   case SpvOpSpecConstant: {
       assert(glsl_type_is_scalar(val->const_type));
       val->constant->values[0].u32[0] = get_specialization(b, val, w[3]);
+      int bit_size = glsl_get_bit_size(val->const_type);
+      if (bit_size == 64)
+         val->constant->values[0].u64[0] =
+            get_specialization64(b, val, vtn_u64_literal(&w[3]));
+      else
+         val->constant->values[0].u32[0] = get_specialization(b, val, w[3]);
       break;
+   }
    case SpvOpSpecConstantComposite:
    case SpvOpConstantComposite: {
       unsigned elem_count = count - 3;
diff --git a/src/compiler/spirv/vtn_private.h b/src/compiler/spirv/vtn_private.h
index ffa00d7f68a..7471ca1e21a 100644
--- a/src/compiler/spirv/vtn_private.h
+++ b/src/compiler/spirv/vtn_private.h
@@ -489,3 +489,9 @@ void vtn_handle_alu(struct vtn_builder *b, SpvOp opcode,
 
 bool vtn_handle_glsl450_instruction(struct vtn_builder *b, uint32_t ext_opcode,
                                     const uint32_t *words, unsigned count);
+
+static inline uint64_t
+vtn_u64_literal(const uint32_t *w)
+{
+   return (uint64_t)w[1] << 32 | w[0];
+}
diff --git a/src/intel/vulkan/anv_pipeline.c b/src/intel/vulkan/anv_pipeline.c
index 17491e34fc0..d5a7406c66c 100644
--- a/src/intel/vulkan/anv_pipeline.c
+++ b/src/intel/vulkan/anv_pipeline.c
@@ -117,7 +117,10 @@ anv_shader_compile_to_nir(struct anv_device *device,
          assert(data + entry.size <= spec_info->pData + spec_info->dataSize);
 
          spec_entries[i].id = spec_info->pMapEntries[i].constantID;
-         spec_entries[i].data = *(const uint32_t *)data;
+         if (spec_info->dataSize == 8)
+            spec_entries[i].data64 = *(const uint64_t *)data;
+         else
+            spec_entries[i].data32 = *(const uint32_t *)data;
       }
    }