Support noarch build variants

This change adds support for noarch build variants. So far we have used
the universal variant for kernels that do not have any AoT-compiled code.
However, the universal variant has two important issues:

1. A kernel without AoT-compiled might still be backend-specific. E.g.
  NVIDIA CuTe-based kernels are not universal in the sense that they
  don't work on non-NVIDIA GPUs.
2. We cannot specify dependencies per backend.

To solve these issues, we introduce the noarch variants to replace
universal kernels. Noarch kernels have variants of the shape
`torch-<backend>` (e.g. `torch-xpu`). This resolves the issues outlined.

This change introduces support for loading noarch kernels. In the
future, we will start emitting deprecation warnings for universal
kernels (to eventually remove support).

Author: danieldk

src/kernels/cli.py

@@ -14,7 +14,7 @@
 from .doc import generate_readme_for_kernel
 from .wheel import build_variant_to_wheel
 
-BUILD_VARIANT_REGEX = re.compile(r"^(torch\d+\d+|torch-universal)")
+BUILD_VARIANT_REGEX = re.compile(r"^(torch\d+\d+|torch-)")
 
 
 def main():

src/kernels/utils.py

@@ -84,11 +84,33 @@ def build_variant() -> str:
     return f"torch{torch_version.major}{torch_version.minor}-{cxxabi}-{compute_framework}-{cpu}-{os}"
 
 
-def universal_build_variant() -> str:
+def build_variant_noarch() -> str:
+    import torch
+
+    if torch.version.cuda is not None:
+        return "torch-cuda"
+    elif torch.version.hip is not None:
+        return "torch-rocm"
+    elif torch.backends.mps.is_available():
+        return "torch-metal"
+    elif hasattr(torch.version, "xpu") and torch.version.xpu is not None:
+        return "torch-xpu"
+    elif _get_privateuse_backend_name() == "npu":
+        return "torch-npu"
+    else:
+        return "torch-cpu"
+
+
+def build_variant_universal() -> str:
     # Once we support other frameworks, detection goes here.
     return "torch-universal"
 
 
+def build_variants() -> List[str]:
+    """Return compatible build variants in preferred order."""
+    return [build_variant(), build_variant_noarch(), build_variant_universal()]
+
+
 def _import_from_path(module_name: str, variant_path: Path) -> ModuleType:
     metadata_path = variant_path / "metadata.json"
     if metadata_path.exists():
@@ -146,13 +168,12 @@ def install_kernel(
         `Tuple[str, Path]`: A tuple containing the package name and the path to the variant directory.
     """
     package_name = package_name_from_repo_id(repo_id)
-    variant = build_variant()
-    universal_variant = universal_build_variant()
+    allow_patterns = [f"build/{variant}/*" for variant in build_variants()]
     user_agent = _get_user_agent(user_agent=user_agent)
     repo_path = Path(
         snapshot_download(
             repo_id,
-            allow_patterns=[f"build/{variant}/*", f"build/{universal_variant}/*"],
+            allow_patterns=allow_patterns,
             cache_dir=CACHE_DIR,
             revision=revision,
             local_files_only=local_files_only,
@@ -173,23 +194,22 @@ def _find_kernel_in_repo_path(
     package_name: str,
     variant_locks: Optional[Dict[str, VariantLock]] = None,
 ) -> Tuple[str, Path]:
-    specific_variant = build_variant()
-    universal_variant = universal_build_variant()
-
-    specific_variant_path = repo_path / "build" / specific_variant
-    universal_variant_path = repo_path / "build" / universal_variant
-
-    if specific_variant_path.exists():
-        variant = specific_variant
-        variant_path = specific_variant_path
-    elif universal_variant_path.exists():
-        variant = universal_variant
-        variant_path = universal_variant_path
-    else:
+    variants = build_variants()
+    variant = None
+    variant_path = None
+    for candidate_variant in variants:
+        variant_path = repo_path / "build" / candidate_variant
+        if variant_path.exists():
+            variant = candidate_variant
+            break
+
+    if variant is None:
         raise FileNotFoundError(
-            f"Kernel at path `{repo_path}` does not have one of build variants: {specific_variant}, {universal_variant}"
+            f"Kernel at path `{repo_path}` does not have one of build variants: {', '.join(variants)}"
         )
 
+    assert variant_path is not None
+
     if variant_locks is not None:
         variant_lock = variant_locks.get(variant)
         if variant_lock is None:
@@ -295,13 +315,9 @@ def get_local_kernel(repo_path: Path, package_name: str) -> ModuleType:
     Returns:
         `ModuleType`: The imported kernel module.
     """
-    variant = build_variant()
-    universal_variant = universal_build_variant()
-
     # Presume we were given the top level path of the kernel repository.
     for base_path in [repo_path, repo_path / "build"]:
-        # Prefer the universal variant if it exists.
-        for v in [universal_variant, variant]:
+        for v in build_variants():
             variant_path = base_path / v
             if variant_path.exists():
                 return _import_from_path(package_name, variant_path)
@@ -337,9 +353,8 @@ def has_kernel(
 
     package_name = package_name_from_repo_id(repo_id)
     variant = build_variant()
-    universal_variant = universal_build_variant()
 
-    for variant in [universal_variant, variant]:
+    for variant in build_variants():
         for init_file in ["__init__.py", f"{package_name}/__init__.py"]:
             if file_exists(
                 repo_id,
@@ -379,13 +394,11 @@ def load_kernel(repo_id: str, *, lockfile: Optional[Path] = None) -> ModuleType:
 
     package_name = package_name_from_repo_id(repo_id)
 
-    variant = build_variant()
-    universal_variant = universal_build_variant()
-
+    allow_patterns = [f"build/{variant}/*" for variant in build_variants()]
     repo_path = Path(
         snapshot_download(
             repo_id,
-            allow_patterns=[f"build/{variant}/*", f"build/{universal_variant}/*"],
+            allow_patterns=allow_patterns,
             cache_dir=CACHE_DIR,
             revision=locked_sha,
             local_files_only=True,
@@ -399,7 +412,7 @@ def load_kernel(repo_id: str, *, lockfile: Optional[Path] = None) -> ModuleType:
         return _import_from_path(package_name, variant_path)
     except FileNotFoundError:
         raise FileNotFoundError(
-            f"Locked kernel `{repo_id}` does not have build `{variant}` or was not downloaded with `kernels download <project>`"
+            f"Locked kernel `{repo_id}` does not have applicable variant or was not downloaded with `kernels download <project>`"
         )
 
 

tests/test_basic.py

@@ -163,6 +163,13 @@ def test_universal_kernel(universal_kernel):
     torch.testing.assert_close(out, out_check, rtol=1e-1, atol=1e-1)
 
 
+def test_noarch_kernel(device):
+    supported_devices = ["cpu", "cuda", "xpu"]
+    if device not in supported_devices:
+        pytest.skip(f"Device is not one of: {','.join(supported_devices)}")
+    get_kernel("kernels-test/silu-and-mul-noarch")
+
+
 @pytest.mark.parametrize(
     "repo_revision",
     [