From 1e7f5418c293d3a9ca91ae4648ca2d522ec2ebe7 Mon Sep 17 00:00:00 2001
From: Yong He <yonghe@outlook.com>
Date: Tue, 19 Nov 2024 14:43:15 -0800
Subject: Fix wgsl legalization around __ref parameter. (#5597)

* Fix wgsl legalization around __ref parameter.

* Add intrinsic and test case.
---
 source/slang/hlsl.meta.slang            | 19 +++++++++++++++++++
 source/slang/slang-ir-wgsl-legalize.cpp |  2 ++
 tests/wgsl/workgroup-uniform-load.slang | 25 +++++++++++++++++++++++++
 3 files changed, 46 insertions(+)
 create mode 100644 tests/wgsl/workgroup-uniform-load.slang

diff --git a/source/slang/hlsl.meta.slang b/source/slang/hlsl.meta.slang
index 14f9f78f2..395ea5012 100644
--- a/source/slang/hlsl.meta.slang
+++ b/source/slang/hlsl.meta.slang
@@ -20645,3 +20645,22 @@ extension<T, L : IBufferDataLayout> RasterizerOrderedStructuredBuffer<T, L> : IR
 {
     int getCount() { uint count; uint stride; this.GetDimensions(count, stride); return count; }
 }
+
+//@public:
+
+/// Mark a variable as being workgroup uniform.
+/// @param v The variable to mark as workgroup uniform.
+/// @return The value of `v`.
+/// @remarks This intrinsic maps to `workgroupUniformLoad` when targeting WGSL and is a no-op on other targets.
+/// WGSL is strict on uniformity, and this intrinsic is needed to mark a variable as workgroup uniform in order
+// to silence uniformity errors in certain cases.
+T workgroupUniformLoad<T>(__ref T v)
+{
+    __target_switch
+    {
+    case wgsl:
+        __intrinsic_asm "workgroupUniformLoad(&($0))";
+    default:
+        return v;
+    }
+}
\ No newline at end of file
diff --git a/source/slang/slang-ir-wgsl-legalize.cpp b/source/slang/slang-ir-wgsl-legalize.cpp
index adb027668..4d617a169 100644
--- a/source/slang/slang-ir-wgsl-legalize.cpp
+++ b/source/slang/slang-ir-wgsl-legalize.cpp
@@ -1398,6 +1398,8 @@ struct LegalizeWGSLEntryPointContext
             {
             case kIROp_Var:
             case kIROp_Param:
+            case kIROp_GlobalParam:
+            case kIROp_GlobalVar:
                 continue;
             default:
                 break;
diff --git a/tests/wgsl/workgroup-uniform-load.slang b/tests/wgsl/workgroup-uniform-load.slang
new file mode 100644
index 000000000..3ace37826
--- /dev/null
+++ b/tests/wgsl/workgroup-uniform-load.slang
@@ -0,0 +1,25 @@
+// groupshared-ref-param.slang
+
+//TEST:SIMPLE(filecheck=CHECK): -target wgsl -entry computeMain -stage compute
+
+
+//TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBuffer
+RWStructuredBuffer<int> outputBuffer;
+
+groupshared uint sharedVal;
+
+// Expect sharedVal to be passed by reference.
+//
+// CHECK: fn computeMain({{.*}}({{.*}})
+// CHECK: {{.*}}workgroupUniformLoad(&((sharedVal_{{[a-zA-Z0-9_]*}}))){{.*}};
+// CHECK: }
+
+
+[numthreads(1, 1, 1)]
+void computeMain(int3 dispatchThreadID: SV_DispatchThreadID)
+{
+    int idx = dispatchThreadID.x;
+
+    sharedVal = 1;
+    outputBuffer[0] = workgroupUniformLoad(sharedVal);
+}
\ No newline at end of file
-- 
cgit v1.2.3