nix: add cmake-xpu-parallel-hook to cap XPU build parallelism

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
   ])

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 { };

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)

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