Parameter binding and gfx fixes. (#3302)

* Parameter binding and gfx fixes.

* Add diagnostics on entry point parameters.

* Fix.

---------

Co-authored-by: Yong He <yhe@nvidia.com>

18 files changed, 90 insertions, 59 deletions

diff --git a/tools/gfx/cpu/cpu-shader-object-layout.cpp b/tools/gfx/cpu/cpu-shader-object-layout.cpp
index 3b969c9a6..2ff89efff 100644
--- a/tools/gfx/cpu/cpu-shader-object-layout.cpp
+++ b/tools/gfx/cpu/cpu-shader-object-layout.cpp
@@ -81,6 +81,7 @@ ShaderObjectLayoutImpl::ShaderObjectLayoutImpl(RendererBase* renderer, slang::Ty
         bindingRangeInfo.baseIndex = baseIndex;
         bindingRangeInfo.uniformOffset = uniformOffset;
         bindingRangeInfo.subObjectIndex = subObjectIndex;
+        bindingRangeInfo.isSpecializable = m_elementTypeLayout->isBindingRangeSpecializable(r);
         m_bindingRanges.add(bindingRangeInfo);
     }
 
diff --git a/tools/gfx/cpu/cpu-shader-object-layout.h b/tools/gfx/cpu/cpu-shader-object-layout.h
index 44c225f83..e421918f1 100644
--- a/tools/gfx/cpu/cpu-shader-object-layout.h
+++ b/tools/gfx/cpu/cpu-shader-object-layout.h
@@ -30,6 +30,8 @@ struct BindingRangeInfo
     // range index and array index.
     //
     Index uniformOffset; // Uniform offset for a resource typed field.
+
+    bool isSpecializable;
 };
 
 struct SubObjectRangeInfo
diff --git a/tools/gfx/cuda/cuda-shader-object-layout.cpp b/tools/gfx/cuda/cuda-shader-object-layout.cpp
index 0cbe23a63..eacae0fc0 100644
--- a/tools/gfx/cuda/cuda-shader-object-layout.cpp
+++ b/tools/gfx/cuda/cuda-shader-object-layout.cpp
@@ -78,7 +78,8 @@ ShaderObjectLayoutImpl::ShaderObjectLayoutImpl(RendererBase* renderer, slang::Ty
         bindingRangeInfo.baseIndex = baseIndex;
         bindingRangeInfo.uniformOffset = uniformOffset;
         bindingRangeInfo.subObjectIndex = subObjectIndex;
-        m_bindingRanges.add(bindingRangeInfo);
+        bindingRangeInfo.isSpecializable = m_elementTypeLayout->isBindingRangeSpecializable(r);
+                m_bindingRanges.add(bindingRangeInfo);
     }
 
     SlangInt subObjectRangeCount = m_elementTypeLayout->getSubObjectRangeCount();
diff --git a/tools/gfx/cuda/cuda-shader-object-layout.h b/tools/gfx/cuda/cuda-shader-object-layout.h
index 305129109..5d3a2d52a 100644
--- a/tools/gfx/cuda/cuda-shader-object-layout.h
+++ b/tools/gfx/cuda/cuda-shader-object-layout.h
@@ -31,6 +31,8 @@ struct BindingRangeInfo
     // range index and array index.
     //
     Index uniformOffset; // Uniform offset for a resource typed field.
+
+    bool isSpecializable;
 };
 
 struct SubObjectRangeInfo
diff --git a/tools/gfx/d3d11/d3d11-shader-object-layout.cpp b/tools/gfx/d3d11/d3d11-shader-object-layout.cpp
index ff1c5d03c..8032719b0 100644
--- a/tools/gfx/d3d11/d3d11-shader-object-layout.cpp
+++ b/tools/gfx/d3d11/d3d11-shader-object-layout.cpp
@@ -51,7 +51,7 @@ Result ShaderObjectLayoutImpl::Builder::setElementTypeLayout(slang::TypeLayoutRe
         BindingRangeInfo bindingRangeInfo;
         bindingRangeInfo.bindingType = slangBindingType;
         bindingRangeInfo.count = count;
-
+        bindingRangeInfo.isSpecializable = typeLayout->isBindingRangeSpecializable(r);
         switch (slangBindingType)
         {
         case slang::BindingType::ConstantBuffer:
diff --git a/tools/gfx/d3d11/d3d11-shader-object-layout.h b/tools/gfx/d3d11/d3d11-shader-object-layout.h
index 717f270bc..e029a7b76 100644
--- a/tools/gfx/d3d11/d3d11-shader-object-layout.h
+++ b/tools/gfx/d3d11/d3d11-shader-object-layout.h
@@ -52,6 +52,9 @@ public:
             /// An index into the sub-object array if this binding range is treated
             /// as a sub-object.
         Index subObjectIndex;
+
+            /// Is this binding range specializable, e.g. an existential value or ParameterBlock<IFoo>.
+        bool isSpecializable;
     };
 
     // Sometimes we just want to iterate over the ranges that represent
diff --git a/tools/gfx/d3d11/d3d11-shader-object.cpp b/tools/gfx/d3d11/d3d11-shader-object.cpp
index 0354b3fdb..285cc60b6 100644
--- a/tools/gfx/d3d11/d3d11-shader-object.cpp
+++ b/tools/gfx/d3d11/d3d11-shader-object.cpp
@@ -334,7 +334,7 @@ Result ShaderObjectImpl::bindAsConstantBuffer(
     // Note that this call will use the `offset` value that might have
     // been modified during `_bindOrindaryDataBufferIfNeeded`.
     //
-    SLANG_RETURN_ON_FAIL(bindAsValue(context, offset, specializedLayout));
+    SLANG_RETURN_ON_FAIL(bindAsValue(context, inOffset, specializedLayout));
 
     return SLANG_OK;
 }
diff --git a/tools/gfx/d3d12/d3d12-device.cpp b/tools/gfx/d3d12/d3d12-device.cpp
index eb9e597bf..44a9faa45 100644
--- a/tools/gfx/d3d12/d3d12-device.cpp
+++ b/tools/gfx/d3d12/d3d12-device.cpp
@@ -98,6 +98,7 @@ Result DeviceImpl::createBuffer(
 
     D3D12_RESOURCE_STATES initialState = finalState;
 
+
     switch (memoryType)
     {
     case MemoryType::ReadBack:
@@ -106,18 +107,17 @@ Result DeviceImpl::createBuffer(
         heapProps.Type = D3D12_HEAP_TYPE_READBACK;
         desc.Flags = D3D12_RESOURCE_FLAG_NONE;
         initialState |= D3D12_RESOURCE_STATE_COPY_DEST;
-
         break;
     case MemoryType::Upload:
 
         heapProps.Type = D3D12_HEAP_TYPE_UPLOAD;
         desc.Flags = D3D12_RESOURCE_FLAG_NONE;
         initialState |= D3D12_RESOURCE_STATE_GENERIC_READ;
-
         break;
     case MemoryType::DeviceLocal:
         heapProps.Type = D3D12_HEAP_TYPE_DEFAULT;
-        initialState = (srcData ? D3D12_RESOURCE_STATE_COPY_DEST : finalState);
+        if (initialState != D3D12_RESOURCE_STATE_RAYTRACING_ACCELERATION_STRUCTURE)
+            initialState = D3D12_RESOURCE_STATE_COMMON;
         break;
     default:
         return SLANG_FAIL;
diff --git a/tools/gfx/d3d12/d3d12-shader-object-layout.cpp b/tools/gfx/d3d12/d3d12-shader-object-layout.cpp
index d90b638ba..c43d3881e 100644
--- a/tools/gfx/d3d12/d3d12-shader-object-layout.cpp
+++ b/tools/gfx/d3d12/d3d12-shader-object-layout.cpp
@@ -124,6 +124,7 @@ Result ShaderObjectLayoutImpl::Builder::setElementTypeLayout(
             static_cast<DeviceImpl*>(m_renderer)->m_extendedDesc.rootParameterShaderAttributeName,
             typeLayout,
             r);
+        bindingRangeInfo.isSpecializable = typeLayout->isBindingRangeSpecializable(r);
         if (bindingRangeInfo.isRootParameter)
         {
             RootParameterInfo rootInfo = {};
@@ -587,7 +588,7 @@ Result RootShaderObjectLayoutImpl::RootSignatureDescBuilder::addDescriptorRange(
         physicalDescriptorSetIndex,
         rangeType,
         (UINT)index + elementOffset[rangeType],
-        (UINT)space + containerOffset.spaceOffset,
+        (UINT)space + elementOffset.spaceOffset,
         (UINT)count,
         isRootParameter);
 }
@@ -646,7 +647,10 @@ void RootShaderObjectLayoutImpl::RootSignatureDescBuilder::addAsValue(
     slang::VariableLayoutReflection* varLayout, Index physicalDescriptorSetIndex)
 {
     BindingRegisterOffsetPair offset(varLayout);
-    addAsValue(varLayout->getTypeLayout(), physicalDescriptorSetIndex, offset, offset);
+    auto elementOffset = offset;
+    elementOffset.primary.spaceOffset = 0;
+    elementOffset.pending.spaceOffset = 0;
+    addAsValue(varLayout->getTypeLayout(), physicalDescriptorSetIndex, offset, elementOffset);
 }
 
 /// Add binding ranges and parameter blocks to the root signature.
@@ -658,30 +662,34 @@ void RootShaderObjectLayoutImpl::RootSignatureDescBuilder::addAsValue(
 /// the descriptor set that binding ranges not belonging to nested
 /// parameter blocks should be added to.
 ///
-/// The `offset` encodes information about space and/or register offsets that
-/// should be applied to descrptor ranges.
+/// The `offsetForChildrenThatNeedNewSpace` and `offsetForOrdinaryChildren` parameters
+/// encode information about space and/or register offsets that should be applied to
+/// descrptor ranges. `offsetForChildrenThatNeedNewSpace` will contain a space offset
+/// for children that requires a new space, such as a ParameterBlock.
+/// `offsetForOrdinaryChildren` contains the space for all direct children that should
+/// be placed in.
 ///
 
 void RootShaderObjectLayoutImpl::RootSignatureDescBuilder::addAsConstantBuffer(
     slang::TypeLayoutReflection* typeLayout,
     Index physicalDescriptorSetIndex,
-    BindingRegisterOffsetPair const& containerOffset,
-    BindingRegisterOffsetPair const& elementOffset)
+    BindingRegisterOffsetPair const& offsetForChildrenThatNeedNewSpace,
+    BindingRegisterOffsetPair const& offsetForOrdinaryChildren)
 {
     if (typeLayout->getSize(SLANG_PARAMETER_CATEGORY_UNIFORM) != 0)
     {
         auto descriptorRangeType = D3D12_DESCRIPTOR_RANGE_TYPE_CBV;
-        auto& offsetForRangeType = containerOffset.primary.offsetForRangeType[descriptorRangeType];
+        auto& offsetForRangeType = offsetForOrdinaryChildren.primary.offsetForRangeType[descriptorRangeType];
         addDescriptorRange(
             physicalDescriptorSetIndex,
             descriptorRangeType,
             offsetForRangeType,
-            containerOffset.primary.spaceOffset,
+            offsetForOrdinaryChildren.primary.spaceOffset,
             1,
             false);
     }
 
-    addAsValue(typeLayout, physicalDescriptorSetIndex, containerOffset, elementOffset);
+    addAsValue(typeLayout, physicalDescriptorSetIndex, offsetForChildrenThatNeedNewSpace, offsetForOrdinaryChildren);
 }
 
 void RootShaderObjectLayoutImpl::RootSignatureDescBuilder::addAsValue(
@@ -743,6 +751,8 @@ void RootShaderObjectLayoutImpl::RootSignatureDescBuilder::addAsValue(
         BindingRegisterOffsetPair subObjectRangeElementOffset = elementOffset;
         subObjectRangeElementOffset +=
             BindingRegisterOffsetPair(typeLayout->getSubObjectRangeOffset(subObjectRangeIndex));
+        subObjectRangeElementOffset.primary.spaceOffset = elementOffset.primary.spaceOffset;
+        subObjectRangeElementOffset.pending.spaceOffset = elementOffset.pending.spaceOffset;
 
         switch (bindingType)
         {
@@ -781,23 +791,30 @@ void RootShaderObjectLayoutImpl::RootSignatureDescBuilder::addAsValue(
 
             BindingRegisterOffsetPair subDescriptorSetOffset;
             subDescriptorSetOffset.primary.spaceOffset =
-                subObjectRangeElementOffset.primary.spaceOffset;
+                subObjectRangeContainerOffset.primary.spaceOffset;
             subDescriptorSetOffset.pending.spaceOffset =
-                subObjectRangeElementOffset.pending.spaceOffset;
+                subObjectRangeContainerOffset.pending.spaceOffset;
 
             auto subPhysicalDescriptorSetIndex = addDescriptorSet();
 
-            BindingRegisterOffsetPair containerOffset = subDescriptorSetOffset;
-            containerOffset += BindingRegisterOffsetPair(containerVarLayout);
-
-            BindingRegisterOffsetPair elementOffset = subDescriptorSetOffset;
-            elementOffset += BindingRegisterOffsetPair(elementVarLayout);
+            // We recursively call `addAsConstantBuffer` to actually generate
+            // the root signature bindings for children in the parameter block.
+            // We must compute `containerOffset`, which include a space offset
+            // that any sub ParameterBlocks should start from, and `elementOffset`
+            // that encodes the space offset of the current parameter block.
+            // The space offset of the current parameter block can be obtained from the
+            // `containerVarLayout`, and the space offset of any sub ParameterBlocks
+            // are obatined from `elementVarLayout`.
+            BindingRegisterOffsetPair offsetForChildrenThatNeedNewSpace = subDescriptorSetOffset;
+            offsetForChildrenThatNeedNewSpace += BindingRegisterOffsetPair(elementVarLayout);
+            BindingRegisterOffsetPair offsetForOrindaryChildren = subDescriptorSetOffset;
+            offsetForOrindaryChildren += BindingRegisterOffsetPair(containerVarLayout);
 
             addAsConstantBuffer(
                 elementTypeLayout,
                 subPhysicalDescriptorSetIndex,
-                containerOffset,
-                elementOffset);
+                offsetForChildrenThatNeedNewSpace,
+                offsetForOrindaryChildren);
         }
         break;
 
diff --git a/tools/gfx/d3d12/d3d12-shader-object-layout.h b/tools/gfx/d3d12/d3d12-shader-object-layout.h
index 8b72067f7..accd30aa9 100644
--- a/tools/gfx/d3d12/d3d12-shader-object-layout.h
+++ b/tools/gfx/d3d12/d3d12-shader-object-layout.h
@@ -67,6 +67,9 @@ public:
         uint32_t subObjectIndex;
 
         bool isRootParameter;
+
+        /// Is this binding range represent a specialization point, such as an existential value, or a `ParameterBlock<IFoo>`.
+        bool isSpecializable;
     };
 
     /// Offset information for a sub-object range
@@ -320,7 +323,7 @@ public:
                 if (varLayout)
                 {
                     spaceOffset =
-                        (UINT)varLayout->getOffset(SLANG_PARAMETER_CATEGORY_REGISTER_SPACE);
+                        (UINT)varLayout->getOffset(SLANG_PARAMETER_CATEGORY_SUB_ELEMENT_REGISTER_SPACE);
                     offsetForRangeType[D3D12_DESCRIPTOR_RANGE_TYPE_CBV] =
                         (UINT)varLayout->getOffset(SLANG_PARAMETER_CATEGORY_CONSTANT_BUFFER);
                     offsetForRangeType[D3D12_DESCRIPTOR_RANGE_TYPE_SRV] =
diff --git a/tools/gfx/d3d12/d3d12-shader-object.cpp b/tools/gfx/d3d12/d3d12-shader-object.cpp
index 74760eabd..5227579ce 100644
--- a/tools/gfx/d3d12/d3d12-shader-object.cpp
+++ b/tools/gfx/d3d12/d3d12-shader-object.cpp
@@ -440,7 +440,7 @@ void ShaderObjectImpl::updateSubObjectsRecursive()
             if (!subObject)
                 continue;
             subObject->updateSubObjectsRecursive();
-            if (m_subObjectVersions[objectIndex] != m_objects[objectIndex]->m_version)
+            if (m_subObjectVersions.getCount() > objectIndex && m_subObjectVersions[objectIndex] != m_objects[objectIndex]->m_version)
             {
                 ShaderOffset offset;
                 offset.bindingRangeIndex = (GfxIndex)subObjectRange.bindingRangeIndex;
diff --git a/tools/gfx/open-gl/render-gl.cpp b/tools/gfx/open-gl/render-gl.cpp
index 16f2b0045..c9eb1d94d 100644
--- a/tools/gfx/open-gl/render-gl.cpp
+++ b/tools/gfx/open-gl/render-gl.cpp
@@ -670,6 +670,7 @@ public:
             Index count;
             Index baseIndex;
             Index subObjectIndex;
+            bool isSpecializable;
         };
 
         struct SubObjectRangeInfo
@@ -721,7 +722,7 @@ public:
                     BindingRangeInfo bindingRangeInfo;
                     bindingRangeInfo.bindingType = slangBindingType;
                     bindingRangeInfo.count = count;
-
+                    bindingRangeInfo.isSpecializable = typeLayout->isBindingRangeSpecializable(r);
                     switch (slangBindingType)
                     {
                     case slang::BindingType::ConstantBuffer:
diff --git a/tools/gfx/renderer-shared.h b/tools/gfx/renderer-shared.h
index f94b909f7..9e208c34c 100644
--- a/tools/gfx/renderer-shared.h
+++ b/tools/gfx/renderer-shared.h
@@ -1524,16 +1524,20 @@ Result ShaderObjectBaseImpl<TShaderObjectImpl, TShaderObjectLayoutImpl, TShaderO
                 case slang::BindingType::ConstantBuffer:
                 case slang::BindingType::RawBuffer:
                 case slang::BindingType::MutableRawBuffer:
-                    // Currently we only handle the case where the field's type is
-                    // `ParameterBlock<SomeStruct>` or `ConstantBuffer<SomeStruct>`, where
-                    // `SomeStruct` is a struct type (not directly an interface type). In this case,
-                    // we just recursively collect the specialization arguments from the bound sub
-                    // object.
+                    // If the field's type is `ParameterBlock<IFoo>`, we want to pull in the type argument
+                    // from the sub object for specialization.
+                    if (bindingRange.isSpecializable)
+                    {
+                        ExtendedShaderObjectType specializedSubObjType;
+                        SLANG_RETURN_ON_FAIL(
+                            subObject->getSpecializedShaderObjectType(&specializedSubObjType));
+                        typeArgs.add(specializedSubObjType);
+                    }
+
+                    // If field's type is `ParameterBlock<SomeStruct>` or `ConstantBuffer<SomeStruct>`, where
+                    // `SomeStruct` is a struct type (not directly an interface type), we need to recursively
+                    // collect the specialization arguments from the bound sub object.
                     SLANG_RETURN_ON_FAIL(subObject->collectSpecializationArgs(typeArgs));
-                    // TODO: we need to handle the case where the field is of the form
-                    // `ParameterBlock<IFoo>`. We should treat this case the same way as the
-                    // `ExistentialValue` case here, but currently we lack a mechanism to
-                    // distinguish the two scenarios.
                     break;
             }
 
diff --git a/tools/gfx/vulkan/vk-command-encoder.cpp b/tools/gfx/vulkan/vk-command-encoder.cpp
index ddb48833f..78d8a751f 100644
--- a/tools/gfx/vulkan/vk-command-encoder.cpp
+++ b/tools/gfx/vulkan/vk-command-encoder.cpp
@@ -130,12 +130,8 @@ Result PipelineCommandEncoder::bindRootShaderObjectImpl(VkPipelineBindPoint bind
     // by the specialized program layout.
     //
     List<VkDescriptorSet> descriptorSetsStorage;
-    auto descriptorSetCount = specializedLayout->getTotalDescriptorSetCount();
 
-    descriptorSetsStorage.setCount(descriptorSetCount);
-    auto descriptorSets = descriptorSetsStorage.getBuffer();
-
-    context.descriptorSets = descriptorSets;
+    context.descriptorSets = &descriptorSetsStorage;
 
     // We kick off recursive binding of shader objects to the pipeline (plus
     // the state in `context`).
@@ -151,15 +147,15 @@ Result PipelineCommandEncoder::bindRootShaderObjectImpl(VkPipelineBindPoint bind
     // Once we've filled in all the descriptor sets, we bind them
     // to the pipeline at once.
     //
-    if (descriptorSetCount > 0)
+    if (descriptorSetsStorage.getCount() > 0)
     {
         m_device->m_api.vkCmdBindDescriptorSets(
             m_commandBuffer->m_commandBuffer,
             bindPoint,
             specializedLayout->m_pipelineLayout,
             0,
-            (uint32_t)descriptorSetCount,
-            descriptorSets,
+            (uint32_t)descriptorSetsStorage.getCount(),
+            descriptorSetsStorage.getBuffer(),
             0,
             nullptr);
     }
diff --git a/tools/gfx/vulkan/vk-helper-functions.h b/tools/gfx/vulkan/vk-helper-functions.h
index e2fae801e..8eab863f4 100644
--- a/tools/gfx/vulkan/vk-helper-functions.h
+++ b/tools/gfx/vulkan/vk-helper-functions.h
@@ -140,12 +140,10 @@ struct RootBindingContext
     DeviceImpl* device;
 
     /// The descriptor sets that are being allocated and bound
-    VkDescriptorSet* descriptorSets;
+    List<VkDescriptorSet>* descriptorSets;
 
     /// Information about all the push-constant ranges that should be bound
     ConstArrayView<VkPushConstantRange> pushConstantRanges;
-
-    uint32_t descriptorSetCounter = 0;
 };
 
 Size calcRowSize(Format format, int width);
diff --git a/tools/gfx/vulkan/vk-shader-object-layout.cpp b/tools/gfx/vulkan/vk-shader-object-layout.cpp
index 03dc1f11a..d7f0d0fd0 100644
--- a/tools/gfx/vulkan/vk-shader-object-layout.cpp
+++ b/tools/gfx/vulkan/vk-shader-object-layout.cpp
@@ -394,7 +394,7 @@ void ShaderObjectLayoutImpl::Builder::addBindingRanges(slang::TypeLayoutReflecti
         bindingRangeInfo.count = count;
         bindingRangeInfo.baseIndex = baseIndex;
         bindingRangeInfo.subObjectIndex = subObjectIndex;
-
+        bindingRangeInfo.isSpecializable = typeLayout->isBindingRangeSpecializable(r);
         // We'd like to extract the information on the GLSL/SPIR-V
         // `binding` that this range should bind into (or whatever
         // other specific kind of offset/index is appropriate to it).
diff --git a/tools/gfx/vulkan/vk-shader-object-layout.h b/tools/gfx/vulkan/vk-shader-object-layout.h
index cf4035f70..e1a01dcf3 100644
--- a/tools/gfx/vulkan/vk-shader-object-layout.h
+++ b/tools/gfx/vulkan/vk-shader-object-layout.h
@@ -59,6 +59,10 @@ public:
         // TODO: Ideally we could refactor so that only the root shader object layout
         // stores a set offset for its binding ranges, and all other objects skip
         // storing a field that never actually matters.
+
+        // Is this binding range representing a specialization point, such as
+        // an existential value or a ParameterBlock<IFoo>.
+        bool isSpecializable;
     };
 
     // Sometimes we just want to iterate over the ranges that represent
diff --git a/tools/gfx/vulkan/vk-shader-object.cpp b/tools/gfx/vulkan/vk-shader-object.cpp
index 3c452c59c..31422429b 100644
--- a/tools/gfx/vulkan/vk-shader-object.cpp
+++ b/tools/gfx/vulkan/vk-shader-object.cpp
@@ -321,7 +321,7 @@ void ShaderObjectImpl::writeBufferDescriptor(
     Offset bufferOffset,
     Size bufferSize)
 {
-    auto descriptorSet = context.descriptorSets[offset.bindingSet];
+    auto descriptorSet = (*context.descriptorSets)[offset.bindingSet];
 
     VkDescriptorBufferInfo bufferInfo = {};
     if (buffer)
@@ -359,7 +359,7 @@ void ShaderObjectImpl::writePlainBufferDescriptor(
     VkDescriptorType descriptorType,
     ArrayView<RefPtr<ResourceViewInternalBase>> resourceViews)
 {
-    auto descriptorSet = context.descriptorSets[offset.bindingSet];
+    auto descriptorSet = (*context.descriptorSets)[offset.bindingSet];
 
     Index count = resourceViews.getCount();
     for (Index i = 0; i < count; ++i)
@@ -398,7 +398,7 @@ void ShaderObjectImpl::writeTexelBufferDescriptor(
     VkDescriptorType descriptorType,
     ArrayView<RefPtr<ResourceViewInternalBase>> resourceViews)
 {
-    auto descriptorSet = context.descriptorSets[offset.bindingSet];
+    auto descriptorSet = (*context.descriptorSets)[offset.bindingSet];
 
     Index count = resourceViews.getCount();
     for (Index i = 0; i < count; ++i)
@@ -432,7 +432,7 @@ void ShaderObjectImpl::writeTextureSamplerDescriptor(
     VkDescriptorType descriptorType,
     ArrayView<CombinedTextureSamplerSlot> slots)
 {
-    auto descriptorSet = context.descriptorSets[offset.bindingSet];
+    auto descriptorSet = (*context.descriptorSets)[offset.bindingSet];
 
     Index count = slots.getCount();
     for (Index i = 0; i < count; ++i)
@@ -473,7 +473,7 @@ void ShaderObjectImpl::writeAccelerationStructureDescriptor(
     VkDescriptorType descriptorType,
     ArrayView<RefPtr<ResourceViewInternalBase>> resourceViews)
 {
-    auto descriptorSet = context.descriptorSets[offset.bindingSet];
+    auto descriptorSet = (*context.descriptorSets)[offset.bindingSet];
 
     Index count = resourceViews.getCount();
     for (Index i = 0; i < count; ++i)
@@ -511,7 +511,7 @@ void ShaderObjectImpl::writeTextureDescriptor(
     VkDescriptorType descriptorType,
     ArrayView<RefPtr<ResourceViewInternalBase>> resourceViews)
 {
-    auto descriptorSet = context.descriptorSets[offset.bindingSet];
+    auto descriptorSet = (*context.descriptorSets)[offset.bindingSet];
 
     Index count = resourceViews.getCount();
     for (Index i = 0; i < count; ++i)
@@ -548,7 +548,7 @@ void ShaderObjectImpl::writeSamplerDescriptor(
     VkDescriptorType descriptorType,
     ArrayView<RefPtr<SamplerStateImpl>> samplers)
 {
-    auto descriptorSet = context.descriptorSets[offset.bindingSet];
+    auto descriptorSet = (*context.descriptorSets)[offset.bindingSet];
 
     Index count = samplers.getCount();
     for (Index i = 0; i < count; ++i)
@@ -865,8 +865,7 @@ Result ShaderObjectImpl::allocateDescriptorSets(
         // we can bind all the descriptor sets to the pipeline when the
         // time comes.
         //
-        context.descriptorSets[context.descriptorSetCounter] = descriptorSetHandle;
-        context.descriptorSetCounter++;
+        (*context.descriptorSets).add(descriptorSetHandle);
     }
 
     return SLANG_OK;
@@ -884,7 +883,7 @@ Result ShaderObjectImpl::bindAsParameterBlock(
     // not the sets for any parent object(s).
     //
     BindingOffset offset = inOffset;
-    offset.bindingSet = context.descriptorSetCounter;
+    offset.bindingSet = (uint32_t)context.descriptorSets->getCount();
     offset.binding = 0;
 
     // TODO: We should also be writing to `offset.pending` here,
@@ -901,7 +900,7 @@ Result ShaderObjectImpl::bindAsParameterBlock(
     //
     SLANG_RETURN_ON_FAIL(allocateDescriptorSets(encoder, context, offset, specializedLayout));
 
-    assert(offset.bindingSet < context.descriptorSetCounter);
+    assert(offset.bindingSet < (uint32_t)context.descriptorSets->getCount());
     SLANG_RETURN_ON_FAIL(bindAsConstantBuffer(encoder, context, offset, specializedLayout));
 
     return SLANG_OK;