Vector & Matrix Prefix Sum & Product (#1272)

* Implement matrix and vector versions of prefixSum and prefix product.

* Comment around how code is organized - where it seems it could be more performant.

4 files changed, 34 insertions, 19 deletions

diff --git a/tests/hlsl-intrinsic/wave-prefix-product.slang b/tests/hlsl-intrinsic/wave-prefix-product.slang
index bc324ed7d..a56912616 100644
--- a/tests/hlsl-intrinsic/wave-prefix-product.slang
+++ b/tests/hlsl-intrinsic/wave-prefix-product.slang
@@ -11,6 +11,14 @@ RWStructuredBuffer<int> outputBuffer;
 void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
 {
     int idx = int(dispatchThreadID.x);
-    int val = WavePrefixProduct(idx + 1);
-    outputBuffer[idx] = val;
+    
+    float2 v1 = float2(1, idx + 1);
+  
+    int r0 = WavePrefixProduct(idx + 1);
+    float2 r1 = WavePrefixProduct(v1);
+    
+    int r2 = int(r1.x) + int(r1.y) - 1;
+    
+    outputBuffer[idx] = r0 + (r2 << 16);
+    
 }
\ No newline at end of file
diff --git a/tests/hlsl-intrinsic/wave-prefix-product.slang.expected.txt b/tests/hlsl-intrinsic/wave-prefix-product.slang.expected.txt
index 03cb63ab9..1b233efaf 100644
--- a/tests/hlsl-intrinsic/wave-prefix-product.slang.expected.txt
+++ b/tests/hlsl-intrinsic/wave-prefix-product.slang.expected.txt
@@ -1,8 +1,8 @@
-1
-1
-2
-6
-18
-78
-2D0
-13B0
+10001
+10001
+20002
+60006
+180018
+780078
+2D002D0
+13B013B0
diff --git a/tests/hlsl-intrinsic/wave-prefix-sum.slang b/tests/hlsl-intrinsic/wave-prefix-sum.slang
index f8d9bb560..343a7afbd 100644
--- a/tests/hlsl-intrinsic/wave-prefix-sum.slang
+++ b/tests/hlsl-intrinsic/wave-prefix-sum.slang
@@ -11,6 +11,13 @@ RWStructuredBuffer<int> outputBuffer;
 void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
 {
     int idx = int(dispatchThreadID.x);
-    int val = WavePrefixSum(1 << idx);
-    outputBuffer[idx] = val;
+    
+    float2 v1 = float2(1, 1 << idx);
+    
+    int r0 = WavePrefixSum(1 << idx);
+    float2 r1 = WavePrefixSum(v1);
+    
+    int r2 = int(r1.x) + int(r1.y) - idx;
+    
+    outputBuffer[idx] = r0 + (r2 << 16);
 }
\ No newline at end of file
diff --git a/tests/hlsl-intrinsic/wave-prefix-sum.slang.expected.txt b/tests/hlsl-intrinsic/wave-prefix-sum.slang.expected.txt
index 6ec6deeea..4b4230415 100644
--- a/tests/hlsl-intrinsic/wave-prefix-sum.slang.expected.txt
+++ b/tests/hlsl-intrinsic/wave-prefix-sum.slang.expected.txt
@@ -1,8 +1,8 @@
 0
-1
-3
-7
-F
-1F
-3F
-7F
+10001
+30003
+70007
+F000F
+1F001F
+3F003F
+7F007F