Merge pull request #335560 from danieldk/xformers-parallel

python3Packages.xformers: enable parallel build

pkgs/development/python-modules/xformers/default.nix

@@ -18,6 +18,7 @@
   fairscale,
   scipy,
   cmake,
+  ninja,
   triton,
   networkx,
   #, apex
@@ -47,13 +48,14 @@ buildPythonPackage {
 
   patches = [ ./0001-fix-allow-building-without-git.patch ];
 
-  preBuild =
-    ''
-      cat << EOF > ./xformers/version.py
-      # noqa: C801
-      __version__ = "${version}"
-      EOF
-    '';
+  preBuild = ''
+    cat << EOF > ./xformers/version.py
+    # noqa: C801
+    __version__ = "${version}"
+    EOF
+
+    export MAX_JOBS=$NIX_BUILD_CORES
+  '';
 
   env = lib.attrsets.optionalAttrs cudaSupport {
     TORCH_CUDA_ARCH_LIST = "${lib.concatStringsSep ";" torch.cudaCapabilities}";
@@ -74,9 +76,10 @@ buildPythonPackage {
     ]
   );
 
-  nativeBuildInputs = [ which ] ++ lib.optionals cudaSupport (with cudaPackages; [
-    cuda_nvcc
-  ]);
+  nativeBuildInputs = [
+    ninja
+    which
+  ] ++ lib.optionals cudaSupport (with cudaPackages; [ cuda_nvcc ]);
 
   propagatedBuildInputs = [
     numpy