tqchen commented on code in PR #20005:
URL: https://github.com/apache/tvm/pull/20005#discussion_r3578335085
##########
python/tvm/tirx/build.py:
##########
@@ -18,6 +18,8 @@
# pylint: disable=invalid-name
"""The build utils in python."""
+from tvm_ffi import DLDeviceType
Review Comment:
Thanks for checking. `tvm_ffi` is the canonical package in the current tree,
and `DLDeviceType` is imported from it elsewhere in the runtime; `tvm._ffi`
does not expose this enum. No change is needed here.
##########
tests/python/target/test_target_target.py:
##########
@@ -27,19 +27,41 @@
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")])
Review Comment:
`Target` is already imported in this module with `from tvm.target import
Target`, so the decorator is evaluated with the name defined. No change is
needed here.
##########
tests/python/target/test_target_target.py:
##########
@@ -27,19 +27,41 @@
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
Review Comment:
Both names are imported at module scope: `Target` from `tvm.target` and the
canonical `tvm_ffi` package. The proposed `tvm._ffi` spelling is not the
current API, so no change is needed here.
##########
python/tvm/runtime/_tensor.py:
##########
@@ -353,6 +353,27 @@ 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
+
+ return device(Target(target).get_target_device_type(), index)
Review Comment:
Good suggestion. Commit `bf539dc094` now reuses an existing `Target` and
normalizes only string/dict inputs before querying the device type.
##########
tests/python/target/test_target_target.py:
##########
@@ -27,19 +27,41 @@
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)
Review Comment:
`Target` is already imported explicitly from `tvm.target` in this module, so
this use is defined and consistent with the neighboring tests. No change is
needed here.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
