Add reflection support for GLSL thread-group-size modifier

Fixes #15

These are the modifiers like:

    layout(local_size_x = 16) in;

Unlike the HLSL case, these don't get attache to the entry point function itself, so there is a bit more work involed in looking them up.

Just to make sure I didn't mess up the HLSL case, I went ahead and added two tests for this capability: one for GLSL and one for HLSL.

7 files changed, 168 insertions, 5 deletions

diff --git a/source/slang/modifier-defs.h b/source/slang/modifier-defs.h
index 665551b35..fadc68695 100644
--- a/source/slang/modifier-defs.h
+++ b/source/slang/modifier-defs.h
@@ -111,6 +111,11 @@ SIMPLE_SYNTAX_CLASS(GLSLSetLayoutModifier        , GLSLParsedLayoutModifier)
 SIMPLE_SYNTAX_CLASS(GLSLLocationLayoutModifier   , GLSLParsedLayoutModifier)
 SIMPLE_SYNTAX_CLASS(GLSLPushConstantLayoutModifier, GLSLParsedLayoutModifier)
 
+SIMPLE_SYNTAX_CLASS(GLSLLocalSizeLayoutModifier,    GLSLUnparsedLayoutModifier)
+SIMPLE_SYNTAX_CLASS(GLSLLocalSizeXLayoutModifier,    GLSLLocalSizeLayoutModifier)
+SIMPLE_SYNTAX_CLASS(GLSLLocalSizeYLayoutModifier,    GLSLLocalSizeLayoutModifier)
+SIMPLE_SYNTAX_CLASS(GLSLLocalSizeZLayoutModifier,    GLSLLocalSizeLayoutModifier)
+
 // A catch-all for single-keyword modifiers
 SIMPLE_SYNTAX_CLASS(SimpleModifier, Modifier)
 
diff --git a/source/slang/parser.cpp b/source/slang/parser.cpp
index e0ff85164..1ad0c6b94 100644
--- a/source/slang/parser.cpp
+++ b/source/slang/parser.cpp
@@ -748,6 +748,9 @@ namespace Slang
                     CASE(set,           GLSLSetLayoutModifier);
                     CASE(location,      GLSLLocationLayoutModifier);
                     CASE(push_constant, GLSLPushConstantLayoutModifier);
+                    CASE(local_size_x,  GLSLLocalSizeXLayoutModifier);
+                    CASE(local_size_y,  GLSLLocalSizeYLayoutModifier);
+                    CASE(local_size_z,  GLSLLocalSizeZLayoutModifier);
 
                 #undef CASE
                     else
diff --git a/source/slang/reflection.cpp b/source/slang/reflection.cpp
index 3c89194c8..aaeae5595 100644
--- a/source/slang/reflection.cpp
+++ b/source/slang/reflection.cpp
@@ -204,7 +204,7 @@ SLANG_API size_t spReflectionType_GetElementCount(SlangReflectionType* inType)
 
     if(auto arrayType = dynamic_cast<ArrayExpressionType*>(type))
     {
-        return (size_t) GetIntVal(arrayType->ArrayLength);
+        return arrayType->ArrayLength ? (size_t) GetIntVal(arrayType->ArrayLength) : 0;
     }
     else if( auto vectorType = dynamic_cast<VectorExpressionType*>(type))
     {
@@ -751,12 +751,43 @@ SLANG_API void spReflectionEntryPoint_getComputeThreadGroupSize(
     auto entryPointFunc = entryPointLayout->entryPoint;
     if(!entryPointFunc) return;
 
+    SlangUInt sizeAlongAxis[3] = { 1, 1, 1 };
+
+    // First look for the HLSL case, where we have an attribute attached to the entry point function
     auto numThreadsAttribute = entryPointFunc->FindModifier<HLSLNumThreadsAttribute>();
-    if(!numThreadsAttribute) return;
+    if (numThreadsAttribute)
+    {
+        sizeAlongAxis[0] = numThreadsAttribute->x;
+        sizeAlongAxis[1] = numThreadsAttribute->y;
+        sizeAlongAxis[2] = numThreadsAttribute->z;
+    }
+    else
+    {
+        // Fall back to the GLSL case, which requires a search over global-scope declarations
+        // to look for anything with the `local_size_*` qualifier
+        auto module = dynamic_cast<ProgramSyntaxNode*>(entryPointFunc->ParentDecl);
+        if (module)
+        {
+            for (auto dd : module->Members)
+            {
+                for (auto mod : dd->GetModifiersOfType<GLSLLocalSizeLayoutModifier>())
+                {
+                    if (auto xMod = dynamic_cast<GLSLLocalSizeXLayoutModifier*>(mod))
+                        sizeAlongAxis[0] = (SlangUInt) getIntegerLiteralValue(xMod->valToken);
+                    else if (auto yMod = dynamic_cast<GLSLLocalSizeYLayoutModifier*>(mod))
+                        sizeAlongAxis[1] = (SlangUInt) getIntegerLiteralValue(yMod->valToken);
+                    else if (auto zMod = dynamic_cast<GLSLLocalSizeZLayoutModifier*>(mod))
+                        sizeAlongAxis[2] = (SlangUInt) getIntegerLiteralValue(zMod->valToken);
+                }
+            }
+        }
+    }
+
+    //
 
-    if(axisCount > 0) outSizeAlongAxis[0] = numThreadsAttribute->x;
-    if(axisCount > 1) outSizeAlongAxis[1] = numThreadsAttribute->y;
-    if(axisCount > 2) outSizeAlongAxis[2] = numThreadsAttribute->z;
+    if(axisCount > 0) outSizeAlongAxis[0] = sizeAlongAxis[0];
+    if(axisCount > 1) outSizeAlongAxis[1] = sizeAlongAxis[1];
+    if(axisCount > 2) outSizeAlongAxis[2] = sizeAlongAxis[2];
     for( SlangUInt aa = 3; aa < axisCount; ++aa )
     {
         outSizeAlongAxis[aa] = 1;
@@ -1526,6 +1557,20 @@ static void emitReflectionEntryPointJSON(
         write(writer, ",\n\"usesAnySampleRateInput\": true");
     }
 
+    if (entryPoint->getStage() == SLANG_STAGE_COMPUTE)
+    {
+        SlangUInt threadGroupSize[3];
+        entryPoint->getComputeThreadGroupSize(3, threadGroupSize);
+
+        write(writer, ",\n\"threadGroupSize\": [");
+        for (int ii = 0; ii < 3; ++ii)
+        {
+            if (ii != 0) write(writer, ", ");
+            write(writer, threadGroupSize[ii]);
+        }
+        write(writer, "]");
+    }
+
     dedent(writer);
     write(writer, "\n}");
 }
diff --git a/tests/reflection/thread-group-size.comp b/tests/reflection/thread-group-size.comp
new file mode 100644
index 000000000..2edf2c2df
--- /dev/null
+++ b/tests/reflection/thread-group-size.comp
@@ -0,0 +1,18 @@
+//TEST:SIMPLE:-no-checking -target reflection-json
+
+// Confirm that we provide reflection data for the `local_size_*` attributes
+
+layout(local_size_x = 3) in;
+
+layout(local_size_y = 5, local_size_z = 7) in;
+
+buffer B
+{
+	float b[];	
+};
+
+void main()
+{
+	uint tid = gl_GlobalInvocationID.x;
+	b[tid] = b[tid + 1] + 1.0f;
+}
\ No newline at end of file
diff --git a/tests/reflection/thread-group-size.comp.expected b/tests/reflection/thread-group-size.comp.expected
new file mode 100644
index 000000000..ea7e97851
--- /dev/null
+++ b/tests/reflection/thread-group-size.comp.expected
@@ -0,0 +1,42 @@
+result code = 0
+standard error = {
+}
+standard output = {
+{
+    "parameters": [
+        {
+            "name": "B",
+            "binding": {"kind": "descriptorTableSlot", "index": 0},
+            "type": {
+                "kind": "shaderStorageBuffer",
+                "elementType": {
+                    "kind": "struct",
+                    "fields": [
+                        {
+                            "name": "b",
+                            "type": {
+                                "kind": "array",
+                                "elementCount": 0,
+                                "elementType": {
+                                    "kind": "scalar",
+                                    "scalarType": "float32"
+                                }
+                            },
+                            "bindings": [
+
+                            ]
+                        }
+                    ]
+                }
+            }
+        }
+    ],
+    "entryPoints": [
+        {
+            "name": "main",
+            "stage:": "compute",
+            "threadGroupSize": [3, 5, 7]
+        }
+    ]
+}
+}
diff --git a/tests/reflection/thread-group-size.hlsl b/tests/reflection/thread-group-size.hlsl
new file mode 100644
index 000000000..650a41e46
--- /dev/null
+++ b/tests/reflection/thread-group-size.hlsl
@@ -0,0 +1,11 @@
+//TEST:SIMPLE:-profile cs_5_0 -target reflection-json
+
+// Confirm that we provide reflection data for the `numthreads` attribute
+
+RWStructuredBuffer<float> b;
+
+[numthreads(3,5,7)]
+void main(uint3 tid : SV_DispatchThreadID)
+{
+	b[tid.x] = b[tid.x + 1] + 1.0f;
+}
\ No newline at end of file
diff --git a/tests/reflection/thread-group-size.hlsl.expected b/tests/reflection/thread-group-size.hlsl.expected
new file mode 100644
index 000000000..60d5e822c
--- /dev/null
+++ b/tests/reflection/thread-group-size.hlsl.expected
@@ -0,0 +1,39 @@
+result code = 0
+standard error = {
+}
+standard output = {
+{
+    "parameters": [
+        {
+            "name": "b",
+            "binding": {"kind": "unorderedAccess", "index": 0},
+            "type": {
+                "kind": "resource",
+                "baseShape": "structuredBuffer",
+                "access": "readWrite"
+            }
+        }
+    ],
+    "entryPoints": [
+        {
+            "name": "main",
+            "stage:": "compute",
+            "parameters": [
+                {
+                    "name": "tid",
+                    "binding": {"kind": "vertexInput", "index": 0},
+                    "type": {
+                        "kind": "vector",
+                        "elementCount": 3,
+                        "elementType": {
+                            "kind": "scalar",
+                            "scalarType": "uint32"
+                        }
+                    }
+                }
+            ],
+            "threadGroupSize": [3, 5, 7]
+        }
+    ]
+}
+}