Support for WaveReadLaneAt with dynamic (but uniform across Wave) on Vk by enabling VK1.4. (#1297)

Fixed wave-lane-at.slang test to test with laneId that is uniform across the Wave.

5 files changed, 87 insertions, 22 deletions

diff --git a/source/slang/hlsl.meta.slang b/source/slang/hlsl.meta.slang
index fe348c60d..739b8579d 100644
--- a/source/slang/hlsl.meta.slang
+++ b/source/slang/hlsl.meta.slang
@@ -2722,14 +2722,17 @@ matrix<T,N,M> WaveReadLaneFirst(matrix<T,N,M> expr);
 
 // NOTE! On GLSL based targets the lane index *must* be a compile time expression!
 // See https://github.com/KhronosGroup/GLSL/blob/master/extensions/khr/GL_KHR_shader_subgroup.txt
+// It is allowed to be 'dynamically uniform within the subgroup' if it's SPIR-V 1.4.
+// TODO(JS): For now we'll use 1.4, but aim for the future for the compiler to determine
+// if the line the is compile constant, and reduce requirement to 1.3
 __generic<T : __BuiltinType>
 __glsl_extension(GL_KHR_shader_subgroup_ballot)
-__spirv_version(1.3)
+__spirv_version(1.4)
 __target_intrinsic(glsl, "subgroupBroadcast($0, $1)")
 __target_intrinsic(cuda, "__shfl_sync(__activemask(), $0, $1)")
 T WaveReadLaneAt(T value, int lane);
 __generic<T : __BuiltinType, let N : int>
-__spirv_version(1.3)
+__spirv_version(1.4)
 __target_intrinsic(glsl, "subgroupBroadcast($0, $1)")
 __target_intrinsic(cuda, "_waveReadLaneAtMultiple($0, $1)")
 vector<T,N> WaveReadLaneAt(vector<T,N> value, int lane);
diff --git a/tests/hlsl-intrinsic/wave-lane-at-vk.slang b/tests/hlsl-intrinsic/wave-lane-at-vk.slang
new file mode 100644
index 000000000..137632f16
--- /dev/null
+++ b/tests/hlsl-intrinsic/wave-lane-at-vk.slang
@@ -0,0 +1,45 @@
+// This is similar to wave-lane-at.slang but tests more limited supported types for vk.
+// We have this 'simple' test, because we can't do matrix (or imat) operations on GLSL/Vk target
+
+//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -dx12 -use-dxil -profile cs_6_0
+// TODO(JS): Disabled for now, as requires upgraded glslang
+//DISABLE_TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute
+
+//TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name outputBuffer
+RWStructuredBuffer<int> outputBuffer;
+
+// `The input lane index must be uniform across the wave.`. 
+// The same restriction applies to glsl/SPIR-V 1.4
+// So we are going to use the input buffer to achieve this.
+
+//TEST_INPUT:ubuffer(data=[1 2 3 0], stride=4):name inputBuffer
+RWStructuredBuffer<int> inputBuffer;
+
+[numthreads(4, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    int idx = int(dispatchThreadID.x);
+    
+    int value = 0;
+    
+    for (int i = 0; i < 4; ++i)
+    {
+        // Scalar
+        
+        // The landId is 'dynamic' but it also uniform across the wave (as required by spec)
+        const int laneId = inputBuffer[i];
+        
+        value += WaveReadLaneAt(idx, laneId);
+        
+        // vector
+        
+        {
+            float2 v = float2(idx + 1, idx + 2);
+            float2 readValue = WaveReadLaneAt(v, (laneId + 1) & 3);
+            
+            value += int(readValue[0] + readValue[1]);
+        }            
+    }
+   
+    outputBuffer[idx] = value;
+}
\ No newline at end of file
diff --git a/tests/hlsl-intrinsic/wave-lane-at-vk.slang.expected.txt b/tests/hlsl-intrinsic/wave-lane-at-vk.slang.expected.txt
new file mode 100644
index 000000000..4e98888c6
--- /dev/null
+++ b/tests/hlsl-intrinsic/wave-lane-at-vk.slang.expected.txt
@@ -0,0 +1,4 @@
+1E
+1E
+1E
+1E
diff --git a/tests/hlsl-intrinsic/wave-lane-at.slang b/tests/hlsl-intrinsic/wave-lane-at.slang
index ca05d985d..8661dbc55 100644
--- a/tests/hlsl-intrinsic/wave-lane-at.slang
+++ b/tests/hlsl-intrinsic/wave-lane-at.slang
@@ -1,13 +1,20 @@
 //DISABLE_TEST(compute):COMPARE_COMPUTE_EX:-cpu -compute 
 //DISABLE_TEST(compute):COMPARE_COMPUTE_EX:-slang -compute
 //TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -dx12 -use-dxil -profile cs_6_0
-// Disabled because on GLSL targets the lane index *must* be a const expr - and in this test it is not.
+// Disabled on VK because glsl can't do WaveReadLaneAt on matrix. 
 //DISABLE_TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute
 //TEST(compute):COMPARE_COMPUTE_EX:-cuda -compute
 
 //TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name outputBuffer
 RWStructuredBuffer<int> outputBuffer;
 
+// Note from HLSL: `The input lane index must be uniform across the wave.`. 
+// The same restriction applies to glsl/SPIR-V 1.4
+// So we are going to use the input buffer to achieve this.
+
+//TEST_INPUT:ubuffer(data=[1 2 3 0], stride=4):name inputBuffer
+RWStructuredBuffer<int> inputBuffer;
+
 [numthreads(4, 1, 1)]
 void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
 {
@@ -15,26 +22,32 @@ void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
     
     int value = 0;
     
-    // Scalar
-    
-    value += WaveReadLaneAt(idx, (idx + 1) & 3);
-    
-    // vector
-    
+    for (int i = 0; i < 4; ++i)
     {
-        float2 v = float2(idx + 1, idx + 2);
-        float2 readValue = WaveReadLaneAt(v, (idx - 1) & 3);
+        // Scalar
         
-        value += int(readValue[0] + readValue[1]);
-    }
+        // The landId is 'dynamic' but it also uniform across the wave (as required by spec)
+        const int laneId = inputBuffer[i];
         
-    // matrix
-    {
-        matrix<int, 2, 2> v = matrix<int, 2, 2>(idx, idx - 1, idx * 3, idx - 2);
+        value += WaveReadLaneAt(idx, laneId);
         
-        matrix<int, 2, 2> readValue = WaveReadLaneAt(v, (idx - 1) & 3);
+        // vector
         
-        value += int(readValue[0][0] + readValue[0][1] + readValue[1][0] + readValue[1][1]);
+        {
+            float2 v = float2(idx + 1, idx + 2);
+            float2 readValue = WaveReadLaneAt(v, (laneId + 1) & 3);
+            
+            value += int(readValue[0] + readValue[1]);
+        }
+            
+        // matrix
+        {
+            matrix<int, 2, 2> v = matrix<int, 2, 2>(idx, idx - 1, idx * 3, idx - 2);
+            
+            matrix<int, 2, 2> readValue = WaveReadLaneAt(v, (laneId - 1) & 3);
+            
+            value += int(readValue[0][0] + readValue[0][1] + readValue[1][0] + readValue[1][1]);
+        }
     }
     
     outputBuffer[idx] = value;
diff --git a/tests/hlsl-intrinsic/wave-lane-at.slang.expected.txt b/tests/hlsl-intrinsic/wave-lane-at.slang.expected.txt
index a327b0804..c6167dbae 100644
--- a/tests/hlsl-intrinsic/wave-lane-at.slang.expected.txt
+++ b/tests/hlsl-intrinsic/wave-lane-at.slang.expected.txt
@@ -1,4 +1,4 @@
-19
-2
-B
-10
+36
+36
+36
+36