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>

4 files changed, 43 insertions, 23 deletions

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;