tidy

NixOS · Jul 5, 2023 · 864c0e9 · 864c0e9
1 parent fe98e88
commit 864c0e9
Showing 1 changed file with 78 additions and 83 deletions.
diff --git a/pkgs/development/libraries/science/math/libtorch/default.nix b/pkgs/development/libraries/science/math/libtorch/default.nix
@@ -9,6 +9,7 @@
   fetchpatch,
   fp16,
   fxdiv,
+  psimd,
   lib,
   magma,
   mpi,
@@ -30,7 +31,6 @@
   useNuma ? true,
   useXnnpack ? true,
   useZstd ? true,
-
 }: let
   inherit (lib) lists strings;
   setBool = bool:
@@ -95,55 +95,40 @@ in
     ];
 
     preConfigure =
-      # Use the system NCCL
-      ''
-        rm -rf third_party/nccl*
-      ''
-      # Use the system XNNPACK
-      + ''
-        rm -rf third_party/{xnnpack,XNNPACK}*
-      ''
-      # Use the system protobuf
-      + ''
-        rm -rf third_party/protobuf*
+      # Enter third_party directory
       ''
-      # Use the system pybind11
-      + ''
-        rm -rf third_party/pybind11*
-      ''
-      # Use the system sleef
-      + ''
-        rm -rf third_party/sleef*
+        pushd third_party
       ''
-      # Use the system CPUINFO
+      # Remove libraries which are unmaintained
       + ''
-        rm -rf third_party/cpuinfo*
+        rm -rf NNPACK*
+        rm -rf QNNPACK*
+        rm -rf tensorpipe
       ''
+      # Remove libraries we've migrated to Nixpkgs
       # NOTE: Cannot remove third_party/ideep* because it is used in the build.
       # However, we can remove third_party/ideep/mkl-dnn.
       + ''
-        rm -rf third_party/ideep/mkl-dnn*
-      ''
-      # Use the system zstd
-      + ''
-        rm -rf third_party/zstd*
-      ''
-      # Don't use FXdiv
-      + ''
-        rm -rf third_party/FXdiv*
-      ''
-      # Don't use FP16
-      + ''
-        rm -rf third_party/FP16*
+        rm -rf cpuinfo*
+        rm -rf fbgemm*
+        rm -rf FP16*
+        rm -rf FXdiv*
+        rm -rf ideep/mkl-dnn*
+        rm -rf protobuf*
+        rm -rf psimd*
+        rm -rf pthreadpool*
+        rm -rf pybind11*
+        rm -rf sleef*
+        rm -rf xnnpack* XNNPACK*
+        rm -rf zstd*
       ''
-      # Libraries which should not be used any more
+      # Return to root directory
       + ''
-        rm -rf third_party/NNPACK*
-        rm -rf third_party/QNNPACK*
-        rm -rf third_party/tensorpipe
-      ''
-
-      ;
+        popd
+      '';
+
+    # TODO(@connorbaker): Use DISABLE_NNPACK_AND_FAMILY for CUDA builds. Allows us to drop
+    # dependencies on pytorch_qnnpack, xnnpack, and pthreadpool.
 
     nativeBuildInputs =
       [
@@ -156,6 +141,7 @@ in
         fxdiv
         ninja
         protobuf
+        psimd
         pthreadpool
         python3
         python3Packages.numpy
@@ -171,48 +157,57 @@ in
       ++ lists.optionals useMpi [mpi]
       ++ lists.optionals useNuma [numactl]
       ++ lists.optionals useXnnpack [xnnpack]
-      ++ lists.optionals useZstd [zstd.dev]
-      ;
+      ++ lists.optionals useZstd [zstd.dev];
+
+    cmakeFlags =
+      # Core configuration options
+      [
+        "-DBUILD_SHARED_LIBS:BOOL=ON"
+        "-DCMAKE_BUILD_TYPE:STRING=Release"
+        "-DATEN_NO_TEST:BOOL=ON"
+        "-DBUILD_PYTHON:BOOL=OFF"
+        "-DUSE_PRECOMPILED_HEADERS:BOOL=ON" # TODO(@connorbaker): Any good?
+      ]
+      # Core dependencies
+      ++ [
+        "-DBUILD_CUSTOM_PROTOBUF:BOOL=OFF"
+        "-DfbgemmLibrary_DIR:PATH=${fbgemm}/share/cmake/fbgemm" # Uses share instead of lib
+        "-DUSE_FBGEMM:BOOL=ON"
+        "-DUSE_SYSTEM_FBGEMM:BOOL=ON"
+        "-DUSE_PYTORCH_QNNPACK:BOOL=ON" # In-tree and maintained, unlike QNNPACK
+        "-DUSE_SYSTEM_CPUINFO:BOOL=ON"
+        "-DUSE_SYSTEM_FP16:BOOL=ON"
+        "-DUSE_SYSTEM_FXDIV:BOOL=ON"
+        "-DUSE_SYSTEM_PSIMD:BOOL=ON"
+        "-DUSE_SYSTEM_PTHREADPOOL:BOOL=ON"
+        "-DUSE_SYSTEM_PYBIND11:BOOL=ON"
+        "-DUSE_SYSTEM_SLEEF:BOOL=ON"
+      ]
+      # Libraries which should no longer be used because they are not maintained.
+      # - NNPACK: https://github.com/Maratyszcza/NNPACK/issues/168#issuecomment-553529745
+      # - QNNPACK: https://github.com/pytorch/QNNPACK
+      # - TensorPipe: https://github.com/pytorch/tensorpipe
+      ++ [
+        "-DUSE_NNPACK:BOOL=OFF"
+        "-DUSE_QNNPACK:BOOL=OFF"
+        "-DUSE_TENSORPIPE:BOOL=OFF"
+      ]
+      # Optional features
+      ++ [
+        "-DUSE_CUDA:BOOL=${setBool useCuda}"
 
-    cmakeFlags = [
-      "-DBUILD_CUSTOM_PROTOBUF:BOOL=OFF"
-      "-DBUILD_SHARED_LIBS:BOOL=ON"
-      "-DCMAKE_BUILD_TYPE:STRING=Release"
-      "-DfbgemmLibrary_DIR:PATH=${fbgemm}/share/cmake/fbgemm" # Uses share instead of lib
-      "-DUSE_CUDA:BOOL=${setBool useCuda}"
-      "-DUSE_FBGEMM:BOOL=ON"
-      "-DUSE_MAGMA:BOOL=${setBool useMagma}"
-      "-DUSE_MKLDNN:BOOL=${setBool useMkldnn}"
-      "-DUSE_MPI:BOOL=${setBool useMpi}"
-      "-DUSE_NCCL:BOOL=${setBool useCuda}"
-      "-DUSE_NUMA:BOOL=${setBool useNuma}"
-      "-DUSE_PYTORCH_QNNPACK:BOOL=ON" # In-tree and maintained, unlike QNNPACK
-      "-DUSE_SYSTEM_CPUINFO:BOOL=ON"
-      "-DUSE_SYSTEM_FBGEMM:BOOL=ON"
-      "-DUSE_SYSTEM_NCCL:BOOL=${setBool useCuda}"
-      "-DUSE_SYSTEM_ONEDNN:BOOL=${setBool useMkldnn}"
-      "-DUSE_SYSTEM_PTHREADPOOL:BOOL=ON"
-      "-DUSE_SYSTEM_PYBIND11:BOOL=ON"
-      "-DUSE_SYSTEM_SLEEF:BOOL=ON"
-      "-DUSE_SYSTEM_FXDIV:BOOL=ON"
-      "-DUSE_SYSTEM_FP16:BOOL=ON"
-      "-DUSE_SYSTEM_XNNPACK:BOOL=${setBool useXnnpack}"
-      "-DUSE_SYSTEM_ZSTD:BOOL=${setBool useZstd}"
-      "-DUSE_XNNPACK:BOOL=${setBool useXnnpack}"
-      "-DUSE_ZSTD:BOOL=${setBool useZstd}"
+        "-DUSE_MAGMA:BOOL=${setBool useMagma}"
+        "-DUSE_MKLDNN:BOOL=${setBool useMkldnn}"
+        "-DUSE_MPI:BOOL=${setBool useMpi}"
+        "-DUSE_NCCL:BOOL=${setBool useCuda}"
+        "-DUSE_NUMA:BOOL=${setBool useNuma}"
 
-      # Misc options
-      "-DATEN_NO_TEST:BOOL=ON"
-      "-DBUILD_PYTHON:BOOL=OFF"
-      "-DUSE_PRECOMPILED_HEADERS:BOOL=ON" # TODO(@connorbaker): Any good?
-    ] 
-    # Libraries which should no longer be used because they are not maintained.
-    # - NNPACK: https://github.com/Maratyszcza/NNPACK/issues/168#issuecomment-553529745
-    # - QNNPACK: https://github.com/pytorch/QNNPACK
-    # - TensorPipe: https://github.com/pytorch/tensorpipe
-    ++ [
-      "-DUSE_NNPACK:BOOL=OFF"
-      "-DUSE_QNNPACK:BOOL=OFF"
-      "-DUSE_TENSORPIPE:BOOL=OFF"
-    ];
+        "-DUSE_SYSTEM_NCCL:BOOL=${setBool useCuda}"
+        "-DUSE_SYSTEM_ONEDNN:BOOL=${setBool useMkldnn}"
+
+        "-DUSE_SYSTEM_XNNPACK:BOOL=${setBool useXnnpack}"
+        "-DUSE_SYSTEM_ZSTD:BOOL=${setBool useZstd}"
+        "-DUSE_XNNPACK:BOOL=${setBool useXnnpack}"
+        "-DUSE_ZSTD:BOOL=${setBool useZstd}"
+      ];
   })