Add a loop analysis step to infer the exit values of loop phi parameters. (#6696)

* Initial loop analysis pass

* More changes for a single-pass implication propagation

* Update slang-ir-autodiff-loop-analysis.cpp

* Cleanup + new system for loop analysis

* Fixup bugs in loop analysis

* Remove some relation types to simplify the analysis. Add test

* Remove unused

* Address comments

* Fix issue with continue loops

* Update reverse-loop-exit-value-inference-1.slang

* Update reverse-continue-loop.slang

5 files changed, 252 insertions, 6 deletions

diff --git a/tests/autodiff/reverse-continue-loop.slang b/tests/autodiff/reverse-continue-loop.slang
index 51f17b611..77bfb358c 100644
--- a/tests/autodiff/reverse-continue-loop.slang
+++ b/tests/autodiff/reverse-continue-loop.slang
@@ -9,14 +9,13 @@ RWStructuredBuffer<float> outputBuffer;
 typedef DifferentialPair<float> dpfloat;
 typedef float.Differential dfloat;
 
-//CHK-DAG: note: checkpointing context of 24 bytes associated with function: 'test_loop_with_continue'
+//CHK-DAG: note: checkpointing context of 20 bytes associated with function: 'test_loop_with_continue'
 [BackwardDifferentiable]
 float test_loop_with_continue(float y)
 {
     //CHK-DAG: note: 20 bytes (FixedArray<float, 5> ) used to checkpoint the following item:
     float t = y;
 
-    //CHK-DAG: note: 4 bytes (int32_t) used to checkpoint the following item:
     for (int i = 0; i < 3; i++)
     {
         if (t > 4.0)
diff --git a/tests/autodiff/reverse-loop-diff-only-2.slang b/tests/autodiff/reverse-loop-diff-only-2.slang
index 2cc33ecca..cc9e14736 100644
--- a/tests/autodiff/reverse-loop-diff-only-2.slang
+++ b/tests/autodiff/reverse-loop-diff-only-2.slang
@@ -32,10 +32,17 @@ float infinitesimal(float x)
 
 // Test that computeLoop's intermediates have no float sitting
 // around (must not cache the outvar from 'compute()')
-// CHECK: struct s_bwd_prop_computeLoop_Intermediates
-// CHECK-NEXT: {
-// CHECK-NOT: {{[A-Za-z0-9_]+}} {{[A-Za-z0-9_]+}}[{{.*}}]
-// CHECK: }
+// 
+// Further, if loop exit value inference is working correctly, 
+// then there should be no context type at all.
+// 
+// CHECK-NOT: struct s_bwd_prop_computeLoop_Intermediates
+//
+// Check that the signature of the s_bwd_prop_computeLoop function only
+// contains an inout DiffPair_float_0 and a float.
+//
+// CHECK: void s_bwd_prop_computeLoop{{[_0-9]*}}(inout DiffPair_float{{[_0-9]*}} dpy{{[_0-9]*}}, float {{[_a-zA-Z0-9]*}})
+//
 
 [BackwardDifferentiable]
 [PreferRecompute]
diff --git a/tests/autodiff/reverse-loop-exit-value-inference-1.slang b/tests/autodiff/reverse-loop-exit-value-inference-1.slang
new file mode 100644
index 000000000..5a8d8c2c3
--- /dev/null
+++ b/tests/autodiff/reverse-loop-exit-value-inference-1.slang
@@ -0,0 +1,240 @@
+//TEST(compute):COMPARE_COMPUTE_EX(filecheck-buffer=CHECK):-slang -compute -shaderobj -output-using-type
+//TEST(compute, vulkan):COMPARE_COMPUTE_EX(filecheck-buffer=CHECK):-vk -compute -shaderobj -output-using-type
+//TEST(compute):COMPARE_COMPUTE_EX(filecheck-buffer=CHECK):-cpu -compute -output-using-type -shaderobj
+//TEST:SIMPLE(filecheck=CHK_REPORT):-target hlsl -stage compute -entry computeMain -report-checkpoint-intermediates
+
+//TEST_INPUT:ubuffer(data=[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0], stride=4):out,name=outputBuffer
+RWStructuredBuffer<float> outputBuffer;
+
+typedef DifferentialPair<float> dpfloat;
+typedef float.Differential dfloat;
+
+// A variety of tests to check for loop exit value inference.
+// For all of these loops, we expect our inference pass to be able to
+// infer the loop exit value correctly.
+//
+// Further, if the optimization pass runs successfully, then there should
+// be absolutely no context stored for any of these tests.
+//
+
+// CHK_REPORT: (0): note: no checkpoint contexts to report
+
+[Differentiable]
+float test_simple(float y)
+{
+    float t = y;
+
+    for (int i = 0; i < 3; i++)
+    {
+        t = t * (i + 1);
+    }
+
+    return t;
+}
+
+[Differentiable]
+float test_strided(float y)
+{
+    float t = y;
+
+    for (int i = 0; i < 5; i+=2)
+    {
+        t = t * (i + 1);
+    }
+
+    return t;
+}
+
+[Differentiable]
+float test_offset(float y)
+{
+    float t = y;
+
+    for (int i = 2; i < 5; i+=2)
+    {
+        t = t * (i + 1);
+    }
+
+    return t;
+}
+
+[Differentiable]
+float test_negative_stride(float y)
+{
+    float t = y;
+
+    for (int i = 7; i >= 1; i-=2)
+    {
+        t = t * (i + 1);
+    }
+
+    return t;
+}
+
+[Differentiable]
+float test_nested(float y)
+{
+    float t = y;
+
+    for (int i = 0; i < 3; i++)
+    {
+        for (int j = 0; j < 3; j++)
+        {
+            t = t * (i + 4 * j + 1);
+        }
+    }
+
+    return t;
+}
+
+[Differentiable]
+float test_nested_with_offset(float y)
+{
+    float t = y;
+
+    for (int i = -3; i < 3; i++)
+    {
+        for (int j = -3; j < 3; j++)
+        {
+            t = t * ((abs(i) % 2) + (abs(j) % 2) + 1);
+        }
+    }
+
+    return t;
+}
+
+[Differentiable]
+float test_nested_with_conditions(float y)
+{
+    float t = y;
+
+    for (int i = 0; i < 3; i++)
+    {
+        if (i % 2 == 0)
+        {
+            for (int j = 0; j < 3; j++)
+            {
+                if (j % 2 == 0)
+                {
+                    t = t * (i + 4 * j + 1);
+                }
+            }
+        }
+    }
+
+    return t;
+}
+
+[Differentiable]
+float test_with_continue(float y)
+{
+    float t = y;
+
+    for (int i = 0; i < 5; i++)
+    {
+        if (i % 2 == 0)
+        {
+            continue;
+        }
+
+        t = t * (i + 1);
+    }
+
+    return t;
+}
+
+[Differentiable]
+float test_nested_with_continue(float y)
+{
+    float t = y;
+
+    for (int i = 0; i < 3; i++)
+    {
+        if (i % 2 == 0)
+            continue;
+    
+        for (int j = 0; j < 3; j++)
+        {
+            if (j % 2 == 0)
+                continue;
+
+            if (j == 0)
+                continue;
+
+            t = t * (i + 4 * j + 1);
+        }
+    }
+
+    return t;
+}
+
+[numthreads(1, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    outputBuffer[0] = 0.0f; // CHECK: 0.000000
+
+    {
+        dpfloat dpa = dpfloat(0.4, 0.0);
+
+        __bwd_diff(test_simple)(dpa, 1.0f);
+        outputBuffer[1] = dpa.d; // CHECK-NEXT: 6.000000
+    }
+
+    {
+        dpfloat dpa = dpfloat(0.4, 0.0);
+
+        __bwd_diff(test_strided)(dpa, 1.0f);
+        outputBuffer[2] = dpa.d; // CHECK-NEXT: 15.000000
+    }
+
+    {
+        dpfloat dpa = dpfloat(0.4, 0.0);
+
+        __bwd_diff(test_offset)(dpa, 1.0f);
+        outputBuffer[3] = dpa.d; // CHECK-NEXT: 15.000000
+    }
+
+    {
+        dpfloat dpa = dpfloat(0.4, 0.0);
+
+        __bwd_diff(test_negative_stride)(dpa, 1.0f);
+        outputBuffer[4] = dpa.d; // CHECK-NEXT: 384.000000
+    }
+
+    {
+        dpfloat dpa = dpfloat(0.4, 0.0);
+
+        __bwd_diff(test_nested)(dpa, 1.0f);
+        outputBuffer[5] = dpa.d;  // CHECK-NEXT: 1247400.000000
+    }
+
+    {
+        dpfloat dpa = dpfloat(0.4, 0.0);
+
+        __bwd_diff(test_nested_with_offset)(dpa, 1.0f);
+        outputBuffer[6] = dpa.d; // CHECK-NEXT: 5159780352.000000
+    }
+
+    {
+        dpfloat dpa = dpfloat(0.4, 0.0);
+
+        __bwd_diff(test_nested_with_conditions)(dpa, 1.0f);
+        outputBuffer[7] = dpa.d; // CHECK-NEXT: 297.000000
+    }
+
+    {
+        dpfloat dpa = dpfloat(0.4, 0.0);
+
+        __bwd_diff(test_with_continue)(dpa, 1.0f);
+        outputBuffer[8] = dpa.d; // CHECK-NEXT: 8.000000
+    }
+
+    {
+        dpfloat dpa = dpfloat(0.4, 0.0);
+
+        __bwd_diff(test_nested_with_continue)(dpa, 1.0f);
+        outputBuffer[9] = dpa.d; // CHECK-NEXT: 6.000000
+    }
+}
+
+//CHK-NOT: note
\ No newline at end of file
diff --git a/tests/autodiff/reverse-loop.slang b/tests/autodiff/reverse-loop-simple.slang
index 18b672860..18b672860 100644
--- a/tests/autodiff/reverse-loop.slang
+++ b/tests/autodiff/reverse-loop-simple.slang
diff --git a/tests/autodiff/reverse-loop.slang.expected.txt b/tests/autodiff/reverse-loop-simple.slang.expected.txt
index 76b7cf779..76b7cf779 100644
--- a/tests/autodiff/reverse-loop.slang.expected.txt
+++ b/tests/autodiff/reverse-loop-simple.slang.expected.txt