From 8ce7c6f6958f9f5ed750ef1a823b9e9ed8c042d8 Mon Sep 17 00:00:00 2001
From: Yong He <yonghe@outlook.com>
Date: Fri, 6 Dec 2024 00:55:35 -0800
Subject: Support specialization constant on WGSL and Metal. (#5780)

---
 source/slang/slang-ir-explicit-global-context.cpp | 24 ++++++++++++++++++++---
 1 file changed, 21 insertions(+), 3 deletions(-)

(limited to 'source/slang/slang-ir-explicit-global-context.cpp')

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
-- 
cgit v1.2.3