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>

Author: fairywreath

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)
     {

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();

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();

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 @@ struct IRBuilder
     {
         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

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