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>

5 files changed, 20 insertions, 23 deletions

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;