Update gfx back-ends to handle static specialization (#1826)

* Update gfx back-ends to handle static specialization

The main goal here is to make the D3D11, D3D12 and Vulkan back-ends support static specialization of interface types in the case where the data for the type won't "fit" in the pre-allocated space for existential values. This includes all cases where the concrete type being specialized to has resources/samplers/etc., as well as any cases where its ordinary/uniform data exceeds the space available.

(Note that the CPU and CUDA targets don't need this work since they can (in theory) support arbitrary-size data in the fixed-size existential payload by using pointer indirection. Actually supporting indirection in those cases should be a distinct change)

The Slang compiler already performs layout for programs that have this kind of data that doesn't "fit," and it lays them out using an idea of "pending" type layouts. Basically, a type that contains some amount of specialized interface-type fields will produce both a "primary" type layout that just covers the data for the unspecialized case, as well as "pending" type layout that describes the layout for all the extra data needed by specialization.
When laying out a `ConstantBuffer<X>` or `ParameterBlocK<X>` ("CB" or "PB"), the front-end will try to place as much of that "pending" data into the layout of the buffer/block itself as is possible. That means that both CBs and PBs will be able to allocate trailing bytes for any ordinary data in the "pending" layout. PBs will be able to allocate any trailing resources/samplers into their layout, but for CBs they will spill out to be part of the pending layout for the buffer itself.

In order for the back-ends to properly handle pending data, they need to *either* assume the exact layout rules used by the front-end and try to reproduce them (e.g., by iterating over binding ranges and sub-objects in the exact same order that front-end layout would enumerate them), *or* they need to respect the reflection information produced by the front-end. This change takes the latter approach, trying to make only minimal assumptions about the layout rules being used. This choice is motivated by wanting to decouple the `gfx` implementation from the compiler front-end, especially insofar as this work has made me question whether the current layout rules are the best ones possible.

A common theme across all the implementations is to have a fixed-size type that can represent "binding offsets" for the chosen back-end. The offset type has fields that depend on the API-specific way bindings are indexed; e.g., for D3D11 it has offsets for CBV, SRV, UAV, and sampler bindings. This fixed-size offset type can be filled in based on Slang reflecton information, and then used to compute derived offsets with just a few add operations.

The simple offset type for each API is then extended to produce an offset type that includes both the offsets for "primary" data and also the offsets for "pending" data. Most logic that traffics in offsets doesn't have to know about this more complicated representation.

Making consistent use of these offsets required that I pretty much rewrite the logic that actually applies shader objects to the API state. Doing so might be lowering the efficiency of the system in the near term, but the increase in clarity was important for getting the work done, and it seems like it will also be important if/when we start trying to perform special-case optimizations around root and entry-point parameter setting.

While there are many API-specific differences, we can identify a repeated pattern where many steps, whether applying parameters to the pipeline stage or constructing signatures / layouts, can be broken down into three main operations on `ShaderObject`s or their layouts:

* `*AsValue()` is the core operation, and is the one used for the `ExistentialValue` case most of the time. It ignores the ordinary data in the object, and instead processes all nested binding ranges (for resources/smaplers) and sub-objects.

* `*AsConstantBuffer()` handles the `ConstntBuffer<X>` case, by dealing with the implicit buffer for ordinary data (if it is needed) and then delegates to the `*AsValue()` case.

* `*AsParameterBlock()` handles the `ParameterBlock<X>` case, by allocating/preparing/etc. any descriptor tables/sets that would be required for the current object/layout and then delegating to `*AsConstantBuffer()` to do the rest

The idea is that by having the parameter block case delegate to the constant buffer case, which delegates to the value/existential case, we can streamline a lot of the logic so that it doesn't seem quite as full of special cases.

Note: When preparing this pull request I spent a reasonable amount of time trying to clean up the D3D11 and Vulkan implementations, so they are probably the easiest to read and understand when it comes to the new code. Doing the cleanup work also helped to work out some weird corner case bugs/issues. In contrast, the D3D12 path hasn't had as much attention given to cleanliness and comments, so it really needs some attention down the line to get things into a state that is easier to understand.

* fixup: remove debugging code spotted in review

6 files changed, 205 insertions, 49 deletions

diff --git a/source/core/slang-string.h b/source/core/slang-string.h
index f42e78ac4..6b9a73bcb 100644
--- a/source/core/slang-string.h
+++ b/source/core/slang-string.h
@@ -279,7 +279,7 @@ namespace Slang
     {
     public:
         UnownedTerminatedStringSlice(char const* b)
-            : UnownedStringSlice(b, b + strlen(b))
+            : UnownedStringSlice(b, b + (b?strlen(b):0))
         {}
     };
 
diff --git a/source/slang/slang-ir-legalize-types.cpp b/source/slang/slang-ir-legalize-types.cpp
index 9b95e069d..1cd94decb 100644
--- a/source/slang/slang-ir-legalize-types.cpp
+++ b/source/slang/slang-ir-legalize-types.cpp
@@ -2435,7 +2435,7 @@ static LegalVal legalizeGlobalVar(
 
     default:
         {
-            context->insertBeforeGlobal = irGlobalVar->getNextInst();
+            context->insertBeforeGlobal = irGlobalVar;
 
             IRGlobalNameInfo globalNameInfo;
             globalNameInfo.globalVar = irGlobalVar;
@@ -2480,7 +2480,7 @@ static LegalVal legalizeGlobalParam(
 
     default:
         {
-            context->insertBeforeGlobal = irGlobalParam->getNextInst();
+            context->insertBeforeGlobal = irGlobalParam;
 
             LegalVarChainLink varChain(LegalVarChain(), varLayout);
 
diff --git a/source/slang/slang-ir-lower-tuple-types.cpp b/source/slang/slang-ir-lower-tuple-types.cpp
index a155ae8d3..429058edf 100644
--- a/source/slang/slang-ir-lower-tuple-types.cpp
+++ b/source/slang/slang-ir-lower-tuple-types.cpp
@@ -121,6 +121,18 @@ namespace Slang
             inst->removeAndDeallocate();
         }
 
+        void processTupleType(IRTupleType* inst)
+        {
+            IRBuilder builderStorage;
+            auto builder = &builderStorage;
+            builder->sharedBuilder = &sharedBuilderStorage;
+            builder->setInsertBefore(inst);
+
+            auto loweredTupleInfo = getLoweredTupleType(builder, inst);
+            SLANG_ASSERT(loweredTupleInfo);
+            SLANG_UNUSED(loweredTupleInfo);
+        }
+
         void processInst(IRInst* inst)
         {
             switch (inst->getOp())
@@ -131,6 +143,9 @@ namespace Slang
             case kIROp_GetTupleElement:
                 processGetTupleElement((IRGetTupleElement*)inst);
                 break;
+            case kIROp_TupleType:
+                processTupleType((IRTupleType*)inst);
+                break;
             default:
                 break;
             }
diff --git a/source/slang/slang-reflection-api.cpp b/source/slang/slang-reflection-api.cpp
index e68efde06..d75920fdb 100644
--- a/source/slang/slang-reflection-api.cpp
+++ b/source/slang/slang-reflection-api.cpp
@@ -1056,22 +1056,104 @@ SLANG_API SlangReflectionVariableLayout* spReflectionTypeLayout_getSpecializedTy
     }
 }
 
+SLANG_API SlangInt spReflectionType_getSpecializedTypeArgCount(SlangReflectionType* inType)
+{
+    auto type = convert(inType);
+    if(!type) return 0;
+
+    auto specializedType = as<ExistentialSpecializedType>(type);
+    if(!specializedType) return 0;
+
+    return specializedType->args.getCount();
+}
+
+SLANG_API SlangReflectionType* spReflectionType_getSpecializedTypeArgType(SlangReflectionType* inType, SlangInt index)
+{
+    auto type = convert(inType);
+    if(!type) return nullptr;
+
+    auto specializedType = as<ExistentialSpecializedType>(type);
+    if(!specializedType) return nullptr;
+
+    if(index < 0) return nullptr;
+    if(index >= specializedType->args.getCount()) return nullptr;
+
+    auto argType = as<Type>(specializedType->args[index].val);
+    return convert(argType);
+}
+
 namespace Slang
 {
+        /// A link in a chain of `VarLayout`s that can be used to compute offset information for a nested field
     struct BindingRangePathLink
     {
+        BindingRangePathLink()
+        {}
+
         BindingRangePathLink(
             BindingRangePathLink*   parent,
             VarLayout*              var)
-            : parent(parent)
-            , var(var)
+            : var(var)
+            , parent(parent)
         {}
 
-        BindingRangePathLink*   parent;
-        VarLayout*              var;
+            /// The inner-most variable that contributes to the offset along this path
+        VarLayout*              var = nullptr;
+
+            /// The next outer link along the path
+        BindingRangePathLink*   parent = nullptr;
+    };
+
+        /// A path leading to some nested field, with both parimary and "pending" data offsets
+    struct BindingRangePath
+    {
+
+            /// The chain of variables that defines the "primary" offset of a nested field
+        BindingRangePathLink* primary = nullptr;
+
+            /// The chain of variables that defines the offset for "pending" data of a nested field
+        BindingRangePathLink* pending = nullptr;
     };
 
+        /// A helper type to construct a `BindingRangePath` that extends an existing path
+    struct ExtendedBindingRangePath : BindingRangePath
+    {
+            /// Construct a path that extends `parent` with offset information from `varLayout`
+        ExtendedBindingRangePath(
+            BindingRangePath const& parent,
+            VarLayout*              varLayout)
+        {
+            SLANG_ASSERT(varLayout);
+
+            // We always add another link to the primary chain.
+            //
+            primaryLink = BindingRangePathLink(parent.primary, varLayout);
+            primary = &primaryLink;
+
+            // If the `varLayout` provided has any offset information
+            // for pending data, then we also add a link to the pending
+            // chain, but otherwise we re-use the pending chain from
+            // the parent path.
+            //
+            if(auto pendingLayout = varLayout->pendingVarLayout)
+            {
+                pendingLink = BindingRangePathLink(parent.pending, pendingLayout);
+                pending = &pendingLink;
+            }
+            else
+            {
+                pending = parent.pending;
+            }
+        }
+
+            /// Storage for a link in the primary chain, if needed
+        BindingRangePathLink primaryLink;
+
+            /// Storage for a link in the pending chain, if needed
+        BindingRangePathLink pendingLink;
+    };
 
+        /// Calculate the offset for resources of the given `kind` in the `path`.
     Int _calcIndexOffset(BindingRangePathLink* path, LayoutResourceKind kind)
     {
         Int result = 0;
@@ -1085,6 +1167,7 @@ namespace Slang
         return result;
     }
 
+        /// Calculate the regsiter space / set for resources of the given `kind` in the `path`.
     Int _calcSpaceOffset(BindingRangePathLink* path, LayoutResourceKind kind)
     {
         Int result = 0;
@@ -1208,12 +1291,24 @@ namespace Slang
         }
     }
 
-
-
     SlangBindingType _calcBindingType(
         Slang::TypeLayout*  typeLayout,
         LayoutResourceKind  kind)
     {
+        // At the type level, a push-constant buffer and a regular constant
+        // buffer are currently not distinct, so we need to detect push
+        // constant buffers/ranges before we inspect the `typeLayout` to
+        // avoid reflecting them all as ordinary constant buffers.
+        //
+        switch(kind)
+        {
+        default:
+            break;
+
+        case LayoutResourceKind::PushConstantBuffer:
+            return SLANG_BINDING_TYPE_PUSH_CONSTANT;
+        }
+
         // If the type or type layout implies a specific binding type
         // (e.g., a `Texture2D` implies a texture binding), then we
         // will always favor the binding type implied.
@@ -1267,7 +1362,42 @@ namespace Slang
             return index;
         }
 
-        void addRangesRec(TypeLayout* typeLayout, BindingRangePathLink* path, LayoutSize multiplier)
+            /// Create a single `VarLayout` for `typeLayout` that summarizes all of the offset information in `path`.
+            ///
+            /// Note: This function does not handle "pending" layout information.
+        RefPtr<VarLayout> _createSimpleOffsetVarLayout(TypeLayout* typeLayout, BindingRangePathLink* path)
+        {
+            SLANG_ASSERT(typeLayout);
+
+            RefPtr<VarLayout> varLayout = new VarLayout();
+            varLayout->typeLayout = typeLayout;
+
+            for(auto typeResInfo : typeLayout->resourceInfos)
+            {
+                auto kind = typeResInfo.kind;
+                auto varResInfo = varLayout->findOrAddResourceInfo(kind);
+                varResInfo->index = _calcIndexOffset(path, kind);
+                varResInfo->space = _calcSpaceOffset(path, kind);
+            }
+
+            return varLayout;
+        }
+
+            /// Create a single `VarLayout` for `typeLayout` that summarizes all of the offset information in `path`.
+        RefPtr<VarLayout> createOffsetVarLayout(TypeLayout* typeLayout, BindingRangePath const& path)
+        {
+            auto primaryVarLayout = _createSimpleOffsetVarLayout(typeLayout, path.primary);
+            SLANG_ASSERT(primaryVarLayout);
+
+            if(auto pendingDataTypeLayout = typeLayout->pendingDataTypeLayout)
+            {
+                primaryVarLayout->pendingVarLayout = _createSimpleOffsetVarLayout(pendingDataTypeLayout, path.pending);
+            }
+
+            return primaryVarLayout;
+        }
+
+        void addRangesRec(TypeLayout* typeLayout, BindingRangePath const& path, LayoutSize multiplier)
         {
             if( auto structTypeLayout = as<StructTypeLayout>(typeLayout) )
             {
@@ -1287,9 +1417,8 @@ namespace Slang
 
                     auto fieldTypeLayout = fieldVarLayout->getTypeLayout();
 
-
-                    BindingRangePathLink fieldLink(path, fieldVarLayout);
-                    addRangesRec(fieldTypeLayout, &fieldLink, multiplier);
+                    ExtendedBindingRangePath fieldPath(path, fieldVarLayout);
+                    addRangesRec(fieldTypeLayout, fieldPath, multiplier);
                 }
                 return;
             }
@@ -1336,7 +1465,7 @@ namespace Slang
                 {
                     if( spaceOffset == -1 )
                     {
-                        spaceOffset = _calcSpaceOffset(path, kind);
+                        spaceOffset = _calcSpaceOffset(path.primary, kind);
                     }
 
                     kind = resInfo.kind;
@@ -1382,24 +1511,17 @@ namespace Slang
                 bindingRange.firstDescriptorRangeIndex = 0;
                 bindingRange.descriptorRangeCount = 0;
 
-                if( kind == LayoutResourceKind::PushConstantBuffer )
-                {
-                    if(auto resInfo = parameterGroupTypeLayout->elementVarLayout->typeLayout->FindResourceInfo(LayoutResourceKind::Uniform))
-                    {
-                        bindingRange.count *= resInfo->count;
-                    }
-                }
-
                 // Every parameter group will introduce a sub-object range,
                 // which will include bindings based on the type of data
                 // inside the sub-object.
                 //
                 TypeLayout::ExtendedInfo::SubObjectRangeInfo subObjectRange;
                 subObjectRange.bindingRangeIndex = bindingRangeIndex;
+                subObjectRange.offsetVarLayout = createOffsetVarLayout(typeLayout, path);
                 subObjectRange.spaceOffset = 0;
-                if (kind == LayoutResourceKind::RegisterSpace && path)
+                if (kind == LayoutResourceKind::RegisterSpace && path.primary)
                 {
-                    auto resInfo = path->var->FindResourceInfo(LayoutResourceKind::RegisterSpace);
+                    auto resInfo = path.primary->var->FindResourceInfo(LayoutResourceKind::RegisterSpace);
                     subObjectRange.spaceOffset = resInfo->index;
                 }
                 // It is possible that the sub-object has descriptor ranges
@@ -1488,16 +1610,7 @@ namespace Slang
                                 descriptorRange.kind = resInfo.kind;
                                 descriptorRange.bindingType = _calcBindingType(typeLayout, resInfo.kind);
                                 descriptorRange.count = multiplier;
-                                descriptorRange.indexOffset = _calcIndexOffset(path, resInfo.kind);
-
-                                if( resInfo.kind == LayoutResourceKind::PushConstantBuffer )
-                                {
-                                    if(auto uniformResInfo = parameterGroupTypeLayout->elementVarLayout->typeLayout->FindResourceInfo(LayoutResourceKind::Uniform))
-                                    {
-                                        descriptorRange.count *= uniformResInfo->count;
-                                    }
-                                }
-
+                                descriptorRange.indexOffset = _calcIndexOffset(path.primary, resInfo.kind);
                                 descriptorSet->descriptorRanges.add(descriptorRange);
                             }
                         }
@@ -1548,12 +1661,12 @@ namespace Slang
                             // up binding information, to ensure that the descriptor ranges
                             // that get enumerated here have correct register/binding offsets.
                             //
-                            BindingRangePathLink elementPath(path, parameterGroupTypeLayout->elementVarLayout);
+                            ExtendedBindingRangePath elementPath(path, parameterGroupTypeLayout->elementVarLayout);
 
                             Index bindingRangeCountBefore = m_extendedInfo->m_bindingRanges.getCount();
                             Index subObjectRangeCountBefore = m_extendedInfo->m_subObjectRanges.getCount();
 
-                            addRangesRec(parameterGroupTypeLayout->elementVarLayout->typeLayout, &elementPath, multiplier);
+                            addRangesRec(parameterGroupTypeLayout->elementVarLayout->typeLayout, elementPath, multiplier);
 
                             m_extendedInfo->m_bindingRanges.setCount(bindingRangeCountBefore);
                             m_extendedInfo->m_subObjectRanges.setCount(subObjectRangeCountBefore);
@@ -1576,10 +1689,6 @@ namespace Slang
                 // An `interface` type should introduce a binding range and a matching
                 // sub-object range.
                 //
-                // We currently do *not* allocate any descriptor ranges to represent
-                // an interface-type field, since the only direct storage required
-                // is all uniform/ordinary data.
-                //
                 TypeLayout::ExtendedInfo::BindingRangeInfo bindingRange;
                 bindingRange.leafTypeLayout = typeLayout;
                 bindingRange.bindingType = SLANG_BINDING_TYPE_EXISTENTIAL_VALUE;
@@ -1590,10 +1699,7 @@ namespace Slang
 
                 TypeLayout::ExtendedInfo::SubObjectRangeInfo subObjectRange;
                 subObjectRange.bindingRangeIndex = m_extendedInfo->m_bindingRanges.getCount();
-
-                // TODO: if we have "pending" layout information that tells us where
-                // data for the sub-object range has been allocated, then we need
-                // a way to reference that data here.
+                subObjectRange.offsetVarLayout = createOffsetVarLayout(typeLayout, path);
 
                 m_extendedInfo->m_bindingRanges.add(bindingRange);
                 m_extendedInfo->m_subObjectRanges.add(subObjectRange);
@@ -1730,8 +1836,8 @@ namespace Slang
                     // `resInfo` representing resouce usage.
                     //
                     auto count = resInfo.count * multiplier;
-                    auto indexOffset = _calcIndexOffset(path, kind);
-                    auto spaceOffset = _calcSpaceOffset(path, kind);
+                    auto indexOffset = _calcIndexOffset(path.primary, kind);
+                    auto spaceOffset = _calcSpaceOffset(path.primary, kind);
 
                     TypeLayout::ExtendedInfo::DescriptorRangeInfo descriptorRange;
                     descriptorRange.kind = kind;
@@ -1767,7 +1873,8 @@ namespace Slang
             context.m_typeLayout = typeLayout;
             context.m_extendedInfo = extendedInfo;
 
-            context.addRangesRec(typeLayout, nullptr, 1);
+            BindingRangePath rootPath;
+            context.addRangesRec(typeLayout, rootPath, 1);
 
             typeLayout->m_extendedInfo = extendedInfo;
         }
@@ -2033,6 +2140,25 @@ SLANG_API SlangInt spReflectionTypeLayout_getSubObjectRangeSpaceOffset(
     return extTypeLayout->m_subObjectRanges[subObjectRangeIndex].spaceOffset;
 }
 
+SLANG_API SlangReflectionVariableLayout* spReflectionTypeLayout_getSubObjectRangeOffset(
+    SlangReflectionTypeLayout* inTypeLayout,
+    SlangInt subObjectRangeIndex)
+{
+    auto typeLayout = convert(inTypeLayout);
+    if (!typeLayout)
+        return 0;
+
+    auto extTypeLayout = Slang::getExtendedTypeLayout(typeLayout);
+
+    if (subObjectRangeIndex < 0)
+        return 0;
+    if (subObjectRangeIndex >= extTypeLayout->m_subObjectRanges.getCount())
+        return 0;
+
+    return convert(extTypeLayout->m_subObjectRanges[subObjectRangeIndex].offsetVarLayout);
+}
+
+
 
 #if 0
 SLANG_API SlangInt spReflectionTypeLayout_getBindingRangeSubObjectRangeIndex(SlangReflectionTypeLayout* inTypeLayout, SlangInt index)
diff --git a/source/slang/slang-type-layout.cpp b/source/slang/slang-type-layout.cpp
index a015bfd78..f21722bd7 100644
--- a/source/slang/slang-type-layout.cpp
+++ b/source/slang/slang-type-layout.cpp
@@ -4013,6 +4013,12 @@ static TypeLayoutResult _createTypeLayout(
         typeLayout->type = type;
         typeLayout->rules = rules;
 
+        for( auto resInfo : baseTypeLayoutResult.layout->resourceInfos )
+        {
+            if(resInfo.kind != LayoutResourceKind::Uniform)
+                typeLayout->addResourceUsage(resInfo);
+        }
+
         RefPtr<VarLayout> pendingDataVarLayout = new VarLayout();
         if(auto pendingDataTypeLayout = baseTypeLayoutResult.layout->pendingDataTypeLayout)
         {
@@ -4033,6 +4039,7 @@ static TypeLayoutResult _createTypeLayout(
                 {
                     if(auto primaryResInfo = baseTypeLayoutResult.layout->FindResourceInfo(kind))
                         index = primaryResInfo->count.getFiniteValue();
+                    typeLayout->addResourceUsage(pendingResInfo);
                 }
                 pendingDataVarLayout->AddResourceInfo(kind)->index = index;
             }
@@ -4041,6 +4048,12 @@ static TypeLayoutResult _createTypeLayout(
         typeLayout->baseTypeLayout = baseTypeLayoutResult.layout;
         typeLayout->pendingDataVarLayout = pendingDataVarLayout;
 
+        typeLayout->uniformAlignment = info.alignment;
+        if( info.size != 0 )
+        {
+            typeLayout->addResourceUsage(LayoutResourceKind::Uniform, info.size);
+        }
+
         return makeTypeLayoutResult(typeLayout);
     }
 
diff --git a/source/slang/slang-type-layout.h b/source/slang/slang-type-layout.h
index 88ba72003..52b575787 100644
--- a/source/slang/slang-type-layout.h
+++ b/source/slang/slang-type-layout.h
@@ -310,6 +310,7 @@ struct SimpleArrayLayoutInfo : SimpleLayoutInfo
 };
 
 struct LayoutRulesImpl;
+class VarLayout;
 
 // Base class for things that store layout info
 class Layout : public RefObject
@@ -436,8 +437,9 @@ public:
 
         struct SubObjectRangeInfo
         {
-            Int     bindingRangeIndex;
-            Int     spaceOffset;
+            Int                 bindingRangeIndex = 0;
+            Int                 spaceOffset = 0;
+            RefPtr<VarLayout>   offsetVarLayout;
         };
 
         List<RefPtr<DescriptorSetInfo>> m_descriptorSets;