From cbdc7e1219e472fd74f7f559d7e417f233e7df39 Mon Sep 17 00:00:00 2001
From: Julius Ikkala <julius.ikkala@gmail.com>
Date: Tue, 14 Jan 2025 20:32:29 +0200
Subject: Implement specialization constant support in numthreads / local_size
 (#5963)

* Allow using specialization constants in numthreads attribute

* Add support for GLSL local_size_x_id syntax

* Fix overeager specialization constant parsing

* Add diagnostics for specialization constant numthreads

* Remove unused variable

* Fix local_size_x_id not finding existing specialization constant

* Allow materializeGetWorkGroupSize to reference specialization constants

* Use SpvOpExecutionModeId for modes that require it

* Cleanup specialization constant numthreads code

* Add tests for specialization constant work group sizes

* Fix implicit Slang::Int -> int32_t cast

* Fix querying thread group size in reflection API

---------

Co-authored-by: Yong He <yonghe@outlook.com>
---
 source/slang/slang-lower-to-ir.cpp | 46 +++++++++++++++++++++++++++++++++-----
 1 file changed, 40 insertions(+), 6 deletions(-)

(limited to 'source/slang/slang-lower-to-ir.cpp')

diff --git a/source/slang/slang-lower-to-ir.cpp b/source/slang/slang-lower-to-ir.cpp
index e82fc03fd..086345719 100644
--- a/source/slang/slang-lower-to-ir.cpp
+++ b/source/slang/slang-lower-to-ir.cpp
@@ -7625,12 +7625,29 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
             {
                 verifyComputeDerivativeGroupModifier = true;
                 getAllEntryPointsNoOverride(entryPoints);
+
+                LoweredValInfo extents[3];
+
+                for (int i = 0; i < 3; ++i)
+                {
+                    extents[i] = layoutLocalSizeAttr->specConstExtents[i]
+                                     ? emitDeclRef(
+                                           context,
+                                           layoutLocalSizeAttr->specConstExtents[i],
+                                           lowerType(
+                                               context,
+                                               getType(
+                                                   context->astBuilder,
+                                                   layoutLocalSizeAttr->specConstExtents[i])))
+                                     : lowerVal(context, layoutLocalSizeAttr->extents[i]);
+                }
+
                 for (auto d : entryPoints)
                     as<IRNumThreadsDecoration>(getBuilder()->addNumThreadsDecoration(
                         d,
-                        getSimpleVal(context, lowerVal(context, layoutLocalSizeAttr->x)),
-                        getSimpleVal(context, lowerVal(context, layoutLocalSizeAttr->y)),
-                        getSimpleVal(context, lowerVal(context, layoutLocalSizeAttr->z))));
+                        getSimpleVal(context, extents[0]),
+                        getSimpleVal(context, extents[1]),
+                        getSimpleVal(context, extents[2])));
             }
             else if (as<GLSLLayoutDerivativeGroupQuadAttribute>(modifier))
             {
@@ -10336,11 +10353,28 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
             }
             else if (auto numThreadsAttr = as<NumThreadsAttribute>(modifier))
             {
+                LoweredValInfo extents[3];
+
+                for (int i = 0; i < 3; ++i)
+                {
+                    extents[i] = numThreadsAttr->specConstExtents[i]
+                                     ? emitDeclRef(
+                                           context,
+                                           numThreadsAttr->specConstExtents[i],
+                                           lowerType(
+                                               context,
+                                               getType(
+                                                   context->astBuilder,
+                                                   numThreadsAttr->specConstExtents[i])))
+                                     : lowerVal(context, numThreadsAttr->extents[i]);
+                }
+
                 numThreadsDecor = as<IRNumThreadsDecoration>(getBuilder()->addNumThreadsDecoration(
                     irFunc,
-                    getSimpleVal(context, lowerVal(context, numThreadsAttr->x)),
-                    getSimpleVal(context, lowerVal(context, numThreadsAttr->y)),
-                    getSimpleVal(context, lowerVal(context, numThreadsAttr->z))));
+                    getSimpleVal(context, extents[0]),
+                    getSimpleVal(context, extents[1]),
+                    getSimpleVal(context, extents[2])));
+                numThreadsDecor->sourceLoc = numThreadsAttr->loc;
             }
             else if (auto waveSizeAttr = as<WaveSizeAttribute>(modifier))
             {
-- 
cgit v1.2.3