valence_workunit.h

// -*-c++-*-
#ifndef STORMM_VALENCE_WORKUNIT_H
#define STORMM_VALENCE_WORKUNIT_H
 
#include <string>
#include <vector>
#include "copyright.h"
#include "Accelerator/hybrid.h"
#include "Accelerator/gpu_details.h"
#include "Constants/behavior.h"
#include "Potential/energy_enumerators.h"
#include "Restraints/restraint_apparatus.h"
#include "Topology/atomgraph.h"
#include "Topology/atomgraph_abstracts.h"
#include "synthesis_enumerators.h"
 
namespace stormm {
namespace synthesis {
 
using card::GpuDetails;
using card::Hybrid;
using constants::PrecisionModel;
using energy::EvaluateForce;
using energy::EvaluateEnergy;
using energy::ValenceKernelSize;
using restraints::RestraintApparatus;
using restraints::RestraintKit;
using topology::AtomGraph;
using topology::ConstraintKit;
using topology::ValenceKit;
using topology::VirtualSiteKit;

constexpr int maximum_valence_work_unit_atoms = 512;
constexpr int half_valence_work_unit_atoms = maximum_valence_work_unit_atoms / 2;
constexpr int quarter_valence_work_unit_atoms = maximum_valence_work_unit_atoms / 4;
constexpr int eighth_valence_work_unit_atoms = maximum_valence_work_unit_atoms / 8;

constexpr int minimum_valence_work_unit_atoms = eighth_valence_work_unit_atoms;

constexpr int max_atom_search_stacks = 8;

constexpr int max_constraint_group_size = 16;
  
class ValenceDelegator {
public:

  ValenceDelegator(const AtomGraph *ag_in, const RestraintApparatus *ra_in = nullptr);

  int getAtomAssignmentCount(int atom_index) const;

  int getFirstUnassignedAtom() const;

  std::vector<int> getBondAffectors(const std::vector<int> &atom_indices) const;
  
  std::vector<int> getAngleAffectors(const std::vector<int> &atom_indices) const;
 
  std::vector<int> getDihedralAffectors(const std::vector<int> &atom_indices) const;
  
  std::vector<int> getUreyBradleyAffectors(const std::vector<int> &atom_indices) const;
   
  std::vector<int> getCharmmImproperAffectors(const std::vector<int> &atom_indices) const;

  std::vector<int> getCmapAffectors(const std::vector<int> &atom_indices) const;

  std::vector<int> getInferred14Affectors(const std::vector<int> &atom_indices) const;

  std::vector<int> getPositionalRestraintAffectors(const std::vector<int> &atom_indices) const;
  
  std::vector<int> getDistanceRestraintAffectors(const std::vector<int> &atom_indices) const;
  
  std::vector<int> getAngleRestraintAffectors(const std::vector<int> &atom_indices) const;

  std::vector<int> getDihedralRestraintAffectors(const std::vector<int> &atom_indices) const;

  std::vector<int> getVirtualSiteAffectors(const std::vector<int> &atom_indices) const;

  std::vector<int> getSettleGroupAffectors(const std::vector<int> &atom_indices) const;

  std::vector<int> getConstraintGroupAffectors(const std::vector<int> &atom_indices) const;

  std::vector<int> findMovementPartners(int atom_idx,
                                        const std::vector<int> &caller_stack = {}) const;

  std::vector<int> findForcePartners(int atom_idx,
                                     const std::vector<int> &caller_stack = {}) const;

  std::vector<int> getUpdateDependencies(const int atom_index) const;
  
  int getUpdateWorkUnit(int atom_index) const;

  int getBondAccumulatorWorkUnit(const int bond_index) const;
  
  int getAngleAccumulatorWorkUnit(const int angl_index) const;

  int getDihedralAccumulatorWorkUnit(const int dihe_index) const;

  int getUreyBradleyAccumulatorWorkUnit(const int ubrd_index) const;

  int getCharmmImproperAccumulatorWorkUnit(const int cimp_index) const;

  int getCmapAccumulatorWorkUnit(const int cmap_index) const;

  int getInferred14AccumulatorWorkUnit(const int infr14_index) const;

  int getPositionalRestraintAccumulatorWorkUnit(const int rposn_index) const;

  int getDistanceRestraintAccumulatorWorkUnit(const int rbond_index) const;

  int getAngleRestraintAccumulatorWorkUnit(const int rangl_index) const;

  int getDihedralRestraintAccumulatorWorkUnit(const int rdihe_index) const;

  const AtomGraph* getTopologyPointer() const;
  
  const RestraintApparatus* getRestraintApparatusPointer() const;
    
  bool checkPresence(int atom_index, int vwu_index) const;

  void markAtomAddition(int vwu_index, int atom_index);

  bool setUpdateWorkUnit(int atom_index, int vwu_index);

  bool setBondAccumulatorWorkUnit(int bond_index, int vwu_index);
  
  bool setAngleAccumulatorWorkUnit(int angl_index, int vwu_index);

  bool setDihedralAccumulatorWorkUnit(int dihe_index, int vwu_index);

  bool setUreyBradleyAccumulatorWorkUnit(int ubrd_index, int vwu_index);

  bool setCharmmImproperAccumulatorWorkUnit(int cimp_index, int vwu_index);

  bool setCmapAccumulatorWorkUnit(int cmap_index, int vwu_index);

  bool setInferred14AccumulatorWorkUnit(int infr14_index, int vwu_index);

  bool setPositionalRestraintAccumulatorWorkUnit(int rposn_index, int vwu_index);

  bool setDistanceRestraintAccumulatorWorkUnit(int rbond_index, int vwu_index);

  bool setAngleRestraintAccumulatorWorkUnit(int rangl_index, int vwu_index);

  bool setDihedralRestraintAccumulatorWorkUnit(int rdihe_index, int vwu_index);
 
private:
  int atom_count;                 
  int first_unassigned_atom;      
  int max_presence_allocation;    
 
  // The following lists are quick to construct and contain a double-, triple-, or even a
  // quintuple-counting of all information in the topology's own "atom assignment" arrays for the
  // same valence terms.  Is there a better way?
  std::vector<int> bond_affector_list;    
  std::vector<int> bond_affector_bounds;  
  std::vector<int> angl_affector_list;    
  std::vector<int> angl_affector_bounds;  
  std::vector<int> dihe_affector_list;    
  std::vector<int> dihe_affector_bounds;  
  std::vector<int> ubrd_affector_list;    
  std::vector<int> ubrd_affector_bounds;  
  std::vector<int> cimp_affector_list;    
  std::vector<int> cimp_affector_bounds;  
  std::vector<int> cmap_affector_list;    
  std::vector<int> cmap_affector_bounds;  
  std::vector<int> infr_affector_list;    
  std::vector<int> infr_affector_bounds;  
 
  // The virtual site affector arrays will be pre-allcoated to hold virtual sites with up to four
  // frame atoms apiece.
  std::vector<int> vste_affector_list;    
  std::vector<int> vste_affector_bounds;  
 
  // SHAKE and RATTLE groups likewise must have all atoms present in a work group in order to
  // evaluate bond constraints.
  std::vector<int> cnst_affector_list;    
  std::vector<int> cnst_affector_bounds;  
 
  // SETTLE groups are special constraint groups for a fast rigid water implementation, kept
  // distinct because of the different treatment they get in the dynamics loop.
  std::vector<int> sett_affector_list;    
  std::vector<int> sett_affector_bounds;  
 
  // Positional restraints apply to a single atom, but multiple restraints could still, in theory,
  // restrain a single atom to a ring or other confined region of space.  A bounds array is still
  // necessary.  Other restraints affect multiple atoms and map like their valence term analogs.
  std::vector<int> rposn_affector_list;    
  std::vector<int> rposn_affector_bounds;  
  std::vector<int> rbond_affector_list;    
  std::vector<int> rbond_affector_bounds;  
  std::vector<int> rangl_affector_list;    
  std::vector<int> rangl_affector_bounds;  
  std::vector<int> rdihe_affector_list;    
  std::vector<int> rdihe_affector_bounds;  
 
  // Individual atoms must leave a record of their whereabouts in the valence work units for rapid
  // retrieval of their locations
  std::vector<int> work_unit_assignment_count;  
  std::vector<int> work_unit_presence;          
  std::vector<int> assigned_update_work_units;  
 
  // Simple energy terms must leave a record of the work units that will be responsible for
  // accumulating their values in the total reported energy.  Multiple work units may evaluate
  // each energy term, but only one must contribute.
  std::vector<int> assigned_bond_acc_work_units;   
  std::vector<int> assigned_angl_acc_work_units;   
  std::vector<int> assigned_dihe_acc_work_units;   
  std::vector<int> assigned_ubrd_acc_work_units;   
  std::vector<int> assigned_cimp_acc_work_units;   
  std::vector<int> assigned_cmap_acc_work_units;   
  std::vector<int> assigned_infr14_acc_work_units; 
  std::vector<int> assigned_rposn_acc_work_units;  
  std::vector<int> assigned_rbond_acc_work_units;  
  std::vector<int> assigned_rangl_acc_work_units;  
  std::vector<int> assigned_rdihe_acc_work_units;  
  
  
  const AtomGraph *ag_pointer;
  const RestraintApparatus *ra_pointer;

  void resizePresenceArrays(int n_units);
  
  void allocate();

  void fillAffectorArrays(const ValenceKit<double> &vk, const VirtualSiteKit<double> &vsk,
                          const ConstraintKit<double> &cnk,
                          const RestraintKit<double, double2, double4> &rar);
};
  
class ValenceWorkUnit {
public:

  ValenceWorkUnit(ValenceDelegator *vdel_in, std::vector<int> *tvwu_coverage, int list_index_in,
                  int seed_atom_in, int max_atoms_in = maximum_valence_work_unit_atoms);

  int getImportedAtomCount() const;

  int getMovedAtomCount() const;

  int getUpdatedAtomCount() const;

  int getListIndex() const;

  int getMinAtomIndex() const;

  int getMaxAtomIndex() const;
  
  int getMaxAtoms() const;

  std::vector<int> getAtomImportList(int atom_offset = 0) const;

  int getImportedAtomIndex(int slot, int atom_offset = 0) const;
  
  std::vector<uint2> getAtomManipulationMasks() const;

  int getPaddedConstraintInstructionCount() const;

  std::vector<int> getTaskCounts() const;

  void storeCompositeBondInstructions(const std::vector<int> &bond_param_map = {},
                                      const std::vector<int> &ubrd_param_map = {});

  void storeAngleInstructions(const std::vector<int> &parameter_map = {});

  void storeCompositeDihedralInstructions(const std::vector<int> &dihe_param_map = {},
                                          const std::vector<int> &dihe14_param_map = {},
                                          const std::vector<int> &cimp_param_map = {});
  
  void storeCmapInstructions(const std::vector<int> &parameter_map = {});
  
  void storeInferred14Instructions(const std::vector<int> &parameter_map = {});

  void storePositionalRestraintInstructions(const std::vector<int> &kr_param_map = {},
                                            const std::vector<int> &xyz_param_map = {});

  void storeDistanceRestraintInstructions(const std::vector<int> &kr_param_map = {});

  void storeAngleRestraintInstructions(const std::vector<int> &kr_param_map = {});

  void storeDihedralRestraintInstructions(const std::vector<int> &kr_param_map = {});

  void storeVirtualSiteInstructions(const std::vector<int> &parameter_map = {});

  void storeSettleGroupInstructions(const std::vector<int> &parameter_map = {});

  void storeConstraintGroupInstructions(const std::vector<int> &parameter_map = {},
                                        const std::vector<int> &group_param_bounds = {});

  const std::vector<uint2>& getCompositeBondInstructions() const;

  const std::vector<uint2>& getAngleInstructions() const;

  const std::vector<uint3>& getCompositeDihedralInstructions() const;

  const std::vector<uint2>& getCmapInstructions() const;

  const std::vector<uint>& getInferred14Instructions() const;

  const std::vector<uint2>& getPositionalRestraintInstructions() const;

  const std::vector<uint2>& getDistanceRestraintInstructions() const;

  const std::vector<uint2>& getAngleRestraintInstructions() const;

  const std::vector<uint2>& getDihedralRestraintInstructions() const;

  const std::vector<uint2>& getVirtualSiteInstructions() const;  
  
  const std::vector<uint2>& getSettleGroupInstructions() const;  
  
  const std::vector<uint2>& getConstraintGroupInstructions() const;  

  uint2 getCompositeBondInstruction(int index) const;

  uint2 getAngleInstruction(int index) const;

  uint3 getCompositeDihedralInstruction(int index) const;

  uint2 getCmapInstruction(int index) const;

  uint getInferred14Instruction(int index) const;

  uint2 getPositionalRestraintInstruction(int index) const;

  uint2 getDistanceRestraintInstruction(int index) const;

  uint2 getAngleRestraintInstruction(int index) const;

  uint2 getDihedralRestraintInstruction(int index) const;

  uint2 getVirtualSiteInstruction(int index) const;  
  
  uint2 getSettleGroupInstruction(int index) const;  
  
  uint2 getConstraintGroupInstruction(int index) const;  

  const std::vector<uint>& getAccumulationFlags(VwuTask vtask) const;

  const std::vector<uint>& getAtomUpdateFlags() const;
  
  const std::vector<int>& getSimpleTaskList(VwuTask vtask) const;

  const std::vector<int>& getCompositeBondTaskList() const;

  const std::vector<int2>& getCompositeDihedralTaskList() const;
  
  ValenceDelegator* getDelegatorPointer();

  const AtomGraph* getTopologyPointer() const;

  const RestraintApparatus* getRestraintApparatusPointer() const;
  
  void setListIndex(int list_index_in);

  void setAtomLimit(int new_limit);
  
  void addNewAtomImport(int atom_index);

  void addNewAtomUpdate(const int atom_index);

  void makeAtomMoveList();

  void sortAtomSets();

  void makeAtomUpdateMask();
  
  void logActivities();
  
private:
  int imported_atom_count;  
  int moved_atom_count;     
  int updated_atom_count;   
  int bond_term_count;      
  int angl_term_count;      
  int dihe_term_count;      
  int ubrd_term_count;      
  int cbnd_term_count;      
  int cimp_term_count;      
  int cdhe_term_count;      
  int cmap_term_count;      
  int infr14_term_count;    
  int rposn_term_count;     
  int rbond_term_count;     
  int rangl_term_count;     
  int rdihe_term_count;     
  int cnst_group_count;     
  int sett_group_count;     
  int vste_count;           
  int list_index;           
  int min_atom_index;       
  int max_atom_index;       
  int atom_limit;           

  std::vector<int> atom_import_list;

  std::vector<int> atom_move_list;

  std::vector<int> atom_update_list;
  
  std::vector<uint> atom_update_mask;
  
  // Valence terms for the work unit, typical force field elements
  std::vector<int> bond_term_list;    
  std::vector<int> angl_term_list;    
  std::vector<int> dihe_term_list;    
  std::vector<int> ubrd_term_list;    
  std::vector<int> cimp_term_list;    
  std::vector<int> cmap_term_list;    
  std::vector<int> infr14_term_list;  
 
  // List of local indices for atoms participating in each task (energy term, virtual site,
  // or constraint group) for which the work unit is responsible
  std::vector<int> bond_i_atoms;      
  std::vector<int> bond_j_atoms;      
  std::vector<int> angl_i_atoms;      
  std::vector<int> angl_j_atoms;      
  std::vector<int> angl_k_atoms;      
  std::vector<int> dihe_i_atoms;      
  std::vector<int> dihe_j_atoms;      
  std::vector<int> dihe_k_atoms;      
  std::vector<int> dihe_l_atoms;      
  std::vector<int> ubrd_i_atoms;      
  std::vector<int> ubrd_k_atoms;      
  std::vector<int> cimp_i_atoms;      
  std::vector<int> cimp_j_atoms;      
  std::vector<int> cimp_k_atoms;      
  std::vector<int> cimp_l_atoms;      
  std::vector<int> cmap_i_atoms;      
  std::vector<int> cmap_j_atoms;      
  std::vector<int> cmap_k_atoms;      
  std::vector<int> cmap_l_atoms;      
  std::vector<int> cmap_m_atoms;      
  std::vector<int> infr14_i_atoms;    
  std::vector<int> infr14_l_atoms;    

  std::vector<int> cbnd_term_list;

  std::vector<bool> cbnd_is_ubrd;  
  std::vector<int> cbnd_i_atoms;   
  std::vector<int> cbnd_jk_atoms;  
  
  std::vector<int2> cdhe_term_list;
 
  // Arrays to support the composite dihedrals
  std::vector<bool> cdhe_is_cimp;     
  std::vector<int> cdhe_i_atoms;      
  std::vector<int> cdhe_j_atoms;      
  std::vector<int> cdhe_k_atoms;      
  std::vector<int> cdhe_l_atoms;      
  
  
  // Restraint terms for this work unit
  std::vector<int> rposn_term_list;  
  std::vector<int> rbond_term_list;  
  std::vector<int> rangl_term_list;  
  std::vector<int> rdihe_term_list;  
  std::vector<int> rposn_atoms;      
  std::vector<int> rbond_i_atoms;    
  std::vector<int> rbond_j_atoms;    
  std::vector<int> rangl_i_atoms;    
  std::vector<int> rangl_j_atoms;    
  std::vector<int> rangl_k_atoms;    
  std::vector<int> rdihe_i_atoms;    
  std::vector<int> rdihe_j_atoms;    
  std::vector<int> rdihe_k_atoms;    
  std::vector<int> rdihe_l_atoms;    
 
  // Bit string directives on whether to contribute energies to the work unit's reported total
  std::vector<uint> acc_bond_energy;    
  std::vector<uint> acc_angl_energy;    
  std::vector<uint> acc_dihe_energy;    
  std::vector<uint> acc_ubrd_energy;    
  std::vector<uint> acc_cimp_energy;    
  std::vector<uint> acc_cmap_energy;    
  std::vector<uint> acc_infr14_energy;  
  std::vector<uint> acc_rposn_energy;   
  std::vector<uint> acc_rbond_energy;   
  std::vector<uint> acc_rangl_energy;   
  std::vector<uint> acc_rdihe_energy;   
  std::vector<uint> acc_cbnd_energy;    
  std::vector<uint> acc_cdhe_energy;    
 
  // Constraint groups for this work unit
  std::vector<int> cnst_group_list;    
  std::vector<int> sett_group_list;    
  std::vector<int> cnst_group_atoms;   
  std::vector<int> cnst_group_bounds;  
  std::vector<int> sett_ox_atoms;      
  std::vector<int> sett_h1_atoms;      
  std::vector<int> sett_h2_atoms;      
 
  // Virtual sites in this work unit
  std::vector<int> virtual_site_list;   
  std::vector<int> vsite_atoms;         
  std::vector<int> vsite_frame1_atoms;  
  std::vector<int> vsite_frame2_atoms;  
  std::vector<int> vsite_frame3_atoms;  
  std::vector<int> vsite_frame4_atoms;  
 
  // Internal arrays of instructions for each energy term, constraint, or virtual site task,
  // stored for fast retrieval.
  std::vector<uint2> cbnd_instructions;   
  std::vector<uint2> angl_instructions;   
  std::vector<uint3> cdhe_instructions;   
  std::vector<uint2> cmap_instructions;   
  std::vector<uint> infr14_instructions;  
  std::vector<uint2> rposn_instructions;  
  std::vector<uint2> rbond_instructions;  
  std::vector<uint2> rangl_instructions;  
  std::vector<uint2> rdihe_instructions;  
  std::vector<uint2> vste_instructions;   
  std::vector<uint2> sett_instructions;   
  std::vector<uint2> cnst_instructions;   
  
  // Pointers to important objects
  ValenceDelegator *vdel_pointer;        
  const AtomGraph *ag_pointer;           
  const RestraintApparatus *ra_pointer;  
};

int2 calculateValenceWorkUnitSize(const int* atom_counts, int system_count);
 
int2 calculateValenceWorkUnitSize(const std::vector<int> &atom_counts);
 
int2 calculateValenceWorkUnitSize(const Hybrid<int> &atom_counts);
 
int2 calculateValenceWorkUnitSize(const int* atom_counts, int system_count, int sm_count,
                                  ValenceKernelSize *kwidth);
 
int2 calculateValenceWorkUnitSize(const std::vector<int> &atom_counts, int sm_count,
                                  ValenceKernelSize *kwidth);
 
int2 calculateValenceWorkUnitSize(const Hybrid<int> &atom_counts, int sm_count,
                                  ValenceKernelSize *kwidth);
 
int2 calculateValenceWorkUnitSize(const std::vector<AtomGraph*> &ag_list,
                                  const std::vector<int> &topology_replicas, const int sm_count,
                                  ValenceKernelSize *kwidth);

std::vector<ValenceWorkUnit>
buildValenceWorkUnits(const AtomGraph *ag, const RestraintApparatus *ra,
                      int2 max_atoms_per_vwu = { maximum_valence_work_unit_atoms,
                                                 maximum_valence_work_unit_atoms });
 
std::vector<ValenceWorkUnit>
buildValenceWorkUnits(ValenceDelegator *vdel,
                      int2 max_atoms_per_vwu = { maximum_valence_work_unit_atoms,
                                                 maximum_valence_work_unit_atoms });
 
} // namespace topology
} // namespace stormm
 
#endif