colvaratoms_gpu.h

Go to the documentation of this file.

#ifndef COLVARATOMS_GPU_H
#define COLVARATOMS_GPU_H
 
#include "colvarmodule.h"
#include "colvartypes.h"
#include "colvar_rotation_derivative.h"
 
namespace colvars_gpu {
 
#if defined (COLVARS_CUDA) || defined (COLVARS_HIP)
struct colvaratoms_gpu_buffer_t {
  int* d_atoms_index;
  cvm::real* d_atoms_pos;
  cvm::real* d_atoms_charge;
  cvm::real* d_atoms_vel;
  cvm::real* d_atoms_mass;
  cvm::real* d_atoms_grad;
  cvm::real* d_atoms_total_force;
  cvm::real* d_atoms_weight;
  cvm::real* d_atoms_applied_force;
  cvm::real* d_fit_gradients;
  cvm::real* d_ref_pos;
  cvm::real* d_atoms_pos_unrotated;
  cvm::rvector* d_com;
  cvm::rvector* d_com_tmp;
  cvm::rvector* d_cog;
  cvm::rvector* d_cog_tmp;
  cvm::rvector* d_cog_orig;
  unsigned int* d_com_cog_tbcount;
  cvm::rvector* h_com;
  cvm::rvector* h_cog;
  cvm::rvector* h_cog_orig;
  cvm::rvector* d_ref_pos_cog;
  cvm::rvector* h_dummy_atom_pos;
};
 
struct colvaratoms_gpu_calc_fit_info_t {
  double3* d_atom_grad;
  double* d_sum_dxdq;
  cvm::rmatrix* d_dxdC;
  unsigned int* d_tbcount;
};
 
class colvaratoms_gpu {
public:
  enum class event_type {
    read_and_calculate = 0,
    calc_fit_gradients,
    apply_force,
    num_event_types,
  };
  colvaratoms_gpu(colvarmodule *cvmodule_in);
  ~colvaratoms_gpu();
  int init_gpu();
  int destroy_gpu();
  int reset_gpu_graphs();
  int read_data_gpu(cvm::atom_group* cpu_atoms);
  int calc_fit_gradients_gpu(cvm::atom_group* cpu_atoms);
  int sync_to_gpu_buffers(const cvm::atom_group* cpu_atoms);
  int clear_gpu_buffers(const cvm::atom_group* cpu_atoms);
  int add_reset_atoms_data_nodes(
    const cvm::atom_group* cpu_atoms, cudaGraph_t& graph,
    std::unordered_map<std::string, cudaGraphNode_t>& nodes_map);
  int add_read_positions_nodes(
    const cvm::atom_group* cpu_atoms, cudaGraph_t& graph,
    std::unordered_map<std::string, cudaGraphNode_t>& nodes_map);
  int add_calc_required_properties_nodes(
    const cvm::atom_group* cpu_atoms, cudaGraph_t& graph,
    std::unordered_map<std::string, cudaGraphNode_t>& nodes_map,
    const std::vector<cudaGraphNode_t>& extra_initial_dependencies = {});
  int add_update_cpu_buffers_nodes(
    cvm::atom_group* cpu_atoms, cudaGraph_t& graph,
    std::unordered_map<std::string, cudaGraphNode_t>& nodes_map);
  int after_read_data_sync(cvm::atom_group* cpu_atoms, bool copy_to_cpu);
  int begin_apply_force_gpu();
  int add_apply_force_nodes(
    const cvm::atom_group* cpu_atoms, cudaGraph_t& graph,
    std::unordered_map<std::string, cudaGraphNode_t>& nodes_map,
    const std::vector<cudaGraphNode_t>& extra_initial_dependencies = {});
  int add_calc_fit_gradients_nodes(
    cvm::atom_group* cpu_atoms, cudaGraph_t& graph,
    std::unordered_map<std::string, cudaGraphNode_t>& nodes_map,
    bool use_cpu_buffers = false);
  int read_positions_gpu_debug(
    cvm::atom_group* cpu_atoms, bool change_fitting_group, size_t change_atom_i,
    int xyz, bool to_cpu, double sign, cudaStream_t stream);
  int calc_required_properties_gpu_debug(
    cvm::atom_group* cpu_atoms, bool to_cpu, cudaStream_t stream);
  void do_feature_side_effects_gpu(
    cvm::atom_group* cpu_atoms, int id);
  int setup_rotation(const cvm::atom_group* cpu_atoms);
  int setup_rotation_derivative(const cvm::atom_group* cpu_atoms);
  colvaratoms_gpu_buffer_t& get_gpu_buffers() { return gpu_buffers; }
  int read_total_forces(cvm::atom_group* cpu_atoms);
  void apply_colvar_force_from_cpu(cvm::real const& cpu_force);
  void set_use_cpu_group_force(bool yesno) { use_group_force = yesno; }
  colvars_gpu::rotation_gpu& get_rot_gpu() { return rot_gpu; }
  const colvars_gpu::rotation_gpu& get_rot_gpu() const { return rot_gpu; }
  colvars_gpu::rotation_derivative_gpu* get_rot_deriv_gpu() { return rot_deriv_gpu; }
  const colvars_gpu::rotation_derivative_gpu* get_rot_deriv_gpu() const { return rot_deriv_gpu; }
  const cudaEvent_t& get_event(event_type type) const {
    return events[static_cast<int>(type)];
  }
  int add_force_to_proxy_gpu(cvm::atom_group* cpu_atoms);
  static int set_weighted_gradient_gpu(
    cvm::atom_group* cpu_atoms,
    const cvm::rvector* d_com_grad,
    cudaStream_t stream);
private:
  colvaratoms_gpu_buffer_t gpu_buffers;
  colvaratoms_gpu_calc_fit_info_t calc_fit_gradients_gpu_info;
  colvaratoms_gpu_calc_fit_info_t calc_fit_forces_gpu_info;
  gpu_graph_t graph_debug;
  cvm::real* h_sum_applied_colvar_force;
  bool use_apply_colvar_force;
  bool use_group_force;
  colvars_gpu::rotation_gpu rot_gpu;
  colvars_gpu::rotation_derivative_gpu* rot_deriv_gpu;
  colvarmodule *cvmodule;
  gpu_graph_t graph_read_compute;
  gpu_graph_t graph_calc_fit_gradients;
  gpu_graph_t graph_apply_force;
  std::array<cudaEvent_t, static_cast<int>(event_type::num_event_types)> events = {};
};
 
#endif // defined (COLVARS_CUDA) || defined (COLVARS_HIP)
}
 
#endif // COLVARATOMS_GPU_H