Support parameter block in metal shader objects. (#4671)

* Support parameter block in metal shader objects.

* Ingore parameter block tests on devices without tier2 argument buffer.

* Fix warning.

* Fix texture subscript test.

---------

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

6 files changed, 43 insertions, 14 deletions

diff --git a/source/slang/slang-compiler.h b/source/slang/slang-compiler.h
index 7fc43d778..4bae6c10d 100755
--- a/source/slang/slang-compiler.h
+++ b/source/slang/slang-compiler.h
@@ -1787,11 +1787,27 @@ namespace Slang
         CodeGenTarget getTarget() { return optionSet.getEnumOption<CodeGenTarget>(CompilerOptionName::Target); }
 
         // TypeLayouts created on the fly by reflection API
-        Dictionary<Type*, RefPtr<TypeLayout>> typeLayouts;
+        struct TypeLayoutKey
+        {
+            Type* type;
+            slang::LayoutRules rules;
+            HashCode getHashCode() const
+            {
+                Hasher hasher;
+                hasher.hashValue(type);
+                hasher.hashValue(rules);
+                return hasher.getResult();
+            }
+            bool operator==(TypeLayoutKey other) const
+            {
+                return type == other.type && rules == other.rules;
+            }
+        };
+        Dictionary<TypeLayoutKey, RefPtr<TypeLayout>> typeLayouts;
 
-        Dictionary<Type*, RefPtr<TypeLayout>>& getTypeLayouts() { return typeLayouts; }
+        Dictionary<TypeLayoutKey, RefPtr<TypeLayout>>& getTypeLayouts() { return typeLayouts; }
 
-        TypeLayout* getTypeLayout(Type* type);
+        TypeLayout* getTypeLayout(Type* type, slang::LayoutRules rules);
 
         CompilerOptionSet& getOptionSet() { return optionSet; }
 
diff --git a/source/slang/slang-parameter-binding.cpp b/source/slang/slang-parameter-binding.cpp
index 8b15126db..fd7fae430 100644
--- a/source/slang/slang-parameter-binding.cpp
+++ b/source/slang/slang-parameter-binding.cpp
@@ -3892,7 +3892,7 @@ RefPtr<ProgramLayout> generateParameterBindings(
     }
 
     // Try to find rules based on the selected code-generation target
-    auto layoutContext = getInitialLayoutContextForTarget(targetReq, programLayout);
+    auto layoutContext = getInitialLayoutContextForTarget(targetReq, programLayout, slang::LayoutRules::Default);
 
     // If there was no target, or there are no rules for the target,
     // then bail out here.
diff --git a/source/slang/slang-reflection-api.cpp b/source/slang/slang-reflection-api.cpp
index 9e08330e4..8dc889c97 100644
--- a/source/slang/slang-reflection-api.cpp
+++ b/source/slang/slang-reflection-api.cpp
@@ -810,13 +810,13 @@ SLANG_API SlangReflectionType * spReflection_FindTypeByName(SlangReflection * re
 SLANG_API SlangReflectionTypeLayout* spReflection_GetTypeLayout(
     SlangReflection* reflection,
     SlangReflectionType* inType,
-    SlangLayoutRules /*rules*/)
+    SlangLayoutRules rules)
 {
     auto context = convert(reflection);
     auto type = convert(inType);
     auto targetReq = context->getTargetReq();
 
-    auto typeLayout = targetReq->getTypeLayout(type);
+    auto typeLayout = targetReq->getTypeLayout(type, (slang::LayoutRules)rules);
     return convert(typeLayout);
 }
 
@@ -1875,6 +1875,7 @@ namespace Slang
                     case LayoutResourceKind::DescriptorTableSlot:
                     case LayoutResourceKind::Uniform:
                     case LayoutResourceKind::ConstantBuffer: // for metal
+                    case LayoutResourceKind::MetalArgumentBufferElement:
                         resInfo = info;
                         break;
                     }
diff --git a/source/slang/slang-type-layout.cpp b/source/slang/slang-type-layout.cpp
index 293dce7c6..eca0d0ec3 100644
--- a/source/slang/slang-type-layout.cpp
+++ b/source/slang/slang-type-layout.cpp
@@ -1811,11 +1811,21 @@ LayoutRulesFamilyImpl* getDefaultLayoutRulesFamilyForTarget(TargetRequest* targe
     }
 }
 
-TypeLayoutContext getInitialLayoutContextForTarget(TargetRequest* targetReq, ProgramLayout* programLayout)
+TypeLayoutContext getInitialLayoutContextForTarget(TargetRequest* targetReq, ProgramLayout* programLayout, slang::LayoutRules rules)
 {
     auto astBuilder = targetReq->getLinkage()->getASTBuilder();
 
-    LayoutRulesFamilyImpl* rulesFamily = getDefaultLayoutRulesFamilyForTarget(targetReq);
+    LayoutRulesFamilyImpl* rulesFamily;
+    switch (rules)
+    {
+    case slang::LayoutRules::Default:
+    default:
+        rulesFamily = getDefaultLayoutRulesFamilyForTarget(targetReq);
+        break;
+    case slang::LayoutRules::MetalArgumentBufferTier2:
+        rulesFamily = &kCPULayoutRulesFamilyImpl;
+        break;
+    }
 
     TypeLayoutContext context;
     context.astBuilder = astBuilder;
diff --git a/source/slang/slang-type-layout.h b/source/slang/slang-type-layout.h
index 37c0cd1e7..701708363 100644
--- a/source/slang/slang-type-layout.h
+++ b/source/slang/slang-type-layout.h
@@ -1288,7 +1288,8 @@ private:
 //
 TypeLayoutContext getInitialLayoutContextForTarget(
     TargetRequest*  targetRequest,
-    ProgramLayout*  programLayout);
+    ProgramLayout* programLayout,
+    slang::LayoutRules rules);
 
     /// Direction(s) of a varying shader parameter
 typedef unsigned int EntryPointParameterDirectionMask;
diff --git a/source/slang/slang.cpp b/source/slang/slang.cpp
index f440ce3cb..c6e460d02 100644
--- a/source/slang/slang.cpp
+++ b/source/slang/slang.cpp
@@ -1368,7 +1368,7 @@ SLANG_NO_THROW slang::TypeLayoutReflection* SLANG_MCALL Linkage::getTypeLayout(
     //
     SLANG_UNUSED(rules);
 
-    auto typeLayout = target->getTypeLayout(type);
+    auto typeLayout = target->getTypeLayout(type, rules);
 
     // TODO: We currently don't have a path for capturing
     // errors that occur during layout (e.g., types that
@@ -1827,7 +1827,7 @@ CapabilitySet TargetRequest::getTargetCaps()
 }
 
 
-TypeLayout* TargetRequest::getTypeLayout(Type* type)
+TypeLayout* TargetRequest::getTypeLayout(Type* type, slang::LayoutRules rules)
 {
     SLANG_AST_BUILDER_RAII(getLinkage()->getASTBuilder());
 
@@ -1841,13 +1841,14 @@ TypeLayout* TargetRequest::getTypeLayout(Type* type)
     // parameter instead (leaving the user to figure out how that
     // maps to the ordering via some API on the program layout).
     //
-    auto layoutContext = getInitialLayoutContextForTarget(this, nullptr);
+    auto layoutContext = getInitialLayoutContextForTarget(this, nullptr, rules);
 
     RefPtr<TypeLayout> result;
-    if (getTypeLayouts().tryGetValue(type, result))
+    auto key = TypeLayoutKey{ type, rules };
+    if (getTypeLayouts().tryGetValue(key, result))
         return result.Ptr();
     result = createTypeLayout(layoutContext, type);
-    getTypeLayouts()[type] = result;
+    getTypeLayouts()[key] = result;
     return result.Ptr();
 }