kernel-builder: ensure that HIP_ARCHITECTURES gets properly cleared

examples/kernels/flake.nix

@@ -27,6 +27,12 @@
       # - drv (system -> flakeOutputs -> derivation): the derivation for the given
       #        system and flake outputs.
       # - torchVersions: optional override for the torchVersions argument
+      # - checkRocmArchs: optional list of ROCm architectures (e.g. "gfx942").
+      #        When set, the kernel dylib must contain exactly this set of
+      #        architectures.
+      # - checkCudaCapabilities: optional list of CUDA capabilities (e.g. "9.0").
+      #        When set, the kernel dylib must contain exactly this set of
+      #        capabilities.
       ciKernels = [
         {
           name = "cpp20-symbols-kernel";
@@ -38,6 +44,28 @@
           path = ./relu;
           drv =
             sys: out: out.packages.${sys}.redistributable.${"torch${torchVersion}-cxx11-${cudaVersion}-${sys}"};
+          checkCudaCapabilities = [
+            "7.0"
+            "7.2"
+            "7.5"
+            "8.0"
+            "8.6"
+            "8.7"
+            "8.9"
+            "9.0"
+          ];
+        }
+        {
+          # Check arch intersection. 5.0 is dropped because it is not supported.
+          name = "relu-archs-subset";
+          path = ./relu-archs-subset;
+          drv =
+            sys: out: out.packages.${sys}.redistributable.${"torch${torchVersion}-cxx11-${cudaVersion}-${sys}"};
+          checkCudaCapabilities = [
+            "7.5"
+            "8.0"
+            "9.0"
+          ];
         }
         {
           name = "relu-torch-stable-abi-kernel";
@@ -175,6 +203,30 @@
           path = ./relu;
           drv =
             sys: out: out.packages.${sys}.redistributable.${"torch${torchVersion}-cxx11-${rocmVersion}-${sys}"};
+          checkRocmArchs = [
+            "gfx906"
+            "gfx908"
+            "gfx90a"
+            "gfx942"
+            "gfx950"
+            "gfx1030"
+            "gfx1100"
+            "gfx1101"
+            "gfx1200"
+            "gfx1201"
+          ];
+        }
+        {
+          # Check arch intersection. gfx940 is dropped because it is not supported.
+          name = "relu-archs-subset";
+          path = ./relu-archs-subset;
+          drv =
+            sys: out: out.packages.${sys}.redistributable.${"torch${torchVersion}-cxx11-${rocmVersion}-${sys}"};
+          checkRocmArchs = [
+            "gfx90a"
+            "gfx942"
+            "gfx1100"
+          ];
         }
         {
           name = "relu-compiler-flags";
@@ -279,6 +331,59 @@
         let
           pkgs = nixpkgs.legacyPackages.${system};
 
+          # Check that every *.so in the output of `drv` contains exactly the
+          # given set of architectures. `listArchs` is a shell command that
+          # prints the architectures of the kernel library `$so`, one per
+          # line. On success, the output is a symlink to `drv`.
+          checkKernelArchs =
+            drv: expectedArchs: listArchs:
+            pkgs.runCommand "${drv.name}-check-archs"
+              {
+                expected = lib.concatLines expectedArchs;
+                passAsFile = [ "expected" ];
+              }
+              ''
+                kernelLibs=$(find ${drv}/ -name '*.so')
+                if [ -z "$kernelLibs" ]; then
+                  echo "no *.so files found in ${drv}" >&2
+                  exit 1
+                fi
+
+                sort -u "$expectedPath" > expected
+                for so in $kernelLibs; do
+                  (
+                    ${listArchs}
+                  ) | sort -u > actual
+                  if ! diff -u expected actual; then
+                    echo "unexpected architectures in $so" >&2
+                    exit 1
+                  fi
+                done
+
+                ln -s ${drv} $out
+              '';
+
+          checkRocmArchs =
+            drv: expectedArchs:
+            checkKernelArchs drv expectedArchs
+              # Use the ROCm LLVM that the kernel was built with. Stock
+              # llvm-readobj does not support `--offloading`.
+              ''
+                ${drv.torch.rocmPackages.rocm-llvm}/llvm/bin/llvm-readobj --offloading "$so" \
+                    | grep -oE 'gfx[0-9]+[a-z]?'
+              '';
+
+          checkCudaCapabilities =
+            drv: expectedCapabilities:
+            let
+              # Convert capabilities like "9.0a" to sm names like "sm_90a".
+              smArchs = map (cap: "sm_" + lib.replaceStrings [ "." ] [ "" ] cap) expectedCapabilities;
+            in
+            checkKernelArchs drv smArchs ''
+              ${drv.torch.cudaPackages.cuda_cuobjdump}/bin/cuobjdump "$so" \
+                  | grep -oE 'sm_[0-9]+[af]?'
+            '';
+
           resolveKernels =
             kernelOutputsList:
             map (kernel: {
@@ -292,6 +397,10 @@
                     drv = baseDrv;
                     expectedBuilderLogEntries = kernel.assertFailLogs or [ ];
                   }
+                else if kernel ? checkRocmArchs then
+                  checkRocmArchs baseDrv kernel.checkRocmArchs
+                else if kernel ? checkCudaCapabilities then
+                  checkCudaCapabilities baseDrv kernel.checkCudaCapabilities
                 else
                   baseDrv;
             }) kernelOutputsList;

examples/kernels/relu-archs-subset/CARD.md

@@ -0,0 +1,61 @@
+---
+library_name: kernels
+{% if license %}license: {{ license }}
+{% endif %}---
+
+This is the repository card of {{ repo_id }} that has been pushed on the Hub. It was built to be used with the [`kernels` library](https://github.com/huggingface/kernels). This card was automatically generated.
+
+## How to use
+{% if functions %}
+
+```python
+# make sure `kernels` is installed: `pip install -U kernels`
+from kernels import get_kernel
+
+kernel_module = get_kernel("{{ repo_id }}", version={{ version }})
+{{ functions[0] }} = kernel_module.{{ functions[0] }}
+
+{{ functions[0] }}(...)
+```
+{% else %}
+
+Usage example not available.
+{% endif %}
+
+## Available functions
+{% if functions %}
+{% for func in functions %}
+- `{{ func }}`
+{% endfor %}
+{% else %}
+
+Function list not available.
+{% endif %}
+{% if layers %}
+
+## Available layers
+{% for layer in layers %}
+- `{{ layer }}`
+{% endfor %}
+{% endif %}
+
+## Benchmarks
+{% if has_benchmark %}
+
+Benchmarking script is available for this kernel. Run `kernels benchmark {{ repo_id }} --version {{ version }}`.
+{% else %}
+
+No benchmark available yet.
+{% endif %}
+{% if upstream %}
+
+## Upstream
+
+The original source code for this kernel comes from {{ upstream }}.
+{% endif %}
+{% if source %}
+
+## Source
+
+The kernel-builder formatted source for this kernel is available at {{ source }}.
+{% endif %}

examples/kernels/relu-archs-subset/build.toml

@@ -0,0 +1,45 @@
+[general]
+name = "relu-archs-subset"
+version = 1
+license = "Apache-2.0"
+backends = [
+    "cuda",
+    "rocm",
+]
+
+[general.hub]
+repo-id = "kernels-test/relu-archs-subset"
+
+[torch]
+src = [
+    "torch-ext/torch_binding.cpp",
+    "torch-ext/torch_binding.h",
+]
+
+# The capability/architecture sets below intentionally contain entries that
+# are not supported by the toolchain versions used in CI (e.g. CUDA 10.0/12.0
+# capabilities require CUDA >= 12.8, gfx940 is not in the supported arch set).
+# CI checks that the built kernels contain exactly the intersection with the
+# supported sets.
+
+[kernel.relu_rocm]
+backend = "rocm"
+depends = ["torch"]
+rocm-archs = [
+    "gfx90a",
+    "gfx940",
+    "gfx942",
+    "gfx1100",
+]
+src = ["relu_cuda/relu.cu"]
+
+[kernel.relu]
+backend = "cuda"
+depends = ["torch"]
+cuda-capabilities = [
+    "5.0",
+    "7.5",
+    "8.0",
+    "9.0",
+]
+src = ["relu_cuda/relu.cu"]

examples/kernels/relu-archs-subset/flake.nix

@@ -0,0 +1,17 @@
+{
+  description = "Flake for ReLU kernel with a subset of CUDA/ROCm archs";
+
+  inputs = {
+    kernel-builder.url = "path:../../..";
+  };
+
+  outputs =
+    {
+      self,
+      kernel-builder,
+    }:
+    kernel-builder.lib.genKernelFlakeOutputs {
+      inherit self;
+      path = ./.;
+    };
+}

examples/kernels/relu-archs-subset/relu_cuda/relu.cu

@@ -0,0 +1,47 @@
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/all.h>
+
+#include <cmath>
+
+__global__ void relu_kernel(float *__restrict__ out,
+                            float const *__restrict__ input, const int d) {
+  const int64_t token_idx = blockIdx.x;
+  for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
+    auto x = input[token_idx * d + idx];
+    out[token_idx * d + idx] = x > 0.0f ? x : 0.0f;
+  }
+}
+
+void relu(torch::Tensor &out, torch::Tensor const &input) {
+  TORCH_CHECK(input.device().is_cuda(), "input must be a CUDA tensor");
+  TORCH_CHECK(input.is_contiguous(), "input must be contiguous");
+  TORCH_CHECK(input.scalar_type() == at::ScalarType::Float &&
+                  input.scalar_type() == at::ScalarType::Float,
+              "relu_kernel only supports float32");
+
+  TORCH_CHECK(input.sizes() == out.sizes(),
+              "Tensors must have the same shape. Got input shape: ",
+              input.sizes(), " and output shape: ", out.sizes());
+
+  TORCH_CHECK(input.scalar_type() == out.scalar_type(),
+              "Tensors must have the same data type. Got input dtype: ",
+              input.scalar_type(), " and output dtype: ", out.scalar_type());
+
+  TORCH_CHECK(input.device() == out.device(),
+              "Tensors must be on the same device. Got input device: ",
+              input.device(), " and output device: ", out.device());
+
+  if (input.numel() == 0) {
+    return;
+  }
+
+  int d = input.size(-1);
+  int64_t num_tokens = input.numel() / d;
+  dim3 grid(num_tokens);
+  dim3 block(std::min(d, 1024));
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  relu_kernel<<<grid, block, 0, stream>>>(out.data_ptr<float>(),
+                                          input.data_ptr<float>(), d);
+}

examples/kernels/relu-archs-subset/torch-ext/relu_archs_subset/__init__.py

@@ -0,0 +1,15 @@
+from typing import Optional
+
+import torch
+
+from ._ops import ops
+
+
+def relu(x: torch.Tensor, out: Optional[torch.Tensor] = None) -> torch.Tensor:
+    if out is None:
+        out = torch.empty_like(x)
+    ops.relu(out, x)
+    return out
+
+
+__all__ = ["relu"]

examples/kernels/relu-archs-subset/torch-ext/torch_binding.cpp

@@ -0,0 +1,19 @@
+#include <torch/library.h>
+
+#include "registration.h"
+#include "torch_binding.h"
+
+TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
+  ops.def("relu(Tensor! out, Tensor input) -> ()");
+#if defined(CPU_KERNEL)
+  ops.impl("relu", torch::kCPU, &relu);
+#elif defined(CUDA_KERNEL) || defined(ROCM_KERNEL)
+  ops.impl("relu", torch::kCUDA, &relu);
+#elif defined(METAL_KERNEL)
+  ops.impl("relu", torch::kMPS, relu);
+#elif defined(XPU_KERNEL)
+  ops.impl("relu", torch::kXPU, &relu);
+#endif
+}
+
+REGISTER_EXTENSION(TORCH_EXTENSION_NAME)

examples/kernels/relu-archs-subset/torch-ext/torch_binding.h

@@ -0,0 +1,5 @@
+#pragma once
+
+#include <torch/torch.h>
+
+void relu(torch::Tensor &out, torch::Tensor const &input);

examples/kernels/relu/build.toml

@@ -36,18 +36,6 @@ src = [
 [kernel.relu_rocm]
 backend = "rocm"
 depends = ["torch"]
-rocm-archs = [
-    "gfx906",
-    "gfx908",
-    "gfx90a",
-    "gfx942",
-    "gfx950",
-    "gfx1030",
-    "gfx1100",
-    "gfx1101",
-    "gfx1200",
-    "gfx1201",
-]
 src = ["relu_cuda/relu.cu"]
 
 [kernel.relu_cpu]

kernel-builder/src/pyproject/templates/torch/preamble.cmake

@@ -198,7 +198,6 @@ elseif(GPU_LANG STREQUAL "SYCL")
   set(sycl_link_flags "-Wl,-z,noexecstack;-fsycl;--offload-compress;-fsycl-targets=spir64_gen,spir64;-Xs;-device pvc,xe-lpg,ats-m150 -options ' -cl-intel-enable-auto-large-GRF-mode -cl-poison-unsupported-fp64-kernels -cl-intel-greater-than-4GB-buffer-required';")
   set(sycl_flags "-fPIC;-fsycl;-fhonor-nans;-fhonor-infinities;-fno-associative-math;-fno-approx-func;-fno-sycl-instrument-device-code;--offload-compress;-fsycl-targets=spir64_gen,spir64;")
   set(GPU_FLAGS "${sycl_flags}")
-  set(GPU_ARCHES "")
 
 
   add_compile_definitions(XPU_KERNEL)

kernel-builder/src/pyproject/templates/torch/torch-extension.cmake

@@ -11,7 +11,6 @@ define_gpu_extension_target(
   LANGUAGE ${GPU_LANG}
   SOURCES ${SRC}
   COMPILE_FLAGS ${GPU_FLAGS}
-  ARCHITECTURES ${GPU_ARCHES}
   USE_SABI 3
   WITH_SOABI)