From b513d0deef521318ad943d820dd37029075a33c4 Mon Sep 17 00:00:00 2001 From: Sai Praveen Bangaru <31557731+saipraveenb25@users.noreply.github.com> Date: Mon, 11 Jul 2022 23:18:06 -0400 Subject: Added support for differentiating calls to basic functions, as well as arithmetic on the float3 type (#2313) * Added support for differentiating calls to basic functions, as well as arithmetic on the float3 type * Added test expected result --- tests/autodiff/local-redecl-custom-jvp.slang | 28 +++++++++++ .../local-redecl-custom-jvp.slang.expected.txt | 5 ++ tests/autodiff/nested-jvp.slang | 54 ++++++++++++++++++++++ tests/autodiff/nested-jvp.slang.expected.txt | 5 ++ tests/autodiff/redecl-custom-jvp.slang | 28 ----------- .../autodiff/redecl-custom-jvp.slang.expected.txt | 5 -- tests/autodiff/test-intrinsics.slang | 5 ++ tests/autodiff/vector-arithmetic-jvp.slang | 30 ++++++++++++ .../vector-arithmetic-jvp.slang.expected.txt | 5 ++ 9 files changed, 132 insertions(+), 33 deletions(-) create mode 100644 tests/autodiff/local-redecl-custom-jvp.slang create mode 100644 tests/autodiff/local-redecl-custom-jvp.slang.expected.txt create mode 100644 tests/autodiff/nested-jvp.slang create mode 100644 tests/autodiff/nested-jvp.slang.expected.txt delete mode 100644 tests/autodiff/redecl-custom-jvp.slang delete mode 100644 tests/autodiff/redecl-custom-jvp.slang.expected.txt create mode 100644 tests/autodiff/vector-arithmetic-jvp.slang create mode 100644 tests/autodiff/vector-arithmetic-jvp.slang.expected.txt (limited to 'tests') diff --git a/tests/autodiff/local-redecl-custom-jvp.slang b/tests/autodiff/local-redecl-custom-jvp.slang new file mode 100644 index 000000000..2bc7cd582 --- /dev/null +++ b/tests/autodiff/local-redecl-custom-jvp.slang @@ -0,0 +1,28 @@ +//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type + +//TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBuffer +RWStructuredBuffer outputBuffer; + +import test_intrinsics; + +float my_pow_jvp(float x, float n, float dx, float dn) +{ + return dx * n * pow(x, n-1) + dn * pow(x, n) * log(x); +} + +[__custom_jvp(my_pow_jvp)] +float _pow(float, float); + +[numthreads(1, 1, 1)] +void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID) +{ + { + float a = 5.0; + float n = 2; + float da = 1.0; + float dn = 0; + + outputBuffer[0] = __jvp(_pow)(a, n, da, dn); // Expect: 10.0 + outputBuffer[1] = __jvp(_pow)(a, n, 0.0, 1.0); // Expect: 40.23595 + } +} diff --git a/tests/autodiff/local-redecl-custom-jvp.slang.expected.txt b/tests/autodiff/local-redecl-custom-jvp.slang.expected.txt new file mode 100644 index 000000000..965f2cb48 --- /dev/null +++ b/tests/autodiff/local-redecl-custom-jvp.slang.expected.txt @@ -0,0 +1,5 @@ +type: float +10.0 +40.23595 +0.0 +0.0 \ No newline at end of file diff --git a/tests/autodiff/nested-jvp.slang b/tests/autodiff/nested-jvp.slang new file mode 100644 index 000000000..c37641c91 --- /dev/null +++ b/tests/autodiff/nested-jvp.slang @@ -0,0 +1,54 @@ +//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type +//TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type + +//TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBuffer +RWStructuredBuffer outputBuffer; + +[__custom_jvp(pow_jvp)] +float pow_(float x, float n) +{ + return pow(x, n); +} + + +[__custom_jvp(max_jvp)] +float max_(float x, float y) +{ + return max(x, y); +} + + +float pow_jvp(float x, float n, float dx, float dn) +{ + return dx * n * pow(x, n-1) + dn * pow(x, n) * log(x); +} + + +float max_jvp(float x, float y, float dx, float dy) +{ + return (x > y) ? dx : dy; +} + + +/* Fresnel Schlick example */ +__differentiate_jvp float3 fresnel(float3 f0, float3 f90, float cosTheta) +{ + return f0 + (f90 - f0) * pow_(max_(1 - cosTheta, 0.0), 5); +} + + +[numthreads(1, 1, 1)] +void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID) +{ + { + float3 f0 = float3(0.2, 0.2, 0.2); + float3 f90 = float3(0.7, 0.7, 0.7); + float cosTheta = 0.5; + + float3 d_f0 = float3(0.1, 0.1, 0.1); + float3 d_f90 = float3(0.9, 0.9, 0.9); + float d_cosTheta = 1.0; + + outputBuffer[0] = __jvp(fresnel)(f0, f90, cosTheta, d_f0, d_f90, d_cosTheta).y; // Expect: -0.031250 + } +} diff --git a/tests/autodiff/nested-jvp.slang.expected.txt b/tests/autodiff/nested-jvp.slang.expected.txt new file mode 100644 index 000000000..107153351 --- /dev/null +++ b/tests/autodiff/nested-jvp.slang.expected.txt @@ -0,0 +1,5 @@ +type: float +-0.031250 +0.0 +0.0 +0.0 \ No newline at end of file diff --git a/tests/autodiff/redecl-custom-jvp.slang b/tests/autodiff/redecl-custom-jvp.slang deleted file mode 100644 index 2bc7cd582..000000000 --- a/tests/autodiff/redecl-custom-jvp.slang +++ /dev/null @@ -1,28 +0,0 @@ -//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type - -//TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBuffer -RWStructuredBuffer outputBuffer; - -import test_intrinsics; - -float my_pow_jvp(float x, float n, float dx, float dn) -{ - return dx * n * pow(x, n-1) + dn * pow(x, n) * log(x); -} - -[__custom_jvp(my_pow_jvp)] -float _pow(float, float); - -[numthreads(1, 1, 1)] -void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID) -{ - { - float a = 5.0; - float n = 2; - float da = 1.0; - float dn = 0; - - outputBuffer[0] = __jvp(_pow)(a, n, da, dn); // Expect: 10.0 - outputBuffer[1] = __jvp(_pow)(a, n, 0.0, 1.0); // Expect: 40.23595 - } -} diff --git a/tests/autodiff/redecl-custom-jvp.slang.expected.txt b/tests/autodiff/redecl-custom-jvp.slang.expected.txt deleted file mode 100644 index 965f2cb48..000000000 --- a/tests/autodiff/redecl-custom-jvp.slang.expected.txt +++ /dev/null @@ -1,5 +0,0 @@ -type: float -10.0 -40.23595 -0.0 -0.0 \ No newline at end of file diff --git a/tests/autodiff/test-intrinsics.slang b/tests/autodiff/test-intrinsics.slang index 3fa53e85a..189004543 100644 --- a/tests/autodiff/test-intrinsics.slang +++ b/tests/autodiff/test-intrinsics.slang @@ -4,3 +4,8 @@ float pow_(float x, float n) { return pow(x, n); } + +float max_(float x, float y) +{ + return max(x, y); +} \ No newline at end of file diff --git a/tests/autodiff/vector-arithmetic-jvp.slang b/tests/autodiff/vector-arithmetic-jvp.slang new file mode 100644 index 000000000..5852b49a7 --- /dev/null +++ b/tests/autodiff/vector-arithmetic-jvp.slang @@ -0,0 +1,30 @@ +//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type +//TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type + +//TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBuffer +RWStructuredBuffer outputBuffer; + +__differentiate_jvp float3 f(float3 x) +{ + return x; +} + +__differentiate_jvp float3 g(float3 x, float3 y) +{ + float3 a = x + y; + float3 b = x - y; + return a * b + 2 * x * y; +} + +[numthreads(1, 1, 1)] +void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID) +{ + { + float3 a = float3(2.0, 2.0, 2.0); + float3 b = float3(1.5, 1.5, 1.5); + float3 da = float3(1.0, 1.0, 1.0); + + outputBuffer[0] = __jvp(f)(a, da).z; // Expect: 1 + outputBuffer[1] = __jvp(g)(a, b, da, float3(2.0, 1.0, 0.0)).y; // Expect: 8 + } +} diff --git a/tests/autodiff/vector-arithmetic-jvp.slang.expected.txt b/tests/autodiff/vector-arithmetic-jvp.slang.expected.txt new file mode 100644 index 000000000..adbb9a448 --- /dev/null +++ b/tests/autodiff/vector-arithmetic-jvp.slang.expected.txt @@ -0,0 +1,5 @@ +type: float +1.0 +8.0 +0.0 +0.0 \ No newline at end of file -- cgit v1.2.3