Metal: Vertex/Fragment builtin and layouts. (#4044)

* Metal: Vertex/Fragment builtin and layouts.

* Fix.

* Fix test.

* Emit user semantic on vertex/fragment attributes.

12 files changed, 590 insertions, 7 deletions

diff --git a/build/visual-studio/slang/slang.vcxproj b/build/visual-studio/slang/slang.vcxproj
index 1926f88cd..6ab460faf 100644
--- a/build/visual-studio/slang/slang.vcxproj
+++ b/build/visual-studio/slang/slang.vcxproj
@@ -442,6 +442,7 @@ IF EXIST ..\..\..\external\slang-glslang\bin\windows-aarch64\release\slang-glsla
     <ClInclude Include="..\..\..\source\slang\slang-ir-lower-witness-lookup.h" />

     <ClInclude Include="..\..\..\source\slang\slang-ir-marshal-native-call.h" />

     <ClInclude Include="..\..\..\source\slang\slang-ir-metadata.h" />

+    <ClInclude Include="..\..\..\source\slang\slang-ir-metal-legalize.h" />

     <ClInclude Include="..\..\..\source\slang\slang-ir-missing-return.h" />

     <ClInclude Include="..\..\..\source\slang\slang-ir-obfuscate-loc.h" />

     <ClInclude Include="..\..\..\source\slang\slang-ir-optix-entry-point-uniforms.h" />

@@ -673,6 +674,7 @@ IF EXIST ..\..\..\external\slang-glslang\bin\windows-aarch64\release\slang-glsla
     <ClCompile Include="..\..\..\source\slang\slang-ir-lower-witness-lookup.cpp" />

     <ClCompile Include="..\..\..\source\slang\slang-ir-marshal-native-call.cpp" />

     <ClCompile Include="..\..\..\source\slang\slang-ir-metadata.cpp" />

+    <ClCompile Include="..\..\..\source\slang\slang-ir-metal-legalize.cpp" />

     <ClCompile Include="..\..\..\source\slang\slang-ir-missing-return.cpp" />

     <ClCompile Include="..\..\..\source\slang\slang-ir-obfuscate-loc.cpp" />

     <ClCompile Include="..\..\..\source\slang\slang-ir-optix-entry-point-uniforms.cpp" />

diff --git a/build/visual-studio/slang/slang.vcxproj.filters b/build/visual-studio/slang/slang.vcxproj.filters
index 64e3af88e..0112bb818 100644
--- a/build/visual-studio/slang/slang.vcxproj.filters
+++ b/build/visual-studio/slang/slang.vcxproj.filters
@@ -414,6 +414,9 @@
     <ClInclude Include="..\..\..\source\slang\slang-ir-metadata.h">

       <Filter>Header Files</Filter>

     </ClInclude>

+    <ClInclude Include="..\..\..\source\slang\slang-ir-metal-legalize.h">

+      <Filter>Header Files</Filter>

+    </ClInclude>

     <ClInclude Include="..\..\..\source\slang\slang-ir-missing-return.h">

       <Filter>Header Files</Filter>

     </ClInclude>

@@ -1103,6 +1106,9 @@
     <ClCompile Include="..\..\..\source\slang\slang-ir-metadata.cpp">

       <Filter>Source Files</Filter>

     </ClCompile>

+    <ClCompile Include="..\..\..\source\slang\slang-ir-metal-legalize.cpp">

+      <Filter>Source Files</Filter>

+    </ClCompile>

     <ClCompile Include="..\..\..\source\slang\slang-ir-missing-return.cpp">

       <Filter>Source Files</Filter>

     </ClCompile>

diff --git a/slang.h b/slang.h
index dad2f1749..20d75e1f7 100644
--- a/slang.h
+++ b/slang.h
@@ -2262,6 +2262,9 @@ extern "C"
         // Metal resource binding points.
         SLANG_PARAMETER_CATEGORY_METAL_ARGUMENT_BUFFER_ELEMENT,
 
+        // Metal [[attribute]] inputs.
+        SLANG_PARAMETER_CATEGORY_METAL_ATTRIBUTE,
+
         //
         SLANG_PARAMETER_CATEGORY_COUNT,
 
@@ -2835,6 +2838,7 @@ namespace slang
         MetalBuffer = SLANG_PARAMETER_CATEGORY_CONSTANT_BUFFER,
         MetalTexture = SLANG_PARAMETER_CATEGORY_METAL_TEXTURE,
         MetalArgumentBufferElement = SLANG_PARAMETER_CATEGORY_METAL_ARGUMENT_BUFFER_ELEMENT,
+        MetalAttribute = SLANG_PARAMETER_CATEGORY_METAL_ATTRIBUTE,
 
         // DEPRECATED:
         VertexInput = SLANG_PARAMETER_CATEGORY_VERTEX_INPUT,
diff --git a/source/slang/slang-emit-metal.cpp b/source/slang/slang-emit-metal.cpp
index 3773b7c2a..7580ed74d 100644
--- a/source/slang/slang-emit-metal.cpp
+++ b/source/slang/slang-emit-metal.cpp
@@ -128,8 +128,13 @@ void MetalSourceEmitter::emitFuncParamLayoutImpl(IRInst* param)
             m_writer->emit(rr->getOffset());
             m_writer->emit(")]]");
             break;
+        case LayoutResourceKind::VaryingInput:
+            m_writer->emit(" [[stage_in]]");
+            break;
         }
     }
+    if (auto sysSemanticAttr = layout->findSystemValueSemanticAttr())
+       _emitSystemSemantic(sysSemanticAttr->getName(), sysSemanticAttr->getIndex());
 }
 
 void MetalSourceEmitter::emitParameterGroupImpl(IRGlobalParam* varDecl, IRUniformParameterGroupType* type)
@@ -610,11 +615,130 @@ void MetalSourceEmitter::emitRateQualifiersAndAddressSpaceImpl(IRRate* rate, [[m
     }
 }
 
+void MetalSourceEmitter::_emitSystemSemantic(UnownedStringSlice semanticName, IRIntegerValue semanticIndex)
+{
+    if (semanticName.caseInsensitiveEquals(toSlice("SV_POSITION")))
+    {
+        m_writer->emit(" [[position]]");
+    }
+    else if (semanticName.caseInsensitiveEquals(toSlice("SV_VERTEXID")))
+    {
+        m_writer->emit(" [[vertex_id]]");
+    }
+    else if (semanticName.caseInsensitiveEquals(toSlice("SV_INSTANCEID")))
+    {
+        m_writer->emit(" [[instance_id]]");
+    }
+    else if (semanticName.caseInsensitiveEquals(toSlice("SV_Target")))
+    {
+        m_writer->emit(" [[color(");
+        m_writer->emit(semanticIndex);
+        m_writer->emit(")]]");
+    }
+    else if (semanticName.caseInsensitiveEquals(toSlice("SV_PRIMITIVEID")))
+    {
+        m_writer->emit(" [[primitive_id]]");
+    }
+    else if (semanticName.caseInsensitiveEquals(toSlice("SV_GROUPID")))
+    {
+        // TODO: not supported by metal.
+        // We need to implement the transformation logic in slang-ir-metal-legalize.cpp
+        // to convert SV_GroupID to something like METAL_threadgroup_position_in_grid.
+    }
+    else if (semanticName.caseInsensitiveEquals(toSlice("SV_GROUPINDEX")))
+    {
+        // TODO: not supported by metal.
+    }
+    else if (semanticName.caseInsensitiveEquals(toSlice("SV_DISPATCHTHREADID")))
+    {
+        m_writer->emit(" [[thread_position_in_grid]]");
+    }
+    else if (semanticName.caseInsensitiveEquals(toSlice("SV_GROUPTHREADID")))
+    {
+        m_writer->emit(" [[thread_position_in_threadgroup]]");
+    }
+    else if (semanticName.caseInsensitiveEquals(toSlice("SV_CLIPDISTANCE")))
+    {
+        m_writer->emit(" [[clip_distance]]");
+    }
+    else if (semanticName.caseInsensitiveEquals(toSlice("SV_RENDERTARGETARRAYINDEX")))
+    {
+        m_writer->emit(" [[render_target_array_index]]");
+    }
+    else if (semanticName.caseInsensitiveEquals(toSlice("SV_VIEWPORTARRAYINDEX")))
+    {
+        m_writer->emit(" [[viewport_array_index]]");
+    }
+    else if (semanticName.caseInsensitiveEquals(toSlice("SV_Depth")))
+    {
+        m_writer->emit(" [[depth(any)]]");
+    }
+    else if (semanticName.caseInsensitiveEquals(toSlice("SV_DepthGreaterEqual")))
+    {
+        m_writer->emit(" [[depth(greater)]]");
+    }
+    else if (semanticName.caseInsensitiveEquals(toSlice("SV_DepthLessEqual")))
+    {
+        m_writer->emit(" [[depth(less)]]");
+    }
+    else if (semanticName.caseInsensitiveEquals(toSlice("SV_Coverage")))
+    {
+        m_writer->emit(" [[sample_mask]]");
+    }
+    else if (semanticName.caseInsensitiveEquals(toSlice("SV_StencilRef")))
+    {
+        m_writer->emit(" [[stencil]]");
+    }
+    else
+    {
+        m_writer->emit(" [[user(");
+        m_writer->emit(semanticName);
+        if (semanticIndex != 0)
+        {
+            m_writer->emit("_");
+            m_writer->emit(semanticIndex);
+        }
+        m_writer->emit(")]]");
+    }
+}
+
 void MetalSourceEmitter::emitSemanticsImpl(IRInst* inst, bool allowOffsets)
 {
-    // Metal does not use semantics.
-    SLANG_UNUSED(inst);
     SLANG_UNUSED(allowOffsets);
+    if (inst->getOp() == kIROp_StructKey)
+    {
+        // Only emit [[attribute(n)]] on struct keys.
+        bool hasSemanticFromLayout = false;
+        if (auto varLayout = findVarLayout(inst))
+        {
+            for (auto attr : varLayout->getAllAttrs())
+            {
+                if (auto offsetAttr = as<IRVarOffsetAttr>(attr))
+                {
+                    if (offsetAttr->getResourceKind() == LayoutResourceKind::MetalAttribute)
+                    {
+                        m_writer->emit(" [[attribute(");
+                        m_writer->emit(offsetAttr->getOffset());
+                        m_writer->emit(")]]");
+                    }
+                }
+                else if (auto semanticAttr = as<IRSemanticAttr>(attr))
+                {
+                    auto semanticName = String(semanticAttr->getName()).toUpper();
+                    _emitSystemSemantic(semanticAttr->getName(), semanticAttr->getIndex());
+                    hasSemanticFromLayout = true;
+                }
+            }
+
+        }
+        if (!hasSemanticFromLayout)
+        {
+            if (auto semanticDecor = inst->findDecoration<IRSemanticDecoration>())
+            {
+                _emitSystemSemantic(semanticDecor->getSemanticName(), semanticDecor->getSemanticIndex());
+            }
+        }
+    }
 }
 
 void MetalSourceEmitter::_emitStageAccessSemantic(IRStageAccessDecoration* decoration, const char* name)
diff --git a/source/slang/slang-emit-metal.h b/source/slang/slang-emit-metal.h
index 38d4c3a2c..d925365da 100644
--- a/source/slang/slang-emit-metal.h
+++ b/source/slang/slang-emit-metal.h
@@ -53,7 +53,6 @@ protected:
     virtual void handleRequiredCapabilitiesImpl(IRInst* inst) SLANG_OVERRIDE;
     
     virtual void emitGlobalInstImpl(IRInst* inst) SLANG_OVERRIDE;
-
     virtual bool doesTargetSupportPtrTypes() SLANG_OVERRIDE { return true; }
 
     void emitFuncParamLayoutImpl(IRInst* param);
@@ -68,6 +67,7 @@ protected:
     void _emitHLSLDecorationSingleInt(const char* name, IRFunc* entryPoint, IRIntLit* val);
 
     void _emitStageAccessSemantic(IRStageAccessDecoration* decoration, const char* name);
+    void _emitSystemSemantic(UnownedStringSlice semanticName, IRIntegerValue semanticIndex);
 };
 
 }
diff --git a/source/slang/slang-emit.cpp b/source/slang/slang-emit.cpp
index 9369afbc5..87f0911e7 100644
--- a/source/slang/slang-emit.cpp
+++ b/source/slang/slang-emit.cpp
@@ -29,6 +29,7 @@
 #include "slang-ir-fuse-satcoop.h"
 #include "slang-ir-glsl-legalize.h"
 #include "slang-ir-hlsl-legalize.h"
+#include "slang-ir-metal-legalize.h"
 #include "slang-ir-insts.h"
 #include "slang-ir-inline.h"
 #include "slang-ir-legalize-array-return-type.h"
@@ -834,7 +835,11 @@ Result linkAndOptimizeIR(
             validateIRModuleIfEnabled(codeGenContext, irModule);
     }
     break;
-
+    case CodeGenTarget::Metal:
+    {
+        legalizeIRForMetal(irModule, sink);
+    }
+    break;
     case CodeGenTarget::CSource:
     case CodeGenTarget::CPPSource:
         {
diff --git a/source/slang/slang-ir-metal-legalize.cpp b/source/slang/slang-ir-metal-legalize.cpp
new file mode 100644
index 000000000..822a1e2f1
--- /dev/null
+++ b/source/slang/slang-ir-metal-legalize.cpp
@@ -0,0 +1,337 @@
+#include "slang-ir-metal-legalize.h"
+
+#include "slang-ir-insts.h"
+#include "slang-ir-util.h"
+#include "slang-ir-clone.h"
+
+namespace Slang
+{
+    struct EntryPointInfo
+    {
+        IRFunc* entryPointFunc;
+        IREntryPointDecoration* entryPointDecor;
+    };
+
+    void hoistEntryPointParameterFromStruct(EntryPointInfo entryPoint)
+    {
+        // If an entry point has a input parameter with a struct type, we want to hoist out
+        // all the fields of the struct type to be individual parameters of the entry point.
+        // This will canonicalize the entry point signature, so we can handle all cases uniformly.
+
+        // For example, given an entry point:
+        // ```
+        // struct VertexInput { float3 pos; float 2 uv; int vertexId : SV_VertexID};
+        // void main(VertexInput vin) { ... }
+        // ```
+        // We will transform it to:
+        // ```
+        // void main(float3 pos, float2 uv, int vertexId : SV_VertexID) {
+        //     VertexInput vin = {pos,uv,vertexId};
+        //     ...
+        // }
+        // ```
+
+        auto func = entryPoint.entryPointFunc;
+        List<IRParam*> paramsToProcess;
+        for (auto param : func->getParams())
+        {
+            if (auto structType = as<IRStructType>(param->getDataType()))
+            {
+                paramsToProcess.add(param);
+            }
+        }
+
+        IRBuilder builder(func);
+        builder.setInsertBefore(func);
+        for (auto param : paramsToProcess)
+        {
+            auto structType = as<IRStructType>(param->getDataType());
+            builder.setInsertBefore(func->getFirstBlock()->getFirstOrdinaryInst());
+            auto varLayout = findVarLayout(param);
+            IRStructTypeLayout* structTypeLayout = nullptr;
+            if (varLayout)
+                structTypeLayout = as<IRStructTypeLayout>(varLayout->getTypeLayout());
+            Index fieldIndex = 0;
+            List<IRInst*> fieldParams;
+            for (auto field : structType->getFields())
+            {
+                auto fieldParam = builder.emitParam(field->getFieldType());
+                
+                IRCloneEnv cloneEnv;
+                cloneInstDecorationsAndChildren(&cloneEnv, builder.getModule(), field->getKey(), fieldParam);
+
+                IRVarLayout* fieldLayout = structTypeLayout ? structTypeLayout->getFieldLayout(fieldIndex) : nullptr;
+                if (varLayout)
+                {
+                    IRVarLayout::Builder varLayoutBuilder(&builder, fieldLayout->getTypeLayout());
+                    varLayoutBuilder.cloneEverythingButOffsetsFrom(fieldLayout);
+                    for (auto offsetAttr : fieldLayout->getOffsetAttrs())
+                    {
+                        auto parentOffsetAttr = varLayout->findOffsetAttr(offsetAttr->getResourceKind());
+                        UInt parentOffset = parentOffsetAttr ? parentOffsetAttr->getOffset() : 0;
+                        UInt parentSpace = parentOffsetAttr ? parentOffsetAttr->getSpace() : 0;
+                        auto resInfo = varLayoutBuilder.findOrAddResourceInfo(offsetAttr->getResourceKind());
+                        resInfo->offset = parentOffset + offsetAttr->getOffset();
+                        resInfo->space = parentSpace + offsetAttr->getSpace();
+                    }
+                    builder.addLayoutDecoration(fieldParam, varLayoutBuilder.build());
+                }
+                param->insertBefore(fieldParam);
+                fieldParams.add(fieldParam);
+                fieldIndex++;
+            }
+            builder.setInsertBefore(func->getFirstBlock()->getFirstOrdinaryInst());
+            auto reconstructedParam = builder.emitMakeStruct(structType, fieldParams.getCount(), fieldParams.getBuffer());
+            param->replaceUsesWith(reconstructedParam);
+            param->removeFromParent();
+        }
+        fixUpFuncType(func);
+    }
+
+    void packStageInParameters(EntryPointInfo entryPoint)
+    {
+        // If the entry point has any parameters whose layout contains VaryingInput,
+        // we need to pack those parameters into a single `struct` type, and decorate
+        // the fields with the appropriate `[[attribute]]` decorations.
+        // For other parameters that are not `VaryingInput`, we need to leave them as is.
+        // 
+        // For example, given this code after `hoistEntryPointParameterFromStruct`:
+         // ```
+        // void main(float3 pos, float2 uv, int vertexId : SV_VertexID) {
+        //     VertexInput vin = {pos,uv,vertexId};
+        //     ...
+        // }
+        // ```
+        // We are going to transform it into:
+        // ```
+        // struct VertexInput {
+        //     float3 pos [[attribute(0)]];
+        //     float2 uv [[attribute(1)]];
+        // };
+        // void main(VertexInput vin, int vertexId : SV_VertexID) {
+        //     let pos = vin.pos;
+        //     let uv = vin.uv;
+        //     ...
+        // }
+
+        auto func = entryPoint.entryPointFunc;
+
+        bool isGeometryStage = false;
+        switch (entryPoint.entryPointDecor->getProfile().getStage())
+        {
+        case Stage::Vertex:
+        case Stage::Amplification:
+        case Stage::Mesh:
+        case Stage::Geometry:
+        case Stage::Domain:
+        case Stage::Hull:
+            isGeometryStage = true;
+            break;
+        }
+
+        List<IRParam*> paramsToPack;
+        for (auto param : func->getParams())
+        {
+            auto layout = findVarLayout(param);
+            if (!layout)
+                continue;
+            if (!layout->findOffsetAttr(LayoutResourceKind::VaryingInput))
+                continue;
+            paramsToPack.add(param);
+        }
+
+        if (paramsToPack.getCount() == 0)
+            return;
+
+        IRBuilder builder(func);
+        builder.setInsertBefore(func);
+        IRStructType* structType = builder.createStructType();
+        auto stageText = getStageText(entryPoint.entryPointDecor->getProfile().getStage());
+        builder.addNameHintDecoration(structType, (String(stageText) + toSlice("Input")).getUnownedSlice());
+        List<IRStructKey*> keys;
+        IRStructTypeLayout::Builder layoutBuilder(&builder);
+        for (auto param : paramsToPack)
+        {
+            auto paramVarLayout = findVarLayout(param);
+            auto key = builder.createStructKey();
+            param->transferDecorationsTo(key);
+            builder.createStructField(structType, key, param->getDataType());
+            if (auto varyingInOffsetAttr = paramVarLayout->findOffsetAttr(LayoutResourceKind::VaryingInput))
+            {
+                if (!key->findDecoration<IRSemanticDecoration>() && !paramVarLayout->findAttr<IRSemanticAttr>())
+                {
+                    // If the parameter doesn't have a semantic, we need to add one for semantic matching.
+                    builder.addSemanticDecoration(key, toSlice("_slang_attr"), (int)varyingInOffsetAttr->getOffset());
+                }
+            }
+            if (isGeometryStage)
+            {
+                // For geometric stages, we need to translate VaryingInput offsets to MetalAttribute offsets.
+                IRVarLayout::Builder elementVarLayoutBuilder(&builder, paramVarLayout->getTypeLayout());
+                elementVarLayoutBuilder.cloneEverythingButOffsetsFrom(paramVarLayout);
+                for (auto offsetAttr : paramVarLayout->getOffsetAttrs())
+                {
+                    auto resourceKind = offsetAttr->getResourceKind();
+                    if (resourceKind == LayoutResourceKind::VaryingInput)
+                    {
+                        resourceKind = LayoutResourceKind::MetalAttribute;
+                    }
+                    auto resInfo = elementVarLayoutBuilder.findOrAddResourceInfo(resourceKind);
+                    resInfo->offset = offsetAttr->getOffset();
+                    resInfo->space = offsetAttr->getSpace();
+                }
+                paramVarLayout = elementVarLayoutBuilder.build();
+            }
+            layoutBuilder.addField(key, paramVarLayout);
+            builder.addLayoutDecoration(key, paramVarLayout);
+            keys.add(key);
+        }
+        builder.setInsertInto(func->getFirstBlock());
+        auto packedParam = builder.emitParamAtHead(structType);
+        auto typeLayout = layoutBuilder.build();
+        IRVarLayout::Builder varLayoutBuilder(&builder, typeLayout);
+
+        // Add a VaryingInput resource info to the packed parameter layout, so that we can emit
+        // the needed `[[stage_in]]` attribute in Metal emitter.
+        varLayoutBuilder.findOrAddResourceInfo(LayoutResourceKind::VaryingInput);
+        auto paramVarLayout = varLayoutBuilder.build();
+        builder.addLayoutDecoration(packedParam, paramVarLayout);
+
+        // Replace the original parameters with the packed parameter
+        builder.setInsertBefore(func->getFirstBlock()->getFirstOrdinaryInst());
+        for (Index paramIndex = 0; paramIndex < paramsToPack.getCount(); paramIndex++)
+        {
+            auto param = paramsToPack[paramIndex];
+            auto key = keys[paramIndex];
+            auto paramField = builder.emitFieldExtract(param->getDataType(), packedParam, key);
+            param->replaceUsesWith(paramField);
+            param->removeFromParent();
+        }
+        fixUpFuncType(func);
+    }
+
+
+    void ensureResultStructHasUserSemantic(IRStructType* structType, IRVarLayout* varLayout)
+    {
+        // Ensure each field in an output struct type has either a system semantic or a user semantic,
+        // so that signature matching can happen correctly.
+        auto typeLayout = as<IRStructTypeLayout>(varLayout->getTypeLayout());
+        Index index = 0;
+        IRBuilder builder(structType);
+        for (auto field : structType->getFields())
+        {
+            auto key = field->getKey();
+            if (key->findDecoration<IRSemanticDecoration>())
+            {
+                index++;
+                continue;
+            }
+            typeLayout->getFieldLayout(index);
+            auto fieldLayout = typeLayout->getFieldLayout(index);
+            if (auto offsetAttr = fieldLayout->findOffsetAttr(LayoutResourceKind::VaryingOutput))
+            {
+                UInt varOffset = 0;
+                if (auto varOffsetAttr = varLayout->findOffsetAttr(LayoutResourceKind::VaryingOutput))
+                    varOffset = varOffsetAttr->getOffset();
+                varOffset += offsetAttr->getOffset();
+                builder.addSemanticDecoration(key, toSlice("_slang_attr"), (int)varOffset);
+            }
+            index++;
+        }
+    }
+
+
+    void wrapReturnValueInStruct(EntryPointInfo entryPoint)
+    {
+        // Wrap return value into a struct if it is not already a struct.
+        // For example, given this entry point:
+        // ```
+        // float4 main() : SV_Target { return float3(1,2,3); }
+        // ```
+        // We are going to transform it into:
+        // ```
+        // struct Output {
+        //     float4 value : SV_Target;
+        // };
+        // Output main() { return {float3(1,2,3)}; }
+
+        auto func = entryPoint.entryPointFunc;
+
+        auto returnType = func->getResultType();
+        if (as<IRVoidType>(returnType))
+            return;
+        auto entryPointLayoutDecor = func->findDecoration<IRLayoutDecoration>();
+        if (!entryPointLayoutDecor)
+            return;
+        auto entryPointLayout = as<IREntryPointLayout>(entryPointLayoutDecor->getLayout());
+        if (!entryPointLayout)
+            return;
+        auto resultLayout = entryPointLayout->getResultLayout();
+
+        // If return type is already a struct, just make sure every field has a semantic.
+        if (auto returnStructType = as<IRStructType>(returnType))
+        {
+            ensureResultStructHasUserSemantic(returnStructType, resultLayout);
+            return;
+        }
+
+        // If not, we need to wrap the result into a struct type.
+        IRBuilder builder(func);
+        builder.setInsertBefore(func);
+        IRStructType* structType = builder.createStructType();
+        auto stageText = getStageText(entryPoint.entryPointDecor->getProfile().getStage());
+        builder.addNameHintDecoration(structType, (String(stageText) + toSlice("Output")).getUnownedSlice());
+        auto key = builder.createStructKey();
+        builder.addNameHintDecoration(key, toSlice("output"));
+        builder.addLayoutDecoration(key, resultLayout);
+        builder.createStructField(structType, key, returnType);
+        IRStructTypeLayout::Builder structTypeLayoutBuilder(&builder);
+        structTypeLayoutBuilder.addField(key, resultLayout);
+        auto typeLayout = structTypeLayoutBuilder.build();
+        IRVarLayout::Builder varLayoutBuilder(&builder, typeLayout);
+        auto varLayout = varLayoutBuilder.build();
+        ensureResultStructHasUserSemantic(structType, varLayout);
+
+        for (auto block : func->getBlocks())
+        {
+            if (auto returnInst = as<IRReturn>(block->getTerminator()))
+            {
+                builder.setInsertBefore(returnInst);
+                auto returnVal = returnInst->getVal();
+                auto newResult = builder.emitMakeStruct(structType, 1, &returnVal);
+                returnInst->setOperand(0, newResult);
+            }
+        }
+        fixUpFuncType(func, structType);
+    }
+
+    void legalizeEntryPointForMetal(EntryPointInfo entryPoint, DiagnosticSink* sink)
+    {
+        SLANG_UNUSED(sink);
+
+        hoistEntryPointParameterFromStruct(entryPoint);
+        packStageInParameters(entryPoint);
+        wrapReturnValueInStruct(entryPoint);
+    }
+
+    void legalizeIRForMetal(IRModule* module, DiagnosticSink* sink)
+    {
+        List<EntryPointInfo> entryPoints;
+        for (auto inst : module->getGlobalInsts())
+        {
+            if (auto func = as<IRFunc>(inst))
+            {
+                if (auto entryPointDecor = func->findDecoration<IREntryPointDecoration>())
+                {
+                    EntryPointInfo info;
+                    info.entryPointDecor = entryPointDecor;
+                    info.entryPointFunc = func;
+                    entryPoints.add(info);
+                }
+            }
+        }
+
+        for (auto entryPoint : entryPoints)
+            legalizeEntryPointForMetal(entryPoint, sink);
+    }
+}
diff --git a/source/slang/slang-ir-metal-legalize.h b/source/slang/slang-ir-metal-legalize.h
new file mode 100644
index 000000000..3eb64438f
--- /dev/null
+++ b/source/slang/slang-ir-metal-legalize.h
@@ -0,0 +1,10 @@
+#pragma once
+
+#include "slang-ir.h"
+
+namespace Slang
+{
+    class DiagnosticSink;
+
+    void legalizeIRForMetal(IRModule* module, DiagnosticSink* sink);
+}
diff --git a/source/slang/slang-options.cpp b/source/slang/slang-options.cpp
index 9d9e8039b..dad69350e 100644
--- a/source/slang/slang-options.cpp
+++ b/source/slang/slang-options.cpp
@@ -2746,6 +2746,9 @@ SlangResult OptionsParser::_parse(
             m_rawTargets[0].format == CodeGenTarget::CUDASource ||
             m_rawTargets[0].format == CodeGenTarget::SPIRV ||
             m_rawTargets[0].format == CodeGenTarget::SPIRVAssembly ||
+            m_rawTargets[0].format == CodeGenTarget::Metal ||
+            m_rawTargets[0].format == CodeGenTarget::MetalLib ||
+            m_rawTargets[0].format == CodeGenTarget::MetalLibAssembly ||
             ArtifactDescUtil::makeDescForCompileTarget(asExternal(m_rawTargets[0].format)).kind == ArtifactKind::HostCallable))
     {
         RawOutput rawOutput;
diff --git a/source/slang/slang.cpp b/source/slang/slang.cpp
index 971d6056f..4a446a351 100644
--- a/source/slang/slang.cpp
+++ b/source/slang/slang.cpp
@@ -6445,6 +6445,9 @@ SlangResult EndToEndCompileRequest::isParameterLocationUsed(Int entryPointIndex,
     if (SLANG_FAILED(_getEntryPointResult(this, static_cast<int>(entryPointIndex), static_cast<int>(targetIndex), artifact)))
         return SLANG_E_INVALID_ARG;
 
+    if (!artifact)
+        return SLANG_E_NOT_AVAILABLE;
+
     // Find a rep
     auto metadata = findAssociatedRepresentation<IArtifactPostEmitMetadata>(artifact);
     if (!metadata)
diff --git a/tests/metal/simple-compute.slang b/tests/metal/simple-compute.slang
index b5c0aaf75..85abaffb4 100644
--- a/tests/metal/simple-compute.slang
+++ b/tests/metal/simple-compute.slang
@@ -1,6 +1,6 @@
 //TEST:SIMPLE(filecheck=CHECK): -target metal
 //TEST:SIMPLE(filecheck=CHECK-ASM): -target metallib
-//TEST:REFLECTION(filecheck=REFLECT):-target metal
+//TEST:REFLECTION(filecheck=REFLECT):-target metal -entry main_kernel -stage compute
 
 uniform RWStructuredBuffer<float> outputBuffer;
 
@@ -26,10 +26,10 @@ ParameterBlock<MyBlock> block2;
 // REFLECT: "binding": {"kind": "metalArgumentBufferElement", "index": 0, "count": 2}
 
 // REFLECT: "name": "outputBuffer",
-// REFLECT-NEXT: "binding": {"kind": "constantBuffer", "index": 0, "used": 0}
+// REFLECT-NEXT: "binding": {"kind": "constantBuffer", "index": 0{{.*}}}
 
 // REFLECT: "name": "block",
-// REFLECT-NEXT: "binding": {"kind": "constantBuffer", "index": 1, "used": 0}
+// REFLECT-NEXT: "binding": {"kind": "constantBuffer", "index": 1{{.*}}}
 
 void func(float v)
 {
diff --git a/tests/metal/stage-in.slang b/tests/metal/stage-in.slang
new file mode 100644
index 000000000..31d224072
--- /dev/null
+++ b/tests/metal/stage-in.slang
@@ -0,0 +1,89 @@
+//TEST:SIMPLE(filecheck=CHECK): -target metal
+//TEST:SIMPLE(filecheck=CHECK-ASM): -target metallib
+
+// CHECK:     struct VOut{{.*}}
+// CHECK-NEXT:{
+// CHECK-NEXT:    float4 position{{.*}} {{\[\[}}position{{\]\]}};
+// CHECK-NEXT:    float4 vertexColor{{.*}} {{\[\[}}user(_slang_attr){{\]\]}};
+// CHECK-NEXT:    float2 vertexUV{{.*}} {{\[\[}}user(_slang_attr_1){{\]\]}};
+// CHECK-NEXT:    float3 vertexNormal{{.*}} {{\[\[}}user(NORMAL){{\]\]}};
+// CHECK-NEXT:};
+
+// CHECK:     struct vertexInput{{.*}}
+// CHECK-NEXT:{
+// CHECK-NEXT:    float4 position{{.*}} {{\[\[}}attribute(0){{\]\]}} {{\[\[}}user(POSITION){{\]\]}};
+// CHECK-NEXT:    float4 color{{.*}} {{\[\[}}attribute(1){{\]\]}} {{\[\[}}user(COLOR){{\]\]}};
+// CHECK-NEXT:};
+
+// CHECK: {{\[\[}}vertex{{\]\]}} VOut{{.*}} main_vertex(vertexInput{{.*}}{{\[\[}}stage_in{{\]\]}}, uint vid{{.*}}{{\[\[}}vertex_id{{\]\]}}, uint instanceID{{.*}} {{\[\[}}instance_id{{\]\]}})
+
+// CHECK:     struct pixelOutput{{.*}}
+// CHECK-NEXT:{
+// CHECK-NEXT:    float4 output{{.*}}{{\[\[}}color(0){{\]\]}};
+// CHECK-NEXT:};
+
+// CHECK:     struct pixelInput{{.*}}
+// CHECK-NEXT:{
+// CHECK-NEXT:    float4 vertexColor{{.*}}  {{\[\[}}user(_slang_attr){{\]\]}};
+// CHECK-NEXT:    float2 vertexUV{{.*}} {{\[\[}}user(_slang_attr_1){{\]\]}};
+// CHECK-NEXT:    float3 vertexNormal{{.*}} {{\[\[}}user(NORMAL){{\]\]}};
+// CHECK-NEXT:};
+
+// CHECK: {{\[\[}}fragment{{\]\]}} pixelOutput{{.*}} main_fragment(pixelInput{{.*}} {{\[\[}}stage_in{{\]\]}}, float4 position{{.*}} {{\[\[}}position{{\]\]}})
+
+// CHECK:     struct FragOut{{.*}}
+// CHECK-NEXT:{
+// CHECK-NEXT:    float4 color{{.*}}{{\[\[}}color(0){{\]\]}};
+// CHECK-NEXT:    float depth{{.*}} {{\[\[}}depth(any){{\]\]}};
+// CHECK-NEXT:};
+
+// CHECK-ASM: define {{.*}} @main_vertex
+// CHECK-ASM: define {{.*}} @main_fragment
+// CHECK-ASM: define {{.*}} @main_fragment1
+
+struct VIn
+{
+    float4 position : POSITION;
+    float4 color : COLOR;
+    uint vid : SV_VertexID;
+    uint instanceID : SV_InstanceID;
+}
+
+struct VOut
+{
+    float4 position : SV_Position;
+    float4 vertexColor;
+    float2 vertexUV;
+    float3 vertexNormal : NORMAL;
+}
+
+[shader("vertex")]
+VOut main_vertex(VIn vertexIn)
+{
+    VOut vertexOut;
+    vertexOut.position = vertexIn.position;
+    vertexOut.vertexColor = vertexIn.color;
+    vertexOut.vertexUV = float2(0.0, 1.0);
+    return vertexOut;
+}
+
+[shader("fragment")]
+float4 main_fragment(VOut fragmentIn) : SV_Target
+{
+    return fragmentIn.vertexColor;
+}
+
+struct FragOut
+{
+    float4 color : SV_Target;
+    float depth : SV_Depth;
+}
+
+[shader("fragment")]
+FragOut main_fragment1(VOut fragmentIn)
+{
+    FragOut fragOut;
+    fragOut.color = fragmentIn.vertexColor;
+    fragOut.depth = 0.5;
+    return fragOut;
+}
\ No newline at end of file