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 29f1183600 [REFACTOR][RELAX] Own global device metadata in Relax
(#20363)
29f1183600 is described below
commit 29f118360074833477eea072cf41b3cb57e6a1bc
Author: Tianqi Chen <[email protected]>
AuthorDate: Wed Sep 16 11:43:56 2026 -0400
[REFACTOR][RELAX] Own global device metadata in Relax (#20363)
Relax owns the concrete global metadata used by its device annotations.
Move `VDevice`, `MemoryScope`, and `DummyGlobalInfo` into the Relax C++
and Python namespaces, with their registrations and consumers, while
keeping the generic `GlobalInfo` base and IRModule storage in IR.
Existing TVMScript global-info syntax is preserved.
Remove the unused legacy `target.VirtualDevice` API and its dedicated
wrappers and tests.
---
include/tvm/ir/expr.h | 3 -
include/tvm/ir/global_info.h | 68 +---
include/tvm/relax/attrs/op.h | 2 +-
include/tvm/{ir => relax}/global_info.h | 39 +--
include/tvm/relax/type.h | 2 +-
include/tvm/target/virtual_device.h | 382 ---------------------
include/tvm/tirx/layout.h | 1 +
python/tvm/ir/__init__.py | 2 +-
python/tvm/ir/global_info.py | 32 +-
python/tvm/relax/__init__.py | 3 +
.../virtual_device.py => relax/global_info.py} | 39 ++-
python/tvm/relax/script/builder/ir.py | 3 +-
python/tvm/relax/transform/transform.py | 5 +-
python/tvm/relax/type.py | 3 +-
python/tvm/relax/utils.py | 3 +-
python/tvm/script/ir_builder/ir/ir.py | 14 +-
python/tvm/target/__init__.py | 1 -
src/{ => relax}/ir/global_info.cc | 10 +-
src/relax/script/printer/call.cc | 2 +-
src/relax/script/printer/utils.h | 4 +-
src/script/ir_builder/ir/ir.cc | 9 +-
src/script/printer/ir/ir.cc | 7 +-
src/target/virtual_device.cc | 201 -----------
tests/cpp/target/virtual_device_test.cc | 129 -------
tests/python/relax/test_analysis_type_analysis.py | 16 +-
.../relax/test_backend_dispatch_sort_scan.py | 2 +-
tests/python/relax/test_op_binary.py | 3 +-
tests/python/relax/test_op_create.py | 3 +-
tests/python/relax/test_op_image.py | 3 +-
tests/python/relax/test_op_index.py | 3 +-
tests/python/relax/test_op_linear_algebra.py | 3 +-
tests/python/relax/test_op_manipulate.py | 3 +-
tests/python/relax/test_op_nn.py | 3 +-
tests/python/relax/test_op_nn_convolution.py | 3 +-
tests/python/relax/test_op_nn_pooling.py | 3 +-
tests/python/relax/test_op_search.py | 3 +-
tests/python/relax/test_op_set.py | 3 +-
tests/python/relax/test_op_sort.py | 3 +-
tests/python/relax/test_op_statistical.py | 3 +-
tests/python/relax/test_op_ternary.py | 3 +-
tests/python/relax/test_op_unary.py | 3 +-
tests/python/relax/test_relax_operators.py | 2 +-
.../python/relax/test_transform_realize_vdevice.py | 2 +-
.../python/relax/test_transform_update_vdevice.py | 2 +-
tests/python/relax/test_tvmscript_parser.py | 2 +-
tests/python/relax/test_vm_build.py | 4 +-
tests/python/target/test_virtual_device.py | 51 ---
47 files changed, 128 insertions(+), 962 deletions(-)
diff --git a/include/tvm/ir/expr.h b/include/tvm/ir/expr.h
index b245795469..c02b804bab 100644
--- a/include/tvm/ir/expr.h
+++ b/include/tvm/ir/expr.h
@@ -43,9 +43,6 @@
namespace tvm {
-// Forward-declare VirtualDevice to avoid circular imports.
-class VirtualDevice;
-
/*! \brief Tuple container */
class TupleNode : public ExprNode {
public:
diff --git a/include/tvm/ir/global_info.h b/include/tvm/ir/global_info.h
index 3533b78687..46d7cce46c 100644
--- a/include/tvm/ir/global_info.h
+++ b/include/tvm/ir/global_info.h
@@ -19,23 +19,15 @@
/*!
* \file tvm/ir/global_info.h
- * \brief GlobalInfo are globally static object that are referred by the IR
itself.
+ * \brief Base for globally static objects referred to by the IR.
*/
-
#ifndef TVM_IR_GLOBAL_INFO_H_
#define TVM_IR_GLOBAL_INFO_H_
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/ir/expr.h>
-#include <tvm/target/target.h>
+#include <tvm/ffi/object.h>
namespace tvm {
-/*!
- * \brief Abstract label for an area of memory.
- */
-using MemoryScope = ffi::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
@@ -56,62 +48,6 @@ class GlobalInfo : public ffi::ObjectRef {
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(GlobalInfo, ffi::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;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<VDeviceNode>()
- .def_ro("target", &VDeviceNode::target)
- .def_ro("vdevice_id", &VDeviceNode::vdevice_id)
- .def_ro("memory_scope", &VDeviceNode::memory_scope);
- }
-
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.VDevice", 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_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(VDevice, GlobalInfo, VDeviceNode);
-};
-
-/*!
- * \brief A dummy global info sub-class for testing purpose.
- */
-class DummyGlobalInfoNode : public GlobalInfoNode {
- public:
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<DummyGlobalInfoNode>();
- }
-
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.DummyGlobalInfo", DummyGlobalInfoNode,
GlobalInfoNode);
-};
-
-/*!
- * \brief Managed reference to DummyGlobalInfoNode.
- * \sa DummyGlobalInfoNode
- */
-class DummyGlobalInfo : public GlobalInfo {
- public:
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(DummyGlobalInfo, GlobalInfo,
DummyGlobalInfoNode);
-};
-
} // namespace tvm
#endif // TVM_IR_GLOBAL_INFO_H_
diff --git a/include/tvm/relax/attrs/op.h b/include/tvm/relax/attrs/op.h
index 4c1451c3dc..54d6bba916 100644
--- a/include/tvm/relax/attrs/op.h
+++ b/include/tvm/relax/attrs/op.h
@@ -24,8 +24,8 @@
#ifndef TVM_RELAX_ATTRS_OP_H_
#define TVM_RELAX_ATTRS_OP_H_
-#include <tvm/ir/global_info.h>
#include <tvm/relax/expr.h>
+#include <tvm/relax/global_info.h>
namespace tvm {
namespace relax {
diff --git a/include/tvm/ir/global_info.h b/include/tvm/relax/global_info.h
similarity index 71%
copy from include/tvm/ir/global_info.h
copy to include/tvm/relax/global_info.h
index 3533b78687..4cb49064b8 100644
--- a/include/tvm/ir/global_info.h
+++ b/include/tvm/relax/global_info.h
@@ -18,44 +18,24 @@
*/
/*!
- * \file tvm/ir/global_info.h
- * \brief GlobalInfo are globally static object that are referred by the IR
itself.
+ * \file tvm/relax/global_info.h
+ * \brief Global information used by Relax.
*/
-
-#ifndef TVM_IR_GLOBAL_INFO_H_
-#define TVM_IR_GLOBAL_INFO_H_
+#ifndef TVM_RELAX_GLOBAL_INFO_H_
+#define TVM_RELAX_GLOBAL_INFO_H_
#include <tvm/ffi/reflection/registry.h>
-#include <tvm/ir/expr.h>
+#include <tvm/ir/global_info.h>
#include <tvm/target/target.h>
namespace tvm {
+namespace relax {
/*!
* \brief Abstract label for an area of memory.
*/
using MemoryScope = ffi::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
- */
-class GlobalInfoNode : public ffi::Object {
- public:
- static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind =
kTVMFFISEqHashKindTreeNode;
-
- TVM_FFI_DECLARE_OBJECT_INFO("ir.GlobalInfo", GlobalInfoNode, ffi::Object);
-};
-
-/*!
- * \brief Managed reference to GlobalInfoNode.
- * \sa GlobalInfoNode
- */
-class GlobalInfo : public ffi::ObjectRef {
- public:
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(GlobalInfo, ffi::ObjectRef,
GlobalInfoNode);
-};
-
/*!
* \brief A global info subclass for virtual devices.
*/
@@ -77,7 +57,7 @@ class VDeviceNode : public GlobalInfoNode {
.def_ro("memory_scope", &VDeviceNode::memory_scope);
}
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.VDevice", VDeviceNode, GlobalInfoNode);
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.VDevice", VDeviceNode,
GlobalInfoNode);
};
/*!
@@ -100,7 +80,7 @@ class DummyGlobalInfoNode : public GlobalInfoNode {
refl::ObjectDef<DummyGlobalInfoNode>();
}
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.DummyGlobalInfo", DummyGlobalInfoNode,
GlobalInfoNode);
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.DummyGlobalInfo",
DummyGlobalInfoNode, GlobalInfoNode);
};
/*!
@@ -112,6 +92,7 @@ class DummyGlobalInfo : public GlobalInfo {
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(DummyGlobalInfo, GlobalInfo,
DummyGlobalInfoNode);
};
+} // namespace relax
} // namespace tvm
-#endif // TVM_IR_GLOBAL_INFO_H_
+#endif // TVM_RELAX_GLOBAL_INFO_H_
diff --git a/include/tvm/relax/type.h b/include/tvm/relax/type.h
index 9fd8ee3b59..cb3c0e8196 100644
--- a/include/tvm/relax/type.h
+++ b/include/tvm/relax/type.h
@@ -28,9 +28,9 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/attrs.h>
#include <tvm/ir/env_func.h>
-#include <tvm/ir/global_info.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/ir/type.h>
+#include <tvm/relax/global_info.h>
#include <string>
#include <utility>
diff --git a/include/tvm/target/virtual_device.h
b/include/tvm/target/virtual_device.h
deleted file mode 100644
index 83c7f5655a..0000000000
--- a/include/tvm/target/virtual_device.h
+++ /dev/null
@@ -1,382 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-/*!
- * \file tvm/target/virtual_device.h
- * \brief A compile time representation for where data is to be stored at
runtime, and how to
- * compile code to compute it.
- */
-
-#ifndef TVM_TARGET_VIRTUAL_DEVICE_H_
-#define TVM_TARGET_VIRTUAL_DEVICE_H_
-
-#include <tvm/ir/transform.h>
-#include <tvm/target/target.h>
-
-#include <string>
-#include <unordered_set>
-#include <utility>
-
-namespace tvm {
-
-/*!
- * \brief Abstract label for an area of memory.
- *
- * Currently uninterpreted and arbitrary. Likely to be replaced by a
structured representation
- * of a memory pool in the future. Please try to use this alias instead of
ffi::String to aid future
- * code migration.
- */
-using MemoryScope = ffi::String;
-
-// NOTE: cannot use enum as they are out of bound of the original enum
-// and results in an undefined behavior
-// A 'null' device type, does not correspond to any DLDeviceType enum.
-// TODO(mbs): This is to help us as we transition away from representing the
'homogenous' case
-// as a singleton target map indexed by the invalid DLDeviceType '0'.
-constexpr int kNullDeviceType = 0;
-
-// An 'invalid' device type, does not correspond to any DLDeviceType enum.
-constexpr int kInvalidDeviceType = -1;
-
-/*!
- * \brief Describes at compile time the constraints on where data is to be
stored at runtime
- * down to the (virtual) device and memory scope level, and how to compile
code to compute that
- * data. Used by the \p PlanDevices pass to collect and solve (virtual) device
constraints for
- * the whole Relax program.
- *
- * Is a quadruple of:
- * - A \p device_type (\p DLDeviceType). May be \p kInvalidDeviceType if
unconstrained.
- * - A \p virtual_device_id (\p int). This allows us to distinguish distinct
devices
- * with the same \p Target, for example in a multi-GPU system. May be -1 if
unconstrained.
- * See "Virtual Devices" below.
- * - A \p target (\p Target) describing how to compile code for the intended
device. May be null
- * if unconstrained.
- * - A \p memory_scope (\p MemoryScope, which is currently just \p String)
describing which memory
- * area is to be used to hold data. May be "" if unconstrained. See "Memory
Scopes and Devices"
- * below.
- *
- * Some or all of these fields may be unconstrained, signaling that device
planning is free to
- * choose a value consistent with the whole program. However if a \p target is
given then the \p
- * device_type must equal \p target->GetTargetDeviceType().
- *
- * Note that currently we assume if a function returns its result on a
particular (virtual) device
- * then the function body is also executed on that device.
- *
- *
- * By 'execution' we include both (fused) primitive operators, and all the
Relax expressions
- * surrounding them which coordinates data and control flow. Again, typically
non-primitive
- * operators must be executed on a 'CPU'-like device with good support for
control flow.
- *
- * Since TVM targets such a wide range of systems it is not possible for \p
VirtualDevice to impose
- * much semantics on these fields, particularly for \p virtual_device_id and
\p memory_scope.
- * Instead we assume downstream passes and codegen will interpret an validate
these fields
- * appropriately.
- *
- * Targets vs Devices
- * ------------------
- * Generally \p Targets (a compile-time only datastructue) describe compiler
options for a specific
- * microarchitecture and toolchain, while \p Devices (a runtime datastructure
also available at
- * compile time) describe a physical device on the target system. Obviously
the target must agree
- * with the device's microarchitecture, but we otherwise don't impose any
constraints between them:
- * - It's ok to use different \p Targets for the same \p Device, eg to
squeeze some extra perf
- * out of a particular primitive using particular compiler flags.
- * - It's ok to use the same \p Target for multiple \p Devices, eg if we have
multiple CPUs.
- *
- * Traditionally TVM assumes at most one \p Target per \p DLDeviceType. We are
moving away from that
- * assumption.
- *
- * Virtual vs Physical Devices
- * ---------------------------
- * The \p virtual_device_id may be used by downstream passes or the runtime to
help decide which
- * \p device_id to use for a particular physical runtime \p Device. For
example:
- * - Some runtimes may support passing in an array of actual `device`
specifications, and the
- * \p virtual_device_id can be used at runtime as an index into that array.
- * - Some runtimes may support dynamically allocating computations to
physical devices. On these
- * systems a large space of \p virtual_device_ids could be used at compile
time, even though
- * at runtime only a few physical devices will be present.
- *
- * The \p virtual_device_id may also be left unconstrained if not needed.
- *
- * Memory Scopes and Devices
- * -------------------------
- * Multi-device systems can have complex memory hierarchies. For example
- * \code
- * (kDLCPU, 0, "llvm", "global")
- * \endcode
- * and
- * \code
- * (kDLCPU, 1, "llvm", "global")
- * \endcode
- * could denote:
- * - The same memory area accessible from two separate CPUs without any CPU
affinity;
- * - Distinct memory areas in a NUMA architecture for which cross-device
access is handled
- * by the memory system;
- * - Outright distinct memory areas, where one device cannot directly address
the memory of
- * another.
- *
- * Similarly:
- * \code
- * (kDLCPU, 0, "llvm", "global")
- * \endcode
- * and
- * \code
- * (kDLCUDA, 0, "cuda", "host")
- * \endcode
- * could denote the same memory area, but with very different access costs.
- *
- * Furthermore, not all memory scopes are accessible to all devices, and it is
possible for
- * a memory scope to only be accessible to a device when code is compiled with
particular
- * \p Target options.
- *
- * \p VirtualDevices themselves have no system-level understanding. Currently
the \p PlanDevices
- * pass will simply insert "device_copy" operators wherever \p VirtualDevices
are not exactly
- * pointwise equal. We may revisit this in the future as the work on memory
pools matures.
- *
- * Joining and Defaulting
- * ----------------------
- * It is possible to 'join' two \p VirtualDevices to yield the most
constrained \p VirtualDevice
- * which agrees with both join arguments. Eg:
- * \code
- * Join((kDLCPU, -1, "llvm", ""), (kInvalidDeviceType, 3, null, "global))
- * => (kDLCPU, 3, "llvm", "global")
- * Join((kDLCPU, -1, "llvm", ""), (kInvalidDeviceType, 3, null, "local))
- * => null (no join possible)
- * \endcode
- *
- * Related to 'join' is 'default', which only takes constrained fields from
the rhs when the
- * lhs is unconstrained:
- * \code
- * Default(kDLCPU, -1, "llvm", "local"), (kDLCPU, 3, null, "global"))
- * => (kDLCPU, 3, "llvm", "local")
- * \endcode
- *
- * These operations are needed during device planning.
- */
-
-class VirtualDeviceNode : public AttrsNode {
- private:
- /*!
- * \brief The \p DLDeviceType (represented as an int) of the virtual device.
If \p target is
- * known then this will be equal to \p target->GetTargetDeviceType(). If \p
target is null then
- * the target is to be determined later.
- *
- * This is needed to support the legacy "on_device" and "device_copy" calls
which only allow
- * a \p DLDeviceTypes (as an integer) to be given.
- *
- * kInvalidDeviceType denotes unconstrained. An int since the DLDeviceType
enum representation
- * is not fixed. Private to discourage further int vs DLDeviceType confusion.
- */
- int /* actually DLDeviceType */ device_type_int;
-
- public:
- DLDeviceType device_type() const { return
static_cast<DLDeviceType>(device_type_int); }
-
- /*!
- * \brief The device identifier for the virtual device. This must be
resolved to a physical
- * device identifier either during compilation or at runtime.
- *
- * -1 denotes unconstrained.
- */
- int virtual_device_id;
-
- /*!
- * \brief The \p Target describing how to compile for the virtual device.
- *
- * Null denotes unconstrained. Note that if a target later becomes known for
this \p VirtualDevice
- * then it must be consistent with the \p device_type if already known. This
is enforced by the
- * Join and Default methods.
- */
- Target target;
-
- /*!
- * \brief The scope of memory w.r.t. the virtual device which holds data.
- *
- * Empty denotes unconstrained.
- */
- MemoryScope memory_scope;
-
- /*!
- * \brief Returns true if virtual device is 'fully unconstrained', ie no
target/device type,
- * device id or memory scope is specified.
- */
- bool IsFullyUnconstrained() const {
- return !target.defined() && device_type() == kInvalidDeviceType &&
virtual_device_id == -1 &&
- memory_scope.empty();
- }
-
- /*!
- * \brief Returns true if virtual device is 'fully constrained', ie target,
device id and memory
- * scope are all specified.
- */
- bool IsFullyConstrained() const {
- return target.defined() && virtual_device_id != -1 &&
!memory_scope.empty();
- }
-
- /*!
- * \brief Returns the (virtual) \p Device implied by this \p VirtualDevice.
Both the \p
- * device_type and \p virtual_device_must be constrained. The returned \p
Device may not
- * correspond to any physical device available at compile time or even
runtime: see "Virtual vs
- * Physical Devices" above.
- */
- Device ToDevice() const {
- TVM_FFI_ICHECK(device_type_int != kInvalidDeviceType);
- TVM_FFI_ICHECK(virtual_device_id != -1);
- Device device;
- device.device_type = device_type();
- device.device_id = virtual_device_id;
- return device;
- }
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<VirtualDeviceNode>()
- .def_ro("device_type_int", &VirtualDeviceNode::device_type_int,
- "The type of the virtual device.",
refl::DefaultValue(kInvalidDeviceType))
- .def_ro("virtual_device_id", &VirtualDeviceNode::virtual_device_id,
- "The device id of the virtual device.", refl::DefaultValue(-1))
- .def_ro("target", &VirtualDeviceNode::target,
- "The target describing how to compile for the virtual device.",
- refl::DefaultValue(Target()))
- .def_ro("memory_scope", &VirtualDeviceNode::memory_scope,
- "The area of memory w.r.t. the virtual device where data is
stored.",
- refl::DefaultValue(""));
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("target.VirtualDevice", VirtualDeviceNode,
AttrsNode);
-
- friend class VirtualDevice;
-};
-
-/*!
- * \brief Managed reference class to \p VirtualDeviceNode.
- */
-class VirtualDevice : public ffi::ObjectRef {
- public:
- /*!
- * \brief Construct a virtual device.
- * \param device_type_int The device type for the virtual device, or \p
kInvalidDeviceType if
- * unconstrained. If \p target is defined then must match its \p
target->GetTargetDeviceType().
- * \param virtual_device_id The device id for the virtual device, or -1 if
unconstrained.
- * \param target The target describing how to compile for the virtual
device, or null if
- * unconstrained.
- * \param memory_scope The memory scope w.r.t. the virtual device which
holds data, or "" if
- * unconstrained.
- */
- TVM_DLL explicit VirtualDevice(int device_type_int = kInvalidDeviceType,
- int virtual_device_id = -1, Target target =
{},
- MemoryScope memory_scope = {});
-
- /*! \brief Returns the unique fully unconstrained \p VirtualDevice. */
- static VirtualDevice FullyUnconstrained();
-
- /*!
- * \brief Returns the \p VirtualDevice for \p device_type and (if not -1) \p
virtual_device_id.
- * The target and memory scope will be unconstrained.
- */
- static VirtualDevice ForDeviceType(DLDeviceType device_type, int
virtual_device_id = -1) {
- TVM_FFI_ICHECK_GT(device_type, 0);
- return VirtualDevice(device_type, virtual_device_id);
- }
- static VirtualDevice ForDeviceType(int device_type, int virtual_device_id =
-1) {
- return ForDeviceType(static_cast<DLDeviceType>(device_type),
virtual_device_id);
- }
- static VirtualDevice ForDeviceType(int64_t device_type, int
virtual_device_id = -1) {
- return ForDeviceType(static_cast<int>(device_type), virtual_device_id);
- }
-
- /*! \brief Returns the \p VirtualDevice for \p device. */
- static VirtualDevice ForDevice(const Device& device) {
- return ForDeviceType(device.device_type, device.device_id);
- }
-
- /*! \brief Returns the \p VirtualDevice for \p device and \p target. */
- static VirtualDevice ForDeviceAndTarget(const Device& device, Target target)
{
- return VirtualDevice(device.device_type, device.device_id,
std::move(target));
- }
-
- /*! \brief Returns the \p VirtualDevice for \p target. */
- static VirtualDevice ForTarget(Target target) {
- DLDeviceType device_type =
static_cast<DLDeviceType>(target->GetTargetDeviceType());
- return VirtualDevice(device_type, /*virtual_device_id=*/0,
std::move(target));
- }
-
- /*! \brief Returns the \p VirtualDevice for \p memory_scope alone. */
- static VirtualDevice ForMemoryScope(MemoryScope memory_scope) {
- return VirtualDevice(kInvalidDeviceType, -1, {}, std::move(memory_scope));
- }
-
- /*! \brief Returns the \p VirtualDevice for \p device, \p target and \p
memory_scope. */
- TVM_DLL static VirtualDevice ForDeviceTargetAndMemoryScope(const Device&
device, Target target,
- MemoryScope
memory_scope) {
- return VirtualDevice(device.device_type, device.device_id,
std::move(target),
- std::move(memory_scope));
- }
-
- /*!
- * \brief Returns the 'join' of \p lhs and \p rhs. The result will agree
pointwise with
- * \p lhs and \p rhs on all their constrained fields. Returns the null
optional if no such
- * join exists, ie there's disagreement on at least one constrained field.
- */
- TVM_DLL static ffi::Optional<VirtualDevice> Join(const VirtualDevice& lhs,
- const VirtualDevice& rhs);
-
- /*!
- * \brief Returns the 'default' of \p lhs and \p rhs. The result will be \p
lhs, except any
- * unconstrained fields in \p lhs will take their value from \p rhs. Always
well-defined.
- */
- TVM_DLL static VirtualDevice Default(const VirtualDevice& lhs, const
VirtualDevice& rhs);
-
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(VirtualDevice, ffi::ObjectRef,
VirtualDeviceNode);
-
- friend class VirtualDeviceCache; // Private implementation helper.
-};
-
-/*!
- * \brief A cache of \p VirtualDevices. This can be used:
- * - To avoid ending up with lots of identical instances, since the space of
VirtualDevices for any
- * one compilation is very small but the number of points they need to be
constructed can
- * be very large (eg during device planning).
- * - So we can assume \p VirtualDevices are pointer equal if and only if they
are structurally
- * equal. This simplifies the unification of 'device domains' which are built
on \p VirtualDevices.
- */
-class TVM_DLL VirtualDeviceCache {
- public:
- /*! \brief Returns the unique \p VirtualDevice representing given fields. */
- VirtualDevice Make(int device_type = kInvalidDeviceType, int
virtual_device_id = -1,
- Target target = {}, MemoryScope memory_scope = {});
-
- /*!
- * \brief Returns the unique \p VirtualDevice structurally equal to the
given \p virtual_device.
- */
- VirtualDevice Unique(const VirtualDevice& virtual_device);
-
- private:
- /*! \brief Already constructed VirtualDevices. */
- std::unordered_set<VirtualDevice, ffi::StructuralHash, ffi::StructuralEqual>
cache_;
-};
-
-/*! brief The attribute key for the virtual device. This key will be promoted
to first class on
- * functions. For use in the parser and printer only.
- *
- * Type: VirtualDevice
- */
-constexpr const char* kVirtualDevice = "virtual_device";
-
-} // namespace tvm
-
-#endif // TVM_TARGET_VIRTUAL_DEVICE_H_
diff --git a/include/tvm/tirx/layout.h b/include/tvm/tirx/layout.h
index 1580338ac4..529ea35823 100644
--- a/include/tvm/tirx/layout.h
+++ b/include/tvm/tirx/layout.h
@@ -29,6 +29,7 @@
#include <tvm/ffi/object.h>
#include <tvm/ir/attr_registry_map.h>
#include <tvm/ir/module.h>
+#include <tvm/target/target.h>
#include <tvm/tirx/exec_scope.h>
#include <tvm/tirx/var.h>
diff --git a/python/tvm/ir/__init__.py b/python/tvm/ir/__init__.py
index 977ec90dcc..bfb378d3bd 100644
--- a/python/tvm/ir/__init__.py
+++ b/python/tvm/ir/__init__.py
@@ -50,7 +50,7 @@ from .expr import (
is_prim_var,
)
from .function import BaseFunc, CallingConv
-from .global_info import GlobalInfo, DummyGlobalInfo, VDevice
+from .global_info import GlobalInfo
from .module import IRModule
from .op import Op, register_intrin_lowering, register_op_attr
diff --git a/python/tvm/ir/global_info.py b/python/tvm/ir/global_info.py
index a9198afac7..a964d4b2e5 100644
--- a/python/tvm/ir/global_info.py
+++ b/python/tvm/ir/global_info.py
@@ -18,10 +18,7 @@
import tvm_ffi
-import tvm
-from tvm.runtime import Device, Object
-
-from . import _ffi_api
+from tvm.runtime import Object
@tvm_ffi.register_object("ir.GlobalInfo")
@@ -38,30 +35,3 @@ class GlobalInfo(Object):
def same_as(self, other):
"""Overload with structural equality."""
return super().__eq__(other)
-
-
-@tvm_ffi.register_object("ir.DummyGlobalInfo")
-class DummyGlobalInfo(GlobalInfo):
- """DummyGlobalInfo"""
-
- def __init__(self) -> None:
- self.__init_handle_by_constructor__(
- _ffi_api.DummyGlobalInfo,
- )
-
-
-@tvm_ffi.register_object("ir.VDevice")
-class VDevice(GlobalInfo):
- """VDevice"""
-
- 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))
- if isinstance(target, Device):
- target = tvm.target.Target.from_device(target)
- self.__init_handle_by_constructor__(_ffi_api.VDevice, target,
vdevice_id, memory_scope)
diff --git a/python/tvm/relax/__init__.py b/python/tvm/relax/__init__.py
index c28ca1b2c0..fa3850c23d 100644
--- a/python/tvm/relax/__init__.py
+++ b/python/tvm/relax/__init__.py
@@ -22,6 +22,9 @@ from tvm.runtime import vm
from tvm.runtime.vm import VirtualMachine, VMInstrumentReturnKind
from tvm.ir import Call
+# Global information
+from .global_info import DummyGlobalInfo, VDevice
+
# Expr
from .expr import (
Expr,
diff --git a/python/tvm/target/virtual_device.py
b/python/tvm/relax/global_info.py
similarity index 51%
rename from python/tvm/target/virtual_device.py
rename to python/tvm/relax/global_info.py
index 37b888196b..2eacc8b072 100644
--- a/python/tvm/target/virtual_device.py
+++ b/python/tvm/relax/global_info.py
@@ -14,26 +14,39 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# ruff: noqa: F821
-"""Python bindings for creating VirtualDevices."""
+"""Global information used by Relax."""
import tvm_ffi
+import tvm
+from tvm.ir import GlobalInfo
+from tvm.runtime import Device
+
from . import _ffi_api
-@tvm_ffi.register_object("target.VirtualDevice")
-class VirtualDevice(tvm_ffi.core.Object):
- """A compile time representation for where data is to be stored at runtime,
- and how to compile code to compute it."""
+@tvm_ffi.register_object("relax.DummyGlobalInfo")
+class DummyGlobalInfo(GlobalInfo):
+ """DummyGlobalInfo"""
- def __init__(self, device=None, target=None, memory_scope="") -> None:
- if device is None:
- # The 'unconstrained' device has device type -1 and device id -1.
- device = tvm.device(-1, -1)
+ def __init__(self) -> None:
self.__init_handle_by_constructor__(
- _ffi_api.VirtualDevice_ForDeviceTargetAndMemoryScope, device,
target, memory_scope
+ _ffi_api.DummyGlobalInfo,
)
- def dlpack_device_type(self) -> int:
- return self.device_type_int
+
+@tvm_ffi.register_object("relax.VDevice")
+class VDevice(GlobalInfo):
+ """VDevice"""
+
+ 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))
+ if isinstance(target, Device):
+ target = tvm.target.Target.from_device(target)
+ self.__init_handle_by_constructor__(_ffi_api.VDevice, target,
vdevice_id, memory_scope)
diff --git a/python/tvm/relax/script/builder/ir.py
b/python/tvm/relax/script/builder/ir.py
index 1f19fbb05b..e3f7c5518a 100644
--- a/python/tvm/relax/script/builder/ir.py
+++ b/python/tvm/relax/script/builder/ir.py
@@ -25,7 +25,7 @@ from typing import Any
import tvm
from tvm import DataType, relax
-from tvm.ir import IRModule, VDevice
+from tvm.ir import IRModule
from tvm.relax import (
Call,
Expr,
@@ -38,6 +38,7 @@ from tvm.relax import (
const,
)
from tvm.relax.dpl import PatternMatchingRewriter
+from tvm.relax.global_info import VDevice
############################### Operators ###############################
from tvm.relax.op import (
diff --git a/python/tvm/relax/transform/transform.py
b/python/tvm/relax/transform/transform.py
index ff4601290a..a28943db1e 100644
--- a/python/tvm/relax/transform/transform.py
+++ b/python/tvm/relax/transform/transform.py
@@ -36,6 +36,7 @@ from tvm.runtime import Object, Tensor
from tvm.tirx import IndexMap, PrimFunc
from ..expr import Var
+from ..global_info import VDevice
from . import _ffi_api
from .legalize_ops.common import LegalizeFunc
@@ -483,12 +484,12 @@ def EliminateCommonSubexpr(call_only=False) ->
FunctionPass:
return _ffi_api.EliminateCommonSubexpr(call_only) # type: ignore
-def UpdateVDevice(new_vdevice: tvm.ir.VDevice, index: int) ->
tvm.ir.transform.Pass:
+def UpdateVDevice(new_vdevice: VDevice, index: int) -> tvm.ir.transform.Pass:
"""Update virtual device.
Parameters
----------
- new_vdevice : tvm.ir.VDevice
+ new_vdevice : tvm.relax.VDevice
The new virtual device.
index : int
The device index indicates the device on which the update will be
performed.
diff --git a/python/tvm/relax/type.py b/python/tvm/relax/type.py
index 98f34d06f0..b5018ec7ea 100644
--- a/python/tvm/relax/type.py
+++ b/python/tvm/relax/type.py
@@ -21,7 +21,8 @@
import tvm_ffi
from tvm_ffi import Array
-from tvm.ir import EnvFunc, PrimType, Span, TupleType, VDevice
+from tvm.ir import EnvFunc, PrimType, Span, TupleType
+from tvm.relax.global_info import VDevice
from . import _ffi_api
from .expr import Expr, ShapeExpr, Type
diff --git a/python/tvm/relax/utils.py b/python/tvm/relax/utils.py
index 78af8ac6b9..91ff545bb7 100644
--- a/python/tvm/relax/utils.py
+++ b/python/tvm/relax/utils.py
@@ -29,9 +29,10 @@ import tvm_ffi
from tvm_ffi import Array, Map
import tvm
+from tvm.relax.global_info import VDevice
from .. import tirx
-from ..ir import Attrs, Type, VDevice
+from ..ir import Attrs, Type
from ..te import Tensor as te_Tensor
from ..te import create_prim_func
from . import _ffi_api
diff --git a/python/tvm/script/ir_builder/ir/ir.py
b/python/tvm/script/ir_builder/ir/ir.py
index a61ad3c344..e2810ce698 100644
--- a/python/tvm/script/ir_builder/ir/ir.py
+++ b/python/tvm/script/ir_builder/ir/ir.py
@@ -18,13 +18,15 @@
from typing import TYPE_CHECKING, Any, TypeVar
-from tvm.ir import BaseFunc, DummyGlobalInfo, GlobalInfo, GlobalVar, VDevice
+from tvm.ir import BaseFunc, GlobalInfo, GlobalVar
from tvm.runtime import Object as tvm_Object
from . import _ffi_api
from .frame import IRModuleFrame
if TYPE_CHECKING:
+ from tvm.relax import DummyGlobalInfo, VDevice
+
T = TypeVar("T")
def meta_var(value: T) -> T:
@@ -162,17 +164,19 @@ def module_global_infos(global_infos: dict[str,
list[GlobalInfo]]) -> None:
############################### GlobalInfo ###############################
-def dummy_global_info() -> DummyGlobalInfo:
+def dummy_global_info() -> "DummyGlobalInfo":
"""Create a dummy global info expression.
Returns
-------
res : DummyGlobalInfo
The result dummy global info.
"""
+ from tvm.relax import DummyGlobalInfo # pylint:
disable=import-outside-toplevel
+
return DummyGlobalInfo() # type: ignore[attr-defined] # pylint:
disable=no-member
-def vdevice(target=None, vdevice_id: int = 0, memory_scope: str = "global") ->
VDevice:
+def vdevice(target=None, vdevice_id: int = 0, memory_scope: str = "global") ->
"VDevice":
"""Create a virtual device global info.
Parameters
----------
@@ -188,10 +192,12 @@ def vdevice(target=None, vdevice_id: int = 0,
memory_scope: str = "global") -> V
res : VDevice
The result virtual device.
"""
+ from tvm.relax import VDevice # pylint: disable=import-outside-toplevel
+
return VDevice(target, vdevice_id, memory_scope) # type:
ignore[attr-defined] # pylint: disable=no-member
-def lookup_vdevice(target_kind: str | None = None, device_index: int = -1) ->
VDevice:
+def lookup_vdevice(target_kind: str | None = None, device_index: int = -1) ->
"VDevice":
"""Retrieve a virtual device from the globalinfo vdevice list.
Parameters
----------
diff --git a/python/tvm/target/__init__.py b/python/tvm/target/__init__.py
index 7303a0d097..20e26bccf3 100644
--- a/python/tvm/target/__init__.py
+++ b/python/tvm/target/__init__.py
@@ -32,7 +32,6 @@ and :py:func:`tvm.target.register_tag` to register new tags.
"""
from .target import Target, TargetKind
-from .virtual_device import VirtualDevice
from .tag import list_tags, register_tag
from . import codegen
from . import tag_registry # registers tags on import
diff --git a/src/ir/global_info.cc b/src/relax/ir/global_info.cc
similarity index 85%
rename from src/ir/global_info.cc
rename to src/relax/ir/global_info.cc
index d8bba04c51..751e6d3186 100644
--- a/src/ir/global_info.cc
+++ b/src/relax/ir/global_info.cc
@@ -18,13 +18,14 @@
*/
/*!
- * \file src/ir/global_info.cc
+ * \file src/relax/ir/global_info.cc
* \brief Module global info.
*/
#include <tvm/ffi/reflection/registry.h>
-#include <tvm/ir/global_info.h>
+#include <tvm/relax/global_info.h>
namespace tvm {
+namespace relax {
TVM_FFI_STATIC_INIT_BLOCK() {
VDeviceNode::RegisterReflection();
@@ -33,7 +34,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("ir.DummyGlobalInfo", []() {
+ refl::GlobalDef().def("relax.DummyGlobalInfo", []() {
auto n = DummyGlobalInfo(ffi::make_object<DummyGlobalInfoNode>());
return n;
});
@@ -49,8 +50,9 @@ VDevice::VDevice(Target tgt, int dev_id, MemoryScope
mem_scope) {
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("ir.VDevice", [](Target tgt, int dev_id, MemoryScope
mem_scope) {
+ refl::GlobalDef().def("relax.VDevice", [](Target tgt, int dev_id,
MemoryScope mem_scope) {
return VDevice(tgt, dev_id, mem_scope);
});
}
+} // namespace relax
} // namespace tvm
diff --git a/src/relax/script/printer/call.cc b/src/relax/script/printer/call.cc
index fd79de5472..628cbfa5f8 100644
--- a/src/relax/script/printer/call.cc
+++ b/src/relax/script/printer/call.cc
@@ -210,7 +210,7 @@ ffi::Optional<ExprDoc> PrintToVDevice(const Call& n, const
AccessPath& n_p, cons
ffi::Array<ExprDoc> kwargs_values;
TVM_FFI_ICHECK(n->attrs.defined());
if (const auto* attrs = n->attrs.as<relax::ToVDeviceAttrs>()) {
- VDevice vdev = attrs->dst_vdevice;
+ relax::VDevice vdev = attrs->dst_vdevice;
std::string dev_kind = vdev->target->kind->name;
int dev_index = FindVDeviceIndexByTargetKind(vdev, d);
kwargs_keys.push_back("dst_vdevice");
diff --git a/src/relax/script/printer/utils.h b/src/relax/script/printer/utils.h
index cb7c8c4b46..c83e15fe2c 100644
--- a/src/relax/script/printer/utils.h
+++ b/src/relax/script/printer/utils.h
@@ -143,11 +143,11 @@ Doc PrintRelaxVar(tvm::Var n, AccessPath p, IRDocsifier
d);
ExprDoc PrintShapeVar(const PrimExpr& e, const AccessPath& e_p, const
IRDocsifier& d);
-inline int FindVDeviceIndexByTargetKind(const VDevice& vdevice, const
IRDocsifier& d) {
+inline int FindVDeviceIndexByTargetKind(const relax::VDevice& vdevice, const
IRDocsifier& d) {
ffi::Array<GlobalInfo> vdevices = d->global_infos["vdevice"];
int kind_index = 0;
for (size_t i = 0; i < vdevices.size(); ++i) {
- auto vdev = vdevices[i].as_or_throw<VDevice>();
+ auto vdev = vdevices[i].as_or_throw<relax::VDevice>();
if (vdev.same_as(vdevice)) {
return kind_index;
}
diff --git a/src/script/ir_builder/ir/ir.cc b/src/script/ir_builder/ir/ir.cc
index ab9c0e935f..872a9f851f 100644
--- a/src/script/ir_builder/ir/ir.cc
+++ b/src/script/ir_builder/ir/ir.cc
@@ -19,6 +19,7 @@
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/module.h>
+#include <tvm/relax/global_info.h>
#include <tvm/runtime/logging.h>
#include <tvm/script/ir_builder/ir/ir.h>
@@ -150,7 +151,7 @@ void ModuleGlobalInfos(ffi::Map<ffi::String,
ffi::Array<GlobalInfo>> global_info
}
}
-VDevice LookupVDevice(ffi::String target_kind, int device_index) {
+relax::VDevice LookupVDevice(ffi::String target_kind, int device_index) {
if (IRBuilder::IsInScope()) {
IRModuleFrame frame = FindModuleFrame();
if (frame->global_infos.empty()) {
@@ -162,11 +163,11 @@ VDevice LookupVDevice(ffi::String target_kind, int
device_index) {
TVM_FFI_THROW(ValueError) << "The target VDevice in the GlobalInfos was
not found.";
}
if (target_kind == "vdevice") {
- return vdevices[device_index].as_or_throw<VDevice>();
+ return vdevices[device_index].as_or_throw<relax::VDevice>();
}
int count = 0;
for (auto vdevice : vdevices) {
- auto vdev = vdevice.as_or_throw<VDevice>();
+ auto vdev = vdevice.as_or_throw<relax::VDevice>();
if (vdev->target->kind->name == target_kind) {
if (count == device_index) {
return vdev;
@@ -176,7 +177,7 @@ VDevice LookupVDevice(ffi::String target_kind, int
device_index) {
}
}
LOG(WARNING) << "The annotated device was not found, please check your
vdevice list.";
- return VDevice();
+ return relax::VDevice();
}
bool LookupName(const ffi::String& name) {
diff --git a/src/script/printer/ir/ir.cc b/src/script/printer/ir/ir.cc
index a2f1237210..11d4811cf8 100644
--- a/src/script/printer/ir/ir.cc
+++ b/src/script/printer/ir/ir.cc
@@ -17,6 +17,7 @@
* under the License.
*/
#include <tvm/ir/type.h>
+#include <tvm/relax/global_info.h>
#include "./utils.h"
@@ -171,15 +172,15 @@ TVM_FFI_STATIC_INIT_BLOCK() {
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<DummyGlobalInfo>(
+ IRDocsifier::vtable().set_dispatch<relax::DummyGlobalInfo>(
"", [](GlobalInfo ginfo, AccessPath p, IRDocsifier d) -> Doc {
return IR(d, "dummy_global_info")->Call({});
});
}
TVM_FFI_STATIC_INIT_BLOCK() {
- IRDocsifier::vtable().set_dispatch<VDevice>(
- "", [](VDevice vdev, AccessPath p, IRDocsifier d) -> Doc {
+ IRDocsifier::vtable().set_dispatch<relax::VDevice>(
+ "", [](relax::VDevice vdev, AccessPath p, IRDocsifier d) -> Doc {
d->AddGlobalInfo("vdevice", vdev);
ffi::Map<ffi::String, ffi::Any> config = vdev->target->ToConfig();
return IR(d, "vdevice")
diff --git a/src/target/virtual_device.cc b/src/target/virtual_device.cc
deleted file mode 100644
index 92803cd143..0000000000
--- a/src/target/virtual_device.cc
+++ /dev/null
@@ -1,201 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-/*!
- * \file tvm/target/virtual_device.cc
- * \brief A compile time representation for where data is to be stored at
runtime, and how to
- * compile code to compute it.
- */
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/runtime/device_api.h>
-#include <tvm/target/virtual_device.h>
-
-#include <sstream>
-
-namespace tvm {
-
-VirtualDevice::VirtualDevice(int device_type_int, int virtual_device_id,
Target target,
- MemoryScope memory_scope) {
- TVM_FFI_ICHECK(!target.defined() || device_type_int ==
target->GetTargetDeviceType())
- << "target " << target->str() << " has device type " <<
target->GetTargetDeviceType()
- << " but virtual device has device type " << device_type_int;
- auto node = ffi::make_object<VirtualDeviceNode>();
- node->device_type_int = device_type_int;
- node->virtual_device_id = virtual_device_id;
- node->target = std::move(target);
- node->memory_scope = std::move(memory_scope);
- data_ = std::move(node);
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- VirtualDeviceNode::RegisterReflection();
- refl::TypeAttrDef<VirtualDeviceNode>().def(
- refl::type_attr::kRepr, [](VirtualDevice vd, ffi::Function fn_repr) ->
ffi::String {
- auto* node = vd.get();
- std::ostringstream os;
- os << "VirtualDevice(";
- if (node->IsFullyUnconstrained()) {
- os << "?";
- } else {
- bool need_sep = false;
- if (node->device_type() != kInvalidDeviceType) {
- os << "device_type=" << node->device_type();
- need_sep = true;
- }
- if (node->virtual_device_id >= 0) {
- if (need_sep) os << ", ";
- os << "virtual_device_id=" << node->virtual_device_id;
- need_sep = true;
- }
- if (node->target.defined()) {
- if (need_sep) os << ", ";
- os << "target=" <<
fn_repr(ffi::AnyView(node->target)).cast<ffi::String>();
- need_sep = true;
- }
- if (!node->memory_scope.empty()) {
- if (need_sep) os << ", ";
- os << "memory_scope='" << node->memory_scope << "'";
- }
- }
- os << ")";
- return os.str();
- });
-}
-
-/* static */ VirtualDevice VirtualDevice::FullyUnconstrained() {
- static const VirtualDevice unconstrained{};
- return unconstrained;
-}
-
-/* static */
-ffi::Optional<VirtualDevice> VirtualDevice::Join(const VirtualDevice& lhs,
- const VirtualDevice& rhs) {
- if (lhs == rhs) {
- return lhs;
- }
- DLDeviceType joined_device_type;
- if (lhs->device_type() != kInvalidDeviceType) {
- joined_device_type = lhs->device_type();
- if (rhs->device_type() != kInvalidDeviceType && lhs->device_type() !=
rhs->device_type()) {
- return {};
- }
- } else {
- joined_device_type = rhs->device_type();
- }
- int joined_virtual_device_id;
- if (lhs->virtual_device_id >= 0) {
- joined_virtual_device_id = lhs->virtual_device_id;
- if (rhs->virtual_device_id >= 0 && lhs->virtual_device_id !=
rhs->virtual_device_id) {
- return {};
- }
- } else {
- joined_virtual_device_id = rhs->virtual_device_id;
- }
- Target joined_target;
- if (lhs->target.defined()) {
- joined_target = lhs->target;
- if (rhs->target.defined() && lhs->target != rhs->target) {
- return {};
- }
- } else {
- joined_target = rhs->target;
- }
- MemoryScope joined_memory_scope;
- if (!lhs->memory_scope.empty()) {
- joined_memory_scope = lhs->memory_scope;
- if (!rhs->memory_scope.empty() && lhs->memory_scope != rhs->memory_scope) {
- return {};
- }
- } else {
- joined_memory_scope = rhs->memory_scope;
- }
- return VirtualDevice(joined_device_type, joined_virtual_device_id,
joined_target,
- joined_memory_scope);
-}
-
-/* static */
-VirtualDevice VirtualDevice::Default(const VirtualDevice& lhs, const
VirtualDevice& rhs) {
- if (lhs == rhs) {
- return lhs;
- }
- DLDeviceType defaulted_device_type;
- if (lhs->device_type() != kInvalidDeviceType) {
- defaulted_device_type = lhs->device_type();
- } else {
- defaulted_device_type = rhs->device_type();
- }
- int defaulted_virtual_device_id;
- if (lhs->virtual_device_id >= 0) {
- defaulted_virtual_device_id = lhs->virtual_device_id;
- } else {
- defaulted_virtual_device_id = rhs->virtual_device_id;
- }
- Target defaulted_target;
- if (lhs->target.defined()) {
- defaulted_target = lhs->target;
- } else {
- // We can only default to the rhs's target if it is consistent with the
device type
- if (rhs->target.defined() && rhs->target->GetTargetDeviceType() ==
defaulted_device_type) {
- defaulted_target = rhs->target;
- }
- // else: leave as null
- }
- MemoryScope defaulted_memory_scope;
- if (!lhs->memory_scope.empty()) {
- defaulted_memory_scope = lhs->memory_scope;
- } else {
- defaulted_memory_scope = rhs->memory_scope;
- }
- return VirtualDevice(defaulted_device_type, defaulted_virtual_device_id,
defaulted_target,
- defaulted_memory_scope);
-}
-
-VirtualDevice VirtualDeviceCache::Make(int device_type, int virtual_device_id,
Target target,
- MemoryScope memory_scope) {
- VirtualDevice prototype(device_type, virtual_device_id, std::move(target),
- std::move(memory_scope));
- if (prototype->IsFullyUnconstrained()) {
- return VirtualDevice::FullyUnconstrained();
- }
- auto itr = cache_.find(prototype);
- if (itr == cache_.end()) {
- cache_.emplace(prototype);
- return prototype;
- } else {
- TVM_FFI_ICHECK_EQ(prototype->target.defined(), (*itr)->target.defined());
- if (prototype->target.defined()) {
- TVM_FFI_ICHECK_EQ(prototype->target->host.has_value(),
(*itr)->target->host.has_value());
- }
- return *itr;
- }
-}
-
-VirtualDevice VirtualDeviceCache::Unique(const VirtualDevice& virtual_device) {
- return Make(virtual_device->device_type(), virtual_device->virtual_device_id,
- virtual_device->target, virtual_device->memory_scope);
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("target.VirtualDevice_ForDeviceTargetAndMemoryScope",
- VirtualDevice::ForDeviceTargetAndMemoryScope);
-}
-
-} // namespace tvm
diff --git a/tests/cpp/target/virtual_device_test.cc
b/tests/cpp/target/virtual_device_test.cc
deleted file mode 100644
index 8e2000852d..0000000000
--- a/tests/cpp/target/virtual_device_test.cc
+++ /dev/null
@@ -1,129 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-#include <gtest/gtest.h>
-#include <tvm/target/target.h>
-#include <tvm/target/virtual_device.h>
-
-namespace tvm {
-namespace {
-
-TEST(VirtualDevice, Join_Defined) {
- {
- Target target_a = Target("cuda");
- VirtualDevice lhs = VirtualDevice(kDLCUDA, 3);
- VirtualDevice rhs = VirtualDevice(kDLCUDA, -1, target_a, "global");
- ffi::Optional<VirtualDevice> actual = VirtualDevice::Join(lhs, rhs);
- EXPECT_TRUE(actual.operator bool());
- VirtualDevice expected = VirtualDevice(kDLCUDA, 3, target_a, "global");
- EXPECT_TRUE(tvm::ffi::StructuralEqual()(actual.value(), expected));
- }
- {
- Target target_a = Target("cuda");
- VirtualDevice lhs = VirtualDevice(kDLCUDA, -1, target_a, "global");
- VirtualDevice rhs = VirtualDevice(kDLCUDA, 3);
- ffi::Optional<VirtualDevice> actual = VirtualDevice::Join(lhs, rhs);
- EXPECT_TRUE(actual.operator bool());
- VirtualDevice expected = VirtualDevice(kDLCUDA, 3, target_a, "global");
- EXPECT_TRUE(tvm::ffi::StructuralEqual()(actual.value(), expected));
- }
- {
- Target target_a = Target("cuda");
- VirtualDevice lhs = VirtualDevice(kDLCUDA);
- VirtualDevice rhs = VirtualDevice(kDLCUDA, 2, target_a);
- ffi::Optional<VirtualDevice> actual = VirtualDevice::Join(lhs, rhs);
- EXPECT_TRUE(actual.operator bool());
- VirtualDevice expected = VirtualDevice(kDLCUDA, 2, target_a);
- EXPECT_TRUE(tvm::ffi::StructuralEqual()(actual.value(), expected));
- }
- {
- Target target_a = Target("cuda");
- VirtualDevice lhs = VirtualDevice();
- VirtualDevice rhs = VirtualDevice(kDLCUDA, 3, target_a, "global");
- ffi::Optional<VirtualDevice> actual = VirtualDevice::Join(lhs, rhs);
- EXPECT_TRUE(actual.operator bool());
- VirtualDevice expected = rhs;
- EXPECT_TRUE(tvm::ffi::StructuralEqual()(actual.value(), expected));
- }
-}
-
-TEST(VirtualDevice, Join_Undefined) {
- {
- VirtualDevice lhs = VirtualDevice(kDLCUDA);
- VirtualDevice rhs = VirtualDevice(kDLCPU);
- ffi::Optional<VirtualDevice> actual = VirtualDevice::Join(lhs, rhs);
- EXPECT_FALSE(actual);
- }
- {
- VirtualDevice lhs = VirtualDevice(kDLCUDA, 3);
- VirtualDevice rhs = VirtualDevice(kDLCUDA, 4);
- ffi::Optional<VirtualDevice> actual = VirtualDevice::Join(lhs, rhs);
- EXPECT_FALSE(actual);
- }
- {
- VirtualDevice lhs = VirtualDevice(kDLCUDA, 3, Target("cuda"));
- VirtualDevice rhs = VirtualDevice(kDLCUDA, 3, Target("cuda"));
- ffi::Optional<VirtualDevice> actual = VirtualDevice::Join(lhs, rhs);
- EXPECT_FALSE(actual);
- }
- {
- VirtualDevice lhs = VirtualDevice(kDLCUDA, 3, Target("cuda"), "local");
- VirtualDevice rhs = VirtualDevice(kDLCUDA, 3, Target("cuda"), "global");
- ffi::Optional<VirtualDevice> actual = VirtualDevice::Join(lhs, rhs);
- EXPECT_FALSE(actual);
- }
-}
-
-TEST(VirtualDevice, Default) {
- Target target_a = Target("cuda");
- VirtualDevice lhs = VirtualDevice(kDLCUDA, -1, Target(), "global");
- VirtualDevice rhs = VirtualDevice(kDLCUDA, 3, target_a, "local");
- VirtualDevice actual = VirtualDevice::Default(lhs, rhs);
- VirtualDevice expected = VirtualDevice(kDLCUDA, 3, target_a, "global");
- EXPECT_TRUE(tvm::ffi::StructuralEqual()(actual, expected));
-}
-
-TEST(VirtualDevice, Constructor_Invalid) {
- EXPECT_ANY_THROW(VirtualDevice(kDLCPU, -1, Target("cuda")));
-}
-
-TEST(VirtualDeviceCache, Memoized) {
- VirtualDeviceCache cache;
- Target target_a = Target("cuda");
- Target target_b = Target("llvm");
- Target target_c = Target("cuda");
- VirtualDevice virtual_device_a = cache.Make(kDLCUDA, 3, target_a, "local");
- VirtualDevice virtual_device_b = cache.Make(kDLCPU, 1, target_b, "global");
-
- EXPECT_EQ(cache.Make(kDLCUDA, 3, target_a, "local"), virtual_device_a);
- EXPECT_EQ(cache.Make(kDLCPU, 1, target_b, "global"), virtual_device_b);
- EXPECT_NE(cache.Make(kDLCUDA, 2, target_a, "local"), virtual_device_a);
- EXPECT_NE(cache.Make(kDLCPU, 3, target_b, "local"), virtual_device_a);
- EXPECT_NE(cache.Make(kDLCUDA, 3, target_a, "global"), virtual_device_a);
- EXPECT_EQ(cache.Make(kDLCUDA, 3, Target("cuda"), "local"), virtual_device_a);
- EXPECT_NE(
- cache.Make(kDLCUDA, 3,
- Target(ffi::Map<ffi::String, ffi::Any>{{"kind",
ffi::String("cuda")},
-
{"max_threads_per_block", int64_t(4096)}}),
- "local"),
- virtual_device_a);
-}
-
-} // namespace
-} // namespace tvm
diff --git a/tests/python/relax/test_analysis_type_analysis.py
b/tests/python/relax/test_analysis_type_analysis.py
index a367de212f..b0b6a54aa0 100644
--- a/tests/python/relax/test_analysis_type_analysis.py
+++ b/tests/python/relax/test_analysis_type_analysis.py
@@ -217,11 +217,11 @@ def test_base_check():
shape3 = rx.ShapeType([1, 2, 3])
shape4 = rx.ShapeType([1, n, 3])
- vdevice0 = ir.VDevice()
- vdevice1 = ir.VDevice("llvm")
- vdevice2 = ir.VDevice("cuda", 0)
- vdevice3 = ir.VDevice("cuda", 2)
- vdevice4 = ir.VDevice("cuda", 0, "")
+ vdevice0 = rx.VDevice()
+ vdevice1 = rx.VDevice("llvm")
+ vdevice2 = rx.VDevice("cuda", 0)
+ vdevice3 = rx.VDevice("cuda", 2)
+ vdevice4 = rx.VDevice("cuda", 0, "")
tensor0 = rx.TensorType(ndim=-1, dtype="int32")
tensor1 = rx.TensorType(ndim=-1, dtype="float32")
@@ -414,7 +414,7 @@ def test_derive_call_ret_type():
_check_derive(bb, func0(2), [obj0], obj0)
# Tensor with vdevice
- vdev = ir.VDevice("llvm")
+ vdev = rx.VDevice("llvm")
def func1(c):
n, m = tirx.Var("n", "int64"), tirx.Var("m", "int64")
@@ -520,8 +520,8 @@ def test_type_lca():
prim0 = tvm.ir.PrimType("int32")
prim1 = tvm.ir.PrimType("float32")
- vdevice0 = ir.VDevice("llvm")
- vdevice1 = ir.VDevice("cuda", 0)
+ vdevice0 = rx.VDevice("llvm")
+ vdevice1 = rx.VDevice("cuda", 0)
shape0 = rx.ShapeType(ndim=-1)
shape1 = rx.ShapeType(ndim=2)
diff --git a/tests/python/relax/test_backend_dispatch_sort_scan.py
b/tests/python/relax/test_backend_dispatch_sort_scan.py
index c2dc7ef66a..795a56537e 100644
--- a/tests/python/relax/test_backend_dispatch_sort_scan.py
+++ b/tests/python/relax/test_backend_dispatch_sort_scan.py
@@ -567,7 +567,7 @@ def test_dispatch_cumsum_webgpu_axes_and_dtypes(
):
"""WebGPU dispatch collapses arbitrary-rank scans to the appropriate
kernel."""
- vdevice = tvm.ir.VDevice("webgpu", 0)
+ vdevice = tvm.relax.VDevice("webgpu", 0)
x = relax.Var("x", relax.TensorType(shape, in_dtype, vdevice=vdevice))
bb = relax.BlockBuilder()
with bb.function("main", (x,)):
diff --git a/tests/python/relax/test_op_binary.py
b/tests/python/relax/test_op_binary.py
index e3befa2a97..1c9da199f6 100644
--- a/tests/python/relax/test_op_binary.py
+++ b/tests/python/relax/test_op_binary.py
@@ -22,7 +22,8 @@ import pytest
import tvm
import tvm.testing
from tvm import relax, tirx
-from tvm.ir import Op, VDevice
+from tvm.ir import Op
+from tvm.relax import VDevice
from tvm.script import relax as R
diff --git a/tests/python/relax/test_op_create.py
b/tests/python/relax/test_op_create.py
index 0e8d403bae..7d4e3899ec 100644
--- a/tests/python/relax/test_op_create.py
+++ b/tests/python/relax/test_op_create.py
@@ -19,7 +19,8 @@ import pytest
import tvm
import tvm.testing
from tvm import relax, tirx
-from tvm.ir import Op, VDevice
+from tvm.ir import Op
+from tvm.relax import VDevice
from tvm.script import relax as R
from tvm.script import tirx as T
diff --git a/tests/python/relax/test_op_image.py
b/tests/python/relax/test_op_image.py
index 1a24b90859..2644e0c3de 100644
--- a/tests/python/relax/test_op_image.py
+++ b/tests/python/relax/test_op_image.py
@@ -24,7 +24,8 @@ pytest.importorskip("scipy") # tvm.topi.testing imports scipy
import tvm.topi.testing
from tvm import relax, tirx
-from tvm.ir import Op, VDevice
+from tvm.ir import Op
+from tvm.relax import VDevice
from tvm.script import relax as R
diff --git a/tests/python/relax/test_op_index.py
b/tests/python/relax/test_op_index.py
index 2894eb98e7..d586471595 100644
--- a/tests/python/relax/test_op_index.py
+++ b/tests/python/relax/test_op_index.py
@@ -21,7 +21,8 @@ import pytest
import tvm
import tvm.testing
from tvm import relax, tirx
-from tvm.ir import Op, VDevice
+from tvm.ir import Op
+from tvm.relax import VDevice
from tvm.script import ir as I
from tvm.script import relax as R
from tvm.script import tirx as T
diff --git a/tests/python/relax/test_op_linear_algebra.py
b/tests/python/relax/test_op_linear_algebra.py
index 4b7b50854d..56c75d77f0 100644
--- a/tests/python/relax/test_op_linear_algebra.py
+++ b/tests/python/relax/test_op_linear_algebra.py
@@ -20,7 +20,8 @@ import pytest
import tvm
import tvm.testing
from tvm import relax, tirx
-from tvm.ir import Op, VDevice
+from tvm.ir import Op
+from tvm.relax import VDevice
from tvm.script import relax as R
diff --git a/tests/python/relax/test_op_manipulate.py
b/tests/python/relax/test_op_manipulate.py
index 694ec82baf..42ddece4c9 100644
--- a/tests/python/relax/test_op_manipulate.py
+++ b/tests/python/relax/test_op_manipulate.py
@@ -20,7 +20,8 @@ import pytest
import tvm
import tvm.testing
from tvm import relax, tirx
-from tvm.ir import Op, VDevice
+from tvm.ir import Op
+from tvm.relax import VDevice
from tvm.script import relax as R
from tvm.script import tirx as T
diff --git a/tests/python/relax/test_op_nn.py b/tests/python/relax/test_op_nn.py
index af4d2c25fc..f605ebf0c6 100644
--- a/tests/python/relax/test_op_nn.py
+++ b/tests/python/relax/test_op_nn.py
@@ -20,7 +20,8 @@ import pytest
import tvm
import tvm.testing
from tvm import relax, tirx
-from tvm.ir import Op, VDevice
+from tvm.ir import Op
+from tvm.relax import VDevice
from tvm.script import relax as R
diff --git a/tests/python/relax/test_op_nn_convolution.py
b/tests/python/relax/test_op_nn_convolution.py
index 31f5a15059..766e32cf16 100644
--- a/tests/python/relax/test_op_nn_convolution.py
+++ b/tests/python/relax/test_op_nn_convolution.py
@@ -19,7 +19,8 @@ import pytest
import tvm
import tvm.testing
from tvm import relax, tirx
-from tvm.ir import Op, VDevice
+from tvm.ir import Op
+from tvm.relax import VDevice
from tvm.script import relax as R
diff --git a/tests/python/relax/test_op_nn_pooling.py
b/tests/python/relax/test_op_nn_pooling.py
index 159a280283..c933f51b4e 100644
--- a/tests/python/relax/test_op_nn_pooling.py
+++ b/tests/python/relax/test_op_nn_pooling.py
@@ -20,7 +20,8 @@ import pytest
import tvm
import tvm.testing
from tvm import relax, tirx
-from tvm.ir import Op, VDevice
+from tvm.ir import Op
+from tvm.relax import VDevice
from tvm.script import relax as R
diff --git a/tests/python/relax/test_op_search.py
b/tests/python/relax/test_op_search.py
index fb3aa91e7e..ebf7eca26e 100644
--- a/tests/python/relax/test_op_search.py
+++ b/tests/python/relax/test_op_search.py
@@ -21,7 +21,8 @@ import pytest
import tvm
import tvm.testing
from tvm import relax, tirx
-from tvm.ir import Op, VDevice
+from tvm.ir import Op
+from tvm.relax import VDevice
from tvm.script import relax as R
diff --git a/tests/python/relax/test_op_set.py
b/tests/python/relax/test_op_set.py
index 8239669103..161dd73b9a 100644
--- a/tests/python/relax/test_op_set.py
+++ b/tests/python/relax/test_op_set.py
@@ -19,7 +19,8 @@ import pytest
import tvm
import tvm.testing
from tvm import relax, tirx
-from tvm.ir import Op, VDevice
+from tvm.ir import Op
+from tvm.relax import VDevice
from tvm.script import relax as R
diff --git a/tests/python/relax/test_op_sort.py
b/tests/python/relax/test_op_sort.py
index c40b58b899..2176bebc80 100644
--- a/tests/python/relax/test_op_sort.py
+++ b/tests/python/relax/test_op_sort.py
@@ -19,7 +19,8 @@ import pytest
import tvm
import tvm.testing
from tvm import relax, tirx
-from tvm.ir import Op, VDevice
+from tvm.ir import Op
+from tvm.relax import VDevice
from tvm.script import relax as R
diff --git a/tests/python/relax/test_op_statistical.py
b/tests/python/relax/test_op_statistical.py
index 06c2dbe055..fb969c3a4e 100644
--- a/tests/python/relax/test_op_statistical.py
+++ b/tests/python/relax/test_op_statistical.py
@@ -21,7 +21,8 @@ import pytest
import tvm
import tvm.testing
from tvm import relax, tirx
-from tvm.ir import Op, VDevice
+from tvm.ir import Op
+from tvm.relax import VDevice
from tvm.script import relax as R
diff --git a/tests/python/relax/test_op_ternary.py
b/tests/python/relax/test_op_ternary.py
index 82b6df52f1..21f2eceffc 100644
--- a/tests/python/relax/test_op_ternary.py
+++ b/tests/python/relax/test_op_ternary.py
@@ -19,7 +19,8 @@ import pytest
import tvm
import tvm.testing
from tvm import relax, tirx
-from tvm.ir import Op, VDevice
+from tvm.ir import Op
+from tvm.relax import VDevice
from tvm.script import relax as R
diff --git a/tests/python/relax/test_op_unary.py
b/tests/python/relax/test_op_unary.py
index 9d2658b32e..8d7e332e91 100644
--- a/tests/python/relax/test_op_unary.py
+++ b/tests/python/relax/test_op_unary.py
@@ -21,7 +21,8 @@ import pytest
import tvm
import tvm.testing
from tvm import relax, tirx
-from tvm.ir import Op, VDevice
+from tvm.ir import Op
+from tvm.relax import VDevice
from tvm.script import relax as R
diff --git a/tests/python/relax/test_relax_operators.py
b/tests/python/relax/test_relax_operators.py
index 437cb33bd3..18f1546ca4 100644
--- a/tests/python/relax/test_relax_operators.py
+++ b/tests/python/relax/test_relax_operators.py
@@ -509,7 +509,7 @@ def test_op_to_vdevice(exec_mode):
@R.function
def to_vdev(x: R.Tensor((3, 4), "float32")):
- dst_vdev = tvm.ir.VDevice("llvm", 0, "global")
+ dst_vdev = tvm.relax.VDevice("llvm", 0, "global")
ret = R.to_vdevice(x, "llvm")
return ret
diff --git a/tests/python/relax/test_transform_realize_vdevice.py
b/tests/python/relax/test_transform_realize_vdevice.py
index fbf4a8a26a..c9d5e2f867 100644
--- a/tests/python/relax/test_transform_realize_vdevice.py
+++ b/tests/python/relax/test_transform_realize_vdevice.py
@@ -19,7 +19,7 @@
import tvm
import tvm.testing
-from tvm.ir import VDevice
+from tvm.relax import VDevice
from tvm.relax.transform import RealizeVDevice
from tvm.script.parser import ir as I
from tvm.script.parser import relax as R
diff --git a/tests/python/relax/test_transform_update_vdevice.py
b/tests/python/relax/test_transform_update_vdevice.py
index 618a286115..48dfa2df7b 100644
--- a/tests/python/relax/test_transform_update_vdevice.py
+++ b/tests/python/relax/test_transform_update_vdevice.py
@@ -18,7 +18,7 @@
import tvm
import tvm.testing
-from tvm.ir import VDevice
+from tvm.relax import VDevice
from tvm.relax.transform import UpdateVDevice
from tvm.script.parser import ir as I
from tvm.script.parser import relax as R
diff --git a/tests/python/relax/test_tvmscript_parser.py
b/tests/python/relax/test_tvmscript_parser.py
index 1e7cddb116..1bff4e99f3 100644
--- a/tests/python/relax/test_tvmscript_parser.py
+++ b/tests/python/relax/test_tvmscript_parser.py
@@ -24,7 +24,7 @@ import tvm
import tvm.script
import tvm.testing
from tvm import IRModule, relax, tirx, topi
-from tvm.ir import DummyGlobalInfo, VDevice
+from tvm.relax import DummyGlobalInfo, VDevice
from tvm.script.parser import ir as I
from tvm.script.parser import relax as R
from tvm.script.parser import tirx as T
diff --git a/tests/python/relax/test_vm_build.py
b/tests/python/relax/test_vm_build.py
index d3537615ff..7a10f81538 100644
--- a/tests/python/relax/test_vm_build.py
+++ b/tests/python/relax/test_vm_build.py
@@ -796,14 +796,14 @@ def test_vm_to_device(exec_mode):
def foo1(
x: R.Tensor((2, 3), "float32"),
) -> R.Tensor((2, 3), "float32"):
- copied = R.to_vdevice(x, tvm.ir.VDevice("cuda", 0, "global"))
+ copied = R.to_vdevice(x, tvm.relax.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"))
+ copied = R.to_vdevice(x, tvm.relax.VDevice("llvm", 0, "global"))
return copied
mod = TestToVDevice
diff --git a/tests/python/target/test_virtual_device.py
b/tests/python/target/test_virtual_device.py
deleted file mode 100644
index b72023b811..0000000000
--- a/tests/python/target/test_virtual_device.py
+++ /dev/null
@@ -1,51 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-# ruff: noqa: F401
-import pytest
-
-import tvm
-import tvm.testing
-
-
-def test_make_virtual_device_for_device():
- virtual_device = tvm.target.VirtualDevice(tvm.cuda())
- assert virtual_device.dlpack_device_type() == 2
- # ie kDLCUDA
- assert virtual_device.virtual_device_id == 0
- assert virtual_device.target is None
- assert virtual_device.memory_scope == ""
-
-
-def test_make_virtual_device_for_device_and_target():
- target = tvm.target.Target("cuda")
- 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 == ""
-
-
-def test_make_virtual_device_for_device_target_and_memory_scope():
- target = tvm.target.Target("cuda")
- scope = "local"
- 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
-
-
-if __name__ == "__main__":
- tvm.testing.main()
