Fix a regression in test shells (#217)

The XPU support PR changed stdenv selection to choose the default stdenv
when one of the kernels is an XPU kernel. However, this does not work
for mixed XPU/CUDA projects, since we need to select the CUDA backend
stdenv when building for CUDA.

Change the logic to always use the stdenv for the backend that we are
building for.

Author: danieldk

.github/workflows/build_kernel.yaml

@@ -47,6 +47,9 @@ jobs:
       - name: Build relu kernel (specific Torch version)
         run: ( cd examples/relu-specific-torch && nix build . )
 
+      - name: Test that we can build a test shell (e.g. that gcc corresponds to CUDA-required)
+        run: ( cd examples/relu && nix build .#devShells.x86_64-linux.test )
+
       - name: Build silu-and-mul-universal kernel
         run: ( cd examples/silu-and-mul-universal && nix build .\#redistributable.torch27-cxx11-cu126-x86_64-linux )
       - name: Copy silu-and-mul-universal kernel

lib/build.nix

@@ -22,6 +22,21 @@ let
     isRocm
     isXpu
     ;
+  mkStdenv =
+    buildSet: oldLinuxCompat:
+    let
+      inherit (buildSet) pkgs torch;
+    in
+    if pkgs.stdenv.hostPlatform.isDarwin then
+      pkgs.stdenv
+    else if oldLinuxCompat then
+      # Uses CUDA stdenv when we are building for CUDA.
+      pkgs.stdenvGlibc_2_27
+    else if torch.cudaSupport then
+      torch.cudaPackages.backendStdenv
+    else
+      pkgs.stdenv;
+
 in
 rec {
   resolveDeps = import ./deps.nix { inherit lib; };
@@ -126,15 +141,7 @@ rec {
           _: buildConfig: builtins.length (buildConfig.cuda-capabilities or supportedCudaCapabilities)
         ) buildConfig.kernel
       );
-      stdenv =
-        if pkgs.stdenv.hostPlatform.isDarwin then
-          pkgs.stdenv
-        else if oldLinuxCompat then
-          pkgs.stdenvGlibc_2_27
-        else if lib.any (k: k.backend == "xpu") (lib.attrValues buildConfig.kernel) then
-          pkgs.stdenv
-        else
-          pkgs.cudaPackages.backendStdenv;
+      stdenv = mkStdenv { inherit pkgs torch; } oldLinuxCompat;
     in
     if buildConfig.general.universal then
       # No torch extension sources? Treat it as a noarch package.
@@ -244,13 +251,7 @@ rec {
         let
           pkgs = buildSet.pkgs;
           rocmSupport = pkgs.config.rocmSupport or false;
-          stdenv =
-            if rocmSupport then
-              pkgs.stdenv
-            else if isXpu buildSet.buildConfig then
-              pkgs.stdenv
-            else
-              pkgs.cudaPackages.backendStdenv;
+          stdenv = mkStdenv buildSet false;
           mkShell = pkgs.mkShell.override { inherit stdenv; };
         in
         {
@@ -291,13 +292,7 @@ rec {
           pkgs = buildSet.pkgs;
           rocmSupport = pkgs.config.rocmSupport or false;
           xpuSupport = pkgs.config.xpuSupport or false;
-          stdenv =
-            if rocmSupport then
-              pkgs.stdenv
-            else if xpuSupport then
-              pkgs.stdenv
-            else
-              pkgs.cudaPackages.backendStdenv;
+          stdenv = mkStdenv buildSet false;
           mkShell = pkgs.mkShell.override { inherit stdenv; };
         in
         {