Enable Kokkos mode from ASE calculator, as well as float32

libsymmetrix/CMakeLists.txt

@@ -11,6 +11,7 @@ set(CMAKE_POSITION_INDEPENDENT_CODE ON)
 
 find_package(BLAS REQUIRED)
 add_subdirectory(external/sphericart)
+option(SYMMETRIX_KOKKOS "Enable Kokkos mode in Symmetrix." ON)
 if(SYMMETRIX_KOKKOS)
     if(TARGET Kokkos::kokkos)
         message(STATUS "Symmetrix: Found existing Kokkos target.")

libsymmetrix/source/mace_kokkos.hpp

@@ -13,6 +13,7 @@
 #include "radial_function_set_kokkos.hpp"
 #include "zbl_kokkos.hpp"
 
+template <typename Precision>
 class MACEKokkos {
 
 public:
@@ -45,31 +46,31 @@ ZBLKokkos zbl;
 
 // R0
 double R0_spline_h;
-Kokkos::View<const double****,Kokkos::LayoutRight> R0_spline_coefficients;
-Kokkos::View<double**,Kokkos::LayoutRight> R0, R0_deriv;
+Kokkos::View<const Precision****,Kokkos::LayoutRight> R0_spline_coefficients;
+Kokkos::View<Precision**,Kokkos::LayoutRight> R0, R0_deriv;
 void compute_R0(const int num_nodes,
                 Kokkos::View<const int*> node_types,
                 Kokkos::View<const int*> num_neigh,
                 Kokkos::View<const int*> neigh_types,
                 Kokkos::View<const double*> r);
 
 // R1
-RadialFunctionSetKokkos radial_1;
-Kokkos::View<double**,Kokkos::LayoutRight> R1, R1_deriv;
+RadialFunctionSetKokkos<Precision> radial_1;
+Kokkos::View<Precision**,Kokkos::LayoutRight> R1, R1_deriv;
 void compute_R1(const int num_nodes,
                 Kokkos::View<const int*> node_types,
                 Kokkos::View<const int*> num_neigh,
                 Kokkos::View<const int*> neigh_types,
                 Kokkos::View<const double*> r);
 
 // Spherical harmonics
-Kokkos::View<double*> xyz_shuffled;
-Kokkos::View<double*> Y, Y_grad;// TODO: make multidimensional
-Kokkos::View<double*> Y_grad_shuffled;
+Kokkos::View<Precision*> xyz_shuffled;
+Kokkos::View<Precision*> Y, Y_grad;// TODO: make multidimensional
+Kokkos::View<Precision*> Y_grad_shuffled;
 void compute_Y(Kokkos::View<const double*> xyz);
 
 // A0
-Kokkos::View<double***,Kokkos::LayoutRight> A0, A0_adj;
+Kokkos::View<Precision***,Kokkos::LayoutRight> A0, A0_adj;
 void compute_A0(const int num_nodes,
                 Kokkos::View<const int*> node_types,
                 Kokkos::View<const int*> num_neigh,
@@ -83,7 +84,7 @@ void reverse_A0(const int num_nodes,
 
 // A0 rescaling
 bool A0_scaled;
-RadialFunctionSetKokkos A0_splines;
+RadialFunctionSetKokkos<double> A0_splines;
 Kokkos::View<double**,Kokkos::LayoutRight> A0_spline_values;
 Kokkos::View<double**,Kokkos::LayoutRight> A0_spline_derivs;
 void compute_A0_scaled(
@@ -101,45 +102,45 @@ void reverse_A0_scaled(
     Kokkos::View<const double*> r);
 
 // M0
-Kokkos::View<double***,Kokkos::LayoutRight> M0, M0_adj;
+Kokkos::View<Precision***,Kokkos::LayoutRight> M0, M0_adj;
 Kokkos::View<Kokkos::View<int**,Kokkos::LayoutRight>*,Kokkos::SharedSpace> M0_monomials;
-Kokkos::View<Kokkos::View<double***,Kokkos::LayoutRight>*,Kokkos::SharedSpace> M0_weights;
+Kokkos::View<Kokkos::View<Precision***,Kokkos::LayoutRight>*,Kokkos::SharedSpace> M0_weights;
 Kokkos::View<Kokkos::View<int**,Kokkos::LayoutRight>*,Kokkos::SharedSpace> M0_poly_spec;
-Kokkos::View<Kokkos::View<double***,Kokkos::LayoutRight>*,Kokkos::SharedSpace> M0_poly_coeff;
-Kokkos::View<Kokkos::View<double***,Kokkos::LayoutRight>*,Kokkos::SharedSpace> M0_poly_values;
-Kokkos::View<Kokkos::View<double***,Kokkos::LayoutRight>*,Kokkos::SharedSpace> M0_poly_adjoints;
+Kokkos::View<Kokkos::View<Precision***,Kokkos::LayoutRight>*,Kokkos::SharedSpace> M0_poly_coeff;
+Kokkos::View<Kokkos::View<Precision***,Kokkos::LayoutRight>*,Kokkos::SharedSpace> M0_poly_values;
+Kokkos::View<Kokkos::View<Precision***,Kokkos::LayoutRight>*,Kokkos::SharedSpace> M0_poly_adjoints;
 void compute_M0(const int num_nodes, Kokkos::View<const int*> node_types);
 void reverse_M0(const int num_nodes, Kokkos::View<const int*> node_types);
 
 // H1
-Kokkos::View<double***,Kokkos::LayoutRight> H1, H1_adj;
-Kokkos::View<double***,Kokkos::LayoutRight> H1_weights;
+Kokkos::View<Precision***,Kokkos::LayoutRight> H1, H1_adj;
+Kokkos::View<Precision***,Kokkos::LayoutRight> H1_weights;
 void compute_H1(const int num_nodes);
 void reverse_H1(const int num_nodes);
 
 // Phi1
 int num_lelm1lm2, num_lme;
 Kokkos::View<int*> Phi1_l, Phi1_l1, Phi1_l2;
 Kokkos::View<int*> Phi1_lme, Phi1_lelm1lm2;
-Kokkos::View<double*> Phi1_clebsch_gordan;
-Kokkos::View<double***,Kokkos::LayoutRight> Phi1r, dPhi1r;
-Kokkos::View<double***,Kokkos::LayoutRight> Phi1, dPhi1;
+Kokkos::View<Precision*> Phi1_clebsch_gordan;
+Kokkos::View<Precision***,Kokkos::LayoutRight> Phi1r, dPhi1r;
+Kokkos::View<Precision***,Kokkos::LayoutRight> Phi1, dPhi1;
 void compute_Phi1(const int num_nodes, Kokkos::View<const int*> num_neigh, Kokkos::View<const int*> neigh_indices);
 void reverse_Phi1(const int num_nodes, Kokkos::View<const int*> num_neigh, Kokkos::View<const int*> neigh_indices, Kokkos::View<const double*> xyz, Kokkos::View<const double*> r, bool zero_dxyz = true, bool zero_H1_adj = true);
 
 // TODO for testing of Phi1 strategies
 Kokkos::View<int*> Phi1_lm1, Phi1_lm2, Phi1_lel1l2;
 
 // A1
-Kokkos::View<double***,Kokkos::LayoutRight> A1, A1_adj;
-Kokkos::View<Kokkos::View<double**,Kokkos::LayoutRight>*,Kokkos::SharedSpace> A1_weights;
-Kokkos::View<Kokkos::View<double**,Kokkos::LayoutRight>*,Kokkos::SharedSpace> A1_weights_trans;
+Kokkos::View<Precision***,Kokkos::LayoutRight> A1, A1_adj;
+Kokkos::View<Kokkos::View<Precision**,Kokkos::LayoutRight>*,Kokkos::SharedSpace> A1_weights;
+Kokkos::View<Kokkos::View<Precision**,Kokkos::LayoutRight>*,Kokkos::SharedSpace> A1_weights_trans;
 void compute_A1(int num_nodes);
 void reverse_A1(int num_nodes);
 
 // A1 rescaling
 bool A1_scaled;
-RadialFunctionSetKokkos A1_splines;
+RadialFunctionSetKokkos<double> A1_splines;
 Kokkos::View<double**,Kokkos::LayoutRight> A1_spline_values;
 Kokkos::View<double**,Kokkos::LayoutRight> A1_spline_derivs;
 void compute_A1_scaled(
@@ -157,13 +158,13 @@ void reverse_A1_scaled(
     Kokkos::View<const double*> r);
 
 // M1
-Kokkos::View<double**,Kokkos::LayoutRight> M1, M1_adj;
+Kokkos::View<Precision**,Kokkos::LayoutRight> M1, M1_adj;
 Kokkos::View<int**,Kokkos::LayoutRight> M1_monomials;
-Kokkos::View<double***,Kokkos::LayoutRight> M1_weights;
+Kokkos::View<Precision***,Kokkos::LayoutRight> M1_weights;
 Kokkos::View<int**,Kokkos::LayoutRight> M1_poly_spec;
-Kokkos::View<double***,Kokkos::LayoutRight> M1_poly_coeff;
-Kokkos::View<double***,Kokkos::LayoutRight> M1_poly_values;
-Kokkos::View<double***,Kokkos::LayoutRight> M1_poly_adjoints;
+Kokkos::View<Precision***,Kokkos::LayoutRight> M1_poly_coeff;
+Kokkos::View<Precision***,Kokkos::LayoutRight> M1_poly_values;
+Kokkos::View<Precision***,Kokkos::LayoutRight> M1_poly_adjoints;
 void compute_M1(int num_nodes, Kokkos::View<const int*> node_types);
 void reverse_M1(int num_nodes, Kokkos::View<const int*> node_types);
 

libsymmetrix/source/radial_function_set_kokkos.cpp

@@ -1,14 +1,15 @@
 #include <vector>
+#include <cmath>
 
 #include "radial_function_set_kokkos.hpp"
 
-
-RadialFunctionSetKokkos::RadialFunctionSetKokkos()
+template <typename Precision>
+RadialFunctionSetKokkos<Precision>::RadialFunctionSetKokkos()
 {
 }
 
-
-RadialFunctionSetKokkos::RadialFunctionSetKokkos(
+template <typename Precision>
+RadialFunctionSetKokkos<Precision>::RadialFunctionSetKokkos(
     double h,
     std::vector<std::vector<std::vector<double>>> node_values,
     std::vector<std::vector<std::vector<double>>> node_derivatives)
@@ -19,18 +20,20 @@ RadialFunctionSetKokkos::RadialFunctionSetKokkos(
     num_functions = node_values[0].size();
     num_nodes = node_values[0][0].size();
 
-    auto c = Kokkos::View<double****, Kokkos::LayoutRight>(
+    auto c = Kokkos::View<Precision****, Kokkos::LayoutRight>(
         "coefficients", num_edge_types, num_nodes-1, 4, num_functions);
     auto h_c = Kokkos::create_mirror_view(c);
     for (int a=0; a<num_edge_types; ++a) {
         for (int i=0; i<num_nodes-1; ++i) {
             for (int j=0; j<num_functions; ++j) {
-                h_c(a,i,0,j) = node_values[a][j][i];
-                h_c(a,i,1,j) = node_derivatives[a][j][i];
-                h_c(a,i,2,j) = (-3*node_values[a][j][i] -2*h*node_derivatives[a][j][i]
-                                + 3*node_values[a][j][i+1] - h*node_derivatives[a][j][i+1]) / (h*h);
-                h_c(a,i,3,j) = (2*node_values[a][j][i] + h*node_derivatives[a][j][i]
-                                - 2*node_values[a][j][i+1] + h*node_derivatives[a][j][i+1]) / (h*h*h);
+                h_c(a,i,0,j) = static_cast<Precision>(node_values[a][j][i]);
+                h_c(a,i,1,j) = static_cast<Precision>(node_derivatives[a][j][i]);
+                h_c(a,i,2,j) = static_cast<Precision>(
+                    (-3*node_values[a][j][i] -2*h*node_derivatives[a][j][i]
+                        + 3*node_values[a][j][i+1] - h*node_derivatives[a][j][i+1]) / (h*h));
+                h_c(a,i,3,j) = static_cast<Precision>(
+                    (2*node_values[a][j][i] + h*node_derivatives[a][j][i]
+                        - 2*node_values[a][j][i+1] + h*node_derivatives[a][j][i+1]) / (h*h*h));
             }
         }
     }
@@ -41,14 +44,15 @@ RadialFunctionSetKokkos::RadialFunctionSetKokkos(
 
 // This older, cleaner version of the function lacks edge-type dependence
 #if 0
-void RadialFunctionSetKokkos::evaluate(
+template <typename Precision>
+void RadialFunctionSetKokkos<Precision>::evaluate(
     const int num_nodes,
     Kokkos::View<const int*> node_types,
     Kokkos::View<const int*> num_neigh,
     Kokkos::View<const int*> neigh_types,
     Kokkos::View<const double*> r,
-    Kokkos::View<double**,Kokkos::LayoutRight> R,
-    Kokkos::View<double**,Kokkos::LayoutRight> R_deriv) const
+    Kokkos::View<Precision**,Kokkos::LayoutRight> R,
+    Kokkos::View<Precision**,Kokkos::LayoutRight> R_deriv) const
 {
     const auto h = this->h;
     const auto num_functions = this->num_functions;
@@ -117,14 +121,15 @@ void RadialFunctionSetKokkos::evaluate(
 #endif
 
 
-void RadialFunctionSetKokkos::evaluate(
+template <typename Precision>
+void RadialFunctionSetKokkos<Precision>::evaluate(
     const int num_nodes,
     Kokkos::View<const int*> node_types,
     Kokkos::View<const int*> num_neigh,
     Kokkos::View<const int*> neigh_types,
     Kokkos::View<const double*> r,
-    Kokkos::View<double**,Kokkos::LayoutRight> R,
-    Kokkos::View<double**,Kokkos::LayoutRight> R_deriv) const
+    Kokkos::View<Precision**,Kokkos::LayoutRight> R,
+    Kokkos::View<Precision**,Kokkos::LayoutRight> R_deriv) const
 {
     const auto h = this->h;
     const auto num_functions = this->num_functions;
@@ -182,13 +187,20 @@ void RadialFunctionSetKokkos::evaluate(
             Kokkos::parallel_for(
                 Kokkos::TeamVectorRange(team_member, num_functions),
                 [&] (const int k) {
-                    const double c0 = c(type_ij,n,0,k);
-                    const double c1 = c(type_ij,n,1,k); 
-                    const double c2 = c(type_ij,n,2,k); 
-                    const double c3 = c(type_ij,n,3,k); 
-                    R(ij,k) = c0 + c1*x + c2*xx + c3*xxx;
-                    R_deriv(ij,k) = c1 + c2*two_x + c3*three_xx;
+                    const Precision c0 = c(type_ij,n,0,k);
+                    const Precision c1 = c(type_ij,n,1,k);
+                    const Precision c2 = c(type_ij,n,2,k);
+                    const Precision c3 = c(type_ij,n,3,k);
+                    R(ij,k) = c0 + c1*static_cast<Precision>(x)
+                        + c2*static_cast<Precision>(xx)
+                        + c3*static_cast<Precision>(xxx);
+                    R_deriv(ij,k) = c1
+                        + c2*static_cast<Precision>(two_x)
+                        + c3*static_cast<Precision>(three_xx);
                 });
         });
         Kokkos::fence();
 }
+
+template class RadialFunctionSetKokkos<float>;
+template class RadialFunctionSetKokkos<double>;

libsymmetrix/source/radial_function_set_kokkos.hpp

@@ -4,6 +4,7 @@
 
 #include <Kokkos_Core.hpp>
 
+template <typename Precision>
 class RadialFunctionSetKokkos
 {
 public:
@@ -19,14 +20,14 @@ class RadialFunctionSetKokkos
         Kokkos::View<const int*> num_neigh,
         Kokkos::View<const int*> neigh_types,
         Kokkos::View<const double*> r,
-        Kokkos::View<double**,Kokkos::LayoutRight> R,
-        Kokkos::View<double**,Kokkos::LayoutRight> R_deriv) const;
+        Kokkos::View<Precision**,Kokkos::LayoutRight> R,
+        Kokkos::View<Precision**,Kokkos::LayoutRight> R_deriv) const;
 
 private:
 
     double h;
     int num_edge_types;
     int num_functions;
     int num_nodes;
-    Kokkos::View<const double****,Kokkos::LayoutRight> coefficients;
+    Kokkos::View<const Precision****,Kokkos::LayoutRight> coefficients;
 };

pair_symmetrix/pair_symmetrix_mace_kokkos.cpp

@@ -134,7 +134,7 @@ void PairSymmetrixMACEKokkos<DeviceType>::coeff(int narg, char **arg)
   if (!allocated) allocate();
 
   utils::logmesg(lmp, "Loading MACEKokkos model from \'{}\' ... ", arg[2]);
-  mace = std::make_unique<MACEKokkos>(arg[2]);
+  mace = std::make_unique<MACEKokkos<double>>(arg[2]);
   utils::logmesg(lmp, "success\n");
 
   // extract atomic numbers from pair_coeff

pair_symmetrix/pair_symmetrix_mace_kokkos.h

@@ -56,7 +56,7 @@ class PairSymmetrixMACEKokkos : public Pair, public KokkosBase {
 
  protected:
   std::string mode;
-  std::unique_ptr<MACEKokkos> mace;
+  std::unique_ptr<MACEKokkos<double>> mace;
   Kokkos::View<int*> mace_types;
   Kokkos::View<double***,Kokkos::LayoutRight> H1, H1_adj;