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

Author: csyonghe

slang.h

@@ -2259,9 +2259,16 @@ extern "C"
         // The input_attachment_index subpass occupancy tracker
         SLANG_PARAMETER_CATEGORY_SUBPASS,
 
+        // Metal resource binding points.
+        SLANG_PARAMETER_CATEGORY_METAL_ARGUMENT_BUFFER_ELEMENT,
+
         //
         SLANG_PARAMETER_CATEGORY_COUNT,
 
+        // Aliases for Metal-specific categories.
+        SLANG_PARAMETER_CATEGORY_METAL_BUFFER = SLANG_PARAMETER_CATEGORY_CONSTANT_BUFFER,
+        SLANG_PARAMETER_CATEGORY_METAL_TEXTURE = SLANG_PARAMETER_CATEGORY_SHADER_RESOURCE,
+        SLANG_PARAMETER_CATEGORY_METAL_SAMPLER = SLANG_PARAMETER_CATEGORY_SAMPLER_STATE,
 
         // DEPRECATED:
         SLANG_PARAMETER_CATEGORY_VERTEX_INPUT = SLANG_PARAMETER_CATEGORY_VARYING_INPUT,
@@ -2825,6 +2832,10 @@ namespace slang
 
         InputAttachmentIndex = SLANG_PARAMETER_CATEGORY_SUBPASS,
 
+        MetalBuffer = SLANG_PARAMETER_CATEGORY_CONSTANT_BUFFER,
+        MetalTexture = SLANG_PARAMETER_CATEGORY_METAL_TEXTURE,
+        MetalArgumentBufferElement = SLANG_PARAMETER_CATEGORY_METAL_ARGUMENT_BUFFER_ELEMENT,
+
         // DEPRECATED:
         VertexInput = SLANG_PARAMETER_CATEGORY_VERTEX_INPUT,
         FragmentOutput = SLANG_PARAMETER_CATEGORY_FRAGMENT_OUTPUT,

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);
 

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)

source/slang/slang-emit-metal.h

@@ -23,7 +23,6 @@ class MetalSourceEmitter : public CLikeSourceEmitter
 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 @@ class MetalSourceEmitter : public CLikeSourceEmitter
     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 @@ class MetalSourceEmitter : public CLikeSourceEmitter
 
     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);
 

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,

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,18 +3988,20 @@ struct IREmptyTypeLegalizationContext : IRTypeLegalizationContext
 // specialized context type to use to get the job done.
 
 void legalizeResourceTypes(
+    TargetProgram* target,
     IRModule*       module,
     DiagnosticSink* sink)
 {
     SLANG_PROFILE;
 
     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);
 }
 

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});
                 }
 

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

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);