From edf5e9f97015a09fa0f2bed58d6a04744992d23f Mon Sep 17 00:00:00 2001
From: Yong He <yonghe@outlook.com>
Date: Wed, 15 Jan 2025 22:47:13 -0800
Subject: Fix argument buffer tier2 layout computation. (#6101)

---
 source/slang/slang-type-layout.cpp | 75 +++++++++++++++++++++++++++++++++++++-
 1 file changed, 73 insertions(+), 2 deletions(-)

(limited to 'source/slang')

diff --git a/source/slang/slang-type-layout.cpp b/source/slang/slang-type-layout.cpp
index 20a18a5ac..a412bf5b2 100644
--- a/source/slang/slang-type-layout.cpp
+++ b/source/slang/slang-type-layout.cpp
@@ -661,7 +661,7 @@ struct MetalLayoutRulesImpl : public CPULayoutRulesImpl
         auto alignedElementCount = 1 << Math::Log2Ceil((uint32_t)elementCount);
 
         // Metal aligns vectors to 2/4 element boundaries.
-        size_t size = elementSize * elementCount;
+        size_t size = alignedElementCount * elementSize;
         size_t alignment = alignedElementCount * elementSize;
 
         SimpleLayoutInfo vectorInfo;
@@ -1147,6 +1147,14 @@ struct MetalLayoutRulesFamilyImpl : LayoutRulesFamilyImpl
     LayoutRulesImpl* getStructuredBufferRules(CompilerOptionSet& compilerOptions) override;
 };
 
+struct MetalArgumentBufferTier2LayoutRulesFamilyImpl : MetalLayoutRulesFamilyImpl
+{
+    virtual LayoutRulesImpl* getConstantBufferRules(
+        CompilerOptionSet& compilerOptions,
+        Type* containerType) override;
+    virtual LayoutRulesImpl* getParameterBlockRules(CompilerOptionSet& compilerOptions) override;
+};
+
 struct WGSLLayoutRulesFamilyImpl : LayoutRulesFamilyImpl
 {
     virtual LayoutRulesImpl* getAnyValueRules() override;
@@ -1175,6 +1183,7 @@ HLSLLayoutRulesFamilyImpl kHLSLLayoutRulesFamilyImpl;
 CPULayoutRulesFamilyImpl kCPULayoutRulesFamilyImpl;
 CUDALayoutRulesFamilyImpl kCUDALayoutRulesFamilyImpl;
 MetalLayoutRulesFamilyImpl kMetalLayoutRulesFamilyImpl;
+MetalArgumentBufferTier2LayoutRulesFamilyImpl kMetalArgumentBufferTier2LayoutRulesFamilyImpl;
 WGSLLayoutRulesFamilyImpl kWGSLLayoutRulesFamilyImpl;
 
 // CPU case
@@ -1969,8 +1978,44 @@ struct MetalArgumentBufferElementLayoutRulesImpl : ObjectLayoutRulesImpl, Defaul
     }
 };
 
+struct MetalTier2ObjectLayoutRulesImpl : ObjectLayoutRulesImpl
+{
+    virtual ObjectLayoutInfo GetObjectLayout(ShaderParameterKind kind, const Options& /* options */)
+        override
+    {
+        switch (kind)
+        {
+        case ShaderParameterKind::ConstantBuffer:
+        case ShaderParameterKind::ParameterBlock:
+        case ShaderParameterKind::StructuredBuffer:
+        case ShaderParameterKind::MutableStructuredBuffer:
+        case ShaderParameterKind::RawBuffer:
+        case ShaderParameterKind::Buffer:
+        case ShaderParameterKind::MutableRawBuffer:
+        case ShaderParameterKind::MutableBuffer:
+        case ShaderParameterKind::ShaderStorageBuffer:
+        case ShaderParameterKind::AccelerationStructure:
+            return SimpleLayoutInfo(LayoutResourceKind::Uniform, 8, 8);
+        case ShaderParameterKind::AppendConsumeStructuredBuffer:
+            return SimpleLayoutInfo(LayoutResourceKind::Uniform, 16, 8);
+        case ShaderParameterKind::MutableTexture:
+        case ShaderParameterKind::TextureUniformBuffer:
+        case ShaderParameterKind::Texture:
+        case ShaderParameterKind::SamplerState:
+            return SimpleLayoutInfo(LayoutResourceKind::Uniform, 8, 8);
+        case ShaderParameterKind::TextureSampler:
+        case ShaderParameterKind::MutableTextureSampler:
+            return SimpleLayoutInfo(LayoutResourceKind::Uniform, 16, 8);
+        default:
+            SLANG_UNEXPECTED("unhandled shader parameter kind");
+            UNREACHABLE_RETURN(SimpleLayoutInfo());
+        }
+    }
+};
+
 static MetalObjectLayoutRulesImpl kMetalObjectLayoutRulesImpl;
 static MetalArgumentBufferElementLayoutRulesImpl kMetalArgumentBufferElementLayoutRulesImpl;
+static MetalTier2ObjectLayoutRulesImpl kMetalTier2ObjectLayoutRulesImpl;
 static MetalLayoutRulesImpl kMetalLayoutRulesImpl;
 
 LayoutRulesImpl kMetalAnyValueLayoutRulesImpl_ = {
@@ -1991,6 +2036,18 @@ LayoutRulesImpl kMetalParameterBlockLayoutRulesImpl_ = {
     &kMetalArgumentBufferElementLayoutRulesImpl,
 };
 
+LayoutRulesImpl kMetalTier2ConstantBufferLayoutRulesImpl_ = {
+    &kMetalLayoutRulesFamilyImpl,
+    &kMetalLayoutRulesImpl,
+    &kMetalTier2ObjectLayoutRulesImpl,
+};
+
+LayoutRulesImpl kMetalTier2ParameterBlockLayoutRulesImpl_ = {
+    &kMetalLayoutRulesFamilyImpl,
+    &kMetalLayoutRulesImpl,
+    &kMetalTier2ObjectLayoutRulesImpl,
+};
+
 LayoutRulesImpl kMetalStructuredBufferLayoutRulesImpl_ = {
     &kMetalLayoutRulesFamilyImpl,
     &kMetalLayoutRulesImpl,
@@ -2079,6 +2136,20 @@ LayoutRulesImpl* MetalLayoutRulesFamilyImpl::getHitAttributesParameterRules()
     return nullptr;
 }
 
+LayoutRulesImpl* MetalArgumentBufferTier2LayoutRulesFamilyImpl::getConstantBufferRules(
+    CompilerOptionSet&,
+    Type*)
+{
+    return &kMetalTier2ConstantBufferLayoutRulesImpl_;
+}
+
+LayoutRulesImpl* MetalArgumentBufferTier2LayoutRulesFamilyImpl::getParameterBlockRules(
+    CompilerOptionSet&)
+{
+    return &kMetalTier2ParameterBlockLayoutRulesImpl_;
+}
+
+
 // WGSL Family
 
 LayoutRulesImpl kWGSLConstantBufferLayoutRulesImpl_ = {
@@ -2229,7 +2300,7 @@ TypeLayoutContext getInitialLayoutContextForTarget(
         rulesFamily = getDefaultLayoutRulesFamilyForTarget(targetReq);
         break;
     case slang::LayoutRules::MetalArgumentBufferTier2:
-        rulesFamily = &kCPULayoutRulesFamilyImpl;
+        rulesFamily = &kMetalArgumentBufferTier2LayoutRulesFamilyImpl;
         break;
     }
 
-- 
cgit v1.2.3