Merge remote-tracking branch 'origin/develop' into lotoFix

.gitmodules

@@ -1,3 +1,6 @@
 [submodule "lib/kokkos"]
 	path = lib/kokkos
 	url = https://github.com/kokkos/kokkos.git
+[submodule "lib/kokkos-kernels"]
+	path = lib/kokkos-kernels
+	url = https://github.com/kokkos/kokkos-kernels

CMakeLists.txt

@@ -30,12 +30,17 @@ option(OMP_AVAIL "Build with OMP" ON)
 option(HDF5_AVAIL "Build with HDF5" ON)
 option(HDF5_SERIAL "Force build to accomodate serial only HDF5, but still use MPI" OFF)
 option(BUILD_DOC "Build documentation" ON)
+
+############## KOKKOS #################
 if(OMP_AVAIL)
   option(Kokkos_ENABLE_OPENMP "Use kokkos with omp" ON)
 else()
   option(Kokkos_ENABLE_OPENMP "Use kokkos with omp" OFF)
   option(Kokkos_ENABLE_SERIAL "Use kokkos in serial" ON)
 endif()
+set(KokkosKernels_ADD_DEFAULT_ETI OFF CACHE BOOL "faster build")
+set(KokkosKernels_ENABLED_COMPONENTS BLAS CACHE STRING "faster build")
+
 ############### SOURCE ###############
 
 FILE(GLOB_RECURSE SOURCE_FILES src src/*.cpp)
@@ -59,9 +64,9 @@ set_target_properties(runTests PROPERTIES EXCLUDE_FROM_ALL TRUE)
 add_dependencies(phoebe spglib_dep pugixml_dep eigen_dep)
 add_dependencies(runTests spglib_dep pugixml_dep eigen_dep)
 
-target_link_libraries(phoebe symspg pugixml kokkos nlohmann_json::nlohmann_json)
+target_link_libraries(phoebe symspg pugixml kokkos Kokkos::kokkoskernels nlohmann_json::nlohmann_json)
 enable_testing()
-target_link_libraries(runTests symspg pugixml gtest_main kokkos nlohmann_json::nlohmann_json)
+target_link_libraries(runTests symspg pugixml gtest_main kokkos Kokkos::kokkoskernels nlohmann_json::nlohmann_json)
 
 gtest_discover_tests(
     runTests

lib/CMakeLists.txt

@@ -1,4 +1,5 @@
 add_subdirectory(lib/kokkos)
+add_subdirectory(lib/kokkos-kernels)
 
 include(ExternalProject)
 include(FetchContent)

src/algebra/PMatrix.cpp

@@ -154,11 +154,15 @@ ParallelMatrix<double>::diagonalize() {
      mpi->time();
   }
 
+  Kokkos::Profiling::pushRegion("pdsyevd");
+
   // call the function to now diagonalize
   pdsyevd_(&jobz, &uplo, &numRows_, mat, &ia, &ja, &descMat_[0], eigenvalues,
           eigenvectors.mat, &ia, &ja, &eigenvectors.descMat_[0],
           work, &lwork, iwork, &liwork, &info);
 
+  Kokkos::Profiling::popRegion();
+
   if(mpi->mpiHead()) {
      std::cout << "Matrix diagonalization completed." << std::endl;
      mpi->time();
@@ -313,11 +317,16 @@ std::tuple<std::vector<double>, ParallelMatrix<double>>
   allocate(work, 1);
   allocate(iwork, 1);
 
+
+  Kokkos::Profiling::pushRegion("pdsyevr");
+
   pdsyevr_(&jobz, &range, &uplo,  &numRows_, mat, &ia, &ja, &descMat_[0],
         &vl, &vu, &il, &iu, &m, &nz, eigenvalues,
         eigenvectors.mat, &iz, &jz, &eigenvectors.descMat_[0],
         work, &lwork, iwork, &liwork, &info);
 
+  Kokkos::Profiling::popRegion();
+
   lwork=int(work[0]);
   delete[] work;
   delete[] iwork;

src/apps/electron_wannier_transport_app.cpp

@@ -12,30 +12,36 @@
 
 void ElectronWannierTransportApp::run(Context &context) {
 
+  Kokkos::Profiling::pushRegion("ETapp.parseHamiltonians");
   auto t2 = Parser::parsePhHarmonic(context);
   auto crystal = std::get<0>(t2);
   auto phononH0 = std::get<1>(t2);
 
   auto t1 = Parser::parseElHarmonicWannier(context, &crystal);
   auto crystalEl = std::get<0>(t1);
   auto electronH0 = std::get<1>(t1);
+  Kokkos::Profiling::popRegion();
 
   // load the elph coupling
   // Note: this file contains the number of electrons
   // which is needed to understand where to place the fermi level
+  Kokkos::Profiling::pushRegion("ETapp.parseInteraction");
   InteractionElPhWan couplingElPh(crystal);
   if (std::isnan(context.getConstantRelaxationTime())) {
     couplingElPh = InteractionElPhWan::parse(context, crystal, &phononH0);
   }
+  Kokkos::Profiling::popRegion();
 
   // compute the band structure on the fine grid
   if (mpi->mpiHead()) {
     std::cout << "\nComputing electronic band structure." << std::endl;
   }
+  Kokkos::Profiling::pushRegion("ETapp.setupElBandstructure");
   Points fullPoints(crystal, context.getKMesh());
   auto t3 = ActiveBandStructure::builder(context, electronH0, fullPoints);
   auto bandStructure = std::get<0>(t3);
   auto statisticsSweep = std::get<1>(t3);
+  Kokkos::Profiling::popRegion();
 
   // print some info about how window and symmetries have reduced things
   if (mpi->mpiHead()) {
@@ -61,10 +67,12 @@ void ElectronWannierTransportApp::run(Context &context) {
   //  StatisticsSweep statisticsSweep(context, &bandStructure);
 
   // build/initialize the scattering matrix and the smearing
+  Kokkos::Profiling::pushRegion("ETapp.setupScatteringMatrix");
   ElScatteringMatrix scatteringMatrix(context, statisticsSweep, bandStructure,
                                       bandStructure, phononH0, &couplingElPh);
   scatteringMatrix.setup();
   scatteringMatrix.outputToJSON("rta_el_relaxation_times.json");
+  Kokkos::Profiling::popRegion();
 
   // solve the BTE at the relaxation time approximation level
   // we always do this, as it's the cheapest solver and is required to know
@@ -78,29 +86,28 @@ void ElectronWannierTransportApp::run(Context &context) {
   // compute the phonon populations in the relaxation time approximation.
   // Note: this is the total phonon population n (n != f(1+f) Delta n)
 
+  Kokkos::Profiling::pushRegion("ETapp.computeRTATransport");
+
   BulkEDrift driftE(statisticsSweep, bandStructure, 3);
   BulkTDrift driftT(statisticsSweep, bandStructure, 3);
   VectorBTE relaxationTimes = scatteringMatrix.getSingleModeTimes();
   VectorBTE nERTA = -driftE * relaxationTimes;
   VectorBTE nTRTA = -driftT * relaxationTimes;
 
   // compute the electrical conductivity
-  OnsagerCoefficients transportCoefficients(statisticsSweep, crystal,
-                                            bandStructure, context);
+  OnsagerCoefficients transportCoefficients(statisticsSweep, crystal, bandStructure, context);
   transportCoefficients.calcFromPopulation(nERTA, nTRTA);
   transportCoefficients.print();
   transportCoefficients.outputToJSON("rta_onsager_coefficients.json");
 
   // compute the Wigner transport coefficients
-  WignerElCoefficients wignerCoefficients(
-      statisticsSweep, crystal, bandStructure, context, relaxationTimes);
+  WignerElCoefficients wignerCoefficients(statisticsSweep, crystal, bandStructure, context, relaxationTimes);
   wignerCoefficients.calcFromPopulation(nERTA, nTRTA);
   wignerCoefficients.print();
   wignerCoefficients.outputToJSON("rta_wigner_coefficients.json");
 
   // compute the electron viscosity
-  ElectronViscosity elViscosity(context, statisticsSweep, crystal,
-                                bandStructure);
+  ElectronViscosity elViscosity(context, statisticsSweep, crystal, bandStructure);
   elViscosity.calcRTA(relaxationTimes);
   elViscosity.print();
   elViscosity.outputToJSON("rta_electron_viscosity.json");
@@ -119,9 +126,13 @@ void ElectronWannierTransportApp::run(Context &context) {
     return; // if we used the constant RTA, we can't solve the BTE exactly
   }
 
+  Kokkos::Profiling::popRegion();
+
   //---------------------------------------------------------------------------
   // start section on exact BTE solvers
 
+  Kokkos::Profiling::pushRegion("ET.exactBTESolvers");
+
   std::vector<std::string> solverBTE = context.getSolverBTE();
 
   bool doIterative = false;
@@ -182,6 +193,9 @@ void ElectronWannierTransportApp::run(Context &context) {
   }
 
   if (doRelaxons) {
+
+    Kokkos::Profiling::pushRegion("ETapp.relaxonsTransport");
+
     if (mpi->mpiHead()) {
       std::cout << "Starting relaxons BTE solver" << std::endl;
     }
@@ -210,7 +224,9 @@ void ElectronWannierTransportApp::run(Context &context) {
       std::cout << "Finished relaxons BTE solver\n\n";
       std::cout << std::string(80, '-') << "\n" << std::endl;
     }
+    Kokkos::Profiling::popRegion();
   }
+  Kokkos::Profiling::popRegion();
 }
 
 void ElectronWannierTransportApp::checkRequirements(Context &context) {

src/bands/active_bandstructure.cpp

@@ -123,11 +123,12 @@ ActiveBandStructure::ActiveBandStructure(const Points &points_,
         size_t ik = iks[start_iik+iik];
         Point point = points.getPoint(ik);
 
-
         for(int i = 0; i < numBands; i++){
           energies(i) = energies_h(iik,i);
-          for(int j = 0; j < numBands; j++){
-            eigenvectors(j,i) = eigenvectors_h(iik,j,i);
+          if (withEigenvectors) {
+            for(int j = 0; j < numBands; j++){
+              eigenvectors(j,i) = eigenvectors_h(iik,j,i);
+            }
           }
         }
 
@@ -145,7 +146,7 @@ ActiveBandStructure::ActiveBandStructure(const Points &points_,
     // the same again, but for velocities
     // TODO: I think this also returns the energies and velocities above, so we could avoid
     // the above calculation entirely (12.5% potential speedup)
-    Kokkos::Profiling::pushRegion("velocity loop");
+    Kokkos::Profiling::pushRegion("bandstructure velocity loop");
 
     int approx_batch_size = h0->estimateBatchSize(true);
 
@@ -442,6 +443,9 @@ void ActiveBandStructure::buildIndices() {
 }
 
 void ActiveBandStructure::buildSymmetries() {
+
+  Kokkos::Profiling::pushRegion("ABS.buildSymmetries");
+
   // ------------------
   // things to use in presence of symmetries
   {
@@ -477,6 +481,7 @@ void ActiveBandStructure::buildSymmetries() {
     }
     ikOld = ik;
   }
+  Kokkos::Profiling::popRegion();
 }
 
 std::tuple<ActiveBandStructure, StatisticsSweep>
@@ -497,6 +502,7 @@ ActiveBandStructure::builder(Context &context, HarmonicHamiltonian &h0,
 
     StatisticsSweep s = activeBandStructure.buildAsPostprocessing(
         context, points_, h0, withEigenvectors, withVelocities);
+
     return std::make_tuple(activeBandStructure, s);
 
   }
@@ -546,9 +552,9 @@ void ActiveBandStructure::buildOnTheFly(Window &window, Points points_,
   std::vector<int> filteredThreadPoints;
   std::vector<std::vector<int>> filteredThreadBands;
 
-  std::vector<size_t> pointsIter = mpi->divideWorkIter(points_.getNumPoints()); 
+  std::vector<size_t> pointsIter = mpi->divideWorkIter(points_.getNumPoints());
   #pragma omp for nowait schedule(static)
-  for (size_t iik = 0; iik < pointsIter.size(); iik++) { 
+  for (size_t iik = 0; iik < pointsIter.size(); iik++) {
 
     int ik = pointsIter[iik];
 

src/bte/helper_el_scattering.cpp

@@ -20,6 +20,7 @@ HelperElScattering::HelperElScattering(BaseBandStructure &innerBandStructure_,
   // 3 - the mesh is complete (if q1 and q2 are only around 0, q3 might be
   //     at the border)
 
+  Kokkos::Profiling::pushRegion("HelperElScattering");
 
   auto t1 = outerBandStructure.getPoints().getMesh();
 //  auto mesh = std::get<0>(t1);
@@ -216,6 +217,7 @@ HelperElScattering::HelperElScattering(BaseBandStructure &innerBandStructure_,
       mappedPolarData.insert({alliq3s[iiq3], allPolarData.col(iiq3)});
     }
   }
+  Kokkos::Profiling::popRegion();
 }
 
 // auto [eigenValues3Minus, nb3Minus, eigenVectors3Minus, v3Minus, bose3]

src/harmonic/phonon_h0.cpp

@@ -18,6 +18,8 @@ PhononH0::PhononH0(Crystal &crystal, const Eigen::Matrix3d &dielectricMatrix_,
   // in this section, we save as class properties a few variables
   // that are needed for the diagonalization of phonon frequencies
 
+  Kokkos::Profiling::pushRegion("phononH0 constructor");
+
   directUnitCell = crystal.getDirectUnitCell();
   Eigen::Matrix3d reciprocalUnitCell = crystal.getReciprocalUnitCell();
   volumeUnitCell = crystal.getVolumeUnitCell();
@@ -98,6 +100,7 @@ PhononH0::PhononH0(Crystal &crystal, const Eigen::Matrix3d &dielectricMatrix_,
       }
     }
 
+    // TODO also add 2d option and move this to a separate function
     longRangeCorrection1.resize(3,3,numAtoms);
     longRangeCorrection1.setZero();
     for (int ig=0; ig<numG; ++ig) {
@@ -236,6 +239,7 @@ PhononH0::PhononH0(Crystal &crystal, const Eigen::Matrix3d &dielectricMatrix_,
     double mem = getDeviceMemoryUsage();
     kokkosDeviceMemory->addDeviceMemoryUsage(mem);
   }
+  Kokkos::Profiling::popRegion();
 }
 
 // copy constructor
@@ -543,6 +547,8 @@ void PhononH0::addLongRangeTerm(Eigen::Tensor<std::complex<double>, 4> &dyn,
   // very rough estimate: geg/4/alpha > gMax = 14
   // (exp (-14) = 10^-6)
 
+  Kokkos::Profiling::pushRegion("PhononH0::addLongRangeTerm");
+
 /*
        IF (loto_2d) THEN
           geg = g1**2 + g2**2 + g3**2
@@ -654,10 +660,14 @@ void PhononH0::addLongRangeTerm(Eigen::Tensor<std::complex<double>, 4> &dyn,
       }
     }
   }
+  Kokkos::Profiling::popRegion();
 }
 
 void PhononH0::reorderDynamicalMatrix(const Eigen::Matrix3d& directUnitCell,
                                       const Eigen::Tensor<double, 7>& forceConstants) {
+
+  Kokkos::Profiling::pushRegion("PhononH0::reorderDynamicalMatrix");
+
   // this part can actually be expensive to execute, so we compute it once
   // at the beginning
 
@@ -752,13 +762,16 @@ void PhononH0::reorderDynamicalMatrix(const Eigen::Matrix3d& directUnitCell,
 
   // wsCache.resize(0, 0, 0, 0, 0);
   // forceConstants.resize(0, 0, 0, 0, 0, 0, 0);
+  Kokkos::Profiling::popRegion();
 }
 
 void PhononH0::shortRangeTerm(Eigen::Tensor<std::complex<double>, 4> &dyn,
                               const Eigen::VectorXd &q) {
   // calculates the dynamical matrix at q from the (short-range part of the)
   // force constants21, by doing the Fourier transform of the force constants
 
+  Kokkos::Profiling::pushRegion("phononH0.shortRangeTerm");
+
   std::vector<std::complex<double>> phases(numBravaisVectors);
   for (int iR = 0; iR < numBravaisVectors; iR++) {
     Eigen::Vector3d r = bravaisVectors.col(iR);
@@ -790,6 +803,7 @@ void PhononH0::shortRangeTerm(Eigen::Tensor<std::complex<double>, 4> &dyn,
     //    }
     //  }
     //}
+  Kokkos::Profiling::popRegion();
 }
 
 std::tuple<Eigen::VectorXd, Eigen::MatrixXcd>
@@ -798,6 +812,8 @@ PhononH0::dynDiagonalize(Eigen::Tensor<std::complex<double>, 4> &dyn) {
   // On input:  speciesMasses = masses, in amu
   // On output: w2 = energies, z = displacements
 
+  Kokkos::Profiling::pushRegion("PhononH0::dynDiagonalize");
+
   // fill the two-indices dynamical matrix
 
   Eigen::MatrixXcd dyn2Tmp(numBands, numBands);
@@ -859,7 +875,7 @@ PhononH0::dynDiagonalize(Eigen::Tensor<std::complex<double>, 4> &dyn) {
   //    printf("old = %.16e %.16e\n", x.real(), x.imag());
   //  }
   //}
-
+  Kokkos::Profiling::popRegion();
   return std::make_tuple(energies, eigenvectors);
 }
 
@@ -872,10 +888,12 @@ PhononH0::diagonalizeVelocity(Point &point) {
 
 Eigen::Tensor<std::complex<double>, 3>
 PhononH0::diagonalizeVelocityFromCoordinates(Eigen::Vector3d &coordinates) {
+
+  Kokkos::Profiling::pushRegion("PhononH0::diagonalizeVelocityFromCoordinates");
+
   Eigen::Tensor<std::complex<double>, 3> velocity(numBands, numBands, 3);
   velocity.setZero();
 
-
   bool withMassScaling = false;
     //for(int i = 0; i < 3; i++){
     //  printf("old = %.16e\n", coordinates(i));
@@ -1041,6 +1059,7 @@ PhononH0::diagonalizeVelocityFromCoordinates(Eigen::Vector3d &coordinates) {
   //    }
   //  }
   //}
+  Kokkos::Profiling::popRegion();
   return velocity;
 }