1 files changed, 117 insertions, 118 deletions

diff --git a/source/slang/slang-ir-link.cpp b/source/slang/slang-ir-link.cpp
index 56b06a499..f46a0db09 100644
--- a/source/slang/slang-ir-link.cpp
+++ b/source/slang/slang-ir-link.cpp
@@ -54,10 +54,6 @@ struct IRSharedSpecContext
 
 struct IRSpecContextBase
 {
-    // A map from the mangled name of a global variable
-    // to the layout to use for it.
-    Dictionary<String, VarLayout*> globalVarLayouts;
-
     IRSharedSpecContext* shared;
 
     IRSharedSpecContext* getShared() { return shared; }
@@ -231,6 +227,7 @@ IRInst* IRSpecContext::maybeCloneValue(IRInst* originalValue)
     case kIROp_StructKey:
     case kIROp_GlobalGenericParam:
     case kIROp_WitnessTable:
+    case kIROp_TaggedUnionType:
         return cloneGlobalValue(this, originalValue);
 
     case kIROp_BoolLit:
@@ -414,6 +411,17 @@ static void cloneExtraDecorations(
     IRBuilder* builder = &builderStorage;
     builder->setInsertInto(clonedInst);
 
+    // If the `clonedInst` already has any non-decoration
+    // children, then we want to insert before those,
+    // to maintain the invariant that decorations always
+    // precede non-decoration instructions in the list of
+    // decorations and children.
+    //
+    if(auto firstChild = clonedInst->getFirstChild())
+    {
+        builder->setInsertBefore(firstChild);
+    }
+
     for(auto sym = originalValues.sym; sym; sym = sym->nextWithSameName)
     {
         for(auto decoration : sym->irGlobalValue->getDecorations())
@@ -424,6 +432,7 @@ static void cloneExtraDecorations(
                 break;
 
             case kIROp_BindExistentialSlotsDecoration:
+            case kIROp_LayoutDecoration:
                 if(!clonedInst->findDecorationImpl(decoration->op))
                 {
                     cloneInst(context, builder, decoration);
@@ -431,6 +440,14 @@ static void cloneExtraDecorations(
                 break;
             }
         }
+
+        // We will also copy over source location information from the alternative
+        // values, in case any of them has it available.
+        //
+        if(sym->irGlobalValue->sourceLoc.isValid() && !clonedInst->sourceLoc.isValid())
+        {
+            clonedInst->sourceLoc = sym->irGlobalValue->sourceLoc;
+        }
     }
 }
 
@@ -472,16 +489,6 @@ IRGlobalParam* cloneGlobalParamImpl(
         cloneType(context, originalVal->getFullType()));
     cloneSimpleGlobalValueImpl(context, originalVal, originalValues, clonedVal);
 
-    if(auto linkage = originalVal->findDecoration<IRLinkageDecoration>())
-    {
-        auto mangledName = String(linkage->getMangledName());
-        VarLayout* layout = nullptr;
-        if (context->globalVarLayouts.TryGetValue(mangledName, layout))
-        {
-            builder->addLayoutDecoration(clonedVal, layout);
-        }
-    }
-
     return clonedVal;
 }
 
@@ -723,9 +730,63 @@ void cloneFunctionCommon(
 IRInst* specializeGeneric(
     IRSpecialize*   specializeInst);
 
+    /// Copy layout information for an entry-point function to its parameters.
+    ///
+    /// When layout information is initially attached to an IR entry point,
+    /// it may be attached to a declaration that would have no `IRParam`s
+    /// to represent the entry-point parameters.
+    ///
+    /// After linking, we expect an entry point to have a full definition,
+    /// so it becomes possible to copy per-parameter layout information
+    /// from the entry-point layout down to the individual parameters,
+    /// which simplifies subsequent IR steps taht want to look for
+    /// per-parameter layout information.
+    ///
+    /// TODO: This step should probably be hoisted out to be an independent
+    /// IR pass that runs after linking, rather than being folded into
+    /// the linking step.
+    ///
+static void maybeCopyLayoutInformationToParameters(
+    IRFunc* func,
+    IRBuilder* builder)
+{
+    auto layoutDecor = func->findDecoration<IRLayoutDecoration>();
+    if(!layoutDecor)
+        return;
+
+    auto entryPointLayout = as<EntryPointLayout>(layoutDecor->getASTLayout());
+    if(!entryPointLayout)
+        return;
+
+    if( auto firstBlock = func->getFirstBlock() )
+    {
+        auto paramsStructLayout = getScopeStructLayout(entryPointLayout);
+        Index paramLayoutCount = paramsStructLayout->fields.getCount();
+        Index paramCounter = 0;
+        for( auto pp = firstBlock->getFirstParam(); pp; pp = pp->getNextParam() )
+        {
+            Index paramIndex = paramCounter++;
+            if( paramIndex < paramLayoutCount )
+            {
+                auto paramLayout = paramsStructLayout->fields[paramIndex];
+
+                auto offsetParamLayout = applyOffsetToVarLayout(paramLayout, entryPointLayout->parametersLayout);
+
+                builder->addLayoutDecoration(
+                    pp,
+                    offsetParamLayout);
+            }
+            else
+            {
+                SLANG_UNEXPECTED("too many parameters");
+            }
+        }
+    }
+}
+
 IRFunc* specializeIRForEntryPoint(
     IRSpecContext*      context,
-    EntryPointLayout*   entryPointLayout)
+    String const&       mangledName)
 {
     // We start by looking up the IR symbol that
     // matches the mangled name given to the
@@ -736,7 +797,6 @@ IRFunc* specializeIRForEntryPoint(
     // so that the mangled name of the decl-ref is
     // not the same as the mangled name of the decl.
     //
-    auto mangledName = getMangledName(entryPointLayout->getFuncDeclRef());
     RefPtr<IRSpecSymbol> sym;
     if (!context->getSymbols().TryGetValue(mangledName, sym))
     {
@@ -744,8 +804,14 @@ IRFunc* specializeIRForEntryPoint(
         return nullptr;
     }
 
-    // TODO: deal with the case where we might
-    // have multiple (profile-overloaded) versions...
+    // Note: it is possible that `sym` shows multiple
+    // definitions, coming from different IR modules that
+    // were input to the linking process. We don't have
+    // to do anything special about that here, because
+    // we can use *any* of the IR values as the starting
+    // point for cloning, and `cloneGlobalValue` will
+    // follow the linkage decoration and discover the
+    // other values on its own.
     //
     auto originalVal = sym->irGlobalValue;
 
@@ -788,26 +854,21 @@ IRFunc* specializeIRForEntryPoint(
         // we don't want to share the definition between
         // an entry point and an ordinary function anyway.
         //
-        clonedVal = specializeGeneric(clonedSpec);
-    }
+        auto specializedClone = specializeGeneric(clonedSpec);
 
-    // TODO: If there is an existential-related decoration
-    // on the entry point, we need to transfer it over
-    // to the specialized function.
-    if( auto bindExistentialSlots = originalVal->findDecorationImpl(kIROp_BindExistentialSlotsDecoration) )
-    {
-        if( !clonedVal->findDecorationImpl(kIROp_BindExistentialSlotsDecoration) )
-        {
-            IRBuilder builderStorage = *context->builder;
-            IRBuilder* builder = &builderStorage;
-            builder->setInsertInto(clonedVal);
+        // One special case we need to be aware of is that there
+        // might be decorations attached to the `specialize`
+        // instruction, which we then want to copy over to the
+        // result of specialization.
+        //
+        cloneExtraDecorations(context, specializedClone, IROriginalValuesForClone(sym));
 
-            auto clonedBind = cloneInst(context, builder, bindExistentialSlots);
-            clonedBind->moveToStart();
-        }
+        // Now we will move to considering the specialized instruction
+        // instead of the unspecialized one for all further steps.
+        //
+        clonedVal = specializedClone;
     }
 
-
     auto clonedFunc = as<IRFunc>(clonedVal);
     if(!clonedFunc)
     {
@@ -820,42 +881,11 @@ IRFunc* specializeIRForEntryPoint(
         context->builder->addKeepAliveDecoration(clonedFunc);
     }
 
-    // We need to attach the layout information for
-    // the entry point to this declaration, so that
-    // we can use it to inform downstream code emit.
-    //
-    context->builder->addLayoutDecoration(
-        clonedFunc,
-        entryPointLayout);
-
     // We will also go on and attach layout information
     // to the function parameters, so that we have it
     // available directly on the parameters, rather
     // than having to look it up on the original entry-point layout.
-    if( auto firstBlock = clonedFunc->getFirstBlock() )
-    {
-        auto paramsStructLayout = getScopeStructLayout(entryPointLayout);
-        Index paramLayoutCount = paramsStructLayout->fields.getCount();
-        Index paramCounter = 0;
-        for( auto pp = firstBlock->getFirstParam(); pp; pp = pp->getNextParam() )
-        {
-            Index paramIndex = paramCounter++;
-            if( paramIndex < paramLayoutCount )
-            {
-                auto paramLayout = paramsStructLayout->fields[paramIndex];
-
-                auto offsetParamLayout = applyOffsetToVarLayout(paramLayout, entryPointLayout->parametersLayout);
-
-                context->builder->addLayoutDecoration(
-                    pp,
-                    offsetParamLayout);
-            }
-            else
-            {
-                SLANG_UNEXPECTED("too many parameters");
-            }
-        }
-    }
+    maybeCopyLayoutInformationToParameters(clonedFunc, context->builder);
 
     return clonedFunc;
 }
@@ -1073,6 +1103,7 @@ IRInst* cloneInst(
     builder->addInst(clonedInst);
     context->builder = oldBuilder;
     cloneDecorations(context, clonedInst, originalInst);
+    cloneExtraDecorations(context, clonedInst, originalValues);
 
     return clonedInst;
 }
@@ -1251,7 +1282,6 @@ void initializeSharedSpecContext(
 
 struct IRSpecializationState
 {
-    ProgramLayout*      programLayout;
     CodeGenTarget       target;
     TargetRequest*      targetReq;
 
@@ -1279,11 +1309,12 @@ struct IRSpecializationState
 
 LinkedIR linkIR(
     BackEndCompileRequest*  compileRequest,
-    EntryPoint*             entryPoint,
-    ProgramLayout*          programLayout,
+    Int                     entryPointIndex,
     CodeGenTarget           target,
-    TargetRequest*          targetReq)
+    TargetProgram*          targetProgram)
 {
+    auto targetReq = targetProgram->getTargetReq();
+
     // TODO: We need to make sure that the program we are being asked
     // to compile has been "resolved" so that it has no outstanding
     // unsatisfied requirements.
@@ -1291,7 +1322,6 @@ LinkedIR linkIR(
     IRSpecializationState stateStorage;
     auto state = &stateStorage;
 
-    state->programLayout = programLayout;
     state->target = target;
     state->targetReq = targetReq;
 
@@ -1316,34 +1346,18 @@ LinkedIR linkIR(
         insertGlobalValueSymbols(sharedContext, irModule);
     });
 
+    // We will also insert the IR global symbols from the IR module
+    // attached to the `TargetProgram`, since this module is
+    // responsible for associating layout information to those
+    // global symbols via decorations.
+    //
+    insertGlobalValueSymbols(sharedContext, targetProgram->getExistingIRModuleForLayout());
+
 
     auto context = state->getContext();
     context->shared = sharedContext;
     context->builder = &sharedContext->builderStorage;
 
-    // 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)
-    //
-    // Note: We are scanning over all the key-value pairs for
-    // entries in the global scope, to account for the fact
-    // that the "same" shader parameter could be declared in
-    // multiple translation units, and thus end up with
-    // multiple mangled names (when the unique translation
-    // unit name gets involved).
-    //
-    auto globalStructLayout = getScopeStructLayout(programLayout);
-    for(auto entry : globalStructLayout->mapVarToLayout)
-    {
-        auto mangledName = getMangledName(entry.Key);
-        auto globalVarLayout = entry.Value;
-        context->globalVarLayouts.AddIfNotExists(mangledName, globalVarLayout);
-    }
-
-    auto entryPointLayout = findEntryPointLayout(programLayout, entryPoint);
-
-    auto offsetEntryPointLayout = entryPointLayout;
-
     context->builder->setInsertInto(context->getModule()->getModuleInst());
 
     // for now, clone all unreferenced witness tables
@@ -1362,7 +1376,15 @@ LinkedIR linkIR(
     // the entry point function itself, and rely on
     // this step to recursively copy over anything else
     // it might reference.
-    auto irEntryPoint = specializeIRForEntryPoint(context, offsetEntryPointLayout);
+    //
+    // Note: We query the mangled name of the entry point from
+    // the `program` instead of the `entryPoint` directly,
+    // because the `program` will include any specialization
+    // arguments which might end up affecting the mangled
+    // entry point name.
+    //
+    auto entryPointMangledName = program->getEntryPointMangledName(entryPointIndex);
+    auto irEntryPoint = specializeIRForEntryPoint(context, entryPointMangledName);
 
     // Bindings for global generic parameters are currently represented
     // as stand-alone global-scope instructions in the IR module for
@@ -1388,29 +1410,6 @@ LinkedIR linkIR(
         }
     });
 
-    // HACK: we need to ensure that any tagged union types
-    // in the IR module have layout information copied over to them.
-    //
-    // Note that we do this *after* cloning the `bindGlobalGenericParam`
-    // instructions, since we expected the tagged union type(s) to
-    // be referenced by them.
-    //
-    for( auto taggedUnionTypeLayout : programLayout->taggedUnionTypeLayouts )
-    {
-        auto taggedUnionType = taggedUnionTypeLayout->getType();
-        auto mangledName = getMangledTypeName(taggedUnionType);
-
-        RefPtr<IRSpecSymbol> sym;
-        if(!context->getSymbols().TryGetValue(mangledName, sym))
-            continue;
-
-        IRInst* clonedType = findClonedValue(context, sym->irGlobalValue);
-        if(!clonedType)
-            continue;
-
-        context->builder->addLayoutDecoration(clonedType, taggedUnionTypeLayout);
-    }
-
     // TODO: *technically* we should consider the case where
     // we have global variables with initializers, since
     // these should get run whether or not the entry point