Support specialization constant on WGSL and Metal. (#5780)

7 files changed, 170 insertions, 71 deletions

diff --git a/source/slang/slang-emit-metal.cpp b/source/slang/slang-emit-metal.cpp
index 3548e9369..e6e2a3ec5 100644
--- a/source/slang/slang-emit-metal.cpp
+++ b/source/slang/slang-emit-metal.cpp
@@ -1331,9 +1331,53 @@ bool MetalSourceEmitter::_emitUserSemantic(
     return false;
 }
 
+bool MetalSourceEmitter::tryEmitGlobalParamImpl(IRGlobalParam* varDecl, IRType* varType)
+{
+    auto layout = getVarLayout(varDecl);
+    if (!layout)
+        return false;
+    if (auto specConstLayout = layout->findOffsetAttr(LayoutResourceKind::SpecializationConstant))
+    {
+        // Emit specialization constant.
+        auto name = getName(varDecl);
+        auto prefixName = "fc_" + name;
+        auto defaultVal = varDecl->findDecoration<IRDefaultValueDecoration>();
+
+        m_writer->emit("constant ");
+        emitType(varType, prefixName);
+        m_writer->emit(" ");
+        m_writer->emit("[[function_constant(");
+        m_writer->emit(specConstLayout->getOffset());
+        m_writer->emit(")]];\n");
+
+        m_writer->emit("constant ");
+        emitType(varType, name);
+        m_writer->emit(" = ");
+        if (defaultVal)
+        {
+            m_writer->emit("is_function_constant_defined(");
+            m_writer->emit(prefixName);
+            m_writer->emit(") ? ");
+            m_writer->emit(prefixName);
+            m_writer->emit(" : ");
+            emitVal(defaultVal->getOperand(0), getInfo(EmitOp::General));
+        }
+        else
+        {
+            m_writer->emit(prefixName);
+        }
+        m_writer->emit(";\n");
+        return true;
+    }
+    return false;
+}
+
 void MetalSourceEmitter::emitSemanticsImpl(IRInst* inst, bool allowOffsets)
 {
     SLANG_UNUSED(allowOffsets);
+
+    auto varLayout = findVarLayout(inst);
+
     if (inst->getOp() == kIROp_StructKey)
     {
         // Only emit [[attribute(n)]] on struct keys.
@@ -1341,7 +1385,6 @@ void MetalSourceEmitter::emitSemanticsImpl(IRInst* inst, bool allowOffsets)
         if (maybeEmitSystemSemantic(inst))
             return;
 
-        auto varLayout = findVarLayout(inst);
         bool hasSemantic = false;
 
         if (varLayout)
@@ -1378,6 +1421,7 @@ void MetalSourceEmitter::emitSemanticsImpl(IRInst* inst, bool allowOffsets)
                     semanticDecor->getSemanticIndex());
             }
         }
+        return;
     }
 }
 
diff --git a/source/slang/slang-emit-metal.h b/source/slang/slang-emit-metal.h
index 7cd01c646..17562b1c4 100644
--- a/source/slang/slang-emit-metal.h
+++ b/source/slang/slang-emit-metal.h
@@ -40,6 +40,7 @@ protected:
 
     virtual void emitRateQualifiersAndAddressSpaceImpl(IRRate* rate, AddressSpace addressSpace)
         SLANG_OVERRIDE;
+    virtual bool tryEmitGlobalParamImpl(IRGlobalParam* varDecl, IRType* varType) SLANG_OVERRIDE;
     virtual void emitSemanticsImpl(IRInst* inst, bool allowOffsets) SLANG_OVERRIDE;
     virtual void emitSimpleFuncParamImpl(IRParam* param) SLANG_OVERRIDE;
     virtual void emitPostDeclarationAttributesForType(IRInst* type) SLANG_OVERRIDE;
diff --git a/source/slang/slang-emit-wgsl.cpp b/source/slang/slang-emit-wgsl.cpp
index cba9c1b54..8d4b33541 100644
--- a/source/slang/slang-emit-wgsl.cpp
+++ b/source/slang/slang-emit-wgsl.cpp
@@ -643,6 +643,21 @@ void WGSLSourceEmitter::emitSimpleTypeImpl(IRType* type)
     }
 }
 
+void WGSLSourceEmitter::emitGlobalParamDefaultVal(IRGlobalParam* varDecl)
+{
+    auto layout = getVarLayout(varDecl);
+    if (!layout)
+        return;
+    if (layout->findOffsetAttr(LayoutResourceKind::SpecializationConstant))
+    {
+        if (auto defaultValDecor = varDecl->findDecoration<IRDefaultValueDecoration>())
+        {
+            m_writer->emit(" = ");
+            emitInstExpr(defaultValDecor->getOperand(0), EmitOpInfo());
+        }
+    }
+}
+
 void WGSLSourceEmitter::emitLayoutQualifiersImpl(IRVarLayout* layout)
 {
 
@@ -670,6 +685,14 @@ void WGSLSourceEmitter::emitLayoutQualifiersImpl(IRVarLayout* layout)
 
             return;
         }
+        else if (kind == LayoutResourceKind::SpecializationConstant)
+        {
+            m_writer->emit("@id(");
+            m_writer->emit(attr->getOffset());
+            m_writer->emit(") ");
+
+            return;
+        }
     }
 }
 
@@ -708,7 +731,15 @@ void WGSLSourceEmitter::emitVarKeywordImpl(IRType* type, IRInst* varDecl)
     case kIROp_GlobalParam:
     case kIROp_GlobalVar:
     case kIROp_Var:
-        m_writer->emit("var");
+        {
+            auto layout = getVarLayout(varDecl);
+            if (layout && layout->findOffsetAttr(LayoutResourceKind::SpecializationConstant))
+            {
+                m_writer->emit("override");
+                break;
+            }
+            m_writer->emit("var");
+        }
         break;
     default:
         if (isStaticConst(varDecl))
diff --git a/source/slang/slang-emit-wgsl.h b/source/slang/slang-emit-wgsl.h
index 0cbedf4eb..28311aa27 100644
--- a/source/slang/slang-emit-wgsl.h
+++ b/source/slang/slang-emit-wgsl.h
@@ -52,6 +52,7 @@ public:
         UnownedStringSlice intrinsicDefinition,
         IRInst* intrinsicInst,
         EmitOpInfo const& inOuterPrec) SLANG_OVERRIDE;
+    virtual void emitGlobalParamDefaultVal(IRGlobalParam* varDecl) SLANG_OVERRIDE;
 
     void emit(const AddressSpace addressSpace);
 
diff --git a/source/slang/slang-ir-explicit-global-context.cpp b/source/slang/slang-ir-explicit-global-context.cpp
index d891263cd..0516d574c 100644
--- a/source/slang/slang-ir-explicit-global-context.cpp
+++ b/source/slang/slang-ir-explicit-global-context.cpp
@@ -3,6 +3,7 @@
 
 #include "slang-ir-clone.h"
 #include "slang-ir-insts.h"
+#include "slang-ir-util.h"
 
 namespace Slang
 {
@@ -39,7 +40,8 @@ struct IntroduceExplicitGlobalContextPass
     class ExplicitContextPolicy
     {
     public:
-        ExplicitContextPolicy(CodeGenTarget target)
+        ExplicitContextPolicy(CodeGenTarget inTarget)
+            : target(inTarget)
         {
             switch (target)
             {
@@ -74,7 +76,7 @@ struct IntroduceExplicitGlobalContextPass
             return (UInt)hoistableGlobalObjectKind & (UInt)hoistable;
         }
 
-        bool canHoistGlobalVar(IRGlobalVar* inst)
+        bool canHoistGlobalVar(IRInst* inst)
         {
             if (!((UInt)hoistGlobalVarOptions & (UInt)HoistGlobalVarOptions::SharedGlobal) &&
                 as<IRGroupSharedRate>(inst->getRate()))
@@ -99,6 +101,19 @@ struct IntroduceExplicitGlobalContextPass
                 }
             }
 
+            // Do not move specialization constants to context.
+            switch (target)
+            {
+            case CodeGenTarget::Metal:
+            case CodeGenTarget::MetalLib:
+            case CodeGenTarget::MetalLibAssembly:
+                {
+                    auto varLayout = findVarLayout(inst);
+                    if (varLayout &&
+                        varLayout->findOffsetAttr(LayoutResourceKind::SpecializationConstant))
+                        return false;
+                }
+            }
             return true;
         }
 
@@ -111,6 +126,7 @@ struct IntroduceExplicitGlobalContextPass
         GlobalObjectKind hoistableGlobalObjectKind = GlobalObjectKind::All;
         bool requiresFuncTypeCorrectionPass = false;
         AddressSpace addressSpaceOfLocals = AddressSpace::ThreadLocal;
+        CodeGenTarget target;
     };
 
     IntroduceExplicitGlobalContextPass(IRModule* module, CodeGenTarget target)
@@ -134,7 +150,7 @@ struct IntroduceExplicitGlobalContextPass
 
     bool canHoistType(GlobalObjectKind hoistable) { return m_options.canHoistType(hoistable); }
 
-    bool canHoistGlobalVar(IRGlobalVar* inst) { return m_options.canHoistGlobalVar(inst); }
+    bool canHoistGlobalVar(IRInst* inst) { return m_options.canHoistGlobalVar(inst); }
 
     void processModule()
     {
@@ -183,6 +199,8 @@ struct IntroduceExplicitGlobalContextPass
                     //
                     auto globalParam = cast<IRGlobalParam>(inst);
 
+                    if (!canHoistGlobalVar(globalParam))
+                        continue;
 
                     // One detail we need to be careful about is that as a result
                     // of legalizing the varying parameters of compute kernels to
diff --git a/source/slang/slang-parameter-binding.cpp b/source/slang/slang-parameter-binding.cpp
index f9c4f0c94..1303a59f1 100644
--- a/source/slang/slang-parameter-binding.cpp
+++ b/source/slang/slang-parameter-binding.cpp
@@ -1081,25 +1081,9 @@ static void addExplicitParameterBindings_GLSL(
     RefPtr<ParameterInfo> parameterInfo,
     RefPtr<VarLayout> varLayout)
 {
-    // We only want to apply GLSL-style layout modifers
-    // when compiling for a Khronos-related target.
-    //
-    // TODO: This should have some finer granularity
-    // so that we are able to distinguish between
-    // Vulkan and OpenGL as targets.
-    //
-    if (!isKhronosTarget(context->getTargetRequest()) && !isWGPUTarget(context->getTargetRequest()))
-        return;
-
     auto typeLayout = varLayout->typeLayout;
     auto varDecl = varLayout->varDecl;
 
-    // The catch in GLSL is that the expected resource type
-    // is implied by the parameter declaration itself, and
-    // the `layout` modifier is only allowed to adjust
-    // the index/offset/etc.
-    //
-
     enum
     {
         kResInfo = 0,
@@ -1120,56 +1104,12 @@ static void addExplicitParameterBindings_GLSL(
     };
     ResAndSemanticInfo info[kMaxResCount] = {};
 
-    if (auto foundInputAttachmentIndex =
-            typeLayout->FindResourceInfo(LayoutResourceKind::InputAttachmentIndex))
-    {
-        foundResInfo = foundInputAttachmentIndex;
-        // Try to find `input_attachment_index`
-        if (auto glslAttachmentIndexAttr =
-                varDecl.getDecl()->findModifier<GLSLInputAttachmentIndexLayoutAttribute>())
-        {
-            info[kSubpassResInfo].resInfo = foundResInfo;
-            // Subpass fills semantic info of a descriptor and subpass
-            info[kSubpassResInfo].semanticInfo.index = (UInt)glslAttachmentIndexAttr->location;
-            info[kSubpassResInfo].semanticInfo.space = 0;
-        }
-    }
+    // First, we want to apply offsets of specialization constants for platforms that supports them.
+    if (!isKhronosTarget(context->getTargetRequest()) &&
+        !isWGPUTarget(context->getTargetRequest()) && !isMetalTarget(context->getTargetRequest()))
+        return;
 
-    if (auto foundDescriptorTableSlot =
-            typeLayout->FindResourceInfo(LayoutResourceKind::DescriptorTableSlot))
-    {
-        foundResInfo = foundDescriptorTableSlot;
-        // Try to find `binding` and `set`
-        if (auto glslBindingAttr = varDecl.getDecl()->findModifier<GLSLBindingAttribute>())
-        {
-            info[kResInfo].resInfo = foundResInfo;
-            info[kResInfo].semanticInfo.index = glslBindingAttr->binding;
-            info[kResInfo].semanticInfo.space = glslBindingAttr->set;
-        }
-    }
-    else if (
-        auto foundSubElementRegisterSpace =
-            typeLayout->FindResourceInfo(LayoutResourceKind::SubElementRegisterSpace))
-    {
-        foundResInfo = foundSubElementRegisterSpace;
-        // Try to find `set`
-        if (auto attr = varDecl.getDecl()->findModifier<GLSLBindingAttribute>())
-        {
-            info[kResInfo].resInfo = foundResInfo;
-            if (attr->binding != 0)
-            {
-                getSink(context)->diagnose(
-                    attr,
-                    Diagnostics::wholeSpaceParameterRequiresZeroBinding,
-                    varDecl.getName(),
-                    attr->binding);
-            }
-            info[kResInfo].semanticInfo.index = attr->set;
-            info[kResInfo].semanticInfo.space = 0;
-        }
-    }
-    else if (
-        auto foundSpecializationConstant =
+    if (auto foundSpecializationConstant =
             typeLayout->FindResourceInfo(LayoutResourceKind::SpecializationConstant))
     {
         info[kResInfo].resInfo = foundSpecializationConstant;
@@ -1180,6 +1120,70 @@ static void addExplicitParameterBindings_GLSL(
             return;
     }
 
+    // For remaining cases, we only want to apply GLSL-style layout modifers
+    // when compiling for Khronos and WGSL targets.
+    //
+    // TODO: This should have some finer granularity
+    // so that we are able to distinguish between
+    // Vulkan and OpenGL as targets.
+    //
+    if (isKhronosTarget(context->getTargetRequest()) || isWGPUTarget(context->getTargetRequest()))
+    {
+        // The catch in GLSL is that the expected resource type
+        // is implied by the parameter declaration itself, and
+        // the `layout` modifier is only allowed to adjust
+        // the index/offset/etc.
+        //
+
+        if (auto foundInputAttachmentIndex =
+                typeLayout->FindResourceInfo(LayoutResourceKind::InputAttachmentIndex))
+        {
+            foundResInfo = foundInputAttachmentIndex;
+            // Try to find `input_attachment_index`
+            if (auto glslAttachmentIndexAttr =
+                    varDecl.getDecl()->findModifier<GLSLInputAttachmentIndexLayoutAttribute>())
+            {
+                info[kSubpassResInfo].resInfo = foundResInfo;
+                // Subpass fills semantic info of a descriptor and subpass
+                info[kSubpassResInfo].semanticInfo.index = (UInt)glslAttachmentIndexAttr->location;
+                info[kSubpassResInfo].semanticInfo.space = 0;
+            }
+        }
+
+        if (auto foundDescriptorTableSlot =
+                typeLayout->FindResourceInfo(LayoutResourceKind::DescriptorTableSlot))
+        {
+            foundResInfo = foundDescriptorTableSlot;
+            // Try to find `binding` and `set`
+            if (auto glslBindingAttr = varDecl.getDecl()->findModifier<GLSLBindingAttribute>())
+            {
+                info[kResInfo].resInfo = foundResInfo;
+                info[kResInfo].semanticInfo.index = glslBindingAttr->binding;
+                info[kResInfo].semanticInfo.space = glslBindingAttr->set;
+            }
+        }
+        else if (
+            auto foundSubElementRegisterSpace =
+                typeLayout->FindResourceInfo(LayoutResourceKind::SubElementRegisterSpace))
+        {
+            foundResInfo = foundSubElementRegisterSpace;
+            // Try to find `set`
+            if (auto attr = varDecl.getDecl()->findModifier<GLSLBindingAttribute>())
+            {
+                info[kResInfo].resInfo = foundResInfo;
+                if (attr->binding != 0)
+                {
+                    getSink(context)->diagnose(
+                        attr,
+                        Diagnostics::wholeSpaceParameterRequiresZeroBinding,
+                        varDecl.getName(),
+                        attr->binding);
+                }
+                info[kResInfo].semanticInfo.index = attr->set;
+                info[kResInfo].semanticInfo.space = 0;
+            }
+        }
+    }
 
     auto varDeclBase = as<VarDeclBase>(varDecl);
     bool hasABinding = false;
diff --git a/source/slang/slang-type-layout.cpp b/source/slang/slang-type-layout.cpp
index da4cd458b..5e63a7cfd 100644
--- a/source/slang/slang-type-layout.cpp
+++ b/source/slang/slang-type-layout.cpp
@@ -2013,7 +2013,7 @@ LayoutRulesImpl* MetalLayoutRulesFamilyImpl::getVaryingOutputRules()
 
 LayoutRulesImpl* MetalLayoutRulesFamilyImpl::getSpecializationConstantRules()
 {
-    return nullptr;
+    return &kGLSLSpecializationConstantLayoutRulesImpl_;
 }
 
 LayoutRulesImpl* MetalLayoutRulesFamilyImpl::getShaderStorageBufferRules(CompilerOptionSet&)