Fusion pass for saturated_cooperation (#2874)

* Fusion pass for saturated_cooperation

* simplify assert

* regenerate vs projects

* missing test output files

* rename shadowing variable to appease msvc

* Fuse calls to sat_coop with differing inputs

* formatting

* add cpu test for hof simple

* Make higher-order functions into compute comparison tests

* comment tests

* remove redundant test

* Add test to confirm inlining in sat_coop fuse

* Add clarifying comment for sat coop fusing

* Add KnownBuiltin decoration

* s/CanUseFuncSignature/TypesFullyResolved for higher order function checking

* Add TODO

* spelling

* Correct detection of sat_coop calls

* Disable tests which are unsupported on testing infra

13 files changed, 308 insertions, 13 deletions

diff --git a/tests/language-feature/higher-order-functions/generic.slang b/tests/language-feature/higher-order-functions/generic.slang
index e7d63eb27..af538bee1 100644
--- a/tests/language-feature/higher-order-functions/generic.slang
+++ b/tests/language-feature/higher-order-functions/generic.slang
@@ -1,14 +1,25 @@
-//TEST:SIMPLE:
+//TEST(compute):COMPARE_COMPUTE:-cpu -shaderobj
 
-func foo(f : functype (float) -> int) -> int{
+//TEST_INPUT:ubuffer(data=[0], stride=4):out,name=outputBuffer
+RWStructuredBuffer<int> outputBuffer;
+[numthreads(1, 1, 1)]
+void computeMain(uint tig : SV_GroupIndex)
+{
+    outputBuffer[tig] = zap();
+}
+
+func foo(f : functype (float) -> int) -> int
+{
     return f(0);
 }
 
-int bap<T>(float) {
+int bap<T>(float)
+{
     return 1;
 }
 
-int zap() {
+int zap()
+{
     // We should be able to specify which foo we want
     return foo(bap<bool>);
 }
diff --git a/tests/language-feature/higher-order-functions/generic.slang.expected.txt b/tests/language-feature/higher-order-functions/generic.slang.expected.txt
new file mode 100644
index 000000000..d00491fd7
--- /dev/null
+++ b/tests/language-feature/higher-order-functions/generic.slang.expected.txt
@@ -0,0 +1 @@
+1
diff --git a/tests/language-feature/higher-order-functions/inference.slang b/tests/language-feature/higher-order-functions/inference.slang
index c01f5b19e..3468395d7 100644
--- a/tests/language-feature/higher-order-functions/inference.slang
+++ b/tests/language-feature/higher-order-functions/inference.slang
@@ -1,21 +1,34 @@
-//TEST:SIMPLE:
+//TEST(compute):COMPARE_COMPUTE:-cpu -shaderobj
+
+//TEST_INPUT:ubuffer(data=[0], stride=4):out,name=outputBuffer
+RWStructuredBuffer<int> outputBuffer;
+[numthreads(1, 1, 1)]
+void computeMain(uint tig : SV_GroupIndex)
+{
+    // Test that we infer the type parameters to X from the type of f
+    x(f);
+    // Testthat we infer the type paramters to from a specialized generic (g)
+    y(g<int>);
+    outputBuffer[tig] = p;
+}
+
+static int p = 0;
 
 func x<A, B>(f : functype (A) -> B)
-{ }
+{
+    ++p;
+}
 
 float f(int)
 {
-    return 0;
+    return 1;
 }
 
 func y<A, B, C>(g : functype (A, B) -> C)
-{ }
+{
+    p += 2;
+}
 
 void g<A>(A, bool)
 { }
 
-void main()
-{
-    x(f);
-    y(g<int>);
-}
diff --git a/tests/language-feature/higher-order-functions/inference.slang.expected.txt b/tests/language-feature/higher-order-functions/inference.slang.expected.txt
new file mode 100644
index 000000000..00750edc0
--- /dev/null
+++ b/tests/language-feature/higher-order-functions/inference.slang.expected.txt
@@ -0,0 +1 @@
+3
diff --git a/tests/language-feature/higher-order-functions/simple.slang b/tests/language-feature/higher-order-functions/simple.slang
index 7c6bf5a26..1443bd60d 100644
--- a/tests/language-feature/higher-order-functions/simple.slang
+++ b/tests/language-feature/higher-order-functions/simple.slang
@@ -1,5 +1,6 @@
 //TEST(compute, vulkan):COMPARE_COMPUTE_EX():-vk -compute -shaderobj -output-using-type
 //TEST(compute):COMPARE_COMPUTE_EX():-slang -compute -shaderobj -output-using-type
+//TEST(compute):COMPARE_COMPUTE:-cpu -shaderobj -output-using-type
 
 //TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBuffer
 RWStructuredBuffer<uint> outputBuffer;
diff --git a/tests/language-feature/saturated-cooperation/fuse-product.slang b/tests/language-feature/saturated-cooperation/fuse-product.slang
new file mode 100644
index 000000000..2d1029140
--- /dev/null
+++ b/tests/language-feature/saturated-cooperation/fuse-product.slang
@@ -0,0 +1,86 @@
+//DISABLED_TEST(compute, vulkan):COMPARE_COMPUTE_EX():-vk -compute -shaderobj -output-using-type
+//TEST(compute):COMPARE_COMPUTE_EX():-dx12 -profile sm_6_5 -use-dxil -compute -shaderobj -output-using-type
+//TEST:SIMPLE(filecheck=CHECK):-target hlsl -profile cs_6_5 -entry computeMain -line-directive-mode none
+//TEST:SIMPLE(filecheck=CHECK):-target glsl -profile cs_6_5 -entry computeMain -line-directive-mode none
+// Disabled because some testing infra (swiftshader) doens't support the necessary wave ops
+
+//
+// This test checks that we fuse calls to saturated cooperation, even with different input values
+//
+
+//TEST_INPUT:ubuffer(data=[0 3 2 2], stride=4):out,name=outputBuffer
+RWStructuredBuffer<int> outputBuffer;
+
+static int count = 0;
+
+int coopAdd(float x, int i)
+{
+    count += i;
+    return int(x) * 2;
+}
+
+int fallback1(float x, int)
+{
+    count = -1;
+    return -1;
+}
+
+int coopMul(float x, float f)
+{
+    count *= int(f);
+    return int(x) * 2;
+}
+
+int fallback2(float x, float)
+{
+    count = -1;
+    return -1;
+}
+
+int coopExp(float x, int i)
+{
+    int c = count;
+    count = 1;
+    for(int j = 0; j < i; ++j)
+        count *= c;
+    return int(x) * 2;
+}
+
+// Make sure that we have enough invocations to saturate the first workgroup
+[numthreads(128, 1, 1)]
+void computeMain(uint tig : SV_GroupIndex)
+{
+    // The values we're cooperating over are {0, 2, 3}
+    // We track the number of sets evaluated in the "count" variable, and write
+    // that at index 0
+    //
+    // If these are not fused, then we'd expect count to be incremented three
+    // times then doubled three times. What we want to see is
+    // 0, +1, *2, **2, +1, *2, **2, +1, *2, **2 = 40804
+    // If they're not interleaved then we'll see
+    // 0, +1, +1, +1, *2, *2, *2, **2, **2, **2 = -1593835520 (having overflowed)
+    let i = tig < 4 ? float(outputBuffer[tig]) : 0;
+    let j = i + 1;
+    let k = j + 1;
+    let y = saturated_cooperation(coopAdd, fallback1, i, 1);
+    let m = 2.f;
+    let x = saturated_cooperation(coopMul, fallback2, j, m);
+    // Put some calculation between the calls to check that the fusion still takes place
+    let e = min(m*100, 2);
+    let z = saturated_cooperation(coopExp, fallback1, k, int(e));
+
+    // check that all the piping functions we invent during fusing is inlined away
+    // CHECK:      = coopAdd
+    // CHECK-NEXT: = coopMul
+    // there is a tuple pack here
+    // CHECK-NEXT: Tuple
+    // CHECK-NEXT: = coopExp
+
+    // CHECK:      = fallback1
+    // CHECK-NEXT: = fallback2
+    // CHECK-NEXT: Tuple
+    // CHECK-NEXT: = fallback1
+
+    if(tig < 4)
+        outputBuffer[tig] = tig == 0 ? count : z;
+}
diff --git a/tests/language-feature/saturated-cooperation/fuse-product.slang.expected.txt b/tests/language-feature/saturated-cooperation/fuse-product.slang.expected.txt
new file mode 100644
index 000000000..cccee985b
--- /dev/null
+++ b/tests/language-feature/saturated-cooperation/fuse-product.slang.expected.txt
@@ -0,0 +1,5 @@
+type: int32_t
+40804
+10
+8
+8
diff --git a/tests/language-feature/saturated-cooperation/fuse.slang b/tests/language-feature/saturated-cooperation/fuse.slang
new file mode 100644
index 000000000..1492be4f3
--- /dev/null
+++ b/tests/language-feature/saturated-cooperation/fuse.slang
@@ -0,0 +1,55 @@
+//DISABLED_TEST(compute, vulkan):COMPARE_COMPUTE_EX():-vk -compute -shaderobj -output-using-type
+//TEST(compute):COMPARE_COMPUTE_EX():-dx12 -profile sm_6_5 -use-dxil -compute -shaderobj -output-using-type
+// Disabled because some testing infra (swiftshader) doens't support the necessary wave ops
+
+//
+// This test checks whether adjacent calls to saturated_cooperation are fused
+//
+
+//TEST_INPUT:ubuffer(data=[0 3 2 2], stride=4):out,name=outputBuffer
+RWStructuredBuffer<uint> outputBuffer;
+
+static int count = 0;
+
+int cooperate(float x, int i)
+{
+    count += i;
+    return int(x) * 2;
+}
+
+int fallback(float x, int)
+{
+    count += 100;
+    return int(x) * 3;
+}
+
+int cooperate2(float x, float f)
+{
+    count *= int(f);
+    return int(x) * 2;
+}
+
+int fallback2(float x, float)
+{
+    count += 100;
+    return int(x) * 3;
+}
+
+// Make sure that we have enough invocations to saturate the first workgroup
+[numthreads(128, 1, 1)]
+void computeMain(uint tig : SV_GroupIndex)
+{
+    // The values we're cooperating over are {0, 2, 3}
+    // We track the number of sets evaluated in the "count" variable, and write
+    // that at index 0
+    //
+    // If these are not fused, then we'd expect count to be incremented three
+    // times then doubled three times. What we want to see is 
+    // 0, +1, *2, +1, *2, +1, *2 = 14
+    let i = tig < 4 ? float(outputBuffer[tig]) : 0;
+    let y = saturated_cooperation(cooperate, fallback, i, 1);
+    let x = saturated_cooperation(cooperate2, fallback2, i, 2.f);
+    if(tig < 4)
+        outputBuffer[tig] = tig == 0 ? count : x;
+
+}
diff --git a/tests/language-feature/saturated-cooperation/fuse.slang.expected.txt b/tests/language-feature/saturated-cooperation/fuse.slang.expected.txt
new file mode 100644
index 000000000..a52613a63
--- /dev/null
+++ b/tests/language-feature/saturated-cooperation/fuse.slang.expected.txt
@@ -0,0 +1,6 @@
+type: uint32_t
+14
+6
+4
+4
+
diff --git a/tests/language-feature/saturated-cooperation/fuse3.slang b/tests/language-feature/saturated-cooperation/fuse3.slang
new file mode 100644
index 000000000..6d5fd9e06
--- /dev/null
+++ b/tests/language-feature/saturated-cooperation/fuse3.slang
@@ -0,0 +1,69 @@
+//DISABLED_TEST(compute, vulkan):COMPARE_COMPUTE_EX():-vk -compute -shaderobj -output-using-type
+//TEST(compute):COMPARE_COMPUTE_EX():-dx12 -profile sm_6_5 -use-dxil -compute -shaderobj -output-using-type
+// Disabled because some testing infra (swiftshader) doens't support the necessary wave ops
+
+//
+// This test checks whether more than 2 adjacent calls to saturated_cooperation
+// are fused, even with operations between them
+//
+
+//TEST_INPUT:ubuffer(data=[0 3 2 2], stride=4):out,name=outputBuffer
+RWStructuredBuffer<int> outputBuffer;
+
+static int count = 0;
+
+int coopAdd(float x, int i)
+{
+    count += i;
+    return int(x) * 2;
+}
+
+int fallback(float x, int)
+{
+    count = -1;
+    return -1;
+}
+
+int coopMul(float x, float f)
+{
+    count *= int(f);
+    return int(x) * 2;
+}
+
+int fallback2(float x, float)
+{
+    count = -1;
+    return -1;
+}
+
+int coopExp(float x, int i)
+{
+    int c = count;
+    count = 1;
+    for(int j = 0; j < i; ++j)
+        count *= c;
+    return int(x) * 2;
+}
+
+// Make sure that we have enough invocations to saturate the first workgroup
+[numthreads(128, 1, 1)]
+void computeMain(uint tig : SV_GroupIndex)
+{
+    // The values we're cooperating over are {0, 2, 3}
+    // We track the number of sets evaluated in the "count" variable, and write
+    // that at index 0
+    //
+    // If these are not fused, then we'd expect count to be incremented three
+    // times then doubled three times. What we want to see is
+    // 0, +1, *2, **2, +1, *2, **2, +1, *2, **2 = 40804
+    let i = tig < 4 ? float(outputBuffer[tig]) : 0;
+    let y = saturated_cooperation(coopAdd, fallback, i, 1);
+    let m = 2.f;
+    let x = saturated_cooperation(coopMul, fallback2, i, m);
+    // Put some calculation between the calls to check that the fusion still takes place
+    let e = min(m*100, 2);
+    let z = saturated_cooperation(coopExp, fallback, i, int(e));
+    if(tig < 4)
+        outputBuffer[tig] = tig == 0 ? count : z;
+
+}
diff --git a/tests/language-feature/saturated-cooperation/fuse3.slang.expected.txt b/tests/language-feature/saturated-cooperation/fuse3.slang.expected.txt
new file mode 100644
index 000000000..c90a6c141
--- /dev/null
+++ b/tests/language-feature/saturated-cooperation/fuse3.slang.expected.txt
@@ -0,0 +1,5 @@
+type: int32_t
+40804
+6
+4
+4
diff --git a/tests/language-feature/saturated-cooperation/simple.slang b/tests/language-feature/saturated-cooperation/simple.slang
new file mode 100644
index 000000000..3b9fb24e7
--- /dev/null
+++ b/tests/language-feature/saturated-cooperation/simple.slang
@@ -0,0 +1,37 @@
+//DISABLED_TEST(compute, vulkan):COMPARE_COMPUTE_EX():-vk -compute -shaderobj -output-using-type
+//TEST(compute):COMPARE_COMPUTE_EX():-dx12 -profile sm_6_5 -use-dxil -compute -shaderobj -output-using-type
+// Disabled because some testing infra (swiftshader) doens't support the necessary wave ops
+
+//TEST_INPUT:ubuffer(data=[0 3 2 2], stride=4):out,name=outputBuffer
+RWStructuredBuffer<uint> outputBuffer;
+
+struct Unit{};
+
+static int count = 0;
+
+int cooperate(float x, Unit)
+{
+    count += 1;
+    return int(x) * 2;
+}
+
+int fallback(float x, Unit)
+{
+    count += 100;
+    return int(x) * 3;
+}
+
+// Make sure that we have enough invocations to saturate the first workgroup
+[numthreads(128, 1, 1)]
+void computeMain(uint tig : SV_GroupIndex)
+{
+    // The values we're cooperating over are {0, 2, 3}
+    // We track the number of sets evaluated in the "count" variable, and write
+    // that at index 0
+    let i = tig < 4 ? float(outputBuffer[tig]) : 0;
+    Unit unit;
+    let x = saturated_cooperation(cooperate, fallback, i, unit);
+    if(tig < 4)
+        outputBuffer[tig] = tig == 0 ? count : x;
+
+}
diff --git a/tests/language-feature/saturated-cooperation/simple.slang.expected.txt b/tests/language-feature/saturated-cooperation/simple.slang.expected.txt
new file mode 100644
index 000000000..53b2894b0
--- /dev/null
+++ b/tests/language-feature/saturated-cooperation/simple.slang.expected.txt
@@ -0,0 +1,5 @@
+type: uint32_t
+3
+6
+4
+4