Commit 4e9c8f49 authored by Christopher Stone's avatar Christopher Stone
Browse files

Update FixRXKokkos for Cuda build. Added inline and other KOKKOS macros. 

- Updated the function prototypes to include the necessary KOKKOS
macros for __host__ and __device__ functions and inlined functions.
- Changed several View definitions to match the disjoint memory spaces
that only come up with Cuda builds.
parent 2f32c1a9

src/KOKKOS/fix_rx_kokkos.cpp

+25 −6
Original line number Diff line number Diff line
@@ -25,13 +25,13 @@ 
#include "neigh_list_kokkos.h"

#include "neigh_request.h"

#include "error.h"

#include "math_special.h"

#include "math_special_kokkos.h"



#include <float.h> // DBL_EPSILON



using namespace LAMMPS_NS;

using namespace FixConst;

using namespace MathSpecial;

using namespace MathSpecialKokkos;



#ifdef DBL_EPSILON

  #define MY_EPSILON (10.0*DBL_EPSILON)
@@ -425,8 +425,8 @@ int FixRxKokkos<DeviceType>::k_rkf45_h0 
      // should we accept this?

      if (hnew_is_ok || iter == max_iters){

         hnew = hg;

         if (iter == max_iters)

            fprintf(stderr, "ERROR_HIN_MAX_ITERS\n");

         //if (iter == max_iters)

         //   fprintf(stderr, "ERROR_HIN_MAX_ITERS\n");

         break;

      }
@@ -1407,6 +1407,14 @@ void FixRxKokkos<DeviceType>::solve_reactions(const int vflag, const bool isPreF 
       );

  }



  // Error flag for any failures.

  DAT::tdual_int_scalar k_error_flag("pair:error_flag");



  // Initialize and sync the device flag.

  k_error_flag.h_view() = 0;

  k_error_flag.template modify<LMPHostType>();

  k_error_flag.template sync<DeviceType>();



  // Create scratch array space.

  const size_t scratchSpaceSize = (8*nspecies + 2*nreactions);

  //double *scratchSpace = new double[ scratchSpaceSize * nlocal ];
@@ -1483,7 +1491,11 @@ void FixRxKokkos<DeviceType>::solve_reactions(const int vflag, const bool isPreF 
        for (int ispecies = 0; ispecies < nspecies; ispecies++)

        {

          if (y[ispecies] < -MY_EPSILON)

            error->one(FLERR,"Computed concentration in RK solver is < -10*DBL_EPSILON");

          {

            //error->one(FLERR,"Computed concentration in RK solver is < -10*DBL_EPSILON");

            k_error_flag.d_view() = 2;

            // This should be an atomic update.

          }

          else if (y[ispecies] < MY_EPSILON)

            y[ispecies] = 0.0;
@@ -1507,6 +1519,12 @@ void FixRxKokkos<DeviceType>::solve_reactions(const int vflag, const bool isPreF 


  TimerType timer_ODE = getTimeStamp();



  // Check the error flag for any failures.

  k_error_flag.template modify<DeviceType>();

  k_error_flag.template sync<LMPHostType>();

  if (k_error_flag.h_view() == 2)

    error->one(FLERR,"Computed concentration in RK solver is < -10*DBL_EPSILON");



  // Signal that dvector has been modified on this execution space.

  atomKK->modified( execution_space, DVECTOR_MASK );
@@ -1815,7 +1833,8 @@ void FixRxKokkos<DeviceType>::computeLocalTemperature() 
        {

          // Create an atomic view of sumWeights and dpdThetaLocal. Only needed

          // for Half/thread scenarios.

          typedef Kokkos::View< E_FLOAT*, typename DAT::t_efloat_1d::array_layout, DeviceType, Kokkos::MemoryTraits< AtomicF< NEIGHFLAG >::value> > AtomicViewType;

          //typedef Kokkos::View< E_FLOAT*, typename DAT::t_efloat_1d::array_layout, DeviceType, Kokkos::MemoryTraits< AtomicF< NEIGHFLAG >::value> > AtomicViewType;

          typedef Kokkos::View< E_FLOAT*, typename DAT::t_efloat_1d::array_layout, typename DAT::t_efloat_1d::device_type, Kokkos::MemoryTraits< AtomicF< NEIGHFLAG >::value> > AtomicViewType;



          AtomicViewType a_dpdThetaLocal = d_dpdThetaLocal;

          AtomicViewType a_sumWeights    = d_sumWeights;

src/KOKKOS/fix_rx_kokkos.h

+25 −10
Original line number Diff line number Diff line
@@ -49,6 +49,7 @@ class FixRxKokkos : public FixRX { 
  {

    int nSteps, nIters, nFuncs, nFails;



    KOKKOS_INLINE_FUNCTION

    CounterType() : nSteps(0), nIters(0), nFuncs(0), nFails(0) {};



    KOKKOS_INLINE_FUNCTION
@@ -72,7 +73,7 @@ class FixRxKokkos : public FixRX { 
    }

  };



 protected:

 //protected:

  PairDPDfdtEnergyKokkos<DeviceType>* pairDPDEKK;

  double VDPD;
@@ -84,13 +85,15 @@ class FixRxKokkos : public FixRX { 


    value_type *m_data;



    KOKKOS_INLINE_FUNCTION

    StridedArrayType() : m_data(NULL) {}

    KOKKOS_INLINE_FUNCTION

    StridedArrayType(value_type *ptr) : m_data(ptr) {}



    inline       value_type& operator()(const int idx)       { return m_data[Stride*idx]; }

    inline const value_type& operator()(const int idx) const { return m_data[Stride*idx]; }

    inline       value_type& operator[](const int idx)       { return m_data[Stride*idx]; }

    inline const value_type& operator[](const int idx) const { return m_data[Stride*idx]; }

    KOKKOS_INLINE_FUNCTION       value_type& operator()(const int idx)       { return m_data[Stride*idx]; }

    KOKKOS_INLINE_FUNCTION const value_type& operator()(const int idx) const { return m_data[Stride*idx]; }

    KOKKOS_INLINE_FUNCTION       value_type& operator[](const int idx)       { return m_data[Stride*idx]; }

    KOKKOS_INLINE_FUNCTION const value_type& operator[](const int idx) const { return m_data[Stride*idx]; }

  };



  template <int stride = 1>
@@ -100,17 +103,22 @@ class FixRxKokkos : public FixRX { 
    StridedArrayType<double,1> rxnRateLaw;

  };



  void solve_reactions(const int vflag, const bool isPreForce = true);

  void solve_reactions(const int vflag, const bool isPreForce);



  int rhs       (double, const double *, double *, void *) const;

  int rhs_dense (double, const double *, double *, void *) const;

  int rhs_sparse(double, const double *, double *, void *) const;



  template <typename VectorType, typename UserDataType>

    KOKKOS_INLINE_FUNCTION

  int k_rhs       (double, const VectorType&, VectorType&, UserDataType& ) const;



  template <typename VectorType, typename UserDataType>

    KOKKOS_INLINE_FUNCTION

  int k_rhs_dense (double, const VectorType&, VectorType&, UserDataType& ) const;



  template <typename VectorType, typename UserDataType>

    KOKKOS_INLINE_FUNCTION

  int k_rhs_sparse(double, const VectorType&, VectorType&, UserDataType& ) const;



  //!< Classic Runge-Kutta 4th-order stepper.
@@ -130,19 +138,23 @@ class FixRxKokkos : public FixRX { 


  //!< Classic Runge-Kutta 4th-order stepper.

  template <typename VectorType, typename UserDataType>

    KOKKOS_INLINE_FUNCTION

  void k_rk4(const double t_stop, VectorType& y, VectorType& rwork, UserDataType& userData) const;



  //!< Runge-Kutta-Fehlberg ODE Solver.

  template <typename VectorType, typename UserDataType>

    KOKKOS_INLINE_FUNCTION

  void k_rkf45(const int neq, const double t_stop, VectorType& y, VectorType& rwork, UserDataType& userData, CounterType& counter) const;



  //!< Runge-Kutta-Fehlberg ODE stepper function.

  template <typename VectorType, typename UserDataType>

    KOKKOS_INLINE_FUNCTION

  void k_rkf45_step (const int neq, const double h, VectorType& y, VectorType& y_out,

                     VectorType& rwk, UserDataType& userData) const;



  //!< Initial step size estimation for the Runge-Kutta-Fehlberg ODE solver.

  template <typename VectorType, typename UserDataType>

    KOKKOS_INLINE_FUNCTION

  int k_rkf45_h0 (const int neq, const double t, const double t_stop,

                  const double hmin, const double hmax,

                  double& h0, VectorType& y, VectorType& rwk, UserDataType& userData) const;
@@ -155,8 +167,10 @@ class FixRxKokkos : public FixRX { 
  int *diagnosticCounterPerODEnFuncs;

  DAT::tdual_int_1d k_diagnosticCounterPerODEnSteps;

  DAT::tdual_int_1d k_diagnosticCounterPerODEnFuncs;

  typename ArrayTypes<DeviceType>::t_int_1d d_diagnosticCounterPerODEnSteps;

  typename ArrayTypes<DeviceType>::t_int_1d d_diagnosticCounterPerODEnFuncs;

  //typename ArrayTypes<DeviceType>::t_int_1d d_diagnosticCounterPerODEnSteps;

  //typename ArrayTypes<DeviceType>::t_int_1d d_diagnosticCounterPerODEnFuncs;

  typename DAT::t_int_1d d_diagnosticCounterPerODEnSteps;

  typename DAT::t_int_1d d_diagnosticCounterPerODEnFuncs;

  typename HAT::t_int_1d h_diagnosticCounterPerODEnSteps;

  typename HAT::t_int_1d h_diagnosticCounterPerODEnFuncs;
@@ -185,7 +199,8 @@ class FixRxKokkos : public FixRX { 


  // Need a dual-view and device-view for dpdThetaLocal and sumWeights since they're used in several callbacks.

  DAT::tdual_efloat_1d k_dpdThetaLocal, k_sumWeights;

  typename ArrayTypes<DeviceType>::t_efloat_1d d_dpdThetaLocal, d_sumWeights;

  //typename ArrayTypes<DeviceType>::t_efloat_1d d_dpdThetaLocal, d_sumWeights;

  typename DAT::t_efloat_1d d_dpdThetaLocal, d_sumWeights;

  typename HAT::t_efloat_1d h_dpdThetaLocal, h_sumWeights;



  template <int WT_FLAG, int LOCAL_TEMP_FLAG, bool NEWTON_PAIR, int NEIGHFLAG>
@@ -196,7 +211,7 @@ class FixRxKokkos : public FixRX { 
  int pack_forward_comm(int , int *, double *, int, int *);

  void unpack_forward_comm(int , int , double *);



 private: // replicate a few from FixRX

 //private: // replicate a few from FixRX

  int my_restartFlag;

};

CRA Git | Maintained and supported by SUSTech CRA and CCSE