Implemented HLSL vertex shader inputs mapping in Vulkan backend using SPV_GOOGLE_hlsl_functionality1 (causes validation layer errors for now)

3 files changed, 110 insertions, 14 deletions

diff --git a/Graphics/GLSLTools/include/SPIRVShaderResources.h b/Graphics/GLSLTools/include/SPIRVShaderResources.h
index 5c72d4a9..56c4253f 100644
--- a/Graphics/GLSLTools/include/SPIRVShaderResources.h
+++ b/Graphics/GLSLTools/include/SPIRVShaderResources.h
@@ -30,7 +30,7 @@
 //
 //   m_MemoryBuffer                                                                                                              m_TotalResources
 //    |                                                                                                                             |
-//    | Uniform Buffers | Storage Buffers | Storage Images | Sampled Images | Atomic Counters | Separate Samplers | Separate Images |  Immutable Samplers  |   Resource Names   |
+//    | Uniform Buffers | Storage Buffers | Storage Images | Sampled Images | Atomic Counters | Separate Samplers | Separate Images |  Immutable Samplers  |   Stage Inputs   |   Resource Names   |
 
 #include <memory>
 #include <vector>
@@ -190,6 +190,18 @@ public:
 };
 static_assert(sizeof(SPIRVShaderResourceAttribs) % sizeof(void*) == 0, "Size of SPIRVShaderResourceAttribs struct must be multiple of sizeof(void*)" );
 
+struct SPIRVShaderStageInputAttribs
+{
+    SPIRVShaderStageInputAttribs(const char* _Semantic, uint32_t _LocationDecorationOffset) :
+        Semantic                (_Semantic),
+        LocationDecorationOffset(_LocationDecorationOffset)
+    {}
+
+    const char* const Semantic;
+    const uint32_t LocationDecorationOffset;
+};
+static_assert(sizeof(SPIRVShaderStageInputAttribs) % sizeof(void*) == 0, "Size of SPIRVShaderStageInputAttribs struct must be multiple of sizeof(void*)" );
+
 /// Diligent::SPIRVShaderResources class
 class SPIRVShaderResources
 {
@@ -198,7 +210,8 @@ public:
                          IRenderDevice*         pRenderDevice,
                          std::vector<uint32_t>  spirv_binary,
                          const ShaderDesc&      shaderDesc,
-                         const char*            CombinedSamplerSuffix);
+                         const char*            CombinedSamplerSuffix,
+                         bool                   LoadShaderStageInputs);
 
     SPIRVShaderResources             (const SPIRVShaderResources&)  = delete;
     SPIRVShaderResources             (      SPIRVShaderResources&&) = delete;
@@ -218,6 +231,7 @@ public:
     Uint32 GetNumSepImgs  ()const noexcept{ return (m_TotalResources        - m_SeparateImageOffset);  }
     Uint32 GetTotalResources()      const noexcept { return m_TotalResources; }
     Uint32 GetNumImmutableSamplers()const noexcept { return m_NumImmutableSamplers; }
+    Uint32 GetNumShaderStageInputs()const noexcept { return m_NumShaderStageInputs; }
 
     const SPIRVShaderResourceAttribs& GetUB      (Uint32 n)const noexcept{ return GetResAttribs(n, GetNumUBs(),         0                      ); }
     const SPIRVShaderResourceAttribs& GetSB      (Uint32 n)const noexcept{ return GetResAttribs(n, GetNumSBs(),         m_StorageBufferOffset  ); }
@@ -239,6 +253,14 @@ public:
         return reinterpret_cast<SamplerPtrType*>(ResourceMemoryEnd)[ImmutableSamplerInd];
     }
 
+    const SPIRVShaderStageInputAttribs& GetShaderStageInputAttribs(Uint32 n)const noexcept
+    {
+        VERIFY(n < m_NumShaderStageInputs, "Shader stage input index (", n, ") is out of range. Total input count: ", m_NumShaderStageInputs);
+        auto* ResourceMemoryEnd = reinterpret_cast<const SPIRVShaderResourceAttribs*>(m_MemoryBuffer.get()) + m_TotalResources;
+        auto* ImmutableSamplerMemoryEnd = reinterpret_cast<const SamplerPtrType*>(ResourceMemoryEnd) + m_NumImmutableSamplers;
+        return reinterpret_cast<const SPIRVShaderStageInputAttribs*>(ImmutableSamplerMemoryEnd)[n];
+    }
+
     struct ResourceCounters
     {
         Uint32 NumUBs       = 0;
@@ -353,18 +375,19 @@ private:
     void Initialize(IMemoryAllocator&       Allocator, 
                     const ResourceCounters& Counters,
                     Uint32                  NumImmutableSamplers,
+                    Uint32                  NumShaderStageInputs,
                     size_t                  ResourceNamesPoolSize);
 
     __forceinline SPIRVShaderResourceAttribs& GetResAttribs(Uint32 n, Uint32 NumResources, Uint32 Offset)noexcept
     {
-        VERIFY(n < NumResources, "Resource index (", n, ") is out of range. Resource array size: ", NumResources);
+        VERIFY(n < NumResources, "Resource index (", n, ") is out of range. Total resource count: ", NumResources);
         VERIFY_EXPR(Offset + n < m_TotalResources);
         return reinterpret_cast<SPIRVShaderResourceAttribs*>(m_MemoryBuffer.get())[Offset + n];
     }
 
     __forceinline const SPIRVShaderResourceAttribs& GetResAttribs(Uint32 n, Uint32 NumResources, Uint32 Offset)const noexcept
     {
-        VERIFY(n < NumResources, "Resource index (", n, ") is out of range. Resource array size: ", NumResources);
+        VERIFY(n < NumResources, "Resource index (", n, ") is out of range. Total resource count: ", NumResources);
         VERIFY_EXPR(Offset + n < m_TotalResources);
         return reinterpret_cast<SPIRVShaderResourceAttribs*>(m_MemoryBuffer.get())[Offset + n];
     }
@@ -380,13 +403,18 @@ private:
 
     SamplerPtrType& GetImmutableSampler(Uint32 n)noexcept
     {
-        VERIFY(n < m_NumImmutableSamplers, "Static sampler index (", n, ") is out of range. Array size: ", m_NumImmutableSamplers);
+        VERIFY(n < m_NumImmutableSamplers, "Immutable sampler index (", n, ") is out of range. Total immutable sampler count: ", m_NumImmutableSamplers);
         auto* ResourceMemoryEnd = reinterpret_cast<SPIRVShaderResourceAttribs*>(m_MemoryBuffer.get()) + m_TotalResources;
         return reinterpret_cast<SamplerPtrType*>(ResourceMemoryEnd)[n];
     }
 
+    SPIRVShaderStageInputAttribs& GetShaderStageInputAttribs(Uint32 n)noexcept
+    {
+        return const_cast<SPIRVShaderStageInputAttribs&>(const_cast<const SPIRVShaderResources*>(this)->GetShaderStageInputAttribs(n));
+    }
+
     // Memory buffer that holds all resources as continuous chunk of memory:
-    // |  UBs  |  SBs  |  StrgImgs  |  SmplImgs  |  ACs  |  SepSamplers  |  SepImgs  |   Immutable Samplers |   Resource Names   |
+    // |  UBs  |  SBs  |  StrgImgs  |  SmplImgs  |  ACs  |  SepSamplers  |  SepImgs  | Immutable Samplers | Stage Inputs | Resource Names |
     std::unique_ptr< void, STDDeleterRawMem<void> > m_MemoryBuffer;
     StringPool m_ResourceNames;
 
@@ -400,7 +428,8 @@ private:
     OffsetType m_SeparateSamplerOffset = 0;
     OffsetType m_SeparateImageOffset   = 0;
     OffsetType m_TotalResources        = 0;
-    OffsetType m_NumImmutableSamplers     = 0;
+    OffsetType m_NumImmutableSamplers  = 0;
+    OffsetType m_NumShaderStageInputs  = 0;
 
     SHADER_TYPE m_ShaderType = SHADER_TYPE_UNKNOWN;
 };
diff --git a/Graphics/GLSLTools/src/SPIRVShaderResources.cpp b/Graphics/GLSLTools/src/SPIRVShaderResources.cpp
index 3e69cae9..5fe9a460 100644
--- a/Graphics/GLSLTools/src/SPIRVShaderResources.cpp
+++ b/Graphics/GLSLTools/src/SPIRVShaderResources.cpp
@@ -95,7 +95,8 @@ SPIRVShaderResources::SPIRVShaderResources(IMemoryAllocator&      Allocator,
                                            IRenderDevice*         pRenderDevice,
                                            std::vector<uint32_t>  spirv_binary,
                                            const ShaderDesc&      shaderDesc,
-                                           const char*            CombinedSamplerSuffix) :
+                                           const char*            CombinedSamplerSuffix,
+                                           bool                   LoadShaderStageInputs) :
     m_ShaderType(shaderDesc.ShaderType)
 {
     // https://github.com/KhronosGroup/SPIRV-Cross/wiki/Reflection-API-user-guide
@@ -103,7 +104,7 @@ SPIRVShaderResources::SPIRVShaderResources(IMemoryAllocator&      Allocator,
 
     // The SPIR-V is now parsed, and we can perform reflection on it.
     spirv_cross::ShaderResources resources = Compiler.get_shader_resources();
-
+    
     size_t ResourceNamesPoolSize = 0;
     for(auto *pResType : 
         {
@@ -125,6 +126,47 @@ SPIRVShaderResources::SPIRVShaderResources(IMemoryAllocator&      Allocator,
         ResourceNamesPoolSize += strlen(CombinedSamplerSuffix) + 1;
     }
 
+    Uint32 NumShaderStageInputs = 0;
+
+    if (resources.stage_inputs.empty())
+        LoadShaderStageInputs = false;
+    if (LoadShaderStageInputs)
+    {
+        const auto& Extensions = Compiler.get_declared_extensions();
+        bool HlslFunctionality1 = false;
+        for (const auto& ext : Extensions )
+        {
+            HlslFunctionality1 = (ext == "SPV_GOOGLE_hlsl_functionality1");
+            if (HlslFunctionality1)
+                break;
+        }
+        
+        if (HlslFunctionality1)
+        {
+            for(const auto& Input : resources.stage_inputs)
+            {
+                if (Compiler.has_decoration(Input.id, spv::Decoration::DecorationHlslSemanticGOOGLE))
+                {
+                    const auto& Semantic = Compiler.get_decoration_string(Input.id, spv::Decoration::DecorationHlslSemanticGOOGLE);
+                    ResourceNamesPoolSize += Semantic.length()+1;
+                    ++NumShaderStageInputs;
+                }
+                else
+                {
+                    LOG_ERROR_MESSAGE("Shader input '", Input.name, "' does not have DecorationHlslSemanticGOOGLE decoration, which is unexpected as the shader declares SPV_GOOGLE_hlsl_functionality1 extension");
+                }
+            }
+        }
+        else
+        {
+            LoadShaderStageInputs = false;
+            LOG_WARNING_MESSAGE("SPIRV byte code of shader '", shaderDesc.Name, "' does not use SPV_GOOGLE_hlsl_functionality1 extension. "
+                                "As a result, it is not possible to get semantics of shader inputs and map them to proper locations. "
+                                "The shader will still work correctly if all attributes are declared in ascending order without any gaps. "
+                                "Enable SPV_GOOGLE_hlsl_functionality1 in your compiler to allow proper mapping of vertex shader inputs.");
+        }
+    }
+
     ResourceCounters ResCounters;
     ResCounters.NumUBs       = static_cast<Uint32>(resources.uniform_buffers.size());
     ResCounters.NumSBs       = static_cast<Uint32>(resources.storage_buffers.size());
@@ -133,7 +175,7 @@ SPIRVShaderResources::SPIRVShaderResources(IMemoryAllocator&      Allocator,
     ResCounters.NumACs       = static_cast<Uint32>(resources.atomic_counters.size());
     ResCounters.NumSepSmplrs = static_cast<Uint32>(resources.separate_samplers.size());
     ResCounters.NumSepImgs   = static_cast<Uint32>(resources.separate_images.size());
-    Initialize(Allocator, ResCounters, shaderDesc.NumStaticSamplers, ResourceNamesPoolSize);
+    Initialize(Allocator, ResCounters, shaderDesc.NumStaticSamplers, NumShaderStageInputs, ResourceNamesPoolSize);
 
     {
         Uint32 CurrUB = 0;
@@ -287,8 +329,6 @@ SPIRVShaderResources::SPIRVShaderResources(IMemoryAllocator&      Allocator,
         m_CombinedSamplerSuffix = m_ResourceNames.CopyString(CombinedSamplerSuffix);
     }
 
-    VERIFY(m_ResourceNames.GetRemainingSize() == 0, "Names pool must be empty");
-
     for (Uint32 s = 0; s < m_NumImmutableSamplers; ++s)
     {
         SamplerPtrType& pStaticSampler = GetImmutableSampler(s);
@@ -296,6 +336,23 @@ SPIRVShaderResources::SPIRVShaderResources(IMemoryAllocator&      Allocator,
         pRenderDevice->CreateSampler(shaderDesc.StaticSamplers[s].Desc, &pStaticSampler);
     }
 
+    if (LoadShaderStageInputs)
+    {
+        Uint32 CurrStageInput = 0;
+        for(const auto& Input : resources.stage_inputs)
+        {
+            if (Compiler.has_decoration(Input.id, spv::Decoration::DecorationHlslSemanticGOOGLE))
+            {
+                const auto& Semantic = Compiler.get_decoration_string(Input.id, spv::Decoration::DecorationHlslSemanticGOOGLE);
+                new (&GetShaderStageInputAttribs(CurrStageInput++)) 
+                    SPIRVShaderStageInputAttribs(m_ResourceNames.CopyString(Semantic), GetDecorationOffset(Compiler, Input, spv::Decoration::DecorationLocation));
+            }
+        }
+        VERIFY_EXPR(CurrStageInput == GetNumShaderStageInputs());
+    }
+
+    VERIFY(m_ResourceNames.GetRemainingSize() == 0, "Names pool must be empty");
+
     //LOG_INFO_MESSAGE(DumpResources());
 
 #ifdef DEVELOPMENT
@@ -376,6 +433,7 @@ SPIRVShaderResources::SPIRVShaderResources(IMemoryAllocator&      Allocator,
 void SPIRVShaderResources::Initialize(IMemoryAllocator&       Allocator, 
                                       const ResourceCounters& Counters,
                                       Uint32                  NumImmutableSamplers,
+                                      Uint32                  NumShaderStageInputs,
                                       size_t                  ResourceNamesPoolSize)
 {
     Uint32 CurrentOffset = 0;
@@ -400,10 +458,14 @@ void SPIRVShaderResources::Initialize(IMemoryAllocator&       Allocator,
     VERIFY(NumImmutableSamplers <= MaxOffset, "Max offset exceeded");
     m_NumImmutableSamplers = static_cast<OffsetType>(NumImmutableSamplers);
 
+    VERIFY(NumShaderStageInputs <= MaxOffset, "Max offset exceeded");
+    m_NumShaderStageInputs = static_cast<OffsetType>(NumShaderStageInputs);
+
     static_assert(sizeof(SPIRVShaderResourceAttribs) % sizeof(void*) == 0, "Size of SPIRVShaderResourceAttribs struct must be multiple of sizeof(void*)");
     static_assert(sizeof(SamplerPtrType)             % sizeof(void*) == 0, "Size of SamplerPtrType must be multiple of sizeof(void*)");
     auto MemorySize = m_TotalResources       * sizeof(SPIRVShaderResourceAttribs) + 
                       m_NumImmutableSamplers * sizeof(SamplerPtrType) +
+                      m_NumShaderStageInputs * sizeof(SPIRVShaderStageInputAttribs) +
                       ResourceNamesPoolSize  * sizeof(char);
 
     VERIFY_EXPR(GetNumUBs()       == Counters.NumUBs);
@@ -420,7 +482,8 @@ void SPIRVShaderResources::Initialize(IMemoryAllocator&       Allocator,
         m_MemoryBuffer = std::unique_ptr<void, STDDeleterRawMem<void>>(pRawMem, Allocator);
         char* NamesPool = reinterpret_cast<char*>(m_MemoryBuffer.get()) + 
                           m_TotalResources       * sizeof(SPIRVShaderResourceAttribs) +
-                          m_NumImmutableSamplers * sizeof(SamplerPtrType);
+                          m_NumImmutableSamplers * sizeof(SamplerPtrType) +
+                          m_NumShaderStageInputs * sizeof(SPIRVShaderStageInputAttribs);
         m_ResourceNames.AssignMemory(NamesPool, ResourceNamesPoolSize);
     }
 }
@@ -450,6 +513,9 @@ SPIRVShaderResources::~SPIRVShaderResources()
 
     for (Uint32 n = 0; n < GetNumImmutableSamplers(); ++n)
         GetImmutableSampler(n).~SamplerPtrType();
+    
+    for (Uint32 n = 0; n < GetNumShaderStageInputs(); ++n)
+        GetShaderStageInputAttribs(n).~SPIRVShaderStageInputAttribs();
 }
 
 SPIRVShaderResources::ResourceCounters  SPIRVShaderResources::CountResources(const SHADER_VARIABLE_TYPE* AllowedVarTypes,
@@ -601,7 +667,7 @@ bool SPIRVShaderResources::IsCompatibleWith(const SPIRVShaderResources& Resource
         GetNumACs()               != Resources.GetNumACs()        ||
         GetNumSepImgs()           != Resources.GetNumSepImgs()    ||
         GetNumSepSmplrs()         != Resources.GetNumSepSmplrs()  ||
-        GetNumImmutableSamplers() != Resources.GetNumImmutableSamplers() )
+        GetNumImmutableSamplers() != Resources.GetNumImmutableSamplers())
         return false;
     VERIFY_EXPR(GetTotalResources() == Resources.GetTotalResources());
 
diff --git a/Graphics/GLSLTools/src/SPIRVUtils.cpp b/Graphics/GLSLTools/src/SPIRVUtils.cpp
index bbd2d656..055b17e2 100644
--- a/Graphics/GLSLTools/src/SPIRVUtils.cpp
+++ b/Graphics/GLSLTools/src/SPIRVUtils.cpp
@@ -424,6 +424,7 @@ std::vector<unsigned int> HLSLtoSPIRV(const ShaderCreationAttribs& Attribs, IDat
     Shader.setEnvTarget(glslang::EShTargetSpv, glslang::EShTargetSpv_1_0);
     Shader.setHlslIoMapping(true);
     Shader.setEntryPoint(Attribs.EntryPoint);
+    Shader.setEnvTargetHlslFunctionality1();
 
     RefCntAutoPtr<IDataBlob> pFileData(MakeNewRCObj<DataBlobImpl>()(0));
     const char* SourceCode = 0;