Report AD checkpoint contexts (#5058)

* Transferring source locations when creating phi instructions

* Tracking for simple variables

* Deriving source locations for loop counters

* Printing checkpoint structure breakdown

* More readable output format

* Special behavior for loop counters

* Writing report to file

* Add slangc option to enable checkpoint reports

* Display types of checkpointed fields

* Message in case there are no checkpointing contexts

* Catch source locations for function calls

* Source cleanup

* Fix compilation warnings

* Remove stray dump()

* Provide the report through diagnostic notes

* Add missing path for sourceLoc during unzip pass

* Add tests for reporting intermediates

* Include more transfer cases for source locations

* Fix ordering in address elimination

* Fill in more holes with source location transfer

* Remove debugging line

* Reverting changes to diagnostic sink

* Simplify address elimination using source location RAII contexts

* Eliminating manual source loc transfers in forward transcription

* Fix local var adaptation to use RAII location setter

* Simplify primal hoisting logic for source location transfer

* Simplify unzipping with RAII location scopes

* Simplify transpose logic

* Cleaning up for rev.cpp

* Reverting spacing changes

* Fix mistake with source loc RAII instantiation

* Fix formatting issues

9 files changed, 48 insertions, 2 deletions

diff --git a/tests/autodiff/reverse-checkpoint-1.slang b/tests/autodiff/reverse-checkpoint-1.slang
index 517297013..3d6e9e702 100644
--- a/tests/autodiff/reverse-checkpoint-1.slang
+++ b/tests/autodiff/reverse-checkpoint-1.slang
@@ -2,6 +2,7 @@
 //TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type
 //TEST(compute):COMPARE_COMPUTE_EX:-cpu -compute -output-using-type -shaderobj
 //TEST:SIMPLE(filecheck=CHECK): -target hlsl -profile cs_5_0 -entry computeMain -line-directive-mode none
+//TEST:SIMPLE(filecheck=CHK):-target glsl -stage compute -entry computeMain -report-checkpoint-intermediates
 
 //TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBuffer
 RWStructuredBuffer<float> outputBuffer;
@@ -16,13 +17,16 @@ float g(float x)
     return log(x);
 }
 
+//CHK: note: checkpointing context of 4 bytes associated with function: 'f'
 [BackwardDifferentiable]
 float f(int p, float x)
 {
     float y = 1.0;
     // Test that phi parameter can be restored.
     if (p == 0)
+        //CHK: note: 4 bytes (float) used to checkpoint the following item:
         y = g(x);
+
     return y * y;
 }
 
@@ -41,3 +45,5 @@ void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
     __bwd_diff(f)(0, dpa, 1.0f);
     outputBuffer[0] = dpa.d; // Expect: 1
 }
+
+//CHK-NOT: note
\ No newline at end of file
diff --git a/tests/autodiff/reverse-checkpoint-2.slang b/tests/autodiff/reverse-checkpoint-2.slang
index 8a7262aa4..1dd3f2963 100644
--- a/tests/autodiff/reverse-checkpoint-2.slang
+++ b/tests/autodiff/reverse-checkpoint-2.slang
@@ -41,3 +41,5 @@ void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
     __bwd_diff(f)(0, dpa, 1.0f);
     outputBuffer[0] = dpa.d; // Expect: 1
 }
+
+//CHK-NOT: note
\ No newline at end of file
diff --git a/tests/autodiff/reverse-continue-loop.slang b/tests/autodiff/reverse-continue-loop.slang
index 0f9502673..0b6e56f78 100644
--- a/tests/autodiff/reverse-continue-loop.slang
+++ b/tests/autodiff/reverse-continue-loop.slang
@@ -1,6 +1,7 @@
 //TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type
 //TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type
 //TEST(compute):COMPARE_COMPUTE_EX:-cpu -compute -output-using-type -shaderobj
+//TEST:SIMPLE(filecheck=CHK):-target glsl -stage compute -entry computeMain -report-checkpoint-intermediates
 
 //TEST_INPUT:ubuffer(data=[0 0 0 0 0], stride=4):out,name=outputBuffer
 RWStructuredBuffer<float> outputBuffer;
@@ -8,11 +9,14 @@ RWStructuredBuffer<float> outputBuffer;
 typedef DifferentialPair<float> dpfloat;
 typedef float.Differential dfloat;
 
+//CHK: note: checkpointing context of 24 bytes associated with function: 'test_loop_with_continue'
 [BackwardDifferentiable]
 float test_loop_with_continue(float y)
 {
+    //CHK: note: 20 bytes (FixedArray<float, 5> ) used to checkpoint the following item:
     float t = y;
 
+    //CHK: note: 4 bytes (int32_t) used for a loop counter here:
     for (int i = 0; i < 3; i++)
     {
         if (t > 4.0)
@@ -41,3 +45,5 @@ void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
         outputBuffer[1] = dpa.d; // Expect: 0.0131072
     }
 }
+
+//CHK-NOT: note
\ No newline at end of file
diff --git a/tests/autodiff/reverse-control-flow-1.slang b/tests/autodiff/reverse-control-flow-1.slang
index 7d2f518be..334de4137 100644
--- a/tests/autodiff/reverse-control-flow-1.slang
+++ b/tests/autodiff/reverse-control-flow-1.slang
@@ -1,5 +1,6 @@
 //TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type
 //TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type
+//TEST:SIMPLE(filecheck=CHK):-target glsl -stage compute -entry computeMain -report-checkpoint-intermediates
 
 //TEST_INPUT:ubuffer(data=[0 0 0 0 0], stride=4):out,name=outputBuffer
 RWStructuredBuffer<float> outputBuffer;
@@ -40,3 +41,5 @@ void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
         outputBuffer[1] = dpa.d; // Expect: 1.0
     }
 }
+
+//CHK: (0): note: no checkpoint contexts to report
\ No newline at end of file
diff --git a/tests/autodiff/reverse-control-flow-2.slang b/tests/autodiff/reverse-control-flow-2.slang
index cde707b4d..c3790367c 100644
--- a/tests/autodiff/reverse-control-flow-2.slang
+++ b/tests/autodiff/reverse-control-flow-2.slang
@@ -1,5 +1,6 @@
 //TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type
 //TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type
+//TEST:SIMPLE(filecheck=CHK):-target glsl -stage compute -entry computeMain -report-checkpoint-intermediates
 
 //TEST_INPUT:ubuffer(data=[0 0 0 0 0], stride=4):out,name=outputBuffer
 RWStructuredBuffer<float> outputBuffer;
@@ -73,3 +74,5 @@ void computeMain(uint3 dispatchThreadID: SV_DispatchThreadID)
         outputBuffer[1] = dpx.d;
     }
 }
+
+//CHK: (0): note: no checkpoint contexts to report
\ No newline at end of file
diff --git a/tests/autodiff/reverse-control-flow-3.slang b/tests/autodiff/reverse-control-flow-3.slang
index 01b533279..b4fa68e3a 100644
--- a/tests/autodiff/reverse-control-flow-3.slang
+++ b/tests/autodiff/reverse-control-flow-3.slang
@@ -1,4 +1,5 @@
 //TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type
+//TEST:SIMPLE(filecheck=CHK):-target glsl -stage compute -entry computeMain -report-checkpoint-intermediates
 
 //TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBuffer
 
@@ -75,7 +76,8 @@ void d_getParam(uint id, MaterialParam.Differential diff)
     outputBuffer[id] += diff.roughness;
 }
 
-
+//CHK-DAG: note: checkpointing context of 8 bytes associated with function: 'updatePathThroughput'
+//CHK-DAG: note: 8 bytes (PathResult_0) used to checkpoint the following item:
 [BackwardDifferentiable]
 void updatePathThroughput(inout PathResult path, const float weight)
 {
@@ -122,9 +124,13 @@ bool generateScatterRay(const BSDFSample bs, const MaterialParam bsdfParams, ino
     \param[in,out] path The path state.
     \return True if a ray was generated, false otherwise.
 */
+
+//CHK-DAG: note: checkpointing context of 16 bytes associated with function: 'generateScatterRay'
 [BackwardDifferentiable]
 bool generateScatterRay(const BSDFSample bs, const MaterialParam bsdfParams, inout PathState path, inout PathResult pathRes)
 {
+    //CHK-DAG: note: 8 bytes (s_bwd_prop_updatePathThroughput_Intermediates_0) used to checkpoint the following item:
+    //CHK-DAG: note: 8 bytes (PathResult_0) used to checkpoint the following item:
     updatePathThroughput(pathRes, bs.val);
     return true;
 }
@@ -215,5 +221,6 @@ void computeMain(uint3 dispatchThreadID: SV_DispatchThreadID)
         var dpx = diffPair(pathRes, pathResD);
         __bwd_diff(tracePath)(1, dpx); // Expect: 5.0 in outputBuffer[3]
     }
-
 }
+
+//CHK-NOT: note
\ No newline at end of file
diff --git a/tests/autodiff/reverse-loop-checkpoint-test.slang b/tests/autodiff/reverse-loop-checkpoint-test.slang
index fc206e128..68ad823ac 100644
--- a/tests/autodiff/reverse-loop-checkpoint-test.slang
+++ b/tests/autodiff/reverse-loop-checkpoint-test.slang
@@ -1,5 +1,6 @@
 //TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type
 //TEST:SIMPLE(filecheck=CHECK): -target hlsl -profile cs_5_0 -entry computeMain -line-directive-mode none
+//TEST:SIMPLE(filecheck=CHK):-target glsl -stage compute -entry computeMain -report-checkpoint-intermediates
 
 //TEST_INPUT:ubuffer(data=[0 0 0 0 0], stride=4):out,name=outputBuffer
 RWStructuredBuffer<float> outputBuffer;
@@ -44,13 +45,18 @@ float3 infinitesimal(float3 x)
     return x - detach(x);
 }
 
+//CHK: note: checkpointing context of 20 bytes associated with function: 'computeLoop'
 [BackwardDifferentiable]
 [PreferRecompute]
 float3 computeLoop(float y)
 {
+    //CHK: note: 4 bytes (float) used to checkpoint the following item:
     float w = 0;
+
+    //CHK: note: 12 bytes (Vector<float, 3> ) used to checkpoint the following item:
     float3 w3 = float3(0, 0, 0);
 
+    //CHK: note: 4 bytes (int32_t) used for a loop counter here:
     for (int i = 0; i < 8; i++)
     {
         float k = compute(i, y);
@@ -93,3 +99,5 @@ void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
 
     outputBuffer[2] = computeLoop(1.0).x;
 }
+
+//CHK-NOT: note
\ No newline at end of file
diff --git a/tests/autodiff/reverse-loop.slang b/tests/autodiff/reverse-loop.slang
index a2c826be9..2ba8535be 100644
--- a/tests/autodiff/reverse-loop.slang
+++ b/tests/autodiff/reverse-loop.slang
@@ -1,6 +1,7 @@
 //TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type
 //TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type
 //TEST(compute):COMPARE_COMPUTE_EX:-cpu -compute -output-using-type -shaderobj
+//TEST:SIMPLE(filecheck=CHK):-target glsl -stage compute -entry computeMain -report-checkpoint-intermediates
 
 //TEST_INPUT:ubuffer(data=[0 0 0 0 0], stride=4):out,name=outputBuffer
 RWStructuredBuffer<float> outputBuffer;
@@ -8,11 +9,14 @@ RWStructuredBuffer<float> outputBuffer;
 typedef DifferentialPair<float> dpfloat;
 typedef float.Differential dfloat;
 
+//CHK: note: checkpointing context of 24 bytes associated with function: 'test_simple_loop'
 [Differentiable]
 float test_simple_loop(float y)
 {
+    //CHK: note: 20 bytes (FixedArray<float, 5> ) used to checkpoint the following item:
     float t = y;
 
+    //CHK: note: 4 bytes (int32_t) used for a loop counter here:
     for (int i = 0; i < 3; i++)
     {
         t = t * t;
@@ -38,3 +42,5 @@ void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
         outputBuffer[1] = dpa.d; // Expect: 0.0131072
     }
 }
+
+//CHK-NOT: note
\ No newline at end of file
diff --git a/tests/autodiff/reverse-nested-calls.slang b/tests/autodiff/reverse-nested-calls.slang
index caf2df6f8..3c1a52c21 100644
--- a/tests/autodiff/reverse-nested-calls.slang
+++ b/tests/autodiff/reverse-nested-calls.slang
@@ -1,6 +1,7 @@
 //TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type
 //TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type
 //TEST(compute):COMPARE_COMPUTE_EX:-cpu -compute -output-using-type -shaderobj
+//TEST:SIMPLE(filecheck=CHK):-target glsl -stage compute -entry computeMain -report-checkpoint-intermediates
 
 //TEST_INPUT:ubuffer(data=[0 0 0 0 0], stride=4):out,name=outputBuffer
 RWStructuredBuffer<float> outputBuffer;
@@ -15,9 +16,11 @@ float g(float y)
     return result * result;
 }
 
+//CHK: note: checkpointing context of 4 bytes associated with function: 'f'
 [BackwardDifferentiable]
 float f(float x)
 {
+    //CHK: note: 4 bytes (float) used to checkpoint the following item:
     return 3.0f * g(2.0f * x);
 }
 
@@ -29,3 +32,5 @@ void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
     __bwd_diff(f)(dpa, 1.0f);
     outputBuffer[0] = dpa.d; // Expect: 96.0
 }
+
+//CHK-NOT: note
\ No newline at end of file