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.git
The following commit(s) were added to refs/heads/main by this push:
new e479a5dbe7 [RUNTIME][PYTHON] Add explicit Target device conversion
(#20005)
e479a5dbe7 is described below
commit e479a5dbe724aa4aa86f301c8d5ad1a5a05a58e3
Author: Tianqi Chen <[email protected]>
AuthorDate: Wed Jul 15 05:34:21 2026 +0800
[RUNTIME][PYTHON] Add explicit Target device conversion (#20005)
## Summary
Compiler Targets can carry device-type semantics that runtime
device-name parsing does not preserve.
- add `tvm.device_from_target` for canonical Target-to-Device
translation
- use explicit runtime constructors where the device kind is fixed
- update target-derived utilities, tests, and documentation to use the
explicit boundary
---
docs/contribute/code_guide.rst | 2 +-
docs/contribute/testing.rst | 8 +-
docs/get_started/tutorials/ir_module.py | 2 +-
docs/how_to/tutorials/customize_opt.py | 2 +-
docs/how_to/tutorials/e2e_opt_model.py | 2 +-
docs/how_to/tutorials/optimize_llm.py | 2 +-
python/tvm/__init__.py | 2 +-
python/tvm/runtime/__init__.py | 2 +-
python/tvm/runtime/_tensor.py | 23 +++++
python/tvm/testing/utils.py | 8 +-
python/tvm/tirx/build.py | 6 +-
.../test_nnapi/test_from_exported_to_cuda.py | 104 ++++++++++-----------
tests/python/codegen/test_gpu_codegen_allreduce.py | 4 +-
tests/python/codegen/test_target_codegen.py | 2 +-
tests/python/codegen/test_target_codegen_bool.py | 2 +-
tests/python/codegen/test_target_codegen_cuda.py | 6 +-
.../python/codegen/test_target_codegen_cuda_fp4.py | 6 +-
.../python/codegen/test_target_codegen_cuda_fp8.py | 16 ++--
tests/python/codegen/test_target_codegen_extern.py | 2 +-
.../codegen/test_target_codegen_gpu_common.py | 2 +-
tests/python/codegen/test_target_codegen_metal.py | 4 +-
tests/python/codegen/test_target_codegen_opencl.py | 12 +--
tests/python/codegen/test_target_codegen_vulkan.py | 22 ++---
.../relax/test_backend_dispatch_sort_scan.py | 2 +-
tests/python/relax/test_codegen_cublas.py | 2 +-
tests/python/relax/test_codegen_cudnn.py | 2 +-
tests/python/relax/test_codegen_cutlass.py | 10 +-
tests/python/relax/test_codegen_hipblas.py | 2 +-
tests/python/relax/test_codegen_tensorrt.py | 4 +-
tests/python/relax/test_contrib_vllm.py | 4 +-
...test_frontend_nn_llm_sequence_prefill_masked.py | 18 ++--
tests/python/relax/test_frontend_nn_op.py | 2 +-
tests/python/relax/test_op_gradient_numeric.py | 104 ++++++++++-----------
tests/python/relax/test_op_take.py | 16 ++--
tests/python/relax/test_op_view.py | 10 +-
.../relax/test_training_optimizer_numeric.py | 6 +-
.../python/relax/test_training_trainer_numeric.py | 8 +-
.../relax/test_transform_gradient_numeric.py | 8 +-
.../relax/test_transform_lazy_transform_params.py | 2 +-
tests/python/relax/test_vm_build.py | 4 +-
tests/python/relax/test_vm_builtin.py | 2 +-
tests/python/relax/test_vm_callback_function.py | 6 +-
tests/python/runtime/test_runtime_dlpack.py | 2 +-
.../test_meta_schedule_mma_tensorize.py | 2 +-
..._tir_schedule_tensorize_ldmatrix_mma_numeric.py | 2 +-
.../test_tir_schedule_tensorize_mfma_numeric.py | 2 +-
...est_s_tir_transform_inject_software_pipeline.py | 2 +-
tests/python/target/test_target_target.py | 48 +++++++---
tests/python/target/test_virtual_device.py | 6 +-
tests/python/tirx-base/test_tir_intrin.py | 2 +-
tests/python/tirx/codegen/test_codegen_ampere.py | 2 +-
tests/python/tirx/codegen/test_codegen_cuda.py | 10 +-
52 files changed, 290 insertions(+), 241 deletions(-)
diff --git a/docs/contribute/code_guide.rst b/docs/contribute/code_guide.rst
index 5136b9d4d6..eb4ba5420d 100644
--- a/docs/contribute/code_guide.rst
+++ b/docs/contribute/code_guide.rst
@@ -140,7 +140,7 @@ If you want your test to run over a variety of targets,
parametrize over ``targe
def test_mytest(target):
if not tvm.testing.device_enabled(target):
pytest.skip(f"{target} not enabled")
- dev = tvm.device(target)
+ dev = tvm.device_from_target(target)
...
will run ``test_mytest`` with ``target="llvm"`` and ``target="cuda"``,
skipping any target whose device is not present. If you only want to test
against a single target, drop the parametrization and hardcode the target. Mark
GPU tests with ``@pytest.mark.gpu`` so the CI can select them, and skip when
the required feature is unavailable with ``@pytest.mark.skipif``. For example,
CUDA tests use:
diff --git a/docs/contribute/testing.rst b/docs/contribute/testing.rst
index f30fd7d03b..8dd71ca4ca 100644
--- a/docs/contribute/testing.rst
+++ b/docs/contribute/testing.rst
@@ -62,7 +62,7 @@ over ``target`` with ``@pytest.mark.parametrize``. Tag each
GPU target
with ``pytest.mark.gpu`` so the CI routes it to a GPU node, skip a target
that cannot run on the current machine with
:py:func:`tvm.testing.device_enabled`, and obtain its device with
-``tvm.device(target)``. The function is run once per target, the
+``tvm.device_from_target(target)``. The function is run once per target, the
success/failure of each is reported separately, and a target whose device
is disabled in ``config.cmake`` or absent from the machine is reported as
skipped.
@@ -76,7 +76,7 @@ skipped.
def test_function(target):
if not tvm.testing.device_enabled(target):
pytest.skip(f"{target} not enabled")
- dev = tvm.device(target)
+ dev = tvm.device_from_target(target)
# Test code goes here
For a test that only applies to a single target, omit the parametrization
@@ -91,7 +91,7 @@ for a GPU target):
)
def test_function():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
# Test code goes here
To exclude a target, leave it out of the parametrize list. To mark a
@@ -113,7 +113,7 @@ as above:
def test_function(target, impl):
if not tvm.testing.device_enabled(target):
pytest.skip(f"{target} not enabled")
- dev = tvm.device(target)
+ dev = tvm.device_from_target(target)
# Test code goes here
diff --git a/docs/get_started/tutorials/ir_module.py
b/docs/get_started/tutorials/ir_module.py
index c5a429461c..beb5e7ed76 100644
--- a/docs/get_started/tutorials/ir_module.py
+++ b/docs/get_started/tutorials/ir_module.py
@@ -265,7 +265,7 @@ with tvm.target.Target("cuda"):
# Now we can compile the IRModule on GPU, the similar way as we did on CPU.
exec = tvm.compile(gpu_mod, target="cuda")
-dev = tvm.device("cuda", 0)
+dev = tvm.cuda(0)
vm = relax.VirtualMachine(exec, dev)
# Need to allocate data and params on GPU device
data = tvm.runtime.tensor(raw_data, dev)
diff --git a/docs/how_to/tutorials/customize_opt.py
b/docs/how_to/tutorials/customize_opt.py
index 4872b324d9..8b236ecd79 100644
--- a/docs/how_to/tutorials/customize_opt.py
+++ b/docs/how_to/tutorials/customize_opt.py
@@ -218,7 +218,7 @@ mod.show()
# We can build and deploy the optimized model to the TVM runtime.
ex = tvm.compile(mod, target="cuda")
-dev = tvm.device("cuda", 0)
+dev = tvm.cuda(0)
vm = relax.VirtualMachine(ex, dev)
# Need to allocate data and params on GPU device
data = tvm.runtime.tensor(np.random.rand(*input_shape).astype("float32"), dev)
diff --git a/docs/how_to/tutorials/e2e_opt_model.py
b/docs/how_to/tutorials/e2e_opt_model.py
index bdb7ac0c91..5de3d50485 100644
--- a/docs/how_to/tutorials/e2e_opt_model.py
+++ b/docs/how_to/tutorials/e2e_opt_model.py
@@ -143,7 +143,7 @@ if not IS_IN_CI:
with target:
mod = tvm.s_tir.transform.DefaultGPUSchedule()(mod)
ex = tvm.compile(mod, target=target)
- dev = tvm.device("cuda", 0)
+ dev = tvm.cuda(0)
vm = relax.VirtualMachine(ex, dev)
# Need to allocate data and params on GPU device
gpu_data = tvm.runtime.tensor(np.random.rand(1, 3, 224,
224).astype("float32"), dev)
diff --git a/docs/how_to/tutorials/optimize_llm.py
b/docs/how_to/tutorials/optimize_llm.py
index 941b67dedd..9e83a0c12a 100644
--- a/docs/how_to/tutorials/optimize_llm.py
+++ b/docs/how_to/tutorials/optimize_llm.py
@@ -92,7 +92,7 @@ class LlamaConfig:
head_dim: int = 64 # hidden_size // num_attention_heads
-dev = tvm.device("cuda", 0)
+dev = tvm.cuda(0)
target = tvm.target.Target.from_device(dev)
diff --git a/python/tvm/__init__.py b/python/tvm/__init__.py
index b6d928edce..4e31f86e0c 100644
--- a/python/tvm/__init__.py
+++ b/python/tvm/__init__.py
@@ -31,7 +31,7 @@ from .base import _RUNTIME_ONLY
# tvm.runtime
from .runtime import Object
-from .runtime._tensor import device, cpu, cuda, opencl, vulkan, metal
+from .runtime._tensor import device, device_from_target, cpu, cuda, opencl,
vulkan, metal
from .runtime._tensor import vpi, rocm, ext_dev, hexagon
from .runtime import DataType, DataTypeCode
diff --git a/python/tvm/runtime/__init__.py b/python/tvm/runtime/__init__.py
index c51cb05dc4..56b904a519 100644
--- a/python/tvm/runtime/__init__.py
+++ b/python/tvm/runtime/__init__.py
@@ -33,7 +33,7 @@ from .module import Module
from .executable import Executable
# function exposures
-from ._tensor import device, cpu, cuda, opencl, vulkan, metal
+from ._tensor import device, device_from_target, cpu, cuda, opencl, vulkan,
metal
from ._tensor import vpi, rocm, ext_dev, from_dlpack
from .module import load_module, enabled, system_lib, load_static_library,
num_threads
from .object_generic import const
diff --git a/python/tvm/runtime/_tensor.py b/python/tvm/runtime/_tensor.py
index 51919c0178..6f1d60b01e 100644
--- a/python/tvm/runtime/_tensor.py
+++ b/python/tvm/runtime/_tensor.py
@@ -353,6 +353,29 @@ def tensor(arr, device=None, mem_scope=None):
return empty(arr.shape, arr.dtype, device, mem_scope).copyfrom(arr)
+def device_from_target(target, index=None):
+ """Construct a runtime device from a compilation target.
+
+ Parameters
+ ----------
+ target : str or dict or tvm.target.Target
+ The compilation target whose device type should be used.
+
+ index : int, optional
+ The integer device index.
+
+ Returns
+ -------
+ dev : Device
+ The created device.
+ """
+ from tvm.target import Target # pylint: disable=import-outside-toplevel
+
+ if not isinstance(target, Target):
+ target = Target(target)
+ return device(target.get_target_device_type(), index)
+
+
def cpu(dev_id=0):
"""Construct a CPU device
diff --git a/python/tvm/testing/utils.py b/python/tvm/testing/utils.py
index f41f06dc07..96420ef50b 100644
--- a/python/tvm/testing/utils.py
+++ b/python/tvm/testing/utils.py
@@ -415,7 +415,7 @@ def _get_targets(target_names=None):
is_runnable = is_enabled and "ANDROID_SERIAL_NUMBER" in os.environ
else:
is_enabled = tvm.runtime.enabled(target_kind)
- is_runnable = is_enabled and tvm.device(target_kind).exist
+ is_runnable = is_enabled and tvm.device_from_target(target).exist
targets.append(
{
@@ -521,7 +521,11 @@ def enabled_targets():
A list of pairs of all enabled devices and the associated context
"""
- return [(t["target"], tvm.device(t["target_kind"])) for t in
_get_targets() if t["is_runnable"]]
+ return [
+ (t["target"], tvm.device_from_target(t["target"]))
+ for t in _get_targets()
+ if t["is_runnable"]
+ ]
def _parse_target_entry(entry):
diff --git a/python/tvm/tirx/build.py b/python/tvm/tirx/build.py
index 10ec096bca..f15650a3ba 100644
--- a/python/tvm/tirx/build.py
+++ b/python/tvm/tirx/build.py
@@ -18,6 +18,8 @@
# pylint: disable=invalid-name
"""The build utils in python."""
+from tvm_ffi import DLDeviceType
+
import tvm
from tvm import ir
from tvm.ir.module import IRModule
@@ -204,9 +206,7 @@ def build(
if target is not None:
if target.host is not None:
target_host = target.host
- elif (
- tvm.device(target.kind.name, 0).dlpack_device_type() ==
tvm.cpu(0).dlpack_device_type()
- ):
+ elif target.get_target_device_type() == DLDeviceType.kDLCPU:
target_host = target
target_host = Target(target_host)
target_to_bind = target_to_bind.with_host(target_host)
diff --git a/tests/nightly/python/test_nnapi/test_from_exported_to_cuda.py
b/tests/nightly/python/test_nnapi/test_from_exported_to_cuda.py
index bf774051ef..438af772b7 100644
--- a/tests/nightly/python/test_nnapi/test_from_exported_to_cuda.py
+++ b/tests/nightly/python/test_nnapi/test_from_exported_to_cuda.py
@@ -73,7 +73,7 @@ def
assert_torch_output_vs_tvm_from_exported_to_cuda(raw_data, torch_module, tar
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_index_tensor():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class IndexModel0(nn.Module):
def __init__(self):
@@ -179,7 +179,7 @@ def test_index_tensor():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_full():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class FullModel(nn.Module):
def __init__(self):
@@ -197,7 +197,7 @@ def test_full():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_full_like():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class FullLike(nn.Module):
def __init__(self):
@@ -216,7 +216,7 @@ def test_full_like():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_ones():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class FullModel(nn.Module):
def __init__(self):
@@ -234,7 +234,7 @@ def test_ones():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_sort():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
raw_data = np.array([[4, 1, 13], [-30, 1, 3], [4, 0,
10]]).astype("float32")
@@ -263,7 +263,7 @@ def test_sort():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_tensor_clamp():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class ClampBothTensor(torch.nn.Module):
def __init__(self):
@@ -348,7 +348,7 @@ def test_tensor_clamp():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_tensor_expand_as():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class ExpandAs0(torch.nn.Module):
def __init__(self):
@@ -399,7 +399,7 @@ def test_tensor_expand_as():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_copy_():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class CopyTester(nn.Module):
def __init__(self, size):
@@ -425,7 +425,7 @@ def test_upsample_with_size():
factor argument but not both. This tests the former.
"""
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
batch_size = 1
channels = 3
@@ -442,7 +442,7 @@ def test_upsample_with_size():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_detach_no_change():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
# In TVM, detach() is just identity
class DetachTester(nn.Module):
@@ -463,7 +463,7 @@ def test_upsample_with_scale_factor():
factor argument but not both. This tests the latter.
"""
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
batch_size = 2
channels = 3
@@ -481,7 +481,7 @@ def test_upsample_with_scale_factor():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_linalg_vector_norm():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class VectorNorm0(torch.nn.Module):
def forward(self, x):
@@ -516,7 +516,7 @@ def test_linalg_vector_norm():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_batch_norm_prog():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
# Default args, in a pytorch program (to ensure output is in proper type
and format)
raw_data = np.random.randn(2, 3, 2, 2).astype(np.float32)
@@ -539,7 +539,7 @@ def test_batch_norm_prog():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_split_size():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
# Test split using the split_size argument such that it is not a divisor
# of the dimension to split (the last tensor will be smaller)
@@ -567,7 +567,7 @@ def test_split_size():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_split_sections_list():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
# Test split using a list of section sizes
batch = 3
@@ -595,7 +595,7 @@ def test_split_sections_list():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_batch_norm0():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
# Eval, no momentum, no affine, no running stats
raw_data = np.random.randn(8, 3, 4, 4).astype(np.float32)
@@ -609,7 +609,7 @@ def test_batch_norm0():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_batch_norm1():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
# Eval, with momentum, no affine, with running stats
raw_data = np.random.randn(1, 4, 2, 2).astype(np.float32)
@@ -623,7 +623,7 @@ def test_batch_norm1():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_batch_norm2():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
# Eval, with momentum, affine, no running stats
raw_data = np.random.randn(3, 4, 2, 2).astype(np.float32)
@@ -637,7 +637,7 @@ def test_batch_norm2():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_batch_norm3():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
# Eval, no momentum, affine, with running stats
raw_data = np.random.randn(1, 2, 2, 2).astype(np.float32)
@@ -651,7 +651,7 @@ def test_batch_norm3():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_chunk_even():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
# Chunks is a divisor of the dimension size
batch = 6
@@ -679,7 +679,7 @@ def test_chunk_even():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_chunk_uneven():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
# Chunks is not a divisor of the dimension size
batch = 2
@@ -707,7 +707,7 @@ def test_chunk_uneven():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_chunk_too_many():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
# If user asks for more chunks than the size of the dim, pytorch simply
splits in sections of size 1
batch = 1
@@ -735,7 +735,7 @@ def test_chunk_too_many():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_arange():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
# arange.default
raw_data = np.array([0, 0, 0, 0, 0])
@@ -772,7 +772,7 @@ def test_arange():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_index_select():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class IndexSelectModel(nn.Module):
def forward(self, x):
@@ -788,7 +788,7 @@ def test_index_select():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_stack():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class StackModel(nn.Module):
def forward(self, x):
@@ -807,7 +807,7 @@ def test_stack():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_sum():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class SumModel(nn.Module):
def forward(self, x):
@@ -823,7 +823,7 @@ def test_sum():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_mul():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class MulModule(nn.Module):
def __init__(self):
@@ -842,7 +842,7 @@ def test_mul():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_concat():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class ConcatFour(nn.Module):
def __init__(self, dim=0):
@@ -864,7 +864,7 @@ def test_concat():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_leakyrelu_module():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class LeakyReLUModule(nn.Module):
def __init__(self):
@@ -883,7 +883,7 @@ def test_leakyrelu_module():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_log_softmax_module():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class LogSoftmaxModule(nn.Module):
def __init__(self):
@@ -902,7 +902,7 @@ def test_log_softmax_module():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_softmax_module():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class SoftmaxModule(nn.Module):
def __init__(self):
@@ -921,7 +921,7 @@ def test_softmax_module():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_adaptive_avg_pool2d_module():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class AdaptiveAvgPool2dModule(nn.Module):
def __init__(self):
@@ -940,7 +940,7 @@ def test_adaptive_avg_pool2d_module():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_avg_pool2d_module():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class AvgPool2dModule(nn.Module):
def __init__(self):
@@ -959,7 +959,7 @@ def test_avg_pool2d_module():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_conv1d_module():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class Conv1dModule(nn.Module):
def __init__(self):
@@ -978,7 +978,7 @@ def test_conv1d_module():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_conv2d_module():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class Conv2dModule(nn.Module):
def __init__(self):
@@ -997,7 +997,7 @@ def test_conv2d_module():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_conv3d_module():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class Conv3dModule(nn.Module):
def __init__(self):
@@ -1016,7 +1016,7 @@ def test_conv3d_module():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_group_norm_module():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class GroupNormModule(nn.Module):
def __init__(self):
@@ -1035,7 +1035,7 @@ def test_group_norm_module():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_layer_norm_module():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class LayerNormModule(nn.Module):
def __init__(self):
@@ -1054,7 +1054,7 @@ def test_layer_norm_module():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_linear_module():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class LinearModule(nn.Module):
def __init__(self):
@@ -1073,7 +1073,7 @@ def test_linear_module():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_max_pool2d_module():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class MaxPool2dModule(nn.Module):
def __init__(self):
@@ -1092,7 +1092,7 @@ def test_max_pool2d_module():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_embedding_module():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class EmbeddingModule(nn.Module):
def __init__(self):
@@ -1111,7 +1111,7 @@ def test_embedding_module():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_flatten_module():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class FlattenModule(nn.Module):
def __init__(self):
@@ -1130,7 +1130,7 @@ def test_flatten_module():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_numel():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class NumelModule(nn.Module):
def forward(self, x):
@@ -1145,7 +1145,7 @@ def test_numel():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_size():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class SizeModule(nn.Module):
def forward(self, x):
@@ -1160,7 +1160,7 @@ def test_size():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_tensor():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class TensorModule(nn.Module):
def forward(self, x):
@@ -1175,7 +1175,7 @@ def test_tensor():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_type():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class TypeModule(nn.Module):
def forward(self, x):
@@ -1190,7 +1190,7 @@ def test_type():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_float():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class FloatModule(nn.Module):
def forward(self, x):
@@ -1205,7 +1205,7 @@ def test_float():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_half():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class HalfModule(nn.Module):
def forward(self, x):
@@ -1220,7 +1220,7 @@ def test_half():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_getattr():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class GetAttrModule(nn.Module):
def forward(self, x):
@@ -1236,7 +1236,7 @@ def test_getattr():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_sym_size_int():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class SymSizeIntModule(nn.Module):
def forward(self, x):
@@ -1251,7 +1251,7 @@ def test_sym_size_int():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_interpolate():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class InterpolateModule(nn.Module):
def forward(self, x):
@@ -1267,7 +1267,7 @@ def test_interpolate():
@pytest.mark.skipif(not tvm.testing.device_enabled("cuda"), reason="cuda not
enabled")
def test_cross_entropy_module():
target = "cuda"
- dev = tvm.device(target)
+ dev = tvm.cuda()
class CrossEntropyModule(nn.Module):
def __init__(self):
diff --git a/tests/python/codegen/test_gpu_codegen_allreduce.py
b/tests/python/codegen/test_gpu_codegen_allreduce.py
index 1012e1e517..119827f99f 100644
--- a/tests/python/codegen/test_gpu_codegen_allreduce.py
+++ b/tests/python/codegen/test_gpu_codegen_allreduce.py
@@ -96,7 +96,7 @@ def test_allreduce_sum(dims, target):
b_np = a_np.sum(axis=-1).astype("float32")
def run_and_check():
- dev = tvm.device(target)
+ dev = tvm.device_from_target(target)
a = tvm.runtime.tensor(a_np, dev)
b = tvm.runtime.tensor(np.zeros_like(b_np), dev)
f(a, b)
@@ -161,7 +161,7 @@ def test_allreduce_max(dims, target):
b_np = a_np.max(axis=-1).astype("float32")
def run_and_check():
- dev = tvm.device(target)
+ dev = tvm.device_from_target(target)
a = tvm.runtime.tensor(a_np, dev)
b = tvm.runtime.tensor(np.zeros_like(b_np), dev)
f(a, b)
diff --git a/tests/python/codegen/test_target_codegen.py
b/tests/python/codegen/test_target_codegen.py
index a3a232d785..7157ae0f69 100644
--- a/tests/python/codegen/test_target_codegen.py
+++ b/tests/python/codegen/test_target_codegen.py
@@ -138,7 +138,7 @@ def test_codegen_loop_step(target):
if target == "c":
assert src.find("for (int32_t i = 3; i < 1024; i += 96)") >= 0
- dev = tvm.device(target, 0)
+ dev = tvm.cpu()
a_np = np.random.rand(1024).astype("float32")
b_np = np.random.rand(1024).astype("float32")
c_np = np.zeros(1024, dtype="float32")
diff --git a/tests/python/codegen/test_target_codegen_bool.py
b/tests/python/codegen/test_target_codegen_bool.py
index c61312d6db..c47e469223 100644
--- a/tests/python/codegen/test_target_codegen_bool.py
+++ b/tests/python/codegen/test_target_codegen_bool.py
@@ -89,7 +89,7 @@ def test_cmp_load_store(target):
b_np = np.random.uniform(size=arr_size).astype("float32")
def run_and_check():
- dev = tvm.device(target)
+ dev = tvm.device_from_target(target)
a = tvm.runtime.tensor(a_np, dev)
b = tvm.runtime.tensor(b_np, dev)
d = tvm.runtime.tensor(np.zeros(arr_size, dtype="float32"), dev)
diff --git a/tests/python/codegen/test_target_codegen_cuda.py
b/tests/python/codegen/test_target_codegen_cuda.py
index 89dc3268f3..a820cf5f67 100644
--- a/tests/python/codegen/test_target_codegen_cuda.py
+++ b/tests/python/codegen/test_target_codegen_cuda.py
@@ -329,7 +329,7 @@ def test_cuda_inf_nan():
fun = tvm.compile(Module, target="cuda")
def run_and_check():
- dev = tvm.device("cuda", 0)
+ dev = tvm.cuda(0)
a = tvm.runtime.empty((n,), dtype, dev)
c = tvm.runtime.empty((n,), dtype, dev)
fun(a, c)
@@ -387,7 +387,7 @@ def test_crossthread_reduction1(target):
vals = [nthd - 1, nthd, nthd + 1]
def run_and_check():
- dev = tvm.device(target)
+ dev = tvm.device_from_target(target)
for kk in vals:
size = (nn, kk)
a =
tvm.runtime.tensor(np.random.uniform(size=size).astype("float32"), dev)
@@ -456,7 +456,7 @@ def test_crossthread_reduction2(target):
vy = [nthdy - 1, nthdy, nthdy + 1]
def run_and_check():
- dev = tvm.device(target)
+ dev = tvm.device_from_target(target)
for kk0, kk1 in [(x, y) for x in vx for y in vy]:
size = (nn, kk0, kk1)
a =
tvm.runtime.tensor(np.random.uniform(size=size).astype("float32"), dev)
diff --git a/tests/python/codegen/test_target_codegen_cuda_fp4.py
b/tests/python/codegen/test_target_codegen_cuda_fp4.py
index d49f10758d..01a0cea9b5 100644
--- a/tests/python/codegen/test_target_codegen_cuda_fp4.py
+++ b/tests/python/codegen/test_target_codegen_cuda_fp4.py
@@ -88,7 +88,7 @@ def test_e2m1_vector_conversions(promoted_dtype):
b_np = np.random.choice(valid_fp4_values,
size=np_shape).astype(np.int8)
def run_and_check():
- dev = tvm.device(target, 0)
+ dev = tvm.cuda(0)
a = tvm.runtime.empty(shape=(vector_length,), dtype=native_dtype,
device=dev)
a.copyfrom(a_np)
b = tvm.runtime.empty(shape=(vector_length,), dtype=native_dtype,
device=dev)
@@ -188,7 +188,7 @@ def _scalar_reinterpret_module(n, num_blocks,
vector_length, num_elem_per_storag
def test_e2m1_dequantize():
n = 128
- dev = tvm.device("cuda", 0)
+ dev = tvm.cuda(0)
target = tvm.target.Target.from_device(dev)
num_elem_per_storage = 32 // 4
@@ -264,7 +264,7 @@ def test_e2m1_scalar_buffer_offset():
expected = fp4_to_fp16[fp4_elements]
def run_and_check():
- dev = tvm.device(target, 0)
+ dev = tvm.cuda(0)
a = tvm.runtime.empty(shape=(n // 2,), dtype="uint8", device=dev)
a.copyfrom(packed)
b = tvm.runtime.empty(shape=(n,), dtype="float16", device=dev)
diff --git a/tests/python/codegen/test_target_codegen_cuda_fp8.py
b/tests/python/codegen/test_target_codegen_cuda_fp8.py
index dd21061a46..bd8fbef978 100644
--- a/tests/python/codegen/test_target_codegen_cuda_fp8.py
+++ b/tests/python/codegen/test_target_codegen_cuda_fp8.py
@@ -77,7 +77,7 @@ def test_fp8_conversions(input):
cuda_src = fadd.imports[0].inspect_source()
assert nv_dtype in cuda_src, f"{nv_dtype} datatype not found in generated
CUDA"
- dev = tvm.device(target, 0)
+ dev = tvm.cuda(0)
a = tvm.runtime.tensor(np.random.uniform(low=0, high=5,
size=64).astype(dtype), dev)
b = tvm.runtime.tensor(np.random.uniform(low=0, high=5,
size=64).astype(dtype), dev)
@@ -130,7 +130,7 @@ def test_fp8_packing(dtype):
mod = _create_mod(native_dtype, packed_dtype, length)
target = "cuda"
f = tvm.compile(mod, target=target)
- dev = tvm.device(target, 0)
+ dev = tvm.cuda(0)
np_shape = (length, vector_length)
a_np = np.random.uniform(low=0, high=5, size=np_shape).astype(dtype)
@@ -192,7 +192,7 @@ def test_fp8_vector_conversions(native_dtype,
promoted_dtype, numpytype):
target = "cuda"
fadd = tvm.tirx.build(mod, target=target)
cuda_src = fadd.imports[0].inspect_source()
- dev = tvm.device(target, 0)
+ dev = tvm.cuda(0)
if "x" in native_dtype:
lanes = int(native_dtype.split("x")[-1])
@@ -242,7 +242,7 @@ def test_half_broadcast(bcast_length):
mod = _create_mod(bcast_length, dtype)
target = "cuda"
func = tvm.compile(mod, target=target)
- dev = tvm.device(target, 0)
+ dev = tvm.cuda(0)
a_np = np.random.uniform(low=0, high=4, size=()).astype(dtype)
a = tvm.runtime.tensor(a_np, device=dev)
@@ -277,7 +277,7 @@ def test_half_misaligned_vector_load(vector_length):
target = "cuda"
f = tvm.compile(vector_load, target=target)
- dev = tvm.device(target, 0)
+ dev = tvm.cuda(0)
a_np = np.random.uniform(low=0, high=1, size=(length,)).astype(dtype)
a = tvm.runtime.tensor(a_np, device=dev)
@@ -321,7 +321,7 @@ def test_half4_vector_add():
target = "cuda"
fadd = tvm.compile(Module, target=target)
- dev = tvm.device(target, 0)
+ dev = tvm.cuda(0)
a_np = np.random.uniform(-1, 1, (length, vector_length)).astype(dtype)
a = tvm.runtime.empty(shape=(length,), dtype=vec_dtype, device=dev)
@@ -780,7 +780,7 @@ class
TestFP8e4x4QuantDequantScale(BaseFP8E4M3QuantScaleOnly):
axis,
target_str,
):
- dev = tvm.device(target_str, 0)
+ dev = tvm.cuda(0)
return self.compile_quant_and_dequant_by_scale(
weight_shape,
scale_shape,
@@ -800,7 +800,7 @@ class
TestFP8e4x4QuantDequantScale(BaseFP8E4M3QuantScaleOnly):
@pytest.mark.skipif(not env.has_cuda_compute(8, 9), reason="need cuda
compute >= 8.9")
def test_main(self, weight_shape, model_dtype, target_str,
compiled_functions):
quant, dequant = compiled_functions
- dev = tvm.device(target_str, 0)
+ dev = tvm.cuda(0)
weight_np = np.random.uniform(-100, 100,
weight_shape).astype(model_dtype)
weight = tvm.runtime.tensor(weight_np, device=dev)
diff --git a/tests/python/codegen/test_target_codegen_extern.py
b/tests/python/codegen/test_target_codegen_extern.py
index 8a67ab8eff..0330b9d441 100644
--- a/tests/python/codegen/test_target_codegen_extern.py
+++ b/tests/python/codegen/test_target_codegen_extern.py
@@ -58,7 +58,7 @@ def test_add_pipeline():
n = nn
def run_and_check():
- dev = tvm.device(target, 0)
+ dev = tvm.device_from_target(target, 0)
a =
tvm.runtime.tensor(np.random.uniform(size=n).astype("float32"), dev)
c = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev)
f(a, c)
diff --git a/tests/python/codegen/test_target_codegen_gpu_common.py
b/tests/python/codegen/test_target_codegen_gpu_common.py
index 145347375b..e2889beebd 100644
--- a/tests/python/codegen/test_target_codegen_gpu_common.py
+++ b/tests/python/codegen/test_target_codegen_gpu_common.py
@@ -66,7 +66,7 @@ def test_int_intrin(target, dtype):
f = tvm.compile(Module, target=target)
def run_and_check():
- dev = tvm.device(target["kind"] if isinstance(target, dict) else
target)
+ dev = tvm.device_from_target(target)
a = tvm.runtime.tensor(np.random.randint(0, 100000,
size=n).astype(dtype), dev)
b = tvm.runtime.tensor(np.zeros(shape=(n,)).astype(dtype), dev)
f(a, b)
diff --git a/tests/python/codegen/test_target_codegen_metal.py
b/tests/python/codegen/test_target_codegen_metal.py
index 794a861c3a..223c670bf7 100644
--- a/tests/python/codegen/test_target_codegen_metal.py
+++ b/tests/python/codegen/test_target_codegen_metal.py
@@ -49,7 +49,7 @@ def test_metal_inf_nan():
fun = tvm.compile(Module, target=target)
def run_and_check():
- dev = tvm.device(target, 0)
+ dev = tvm.metal(0)
a = tvm.runtime.empty((n,), dtype, dev)
c = tvm.runtime.empty((n,), dtype, dev)
fun(a, c)
@@ -116,7 +116,7 @@ def test_metal_erf():
fun = tvm.compile(Module, target=target)
def run_and_check():
- dev = tvm.device(target, 0)
+ dev = tvm.metal(0)
a = tvm.runtime.empty((n,), dtype, dev)
c = tvm.runtime.empty((n,), dtype, dev)
fun(a, c)
diff --git a/tests/python/codegen/test_target_codegen_opencl.py
b/tests/python/codegen/test_target_codegen_opencl.py
index 16da13039c..82e46dd779 100644
--- a/tests/python/codegen/test_target_codegen_opencl.py
+++ b/tests/python/codegen/test_target_codegen_opencl.py
@@ -54,7 +54,7 @@ def test_opencl_ternary_expression():
fun = tvm.tirx.build(Module, target=target)
def run_and_check():
- dev = tvm.device(target, 0)
+ dev = tvm.opencl(0)
a = tvm.runtime.empty((n,), dtype, dev)
c = tvm.runtime.empty((n,), dtype, dev)
fun(a, c)
@@ -84,7 +84,7 @@ def test_opencl_ternary_expression():
fun = tvm.tirx.build(Module, target=target)
def run_and_check():
- dev = tvm.device(target, 0)
+ dev = tvm.opencl(0)
a = tvm.runtime.empty((n,), dtype, dev)
c = tvm.runtime.empty((n,), dtype, dev)
fun(a, c)
@@ -120,7 +120,7 @@ def test_opencl_inf_nan():
fun = tvm.tirx.build(Module, target=target)
def run_and_check():
- dev = tvm.device(target, 0)
+ dev = tvm.opencl(0)
a = tvm.runtime.empty((n,), dtype, dev)
c = tvm.runtime.empty((n,), dtype, dev)
fun(a, c)
@@ -154,7 +154,7 @@ def test_opencl_max():
fun = tvm.tirx.build(Module, target=target)
def run_and_check():
- dev = tvm.device(target, 0)
+ dev = tvm.opencl(0)
a = tvm.runtime.empty((n,), dtype, dev)
c = tvm.runtime.empty((n,), dtype, dev)
fun(a, c)
@@ -221,7 +221,7 @@ def test_opencl_type_casting():
assert assembly.count(pattern_cond) != 0
def run_and_check():
- dev = tvm.device(target, 0)
+ dev = tvm.opencl(0)
c = tvm.runtime.empty((n,), dtype, dev)
fun(c)
@@ -321,7 +321,7 @@ def test_export_load_with_fallback(monkeypatch, tmp_path):
b_np = np.zeros((n,), dtype="float32")
def run_and_check():
- dev = tvm.device(target, 0)
+ dev = tvm.opencl(0)
a = tvm.runtime.tensor(a_np, dev)
b = tvm.runtime.tensor(b_np, dev)
reloaded["main"](a, b)
diff --git a/tests/python/codegen/test_target_codegen_vulkan.py
b/tests/python/codegen/test_target_codegen_vulkan.py
index 42f12178bf..7ff86da9dc 100644
--- a/tests/python/codegen/test_target_codegen_vulkan.py
+++ b/tests/python/codegen/test_target_codegen_vulkan.py
@@ -108,7 +108,7 @@ def test_array_copy(target, dtype, fuzz_seed):
a_np = np.random.uniform(size=(arr_size,)).astype(dtype)
def run_and_check():
- dev = tvm.device(target)
+ dev = tvm.device_from_target(target)
a = tvm.runtime.empty((arr_size,), dtype, dev).copyfrom(a_np)
tvm.testing.assert_allclose(a_np, a.numpy())
@@ -145,7 +145,7 @@ def test_array_vectorize_add(dtype):
f = tvm.compile(Module, target=target)
def run_and_check():
- dev = tvm.device(target.kind.name)
+ dev = tvm.vulkan()
a = tvm.runtime.empty((arr_size,), vec_dtype, dev).copyfrom(
np.random.uniform(size=(arr_size, lanes))
)
@@ -183,7 +183,7 @@ def test_vulkan_bool_load():
ref = a_np.astype(np.int32)
def run_and_check():
- dev = tvm.device(target.kind.name)
+ dev = tvm.vulkan()
a = tvm.runtime.tensor(a_np, dev)
b = tvm.runtime.tensor(b_np, dev)
f(a, b)
@@ -260,7 +260,7 @@ def test_vulkan_constant_passing(vulkan_parameter_impl,
vulkan_parameter_dtype):
scalars = np.array([1 for _ in range(num_int_params)]).astype(dtype)
def run_and_check():
- dev = tvm.device(target.kind.name)
+ dev = tvm.vulkan()
a = tvm.runtime.tensor(np.random.uniform(size=n).astype(dtype), dev)
b = tvm.runtime.tensor(np.zeros(n, dtype=dtype), dev)
f_add(*scalars, a, b)
@@ -294,7 +294,7 @@ def test_vulkan_while_if():
compiled_func = tvm.compile(mod, target=target)
def run_and_check():
- dev = tvm.device(target.kind.name)
+ dev = tvm.vulkan()
for input_value, expected in [(5, [55]), (-5, [210])]:
a = tvm.runtime.tensor(np.array([input_value], dtype=dtype), dev)
b = tvm.runtime.tensor(np.zeros(n, dtype=dtype), dev)
@@ -332,7 +332,7 @@ def test_vulkan_local_threadidx():
b_np = np.zeros((n,), dtype="int32")
def run_and_check():
- dev = tvm.device(target.kind.name)
+ dev = tvm.vulkan()
a = tvm.runtime.tensor(a_np, dev)
b = tvm.runtime.tensor(b_np, dev)
func(a, b)
@@ -371,7 +371,7 @@ def test_vectorized_index_ramp():
b_np = np.zeros(n, dtype="int32")
def run_and_check():
- dev = tvm.device(target["kind"])
+ dev = tvm.vulkan()
a = tvm.runtime.tensor(a_np, dev)
b = tvm.runtime.tensor(b_np, dev)
f(a, b)
@@ -412,7 +412,7 @@ def test_vectorized_index_broadcast():
b_np = np.zeros(n, dtype="int32")
def run_and_check():
- dev = tvm.device(target["kind"])
+ dev = tvm.vulkan()
a = tvm.runtime.tensor(a_np, dev)
b = tvm.runtime.tensor(b_np, dev)
f(a, b)
@@ -457,7 +457,7 @@ def test_negative_operand_divmod():
built = tvm.compile(func, target=target)
def run_and_check():
- dev = tvm.device(target["kind"])
+ dev = tvm.vulkan()
a_dev = tvm.runtime.empty([N, 2], "int32", dev)
built(a_dev)
a = a_dev.numpy()
@@ -556,7 +556,7 @@ def test_cooperative_matrix(out_dtype):
f = tvm.compile(Module, target=target)
def run_and_check():
- dev = tvm.device("vulkan", 0)
+ dev = tvm.vulkan(0)
A = tvm.runtime.tensor(np.random.randn(M, K).astype("float16"),
dev)
B = tvm.runtime.tensor(np.random.randn(K, N).astype("float16"),
dev)
C = tvm.runtime.tensor(np.random.randn(M, N).astype(out_dtype),
dev)
@@ -656,7 +656,7 @@ def test_unary():
data = np.random.uniform(0.1, 0.9, size=n)
def run_and_check():
- dev = tvm.device(target.kind.name, 0)
+ dev = tvm.vulkan(0)
a = tvm.runtime.tensor(data.astype("float32"), dev)
b = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev)
func(a, b)
diff --git a/tests/python/relax/test_backend_dispatch_sort_scan.py
b/tests/python/relax/test_backend_dispatch_sort_scan.py
index 6b44c597a5..2cab14f36a 100644
--- a/tests/python/relax/test_backend_dispatch_sort_scan.py
+++ b/tests/python/relax/test_backend_dispatch_sort_scan.py
@@ -439,7 +439,7 @@ def test_dispatch_cumsum_gpu(target):
ex = tvm.compile(mod, target)
def run_and_check():
- dev = tvm.device(target["kind"] if isinstance(target, dict) else
target)
+ dev = tvm.device_from_target(target)
vm = tvm.relax.VirtualMachine(ex, dev)
tvm_data = tvm.runtime.tensor(np_data, dev)
cumsum = vm["main"](tvm_data)
diff --git a/tests/python/relax/test_codegen_cublas.py
b/tests/python/relax/test_codegen_cublas.py
index 08edb8ea6c..f273ebf8ef 100644
--- a/tests/python/relax/test_codegen_cublas.py
+++ b/tests/python/relax/test_codegen_cublas.py
@@ -59,7 +59,7 @@ def build_and_run(mod, inputs_np, target, legalize=False,
cuda_graph=False):
ex = tvm.compile(mod, target)
def run_and_check():
- dev = tvm.device(target, 0)
+ dev = tvm.device_from_target(target, 0)
vm = relax.VirtualMachine(ex, dev)
f = vm["main"]
inputs = [tvm.runtime.tensor(inp, dev) for inp in inputs_np]
diff --git a/tests/python/relax/test_codegen_cudnn.py
b/tests/python/relax/test_codegen_cudnn.py
index b44a9e15d1..419c805fdc 100644
--- a/tests/python/relax/test_codegen_cudnn.py
+++ b/tests/python/relax/test_codegen_cudnn.py
@@ -120,7 +120,7 @@ def build_and_run(mod, inputs_np, target, legalize=False,
cuda_graph=False):
ex = tvm.compile(mod, target)
def run_and_check():
- dev = tvm.device(target, 0)
+ dev = tvm.device_from_target(target, 0)
vm = relax.VirtualMachine(ex, dev)
f = vm["main"]
inputs = [tvm.runtime.tensor(inp, dev) for inp in inputs_np]
diff --git a/tests/python/relax/test_codegen_cutlass.py
b/tests/python/relax/test_codegen_cutlass.py
index 3fda98546e..0871b59373 100644
--- a/tests/python/relax/test_codegen_cutlass.py
+++ b/tests/python/relax/test_codegen_cutlass.py
@@ -99,7 +99,7 @@ def build_and_run(mod, inputs_np, target, legalize=True,
cuda_graph=False):
ex = tvm.compile(mod, target)
def run_and_check():
- dev = tvm.device(target, 0)
+ dev = tvm.device_from_target(target, 0)
vm = relax.VirtualMachine(ex, dev)
f = vm["main"]
inputs = [tvm.runtime.tensor(inp, dev) for inp in inputs_np]
@@ -1497,7 +1497,7 @@ def test_fp16A_int4B_gemm():
ex_cuda = tvm.compile(mod_deploy, target="cuda")
def run_and_check():
- dev = tvm.device("cuda", 0)
+ dev = tvm.cuda(0)
vm = relax.vm.VirtualMachine(ex_cuda, dev)
x_nd = tvm.runtime.tensor(x, dev)
residual_nd = tvm.runtime.tensor(residual, dev)
@@ -1655,7 +1655,7 @@ def test_fp16A_int8B_gemm():
return x * 0.5 * (1.0 + erf_out)
def run_and_check():
- dev = tvm.device("cuda", 0)
+ dev = tvm.cuda(0)
vm = relax.vm.VirtualMachine(ex_cuda, dev)
x_nd = tvm.runtime.tensor(x, dev)
inp = [x_nd, packed_weight.copyto(dev), scales.copyto(dev),
bias_trans.copyto(dev)]
@@ -1922,7 +1922,7 @@ def test_fp16A_int8B_gemm_batched():
ex_cuda = tvm.compile(mod_deploy, target="cuda")
def run_and_check():
- dev = tvm.device("cuda", 0)
+ dev = tvm.cuda(0)
vm = relax.vm.VirtualMachine(ex_cuda, dev)
x_nd = tvm.runtime.tensor(x, dev)
inp = [x_nd, packed_weight.copyto(dev), scales.copyto(dev)]
@@ -2080,7 +2080,7 @@ def test_fp16A_int8B_gemm_batched_finegrained():
ex_cuda = tvm.compile(mod_deploy, target="cuda")
def run_and_check():
- dev = tvm.device("cuda", 0)
+ dev = tvm.cuda(0)
vm = relax.vm.VirtualMachine(ex_cuda, dev)
x_nd = tvm.runtime.tensor(x, dev)
inp = [x_nd, packed_weight.copyto(dev), scales.copyto(dev)]
diff --git a/tests/python/relax/test_codegen_hipblas.py
b/tests/python/relax/test_codegen_hipblas.py
index b61578e172..77ed14fbac 100644
--- a/tests/python/relax/test_codegen_hipblas.py
+++ b/tests/python/relax/test_codegen_hipblas.py
@@ -51,7 +51,7 @@ def build_and_run(mod, inputs_np, target, legalize=False):
ex = tvm.compile(mod, target)
def run_and_check():
- dev = tvm.device(target, 0)
+ dev = tvm.device_from_target(target, 0)
vm = relax.VirtualMachine(ex, dev)
f = vm["main"]
inputs = [tvm.runtime.tensor(inp, dev) for inp in inputs_np]
diff --git a/tests/python/relax/test_codegen_tensorrt.py
b/tests/python/relax/test_codegen_tensorrt.py
index f6f93aca4d..73e3471b20 100644
--- a/tests/python/relax/test_codegen_tensorrt.py
+++ b/tests/python/relax/test_codegen_tensorrt.py
@@ -69,7 +69,7 @@ def build_and_run(mod, inputs_np, target, legalize=False):
ex = tvm.compile(mod, target)
def run_and_check():
- dev = tvm.device(target, 0)
+ dev = tvm.device_from_target(target, 0)
vm = relax.VirtualMachine(ex, dev)
f = vm["main"]
inputs = [tvm.runtime.tensor(inp, dev) for inp in inputs_np]
@@ -318,7 +318,7 @@ def test_tensorrt_int8_calibration(monkeypatch):
ex = tvm.compile(offloaded, "cuda")
def run_and_check():
- dev = tvm.device("cuda", 0)
+ dev = tvm.cuda(0)
vm = relax.VirtualMachine(ex, dev)
data_trt = tvm.runtime.tensor(data, dev)
out = None
diff --git a/tests/python/relax/test_contrib_vllm.py
b/tests/python/relax/test_contrib_vllm.py
index 607c1fde67..096efa6f11 100644
--- a/tests/python/relax/test_contrib_vllm.py
+++ b/tests/python/relax/test_contrib_vllm.py
@@ -50,7 +50,7 @@ def build_and_run(mod, inputs_np, target, legalize=True):
with tvm.transform.PassContext():
ex = tvm.compile(mod, target)
- dev = tvm.device(target, 0)
+ dev = tvm.cuda(0)
vm = relax.VirtualMachine(ex, dev)
f = vm["main"]
inputs = [tvm.runtime.tensor(inp, dev) for inp in inputs_np]
@@ -759,7 +759,7 @@ def test_reconstruct_from_cache():
reconstruct_from_cache_func =
tvm.get_global_func("tvm.contrib.vllm.reconstruct_from_cache")
def run_and_check():
- dev = tvm.device("cuda", 0)
+ dev = tvm.cuda(0)
key = tvm.runtime.tensor(
np.random.randn(num_tokens, num_heads,
head_dim).astype("float16"), dev
)
diff --git a/tests/python/relax/test_frontend_nn_llm_sequence_prefill_masked.py
b/tests/python/relax/test_frontend_nn_llm_sequence_prefill_masked.py
index af4f983742..b04d0a966c 100644
--- a/tests/python/relax/test_frontend_nn_llm_sequence_prefill_masked.py
+++ b/tests/python/relax/test_frontend_nn_llm_sequence_prefill_masked.py
@@ -193,7 +193,7 @@ def test_valid_len_zero(target):
"""All samples are fully padded: kernel must not crash and must stay
bounded."""
if not tvm.testing.device_enabled(target):
pytest.skip(f"{target} not enabled")
- dev = tvm.device(target)
+ dev = tvm.device_from_target(target)
_run_case(
target=target,
dev=dev,
@@ -216,7 +216,7 @@ def test_valid_len_full(target):
"""All samples are fully valid: must match a plain unmasked attention."""
if not tvm.testing.device_enabled(target):
pytest.skip(f"{target} not enabled")
- dev = tvm.device(target)
+ dev = tvm.device_from_target(target)
_run_case(
target=target,
dev=dev,
@@ -239,7 +239,7 @@ def test_valid_len_mixed(target):
"""Typical encoder batch with different valid lengths per sample."""
if not tvm.testing.device_enabled(target):
pytest.skip(f"{target} not enabled")
- dev = tvm.device(target)
+ dev = tvm.device_from_target(target)
_run_case(
target=target,
dev=dev,
@@ -262,7 +262,7 @@ def test_valid_len_mixed_gqa(target):
"""Grouped-query attention: ``group_size = h_q / h_kv > 1``."""
if not tvm.testing.device_enabled(target):
pytest.skip(f"{target} not enabled")
- dev = tvm.device(target)
+ dev = tvm.device_from_target(target)
_run_case(
target=target,
dev=dev,
@@ -285,7 +285,7 @@ def test_causal_padded_left_valid_len_zero(target):
"""Causal left-pad: all samples are fully padded."""
if not tvm.testing.device_enabled(target):
pytest.skip(f"{target} not enabled")
- dev = tvm.device(target)
+ dev = tvm.device_from_target(target)
_run_case(
target=target,
dev=dev,
@@ -309,7 +309,7 @@ def test_causal_padded_left_valid_len_full(target):
"""Causal left-pad: all samples are fully valid — degenerates to plain
causal attention."""
if not tvm.testing.device_enabled(target):
pytest.skip(f"{target} not enabled")
- dev = tvm.device(target)
+ dev = tvm.device_from_target(target)
_run_case(
target=target,
dev=dev,
@@ -333,7 +333,7 @@ def test_causal_padded_left_valid_len_mixed(target):
"""Causal left-pad: typical decoder-embedding batch with mixed lengths."""
if not tvm.testing.device_enabled(target):
pytest.skip(f"{target} not enabled")
- dev = tvm.device(target)
+ dev = tvm.device_from_target(target)
_run_case(
target=target,
dev=dev,
@@ -357,7 +357,7 @@ def test_causal_padded_left_valid_len_mixed_gqa(target):
"""Causal left-pad: grouped-query attention with mixed lengths."""
if not tvm.testing.device_enabled(target):
pytest.skip(f"{target} not enabled")
- dev = tvm.device(target)
+ dev = tvm.device_from_target(target)
_run_case(
target=target,
dev=dev,
@@ -381,7 +381,7 @@ def
test_causal_padded_left_qo_len_differs_from_kv_len(target):
"""Causal left-pad: Q and K/V may have different padded lengths."""
if not tvm.testing.device_enabled(target):
pytest.skip(f"{target} not enabled")
- dev = tvm.device(target)
+ dev = tvm.device_from_target(target)
_run_case(
target=target,
dev=dev,
diff --git a/tests/python/relax/test_frontend_nn_op.py
b/tests/python/relax/test_frontend_nn_op.py
index af2700ae9f..2877ef11c7 100644
--- a/tests/python/relax/test_frontend_nn_op.py
+++ b/tests/python/relax/test_frontend_nn_op.py
@@ -981,7 +981,7 @@ def test_multinomial_from_uniform():
mod = relax.backend.DispatchSampling()(mod)
mod = s_tir.transform.DefaultGPUSchedule()(mod)
ex = tvm.compile(mod, target)
- dev = tvm.device(target.kind.name, 0)
+ dev = tvm.cuda(0)
vm = relax.VirtualMachine(ex, dev)
effects = vm["_initialize_effect"]()
diff --git a/tests/python/relax/test_op_gradient_numeric.py
b/tests/python/relax/test_op_gradient_numeric.py
index bce591c798..879ffbe7e5 100644
--- a/tests/python/relax/test_op_gradient_numeric.py
+++ b/tests/python/relax/test_op_gradient_numeric.py
@@ -224,7 +224,7 @@ def relax_check_gradients(
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_unary(unary_op_func, can_be_neg):
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
(low, high) = (-1, 1) if can_be_neg else (0.1, 1)
data_numpy = np.random.uniform(low, high, (3, 3)).astype(np.float32)
relax_check_gradients(unary_op_func, [data_numpy], target, dev)
@@ -246,7 +246,7 @@ def test_unary(unary_op_func, can_be_neg):
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_binary_arith(binary_arith_op_func):
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data1_numpy = np.random.uniform(1, 2, (3, 3)).astype(np.float32)
data2_numpy = np.random.uniform(1, 2, (3, 3)).astype(np.float32)
relax_check_gradients(binary_arith_op_func, [data1_numpy, data2_numpy],
target, dev)
@@ -256,7 +256,7 @@ def test_binary_arith(binary_arith_op_func):
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_binary_minmax(binary_minmax_op_func):
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
# Checking numerical gradient of min and max requires data1_numpy[i] !=
data2_numpy[i]
# for all possible i.
# If data1_numpy[i] == data2_numpy[i], the operator is not differentiable
w.r.t. place i
@@ -281,7 +281,7 @@ def test_binary_minmax(binary_minmax_op_func):
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_binary_cmp(binary_cmp_op_func):
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data1_numpy = np.random.uniform(1, 2, (3, 3)).astype(np.float32)
data2_numpy = np.random.uniform(1, 2, (3, 3)).astype(np.float32)
relax_check_gradients(
@@ -296,7 +296,7 @@ def test_binary_cmp(binary_cmp_op_func):
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_ones_zeros_like(like_op_func):
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data_numpy = np.random.uniform(-1, 1, (3, 3)).astype(np.float32)
relax_check_gradients(like_op_func, [data_numpy], target, dev,
ignore_grads=[0])
@@ -304,7 +304,7 @@ def test_ones_zeros_like(like_op_func):
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_full_like():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data_numpy = np.random.uniform(-1, 1, (3, 3)).astype(np.float32)
fill_value = np.random.uniform(-1, 1, ()).astype(np.float32)
relax_check_gradients(
@@ -316,7 +316,7 @@ def test_full_like():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_ones_zeros(create_op_func):
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
relax_check_gradients(
create_op_func, [], target, dev, ignore_grads=[0], shape=(3, 3),
dtype="float32"
)
@@ -325,7 +325,7 @@ def test_ones_zeros(create_op_func):
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_triu():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data_numpy = np.random.uniform(-1, 1, (3, 3)).astype(np.float32)
relax_check_gradients(relax.op.triu, [data_numpy], target, dev, k=0)
@@ -336,7 +336,7 @@ def test_triu():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_sum():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data1_numpy = np.random.uniform(0, 16, (3, 3)).astype(np.float32)
relax_check_gradients(relax.op.sum, [data1_numpy], target, dev)
@@ -344,7 +344,7 @@ def test_sum():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_sum_with_axis():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data1_numpy = np.random.uniform(0, 16, (2, 3, 4, 5)).astype(np.float32)
relax_check_gradients(relax.op.sum, [data1_numpy], target, dev, axis=[1,
3])
@@ -352,7 +352,7 @@ def test_sum_with_axis():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_sum_keepdims():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data1_numpy = np.random.uniform(0, 16, (3, 3)).astype(np.float32)
relax_check_gradients(relax.op.sum, [data1_numpy], target, dev,
keepdims=True, axis=1)
@@ -360,7 +360,7 @@ def test_sum_keepdims():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_mean():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data1_numpy = np.random.uniform(0, 16, (3, 3)).astype(np.float32)
relax_check_gradients(relax.op.mean, [data1_numpy], target, dev)
@@ -368,7 +368,7 @@ def test_mean():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_mean_with_axis():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data1_numpy = np.random.uniform(0, 16, (2, 3, 4, 5)).astype(np.float32)
relax_check_gradients(relax.op.mean, [data1_numpy], target, dev, axis=[1,
3])
@@ -376,7 +376,7 @@ def test_mean_with_axis():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_mean_keepdims():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data1_numpy = np.random.uniform(0, 16, (3, 3)).astype(np.float32)
relax_check_gradients(relax.op.mean, [data1_numpy], target, dev,
keepdims=True, axis=1)
@@ -384,7 +384,7 @@ def test_mean_keepdims():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_variance():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data1_numpy = np.random.uniform(0, 16, (3, 3)).astype(np.float32)
relax_check_gradients(relax.op.variance, [data1_numpy], target, dev)
@@ -392,7 +392,7 @@ def test_variance():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_variance_with_axis():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data1_numpy = np.random.uniform(0, 16, (2, 3, 4, 5)).astype(np.float32)
relax_check_gradients(relax.op.variance, [data1_numpy], target, dev,
axis=[1, 3])
@@ -400,7 +400,7 @@ def test_variance_with_axis():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_variance_keepdims():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data1_numpy = np.random.uniform(0, 16, (3, 3)).astype(np.float32)
relax_check_gradients(relax.op.variance, [data1_numpy], target, dev,
keepdims=True, axis=1)
@@ -411,7 +411,7 @@ def test_variance_keepdims():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_reshape():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data_numpy = np.random.uniform(0, 16, (2, 3, 5)).astype(np.float32)
relax_check_gradients(
relax.op.reshape, [data_numpy], target, dev, ignore_grads=[1],
shape=(5, 6)
@@ -421,7 +421,7 @@ def test_reshape():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_reshape_infer_dim():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data_numpy = np.random.uniform(0, 16, (2, 3, 5)).astype(np.float32)
relax_check_gradients(
relax.op.reshape, [data_numpy], target, dev, ignore_grads=[1],
shape=(5, 2, 1, -1)
@@ -431,7 +431,7 @@ def test_reshape_infer_dim():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_permute_dims():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data_numpy = np.random.uniform(0, 16, (2, 3, 4, 5)).astype(np.float32)
relax_check_gradients(relax.op.permute_dims, [data_numpy], target, dev)
@@ -439,7 +439,7 @@ def test_permute_dims():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_permute_dims_with_axes():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data_numpy = np.random.uniform(0, 16, (2, 3, 4, 5)).astype(np.float32)
relax_check_gradients(
relax.op.permute_dims,
@@ -453,7 +453,7 @@ def test_permute_dims_with_axes():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_concat():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data_numpy1 = np.random.uniform(1, 16, (3, 3)).astype(np.float32)
data_numpy2 = np.random.uniform(1, 16, (3, 4)).astype(np.float32)
data_numpy3 = np.random.uniform(1, 16, (3, 5)).astype(np.float32)
@@ -470,7 +470,7 @@ def test_concat():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_split_indices():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data_numpy = np.random.uniform(1, 16, (3, 12)).astype(np.float32)
relax_check_gradients(
relax.op.split,
@@ -485,7 +485,7 @@ def test_split_indices():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_split_section():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data_numpy = np.random.uniform(1, 16, (3, 12)).astype(np.float32)
relax_check_gradients(
relax.op.split,
@@ -500,7 +500,7 @@ def test_split_section():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_reshape():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data_numpy = np.random.uniform(1, 16, (3, 4)).astype(np.float32)
relax_check_gradients(
@@ -516,7 +516,7 @@ def test_reshape():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_cumsum():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data_numpy1 = np.random.uniform(1, 16, (3, 3)).astype(np.float32)
relax_check_gradients(
relax.op.cumsum,
@@ -530,7 +530,7 @@ def test_cumsum():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_cumsum_no_axis():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data_numpy1 = np.random.uniform(1, 16, (3, 3)).astype(np.float32)
relax_check_gradients(
relax.op.cumsum,
@@ -543,7 +543,7 @@ def test_cumsum_no_axis():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_expand_dims():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data_numpy = np.random.uniform(1, 16, (3, 12)).astype(np.float32)
relax_check_gradients(relax.op.expand_dims, [data_numpy], target, dev,
axis=1)
@@ -551,7 +551,7 @@ def test_expand_dims():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_expand_dims_list():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data_numpy = np.random.uniform(1, 16, (3, 12)).astype(np.float32)
relax_check_gradients(relax.op.expand_dims, [data_numpy], target, dev,
axis=(0, 2, 3))
@@ -559,7 +559,7 @@ def test_expand_dims_list():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_broadcast_to():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data_numpy = np.random.uniform(1, 16, (3, 4)).astype(np.float32)
relax_check_gradients(
relax.op.broadcast_to,
@@ -577,7 +577,7 @@ def test_broadcast_to():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_take():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data_numpy = np.random.uniform(0, 16, size=(2, 3, 4)).astype(np.float32)
indices = np.array([0, 1])
relax_check_gradients(
@@ -593,7 +593,7 @@ def test_take():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_take_no_axis():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data_numpy = np.random.uniform(0, 16, size=(5,)).astype(np.float32)
indices = np.array([1, 3])
relax_check_gradients(
@@ -611,7 +611,7 @@ def test_take_no_axis():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_where():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data1_numpy = np.random.uniform(0, 1, size=(3, 3)) > 0.5
data2_numpy = np.random.uniform(0, 16, size=(3, 3)).astype(np.float32)
data3_numpy = np.random.uniform(0, 16, size=(3, 3)).astype(np.float32)
@@ -631,7 +631,7 @@ def test_where():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_matmul_2_2():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data1_numpy = np.random.uniform(0, 16, (2, 3)).astype(np.float32)
data2_numpy = np.random.uniform(0, 16, (3, 4)).astype(np.float32)
relax_check_gradients(relax.op.matmul, [data1_numpy, data2_numpy], target,
dev)
@@ -640,7 +640,7 @@ def test_matmul_2_2():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_matmul_1_1():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data1_numpy = np.random.uniform(0, 16, (4,)).astype(np.float32)
data2_numpy = np.random.uniform(0, 16, (4,)).astype(np.float32)
relax_check_gradients(relax.op.matmul, [data1_numpy, data2_numpy], target,
dev)
@@ -649,7 +649,7 @@ def test_matmul_1_1():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_matmul_1_4():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data1_numpy = np.random.uniform(0, 16, (4,)).astype(np.float32)
data2_numpy = np.random.uniform(0, 16, (2, 3, 4, 5)).astype(np.float32)
relax_check_gradients(relax.op.matmul, [data1_numpy, data2_numpy], target,
dev)
@@ -658,7 +658,7 @@ def test_matmul_1_4():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_matmul_4_1():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data1_numpy = np.random.uniform(0, 16, (2, 3, 4, 5)).astype(np.float32)
data2_numpy = np.random.uniform(0, 16, (5,)).astype(np.float32)
relax_check_gradients(relax.op.matmul, [data1_numpy, data2_numpy], target,
dev)
@@ -667,7 +667,7 @@ def test_matmul_4_1():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_matmul_5_4():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data1_numpy = np.random.uniform(0, 16, (2, 3, 1, 4, 5)).astype(np.float32)
data2_numpy = np.random.uniform(0, 16, (3, 2, 5, 4)).astype(np.float32)
relax_check_gradients(
@@ -684,7 +684,7 @@ def test_matmul_5_4():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_astype():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data_numpy = np.random.uniform(0, 16, size=(3, 3)).astype(np.float64)
relax_check_gradients(relax.op.astype, [data_numpy], target, dev,
dtype="float32")
@@ -695,7 +695,7 @@ def test_astype():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_relu():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data1_numpy = np.random.uniform(0.2, 1, (3, 3)).astype(np.float32)
sign = np.random.randint(0, 2, (3, 3)).astype(np.float32) * 2 - 1
data1_numpy *= sign
@@ -705,7 +705,7 @@ def test_relu():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_silu():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data1_numpy = np.random.uniform(0, 16, (3, 3)).astype(np.float32)
relax_check_gradients(relax.op.nn.silu, [data1_numpy], target, dev)
@@ -713,7 +713,7 @@ def test_silu():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_softmax():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data1_numpy = np.random.uniform(0, 16, (3, 3)).astype(np.float32)
relax_check_gradients(relax.op.nn.softmax, [data1_numpy], target, dev)
@@ -721,7 +721,7 @@ def test_softmax():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_softmax_with_axis():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data1_numpy = np.random.uniform(0, 16, (3, 3)).astype(np.float32)
relax_check_gradients(relax.op.nn.softmax, [data1_numpy], target, dev,
axis=1)
@@ -729,7 +729,7 @@ def test_softmax_with_axis():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_log_softmax():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data1_numpy = np.random.uniform(0, 16, (3, 3)).astype(np.float32)
relax_check_gradients(relax.op.nn.log_softmax, [data1_numpy], target, dev)
@@ -737,7 +737,7 @@ def test_log_softmax():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_log_softmax_with_axis():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data1_numpy = np.random.uniform(0, 16, (3, 3)).astype(np.float32)
relax_check_gradients(relax.op.nn.log_softmax, [data1_numpy], target, dev,
axis=1)
@@ -745,7 +745,7 @@ def test_log_softmax_with_axis():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_cross_entropy_with_logits():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data_numpy1 = np.random.uniform(1, 16, (3,)).astype(np.float32)
data_numpy2 = np.random.uniform(1, 16, (3,)).astype(np.float32)
relax_check_gradients(
@@ -759,7 +759,7 @@ def test_cross_entropy_with_logits():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_cross_entropy_with_logits_batch():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data_numpy1 = np.random.uniform(1, 16, (2, 3)).astype(np.float32)
data_numpy2 = np.random.uniform(1, 16, (2, 3)).astype(np.float32)
relax_check_gradients(
@@ -784,7 +784,7 @@ def test_cross_entropy_with_logits_batch():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_nll_loss(nll_reduction, nll_weighted, nll_ignore_index):
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data1_numpy = np.random.uniform(0, 16, (2, 3, 4)).astype(np.float32)
data2_numpy = np.random.randint(0, 3, (2, 4)).astype(np.int64)
# force a position in targets it not ignore_index, to avoid zero total
weight
@@ -817,7 +817,7 @@ def test_nll_loss(nll_reduction, nll_weighted,
nll_ignore_index):
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_nll_loss_no_batch(nll_reduction1, nll_weighted1, nll_ignore_index1):
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data1_numpy = np.random.uniform(0, 16, (3,)).astype(np.float32)
data2_numpy = np.random.randint(0, 3, ()).astype(np.int64)
# weight > 0
@@ -875,7 +875,7 @@ def test_nll_loss_no_batch(nll_reduction1, nll_weighted1,
nll_ignore_index1):
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_conv2d(c2d_shape1, c2d_shape2, c2d_kwargs):
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
import pytest
# Use smaller range to reduce numerical errors in gradient check
@@ -916,7 +916,7 @@ pool_params = [
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_max_pool2d(pool_size, pool_kwargs):
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data_numpy = np.random.uniform(0, 3, size=(3, 2, 10,
10)).astype(np.float32)
relax_check_gradients(
relax.op.nn.max_pool2d,
@@ -932,7 +932,7 @@ def test_max_pool2d(pool_size, pool_kwargs):
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_avg_pool2d(pool_size, pool_kwargs):
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
data_numpy = np.random.uniform(0, 3, size=(3, 2, 10,
10)).astype(np.float32)
relax_check_gradients(
relax.op.nn.avg_pool2d,
diff --git a/tests/python/relax/test_op_take.py
b/tests/python/relax/test_op_take.py
index 18006e71fb..eab0d14836 100644
--- a/tests/python/relax/test_op_take.py
+++ b/tests/python/relax/test_op_take.py
@@ -36,7 +36,7 @@ def test_take_scalar_tensor_as_index(axis):
"""
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
@I.ir_module
class Module:
@@ -64,7 +64,7 @@ def test_take_1d_tensor_as_index(axis):
`data.ndim + indices.ndim - 1`.
"""
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
@I.ir_module
class Module:
@@ -88,7 +88,7 @@ def test_take_1d_tensor_as_index(axis):
def test_take_2d_tensor_as_index(axis):
"""The index of R.take may be a 2-d tensor"""
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
@I.ir_module
class Module:
@@ -117,7 +117,7 @@ def test_take_constant_prim_value_as_index(axis):
"""
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
@I.ir_module
class Module:
@@ -146,7 +146,7 @@ def test_take_dynamic_prim_value_as_index(axis):
"""
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
@I.ir_module
class Module:
@@ -173,7 +173,7 @@ def test_take_nan_mode_OOB_indices(axis):
This test checks that out-of-bounds indices produce NaN values in the
output tensor.
"""
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
@I.ir_module
class Module:
@@ -208,7 +208,7 @@ def test_take_wrap_mode_OOB_indices(axis):
This test checks that out-of-bounds indices wrap around to the valid range.
"""
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
@I.ir_module
class Module:
@@ -234,7 +234,7 @@ def test_take_clip_mode_OOB_indices(axis):
This test checks that out-of-bounds indices are clipped to the valid range.
"""
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
@I.ir_module
class Module:
diff --git a/tests/python/relax/test_op_view.py
b/tests/python/relax/test_op_view.py
index 9bf55ece1b..dc7e658ab9 100644
--- a/tests/python/relax/test_op_view.py
+++ b/tests/python/relax/test_op_view.py
@@ -682,7 +682,7 @@ def test_execute_no_op_view(target):
np_expected = np_input
def run_and_check():
- dev = tvm.device(target)
+ dev = tvm.device_from_target(target)
vm = tvm.relax.VirtualMachine(built, device=dev)
tvm_input = tvm.runtime.tensor(np_input, dev)
tvm_output = vm["main"](tvm_input)
@@ -711,7 +711,7 @@ def test_execute_view_with_new_shape(target):
np_expected = np_input.reshape(64, 64)
def run_and_check():
- dev = tvm.device(target)
+ dev = tvm.device_from_target(target)
vm = tvm.relax.VirtualMachine(built, device=dev)
tvm_input = tvm.runtime.tensor(np_input, dev)
tvm_output = vm["main"](tvm_input)
@@ -744,7 +744,7 @@ def test_execute_view_with_new_byte_offset(target):
np_expected = np_input.reshape(64, 64)[32:48, :]
def run_and_check():
- dev = tvm.device(target)
+ dev = tvm.device_from_target(target)
vm = tvm.relax.VirtualMachine(built, device=dev)
tvm_input = tvm.runtime.tensor(np_input, dev)
tvm_output = vm["main"](tvm_input)
@@ -773,7 +773,7 @@ def test_execute_view_with_new_dtype(target):
np_expected = np_input.view("uint32")
def run_and_check():
- dev = tvm.device(target)
+ dev = tvm.device_from_target(target)
vm = tvm.relax.VirtualMachine(built, device=dev)
tvm_input = tvm.runtime.tensor(np_input, dev)
tvm_output = vm["main"](tvm_input)
@@ -815,7 +815,7 @@ def test_execute_view_with_multiple_updated_fields(target):
]
def run_and_check():
- dev = tvm.device(target)
+ dev = tvm.device_from_target(target)
vm = tvm.relax.VirtualMachine(built, device=dev)
tvm_input = tvm.runtime.tensor(np_input, dev)
tvm_output = vm["main"](tvm_input)
diff --git a/tests/python/relax/test_training_optimizer_numeric.py
b/tests/python/relax/test_training_optimizer_numeric.py
index fc5264b796..1b3c9a6fbf 100644
--- a/tests/python/relax/test_training_optimizer_numeric.py
+++ b/tests/python/relax/test_training_optimizer_numeric.py
@@ -89,7 +89,7 @@ def _test_optimizer(target, dev, np_func, opt_type, *args,
**kwargs):
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_sgd(lr, weight_decay):
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
def np_func(param_tuple, grad_tuple, state_tuple):
num_steps = state_tuple[0]
@@ -115,7 +115,7 @@ def test_sgd(lr, weight_decay):
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_momentum_sgd(lr, momentum, dampening, weight_decay, nesterov):
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
def np_func(param_tuple, grad_tuple, state_tuple):
num_steps = state_tuple[0]
@@ -152,7 +152,7 @@ def test_momentum_sgd(lr, momentum, dampening,
weight_decay, nesterov):
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_adam(lr, betas, eps, weight_decay):
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
def np_func(param_tuple, grad_tuple, state_tuple):
num_steps = state_tuple[0]
diff --git a/tests/python/relax/test_training_trainer_numeric.py
b/tests/python/relax/test_training_trainer_numeric.py
index 1164e704f0..a8ec4740e2 100644
--- a/tests/python/relax/test_training_trainer_numeric.py
+++ b/tests/python/relax/test_training_trainer_numeric.py
@@ -56,7 +56,7 @@ def _make_dataset():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_execute():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
backbone = _get_backbone()
pred_ty = relax.TensorType((1, 5), "float32")
@@ -82,7 +82,7 @@ def test_execute():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_execute_numeric():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
backbone = _get_backbone()
pred_ty = relax.TensorType((1, 5), "float32")
@@ -113,7 +113,7 @@ def test_execute_numeric():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_load_export_params():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
backbone = _get_backbone()
pred_ty = relax.TensorType((1, 5), "float32")
@@ -150,7 +150,7 @@ def test_load_export_params():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_setting_error():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
backbone = _get_backbone()
pred_ty = relax.TensorType((1, 5), "float32")
diff --git a/tests/python/relax/test_transform_gradient_numeric.py
b/tests/python/relax/test_transform_gradient_numeric.py
index f3d9b357b2..beb57cd8b2 100644
--- a/tests/python/relax/test_transform_gradient_numeric.py
+++ b/tests/python/relax/test_transform_gradient_numeric.py
@@ -40,7 +40,7 @@ def _legalize_and_build(mod, target, dev):
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_manual_gradient():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
# The expression computed is sum((2x - 2y) * (y + z))
# the gradient of x is broadcast_to(2y + 2z, x.shape)
@@ -89,7 +89,7 @@ def test_manual_gradient():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_mlp_blockbuilder():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
layers, in_size, out_size, hidden_size, batch_size = 3, 5, 5, 5, 4
input_list = [relax.Var("x", R.Tensor((batch_size, in_size), "float32"))]
@@ -147,7 +147,7 @@ def test_mlp_blockbuilder():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_complex():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
cst = relax.const(np.ones((6,)), dtype="float32")
cst1 = relax.const(np.array(3), dtype="int64")
@@ -205,7 +205,7 @@ def test_complex():
@pytest.mark.skipif(not tvm.testing.device_enabled("llvm"), reason="llvm not
enabled")
def test_matmul():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
@tvm.script.ir_module
class Before:
diff --git a/tests/python/relax/test_transform_lazy_transform_params.py
b/tests/python/relax/test_transform_lazy_transform_params.py
index 23b5d61c70..0ff259b948 100644
--- a/tests/python/relax/test_transform_lazy_transform_params.py
+++ b/tests/python/relax/test_transform_lazy_transform_params.py
@@ -617,7 +617,7 @@ def test_output_with_use_site():
def test_output():
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
@I.ir_module(s_tir=True)
class TransformModule:
diff --git a/tests/python/relax/test_vm_build.py
b/tests/python/relax/test_vm_build.py
index 5156202726..ad8f94e8a5 100644
--- a/tests/python/relax/test_vm_build.py
+++ b/tests/python/relax/test_vm_build.py
@@ -1215,7 +1215,7 @@ def test_relax_module_with_multiple_targets(exec_mode):
np_A = np.random.random([32, 32]).astype("float32")
np_B = np.random.random([32, 32]).astype("float32")
- dev_llvm = tvm.device("llvm")
+ dev_llvm = tvm.cpu()
vm_llvm = tvm.relax.VirtualMachine(built, device=dev_llvm)
llvm_output = vm_llvm["func_llvm"](
tvm.runtime.tensor(np_A, dev_llvm),
@@ -1227,7 +1227,7 @@ def test_relax_module_with_multiple_targets(exec_mode):
tvm.testing.assert_allclose(llvm_output.numpy(), np_C)
def run_and_check():
- dev_cuda = tvm.device("cuda")
+ dev_cuda = tvm.cuda()
vm_cuda = tvm.relax.VirtualMachine(built, device=dev_cuda)
cuda_output = vm_cuda["func_cuda"](
tvm.runtime.tensor(np_A, dev_cuda),
diff --git a/tests/python/relax/test_vm_builtin.py
b/tests/python/relax/test_vm_builtin.py
index c3161e7855..b05e474d69 100644
--- a/tests/python/relax/test_vm_builtin.py
+++ b/tests/python/relax/test_vm_builtin.py
@@ -79,7 +79,7 @@ def test_alloc_tensor_raises_out_of_memory():
built = tvm.compile(Module, target=target)
def run_and_check():
- dev = tvm.device(target)
+ dev = tvm.cuda()
vm = relax.VirtualMachine(built, dev)
with pytest.raises(Exception, match="CUDA.*out of memory"):
vm["main"]()
diff --git a/tests/python/relax/test_vm_callback_function.py
b/tests/python/relax/test_vm_callback_function.py
index 78d314abc5..ef16d8fe58 100644
--- a/tests/python/relax/test_vm_callback_function.py
+++ b/tests/python/relax/test_vm_callback_function.py
@@ -30,7 +30,7 @@ pytestmark = pytest.mark.skipif(not
tvm.testing.device_enabled("llvm"), reason="
def test_pass_tensor_to_function(exec_mode):
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
@R.function
def relax_func(
@@ -65,7 +65,7 @@ def test_pass_tensor_to_function(exec_mode):
def test_generate_tensor_in_function(exec_mode):
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
@R.function
def relax_func(
@@ -94,7 +94,7 @@ def test_generate_tensor_in_function(exec_mode):
def test_catch_exception_with_full_stack_trace(exec_mode):
target = "llvm"
- dev = tvm.device(target)
+ dev = tvm.cpu()
@R.function
def relax_func(
diff --git a/tests/python/runtime/test_runtime_dlpack.py
b/tests/python/runtime/test_runtime_dlpack.py
index b1fcc83dcf..53ebccb823 100644
--- a/tests/python/runtime/test_runtime_dlpack.py
+++ b/tests/python/runtime/test_runtime_dlpack.py
@@ -42,7 +42,7 @@ def test_from_dlpack_shape_one():
fadd = tvm.compile(te.create_prim_func([A, B, C]), target=tgt)
- dev = tvm.device(tgt.kind.name, 0)
+ dev = tvm.cpu()
b = tvm.runtime.tensor(np.random.uniform(size=(rows, 16)).astype(B.dtype),
dev)
c = tvm.runtime.tensor(np.zeros((rows, 16), dtype=C.dtype), dev)
diff --git
a/tests/python/s_tir/meta_schedule/test_meta_schedule_mma_tensorize.py
b/tests/python/s_tir/meta_schedule/test_meta_schedule_mma_tensorize.py
index de24f0f6cc..efa260d298 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_mma_tensorize.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_mma_tensorize.py
@@ -82,7 +82,7 @@ def test_run_target(mod=None, tgt_str=None,
in_dtype="float16", out_dtype="float
f = lib["main"]
def run_and_check():
- dev = tvm.device(tgt_str, 0)
+ dev = tvm.cuda(0)
a = tvm.runtime.tensor(a_np, dev)
b = tvm.runtime.tensor(b_np, dev)
c = tvm.runtime.tensor(c_np, dev)
diff --git
a/tests/python/s_tir/schedule/test_tir_schedule_tensorize_ldmatrix_mma_numeric.py
b/tests/python/s_tir/schedule/test_tir_schedule_tensorize_ldmatrix_mma_numeric.py
index c7fa9ee40a..6b0c521f18 100644
---
a/tests/python/s_tir/schedule/test_tir_schedule_tensorize_ldmatrix_mma_numeric.py
+++
b/tests/python/s_tir/schedule/test_tir_schedule_tensorize_ldmatrix_mma_numeric.py
@@ -169,7 +169,7 @@ def run_test(
c_np = np.dot(a_np.astype("float32"),
b_np.astype("float32")).astype("int32")
def run_and_check(measure=False):
- dev = tvm.device("cuda", 0)
+ dev = tvm.cuda(0)
a = tvm.runtime.tensor(a_np, dev)
b = tvm.runtime.tensor(b_np, dev)
c = tvm.runtime.tensor(np.zeros((M, N), dtype=out_dtype), dev)
diff --git
a/tests/python/s_tir/schedule/test_tir_schedule_tensorize_mfma_numeric.py
b/tests/python/s_tir/schedule/test_tir_schedule_tensorize_mfma_numeric.py
index 537c1c3919..969141f8ba 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_tensorize_mfma_numeric.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_tensorize_mfma_numeric.py
@@ -147,7 +147,7 @@ def run_test(
c_np = np.dot(a_np.astype("float32"),
b_np.astype("float32")).astype("int32")
def run_and_check(measure=False):
- dev = tvm.device("rocm", 0)
+ dev = tvm.rocm(0)
a = tvm.runtime.tensor(a_np, dev)
b = tvm.runtime.tensor(b_np, dev)
c = tvm.runtime.tensor(np.zeros((M, N), dtype=out_dtype), dev)
diff --git
a/tests/python/s_tir/transform/test_s_tir_transform_inject_software_pipeline.py
b/tests/python/s_tir/transform/test_s_tir_transform_inject_software_pipeline.py
index 4df1bd14b7..56633f3eb1 100644
---
a/tests/python/s_tir/transform/test_s_tir_transform_inject_software_pipeline.py
+++
b/tests/python/s_tir/transform/test_s_tir_transform_inject_software_pipeline.py
@@ -1542,7 +1542,7 @@ def build_and_run(sch):
c_np = np.dot(a_np.astype("float32"), b_np.astype("float32"))
def run_and_check():
- dev = tvm.device("cuda", 0)
+ dev = tvm.cuda(0)
a = tvm.runtime.tensor(a_np, dev)
b = tvm.runtime.tensor(b_np, dev)
c = tvm.runtime.tensor(np.zeros((N, M), dtype="float32"), dev)
diff --git a/tests/python/target/test_target_target.py
b/tests/python/target/test_target_target.py
index f6dd1be8de..4421a57ce5 100644
--- a/tests/python/target/test_target_target.py
+++ b/tests/python/target/test_target_target.py
@@ -27,19 +27,41 @@ from tvm.testing import env
def test_all_targets_device_type_verify():
"""Consistency verification for all targets' device type"""
- target_kind_set = set(tvm.target.Target.list_kinds())
- target_kind_set.remove("composite")
- all_targets = [tvm.target.Target(t) for t in target_kind_set]
-
- for tgt in all_targets:
- if tgt.kind.name not in tvm.runtime.Device._DEVICE_NAME_TO_TYPE:
- raise KeyError(
- f"Cannot find target kind: {tgt.kind.name} in
Device._DEVICE_NAME_TO_TYPE"
- )
-
- assert (
- tgt.get_target_device_type() ==
tvm.runtime.Device._DEVICE_NAME_TO_TYPE[tgt.kind.name]
- )
+ for target_kind in tvm.target.Target.list_kinds():
+ target = Target(target_kind)
+ device = tvm.device_from_target(target)
+
+ assert device.dlpack_device_type() == target.get_target_device_type()
+
+
[email protected]("target", ["llvm", {"kind": "llvm"}, Target("llvm")])
+def test_device_from_target_input_forms(target):
+ device = tvm.device_from_target(target)
+
+ assert device == tvm.cpu()
+ assert isinstance(device, tvm.runtime.Device)
+ assert tvm.runtime.device_from_target(target) == tvm.cpu()
+
+
+def test_device_from_target_compiler_only_kind():
+ assert tvm.device_from_target("composite") == tvm.cpu()
+
+
+def test_device_from_target_index():
+ assert tvm.device_from_target("llvm").index == 0
+ assert tvm.device_from_target("llvm", None).index == 0
+ assert tvm.device_from_target("llvm", 3).index == 3
+
+
+def test_device_from_target_override():
+ target = Target(
+ {
+ "kind": "llvm",
+ "target_device_type": int(tvm_ffi.DLDeviceType.kDLCUDA),
+ }
+ )
+
+ assert tvm.device_from_target(target).dlpack_device_type() ==
tvm_ffi.DLDeviceType.kDLCUDA
def test_target_string_parse():
diff --git a/tests/python/target/test_virtual_device.py
b/tests/python/target/test_virtual_device.py
index d1fd97b0f2..b72023b811 100644
--- a/tests/python/target/test_virtual_device.py
+++ b/tests/python/target/test_virtual_device.py
@@ -22,7 +22,7 @@ import tvm.testing
def test_make_virtual_device_for_device():
- virtual_device = tvm.target.VirtualDevice(tvm.device("cuda"))
+ virtual_device = tvm.target.VirtualDevice(tvm.cuda())
assert virtual_device.dlpack_device_type() == 2
# ie kDLCUDA
assert virtual_device.virtual_device_id == 0
@@ -32,7 +32,7 @@ def test_make_virtual_device_for_device():
def test_make_virtual_device_for_device_and_target():
target = tvm.target.Target("cuda")
- virtual_device = tvm.target.VirtualDevice(tvm.device("cuda"), target)
+ virtual_device = tvm.target.VirtualDevice(tvm.cuda(), target)
assert virtual_device.dlpack_device_type() == 2 # ie kDLCUDA
assert virtual_device.target == target
assert virtual_device.memory_scope == ""
@@ -41,7 +41,7 @@ def test_make_virtual_device_for_device_and_target():
def test_make_virtual_device_for_device_target_and_memory_scope():
target = tvm.target.Target("cuda")
scope = "local"
- virtual_device = tvm.target.VirtualDevice(tvm.device("cuda"), target,
scope)
+ virtual_device = tvm.target.VirtualDevice(tvm.cuda(), target, scope)
assert virtual_device.dlpack_device_type() == 2 # ie kDLCUDA
assert virtual_device.target == target
assert virtual_device.memory_scope == scope
diff --git a/tests/python/tirx-base/test_tir_intrin.py
b/tests/python/tirx-base/test_tir_intrin.py
index 86b84d94c5..4765bd2f94 100644
--- a/tests/python/tirx-base/test_tir_intrin.py
+++ b/tests/python/tirx-base/test_tir_intrin.py
@@ -318,7 +318,7 @@ def test_clz(target, dtype):
func = tvm.compile(sch.mod, target=target)
def run_and_check():
- dev = tvm.device(target.kind.name)
+ dev = tvm.device_from_target(target)
n = 10
highs = [10, 100, 1000, 10000, 100000, 1000000]
diff --git a/tests/python/tirx/codegen/test_codegen_ampere.py
b/tests/python/tirx/codegen/test_codegen_ampere.py
index ae3ab17f59..279ca71b2e 100644
--- a/tests/python/tirx/codegen/test_codegen_ampere.py
+++ b/tests/python/tirx/codegen/test_codegen_ampere.py
@@ -63,7 +63,7 @@ def _run_mma(mod, K, no_c_ptr, np_in):
ref = ref + C_np
def run_and_check():
- dev = tvm.device("cuda")
+ dev = tvm.cuda()
D = tvm.runtime.tensor(np.zeros((16, 8), np.float32), device=dev)
A = tvm.runtime.tensor(A_np, device=dev)
B = tvm.runtime.tensor(B_np, device=dev)
diff --git a/tests/python/tirx/codegen/test_codegen_cuda.py
b/tests/python/tirx/codegen/test_codegen_cuda.py
index 6358867ed2..f5fc621c69 100644
--- a/tests/python/tirx/codegen/test_codegen_cuda.py
+++ b/tests/python/tirx/codegen/test_codegen_cuda.py
@@ -139,7 +139,7 @@ def test_cuda_atomic_add():
B_np = np.zeros(1, dtype="float32")
def run_and_check():
- dev = tvm.device("cuda")
+ dev = tvm.cuda()
A_tvm = tvm.runtime.tensor(A_np, device=dev)
B_tvm = tvm.runtime.tensor(B_np, device=dev)
mod["main"](A_tvm, B_tvm)
@@ -479,7 +479,7 @@ __device__ int32_t add_one(int32_t a) {
B = np.zeros((16, 16), dtype="int32")
def run_and_check():
- dev = tvm.device("cuda")
+ dev = tvm.cuda()
A_tvm = tvm.runtime.tensor(A, device=dev)
B_tvm = tvm.runtime.tensor(B, device=dev)
mod["main"](A_tvm, B_tvm)
@@ -510,7 +510,7 @@ __device__ void print(int32_t a) {
A = np.random.randint(0, 10, (16, 16)).astype("int32")
def run_and_check():
- dev = tvm.device("cuda")
+ dev = tvm.cuda()
A_tvm = tvm.runtime.tensor(A, device=dev)
mod["main"](A_tvm)
dev.sync()
@@ -602,7 +602,7 @@ def test_ptx_cp_async(cp_size, cache_hint, prefetch_size,
predicate, fill_mode):
A_ref = np.ones(N, dtype="float16") * 6
def run_and_check():
- dev = tvm.device("cuda")
+ dev = tvm.cuda()
A = tvm.runtime.tensor(A_np, device=dev)
mod(A)
np.testing.assert_allclose(A.numpy(), A_ref)
@@ -678,7 +678,7 @@ def test_ptx_ldmatrix(trans, num):
B_ref[8:16, 8:16] = A_np[8:16, 8:16] if not trans else A_np[8:16,
8:16].T
def run_and_check():
- dev = tvm.device("cuda")
+ dev = tvm.cuda()
A = tvm.runtime.tensor(A_np, device=dev)
B = tvm.runtime.tensor(B_np, device=dev)
mod(A, B)
