1 files changed, 37 insertions, 3 deletions
diff --git a/source/slang/slang-emit-c-like.cpp b/source/slang/slang-emit-c-like.cpp
index 7b51495e2..d3a9359ff 100644
--- a/source/slang/slang-emit-c-like.cpp
+++ b/source/slang/slang-emit-c-like.cpp
@@ -295,14 +295,48 @@ void CLikeSourceEmitter::emitSimpleType(IRType* type)
 }
 
 
-/* static */ IRNumThreadsDecoration* CLikeSourceEmitter::getComputeThreadGroupSize(
+IRNumThreadsDecoration* CLikeSourceEmitter::getComputeThreadGroupSize(
     IRFunc* func,
     Int outNumThreads[kThreadGroupAxisCount])
 {
+    Int specializationConstantIds[kThreadGroupAxisCount];
+    IRNumThreadsDecoration* decor =
+        getComputeThreadGroupSize(func, outNumThreads, specializationConstantIds);
+
+    for (auto id : specializationConstantIds)
+    {
+        if (id >= 0)
+        {
+            getSink()->diagnose(decor, Diagnostics::unsupportedSpecializationConstantForNumThreads);
+            break;
+        }
+    }
+    return decor;
+}
+
+/* static */ IRNumThreadsDecoration* CLikeSourceEmitter::getComputeThreadGroupSize(
+    IRFunc* func,
+    Int outNumThreads[kThreadGroupAxisCount],
+    Int outSpecializationConstantIds[kThreadGroupAxisCount])
+{
     IRNumThreadsDecoration* decor = func->findDecoration<IRNumThreadsDecoration>();
-    for (int i = 0; i < 3; ++i)
+    for (int i = 0; i < kThreadGroupAxisCount; ++i)
     {
-        outNumThreads[i] = decor ? Int(getIntVal(decor->getOperand(i))) : 1;
+        if (!decor)
+        {
+            outNumThreads[i] = 1;
+            outSpecializationConstantIds[i] = -1;
+        }
+        else if (auto specConst = as<IRGlobalParam>(decor->getOperand(i)))
+        {
+            outNumThreads[i] = 1;
+            outSpecializationConstantIds[i] = getSpecializationConstantId(specConst);
+        }
+        else
+        {
+            outNumThreads[i] = Int(getIntVal(decor->getOperand(i)));
+            outSpecializationConstantIds[i] = -1;
+        }
     }
     return decor;
 }