Parameter layout and reflection for Metal bindings. (#4022)

11 files changed, 336 insertions, 101 deletions

diff --git a/source/slang/slang-compiler.h b/source/slang/slang-compiler.h
index 105db5789..f9fa24622 100755
--- a/source/slang/slang-compiler.h
+++ b/source/slang/slang-compiler.h
@@ -1688,6 +1688,9 @@ namespace Slang
     /// Are we generating code for a D3D API?
     bool isD3DTarget(TargetRequest* targetReq);
 
+    // Are we generating code for Metal?
+    bool isMetalTarget(TargetRequest* targetReq);
+
     /// Are we generating code for a Khronos API (OpenGL or Vulkan)?
     bool isKhronosTarget(TargetRequest* targetReq);
 
diff --git a/source/slang/slang-emit-metal.cpp b/source/slang/slang-emit-metal.cpp
index 459104ac9..3773b7c2a 100644
--- a/source/slang/slang-emit-metal.cpp
+++ b/source/slang/slang-emit-metal.cpp
@@ -37,44 +37,6 @@ void MetalSourceEmitter::_emitHLSLDecorationSingleInt(const char* name, IRFunc*
     m_writer->emit(")]]\n");
 }
 
-void MetalSourceEmitter::_emitHLSLRegisterSemantic(LayoutResourceKind kind, EmitVarChain* chain, IRInst* inst, char const* uniformSemanticSpelling)
-{
-    // Metal does not use explicit binding.
-    SLANG_UNUSED(kind);
-    SLANG_UNUSED(chain);
-    SLANG_UNUSED(inst);
-    SLANG_UNUSED(uniformSemanticSpelling);
-}
-
-void MetalSourceEmitter::_emitHLSLRegisterSemantics(EmitVarChain* chain, IRInst* inst, char const* uniformSemanticSpelling)
-{
-    // TODO: implement.
-    SLANG_UNUSED(chain);
-    SLANG_UNUSED(inst);
-    SLANG_UNUSED(uniformSemanticSpelling);
-}
-
-void MetalSourceEmitter::_emitHLSLRegisterSemantics(IRVarLayout* varLayout, IRInst* inst, char const* uniformSemanticSpelling)
-{
-    // TODO: implement.
-    SLANG_UNUSED(varLayout);
-    SLANG_UNUSED(inst);
-    SLANG_UNUSED(uniformSemanticSpelling);
-}
-
-void MetalSourceEmitter::_emitHLSLParameterGroupFieldLayoutSemantics(EmitVarChain* chain)
-{
-    // TODO: implement.
-    SLANG_UNUSED(chain);
-}
-
-void MetalSourceEmitter::_emitHLSLParameterGroupFieldLayoutSemantics(IRVarLayout* fieldLayout, EmitVarChain* inChain)
-{
-    // TODO: implement.
-    SLANG_UNUSED(fieldLayout);
-    SLANG_UNUSED(inChain);
-}
-
 void MetalSourceEmitter::_emitHLSLParameterGroup(IRGlobalParam* varDecl, IRUniformParameterGroupType* type)
 {
     // Metal does not allow shader parameters declared as global variables, so we shouldn't see this.
@@ -139,17 +101,34 @@ void MetalSourceEmitter::_emitHLSLTextureType(IRTextureTypeBase* texType)
     m_writer->emit(">");
 }
 
-void MetalSourceEmitter::_emitHLSLSubpassInputType(IRSubpassInputType* subpassType)
-{
-    SLANG_UNUSED(subpassType);
-}
-
-void MetalSourceEmitter::emitLayoutSemanticsImpl(IRInst* inst, char const* uniformSemanticSpelling)
+void MetalSourceEmitter::emitFuncParamLayoutImpl(IRInst* param)
 {
-    auto layout = getVarLayout(inst); 
-    if (layout)
+    auto layoutDecoration = param->findDecoration<IRLayoutDecoration>();
+    if (!layoutDecoration)
+        return;
+    auto layout = as<IRVarLayout>(layoutDecoration->getLayout());
+    if (!layout)
+        return;
+    for (auto rr : layout->getOffsetAttrs())
     {
-        _emitHLSLRegisterSemantics(layout, inst, uniformSemanticSpelling);
+        switch (rr->getResourceKind())
+        {
+        case LayoutResourceKind::MetalTexture:
+            m_writer->emit(" [[texture(");
+            m_writer->emit(rr->getOffset());
+            m_writer->emit(")]]");
+            break;
+        case LayoutResourceKind::MetalBuffer:
+            m_writer->emit(" [[buffer(");
+            m_writer->emit(rr->getOffset());
+            m_writer->emit(")]]");
+            break;
+        case LayoutResourceKind::SamplerState:
+            m_writer->emit(" [[sampler(");
+            m_writer->emit(rr->getOffset());
+            m_writer->emit(")]]");
+            break;
+        }
     }
 }
 
@@ -647,6 +626,7 @@ void MetalSourceEmitter::_emitStageAccessSemantic(IRStageAccessDecoration* decor
 void MetalSourceEmitter::emitSimpleFuncParamImpl(IRParam* param)
 {
     Super::emitSimpleFuncParamImpl(param);
+    emitFuncParamLayoutImpl(param);
 }
 
 static UnownedStringSlice _getInterpolationModifierText(IRInterpolationMode mode)
diff --git a/source/slang/slang-emit-metal.h b/source/slang/slang-emit-metal.h
index 6c3de04c4..38d4c3a2c 100644
--- a/source/slang/slang-emit-metal.h
+++ b/source/slang/slang-emit-metal.h
@@ -23,7 +23,6 @@ public:
 protected:
     RefPtr<MetalExtensionTracker> m_extensionTracker;
 
-    virtual void emitLayoutSemanticsImpl(IRInst* inst, char const* uniformSemanticSpelling) SLANG_OVERRIDE;
     virtual void emitParameterGroupImpl(IRGlobalParam* varDecl, IRUniformParameterGroupType* type) SLANG_OVERRIDE;
     virtual void emitEntryPointAttributesImpl(IRFunc* irFunc, IREntryPointDecoration* entryPointDecor) SLANG_OVERRIDE;
     
@@ -32,6 +31,7 @@ protected:
     virtual void emitRateQualifiersAndAddressSpaceImpl(IRRate* rate, IRIntegerValue addressSpace) SLANG_OVERRIDE;
     virtual void emitSemanticsImpl(IRInst* inst, bool allowOffsets) SLANG_OVERRIDE;
     virtual void emitSimpleFuncParamImpl(IRParam* param) SLANG_OVERRIDE;
+
     virtual void emitInterpolationModifiersImpl(IRInst* varInst, IRType* valueType, IRVarLayout* layout) SLANG_OVERRIDE;
     virtual void emitPackOffsetModifier(IRInst* varInst, IRType* valueType, IRPackOffsetDecoration* decoration) SLANG_OVERRIDE;
 
@@ -56,16 +56,7 @@ protected:
 
     virtual bool doesTargetSupportPtrTypes() SLANG_OVERRIDE { return true; }
 
-        // Emit a single `register` semantic, as appropriate for a given resource-type-specific layout info
-        // Keyword to use in the uniform case (`register` for globals, `packoffset` inside a `cbuffer`)
-    void _emitHLSLRegisterSemantic(LayoutResourceKind kind, EmitVarChain* chain, IRInst* inst, char const* uniformSemanticSpelling = "register");
-
-        // Emit all the `register` semantics that are appropriate for a particular variable layout
-    void _emitHLSLRegisterSemantics(EmitVarChain* chain, IRInst* inst, char const* uniformSemanticSpelling = "register");
-    void _emitHLSLRegisterSemantics(IRVarLayout* varLayout, IRInst* inst, char const* uniformSemanticSpelling = "register");
-
-    void _emitHLSLParameterGroupFieldLayoutSemantics(EmitVarChain* chain);
-    void _emitHLSLParameterGroupFieldLayoutSemantics(IRVarLayout* fieldLayout, EmitVarChain* inChain);
+    void emitFuncParamLayoutImpl(IRInst* param);
 
     void _emitHLSLParameterGroup(IRGlobalParam* varDecl, IRUniformParameterGroupType* type);
 
diff --git a/source/slang/slang-emit.cpp b/source/slang/slang-emit.cpp
index 77164dae3..5cfcde0c9 100644
--- a/source/slang/slang-emit.cpp
+++ b/source/slang/slang-emit.cpp
@@ -583,6 +583,7 @@ Result linkAndOptimizeIR(
         //  will have (more) legal shader code.
         //
         legalizeExistentialTypeLayout(
+            targetProgram,
             irModule,
             sink);
 
@@ -603,6 +604,7 @@ Result linkAndOptimizeIR(
         // then become multiple variables/parameters/arguments/etc.
         //
         legalizeResourceTypes(
+            targetProgram,
             irModule,
             sink);
 
@@ -617,6 +619,7 @@ Result linkAndOptimizeIR(
         // On CPU/CUDA targets, we simply elminate any empty types if
         // they are not part of public interface.
         legalizeEmptyTypes(
+            targetProgram,
             irModule,
             sink);
     }
@@ -876,6 +879,7 @@ Result linkAndOptimizeIR(
     case CodeGenTarget::GLSL:
     case CodeGenTarget::SPIRV:
     case CodeGenTarget::SPIRVAssembly:
+    case CodeGenTarget::Metal:
         moveGlobalVarInitializationToEntryPoints(irModule);
         break;
     case CodeGenTarget::CPPSource:
@@ -941,9 +945,12 @@ Result linkAndOptimizeIR(
 
     if (options.shouldLegalizeExistentialAndResourceTypes)
     {
-        // We need to lower any types used in a buffer resource (e.g. ContantBuffer or StructuredBuffer) into
-        // a simple storage type that has target independent layout based on the kind of buffer resource.
-        lowerBufferElementTypeToStorageType(targetProgram, irModule);
+        if (!isMetalTarget(targetRequest))
+        {
+            // We need to lower any types used in a buffer resource (e.g. ContantBuffer or StructuredBuffer) into
+            // a simple storage type that has target independent layout based on the kind of buffer resource.
+            lowerBufferElementTypeToStorageType(targetProgram, irModule);
+        }
     }
 
     // Rewrite functions that return arrays to return them via `out` parameter,
diff --git a/source/slang/slang-ir-legalize-types.cpp b/source/slang/slang-ir-legalize-types.cpp
index fb914f7c0..f36e02066 100644
--- a/source/slang/slang-ir-legalize-types.cpp
+++ b/source/slang/slang-ir-legalize-types.cpp
@@ -95,8 +95,11 @@ LegalVal LegalVal::wrappedBuffer(
 //
 
 IRTypeLegalizationContext::IRTypeLegalizationContext(
+    TargetProgram* target,
     IRModule* inModule)
 {
+    targetProgram = target;
+
     session = inModule->getSession();
     module = inModule;
 
@@ -3868,8 +3871,8 @@ static void legalizeTypes(
 //
 struct IRResourceTypeLegalizationContext : IRTypeLegalizationContext
 {
-    IRResourceTypeLegalizationContext(IRModule* module)
-        : IRTypeLegalizationContext(module)
+    IRResourceTypeLegalizationContext(TargetProgram* target, IRModule* module)
+        : IRTypeLegalizationContext(target, module)
     {}
 
     bool isSpecialType(IRType* type) override
@@ -3903,8 +3906,8 @@ struct IRResourceTypeLegalizationContext : IRTypeLegalizationContext
 //
 struct IRExistentialTypeLegalizationContext : IRTypeLegalizationContext
 {
-    IRExistentialTypeLegalizationContext(IRModule* module)
-        : IRTypeLegalizationContext(module)
+    IRExistentialTypeLegalizationContext(TargetProgram* target, IRModule* module)
+        : IRTypeLegalizationContext(target, module)
     {}
 
     bool isSpecialType(IRType* inType) override
@@ -3944,8 +3947,8 @@ struct IRExistentialTypeLegalizationContext : IRTypeLegalizationContext
 // a public function signature.
 struct IREmptyTypeLegalizationContext : IRTypeLegalizationContext
 {
-    IREmptyTypeLegalizationContext(IRModule* module)
-        : IRTypeLegalizationContext(module)
+    IREmptyTypeLegalizationContext(TargetProgram* target, IRModule* module)
+        : IRTypeLegalizationContext(target, module)
     {}
 
     bool isSpecialType(IRType*) override
@@ -3985,6 +3988,7 @@ struct IREmptyTypeLegalizationContext : IRTypeLegalizationContext
 // specialized context type to use to get the job done.
 
 void legalizeResourceTypes(
+    TargetProgram* target,
     IRModule*       module,
     DiagnosticSink* sink)
 {
@@ -3992,11 +3996,12 @@ void legalizeResourceTypes(
 
     SLANG_UNUSED(sink);
 
-    IRResourceTypeLegalizationContext context(module);
+    IRResourceTypeLegalizationContext context(target, module);
     legalizeTypes(&context);
 }
 
 void legalizeExistentialTypeLayout(
+    TargetProgram* target,
     IRModule*       module,
     DiagnosticSink* sink)
 {
@@ -4005,15 +4010,15 @@ void legalizeExistentialTypeLayout(
     SLANG_UNUSED(module);
     SLANG_UNUSED(sink);
 
-    IRExistentialTypeLegalizationContext context(module);
+    IRExistentialTypeLegalizationContext context(target, module);
     legalizeTypes(&context);
 }
 
-void legalizeEmptyTypes(IRModule* module, DiagnosticSink* sink)
+void legalizeEmptyTypes(TargetProgram* target, IRModule* module, DiagnosticSink* sink)
 {
     SLANG_UNUSED(sink);
 
-    IREmptyTypeLegalizationContext context(module);
+    IREmptyTypeLegalizationContext context(target, module);
     legalizeTypes(&context);
 }
 
diff --git a/source/slang/slang-ir-wrap-global-context.cpp b/source/slang/slang-ir-wrap-global-context.cpp
index 01a343a02..32bf7995c 100644
--- a/source/slang/slang-ir-wrap-global-context.cpp
+++ b/source/slang/slang-ir-wrap-global-context.cpp
@@ -65,8 +65,7 @@ namespace Slang
                 {
                     auto newParam = builder.createParam(globalParam.first->getFullType());
                     newParam->insertBefore(paramInsertPoint);
-                    if (auto name = findNameHint(globalParam.first))
-                        builder.addNameHintDecoration(newParam, name);
+                    globalParam.first->transferDecorationsTo(newParam);
                     newParams.add({newParam, globalParam.second});
                 }
 
diff --git a/source/slang/slang-legalize-types.cpp b/source/slang/slang-legalize-types.cpp
index fc0947c62..393b34e68 100644
--- a/source/slang/slang-legalize-types.cpp
+++ b/source/slang/slang-legalize-types.cpp
@@ -1156,23 +1156,34 @@ LegalType legalizeTypeImpl(
         auto originalElementType = uniformBufferType->getElementType();
 
         // Legalize the element type to see what we are working with.
-        auto legalElementType = legalizeType(context,
-            originalElementType);
-
-        // As a bit of a corner case, if the user requested something
-        // like `ConstantBuffer<Texture2D>` the element type would
-        // legalize to a "simple" type, and that would be interpreted
-        // as an *ordinary* type, but we really need to notice the
-        // case when the element type is simple, but *special*.
-        //
-        if( context->isSpecialType(originalElementType) )
+        LegalType legalElementType;
+        
+        if (isMetalTarget(context->targetProgram->getTargetReq()) &&
+            as<IRParameterBlockType>(uniformBufferType))
+        {
+            // On Metal, we do not need to legalize the element type of
+            // a parameter block because we can translate it directly into
+            // an argument buffer.
+            legalElementType = LegalType::simple(originalElementType);
+        }
+        else
         {
-            // Anything that has a special element type needs to
-            // be handled by the pass-specific logic in the context.
+            legalElementType = legalizeType(context, originalElementType);
+            // As a bit of a corner case, if the user requested something
+            // like `ConstantBuffer<Texture2D>` the element type would
+            // legalize to a "simple" type, and that would be interpreted
+            // as an *ordinary* type, but we really need to notice the
+            // case when the element type is simple, but *special*.
             //
-            return context->createLegalUniformBufferType(
-                uniformBufferType->getOp(),
-                legalElementType);
+            if( context->isSpecialType(originalElementType) )
+            {
+                // Anything that has a special element type needs to
+                // be handled by the pass-specific logic in the context.
+                //
+                return context->createLegalUniformBufferType(
+                    uniformBufferType->getOp(),
+                    legalElementType);
+            }
         }
 
         // Note that even when legalElementType.flavor == Simple
diff --git a/source/slang/slang-legalize-types.h b/source/slang/slang-legalize-types.h
index 4909433e3..f954e2b25 100644
--- a/source/slang/slang-legalize-types.h
+++ b/source/slang/slang-legalize-types.h
@@ -602,10 +602,11 @@ struct IRTypeLegalizationContext
     Session*    session;
     IRModule*   module;
     IRBuilder*  builder;
-
+    TargetProgram* targetProgram;
     IRBuilder builderStorage;
 
     IRTypeLegalizationContext(
+        TargetProgram* target,
         IRModule* inModule);
 
     // When inserting new globals, put them before this one.
@@ -710,14 +711,17 @@ LegalType createLegalUniformBufferTypeForExistentials(
 
 
 void legalizeExistentialTypeLayout(
+    TargetProgram* target,
     IRModule*       module,
     DiagnosticSink* sink);
 
 void legalizeResourceTypes(
+    TargetProgram* target,
     IRModule*       module,
     DiagnosticSink* sink);
 
 void legalizeEmptyTypes(
+    TargetProgram* target,
     IRModule* module,
     DiagnosticSink* sink);
 
diff --git a/source/slang/slang-parameter-binding.cpp b/source/slang/slang-parameter-binding.cpp
index f702efb94..56c85e9cb 100644
--- a/source/slang/slang-parameter-binding.cpp
+++ b/source/slang/slang-parameter-binding.cpp
@@ -973,7 +973,7 @@ static void addExplicitParameterBindings_HLSL(
     RefPtr<VarLayout>           varLayout)
 {
     // We only want to apply D3D `register` modifiers when compiling for
-    // D3D targets.
+    // D3D and Metal targets.
     //
     // TODO: Nominally, the `register` keyword allows for a shader
     // profile to be specified, so that a given binding only
@@ -989,7 +989,7 @@ static void addExplicitParameterBindings_HLSL(
     //
     // For now we do the filtering on target in a very direct fashion:
     //
-    if(!isD3DTarget(context->getTargetRequest()))
+    if(!isD3DTarget(context->getTargetRequest()) && !isMetalTarget(context->getTargetRequest()))
         return;
 
     auto typeLayout = varLayout->typeLayout;
diff --git a/source/slang/slang-type-layout.cpp b/source/slang/slang-type-layout.cpp
index ebb789876..d81cc92cd 100644
--- a/source/slang/slang-type-layout.cpp
+++ b/source/slang/slang-type-layout.cpp
@@ -673,6 +673,13 @@ struct HLSLVaryingLayoutRulesImpl : DefaultVaryingLayoutRulesImpl
     {}
 };
 
+struct MetalVaryingLayoutRulesImpl : DefaultVaryingLayoutRulesImpl
+{
+    MetalVaryingLayoutRulesImpl(LayoutResourceKind kind)
+        : DefaultVaryingLayoutRulesImpl(kind)
+    {}
+};
+
 //
 
 struct GLSLSpecializationConstantLayoutRulesImpl : DefaultLayoutRulesImpl
@@ -744,7 +751,7 @@ struct GLSLObjectLayoutRulesImpl : ObjectLayoutRulesImpl
     virtual ObjectLayoutInfo GetObjectLayout(ShaderParameterKind kind, const Options& options) override
     {
         int slotCount = 1;
-
+        
         // In Vulkan GLSL, pretty much every object is just a descriptor-table slot.
         // Except for AppendConsumeStructuredBuffer, which takes two slots.
         // This, however, is added in 'createStructuredBufferWithCounterTypeLayout'
@@ -902,6 +909,9 @@ CUDARayTracingLayoutRulesImpl kCUDARayPayloadParameterLayoutRulesImpl(LayoutReso
 //CUDARayTracingLayoutRulesImpl kCUDACallablePayloadParameterLayoutRulesImpl(LayoutResourceKind::CallablePayload);
 CUDARayTracingLayoutRulesImpl kCUDAHitAttributesParameterLayoutRulesImpl(LayoutResourceKind::HitAttributes);
 
+MetalVaryingLayoutRulesImpl kMetalVaryingInputLayoutRulesImpl(LayoutResourceKind::VertexInput);
+MetalVaryingLayoutRulesImpl kMetalVaryingOutputLayoutRulesImpl(LayoutResourceKind::FragmentOutput);
+
 struct GLSLLayoutRulesFamilyImpl : LayoutRulesFamilyImpl
 {
     virtual LayoutRulesImpl* getAnyValueRules() override;
@@ -984,10 +994,31 @@ struct CUDALayoutRulesFamilyImpl : LayoutRulesFamilyImpl
     LayoutRulesImpl* getStructuredBufferRules(CompilerOptionSet& compilerOptions) override;
 };
 
+struct MetalLayoutRulesFamilyImpl : LayoutRulesFamilyImpl
+{
+    virtual LayoutRulesImpl* getAnyValueRules() override;
+    virtual LayoutRulesImpl* getConstantBufferRules(CompilerOptionSet& compilerOptions) override;
+    virtual LayoutRulesImpl* getPushConstantBufferRules() override;
+    virtual LayoutRulesImpl* getTextureBufferRules() override;
+    virtual LayoutRulesImpl* getVaryingInputRules() override;
+    virtual LayoutRulesImpl* getVaryingOutputRules() override;
+    virtual LayoutRulesImpl* getSpecializationConstantRules() override;
+    virtual LayoutRulesImpl* getShaderStorageBufferRules(CompilerOptionSet& compilerOptions) override;
+    virtual LayoutRulesImpl* getParameterBlockRules(CompilerOptionSet& compilerOptions) override;
+
+    LayoutRulesImpl* getRayPayloadParameterRules()      override;
+    LayoutRulesImpl* getCallablePayloadParameterRules() override;
+    LayoutRulesImpl* getHitAttributesParameterRules()   override;
+
+    LayoutRulesImpl* getShaderRecordConstantBufferRules() override;
+    LayoutRulesImpl* getStructuredBufferRules(CompilerOptionSet& compilerOptions) override;
+};
+
 GLSLLayoutRulesFamilyImpl kGLSLLayoutRulesFamilyImpl;
 HLSLLayoutRulesFamilyImpl kHLSLLayoutRulesFamilyImpl;
 CPULayoutRulesFamilyImpl kCPULayoutRulesFamilyImpl;
 CUDALayoutRulesFamilyImpl kCUDALayoutRulesFamilyImpl;
+MetalLayoutRulesFamilyImpl kMetalLayoutRulesFamilyImpl;
 
 // CPU case
 
@@ -1023,7 +1054,7 @@ struct CPUObjectLayoutRulesImpl : ObjectLayoutRulesImpl
             case ShaderParameterKind::SamplerState:
                 // It's a pointer
                 return SimpleLayoutInfo(LayoutResourceKind::Uniform, sizeof(void*), SLANG_ALIGN_OF(void*));
- 
+
             case ShaderParameterKind::TextureSampler:
             case ShaderParameterKind::MutableTextureSampler:
             case ShaderParameterKind::InputRenderTarget:
@@ -1525,6 +1556,185 @@ LayoutRulesImpl* CUDALayoutRulesFamilyImpl::getStructuredBufferRules(CompilerOpt
     return &kCUDALayoutRulesImpl_;
 }
 
+// Metal Family
+
+struct MetalObjectLayoutRulesImpl : ObjectLayoutRulesImpl
+{
+    virtual ObjectLayoutInfo GetObjectLayout(ShaderParameterKind kind, const Options& /* options */) override
+    {
+        switch (kind)
+        {
+        case ShaderParameterKind::ConstantBuffer:
+        case ShaderParameterKind::StructuredBuffer:
+        case ShaderParameterKind::MutableStructuredBuffer:
+        case ShaderParameterKind::RawBuffer:
+        case ShaderParameterKind::Buffer:
+        case ShaderParameterKind::MutableRawBuffer:
+        case ShaderParameterKind::MutableBuffer:
+            return SimpleLayoutInfo(LayoutResourceKind::MetalBuffer, 1);
+        case ShaderParameterKind::AppendConsumeStructuredBuffer:
+            return SimpleLayoutInfo(LayoutResourceKind::MetalBuffer, 2);
+        case ShaderParameterKind::MutableTexture:
+        case ShaderParameterKind::TextureUniformBuffer:
+        case ShaderParameterKind::Texture:
+            return SimpleLayoutInfo(LayoutResourceKind::MetalTexture, 1);
+
+        case ShaderParameterKind::SamplerState:
+            return SimpleLayoutInfo(LayoutResourceKind::SamplerState, 1);
+
+        case ShaderParameterKind::TextureSampler:
+        case ShaderParameterKind::MutableTextureSampler:
+        {
+            ObjectLayoutInfo info;
+            info.layoutInfos.add(SimpleLayoutInfo(LayoutResourceKind::MetalTexture, 1));
+            info.layoutInfos.add(SimpleLayoutInfo(LayoutResourceKind::SamplerState, 1));
+            return info;
+        }
+        default:
+            SLANG_UNEXPECTED("unhandled shader parameter kind");
+            UNREACHABLE_RETURN(SimpleLayoutInfo());
+        }
+    }
+};
+
+struct MetalArgumentBufferElementLayoutRulesImpl : ObjectLayoutRulesImpl, DefaultLayoutRulesImpl
+{
+    SimpleLayoutInfo GetScalarLayout(BaseType baseType) override
+    {
+        SLANG_UNUSED(baseType);
+        // Everything in a metal argument buffer, including basic scalar values, occupy one `[[id]]` slot.
+        return SimpleLayoutInfo(LayoutResourceKind::MetalArgumentBufferElement, 1);
+    }
+
+    virtual ObjectLayoutInfo GetObjectLayout(ShaderParameterKind kind, const Options& /* options */) override
+    {
+        switch (kind)
+        {
+        case ShaderParameterKind::ConstantBuffer:
+        case ShaderParameterKind::StructuredBuffer:
+        case ShaderParameterKind::MutableStructuredBuffer:
+        case ShaderParameterKind::RawBuffer:
+        case ShaderParameterKind::Buffer:
+        case ShaderParameterKind::MutableRawBuffer:
+        case ShaderParameterKind::MutableBuffer:
+        case ShaderParameterKind::MutableTexture:
+        case ShaderParameterKind::TextureUniformBuffer:
+        case ShaderParameterKind::Texture:
+        case ShaderParameterKind::SamplerState:
+        {
+            return SimpleLayoutInfo(LayoutResourceKind::MetalArgumentBufferElement, 1);
+        }
+        case ShaderParameterKind::TextureSampler:
+        case ShaderParameterKind::AppendConsumeStructuredBuffer:
+        case ShaderParameterKind::MutableTextureSampler:
+        {
+            return SimpleLayoutInfo(LayoutResourceKind::MetalArgumentBufferElement, 2);
+        }
+        default:
+            SLANG_UNEXPECTED("unhandled shader parameter kind");
+            UNREACHABLE_RETURN(SimpleLayoutInfo());
+        }
+    }
+};
+
+static MetalObjectLayoutRulesImpl kMetalObjectLayoutRulesImpl;
+static MetalArgumentBufferElementLayoutRulesImpl kMetalArgumentBufferElementLayoutRulesImpl;
+
+LayoutRulesImpl kMetalAnyValueLayoutRulesImpl_ = {
+    &kMetalLayoutRulesFamilyImpl,
+    &kDefaultLayoutRulesImpl,
+    &kMetalObjectLayoutRulesImpl,
+};
+
+LayoutRulesImpl kMetalConstantBufferLayoutRulesImpl_ = {
+    &kMetalLayoutRulesFamilyImpl, & kCPULayoutRulesImpl, &kMetalObjectLayoutRulesImpl,
+};
+
+LayoutRulesImpl kMetalParameterBlockLayoutRulesImpl_ = {
+    &kMetalLayoutRulesFamilyImpl, &kMetalArgumentBufferElementLayoutRulesImpl, &kMetalArgumentBufferElementLayoutRulesImpl,
+};
+
+LayoutRulesImpl kMetalStructuredBufferLayoutRulesImpl_ = {
+    &kMetalLayoutRulesFamilyImpl, &kCPULayoutRulesImpl, & kMetalObjectLayoutRulesImpl,
+};
+
+LayoutRulesImpl kMetalVaryingInputLayoutRulesImpl_ = {
+    &kMetalLayoutRulesFamilyImpl, &kMetalVaryingInputLayoutRulesImpl, &kHLSLObjectLayoutRulesImpl,
+};
+
+LayoutRulesImpl kMetalVaryingOutputLayoutRulesImpl_ = {
+    &kMetalLayoutRulesFamilyImpl, &kMetalVaryingOutputLayoutRulesImpl, &kHLSLObjectLayoutRulesImpl,
+};
+
+LayoutRulesImpl* MetalLayoutRulesFamilyImpl::getAnyValueRules()
+{
+    return &kHLSLAnyValueLayoutRulesImpl_;
+}
+
+LayoutRulesImpl* MetalLayoutRulesFamilyImpl::getConstantBufferRules(CompilerOptionSet&)
+{
+    return &kMetalConstantBufferLayoutRulesImpl_;
+}
+
+LayoutRulesImpl* MetalLayoutRulesFamilyImpl::getParameterBlockRules(CompilerOptionSet&)
+{
+    return &kMetalParameterBlockLayoutRulesImpl_;
+}
+
+LayoutRulesImpl* MetalLayoutRulesFamilyImpl::getPushConstantBufferRules()
+{
+    return &kMetalConstantBufferLayoutRulesImpl_;
+}
+
+LayoutRulesImpl* MetalLayoutRulesFamilyImpl::getShaderRecordConstantBufferRules()
+{
+    return &kMetalConstantBufferLayoutRulesImpl_;
+}
+
+LayoutRulesImpl* MetalLayoutRulesFamilyImpl::getStructuredBufferRules(CompilerOptionSet&)
+{
+    return &kMetalStructuredBufferLayoutRulesImpl_;
+}
+
+LayoutRulesImpl* MetalLayoutRulesFamilyImpl::getTextureBufferRules()
+{
+    return nullptr;
+}
+
+LayoutRulesImpl* MetalLayoutRulesFamilyImpl::getVaryingInputRules()
+{
+    return &kMetalVaryingInputLayoutRulesImpl_;
+}
+
+LayoutRulesImpl* MetalLayoutRulesFamilyImpl::getVaryingOutputRules()
+{
+    return &kHLSLVaryingOutputLayoutRulesImpl_;
+}
+
+LayoutRulesImpl* MetalLayoutRulesFamilyImpl::getSpecializationConstantRules()
+{
+    return nullptr;
+}
+
+LayoutRulesImpl* MetalLayoutRulesFamilyImpl::getShaderStorageBufferRules(CompilerOptionSet&)
+{
+    return nullptr;
+}
+
+LayoutRulesImpl* MetalLayoutRulesFamilyImpl::getRayPayloadParameterRules()
+{
+    return nullptr;
+}
+
+LayoutRulesImpl* MetalLayoutRulesFamilyImpl::getCallablePayloadParameterRules()
+{
+    return nullptr;
+}
+
+LayoutRulesImpl* MetalLayoutRulesFamilyImpl::getHitAttributesParameterRules()
+{
+    return nullptr;
+}
 
 
 LayoutRulesFamilyImpl* getDefaultLayoutRulesFamilyForTarget(TargetRequest* targetReq)
@@ -1549,9 +1759,6 @@ LayoutRulesFamilyImpl* getDefaultLayoutRulesFamilyForTarget(TargetRequest* targe
     case CodeGenTarget::ShaderSharedLibrary:
     case CodeGenTarget::CPPSource:
     case CodeGenTarget::CSource:
-    case CodeGenTarget::Metal:
-    case CodeGenTarget::MetalLib:
-    case CodeGenTarget::MetalLibAssembly:
     {
         // For now lets use some fairly simple CPU binding rules
 
@@ -1563,6 +1770,10 @@ LayoutRulesFamilyImpl* getDefaultLayoutRulesFamilyForTarget(TargetRequest* targe
 
         return &kCPULayoutRulesFamilyImpl;
     }
+    case CodeGenTarget::Metal:
+    case CodeGenTarget::MetalLib:
+    case CodeGenTarget::MetalLibAssembly:
+        return &kMetalLayoutRulesFamilyImpl;
 
     case CodeGenTarget::PTX:
     case CodeGenTarget::CUDASource:
@@ -1761,6 +1972,20 @@ bool isD3DTarget(TargetRequest* targetReq)
     }
 }
 
+bool isMetalTarget(TargetRequest* targetReq)
+{
+    switch (targetReq->getTarget())
+    {
+    default:
+        return false;
+
+    case CodeGenTarget::Metal:
+    case CodeGenTarget::MetalLib:
+    case CodeGenTarget::MetalLibAssembly:
+        return true;
+    }
+}
+
 bool isKhronosTarget(TargetRequest* targetReq)
 {
     switch( targetReq->getTarget() )
@@ -2242,7 +2467,9 @@ static void _addUnmaskedResourceUsage(
                 dstTypeLayout->addResourceUsage(resInfo);
             }
             break;
-
+        case LayoutResourceKind::MetalArgumentBufferElement:
+            // A metal argument buffer element will always be masked.
+            break;
         case LayoutResourceKind::SubElementRegisterSpace:
         case LayoutResourceKind::ExistentialTypeParam:
             // A parameter group will always pay for full registers
@@ -2390,7 +2617,8 @@ static RefPtr<TypeLayout> _createParameterGroupTypeLayout(
     bool wantConstantBuffer = _usesOrdinaryData(rawElementTypeLayout)
         || _usesExistentialData(rawElementTypeLayout)
         || isCUDATarget(context.targetReq)
-        || isCPUTarget(context.targetReq);
+        || isCPUTarget(context.targetReq)
+        || isMetalTarget(context.targetReq);
     if( wantConstantBuffer )
     {
         // If there is any ordinary data, then we'll need to
@@ -2489,8 +2717,11 @@ static RefPtr<TypeLayout> _createParameterGroupTypeLayout(
     {
         auto kind = elementTypeResInfo.kind;
 
-        // TODO: Added to make layout work correctly for CPU target 
-        if(kind == LayoutResourceKind::Uniform)
+        // Uniforms and MetalArgumentBUfferElements are private
+        // to the element layout and do not share the binding space
+        // with the container.
+        if (kind == LayoutResourceKind::Uniform ||
+            kind == LayoutResourceKind::MetalArgumentBufferElement)
         {
             continue;
         }
diff --git a/source/slang/slang-type-layout.h b/source/slang/slang-type-layout.h
index acf6088eb..6ca756e8c 100644
--- a/source/slang/slang-type-layout.h
+++ b/source/slang/slang-type-layout.h
@@ -272,7 +272,11 @@ typedef slang::ParameterCategory LayoutResourceKind;
     x(SubElementRegisterSpace) \
     x(VertexInput) \
     x(FragmentOutput) \
-    x(InputAttachmentIndex)
+    x(InputAttachmentIndex) \
+    \
+    x(MetalBuffer) \
+    x(MetalTexture) \
+    x(MetalArgumentBufferElement)
 
 #define SLANG_PARAMETER_CATEGORY_FLAG(x) x = ParameterCategoryFlags(1) << int(slang::x),