Fix push constant on global variables. (#3034)

Co-authored-by: Yong He <yhe@nvidia.com>

4 files changed, 31 insertions, 0 deletions

diff --git a/source/slang/slang-ast-builder.cpp b/source/slang/slang-ast-builder.cpp
index 72cb647d5..64a7abd8c 100644
--- a/source/slang/slang-ast-builder.cpp
+++ b/source/slang/slang-ast-builder.cpp
@@ -304,6 +304,19 @@ ArrayExpressionType* ASTBuilder::getArrayType(Type* elementType, IntVal* element
     return result;
 }
 
+ConstantBufferType* ASTBuilder::getConstantBufferType(Type* elementType)
+{
+    auto result = getOrCreate<ConstantBufferType>(elementType);
+    if (!result->declRef.getDecl())
+    {
+        auto genericDecl = as<GenericDecl>(m_sharedASTBuilder->findMagicDecl("ConstantBuffer"));
+        auto typeDecl = genericDecl->inner;
+        auto substitutions = getOrCreateGenericSubstitution(nullptr, genericDecl, elementType);
+        result->declRef = getSpecializedDeclRef<Decl>(typeDecl, substitutions);
+    }
+    return result;
+}
+
 VectorExpressionType* ASTBuilder::getVectorType(
     Type*    elementType,
     IntVal*  elementCount)
diff --git a/source/slang/slang-ast-builder.h b/source/slang/slang-ast-builder.h
index 479ad6540..cf0975cdd 100644
--- a/source/slang/slang-ast-builder.h
+++ b/source/slang/slang-ast-builder.h
@@ -405,6 +405,8 @@ public:
 
     VectorExpressionType* getVectorType(Type* elementType, IntVal* elementCount);
 
+    ConstantBufferType* getConstantBufferType(Type* elementType);
+
     DifferentialPairType* getDifferentialPairType(
         Type* valueType,
         Witness* primalIsDifferentialWitness);
diff --git a/source/slang/slang-ast-type.h b/source/slang/slang-ast-type.h
index 53544edfa..b32d62404 100644
--- a/source/slang/slang-ast-type.h
+++ b/source/slang/slang-ast-type.h
@@ -396,6 +396,7 @@ class VaryingParameterGroupType : public ParameterGroupType
 class ConstantBufferType : public UniformParameterGroupType 
 {
     SLANG_AST_CLASS(ConstantBufferType)
+    ConstantBufferType(Type* elementType) { SLANG_UNUSED(elementType); }
 };
 
 
diff --git a/source/slang/slang-check-decl.cpp b/source/slang/slang-check-decl.cpp
index 2d009c28c..b1dd2d533 100644
--- a/source/slang/slang-check-decl.cpp
+++ b/source/slang/slang-check-decl.cpp
@@ -1300,6 +1300,21 @@ namespace Slang
                 addModifier(varDecl, m_astBuilder->getOrCreate<NoDiffModifier>());
             }
         }
+
+
+        if (as<NamespaceDeclBase>(varDecl->parentDecl))
+        {
+            // If this is a global variable with [vk::push_constant] attribute,
+            // we need to make sure to wrap it in a `ConstantBuffer`.
+            
+            if (!as<ConstantBufferType>(varDecl->type))
+            {
+                if (varDecl->findModifier<PushConstantAttribute>())
+                {
+                    varDecl->type.type = m_astBuilder->getConstantBufferType(varDecl->type);
+                }
+            }
+        }
     }
 
     void SemanticsDeclHeaderVisitor::visitStructDecl(StructDecl* structDecl)