Add WaveGetLane* support for Metal and WGSL (#6371)

* support WaveGetLane* for WGSL and Metal

* update test and glsl support

* address review comments and fix metal test

* add missing pragma guard

* update test

* Revert "update test"

This reverts commit f2b97e91c29de154190710580c343bd0764aedbb.

* update failing glsl metal test and added new test

* make hlsl and glsl outputs similar

* update test

* disable tests for Metal and cleanup

* comment fix

* add expected failures

* correct expected failures list

* remove expected failure

* add tests to expected failure

---------

Co-authored-by: Yong He <yonghe@outlook.com>

5 files changed, 56 insertions, 12 deletions

diff --git a/tests/expected-failure-github.txt b/tests/expected-failure-github.txt
index e9f634ea8..10897b31e 100644
--- a/tests/expected-failure-github.txt
+++ b/tests/expected-failure-github.txt
@@ -10,4 +10,6 @@ tests/autodiff/custom-intrinsic.slang.2 syn (wgpu)
 tests/bugs/buffer-swizzle-store.slang.3 syn (wgpu)
 tests/compute/interface-shader-param-in-struct.slang.4 syn (wgpu)
 tests/compute/interface-shader-param.slang.5 syn (wgpu)
-tests/language-feature/shader-params/interface-shader-param-ordinary.slang.4 syn (wgpu)
\ No newline at end of file
+tests/language-feature/shader-params/interface-shader-param-ordinary.slang.4 syn (wgpu)
+tests/glsl-intrinsic/shader-subgroup/shader-subgroup-builtin-variables.slang.8 (mtl)
+tests/glsl-intrinsic/shader-subgroup/shader-subgroup-builtin-variables-2.slang.3 (mtl)
diff --git a/tests/glsl-intrinsic/shader-subgroup/shader-subgroup-builtin-variables-2.slang b/tests/glsl-intrinsic/shader-subgroup/shader-subgroup-builtin-variables-2.slang
new file mode 100644
index 000000000..2e3896cc5
--- /dev/null
+++ b/tests/glsl-intrinsic/shader-subgroup/shader-subgroup-builtin-variables-2.slang
@@ -0,0 +1,36 @@
+//TEST(compute, vulkan):COMPARE_COMPUTE_EX(filecheck-buffer=BUF):-vk -compute -entry computeMain -allow-glsl
+//TEST(compute, vulkan):COMPARE_COMPUTE_EX(filecheck-buffer=BUF):-vk -compute -entry computeMain -allow-glsl -emit-spirv-directly
+//TEST(compute):COMPARE_COMPUTE_EX(filecheck-buffer=BUF):-wgpu -compute -entry computeMain -allow-glsl
+
+
+// There are some issues with the Metal backend when using glsl-style syntax - this should be fixed.
+//TEST(compute):COMPARE_COMPUTE_EX(filecheck-buffer=BUF):-metal -compute -entry computeMain -allow-glsl
+
+#version 430
+
+//TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBuffer
+buffer MyBlockName2
+{
+    int data[];
+} outputBuffer;
+
+layout(local_size_x = 4) in;
+
+void computeMain(int3 dispatchThreadID : SV_DispatchThreadID)
+{
+    // There may be some issues with structure padding for global context containing
+    // global builtin variables.
+    // int idx = gl_GlobalInvocationID.x;
+    int idx = dispatchThreadID.x;
+
+    uint laneId = gl_SubgroupInvocationID;
+    // The laneCount will be dependent on target hardware. It seems a count of 1 is valid in spec. 
+    // For now we'll just check it's not 0.
+    uint laneCount = gl_SubgroupSize;
+    outputBuffer.data[idx] = int(((laneCount > 0) ? 0x100 : 0) + laneId);
+
+    // BUF: 100
+    // BUF-NEXT: 101
+    // BUF-NEXT: 102
+    // BUF-NEXT: 103
+}
diff --git a/tests/glsl-intrinsic/shader-subgroup/shader-subgroup-builtin-variables.slang b/tests/glsl-intrinsic/shader-subgroup/shader-subgroup-builtin-variables.slang
index 21b533178..2d11ca5fb 100644
--- a/tests/glsl-intrinsic/shader-subgroup/shader-subgroup-builtin-variables.slang
+++ b/tests/glsl-intrinsic/shader-subgroup/shader-subgroup-builtin-variables.slang
@@ -1,5 +1,5 @@
-//TEST:SIMPLE(filecheck=CHECK_GLSL):  -allow-glsl -stage compute -entry computeMain -target glsl
-//TEST:SIMPLE(filecheck=CHECK_SPV):  -allow-glsl -stage compute -entry computeMain -target spirv -emit-spirv-directly
+//TEST:SIMPLE(filecheck=CHECK_GLSL):  -allow-glsl -stage compute -entry computeMain -target glsl -DTARGET_VK
+//TEST:SIMPLE(filecheck=CHECK_SPV):  -allow-glsl -stage compute -entry computeMain -target spirv -emit-spirv-directly -DTARGET_VK
 
 // missing implementation of most builtin values due to non trivial translation
 //DISABLE_TEST:SIMPLE(filecheck=CHECK_HLSL): -allow-glsl -stage compute -entry computeMain -target hlsl -DTARGET_HLSL
@@ -8,8 +8,13 @@
 //missing implementation of system (varying?) values
 //DISABLE_TEST:SIMPLE(filecheck=CHECK_CPP):  -allow-glsl -stage compute -entry computeMain -target cpp -DTARGET_CPP
 
-//TEST(compute, vulkan):COMPARE_COMPUTE(filecheck-buffer=BUF):-vk -compute -entry computeMain -allow-glsl
-//TEST(compute, vulkan):COMPARE_COMPUTE(filecheck-buffer=BUF):-vk -compute -entry computeMain -allow-glsl -emit-spirv-directly
+//TEST(compute, vulkan):COMPARE_COMPUTE_EX(filecheck-buffer=BUF):-vk -compute -entry computeMain -allow-glsl -xslang -DTARGET_VK
+//TEST(compute, vulkan):COMPARE_COMPUTE_EX(filecheck-buffer=BUF):-vk -compute -entry computeMain -allow-glsl -emit-spirv-directly -xslang -DTARGET_VK
+//TEST(compute):COMPARE_COMPUTE_EX(filecheck-buffer=BUF):-wgpu -compute -entry computeMain -allow-glsl
+
+// There are some issues with the Metal backend when using glsl-style syntax - this should be fixed.
+//TEST(compute):COMPARE_COMPUTE_EX(filecheck-buffer=BUF):-metal -compute -entry computeMain -allow-glsl -xslang -DTARGET_METAL
+
 #version 430
 
 //TEST_INPUT:ubuffer(data=[0], stride=4):out,name=outputBuffer
@@ -24,15 +29,18 @@ void computeMain()
 {
     if (gl_GlobalInvocationID.x == 3) {
         outputBuffer.data[0] = true
-            && gl_NumSubgroups == 1
-            && gl_SubgroupID  == 0 //1 subgroup, 0 based indexing
             && gl_SubgroupSize == 32
             && gl_SubgroupInvocationID == 3
+            // These intrinsics are only available on Vulkan(SPIRV and GLSL).
+#if defined(TARGET_VK)
+            && gl_SubgroupID  == 0 //1 subgroup, 0 based indexing
+            && gl_NumSubgroups == 1
             && gl_SubgroupEqMask == uvec4(0b1000,0,0,0)
             && gl_SubgroupGeMask == uvec4(0xFFFFFFF8,0,0,0)
             && gl_SubgroupGtMask == uvec4(0xFFFFFFF0,0,0,0)
             && gl_SubgroupLeMask == uvec4(0b1111,0,0,0)
             && gl_SubgroupLtMask == uvec4(0b111,0,0,0)
+#endif
             ;
     }
     // CHECK_GLSL: void main(
diff --git a/tests/hlsl-intrinsic/quad-control/quad-control-comp-functionality.slang b/tests/hlsl-intrinsic/quad-control/quad-control-comp-functionality.slang
index 01771f7a4..77ca03178 100644
--- a/tests/hlsl-intrinsic/quad-control/quad-control-comp-functionality.slang
+++ b/tests/hlsl-intrinsic/quad-control/quad-control-comp-functionality.slang
@@ -8,11 +8,7 @@ RWStructuredBuffer<uint> outputBuffer;
 [numthreads(16, 1, 1)]
 void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
 {
-#if !defined(METAL)
     uint index = WaveGetLaneIndex();
-#else
-    uint index = dispatchThreadID.x;
-#endif
 
     if (index < 4)
     {
diff --git a/tests/hlsl-intrinsic/wave-get-lane-index.slang b/tests/hlsl-intrinsic/wave-get-lane-index.slang
index fb09022c2..e1a9262a9 100644
--- a/tests/hlsl-intrinsic/wave-get-lane-index.slang
+++ b/tests/hlsl-intrinsic/wave-get-lane-index.slang
@@ -4,6 +4,8 @@
 //TEST:COMPARE_COMPUTE_EX:-slang -compute -dx12 -use-dxil -profile cs_6_0 -shaderobj
 //TEST(vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj
 //TEST:COMPARE_COMPUTE_EX:-cuda -compute -shaderobj
+//TEST:COMPARE_COMPUTE_EX:-wgpu -compute -shaderobj
+//TEST:COMPARE_COMPUTE_EX:-metal -compute -shaderobj
 
 //TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name outputBuffer
 RWStructuredBuffer<int> outputBuffer;
@@ -17,4 +19,4 @@ void computeMain(int3 dispatchThreadID : SV_DispatchThreadID)
     // For now we'll just check it's not 0.
     uint laneCount = WaveGetLaneCount();
     outputBuffer[idx] = int(((laneCount > 0) ? 0x100 : 0) + laneId);
-}
\ No newline at end of file
+}