(tvm-ffi) branch main updated: fix(device): require canonical Python device names (#671)

[tqchen]

This is an automated email from the ASF dual-hosted git repository.

tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git


The following commit(s) were added to refs/heads/main by this push:
     new 5fbad39d fix(device): require canonical Python device names (#671)
5fbad39d is described below

commit 5fbad39d7d0a83300854d3a60a54117ebe9c1e30
Author: Tianqi Chen <tqc...@users.noreply.github.com>
AuthorDate: Thu Jul 16 13:52:46 2026 +0800

    fix(device): require canonical Python device names (#671)
    
    ## Summary
    
    - remove the Python-only `llvm`, `c`, `test`, `nvptx`, and `cl` device
    aliases
    - retain canonical device names and their existing error behavior
    - cover direct construction and typed FFI conversion rejection
    
    ## Testing
    
    - `pytest -o addopts= tests/python/test_device.py
    tests/python/test_function.py tests/python/test_type_converter.py`
    - scoped pre-commit hooks, including Ruff, Cython lint, and ty
---
 python/tvm_ffi/cython/device.pxi    |  7 ++-----
 tests/python/test_device.py         | 17 +++++++++++++++++
 tests/python/test_type_converter.py |  5 +++--
 3 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/python/tvm_ffi/cython/device.pxi b/python/tvm_ffi/cython/device.pxi
index f7ae3340..0f3fa863 100644
--- a/python/tvm_ffi/cython/device.pxi
+++ b/python/tvm_ffi/cython/device.pxi
@@ -115,20 +115,17 @@ cdef class Device:
     }
 
     _DEVICE_NAME_TO_TYPE = {
-        "llvm": DLDeviceType.kDLCPU,
         "cpu": DLDeviceType.kDLCPU,
-        "c": DLDeviceType.kDLCPU,
-        "test": DLDeviceType.kDLCPU,
         "cuda": DLDeviceType.kDLCUDA,
-        "nvptx": DLDeviceType.kDLCUDA,
-        "cl": DLDeviceType.kDLOpenCL,
         "opencl": DLDeviceType.kDLOpenCL,
         "vulkan": DLDeviceType.kDLVulkan,
         "metal": DLDeviceType.kDLMetal,
+        "mps": DLDeviceType.kDLMetal,
         "vpi": DLDeviceType.kDLVPI,
         "rocm": DLDeviceType.kDLROCM,
         "ext_dev": DLDeviceType.kDLExtDev,
         "hexagon": DLDeviceType.kDLHexagon,
+        "wgpu": DLDeviceType.kDLWebGPU,
         "webgpu": DLDeviceType.kDLWebGPU,
         "maia": DLDeviceType.kDLMAIA,
         "trn": DLDeviceType.kDLTrn,
diff --git a/tests/python/test_device.py b/tests/python/test_device.py
index 6575ec90..0b622c72 100644
--- a/tests/python/test_device.py
+++ b/tests/python/test_device.py
@@ -87,6 +87,23 @@ def test_device_with_dev_id(
     assert dev.index == expect_device_id
 
 
+@pytest.mark.parametrize(
+    "alias, canonical_name",
+    [
+        ("llvm", "cpu"),
+        ("c", "cpu"),
+        ("test", "cpu"),
+        ("nvptx", "cuda"),
+        ("cl", "opencl"),
+    ],
+)
+def test_device_rejects_noncanonical_aliases(alias: str, canonical_name: str) 
-> None:
+    with pytest.raises(ValueError, match=rf"^Unknown device: {alias}$"):
+        tvm_ffi.device(alias)
+
+    assert tvm_ffi.device(canonical_name).type == canonical_name
+
+
 @pytest.mark.parametrize("dev_type, dev_id", [("cpu:0:0", None), ("cpu:?", 
None), ("cpu:", None)])
 def test_deive_type_error(dev_type: str, dev_id: int | None) -> None:
     with pytest.raises(ValueError):
diff --git a/tests/python/test_type_converter.py 
b/tests/python/test_type_converter.py
index 0290a07b..ab582990 100644
--- a/tests/python/test_type_converter.py
+++ b/tests/python/test_type_converter.py
@@ -855,10 +855,11 @@ class TestConvertSpecialTypes:
         result = _to_py_class_value(A(tvm_ffi.Device).convert("cuda:1"))
         assert result == tvm_ffi.device("cuda", 1)
 
-    def test_device_bad_str_rejected(self) -> None:
+    @pytest.mark.parametrize("value", ["not_a_device", "llvm", "c", "test", 
"nvptx", "cl"])
+    def test_device_bad_str_rejected(self, value: str) -> None:
         """Invalid device strings remain type-conversion errors."""
         with pytest.raises(TypeError, match="device"):
-            A(tvm_ffi.Device).convert("not_a_device")
+            A(tvm_ffi.Device).convert(value)
 
     def test_device_base_class_passthrough_after_public_class_override(self) 
-> None:
         """The base Device cdef class remains accepted if _CLASS_DEVICE is 
overridden."""