Merge pull request #572 from steve-lunarg/numthreads

HLSL: implement numthreads for compute shaders

10 files changed, 352 insertions, 24 deletions

diff --git a/Test/baseResults/hlsl.numthreads.comp.out b/Test/baseResults/hlsl.numthreads.comp.out
new file mode 100644
index 00000000..76e95c8f
--- /dev/null
+++ b/Test/baseResults/hlsl.numthreads.comp.out
@@ -0,0 +1,60 @@
+hlsl.numthreads.comp
+Shader version: 450
+local_size = (4, 4, 2)
+0:? Sequence
+0:4  Function Definition: main(vu3; (temp void)
+0:4    Function Parameters: 
+0:4      'tid' (in 3-component vector of uint)
+0:9  Function Definition: main_aux1(vu3; (temp void)
+0:9    Function Parameters: 
+0:9      'tid' (in 3-component vector of uint LocalInvocationID)
+0:?   Linker Objects
+0:?     'tid' (in 3-component vector of uint LocalInvocationID)
+
+
+Linked compute stage:
+
+
+Shader version: 450
+local_size = (4, 4, 2)
+0:? Sequence
+0:4  Function Definition: main(vu3; (temp void)
+0:4    Function Parameters: 
+0:4      'tid' (in 3-component vector of uint)
+0:9  Function Definition: main_aux1(vu3; (temp void)
+0:9    Function Parameters: 
+0:9      'tid' (in 3-component vector of uint LocalInvocationID)
+0:?   Linker Objects
+0:?     'tid' (in 3-component vector of uint LocalInvocationID)
+
+// Module Version 10000
+// Generated by (magic number): 80001
+// Id's are bound by 15
+
+                              Capability Shader
+               1:             ExtInstImport  "GLSL.std.450"
+                              MemoryModel Logical GLSL450
+                              EntryPoint GLCompute 4  "main_aux1" 14
+                              ExecutionMode 4 LocalSize 4 4 2
+                              Name 4  "main_aux1"
+                              Name 11  "main(vu3;"
+                              Name 10  "tid"
+                              Name 14  "tid"
+                              Decorate 14(tid) BuiltIn LocalInvocationId
+               2:             TypeVoid
+               3:             TypeFunction 2
+               6:             TypeInt 32 0
+               7:             TypeVector 6(int) 3
+               8:             TypePointer Function 7(ivec3)
+               9:             TypeFunction 2 8(ptr)
+              13:             TypePointer Input 7(ivec3)
+         14(tid):     13(ptr) Variable Input
+    4(main_aux1):           2 Function None 3
+               5:             Label
+                              Return
+                              FunctionEnd
+   11(main(vu3;):           2 Function None 9
+         10(tid):      8(ptr) FunctionParameter
+              12:             Label
+                              Return
+                              FunctionEnd
diff --git a/Test/hlsl.numthreads.comp b/Test/hlsl.numthreads.comp
new file mode 100644
index 00000000..fcc97f37
--- /dev/null
+++ b/Test/hlsl.numthreads.comp
@@ -0,0 +1,14 @@
+
+[numthreads(8,8,1)]
+void main(uint3 tid : SV_DispatchThreadID )
+{
+}
+
+[numTHreaDs(4,4,2)]  // case insensitive
+void main_aux1(uint3 tid : SV_DispatchThreadID )
+{
+}
+
+[numthreads(1,4,8)]
+void main_aux2(uint3 tid : SV_DispatchThreadID );
+
diff --git a/gtests/Hlsl.FromFile.cpp b/gtests/Hlsl.FromFile.cpp
index 9a6072a3..66fba2e7 100644
--- a/gtests/Hlsl.FromFile.cpp
+++ b/gtests/Hlsl.FromFile.cpp
@@ -151,6 +151,7 @@ INSTANTIATE_TEST_CASE_P(
         {"hlsl.multiReturn.frag", "main"},
         {"hlsl.matrixindex.frag", "main"},
         {"hlsl.numericsuffixes.frag", "main"},
+        {"hlsl.numthreads.comp", "main_aux1"},
         {"hlsl.overload.frag", "PixelShaderFunction"},
         {"hlsl.pp.line.frag", "main"},
         {"hlsl.precise.frag", "main"},
diff --git a/hlsl/CMakeLists.txt b/hlsl/CMakeLists.txt
index c7537e27..ec5f1a56 100755
--- a/hlsl/CMakeLists.txt
+++ b/hlsl/CMakeLists.txt
@@ -1,4 +1,5 @@
 set(SOURCES
+    hlslAttributes.cpp
     hlslParseHelper.cpp
     hlslScanContext.cpp
     hlslOpMap.cpp
@@ -6,7 +7,8 @@ set(SOURCES
     hlslGrammar.cpp
     hlslParseables.cpp)
 
-set(HEADERS
+  set(HEADERS
+    hlslAttributes.h
     hlslParseHelper.h
     hlslTokens.h
     hlslScanContext.h
diff --git a/hlsl/hlslAttributes.cpp b/hlsl/hlslAttributes.cpp
new file mode 100644
index 00000000..957f282e
--- /dev/null
+++ b/hlsl/hlslAttributes.cpp
@@ -0,0 +1,108 @@
+//
+//Copyright (C) 2016 LunarG, Inc.
+//
+//All rights reserved.
+//
+//Redistribution and use in source and binary forms, with or without
+//modification, are permitted provided that the following conditions
+//are met:
+//
+//    Redistributions of source code must retain the above copyright
+//    notice, this list of conditions and the following disclaimer.
+//
+//    Redistributions in binary form must reproduce the above
+//    copyright notice, this list of conditions and the following
+//    disclaimer in the documentation and/or other materials provided
+//    with the distribution.
+//
+//    Neither the name of Google, Inc., nor the names of its
+//    contributors may be used to endorse or promote products derived
+//    from this software without specific prior written permission.
+//
+//THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+//"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+//LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
+//FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
+//COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
+//INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+//BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+//LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+//CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
+//LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
+//ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+//POSSIBILITY OF SUCH DAMAGE.
+//
+
+#include "hlslAttributes.h"
+#include <cstdlib>
+#include <cctype>
+
+namespace glslang {
+    // Map the given string to an attribute enum from TAttributeType,
+    // or EatNone if invalid.
+    TAttributeType TAttributeMap::attributeFromName(const TString& name)
+    {
+        // These are case insensitive.
+        TString lowername(name);
+        std::transform(lowername.begin(), lowername.end(), lowername.begin(), ::tolower);
+
+        if (lowername == "allow_uav_condition")
+            return EatAllow_uav_condition;
+        else if (lowername == "branch")
+            return EatBranch;
+        else if (lowername == "call")
+            return EatCall;
+        else if (lowername == "domain")
+            return EatDomain;
+        else if (lowername == "earlydepthstencil")
+            return EatEarlydepthstencil;
+        else if (lowername == "fastopt")
+            return EatFastopt;
+        else if (lowername == "flatten")
+            return EatFlatten;
+        else if (lowername == "forcecase")
+            return EatForcecase;
+        else if (lowername == "instance")
+            return EatInstance;
+        else if (lowername == "maxtessfactor")
+            return EatMaxtessfactor;
+        else if (lowername == "numthreads")
+            return EatNumthreads;
+        else if (lowername == "outputcontrolpoints")
+            return EatOutputcontrolpoints;
+        else if (lowername == "outputtopology")
+            return EatOutputtopology;
+        else if (lowername == "partitioning")
+            return EatPartitioning;
+        else if (lowername == "patchconstantfunc")
+            return EatPatchconstantfunc;
+        else if (lowername == "unroll")
+            return EatUnroll;
+        else
+            return EatNone;
+    }
+
+    // Look up entry, inserting if it's not there, and if name is a valid attribute name
+    // as known by attributeFromName.
+    TAttributeType TAttributeMap::setAttribute(const TString* name, TIntermAggregate* value)
+    {
+        if (name == nullptr)
+            return EatNone;
+
+        const TAttributeType attr = attributeFromName(*name);
+
+        if (attr != EatNone)
+            attributes[attr] = value;
+ 
+        return attr;
+    }
+
+    // Look up entry (const version), and return aggregate node.  This cannot change the map.
+    const TIntermAggregate* TAttributeMap::operator[](TAttributeType attr) const
+    {
+        const auto entry = attributes.find(attr);
+
+        return (entry == attributes.end()) ? nullptr : entry->second;
+    }
+    
+} // end namespace glslang
diff --git a/hlsl/hlslAttributes.h b/hlsl/hlslAttributes.h
new file mode 100644
index 00000000..da5ee5ef
--- /dev/null
+++ b/hlsl/hlslAttributes.h
@@ -0,0 +1,96 @@
+//
+//Copyright (C) 2016 LunarG, Inc.
+//
+//All rights reserved.
+//
+//Redistribution and use in source and binary forms, with or without
+//modification, are permitted provided that the following conditions
+//are met:
+//
+//    Redistributions of source code must retain the above copyright
+//    notice, this list of conditions and the following disclaimer.
+//
+//    Redistributions in binary form must reproduce the above
+//    copyright notice, this list of conditions and the following
+//    disclaimer in the documentation and/or other materials provided
+//    with the distribution.
+//
+//    Neither the name of Google, Inc., nor the names of its
+//    contributors may be used to endorse or promote products derived
+//    from this software without specific prior written permission.
+//
+//THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+//"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+//LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
+//FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
+//COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
+//INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+//BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+//LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+//CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
+//LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
+//ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+//POSSIBILITY OF SUCH DAMAGE.
+//
+
+#ifndef HLSLATTRIBUTES_H_
+#define HLSLATTRIBUTES_H_
+
+#include <unordered_map>
+#include <functional>
+#include "hlslScanContext.h"
+#include "../glslang/Include/Common.h"
+
+namespace glslang {
+    enum TAttributeType {
+        EatNone,
+        EatAllow_uav_condition,
+        EatBranch,
+        EatCall,
+        EatDomain,
+        EatEarlydepthstencil,
+        EatFastopt,
+        EatFlatten,
+        EatForcecase,
+        EatInstance,
+        EatMaxtessfactor,
+        EatNumthreads,
+        EatOutputcontrolpoints,
+        EatOutputtopology,
+        EatPartitioning,
+        EatPatchconstantfunc,
+        EatUnroll,
+    };
+}
+
+namespace std {
+    // Allow use of TAttributeType enum in hash_map without calling code having to cast.
+    template <> struct hash<glslang::TAttributeType> {
+        std::size_t operator()(glslang::TAttributeType attr) const {
+            return std::hash<int>()(int(attr));
+        }
+    };
+} // end namespace std
+
+namespace glslang {
+    class TIntermAggregate;
+
+    class TAttributeMap {
+    public:
+        // Search for and potentially add the attribute into the map.  Return the
+        // attribute type enum for it, if found, else EatNone.
+        TAttributeType setAttribute(const TString* name, TIntermAggregate* value);
+
+        // Const lookup: search for (but do not modify) the attribute in the map.
+        const TIntermAggregate* operator[](TAttributeType) const;
+
+    protected:
+        // Find an attribute enum given its name.
+        static TAttributeType attributeFromName(const TString&);
+
+        std::unordered_map<TAttributeType, TIntermAggregate*> attributes;
+    };
+} // end namespace glslang
+
+
+#endif // HLSLATTRIBUTES_H_
diff --git a/hlsl/hlslGrammar.cpp b/hlsl/hlslGrammar.cpp
index 4b2d6a23..e16c4b50 100755
--- a/hlsl/hlslGrammar.cpp
+++ b/hlsl/hlslGrammar.cpp
@@ -53,6 +53,7 @@
 
 #include "hlslTokens.h"
 #include "hlslGrammar.h"
+#include "hlslAttributes.h"
 
 namespace glslang {
 
@@ -268,6 +269,10 @@ bool HlslGrammar::acceptDeclaration(TIntermNode*& node)
     node = nullptr;
     bool list = false;
 
+    // attributes
+    TAttributeMap attributes;
+    acceptAttributes(attributes);
+
     // typedef
     bool typedefDecl = acceptTokenClass(EHTokTypedef);
 
@@ -302,7 +307,7 @@ bool HlslGrammar::acceptDeclaration(TIntermNode*& node)
                     parseContext.error(idToken.loc, "function body can't be in a declarator list", "{", "");
                 if (typedefDecl)
                     parseContext.error(idToken.loc, "function body can't be in a typedef", "{", "");
-                return acceptFunctionDefinition(function, node);
+                return acceptFunctionDefinition(function, node, attributes);
             } else {
                 if (typedefDecl)
                     parseContext.error(idToken.loc, "function typedefs not implemented", "{", "");
@@ -1687,13 +1692,13 @@ bool HlslGrammar::acceptParameterDeclaration(TFunction& function)
 
 // Do the work to create the function definition in addition to
 // parsing the body (compound_statement).
-bool HlslGrammar::acceptFunctionDefinition(TFunction& function, TIntermNode*& node)
+bool HlslGrammar::acceptFunctionDefinition(TFunction& function, TIntermNode*& node, const TAttributeMap& attributes)
 {
     TFunction& functionDeclarator = parseContext.handleFunctionDeclarator(token.loc, function, false /* not prototype */);
     TSourceLoc loc = token.loc;
 
     // This does a pushScope()
-    node = parseContext.handleFunctionDefinition(loc, functionDeclarator);
+    node = parseContext.handleFunctionDefinition(loc, functionDeclarator, attributes);
 
     // compound_statement
     TIntermNode* functionBody = nullptr;
@@ -2344,7 +2349,8 @@ bool HlslGrammar::acceptStatement(TIntermNode*& statement)
     statement = nullptr;
 
     // attributes
-    acceptAttributes();
+    TAttributeMap attributes;
+    acceptAttributes(attributes);
 
     // attributed_statement
     switch (peek()) {
@@ -2417,42 +2423,68 @@ bool HlslGrammar::acceptStatement(TIntermNode*& statement)
 //      | FLATTEN
 //      | FORCECASE
 //      | CALL
-//
-void HlslGrammar::acceptAttributes()
+//      | DOMAIN
+//      | EARLYDEPTHSTENCIL
+//      | INSTANCE
+//      | MAXTESSFACTOR
+//      | OUTPUTCONTROLPOINTS
+//      | OUTPUTTOPOLOGY
+//      | PARTITIONING
+//      | PATCHCONSTANTFUNC
+//      | NUMTHREADS LEFT_PAREN x_size, y_size,z z_size RIGHT_PAREN
+//
+void HlslGrammar::acceptAttributes(TAttributeMap& attributes)
 {
-    // For now, accept the [ XXX(X) ] syntax, but drop.
+    // For now, accept the [ XXX(X) ] syntax, but drop all but
+    // numthreads, which is used to set the CS local size.
     // TODO: subset to correct set?  Pass on?
     do {
+        HlslToken idToken;
+
         // LEFT_BRACKET?
         if (! acceptTokenClass(EHTokLeftBracket))
             return;
 
         // attribute
-        if (peekTokenClass(EHTokIdentifier)) {
-            // 'token.string' is the attribute
-            advanceToken();
+        if (acceptIdentifier(idToken)) {
+            // 'idToken.string' is the attribute
         } else if (! peekTokenClass(EHTokRightBracket)) {
             expected("identifier");
             advanceToken();
         }
 
-        // (x)
+        TIntermAggregate* literals = nullptr;
+
+        // (x, ...)
         if (acceptTokenClass(EHTokLeftParen)) {
+            literals = new TIntermAggregate;
+
             TIntermTyped* node;
-            if (! acceptLiteral(node))
-                expected("literal");
-            // 'node' has the literal in it
+            bool expectingLiteral = false;
+            
+            while (acceptLiteral(node)) {
+                expectingLiteral = false;
+                literals->getSequence().push_back(node);
+                if (acceptTokenClass(EHTokComma))
+                    expectingLiteral = true;
+            }
+
+            // 'literals' is an aggregate with the literals in it
             if (! acceptTokenClass(EHTokRightParen))
                 expected(")");
+            if (expectingLiteral || literals->getSequence().empty())
+                expected("literal");
         }
 
         // RIGHT_BRACKET
-        if (acceptTokenClass(EHTokRightBracket))
-            continue;
-
-        expected("]");
-        return;
+        if (!acceptTokenClass(EHTokRightBracket)) {
+            expected("]");
+            return;
+        }
 
+        // Add any values we found into the attribute map.  This accepts
+        // (and ignores) values not mapping to a known TAttributeType;
+        attributes.setAttribute(idToken.string, literals);
     } while (true);
 }
 
diff --git a/hlsl/hlslGrammar.h b/hlsl/hlslGrammar.h
index 992eb5e9..c3e42f67 100755
--- a/hlsl/hlslGrammar.h
+++ b/hlsl/hlslGrammar.h
@@ -43,6 +43,8 @@
 
 namespace glslang {
 
+    class TAttributeMap; // forward declare
+    
     // Should just be the grammar aspect of HLSL.
     // Described in more detail in hlslGrammar.cpp.
 
@@ -80,7 +82,7 @@ namespace glslang {
         bool acceptStructDeclarationList(TTypeList*&);
         bool acceptFunctionParameters(TFunction&);
         bool acceptParameterDeclaration(TFunction&);
-        bool acceptFunctionDefinition(TFunction&, TIntermNode*&);
+        bool acceptFunctionDefinition(TFunction&, TIntermNode*&, const TAttributeMap&);
         bool acceptParenExpression(TIntermTyped*&);
         bool acceptExpression(TIntermTyped*&);
         bool acceptInitializer(TIntermTyped*&);
@@ -98,7 +100,7 @@ namespace glslang {
         bool acceptScopedStatement(TIntermNode*&);
         bool acceptScopedCompoundStatement(TIntermNode*&);
         bool acceptNestedStatement(TIntermNode*&);
-        void acceptAttributes();
+        void acceptAttributes(TAttributeMap&);
         bool acceptSelectionStatement(TIntermNode*&);
         bool acceptSwitchStatement(TIntermNode*&);
         bool acceptIterationStatement(TIntermNode*&);
diff --git a/hlsl/hlslParseHelper.cpp b/hlsl/hlslParseHelper.cpp
index a9e3a107..cf97ec7e 100755
--- a/hlsl/hlslParseHelper.cpp
+++ b/hlsl/hlslParseHelper.cpp
@@ -37,6 +37,7 @@
 #include "hlslParseHelper.h"
 #include "hlslScanContext.h"
 #include "hlslGrammar.h"
+#include "hlslAttributes.h"
 
 #include "../glslang/MachineIndependent/Scan.h"
 #include "../glslang/MachineIndependent/preprocessor/PpContext.h"
@@ -1045,7 +1046,8 @@ TFunction& HlslParseContext::handleFunctionDeclarator(const TSourceLoc& loc, TFu
 // Handle seeing the function prototype in front of a function definition in the grammar.  
 // The body is handled after this function returns.
 //
-TIntermAggregate* HlslParseContext::handleFunctionDefinition(const TSourceLoc& loc, TFunction& function)
+TIntermAggregate* HlslParseContext::handleFunctionDefinition(const TSourceLoc& loc, TFunction& function, 
+                                                             const TAttributeMap& attributes)
 {
     currentCaller = function.getMangledName();
     TSymbol* symbol = symbolTable.find(function.getMangledName());
@@ -1134,6 +1136,15 @@ TIntermAggregate* HlslParseContext::handleFunctionDefinition(const TSourceLoc& l
     controlFlowNestingLevel = 0;
     postMainReturn = false;
 
+    // Handle function attributes
+    const TIntermAggregate* numThreadliterals = attributes[EatNumthreads];
+    if (numThreadliterals != nullptr && inEntryPoint) {
+        const TIntermSequence& sequence = numThreadliterals->getSequence();
+ 
+        for (int lid = 0; lid < int(sequence.size()); ++lid)
+            intermediate.setLocalSize(lid, sequence[lid]->getAsConstantUnion()->getConstArray()[0].getIConst());
+    }
+
     return paramNodes;
 }
 
diff --git a/hlsl/hlslParseHelper.h b/hlsl/hlslParseHelper.h
index 2eeba5ac..8410c344 100755
--- a/hlsl/hlslParseHelper.h
+++ b/hlsl/hlslParseHelper.h
@@ -41,6 +41,8 @@
 
 namespace glslang {
 
+class TAttributeMap; // forward declare
+
 class HlslParseContext : public TParseContextBase {
 public:
     HlslParseContext(TSymbolTable&, TIntermediate&, bool parsingBuiltins,
@@ -69,7 +71,7 @@ public:
     TIntermTyped* handleDotDereference(const TSourceLoc&, TIntermTyped* base, const TString& field);
     void assignLocations(TVariable& variable);
     TFunction& handleFunctionDeclarator(const TSourceLoc&, TFunction& function, bool prototype);
-    TIntermAggregate* handleFunctionDefinition(const TSourceLoc&, TFunction&);
+    TIntermAggregate* handleFunctionDefinition(const TSourceLoc&, TFunction&, const TAttributeMap&);
     void handleFunctionBody(const TSourceLoc&, TFunction&, TIntermNode* functionBody, TIntermNode*& node);
     void remapEntryPointIO(TFunction& function);
     void remapNonEntryPointIO(TFunction& function);