This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch unity in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/unity by this push: new 2226a1f558 [Unity] Multi-device support for Relax (#15447) 2226a1f558 is described below commit 2226a1f5581cb77e599ce16d61227ed1088ff058 Author: Yong Wu <yongc...@gmail.com> AuthorDate: Fri Aug 11 12:25:29 2023 -0700 [Unity] Multi-device support for Relax (#15447) --- include/tvm/ir/global_info.h | 50 ++++++++++- include/tvm/relax/attrs/op.h | 8 ++ include/tvm/relax/struct_info.h | 17 +++- include/tvm/relay/transform.h | 12 +++ include/tvm/script/printer/ir_docsifier.h | 7 ++ include/tvm/target/target.h | 10 --- include/tvm/target/target_kind.h | 27 +----- include/tvm/tir/usmp/utils.h | 1 + python/tvm/ir/__init__.py | 2 +- python/tvm/ir/global_info.py | 12 +++ python/tvm/ir/json_compact.py | 17 ++++ python/tvm/relax/op/base.py | 21 +++++ python/tvm/relax/struct_info.py | 14 ++-- python/tvm/script/ir_builder/ir/__init__.py | 2 + python/tvm/script/ir_builder/ir/ir.py | 38 ++++++++- python/tvm/script/ir_builder/relax/ir.py | 2 + python/tvm/script/parser/ir/__init__.py | 2 + python/tvm/script/parser/relax/dist.py | 8 +- python/tvm/script/parser/relax/entry.py | 21 ++++- src/driver/driver_api.cc | 17 ++++ src/ir/global_info.cc | 13 +++ src/relax/analysis/struct_info_analysis.cc | 25 ++++-- src/relax/backend/vm/vm_builtin_lower.cc | 21 +++++ src/relax/ir/expr.cc | 2 +- src/relax/ir/struct_info.cc | 12 +-- src/relax/ir/struct_info_functor.cc | 7 +- src/relax/op/op.cc | 28 +++++++ src/relax/transform/convert_layout.cc | 6 +- src/relax/transform/to_mixed_precision.cc | 6 +- src/relay/backend/contrib/cmsisnn/target.cc | 4 +- src/relay/backend/contrib/codegen_c/target.cc | 2 +- src/relay/backend/contrib/cutlass/target.cc | 2 +- src/relay/backend/contrib/ethosu/codegen.cc | 4 +- .../backend/contrib/example_target_hooks/target.cc | 6 +- src/relay/backend/contrib/tensorrt/target.cc | 2 +- src/relay/backend/contrib/uma/targets.cc | 28 +++---- src/runtime/relax_vm/builtin.cc | 6 ++ src/script/ir_builder/ir/ir.cc | 30 +++++++ src/script/printer/ir/ir.cc | 10 +++ src/script/printer/ir_docsifier.cc | 6 ++ src/script/printer/relax/struct_info.cc | 7 ++ src/script/printer/relax/utils.h | 16 ++++ src/target/codegen.cc | 3 + src/target/target.cc | 15 +--- tests/cpp/target_test.cc | 2 +- .../python/relax/test_json_compact.py | 50 ++++++----- tests/python/relax/test_relax_operators.py | 30 +++++++ tests/python/relax/test_tvmscript_parser.py | 97 +++++++++++++++++++++- .../relax/test_tvmscript_parser_op_manipulate.py | 15 ++++ tests/python/relax/test_vm_build.py | 35 ++++++++ tests/python/relax/test_vm_codegen_only.py | 27 ++++++ 51 files changed, 680 insertions(+), 125 deletions(-) diff --git a/include/tvm/ir/global_info.h b/include/tvm/ir/global_info.h index bd6006864a..57bc57fd09 100644 --- a/include/tvm/ir/global_info.h +++ b/include/tvm/ir/global_info.h @@ -25,10 +25,16 @@ #ifndef TVM_IR_GLOBAL_INFO_H_ #define TVM_IR_GLOBAL_INFO_H_ -#include "tvm/ir/expr.h" +#include <tvm/ir/expr.h> +#include <tvm/target/target.h> namespace tvm { +/*! + * \brief Abstract label for an area of memory. + */ +using MemoryScope = String; + /*! * \brief GlobalInfo are globally static object that are referred by the IR itself. * Base node for all global info that can appear in the IR @@ -50,6 +56,48 @@ class GlobalInfo : public ObjectRef { TVM_DEFINE_OBJECT_REF_METHODS(GlobalInfo, ObjectRef, GlobalInfoNode); }; +/*! + * \brief A global info subclass for virtual devices. + */ +class VDeviceNode : public GlobalInfoNode { + public: + /*! \brief The \p Target describing how to compile for the virtual device. */ + Target target; + /*! \brief The device identifier for the virtual device. This enables us to + * differentiate between distinct devices with same Target, such as multiple GPUs. + */ + int vdevice_id; + MemoryScope memory_scope; + void VisitAttrs(tvm::AttrVisitor* v) { + v->Visit("target", &target); + v->Visit("vdevice_id", &vdevice_id); + v->Visit("memory_scope", &memory_scope); + } + + TVM_DLL bool SEqualReduce(const VDeviceNode* other, SEqualReducer equal) const { + return equal(target, other->target) && equal(vdevice_id, other->vdevice_id) && + equal(memory_scope, other->memory_scope); + } + + TVM_DLL void SHashReduce(SHashReducer hash_reduce) const { + hash_reduce(target); + hash_reduce(vdevice_id); + hash_reduce(memory_scope); + } + static constexpr const char* _type_key = "VDevice"; + TVM_DECLARE_FINAL_OBJECT_INFO(VDeviceNode, GlobalInfoNode); +}; + +/*! + * \brief Managed reference to VDeviceNode. + * \sa VDeviceNode + */ +class VDevice : public GlobalInfo { + public: + TVM_DLL explicit VDevice(Target tgt, int dev_id, MemoryScope mem_scope); + TVM_DEFINE_OBJECT_REF_METHODS(VDevice, GlobalInfo, VDeviceNode); +}; + /*! * \brief A dummy global info sub-class for testing purpose. */ diff --git a/include/tvm/relax/attrs/op.h b/include/tvm/relax/attrs/op.h index 8c0ee6abc6..f33f3a9701 100644 --- a/include/tvm/relax/attrs/op.h +++ b/include/tvm/relax/attrs/op.h @@ -57,6 +57,14 @@ struct CallTIRInplaceAttrs : public tvm::AttrsNode<CallTIRInplaceAttrs> { } }; // struct CallTIRInplaceAttrs +/*! \brief Attributes used in to_vdevice */ +struct ToVDeviceAttrs : public tvm::AttrsNode<ToVDeviceAttrs> { + VDevice dst_vdevice; + TVM_DECLARE_ATTRS(ToVDeviceAttrs, "relax.attrs.ToVDeviceAttrs") { + TVM_ATTR_FIELD(dst_vdevice).describe("The destination device where the data is copied to."); + } +}; // struct ToVDeviceAttrs + } // namespace relax } // namespace tvm diff --git a/include/tvm/relax/struct_info.h b/include/tvm/relax/struct_info.h index 385c320db1..d2bf525225 100644 --- a/include/tvm/relax/struct_info.h +++ b/include/tvm/relax/struct_info.h @@ -156,6 +156,10 @@ class TensorStructInfoNode : public StructInfoNode { * \note shape must be normalized: it can only be NullOpt or ShapeExpr or Var. */ Optional<Expr> shape; + /*! \brief The virtual device, indicates where the tensor + * is expected to be executed. + */ + Optional<VDevice> vdevice; /*! \brief The content data type, use void to denote the dtype is unknown. */ DataType dtype; /*! @@ -180,17 +184,20 @@ class TensorStructInfoNode : public StructInfoNode { void VisitAttrs(AttrVisitor* v) { v->Visit("shape", &shape); v->Visit("dtype", &dtype); + v->Visit("vdevice", &vdevice); v->Visit("ndim", &ndim); v->Visit("span", &span); } bool SEqualReduce(const TensorStructInfoNode* other, SEqualReducer equal) const { - return equal(shape, other->shape) && equal(ndim, other->ndim) && equal(dtype, other->dtype); + return equal(shape, other->shape) && equal(ndim, other->ndim) && + equal(vdevice, other->vdevice) && equal(dtype, other->dtype); } void SHashReduce(SHashReducer hash_reduce) const { hash_reduce(shape); hash_reduce(dtype); + hash_reduce(vdevice); hash_reduce(ndim); } @@ -208,19 +215,23 @@ class TensorStructInfo : public StructInfo { * \brief Construction with a known shape expression. * \param shape The shape of the tensor. * \param dtype The data type of tensor's elements. + * \param vdevice The virtual device. * \param span The span of the AST. * * \note shape must already be normalized. */ - TVM_DLL TensorStructInfo(Expr shape, DataType dtype, Span span = Span()); + TVM_DLL TensorStructInfo(Expr shape, DataType dtype, VDevice vdevice = VDevice(), + Span span = Span()); /*! * \brief Construction with an unknown shape expression. * \param dtype The data type of tensor's elements. * \param ndim The number of dimensions + * \param vdevice The virtual device. * \param span The span of the AST. */ - TVM_DLL TensorStructInfo(DataType dtype, int ndim, Span span = Span()); + TVM_DLL TensorStructInfo(DataType dtype, int ndim, VDevice vdevice = VDevice(), + Span span = Span()); TVM_DEFINE_OBJECT_REF_METHODS(TensorStructInfo, StructInfo, TensorStructInfoNode); }; diff --git a/include/tvm/relay/transform.h b/include/tvm/relay/transform.h index 7a0e003038..f4286512e5 100644 --- a/include/tvm/relay/transform.h +++ b/include/tvm/relay/transform.h @@ -47,6 +47,18 @@ using PassInfoNode = tvm::transform::PassInfoNode; using PassContext = tvm::transform::PassContext; using PassContextNode = tvm::transform::PassContextNode; using Sequential = tvm::transform::Sequential; +using FTVMRelayToTIR = tvm::transform::Pass; +/*! + * \brief TIRToRuntime conversion specific to a TargetKind + * + * This function is responsible for scanning an IRModule for appropriate Target-specific functions + and generating a Runtime module representing the compiled output + * + * \param ir_module Unified IRModule + * \param target Target to filter on or retrieve arguments from + * \return Runtime Module containing compiled functions + */ +using FTVMTIRToRuntime = tvm::runtime::TypedPackedFunc<runtime::Module(IRModule, Target)>; /* * \brief Create a function pass. diff --git a/include/tvm/script/printer/ir_docsifier.h b/include/tvm/script/printer/ir_docsifier.h index 156daebf00..1163464738 100644 --- a/include/tvm/script/printer/ir_docsifier.h +++ b/include/tvm/script/printer/ir_docsifier.h @@ -145,6 +145,8 @@ class IRDocsifierNode : public Object { std::unordered_map<ObjectRef, VariableInfo, ObjectPtrHash, ObjectPtrEqual> obj2info; /*! \brief Metadata printing */ std::unordered_map<String, Array<ObjectRef>> metadata; + /*! \brief GlobalInfo printing */ + std::unordered_map<String, Array<GlobalInfo>> global_infos; /*! \brief The variable names used already */ std::unordered_set<String> defined_names; /*! \brief Common prefixes of variable usages */ @@ -206,6 +208,11 @@ class IRDocsifierNode : public Object { Optional<ExprDoc> GetVarDoc(const ObjectRef& obj) const; /*! \brief Add a TVM object to the metadata section*/ ExprDoc AddMetadata(const ObjectRef& obj); + /*! \brief Add a GlobalInfo to the global_infos map. + * \param name The name of key of global_infos. + * \param ginfo The GlobalInfo to be added. + */ + void AddGlobalInfo(const String& name, const GlobalInfo& ginfo); /*! * \brief Check if a variable exists in the table. * \param obj The variable object. diff --git a/include/tvm/target/target.h b/include/tvm/target/target.h index 56d6a596b9..d47ac94e06 100644 --- a/include/tvm/target/target.h +++ b/include/tvm/target/target.h @@ -25,7 +25,6 @@ #define TVM_TARGET_TARGET_H_ #include <tvm/ir/expr.h> -#include <tvm/ir/module.h> #include <tvm/node/node.h> #include <tvm/support/with.h> #include <tvm/target/target_kind.h> @@ -284,14 +283,5 @@ class Target : public ObjectRef { */ void CheckAndUpdateHostConsistency(Target* target, Target* host); -/*! - * \brief Check and update host field of the given legacy heterogeneous targets and - * target host.Note that this function is for legacy target api compatibility issue only, - * not recommended for other use. - * \param ir_modules The pointer to a Map objects with keys being Target objects - * \param host The Target typed object for target host to be updated - */ -void CheckAndUpdateHostConsistency(Map<Target, IRModule>* ir_modules, Target* host); - } // namespace tvm #endif // TVM_TARGET_TARGET_H_ diff --git a/include/tvm/target/target_kind.h b/include/tvm/target/target_kind.h index 19bcce3116..10808fd12d 100644 --- a/include/tvm/target/target_kind.h +++ b/include/tvm/target/target_kind.h @@ -24,7 +24,6 @@ #ifndef TVM_TARGET_TARGET_KIND_H_ #define TVM_TARGET_TARGET_KIND_H_ -#include <tvm/ir/transform.h> #include <tvm/node/attr_registry_map.h> #include <tvm/node/node.h> @@ -50,31 +49,7 @@ using TargetFeatures = Map<String, ObjectRef>; * \return The transformed Target JSON object. */ using TargetJSON = Map<String, ObjectRef>; -using FTVMTargetParser = TypedPackedFunc<TargetJSON(TargetJSON)>; - -/*! - * \brief RelayToTIR tvm::transform::Pass specific to a TargetKind - * - * Called before the default lowering passes. - * - * \param mod The module that an optimization pass runs on. - * \param pass_ctx The pass context that can provide information for the optimization. - * - * \return The transformed module. - */ -using FTVMRelayToTIR = transform::Pass; - -/*! - * \brief TIRToRuntime conversion specific to a TargetKind - * - * This function is responsible for scanning an IRModule for appropriate Target-specific functions - and generating a Runtime module representing the compiled output - * - * \param ir_module Unified IRModule - * \param target Target to filter on or retrieve arguments from - * \return Runtime Module containing compiled functions - */ -using FTVMTIRToRuntime = runtime::TypedPackedFunc<runtime::Module(IRModule, Target)>; +using FTVMTargetParser = tvm::runtime::TypedPackedFunc<TargetJSON(TargetJSON)>; namespace detail { template <typename, typename, typename> diff --git a/include/tvm/tir/usmp/utils.h b/include/tvm/tir/usmp/utils.h index f49e9ceef7..a67350a2bb 100644 --- a/include/tvm/tir/usmp/utils.h +++ b/include/tvm/tir/usmp/utils.h @@ -27,6 +27,7 @@ #include <tvm/ir/expr.h> #include <tvm/ir/memory_pools.h> +#include <tvm/ir/module.h> #include <tvm/runtime/device_api.h> #include <tvm/target/target.h> #include <tvm/tir/stmt.h> diff --git a/python/tvm/ir/__init__.py b/python/tvm/ir/__init__.py index 49c2cf6348..939a5f6383 100644 --- a/python/tvm/ir/__init__.py +++ b/python/tvm/ir/__init__.py @@ -35,7 +35,7 @@ from .base import ( from .container import Array, Map from .expr import BaseExpr, GlobalVar, PrimExpr, Range, RelayExpr from .function import BaseFunc, CallingConv -from .global_info import GlobalInfo, DummyGlobalInfo +from .global_info import GlobalInfo, DummyGlobalInfo, VDevice from .memory_pools import ( ConstantMemoryPools, ConstantPoolInfo, diff --git a/python/tvm/ir/global_info.py b/python/tvm/ir/global_info.py index 17011e76a6..458a16717b 100644 --- a/python/tvm/ir/global_info.py +++ b/python/tvm/ir/global_info.py @@ -40,3 +40,15 @@ class DummyGlobalInfo(GlobalInfo): self.__init_handle_by_constructor__( _ffi_api.DummyGlobalInfo, ) + + +class VDevice(GlobalInfo): + def __init__( + self, + target=None, + vdevice_id: int = 0, + memory_scope: str = "global", + ) -> None: + if isinstance(target, (dict, str)): + target = tvm.target.Target(tvm.runtime.convert(target)) + self.__init_handle_by_constructor__(_ffi_api.VDevice, target, vdevice_id, memory_scope) diff --git a/python/tvm/ir/json_compact.py b/python/tvm/ir/json_compact.py index 6ce2a8b9e2..224932b00c 100644 --- a/python/tvm/ir/json_compact.py +++ b/python/tvm/ir/json_compact.py @@ -57,6 +57,21 @@ def create_updater(node_map, from_ver, to_ver): return _updater +def create_updater_13_to_14(): + """Create an update to upgrade json from v0.13 to v0.14 for TVM Unity""" + + def _update_vdevice(item, _): + if "vdevice" not in item["attrs"]: + item["attrs"]["vdevice"] = "0" + return item + + node_map = { + "relax.TensorStructInfo": _update_vdevice, + } + + return create_updater(node_map, "0.13", "0.14") + + def create_updater_08_to_09(): """ Create an update to upgrade json from v0.8 to v0.9 @@ -259,6 +274,8 @@ def upgrade_json(json_str): data = create_updater_08_to_09()(create_updater_07_to_08()(data)) elif from_version.startswith("0.8"): data = create_updater_08_to_09()(data) + elif from_version.startswith("0.13"): + data = create_updater_13_to_14()(data) else: raise ValueError(f"Cannot update from version {from_version}") return json.dumps(data, indent=2) diff --git a/python/tvm/relax/op/base.py b/python/tvm/relax/op/base.py index 25c70e0493..1d49c00ea8 100644 --- a/python/tvm/relax/op/base.py +++ b/python/tvm/relax/op/base.py @@ -690,3 +690,24 @@ def invoke_pure_closure( sinfo_args = [sinfo_args] return _ffi_api.invoke_pure_closure(closure, args, sinfo_args) # type: ignore + + +def to_vdevice(data, dst_vdevice) -> Expr: + """Copy data to the destination device. This + operator helps data transferring between difference devices for + heterogeneous execution. + + Parameters + ---------- + data : Expr + The tensor to be copied. + + dst_device : Union[:py:class:`Device`, str] + The destination device where the data is copied to. + + Returns + ------- + result : Expr + The copied result. + """ + return _ffi_api.to_vdevice(data, dst_vdevice) # type: ignore diff --git a/python/tvm/relax/struct_info.py b/python/tvm/relax/struct_info.py index 3dcc3dc9a0..e78e1cf69a 100644 --- a/python/tvm/relax/struct_info.py +++ b/python/tvm/relax/struct_info.py @@ -16,14 +16,14 @@ # under the License. # pylint: disable=invalid-name, unused-import """The struct info nodes of the Relax language.""" -from typing import List, Optional, Tuple, Union +from typing import List, Optional, Union import tvm._ffi import tvm -from tvm.ir import Span, Node, EnvFunc, Array, Type +from tvm.ir import Span, EnvFunc, Array, VDevice from tvm.tir import PrimExpr -from .expr import StructInfo, Var, Expr, ShapeExpr +from .expr import StructInfo, Expr, ShapeExpr from . import _ffi_api, ty, expr @@ -93,6 +93,9 @@ class TensorStructInfo(StructInfo): dtype : Optional[str] The content data type. + vdevice : Optional[Vdevice] + The virtual device. + ndim : Optional[int] The number of dimensions of the tensor. @@ -103,6 +106,7 @@ class TensorStructInfo(StructInfo): shape: Optional[Expr] dtype: str + vdevice: Optional[VDevice] ndim: int span: Span @@ -110,14 +114,14 @@ class TensorStructInfo(StructInfo): self, shape: Union[Optional[Expr], List[PrimExpr]] = None, dtype: str = "float32", + vdevice: Union[Optional[VDevice], str] = None, ndim: int = -1, span: Span = None, ) -> None: if isinstance(shape, (list, tuple, Array)): shape = ShapeExpr(shape) - self.__init_handle_by_constructor__( - _ffi_api.TensorStructInfo, shape, dtype, ndim, span # type: ignore + _ffi_api.TensorStructInfo, shape, dtype, ndim, vdevice, span # type: ignore ) diff --git a/python/tvm/script/ir_builder/ir/__init__.py b/python/tvm/script/ir_builder/ir/__init__.py index 68eda2cfee..fdf44b2b79 100644 --- a/python/tvm/script/ir_builder/ir/__init__.py +++ b/python/tvm/script/ir_builder/ir/__init__.py @@ -22,5 +22,7 @@ from .ir import ( ir_module, module_attrs, module_global_infos, + lookup_vdevice, + vdevice, dummy_global_info, ) diff --git a/python/tvm/script/ir_builder/ir/ir.py b/python/tvm/script/ir_builder/ir/ir.py index 53c48b4cc5..0d3523ec7d 100644 --- a/python/tvm/script/ir_builder/ir/ir.py +++ b/python/tvm/script/ir_builder/ir/ir.py @@ -18,7 +18,7 @@ from typing import Dict, List -from tvm.ir import BaseFunc, GlobalVar, GlobalInfo, DummyGlobalInfo +from tvm.ir import BaseFunc, GlobalVar, GlobalInfo, VDevice, DummyGlobalInfo from tvm.runtime import Object as tvm_Object @@ -104,3 +104,39 @@ def dummy_global_info() -> DummyGlobalInfo: The result dummy global info. """ return DummyGlobalInfo() # type: ignore[attr-defined] # pylint: disable=no-member + + +def vdevice(target=None, vdevice_id: int = 0, memory_scope: str = "global") -> VDevice: + """Create a virtual device global info. + Parameters + ---------- + target + The target. + vdevice_id: int + The virtual device index. + memory_scope: str + The memory scope, default is "global" + + Returns + ------- + res : VDevice + The result virtual device. + """ + return VDevice(target, vdevice_id, memory_scope) # type: ignore[attr-defined] # pylint: disable=no-member + + +def lookup_vdevice(target_kind: str = None, device_index: int = -1) -> VDevice: + """Retrieve a virtual device from the globalinfo vdevice list. + Parameters + ---------- + target_kind: str + The target device kind, for example 'llvm' or 'cuda'. + device_index: int + The virtual device index. + + Returns + ------- + res : VDevice + The result virtual device. + """ + return _ffi_api.LookupVDevice(target_kind, device_index) # type: ignore[attr-defined] # pylint: disable=no-member diff --git a/python/tvm/script/ir_builder/relax/ir.py b/python/tvm/script/ir_builder/relax/ir.py index 5bb0374d35..8a538c1868 100644 --- a/python/tvm/script/ir_builder/relax/ir.py +++ b/python/tvm/script/ir_builder/relax/ir.py @@ -140,6 +140,7 @@ from tvm.relax.op import ( tanh, erf, tile, + to_vdevice, tril, triu, unique, @@ -693,6 +694,7 @@ __all__ = [ "tan", "tanh", "tile", + "to_vdevice", "tril", "triu", "tuple", diff --git a/python/tvm/script/parser/ir/__init__.py b/python/tvm/script/parser/ir/__init__.py index ec518f8573..3a8196288d 100644 --- a/python/tvm/script/parser/ir/__init__.py +++ b/python/tvm/script/parser/ir/__init__.py @@ -27,4 +27,6 @@ __all__ = [ "module_global_infos", "dummy_global_info", "Range", + "lookup_vdevice", + "vdevice", ] diff --git a/python/tvm/script/parser/relax/dist.py b/python/tvm/script/parser/relax/dist.py index 120d57ca56..f9c78f980f 100644 --- a/python/tvm/script/parser/relax/dist.py +++ b/python/tvm/script/parser/relax/dist.py @@ -78,7 +78,11 @@ def DTensor( raise ValueError(f"shape must be a list or tuple, but got: {shape}") if isinstance(device_mesh, str): if not IRBuilder.is_in_scope(): - return (DTensorProxy(TensorProxy(shape, dtype, ndim), DeviceMesh([], Range(0, 1)), ""),) + return ( + DTensorProxy( + TensorProxy(shape, dtype, None, ndim), DeviceMesh([], Range(0, 1)), "" + ), + ) name, index = device_mesh.split("[") index = int(index[:-1]) frames = IRBuilder.current().frames @@ -89,7 +93,7 @@ def DTensor( assert isinstance(device_mesh, DeviceMesh) if isinstance(placement, str): placement = Placement.from_text(placement) - return DTensorProxy(TensorProxy(shape, dtype, ndim), device_mesh, placement) + return DTensorProxy(TensorProxy(shape, dtype, None, ndim), device_mesh, placement) __all__ = ["DTensor", "device_mesh"] diff --git a/python/tvm/script/parser/relax/entry.py b/python/tvm/script/parser/relax/entry.py index ff237a5600..1c18d75be4 100644 --- a/python/tvm/script/parser/relax/entry.py +++ b/python/tvm/script/parser/relax/entry.py @@ -37,6 +37,7 @@ from tvm.runtime import ObjectGeneric from tvm.tir import PrimExpr from .._core import parse, utils +from ..ir import lookup_vdevice FType = TypeVar("FType", bound=_Callable) @@ -103,12 +104,14 @@ def _eval_shape(expr: Union[str, PrimExpr], dict_globals: Optional[Dict[str, Any class TensorProxy(StructInfoProxy): shape: Optional[List[Union[str, PrimExpr]]] dtype: str + vdevice: Optional[str] ndim: int def __init__( self, shape: Optional[Union[List[Union[PrimExpr, str]], Expr]] = None, dtype: Optional[str] = None, + vdevice: Optional[str] = None, ndim: int = -1, ) -> None: if isinstance(shape, Expr): @@ -124,6 +127,7 @@ class TensorProxy(StructInfoProxy): ) self.shape = shape self.dtype = dtype + self.vdevice = vdevice self.ndim = ndim def get_symbolic_vars(self) -> Set[str]: @@ -133,10 +137,18 @@ class TensorProxy(StructInfoProxy): return {s for s in self.shape if isinstance(s, str) and s.isidentifier()} def as_struct_info(self, dict_globals: Optional[Dict[str, Any]] = None) -> TensorStructInfo: + vdev = self.vdevice + if isinstance(self.vdevice, str): + if ":" in self.vdevice: + split_vdev = self.vdevice.split(":") + vdev = lookup_vdevice(split_vdev[0], int(split_vdev[1])) + else: + vdev = lookup_vdevice(self.vdevice, 0) + if self.shape is None: - return TensorStructInfo(None, self.dtype, self.ndim) + return TensorStructInfo(None, self.dtype, vdev, self.ndim) elif isinstance(self.shape, (ShapeExpr, Var)): - return TensorStructInfo(self.shape, self.dtype, self.ndim) + return TensorStructInfo(self.shape, self.dtype, vdev, self.ndim) else: if dict_globals is None and any([isinstance(s, str) for s in self.shape]): raise ValueError( @@ -144,12 +156,13 @@ class TensorProxy(StructInfoProxy): "and return annotations for TVMScript." ) shape = [_eval_shape(s, dict_globals) for s in self.shape] - return TensorStructInfo(shape, self.dtype, self.ndim) + return TensorStructInfo(shape, self.dtype, vdev, self.ndim) def Tensor( shape: Optional[Union[List[Union[PrimExpr, str]], Expr]] = None, dtype: Optional[str] = None, + vdevice: Optional[str] = None, ndim: int = -1, ) -> TensorProxy: # scalar tensor case @@ -161,7 +174,7 @@ def Tensor( if shape is not None and not isinstance(shape, (tuple, list)) and not isinstance(shape, Expr): raise ValueError(f"shape must be a list/tuple or an Expr, but got: {shape}") - return TensorProxy(shape, dtype, ndim) + return TensorProxy(shape, dtype, vdevice, ndim) ############################## R.Callable ############################## diff --git a/src/driver/driver_api.cc b/src/driver/driver_api.cc index d46fab7168..b7ba0ffe44 100644 --- a/src/driver/driver_api.cc +++ b/src/driver/driver_api.cc @@ -431,6 +431,23 @@ std::pair<IRModule, IRModule> SplitMixedModule(IRModule mod_mixed, const Target& return {host_mod, device_mod}; } +/*! + * \brief Check and update host field of the given legacy heterogeneous targets and + * target host.Note that this function is for legacy target api compatibility issue only, + * not recommended for other use. + * \param ir_modules The pointer to a Map objects with keys being Target objects + * \param host The Target typed object for target host to be updated + */ +void CheckAndUpdateHostConsistency(Map<Target, IRModule>* targets, Target* host) { + Map<Target, IRModule> new_targets; + for (auto& it : *targets) { + auto target = it.first; + CheckAndUpdateHostConsistency(&target, host); + new_targets.Set(target, it.second); + } + *targets = new_targets; +} + runtime::Module TIRToRuntime(const Map<Target, IRModule>& inputs_arg, const Target& target_host_arg) { std::vector<runtime::Module> device_modules; diff --git a/src/ir/global_info.cc b/src/ir/global_info.cc index 48f56d60d6..f1ecc8cd04 100644 --- a/src/ir/global_info.cc +++ b/src/ir/global_info.cc @@ -29,4 +29,17 @@ TVM_REGISTER_GLOBAL("ir.DummyGlobalInfo").set_body_typed([]() { auto n = DummyGlobalInfo(make_object<DummyGlobalInfoNode>()); return n; }); + +VDevice::VDevice(Target tgt = {}, int dev_id = -1, MemoryScope mem_scope = {}) { + ObjectPtr<VDeviceNode> n = make_object<VDeviceNode>(); + n->target = std::move(tgt); + n->vdevice_id = std::move(dev_id); + n->memory_scope = std::move(mem_scope); + data_ = std::move(n); +} + +TVM_REGISTER_NODE_TYPE(VDeviceNode); +TVM_REGISTER_GLOBAL("ir.VDevice").set_body_typed([](Target tgt, int dev_id, MemoryScope mem_scope) { + return VDevice(tgt, dev_id, mem_scope); +}); } // namespace tvm diff --git a/src/relax/analysis/struct_info_analysis.cc b/src/relax/analysis/struct_info_analysis.cc index 7006f71198..82ccdf33ea 100644 --- a/src/relax/analysis/struct_info_analysis.cc +++ b/src/relax/analysis/struct_info_analysis.cc @@ -149,19 +149,24 @@ class WellDefinedEraser : public StructInfoMutator, std::swap(has_undefined_, has_undefined); } + VDevice vdev = VDevice(); + if (op->vdevice.defined()) { + vdev = op->vdevice.value(); + } + // erase symbolic shape if we have undefined. if (!has_undefined) { if (shape.same_as(op->shape)) { return GetRef<StructInfo>(op); } else { if (shape.defined()) { - return TensorStructInfo(shape.value(), op->dtype, op->span); + return TensorStructInfo(shape.value(), op->dtype, vdev, op->span); } else { - return TensorStructInfo(op->dtype, op->ndim, op->span); + return TensorStructInfo(op->dtype, op->ndim, vdev, op->span); } } } else { - return TensorStructInfo(op->dtype, op->ndim, op->span); + return TensorStructInfo(op->dtype, op->ndim, vdev, op->span); } } @@ -767,6 +772,16 @@ class StructInfoLCAFinder // find the target dtype and ndim. DataType dtype = lhs->dtype == rhs->dtype ? lhs->dtype : DataType::Void(); int ndim = lhs->ndim == rhs->ndim ? lhs->ndim : kUnknownNDim; + VDevice vdev = VDevice(); + if (lhs->vdevice.defined() && rhs->vdevice.defined()) { + if (lhs->vdevice.value().same_as(lhs->vdevice.value())) { + vdev = lhs->vdevice.value(); + } + } else if (lhs->vdevice.defined()) { + vdev = lhs->vdevice.value(); + } else if (rhs->vdevice.defined()) { + vdev = rhs->vdevice.value(); + } // if ndim mismatch or one side of shape is missing // then we cannot keep in symbolic shape if (lhs->ndim != rhs->ndim || !lhs->shape.defined() || !rhs->shape.defined() || @@ -775,12 +790,12 @@ class StructInfoLCAFinder if (!lhs->shape.defined() && lhs->dtype == dtype && lhs->ndim == ndim) { return GetRef<StructInfo>(lhs); } else { - return TensorStructInfo(dtype, ndim, lhs->span); + return TensorStructInfo(dtype, ndim, vdev, lhs->span); } } // symbolic shape match but dtype mismatch if (lhs->dtype != dtype) { - return TensorStructInfo(lhs->shape.value(), dtype, lhs->span); + return TensorStructInfo(lhs->shape.value(), dtype, vdev, lhs->span); } else { return GetRef<StructInfo>(lhs); } diff --git a/src/relax/backend/vm/vm_builtin_lower.cc b/src/relax/backend/vm/vm_builtin_lower.cc index 6087c2bb25..784b3c9fd5 100644 --- a/src/relax/backend/vm/vm_builtin_lower.cc +++ b/src/relax/backend/vm/vm_builtin_lower.cc @@ -21,6 +21,7 @@ * \brief Lowers most builtin functions and packed calls. */ #include <tvm/relax/analysis.h> +#include <tvm/relax/attrs/op.h> #include <tvm/relax/backend.h> #include <tvm/relax/expr_functor.h> #include <tvm/relax/type.h> @@ -46,6 +47,8 @@ class VMBuiltinLowerMutator : public ExprMutator { return Reshape(call); } else if (call->op == shape_of_op_) { return ShapeOf(call); + } else if (call->op == to_vdevice_op_) { + return ToDevice(call); } else if (call->op == make_closure_op_) { return MakeClosure(call); } else if (call->op == invoke_closure_op_) { @@ -156,6 +159,22 @@ class VMBuiltinLowerMutator : public ExprMutator { return Call(builtin_shape_of_, call_node->args, Attrs(), {GetStructInfo(call_node)}); } + Expr ToDevice(const Call& call_node) { + // TODO(yongwww): replace ToVDeviceAttrs with related Expr + ICHECK(call_node->args.size() == 1); + ICHECK(call_node->struct_info_.defined()); + auto attrs = call_node->attrs.as<ToVDeviceAttrs>(); + Array<Expr> args; + args.push_back(call_node->args[0]); + // Get the DLDeviceType and device_id from VDevice + VDevice vdev = attrs->dst_vdevice; + int dev_type = vdev->target->GetTargetDeviceType(); + int dev_id = vdev->vdevice_id; + args.push_back(PrimValue::Int64(dev_type)); + args.push_back(PrimValue::Int64(dev_id)); + return Call(builtin_to_device_, args, call_node->attrs, {GetStructInfo(call_node)}); + } + Expr MakeClosure(const Call& call_node) { ICHECK(call_node->args.size() == 2); ICHECK(call_node->args[0]->IsInstance<GlobalVarNode>()); @@ -198,6 +217,7 @@ class VMBuiltinLowerMutator : public ExprMutator { const Op& call_tir_dyn_op_ = Op::Get("relax.vm.call_tir_dyn"); const Op& reshape_op_ = Op::Get("relax.reshape"); const Op& shape_of_op_ = Op::Get("relax.shape_of"); + const Op& to_vdevice_op_ = Op::Get("relax.to_vdevice"); const Op& make_closure_op_ = Op::Get("relax.make_closure"); const Op& invoke_closure_op_ = Op::Get("relax.invoke_closure"); const Op& alloc_tensor_op_ = Op::Get("relax.builtin.alloc_tensor"); @@ -214,6 +234,7 @@ class VMBuiltinLowerMutator : public ExprMutator { const ExternFunc builtin_call_tir_dyn_{"vm.builtin.call_tir_dyn"}; const ExternFunc builtin_reshape_{"vm.builtin.reshape"}; const ExternFunc builtin_shape_of_{"vm.builtin.shape_of"}; + const ExternFunc builtin_to_device_{"vm.builtin.to_device"}; const ExternFunc builtin_make_closure_{"vm.builtin.make_closure"}; const ExternFunc builtin_invoke_closure_{"vm.builtin.invoke_closure"}; }; diff --git a/src/relax/ir/expr.cc b/src/relax/ir/expr.cc index ccff18cd40..ac04096aaf 100644 --- a/src/relax/ir/expr.cc +++ b/src/relax/ir/expr.cc @@ -292,7 +292,7 @@ Constant::Constant(runtime::NDArray data, Optional<StructInfo> struct_info_annot n->struct_info_ = struct_info_annotation.value(); n->checked_type_ = GetStaticType(struct_info_annotation.value()); } else { - TensorStructInfo tinfo(ShapeExpr(values), n->data.DataType(), span); + TensorStructInfo tinfo(ShapeExpr(values), n->data.DataType(), VDevice(), span); n->struct_info_ = tinfo; n->checked_type_ = DynTensorType(tinfo->ndim, tinfo->dtype); } diff --git a/src/relax/ir/struct_info.cc b/src/relax/ir/struct_info.cc index c290711dcd..31784af000 100644 --- a/src/relax/ir/struct_info.cc +++ b/src/relax/ir/struct_info.cc @@ -92,7 +92,7 @@ TVM_REGISTER_GLOBAL("relax.ShapeStructInfo") }); // Tensor -TensorStructInfo::TensorStructInfo(Expr shape, DataType dtype, Span span) { +TensorStructInfo::TensorStructInfo(Expr shape, DataType dtype, VDevice vdevice, Span span) { ObjectPtr<TensorStructInfoNode> n = make_object<TensorStructInfoNode>(); // assign ndim before move Optional<ShapeStructInfo> sinfo = MatchStructInfo<ShapeStructInfo>(shape); @@ -104,15 +104,17 @@ TensorStructInfo::TensorStructInfo(Expr shape, DataType dtype, Span span) { // assign rest of the fields. n->shape = std::move(shape); n->dtype = dtype; + n->vdevice = vdevice; n->span = span; data_ = std::move(n); } -TensorStructInfo::TensorStructInfo(DataType dtype, int ndim, Span span) { +TensorStructInfo::TensorStructInfo(DataType dtype, int ndim, VDevice vdevice, Span span) { ObjectPtr<TensorStructInfoNode> n = make_object<TensorStructInfoNode>(); CHECK_GE(ndim, -1) << "ndim of TensorStructInfo must be >= -1, but got " << ndim; n->ndim = ndim; n->dtype = dtype; + n->vdevice = vdevice; n->span = span; data_ = std::move(n); } @@ -120,12 +122,12 @@ TensorStructInfo::TensorStructInfo(DataType dtype, int ndim, Span span) { TVM_REGISTER_NODE_TYPE(TensorStructInfoNode); TVM_REGISTER_GLOBAL("relax.TensorStructInfo") - .set_body_typed([](Optional<Expr> shape, DataType dtype, int ndim, Span span) { + .set_body_typed([](Optional<Expr> shape, DataType dtype, int ndim, VDevice vdevice, Span span) { if (shape.defined()) { CHECK_EQ(ndim, kUnknownNDim) << "ValueError: Cannot both specify shape and ndim"; - return TensorStructInfo(shape.value(), dtype, span); + return TensorStructInfo(shape.value(), dtype, vdevice, span); } else { - return TensorStructInfo(dtype, ndim, span); + return TensorStructInfo(dtype, ndim, vdevice, span); } }); diff --git a/src/relax/ir/struct_info_functor.cc b/src/relax/ir/struct_info_functor.cc index 72ea623e07..c998d8c0b2 100644 --- a/src/relax/ir/struct_info_functor.cc +++ b/src/relax/ir/struct_info_functor.cc @@ -94,10 +94,15 @@ StructInfo StructInfoMutator::VisitStructInfo_(const TensorStructInfoNode* op) { shape = this->VisitStructInfoExprField(op->shape.value()); } + VDevice vdev = VDevice(); + if (op->vdevice.defined()) { + vdev = op->vdevice.value(); + } + if (shape.same_as(op->shape)) { return GetRef<StructInfo>(op); } else { - return TensorStructInfo(shape.value(), op->dtype, op->span); + return TensorStructInfo(shape.value(), op->dtype, vdev, op->span); } } diff --git a/src/relax/op/op.cc b/src/relax/op/op.cc index af93b43dcf..1f4f7d5c34 100644 --- a/src/relax/op/op.cc +++ b/src/relax/op/op.cc @@ -866,5 +866,33 @@ Expr MakeStopLiftParams(Expr x) { TVM_REGISTER_GLOBAL("relax.op.builtin.stop_lift_params").set_body_typed(MakeStopLiftParams); +// to_vdevice +TVM_REGISTER_NODE_TYPE(ToVDeviceAttrs); + +StructInfo InferToVDeviceStructInfo(const Call& call, const BlockBuilder& ctx) { + ICHECK(call->args.size() == 1); + ICHECK(call->args[0]->struct_info_.defined()); + TensorStructInfo data_sinfo = GetUnaryInputTensorStructInfo(call, ctx); + return data_sinfo; +} + +RELAY_REGISTER_OP("relax.to_vdevice") + .set_num_inputs(1) + .set_attrs_type<ToVDeviceAttrs>() + .add_argument("data", "Expr", "The input expression to be copied") + .set_attr<FInferStructInfo>("FInferStructInfo", InferToVDeviceStructInfo) + .set_attr<Bool>("FPurity", Bool(true)); + +Expr MakeToVDevice(Expr data, VDevice dst_vdevice) { + static const Op& op = Op::Get("relax.to_vdevice"); + // TODO(@yongwww): replace Attr with TensorStructInfo + ObjectPtr<ToVDeviceAttrs> attrs = make_object<ToVDeviceAttrs>(); + attrs->dst_vdevice = dst_vdevice; + + return Call(op, {data}, Attrs(attrs), {}); +} + +TVM_REGISTER_GLOBAL("relax.op.to_vdevice").set_body_typed(MakeToVDevice); + } // namespace relax } // namespace tvm diff --git a/src/relax/transform/convert_layout.cc b/src/relax/transform/convert_layout.cc index 91dcd5d8e8..dd09dd67b8 100644 --- a/src/relax/transform/convert_layout.cc +++ b/src/relax/transform/convert_layout.cc @@ -267,7 +267,11 @@ class LayoutConvertMutator : public ExprMutator { new_shape.push_back( shape->values[from.LeafValue()->layout.IndexOf(to.LeafValue()->layout[i])]); } - return TensorStructInfo(ShapeExpr(new_shape), tsinfo->dtype, tsinfo->span); + VDevice vdev = VDevice(); + if (tsinfo->vdevice.defined()) { + vdev = tsinfo->vdevice.value(); + } + return TensorStructInfo(ShapeExpr(new_shape), tsinfo->dtype, vdev, tsinfo->span); }; StructInfo new_struct_info = TransformTupleLeaf<LayoutDecision>( binding->struct_info, std::array<NLayout, 2>({from_layout, input_layout}), fvisitleaf); diff --git a/src/relax/transform/to_mixed_precision.cc b/src/relax/transform/to_mixed_precision.cc index 64763276d0..d12d1080b9 100644 --- a/src/relax/transform/to_mixed_precision.cc +++ b/src/relax/transform/to_mixed_precision.cc @@ -289,7 +289,11 @@ class ToMixedPrecisionRewriter : public ExprMutator { if (fp16_input_names_.count(var->name_hint())) { auto sinfo = GetStructInfo(var); if (auto tensor_sinfo = sinfo.as<TensorStructInfoNode>()) { - TensorStructInfo fp16_sinfo(tensor_sinfo->shape.value(), DataType::Float(16), + VDevice vdev = VDevice(); + if (tensor_sinfo->vdevice.defined()) { + vdev = tensor_sinfo->vdevice.value(); + } + TensorStructInfo fp16_sinfo(tensor_sinfo->shape.value(), DataType::Float(16), vdev, tensor_sinfo->span); Var fp16_var(var->vid, fp16_sinfo, var->span); var_remap_[var->vid] = fp16_var; diff --git a/src/relay/backend/contrib/cmsisnn/target.cc b/src/relay/backend/contrib/cmsisnn/target.cc index f14c106703..527fba98c0 100644 --- a/src/relay/backend/contrib/cmsisnn/target.cc +++ b/src/relay/backend/contrib/cmsisnn/target.cc @@ -37,8 +37,8 @@ TVM_REGISTER_TARGET_KIND("cmsis-nn", kDLCPU) .add_attr_option<Array<String>>("mattr") .add_attr_option<String>("mcpu") .add_attr_option<Bool>("debug_last_error") - .set_attr<FTVMRelayToTIR>(tvm::attr::kRelayToTIR, RelayToTIR()) - .set_attr<FTVMTIRToRuntime>("TIRToRuntime", TIRToRuntime) + .set_attr<relay::transform::FTVMRelayToTIR>(tvm::attr::kRelayToTIR, RelayToTIR()) + .set_attr<relay::transform::FTVMTIRToRuntime>("TIRToRuntime", TIRToRuntime) .set_target_parser(tvm::target::parsers::cpu::ParseTarget); } // namespace cmsisnn diff --git a/src/relay/backend/contrib/codegen_c/target.cc b/src/relay/backend/contrib/codegen_c/target.cc index 623057ac17..cd1e0283df 100644 --- a/src/relay/backend/contrib/codegen_c/target.cc +++ b/src/relay/backend/contrib/codegen_c/target.cc @@ -34,7 +34,7 @@ namespace contrib { */ TVM_REGISTER_TARGET_KIND("ccompiler", kDLCPU) .set_attr<Bool>(tvm::attr::kIsExternalCodegen, Bool(true)) - .set_attr<FTVMRelayToTIR>(tvm::attr::kRelayToTIR, CCompilerPass()) + .set_attr<relay::transform::FTVMRelayToTIR>(tvm::attr::kRelayToTIR, CCompilerPass()) // Value is prepended to every output CModule. .add_attr_option<String>("header", String("")); diff --git a/src/relay/backend/contrib/cutlass/target.cc b/src/relay/backend/contrib/cutlass/target.cc index 7b377f340a..50c8b84a90 100644 --- a/src/relay/backend/contrib/cutlass/target.cc +++ b/src/relay/backend/contrib/cutlass/target.cc @@ -40,7 +40,7 @@ namespace cutlass { */ TVM_REGISTER_TARGET_KIND("cutlass", kDLCUDA) .set_attr<Bool>(tvm::attr::kIsExternalCodegen, Bool(true)) - .set_attr<FTVMRelayToTIR>("RelayToTIR", CompileForCutlass()) + .set_attr<tvm::transform::Pass>("RelayToTIR", CompileForCutlass()) // An integer specifying the compute capability. For example, 75 for Turing and // 80 or 86 for Ampere. .add_attr_option<Integer>("sm", Integer(80)) diff --git a/src/relay/backend/contrib/ethosu/codegen.cc b/src/relay/backend/contrib/ethosu/codegen.cc index f35d4c6d48..2e635455e9 100644 --- a/src/relay/backend/contrib/ethosu/codegen.cc +++ b/src/relay/backend/contrib/ethosu/codegen.cc @@ -320,8 +320,8 @@ runtime::Module TIRToRuntime(IRModule mod, Target target) { TVM_REGISTER_TARGET_KIND("ethos-u", kDLCPU) .set_attr<Bool>("use_device_api", Bool(true)) - .set_attr<FTVMRelayToTIR>(tvm::attr::kRelayToTIR, RelayToTIR()) - .set_attr<FTVMTIRToRuntime>("TIRToRuntime", TIRToRuntime); + .set_attr<relay::transform::FTVMRelayToTIR>(tvm::attr::kRelayToTIR, RelayToTIR()) + .set_attr<relay::transform::FTVMTIRToRuntime>("TIRToRuntime", TIRToRuntime); } // namespace ethosu } // namespace contrib diff --git a/src/relay/backend/contrib/example_target_hooks/target.cc b/src/relay/backend/contrib/example_target_hooks/target.cc index b01c23ed80..275efaa933 100644 --- a/src/relay/backend/contrib/example_target_hooks/target.cc +++ b/src/relay/backend/contrib/example_target_hooks/target.cc @@ -33,8 +33,10 @@ runtime::Module TIRToRuntime(IRModule mod, Target target); TVM_REGISTER_TARGET_KIND("example_target_hook", kDLCPU) .set_attr<Bool>("use_device_api", Bool(true)) - .set_attr<FTVMRelayToTIR>(attr::kRelayToTIR, relay::contrib::example_target_hooks::RelayToTIR()) - .set_attr<FTVMTIRToRuntime>("TIRToRuntime", relay::contrib::example_target_hooks::TIRToRuntime) + .set_attr<relay::transform::FTVMRelayToTIR>(attr::kRelayToTIR, + relay::contrib::example_target_hooks::RelayToTIR()) + .set_attr<relay::transform::FTVMTIRToRuntime>( + "TIRToRuntime", relay::contrib::example_target_hooks::TIRToRuntime) .add_attr_option<Integer>("example_attribute", Integer(0)); } // namespace tvm diff --git a/src/relay/backend/contrib/tensorrt/target.cc b/src/relay/backend/contrib/tensorrt/target.cc index 2e4581d30a..0277787a8c 100644 --- a/src/relay/backend/contrib/tensorrt/target.cc +++ b/src/relay/backend/contrib/tensorrt/target.cc @@ -39,7 +39,7 @@ namespace tensorrt { */ TVM_REGISTER_TARGET_KIND("tensorrt", kDLCUDA) .set_attr<Bool>(tvm::attr::kIsExternalCodegen, Bool(true)) - .set_attr<FTVMRelayToTIR>("RelayToTIR", CompileForTensorRT()) + .set_attr<tvm::transform::Pass>("RelayToTIR", CompileForTensorRT()) // A array of three integers given the major, minor, and patch numbers for the supported // TensorRT compiler version. If empty will be auto-detected from linked library. Default empty. .add_attr_option<Array<Integer>>("tensorrt_version", Array<Integer>()) diff --git a/src/relay/backend/contrib/uma/targets.cc b/src/relay/backend/contrib/uma/targets.cc index e2fe644cb9..d01f5b4c73 100644 --- a/src/relay/backend/contrib/uma/targets.cc +++ b/src/relay/backend/contrib/uma/targets.cc @@ -46,20 +46,20 @@ TVM_REGISTER_GLOBAL("relay.backend.contrib.uma.RegisterTarget") } } - auto target_kind = - TargetKindRegEntry::RegisterOrGet(target_name) - .set_name() - .set_default_device_type(kDLCPU) - .add_attr_option<Array<String>>("keys") - .add_attr_option<String>("tag") - .add_attr_option<String>("device") - .add_attr_option<String>("model") - .add_attr_option<Array<String>>("libs") - .add_attr_option<Target>("host") - .add_attr_option<Integer>("from_device") - .set_attr<FTVMRelayToTIR>(attr::kRelayToTIR, - relay::contrib::uma::RelayToTIR(target_name)) - .set_attr<FTVMTIRToRuntime>("TIRToRuntime", relay::contrib::uma::TIRToRuntime); + auto target_kind = TargetKindRegEntry::RegisterOrGet(target_name) + .set_name() + .set_default_device_type(kDLCPU) + .add_attr_option<Array<String>>("keys") + .add_attr_option<String>("tag") + .add_attr_option<String>("device") + .add_attr_option<String>("model") + .add_attr_option<Array<String>>("libs") + .add_attr_option<Target>("host") + .add_attr_option<Integer>("from_device") + .set_attr<relay::transform::FTVMRelayToTIR>( + attr::kRelayToTIR, relay::contrib::uma::RelayToTIR(target_name)) + .set_attr<relay::transform::FTVMTIRToRuntime>( + "TIRToRuntime", relay::contrib::uma::TIRToRuntime); // target kind attrs inventory auto kind = TargetKind::Get(target_name).value(); diff --git a/src/runtime/relax_vm/builtin.cc b/src/runtime/relax_vm/builtin.cc index 86f0152ce7..8b27bb2d9e 100644 --- a/src/runtime/relax_vm/builtin.cc +++ b/src/runtime/relax_vm/builtin.cc @@ -327,6 +327,12 @@ TVM_REGISTER_GLOBAL("vm.builtin.null_value").set_body([](TVMArgs args, TVMRetVal *rv = nullptr; }); +TVM_REGISTER_GLOBAL("vm.builtin.to_device") + .set_body_typed([](NDArray data, int dev_type, int dev_id) { + Device dst_device = {(DLDeviceType)dev_type, dev_id}; + return data.CopyTo(dst_device); + }); + /*! * \brief Load the scalar value in cond and return the result value. * \param cond The condition diff --git a/src/script/ir_builder/ir/ir.cc b/src/script/ir_builder/ir/ir.cc index fb51886a7d..2f2785ca44 100644 --- a/src/script/ir_builder/ir/ir.cc +++ b/src/script/ir_builder/ir/ir.cc @@ -110,11 +110,41 @@ void ModuleGlobalInfos(Map<String, Array<GlobalInfo>> global_infos) { } } +VDevice LookupVDevice(String target_kind, int device_index) { + if (IRBuilder::IsInScope()) { + IRModuleFrame frame = FindModuleFrame(); + if (frame->global_infos.empty()) { + LOG(FATAL) << "ValueError: The GlobalInfos in the IRModule is not defined."; + } + Array<GlobalInfo> vdevices = frame->global_infos["vdevice"]; + if (vdevices.empty() || device_index < 0 || + static_cast<size_t>(device_index) >= vdevices.size()) { + LOG(FATAL) << "ValueError: The target VDevice in the GlobalInfos was not found."; + } + if (target_kind == "vdevice") { + return Downcast<VDevice>(vdevices[device_index]); + } + int count = 0; + for (auto vdevice : vdevices) { + auto vdev = Downcast<VDevice>(vdevice); + if (vdev->target->kind->name == target_kind) { + if (count == device_index) { + return vdev; + } + count++; + } + } + } + LOG(WARNING) << "The annotated device was not found, please check your vdevice list."; + return VDevice(); +} + TVM_REGISTER_GLOBAL("script.ir_builder.ir.IRModule").set_body_typed(IRModule); TVM_REGISTER_GLOBAL("script.ir_builder.ir.DeclFunction").set_body_typed(DeclFunction); TVM_REGISTER_GLOBAL("script.ir_builder.ir.DefFunction").set_body_typed(DefFunction); TVM_REGISTER_GLOBAL("script.ir_builder.ir.ModuleAttrs").set_body_typed(ModuleAttrs); TVM_REGISTER_GLOBAL("script.ir_builder.ir.ModuleGlobalInfos").set_body_typed(ModuleGlobalInfos); +TVM_REGISTER_GLOBAL("script.ir_builder.ir.LookupVDevice").set_body_typed(LookupVDevice); } // namespace ir } // namespace ir_builder diff --git a/src/script/printer/ir/ir.cc b/src/script/printer/ir/ir.cc index ecf92897a5..a239481d03 100644 --- a/src/script/printer/ir/ir.cc +++ b/src/script/printer/ir/ir.cc @@ -118,6 +118,16 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) return IR(d, "dummy_global_info")->Call({}); }); +TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) + .set_dispatch<VDevice>("", [](VDevice vdev, ObjectPath p, IRDocsifier d) -> Doc { + d->AddGlobalInfo("vdevice", vdev); + Map<String, ObjectRef> config = vdev->target->Export(); + return IR(d, "vdevice") + ->Call({d->AsDoc<ExprDoc>(config, p), + LiteralDoc::Int(vdev->vdevice_id, p->Attr("vdevice_id")), + LiteralDoc::Str(vdev->memory_scope, p->Attr("memory_scope"))}); + }); + TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) .set_dispatch<Op>("", [](Op op, ObjectPath p, IRDocsifier d) -> Doc { return IR(d, "Op")->Call({LiteralDoc::Str(op->name, p->Attr("name"))}); diff --git a/src/script/printer/ir_docsifier.cc b/src/script/printer/ir_docsifier.cc index 521ab07359..a424863495 100644 --- a/src/script/printer/ir_docsifier.cc +++ b/src/script/printer/ir_docsifier.cc @@ -63,6 +63,12 @@ ExprDoc IRDocsifierNode::AddMetadata(const ObjectRef& obj) { return IdDoc("metadata")[{LiteralDoc::Str(key, NullOpt)}][{LiteralDoc::Int(index, NullOpt)}]; } +void IRDocsifierNode::AddGlobalInfo(const String& name, const GlobalInfo& ginfo) { + ICHECK(ginfo.defined()) << "TypeError: Cannot add nullptr to global_infos"; + Array<GlobalInfo>& array = global_infos[name]; + array.push_back(ginfo); +} + bool IRDocsifierNode::IsVarDefined(const ObjectRef& obj) const { return obj2info.count(obj); } void IRDocsifierNode::RemoveVar(const ObjectRef& obj) { diff --git a/src/script/printer/relax/struct_info.cc b/src/script/printer/relax/struct_info.cc index 49162bb824..7fab5b59a2 100644 --- a/src/script/printer/relax/struct_info.cc +++ b/src/script/printer/relax/struct_info.cc @@ -111,6 +111,13 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) kwargs_keys.push_back("ndim"); kwargs_values.push_back(LiteralDoc::Int(n->ndim, n_p->Attr("ndim"))); } + if (n->vdevice.defined()) { + kwargs_keys.push_back("vdevice"); + std::string dev_kind = n->vdevice.value()->target->kind->name; + int dev_index = FindVDeviceIndexByTargetKind(n->vdevice.value(), d); + kwargs_values.push_back( + LiteralDoc::Str(dev_kind + ":" + std::to_string(dev_index), n_p->Attr("vdevice"))); + } if (args.empty() && kwargs_keys.empty()) { return Relax(d, "Tensor"); } diff --git a/src/script/printer/relax/utils.h b/src/script/printer/relax/utils.h index 88fc7491c2..e0b5348d73 100644 --- a/src/script/printer/relax/utils.h +++ b/src/script/printer/relax/utils.h @@ -97,6 +97,22 @@ Array<StmtDoc> PrintSeqExpr(const relax::SeqExpr& n, const ObjectPath& n_p, cons ExprDoc PrintShapeVar(const PrimExpr& e, const ObjectPath& e_p, const IRDocsifier& d); +inline int FindVDeviceIndexByTargetKind(const VDevice& vdevice, const IRDocsifier& d) { + Array<GlobalInfo> vdevices = d->global_infos["vdevice"]; + int kind_index = 0; + for (size_t i = 0; i < vdevices.size(); ++i) { + auto vdev = Downcast<VDevice>(vdevices[i]); + if (vdev.same_as(vdevice)) { + return kind_index; + } + if (vdev->target->kind->name == vdevice->target->kind->name) { + kind_index++; + } + } + LOG(WARNING) << "The VDevice was not found in the global_infos map: " << vdevice; + return -1; +} + } // namespace printer } // namespace script } // namespace tvm diff --git a/src/target/codegen.cc b/src/target/codegen.cc index bbb2c15a64..55af8889e1 100644 --- a/src/target/codegen.cc +++ b/src/target/codegen.cc @@ -38,6 +38,9 @@ #include <vector> namespace tvm { + +using FTVMTIRToRuntime = runtime::TypedPackedFunc<runtime::Module(IRModule, Target)>; + namespace codegen { runtime::Module Build(IRModule mod, Target target) { diff --git a/src/target/target.cc b/src/target/target.cc index 2f585188d0..cd2e3714e4 100644 --- a/src/target/target.cc +++ b/src/target/target.cc @@ -21,6 +21,7 @@ * \file src/target/target.cc */ #include <dmlc/thread_local.h> +#include <tvm/ir/transform.h> #include <tvm/runtime/device_api.h> #include <tvm/runtime/logging.h> #include <tvm/runtime/registry.h> @@ -91,16 +92,6 @@ void CheckAndUpdateHostConsistency(Target* target, Target* host) { *host = (*target)->GetHost().value_or(Target()); } -void CheckAndUpdateHostConsistency(Map<Target, IRModule>* targets, Target* host) { - Map<Target, IRModule> new_targets; - for (auto& it : *targets) { - auto target = it.first; - CheckAndUpdateHostConsistency(&target, host); - new_targets.Set(target, it.second); - } - *targets = new_targets; -} - static std::vector<String> DeduplicateKeys(const std::vector<String>& keys) { std::vector<String> new_keys; for (size_t i = 0; i < keys.size(); ++i) { @@ -614,8 +605,8 @@ Target::Target(TargetKind kind, Optional<ObjectRef> host, String tag, Array<Stri bool Target::IsExternalCodegen() const { TargetKindAttrMap<Bool> is_external_codegen_map = TargetKind::GetAttrMap<Bool>(tvm::attr::kIsExternalCodegen); - TargetKindAttrMap<FTVMRelayToTIR> relay_to_tir_map = - TargetKind::GetAttrMap<FTVMRelayToTIR>(tvm::attr::kRelayToTIR); + TargetKindAttrMap<tvm::transform::Pass> relay_to_tir_map = + TargetKind::GetAttrMap<tvm::transform::Pass>(tvm::attr::kRelayToTIR); return is_external_codegen_map.get(get()->kind, Bool(false)) || relay_to_tir_map.count(get()->kind); } diff --git a/tests/cpp/target_test.cc b/tests/cpp/target_test.cc index 37a8eeb448..9d05023e79 100644 --- a/tests/cpp/target_test.cc +++ b/tests/cpp/target_test.cc @@ -458,7 +458,7 @@ TVM_REGISTER_TARGET_KIND("test_external_codegen_2", kDLMetal) .set_attr<Bool>(tvm::attr::kIsExternalCodegen, Bool(true)); TVM_REGISTER_TARGET_KIND("test_external_codegen_3", kDLCPU) - .set_attr<FTVMRelayToTIR>(tvm::attr::kRelayToTIR, tvm::relay::transform::InferType()); + .set_attr<tvm::transform::Pass>(tvm::attr::kRelayToTIR, tvm::relay::transform::InferType()); TEST(Target, ExternalCodegen) { Target regular("cuda"); diff --git a/python/tvm/ir/global_info.py b/tests/python/relax/test_json_compact.py similarity index 53% copy from python/tvm/ir/global_info.py copy to tests/python/relax/test_json_compact.py index 17011e76a6..1320ff1cd6 100644 --- a/python/tvm/ir/global_info.py +++ b/tests/python/relax/test_json_compact.py @@ -14,29 +14,37 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""Global Info.""" -import tvm -from tvm.runtime.object import Object -from . import _ffi_api - -class GlobalInfo(Object): - """Base node for all global info that can appear in the IR""" - - def __eq__(self, other): - """Compare two struct info for structural equivalence.""" - return tvm.ir.structural_equal(self, other) +import tvm +import tvm.testing +from tvm import relax +import json - def __ne__(self, other): - return not self.__eq__(other) - def same_as(self, other): - """Overload with structural equality.""" - return super().__eq__(other) +# 0.13 BACKWARDS COMPATIBILITY TESTS +def test_vdevice(): + nodes = [ + {"type_key": ""}, + { + "type_key": "relax.TensorStructInfo", + "attrs": { + "dtype": "float32", + "ndim": "-1", + "shape": "0", + "span": "0", + }, + }, + ] + data = { + "root": 1, + "nodes": nodes, + "attrs": {"tvm_version": "0.13.0"}, + "b64ndarrays": [], + } + tsinfo = tvm.ir.load_json(json.dumps(data)) + assert isinstance(tsinfo, relax.TensorStructInfo) + assert not tsinfo.vdevice -class DummyGlobalInfo(GlobalInfo): - def __init__(self) -> None: - self.__init_handle_by_constructor__( - _ffi_api.DummyGlobalInfo, - ) +if __name__ == "__main__": + tvm.testing.main() diff --git a/tests/python/relax/test_relax_operators.py b/tests/python/relax/test_relax_operators.py index 90608df4b6..e0904477d4 100644 --- a/tests/python/relax/test_relax_operators.py +++ b/tests/python/relax/test_relax_operators.py @@ -247,5 +247,35 @@ def test_op_call_pure_packed(): assert (copy_found.numpy() == arr).all() +def test_op_to_device(): + @tvm.script.ir_module + class CallToDevice: + @R.function + def to_dev(x: R.Tensor((3, 4), "float32")): + z = R.call_pure_packed( + "vm.builtin.to_device", + x, + 1, + 0, + sinfo_args=(R.Tensor((3, 4), dtype="float32")), + ) + return z + + np.random.seed(0) # to avoid flakiness + arr = np.random.rand(3, 4).astype("float32") + copy_found = run_cpu(CallToDevice, "to_dev", tvm.nd.array(arr)) + assert (copy_found.numpy() == arr).all() + + +def test_op_to_vdevice(): + @tvm.script.ir_module + class ToVDevice: + @R.function + def to_vdev(x: R.Tensor((3, 4), "float32")): + dst_vdev = tvm.ir.VDevice("llvm", 0, "global") + ret = R.to_vdevice(x, dst_vdev) + return ret + + if __name__ == "__main__": tvm.testing.main() diff --git a/tests/python/relax/test_tvmscript_parser.py b/tests/python/relax/test_tvmscript_parser.py index bc324fe364..39a4d33ca6 100644 --- a/tests/python/relax/test_tvmscript_parser.py +++ b/tests/python/relax/test_tvmscript_parser.py @@ -22,7 +22,7 @@ import tvm import tvm.script import tvm.testing from tvm import IRModule, relax, tir, topi -from tvm.ir import DummyGlobalInfo +from tvm.ir import VDevice, DummyGlobalInfo from tvm.script.parser import ir as I from tvm.script.parser import relax as R from tvm.script.parser import tir as T @@ -303,6 +303,56 @@ def test_module_with_attr_and_global_info(): _check(TestModule, mod) +def test_global_info_vdevice(): + vdevices = [ + VDevice("llvm"), + VDevice("cuda", 0), + VDevice("cuda -arch=sm_80", 0), + VDevice("metal", 0, "global"), + ] + + @I.ir_module + class TestModule: + I.module_attrs({"attr": 10}) + I.module_global_infos( + { + "vdevice": [ + I.vdevice("llvm"), + I.vdevice("cuda", 0), + I.vdevice("cuda -arch=sm_80", 0), + I.vdevice("metal", 0, "global"), + ] + } + ) + + @T.prim_func(private=True) + def tir_func( + x: T.Buffer((T.int64(128), T.int64(128)), "float32"), + y: T.Buffer((T.int64(128), T.int64(128)), "float32"), + ): + T.func_attr({"tir.noalias": True}) + for i, j in T.grid(T.int64(128), T.int64(128)): + with T.block(): + vi, vj = T.axis.remap("SS", [i, j]) + y[vi, vj] = x[vi, vj] + 1.0 + + @R.function + def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor((128, 128), "float32"): + cls = TestModule + gv0 = R.call_tir(cls.tir_func, x, R.Tensor((128, 128), dtype="float32")) + return gv0 + + x = relax.Var("x", R.Tensor((128, 128), "float32")) + bb = relax.BlockBuilder() + with bb.function("foo", (x,)): + out = bb.emit_te(lambda x: x + 1, x, primfunc_name_hint="tir_func") + bb.emit_func_output(out) + mod = bb.get() + mod.update_global_info("vdevice", vdevices) + mod = mod.with_attr("attr", tvm.tir.IntImm("int32", 10)) + _check(TestModule, mod) + + def test_relax_tensor_op(): @R.function def foo(x: R.Tensor((4, 4), "float32")) -> R.Tensor((4, 4), "float32"): @@ -714,6 +764,51 @@ def test_tensor_type_without_args(): _check(foo, bb.get()["foo"]) +def test_tensor_with_vdevice(): + vdevices = [ + VDevice("llvm"), + VDevice("cuda", 0), + VDevice("metal", 0, "global"), + VDevice("cuda -arch=sm_80", 0), + ] + + @I.ir_module + class TestModule: + I.module_attrs({"attr": 10}) + I.module_global_infos( + { + "vdevice": [ + I.vdevice("llvm"), + I.vdevice("cuda", 0), + I.vdevice("metal", 0, "global"), + I.vdevice("cuda -arch=sm_80", 0), + ] + } + ) + + @R.function + def foo( + a: R.Tensor((128, 128), "float32", "cuda:1"), # noqa: F722 + b: R.Tensor((128, 128), "float32", "llvm"), + c: R.Tensor((128, 128), "float32", "vdevice:3"), # noqa: F722 + ) -> R.Tensor((128, 128), "float32"): + s = R.add(a, c) + return s + + a = relax.Var("a", R.Tensor((128, 128), "float32", vdevices[3])) + b = relax.Var("b", R.Tensor((128, 128), "float32", vdevices[0])) + c = relax.Var("c", R.Tensor((128, 128), "float32", vdevices[3])) + bb = relax.BlockBuilder() + with bb.function("foo", (a, b, c)): + out = bb.emit(relax.op.add(a, c)) + bb.emit_func_output(out) + mod = bb.get() + mod = mod.with_attr("attr", tvm.tir.IntImm("int32", 10)) + mod.update_global_info("vdevice", vdevices) + + _check(TestModule, mod) + + def test_direct_return(): @R.function def foo(x: R.Tensor((32, 32), "float32")) -> R.Tensor((32, 32), "float32"): diff --git a/tests/python/relax/test_tvmscript_parser_op_manipulate.py b/tests/python/relax/test_tvmscript_parser_op_manipulate.py index a80e8aad37..25f4f08520 100644 --- a/tests/python/relax/test_tvmscript_parser_op_manipulate.py +++ b/tests/python/relax/test_tvmscript_parser_op_manipulate.py @@ -403,5 +403,20 @@ def test_flip(): _check(foo, bb.get()["foo"]) +def test_to_vdevice(): + @R.function + def foo(x: R.Tensor((), "int32")) -> R.Tensor((), "int32"): + tensor = R.to_vdevice(x, tvm.ir.VDevice("llvm", 0, "global")) + return tensor + + x = relax.Var("x", R.Tensor((), "int32")) + bb = relax.BlockBuilder() + with bb.function("foo", (x,)): + tensor = bb.emit(relax.op.to_vdevice(x, tvm.ir.VDevice("llvm", 0, "global"))) + bb.emit_func_output(tensor) + + _check(foo, bb.get()["foo"]) + + if __name__ == "__main__": tvm.testing.main() diff --git a/tests/python/relax/test_vm_build.py b/tests/python/relax/test_vm_build.py index 0b34d24540..085b6137ac 100644 --- a/tests/python/relax/test_vm_build.py +++ b/tests/python/relax/test_vm_build.py @@ -725,6 +725,41 @@ def test_recursion(exec_mode): tvm.testing.assert_allclose(res.numpy(), np.power(2.0, recursion_runs), rtol=1e-7, atol=1e-7) +@tvm.testing.requires_gpu +@pytest.mark.parametrize("exec_mode", EXEC_MODE) +def test_vm_to_device(exec_mode): + @tvm.script.ir_module + class TestToVDevice: + @R.function + def foo1( + x: R.Tensor((2, 3), "float32"), + ) -> R.Tensor((2, 3), "float32"): + copied = R.to_vdevice(x, tvm.ir.VDevice("cuda", 0, "global")) + return copied + + @R.function + def foo2( + x: R.Tensor((2, 3), "float32"), + ) -> R.Tensor((2, 3), "float32"): + copied = R.to_vdevice(x, tvm.ir.VDevice("llvm", 0, "global")) + return copied + + mod = TestToVDevice + target = tvm.target.Target("llvm", host="llvm") + ex = relax.build(mod, target, exec_mode=exec_mode) + vm = relax.VirtualMachine(ex, tvm.cpu()) + x_inp = tvm.nd.array(np.random.rand(2, 3).astype("float32")) + res_1 = check_saved_func(vm, "foo1", x_inp) + res_2 = check_saved_func(vm, "foo2", x_inp) + + # check the copied tensor's device + assert str(res_1.device) == "cuda(0)" + assert str(res_2.device) == "cpu(0)" + + tvm.testing.assert_allclose(res_1.numpy(), x_inp.numpy()) + tvm.testing.assert_allclose(res_2.numpy(), x_inp.numpy()) + + @pytest.mark.parametrize("exec_mode", EXEC_MODE) def test_vm_closure(exec_mode): @tvm.script.ir_module diff --git a/tests/python/relax/test_vm_codegen_only.py b/tests/python/relax/test_vm_codegen_only.py index ffa9837d02..d9fb130f3c 100644 --- a/tests/python/relax/test_vm_codegen_only.py +++ b/tests/python/relax/test_vm_codegen_only.py @@ -57,6 +57,33 @@ def test_vm_copy(exec_mode): tvm.testing.assert_allclose(res.numpy(), inp.numpy(), rtol=1e-7, atol=1e-7) +@pytest.mark.parametrize("exec_mode", EXEC_MODE) +def test_vm_to_device(exec_mode): + @tvm.script.ir_module + class TestVMToDevice: + @R.function + def foo(x: R.Tensor((3, 4), "float32")): + R.func_attr({"global_symbol": "foo"}) + # Copy x to the first cpu: device_type=1 and device_id=0. + # More device info. please take a look at python/tvm/_ffi/runtime_ctypes.py + z = R.call_packed( + "vm.builtin.to_device", x, 1, 0, sinfo_args=(R.Tensor((3, 4), dtype="float32")) + ) + return z + + mod = TestVMToDevice + target = tvm.target.Target("llvm", host="llvm") + ex = codegen(mod, target, exec_mode) + inp = tvm.nd.array(np.random.rand(3, 4).astype(np.float32)) + vm = relax.VirtualMachine(ex, tvm.cpu()) + res = check_saved_func(vm, "foo", inp) + tvm.testing.assert_allclose(res.numpy(), inp.numpy(), rtol=1e-7, atol=1e-7) + # check the resulting tensor is on cpu:0 + assert str(res.device) == "cpu(0)" + assert res.device.device_type == 1 + assert res.device.device_id == 0 + + @pytest.mark.parametrize("exec_mode", EXEC_MODE) def test_if_cond_const(exec_mode): @tvm.script.ir_module