Eliminate the specializeProgramLayout() function (#869)

The `specializeProgramLayout()` primarily existed to support global generic parameters.
The guiding idea of the design had been that plugging in concrete types for global generic parameters should only affect the layout/binding of shader parameters that depend on the global generic type parameters.
All other shader parameters should keep the same layout/location across all specializations.

This idea was implemented by conceptually having two phases of layout:

* A generic-argument-independent phase would do layout on all the shader parameters that *don't* depend on global generic type parameters.

* A second phase would then pick up where the other one left off (re-using existing parameter layouts to guarantee a match), and just layout out the shader parameters that *do* depend on global generic type parameters. Because the other parameters were already laid out, these new parameters would only ever fill in the gaps in layout (or come after all the other parameters, if explicit bindings aren't used).

This implementation strategy proved to be a bit of a mess, since we had to duplicate most of the code between the two passes anyway.

This commit eliminates `specializeProgramLayout()` entirely, and instead threads through global generic type arguments as part of the main layout pass. It is almost strictly a cleanup pass, now that the refactored logic for `Program` means the same layout algorithm can apply to specialized and unspecialized programs.

This change has one important semantic consequence (which is technically a break in backwards compatiblity for anybody using global generic parameters). Parameters that depend on global generic type parameters now get laid out in declaration order, just like all other shader parameters. This simplifies the rules, and in my experience actually makes application code *easier* to write in a systematic way (whereas the original design was motivated by the idea that giving more things stable locations would be beneficial).

A future improvement could be made so that we don't thread through the global generic substitution as part of layout. Instead, we could just attach a list of the global generic type arguments directly to the `TypeLayoutContext` and look those up on-demand when we encounter a global generic parameter type during layout. This would actually eliminate the need for global generics to appear as a `Substitution` entirely.

3 files changed, 29 insertions, 338 deletions

diff --git a/docs/layout.md b/docs/layout.md
index aecb22d70..12144c155 100644
--- a/docs/layout.md
+++ b/docs/layout.md
@@ -210,19 +210,19 @@ Generics
 Generic type parameters complicate these layout rules.
 For example, we cannot compute the exact resource requirements for a `vector<T,3>` without knowing what the type `T` is.
 
-When generic parameters are used in a program, the layout rules are ammended as follows:
+When computing layouts for fully specialized types or programs, no special considerations are needed: the rules as described in this document still apply.
+One important consequence to understand is that given a type like:
 
-* Identify those parameters with types that *depend on* a generic parameter
-  * A generic type parameter depends on itself
-  * An instantiation of a generic type `F<A>` depends on a generic parameter if `A` does
-  * An aggregate (`struct`) type depends on a generic type if the type of any of its fields does
-
-* Split the global parameter order into two lists: one for parameters that do not depend on generic parameters, and one for parameters that do
-
-* When enumerating shader parameters, always enumerate the non-generic-dependent parameter first
+```hlsl
+struct MyStuff<T>
+{
+	int a;
+	T b;
+	int c;
+}
+```
 
-This choie ensures that plugging in different types for the generic type parameters of an entry point will not change the layout that was computed for any "ordinary" parameters that did not depend on the generics.
+the offset computed for the `c` field depends on the concrete type that gets plugged in for `T`.
+We think this is the least surprising behavior for programmers who might be familiar with things like C++ template specialization.
 
-Note that we do *not* currently apply this same logic to `struct` types (e.g., lay out all the non-generic-dependent fields before the generic-dependent ones).
-At present, we think it is valuable for a specialization of a generic type to lay out like an "equivalent" non-generic type, but we may consider changing that decision.
-To avoid problems, users are encourage to declare types so that all non-generic-dependent fields come before generic-dependent ones.
+In cases where confusion about a field like `c` getting different offsets in different specializations is a concern, users are encouraged to declare types so that all non-generic-dependent fields come before generic-dependent ones.
diff --git a/source/slang/ir-link.cpp b/source/slang/ir-link.cpp
index fd9f1222e..60bafb9d1 100644
--- a/source/slang/ir-link.cpp
+++ b/source/slang/ir-link.cpp
@@ -1180,13 +1180,6 @@ void initializeSharedSpecContext(
     sharedContext->target = target;
 }
 
-// implementation provided in parameter-binding.cpp
-RefPtr<ProgramLayout> specializeProgramLayout(
-    TargetRequest * targetReq,
-    ProgramLayout*  programLayout,
-    SubstitutionSet typeSubst,
-    DiagnosticSink* sink);
-
 struct IRSpecializationState
 {
     ProgramLayout*      programLayout;
@@ -1194,7 +1187,6 @@ struct IRSpecializationState
     TargetRequest*      targetReq;
 
     IRModule* irModule = nullptr;
-    RefPtr<ProgramLayout> newProgramLayout;
 
     IRSharedSpecContext sharedContextStorage;
     IRSpecContext contextStorage;
@@ -1211,7 +1203,6 @@ struct IRSpecializationState
 
     ~IRSpecializationState()
     {
-        newProgramLayout = nullptr;
         contextStorage = IRSpecContext();
         sharedContextStorage = IRSharedSpecContext();
     }
@@ -1260,25 +1251,6 @@ LinkedIR linkIR(
     context->shared = sharedContext;
     context->builder = &sharedContext->builderStorage;
 
-    // Now specialize the program layout using the substitution
-    //
-    // TODO: The specialization of the layout is conceptually an AST-level operations,
-    // and shouldn't be done here in the IR at all.
-    //
-    RefPtr<ProgramLayout> newProgramLayout = specializeProgramLayout(
-        targetReq,
-        programLayout,
-        SubstitutionSet(program->getGlobalGenericSubstitution()),
-        compileRequest->getSink());
-
-    // TODO: we need to register the (IR-level) arguments of the global generic parameters as the
-    // substitutions for the generic parameters in the original IR.
-
-    // applyGlobalGenericParamSubsitution(...);
-
-
-    state->newProgramLayout = newProgramLayout;
-
     // Next, we want to optimize lookup for layout information
     // associated with global declarations, so that we can
     // look things up based on the IR values (using mangled names)
@@ -1290,7 +1262,7 @@ LinkedIR linkIR(
     // multiple mangled names (when the unique translation
     // unit name gets involved).
     //
-    auto globalStructLayout = getScopeStructLayout(newProgramLayout);
+    auto globalStructLayout = getScopeStructLayout(programLayout);
     for(auto entry : globalStructLayout->mapVarToLayout)
     {
         auto mangledName = getMangledName(entry.Key);
@@ -1311,7 +1283,7 @@ LinkedIR linkIR(
             cloneGlobalValue(context, (IRWitnessTable*)sym.Value->irGlobalValue);
     }
 
-    auto entryPointLayout = findEntryPointLayout(newProgramLayout, entryPoint);
+    auto entryPointLayout = findEntryPointLayout(programLayout, entryPoint);
 
     // Next, we make sure to clone the global value for
     // the entry point function itself, and rely on
diff --git a/source/slang/parameter-binding.cpp b/source/slang/parameter-binding.cpp
index c56ace9fe..f0256164c 100644
--- a/source/slang/parameter-binding.cpp
+++ b/source/slang/parameter-binding.cpp
@@ -1149,13 +1149,6 @@ RefPtr<TypeLayout> getTypeLayoutForGlobalShaderParameter(
         type);
 }
 
-RefPtr<TypeLayout> getTypeLayoutForGlobalShaderParameter(
-    ParameterBindingContext*    context,
-    VarDeclBase*                varDecl)
-{
-    return getTypeLayoutForGlobalShaderParameter(context, varDecl, varDecl->getType());
-}
-
 //
 
 struct EntryPointParameterState
@@ -1191,14 +1184,19 @@ static void collectGlobalGenericParameter(
 // Collect a single declaration into our set of parameters
 static void collectGlobalScopeParameter(
     ParameterBindingContext*    context,
-    RefPtr<VarDeclBase>         varDecl)
+    RefPtr<VarDeclBase>         varDecl,
+    SubstitutionSet             globalGenericSubst)
 {
+    // We apply any substitutions for global generic parameters here.
+    auto type = varDecl->getType()->Substitute(globalGenericSubst).as<Type>();
+
     // We use a single operation to both check whether the
     // variable represents a shader parameter, and to compute
     // the layout for that parameter's type.
     auto typeLayout = getTypeLayoutForGlobalShaderParameter(
         context,
-        varDecl.Ptr());
+        varDecl,
+        type);
 
     // If we did not find appropriate layout rules, then it
     // must mean that this global variable is *not* a shader
@@ -1789,7 +1787,8 @@ static void completeBindingsForParameter(
 
 static void collectGlobalScopeParameters(
     ParameterBindingContext*    context,
-    ModuleDecl*          program)
+    ModuleDecl*                 program,
+    SubstitutionSet             globalGenericSubst)
 {
     // First enumerate parameters at global scope
     // We collect two things here:
@@ -1807,7 +1806,7 @@ static void collectGlobalScopeParameters(
     for (auto decl : program->Members)
     {
         if (auto varDecl = as<VarDeclBase>(decl))
-            collectGlobalScopeParameter(context, varDecl);
+            collectGlobalScopeParameter(context, varDecl, globalGenericSubst);
     }
 
     // Next, we need to enumerate the parameters of
@@ -2663,19 +2662,21 @@ static void collectParameters(
     ParameterBindingContext contextData = *inContext;
     auto context = &contextData;
 
+    auto globalGenericSubst = program->getGlobalGenericSubstitution();
+
     for(RefPtr<Module> module : program->getModuleDependencies())
     {
         context->stage = Stage::Unknown;
 
         // First look at global-scope parameters
-        collectGlobalScopeParameters(context, module->getModuleDecl());
+        collectGlobalScopeParameters(context, module->getModuleDecl(), globalGenericSubst);
     }
 
     // Next consider parameters for entry points
     for(auto entryPoint : program->getEntryPoints())
     {
         context->stage = entryPoint->getStage();
-        collectEntryPointParameters(context, entryPoint, SubstitutionSet());
+        collectEntryPointParameters(context, entryPoint, globalGenericSubst);
     }
     context->entryPointLayout = nullptr;
 }
@@ -2949,286 +2950,4 @@ void generateParameterBindings(
     program->getTargetProgram(targetReq)->getOrCreateLayout(sink);
 }
 
-RefPtr<ProgramLayout> specializeProgramLayout(
-    TargetRequest*  targetReq,
-    ProgramLayout*  oldProgramLayout, 
-    SubstitutionSet typeSubst,
-    DiagnosticSink* sink)
-{
-    // The goal of the layout specialization step is to take an existing `ProgramLayout`,
-    // and add a layout to any parameter(s) that could not be laid out previously, because
-    // they had a dependence on generic type parameters that made layout impossible at
-    // the time.
-    //
-    // TODO: It would be far simpler to just "re-do" the entire layout process, just
-    // with knowledge of what the global type substitution is, but that would mean that
-    // global parameters that come after a generic-dependent parameter might change
-    // their location/binding/register depending on what types are plugged in.
-    // Our current design preserves the layout for any global parameter that was placed during
-    // the initial layout of a program (before the generic arguments were know).
-    // It isn't clear that this design choice pays off in practice, since there is  lot
-    // of complexity in this function.
-
-    RefPtr<ProgramLayout> newProgramLayout;
-    newProgramLayout = new ProgramLayout();
-    newProgramLayout->targetProgram = oldProgramLayout->targetProgram;
-    newProgramLayout->globalGenericParams = oldProgramLayout->globalGenericParams;
-
-    // The basic idea will be to iterate over the parameters in the old layout,
-    // and "pick up where we left off" in terms of allocating registers to things.
-    //
-    // That means we will look at the existing parameters (that were laid out already)
-    // and mark any registers/bytes/bindings/etc. that they occupy as "used" so
-    // that the subsequent layout of the generic-dependency parameters will not
-    // collide with them.
-    //
-    // We will use the same kind of context type as the original parameter binding
-    // step did, so we initialize its state here:
-
-    auto layoutContext = getInitialLayoutContextForTarget(targetReq, newProgramLayout);
-    SLANG_ASSERT(layoutContext.rules);
-
-    SharedParameterBindingContext sharedContext(
-        layoutContext.getRulesFamily(),
-        newProgramLayout,
-        targetReq,
-        sink);
-
-    ParameterBindingContext context;
-    context.shared = &sharedContext;
-    context.layoutContext = layoutContext;
-
-    // We will also need state for laying out any global-scope parameters
-    // that include ordinary/uniform data.
-    //
-    auto oldGlobalStructLayout = getGlobalStructLayout(oldProgramLayout);
-    SLANG_ASSERT(oldGlobalStructLayout);
-
-    ScopeLayoutBuilder newGlobalScopeLayoutBuilder;
-    newGlobalScopeLayoutBuilder.beginLayout(&context);
-    auto& newGlobalStructLayoutInfo = newGlobalScopeLayoutBuilder.m_structLayoutInfo;
-    auto newGlobalStructLayout = newGlobalScopeLayoutBuilder.m_structLayout;
-
-    // The initial state for uniform layout will be based on whatever
-    // global-scope ordinary/uniform parameters were laid out before.
-    // The alignment can be read directly from the old global layout.
-    //
-    newGlobalStructLayoutInfo.alignment = oldGlobalStructLayout->uniformAlignment;
-    newGlobalStructLayoutInfo.size = 0;
-
-    // The remaining information needs to be collected by looking at
-    // the individual parameters in the existing layout.
-    //
-    bool oldAnyUniforms = false;
-    for(auto oldVarLayout : oldGlobalStructLayout->fields)
-    {
-        // If a parameter made use of a global generic parameter, then we would
-        // have skipped applying layout to it in the original layout process,
-        // and so we should skip it for the process of recovering the existing
-        // layout information.
-        //
-        if (oldVarLayout->FindResourceInfo(LayoutResourceKind::GenericResource))
-            continue;
-
-        // Otherwise, we will "reserve" any resources that the parameter was
-        // determined to consume.
-        //
-        // The easy case is any registers/bindings used for textures/sampler/etc.
-        // We iterate over the kinds of resources consumed by teh parameter.
-        //
-        for( auto varResInfo : oldVarLayout->resourceInfos )
-        {
-            // For each kind of resource consumed the `varResInfo` will tell us
-            // the start of the consumed range, whle the type will be needed
-            // to tell us the amount of resources consumed.
-            //
-            if( auto typeResInfo = oldVarLayout->typeLayout->FindResourceInfo(varResInfo.kind) )
-            {
-                // We will mark the range of resources consumed by theis parameter
-                // as "used" so that it cannot be claimed by later parameters.
-                //
-                auto usedRangeSet = findUsedRangeSetForSpace(&context, varResInfo.space);
-                markSpaceUsed(&context, varResInfo.space);
-                usedRangeSet->usedResourceRanges[(int)varResInfo.kind].Add(
-                    nullptr, // we don't need to track parameter info here
-                    varResInfo.index,
-                    varResInfo.index + typeResInfo->count);
-            }
-        }
-
-        // The more subtle case is when the parameter consumes ordinary bytes
-        // of uniform (constant buffer) memory, because we do not use the
-        // same "used range" model to allocate space for ordinary data.
-        //
-        // Instead, we simply track the highest byte offset covered by any parameter.
-        //
-        if (auto varUniformInfo = oldVarLayout->FindResourceInfo(LayoutResourceKind::Uniform))
-        {
-            oldAnyUniforms = true;
-
-            if( auto typeUniformInfo = oldVarLayout->typeLayout->FindResourceInfo(LayoutResourceKind::Uniform) )
-            {
-                newGlobalStructLayoutInfo.size = maximum(
-                    newGlobalStructLayoutInfo.size,
-                    varUniformInfo->index + typeUniformInfo->count);
-            }
-        }
-    }
-
-    // Rather than attempt to re-use the entry-point layout information
-    // that was collected in the first pass, we will re-collect the
-    // information for entry points from scratch.
-    //
-    // This ensures that when an entry point makes use of a generic type
-    // parameter, the layout of its parameter list strictly follows
-    // the declaration order.
-    //
-    for( auto entryPoint : oldProgramLayout->getProgram()->getEntryPoints() )
-    {
-        collectEntryPointParameters(&context, entryPoint, typeSubst);
-        context.entryPointLayout = nullptr;
-    }
-
-    // Now that we've marked thing as being used, we can make a second
-    // sweep to compute the requirements of any generic-dependent parameters.
-    //
-    // Along the way we will build up the new layout for the global-scope
-    // structure type, including the offsets of all ordinary/uniform fields.
-    //
-
-    bool newAnyUniforms = oldAnyUniforms;
-    List<RefPtr<VarLayout>> newVarLayouts;
-    Dictionary<RefPtr<VarLayout>, RefPtr<VarLayout>> mapOldLayoutToNew;
-    for(auto oldVarLayout : oldGlobalStructLayout->fields)
-    {
-        // In this pass, the variables that *don't* depend on generic parameters
-        // are the easy ones to handle. We can just copy them over to the new layout.
-        //
-        if(!oldVarLayout->FindResourceInfo(LayoutResourceKind::GenericResource))
-        {
-            newGlobalStructLayout->fields.Add(oldVarLayout);
-            continue;
-        }
-
-        // In the case where things are generic-dependent, we need to re-do
-        // the type layout process on the type that results from doing
-        // substitution with the global generic arguments.
-        //
-        RefPtr<Type> oldType = oldVarLayout->getTypeLayout()->getType();
-        SLANG_ASSERT(oldType);
-        RefPtr<Type> newType = oldType->Substitute(typeSubst).as<Type>();
-
-        RefPtr<TypeLayout> newTypeLayout = getTypeLayoutForGlobalShaderParameter(
-            &context,
-            oldVarLayout->varDecl,
-            newType);
-
-        RefPtr<VarLayout> newVarLayout = new VarLayout();
-        newVarLayout->varDecl = oldVarLayout->varDecl;
-        newVarLayout->stage = oldVarLayout->stage;
-        newVarLayout->typeLayout = newTypeLayout;
-
-        newGlobalScopeLayoutBuilder.addParameter(newVarLayout);
-        newVarLayouts.Add(newVarLayout);
-        mapOldLayoutToNew.Add(oldVarLayout, newVarLayout);
-
-        if(auto uniformInfo = newTypeLayout->FindResourceInfo(LayoutResourceKind::Uniform))
-        {
-            if(uniformInfo->count != 0)
-            {
-                newAnyUniforms = true;
-                diagnoseGlobalUniform(&sharedContext, newVarLayout->varDecl);
-            }
-        }
-
-    }
-    auto newGlobalScopeVarLayout = newGlobalScopeLayoutBuilder.endLayout();
-
-    // We had better have made a copy of every field in the original layout.
-    //
-    SLANG_ASSERT(oldGlobalStructLayout->fields.Count() == newGlobalStructLayout->fields.Count());
-
-    // If there were no global-scope uniforms before, but there
-    // are now that we've done global substitution, then we
-    // need to allocate a global constant buffer to hold them.
-    //
-    auto newGlobalConstantBufferBinding = maybeAllocateConstantBufferBinding(&context, newAnyUniforms && !oldAnyUniforms);
-
-    // Now we need to "complete" finding for each of the new parameters,
-    // which is the step that actually allocates resource to them.
-    //
-    // Note: we don't support generic-dependent parameters with explicit bindings,
-    // so we should probably emit an error message about that in the original
-    // layout step.
-    //
-    for(auto newVarLayout : newVarLayouts)
-    {
-        completeBindingsForParameter(&context, newVarLayout);
-    }
-
-    // One remaining missing step is that the `StructLayout` type maintains
-    // a map from variable declarations to their layouts, and in some cases
-    // multiple declarations will map to the same layout (because, e.g., the
-    // same `cbuffer` was declared in both a vertex and fragment shader file).
-    //
-    // We need to clone that remapping information over from the old program
-    // layout. This is why we created the `mapOldLayoutToNew` mapping in
-    // the preceding loop.
-    //
-    // TODO: This step would be easier if the `StructLayout::mapVarToLayout`
-    // dictionary were instead a mapping from variable declaration to the
-    // *index* of the corresponding layout in the `fields` array.
-    //
-    for(auto entry : oldGlobalStructLayout->mapVarToLayout)
-    {
-        RefPtr<VarLayout> varLayout = entry.Value;
-        mapOldLayoutToNew.TryGetValue(varLayout, varLayout);
-        newGlobalStructLayout->mapVarToLayout[entry.Key] = varLayout;
-    }
-
-    // Just as for the initial computation of layout, we will complete
-    // binding for entry-point parameters *after* we have laid out
-    // all the global-scope parameters.
-    //
-    // Note that this includes layout of generic-dependent global scope
-    // parameters, so it is possible for entry point uniform parameters
-    // to end up with a different register/binding after generic specialization.
-    // (There really isn't a great way around that)
-    //
-    for( auto entryPoint : sharedContext.programLayout->entryPoints )
-    {
-        auto entryPointParamsLayout = entryPoint->parametersLayout;
-        completeBindingsForParameter(&context, entryPointParamsLayout);
-    }
-
-    // As a last step we need to set up the binding/offset information
-    // for the global scope itself.
-    //
-    // We will start by copying whatever information was in the old layout.
-    //
-    {
-        auto oldGlobalScopeVarLayout = oldProgramLayout->parametersLayout;
-        for( auto oldResInfo : oldGlobalScopeVarLayout->resourceInfos )
-        {
-            auto newResInfo = newGlobalScopeVarLayout->findOrAddResourceInfo(oldResInfo.kind);
-            newResInfo->space = oldResInfo.space;
-            newResInfo->kind = oldResInfo.kind;
-        }
-    }
-
-    // If we had to create a constant buffer to house the global-scope
-    // ordinary/uniform data, then we need to make sure to set that
-    // information on the global scope.
-    //
-    if(newGlobalConstantBufferBinding.kind != LayoutResourceKind::None )
-    {
-        auto resInfo = newGlobalScopeVarLayout->findOrAddResourceInfo(newGlobalConstantBufferBinding.kind);
-        resInfo->space = newGlobalConstantBufferBinding.space;
-        resInfo->index = newGlobalConstantBufferBinding.index;
-    }
-
-    newProgramLayout->parametersLayout = newGlobalScopeVarLayout;
-    return newProgramLayout;
-}
-
 } // namespace Slang