From 91430870a8a6ec2825969823b892a9c8a0d588ab Mon Sep 17 00:00:00 2001
From: Darren Wihandi <65404740+fairywreath@users.noreply.github.com>
Date: Fri, 17 Jan 2025 11:57:06 -0500
Subject: [PATCH] hoist entry point params for wgsl (#6116)

Co-authored-by: Yong He <yonghe@outlook.com>
---
 .../slang-ir-legalize-varying-params.cpp      | 20 ++++++-------------
 .../metal/nested-struct-fragment-input.slang  |  4 ++--
 tests/metal/stage-in-2.slang                  |  6 ++----
 tests/wgsl/nested-varying-input.slang         | 20 ++++++++++---------
 4 files changed, 21 insertions(+), 29 deletions(-)

diff --git a/source/slang/slang-ir-legalize-varying-params.cpp b/source/slang/slang-ir-legalize-varying-params.cpp
index 9720084ba3..3b65ee59af 100644
--- a/source/slang/slang-ir-legalize-varying-params.cpp
+++ b/source/slang/slang-ir-legalize-varying-params.cpp
@@ -1574,8 +1574,8 @@ class LegalizeShaderEntryPointContext
     }
 
 protected:
-    LegalizeShaderEntryPointContext(IRModule* module, DiagnosticSink* sink, bool hoistParameters)
-        : m_module(module), m_sink(sink), hoistParameters(hoistParameters)
+    LegalizeShaderEntryPointContext(IRModule* module, DiagnosticSink* sink)
+        : m_module(module), m_sink(sink)
     {
     }
 
@@ -1758,7 +1758,6 @@ class LegalizeShaderEntryPointContext
     }
 
 private:
-    const bool hoistParameters;
     HashSet<IRStructField*> semanticInfoToRemove;
 
     void removeSemanticLayoutsFromLegalizedStructs()
@@ -2985,16 +2984,9 @@ class LegalizeShaderEntryPointContext
         // If the entrypoint is receiving varying inputs as a pointer, turn it into a value.
         depointerizeInputParams(entryPoint.entryPointFunc);
 
-        // TODO FIXME: Enable these for WGSL and remove the `hoistParemeters` member field.
-        // WGSL entry point legalization currently only applies attributes to struct parameters,
-        // apply the same hoisting from Metal to WGSL to fix it.
-        if (hoistParameters)
-        {
-            hoistEntryPointParameterFromStruct(entryPoint);
-            packStageInParameters(entryPoint);
-        }
-
         // Input Parameter Legalize
+        hoistEntryPointParameterFromStruct(entryPoint);
+        packStageInParameters(entryPoint);
         flattenInputParameters(entryPoint);
 
         // System Value Legalize
@@ -3023,7 +3015,7 @@ class LegalizeMetalEntryPointContext : public LegalizeShaderEntryPointContext
 {
 public:
     LegalizeMetalEntryPointContext(IRModule* module, DiagnosticSink* sink)
-        : LegalizeShaderEntryPointContext(module, sink, true)
+        : LegalizeShaderEntryPointContext(module, sink)
     {
         generatePermittedTypes_sv_target();
     }
@@ -3681,7 +3673,7 @@ class LegalizeWGSLEntryPointContext : public LegalizeShaderEntryPointContext
 {
 public:
     LegalizeWGSLEntryPointContext(IRModule* module, DiagnosticSink* sink)
-        : LegalizeShaderEntryPointContext(module, sink, false)
+        : LegalizeShaderEntryPointContext(module, sink)
     {
     }
 
diff --git a/tests/metal/nested-struct-fragment-input.slang b/tests/metal/nested-struct-fragment-input.slang
index f70303306f..77530f9fd5 100644
--- a/tests/metal/nested-struct-fragment-input.slang
+++ b/tests/metal/nested-struct-fragment-input.slang
@@ -22,7 +22,7 @@
 // METAL-NOT: [[ATTR3]]
 // METAL-DAG: [[ATTR4:COARSEVERTEX_(1|2|3|4)]]
 
-// WGSL:struct FragmentStageInput
+// WGSL:struct pixelInput
 // WGSL-DAG:@location(0) [[VAR0:[A-Za-z_0-9]+]]
 // WGSL-DAG:@location(1) [[VAR1:[A-Za-z_0-9]+]]
 // WGSL-DAG:@location(2) [[VAR2:[A-Za-z_0-9]+]]
@@ -78,7 +78,7 @@ float4 fragmentMain(FragmentStageInput input)
     // METAL-DAG: {{.*}}->p3{{.*}}->p2{{.*}}->p1{{.*}}=
     // METAL-DAG: {{.*}}->p3{{.*}}->p3{{.*}}->p1{{.*}}=
 
-    // WGSL: var [[UnpackedInput:[A-Za-z_0-9]+]] : FragmentStageInput
+    // WGSL: var [[UnpackedInput:[A-Za-z_0-9]+]] : pixelInput
     // WGSL-DAG: [[UnpackedInput]].{{[A-Za-z_0-9]+}}.{{[A-Za-z_0-9]+}} = [[InputVar]].[[VAR7]];
 
     // WGSL-DAG: [[UnpackedInput]].{{[A-Za-z_0-9]+}}.{{[A-Za-z_0-9]+}}.{{[A-Za-z_0-9]+}} = [[InputVar]].[[VAR6]];
diff --git a/tests/metal/stage-in-2.slang b/tests/metal/stage-in-2.slang
index 5b885af492..6ee3b8c512 100644
--- a/tests/metal/stage-in-2.slang
+++ b/tests/metal/stage-in-2.slang
@@ -15,10 +15,8 @@
 //WGSLSPIRV: %vertexMain = OpFunction %
 //WGSLSPIRV: %fragmentMain = OpFunction %
 
-//WGSL: struct [[CoarseVertex:CoarseVertex[_0-9]*]]
-//WGSL-NEXT: {
-//WGSL-NEXT:     @location(0) color
-//WGSL: fn fragmentMain({{.*}}[[CoarseVertex]]
+//WGSL:     @location(0) output
+//WGSL:     @location(0) color
 
 // Uniform data to be passed from application -> shader.
 cbuffer Uniforms
diff --git a/tests/wgsl/nested-varying-input.slang b/tests/wgsl/nested-varying-input.slang
index 2cdf4f7ebd..d2566a6b90 100644
--- a/tests/wgsl/nested-varying-input.slang
+++ b/tests/wgsl/nested-varying-input.slang
@@ -22,17 +22,17 @@ struct NestedVertexOutput
 
 struct VertexOutput
 {
-    //VERT: @builtin(position) position
-    //FRAG: @builtin(position) position
-    float4 position : SV_Position;
-
     //VERT: @location(0) uv
-    //FRAG: @location(0) uv
+    //FRAG-DAG: @location(0) uv
     float2 uv : TEXCOORD0;
 
     //VERT: @location(1) color
-    //FRAG: @location(1) color
+    //FRAG-DAG: @location(1) color
     NestedVertexOutput nested;
+
+    //VERT: @builtin(position) position
+    //FRAG-DAG: @builtin(position) position
+    float4 position : SV_Position;
 };
 
 VertexOutput vertexMain()
@@ -44,10 +44,12 @@ VertexOutput vertexMain()
     return out;
 }
 
-FragmentOutput fragmentMain(VertexOutput input)
+//FRAG-DAG: @location(3) color3
+//FRAG-DAG: @location(6) color6
+FragmentOutput fragmentMain(VertexOutput input, float4 color3: COLOR3, float4 color6: COLOR6)
 {
     FragmentOutput out;
-    out.color0 = input.nested.color;
-    out.color1 = input.nested.color;
+    out.color0 = input.nested.color + color3;
+    out.color1 = input.nested.color + color6;
     return out;
 }