Rework shader object specialization control interface. (#1857)

3 files changed, 101 insertions, 67 deletions

diff --git a/tools/gfx/debug-layer.cpp b/tools/gfx/debug-layer.cpp
index 3fa2eee9d..26e55ca7e 100644
--- a/tools/gfx/debug-layer.cpp
+++ b/tools/gfx/debug-layer.cpp
@@ -955,23 +955,12 @@ Result DebugShaderObject::setCombinedTextureSampler(
 }
 
 Result DebugShaderObject::setSpecializationArgs(
+    ShaderOffset const& offset,
     const slang::SpecializationArg* args,
     uint32_t count)
 {
-    ComPtr<slang::ISession> session;
-    m_device->getSlangSession(session.writeRef());
-    auto expectedCount = (uint32_t)session->getTypeLayout(m_slangType)
-                             ->getSize(SLANG_PARAMETER_CATEGORY_EXISTENTIAL_TYPE_PARAM);
-    if (expectedCount != count)
-    {
-        GFX_DIAGNOSE_ERROR_FORMAT(
-            "specialization argument count for shader object type %s mismatch: expecting %d but %d "
-            "provided.",
-            m_typeName.getBuffer(),
-            expectedCount,
-            count);
-    };
-    return baseObject->setSpecializationArgs(args, count);
+
+    return baseObject->setSpecializationArgs(offset, args, count);
 }
 
 DebugObjectBase::DebugObjectBase()
@@ -981,11 +970,11 @@ DebugObjectBase::DebugObjectBase()
 }
 
 Result DebugRootShaderObject::setSpecializationArgs(
+    ShaderOffset const& offset,
     const slang::SpecializationArg* args,
     uint32_t count)
 {
-    GFX_DIAGNOSE_ERROR("`setSpecializationArgs` should not be called directly on root objects.");
-    return baseObject->setSpecializationArgs(args, count);
+    return baseObject->setSpecializationArgs(offset, args, count);
 }
 
 } // namespace gfx
diff --git a/tools/gfx/debug-layer.h b/tools/gfx/debug-layer.h
index f9ecc7dfe..89ee9d837 100644
--- a/tools/gfx/debug-layer.h
+++ b/tools/gfx/debug-layer.h
@@ -168,6 +168,7 @@ public:
         IResourceView* textureView,
         ISamplerState* sampler) override;
     virtual SLANG_NO_THROW Result SLANG_MCALL setSpecializationArgs(
+        ShaderOffset const& offset,
         const slang::SpecializationArg* args,
         uint32_t count) override;
 
@@ -204,8 +205,10 @@ class DebugRootShaderObject : public DebugShaderObject
 public:
     virtual SLANG_NO_THROW uint32_t SLANG_MCALL addRef() override { return 1; }
     virtual SLANG_NO_THROW uint32_t SLANG_MCALL release() override { return 1; }
-    virtual SLANG_NO_THROW Result SLANG_MCALL
-        setSpecializationArgs(const slang::SpecializationArg* args, uint32_t count) override;
+    virtual SLANG_NO_THROW Result SLANG_MCALL setSpecializationArgs(
+        ShaderOffset const& offset,
+        const slang::SpecializationArg* args,
+        uint32_t count) override;
 };
 
 class DebugCommandBuffer;
diff --git a/tools/gfx/renderer-shared.h b/tools/gfx/renderer-shared.h
index 83799ac25..0dc0f75ae 100644
--- a/tools/gfx/renderer-shared.h
+++ b/tools/gfx/renderer-shared.h
@@ -473,10 +473,10 @@ class ShaderObjectBaseImpl : public ShaderObjectBase
 protected:
     TShaderObjectData m_data;
     Slang::List<Slang::RefPtr<TShaderObjectImpl>> m_objects;
+    Slang::List<Slang::RefPtr<ExtendedShaderObjectTypeListObject>> m_userProvidedSpecializationArgs;
 
     // Specialization args for a StructuredBuffer object.
     ExtendedShaderObjectTypeList m_structuredBufferSpecializationArgs;
-    Slang::RefPtr<ExtendedShaderObjectTypeListObject> m_userProvidedSpecializationArgs;
 
 public:
     TShaderObjectLayoutImpl* getLayout()
@@ -502,6 +502,39 @@ public:
         return SLANG_OK;
     }
 
+    void setSpecializationArgsForContainerElement(ExtendedShaderObjectTypeList& specializationArgs)
+    {
+        // Compute specialization args for the structured buffer object.
+        // If we haven't filled anything to `m_structuredBufferSpecializationArgs` yet,
+        // use `specializationArgs` directly.
+        if (m_structuredBufferSpecializationArgs.getCount() == 0)
+        {
+            m_structuredBufferSpecializationArgs = Slang::_Move(specializationArgs);
+        }
+        else
+        {
+            // If `m_structuredBufferSpecializationArgs` already contains some arguments, we
+            // need to check if they are the same as `specializationArgs`, and replace
+            // anything that is different with `__Dynamic` because we cannot specialize the
+            // buffer type if the element types are not the same.
+            SLANG_ASSERT(
+                m_structuredBufferSpecializationArgs.getCount() == specializationArgs.getCount());
+            auto device = getRenderer();
+            for (Slang::Index i = 0; i < m_structuredBufferSpecializationArgs.getCount(); i++)
+            {
+                if (m_structuredBufferSpecializationArgs[i].componentID !=
+                    specializationArgs[i].componentID)
+                {
+                    auto dynamicType = device->slangContext.session->getDynamicType();
+                    m_structuredBufferSpecializationArgs.componentIDs[i] =
+                        device->shaderCache.getComponentId(dynamicType);
+                    m_structuredBufferSpecializationArgs.components[i] =
+                        slang::SpecializationArg::fromType(dynamicType);
+                }
+            }
+        }
+    }
+
     virtual SLANG_NO_THROW Result SLANG_MCALL
         setObject(ShaderOffset const& offset, IShaderObject* object) SLANG_OVERRIDE
     {
@@ -563,36 +596,7 @@ public:
                 subObject->m_data.getBuffer(),
                 (size_t)subObject->m_data.getCount()));
 
-            // Compute specialization args for the structured buffer object.
-            // If we haven't filled anything to `m_structuredBufferSpecializationArgs` yet,
-            // use `specializationArgs` directly.
-            if (m_structuredBufferSpecializationArgs.getCount() == 0)
-            {
-                m_structuredBufferSpecializationArgs = Slang::_Move(specializationArgs);
-            }
-            else
-            {
-                // If `m_structuredBufferSpecializationArgs` already contains some arguments, we
-                // need to check if they are the same as `specializationArgs`, and replace
-                // anything that is different with `__Dynamic` because we cannot specialize the
-                // buffer type if the element types are not the same.
-                SLANG_ASSERT(
-                    m_structuredBufferSpecializationArgs.getCount() ==
-                    specializationArgs.getCount());
-                auto device = getRenderer();
-                for (Slang::Index i = 0; i < m_structuredBufferSpecializationArgs.getCount(); i++)
-                {
-                    if (m_structuredBufferSpecializationArgs[i].componentID !=
-                        specializationArgs[i].componentID)
-                    {
-                        auto dynamicType = device->slangContext.session->getDynamicType();
-                        m_structuredBufferSpecializationArgs.componentIDs[i] =
-                            device->shaderCache.getComponentId(dynamicType);
-                        m_structuredBufferSpecializationArgs.components[i] =
-                            slang::SpecializationArg::fromType(dynamicType);
-                    }
-                }
-            }
+            setSpecializationArgsForContainerElement(specializationArgs);
             return SLANG_OK;
         }
 
@@ -693,17 +697,11 @@ public:
         return SLANG_OK;
     }
 
-    virtual SLANG_NO_THROW Result SLANG_MCALL
-        setSpecializationArgs(const slang::SpecializationArg* args, uint32_t count) override
+    Result getExtendedShaderTypeListFromSpecializationArgs(
+        ExtendedShaderObjectTypeList& list,
+        const slang::SpecializationArg* args,
+        uint32_t count)
     {
-        if (!m_userProvidedSpecializationArgs)
-        {
-            m_userProvidedSpecializationArgs = new ExtendedShaderObjectTypeListObject();
-        }
-        else
-        {
-            m_userProvidedSpecializationArgs->clear();
-        }
         auto device = getRenderer();
         for (uint32_t i = 0; i < count; i++)
         {
@@ -718,20 +716,57 @@ public:
                 SLANG_ASSERT(false && "Unexpected specialization argument kind.");
                 return SLANG_FAIL;
             }
-            m_userProvidedSpecializationArgs->add(extendedType);
+            list.add(extendedType);
         }
         return SLANG_OK;
     }
 
-    // Appends all types that are used to specialize the element type of this shader object in
-    // `args` list.
-    virtual Result collectSpecializationArgs(ExtendedShaderObjectTypeList& args) override
+    virtual SLANG_NO_THROW Result SLANG_MCALL setSpecializationArgs(
+        ShaderOffset const& offset,
+        const slang::SpecializationArg* args,
+        uint32_t count) override
     {
-        if (m_userProvidedSpecializationArgs)
+        auto layout = getLayout();
+
+        // If the shader object is a container, delegate the processing to
+        // `setSpecializationArgsForContainerElements`.
+        if (layout->getContainerType() != ShaderObjectContainerType::None)
         {
-            args.addRange(*m_userProvidedSpecializationArgs);
+            ExtendedShaderObjectTypeList argList;
+            SLANG_RETURN_ON_FAIL(
+                getExtendedShaderTypeListFromSpecializationArgs(argList, args, count));
+            setSpecializationArgsForContainerElement(argList);
             return SLANG_OK;
         }
+
+        if (offset.bindingRangeIndex < 0)
+            return SLANG_E_INVALID_ARG;
+        if (offset.bindingRangeIndex >= layout->getBindingRangeCount())
+            return SLANG_E_INVALID_ARG;
+
+        auto bindingRangeIndex = offset.bindingRangeIndex;
+        auto bindingRange = layout->getBindingRange(bindingRangeIndex);
+        auto objectIndex = bindingRange.subObjectIndex + offset.bindingArrayIndex;
+        if (objectIndex >= m_userProvidedSpecializationArgs.getCount())
+            m_userProvidedSpecializationArgs.setCount(objectIndex + 1);
+        if (!m_userProvidedSpecializationArgs[objectIndex])
+        {
+            m_userProvidedSpecializationArgs[objectIndex] =
+                new ExtendedShaderObjectTypeListObject();
+        }
+        else
+        {
+            m_userProvidedSpecializationArgs[objectIndex]->clear();
+        }
+        SLANG_RETURN_ON_FAIL(getExtendedShaderTypeListFromSpecializationArgs(
+            *m_userProvidedSpecializationArgs[objectIndex], args, count));
+        return SLANG_OK;
+    }
+
+    // Appends all types that are used to specialize the element type of this shader object in
+    // `args` list.
+    virtual Result collectSpecializationArgs(ExtendedShaderObjectTypeList& args) override
+    {
         if (m_layout->getContainerType() != ShaderObjectContainerType::None)
         {
             args.addRange(m_structuredBufferSpecializationArgs);
@@ -761,12 +796,19 @@ public:
                  subObjectIndexInRange++)
             {
                 ExtendedShaderObjectTypeList typeArgs;
-
-                auto subObject = m_objects[bindingRange.subObjectIndex + subObjectIndexInRange];
+                auto objectIndex = bindingRange.subObjectIndex + subObjectIndexInRange;
+                auto subObject = m_objects[objectIndex];
 
                 if (!subObject)
                     continue;
 
+                if (objectIndex < m_userProvidedSpecializationArgs.getCount() &&
+                    m_userProvidedSpecializationArgs[objectIndex])
+                {
+                    args.addRange(*m_userProvidedSpecializationArgs[objectIndex]);
+                    continue;
+                }
+
                 switch (bindingRange.bindingType)
                 {
                 case slang::BindingType::ExistentialValue: