Add Kokkos unittests

vgvassilev · Mar 28, 2024 · bd74b58 · bd74b58
1 parent 2cd0b32
commit bd74b58
Show file tree

Hide file tree

Showing 7 changed files with 713 additions and 22 deletions.
diff --git a/unittests/Kokkos/CMakeLists.txt b/unittests/Kokkos/CMakeLists.txt
@@ -1,6 +1,8 @@
 add_clad_unittest(KokkosTests
-  parallel_for.cpp
   main.cpp
+  ViewBasics.cpp
+  ParallelReduce.cpp
+  ParallelFor.cpp
   )
 
 # If llvm does not require rtti, kokkos does.

diff --git a/unittests/Kokkos/ParallelFor.cpp b/unittests/Kokkos/ParallelFor.cpp
@@ -0,0 +1,95 @@
+#include <Kokkos_Core.hpp>
+#include "clad/Differentiator/Differentiator.h"
+#include "gtest/gtest.h"
+// #include "TestUtils.hpp"
+#include "parallel_add.hpp"
+
+TEST(ParallelFor, HelloWorldLambdaLoopForward) {
+    // // check finite difference and forward mode similarity
+    // const double eps = 1e-5;
+    // const double tau = 1e-6; // tolerance
+
+    std::function<double(double)> _f = [](double x) {
+        using Policy = Kokkos::RangePolicy<Kokkos::DefaultHostExecutionSpace>;
+        double res[5] = {0};
+        Kokkos::parallel_for("HelloWorld-forward", Policy(0, 5),
+        [&res, x](const int i) {
+            res[i] = x*x;
+        });
+        // all elements of res should be the same, so return one of them arbitrarily
+        return res[2];
+    };
+
+    // // TODO: uncomment this once it has been implemented
+    // auto f_diff = clad::differentiate(_f, 0/*x*/);
+    // for (double x = -2; x <= 2; x += 1) {
+    //   double f_diff_ex = f_diff.execute(x);
+    //   double dx_f_FD = finite_difference_tangent(_f, x, eps);
+    //   EXPECT_NEAR(f_diff_ex, dx_f_FD, abs(tau*dx_f_FD));  
+    // }
+}
+
+TEST(ParallelFor, HelloWorldLambdaLoopReverse) {
+    // // check finite difference and reverse mode similarity
+    // const double eps = 1e-5;
+    // const double tau = 1e-6; // tolerance
+
+    std::function<double(double)> _f = [](double x) {
+        using Policy = Kokkos::RangePolicy<Kokkos::DefaultHostExecutionSpace>;
+        double res[5] = {0};
+        Kokkos::parallel_for("HelloWorld-reverse", Policy(0, 5),
+        [&res, x](const int i) {
+            res[i] = x*x;
+        });
+        // all elements of res should be the same, so return one of them arbitrarily
+        return res[2];
+    };
+
+    // // TODO: uncomment this once it has been implemented
+    // auto f_grad = clad::gradient(_f);
+    // for (double x = -2; x <= 2; x += 1) {
+    //   double dx_f_FD = finite_difference_tangent(_f, x, eps);
+    //   double dx;
+    //   f_grad.execute(x, &dx);
+    //   EXPECT_NEAR(dx_f_FD, dx, abs(tau*dx));
+    // }
+}
+
+double parallel_polynomial_for(double x) { // data races
+    Kokkos::View <double[1], Kokkos::HostSpace> res("res");
+    res(0) = 0;
+    Kokkos::parallel_for("polycalc", 5, KOKKOS_LAMBDA(const int i) {
+        res(0) += pow(x, i+1)/(i+1);
+    });
+    // for (int i = 0; i < 6; ++i) {
+    //   res(0) += pow(x, i+1)/(i+1);
+    // }
+    return res(0);
+}
+
+TEST(ParallelFor, ParallelPolynomialForward) {
+    // // check true derivative and forward mode similarity
+    // const double tau = 1e-5; // tolerance
+
+    // // TODO: uncomment this once it has been implemented
+    // auto f_diff = clad::differentiate(parallel_polynomial_for, "x");
+    // for (double x = -2; x <= 2; x += 1) {
+    //   double f_diff_ex = f_diff.execute(x);
+    //   double dx_f_true = parallel_polynomial_true_derivative(x);
+    //   EXPECT_NEAR(f_diff_ex, dx_f_true, abs(tau*dx_f_true)); 
+    // }
+}
+
+TEST(ParallelFor, ParallelPolynomialReverse) {
+    // // check true derivative and reverse mode similarity
+    // const double tau = 1e-5; // tolerance
+
+    // // TODO: uncomment this once it has been implemented
+    // auto f_grad = clad::gradient(parallel_polynomial_for);
+    // for (double x = -2; x <= 2; x += 1) {
+    //   double dx_f_true = parallel_polynomial_true_derivative(x);
+    //   double dx = 0;
+    //   f_grad.execute(x, &dx);
+    //   EXPECT_NEAR(dx_f_true, dx, abs(tau*dx));
+    // }
+}
diff --git a/unittests/Kokkos/ParallelReduce.cpp b/unittests/Kokkos/ParallelReduce.cpp
@@ -0,0 +1,87 @@
+#include <Kokkos_Core.hpp>
+#include "clad/Differentiator/Differentiator.h"
+#include "gtest/gtest.h"
+// #include "TestUtils.hpp"
+#include "parallel_add.hpp"
+
+TEST(ParallelReduce, HelloWorldLambdaLoopForward) {
+    // // check finite difference and forward mode similarity
+    // const double eps = 1e-5;
+    // const double tau = 1e-6; // tolerance
+
+    std::function<double(double)> _f = [](double x) {
+      double res = 0.;
+      Kokkos::parallel_reduce("HelloWorld-forward", 5, KOKKOS_LAMBDA(const int& i, double& _res){ _res += x; }, res);
+      // res = 5*x;
+      return res;
+    };
+
+    // // TODO: uncomment this once it has been implemented
+    // auto f_diff = clad::differentiate(_f, 0/*x*/);
+    // for (double x = -2; x <= 2; x += 1) {
+    //   double f_diff_ex = f_diff.execute(x);
+    //   double dx_f_FD = finite_difference_tangent(_f, x, eps);
+    //   EXPECT_NEAR(f_diff_ex, dx_f_FD, abs(tau*dx_f_FD));  
+    // }
+}
+
+TEST(ParallelReduce, HelloWorldLambdaLoopReverse) {
+    // // check finite difference and reverse mode similarity
+    // const double eps = 1e-5;
+    // const double tau = 1e-6; // tolerance
+
+    std::function<double(double)> _f = [](double x) {
+      double res = 0.;
+      Kokkos::parallel_reduce("HelloWorld-reverse", 5, KOKKOS_LAMBDA(const int& i, double& _res){ _res += x; }, res);
+      // res = 5*x;
+      return res;
+    };
+
+    // // TODO: uncomment this once it has been implemented
+    // auto f_grad = clad::gradient(_f);
+    // for (double x = -2; x <= 2; x += 1) {
+    //   double dx_f_FD = finite_difference_tangent(_f, x, eps);
+    //   double dx;
+    //   f_grad.execute(x, &dx);
+    //   EXPECT_NEAR(dx_f_FD, dx, abs(tau*dx));
+    // }
+}
+
+double parallel_polynomial_reduce(double x) {
+    Kokkos::View <double[1], Kokkos::HostSpace> res("res");
+    res(0) = 0;
+    Kokkos::parallel_reduce("polycalc", 5, KOKKOS_LAMBDA(const int& i, double& _res) {
+        _res += pow(x, i+1)/(i+1);
+    }, res(0));
+    // for (int i = 0; i < 6; ++i) {
+    //   res(0) += pow(x, i+1)/(i+1);
+    // }
+    return res(0);
+}
+
+TEST(ParallelReduce, ParallelPolynomialForward) {
+    // // check true derivative and forward mode similarity
+    // const double tau = 1e-5; // tolerance
+
+    // // TODO: uncomment this once it has been implemented
+    // auto f_diff = clad::differentiate(parallel_polynomial_reduce, "x");
+    // for (double x = -2; x <= 2; x += 1) {
+    //   double f_diff_ex = f_diff.execute(x);
+    //   double dx_f_true = parallel_polynomial_true_derivative(x);
+    //   EXPECT_NEAR(f_diff_ex, dx_f_true, abs(tau*dx_f_true)); 
+    // }
+}
+
+TEST(ParallelReduce, ParallelPolynomialReverse) {
+    // // check true derivative and reverse mode similarity
+    // const double tau = 1e-5; // tolerance
+
+    // // TODO: uncomment this once it has been implemented
+    // auto f_grad = clad::gradient(parallel_polynomial_reduce);
+    // for (double x = -2; x <= 2; x += 1) {
+    //   double dx_f_true = parallel_polynomial_true_derivative(x);
+    //   double dx = 0;
+    //   f_grad.execute(x, &dx);
+    //   EXPECT_NEAR(dx_f_true, dx, abs(tau*dx));
+    // }
+}
diff --git a/unittests/Kokkos/TestUtils.hpp b/unittests/Kokkos/TestUtils.hpp
@@ -0,0 +1,21 @@
+// Useful things
+
+#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) {
+  return (func(x+epsilon)-func(x-epsilon)) / (2 * epsilon);
+}
+
+double parallel_polynomial_true_derivative(double x) { // the true derivative of the polynomial tested in ParallelFor.cpp and ParallelReduce.cpp
+    double res = 0;
+    double x_c = 1;
+    for (unsigned i = 0; i < 6; ++i) {
+        res += x_c;
+        x_c *= x;
+    }
+    return res;
+}
+
+#endif
diff --git a/unittests/Kokkos/ViewBasics.cpp b/unittests/Kokkos/ViewBasics.cpp
@@ -0,0 +1,159 @@
+// 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 <Kokkos_Core.hpp>
+#include "clad/Differentiator/Differentiator.h"
+#include "gtest/gtest.h"
+#include "TestUtils.hpp"
+#include "parallel_add.hpp"
+
+double f(double x, double y) {
+  const int N = 2;
+
+  Kokkos::View<double*[N], Kokkos::HostSpace> a("a", N);
+  Kokkos::View<double*[N], Kokkos::HostSpace> 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<double(double)> 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<double(double)> 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<double*[4], Kokkos::LayoutLeft, Kokkos::HostSpace> a("a", N);
+  Kokkos::View<double*[4], Kokkos::LayoutLeft, Kokkos::HostSpace> 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<double(double)> 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<double(double)> 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<double*, Kokkos::LayoutLeft, Kokkos::HostSpace> a("a", N);
+  Kokkos::View<double*, Kokkos::LayoutLeft, Kokkos::HostSpace> 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<double(double)> 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<double(double)> 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));
+  //   }
+  // }
+}