From 5085b2b443ea15050db626811eff5601bcb68f53 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= <me@danieldk.eu>
Date: Fri, 29 May 2026 06:51:27 +0000
Subject: [PATCH] CI: do concurrent builds using Nix for XPU and Metal

This was already done previously for CUDA/ROCm.
---
 .github/workflows/build_kernel_macos.yaml |  7 +--
 .github/workflows/build_kernel_xpu.yaml   | 13 +-----
 examples/kernels/flake.nix                | 56 ++++++++++++++++++++++-
 3 files changed, 59 insertions(+), 17 deletions(-)

diff --git a/.github/workflows/build_kernel_macos.yaml b/.github/workflows/build_kernel_macos.yaml
index 905a05bd..03fd552a 100644
--- a/.github/workflows/build_kernel_macos.yaml
+++ b/.github/workflows/build_kernel_macos.yaml
@@ -27,8 +27,5 @@ jobs:
           #authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
       # For now we only test that there are no regressions in building macOS
       # kernels. Also run tests once we have a macOS runner.
-      - name: Build relu kernel
-        run: ( cd examples/kernels/relu && nix build .\#redistributable.torch211-metal-aarch64-darwin -L )
-
-      - name: Build relu metal cpp kernel
-        run: ( cd examples/kernels/relu-metal-cpp && nix build .\#redistributable.torch211-metal-aarch64-darwin -L )
+      - name: Build all Metal example kernels
+        run: nix build -L ./examples/kernels#ci-build-metal
diff --git a/.github/workflows/build_kernel_xpu.yaml b/.github/workflows/build_kernel_xpu.yaml
index 507fa7bc..7a284846 100644
--- a/.github/workflows/build_kernel_xpu.yaml
+++ b/.github/workflows/build_kernel_xpu.yaml
@@ -33,14 +33,5 @@ jobs:
         run: nix-shell -p nix-info --run "nix-info -m"
       # For now we only test that there are no regressions in building XPU
       # kernels. Also run tests once we have a XPU runner.
-      - name: Build relu kernel
-        run: ( cd examples/kernels/relu && nix build .\#redistributable.torch211-cxx11-xpu20253-x86_64-linux -L )
-
-      - name: Build relu tvm-ffi kernel
-        run: ( cd examples/kernels/relu-tvm-ffi && nix build .\#redistributable.tvm-ffi01-xpu20253-x86_64-linux -L )
-
-      - name: Build relu kernel (compiler flags)
-        run: ( cd examples/kernels/relu-compiler-flags && nix build .\#redistributable.torch211-cxx11-xpu20253-x86_64-linux )
-
-      - name: Build cutlass-gemm kernel
-        run: ( cd examples/kernels/cutlass-gemm && nix build .\#redistributable.torch211-cxx11-xpu20253-x86_64-linux -L )
+      - name: Build all XPU example kernels
+        run: nix build -L ./examples/kernels#ci-build-xpu
diff --git a/examples/kernels/flake.nix b/examples/kernels/flake.nix
index e1a351a6..82def60a 100644
--- a/examples/kernels/flake.nix
+++ b/examples/kernels/flake.nix
@@ -16,6 +16,7 @@
 
       cudaVersion = "cu126";
       rocmVersion = "rocm71";
+      xpuVersion = "xpu20253";
       torchVersion = "211";
       tvmFfiVersion = "01";
 
@@ -194,11 +195,57 @@
 
       ciKernelOutputs = mkKernelOutputs' ciKernels;
       ciRocmKernelOutputs = mkKernelOutputs' ciRocmKernels;
+
+      # XPU kernels to build in CI.
+      ciXpuKernels = [
+        {
+          name = "relu-kernel";
+          path = ./relu;
+          drv =
+            sys: out: out.packages.${sys}.redistributable.${"torch${torchVersion}-cxx11-${xpuVersion}-${sys}"};
+        }
+        {
+          name = "relu-tvm-ffi-kernel";
+          path = ./relu-tvm-ffi;
+          drv =
+            sys: out: out.packages.${sys}.redistributable.${"tvm-ffi${tvmFfiVersion}-${xpuVersion}-${sys}"};
+        }
+        {
+          name = "relu-compiler-flags";
+          path = ./relu-compiler-flags;
+          drv =
+            sys: out: out.packages.${sys}.redistributable.${"torch${torchVersion}-cxx11-${xpuVersion}-${sys}"};
+        }
+        {
+          name = "cutlass-gemm-kernel";
+          path = ./cutlass-gemm;
+          drv =
+            sys: out: out.packages.${sys}.redistributable.${"torch${torchVersion}-cxx11-${xpuVersion}-${sys}"};
+        }
+      ];
+
+      # Metal kernels to build in CI.
+      ciMetalKernels = [
+        {
+          name = "relu-kernel";
+          path = ./relu;
+          drv = sys: out: out.packages.${sys}.redistributable.${"torch${torchVersion}-metal-${sys}"};
+        }
+        {
+          name = "relu-metal-cpp-kernel";
+          path = ./relu-metal-cpp;
+          drv = sys: out: out.packages.${sys}.redistributable.${"torch${torchVersion}-metal-${sys}"};
+        }
+      ];
+
+      ciXpuKernelOutputs = mkKernelOutputs' ciXpuKernels;
+      ciMetalKernelOutputs = mkKernelOutputs' ciMetalKernels;
     in
     flake-utils.lib.eachSystem
       [
         "x86_64-linux"
         "aarch64-linux"
+        "aarch64-darwin"
       ]
       (
         system:
@@ -233,10 +280,17 @@
 
           ci-build-cuda = mkCiBuild "ci-kernels-cuda" ciKernelOutputs;
           ci-build-rocm = mkCiBuild "ci-kernels-rocm" ciRocmKernelOutputs;
+          ci-build-xpu = mkCiBuild "ci-kernels-xpu" ciXpuKernelOutputs;
+          ci-build-metal = mkCiBuild "ci-kernels-metal" ciMetalKernelOutputs;
         in
         {
           packages = {
-            inherit ci-build-cuda ci-build-rocm;
+            inherit
+              ci-build-cuda
+              ci-build-rocm
+              ci-build-xpu
+              ci-build-metal
+              ;
             default = ci-build-cuda;
           };
         }