From c6ae5a9029f65a42467f34ec5f4d7e3476c4ae12 Mon Sep 17 00:00:00 2001 From: tqchen Date: Thu, 16 Jul 2026 00:21:25 +0000 Subject: [PATCH 1/3] fix(device): require canonical Python device names --- python/tvm_ffi/cython/device.pxi | 5 ----- tests/python/test_device.py | 17 +++++++++++++++++ tests/python/test_type_converter.py | 5 +++-- 3 files changed, 20 insertions(+), 7 deletions(-) diff --git a/python/tvm_ffi/cython/device.pxi b/python/tvm_ffi/cython/device.pxi index f7ae33402..ec6d0e6d8 100644 --- a/python/tvm_ffi/cython/device.pxi +++ b/python/tvm_ffi/cython/device.pxi @@ -115,13 +115,8 @@ 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, diff --git a/tests/python/test_device.py b/tests/python/test_device.py index 6575ec905..0b622c725 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 64a9bb195..a8e83db98 100644 --- a/tests/python/test_type_converter.py +++ b/tests/python/test_type_converter.py @@ -855,10 +855,11 @@ def test_device_str_converts(self) -> None: 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.""" From 68cc76b635e2a0f5063372accbe77d8fa6a350c6 Mon Sep 17 00:00:00 2001 From: tqchen Date: Thu, 16 Jul 2026 00:58:58 +0000 Subject: [PATCH 2/3] fix(device): remove remaining noncanonical aliases --- include/tvm/ffi/device.h | 4 ++-- tests/cpp/test_device.cc | 9 +++++---- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/include/tvm/ffi/device.h b/include/tvm/ffi/device.h index 401657895..ff5ad40c9 100644 --- a/include/tvm/ffi/device.h +++ b/include/tvm/ffi/device.h @@ -43,12 +43,12 @@ TVM_FFI_INLINE static std::optional TryParseDLDeviceType(std::stri if (name == "cuda") return kDLCUDA; if (name == "opencl") return kDLOpenCL; if (name == "vulkan") return kDLVulkan; - if (name == "metal" || name == "mps") return kDLMetal; + if (name == "metal") return kDLMetal; if (name == "vpi") return kDLVPI; if (name == "rocm") return kDLROCM; if (name == "ext_dev") return kDLExtDev; if (name == "hexagon") return kDLHexagon; - if (name == "wgpu" || name == "webgpu") return kDLWebGPU; + if (name == "webgpu") return kDLWebGPU; if (name == "maia") return kDLMAIA; if (name == "trn") return kDLTrn; return std::nullopt; diff --git a/tests/cpp/test_device.cc b/tests/cpp/test_device.cc index 3eb80859f..612f930d9 100644 --- a/tests/cpp/test_device.cc +++ b/tests/cpp/test_device.cc @@ -39,10 +39,9 @@ TEST(Device, AnyConversionWithString) { DLDeviceType device_type; }; const TestCase test_cases[] = { - {"cpu", kDLCPU}, {"cuda", kDLCUDA}, {"opencl", kDLOpenCL}, {"vulkan", kDLVulkan}, - {"metal", kDLMetal}, {"mps", kDLMetal}, {"vpi", kDLVPI}, {"rocm", kDLROCM}, - {"ext_dev", kDLExtDev}, {"hexagon", kDLHexagon}, {"wgpu", kDLWebGPU}, {"webgpu", kDLWebGPU}, - {"maia", kDLMAIA}, {"trn", kDLTrn}, + {"cpu", kDLCPU}, {"cuda", kDLCUDA}, {"opencl", kDLOpenCL}, {"vulkan", kDLVulkan}, + {"metal", kDLMetal}, {"vpi", kDLVPI}, {"rocm", kDLROCM}, {"ext_dev", kDLExtDev}, + {"hexagon", kDLHexagon}, {"webgpu", kDLWebGPU}, {"maia", kDLMAIA}, {"trn", kDLTrn}, }; for (const TestCase& test_case : test_cases) { SCOPED_TRACE(test_case.name); @@ -74,6 +73,8 @@ TEST(Device, RejectInvalidStrings) { EXPECT_FALSE(AnyView("test").try_cast().has_value()); EXPECT_FALSE(AnyView("nvptx").try_cast().has_value()); EXPECT_FALSE(AnyView("cl").try_cast().has_value()); + EXPECT_FALSE(AnyView("mps").try_cast().has_value()); + EXPECT_FALSE(AnyView("wgpu").try_cast().has_value()); EXPECT_FALSE(AnyView("cuda:").try_cast().has_value()); EXPECT_FALSE(AnyView("cuda:-1").try_cast().has_value()); EXPECT_FALSE(AnyView("cuda:1:2").try_cast().has_value()); From 064908bbb55a8bc1b4bf3602c1958e64daa41937 Mon Sep 17 00:00:00 2001 From: tqchen Date: Thu, 16 Jul 2026 05:01:54 +0000 Subject: [PATCH 3/3] fix(device): accept standard mps and wgpu spellings --- include/tvm/ffi/device.h | 4 ++-- python/tvm_ffi/cython/device.pxi | 2 ++ tests/cpp/test_device.cc | 9 ++++----- 3 files changed, 8 insertions(+), 7 deletions(-) diff --git a/include/tvm/ffi/device.h b/include/tvm/ffi/device.h index ff5ad40c9..401657895 100644 --- a/include/tvm/ffi/device.h +++ b/include/tvm/ffi/device.h @@ -43,12 +43,12 @@ TVM_FFI_INLINE static std::optional TryParseDLDeviceType(std::stri if (name == "cuda") return kDLCUDA; if (name == "opencl") return kDLOpenCL; if (name == "vulkan") return kDLVulkan; - if (name == "metal") return kDLMetal; + if (name == "metal" || name == "mps") return kDLMetal; if (name == "vpi") return kDLVPI; if (name == "rocm") return kDLROCM; if (name == "ext_dev") return kDLExtDev; if (name == "hexagon") return kDLHexagon; - if (name == "webgpu") return kDLWebGPU; + if (name == "wgpu" || name == "webgpu") return kDLWebGPU; if (name == "maia") return kDLMAIA; if (name == "trn") return kDLTrn; return std::nullopt; diff --git a/python/tvm_ffi/cython/device.pxi b/python/tvm_ffi/cython/device.pxi index ec6d0e6d8..0f3fa863c 100644 --- a/python/tvm_ffi/cython/device.pxi +++ b/python/tvm_ffi/cython/device.pxi @@ -120,10 +120,12 @@ cdef class Device: "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/cpp/test_device.cc b/tests/cpp/test_device.cc index 612f930d9..3eb80859f 100644 --- a/tests/cpp/test_device.cc +++ b/tests/cpp/test_device.cc @@ -39,9 +39,10 @@ TEST(Device, AnyConversionWithString) { DLDeviceType device_type; }; const TestCase test_cases[] = { - {"cpu", kDLCPU}, {"cuda", kDLCUDA}, {"opencl", kDLOpenCL}, {"vulkan", kDLVulkan}, - {"metal", kDLMetal}, {"vpi", kDLVPI}, {"rocm", kDLROCM}, {"ext_dev", kDLExtDev}, - {"hexagon", kDLHexagon}, {"webgpu", kDLWebGPU}, {"maia", kDLMAIA}, {"trn", kDLTrn}, + {"cpu", kDLCPU}, {"cuda", kDLCUDA}, {"opencl", kDLOpenCL}, {"vulkan", kDLVulkan}, + {"metal", kDLMetal}, {"mps", kDLMetal}, {"vpi", kDLVPI}, {"rocm", kDLROCM}, + {"ext_dev", kDLExtDev}, {"hexagon", kDLHexagon}, {"wgpu", kDLWebGPU}, {"webgpu", kDLWebGPU}, + {"maia", kDLMAIA}, {"trn", kDLTrn}, }; for (const TestCase& test_case : test_cases) { SCOPED_TRACE(test_case.name); @@ -73,8 +74,6 @@ TEST(Device, RejectInvalidStrings) { EXPECT_FALSE(AnyView("test").try_cast().has_value()); EXPECT_FALSE(AnyView("nvptx").try_cast().has_value()); EXPECT_FALSE(AnyView("cl").try_cast().has_value()); - EXPECT_FALSE(AnyView("mps").try_cast().has_value()); - EXPECT_FALSE(AnyView("wgpu").try_cast().has_value()); EXPECT_FALSE(AnyView("cuda:").try_cast().has_value()); EXPECT_FALSE(AnyView("cuda:-1").try_cast().has_value()); EXPECT_FALSE(AnyView("cuda:1:2").try_cast().has_value());