add xpu support in nix

Signed-off-by: Wang, Yi A <yi.a.wang@intel.com>

Author: sywangyi

build2cmake/src/templates/xpu/preamble.cmake

@@ -33,7 +33,7 @@ endif()
 
 # Check for Intel XPU support in PyTorch
 run_python(XPU_AVAILABLE
-  "import torch; print('true' if hasattr(torch, 'xpu') and torch.xpu.is_available() else 'false')"
+  "import torch; print('true' if hasattr(torch, 'xpu') else 'false')"
   "Failed to check XPU availability")
 
 if(NOT XPU_AVAILABLE STREQUAL "true")

lib/build-sets.nix

@@ -15,6 +15,7 @@ let
     isCuda
     isMetal
     isRocm
+    isXpu
     ;
 
   # All build configurations supported by Torch.
@@ -39,6 +40,12 @@ let
     in
     lib.unique (builtins.map (torchVersion: torchVersion.rocmVersion) withRocm);
 
+  xpuVersions =
+    let
+      withXpu = builtins.filter (torchVersion: torchVersion ? xpuVersion) torchVersions;
+    in
+    lib.unique (builtins.map (torchVersion: torchVersion.xpuVersion) withXpu);
+
   flattenVersion = version: lib.replaceStrings [ "." ] [ "_" ] (lib.versions.pad 2 version);
 
   # An overlay that overides CUDA to the given version.
@@ -50,12 +57,16 @@ let
     rocmPackages = super."rocmPackages_${flattenVersion rocmVersion}";
   };
 
+  overlayForXpuVersion = xpuVersion: self: super: {
+    xpuPackages = super."xpuPackages_${flattenVersion xpuVersion}";
+  };
   # Construct the nixpkgs package set for the given versions.
   pkgsForVersions =
     buildConfig@{
       cudaVersion ? null,
       metal ? false,
       rocmVersion ? null,
+      xpuVersion ? null,
       torchVersion,
       cxx11Abi,
       system,
@@ -69,11 +80,18 @@ let
           pkgsByRocmVer.${rocmVersion}
         else if isMetal buildConfig then
           pkgsForMetal
+        else if isXpu buildConfig then
+          pkgsByXpuVer.${xpuVersion}
         else
           throw "No compute framework set in Torch version";
-      torch = pkgs.python3.pkgs."torch_${flattenVersion torchVersion}".override {
-        inherit cxx11Abi;
-      };
+      torch = pkgs.python3.pkgs."torch_${flattenVersion torchVersion}".overrideAttrs (old: {
+        passthru = (old.passthru or { }) // {
+          xpuPackages = pkgs.xpuPackages or { };
+        };
+        cxx11Abi = cxx11Abi;
+        xpuSupport =
+          (pkgs ? xpuPackages) && ((builtins.length (builtins.attrNames pkgs.xpuPackages or { })) > 0);
+      });
     in
     {
       inherit
@@ -83,6 +101,26 @@ let
         bundleBuild
         ;
     };
+  pkgsForXpuVersions =
+    xpuVersions:
+    builtins.listToAttrs (
+      map (xpuVersion: {
+        name = xpuVersion;
+        value = import nixpkgs {
+          inherit system;
+          config = {
+            allowUnfree = true;
+            xpuSupport = true;
+          };
+          overlays = [
+            hf-nix
+            overlay
+            (overlayForXpuVersion xpuVersion)
+          ];
+        };
+      }) xpuVersions
+    );
+  pkgsByXpuVer = pkgsForXpuVersions xpuVersions;
 
   pkgsForMetal = import nixpkgs {
     inherit system;

lib/build-variants.nix

@@ -5,6 +5,7 @@ let
     isCuda
     isMetal
     isRocm
+    isXpu
     ;
 in
 rec {
@@ -16,8 +17,10 @@ rec {
       "metal"
     else if buildConfig ? "rocmVersion" then
       "rocm"
+    else if buildConfig ? xpuVersion then
+      "xpu"
     else
-      throw "Could not find compute framework: no CUDA or ROCm version specified and Metal is not enabled";
+      throw "Could not find compute framework: no CUDA, ROCm, XPU version specified and Metal is not enabled";
 
   # Build variants included in bundle builds.
   buildVariants =
@@ -31,6 +34,8 @@ rec {
           "rocm${flattenVersion (lib.versions.majorMinor version.rocmVersion)}"
         else if isMetal version then
           "metal"
+        else if isXpu version then
+          "xpu${flattenVersion (lib.versions.majorMinor version.xpuVersion)}"
         else
           throw "No compute framework set in Torch version";
       buildName =

lib/build-version.nix

@@ -12,10 +12,24 @@ let
   cudaVersion = torch: "cu${flattenVersion torch.cudaPackages.cudaMajorMinorVersion}";
   rocmVersion =
     torch: "rocm${flattenVersion (lib.versions.majorMinor torch.rocmPackages.rocm.version)}";
-  gpuVersion = torch: (if torch.cudaSupport then cudaVersion else rocmVersion) torch;
   torchVersion = torch: flattenVersion torch.version;
+  xpuVersion =
+    torch:
+    "xpu${flattenVersion (lib.versions.majorMinor torch.passthru.xpuPackages.intel-oneapi-dpcpp-cpp.version)}";
+  gpuVersion =
+    torch:
+    if torch.cudaSupport then
+      cudaVersion torch
+    else if (torch ? rocmPackages) && (torch.rocmSupport or false) then
+      rocmVersion torch
+    else if (torch ? passthru) && (torch.passthru ? xpuPackages) && (torch.xpuSupport or false) then
+      xpuVersion torch
+    else
+      null;
 in
 if pkgs.stdenv.hostPlatform.isDarwin then
   "torch${torchVersion torch}-metal-${targetPlatform}"
-else
+else if gpuVersion torch != null then
   "torch${torchVersion torch}-${abi torch}-${gpuVersion torch}-${targetPlatform}"
+else
+  throw "No supported GPU framework (CUDA, ROCm, XPU, Metal) detected for build-version.nix"

lib/build.nix

@@ -16,7 +16,12 @@ let
   supportedCudaCapabilities = builtins.fromJSON (
     builtins.readFile ../build2cmake/src/cuda_supported_archs.json
   );
-  inherit (import ./torch-version-utils.nix { inherit lib; }) isCuda isMetal isRocm;
+  inherit (import ./torch-version-utils.nix { inherit lib; })
+    isCuda
+    isMetal
+    isRocm
+    isXpu
+    ;
 in
 rec {
   resolveDeps = import ./deps.nix { inherit lib; };
@@ -45,11 +50,13 @@ rec {
         cuda = false;
         metal = false;
         rocm = false;
+        xpu = false;
       };
     in
     lib.foldl (backends: kernel: backends // { ${kernelBackend kernel} = true; }) init kernels;
 
   readBuildConfig = path: validateBuildConfig (readToml (path + "/build.toml"));
+  tracedReadBuildConfig = path: readBuildConfig path;
 
   srcFilter =
     src: name: type:
@@ -75,6 +82,7 @@ rec {
             (isCuda buildSet.buildConfig && backends'.cuda)
             || (isRocm buildSet.buildConfig && backends'.rocm)
             || (isMetal buildSet.buildConfig && backends'.metal)
+            || (isXpu buildSet.buildConfig && backends'.xpu)
             || (buildConfig.general.universal or false);
           cudaVersionSupported =
             !(isCuda buildSet.buildConfig)
@@ -121,6 +129,8 @@ rec {
       stdenv =
         if pkgs.stdenv.hostPlatform.isDarwin then
           pkgs.stdenv
+        else if lib.any (k: k.backend == "xpu") (lib.attrValues buildConfig.kernel) then
+          pkgs.stdenv
         else if oldLinuxCompat then
           pkgs.stdenvGlibc_2_27
         else
@@ -234,7 +244,13 @@ rec {
         let
           pkgs = buildSet.pkgs;
           rocmSupport = pkgs.config.rocmSupport or false;
-          stdenv = if rocmSupport then pkgs.stdenv else pkgs.cudaPackages.backendStdenv;
+          stdenv =
+            if rocmSupport then
+              pkgs.stdenv
+            else if isXpu buildSet.buildConfig then
+              pkgs.stdenv
+            else
+              pkgs.cudaPackages.backendStdenv;
           mkShell = pkgs.mkShell.override { inherit stdenv; };
         in
         {
@@ -274,7 +290,14 @@ rec {
         let
           pkgs = buildSet.pkgs;
           rocmSupport = pkgs.config.rocmSupport or false;
-          stdenv = if rocmSupport then pkgs.stdenv else pkgs.cudaPackages.backendStdenv;
+          xpuSupport = pkgs.config.xpuSupport or false;
+          stdenv =
+            if rocmSupport then
+              pkgs.stdenv
+            else if xpuSupport then
+              pkgs.stdenv
+            else
+              pkgs.cudaPackages.backendStdenv;
           mkShell = pkgs.mkShell.override { inherit stdenv; };
         in
         {

lib/torch-extension/default.nix

@@ -14,6 +14,7 @@
   config,
   cudaSupport ? config.cudaSupport,
   rocmSupport ? config.rocmSupport,
+  xpuSupport ? config.xpuSupport,
 
   lib,
   stdenv,
@@ -70,6 +71,8 @@ stdenv.mkDerivation (prevAttrs: {
         "cuda"
       else if rocmSupport then
         "rocm"
+      else if xpuSupport then
+        "xpu"
       else
         "metal"
     } --ops-id ${rev} build.toml
@@ -98,6 +101,13 @@ stdenv.mkDerivation (prevAttrs: {
   ++ lib.optionals rocmSupport [
     clr
   ]
+  ++ lib.optionals xpuSupport (
+    with torch.passthru.xpuPackages;
+    [
+      ocloc
+      oneapi-torch-dev
+    ]
+  )
   ++ lib.optionals stdenv.hostPlatform.isDarwin [
     rewrite-nix-paths-macho
   ];
@@ -120,6 +130,12 @@ stdenv.mkDerivation (prevAttrs: {
     ]
   )
   ++ lib.optionals rocmSupport (with rocmPackages; [ hipsparselt ])
+  ++ lib.optionals xpuSupport (
+    with torch.passthru.xpuPackages;
+    [
+      oneapi-torch-dev
+    ]
+  )
   ++ lib.optionals stdenv.hostPlatform.isDarwin [
     apple-sdk_15
   ]
@@ -136,6 +152,10 @@ stdenv.mkDerivation (prevAttrs: {
     }
     // lib.optionalAttrs rocmSupport {
       PYTORCH_ROCM_ARCH = lib.concatStringsSep ";" torch.rocmArchs;
+    }
+    // lib.optionalAttrs xpuSupport {
+      MKLROOT = torch.passthru.xpuPackages.oneapi-torch-dev;
+      SYCL_ROOT = torch.passthru.xpuPackages.oneapi-torch-dev;
     };
 
   # If we use the default setup, CMAKE_CUDA_HOST_COMPILER gets set to nixpkgs g++.

lib/torch-version-utils.nix

@@ -10,4 +10,5 @@
   isCuda = version: version ? cudaVersion;
   isMetal = version: version.metal or false;
   isRocm = version: version ? rocmVersion;
+  isXpu = version: version ? xpuVersion;
 }

versions.nix

@@ -40,6 +40,20 @@
     bundleBuild = false;
   }
 
+  {
+    torchVersion = "2.7";
+    xpuVersion = "2025.0.2";
+    cxx11Abi = true;
+    systems = [ "x86_64-linux" ];
+    bundleBuild = true;
+  }
+  {
+    torchVersion = "2.8";
+    xpuVersion = "2025.1.3";
+    cxx11Abi = true;
+    systems = [ "x86_64-linux" ];
+    bundleBuild = true;
+  }
   {
     torchVersion = "2.7";
     cxx11Abi = true;