From 35000d0cd66f3d67f4593348597637893df6daf5 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 | 74 +++++++++++++++++++ unittests/Kokkos/CMakeLists.txt | 1 + unittests/Kokkos/ViewAccess.cpp | 76 ++++++++++++++++++++ unittests/Kokkos/ViewBasics.cpp | 32 ++++----- 4 files changed, 167 insertions(+), 16 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..4f577a7eb --- /dev/null +++ b/include/clad/Differentiator/KokkosBuiltins.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. Please include it manually to use Clad on Kokkos code. + +#ifndef CLAD_KOKKOS_BUILTINS_H +#define CLAD_KOKKOS_BUILTINS_H + +#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 CLAD_KOKKOS_BUILTINS_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/ViewAccess.cpp b/unittests/Kokkos/ViewAccess.cpp new file mode 100644 index 000000000..43dd17cfb --- /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 "clad/Differentiator/KokkosBuiltins.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 index d08c94a61..629826033 100644 --- a/unittests/Kokkos/ViewBasics.cpp +++ b/unittests/Kokkos/ViewBasics.cpp @@ -4,12 +4,12 @@ // it has been modified to match gtest guidelines and improve readability #include "ParallelAdd.h" -#include "TestUtils.h" +// #include "TestUtils.h" #include #include "clad/Differentiator/Differentiator.h" #include "gtest/gtest.h" -double f(double x, double y) { +double ff(double x, double y) { const int N = 2; Kokkos::View a("a", N); @@ -27,9 +27,9 @@ 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(ff, "x"); // for (double y = 3; y <= 5; y += 1) { - // std::function f_tmp = [y](double t){ return f(t, y); }; + // std::function f_tmp = [y](double t){ return ff(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); @@ -44,9 +44,9 @@ 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(ff); // for (double y = 3; y <= 5; y += 1) { - // std::function f_tmp = [y](double t){ return f(t, y); }; + // std::function f_tmp = [y](double t){ return ff(t, y); }; // for (double x = 3; x <= 5; x += 1) { // double dx_f_FD = finite_difference_tangent(f_tmp, x, eps); // double dx, dy; @@ -56,7 +56,7 @@ TEST(ViewBasics, TestAccessReverse) { // } } -double f_2(double x, double y) { +double ff_2(double x, double y) { const int N = 2; Kokkos::View a("a", N); @@ -77,9 +77,9 @@ 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(ff_2, "x"); // for (double y = 3; y <= 5; y += 1) { - // std::function f_tmp = [y](double t){ return f_2(t, y); }; + // std::function f_tmp = [y](double t){ return ff_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); @@ -94,9 +94,9 @@ 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(ff_2); // for (double y = 3; y <= 5; y += 1) { - // std::function f_tmp = [y](double t){ return f_2(t, y); }; + // std::function f_tmp = [y](double t){ return ff_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; @@ -106,7 +106,7 @@ TEST(ViewBasics, TestDeepCopyReverse) { // } } -double f_3(double x, double y) { +double ff_3(double x, double y) { const int N = 2; Kokkos::View a("a", N); @@ -130,9 +130,9 @@ 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(ff_3, "y"); // for (double x = 3; x <= 5; x += 1) { - // std::function f_tmp = [x](double t){ return f_3(x, t); }; + // std::function f_tmp = [x](double t){ return ff_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); @@ -147,9 +147,9 @@ 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(ff_3); // for (double x = 3; x <= 5; x += 1) { - // std::function f_tmp = [x](double t){ return f_3(x, t); }; + // std::function f_tmp = [x](double t){ return ff_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;