From b808aa4df50d46eaa569561f7e464c55c1c2d72a Mon Sep 17 00:00:00 2001 From: venkataram-nv Date: Wed, 18 Sep 2024 20:42:07 -0700 Subject: 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 --- tests/autodiff/reverse-checkpoint-1.slang | 6 ++++++ tests/autodiff/reverse-checkpoint-2.slang | 2 ++ tests/autodiff/reverse-continue-loop.slang | 6 ++++++ tests/autodiff/reverse-control-flow-1.slang | 3 +++ tests/autodiff/reverse-control-flow-2.slang | 3 +++ tests/autodiff/reverse-control-flow-3.slang | 11 +++++++++-- tests/autodiff/reverse-loop-checkpoint-test.slang | 8 ++++++++ tests/autodiff/reverse-loop.slang | 6 ++++++ tests/autodiff/reverse-nested-calls.slang | 5 +++++ 9 files changed, 48 insertions(+), 2 deletions(-) (limited to 'tests') 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 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 outputBuffer; @@ -8,11 +9,14 @@ RWStructuredBuffer outputBuffer; typedef DifferentialPair 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 ) 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 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 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 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 ) 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 outputBuffer; @@ -8,11 +9,14 @@ RWStructuredBuffer outputBuffer; typedef DifferentialPair 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 ) 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 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 -- cgit v1.2.3