Clean up the way we stride over sub-object ranges for D3D11 and Vulkan (#1829)

This change applies to the case where we have a sub-object range with more than one object in it (so arrays of constant buffers, parameter blocks, or existentials). It is worth noting that we don't really seem to have any tests covering this case right now, so it is entirely possible that things are busted already, and/or that this change doesn't work the way I think it does.

When we go to actually bind the state from a sub-object range into the pipeline, we care about:

* The offset to where the first object should be written
* The "stride" between consecutive objects

We were already capturing the offset information from Slang reflection data, and the same was true for the part of the offset that pertains to "pending" ordinary data. For other cases, though, we were manually incrementing the offset by values computed manually in `gfx` for Vulkan, and we were just skipping the offsetting step entirely for D3D11.

With this change we extract more complete stride information from the Slang reflection data and use that instead of ad hoc computations. Hopefully this data is correct, and if it isn't we can consider whether it needs fixing at the Slang reflection level rather than in `gfx`.

This change doesn't apply to the OpenGL back-end (which hasn't been updated to match the new approach to static specialization at all) or to the D3D12 back end (which has been updated a bit, but probably needs other cleanup work to be done first to bring it in line).

4 files changed, 86 insertions, 37 deletions

diff --git a/slang.h b/slang.h
index f8f5c3073..a5982a731 100644
--- a/slang.h
+++ b/slang.h
@@ -1968,6 +1968,7 @@ extern "C"
     SLANG_API SlangReflectionType* spReflectionTypeLayout_GetType(SlangReflectionTypeLayout* type);
     SLANG_API SlangTypeKind spReflectionTypeLayout_getKind(SlangReflectionTypeLayout* type);
     SLANG_API size_t spReflectionTypeLayout_GetSize(SlangReflectionTypeLayout* type, SlangParameterCategory category);
+    SLANG_API size_t spReflectionTypeLayout_GetStride(SlangReflectionTypeLayout* type, SlangParameterCategory category);
     SLANG_API int32_t spReflectionTypeLayout_getAlignment(SlangReflectionTypeLayout* type, SlangParameterCategory category);
 
     SLANG_API SlangReflectionVariableLayout* spReflectionTypeLayout_GetFieldByIndex(SlangReflectionTypeLayout* type, unsigned index);
@@ -2419,6 +2420,11 @@ namespace slang
             return spReflectionTypeLayout_GetSize((SlangReflectionTypeLayout*) this, category);
         }
 
+        size_t getStride(SlangParameterCategory category = SLANG_PARAMETER_CATEGORY_UNIFORM)
+        {
+            return spReflectionTypeLayout_GetStride((SlangReflectionTypeLayout*) this, category);
+        }
+
         int32_t getAlignment(SlangParameterCategory category = SLANG_PARAMETER_CATEGORY_UNIFORM)
         {
             return spReflectionTypeLayout_getAlignment((SlangReflectionTypeLayout*) this, category);
diff --git a/source/slang/slang-reflection-api.cpp b/source/slang/slang-reflection-api.cpp
index d75920fdb..b74a019e3 100644
--- a/source/slang/slang-reflection-api.cpp
+++ b/source/slang/slang-reflection-api.cpp
@@ -819,6 +819,35 @@ namespace
 
         return SLANG_UNBOUNDED_SIZE;
     }
+
+    static int32_t getAlignment(TypeLayout* typeLayout, SlangParameterCategory category)
+    {
+        if( category == SLANG_PARAMETER_CATEGORY_UNIFORM )
+        {
+            return int32_t(typeLayout->uniformAlignment);
+        }
+        else
+        {
+            return 1;
+        }
+    }
+
+    static size_t getStride(TypeLayout* typeLayout, SlangParameterCategory category)
+    {
+        auto info = typeLayout->FindResourceInfo(LayoutResourceKind(category));
+        if(!info) return 0;
+
+        auto size = info->count;
+        if(size.isInfinite())
+            return SLANG_UNBOUNDED_SIZE;
+
+        size_t finiteSize = size.getFiniteValue();
+        size_t alignment = getAlignment(typeLayout, category);
+        SLANG_ASSERT(alignment >= 1);
+
+        auto stride = (finiteSize + (alignment-1)) & ~(alignment-1);
+        return stride;
+    }
 }
 
 SLANG_API size_t spReflectionTypeLayout_GetSize(SlangReflectionTypeLayout* inTypeLayout, SlangParameterCategory category)
@@ -832,19 +861,20 @@ SLANG_API size_t spReflectionTypeLayout_GetSize(SlangReflectionTypeLayout* inTyp
     return getReflectionSize(info->count);
 }
 
+SLANG_API size_t spReflectionTypeLayout_GetStride(SlangReflectionTypeLayout* inTypeLayout, SlangParameterCategory category)
+{
+    auto typeLayout = convert(inTypeLayout);
+    if(!typeLayout) return 0;
+
+    return getStride(typeLayout, category);
+}
+
 SLANG_API int32_t spReflectionTypeLayout_getAlignment(SlangReflectionTypeLayout* inTypeLayout, SlangParameterCategory category)
 {
     auto typeLayout = convert(inTypeLayout);
     if(!typeLayout) return 0;
 
-    if( category == SLANG_PARAMETER_CATEGORY_UNIFORM )
-    {
-        return int32_t(typeLayout->uniformAlignment);
-    }
-    else
-    {
-        return 1;
-    }
+    return getAlignment(typeLayout, category);
 }
 
 SLANG_API SlangReflectionVariableLayout* spReflectionTypeLayout_GetFieldByIndex(SlangReflectionTypeLayout* inTypeLayout, unsigned index)
diff --git a/tools/gfx/d3d11/render-d3d11.cpp b/tools/gfx/d3d11/render-d3d11.cpp
index 415e0b833..6f6ba4faf 100644
--- a/tools/gfx/d3d11/render-d3d11.cpp
+++ b/tools/gfx/d3d11/render-d3d11.cpp
@@ -550,6 +550,18 @@ protected:
             }
         }
 
+            /// Create an offset based on size/stride information in the given Slang `typeLayout`
+        SimpleBindingOffset(slang::TypeLayoutReflection* typeLayout)
+        {
+            if(typeLayout)
+            {
+                cbv     = (uint32_t) typeLayout->getSize(SLANG_PARAMETER_CATEGORY_CONSTANT_BUFFER);
+                srv     = (uint32_t) typeLayout->getSize(SLANG_PARAMETER_CATEGORY_SHADER_RESOURCE);
+                uav     = (uint32_t) typeLayout->getSize(SLANG_PARAMETER_CATEGORY_UNORDERED_ACCESS);
+                sampler = (uint32_t) typeLayout->getSize(SLANG_PARAMETER_CATEGORY_SAMPLER_STATE);
+            }
+        }
+
             /// Add any values in the given `offset`
         void operator+=(SimpleBindingOffset const& offset)
         {
@@ -594,6 +606,12 @@ protected:
             , pending(varLayout->getPendingDataLayout())
         {}
 
+            /// Create an offset based on size/stride information in the given Slang `typeLayout`
+        BindingOffset(slang::TypeLayoutReflection* typeLayout)
+            : SimpleBindingOffset(typeLayout)
+            , pending(typeLayout->getPendingDataTypeLayout())
+        {}
+
             /// Add any values in the given `offset`
         void operator+=(SimpleBindingOffset const& offset)
         {
@@ -673,16 +691,17 @@ protected:
         };
 
             /// Stride information for a sub-object range
-        struct SubObjectRangeStride
+        struct SubObjectRangeStride : BindingOffset
         {
             SubObjectRangeStride()
             {}
 
             SubObjectRangeStride(slang::TypeLayoutReflection* typeLayout)
+                : BindingOffset(typeLayout)
             {
                 if(auto pendingLayout = typeLayout->getPendingDataTypeLayout())
                 {
-                    pendingOrdinaryData = (uint32_t) pendingLayout->getSize(SLANG_PARAMETER_CATEGORY_UNIFORM);
+                    pendingOrdinaryData = (uint32_t) typeLayout->getStride();
                 }
             }
 
@@ -1747,6 +1766,11 @@ protected:
                 BindingOffset rangeOffset = offset;
                 rangeOffset += subObjectRange.offset;
 
+                // Similarly, the "stride" between consecutive objects in
+                // the range was also pre-computed.
+                //
+                BindingOffset rangeStride = subObjectRange.stride;
+
                 switch(bindingRange.bindingType)
                 {
                 // For D3D11-compatible compilation targets, the Slang compiler
@@ -1765,8 +1789,7 @@ protected:
                             //
                             subObject->bindAsConstantBuffer(context, objOffset, subObjectLayout);
 
-                            // TODO: We need to update `objOffset` by the stride for
-                            // this range.
+                            objOffset += rangeStride;
                         }
                     }
                     break;
@@ -1782,15 +1805,15 @@ protected:
                         // As a result, the offset for the first object in the range
                         // will come from the `pending` part of the range's offset.
                         //
-                        BindingOffset objOffset = BindingOffset(rangeOffset.pending);
+                        SimpleBindingOffset objOffset = rangeOffset.pending;
+                        SimpleBindingOffset objStride = rangeStride.pending;
 
                         for(Index i = 0; i < count; ++i)
                         {
                             auto subObject = m_objects[ baseIndex + i ];
-                            subObject->bindAsValue(context, objOffset, subObjectLayout);
+                            subObject->bindAsValue(context, BindingOffset(objOffset), subObjectLayout);
 
-                            // TODO: We need to update `objOffset` by the stride for
-                            // this range.
+                            objOffset += objStride;
                         }
                     }
                     break;
diff --git a/tools/gfx/vulkan/render-vk.cpp b/tools/gfx/vulkan/render-vk.cpp
index a0860f0e5..f330d95a8 100644
--- a/tools/gfx/vulkan/render-vk.cpp
+++ b/tools/gfx/vulkan/render-vk.cpp
@@ -853,7 +853,7 @@ public:
         };
 
             /// Stride information for a sub-object range
-        struct SubObjectRangeStride
+        struct SubObjectRangeStride : BindingOffset
         {
             SubObjectRangeStride()
             {}
@@ -862,7 +862,7 @@ public:
             {
                 if(auto pendingLayout = typeLayout->getPendingDataTypeLayout())
                 {
-                    pendingOrdinaryData = (uint32_t) pendingLayout->getSize(SLANG_PARAMETER_CATEGORY_UNIFORM);
+                    pendingOrdinaryData = (uint32_t) pendingLayout->getStride();
                 }
             }
 
@@ -3095,6 +3095,8 @@ public:
                 BindingOffset rangeOffset = offset;
                 rangeOffset += subObjectRange.offset;
 
+                BindingOffset rangeStride = subObjectRange.stride;
+
                 switch( bindingRangeInfo.bindingType )
                 {
                 case slang::BindingType::ConstantBuffer:
@@ -3114,12 +3116,7 @@ public:
                             // sure to increment the offset for each subsequent object
                             // by the appropriate stride.
                             //
-                            // TODO: We should pre-compute these and simply have
-                            // `subObjectRange.stride` to go with `subObjectRange.offset`.
-                            //
-                            objOffset.binding += subObjectLayout->getTotalBindingCount();
-                            objOffset.childSet += subObjectLayout->getChildDescriptorSetCount();
-                            objOffset.pushConstantRange += subObjectLayout->getTotalPushConstantRangeCount();
+                            objOffset += rangeStride;
                         }
                     }
                     break;
@@ -3135,12 +3132,7 @@ public:
                             ShaderObjectImpl* subObject = m_objects[baseIndex + i];
                             subObject->bindAsParameterBlock(encoder, context, objOffset, subObjectLayout);
 
-                            // The logic for striding from one sub-object to another is also
-                            // different, given the differences in how constant buffers and
-                            // parameter blocks are allocated.
-                            //
-                            objOffset.childSet += subObjectLayout->getChildDescriptorSetCount();
-                            objOffset.pushConstantRange += subObjectLayout->getTotalPushConstantRangeCount();
+                            objOffset += rangeStride;
                         }
                     }
                     break;
@@ -3153,14 +3145,15 @@ public:
                     //
                     if( subObjectLayout )
                     {
-                        // Second, the offset where we want to start bindign for existential-type
+                        // Second, the offset where we want to start binding for existential-type
                         // ranges is a bit different, because we don't wnat to bind at the "primary"
                         // offset that got passed down, but instead at the "pending" offset.
                         //
                         // For the purposes of nested binding, what used to be the pending offset
                         // will now be used as the primary offset.
                         //
-                        BindingOffset objOffset = BindingOffset(rangeOffset.pending);
+                        SimpleBindingOffset objOffset = rangeOffset.pending;
+                        SimpleBindingOffset objStride = rangeStride.pending;
                         for (uint32_t i = 0; i < count; ++i)
                         {
                             // An existential-type sub-object is always bound just as a value,
@@ -3170,11 +3163,8 @@ public:
                             // already.
                             //
                             ShaderObjectImpl* subObject = m_objects[baseIndex + i];
-                            subObject->bindAsValue(encoder, context, objOffset, subObjectLayout);
-
-                            objOffset.binding += subObjectLayout->getTotalBindingCount();
-                            objOffset.childSet += subObjectLayout->getChildDescriptorSetCount();
-                            objOffset.pushConstantRange += subObjectLayout->getTotalPushConstantRangeCount();
+                            subObject->bindAsValue(encoder, context, BindingOffset(objOffset), subObjectLayout);
+                            objOffset += objStride;
                         }
                     }
                     break;