From 945409c4c6871c18aad24086c594cc66b5913733 Mon Sep 17 00:00:00 2001 From: Sai Praveen Bangaru <31557731+saipraveenb25@users.noreply.github.com> Date: Thu, 17 Aug 2023 14:45:13 -0400 Subject: Initial support for differentiating existential types (#3111) * Merge * WIP: Complete auto-diff logic for existential types * Revert "Add compiler option for generating representative hash" This reverts commit 13b09ef4621e73844c96d64d9c111a8ed0d45aae. * More fixes for fwd-mode AD on existential types * Add anyValueSize inference pass * Fix checking of `Differential.Differential==Differential` * In-progress: infer any-value-size for existential types * Existentials now work in forward-mode * Overhaul handling of existential AD types. Fwd-mode works, reverse-mode requires front-end changes * Reverse-mode now works on existentials * Cleanup * Remove diff rules for create existential object for now * Revert treat-as-differentiable changes * Fixes * More fixes * Cleanup * more cleanup * signed/unsigned * Revert "Cleanup" This reverts commit e4f7d71f07bb207736f90708961eeecd09a1b652. * Cleanup (again) * Remove public/export/keep-alive on null differential after AD pass * Minor fix * Update dictionary accessors * Keep export decoration * More fixes + Support for `kIROp_PackAnyValue` * Merge upstream * Update expected-failure.txt --- tests/autodiff/existential-1.slang | 76 ++++++++++++++++++++++ tests/autodiff/existential-1.slang.expected.txt | 6 ++ tests/autodiff/existential-2.slang | 75 +++++++++++++++++++++ tests/autodiff/existential-2.slang.expected.txt | 6 ++ tests/autodiff/null-differential.slang | 68 +++++++++++++++++++ tests/autodiff/treat-as-differentiable-1.slang | 37 +++++++++++ .../treat-as-differentiable-1.slang.expected.txt | 2 + tests/autodiff/treat-as-differentiable.slang | 37 ----------- .../treat-as-differentiable.slang.expected.txt | 2 - tests/expected-failure.txt | 6 +- 10 files changed, 275 insertions(+), 40 deletions(-) create mode 100644 tests/autodiff/existential-1.slang create mode 100644 tests/autodiff/existential-1.slang.expected.txt create mode 100644 tests/autodiff/existential-2.slang create mode 100644 tests/autodiff/existential-2.slang.expected.txt create mode 100644 tests/autodiff/null-differential.slang create mode 100644 tests/autodiff/treat-as-differentiable-1.slang create mode 100644 tests/autodiff/treat-as-differentiable-1.slang.expected.txt delete mode 100644 tests/autodiff/treat-as-differentiable.slang delete mode 100644 tests/autodiff/treat-as-differentiable.slang.expected.txt (limited to 'tests') diff --git a/tests/autodiff/existential-1.slang b/tests/autodiff/existential-1.slang new file mode 100644 index 000000000..4f93dd04d --- /dev/null +++ b/tests/autodiff/existential-1.slang @@ -0,0 +1,76 @@ +// Test calling differentiable function through dynamic dispatch. + +//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 0], stride=4):out,name=outputBuffer +RWStructuredBuffer outputBuffer; + +[anyValueSize(16)] +interface IInterface : IDifferentiable +{ + [Differentiable] + float calc(float x); +} + +struct A : IInterface +{ + int data1; + + __init(int data1) { this.data1 = data1; } + + [Differentiable] + float calc(float x) { return x * x * x * data1; } +}; + +struct B : IInterface +{ + int data1; + int data2; + + __init(int data1, int data2) { this.data1 = data1; this.data2 = data2; } + + [Differentiable] + float calc(float x) { return x * x * data1 * data2; } +}; + +[Differentiable] +float doThing(IInterface obj, float x) +{ + return obj.calc(x); +} + +[Differentiable] +float f(uint id, float x) +{ + IInterface obj; + + if (id == 0) + obj = no_diff(A(2)); + else + obj = no_diff(B(2, 3)); + + return doThing(obj, x); +} + +//TEST_INPUT: type_conformance A:IInterface = 0 +//TEST_INPUT: type_conformance B:IInterface = 1 + +[numthreads(1, 1, 1)] +void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID) +{ + outputBuffer[0] = fwd_diff(f)(dispatchThreadID.x, DifferentialPair(1.0, 2.0)).d; // A.calc, expect 12 + outputBuffer[1] = fwd_diff(f)(dispatchThreadID.x + 1, DifferentialPair(1.5, 1.0)).d; // B.calc, expect 18 + + { + var dpx = diffPair(1.0); + bwd_diff(f)(dispatchThreadID.x, dpx, 2.0); // A.calc, expect 12 + outputBuffer[2] = dpx.d; + } + + { + var dpx = diffPair(1.5); + bwd_diff(f)(dispatchThreadID.x + 1, dpx, 1.0); // B.calc, expect 18 + outputBuffer[3] = dpx.d; + } +} diff --git a/tests/autodiff/existential-1.slang.expected.txt b/tests/autodiff/existential-1.slang.expected.txt new file mode 100644 index 000000000..c5fce63b1 --- /dev/null +++ b/tests/autodiff/existential-1.slang.expected.txt @@ -0,0 +1,6 @@ +type: float +12.000000 +18.000000 +12.000000 +18.000000 +0.000000 diff --git a/tests/autodiff/existential-2.slang b/tests/autodiff/existential-2.slang new file mode 100644 index 000000000..dd68c2972 --- /dev/null +++ b/tests/autodiff/existential-2.slang @@ -0,0 +1,75 @@ +// Test calling differentiable function through dynamic dispatch. + +//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 0], stride=4):out,name=outputBuffer +RWStructuredBuffer outputBuffer; + +[anyValueSize(16)] +interface IInterface : IDifferentiable +{ + [Differentiable] + float calc(float x); +} + +struct A : IInterface +{ + float data1; + + [Differentiable] + __init(float data1) { this.data1 = data1; } + + [Differentiable] + float calc(float x) { return x * x * x * data1; } +}; + +struct B : IInterface +{ + float data1; + float data2; + + [Differentiable] + __init(float data1, float data2) { this.data1 = data1; this.data2 = data2; } + + [Differentiable] + float calc(float x) { return x * x * data1 * data2; } +}; + +[Differentiable] +float doThing(IInterface obj, float x) +{ + return obj.calc(x); +} + +[Differentiable] +float f(uint id, float x) +{ + IInterface obj; + + if (id == 0) + obj = A(x); // x^4 + else + obj = B(x, x); // x^4 + + return doThing(obj, x) + doThing(obj, x); // 2 * x^4 +} + +//TEST_INPUT: type_conformance A:IInterface = 0 +//TEST_INPUT: type_conformance B:IInterface = 1 + +[numthreads(1, 1, 1)] +void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID) +{ + { + var dpx = diffPair(1.0); + bwd_diff(f)(dispatchThreadID.x, dpx, 2.0); + outputBuffer[0] = dpx.d; // Expect: 2 * 4 * x^3 * dx = 8 * x^3 * dx = 16 + } + + { + var dpx = diffPair(1.5); + bwd_diff(f)(dispatchThreadID.x + 1, dpx, 1.0); + outputBuffer[1] = dpx.d; // Expect: 2 * 4 * x^3 * dx = 8 * (1.5)^3 * 1.0 = 27 + } +} diff --git a/tests/autodiff/existential-2.slang.expected.txt b/tests/autodiff/existential-2.slang.expected.txt new file mode 100644 index 000000000..98cc988a3 --- /dev/null +++ b/tests/autodiff/existential-2.slang.expected.txt @@ -0,0 +1,6 @@ +type: float +16.000000 +27.000000 +0.000000 +0.000000 +0.000000 diff --git a/tests/autodiff/null-differential.slang b/tests/autodiff/null-differential.slang new file mode 100644 index 000000000..71e0f1046 --- /dev/null +++ b/tests/autodiff/null-differential.slang @@ -0,0 +1,68 @@ +// Test calling differentiable function through dynamic dispatch. + +//DISABLE_TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type +//DISABLE_TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type + +//TEST_INPUT:ubuffer(data=[0 0 0 0 0], stride=4):out,name=outputBuffer +RWStructuredBuffer outputBuffer; + +[anyValueSize(16)] +interface IInterface : IDifferentiable +{ + [Differentiable] + float calc(float x); +} + +struct A : IInterface +{ + int data1; + + __init(int data1) { this.data1 = data1; } + + [Differentiable] + float calc(float x) { return x * x * x * data1; } +}; + +struct B : IInterface +{ + int data1; + int data2; + + __init(int data1, int data2) { this.data1 = data1; this.data2 = data2; } + + [Differentiable] + float calc(float x) { return x * x * data1 * data2; } +}; + +[Differentiable] +float doThing(IInterface obj, IInterface obj2, float x) +{ + return obj.calc(x); +} + +//TEST_INPUT: type_conformance A:IInterface = 0 +//TEST_INPUT: type_conformance B:IInterface = 1 + +[numthreads(1, 1, 1)] +void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID) +{ + { + IInterface obj; + if (dispatchThreadID.x == 0) + obj = A(2); + else + obj = B(2, 3); + + IInterface obj2 = A(3); + + var dpx = diffPair(1.0); + DifferentialPair dpobj = diffPair(obj); + DifferentialPair dpobj2 = diffPair(obj2); + + bwd_diff(doThing)(dpobj, dpobj2, dpx, 2.0); // A.calc, expect 12 + + IDifferentiable objd = dpobj.d; + + outputBuffer[0] = (isNullDifferential((objd)) ? 1.f : 0.f); + } +} diff --git a/tests/autodiff/treat-as-differentiable-1.slang b/tests/autodiff/treat-as-differentiable-1.slang new file mode 100644 index 000000000..95423d978 --- /dev/null +++ b/tests/autodiff/treat-as-differentiable-1.slang @@ -0,0 +1,37 @@ +// Tests automatic synthesis of Differential type and method requirements. + +//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], stride=4):out,name=outputBuffer +RWStructuredBuffer outputBuffer; + +interface IFoo +{ + [BackwardDifferentiable] + float f(float v); +} + +struct B : IFoo +{ + [TreatAsDifferentiable] + float f(float v) + { + return v * v; + } +} + +[BackwardDifferentiable] +float use(IFoo o, float x) +{ + return o.f(x); +} + +[numthreads(1, 1, 1)] +void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID) +{ + B b; + var p = diffPair(1.0); + __bwd_diff(use)(b, p, 1.0); + outputBuffer[0] = p.d; +} diff --git a/tests/autodiff/treat-as-differentiable-1.slang.expected.txt b/tests/autodiff/treat-as-differentiable-1.slang.expected.txt new file mode 100644 index 000000000..9d11e5c94 --- /dev/null +++ b/tests/autodiff/treat-as-differentiable-1.slang.expected.txt @@ -0,0 +1,2 @@ +type: float +0.0 \ No newline at end of file diff --git a/tests/autodiff/treat-as-differentiable.slang b/tests/autodiff/treat-as-differentiable.slang deleted file mode 100644 index 95423d978..000000000 --- a/tests/autodiff/treat-as-differentiable.slang +++ /dev/null @@ -1,37 +0,0 @@ -// Tests automatic synthesis of Differential type and method requirements. - -//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], stride=4):out,name=outputBuffer -RWStructuredBuffer outputBuffer; - -interface IFoo -{ - [BackwardDifferentiable] - float f(float v); -} - -struct B : IFoo -{ - [TreatAsDifferentiable] - float f(float v) - { - return v * v; - } -} - -[BackwardDifferentiable] -float use(IFoo o, float x) -{ - return o.f(x); -} - -[numthreads(1, 1, 1)] -void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID) -{ - B b; - var p = diffPair(1.0); - __bwd_diff(use)(b, p, 1.0); - outputBuffer[0] = p.d; -} diff --git a/tests/autodiff/treat-as-differentiable.slang.expected.txt b/tests/autodiff/treat-as-differentiable.slang.expected.txt deleted file mode 100644 index 9d11e5c94..000000000 --- a/tests/autodiff/treat-as-differentiable.slang.expected.txt +++ /dev/null @@ -1,2 +0,0 @@ -type: float -0.0 \ No newline at end of file diff --git a/tests/expected-failure.txt b/tests/expected-failure.txt index ad1901ca3..a04eb90c9 100644 --- a/tests/expected-failure.txt +++ b/tests/expected-failure.txt @@ -173,4 +173,8 @@ tests/slang-extension/cas-int64-byte-address-buffer.slang.4 (vk) tests/slang-extension/exchange-int64-byte-address-buffer.slang.4 (vk) tests/slang-extension/realtime-clock.slang.2 (vk) tests/spirv/spirv-instruction.slang (vk) -tests/type/texture-sampler/texture-sampler-2d.slang (vk) \ No newline at end of file +tests/type/texture-sampler/texture-sampler-2d.slang (vk) +tests/autodiff/existential-1.slang.1 (vk) +tests/autodiff/existential-2.slang.1 (vk) +tests/ir/loop-inversion.slang.4 (vk) +tests/spirv/direct-spirv-control-flow-2.slang (vk) \ No newline at end of file -- cgit v1.2.3