Preconditioning precision

src/CMakeLists.txt

@@ -44,7 +44,7 @@ set(SOURCES
  HamiltonianMVP_DMStrategy.cc 
  MVPSolver.cc 
  HamiltonianMVPSolver.cc 
- OrbitalsPreconditioning.cc 
+ MGOrbitalsPreconditioning.cc
  DFTsolver.cc 
  NonOrthoDMStrategy.cc 
  FullyOccupiedNonOrthoDMStrategy.cc 

src/Control.cc

@@ -263,7 +263,8 @@ void Control::print(std::ostream& os)
         os << " Localization radius       = " << cut_radius << std::endl;
     os << std::endl;
 
-    os << " preconditioner factor:" << precond_factor << std::endl;
+    os << " preconditioner factor: " << precond_factor << std::endl;
+    os << " preconditioner precision: " << precond_precision_ << std::endl;
     if (precond_type_ == 10)
     {
         os << " Multigrid preconditioning for wave functions:" << std::endl;
@@ -332,7 +333,7 @@ void Control::sync(void)
     if (onpe0 && verbose > 0)
         (*MPIdata::sout) << "Control::sync()" << std::endl;
     // pack
-    const short size_short_buffer = 92;
+    const short size_short_buffer = 93;
     short* short_buffer           = new short[size_short_buffer];
     if (mype_ == 0)
     {
@@ -423,6 +424,7 @@ void Control::sync(void)
         short_buffer[89] = MD_last_step_;
         short_buffer[90] = (short)static_cast<int>(poisson_lap_type_);
         short_buffer[91] = poisson_pc_data_;
+        short_buffer[92] = precond_precision_;
     }
     else
     {
@@ -635,8 +637,9 @@ void Control::sync(void)
     max_electronic_steps_tight_      = short_buffer[86];
     hartree_reset_                   = short_buffer[88];
     MD_last_step_                    = short_buffer[89];
-    poisson_lap_type_ = static_cast<PoissonFDtype>(short_buffer[90]);
-    poisson_pc_data_  = short_buffer[91];
+    poisson_lap_type_  = static_cast<PoissonFDtype>(short_buffer[90]);
+    poisson_pc_data_   = short_buffer[91];
+    precond_precision_ = short_buffer[92];
 
     numst    = int_buffer[0];
     nel_     = int_buffer[1];
@@ -1479,8 +1482,9 @@ void Control::setOptions(const boost::program_options::variables_map& vm)
         std::cout << "Outer solver type: " << str << std::endl;
         assert(it_algo_type_ >= 0);
 
-        mg_levels_     = vm["Quench.preconditioner_num_levels"].as<short>() - 1;
-        precond_factor = vm["Quench.step_length"].as<float>();
+        mg_levels_ = vm["Quench.preconditioner_num_levels"].as<short>() - 1;
+        precond_precision_ = vm["Quench.preconditioner_precision"].as<short>();
+        precond_factor     = vm["Quench.step_length"].as<float>();
         if (precond_factor < 0.)
         {
             switch (lap_type)

src/Control.h

@@ -414,6 +414,8 @@ class Control
 
     short orthof; // orthogonalization frequency
 
+    short precond_precision_;
+
     // screening constant for potential mixing
     float screening_const;
 

src/OrbitalsPreconditioning.cc → src/MGOrbitalsPreconditioning.cc

@@ -7,7 +7,7 @@
 // This file is part of MGmol. For details, see https://github.com/llnl/mgmol.
 // Please also read this link https://github.com/llnl/mgmol/LICENSE
 
-#include "OrbitalsPreconditioning.h"
+#include "MGOrbitalsPreconditioning.h"
 
 #include "Control.h"
 #include "ExtendedGridOrbitals.h"
@@ -20,27 +20,30 @@
 #include "ProjectedMatricesInterface.h"
 
 template <class OrbitalsType, typename PDataType>
-OrbitalsPreconditioning<OrbitalsType, PDataType>::~OrbitalsPreconditioning()
+MGOrbitalsPreconditioning<OrbitalsType, PDataType>::MGOrbitalsPreconditioning(
+    const short mg_levels, const short lap_type)
+    : mg_levels_(mg_levels), lap_type_(lap_type), is_set_(false){};
+
+template <class OrbitalsType, typename PDataType>
+MGOrbitalsPreconditioning<OrbitalsType, PDataType>::~MGOrbitalsPreconditioning()
 {
     assert(is_set_);
     assert(precond_);
 }
 
 template <class OrbitalsType, typename PDataType>
-void OrbitalsPreconditioning<OrbitalsType, PDataType>::setup(
-    OrbitalsType& orbitals, const short mg_levels, const short lap_type,
-    MasksSet* currentMasks, const std::shared_ptr<LocalizationRegions>& lrs)
+void MGOrbitalsPreconditioning<OrbitalsType, PDataType>::setup(
+    OrbitalsType& orbitals, MasksSet* currentMasks,
+    const std::shared_ptr<LocalizationRegions>& lrs)
 {
     assert(!is_set_);
 
-    lap_type_ = lap_type;
-
     Control& ct(*(Control::instance()));
     Mesh* mymesh = Mesh::instance();
     const pb::Grid& mygrid(mymesh->grid());
 
     precond_ = std::make_shared<Preconditioning<PDataType>>(
-        lap_type, mg_levels, mygrid, ct.bcWF);
+        lap_type_, mg_levels_, mygrid, ct.bcWF);
 
     if (currentMasks != nullptr)
     {
@@ -78,7 +81,7 @@ void OrbitalsPreconditioning<OrbitalsType, PDataType>::setup(
 }
 
 template <class OrbitalsType, typename PDataType>
-void OrbitalsPreconditioning<OrbitalsType, PDataType>::precond_mg(
+void MGOrbitalsPreconditioning<OrbitalsType, PDataType>::precond(
     OrbitalsType& orbitals)
 {
     assert(is_set_);
@@ -127,15 +130,16 @@ void OrbitalsPreconditioning<OrbitalsType, PDataType>::precond_mg(
 
 #ifdef PRINT_OPERATIONS
     if (onpe0)
-        (*MPIdata::sout) << "OrbitalsPreconditioning<OrbitalsType,PDataType>::"
-                            "precond_mg() done"
-                         << endl;
+        (*MPIdata::sout)
+            << "MGOrbitalsPreconditioning<OrbitalsType,PDataType>::"
+               "precond_mg() done"
+            << endl;
 #endif
     precond_tm_.stop();
 }
 
 template <class OrbitalsType, typename PDataType>
-void OrbitalsPreconditioning<OrbitalsType, PDataType>::setGamma(
+void MGOrbitalsPreconditioning<OrbitalsType, PDataType>::setGamma(
     const pb::Lap<ORBDTYPE>& lapOper, const Potentials& pot,
     const short mg_levels, ProjectedMatricesInterface* proj_matrices)
 {
@@ -163,11 +167,13 @@ void OrbitalsPreconditioning<OrbitalsType, PDataType>::setGamma(
 }
 
 template <class OrbitalsType, typename PDataType>
-void OrbitalsPreconditioning<OrbitalsType, PDataType>::printTimers(
+void MGOrbitalsPreconditioning<OrbitalsType, PDataType>::printTimers(
     std::ostream& os)
 {
     precond_tm_.print(os);
 }
 
-template class OrbitalsPreconditioning<LocGridOrbitals, MGPRECONDTYPE>;
-template class OrbitalsPreconditioning<ExtendedGridOrbitals, MGPRECONDTYPE>;
+template class MGOrbitalsPreconditioning<LocGridOrbitals, float>;
+template class MGOrbitalsPreconditioning<LocGridOrbitals, double>;
+template class MGOrbitalsPreconditioning<ExtendedGridOrbitals, float>;
+template class MGOrbitalsPreconditioning<ExtendedGridOrbitals, double>;

src/MGOrbitalsPreconditioning.h

@@ -0,0 +1,79 @@
+// Copyright (c) 2017, Lawrence Livermore National Security, LLC and
+// UT-Battelle, LLC.
+// Produced at the Lawrence Livermore National Laboratory and the Oak Ridge
+// National Laboratory.
+// LLNL-CODE-743438
+// All rights reserved.
+// This file is part of MGmol. For details, see https://github.com/llnl/mgmol.
+// Please also read this link https://github.com/llnl/mgmol/LICENSE
+
+#ifndef MGMOL_MGOrbitalsPreconditioning_H
+#define MGMOL_MGOrbitalsPreconditioning_H
+
+#include "GridFuncVector.h"
+#include "Lap.h"
+#include "Map2Masks.h"
+#include "OrbitalsPreconditioning.h"
+#include "Preconditioning.h"
+
+#include <memory>
+
+// class Masks4Orbitals;
+// class MasksSet;
+class ProjectedMatricesInterface;
+class Potentials;
+// class LocalizationRegions;
+
+template <class OrbitalsType, typename PDataType>
+class MGOrbitalsPreconditioning : public OrbitalsPreconditioning<OrbitalsType>
+{
+private:
+#ifdef HAVE_MAGMA
+    using memory_space_type = MemorySpace::Device;
+#else
+    using memory_space_type = MemorySpace::Host;
+#endif
+
+    std::shared_ptr<Preconditioning<PDataType>> precond_;
+
+    // work arrays with preconditioner precision
+    std::shared_ptr<pb::GridFuncVector<PDataType, memory_space_type>>
+        gfv_work1_;
+    std::shared_ptr<pb::GridFuncVector<PDataType, memory_space_type>>
+        gfv_work2_;
+
+    // tmp work array for case ORBDTYPE!=PDataType
+    std::shared_ptr<pb::GridFuncVector<ORBDTYPE, memory_space_type>> gfv_work3_;
+
+    short mg_levels_;
+
+    short lap_type_;
+
+    bool is_set_;
+
+    // coefficient for preconditioning
+    double gamma_;
+
+    // timers
+    static Timer precond_tm_;
+
+    std::shared_ptr<Map2Masks> map2masks_;
+
+public:
+    MGOrbitalsPreconditioning(const short mg_levels, const short lap_type);
+
+    ~MGOrbitalsPreconditioning();
+
+    void setup(OrbitalsType& orbitals, MasksSet*,
+        const std::shared_ptr<LocalizationRegions>&) override;
+    void precond(OrbitalsType& orbitals) override;
+    void setGamma(const pb::Lap<ORBDTYPE>& lapOper, const Potentials& pot,
+        const short mg_levels, ProjectedMatricesInterface* proj_matrices);
+    static void printTimers(std::ostream& os);
+};
+
+template <class OrbitalsType, typename PDataType>
+Timer MGOrbitalsPreconditioning<OrbitalsType, PDataType>::precond_tm_(
+    "MGOrbitalsPreconditioning::precond");
+
+#endif

src/MGmol.cc

@@ -42,14 +42,14 @@
 #include "LocalMatrices2ReplicatedMatrix.h"
 #include "LocalizationRegions.h"
 #include "MDfiles.h"
+#include "MGOrbitalsPreconditioning.h"
 #include "MGkernels.h"
 #include "MGmol.h"
 #include "MLWFTransform.h"
 #include "MPIdata.h"
 #include "MVPSolver.h"
 #include "MasksSet.h"
 #include "Mesh.h"
-#include "OrbitalsPreconditioning.h"
 #include "PackedCommunicationBuffer.h"
 #include "PoissonInterface.h"
 #include "Potentials.h"
@@ -960,7 +960,8 @@ void MGmol<OrbitalsType>::printTimers()
         ChebyshevApproximation<
             dist_matrix::DistMatrix<DISTMATDTYPE>>::printTimers(os_);
     }
-    OrbitalsPreconditioning<OrbitalsType, MGPRECONDTYPE>::printTimers(os_);
+    MGOrbitalsPreconditioning<OrbitalsType, float>::printTimers(os_);
+    MGOrbitalsPreconditioning<OrbitalsType, double>::printTimers(os_);
     MDfiles::printTimers(os_);
     ChebyshevApproximationInterface::printTimers(os_);
 }
@@ -1180,10 +1181,35 @@ void MGmol<OrbitalsType>::precond_mg(OrbitalsType& phi)
     Potentials& pot            = hamiltonian_->potential();
     pb::Lap<ORBDTYPE>* lapOper = hamiltonian_->lapOper();
 
-    orbitals_precond_->setGamma(
-        *lapOper, pot, ct.getMGlevels(), proj_matrices_.get());
+    const short precision = ct.precond_precision_;
+    if (precision == 32)
+    {
+        using OrbitalsPrecond = MGOrbitalsPreconditioning<OrbitalsType, float>;
+
+        std::shared_ptr<OrbitalsPrecond> orbitals_precond
+            = std::dynamic_pointer_cast<OrbitalsPrecond>(orbitals_precond_);
+
+        orbitals_precond->setGamma(
+            *lapOper, pot, ct.getMGlevels(), proj_matrices_.get());
+    }
+    else if (precision == 64)
+    {
+        using OrbitalsPrecond = MGOrbitalsPreconditioning<OrbitalsType, double>;
+
+        std::shared_ptr<OrbitalsPrecond> orbitals_precond
+            = std::dynamic_pointer_cast<OrbitalsPrecond>(orbitals_precond_);
+
+        orbitals_precond->setGamma(
+            *lapOper, pot, ct.getMGlevels(), proj_matrices_.get());
+    }
+    else
+    {
+        std::cerr << "Precision " << precision
+                  << " not supported for orbitals preconditioner!!!"
+                  << std::endl;
+    }
 
-    orbitals_precond_->precond_mg(phi);
+    orbitals_precond_->precond(phi);
 }
 
 template <class OrbitalsType>

src/MGmol.h

@@ -46,9 +46,9 @@ class IonicAlgorithm;
 #include "Forces.h"
 #include "Ions.h"
 #include "LocGridOrbitals.h"
+#include "MGOrbitalsPreconditioning.h"
 #include "MGmolInterface.h"
 #include "OrbitalsExtrapolation.h"
-#include "OrbitalsPreconditioning.h"
 #include "Rho.h"
 #include "SpreadPenaltyInterface.h"
 #include "SpreadsAndCenters.h"
@@ -100,8 +100,7 @@ class MGmol : public MGmolInterface
 
     std::shared_ptr<HDFrestart> h5f_file_;
 
-    std::shared_ptr<OrbitalsPreconditioning<OrbitalsType, MGPRECONDTYPE>>
-        orbitals_precond_;
+    std::shared_ptr<OrbitalsPreconditioning<OrbitalsType>> orbitals_precond_;
 
     double total_energy_;
     std::shared_ptr<ConstraintSet> constraints_;