add tvm-ffi

Author: ved1beta

kernels/src/kernels/variants.py

@@ -23,7 +23,12 @@
 )
 from kernels.compat import has_torch, has_tvm_ffi
 
-BUILD_VARIANT_REGEX = re.compile(r"^(torch\d+\d+|torch-(cpu|cuda|metal|neuron|rocm|xpu)|tvm-ffi\d+\d+)")
+BUILD_VARIANT_REGEX = re.compile(
+    r"^(torch\d+\d+"
+    r"|torch-(cpu|cuda|metal|neuron|rocm|xpu|metal)"
+    r"|tvm-ffi\d+\d+"
+    r"|tvm-ffi-(cpu|cuda|metal|neuron|rocm|xpu))"
+)
 
 
 @dataclass(unsafe_hash=True)
@@ -134,6 +139,23 @@ def variant_str(self) -> str:
         return "torch"
 
 
+@strict
+@dataclass(unsafe_hash=True)
+class TvmFfiNoarch:
+    """Versionless tvm-ffi framework (noarch variants)."""
+
+    @staticmethod
+    def possible_variants() -> list["TvmFfiNoarch"]:
+        if has_tvm_ffi:
+            return [TvmFfiNoarch()]
+        else:
+            return []
+
+    @property
+    def variant_str(self) -> str:
+        return "tvm-ffi"
+
+
 @strict
 @dataclass(unsafe_hash=True)
 class Arch:
@@ -220,12 +242,14 @@ def variant_str(self) -> str:
 class NoarchVariant:
     """Noarch kernel build variant."""
 
-    framework: TorchNoarch
+    framework: TorchNoarch | TvmFfiNoarch
     arch: Noarch
 
     @staticmethod
     def possible_variants() -> list["NoarchVariant"]:
-        frameworks = TorchNoarch.possible_variants()
+        frameworks: list[TorchNoarch | TvmFfiNoarch] = (
+            TorchNoarch.possible_variants() + TvmFfiNoarch.possible_variants()
+        )
         archs = Noarch.possible_variants()
         return [NoarchVariant(framework=fw, arch=arch) for fw, arch in itertools.product(frameworks, archs)]
 
@@ -264,6 +288,9 @@ def parse_variant(variant_str: str) -> Variant:
             framework_str = parts[0]
             arch_parts = parts[1:]
         return ArchVariant(framework=Torch.parse(framework_str), arch=Arch.parse(arch_parts))
+    elif parts[0] == "tvm" and len(parts) >= 2 and parts[1] == "ffi":
+        # noarch: e.g. "tvm-ffi-metal"
+        return NoarchVariant(framework=TvmFfiNoarch(), arch=Noarch.parse("-".join(parts[2:])))
     elif parts[0] == "tvm" and len(parts) >= 2 and parts[1].startswith("ffi"):
         return ArchVariant(framework=TvmFfi.parse(f"tvm-{parts[1]}"), arch=Arch.parse(parts[2:]))
     else:
@@ -432,7 +459,8 @@ def _sort_variants(
     1. Torch arch kernels with with the highest compatible CUDA version.
     2. tvm-ffi arch kernels with with the highest compatible CUDA version.
     3. Torch noarch kernels.
-    4. Old Torch universal kernels.
+    4. tvm-ffi noarch kernels.
+    5. Old universal kernels.
     """
 
     def sort_key(v: Variant) -> tuple:
@@ -447,6 +475,7 @@ def sort_key(v: Variant) -> tuple:
         else:
             assert isinstance(v, NoarchVariant)
             universal_order = 1 if v.arch.backend_name == "universal" else 0
-            return (2, universal_order)
+            framework_order = 0 if isinstance(v.framework, TorchNoarch) else 1
+            return (2, universal_order, framework_order)
 
     return sorted(variants, key=sort_key)

kernels/tests/test_variants.py

@@ -58,6 +58,12 @@
     "torch-xpu",
     "torch-npu",
     "torch-universal",
+    "tvm-ffi-cpu",
+    "tvm-ffi-cuda",
+    "tvm-ffi-metal",
+    "tvm-ffi-rocm",
+    "tvm-ffi-xpu",
+    "tvm-ffi-universal",
 ]
 
 SUPERSET_VARIANT_STRINGS = [
@@ -259,6 +265,73 @@ def test_resolve_noarch_fallback():
     assert result[0].variant_str == "torch-cuda"
 
 
+RESOLVE_VARIANTS_TVM_FFI = [
+    parse_variant(s)
+    for s in [
+        "tvm-ffi01-metal-aarch64-darwin",
+        "tvm-ffi01-cu128-x86_64-linux",
+        "tvm-ffi-metal",
+        "tvm-ffi-cuda",
+    ]
+]
+
+
+def test_resolve_tvm_ffi_metal_darwin():
+    # tvm-ffi arch metal build should be picked on Darwin/aarch64 when
+    # tvm-ffi 0.1 is installed.
+    from kernels.backends import Metal
+
+    result = _resolve_variant_for_system(
+        variants=RESOLVE_VARIANTS_TVM_FFI,
+        selected_backend=Metal(),
+        cpu="aarch64",
+        os="darwin",
+        torch_version=None,
+        torch_cxx11_abi=None,
+        tvm_ffi_version=Version("0.1"),
+    )
+    assert result != []
+    assert result[0].variant_str == "tvm-ffi01-metal-aarch64-darwin"
+
+
+def test_resolve_tvm_ffi_metal_noarch_fallback():
+    # With no matching arch variant for the system, should fall back to
+    # the tvm-ffi metal noarch variant.
+    from kernels.backends import Metal
+
+    variants = [parse_variant(s) for s in ["tvm-ffi01-cu128-x86_64-linux", "tvm-ffi-metal"]]
+    result = _resolve_variant_for_system(
+        variants=variants,
+        selected_backend=Metal(),
+        cpu="aarch64",
+        os="darwin",
+        torch_version=None,
+        torch_cxx11_abi=None,
+        tvm_ffi_version=Version("0.1"),
+    )
+    assert result != []
+    assert result[0].variant_str == "tvm-ffi-metal"
+
+
+def test_resolve_torch_noarch_preferred_over_tvm_ffi_noarch():
+    # When both are available, Torch noarch is preferred over tvm-ffi noarch
+    # (mirrors the arch precedence: Torch > tvm-ffi).
+    variants = [parse_variant(s) for s in ["tvm-ffi-metal", "torch-metal"]]
+    from kernels.backends import Metal
+
+    result = _resolve_variant_for_system(
+        variants=variants,
+        selected_backend=Metal(),
+        cpu="aarch64",
+        os="darwin",
+        torch_version=Version("2.10"),
+        torch_cxx11_abi=None,
+        tvm_ffi_version=Version("0.1"),
+    )
+    assert result != []
+    assert result[0].variant_str == "torch-metal"
+
+
 def test_resolve_no_match():
     result = _resolve_variant_for_system(
         variants=RESOLVE_VARIANTS,