Better handling of bindings with multiple resource kind "aliases" for GLSL emit (#3009)

* A more way robust way to handle resource consumption might use multiple `kind`s on GLSL emit.

* Improve method naming and some comments.

* Small consistency fix.

7 files changed, 159 insertions, 44 deletions

diff --git a/source/slang/slang-emit-c-like.cpp b/source/slang/slang-emit-c-like.cpp
index 718653545..b05174b0a 100644
--- a/source/slang/slang-emit-c-like.cpp
+++ b/source/slang/slang-emit-c-like.cpp
@@ -605,6 +605,48 @@ void CLikeSourceEmitter::emitVal(IRInst* val, EmitOpInfo const& outerPrec)
     }
 }
 
+UInt CLikeSourceEmitter::getBindingOffsetForKinds(EmitVarChain* chain, LayoutResourceKindFlags kindFlags)
+{
+    UInt offset = 0;
+    for (auto cc = chain; cc; cc = cc->next)
+    {
+        for (auto offsetAttr : cc->varLayout->getOffsetAttrs())
+        {
+            // Accumulate offset for all matching kind
+            if (LayoutResourceKindFlag::make(offsetAttr->getResourceKind()) & kindFlags)
+            {
+                offset += offsetAttr->getOffset();
+            }
+        }
+    }
+
+    return offset;
+}
+
+UInt CLikeSourceEmitter::getBindingSpaceForKinds(EmitVarChain* chain, LayoutResourceKindFlags kindFlags)
+{
+    UInt space = 0;
+    for (auto cc = chain; cc; cc = cc->next)
+    {
+        auto varLayout = cc->varLayout;
+
+        for (auto offsetAttr : cc->varLayout->getOffsetAttrs())
+        {
+            // Accumulate offset for all matching kinds
+            if (LayoutResourceKindFlag::make(offsetAttr->getResourceKind()) & kindFlags)
+            {
+                space += offsetAttr->getSpace();
+            }
+        }
+
+        if (auto resInfo = varLayout->findOffsetAttr(LayoutResourceKind::RegisterSpace))
+        {
+            space += resInfo->getOffset();
+        }
+    }
+    return space;
+}
+
 UInt CLikeSourceEmitter::getBindingOffset(EmitVarChain* chain, LayoutResourceKind kind)
 {
     UInt offset = 0;
diff --git a/source/slang/slang-emit-c-like.h b/source/slang/slang-emit-c-like.h
index 47eff1ac4..4f8d23a0d 100644
--- a/source/slang/slang-emit-c-like.h
+++ b/source/slang/slang-emit-c-like.h
@@ -316,6 +316,13 @@ public:
     UInt getBindingOffset(EmitVarChain* chain, LayoutResourceKind kind);
     UInt getBindingSpace(EmitVarChain* chain, LayoutResourceKind kind);
 
+        /// Finds the binding offset for *all* the kinds that match the kindFlags
+        /// Thus only meaningful if multiple kinds can be treated as the same as far as binding is concerned.
+        /// In particular is useful for GLSL binding emit, where some HLSL resource kinds can appear but are in effect the 
+        /// same as DescriptorSlot
+    UInt getBindingOffsetForKinds(EmitVarChain* chain, LayoutResourceKindFlags kindFlags);   
+    UInt getBindingSpaceForKinds(EmitVarChain* chain, LayoutResourceKindFlags kindFlags);
+    
         // Utility code for generating unique IDs as needed
         // during the emit process (e.g., for declarations
         // that didn't originally have names, but now need to).
diff --git a/source/slang/slang-emit-glsl.cpp b/source/slang/slang-emit-glsl.cpp
index b378f69dc..c377b40c6 100644
--- a/source/slang/slang-emit-glsl.cpp
+++ b/source/slang/slang-emit-glsl.cpp
@@ -145,28 +145,6 @@ void GLSLSourceEmitter::_requireGLSLVersion(int version)
     }
 }
 
-
-// This function was needed because the vk-shift-* change allows kinds which are "in effect" DescriptorSlots but appear as HLSL
-// kinds. 
-// 
-// This seems safe as a test, because in the description of `usesResourceKind`, it appears that VarLayout require offsets even if
-// it's zero. 
-static LayoutResourceKind _findUsedDescriptorSlotLikeKind(IRVarLayout* layout, LayoutResourceKind kind)
-{
-    // DescriptorTableSlot is the most likely, so look for that first
-    if (layout->usesResourceKind(LayoutResourceKind::DescriptorTableSlot))
-    {
-        return LayoutResourceKind::DescriptorTableSlot;
-    }
-    else if (layout->usesResourceKind(kind))
-    {
-        // If we find the optional kind use that
-        return kind;
-    }
-    // We'll just assume descriptor slot then
-    return LayoutResourceKind::DescriptorTableSlot;
-}
-
 void GLSLSourceEmitter::_emitGLSLStructuredBuffer(IRGlobalParam* varDecl, IRHLSLStructuredBufferTypeBase* structuredBufferType)
 {
     // Shader storage buffer is an OpenGL 430 feature
@@ -177,14 +155,18 @@ void GLSLSourceEmitter::_emitGLSLStructuredBuffer(IRGlobalParam* varDecl, IRHLSL
     m_writer->emit("layout(");
     m_writer->emit(getTargetReq()->getForceGLSLScalarBufferLayout() ? "scalar" : "std430");
 
+    
     auto layout = getVarLayout(varDecl);
     if (layout)
     {
-        const LayoutResourceKind usedKind = _findUsedDescriptorSlotLikeKind(layout, LayoutResourceKind::ShaderResource);
+        // We can use ShaderResource/DescriptorSlot interchangably here. 
+       // This is possible because vk-shift-*
+        const LayoutResourceKindFlags kinds = LayoutResourceKindFlag::ShaderResource | LayoutResourceKindFlag::DescriptorTableSlot;
+
         EmitVarChain chain(layout);
 
-        const UInt index = getBindingOffset(&chain, usedKind);
-        const UInt space = getBindingSpace(&chain, usedKind);
+        const UInt index = getBindingOffsetForKinds(&chain, kinds);
+        const UInt space = getBindingSpaceForKinds(&chain, kinds);
 
         m_writer->emit(", binding = ");
         m_writer->emit(index);
@@ -256,12 +238,14 @@ void GLSLSourceEmitter::_emitGLSLByteAddressBuffer(IRGlobalParam* varDecl, IRByt
     auto layout = getVarLayout(varDecl);
     if (layout)
     {
-        const LayoutResourceKind usedKind = _findUsedDescriptorSlotLikeKind(layout, LayoutResourceKind::ShaderResource);
+        // We can use ShaderResource/DescriptorSlot interchangably here. 
+        // This is possible because vk-shift-*
+        const LayoutResourceKindFlags kinds = LayoutResourceKindFlag::ShaderResource | LayoutResourceKindFlag::DescriptorTableSlot;
 
         EmitVarChain chain(layout);
 
-        const UInt index = getBindingOffset(&chain, usedKind);
-        const UInt space = getBindingSpace(&chain, usedKind);
+        const UInt index = getBindingOffsetForKinds(&chain, kinds);
+        const UInt space = getBindingSpaceForKinds(&chain, kinds);
 
         m_writer->emit(", binding = ");
         m_writer->emit(index);
@@ -335,15 +319,10 @@ void GLSLSourceEmitter::_emitGLSLParameterGroup(IRGlobalParam* varDecl, IRUnifor
     */
 
     {
-        // TODO(JS):
-        // The *assumption* here is that we only need to detect the LayoutResourceKind as used on the initial containerChain.varLayout
-        // rather than search up the chain. 
-        //
-        // Perhaps there is still an issue here, because perhaps it's possible (?) to have a mixture of ConstantBuffer/DescriptorSlot
-        // and in that case the usedKind would just restrict to one and so produce the wrong answer.
-        const auto usedKind = _findUsedDescriptorSlotLikeKind(containerChain.varLayout, LayoutResourceKind::ConstantBuffer);
-        _emitGLSLLayoutQualifier(usedKind, &containerChain);
+        const LayoutResourceKindFlags kinds = LayoutResourceKindFlag::ConstantBuffer | LayoutResourceKindFlag::DescriptorTableSlot;
+        _emitGLSLLayoutQualifierWithBindingKinds(LayoutResourceKind::DescriptorTableSlot, &containerChain, kinds);
     }
+
     _emitGLSLLayoutQualifier(LayoutResourceKind::PushConstantBuffer, &containerChain);
     bool isShaderRecord = _emitGLSLLayoutQualifier(LayoutResourceKind::ShaderRecord, &containerChain);
 
@@ -567,15 +546,37 @@ void GLSLSourceEmitter::_emitGLSLImageFormatModifier(IRInst* var, IRTextureType*
     }
 }
 
-bool GLSLSourceEmitter::_emitGLSLLayoutQualifier(LayoutResourceKind kind, EmitVarChain* chain)
+bool GLSLSourceEmitter::_emitGLSLLayoutQualifierWithBindingKinds(LayoutResourceKind kind, EmitVarChain* chain, LayoutResourceKindFlags bindingKinds)
 {
     if (!chain)
         return false;
-    if (!chain->varLayout->findOffsetAttr(kind))
-        return false;
 
-    UInt index = getBindingOffset(chain, kind);
-    UInt space = getBindingSpace(chain, kind);
+    UInt index, space; 
+    auto varLayout = chain->varLayout;
+
+    // If bindingKinds are set, we use that for binding lookup
+    if (bindingKinds != 0)
+    {
+        if (!varLayout->usesResourceFromKinds(bindingKinds))
+        {
+            return false;
+        }
+
+        index = getBindingOffsetForKinds(chain, bindingKinds);
+        space = getBindingSpaceForKinds(chain, bindingKinds);
+    }
+    else
+    {
+        // Otherwise we just use kind
+        if (!varLayout->usesResourceKind(kind))
+        {
+            return false;
+        }
+        
+        index = getBindingOffset(chain, kind);
+        space = getBindingSpace(chain, kind);
+    }
+
     switch (kind)
     {
         case LayoutResourceKind::Uniform:
@@ -613,7 +614,7 @@ bool GLSLSourceEmitter::_emitGLSLLayoutQualifier(LayoutResourceKind kind, EmitVa
         case LayoutResourceKind::VaryingOutput:
             m_writer->emit("layout(location = ");
             m_writer->emit(index);
-            if( space )
+            if (space)
             {
                 m_writer->emit(", index = ");
                 m_writer->emit(space);
@@ -647,7 +648,7 @@ bool GLSLSourceEmitter::_emitGLSLLayoutQualifier(LayoutResourceKind kind, EmitVa
             m_writer->emit("layout(push_constant)\n");
             break;
         case LayoutResourceKind::ShaderRecord:
-            if( getTargetCaps().implies(CapabilityAtom::GL_NV_ray_tracing) )
+            if (getTargetCaps().implies(CapabilityAtom::GL_NV_ray_tracing))
             {
                 m_writer->emit("layout(shaderRecordNV)\n");
             }
diff --git a/source/slang/slang-emit-glsl.h b/source/slang/slang-emit-glsl.h
index 757955677..c0f3336f0 100644
--- a/source/slang/slang-emit-glsl.h
+++ b/source/slang/slang-emit-glsl.h
@@ -66,7 +66,10 @@ protected:
     void _emitGLSLImageFormatModifier(IRInst* var, IRTextureType* resourceType);
 
     void _emitGLSLLayoutQualifiers(IRVarLayout* layout, EmitVarChain* inChain, LayoutResourceKind filter = LayoutResourceKind::None);
-    bool _emitGLSLLayoutQualifier(LayoutResourceKind kind, EmitVarChain* chain);
+
+        /// If bindingKinds is set, it is used for binding index/set lookup. Passing in 0 is equivalent to using the kind only.
+    bool _emitGLSLLayoutQualifierWithBindingKinds(LayoutResourceKind kind, EmitVarChain* chain, LayoutResourceKindFlags bindingKinds);
+    bool _emitGLSLLayoutQualifier(LayoutResourceKind kind, EmitVarChain* chain) { return _emitGLSLLayoutQualifierWithBindingKinds(kind, chain, 0); }
 
     void _emitGLSLTypePrefix(IRType* type, bool promoteHalfToFloat = false);
 
diff --git a/source/slang/slang-ir-insts.h b/source/slang/slang-ir-insts.h
index a4563c254..f933697a2 100644
--- a/source/slang/slang-ir-insts.h
+++ b/source/slang/slang-ir-insts.h
@@ -1851,6 +1851,8 @@ struct IRVarLayout : IRLayout
 
         /// Does this variable use any resources of the given `kind`?
     bool usesResourceKind(LayoutResourceKind kind);
+        /// Returns true if there is use of one or more of the kinds
+    bool usesResourceFromKinds(LayoutResourceKindFlags kindFlags);
 
         /// Get the fixed/known stage that this variable is associated with.
         ///
diff --git a/source/slang/slang-ir.cpp b/source/slang/slang-ir.cpp
index 2fc1a9466..a27bf8658 100644
--- a/source/slang/slang-ir.cpp
+++ b/source/slang/slang-ir.cpp
@@ -1047,6 +1047,20 @@ namespace Slang
         return findOffsetAttr(kind) != nullptr;
     }
 
+    bool IRVarLayout::usesResourceFromKinds(LayoutResourceKindFlags kindFlags)
+    {
+        // Like usesResourceKind this works because there is an offset stored even if it's 0.
+        if (kindFlags)
+        {
+            for (auto offsetAttr : getOffsetAttrs())
+            {
+                if (LayoutResourceKindFlag::make(offsetAttr->getResourceKind()) & kindFlags )
+                    return true;
+            }
+        }
+        return false;
+    }
+
     IRSystemValueSemanticAttr* IRVarLayout::findSystemValueSemanticAttr()
     {
         return findAttr<IRSystemValueSemanticAttr>();
diff --git a/source/slang/slang-type-layout.h b/source/slang/slang-type-layout.h
index dc769d1c7..7b822eac4 100644
--- a/source/slang/slang-type-layout.h
+++ b/source/slang/slang-type-layout.h
@@ -243,6 +243,52 @@ struct UniformArrayLayoutInfo : UniformLayoutInfo
 
 typedef slang::ParameterCategory LayoutResourceKind;
 
+// Any change to slang::ParameterCategory, requires a change to this macro.
+#define SLANG_PARAMETER_CATEGORIES(x) \
+    x(None) \
+    x(Mixed) \
+    x(ConstantBuffer) \
+    x(ShaderResource) \
+    x(UnorderedAccess) \
+    x(VaryingInput) \
+    x(VaryingOutput) \
+    x(SamplerState) \
+    x(Uniform) \
+    x(DescriptorTableSlot) \
+    x(SpecializationConstant) \
+    x(PushConstantBuffer) \
+    x(RegisterSpace) \
+    x(GenericResource) \
+    \
+    x(RayPayload) \
+    x(HitAttributes) \
+    x(CallablePayload) \
+    \
+    x(ShaderRecord) \
+    \
+    x(ExistentialTypeParam) \
+    x(ExistentialObjectParam) \
+    \
+    x(VertexInput) \
+    x(FragmentOutput)
+
+#define SLANG_PARAMETER_CATEGORY_FLAG(x) x = ParameterCategoryFlags(1) << int(slang::x), 
+
+typedef uint32_t ParameterCategoryFlags;
+struct ParameterCategoryFlag
+{
+    enum Enum : ParameterCategoryFlags
+    {
+        SLANG_PARAMETER_CATEGORIES(SLANG_PARAMETER_CATEGORY_FLAG)
+    };
+
+        /// Make a flag from a category
+    SLANG_FORCE_INLINE static Enum make(slang::ParameterCategory cat) { return Enum(ParameterCategoryFlags(1) << int(cat)); }
+};
+
+typedef ParameterCategoryFlags LayoutResourceKindFlags;
+typedef ParameterCategoryFlag LayoutResourceKindFlag;
+
 // Layout information for a value that only consumes
 // a single resource kind.
 struct SimpleLayoutInfo