-
Notifications
You must be signed in to change notification settings - Fork 123
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
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.
- Loading branch information
1 parent
ed26dbd
commit f03fc98
Showing
5 changed files
with
157 additions
and
27 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 | ||
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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); | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters