4 files changed, 93 insertions, 6 deletions

diff --git a/source/slang/slang-ir-wgsl-legalize.cpp b/source/slang/slang-ir-wgsl-legalize.cpp
index 8eb820b4d..4d5f6e4ba 100644
--- a/source/slang/slang-ir-wgsl-legalize.cpp
+++ b/source/slang/slang-ir-wgsl-legalize.cpp
@@ -147,6 +147,9 @@ struct LegalizeWGSLEntryPointContext
                     semanticInfoToRemove);
                 // Validate/rearange all semantics which overlap in our flat struct.
                 fixFieldSemanticsOfFlatStruct(flattenedStruct);
+                ensureStructHasUserSemantic<LayoutResourceKind::VaryingInput>(
+                    flattenedStruct,
+                    layout);
                 if (flattenedStruct != structType)
                 {
                     // Replace the 'old IRParam type' with a 'new IRParam type'
@@ -381,7 +384,8 @@ struct LegalizeWGSLEntryPointContext
             semanticName);
     }
 
-    void ensureResultStructHasUserSemantic(IRStructType* structType, IRVarLayout* varLayout)
+    template<LayoutResourceKind K>
+    void ensureStructHasUserSemantic(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.
@@ -415,11 +419,10 @@ struct LegalizeWGSLEntryPointContext
             }
             typeLayout->getFieldLayout(index);
             auto fieldLayout = typeLayout->getFieldLayout(index);
-            if (auto offsetAttr = fieldLayout->findOffsetAttr(LayoutResourceKind::VaryingOutput))
+            if (auto offsetAttr = fieldLayout->findOffsetAttr(K))
             {
                 UInt varOffset = 0;
-                if (auto varOffsetAttr =
-                        varLayout->findOffsetAttr(LayoutResourceKind::VaryingOutput))
+                if (auto varOffsetAttr = varLayout->findOffsetAttr(K))
                     varOffset = varOffsetAttr->getOffset();
                 varOffset += offsetAttr->getOffset();
                 builder.addSemanticDecoration(key, toSlice("_slang_attr"), (int)varOffset);
@@ -1101,7 +1104,9 @@ struct LegalizeWGSLEntryPointContext
             }
             // Ensure non-overlapping semantics
             fixFieldSemanticsOfFlatStruct(flattenedStruct);
-            ensureResultStructHasUserSemantic(flattenedStruct, resultLayout);
+            ensureStructHasUserSemantic<LayoutResourceKind::VaryingOutput>(
+                flattenedStruct,
+                resultLayout);
             return;
         }
 
@@ -1121,7 +1126,7 @@ struct LegalizeWGSLEntryPointContext
         auto typeLayout = structTypeLayoutBuilder.build();
         IRVarLayout::Builder varLayoutBuilder(&builder, typeLayout);
         auto varLayout = varLayoutBuilder.build();
-        ensureResultStructHasUserSemantic(structType, varLayout);
+        ensureStructHasUserSemantic<LayoutResourceKind::VaryingOutput>(structType, varLayout);
 
         _replaceAllReturnInst(
             builder,
diff --git a/tests/render/multiple-stage-io-locations-without-user-semantics.slang b/tests/render/multiple-stage-io-locations-without-user-semantics.slang
new file mode 100644
index 000000000..af0e3e39f
--- /dev/null
+++ b/tests/render/multiple-stage-io-locations-without-user-semantics.slang
@@ -0,0 +1,77 @@
+// TODO: Investigate failures on non-WebGPU backends
+//TEST(smoke,render):COMPARE_HLSL_RENDER: -wgpu
+
+cbuffer Uniforms
+{
+    float4x4 modelViewProjection;
+}
+
+struct AssembledVertex
+{
+    float3 position;
+    float3 color;
+};
+
+struct Fragment
+{
+    float4 color;
+};
+
+// Vertex  Shader
+
+struct VertexStageInput
+{
+    AssembledVertex assembledVertex : A;
+};
+
+struct VertexStageOutput
+{
+    float3 color;
+    float3 localPosition;
+    float4 sv_position     : SV_Position;
+};
+
+VertexStageOutput vertexMain(VertexStageInput input)
+{
+    VertexStageOutput output;
+
+    float3 position = input.assembledVertex.position;
+    float3 color    = input.assembledVertex.color;
+
+    output.color         = color;
+    output.sv_position   = mul(modelViewProjection, float4(position, 1.0));
+    output.localPosition = position;
+
+    return output;
+}
+
+// Fragment Shader
+
+struct FragmentStageInput
+{
+    float3 color;
+    float3 localPosition;
+};
+
+struct FragmentStageOutput
+{
+    Fragment fragment : SV_Target;
+};
+
+FragmentStageOutput fragmentMain(FragmentStageInput input)
+{
+    FragmentStageOutput output;
+
+    float3 color = input.color;
+
+    if (input.color.y < input.color.z)
+    {
+        output.fragment.color = float4(input.localPosition, 1.0);
+    }
+    else
+    {
+        output.fragment.color = float4(input.color, 1.0);
+    }
+
+    return output;
+}
diff --git a/tests/render/multiple-stage-io-locations-without-user-semantics.slang.expected b/tests/render/multiple-stage-io-locations-without-user-semantics.slang.expected
new file mode 100644
index 000000000..4c32e2510
--- /dev/null
+++ b/tests/render/multiple-stage-io-locations-without-user-semantics.slang.expected
@@ -0,0 +1,5 @@
+result code = 0
+standard error = {
+}
+standard output = {
+}
diff --git a/tests/render/multiple-stage-io-locations-without-user-semantics.slang.expected.png b/tests/render/multiple-stage-io-locations-without-user-semantics.slang.expected.png
new file mode 100644
index 000000000..3333a12d7
--- /dev/null
+++ b/tests/render/multiple-stage-io-locations-without-user-semantics.slang.expected.png