From 2d0912bfe2de7799b32e80722fa5c8dc279a339b Mon Sep 17 00:00:00 2001
From: Ellie Hermaszewska <ellieh@nvidia.com>
Date: Wed, 31 Jan 2024 03:28:04 +0800
Subject: Correctly apply glsl local size layout to entry points during
 lowering (#3528)

* Correctly apply glsl local size layout to entry points during lowering

* Test for glsl layout correctness
---
 source/slang/slang-lower-to-ir.cpp | 40 +++++++++++++++++++++++++++-----------
 1 file changed, 29 insertions(+), 11 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 00db77511..e0bf72f7d 100644
--- a/source/slang/slang-lower-to-ir.cpp
+++ b/source/slang/slang-lower-to-ir.cpp
@@ -6938,7 +6938,6 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
     IGNORED_CASE(IncludeDecl)
     IGNORED_CASE(ImplementingDecl)
     IGNORED_CASE(UsingDecl)
-    IGNORED_CASE(EmptyDecl)
     IGNORED_CASE(SyntaxDecl)
     IGNORED_CASE(AttributeDecl)
     IGNORED_CASE(NamespaceDecl)
@@ -6947,6 +6946,29 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
 
 #undef IGNORED_CASE
 
+    LoweredValInfo visitEmptyDecl(EmptyDecl* decl)
+    {
+        for(const auto modifier : decl->modifiers)
+        {
+            if(const auto layoutLocalSizeAttr = as<GLSLLayoutLocalSizeAttribute>(modifier))
+            {
+                for(const auto d : context->irBuilder->getModule()->getModuleInst()->getGlobalInsts())
+                {
+                    if(d->findDecoration<IREntryPointDecoration>())
+                    {
+                        getBuilder()->addNumThreadsDecoration(
+                            d,
+                            layoutLocalSizeAttr->x,
+                            layoutLocalSizeAttr->y,
+                            layoutLocalSizeAttr->z
+                        );
+                    }
+                }
+            }
+        }
+        return LoweredValInfo();
+    }
+
     void ensureInsertAtGlobalScope(IRBuilder* builder)
     {
         auto inst = builder->getInsertLoc().getInst();
@@ -9325,16 +9347,12 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
             }
             else if (auto numThreadsAttr = as<NumThreadsAttribute>(modifier))
             {
-                auto builder = getBuilder();
-                IRType* intType = builder->getIntType();
-
-                IRInst* operands[3] = {
-                    builder->getIntValue(intType, numThreadsAttr->x),
-                    builder->getIntValue(intType, numThreadsAttr->y),
-                    builder->getIntValue(intType, numThreadsAttr->z)
-                };
-
-                builder->addDecoration(irFunc, kIROp_NumThreadsDecoration, operands, 3);
+                getBuilder()->addNumThreadsDecoration(
+                    irFunc,
+                    numThreadsAttr->x,
+                    numThreadsAttr->y,
+                    numThreadsAttr->z
+                );
             }
             else if (as<ReadNoneAttribute>(modifier))
             {
-- 
cgit v1.2.3