pmigrid.h

// -*-c++-*-
#ifndef STORMM_PMEGRID_H
#define STORMM_PMEGRID_H
 
#include <climits>
#include "copyright.h"
#include "Accelerator/gpu_details.h"
#include "Accelerator/hybrid.h"
#include "Constants/behavior.h"
#include "Constants/hpc_bounds.h"
#include "DataTypes/stormm_vector_types.h"
#include "Math/math_enumerators.h"
#include "Numerics/split_fixed_precision.h"
#include "Synthesis/phasespace_synthesis.h"
#include "cellgrid.h"
 
namespace stormm {
namespace energy {
 
using card::GpuDetails;
using card::Hybrid;
using card::HybridKind;
using card::HybridTargetLevel;
using constants::CartesianDimension;
using constants::PrecisionModel;
using constants::UnitCellAxis;
using stmath::FFTMode;
using synthesis::PhaseSpaceSynthesis;

constexpr int minimum_qspread_fp_bits = 24;
constexpr int minimum_ljspread_fp_bits = 24;
constexpr int maximum_qspread_fp_bits_sp = 56;
constexpr int maximum_qspread_fp_bits_dp = 88;
constexpr int default_qspread_fp_bits_sp = 32;
constexpr int default_qspread_fp_bits_dp = 64;
constexpr int maximum_ljspread_fp_bits_sp = 50;
constexpr int maximum_ljspread_fp_bits_dp = 82;
constexpr int default_ljspread_fp_bits_sp = 24;
constexpr int default_ljspread_fp_bits_dp = 56;
constexpr int mapping_nonoverflow_bits_dp = 61;
constexpr int mapping_nonoverflow_bits_sp = 29;
constexpr int default_bspline_order = 5;
constexpr int density_mapping_wu_size = 32;

#ifdef STORMM_USE_CUDA
constexpr int max_shared_acc_atom_bearing_region_adim = 28;
#else
constexpr int max_shared_acc_atom_bearing_region_adim = 9;
#endif
constexpr int max4s_grid_mapping_volume = 48;
constexpr int max4s_atom_bearing_region_adim = 7;
constexpr int max4s_atom_bearing_cross_section = 16;
constexpr int max5s_grid_mapping_volume = 48;
constexpr int max5s_atom_bearing_region_adim = 7;
constexpr int max5s_atom_bearing_cross_section = 16;
constexpr int max6s_grid_mapping_volume = 48;
constexpr int max6s_atom_bearing_region_adim = 4;
constexpr int max6s_atom_bearing_cross_section = 12;

struct PMIGridWriter {

  PMIGridWriter(NonbondedTheme theme_in, PrecisionModel mode_in, FFTMode fftm_in, int fp_bits_in,
                int nsys_in, int order_in, int wu_count_in, int max_grid_points_in,
                const uint4* dims_in, double* ddata_in, float* fdata_in,
                const uint* work_units_in);

  PMIGridWriter(const PMIGridWriter &original) = default;
  PMIGridWriter(PMIGridWriter &&original) = default;
 
  const NonbondedTheme theme;  
  const PrecisionModel mode;   
  const FFTMode fftm;          
  const float shacc_fp_scale;  
  const int nsys;              
  const int order;             
  const int wu_count;          
  const int max_grid_points;   
  const uint4* dims;           
  double* ddata;               
  float* fdata;                
  const uint* work_units;      
};

struct PMIGridReader {

  PMIGridReader(NonbondedTheme theme_in, PrecisionModel mode_in, FFTMode fftm_in, int fp_bits_in,
                int nsys_in, int order_in, const uint4* dims_in, const double* ddata_in,
                const float* fdata_in);

  PMIGridReader(const PMIGridWriter &w);
  PMIGridReader(const PMIGridWriter *w);
  
  PMIGridReader(const PMIGridReader &original) = default;
  PMIGridReader(PMIGridReader &&original) = default;
 
  const NonbondedTheme theme;  
  const PrecisionModel mode;   
  const FFTMode fftm;          
  const float shacc_fp_scale;  
  const int nsys;              
  const int order;             
  const uint4* dims;           
  const double* ddata;         
  const float* fdata;          
};
  
struct PMIGridAccumulator {

  PMIGridAccumulator(NonbondedTheme theme_in, PrecisionModel mode_in, FFTMode fftm_in,
                     bool use_overflow_in, int fp_bits_in, int nsys_in, int order_in,
                     int wu_count_in, const uint4* dims_in, double* ddata_in, float* fdata_in,
                     int* overflow_in, const uint* work_units_in);

  PMIGridAccumulator(const PMIGridAccumulator &original) = default;
  PMIGridAccumulator(PMIGridAccumulator &&original) = default;
  
  const NonbondedTheme theme;  
  const PrecisionModel mode;   
  const FFTMode fftm;          
  const bool use_overflow;     
  const int fp_bits;           
  const float fp_scale;        
  const int nsys;              
  const int order;             
  const int order_squared;     
  const int order_cubed;       
  const int wu_count;          
  const uint4* dims;           
  llint* lldata;               
  int* idata;                  
  int* overflow;               
  const uint* work_units;      
};

struct PMIGridFPReader {

  PMIGridFPReader(NonbondedTheme theme_in, PrecisionModel mode_in, FFTMode fftm_in,
                  bool use_overflow_in, int fp_bits_in, int nsys_in, int order_in,
                  const uint4* dims_in, const double* ddata_in, const float* fdata_in,
                  const int* overflow_in);

  PMIGridFPReader(const PMIGridAccumulator &w);
  PMIGridFPReader(const PMIGridAccumulator *w);

  PMIGridFPReader(const PMIGridFPReader &original) = default;
  PMIGridFPReader(PMIGridFPReader &&original) = default;
 
  const NonbondedTheme theme;  
  const PrecisionModel mode;   
  const FFTMode fftm;          
  const bool use_overflow;     
  const int fp_bits;           
  const float fp_scale;        
  const int nsys;              
  const int order;             
  const uint4* dims;           
  const llint* lldata;         
  const int* idata;            
  const int* overflow;         
};

class PMIGrid {
public:

  template <typename T, typename Tacc, typename Tcalc, typename T4>
  PMIGrid(const CellGrid<T, Tacc, Tcalc, T4> *cg_in, NonbondedTheme theme_in,
          int b_spline_order_in = default_bspline_order,
          PrecisionModel mode_in = PrecisionModel::SINGLE,
          FFTMode fft_staging_in = FFTMode::IN_PLACE, int fp_accumulation_bits_in = 0,
          int shared_fp_accumulation_bits_in = -1, const GpuDetails &gpu = null_gpu,
          const QMapMethod work_unit_configuration_in = QMapMethod::GENERAL_PURPOSE);
 
  template <typename T, typename Tacc, typename Tcalc, typename T4>
  PMIGrid(const CellGrid<T, Tacc, Tcalc, T4> &cg_in, NonbondedTheme theme_in,
          int b_spline_order_in = default_bspline_order,
          PrecisionModel mode_in = PrecisionModel::SINGLE,
          FFTMode fft_staging_in = FFTMode::IN_PLACE, int fp_accumulation_bits_in = 0,
          int shared_fp_accumulation_bits_in = -1, const GpuDetails &gpu = null_gpu,
          const QMapMethod work_unit_configuration_in = QMapMethod::GENERAL_PURPOSE);

  PMIGrid(const PMIGrid &original) = default;
  PMIGrid(PMIGrid &&original) = default;
  PMIGrid& operator=(const PMIGrid &original) = default;
  PMIGrid& operator=(PMIGrid &&original) = default;

  NonbondedTheme getTheme() const;
  
  PrecisionModel getMode() const;

  FFTMode getFFTStaging() const;

  bool fixedPrecisionEnabled() const;

  QMapMethod getRecommendedMappingMethod() const;
  
  QMapMethod getWorkUnitConfiguration() const;
  
  bool dataIsReal() const;
  
  int getFixedPrecisionBits() const;

  int getSharedFixedPrecisionBits() const;
  
  int getSystemCount() const;

  int getInterpolationOrder() const;

  size_t getTotalCapacity() const;
  
  uint4 getGridDimensions(int system_index) const;
  int getGridDimensions(int system_index, UnitCellAxis uc_axis) const;
  int getGridDimensions(int system_index, CartesianDimension uc_axis) const;

  int getWorkUnitCount() const;

  int getLargestWorkUnitGridPoints() const;

  bool shortFormatAccumulation() const;

  bool useOverflowAccumulation() const;

  PMIGridWriter data(HybridTargetLevel tier = HybridTargetLevel::HOST);
  const PMIGridReader data(HybridTargetLevel tier = HybridTargetLevel::HOST) const;

  PMIGridAccumulator fpData(HybridTargetLevel tier = HybridTargetLevel::HOST);
  const PMIGridFPReader fpData(HybridTargetLevel tier = HybridTargetLevel::HOST) const;
  
  template <typename T, typename Tacc, typename Tcalc, typename T4>
  const CellGrid<T, Tacc, Tcalc, T4>* getCellGridPointer() const;

  const CellGridReader<void, void, void, void>
  getTemplateFreeCellGridReader(HybridTargetLevel tier = HybridTargetLevel::HOST) const;
  
  size_t getCellGridMatrixTypeID() const;

  size_t getCellGridAccumulatorTypeID() const;

  size_t getCellGridCalculationTypeID() const;

  size_t getCellGridCoordinateTypeID() const;
  
  double getTotalOnGrid(int system_index) const;

  std::vector<double> getGrid(int system_index) const;

  const PhaseSpaceSynthesis* getCoordinateSynthesisPointer() const;

  const AtomGraphSynthesis* getTopologySynthesisPointer() const;
  
  const PMIGrid* getSelfPointer() const;
  
#if STORMM_USE_HPC
  void upload();

  void download();
#endif

  void setMode(PrecisionModel mode_in);

  void setRealDataFormat(bool real_in = true);
  
  void prepareFixedPrecisionModel(int fp_accumulation_bits_in, int shared_fp_bits_in);

  void setRecommendedMappingMethod(const GpuDetails &gpu);

  void prepareWorkUnits(QMapMethod approach, const GpuDetails &gpu);
  void prepareWorkUnits(const GpuDetails &gpu);
  
  void initialize(HybridTargetLevel tier = HybridTargetLevel::HOST,
                  const GpuDetails &gpu = null_gpu);

  void convertToReal(HybridTargetLevel tier = HybridTargetLevel::HOST,
                     const GpuDetails &gpu = null_gpu);
  
private:

  NonbondedTheme theme;
  
  PrecisionModel mode; 

  FFTMode fft_staging;
  
  int fp_accumulation_bits;

  int shared_fp_accumulation_bits;
  
  bool data_is_real;

  bool use_short_format_accumulation;
  
  int system_count;

  int b_spline_order;

  int shared_acc_buffer_size;
  
  QMapMethod recommendation;

  QMapMethod work_unit_configuration;
  
  Hybrid<uint4> grid_dimensions;

  size_t capacity;

  int work_unit_count;

  int largest_work_unit_grid_points;
  
  Hybrid<double> dgrid_stack;

  Hybrid<float> fgrid_stack;

  Hybrid<int> overflow_stack;

  Hybrid<uint> work_units;

  CellGrid<double, double, double, double4> *cg_pointer;
 
  // The following integers provide type ID numbers for the templated characteristics of the
  // attached CellGrid object.
  size_t cg_tmat;   
  size_t cg_tacc;   
  size_t cg_tcalc;  
  size_t cg_tcrd;   
  
  PhaseSpaceSynthesis *poly_ps_pointer;

  void validateFixedPrecisionBits(int fp_bits) const;
  
  void validateSystemIndex(int system_index) const;

  int findSharedBufferSize(const GpuDetails &gpu) const;
  
  int computeMappingHalo() const;
  
  void addWorkUnit(std::vector<uint> *result, int sysid, int pos_a, int pos_b, int pos_c,
                   int gmap_a, int gmap_b, int gmap_c,
                   const CellGridReader<void, void, void, void> &cgr_vc, int halo = 1);

  void checkShortFormatViability();

  void computeLargestWorkUnitGridPoints();
  
  template <typename T, typename T4> const CellGridReader<void, void, void, void>
  unrollTemplateFreeCGReader(HybridTargetLevel tier = HybridTargetLevel::HOST) const;

  template <typename T, typename T4> const AtomGraphSynthesis* unrollCgAgsPtrOne() const;

  template <typename T, typename Tacc, typename T4>
  const AtomGraphSynthesis* unrollCgAgsPtrTwo() const;
};
 
#ifdef STORMM_USE_HPC
void launchPMIGridInitialization(PMIGridAccumulator *pm_acc, const GpuDetails &gpu);
 
void launchPMIGridInitialization(PMIGridAccumulator *pm_acc, int block_count);

void launchPMIGridRealConversion(PMIGridWriter *pm_wrt, const PMIGridAccumulator &pm_acc,
                                 const GpuDetails &gpu);
 
void launchPMIGridRealConversion(PMIGridWriter *pm_wrt, const PMIGridAccumulator &pm_acc,
                                 int block_count);
#endif
 
} // namespace energy
} // namespace stormm
 
#include "pmigrid.tpp"
 
#endif