Remove unnecessary parameters from Metal entry point signature (#6131)

* fix metal entry point global params

* address review comments, cleanup and test

* remove dead code

* undo accidental change

* address review comments and cleanup

* minor fix and cleanup

---------

Co-authored-by: Yong He <yonghe@outlook.com>

12 files changed, 277 insertions, 160 deletions

diff --git a/source/slang/slang-emit.cpp b/source/slang/slang-emit.cpp
index d5c13121a..ea209b598 100644
--- a/source/slang/slang-emit.cpp
+++ b/source/slang/slang-emit.cpp
@@ -93,10 +93,8 @@
 #include "slang-ir-ssa-simplification.h"
 #include "slang-ir-ssa.h"
 #include "slang-ir-string-hash.h"
-#include "slang-ir-strip-cached-dict.h"
 #include "slang-ir-strip-default-construct.h"
-#include "slang-ir-strip-witness-tables.h"
-#include "slang-ir-strip.h"
+#include "slang-ir-strip-legalization-insts.h"
 #include "slang-ir-synthesize-active-mask.h"
 #include "slang-ir-translate-glsl-global-var.h"
 #include "slang-ir-uniformity.h"
@@ -1501,12 +1499,10 @@ Result linkAndOptimizeIR(
         break;
     }
 
-    stripCachedDictionaries(irModule);
-
     // TODO: our current dynamic dispatch pass will remove all uses of witness tables.
     // If we are going to support function-pointer based, "real" modular dynamic dispatch,
     // we will need to disable this pass.
-    stripWitnessTables(irModule);
+    stripLegalizationOnlyInstructions(irModule);
 
     switch (target)
     {
diff --git a/source/slang/slang-ir-entry-point-uniforms.cpp b/source/slang/slang-ir-entry-point-uniforms.cpp
index b3073a97d..12dcd187a 100644
--- a/source/slang/slang-ir-entry-point-uniforms.cpp
+++ b/source/slang/slang-ir-entry-point-uniforms.cpp
@@ -494,19 +494,27 @@ struct MoveEntryPointUniformParametersToGlobalScope : PerEntryPointPass
             // for CPU/CUDA) that might want to treat entry-point parameters
             // different from other cases.
             //
-            // TODO: Once we have support for multiple entry points to be emitted
-            // at once, we need a way to associate these per-entry-point parameters
-            // more closely with the original entry point. The two easiest options
-            // are:
+            // We need a way to associate these per-entry-point parameters
+            // more closely with the original entry point. The two current
+            // methods are:
             //
             // 1. Don't move the new aggregate parameter to the global scope
             // on those targets, and instead keep it as a parameter of the
-            // entry point.
+            // entry point. This is used for CPU/CUDA targets.
             //
-            // 2. Use a decoration on the entry point itself to point at the
-            // global parameter for its per-entry-point parameter data.
+            // 2. Use a decoration on the global param itself to point at the
+            // entry point for its per-entry-point parameter data, without moving
+            // the parameter to the global scope. This is used for Metal targets, as
+            // Metal does not have global parameters at the global scope.
             //
-            builder->addDecoration(globalParam, kIROp_EntryPointParamDecoration);
+            // Method (1) is not used because Metal contains shading language concepts
+            // such as binding offets, similar to other shading language targets.
+            // We want to reuse code from other shading language targets for Metal, hence
+            // we move parameters to the global scope, and then move the parameters back to
+            // the entry points that they originate from. The originating entry points are
+            // tracked through this decoration.
+            //
+            builder->addEntryPointParamDecoration(globalParam, entryPointFunc);
 
             param->replaceUsesWith(globalParam);
             param->removeAndDeallocate();
diff --git a/source/slang/slang-ir-explicit-global-context.cpp b/source/slang/slang-ir-explicit-global-context.cpp
index 0516d574c..f235ba3e4 100644
--- a/source/slang/slang-ir-explicit-global-context.cpp
+++ b/source/slang/slang-ir-explicit-global-context.cpp
@@ -140,7 +140,19 @@ struct IntroduceExplicitGlobalContextPass
     IRStructType* m_contextStructType = nullptr;
     IRPtrType* m_contextStructPtrType = nullptr;
 
-    List<IRGlobalParam*> m_globalParams;
+    struct GlobalParamInfo
+    {
+        // Original global param inst.
+        IRGlobalParam* globalParam = nullptr;
+
+        // New entry point param that is created by this pass.
+        IRParam* entryPointParam = nullptr;
+
+        // Orignating entry point obtained from entry point param decoration, if it exists.
+        IRFunc* originatingEntryPoint = nullptr;
+    };
+
+    List<GlobalParamInfo> m_globalParams;
     List<IRGlobalVar*> m_globalVars;
     List<IRFunc*> m_entryPoints;
 
@@ -237,7 +249,22 @@ struct IntroduceExplicitGlobalContextPass
                     if (m_target == CodeGenTarget::CUDASource)
                         continue;
 
-                    m_globalParams.add(globalParam);
+                    GlobalParamInfo globalParamInfo;
+                    globalParamInfo.globalParam = globalParam;
+
+                    // Entry point param decorations are not required anymore after this pass and
+                    // must be removed for entry point param emit. Remoeving it here prevents the
+                    // decoration from being cloned when creating struct keys and entry point
+                    // parameters.
+                    if (const auto entryPointParamDecoration =
+                            globalParam->findDecoration<IREntryPointParamDecoration>())
+                    {
+                        globalParamInfo.originatingEntryPoint =
+                            entryPointParamDecoration->getEntryPoint();
+                        entryPointParamDecoration->removeAndDeallocate();
+                    }
+
+                    m_globalParams.add(globalParamInfo);
                 }
                 break;
 
@@ -305,11 +332,10 @@ struct IntroduceExplicitGlobalContextPass
             // For the parameter representing all the global uniform shader
             // parameters, we create a field that exactly matches its type.
             //
-
             createContextStructField(
-                globalParam,
+                globalParam.globalParam,
                 GlobalObjectKind::GlobalParam,
-                globalParam->getFullType());
+                globalParam.globalParam->getFullType());
         }
         for (auto globalVar : m_globalVars)
         {
@@ -347,7 +373,7 @@ struct IntroduceExplicitGlobalContextPass
         //
         for (auto globalParam : m_globalParams)
         {
-            replaceUsesOfGlobalParam(globalParam);
+            replaceUsesOfGlobalParam(globalParam.globalParam);
         }
         for (auto globalVar : m_globalVars)
         {
@@ -444,23 +470,32 @@ struct IntroduceExplicitGlobalContextPass
         // then we need to introduce an explicit parameter onto
         // each entry-point function to represent it.
         //
-        struct GlobalParamInfo
-        {
-            IRGlobalParam* globalParam;
-            IRParam* entryPointParam;
-        };
-        List<GlobalParamInfo> entryPointParams;
+
+        List<GlobalParamInfo> entryPointParamsToAdd;
         for (auto globalParam : m_globalParams)
         {
-            auto entryPointParam = builder.createParam(globalParam->getFullType());
+            // Do not add global param to current entry point if global param
+            // explicitly originates from a different entry point.
+            if (globalParam.originatingEntryPoint &&
+                globalParam.originatingEntryPoint != entryPointFunc)
+            {
+                continue;
+            }
+
+            globalParam.entryPointParam =
+                builder.createParam(globalParam.globalParam->getFullType());
             IRCloneEnv cloneEnv;
-            cloneInstDecorationsAndChildren(&cloneEnv, m_module, globalParam, entryPointParam);
-            entryPointParams.add({globalParam, entryPointParam});
+            cloneInstDecorationsAndChildren(
+                &cloneEnv,
+                m_module,
+                globalParam.globalParam,
+                globalParam.entryPointParam);
+            entryPointParamsToAdd.add(globalParam);
 
             // The new parameter will be the last one in the
             // parameter list of the entry point.
             //
-            entryPointParam->insertBefore(firstOrdinary);
+            globalParam.entryPointParam->insertBefore(firstOrdinary);
         }
 
         if (m_target == CodeGenTarget::CPPSource && m_globalParams.getCount() == 0)
@@ -485,7 +520,7 @@ struct IntroduceExplicitGlobalContextPass
         // to inialize the corresponding field of the `KernelContext`
         // before moving on with execution of the kernel body.
         //
-        for (auto entryPointParam : entryPointParams)
+        for (auto entryPointParam : entryPointParamsToAdd)
         {
             auto fieldInfo = m_mapInstToContextFieldInfo[entryPointParam.globalParam];
             auto fieldType = entryPointParam.globalParam->getFullType();
diff --git a/source/slang/slang-ir-insts.h b/source/slang/slang-ir-insts.h
index eebb8f119..a311e024c 100644
--- a/source/slang/slang-ir-insts.h
+++ b/source/slang/slang-ir-insts.h
@@ -812,6 +812,14 @@ struct IRKnownBuiltinDecoration : IRDecoration
     UnownedStringSlice getName() { return getNameOperand()->getStringSlice(); }
 };
 
+struct IREntryPointParamDecoration : IRDecoration
+{
+    IR_LEAF_ISA(EntryPointParamDecoration)
+
+    /// Get the entry point that this parameter orignates from.
+    IRFunc* getEntryPoint() { return cast<IRFunc>(getOperand(0)); }
+};
+
 struct IRFormatDecoration : IRDecoration
 {
     enum
@@ -5226,6 +5234,11 @@ public:
     {
         addDecoration(inst, kIROp_CheckpointIntermediateDecoration, func);
     }
+
+    void addEntryPointParamDecoration(IRInst* inst, IRFunc* entryPointFunc)
+    {
+        addDecoration(inst, kIROp_EntryPointParamDecoration, entryPointFunc);
+    }
 };
 
 // Helper to establish the source location that will be used
diff --git a/source/slang/slang-ir-legalize-types.cpp b/source/slang/slang-ir-legalize-types.cpp
index 8eb70596f..a9d9d65f3 100644
--- a/source/slang/slang-ir-legalize-types.cpp
+++ b/source/slang/slang-ir-legalize-types.cpp
@@ -110,10 +110,11 @@ static void registerLegalizedValue(
     context->mapValToLegalVal[irValue] = legalVal;
 }
 
-struct IRGlobalNameInfo
+/// Structure to pass information from the original/old global param to
+/// composite members during tuple flavored global param legalization.
+struct IRGlobalParamInfo
 {
-    IRInst* globalVar;
-    UInt counter;
+    IRFunc* originatingEntryPoint = nullptr;
 };
 
 static LegalVal declareVars(
@@ -124,7 +125,7 @@ static LegalVal declareVars(
     LegalVarChain const& varChain,
     UnownedStringSlice nameHint,
     IRInst* leafVar,
-    IRGlobalNameInfo* globalNameInfo,
+    IRGlobalParamInfo* globalParamInfo,
     bool isSpecial);
 
 /// Unwrap a value with flavor `wrappedBuffer`
@@ -2727,10 +2728,8 @@ static LegalVal declareSimpleVar(
     LegalVarChain const& varChain,
     UnownedStringSlice nameHint,
     IRInst* leafVar,
-    IRGlobalNameInfo* globalNameInfo)
+    IRGlobalParamInfo* globalParamInfo)
 {
-    SLANG_UNUSED(globalNameInfo);
-
     IRVarLayout* varLayout = createVarLayout(context->builder, varChain, typeLayout);
 
     IRBuilder* builder = context->builder;
@@ -2757,6 +2756,19 @@ static LegalVal declareSimpleVar(
             globalParam->removeFromParent();
             globalParam->insertBefore(context->insertBeforeGlobal);
 
+            // Add originating entry point decoration if original global param
+            // comes from an entry point parameter. This is required in cases where the global
+            // param has to be linked back to the originating entry point, such as when
+            // emitting Metal where there global params have to be moved back to the
+            // entry point parameter.
+            SLANG_ASSERT(globalParamInfo);
+            if (globalParamInfo->originatingEntryPoint)
+            {
+                builder->addEntryPointParamDecoration(
+                    globalParam,
+                    globalParamInfo->originatingEntryPoint);
+            }
+
             irVar = globalParam;
             legalVarVal = LegalVal::simple(globalParam);
         }
@@ -3416,7 +3428,7 @@ static LegalVal declareVars(
     LegalVarChain const& inVarChain,
     UnownedStringSlice nameHint,
     IRInst* leafVar,
-    IRGlobalNameInfo* globalNameInfo,
+    IRGlobalParamInfo* globalParamInfo,
     bool isSpecial)
 {
     LegalVarChain varChain = inVarChain;
@@ -3451,7 +3463,7 @@ static LegalVal declareVars(
             varChain,
             nameHint,
             leafVar,
-            globalNameInfo);
+            globalParamInfo);
         break;
 
     case LegalType::Flavor::implicitDeref:
@@ -3466,7 +3478,7 @@ static LegalVal declareVars(
                 varChain,
                 nameHint,
                 leafVar,
-                globalNameInfo,
+                globalParamInfo,
                 isSpecial);
             return LegalVal::implicitDeref(val);
         }
@@ -3483,7 +3495,7 @@ static LegalVal declareVars(
                 varChain,
                 nameHint,
                 leafVar,
-                globalNameInfo,
+                globalParamInfo,
                 false);
             auto specialVal = declareVars(
                 context,
@@ -3493,7 +3505,7 @@ static LegalVal declareVars(
                 varChain,
                 nameHint,
                 leafVar,
-                globalNameInfo,
+                globalParamInfo,
                 true);
             return LegalVal::pair(ordinaryVal, specialVal, pairType->pairInfo);
         }
@@ -3545,7 +3557,7 @@ static LegalVal declareVars(
                     newVarChain,
                     fieldNameHint,
                     ee.key,
-                    globalNameInfo,
+                    globalParamInfo,
                     true);
 
                 TuplePseudoVal::Element element;
@@ -3600,7 +3612,7 @@ static LegalVal declareVars(
                 varChain,
                 nameHint,
                 leafVar,
-                globalNameInfo);
+                globalParamInfo);
 
             return LegalVal::wrappedBuffer(innerVal, wrappedBuffer->elementInfo);
         }
@@ -3634,10 +3646,6 @@ static LegalVal legalizeGlobalVar(IRTypeLegalizationContext* context, IRGlobalVa
         {
             context->insertBeforeGlobal = irGlobalVar;
 
-            IRGlobalNameInfo globalNameInfo;
-            globalNameInfo.globalVar = irGlobalVar;
-            globalNameInfo.counter = 0;
-
             UnownedStringSlice nameHint = findNameHint(irGlobalVar);
             context->builder->setInsertBefore(irGlobalVar);
             LegalVal newVal = declareVars(
@@ -3648,7 +3656,7 @@ static LegalVal legalizeGlobalVar(IRTypeLegalizationContext* context, IRGlobalVa
                 LegalVarChain(),
                 nameHint,
                 irGlobalVar,
-                &globalNameInfo,
+                nullptr,
                 context->isSpecialType(originalValueType));
 
             // Register the new value as the replacement for the old
@@ -3689,9 +3697,12 @@ static LegalVal legalizeGlobalParam(
 
             LegalVarChainLink varChain(LegalVarChain(), varLayout);
 
-            IRGlobalNameInfo globalNameInfo;
-            globalNameInfo.globalVar = irGlobalParam;
-            globalNameInfo.counter = 0;
+            IRGlobalParamInfo globalParamInfo;
+            if (auto entryPointParamDecoration =
+                    irGlobalParam->findDecoration<IREntryPointParamDecoration>())
+            {
+                globalParamInfo.originatingEntryPoint = entryPointParamDecoration->getEntryPoint();
+            }
 
             // TODO: need to handle initializer here!
 
@@ -3705,7 +3716,7 @@ static LegalVal legalizeGlobalParam(
                 varChain,
                 nameHint,
                 irGlobalParam,
-                &globalNameInfo,
+                &globalParamInfo,
                 context->isSpecialType(irGlobalParam->getDataType()));
 
             // Register the new value as the replacement for the old
diff --git a/source/slang/slang-ir-strip-cached-dict.cpp b/source/slang/slang-ir-strip-cached-dict.cpp
deleted file mode 100644
index 202b3a4a2..000000000
--- a/source/slang/slang-ir-strip-cached-dict.cpp
+++ /dev/null
@@ -1,29 +0,0 @@
-// slang-ir-strip-cached-dict.cpp
-#include "slang-ir-strip-cached-dict.h"
-
-#include "slang-ir-insts.h"
-
-namespace Slang
-{
-
-void stripCachedDictionaries(IRModule* module)
-{
-    List<IRInst*> toRemove;
-    for (auto inst : module->getGlobalInsts())
-    {
-        switch (inst->getOp())
-        {
-        case kIROp_GenericSpecializationDictionary:
-        case kIROp_ExistentialFuncSpecializationDictionary:
-        case kIROp_ExistentialTypeSpecializationDictionary:
-            toRemove.add(inst);
-            break;
-        default:
-            continue;
-        }
-    }
-    for (auto inst : toRemove)
-        inst->removeAndDeallocate();
-}
-
-} // namespace Slang
diff --git a/source/slang/slang-ir-strip-cached-dict.h b/source/slang/slang-ir-strip-cached-dict.h
deleted file mode 100644
index c018ab4dd..000000000
--- a/source/slang/slang-ir-strip-cached-dict.h
+++ /dev/null
@@ -1,11 +0,0 @@
-// slang-ir-strip-cached-dict.h
-#pragma once
-
-namespace Slang
-{
-struct IRModule;
-struct IRCall;
-
-/// Removes specialization dictionaries from module.
-void stripCachedDictionaries(IRModule* module);
-} // namespace Slang
diff --git a/source/slang/slang-ir-strip-legalization-insts.cpp b/source/slang/slang-ir-strip-legalization-insts.cpp
new file mode 100644
index 000000000..f55d3f119
--- /dev/null
+++ b/source/slang/slang-ir-strip-legalization-insts.cpp
@@ -0,0 +1,80 @@
+// slang-ir-strip-legalization-insts.cpp
+#include "slang-ir-strip-legalization-insts.h"
+
+#include "slang-ir-insts.h"
+#include "slang-ir.h"
+
+namespace Slang
+{
+
+void stripLegalizationOnlyInstructions(IRModule* module)
+{
+    for (auto inst : module->getGlobalInsts())
+    {
+        switch (inst->getOp())
+        {
+        // Remove cached dictionaries.
+        case kIROp_GenericSpecializationDictionary:
+        case kIROp_ExistentialFuncSpecializationDictionary:
+        case kIROp_ExistentialTypeSpecializationDictionary:
+            {
+                inst->removeAndDeallocate();
+                break;
+            }
+
+        // Remove global param entry point param decoration.
+        case kIROp_GlobalParam:
+            {
+                if (const auto entryPointParamDecoration =
+                        inst->findDecoration<IREntryPointParamDecoration>())
+                    entryPointParamDecoration->removeAndDeallocate();
+                break;
+            }
+
+        // Remove witness tables.
+        // Our goal here is to empty out any witness tables in
+        // the IR so that they don't keep other symbols alive
+        // further into compilation. Luckily we expect all
+        // witness tables to live directly at the global scope
+        // (or inside of a generic, which we can ignore for
+        // now because the emit logic also ignores generics),
+        // and there is a single function we can call to
+        // remove all of the content from the witness tables
+        // (since the key-value associations are stored as
+        // children of each table).
+        case kIROp_WitnessTable:
+            {
+                auto witnessTable = as<IRWitnessTable>(inst);
+                auto conformanceType = witnessTable->getConformanceType();
+                if (!conformanceType ||
+                    !conformanceType->findDecoration<IRComInterfaceDecoration>())
+                {
+                    witnessTable->removeAndDeallocateAllDecorationsAndChildren();
+                }
+                break;
+            }
+
+        default:
+            break;
+        }
+    }
+}
+
+void unpinWitnessTables(IRModule* module)
+{
+    for (auto inst : module->getGlobalInsts())
+    {
+        auto witnessTable = as<IRWitnessTable>(inst);
+        if (!witnessTable)
+            continue;
+
+        // If a witness table is not used for dynamic dispatch, unpin it.
+        if (!witnessTable->findDecoration<IRDynamicDispatchWitnessDecoration>())
+        {
+            while (auto decor = witnessTable->findDecoration<IRKeepAliveDecoration>())
+                decor->removeAndDeallocate();
+        }
+    }
+}
+
+} // namespace Slang
diff --git a/source/slang/slang-ir-strip-legalization-insts.h b/source/slang/slang-ir-strip-legalization-insts.h
new file mode 100644
index 000000000..0e5883118
--- /dev/null
+++ b/source/slang/slang-ir-strip-legalization-insts.h
@@ -0,0 +1,23 @@
+// slang-ir-strip-legalization-insts.h
+#pragma once
+
+namespace Slang
+{
+
+struct IRModule;
+
+/// Removes global instructions from the module that are only required for legalization.
+/// These instructions are safe to or must be removed for instruction emitting.
+///
+/// Currently does the following:
+/// - Removes specialization dictionaries.
+/// - Removes the contents of all witness table instructions.
+/// - Removes global param entry point decorations, as they will cause false code-gen circularity
+///   alerts during pre-emit actions.
+///
+void stripLegalizationOnlyInstructions(IRModule* module);
+
+/// Remove [KeepAlive] decorations from witness tables.
+void unpinWitnessTables(IRModule* module);
+
+} // namespace Slang
diff --git a/source/slang/slang-ir-strip-witness-tables.cpp b/source/slang/slang-ir-strip-witness-tables.cpp
deleted file mode 100644
index 038f8c312..000000000
--- a/source/slang/slang-ir-strip-witness-tables.cpp
+++ /dev/null
@@ -1,53 +0,0 @@
-// slang-ir-strip-witness-tables.cpp
-#include "slang-ir-strip-witness-tables.h"
-
-#include "slang-ir-insts.h"
-#include "slang-ir.h"
-
-namespace Slang
-{
-
-void stripWitnessTables(IRModule* module)
-{
-    // Our goal here is to empty out any witness tables in
-    // the IR so that they don't keep other symbols alive
-    // further into compilation. Luckily we expect all
-    // witness tables to live directly at the global scope
-    // (or inside of a generic, which we can ignore for
-    // now because the emit logic also ignores generics),
-    // and there is a single function we can call to
-    // remove all of the content from the witness tables
-    // (since the key-value associations are stored as
-    // children of each table).
-
-    for (auto inst : module->getGlobalInsts())
-    {
-        auto witnessTable = as<IRWitnessTable>(inst);
-        if (!witnessTable)
-            continue;
-        auto conformanceType = witnessTable->getConformanceType();
-        if (conformanceType && conformanceType->findDecoration<IRComInterfaceDecoration>())
-            continue;
-
-        witnessTable->removeAndDeallocateAllDecorationsAndChildren();
-    }
-}
-
-void unpinWitnessTables(IRModule* module)
-{
-    for (auto inst : module->getGlobalInsts())
-    {
-        auto witnessTable = as<IRWitnessTable>(inst);
-        if (!witnessTable)
-            continue;
-
-        // If a witness table is not used for dynamic dispatch, unpin it.
-        if (!witnessTable->findDecoration<IRDynamicDispatchWitnessDecoration>())
-        {
-            while (auto decor = witnessTable->findDecoration<IRKeepAliveDecoration>())
-                decor->removeAndDeallocate();
-        }
-    }
-}
-
-} // namespace Slang
diff --git a/source/slang/slang-ir-strip-witness-tables.h b/source/slang/slang-ir-strip-witness-tables.h
deleted file mode 100644
index 835ae8fc3..000000000
--- a/source/slang/slang-ir-strip-witness-tables.h
+++ /dev/null
@@ -1,13 +0,0 @@
-// slang-ir-strip-witness-tables.cpp
-#pragma once
-
-namespace Slang
-{
-struct IRModule;
-
-/// Strip the contents of all witness table instructions from the given IR `module`
-void stripWitnessTables(IRModule* module);
-
-/// Remove [KeepAlive] decorations from witness tables.
-void unpinWitnessTables(IRModule* module);
-} // namespace Slang
diff --git a/tests/metal/multi-entry-point-params.slang b/tests/metal/multi-entry-point-params.slang
new file mode 100644
index 000000000..11aefe267
--- /dev/null
+++ b/tests/metal/multi-entry-point-params.slang
@@ -0,0 +1,57 @@
+//TEST:SIMPLE(filecheck=CHECK): -target metal -fvk-use-entrypoint-name
+
+struct FirstStruct {
+    uint a;
+}
+
+struct SecondStruct {
+    uint a;
+}
+
+struct ThirdStruct {
+    uint a;
+}
+
+struct FourthStruct {
+    uint a;
+}
+
+StructuredBuffer<uint> globalScopeBuffer;
+
+RWStructuredBuffer<uint> outBuffer;
+
+//
+// Checks for the following:
+//  - Output entry points will only contain parameters that they originally have from the Slang source. 
+//  - Binding offset calculation for global params originating from entry point parameters are done per-whole file.
+//  - Binding offset for global params in global scope are the same for each generated entry point.
+//
+
+// CHECK: main1({{.*}}globalScopeBuffer{{.*}}buffer(0)]], FirstStruct{{.*}}buffer(2)]], SecondStruct{{.*}}buffer(3)]], float{{.*}}buffer(4)]])
+[shader("compute")]
+[numthreads(5, 1, 1)]
+void main1(
+    uint3 dispatchThreadID: SV_DispatchThreadID,
+    RWStructuredBuffer<FirstStruct> custom,
+    RWStructuredBuffer<SecondStruct> other,
+    ConstantBuffer<float> factor
+) {
+    uint index = dispatchThreadID.x;
+    outBuffer[index] = globalScopeBuffer[index] * custom[index].a * other[index].a * uint(factor);
+}
+
+
+// CHECK-NOT: FirstStruct
+// CHECK-NOT: SecondStruct
+// CHECK: main2({{.*}}globalScopeBuffer{{.*}}buffer(0)]], ThirdStruct{{.*}}buffer(5)]], FourthStruct{{.*}}buffer(6)]])
+[shader("compute")]
+[numthreads(5, 1, 1)]
+void main2(
+    uint3 dispatchThreadID: SV_DispatchThreadID,
+    RWStructuredBuffer<ThirdStruct> custom,
+    RWStructuredBuffer<FourthStruct> other,
+) {
+    uint index = dispatchThreadID.x;
+    outBuffer[index] = globalScopeBuffer[index] * custom[index].a * other[index].a;
+}
+