From a88055c6f5190ca62bb4aa853b4f0fa11546278f Mon Sep 17 00:00:00 2001
From: Yong He <yonghe@outlook.com>
Date: Thu, 5 Sep 2024 10:26:59 -0700
Subject: Respect matrix layout in uniform and in/out parameters for HLSL
 target. (#5013)

* Respect matrix layout in uniform and in/out parameters for HLSL target.

* Update test.

* Fix test.

* fix test.

* Fix metal layout calculation.

* Fix compile error.

* Fix compiler error.

---------

Co-authored-by: Yong He <yhe@nvidia.com>
---
 source/slang/slang-type-layout.cpp | 27 +++++++++++++++++++++++++--
 1 file changed, 25 insertions(+), 2 deletions(-)

(limited to 'source/slang/slang-type-layout.cpp')

diff --git a/source/slang/slang-type-layout.cpp b/source/slang/slang-type-layout.cpp
index 7d6d047d3..f654135a1 100644
--- a/source/slang/slang-type-layout.cpp
+++ b/source/slang/slang-type-layout.cpp
@@ -606,6 +606,28 @@ struct CUDALayoutRulesImpl : DefaultLayoutRulesImpl
     }
 };
 
+struct MetalLayoutRulesImpl : public CPULayoutRulesImpl
+{
+    SimpleLayoutInfo GetVectorLayout(BaseType elementType, SimpleLayoutInfo elementInfo, size_t elementCount) override
+    {
+        SLANG_UNUSED(elementType);
+
+        const auto elementSize = elementInfo.size.getFiniteValue();
+        auto alignedElementCount = 1 << Math::Log2Ceil((uint32_t)elementCount);
+
+        // Metal aligns vectors to 2/4 element boundaries.
+        size_t size = elementSize * elementCount;
+        size_t alignment = alignedElementCount * elementSize;
+
+        SimpleLayoutInfo vectorInfo;
+        vectorInfo.kind = elementInfo.kind;
+        vectorInfo.size = size;
+        vectorInfo.alignment = alignment;
+     
+        return vectorInfo;
+    }
+};
+
 struct HLSLStructuredBufferLayoutRulesImpl : DefaultLayoutRulesImpl
 {
     // HLSL structured buffers drop the restrictions added for constant buffers,
@@ -1696,6 +1718,7 @@ struct MetalArgumentBufferElementLayoutRulesImpl : ObjectLayoutRulesImpl, Defaul
 
 static MetalObjectLayoutRulesImpl kMetalObjectLayoutRulesImpl;
 static MetalArgumentBufferElementLayoutRulesImpl kMetalArgumentBufferElementLayoutRulesImpl;
+static MetalLayoutRulesImpl kMetalLayoutRulesImpl;
 
 LayoutRulesImpl kMetalAnyValueLayoutRulesImpl_ = {
     &kMetalLayoutRulesFamilyImpl,
@@ -1704,7 +1727,7 @@ LayoutRulesImpl kMetalAnyValueLayoutRulesImpl_ = {
 };
 
 LayoutRulesImpl kMetalConstantBufferLayoutRulesImpl_ = {
-    &kMetalLayoutRulesFamilyImpl, & kCPULayoutRulesImpl, &kMetalObjectLayoutRulesImpl,
+    &kMetalLayoutRulesFamilyImpl, & kMetalLayoutRulesImpl, &kMetalObjectLayoutRulesImpl,
 };
 
 LayoutRulesImpl kMetalParameterBlockLayoutRulesImpl_ = {
@@ -1712,7 +1735,7 @@ LayoutRulesImpl kMetalParameterBlockLayoutRulesImpl_ = {
 };
 
 LayoutRulesImpl kMetalStructuredBufferLayoutRulesImpl_ = {
-    &kMetalLayoutRulesFamilyImpl, &kCPULayoutRulesImpl, & kMetalObjectLayoutRulesImpl,
+    &kMetalLayoutRulesFamilyImpl, &kMetalLayoutRulesImpl, & kMetalObjectLayoutRulesImpl,
 };
 
 LayoutRulesImpl kMetalVaryingInputLayoutRulesImpl_ = {
-- 
cgit v1.2.3