From dcd7efd1f993419527a63313d7158cef746b0106 Mon Sep 17 00:00:00 2001
From: Phillip Seeber <phillip.seeber@uni-jena.de>
Date: Thu, 5 Oct 2023 14:24:58 +0200
Subject: [PATCH] cp2k: enable GPU offloading, sirius planewaves and
 optimisations

---
 .../science/chemistry/cp2k/default.nix        | 98 ++++++++++++++++---
 1 file changed, 85 insertions(+), 13 deletions(-)

diff --git a/pkgs/applications/science/chemistry/cp2k/default.nix b/pkgs/applications/science/chemistry/cp2k/default.nix
index e34810a3982b..052d791c0bb7 100644
--- a/pkgs/applications/science/chemistry/cp2k/default.nix
+++ b/pkgs/applications/science/chemistry/cp2k/default.nix
@@ -1,15 +1,57 @@
-{ lib, stdenv, fetchFromGitHub, mpiCheckPhaseHook, python3, gfortran, blas, lapack
-, fftw, libint, libvori, libxc, mpi, gsl, scalapack, openssh, makeWrapper
-, libxsmm, spglib, which, pkg-config, plumed, zlib
+{ lib
+, stdenv
+, fetchFromGitHub
+, mpiCheckPhaseHook
+, python3
+, gfortran
+, blas
+, lapack
+, fftw
+, libint
+, libvori
+, libxc
+, mpi
+, gsl
+, scalapack
+, openssh
+, makeWrapper
+, libxsmm
+, spglib
+, which
+, pkg-config
+, plumed
+, zlib
+, hdf5-fortran
+, sirius
+, libvdwxc
+, spla
+, spfft
 , enableElpa ? false
 , elpa
-} :
+, gpuBackend ? "none"
+, cudaPackages
+# gpuVersion needs to be set for both CUDA as well as ROCM hardware.
+# gpuArch is only required for the ROCM stack.
+# Change to a value suitable for your target GPU.
+# For AMD values see https://github.com/cp2k/cp2k/blob/master/INSTALL.md#2v-rocmhip-support-for-amd-gpu
+# and for Nvidia see https://github.com/cp2k/cp2k/blob/master/INSTALL.md#2i-cuda-optional-improved-performance-on-gpu-systems
+, gpuVersion ? "Mi100"
+, gpuArch ? "gfx908"
+, rocm-core
+, hip
+, hipblas
+, hipfft
+, rocblas
+}:
+
+assert builtins.elem gpuBackend [ "none" "cuda" "rocm" ];
 
 let
   cp2kVersion = "psmp";
   arch = "Linux-x86-64-gfortran";
 
-in stdenv.mkDerivation rec {
+in
+stdenv.mkDerivation rec {
   pname = "cp2k";
   version = "2023.2";
 
@@ -36,7 +78,16 @@ in stdenv.mkDerivation rec {
     lapack
     plumed
     zlib
-  ] ++ lib.optional enableElpa elpa;
+    hdf5-fortran
+    sirius
+    spla
+    spfft
+    libvdwxc
+  ]
+  ++ lib.optional enableElpa elpa
+  ++ lib.optional (gpuBackend == "cuda") cudaPackages.cudatoolkit
+  ++ lib.optional (gpuBackend == "rocm") [hip rocm-core hipblas hipfft rocblas]
+  ;
 
   propagatedBuildInputs = [ mpi ];
   propagatedUserEnvPkgs = [ mpi ];
@@ -46,7 +97,7 @@ in stdenv.mkDerivation rec {
     "VERSION=${cp2kVersion}"
   ];
 
-  doCheck = true;
+  doCheck = gpuBackend == "none";
 
   enableParallelBuilding = true;
 
@@ -64,25 +115,46 @@ in stdenv.mkDerivation rec {
     FC         = mpif90
     LD         = mpif90
     AR         = ar -r
+    ${lib.strings.optionalString (gpuBackend == "cuda") ''
+    OFFLOAD_CC = nvcc
+    OFFLOAD_FLAGS = -O3 -g -w --std=c++11
+    OFFLOAD_TARGET = cuda
+    GPUVER = ${gpuVersion}
+    CXX = mpicxx
+    CXXFLAGS = -std=c++11 -fopenmp
+    ''}
+    ${lib.strings.optionalString (gpuBackend == "rocm") ''
+    GPUVER = ${gpuVersion}
+    OFFLOAD_CC = hipcc
+    OFFLOAD_FLAGS = -fopenmp -m64 -pthread -fPIC -D__GRID_HIP -O2 --offload-arch=${gpuArch} --rocm-path=${rocm-core}
+    OFFLOAD_TARGET = hip
+    CXX = mpicxx
+    CXXFLAGS = -std=c++11 -fopenmp -D__HIP_PLATFORM_AMD__
+    ''}
     DFLAGS     = -D__FFTW3 -D__LIBXC -D__LIBINT -D__parallel -D__SCALAPACK \
                  -D__MPI_VERSION=3 -D__F2008 -D__LIBXSMM -D__SPGLIB \
                  -D__MAX_CONTR=4 -D__LIBVORI ${lib.optionalString enableElpa "-D__ELPA"} \
-                 -D__PLUMED2
-    CFLAGS    = -fopenmp
+                 -D__PLUMED2 -D__HDF5 -D__GSL -D__SIRIUS -D__LIBVDWXC -D__SPFFT -D__SPLA \
+                 ${lib.strings.optionalString (gpuBackend == "cuda") "-D__OFFLOAD_CUDA -D__DBCSR_ACC"} \
+                 ${lib.strings.optionalString (gpuBackend == "rocm") "-D__OFFLOAD_HIP -D__DBCSR_ACC -D__NO_OFFLOAD_PW"}
+    CFLAGS    = -fopenmp -I${lib.getDev hdf5-fortran}/include -I${lib.getDev gsl}/include
     FCFLAGS    = \$(DFLAGS) -O2 -ffree-form -ffree-line-length-none \
                  -ftree-vectorize -funroll-loops -msse2 \
                  -std=f2008 \
                  -fopenmp -ftree-vectorize -funroll-loops \
-                 -I${lib.getDev libxc}/include -I${lib.getDev libxsmm}/include \
-                 -I${libint}/include ${lib.optionalString enableElpa "$(pkg-config --variable=fcflags elpa)"}
+                 -I${lib.getDev libint}/include ${lib.optionalString enableElpa "$(pkg-config --variable=fcflags elpa)"} \
+                 -I${lib.getDev sirius}/include/sirius \
+                 -I${lib.getDev libxc}/include -I${lib.getDev libxsmm}/include
     LIBS       = -lfftw3 -lfftw3_threads \
                  -lscalapack -lblas -llapack \
                  -lxcf03 -lxc -lxsmmf -lxsmm -lsymspg \
                  -lint2 -lstdc++ -lvori \
                  -lgomp -lpthread -lm \
                  -fopenmp ${lib.optionalString enableElpa "$(pkg-config --libs elpa)"} \
-                 -lz -ldl -lstdc++ ${lib.optionalString (mpi.pname == "openmpi") "$(mpicxx --showme:link)"} \
-                 -lplumed
+                 -lz -ldl ${lib.optionalString (mpi.pname == "openmpi") "$(mpicxx --showme:link)"} \
+                 -lplumed -lhdf5_fortran -lhdf5_hl -lhdf5 -lgsl -lsirius -lspla -lspfft -lvdwxc \
+                 ${lib.strings.optionalString (gpuBackend == "cuda") "-lcudart -lnvrtc -lcuda -lcublas"} \
+                 ${lib.strings.optionalString (gpuBackend == "rocm") "-lamdhip64 -lhipfft -lhipblas -lrocblas"}
     LDFLAGS    = \$(FCFLAGS) \$(LIBS)
     include ${plumed}/lib/plumed/src/lib/Plumed.inc
     EOF