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.

4 files changed, 82 insertions, 20 deletions

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