Fix push constant on global variables. (#3034)

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

5 files changed, 51 insertions, 12 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)
diff --git a/tests/bugs/vk-shift-uniform-issue.slang b/tests/bugs/vk-shift-uniform-issue.slang
index fa49c3561..bc5963fe0 100644
--- a/tests/bugs/vk-shift-uniform-issue.slang
+++ b/tests/bugs/vk-shift-uniform-issue.slang
@@ -1,31 +1,39 @@
 //TEST:SIMPLE(filecheck=CHECK):-target glsl -profile ps_4_0 -entry main -fvk-t-shift 10 all  -fvk-s-shift 100 all -fvk-u-shift 100 all -fvk-b-shift 1000 all
 
 // CHECK:layout(binding = 10)
-// CHECK:uniform texture2D texture0_0;
+// CHECK-NEXT:uniform texture2D texture0_0;
 
 // CHECK:layout(binding = 100)
-// CHECK:uniform sampler sampler0_0;
+// CHECK-NEXT:uniform sampler sampler0_0;
 
 // CHECK:layout(binding = 11, set = 2)
-// CHECK:uniform texture2D texture1_0;
+// CHECK-NEXT:uniform texture2D texture1_0;
 
 // CHECK:layout(binding = 101, set = 2)
-// CHECK:uniform sampler sampler1_0;
+// CHECK-NEXT:uniform sampler sampler1_0;
+
+// CHECK: layout(push_constant)
+// CHECK-NEXT: layout(std140) uniform
 
 // CHECK:layout(binding = 1004)
-// CHECK:layout(std140) uniform _S1
+// CHECK-NEXT:layout(std140) uniform
 
 // CHECK:layout(binding = 1003)
-// CHECK:layout(std140) uniform _S2
+// CHECK-NEXT:layout(std140) uniform
 
 // CHECK:layout(binding = 1002)
-// CHECK:layout(std140) uniform _S3
+// CHECK-NEXT:layout(std140) uniform
 
 // CHECK:layout(binding = 1001)
-// CHECK:layout(std140) uniform _S4
+// CHECK-NEXT:layout(std140) uniform
+
+// CHECK: struct GlobalParams
+// CHECK-NEXT: {
+// CHECK-NEXT: float g_value
+// CHECK-NEXT: }
 
 // CHECK:layout(binding = 1000)
-// CHECK:layout(std140) uniform _S5
+// CHECK-NEXT:layout(std140) uniform
 
 Texture2D texture0; 
 SamplerState sampler0;
@@ -59,10 +67,10 @@ struct StructA
 { 
     float a;
 };
- 
-[[ vk::push_constant ]] 
+
+[[ vk::push_constant ]]
 StructA pushConstantA;
-  
+
 struct PixelInput
 { 
     float4 t : TEXCOORD0;