From 41e7e565eb3dfa13562cbfa3e8641874c2c6d66c Mon Sep 17 00:00:00 2001 From: Sai Praveen Bangaru <31557731+saipraveenb25@users.noreply.github.com> Date: Fri, 4 Apr 2025 14:12:27 -0400 Subject: 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 --- tests/autodiff/reverse-continue-loop.slang | 3 +- tests/autodiff/reverse-loop-diff-only-2.slang | 15 +- .../reverse-loop-exit-value-inference-1.slang | 240 +++++++++++++++++++++ tests/autodiff/reverse-loop-simple.slang | 46 ++++ .../reverse-loop-simple.slang.expected.txt | 6 + tests/autodiff/reverse-loop.slang | 46 ---- tests/autodiff/reverse-loop.slang.expected.txt | 6 - 7 files changed, 304 insertions(+), 58 deletions(-) create mode 100644 tests/autodiff/reverse-loop-exit-value-inference-1.slang create mode 100644 tests/autodiff/reverse-loop-simple.slang create mode 100644 tests/autodiff/reverse-loop-simple.slang.expected.txt delete mode 100644 tests/autodiff/reverse-loop.slang delete mode 100644 tests/autodiff/reverse-loop.slang.expected.txt (limited to 'tests') 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 outputBuffer; typedef DifferentialPair 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 ) 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 outputBuffer; + +typedef DifferentialPair 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-simple.slang b/tests/autodiff/reverse-loop-simple.slang new file mode 100644 index 000000000..18b672860 --- /dev/null +++ b/tests/autodiff/reverse-loop-simple.slang @@ -0,0 +1,46 @@ +//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 +//DISABLE_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; + +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; + } + + return t; +} + +[numthreads(1, 1, 1)] +void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID) +{ + { + dpfloat dpa = dpfloat(1.0, 0.0); + + __bwd_diff(test_simple_loop)(dpa, 1.0f); + outputBuffer[0] = dpa.d; // Expect: 8.0 + } + + { + dpfloat dpa = dpfloat(0.4, 0.0); + + __bwd_diff(test_simple_loop)(dpa, 1.0f); + outputBuffer[1] = dpa.d; // Expect: 0.0131072 + } +} + +//CHK-NOT: note \ No newline at end of file diff --git a/tests/autodiff/reverse-loop-simple.slang.expected.txt b/tests/autodiff/reverse-loop-simple.slang.expected.txt new file mode 100644 index 000000000..76b7cf779 --- /dev/null +++ b/tests/autodiff/reverse-loop-simple.slang.expected.txt @@ -0,0 +1,6 @@ +type: float +8.000000 +0.013107 +0.000000 +0.000000 +0.000000 diff --git a/tests/autodiff/reverse-loop.slang b/tests/autodiff/reverse-loop.slang deleted file mode 100644 index 18b672860..000000000 --- a/tests/autodiff/reverse-loop.slang +++ /dev/null @@ -1,46 +0,0 @@ -//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 -//DISABLE_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; - -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; - } - - return t; -} - -[numthreads(1, 1, 1)] -void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID) -{ - { - dpfloat dpa = dpfloat(1.0, 0.0); - - __bwd_diff(test_simple_loop)(dpa, 1.0f); - outputBuffer[0] = dpa.d; // Expect: 8.0 - } - - { - dpfloat dpa = dpfloat(0.4, 0.0); - - __bwd_diff(test_simple_loop)(dpa, 1.0f); - outputBuffer[1] = dpa.d; // Expect: 0.0131072 - } -} - -//CHK-NOT: note \ No newline at end of file diff --git a/tests/autodiff/reverse-loop.slang.expected.txt b/tests/autodiff/reverse-loop.slang.expected.txt deleted file mode 100644 index 76b7cf779..000000000 --- a/tests/autodiff/reverse-loop.slang.expected.txt +++ /dev/null @@ -1,6 +0,0 @@ -type: float -8.000000 -0.013107 -0.000000 -0.000000 -0.000000 -- cgit v1.2.3