diff options
| -rw-r--r-- | source/slang/diff.meta.slang | 82 | ||||
| -rw-r--r-- | tests/autodiff/dstdlib-vector.slang | 26 | ||||
| -rw-r--r-- | tests/autodiff/dstdlib-vector.slang.expected.txt | 7 |
3 files changed, 66 insertions, 49 deletions
diff --git a/source/slang/diff.meta.slang b/source/slang/diff.meta.slang index 674531048..1c3066e1d 100644 --- a/source/slang/diff.meta.slang +++ b/source/slang/diff.meta.slang @@ -59,60 +59,15 @@ extension float : IDifferentiable } } -extension vector<float, 3> : IDifferentiable +__generic<let N:int> +extension vector<float, N> : IDifferentiable { - typedef vector<float, 3> Differential; + typedef vector<float, N> Differential; [__unsafeForceInlineEarly] static Differential dzero() { - return vector<float, 3>(0.f); - } - - [__unsafeForceInlineEarly] - static Differential dadd(Differential a, Differential b) - { - return a + b; - } - - [__unsafeForceInlineEarly] - static Differential dmul(This a, Differential b) - { - return a * b; - } -} - -extension vector<float, 2> : IDifferentiable -{ - typedef vector<float, 2> Differential; - - [__unsafeForceInlineEarly] - static Differential dzero() - { - return vector<float, 2>(0.f); - } - - [__unsafeForceInlineEarly] - static Differential dadd(Differential a, Differential b) - { - return a + b; - } - - [__unsafeForceInlineEarly] - static Differential dmul(This a, Differential b) - { - return a * b; - } -} - -extension vector<float, 4> : IDifferentiable -{ - typedef vector<float, 4> Differential; - - [__unsafeForceInlineEarly] - static Differential dzero() - { - return vector<float, 4>(0.f); + return vector<float, N>(0.f); } [__unsafeForceInlineEarly] @@ -207,6 +162,9 @@ struct DifferentialPair : IDifferentiable typealias IDFloat = IFloat & IDifferentiable; +#define VECTOR_MAP_UNARY(TYPE, COUNT, FUNC, VALUE) \ + vector<TYPE,COUNT> result; for(int i = 0; i < COUNT; ++i) { result[i] = FUNC(VALUE[i]); } return result + namespace dstd { // Natural Exponent @@ -262,4 +220,30 @@ namespace dstd cos(dpx.p()), T.dmul(-sin(dpx.p()), dpx.d())); } + + __generic<let N : int> + __target_intrinsic(hlsl) + __target_intrinsic(glsl) + __target_intrinsic(spirv_direct, "12 resultType resultId glsl450 27 _0") + [ForwardDerivative(d_exp_vector)] + vector<float, N> exp(vector<float, N> x) + { + VECTOR_MAP_UNARY(float, N, dstd.exp, x); + } + + __generic<let N : int> + DifferentialPair<vector<float, N>> d_exp_vector(DifferentialPair<vector<float, N>> dpx) + { + vector<float, N> result; + vector<float, N>.Differential d_result; + for(int i = 0; i < N; ++i) + { + DifferentialPair<float> dpexp = dstd.d_exp(DifferentialPair<float>(dpx.p()[i], dpx.d()[i])); + result[i] = dpexp.p(); + d_result[i] = dpexp.d(); + } + + return DifferentialPair<vector<float, N>>(result, d_result); + } + }; diff --git a/tests/autodiff/dstdlib-vector.slang b/tests/autodiff/dstdlib-vector.slang new file mode 100644 index 000000000..44315daf6 --- /dev/null +++ b/tests/autodiff/dstdlib-vector.slang @@ -0,0 +1,26 @@ +//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<float> outputBuffer; + +typedef DifferentialPair<float> dpfloat; + +[ForwardDifferentiable] +float f(float x) +{ + float3 vx = float3(x, 2*x, 3*x); + float3 vexpx = dstd.exp(vx); + return vexpx.x + vexpx.y + vexpx.z; +} + +[numthreads(1, 1, 1)] +void computeMain(uint3 dispatchThreadID: SV_DispatchThreadID) +{ + { + dpfloat dpa = dpfloat(2.0, 1.0); + + outputBuffer[0] = f(dpa.p()); // Expect: 465.415999 + outputBuffer[1] = __fwd_diff(f)(dpa).d(); // Expect: 1326.871736 + } +}
\ No newline at end of file diff --git a/tests/autodiff/dstdlib-vector.slang.expected.txt b/tests/autodiff/dstdlib-vector.slang.expected.txt new file mode 100644 index 000000000..3ab3a05e5 --- /dev/null +++ b/tests/autodiff/dstdlib-vector.slang.expected.txt @@ -0,0 +1,7 @@ +type: float +465.415999 +1326.871736 +0.000000 +0.000000 +0.000000 +0.000000 |