wgsl: Fix field semantics even if input struct doesn't need flattening (#5642)

* wgsl: Fix field semantics even if input struct doesn't need flattening

Helps to address issue #5633.

* Add test for multiple stage IO locations

This verifies part of issue #5633.

9 files changed, 96 insertions, 3 deletions

diff --git a/source/slang/slang-ir-wgsl-legalize.cpp b/source/slang/slang-ir-wgsl-legalize.cpp
index 4d617a169..8eb820b4d 100644
--- a/source/slang/slang-ir-wgsl-legalize.cpp
+++ b/source/slang/slang-ir-wgsl-legalize.cpp
@@ -145,11 +145,10 @@ struct LegalizeWGSLEntryPointContext
                     structType,
                     mapOldFieldToNewField,
                     semanticInfoToRemove);
+                // Validate/rearange all semantics which overlap in our flat struct.
+                fixFieldSemanticsOfFlatStruct(flattenedStruct);
                 if (flattenedStruct != structType)
                 {
-                    // Validate/rearange all semantics which overlap in our flat struct
-                    fixFieldSemanticsOfFlatStruct(flattenedStruct);
-
                     // Replace the 'old IRParam type' with a 'new IRParam type'
                     param->setFullType(flattenedStruct);
 
diff --git a/tests/expected-failure-github.txt b/tests/expected-failure-github.txt
index fdb6127a2..671b8dafd 100644
--- a/tests/expected-failure-github.txt
+++ b/tests/expected-failure-github.txt
@@ -5,6 +5,7 @@ tests/language-feature/saturated-cooperation/fuse-product.slang (vk)
 tests/language-feature/saturated-cooperation/fuse.slang (vk)
 tests/bugs/byte-address-buffer-interlocked-add-f32.slang (vk)
 tests/render/render0.hlsl (mtl)
+tests/render/multiple-stage-io-locations.slang (mtl)
 tests/render/nointerpolation.hlsl (mtl)
 tests/serialization/obfuscated-serialized-module-test.slang.2 syn (mtl)
 tests/autodiff/custom-intrinsic.slang.2 syn (wgpu)
diff --git a/tests/render/multiple-stage-io-locations.slang b/tests/render/multiple-stage-io-locations.slang
new file mode 100644
index 000000000..9f74d1398
--- /dev/null
+++ b/tests/render/multiple-stage-io-locations.slang
@@ -0,0 +1,78 @@
+//TEST(smoke,render):COMPARE_HLSL_RENDER:
+// TODO: Investigate Metal failure
+//DISABLE_TEST(smoke,render):COMPARE_HLSL_RENDER: -mtl
+
+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           : VERTEX_COLOR;
+    float3 localPosition   : VERTEX_LOCAL_POSITION;
+    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         : VERTEX_COLOR;
+    float3 localPosition : VERTEX_LOCAL_POSITION;
+};
+
+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.slang.1.expected b/tests/render/multiple-stage-io-locations.slang.1.expected
new file mode 100644
index 000000000..4c32e2510
--- /dev/null
+++ b/tests/render/multiple-stage-io-locations.slang.1.expected
@@ -0,0 +1,5 @@
+result code = 0
+standard error = {
+}
+standard output = {
+}
diff --git a/tests/render/multiple-stage-io-locations.slang.1.expected.png b/tests/render/multiple-stage-io-locations.slang.1.expected.png
new file mode 100644
index 000000000..3333a12d7
--- /dev/null
+++ b/tests/render/multiple-stage-io-locations.slang.1.expected.png
diff --git a/tests/render/multiple-stage-io-locations.slang.2.expected b/tests/render/multiple-stage-io-locations.slang.2.expected
new file mode 100644
index 000000000..4c32e2510
--- /dev/null
+++ b/tests/render/multiple-stage-io-locations.slang.2.expected
@@ -0,0 +1,5 @@
+result code = 0
+standard error = {
+}
+standard output = {
+}
diff --git a/tests/render/multiple-stage-io-locations.slang.2.expected.png b/tests/render/multiple-stage-io-locations.slang.2.expected.png
new file mode 100644
index 000000000..3333a12d7
--- /dev/null
+++ b/tests/render/multiple-stage-io-locations.slang.2.expected.png
diff --git a/tests/render/multiple-stage-io-locations.slang.expected b/tests/render/multiple-stage-io-locations.slang.expected
new file mode 100644
index 000000000..4c32e2510
--- /dev/null
+++ b/tests/render/multiple-stage-io-locations.slang.expected
@@ -0,0 +1,5 @@
+result code = 0
+standard error = {
+}
+standard output = {
+}
diff --git a/tests/render/multiple-stage-io-locations.slang.expected.png b/tests/render/multiple-stage-io-locations.slang.expected.png
new file mode 100644
index 000000000..3333a12d7
--- /dev/null
+++ b/tests/render/multiple-stage-io-locations.slang.expected.png