From faec3f79464c72833876f7b2a84114955beee7fb Mon Sep 17 00:00:00 2001 From: Atell Krasnopolski Date: Wed, 10 Jul 2024 14:06:44 +0200 Subject: [PATCH] Add basic Kokkos support for the original tests of #783 This commit adds a test from the original Kokkos PR #783 and provides the implementation of the simplest Kokkos features in the form of custom pushforwards. This commit only addresses the forward mode. --- include/clad/Differentiator/KokkosBuiltins.h | 54 +++++++++++++++ unittests/Kokkos/CMakeLists.txt | 1 + unittests/Kokkos/TestUtils.h | 9 ++- unittests/Kokkos/ViewAccess.cpp | 73 ++++++++++++++++++++ unittests/Kokkos/ViewBasics.cpp | 47 +++++++------ 5 files changed, 157 insertions(+), 27 deletions(-) create mode 100644 include/clad/Differentiator/KokkosBuiltins.h create mode 100644 unittests/Kokkos/ViewAccess.cpp diff --git a/include/clad/Differentiator/KokkosBuiltins.h b/include/clad/Differentiator/KokkosBuiltins.h new file mode 100644 index 000000000..8d24adcf0 --- /dev/null +++ b/include/clad/Differentiator/KokkosBuiltins.h @@ -0,0 +1,54 @@ +// This header file contains the implementation of the Kokkos framework +// differentiation support in Clad in the form of custom pushforwards and +// pullbacks. Please include it manually to enable Clad for Kokkos code. + +#ifndef CLAD_DIFFERENTIATOR_KOKKOSBUILTINS_H +#define CLAD_DIFFERENTIATOR_KOKKOSBUILTINS_H + +#include +#include "clad/Differentiator/Differentiator.h" + +namespace clad::custom_derivatives { +namespace class_functions { +/// Kokkos arrays +template +clad::ValueAndPushforward, + Kokkos::View> +constructor_pushforward( + clad::ConstructorPushforwardTag>, + const ::std::string& name, const size_t& idx0, const size_t& idx1, + const size_t& idx2, const size_t& idx3, const size_t& idx4, + const size_t& idx5, const size_t& idx6, const size_t& idx7, + const ::std::string& d_name, const size_t& d_idx0, const size_t& d_idx1, + const size_t& d_idx2, const size_t& d_idx3, const size_t& d_idx4, + const size_t& d_idx5, const size_t& d_idx6, const size_t& d_idx7) { + return {Kokkos::View(name, idx0, idx1, idx2, idx3, + idx4, idx5, idx6, idx7), + Kokkos::View( + "_diff_" + name, idx0, idx1, idx2, idx3, idx4, idx5, idx6, idx7)}; +} +} // namespace class_functions + +/// Kokkos functions +namespace Kokkos { +template +inline void deep_copy_pushforward(const View1& dst, const View2& src, T param, + const View1& d_dst, const View2& d_src, + T d_param) { + deep_copy(dst, src); + deep_copy(d_dst, d_src); +} + +template +inline void +parallel_for_pushforward(const ::std::string& str, const ExecPolicy& policy, + const FunctorType& functor, const ::std::string& d_str, + const ExecPolicy& d_policy, + const FunctorType& d_functor) { + // TODO: implement parallel_for_pushforward + return; +} +} // namespace Kokkos +} // namespace clad::custom_derivatives + +#endif // CLAD_DIFFERENTIATOR_KOKKOSBUILTINS_H \ No newline at end of file diff --git a/unittests/Kokkos/CMakeLists.txt b/unittests/Kokkos/CMakeLists.txt index bd89f86f5..757d55b52 100644 --- a/unittests/Kokkos/CMakeLists.txt +++ b/unittests/Kokkos/CMakeLists.txt @@ -1,5 +1,6 @@ add_clad_unittest(KokkosTests main.cpp + ViewAccess.cpp ViewBasics.cpp ParallelReduce.cpp ParallelFor.cpp diff --git a/unittests/Kokkos/TestUtils.h b/unittests/Kokkos/TestUtils.h index ab55219aa..239aa2f60 100644 --- a/unittests/Kokkos/TestUtils.h +++ b/unittests/Kokkos/TestUtils.h @@ -3,11 +3,10 @@ #ifndef KOKKOS_UNITTEST_UTILS #define KOKKOS_UNITTEST_UTILS -template // comparison with the finite difference approx. has been - // tested in the initial PR for Kokkos-aware Clad by Kim - // Liegeois -T finite_difference_tangent(std::function func, const T& x, - const T& epsilon) { +template // comparison with the finite difference + // approx. has been tested in the initial PR + // for Kokkos-aware Clad by Kim Liegeois +T finite_difference_tangent(F func, const T& x, const T& epsilon) { return (func(x + epsilon) - func(x - epsilon)) / (2 * epsilon); } diff --git a/unittests/Kokkos/ViewAccess.cpp b/unittests/Kokkos/ViewAccess.cpp new file mode 100644 index 000000000..e77b278f0 --- /dev/null +++ b/unittests/Kokkos/ViewAccess.cpp @@ -0,0 +1,73 @@ +// Tests originally created by Kim Liegeois + +#include "TestUtils.h" +#include +#include "clad/Differentiator/Differentiator.h" +#include "clad/Differentiator/KokkosBuiltins.h" +#include "gtest/gtest.h" + +double f(double x, double y) { + + const int N1 = 4; + + Kokkos::View a("a", N1); + Kokkos::View b("b", N1); + + a(0, 0) = x; + b(0, 0) = y; + + b(0, 0) += a(0, 0) * b(0, 0); + + return a(0, 0) * a(0, 0) * b(0, 0) + b(0, 0); +} + +double f_2(double x, double y) { + + const int N1 = 4; + + Kokkos::View a("a", N1); + Kokkos::View b("b", N1); + + Kokkos::deep_copy(a, 3 * x + y); + b(0, 0) = x; + Kokkos::deep_copy(b, a); + + b(0, 0) += a(0, 0) * b(0, 0); + + return a(0, 0); +} + +TEST(ViewAccess, Test1) { + EXPECT_NEAR(f(0, 1), 1, 1e-8); + EXPECT_NEAR(f(0, 2), 2, 1e-8); +} + +TEST(ViewAccess, Test2) { + + double tolerance = 1e-8; + double epsilon = 1e-6; + + auto f_x = clad::differentiate(f, "x"); + + std::function f_tmp = [](double x) { return f(x, 4.); }; + double dx_f_FD = finite_difference_tangent(f_tmp, 3., epsilon); + + EXPECT_NEAR(f_x.execute(3, 4), dx_f_FD, tolerance * dx_f_FD); + + auto f_2_x = clad::differentiate(f_2, "x"); + + std::function f_2_tmp = [](double x) { return f_2(x, 4.); }; + double dx_f_2_FD = finite_difference_tangent(f_2_tmp, 3., epsilon); + EXPECT_NEAR(f_2_x.execute(3, 4), dx_f_2_FD, tolerance * dx_f_2_FD); + + // TODO: uncomment this once it has been implemented + // auto f_grad_exe = clad::gradient(f); + // double dx, dy; + // f_grad_exe.execute(3., 4., &dx, &dy); + // EXPECT_NEAR(f_x.execute(3, 4),dx,tolerance*dx); + + // double dx_2, dy_2; + // auto f_2_grad_exe = clad::gradient(f_2); + // f_2_grad_exe.execute(3., 4., &dx_2, &dy_2); + // EXPECT_NEAR(f_2_x.execute(3, 4),dx_2,tolerance*dx_2); +} \ No newline at end of file diff --git a/unittests/Kokkos/ViewBasics.cpp b/unittests/Kokkos/ViewBasics.cpp index d08c94a61..8773f5451 100644 --- a/unittests/Kokkos/ViewBasics.cpp +++ b/unittests/Kokkos/ViewBasics.cpp @@ -4,12 +4,11 @@ // it has been modified to match gtest guidelines and improve readability #include "ParallelAdd.h" -#include "TestUtils.h" #include #include "clad/Differentiator/Differentiator.h" #include "gtest/gtest.h" -double f(double x, double y) { +double f_basics(double x, double y) { const int N = 2; Kokkos::View a("a", N); @@ -27,10 +26,10 @@ TEST(ViewBasics, TestAccessForward) { // const double tau = 1e-6; // tolerance // // TODO: uncomment this once it has been implemented - // auto f_x = clad::differentiate(f, "x"); + // auto f_x = clad::differentiate(f_basics, "x"); // for (double y = 3; y <= 5; y += 1) { - // std::function f_tmp = [y](double t){ return f(t, y); }; - // for (double x = 3; x <= 5; x += 1) { + // std::function f_tmp = [y](double t){ return f_basics(t, + // y); }; for (double x = 3; x <= 5; x += 1) { // double f_x_ex = f_x.execute(x, y); // double dx_f_FD = finite_difference_tangent(f_tmp, x, eps); // EXPECT_NEAR(f_x_ex, dx_f_FD, abs(tau*dx_f_FD)); @@ -44,10 +43,10 @@ TEST(ViewBasics, TestAccessReverse) { // const double tau = 1e-6; // tolerance // // TODO: uncomment this once it has been implemented - // auto f_grad_exe = clad::gradient(f); + // auto f_grad_exe = clad::gradient(f_basics); // for (double y = 3; y <= 5; y += 1) { - // std::function f_tmp = [y](double t){ return f(t, y); }; - // for (double x = 3; x <= 5; x += 1) { + // std::function f_tmp = [y](double t){ return f_basics(t, + // y); }; for (double x = 3; x <= 5; x += 1) { // double dx_f_FD = finite_difference_tangent(f_tmp, x, eps); // double dx, dy; // f_grad_exe.execute(x, y, &dx, &dy); @@ -56,7 +55,7 @@ TEST(ViewBasics, TestAccessReverse) { // } } -double f_2(double x, double y) { +double f_basics_deep_copy(double x, double y) { const int N = 2; Kokkos::View a("a", N); @@ -77,10 +76,11 @@ TEST(ViewBasics, TestDeepCopyForward) { // const double tau = 1e-6; // tolerance // // TODO: uncomment this once it has been implemented - // auto f_x = clad::differentiate(f_2, "x"); + // auto f_x = clad::differentiate(f_basics_deep_copy, "x"); // for (double y = 3; y <= 5; y += 1) { - // std::function f_tmp = [y](double t){ return f_2(t, y); }; - // for (double x = 3; x <= 5; x += 1) { + // std::function f_tmp = [y](double t){ return + // f_basics_deep_copy(t, y); + // }; for (double x = 3; x <= 5; x += 1) { // double f_x_ex = f_x.execute(x, y); // double dx_f_FD = finite_difference_tangent(f_tmp, x, eps); // EXPECT_NEAR(f_x_ex, dx_f_FD, abs(tau*dx_f_FD)); @@ -94,10 +94,11 @@ TEST(ViewBasics, TestDeepCopyReverse) { // const double tau = 1e-6; // tolerance // // TODO: uncomment this once it has been implemented - // auto f_grad_exe = clad::gradient(f_2); + // auto f_grad_exe = clad::gradient(f_basics_deep_copy); // for (double y = 3; y <= 5; y += 1) { - // std::function f_tmp = [y](double t){ return f_2(t, y); }; - // for (double x = 3; x <= 5; x += 1) { + // std::function f_tmp = [y](double t){ return + // f_basics_deep_copy(t, y); + // }; for (double x = 3; x <= 5; x += 1) { // double dx_f_FD = finite_difference_tangent(f_tmp, x, eps); // double dx, dy; // f_grad_exe.execute(x, y, &dx, &dy); @@ -106,7 +107,7 @@ TEST(ViewBasics, TestDeepCopyReverse) { // } } -double f_3(double x, double y) { +double f_basics_deep_copy_2(double x, double y) { const int N = 2; Kokkos::View a("a", N); @@ -130,10 +131,11 @@ TEST(ViewBasics, TestDeepCopy2Forward) { // const double tau = 1e-6; // tolerance // // TODO: uncomment this once it has been implemented - // auto f_y = clad::differentiate(f_3, "y"); + // auto f_y = clad::differentiate(f_basics_deep_copy_2, "y"); // for (double x = 3; x <= 5; x += 1) { - // std::function f_tmp = [x](double t){ return f_3(x, t); }; - // for (double y = 3; y <= 5; y += 1) { + // std::function f_tmp = [x](double t){ return + // f_basics_deep_copy_2(x, t); + // }; for (double y = 3; y <= 5; y += 1) { // double f_y_ex = f_y.execute(x, y); // double dy_f_FD = finite_difference_tangent(f_tmp, y, eps); // EXPECT_NEAR(f_y_ex, dy_f_FD, abs(tau*dy_f_FD)); @@ -147,10 +149,11 @@ TEST(ViewBasics, TestDeepCopy2Reverse) { // const double tau = 1e-6; // tolerance // // TODO: uncomment this once it has been implemented - // auto f_grad_exe = clad::gradient(f_3); + // auto f_grad_exe = clad::gradient(f_basics_deep_copy_2); // for (double x = 3; x <= 5; x += 1) { - // std::function f_tmp = [x](double t){ return f_3(x, t); }; - // for (double y = 3; y <= 5; y += 1) { + // std::function f_tmp = [x](double t){ return + // f_basics_deep_copy_2(x, t); + // }; for (double y = 3; y <= 5; y += 1) { // double dy_f_FD = finite_difference_tangent(f_tmp, y, eps); // double dx, dy; // f_grad_exe.execute(x, y, &dx, &dy);