bscpkgs/pkgs/intel-oneapi/patch_intel.nix

{ stdenv
, stdenvNoCC
, lib
, symlinkJoin
, autoPatchelfHook
, wrapCCWith
, overrideCC
, gcc
, hwloc
, libelf
, libffi_3_3
, libpsm2
, libuuid
, libxml2
, numactl
, ocl-icd
, openssl
, python3
, rdma-core
, ucx
, zlib
, makeOverridable
, recurseIntoAttrs
}:

makeOverridable (
{
  unpatched,
  components ? { },
  extraPackages ? components.extraPackages or []
}:

let
  inherit (builtins) attrValues filter mapAttrs removeAttrs;

  __components = removeAttrs components ["extraPackages"];
  _components = __components;
  # _components = lib.traceSeqN 2 {
  #   inherit unpatched __components;
  #   deps = builtins.map (x: "${x.pname}-${x.version}") unpatched.deps;
  # } __components;

  wrapIntel = { cc, extraBuildCommands ? "", extraInstall ? "", ... }@args:
    let
      targetConfig = stdenv.targetPlatform.config;
    in (wrapCCWith {
      inherit cc;
      extraBuildCommands = ''
        echo "-isystem ${cc}/include" >> $out/nix-support/cc-cflags
        echo "-isystem ${cc}/include/intel64" >> $out/nix-support/cc-cflags

        echo "-L${gcc.cc}/lib/gcc/${targetConfig}/${gcc.version}" >> $out/nix-support/cc-ldflags
        echo "-L${gcc.cc.lib}/lib" >> $out/nix-support/cc-ldflags
        echo "-L${cc}/lib" >> $out/nix-support/cc-ldflags

        echo "--gcc-toolchain=${gcc.cc}" >> $out/nix-support/libcxx-cxxflags

        for dir in ${gcc.cc}/include/c++/*; do
          echo "-isystem $dir" >> $out/nix-support/libcxx-cxxflags
        done
        for dir in ${gcc.cc}/include/c++/*/${targetConfig}; do
          echo "-isystem $dir" >> $out/nix-support/libcxx-cxxflags
        done

        # Need the gcc in the path
        echo 'export "PATH=${gcc}/bin:$PATH"' >> $out/nix-support/cc-wrapper-hook

        # Disable hardening by default
        echo "" > $out/nix-support/add-hardening.sh

        # For some reason, If we don't resolve the realpath things go wrong
        for stddef in ${cc}/lib/clang/*/include/stddef.h ; do
          dir=$(dirname $(realpath "$stddef"))
          echo "-isystem $dir" >> $out/nix-support/cc-cflags
        done

      '' + extraBuildCommands;
    } // (removeAttrs args ["cc" "extraBuildCommands" "extraInstall"])
    ).overrideAttrs (old: {
      installPhase = old.installPhase + extraInstall;
    });

in
stdenvNoCC.mkDerivation (finalAttrs: {
  pname = lib.removeSuffix "-extracted" unpatched.pname;
  inherit (unpatched) version;
  src = unpatched;

  phases = [ "installPhase" "fixupPhase" ];
  buildInputs = [
    libffi_3_3
    libelf
    libxml2
    hwloc
    numactl
    libuuid
    libpsm2
    zlib
    ocl-icd
    rdma-core
    ucx
    openssl
    python3
    stdenv.cc.cc.lib
  ] ++ extraPackages;

  autoPatchelfIgnoreMissingDeps = [ "libhwloc.so.5" "libcuda.so.1" "libze_loader.so.1" ];

  nativeBuildInputs = [ autoPatchelfHook ];
  installPhase = ''
    cp -r $src/opt/intel/oneapi/ $out
  '';

  passthru = let
    pkgs = mapAttrs
      (folder: version: symlinkJoin {
        pname = "intel-${folder}";
        inherit version;
        paths = ["${finalAttrs.finalPackage}/${folder}/${version}"];
      })
      _components;
  in pkgs // {
    inherit unpatched;
    pkgs = recurseIntoAttrs pkgs;
    components = _components;

    # This contains all packages properly symlinked into toplevel directories
    # in $out.
    #
    # NOTE: there are clashes with packages that have symlinks outside their
    # scope (libtcm and env/vars.sh)
    all = symlinkJoin {
      pname = finalAttrs.finalPackage + "-symlinked";
      inherit (finalAttrs.finalPackage) version;
      paths = filter lib.isDerivation (attrValues finalAttrs.finalPackage.pkgs);
    };

    stdenv = overrideCC stdenv finalAttrs.finalPackage.cc;

    cc = wrapIntel {
      cc = finalAttrs.finalPackage.pkgs.compiler;
      extraBuildCommands = ''
        wrap icx  $wrapper $ccPath/icx
        wrap icpx $wrapper $ccPath/icpx
        wrap ifx  $wrapper $ccPath/ifx

        ln -s $out/bin/icpx $out/bin/c++
        ln -s $out/bin/icx  $out/bin/cc

        # Use this to detect when a compiler subprocess is called
        # from icpx (--fsycl-host-compiler)
        echo 'export "NIX_CC_WRAPPER_INTEL=1"' >> $out/nix-support/cc-wrapper-hook

        sed -i 's/.*isCxx=0/isCxx=1/' $out/bin/icpx

        # oneMath looks for sycl libraries in bin/../lib
        ln -s ${finalAttrs.finalPackage.pkgs.compiler}/lib $out/lib
        ln -s ${finalAttrs.finalPackage.pkgs.compiler}/include $out/include
      '';
      extraInstall = ''
        export named_cc="icx"
        export named_cxx="icpx"
        export named_fc="ifx"
      '';
    };
  };

}))