Cache extra ROCm/XPU build dependencies

Author: danieldk

flake.nix

@@ -146,6 +146,7 @@
                   with buildSet.pkgs;
                   (
                     allOutputs buildSet.torch
+                    ++ lib.concatMap allOutputs buildSet.extension.extraBuildDeps
                     ++ allOutputs build2cmake
                     ++ allOutputs kernel-abi-check
                     ++ allOutputs python3Packages.kernels

lib/torch-extension/default.nix

@@ -1,4 +1,8 @@
 {
+  rocmSupport ? torch.rocmSupport,
+  xpuSupport ? torch.xpuSupport,
+
+  lib,
   callPackage,
   stdenv,
   stdenvGlibc_2_27,
@@ -31,6 +35,13 @@ let
   };
 in
 {
+  extraBuildDeps =
+    lib.optionals xpuSupport [
+      oneapi-torch-dev
+      onednn-xpu
+    ]
+    ++ lib.optionals rocmSupport [ clr ];
+
   mkExtension = callPackage ./arch.nix {
     inherit
       clr