5 files changed, 48 insertions, 1 deletions
diff --git a/source/slang/ast-legalize.cpp b/source/slang/ast-legalize.cpp
index afc8b31c8..1b045cbc9 100644
--- a/source/slang/ast-legalize.cpp
+++ b/source/slang/ast-legalize.cpp
@@ -3309,6 +3309,7 @@ struct LoweringVisitor
                     RefPtr<VarLayout> newFieldLayout = new VarLayout();
                     newFieldLayout->typeLayout = fieldLayout->typeLayout;
                     newFieldLayout->flags = fieldLayout->flags;
+                    newFieldLayout->stage = fieldLayout->stage;
                     newFieldLayout->varDecl = fieldLayout->varDecl;
                     newFieldLayout->systemValueSemantic = fieldLayout->systemValueSemantic;
                     newFieldLayout->systemValueSemanticIndex = fieldLayout->systemValueSemanticIndex;
diff --git a/source/slang/ir-legalize-types.cpp b/source/slang/ir-legalize-types.cpp
index 37df5f698..7fad2ccc5 100644
--- a/source/slang/ir-legalize-types.cpp
+++ b/source/slang/ir-legalize-types.cpp
@@ -1620,6 +1620,8 @@ static LegalVal declareSimpleVar(
         // well as all the offset information that has accumulated
         // along the chain of parent variables.
 
+        // TODO: this logic needs to propagate through semantics...
+
         varLayout = new VarLayout();
         varLayout->typeLayout = typeLayout;
 
diff --git a/source/slang/parameter-binding.cpp b/source/slang/parameter-binding.cpp
index b97c174fc..66f7d437c 100644
--- a/source/slang/parameter-binding.cpp
+++ b/source/slang/parameter-binding.cpp
@@ -627,6 +627,7 @@ struct EntryPointParameterState
     int*                                ioSemanticIndex = nullptr;
     EntryPointParameterDirectionMask    directionMask;
     int                                 semanticSlotCount;
+    Stage                               stage = Stage::Unknown;
 };
 
 
@@ -639,7 +640,7 @@ static RefPtr<TypeLayout> processEntryPointParameter(
 static void collectGlobalScopeGLSLVaryingParameter(
     ParameterBindingContext*        context,
     RefPtr<VarDeclBase>             varDecl,
-    RefPtr<Type>          effectiveType,
+    RefPtr<Type>                    effectiveType,
     EntryPointParameterDirection    direction)
 {
     int defaultSemanticIndex = 0;
@@ -647,6 +648,7 @@ static void collectGlobalScopeGLSLVaryingParameter(
     EntryPointParameterState state;
     state.directionMask = direction;
     state.ioSemanticIndex = &defaultSemanticIndex;
+    state.stage = context->stage;
 
     RefPtr<VarLayout> varLayout = new VarLayout();
     varLayout->varDecl = makeDeclRef(varDecl.Ptr());
@@ -1333,6 +1335,10 @@ static RefPtr<TypeLayout> processEntryPointParameter(
         varLayout->flags |= VarLayoutFlag::HasSemantic;
     }
 
+    if (varLayout)
+    {
+        varLayout->stage = state.stage;
+    }
 
     // Scalar and vector types are treated as outputs directly
     if(auto basicType = type->As<BasicExpressionType>())
@@ -1485,6 +1491,7 @@ static void collectEntryPointParameters(
     state.ioSemanticIndex = &defaultSemanticIndex;
     state.optSemanticName = nullptr;
     state.semanticSlotCount = 0;
+    state.stage = entryPoint->profile.GetStage();
 
     for( auto m : entryPointFuncDecl->Members )
     {
@@ -2020,6 +2027,7 @@ RefPtr<ProgramLayout> specializeProgramLayout(
             RefPtr<VarLayout> newVarLayout = new VarLayout();
             RefPtr<ParameterInfo> paramInfo = new ParameterInfo();
             newVarLayout->varDecl = varLayout->varDecl;
+            newVarLayout->stage = varLayout->stage;
             newVarLayout->typeLayout = newTypeLayout;
             paramInfo->varLayouts.Add(newVarLayout);
             completeBindingsForParameter(&context, paramInfo);
diff --git a/source/slang/reflection.cpp b/source/slang/reflection.cpp
index 14199f126..a1ea3ff72 100644
--- a/source/slang/reflection.cpp
+++ b/source/slang/reflection.cpp
@@ -666,6 +666,36 @@ SLANG_API size_t spReflectionVariableLayout_GetSemanticIndex(SlangReflectionVari
     return varLayout->semanticIndex;
 }
 
+SLANG_API SlangStage spReflectionVariableLayout_getStage(
+    SlangReflectionVariableLayout* inVarLayout)
+{
+    auto varLayout = convert(inVarLayout);
+    if(!varLayout) return SLANG_STAGE_NONE;
+
+    // A parameter that is not a varying input or output is
+    // not considered to belong to a single stage.
+    //
+    // TODO: We might need to reconsider this for, e.g., entry
+    // point parameters, where they might be stage-specific even
+    // if they are uniform.
+    if (!varLayout->FindResourceInfo(Slang::LayoutResourceKind::VaryingInput)
+        && !varLayout->FindResourceInfo(Slang::LayoutResourceKind::VaryingOutput))
+    {
+        return SLANG_STAGE_NONE;
+    }
+
+    // TODO: We should find the stage for a variable layout by
+    // walking up the tree of layout information, until we find
+    // something that has a definitive stage attached to it (e.g.,
+    // either an entry point or a GLSL translation unit).
+    //
+    // We don't currently have parent links in the reflection layout
+    // information, so doing that walk would be tricky right now, so
+    // it is easier to just bloat the representation and store yet another
+    // field on every variable layout.
+    return (SlangStage) varLayout->stage;
+}
+
 
 // Shader Parameter Reflection
 
diff --git a/source/slang/type-layout.h b/source/slang/type-layout.h
index 07530bdfc..66763c8a7 100644
--- a/source/slang/type-layout.h
+++ b/source/slang/type-layout.h
@@ -251,6 +251,12 @@ public:
     String  semanticName;
     int     semanticIndex;
 
+    // The stage this variable belongs to, in case it is
+    // stage-specific.
+    // TODO: This is wasteful to be storing on every single
+    // variable layout.
+    Stage stage = Stage::Unknown;
+
     // The start register(s) for any resources
     struct ResourceInfo
     {