Skip to content

Commit

Permalink
Add basic Kokkos support for the original tests of #783
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
gojakuch authored and vgvassilev committed Jul 22, 2024
1 parent ed26dbd commit f03fc98
Show file tree
Hide file tree
Showing 5 changed files with 157 additions and 27 deletions.
54 changes: 54 additions & 0 deletions include/clad/Differentiator/KokkosBuiltins.h
Original file line number Diff line number Diff line change
@@ -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 <Kokkos_Core.hpp>
#include "clad/Differentiator/Differentiator.h"

namespace clad::custom_derivatives {
namespace class_functions {
/// Kokkos arrays
template <class DataType, class... ViewParams>
clad::ValueAndPushforward<Kokkos::View<DataType, ViewParams...>,
Kokkos::View<DataType, ViewParams...>>
constructor_pushforward(
clad::ConstructorPushforwardTag<Kokkos::View<DataType, ViewParams...>>,
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<DataType, ViewParams...>(name, idx0, idx1, idx2, idx3,
idx4, idx5, idx6, idx7),
Kokkos::View<DataType, ViewParams...>(
"_diff_" + name, idx0, idx1, idx2, idx3, idx4, idx5, idx6, idx7)};
}
} // namespace class_functions

/// Kokkos functions
namespace Kokkos {
template <typename View1, typename View2, typename T>
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 <class ExecPolicy, class FunctorType>
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
1 change: 1 addition & 0 deletions unittests/Kokkos/CMakeLists.txt
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
add_clad_unittest(KokkosTests
main.cpp
ViewAccess.cpp
ViewBasics.cpp
ParallelReduce.cpp
ParallelFor.cpp
Expand Down
9 changes: 4 additions & 5 deletions unittests/Kokkos/TestUtils.h
Original file line number Diff line number Diff line change
Expand Up @@ -3,11 +3,10 @@
#ifndef KOKKOS_UNITTEST_UTILS
#define KOKKOS_UNITTEST_UTILS

template <typename T> // 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<T(T)> func, const T& x,
const T& epsilon) {
template <typename F, typename T> // 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);
}

Expand Down
73 changes: 73 additions & 0 deletions unittests/Kokkos/ViewAccess.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,73 @@
// Tests originally created by Kim Liegeois

#include "TestUtils.h"
#include <Kokkos_Core.hpp>
#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<double* [4], Kokkos::LayoutLeft, Kokkos::HostSpace> a("a", N1);
Kokkos::View<double* [4], Kokkos::LayoutLeft, Kokkos::HostSpace> 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<double* [4], Kokkos::LayoutLeft, Kokkos::HostSpace> a("a", N1);
Kokkos::View<double* [4], Kokkos::LayoutLeft, Kokkos::HostSpace> 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<double(double)> 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<double(double)> 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);
}
47 changes: 25 additions & 22 deletions unittests/Kokkos/ViewBasics.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -4,12 +4,11 @@
// it has been modified to match gtest guidelines and improve readability

#include "ParallelAdd.h"
#include "TestUtils.h"
#include <Kokkos_Core.hpp>
#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<double* [N], Kokkos::HostSpace> a("a", N);
Expand All @@ -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<double(double)> f_tmp = [y](double t){ return f(t, y); };
// for (double x = 3; x <= 5; x += 1) {
// std::function<double(double)> 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));
Expand All @@ -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<double(double)> f_tmp = [y](double t){ return f(t, y); };
// for (double x = 3; x <= 5; x += 1) {
// std::function<double(double)> 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);
Expand All @@ -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<double* [4], Kokkos::LayoutLeft, Kokkos::HostSpace> a("a", N);
Expand All @@ -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<double(double)> f_tmp = [y](double t){ return f_2(t, y); };
// for (double x = 3; x <= 5; x += 1) {
// std::function<double(double)> 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));
Expand All @@ -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<double(double)> f_tmp = [y](double t){ return f_2(t, y); };
// for (double x = 3; x <= 5; x += 1) {
// std::function<double(double)> 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);
Expand All @@ -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<double*, Kokkos::LayoutLeft, Kokkos::HostSpace> a("a", N);
Expand All @@ -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<double(double)> f_tmp = [x](double t){ return f_3(x, t); };
// for (double y = 3; y <= 5; y += 1) {
// std::function<double(double)> 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));
Expand All @@ -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<double(double)> f_tmp = [x](double t){ return f_3(x, t); };
// for (double y = 3; y <= 5; y += 1) {
// std::function<double(double)> 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);
Expand Down

0 comments on commit f03fc98

Please sign in to comment.