From 7e38ee602ed1cae15238d64abd74e48e676b992c Mon Sep 17 00:00:00 2001
From: chn <chn@chn.moe>
Date: Sun, 22 Sep 2024 17:18:24 +0800
Subject: [PATCH] packages.vasp.intel: fix

---
 modules/packages/vasp.nix       | 19 +++++---
 modules/services/slurm.nix      |  9 ++++
 packages/vasp/intel/default.nix | 78 ++++++++++++++++++++-------------
 3 files changed, 70 insertions(+), 36 deletions(-)

diff --git a/modules/packages/vasp.nix b/modules/packages/vasp.nix
index 8184fe19..37b66bc5 100644
--- a/modules/packages/vasp.nix
+++ b/modules/packages/vasp.nix
@@ -8,10 +8,19 @@ inputs:
   # TODO: add more options to correctly configure VASP
   config = let inherit (inputs.config.nixos.packages) vasp; in inputs.lib.mkIf (vasp != null)
   {
-    nixos.packages.packages._packages = (with inputs.pkgs.localPackages.vasp; [ intel vtstscripts ])
-      ++ (with inputs.pkgs.localPackages; [ py4vasp vaspkit ])
-      ++ (inputs.lib.optional
-        (let inherit (inputs.config.nixos.system.nixpkgs) cuda; in cuda.enable && cuda.capabilities != null)
-        inputs.pkgs.localPackages.vasp.nvidia);
+    nixos.packages.packages._packages =
+      (with inputs.pkgs.localPackages.vasp;
+      [
+        (intel.override
+        {
+          integratedWithSlurm = inputs.config.nixos.services.slurm.enable;
+          slurm = inputs.config.services.slurm.package;
+        })
+        vtstscripts
+      ])
+        ++ (with inputs.pkgs.localPackages; [ py4vasp vaspkit ])
+        ++ (inputs.lib.optional
+          (let inherit (inputs.config.nixos.system.nixpkgs) cuda; in cuda.enable && cuda.capabilities != null)
+          inputs.pkgs.localPackages.vasp.nvidia);
   };
 }
diff --git a/modules/services/slurm.nix b/modules/services/slurm.nix
index 794fc505..22941115 100644
--- a/modules/services/slurm.nix
+++ b/modules/services/slurm.nix
@@ -52,6 +52,15 @@ inputs:
                 buildInputs = prev.buildInputs or [] ++ additionalInputs;
                 LDFLAGS = prev.LDFLAGS or [] ++ additionalFlags;
                 nativeBuildInputs = prev.nativeBuildInputs ++ [ inputs.pkgs.wrapGAppsHook ];
+                postInstall =
+                ''
+                  pushd contribs/pmi2
+                  make install
+                  popd
+                  pushd contribs/pmi
+                  make install
+                  popd
+                '' + prev.postInstall;
               }
             );
           client.enable = true;
diff --git a/packages/vasp/intel/default.nix b/packages/vasp/intel/default.nix
index 0a8801f2..34b2cb26 100644
--- a/packages/vasp/intel/default.nix
+++ b/packages/vasp/intel/default.nix
@@ -1,35 +1,51 @@
 {
-  stdenv, src, rsync, which, wannier90, hdf5, vtst, mpi, mkl
-}: stdenv.mkDerivation
-{
-  name = "vasp-intel";
-  inherit src;
-  # patches = [ ../vtst.patch ];
-  configurePhase =
-  ''
-    cp ${./makefile.include} makefile.include
-    chmod +w makefile.include
-    cp ${../constr_cell_relax.F} src/constr_cell_relax.F
-    # cp -r ${vtst}/* src
-    chmod -R +w src
-  '';
-  buildInputs = [ hdf5 wannier90 mkl ];
-  nativeBuildInputs = [ rsync which mpi ];
-  installPhase =
-  ''
-    mkdir -p $out/bin
-    for i in std gam ncl; do cp bin/vasp_$i $out/bin/vasp-$i; done
-    mkdir $out/src
-    ln -s ${src} $out/src/vasp
-    ln -s ${vtst} $out/src/vtst
-  '';
+  stdenv, src, writeShellScriptBin, lib,
+  rsync, which, wannier90, hdf5, vtst, mpi, mkl, libfabric,
+  integratedWithSlurm ? false, slurm
+}:
+let vasp = stdenv.mkDerivation
+  {
+    name = "vasp-intel";
+    inherit src;
+    # patches = [ ../vtst.patch ];
+    configurePhase =
+    ''
+      cp ${./makefile.include} makefile.include
+      chmod +w makefile.include
+      cp ${../constr_cell_relax.F} src/constr_cell_relax.F
+      # cp -r ${vtst}/* src
+      chmod -R +w src
+    '';
+    buildInputs = [ hdf5 wannier90 mkl ];
+    nativeBuildInputs = [ rsync which mpi ];
+    installPhase =
+    ''
+      mkdir -p $out/bin
+      for i in std gam ncl; do cp bin/vasp_$i $out/bin/vasp-$i; done
+      mkdir $out/src
+      ln -s ${src} $out/src/vasp
+      ln -s ${vtst} $out/src/vtst
+    '';
 
-  # NIX_DEBUG = "7";
+    # NIX_DEBUG = "7";
 
-  # enable parallel build
-  enableParallelBuilding = true;
-  DEPS = "1";
+    # enable parallel build
+    enableParallelBuilding = true;
+    DEPS = "1";
 
-  # vasp directly include headers under ${mkl}/include/fftw
-  MKLROOT = mkl;
-}
+    # vasp directly include headers under ${mkl}/include/fftw
+    MKLROOT = mkl;
+  };
+in writeShellScriptBin "vasp-intel"
+''
+  # not sure why mpi could not find libfabric.so
+  export LD_LIBRARY_PATH=${libfabric}/lib''${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}
+
+  # intel mpi need this to talk with slurm
+  ${lib.optionalString integratedWithSlurm "export I_MPI_MPI_LIBRARY=${slurm}/lib/libpmi2.so"}
+
+  # add vasp and intel mpi in PATH
+  export PATH=${vasp}/bin:${mpi}/bin''${PATH:+:$PATH}
+
+  exec "$@"
+''