Change the policy for entry-point uniform parameters on Vulkan (#1476)

Entry point `uniform` parameters were a feature of the original Cg and HLSL, but have not been used much in production shader code. One of our goals on Slang is to reduce the (ab)use of the global scope, so bringing entry point `uniform` parameters up to a greater level of usability is an important goal.

Some policy choices about how global vs. entry-point `uniform` parameters behave have already been made, that shape decisions looking forward:

* For DXBC/DXIL, it makes the most sense to follow the lead of fxc/dxc, by treating entry point `uniform` parameters as a kind of syntax sugar for global shader parameters. Any parameters of "ordinary" types are bundles up into an implicit constant buffer, and all the resources (including the implicit constant buffer) are assigned `register`s just as for globals. It is up to the application to decide how to bind those parameters via a root signature (using root descriptors, root constants, descriptor tables, local vs. global root signature, etc.)

* For CPU, it makes sense to pass global vs. entry-point parameters as two different pointers, although the details of what we do for CPU are the least constrained across all current targets.

* For CUDA compute, it makes the most sense to map global shader parameters to `__constant__` global data, and entry-point `uniform` parameters to kernel parameters. This choice ensures that the signature of a kernel when translated from Slang->CUDA follows the Principle of Least Surprise, at the cost of making entry-point vs. global parameters be passed via different mechanisms.

* For OptiX ray tracing, it makes sense to expand on the precedent from CUDA compute: pass global parameters via global `__constant__` data (as is already expected by OptiX for whole-launch parameters), and pass entry-point `uniform` parameters via the "shader record." This establishes a precedent that for ray-tracing shaders, global-scope parameters map to the "global root signature" concept from DXR, while entry-point `uniform` parameters map to a "local root signature" or "shader record."

* For Vulkan ray tracing, the precedent from OptiX then argues that entry-point `uniform` parameters should map to the Vulkan "shader record" concept (and thus cannot support things like resource types).

* The remaining interesting case is what to do for non-ray-tracing shaders on Vulkan.

The dev team agrees that the most reasonable choice to make for non-ray-tracing Vulkan shaders is to map entry-point `uniform` parameters to "push constants." In particular, this makes it easy to express the case of a compute kernel with direct parameters of ordinary/value types in the way that will be implemented most efficiently.

The big picture is then that a kernel like:

```hlsl
void computeMain(uniform float someValue) { ... }
```

will map to output GLSL like:

```glsl
layout(push_constant)
uniform
{
    float someValue;
} U;

void main() { ... }
```

If the user really wanted a constant-buffer binding to be created instead, they can easily change their input to make the buffer explicit:

```hlsl
struct Params { float someValue; }
void computeMain(uniform ConstantBuffer<Params> params) { ... }
```

(Forcing the user to be explicit about the desire for a buffer here creates a nice symmetry between Vulkan and CUDA; in the first case the user sets up the data in host memory and passes it to the GPU by copy, while in the second case the user must allocate and set up a device-memory buffer for the data. This symmetry extends to D3D if the application chooses to map entry-point `uniform` parameters to root constants.)

This change implements logic in the "parameter binding" part of the Slang compiler to make sure that entry-point `uniform` parameters are wrapped up in a push-constant buffer rather than an ordinary constant buffer for non-ray-tracing shaders on Vulkan (and in a shader record "buffer" for the ray-tracing case).

The majority of the actual work was in adding support for root/push constants to the test framework and the graphics API abstraction it uses. To be clear about that support:

* Root constant ranges are (perhaps confusingly) treated as a new kind of "slot" that can appear on a descriptor set. This choice ensures that the implicit numbering of registers/spaces used by the back-ends can account for these ranges correctly.

* The `TEST_INPUT` lines are extended to allow a `root_constants` case that behaves more or less like `cbuffer`

* The CPU and CUDA paths can treat a `root_constants` input identically to a `cbuffer`. They already allocate the actual buffers based on reflection, and just use `cbuffer` as a directive that causes bytes to be copied in.

* On D3D12 and Vulkan, a descriptor set allocates a `List<char>` to hold the bytes of root constant data assigned into it, and these bytes are flushed to the command list when the table is actually bound (usually right before rendering).

* On D3D11, a descriptor set treats a root constant range more or less like a constant buffer range (with a single buffer), except that it also automatically allocates a buffer to hold the data. Assigning "root constant" data automatically copies it into that buffer.

The small number of tests that used entry-point `uniform` parameters of ordinary types were updated to use the new `root_constant` input type, and the bugs that surfaced were fixed.

A new test to confirm that entry-point `uniform` parameters map to the shader record for VK ray tracing was added.

An important but technically unrelated change is the removal of the `DescriptorSetImpl::Binding` type and related function from the Vulkan implementation of `Renderer`. That type was created to ensure that objects that are bound into a descriptor set don't get released while the descriptor set is still alive, but the implementation relied on a complicated linear search to check for existing bindings, which could create a performance issue for descriptor sets that include large arrays of descriptors. The new implementation makes use of the approach already present in the various `Renderer` implementations (including the Vulkan one) for assigning ranges in a descriptor set a flat/linear index for where their pertinent data is to be bound. As a result, the Vulkan `DescriptorSetImpl` now uses a single flat array of `RefPtr`s to track bound objects, and has no need for linear search when binding.

Co-authored-by: Yong He <yonghe@outlook.com>

2 files changed, 283 insertions, 113 deletions

diff --git a/tools/gfx/vulkan/render-vk.cpp b/tools/gfx/vulkan/render-vk.cpp
index 53c210a20..b9cc82469 100644
--- a/tools/gfx/vulkan/render-vk.cpp
+++ b/tools/gfx/vulkan/render-vk.cpp
@@ -91,6 +91,9 @@ public:
 
     protected:
 
+        /// Flush state from descriptor set bindings into `commandBuffer`
+    void _flushBindingState(VkCommandBuffer commandBuffer, VkPipelineBindPoint pipelineBindPoint);
+
     class Buffer
     {
         public:
@@ -280,11 +283,57 @@ public:
         VkDescriptorSetLayout m_descriptorSetLayout = VK_NULL_HANDLE;
         VkDescriptorPool m_descriptorPool = VK_NULL_HANDLE;
 
+        // Vulkan descriptor sets are the closest in design to what
+        // the `Renderer` abstraction exposes as a `DescriptorSet`.
+        // The main difference is that a `DescriptorSet` can include
+        // root constant ranges, while under Vulkan push constant
+        // ranges are part of the `VkPipelineLayout`, but not part
+        // of any `VkDescriptorSetLayout`.
+        //
+        // Information about each descriptor slot range in the
+        // original `Desc` will be stored as `RangeInfo` values,
+        // which store the relevant information from the `Desc`
+        // as well as additional information specific to the
+        // Vulkan implementation path.
+        //
         struct RangeInfo
         {
-            VkDescriptorType descriptorType;
+                /// The type of descriptor slot range from the original `Desc`
+            DescriptorSlotType  type;
+
+                /// The start index of the range in the appropriate type-specific array
+            Index               arrayIndex;
+
+                /// The equivalent Vulkan descriptor type, where applicable
+            VkDescriptorType    vkDescriptorType;
+
+                /// The Vulkan `binding` index for this range
+            uint32_t            vkBindingIndex;
         };
         List<RangeInfo> m_ranges;
+
+        // Because root constant ranges aren't part of a `VkDescriptorSetLayout`,
+        // we store additional data to represent the ranges so that
+        // we can store their data on a `DescriptorSetImpl` and then
+        // bind it to the API later.
+        //
+        struct RootConstantRangeInfo
+        {
+                /// The offset of the range's data in the backing storage.
+            Index offset;
+
+                /// The size of the range's data.
+            Index size;
+        };
+        Slang::List<RootConstantRangeInfo> m_rootConstantRanges;
+
+            /// The total size, in bytes, or root constant data for this descriptor set.
+        uint32_t m_rootConstantDataSize = 0;
+
+            /// The total number of reference counted objects that can be bound
+            /// to descriptor sets described by this layout.
+            ///
+        Index m_totalBoundObjectCount = 0;
     };
 
     class PipelineLayoutImpl : public PipelineLayout
@@ -306,28 +355,14 @@ public:
         VulkanApi const*    m_api;
         VkPipelineLayout    m_pipelineLayout        = VK_NULL_HANDLE;
         UInt                m_descriptorSetCount    = 0;
+
+            /// For each descriptor set, stores the start offset of that set's root constant data in the pipeline layout
+        List<uint32_t>      m_descriptorSetRootConstantOffsets;
     };
 
     class DescriptorSetImpl : public DescriptorSet
     {
     public:
-        // Record the view binding
-        struct Binding
-        {
-            enum class Type : uint8_t
-            {
-                Unknown,
-                ResourceView,
-                SamplerState,
-                BufferResource,
-                CountOf,
-            };
-            Type type;
-            uint32_t range;
-            uint32_t index;
-            RefPtr<RefObject> obj;
-        };
-
         DescriptorSetImpl(VKRenderer* renderer)
             : m_renderer(renderer)
         {
@@ -345,15 +380,21 @@ public:
             UInt index,
             ResourceView*   textureView,
             SamplerState*   sampler) override;
-
-        static Binding::Type _getBindingType(RefObject* ptr);
-        void _setBinding(Binding::Type type, UInt range, UInt index, RefObject* ptr);
+        virtual void setRootConstants(
+            UInt range,
+            UInt offset,
+            UInt size,
+            void const* data) override;
 
         VKRenderer*                         m_renderer = nullptr;   ///< Weak pointer, can't be strong, because if set will become circular reference
         RefPtr<DescriptorSetLayoutImpl>     m_layout;
         VkDescriptorSet                     m_descriptorSet = VK_NULL_HANDLE;
 
-        List<Binding>                       m_bindings;             ///< Records entities are bound to this descriptor set, and keeps the associated resources/views/state in scope
+            /// Records entities that are bound to this descriptor set, and keeps the associated resources/views/state in scope
+        List<RefPtr<RefObject>>             m_boundObjects;
+
+            /// Backing storage for root constant ranges belonging to this descriptor set
+        List<char>                          m_rootConstantData;
     };
 
     struct BoundVertexBuffer
@@ -2019,6 +2060,46 @@ void VKRenderer::setPipelineState(PipelineType pipelineType, PipelineState* stat
     m_currentPipeline = (PipelineStateImpl*)state;
 }
 
+void VKRenderer::_flushBindingState(VkCommandBuffer commandBuffer, VkPipelineBindPoint pipelineBindPoint)
+{
+    auto pipeline = m_currentPipeline;
+
+    // We start by binding the pipeline state.
+    //
+    m_api.vkCmdBindPipeline(commandBuffer, pipelineBindPoint, pipeline->m_pipeline);
+
+    // Next we bind all the descriptor sets that were set in the `VKRenderer`.
+    //
+    auto pipelineLayoutImpl = pipeline->m_pipelineLayout.Ptr();
+    auto vkPipelineLayout = pipelineLayoutImpl->m_pipelineLayout;
+    auto descriptorSetCount = pipelineLayoutImpl->m_descriptorSetCount;
+    m_api.vkCmdBindDescriptorSets(commandBuffer, pipelineBindPoint, vkPipelineLayout,
+        0, uint32_t(descriptorSetCount),
+        &m_currentDescriptorSets[0],
+        0, nullptr);
+
+    // For any descriptor sets with root-constant ranges, we need to
+    // bind the relevant data to the context.
+    //
+    for(gfx::UInt ii = 0; ii < descriptorSetCount; ++ii)
+    {
+        auto descriptorSet = m_currentDescriptorSetImpls[ii];
+        auto descriptorSetLayout = descriptorSet->m_layout;
+        auto size = descriptorSetLayout->m_rootConstantDataSize;
+        if(size == 0)
+            continue;
+        auto data = descriptorSet->m_rootConstantData.getBuffer();
+
+        // The absolute offset of the descriptor set's data in
+        // the push-constant data for the entire pipeline was
+        // computed and cached in the pipeline layout.
+        //
+        uint32_t offset = pipelineLayoutImpl->m_descriptorSetRootConstantOffsets[ii];
+
+        m_api.vkCmdPushConstants(commandBuffer, vkPipelineLayout, VK_SHADER_STAGE_ALL, offset, size, data);
+    }
+}
+
 void VKRenderer::draw(UInt vertexCount, UInt startVertex = 0)
 {
     auto pipeline = m_currentPipeline;
@@ -2033,13 +2114,7 @@ void VKRenderer::draw(UInt vertexCount, UInt startVertex = 0)
     // Also create descriptor sets based on the given pipeline layout
     VkCommandBuffer commandBuffer = m_deviceQueue.getCommandBuffer();
 
-    m_api.vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, pipeline->m_pipeline);
-
-    auto pipelineLayoutImpl = pipeline->m_pipelineLayout.Ptr();
-    m_api.vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, pipelineLayoutImpl->m_pipelineLayout,
-        0, uint32_t(pipelineLayoutImpl->m_descriptorSetCount),
-        &m_currentDescriptorSets[0],
-        0, nullptr);
+    _flushBindingState(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS);
 
     // Bind the vertex buffer
     if (m_boundVertexBuffers.getCount() > 0 && m_boundVertexBuffers[0].m_buffer)
@@ -2073,13 +2148,7 @@ void VKRenderer::dispatchCompute(int x, int y, int z)
     // Also create descriptor sets based on the given pipeline layout
     VkCommandBuffer commandBuffer = m_deviceQueue.getCommandBuffer();
 
-    m_api.vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline->m_pipeline);
-
-    auto pipelineLayoutImpl = pipeline->m_pipelineLayout.Ptr();
-    m_api.vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_COMPUTE, pipelineLayoutImpl->m_pipelineLayout,
-        0, uint32_t(pipelineLayoutImpl->m_descriptorSetCount),
-        &m_currentDescriptorSets[0],
-        0, nullptr);
+    _flushBindingState(commandBuffer, VK_PIPELINE_BIND_POINT_COMPUTE);
 
     m_api.vkCmdDispatch(commandBuffer, x, y, z);
 }
@@ -2284,10 +2353,61 @@ Result VKRenderer::createDescriptorSetLayout(const DescriptorSetLayout::Desc& de
     {
         auto& srcRange = desc.slotRanges[rr];
 
+        if(srcRange.type == DescriptorSlotType::RootConstant)
+        {
+            // Root constant ranges are a special case, since they
+            // don't actually map to `VkDescriptorSetLayoutBinding`s
+            // like the other cases.
+
+            // We start by computing the offset of the range within
+            // the backing storage for the descriptor set, while
+            // also updating the computed total size of root constant
+            // data needed by the set.
+            //
+            auto size = uint32_t(srcRange.count);
+            auto offset = descriptorSetLayoutImpl->m_rootConstantDataSize;
+            descriptorSetLayoutImpl->m_rootConstantDataSize += size;
+
+            // We will keep track of the information for this
+            // range as part of the descriptor set layout.
+            //
+            DescriptorSetLayoutImpl::RootConstantRangeInfo rootConstantRangeInfo;
+            rootConstantRangeInfo.offset = offset;
+            rootConstantRangeInfo.size = size;
+
+            auto rootConstantRangeIndex = descriptorSetLayoutImpl->m_rootConstantRanges.getCount();
+            descriptorSetLayoutImpl->m_rootConstantRanges.add(rootConstantRangeInfo);
+
+            // We will also add a `RangeInfo` to reprsent this
+            // range, even though it doesn't map to a VK-level
+            // descriptor range.
+            //
+            DescriptorSetLayoutImpl::RangeInfo rangeInfo;
+            rangeInfo.type = srcRange.type;
+            rangeInfo.vkDescriptorType = VkDescriptorType(-1);
+            rangeInfo.arrayIndex = rootConstantRangeIndex;
+            descriptorSetLayoutImpl->m_ranges.add(rangeInfo);
+
+            // Finally, we bail out instead of performing
+            // the logic that applies to the other descriptor
+            // range types.
+            //
+            continue;
+        }
+
+        // Note: Because of the existence of root constant ranges,
+        // we cannot assume that the `binding` for a range is
+        // the same as its index in the input array of ranges.
+        //
+        // Instead, the `binding` for a range is its index in
+        // the output array of `VkDescriptorSetLayoutBinding`s.
+        //
+        uint32_t bindingIndex = uint32_t(dstBindings.getCount());
+
         VkDescriptorType dstDescriptorType = translateDescriptorType(srcRange.type);
 
         VkDescriptorSetLayoutBinding dstBinding;
-        dstBinding.binding = uint32_t(rr);
+        dstBinding.binding = uint32_t(bindingIndex);
         dstBinding.descriptorType = dstDescriptorType;
         dstBinding.descriptorCount = uint32_t(srcRange.count);
         dstBinding.stageFlags = VK_SHADER_STAGE_ALL;
@@ -2297,8 +2417,20 @@ Result VKRenderer::createDescriptorSetLayout(const DescriptorSetLayout::Desc& de
 
         dstBindings.add(dstBinding);
 
+        UInt boundObjectCount = srcRange.count;
+        if( srcRange.type == DescriptorSlotType::CombinedImageSampler )
+        {
+            boundObjectCount = 2 * srcRange.count;
+        }
+
+        auto boundObjectArrayIndex = descriptorSetLayoutImpl->m_totalBoundObjectCount;
+        descriptorSetLayoutImpl->m_totalBoundObjectCount += boundObjectCount;
+
         DescriptorSetLayoutImpl::RangeInfo rangeInfo;
-        rangeInfo.descriptorType = dstDescriptorType;
+        rangeInfo.type = srcRange.type;
+        rangeInfo.vkDescriptorType = dstDescriptorType;
+        rangeInfo.vkBindingIndex = bindingIndex;
+        rangeInfo.arrayIndex = boundObjectArrayIndex;
         descriptorSetLayoutImpl->m_ranges.add(rangeInfo);
     }
 
@@ -2344,15 +2476,47 @@ Result VKRenderer::createPipelineLayout(const PipelineLayout::Desc& desc, Pipeli
     UInt descriptorSetCount = desc.descriptorSetCount;
 
     VkDescriptorSetLayout descriptorSetLayouts[kMaxDescriptorSets];
+    uint32_t descriptorSetRootConstantOffsets[kMaxDescriptorSets];
+    uint32_t totalRootConstantSize = 0;
     for(UInt ii = 0; ii < descriptorSetCount; ++ii)
     {
-        descriptorSetLayouts[ii] = ((DescriptorSetLayoutImpl*) desc.descriptorSets[ii].layout)->m_descriptorSetLayout;
+        auto descriptorSetLayoutImpl = (DescriptorSetLayoutImpl*) desc.descriptorSets[ii].layout;
+        descriptorSetLayouts[ii] = descriptorSetLayoutImpl->m_descriptorSetLayout;
+
+        descriptorSetRootConstantOffsets[ii] = totalRootConstantSize;
+        totalRootConstantSize += descriptorSetLayoutImpl->m_rootConstantDataSize;
     }
 
     VkPipelineLayoutCreateInfo pipelineLayoutInfo = { VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO };
     pipelineLayoutInfo.setLayoutCount = uint32_t(desc.descriptorSetCount);
     pipelineLayoutInfo.pSetLayouts = &descriptorSetLayouts[0];
 
+    // Our abstraction allows the user to specify any number of root-constant
+    // ranges across all of their descriptor sets, but Vulkan has a restriction
+    // that a pipeline layout may only include a single push constant range
+    // accessible from a given stage. (In other words, the only situation where
+    // multiple push-constant ranges are allowed is if you want to have, say,
+    // distinct ranges for the vertex and fragment stages to access).
+    //
+    // We handle this by declaring at most one push constant range, which
+    // represents the concatenation of the data from all ranges that the
+    // user might have asked for.
+    //
+    // Note: The Slang compiler doesn't yet have logic to concatenate multiple
+    // push-constant ranges in this way, but if/when it does, it should hopefully
+    // Just Work with this logic.
+    //
+    VkPushConstantRange pushConstantRange;
+    if( totalRootConstantSize )
+    {
+        pushConstantRange.offset = 0;
+        pushConstantRange.size = totalRootConstantSize;
+        pushConstantRange.stageFlags = VK_SHADER_STAGE_ALL;
+
+        pipelineLayoutInfo.pushConstantRangeCount = 1;
+        pipelineLayoutInfo.pPushConstantRanges = &pushConstantRange;
+    }
+
     VkPipelineLayout pipelineLayout;
     SLANG_VK_CHECK(m_api.vkCreatePipelineLayout(m_device, &pipelineLayoutInfo, nullptr, &pipelineLayout));
 
@@ -2360,6 +2524,12 @@ Result VKRenderer::createPipelineLayout(const PipelineLayout::Desc& desc, Pipeli
     pipelineLayoutImpl->m_pipelineLayout = pipelineLayout;
     pipelineLayoutImpl->m_descriptorSetCount = descriptorSetCount;
 
+    for(UInt ii = 0; ii < descriptorSetCount; ++ii)
+    {
+        pipelineLayoutImpl->m_descriptorSetRootConstantOffsets.add(
+            descriptorSetRootConstantOffsets[ii]);
+    }
+
     *outLayout = pipelineLayoutImpl.detach();
     return SLANG_OK;
 }
@@ -2379,73 +2549,23 @@ Result VKRenderer::createDescriptorSet(DescriptorSetLayout* layout, DescriptorSe
     RefPtr<DescriptorSetImpl> descriptorSetImpl = new DescriptorSetImpl(this);
     descriptorSetImpl->m_layout = layoutImpl;
     descriptorSetImpl->m_descriptorSet = descriptorSet;
-    *outDescriptorSet = descriptorSetImpl.detach();
-    return SLANG_OK;
-}
 
-/* static */VKRenderer::DescriptorSetImpl::Binding::Type VKRenderer::DescriptorSetImpl::_getBindingType(RefObject* ptr)
-{
-    typedef Binding::Type Type;
+    descriptorSetImpl->m_rootConstantData.setCount(layoutImpl->m_rootConstantDataSize);
+    descriptorSetImpl->m_boundObjects.setCount(layoutImpl->m_totalBoundObjectCount);
 
-    if (ptr)
-    {
-        if (dynamic_cast<ResourceView*>(ptr))
-        {
-            return Type::ResourceView;
-        }
-        else if (dynamic_cast<BufferResource*>(ptr))
-        {
-            return Type::BufferResource;
-        }
-        else if (dynamic_cast<SamplerState*>(ptr))
-        {
-            return Type::SamplerState;
-        }
-    }
-    return Type::Unknown;
-}
-
-void VKRenderer::DescriptorSetImpl::_setBinding(Binding::Type type, UInt range, UInt index, RefObject* ptr)
-{
-    SLANG_ASSERT(ptr == nullptr || _getBindingType(ptr) == type);
-
-    const Index numBindings = m_bindings.getCount();
-    for (Index i = 0; i < numBindings; ++i)
-    {
-        Binding& binding = m_bindings[i];
-
-        if (binding.type == type && binding.range == uint32_t(range) && binding.index == uint32_t(index))
-        {
-            if (ptr)
-            {
-                binding.obj = ptr;
-            }
-            else
-            {
-                m_bindings.removeAt(i);
-            }
-
-            return;
-        }
-    }
-
-    // If an entry is not found, and we have a pointer, create an entry
-    if (ptr)
-    {
-        Binding binding;
-        binding.type = type;
-        binding.range = uint32_t(range);
-        binding.index = uint32_t(index);
-        binding.obj = ptr;
-
-        m_bindings.add(binding);
-    }
+    *outDescriptorSet = descriptorSetImpl.detach();
+    return SLANG_OK;
 }
 
 void VKRenderer::DescriptorSetImpl::setConstantBuffer(UInt range, UInt index, BufferResource* buffer)
 {
     auto bufferImpl = (BufferResourceImpl*)buffer;
 
+    SLANG_ASSERT(range < UInt(m_layout->m_ranges.getCount()));
+    auto& rangeInfo = m_layout->m_ranges[range];
+    auto bindingIndex = rangeInfo.vkBindingIndex;
+    auto boundObjectIndex = rangeInfo.arrayIndex + index;
+
     VkDescriptorBufferInfo bufferInfo = {};
     bufferInfo.buffer = bufferImpl->m_buffer.m_buffer;
     bufferInfo.offset = 0;
@@ -2453,19 +2573,24 @@ void VKRenderer::DescriptorSetImpl::setConstantBuffer(UInt range, UInt index, Bu
 
     VkWriteDescriptorSet writeInfo = { VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET };
     writeInfo.dstSet = m_descriptorSet;
-    writeInfo.dstBinding = uint32_t(range);
+    writeInfo.dstBinding = uint32_t(bindingIndex);
     writeInfo.dstArrayElement = uint32_t(index);
     writeInfo.descriptorCount = 1;
-    writeInfo.descriptorType = m_layout->m_ranges[range].descriptorType;
+    writeInfo.descriptorType = rangeInfo.vkDescriptorType;
     writeInfo.pBufferInfo = &bufferInfo;
 
     m_renderer->m_api.vkUpdateDescriptorSets(m_renderer->m_device, 1, &writeInfo, 0, nullptr);
-
-    _setBinding(Binding::Type::BufferResource, range, index, buffer);
+    m_boundObjects[boundObjectIndex] = buffer;
 }
 
 void VKRenderer::DescriptorSetImpl::setResource(UInt range, UInt index, ResourceView* view)
 {
+    SLANG_ASSERT(range < UInt(m_layout->m_ranges.getCount()));
+    auto& rangeInfo = m_layout->m_ranges[range];
+    auto bindingIndex = rangeInfo.vkBindingIndex;
+    auto boundObjectIndex = rangeInfo.arrayIndex + index;
+    auto descriptorType = rangeInfo.vkDescriptorType;
+
     auto viewImpl = (ResourceViewImpl*)view;
     switch (viewImpl->m_type)
     {
@@ -2479,10 +2604,10 @@ void VKRenderer::DescriptorSetImpl::setResource(UInt range, UInt index, Resource
 
             VkWriteDescriptorSet writeInfo = { VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET };
             writeInfo.dstSet = m_descriptorSet;
-            writeInfo.dstBinding = uint32_t(range);
+            writeInfo.dstBinding = uint32_t(bindingIndex);
             writeInfo.dstArrayElement = uint32_t(index);
             writeInfo.descriptorCount = 1;
-            writeInfo.descriptorType = m_layout->m_ranges[range].descriptorType;
+            writeInfo.descriptorType = descriptorType;
             writeInfo.pImageInfo = &imageInfo;
 
             m_renderer->m_api.vkUpdateDescriptorSets(m_renderer->m_device, 1, &writeInfo, 0, nullptr);
@@ -2495,10 +2620,10 @@ void VKRenderer::DescriptorSetImpl::setResource(UInt range, UInt index, Resource
 
             VkWriteDescriptorSet writeInfo = { VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET };
             writeInfo.dstSet = m_descriptorSet;
-            writeInfo.dstBinding = uint32_t(range);
+            writeInfo.dstBinding = uint32_t(bindingIndex);
             writeInfo.dstArrayElement = uint32_t(index);
             writeInfo.descriptorCount = 1;
-            writeInfo.descriptorType = m_layout->m_ranges[range].descriptorType;
+            writeInfo.descriptorType = descriptorType;
             writeInfo.pTexelBufferView = &bufferViewImpl->m_view;
 
             m_renderer->m_api.vkUpdateDescriptorSets(m_renderer->m_device, 1, &writeInfo, 0, nullptr);
@@ -2516,25 +2641,32 @@ void VKRenderer::DescriptorSetImpl::setResource(UInt range, UInt index, Resource
 
             VkWriteDescriptorSet writeInfo = { VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET };
             writeInfo.dstSet = m_descriptorSet;
-            writeInfo.dstBinding = uint32_t(range);
+            writeInfo.dstBinding = uint32_t(bindingIndex);
             writeInfo.dstArrayElement = uint32_t(index);
             writeInfo.descriptorCount = 1;
-            writeInfo.descriptorType = m_layout->m_ranges[range].descriptorType;
+            writeInfo.descriptorType = descriptorType;
             writeInfo.pBufferInfo = &bufferInfo;
 
             m_renderer->m_api.vkUpdateDescriptorSets(m_renderer->m_device, 1, &writeInfo, 0, nullptr);
         }
         break;
-
     }
 
-    _setBinding(Binding::Type::ResourceView, range, index, view);
+    m_boundObjects[boundObjectIndex] = view;
 }
 
 void VKRenderer::DescriptorSetImpl::setSampler(UInt range, UInt index, SamplerState* sampler)
 {
+    SLANG_ASSERT(range < UInt(m_layout->m_ranges.getCount()));
+    auto& rangeInfo = m_layout->m_ranges[range];
+    SLANG_ASSERT(rangeInfo.type == DescriptorSlotType::Sampler);
+    auto bindingIndex = rangeInfo.vkBindingIndex;
+    auto boundObjectIndex = rangeInfo.arrayIndex + index;
+    auto descriptorType = rangeInfo.vkDescriptorType;
+
+    // TODO: Actually bind it!
 
-    _setBinding(Binding::Type::SamplerState, range, index, sampler);
+    m_boundObjects[boundObjectIndex] = sampler;
 }
 
 void VKRenderer::DescriptorSetImpl::setCombinedTextureSampler(
@@ -2543,9 +2675,46 @@ void VKRenderer::DescriptorSetImpl::setCombinedTextureSampler(
     ResourceView*   textureView,
     SamplerState*   sampler)
 {
+    SLANG_ASSERT(range < UInt(m_layout->m_ranges.getCount()));
+    auto& rangeInfo = m_layout->m_ranges[range];
+    SLANG_ASSERT(rangeInfo.type == DescriptorSlotType::CombinedImageSampler);
+    auto bindingIndex = rangeInfo.vkBindingIndex;
+    auto descriptorType = rangeInfo.vkDescriptorType;
+
+    // TODO: Actually bind it!
+
+    // Note: Each entry in a combined texture/sampler range consumes
+    // two entries in the `m_boundObjects` array, since we have
+    // to keep both the texture view and the sampler object live.
+    //
+    auto boundObjectIndex = rangeInfo.arrayIndex + 2 * index;
+    m_boundObjects[boundObjectIndex + 0] = textureView;
+    m_boundObjects[boundObjectIndex + 1] = sampler;
+}
+
+void VKRenderer::DescriptorSetImpl::setRootConstants(
+    UInt range,
+    UInt offset,
+    UInt size,
+    void const* data)
+{
+    // The `range` variabel is the index of one of the descriptor
+    // slot ranges, which had better be a `RootConstant` range.
+    //
+    SLANG_ASSERT(range < UInt(m_layout->m_ranges.getCount()));
+    auto& rangeInfo = m_layout->m_ranges[range];
+    SLANG_ASSERT(rangeInfo.type == DescriptorSlotType::RootConstant);
+
+    // The `arrayIndex` for the descriptor slot range will refer
+    // to a root constant range, which is the range to be set.
+    //
+    auto rootConstantIndex = rangeInfo.arrayIndex;
+    SLANG_ASSERT(rootConstantIndex >= 0);
+    SLANG_ASSERT(rootConstantIndex < m_layout->m_rootConstantRanges.getCount());
+    auto& rootConstantRangeInfo = m_layout->m_rootConstantRanges[rootConstantIndex];
+    SLANG_ASSERT(offset + size <= rootConstantRangeInfo.size);
 
-    _setBinding(Binding::Type::SamplerState, range, index, sampler);
-    _setBinding(Binding::Type::ResourceView, range, index, textureView);
+    memcpy(m_rootConstantData.getBuffer() + rootConstantRangeInfo.offset + offset, data, size);
 }
 
 void VKRenderer::setDescriptorSet(PipelineType pipelineType, PipelineLayout* layout, UInt index, DescriptorSet* descriptorSet)
diff --git a/tools/gfx/vulkan/vk-api.h b/tools/gfx/vulkan/vk-api.h
index 6f12d5ddb..58bcc89fd 100644
--- a/tools/gfx/vulkan/vk-api.h
+++ b/tools/gfx/vulkan/vk-api.h
@@ -84,6 +84,7 @@ namespace gfx {
     x(vkCmdEndRenderPass) \
     x(vkCmdPipelineBarrier) \
     x(vkCmdCopyBufferToImage)\
+    x(vkCmdPushConstants) \
     \
     x(vkCreateFence) \
     x(vkDestroyFence) \