From 85f0bf2764a2fd52c2292e14f1f2d35892b31593 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. --- unittests/Kokkos/CMakeLists.txt | 2 +- unittests/Kokkos/KokkosImplementation.h | 74 +++++++++++ unittests/Kokkos/ViewAccess.cpp | 76 +++++++++++ unittests/Kokkos/ViewBasics.cpp | 160 ------------------------ 4 files changed, 151 insertions(+), 161 deletions(-) create mode 100644 unittests/Kokkos/KokkosImplementation.h create mode 100644 unittests/Kokkos/ViewAccess.cpp delete mode 100644 unittests/Kokkos/ViewBasics.cpp diff --git a/unittests/Kokkos/CMakeLists.txt b/unittests/Kokkos/CMakeLists.txt index bd89f86f5..1964cbc95 100644 --- a/unittests/Kokkos/CMakeLists.txt +++ b/unittests/Kokkos/CMakeLists.txt @@ -1,6 +1,6 @@ add_clad_unittest(KokkosTests main.cpp - ViewBasics.cpp + ViewAccess.cpp ParallelReduce.cpp ParallelFor.cpp ) diff --git a/unittests/Kokkos/KokkosImplementation.h b/unittests/Kokkos/KokkosImplementation.h new file mode 100644 index 000000000..6162a6a34 --- /dev/null +++ b/unittests/Kokkos/KokkosImplementation.h @@ -0,0 +1,74 @@ +// This header file contains the implementation of the Kokkos framework +// differentiation support in Clad in the form of custom pushforwards and +// pullbacks. + +#ifndef KOKKOS_IMPLEMENTATION +#define KOKKOS_IMPLEMENTATION + +#include +#include "clad/Differentiator/Differentiator.h" + +namespace clad { +namespace 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)}; +}; + +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 custom_derivatives +} // namespace clad + +#endif // #ifndef KOKKOS_IMPLEMENTATION \ No newline at end of file diff --git a/unittests/Kokkos/ViewAccess.cpp b/unittests/Kokkos/ViewAccess.cpp new file mode 100644 index 000000000..0be2aa5f0 --- /dev/null +++ b/unittests/Kokkos/ViewAccess.cpp @@ -0,0 +1,76 @@ +// Source: +// https://github.com/vgvassilev/clad/pull/783/files#diff-c919937640674e1911d1f17d3df2b65426d152400589873ff9ba2f0391048d78 +// Source: https://github.com/kliegeois/clad/tree/kokkos-PR +// Originally created by Kim Liegeois + +#include "KokkosImplementation.h" +#include "TestUtils.h" +#include +#include "clad/Differentiator/Differentiator.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 deleted file mode 100644 index d08c94a61..000000000 --- a/unittests/Kokkos/ViewBasics.cpp +++ /dev/null @@ -1,160 +0,0 @@ -// Very basic Kokkos::View usage test that should work by all means -// inspired by -// https://github.com/kliegeois/clad/blob/kokkos-PR/unittests/Kokkos/view_access.cpp -// 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) { - const int N = 2; - - Kokkos::View a("a", N); - Kokkos::View b("b", N); - - a(0, 0) = x; - b(0, 0) = y * x; - - return a(0, 0) + a(0, 0) * b(0, 0) + b(0, 0); -} - -TEST(ViewBasics, TestAccessForward) { - // // check finite difference and forward mode similarity - // const double eps = 1e-5; - // const double tau = 1e-6; // tolerance - - // // TODO: uncomment this once it has been implemented - // auto f_x = clad::differentiate(f, "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) { - // 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)); - // } - // } -} - -TEST(ViewBasics, TestAccessReverse) { - // // check reverse mode and forward mode similarity - // const double eps = 1e-5; - // const double tau = 1e-6; // tolerance - - // // TODO: uncomment this once it has been implemented - // auto f_grad_exe = clad::gradient(f); - // 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) { - // double dx_f_FD = finite_difference_tangent(f_tmp, x, eps); - // double dx, dy; - // f_grad_exe.execute(x, y, &dx, &dy); - // EXPECT_NEAR(dx_f_FD, dx, abs(tau*dx)); - // } - // } -} - -double f_2(double x, double y) { - const int N = 2; - - Kokkos::View a("a", N); - Kokkos::View b("b", N); - - Kokkos::deep_copy(a, 3 * x + y); - b(0, 0) = x; - Kokkos::deep_copy(b, a); - - b(0, 0) = b(0, 0) + a(0, 0) * b(0, 0); - - return a(0, 0); // derivative wrt x is constantly 3 -} - -TEST(ViewBasics, TestDeepCopyForward) { - // // check finite difference and forward mode similarity - // const double eps = 1e-5; - // const double tau = 1e-6; // tolerance - - // // TODO: uncomment this once it has been implemented - // auto f_x = clad::differentiate(f_2, "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) { - // 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)); - // } - // } -} - -TEST(ViewBasics, TestDeepCopyReverse) { - // // check reverse mode and forward mode similarity - // const double eps = 1e-5; - // const double tau = 1e-6; // tolerance - - // // TODO: uncomment this once it has been implemented - // auto f_grad_exe = clad::gradient(f_2); - // 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) { - // double dx_f_FD = finite_difference_tangent(f_tmp, x, eps); - // double dx, dy; - // f_grad_exe.execute(x, y, &dx, &dy); - // EXPECT_NEAR(dx_f_FD, dx, abs(tau*dx)); - // } - // } -} - -double f_3(double x, double y) { - const int N = 2; - - Kokkos::View a("a", N); - Kokkos::View b("b", N); - - Kokkos::deep_copy(a, 3 * y + x + 50); - b(1) = x * y; - Kokkos::deep_copy(b, a); - - b(1) = b(1) + a(0) * b(1); - - a(1) = x * x * x; - a(0) += a(1); - - return a(0); // derivative of this wrt y is constantly 3 -} - -TEST(ViewBasics, TestDeepCopy2Forward) { - // // check finite difference and forward mode similarity - // const double eps = 1e-5; - // const double tau = 1e-6; // tolerance - - // // TODO: uncomment this once it has been implemented - // auto f_y = clad::differentiate(f_3, "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) { - // 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)); - // } - // } -} - -TEST(ViewBasics, TestDeepCopy2Reverse) { - // // check reverse mode and forward mode similarity - // const double eps = 1e-5; - // const double tau = 1e-6; // tolerance - - // // TODO: uncomment this once it has been implemented - // auto f_grad_exe = clad::gradient(f_3); - // 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) { - // double dy_f_FD = finite_difference_tangent(f_tmp, y, eps); - // double dx, dy; - // f_grad_exe.execute(x, y, &dx, &dy); - // EXPECT_NEAR(dy_f_FD, dy, abs(tau*dy)); - // } - // } -} \ No newline at end of file