diff --git a/nix-builder/lib/extension/torch/arch.nix b/nix-builder/lib/extension/torch/arch.nix index f7c64fac..9535ad9a 100644 --- a/nix-builder/lib/extension/torch/arch.nix +++ b/nix-builder/lib/extension/torch/arch.nix @@ -12,6 +12,7 @@ kernel-builder, cmake, cmakeNvccThreadsHook, + cmakeXpuParallelHook, cuda_nvcc, get-kernel-check, kernel-abi-check, @@ -54,6 +55,9 @@ nvccThreads, + # TODO: Maximum parallel XPU compilation jobs. (for FA2) + xpuParallelJobs ? 12, + # A stringly-typed list of Python dependencies. Ideally we'd take a # list of derivations, but we also need to write the dependencies to # the output. @@ -109,6 +113,7 @@ stdenv.mkDerivation (prevAttrs: { doKernelBuildCheck moduleName nvccThreads + xpuParallelJobs ; framework = "torch"; @@ -167,6 +172,7 @@ stdenv.mkDerivation (prevAttrs: { clr ] ++ lib.optionals xpuSupport ([ + cmakeXpuParallelHook xpuPackages.ocloc oneapi-torch-dev ]) diff --git a/nix-builder/overlay.nix b/nix-builder/overlay.nix index a1a69755..746b9c90 100644 --- a/nix-builder/overlay.nix +++ b/nix-builder/overlay.nix @@ -12,6 +12,8 @@ in cmakeNvccThreadsHook = final.callPackage ./pkgs/cmake-nvcc-threads-hook { }; + cmakeXpuParallelHook = final.callPackage ./pkgs/cmake-xpu-parallel-hook { }; + get-kernel-check = final.callPackage ./pkgs/get-kernel-check { }; kernel-abi-check = final.callPackage ./pkgs/kernel-abi-check { }; diff --git a/nix-builder/pkgs/cmake-xpu-parallel-hook/cmake-xpu-parallel-hook.sh b/nix-builder/pkgs/cmake-xpu-parallel-hook/cmake-xpu-parallel-hook.sh new file mode 100644 index 00000000..1c726f69 --- /dev/null +++ b/nix-builder/pkgs/cmake-xpu-parallel-hook/cmake-xpu-parallel-hook.sh @@ -0,0 +1,18 @@ +#!/bin/sh + +_setXpuParallelHook() { + if [ -z "${xpuParallelJobs}" ] || [ "${xpuParallelJobs}" -ne "${xpuParallelJobs}" ] 2>/dev/null; then + >&2 echo "Number of XPU parallel jobs is not (correctly) set, setting to 12" + xpuParallelJobs=12 + fi + + # Cap parallel jobs to available cores. + xpuParallelJobs=$((NIX_BUILD_CORES < xpuParallelJobs ? NIX_BUILD_CORES : xpuParallelJobs)) + + # Reduce NIX_BUILD_CORES to limit ninja parallelism for XPU builds, + export NIX_BUILD_CORES="${xpuParallelJobs}" + + >&2 echo "XPU parallel hook: limiting ninja to -j${NIX_BUILD_CORES}" +} + +preConfigureHooks+=(_setXpuParallelHook) diff --git a/nix-builder/pkgs/cmake-xpu-parallel-hook/default.nix b/nix-builder/pkgs/cmake-xpu-parallel-hook/default.nix new file mode 100644 index 00000000..73211f0e --- /dev/null +++ b/nix-builder/pkgs/cmake-xpu-parallel-hook/default.nix @@ -0,0 +1,5 @@ +{ makeSetupHook }: + +makeSetupHook { + name = "cmake-xpu-parallel-hook"; +} ./cmake-xpu-parallel-hook.sh