From f7f0dcadd3b2aca4c0bcd03a96e11c617cf69fc2 Mon Sep 17 00:00:00 2001 From: Sai Praveen Bangaru <31557731+saipraveenb25@users.noreply.github.com> Date: Wed, 26 Oct 2022 22:21:29 -0400 Subject: Adding a differentiable standard library (#2465) --- tests/autodiff/custom-intrinsic.slang | 119 +++++++++++++++++++++ tests/autodiff/custom-intrinsic.slang.expected.txt | 6 ++ tests/autodiff/dstdlib.slang | 37 +++++++ tests/autodiff/dstdlib.slang.expected.txt | 7 ++ 4 files changed, 169 insertions(+) create mode 100644 tests/autodiff/custom-intrinsic.slang create mode 100644 tests/autodiff/custom-intrinsic.slang.expected.txt create mode 100644 tests/autodiff/dstdlib.slang create mode 100644 tests/autodiff/dstdlib.slang.expected.txt (limited to 'tests') diff --git a/tests/autodiff/custom-intrinsic.slang b/tests/autodiff/custom-intrinsic.slang new file mode 100644 index 000000000..8ce354edc --- /dev/null +++ b/tests/autodiff/custom-intrinsic.slang @@ -0,0 +1,119 @@ +//TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type +//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type + +//TEST_INPUT:ubuffer(data=[0 0 0 0 0], stride=4):out,name=outputBuffer +RWStructuredBuffer outputBuffer; + +typedef __DifferentialPair dpfloat; + +typealias IDFloat = IFloat & IDifferentiable; + +namespace myintrinsiclib +{ + __generic + __target_intrinsic(hlsl) + __target_intrinsic(glsl) + __target_intrinsic(cuda, "$P_exp($0)") + __target_intrinsic(cpp, "$P_exp($0)") + __target_intrinsic(spirv_direct, "12 resultType resultId glsl450 27 _0") + [__custom_jvp(d_exp)] + T exp(T x); + + __generic + __DifferentialPair d_exp(__DifferentialPair dpx) + { + return __DifferentialPair( + exp(dpx.p()), + T.dmul(exp(dpx.p()), dpx.d())); + } + + + // Sine + __generic + __target_intrinsic(hlsl) + __target_intrinsic(glsl) + __target_intrinsic(cuda, "$P_sin($0)") + __target_intrinsic(cpp, "$P_sin($0)") + __target_intrinsic(spirv_direct, "12 resultType resultId glsl450 13 _0") + [__custom_jvp(d_sin)] + T sin(T x); + + __generic + __DifferentialPair d_sin(__DifferentialPair dpx) + { + return __DifferentialPair( + sin(dpx.p()), + T.dmul(cos(dpx.p()), dpx.d())); + } + + // Cosine + __generic + __target_intrinsic(hlsl) + __target_intrinsic(glsl) + __target_intrinsic(cuda, "$P_cos($0)") + __target_intrinsic(cpp, "$P_cos($0)") + __target_intrinsic(spirv_direct, "12 resultType resultId glsl450 14 _0") + [__custom_jvp(d_cos)] + T cos(T x); + + __generic + __DifferentialPair d_cos(__DifferentialPair dpx) + { + return __DifferentialPair( + cos(dpx.p()), + T.dmul(-sin(dpx.p()), dpx.d())); + } + + // Sine and cosine + __generic + __target_intrinsic(hlsl) + __target_intrinsic(cuda, "$P_sincos($0, $1, $2)") + [__custom_jvp(d_sincos)] + void sincos(T x, out T s, out T c) + { + s = sin(x); + c = cos(x); + } + + __generic + void d_sincos(__DifferentialPair x, out __DifferentialPair s, out __DifferentialPair c) + { + T _s; + T _c; + sincos(x.p(), _s, _c); + + s = __DifferentialPair(_s, T.dmul(_c, x.d())); + c = __DifferentialPair(_c, T.dmul(-_s, x.d())); + } +}; + +__differentiate_jvp float f(float x) +{ + return myintrinsiclib.exp(x); +} + +__differentiate_jvp float g(float x) +{ + float s; + float t; + myintrinsiclib.sincos(x, s, t); + + return s + t; +} + +[numthreads(1, 1, 1)] +void computeMain(uint3 dispatchThreadID: SV_DispatchThreadID) +{ + { + dpfloat dpa = dpfloat(2.0, 1.0); + + outputBuffer[0] = f(dpa.p()); // Expect: 7.389056 + outputBuffer[1] = __jvp(f)(dpa).d(); // Expect: 7.389056 + + // g() needs additional handling of IRMakeDifferentialPair(PtrType). This needs to + // generate a new var, load from the individual vars and store into the pair var. + + //outputBuffer[2] = g(dpa.p()); // Expect: 1.381773 + //outputBuffer[3] = __jvp(g)(dpa).d(); // Expect: -0.301168 + } +} \ No newline at end of file diff --git a/tests/autodiff/custom-intrinsic.slang.expected.txt b/tests/autodiff/custom-intrinsic.slang.expected.txt new file mode 100644 index 000000000..ce22a5b95 --- /dev/null +++ b/tests/autodiff/custom-intrinsic.slang.expected.txt @@ -0,0 +1,6 @@ +type: float +7.389056 +7.389056 +0.0 +0.0 +0.0 \ No newline at end of file diff --git a/tests/autodiff/dstdlib.slang b/tests/autodiff/dstdlib.slang new file mode 100644 index 000000000..6c7ecffbe --- /dev/null +++ b/tests/autodiff/dstdlib.slang @@ -0,0 +1,37 @@ +//TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type +//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type + +//TEST_INPUT:ubuffer(data=[0 0 0 0 0 0], stride=4):out,name=outputBuffer +RWStructuredBuffer outputBuffer; + +typedef __DifferentialPair dpfloat; + +__differentiate_jvp float f(float x) +{ + return dstd.exp(x); +} + +__differentiate_jvp float g(float x) +{ + return dstd.sin(x); +} + +__differentiate_jvp float h(float x) +{ + return dstd.cos(x); +} + +[numthreads(1, 1, 1)] +void computeMain(uint3 dispatchThreadID: SV_DispatchThreadID) +{ + { + dpfloat dpa = dpfloat(2.0, 1.0); + + outputBuffer[0] = f(dpa.p()); // Expect: 7.389056 + outputBuffer[1] = __jvp(f)(dpa).d(); // Expect: 7.389056 + outputBuffer[2] = g(dpa.p()); // Expect: 0.909297 + outputBuffer[3] = __jvp(g)(dpa).d(); // Expect: -0.416146 + outputBuffer[4] = h(dpa.p()); // Expect: -0.416146 + outputBuffer[5] = __jvp(h)(dpa).d(); // Expect: -0.909297 + } +} \ No newline at end of file diff --git a/tests/autodiff/dstdlib.slang.expected.txt b/tests/autodiff/dstdlib.slang.expected.txt new file mode 100644 index 000000000..82053b379 --- /dev/null +++ b/tests/autodiff/dstdlib.slang.expected.txt @@ -0,0 +1,7 @@ +type: float +7.389056 +7.389056 +0.909297 +-0.416147 +-0.416147 +-0.909297 \ No newline at end of file -- cgit v1.2.3