Fix Metal type layout reflection for nested parameter blocks. (#6042)

2 files changed, 33 insertions, 3 deletions

diff --git a/source/slang/slang-type-layout.cpp b/source/slang/slang-type-layout.cpp
index d7b303535..74528c61a 100644
--- a/source/slang/slang-type-layout.cpp
+++ b/source/slang/slang-type-layout.cpp
@@ -845,6 +845,8 @@ struct GLSLObjectLayoutRulesImpl : ObjectLayoutRulesImpl
         {
         case ShaderParameterKind::SubpassInput:
             return SimpleLayoutInfo(LayoutResourceKind::InputAttachmentIndex, slotCount);
+        case ShaderParameterKind::ParameterBlock:
+            return SimpleLayoutInfo(LayoutResourceKind::SubElementRegisterSpace, 1);
         default:
             break;
         }
@@ -891,7 +893,8 @@ struct HLSLObjectLayoutRulesImpl : ObjectLayoutRulesImpl
         {
         case ShaderParameterKind::ConstantBuffer:
             return SimpleLayoutInfo(LayoutResourceKind::ConstantBuffer, 1);
-
+        case ShaderParameterKind::ParameterBlock:
+            return SimpleLayoutInfo(LayoutResourceKind::SubElementRegisterSpace, 1);
         case ShaderParameterKind::TextureUniformBuffer:
         case ShaderParameterKind::StructuredBuffer:
         case ShaderParameterKind::RawBuffer:
@@ -1180,6 +1183,7 @@ struct CPUObjectLayoutRulesImpl : ObjectLayoutRulesImpl
         switch (kind)
         {
         case ShaderParameterKind::ConstantBuffer:
+        case ShaderParameterKind::ParameterBlock:
             // It's a pointer to the actual uniform data
             return SimpleLayoutInfo(
                 LayoutResourceKind::Uniform,
@@ -1262,6 +1266,7 @@ struct CUDAObjectLayoutRulesImpl : CPUObjectLayoutRulesImpl
         switch (kind)
         {
         case ShaderParameterKind::ConstantBuffer:
+        case ShaderParameterKind::ParameterBlock:
             {
                 // It's a pointer to the actual uniform data
                 return SimpleLayoutInfo(
@@ -1856,6 +1861,7 @@ struct MetalObjectLayoutRulesImpl : ObjectLayoutRulesImpl
         switch (kind)
         {
         case ShaderParameterKind::ConstantBuffer:
+        case ShaderParameterKind::ParameterBlock:
         case ShaderParameterKind::StructuredBuffer:
         case ShaderParameterKind::MutableStructuredBuffer:
         case ShaderParameterKind::RawBuffer:
@@ -1917,6 +1923,7 @@ struct MetalArgumentBufferElementLayoutRulesImpl : ObjectLayoutRulesImpl, Defaul
         switch (kind)
         {
         case ShaderParameterKind::ConstantBuffer:
+        case ShaderParameterKind::ParameterBlock:
         case ShaderParameterKind::StructuredBuffer:
         case ShaderParameterKind::MutableStructuredBuffer:
         case ShaderParameterKind::RawBuffer:
@@ -2341,6 +2348,10 @@ static SimpleLayoutInfo _getParameterGroupLayoutInfo(
     }
     else if (as<ParameterBlockType>(type))
     {
+        auto info =
+            rules->GetObjectLayout(ShaderParameterKind::ParameterBlock, context.objectLayoutOptions)
+                .getSimple();
+
         // Note: we default to consuming zero register spces here, because
         // a parameter block might not contain anything (or all it contains
         // is other blocks), and so it won't get a space allocated.
@@ -2351,7 +2362,9 @@ static SimpleLayoutInfo _getParameterGroupLayoutInfo(
         //
         // TODO: wouldn't it be any different to just allocate this
         // as an empty `SimpleLayoutInfo` of any other kind?
-        return SimpleLayoutInfo(LayoutResourceKind::SubElementRegisterSpace, 0);
+        if (info.kind == LayoutResourceKind::SubElementRegisterSpace)
+            info.size = 0;
+        return info;
     }
 
     // TODO: the vertex-input and fragment-output cases should
@@ -4049,6 +4062,21 @@ RefPtr<VarLayout> StructTypeLayoutBuilder::addField(
     RefPtr<TypeLayout> fieldTypeLayout = fieldResult.layout;
     UniformLayoutInfo fieldInfo = fieldResult.info.getUniformLayout();
 
+    if (fieldTypeLayout->resourceInfos.getCount() == 0)
+    {
+        if (auto paramGroupTypeLayout = as<ParameterGroupTypeLayout>(fieldTypeLayout))
+        {
+            // If field type layout is a parameter block and it has a size that is not just a space,
+            // we need to count for it in the struct layout.
+            auto containerTypeLayout = paramGroupTypeLayout->containerVarLayout->getTypeLayout();
+            if (containerTypeLayout->FindResourceInfo(
+                    LayoutResourceKind::SubElementRegisterSpace) == nullptr)
+            {
+                fieldTypeLayout = containerTypeLayout;
+            }
+        }
+    }
+
     // Note: we don't add any zero-size fields
     // when computing structure layout, just
     // to avoid having a resource type impact
@@ -4079,7 +4107,7 @@ RefPtr<VarLayout> StructTypeLayoutBuilder::addField(
     // for each field of the structure.
     RefPtr<VarLayout> fieldLayout = new VarLayout();
     fieldLayout->varDecl = field;
-    fieldLayout->typeLayout = fieldTypeLayout;
+    fieldLayout->typeLayout = fieldResult.layout;
     m_typeLayout->fields.add(fieldLayout);
 
     if (field)
diff --git a/source/slang/slang-type-layout.h b/source/slang/slang-type-layout.h
index 2cdb4b386..91132c76a 100644
--- a/source/slang/slang-type-layout.h
+++ b/source/slang/slang-type-layout.h
@@ -975,6 +975,8 @@ enum class ShaderParameterKind
     AtomicUint,
 
     SubpassInput,
+
+    ParameterBlock,
 };
 
 struct SimpleLayoutRulesImpl