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 8697437586 [REFACTOR][Relax] Remove compiled VM TIR path and AnyList
intrinsics (#20464)
8697437586 is described below
commit 86974375869e476a6130d814ce912909f89b66ce
Author: Tianqi Chen <[email protected]>
AuthorDate: Sat Sep 26 11:47:57 2026 -0400
[REFACTOR][Relax] Remove compiled VM TIR path and AnyList intrinsics
(#20464)
Relax functions now execute through VM bytecode only. This removes the
unused TIR codegen and its runtime dispatch, along with the AnyList
intrinsics and lowering hooks that existed solely for that path. The
Relax build API and VM documentation describe the single mode, and
former dual-mode tests retain their bytecode behavior coverage.
---
docs/arch/relax_vm.rst | 20 +-
docs/tirx/api/script.rst | 2 +-
include/tvm/runtime/vm/executable.h | 2 -
include/tvm/tirx/builtin.h | 44 --
python/tvm/relax/vm_build.py | 16 +-
python/tvm/tirx/op.py | 84 ---
python/tvm/tirx/script/ir_builder/op.py | 16 -
src/relax/backend/vm/codegen_vm_tir.cc | 608 ---------------------
src/relax/backend/vm/exec_builder.cc | 6 -
src/relax/transform/kill_after_last_use.cc | 11 +-
src/runtime/vm/builtin.cc | 62 ---
src/runtime/vm/executable.cc | 7 -
src/runtime/vm/vm.cc | 57 +-
src/tirx/op/builtin.cc | 13 -
src/tirx/op/runtime.cc | 43 --
src/tirx/transform/lower_tvm_builtin.cc | 107 ++--
tests/python/relax/test_kill_after_last_use.py | 10 +-
tests/python/relax/test_relax_operators.py | 119 ++--
.../relax/test_vm_alloc_storage_with_scope.py | 2 +-
tests/python/relax/test_vm_build.py | 178 +++---
tests/python/relax/test_vm_callback_function.py | 11 +-
tests/python/relax/test_vm_codegen_only.py | 74 +--
tests/python/relax/test_vm_codegen_tir.py | 225 --------
tests/python/relax/test_vm_cuda_graph.py | 4 +-
24 files changed, 243 insertions(+), 1478 deletions(-)
diff --git a/docs/arch/relax_vm.rst b/docs/arch/relax_vm.rst
index 222329ce0b..ec3cbfcab5 100644
--- a/docs/arch/relax_vm.rst
+++ b/docs/arch/relax_vm.rst
@@ -91,13 +91,7 @@ Internally, ``relax.build()`` performs these steps:
4. **Link** the bytecode executable with the compiled native module using
``VMLink``, producing
a ``VMExecutable``.
-Two execution modes are supported:
-
-- ``exec_mode="bytecode"`` (default): Relax functions are interpreted by the
VM's bytecode
- dispatch loop.
-- ``exec_mode="compiled"``: Relax functions are compiled into TIR functions
(``VMTIRCodeGen``)
- that directly manipulate the register file, bypassing the interpreter loop.
This avoids
- dispatch overhead but produces more code.
+Relax functions execute through the VM bytecode dispatch loop.
Bytecode generation
~~~~~~~~~~~~~~~~~~~
@@ -180,7 +174,7 @@ execution:
Function kinds
~~~~~~~~~~~~~~
-The VM recognizes three function kinds (``VMFuncInfo::FuncKind``):
+The VM recognizes two function kinds (``VMFuncInfo::FuncKind``):
.. list-table::
:header-rows: 1
@@ -193,10 +187,6 @@ The VM recognizes three function kinds
(``VMFuncInfo::FuncKind``):
registry. Examples: ``vm.builtin.alloc_shape_heap``,
``vm.builtin.match_shape``.
* - ``kVMFunc``
- A bytecode-interpreted Relax function. The VM interprets its
instructions in ``RunLoop()``.
- * - ``kVMTIRFunc``
- - A Relax function compiled to a TIR function (``exec_mode="compiled"``).
Found in
- imports under the name ``__vmtir__<func_name>``. Called directly with
register file
- pointers, bypassing the interpreter loop.
Serialization
~~~~~~~~~~~~~
@@ -239,7 +229,7 @@ Under the hood:
(either ``NAIVE_ALLOCATOR`` or ``POOLED_ALLOCATOR``, defaulting to pooled).
A CPU device
is always added for shape computations.
3. **InitFuncPool**: the function pool is populated — ``kPackedFunc`` entries
are resolved from
- imports or the global registry; ``kVMFunc`` and ``kVMTIRFunc`` entries are
wrapped in
+ imports or the global registry; ``kVMFunc`` entries are wrapped in
``VMClosure`` objects.
4. **Constant pool**: model constants are loaded and optionally transferred to
the target device.
@@ -286,7 +276,7 @@ VMClosure and function dispatch
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Functions in the VM are stored in a ``func_pool_`` indexed by function table
position.
-``kVMFunc`` and ``kVMTIRFunc`` entries are wrapped as ``VMClosure`` objects,
while ``kPackedFunc``
+``kVMFunc`` entries are wrapped as ``VMClosure`` objects, while ``kPackedFunc``
entries are stored as plain ``ffi::Function``. A ``VMClosure`` stores:
- ``func_name``: the function's string name.
@@ -419,8 +409,6 @@ Source Code Map
- Built-in operations (shape matching, allocation)
* - ``src/relax/backend/vm/codegen_vm.cc``
- CodeGenVM: Relax IR → bytecode
- * - ``src/relax/backend/vm/codegen_vm_tir.cc``
- - VMTIRCodeGen: Relax IR → compiled TIR
* - ``python/tvm/runtime/vm.py``
- Python VirtualMachine wrapper
* - ``python/tvm/relax/vm_build.py``
diff --git a/docs/tirx/api/script.rst b/docs/tirx/api/script.rst
index b3d50539b4..5f0d50c76f 100644
--- a/docs/tirx/api/script.rst
+++ b/docs/tirx/api/script.rst
@@ -40,7 +40,7 @@ Core IR builder
.. automodule:: tvm.tirx.script.ir_builder.ir
:members:
:no-index:
- :exclude-members: LetAnnotation, alloc_tcgen05_ldst_frag, anylist_getitem,
anylist_resetitem, anylist_setitem_call_packed, anylist_setitem_call_cpacked
+ :exclude-members: LetAnnotation, alloc_tcgen05_ldst_frag
.. currentmodule:: tvm.tirx.script.ir_builder.ir
diff --git a/include/tvm/runtime/vm/executable.h
b/include/tvm/runtime/vm/executable.h
index 29ee7e420c..eff16e8ff8 100644
--- a/include/tvm/runtime/vm/executable.h
+++ b/include/tvm/runtime/vm/executable.h
@@ -57,8 +57,6 @@ struct VMFuncInfo {
kPackedFunc = 0,
/*! \brief VM function. */
kVMFunc = 1,
- /*! \brief VMTIR function. */
- kVMTIRFunc = 2,
};
/*! \brief The kind of function. */
FuncKind kind;
diff --git a/include/tvm/tirx/builtin.h b/include/tvm/tirx/builtin.h
index 146d5af881..eba9f21fed 100644
--- a/include/tvm/tirx/builtin.h
+++ b/include/tvm/tirx/builtin.h
@@ -552,50 +552,6 @@ TVM_DLL const Op& assume();
*/
TVM_DLL const Op& undef();
-/*!
- * \brief Get a item from any list and return it.
- *
- * Any anylist_getitem(Handle anylist,
- * int index)
- * return anylist[index];
- * }
- *
- * \note This intrinsic is only applicable when appearing
- * in call_packed and anylist_setitem_call_packed.
- */
-TVM_DLL const Op& anylist_getitem();
-
-/*!
- * \brief Reset and clear a item in any list.
- *
- * void anylist_resetitem(Handle anylist,
- * int index)
- * anylist[index] = nullptr;
- * }
- *
- * \note This intrinsic is only applicable when appearing
- * in call_packed and anylist_setitem_call_packed.
- */
-TVM_DLL const Op& anylist_resetitem();
-
-/*!
- * \brief Set an item into any list by running packed function call.
- *
- * void anylist_setitem_call_packed(Handle anylist,
- * int index,
- * name, *args)
- *
- * anylist[index] = call_packed(name, *args)
- * }
- * \note This intrinsic can be used in combination with anylist_getitem.
- */
-TVM_DLL const Op& anylist_setitem_call_packed();
-
-/*!
- * \brief Same as anylist_setitem_call_packed but use C calling convention.
- */
-TVM_DLL const Op& anylist_setitem_call_cpacked();
-
/*!
* \brief Calculate a predicate mask given an upper bound (limit) and a
current value (base).
*
diff --git a/python/tvm/relax/vm_build.py b/python/tvm/relax/vm_build.py
index adc0f7ad83..f146b51ab1 100644
--- a/python/tvm/relax/vm_build.py
+++ b/python/tvm/relax/vm_build.py
@@ -51,7 +51,6 @@ class VMExecutable(Executable):
def _vmcodegen(
builder: "relax.ExecBuilder",
mod: tvm.IRModule,
- exec_mode: str = "bytecode",
) -> tvm.IRModule:
"""Running VM codegen.
@@ -63,20 +62,13 @@ def _vmcodegen(
mod: IRModule
The input IRModule to be built.
- exec_mode: {"bytecode", "compiled"}
- The execution mode.
-
Return
------
leftover: IRModule
Left over IRModule that may contain extra functions.
"""
- if exec_mode == "bytecode":
- return _ffi_api.VMCodeGen(builder, mod) # type:ignore
- if exec_mode == "compiled":
- return _ffi_api.VMTIRCodeGen(builder, mod) # type: ignore
- raise ValueError(f"Unknown exec_mode {exec_mode}")
+ return _ffi_api.VMCodeGen(builder, mod) # type: ignore
def _auto_attach_system_lib_prefix(
@@ -175,7 +167,6 @@ def build(
params: dict[str, list] | None = None,
relax_pipeline: None | str | tvm.transform.Pass = "default",
tir_pipeline: None | str | tvm.transform.Pass = "default",
- exec_mode: str = "bytecode",
*,
system_lib: bool | None = None,
) -> Executable:
@@ -206,9 +197,6 @@ def build(
tir_pipelinie : str = "default"
The TIR compilation pipeline to use.
- exec_mode: {"bytecode", "compiled"}
- The execution mode.
-
system_lib: Optional[bool]
Whether to build system lib that is being packed statically and
auto registers generated functions to the system.
@@ -272,7 +260,7 @@ def build(
ext_libs, constants = _extract_attrs(mod)
params.update(dict(constants))
builder = relax.ExecBuilder()
- mod = _vmcodegen(builder, mod, exec_mode)
+ mod = _vmcodegen(builder, mod)
return _vmlink(
builder=builder,
target=target,
diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op.py
index c93e0da604..49fc608891 100644
--- a/python/tvm/tirx/op.py
+++ b/python/tvm/tirx/op.py
@@ -2863,90 +2863,6 @@ def TVMBackendFreeWorkspace(device_type, device_id, ptr):
return call_intrin("int32", "tirx.TVMBackendFreeWorkspace", device_type,
device_id, ptr)
-def anylist_getitem(list_handle, index):
- """Returns an item from any list.
-
- Parameters
- ----------
- list_handle: Var
- The handle to anylist
- index : int
- The index
-
- Returns
- -------
- call : Expr
- The call expression.
- """
- return call_intrin("handle", "tirx.anylist_getitem", list_handle, index)
-
-
-def anylist_resetitem(list_handle, index):
- """Reset an item from any list.
-
- Parameters
- ----------
- list_handle: Var
- The handle to anylist
- index : int
- The index
-
- Returns
- -------
- call : Expr
- The call expression.
- """
- return call_intrin("int", "tirx.anylist_resetitem", list_handle, index)
-
-
-def anylist_setitem_call_packed(list_handle, index, func_name, *args):
- """Set anylist item by result of packed call.
-
- Parameters
- ----------
- list_handle: Var
- The handle to anylist
- index : int
- The index
- func_name: str
- The name of the function to be called.
- args:
- Extra arguments
-
- Returns
- -------
- call : Expr
- The call expression.
- """
- return call_intrin(
- "int", "tirx.anylist_setitem_call_packed", list_handle, index,
func_name, *args
- )
-
-
-def anylist_setitem_call_cpacked(list_handle, index, func_name, *args):
- """Set anylist item by result of packed call.
-
- Parameters
- ----------
- list_handle: Var
- The handle to anylist
- index : int
- The index
- func_name: str
- The name of the function to be called.
- args:
- Extra arguments
-
- Returns
- -------
- call : Expr
- The call expression.
- """
- return call_intrin(
- "int", "tirx.anylist_setitem_call_cpacked", list_handle, index,
func_name, *args
- )
-
-
def vscale():
"""Get the target's vscale value. It will be lowered to llvm.vscale
intrinsic
(https://llvm.org/docs/LangRef.html#llvm-vscale-intrinsic)
diff --git a/python/tvm/tirx/script/ir_builder/op.py
b/python/tvm/tirx/script/ir_builder/op.py
index 97ab45c37b..946a8fcd3e 100644
--- a/python/tvm/tirx/script/ir_builder/op.py
+++ b/python/tvm/tirx/script/ir_builder/op.py
@@ -644,18 +644,6 @@ TVMBackendAllocWorkspace =
_op_wrapper(_tir_op.TVMBackendAllocWorkspace)
TVMBackendFreeWorkspace = _op_wrapper(_tir_op.TVMBackendFreeWorkspace)
-anylist_getitem = _op_wrapper(_tir_op.anylist_getitem)
-
-
-anylist_resetitem = _op_wrapper(_tir_op.anylist_resetitem)
-
-
-anylist_setitem_call_packed = _op_wrapper(_tir_op.anylist_setitem_call_packed)
-
-
-anylist_setitem_call_cpacked =
_op_wrapper(_tir_op.anylist_setitem_call_cpacked)
-
-
vscale = _op_wrapper(_tir_op.vscale)
@@ -894,10 +882,6 @@ __all__ = [
"acos",
"acosh",
"address_of",
- "anylist_getitem",
- "anylist_resetitem",
- "anylist_setitem_call_cpacked",
- "anylist_setitem_call_packed",
"asin",
"asinh",
"assume",
diff --git a/src/relax/backend/vm/codegen_vm_tir.cc
b/src/relax/backend/vm/codegen_vm_tir.cc
deleted file mode 100644
index 2537f9b440..0000000000
--- a/src/relax/backend/vm/codegen_vm_tir.cc
+++ /dev/null
@@ -1,608 +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 src/relax/backend/vm/codegen_tir.cc
- * \brief A codegen to generate VMTIR function(that can be compiled) from
executable.
- */
-#include <tvm/ffi/cast.h>
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/ir/module.h>
-#include <tvm/ir/prim/builtin.h>
-#include <tvm/ir/prim/expr.h>
-#include <tvm/relax/exec_builder.h>
-#include <tvm/relax/expr_functor.h>
-#include <tvm/relax/op_attr_types.h>
-#include <tvm/runtime/logging.h>
-#include <tvm/runtime/vm/executable.h>
-#include <tvm/target/target.h>
-#include <tvm/tirx/builtin.h>
-#include <tvm/tirx/function.h>
-#include <tvm/tirx/stmt.h>
-
-#include <cctype>
-#include <string>
-#include <unordered_map>
-#include <vector>
-
-#include "../../transform/utils.h"
-
-namespace tvm {
-namespace relax {
-namespace codegen_vm {
-
-using vm::VMFuncInfo;
-
-/*!
- * \brief A class to generate VMTIR for Relax functions.
- *
- * \note Skip CallPacked with special attrs for now, as they can be
- * further simplified with PrimExpr.
- */
-class CodeGenVMTIR : public ExprFunctor<ffi::Optional<Expr>(const Expr&)> {
- public:
- explicit CodeGenVMTIR(relax::ExecBuilder builder, IRModule ctx_mod)
- : builder_(builder), ctx_mod_(ctx_mod) {
- system_lib_prefix_ =
ctx_mod_->GetAttr<ffi::String>(tvm::attr::kSystemLibPrefix);
- }
-
- static IRModule Run(relax::ExecBuilder builder, IRModule mod) {
- // create a new copy
- IRModule res_mod = mod;
- res_mod.CopyOnWrite();
-
- CodeGenVMTIR codegen(builder, mod);
- // Remove relax function and turn into TIR func.
- for (auto& p : mod->functions) {
- if (auto* func = p.second.as<FunctionNode>()) {
- auto tir_func = codegen.Codegen(ffi::GetRef<Function>(func));
- auto gsymbol =
tir_func->GetAttr<ffi::String>(tvm::attr::kGlobalSymbol);
- res_mod->Add(GlobalVar(gsymbol.value()), tir_func);
- res_mod->Remove(p.first);
- }
- }
- return res_mod;
- }
-
- private:
- int64_t NewRegister() { return registers_num_++; }
-
- static IntImm ConstInt64(int64_t value) { return IntImm::Int64(value); }
-
- static IntImm ConstInt32(int64_t value) { return IntImm::Int32(value); }
-
- Expr RegListGet(int64_t slot) const {
- // use 128 bits to represent any
- return tvm::Call(tvm::PointerType::VoidPointerTy(),
tirx::builtin::anylist_getitem(),
- {reg_anylist_handle_, ConstInt32(slot)});
- }
-
- Expr ConstListGet(int64_t slot) const {
- // use 128 bits to represent any
- return tvm::Call(tvm::PointerType::VoidPointerTy(),
tirx::builtin::anylist_getitem(),
- {const_anylist_handle_, ConstInt32(slot)});
- }
-
- Expr FuncListGet(int64_t slot) const {
- // use 128 bits to represent any
- return tvm::Call(tvm::PointerType::VoidPointerTy(),
tirx::builtin::anylist_getitem(),
- {func_anylist_handle_, ConstInt32(slot)});
- }
-
- void EmitStmt(tirx::Stmt stmt) {
- TVM_FFI_ICHECK(!stmt_stack_.empty());
- stmt_stack_.back().emplace_back(stmt);
- }
-
- void EmitCallPacked(ffi::String name, const ffi::Array<Expr>& args,
- int64_t dst_anylist_slot = -1) {
- ffi::Array<Expr> all_args;
- // negative index indicate return value can be discarded, emit call_packed
- if (dst_anylist_slot >= 0) {
- all_args = {reg_anylist_handle_, ConstInt32(dst_anylist_slot)};
- }
- all_args.push_back(StringImm(name));
- for (Expr arg : args) {
- all_args.push_back(arg);
- }
- if (dst_anylist_slot >= 0) {
- this->EmitStmt(tirx::Evaluate(
- tvm::Call(tvm::PrimType::Int(32),
tirx::builtin::anylist_setitem_call_packed(), all_args)
- .as_or_throw<PrimExpr>()));
- } else {
- this->EmitStmt(tirx::Evaluate(
- tvm::Call(tvm::PrimType::Int(32), tirx::builtin::tvm_call_packed(),
all_args)
- .as_or_throw<PrimExpr>()));
- }
- }
-
- void EmitCallCPacked(const tirx::PrimFunc& prim_func, const
ffi::Array<Expr>& args,
- int64_t dst_anylist_slot = -1) {
- ffi::Optional<ffi::String> gsymbol =
prim_func->GetAttr<ffi::String>(tvm::attr::kGlobalSymbol);
- TVM_FFI_ICHECK(gsymbol.has_value()) << "All functions must have global
symbol at this phase";
- ffi::Array<Expr> all_args;
- // negative index indicate return value can be discarded, emit call_packed
- if (dst_anylist_slot >= 0) {
- all_args = {reg_anylist_handle_, ConstInt32(dst_anylist_slot)};
- }
- all_args.push_back(StringImm(gsymbol.value()));
- for (Expr arg : args) {
- all_args.push_back(arg);
- }
- if (dst_anylist_slot >= 0) {
- this->EmitStmt(tirx::Evaluate(
- tvm::Call(tvm::PrimType::Int(32),
tirx::builtin::anylist_setitem_call_cpacked(), all_args)
- .as_or_throw<PrimExpr>()));
- } else {
- this->EmitStmt(tirx::Evaluate(
- tvm::Call(tvm::PrimType::Int(32), tirx::builtin::tvm_call_cpacked(),
all_args)
- .as_or_throw<PrimExpr>()));
- }
- }
-
- tirx::PrimFunc Codegen(const Function& func) {
- ffi::Optional<ffi::String> gsymbol =
func->GetAttr<ffi::String>(tvm::attr::kGlobalSymbol);
- TVM_FFI_ICHECK(gsymbol.has_value())
- << "there should be no local functions in Relax VM codegen phase. "
- "Did you forget to apply LambdaLift or AttachGlobalSymbol Pass?";
- // initialize the state
- stmt_stack_ = {};
- registers_num_ = 0;
- var_map_.clear();
- ctx_ptr_ = tirx::Var("ctx_ptr", PointerType::VoidPointerTy());
- reg_anylist_handle_ = tirx::Var("r", PointerType::VoidPointerTy());
- func_anylist_handle_ = tirx::Var("f", PointerType::VoidPointerTy());
- const_anylist_handle_ = tirx::Var("c", PointerType::VoidPointerTy());
-
- ffi::Array<ffi::String> param_names;
- for (Var param : func->params) {
- param_names.push_back(param->name);
- }
- // declare this function.
- builder_->DeclareFunction(gsymbol.value(),
vm::VMFuncInfo::FuncKind::kVMTIRFunc);
-
- for (size_t i = 0; i < func->params.size(); ++i) {
- int64_t r = NewRegister();
- TVM_FFI_ICHECK_EQ(static_cast<size_t>(r), i);
- this->var_map_.insert({func->params[i], RegListGet(r)});
- }
- size_t ret_reg = NewRegister();
-
- tirx::Stmt body = WithNewScope([&]() {
- ffi::Optional<Expr> ret = ExprFunctor::VisitExpr(func->body);
- if (ret.has_value()) {
- this->EmitCallPacked("vm.builtin.copy", {ret.value()}, ret_reg);
- }
- });
-
- // Mark the function entry internally.
- builder_->EmitFunction(gsymbol.value(), param_names.size(), param_names,
- VMFuncInfo::FuncKind::kVMTIRFunc, registers_num_);
- builder_->EndFunction(gsymbol.value());
-
- Type ret_type = VoidType();
- ffi::Array<tirx::Var> tir_params = {ctx_ptr_, reg_anylist_handle_,
const_anylist_handle_,
- func_anylist_handle_};
- ffi::String tir_func_name = system_lib_prefix_.value_or("") + "__vmtir__"
+ gsymbol.value();
- tirx::PrimFunc tir_func(tir_params, body, ret_type);
- tir_func = WithAttr(tir_func, "global_symbol", tir_func_name);
- tir_func = WithAttr(tir_func, tvm::attr::kSTir, true);
- registers_num_ = 0;
- var_map_.clear();
- stmt_stack_.clear();
- return tir_func;
- }
-
- ffi::Optional<Expr> VisitExpr_(const SeqExprNode* op) final {
- for (auto block : op->blocks) {
- for (Binding binding : block->bindings) {
- Expr expr = GetBoundValue(binding);
- ffi::Optional<Expr> value = VisitExpr(expr);
-
- if (expr.as<Var>() && value.has_value()) {
- // For a normalized relax module, there should be one
- // register for each relax::Binding. This makes the Relax
- // semantics of R.vm.kill_* operate the same as the Python
- // "del" operator. These bindings may be removable by using
- // relax.transform.CanonicalizeBindings earlier in lowering.
- auto new_reg = NewRegister();
- EmitCallPacked("vm.builtin.copy", {value.value()}, new_reg);
- value = RegListGet(new_reg);
- }
-
- this->var_map_.insert({binding->var, value});
- }
- }
- return this->VisitExpr(op->body);
- }
-
- ffi::Optional<Expr> VisitExpr_(const CallNode* call_node) final {
- Call call = ffi::GetRef<Call>(call_node);
-
- if (call_node->op.same_as(null_value_op_)) {
- return tvm::Call(tvm::PointerType::VoidPointerTy(),
tirx::builtin::reinterpret(),
- {IntImm::Int64(0)});
- }
- int64_t dst_reg = HasVoidType(call) ? -1 : NewRegister();
- if (call->op.as<OpNode>()) {
- if (call_node->op.same_as(call_builtin_with_ctx_op_)) {
- EmitCallBuiltinWithCtx(call, dst_reg);
- } else if (call_node->op.same_as(alloc_storage_op_)) {
- EmitAllocStorage(call, dst_reg);
- } else if (call_node->op.same_as(alloc_tensor_op_)) {
- EmitAllocTensor(call, dst_reg);
- } else if (call_node->op.same_as(kill_object_op_)) {
- dst_reg = EmitKillObject(call);
- } else {
- // every "normal" operator is lowered to a global var in the IRModule.
The Attrs for those
- // ops are handled in a pass when lowering them to TIR.
- TVM_FFI_THROW(InternalError)
- << "CodeGenVMTIR cannot emit this Relax operator directly. "
- << "Run the appropriate lowering pass, or route the operator to an
external "
- << "codegen before VM codegen.\nOffending call:\n"
- << call;
- }
- } else {
- EmitNormalCall(call, dst_reg);
- }
- if (dst_reg >= 0) {
- return RegListGet(dst_reg);
- } else {
- return std::nullopt;
- }
- }
-
- ffi::Optional<Expr> VisitExpr_(const IfNode* op) final {
- // Reserve a register for return
- size_t merge_register = NewRegister();
- Expr cond_value = this->VisitExpr(op->cond).value();
-
- PrimExpr condition = tvm::Call(tvm::PrimType::Bool(),
tirx::builtin::tvm_call_packed(),
- {StringImm("vm.builtin.read_if_cond"),
cond_value})
- .as_or_throw<PrimExpr>();
-
- tirx::Stmt true_branch = WithNewScope([&]() {
- Expr true_value = this->VisitExpr(op->true_branch).value();
- this->EmitCallPacked("vm.builtin.copy", {true_value}, merge_register);
- });
- tirx::Stmt false_branch = WithNewScope([&]() {
- Expr false_value = this->VisitExpr(op->false_branch).value();
- this->EmitCallPacked("vm.builtin.copy", {false_value}, merge_register);
- });
- this->EmitStmt(tirx::IfThenElse(condition, true_branch, false_branch));
- return RegListGet(merge_register);
- }
-
- ffi::Optional<Expr> VisitExpr_(const VarNode* op) final {
- Var var = ffi::GetRef<Var>(op);
- auto it = this->var_map_.find(var);
- TVM_FFI_ICHECK(it != this->var_map_.end()) << "Var " << var << " is not
defined";
- return it->second;
- }
-
- ffi::Optional<Expr> VisitExpr_(const DataflowVarNode* op) final {
- return VisitExpr_(static_cast<const VarNode*>(op));
- }
-
-#define VM_TIR_PRIM_EXPR(OP) \
- ffi::Optional<Expr> VisitExpr_(const OP* op) final { return
ffi::GetRef<Expr>(op); }
-
- VM_TIR_PRIM_EXPR(TensorLoadNode);
- VM_TIR_PRIM_EXPR(prim::AddNode);
- VM_TIR_PRIM_EXPR(prim::LShiftNode);
- VM_TIR_PRIM_EXPR(prim::RShiftNode);
- VM_TIR_PRIM_EXPR(prim::BitwiseAndNode);
- VM_TIR_PRIM_EXPR(prim::BitwiseOrNode);
- VM_TIR_PRIM_EXPR(prim::BitwiseXorNode);
- VM_TIR_PRIM_EXPR(prim::BitwiseNotNode);
- VM_TIR_PRIM_EXPR(prim::SubNode);
- VM_TIR_PRIM_EXPR(prim::MulNode);
- VM_TIR_PRIM_EXPR(prim::DivNode);
- VM_TIR_PRIM_EXPR(prim::ModNode);
- VM_TIR_PRIM_EXPR(prim::FloorDivNode);
- VM_TIR_PRIM_EXPR(prim::FloorModNode);
- VM_TIR_PRIM_EXPR(prim::MinNode);
- VM_TIR_PRIM_EXPR(prim::MaxNode);
- VM_TIR_PRIM_EXPR(prim::EQNode);
- VM_TIR_PRIM_EXPR(prim::NENode);
- VM_TIR_PRIM_EXPR(prim::LTNode);
- VM_TIR_PRIM_EXPR(prim::LENode);
- VM_TIR_PRIM_EXPR(prim::GTNode);
- VM_TIR_PRIM_EXPR(prim::GENode);
- VM_TIR_PRIM_EXPR(prim::AndNode);
- VM_TIR_PRIM_EXPR(prim::OrNode);
- VM_TIR_PRIM_EXPR(prim::CastNode);
- VM_TIR_PRIM_EXPR(prim::NotNode);
- VM_TIR_PRIM_EXPR(prim::SelectNode);
- VM_TIR_PRIM_EXPR(prim::RampNode);
- VM_TIR_PRIM_EXPR(prim::BroadcastNode);
- VM_TIR_PRIM_EXPR(prim::ShuffleNode);
- VM_TIR_PRIM_EXPR(tvm::IntImmNode);
- VM_TIR_PRIM_EXPR(tvm::FloatImmNode);
-
-#undef VM_TIR_PRIM_EXPR
-
- ffi::Optional<Expr> VisitExpr_(const GenericConstNode* op) final {
- return ConstListGet(builder_->ConvertConstant(op->value).value());
- }
-
- ffi::Optional<Expr> VisitExpr_(const ShapeExprNode* op) final {
- std::vector<int64_t> shape;
- for (PrimExpr e : op->values) {
- if (auto* int_value = e.as<IntImmNode>()) {
- shape.push_back(static_cast<int64_t>(int_value->value));
- } else {
- TVM_FFI_THROW(InternalError)
- << "Should only use constant shape after shape lowering: " <<
op->values;
- }
- }
- return ConstListGet(builder_->ConvertConstant(ffi::Shape(shape)).value());
- }
-
- ffi::Optional<Expr> VisitExpr_(const StringImmNode* op) final {
- return ConstListGet(builder_->ConvertConstant(op->value).value());
- }
- ffi::Optional<Expr> VisitExpr_(const DataTypeImmNode* op) final {
- return ConstListGet(builder_->ConvertConstant(op->value).value());
- }
-
- ffi::Optional<Expr> VisitExpr_(const TupleNode* op) final {
- Tuple tuple = ffi::GetRef<Tuple>(op);
- ffi::Array<Expr> args;
- for (auto arg : tuple->fields) {
- args.push_back(this->VisitExpr(arg).value());
- }
- int32_t dst_register = NewRegister();
- this->EmitCallPacked("vm.builtin.make_tuple", args, dst_register);
- return RegListGet(dst_register);
- }
-
- ffi::Optional<Expr> VisitExpr_(const TupleGetItemNode* op) final {
- TupleGetItem expr = ffi::GetRef<TupleGetItem>(op);
- ffi::Array<Expr> args = {this->VisitExpr(expr->tuple).value()};
-
- args.push_back(ConstInt64(expr->index));
-
- int64_t dst_register = NewRegister();
- this->EmitCallPacked("vm.builtin.tuple_getitem", args, dst_register);
- return RegListGet(dst_register);
- }
-
- // Lookup the function and see if it matches
- ffi::Optional<ffi::String> LookupFunction(const Expr& expr,
VMFuncInfo::FuncKind* kind) {
- if (auto* ext_func = expr.as<ExternFuncNode>()) {
- *kind = VMFuncInfo::FuncKind::kPackedFunc;
- return ext_func->global_symbol;
- } else if (auto* gvar_ptr = expr.as<GlobalVarNode>()) {
- GlobalVar gvar = ffi::GetRef<GlobalVar>(gvar_ptr);
- // Run a look up in the env to see if it maps to an extern func.
- auto it = ctx_mod_->functions.find(gvar);
- if (it != ctx_mod_->functions.end()) {
- BaseFunc func = (*it).second;
- if (auto* efunc = func.as<ExternFuncNode>()) {
- *kind = VMFuncInfo::FuncKind::kPackedFunc;
- return efunc->global_symbol;
- } else if (func.as<FunctionNode>()) {
- *kind = VMFuncInfo::FuncKind::kVMTIRFunc;
- return gvar->name_hint;
- } else if (func.as<tirx::PrimFuncNode>()) {
- *kind = VMFuncInfo::FuncKind::kPackedFunc;
- return gvar->name_hint;
- } else {
- *kind = VMFuncInfo::FuncKind::kPackedFunc;
- return gvar->name_hint;
- }
- }
- LOG(WARNING) << "Undefined global var " << gvar->name_hint;
- // undefined global var, consider eliminate later.
- *kind = VMFuncInfo::FuncKind::kPackedFunc;
- return gvar->name_hint;
- } else {
- return std::nullopt;
- }
- }
- // Lookup PrimFunc in the same module
- // We can do direct PrimFunc call in such cases
- ffi::Optional<tirx::PrimFunc> LookupPrimFunc(const ffi::String& name) {
- if (!ctx_mod_->ContainGlobalVar(name)) return std::nullopt;
-
- GlobalVar gvar = ctx_mod_->GetGlobalVar(name);
- auto it = ctx_mod_->functions.find(gvar);
- if (it != ctx_mod_->functions.end()) {
- BaseFunc func = (*it).second;
- if (auto* prim_func = func.as<tirx::PrimFuncNode>()) {
- return ffi::GetRef<tirx::PrimFunc>(prim_func);
- }
- }
- return std::nullopt;
- }
-
- ffi::Optional<Expr> VisitExpr_(const GlobalVarNode* op) final {
- VMFuncInfo::FuncKind kind;
- auto symbol = LookupFunction(ffi::GetRef<Expr>(op), &kind);
- TVM_FFI_ICHECK(symbol.has_value());
- builder_->DeclareFunction(symbol.value(), kind);
- return FuncListGet(builder_->GetFunction(symbol.value()).value());
- }
-
- ffi::Optional<Expr> VisitExpr_(const ExternFuncNode* op) final {
- builder_->DeclareFunction(op->global_symbol,
VMFuncInfo::FuncKind::kPackedFunc);
- return FuncListGet(builder_->GetFunction(op->global_symbol).value());
- }
-
- void EmitAllocStorage(const Call& call_node, int64_t dst_reg) {
- // Handle args of the call
- ffi::Array<Expr> args;
- args.push_back(ctx_ptr_);
- for (Expr arg : call_node->args) {
- args.push_back(this->VisitExpr(arg).value());
- }
- this->EmitCallPacked("vm.builtin.alloc_storage", args, dst_reg);
- }
-
- void EmitAllocTensor(const Call& call_node, int64_t dst_reg) {
- TVM_FFI_ICHECK_EQ(call_node->args.size(), 5);
- ffi::Array<Expr> args;
- for (int i = 0; i < 4; ++i) {
- args.push_back(this->VisitExpr(call_node->args[i]).value());
- }
- int64_t vdevice_index = -1;
- if (const auto* int_imm = call_node->args[4].as<IntImmNode>()) {
- vdevice_index = int_imm->value.as<int>().value();
- }
- auto vdevice = GetGlobalVDevice(ctx_mod_, vdevice_index);
-
- if (vdevice.has_value()) {
- args.push_back(StringImm(vdevice.value()->memory_scope));
- }
-
- this->EmitCallPacked("vm.builtin.alloc_tensor", args, dst_reg);
- }
-
- int64_t EmitKillObject(const Call& call_node) {
- TVM_FFI_ICHECK_EQ(call_node->args.size(), 1);
- Expr arg = this->VisitExpr(call_node->args[0]).value();
-
- // Check the arg is a register.
- const auto* tir_call = arg.as<CallNode>();
- TVM_FFI_ICHECK(tir_call != nullptr);
- TVM_FFI_ICHECK(tir_call->op.same_as(tirx::builtin::anylist_getitem()));
- TVM_FFI_ICHECK(tir_call->args.size() == 2);
- TVM_FFI_ICHECK(tir_call->args[0].same_as(reg_anylist_handle_));
- const auto* p_dst_reg = tir_call->args[1].as<IntImmNode>();
- TVM_FFI_ICHECK(p_dst_reg != nullptr);
- TVM_FFI_ICHECK(
-
p_dst_reg->ty.as_or_throw<PrimType>().MatchesElementType(DLDataTypeCode::kDLInt,
32));
-
- int64_t dst_reg = static_cast<int64_t>(p_dst_reg->value);
- this->EmitCallPacked("vm.builtin.null_value", {}, dst_reg);
- return dst_reg;
- }
-
- void EmitCallBuiltinWithCtx(const Call& call_node, int64_t dst_reg) {
- ffi::Array<Expr> args;
- // if context is required, pass as first argument.
- args.push_back(ctx_ptr_);
- auto* func = call_node->args[0].as<ExternFuncNode>();
- TVM_FFI_ICHECK(func) << "CallBuiltin comes with extern func";
-
- auto tuple_arg = call_node->args[1].as_or_throw<Tuple>();
-
- // Handle args of the call
- for (Expr arg : tuple_arg->fields) {
- args.push_back(this->VisitExpr(arg).value());
- }
-
- this->EmitCallPacked(func->global_symbol, args, dst_reg);
- }
-
- void EmitNormalCall(const Call& call_node, int64_t dst_reg) {
- ffi::Array<Expr> args = VisitArray(call_node->args);
- // A function can be a closure that comes from parent
- // Do call closure to be safe.
- VMFuncInfo::FuncKind kind;
- auto symbol = LookupFunction(call_node->op, &kind);
-
- if (symbol.has_value() && kind == VMFuncInfo::FuncKind::kPackedFunc) {
- // primfunc in the same module.
- // use cpacked to directly invoke without named based lookup
- if (ffi::Optional<tirx::PrimFunc> prim_func =
LookupPrimFunc(symbol.value())) {
- this->EmitCallCPacked(prim_func.value(), args, dst_reg);
- } else {
- this->EmitCallPacked(symbol.value(), args, dst_reg);
- }
- } else {
- // Default path, leverage function table and invoke as closure
- ffi::Array<Expr> all_args;
- all_args.push_back(ctx_ptr_);
- all_args.push_back(this->VisitExpr(call_node->op).value());
- for (auto arg : args) {
- all_args.push_back(arg);
- }
- this->EmitCallPacked("vm.builtin.invoke_closure", all_args, dst_reg);
- }
- }
-
- template <typename FLambda>
- tirx::Stmt WithNewScope(const FLambda& callback) {
- stmt_stack_.push_back({});
- callback();
- tirx::Stmt stmt = tirx::SeqStmt::Flatten(stmt_stack_.back());
- stmt_stack_.pop_back();
- return stmt;
- }
-
- ffi::Array<Expr> VisitArray(const ffi::Array<Expr>& arr) {
- ffi::Array<Expr> ret;
- for (size_t i = 0; i < arr.size(); ++i) {
- ret.push_back(this->VisitExpr(arr[i]).value());
- }
- return ret;
- }
- /*! \brief Internal ExecBuilder. */
- relax::ExecBuilder builder_;
- /*! \brief List to ctx_ptr */
- tirx::Var ctx_ptr_;
- /*! \brief List to store temp object registers */
- tirx::Var reg_anylist_handle_;
- /*! \brief List to store closures */
- tirx::Var func_anylist_handle_;
- /*! \brief List to store constants */
- tirx::Var const_anylist_handle_;
- /*!
- * \brief Total number of virtual registers allocated.
- * \note The first two registers are reserved for special registers.
- */
- int64_t registers_num_ = 0;
- /*! \brief Stack to build up statements */
- std::vector<std::vector<tirx::Stmt>> stmt_stack_;
- /*! \brief Map from var to Expr. */
- std::unordered_map<Var, ffi::Optional<Expr>> var_map_;
- /*! \brief the context module. */
- IRModule ctx_mod_;
- /*! \brief system lib prefix */
- ffi::Optional<ffi::String> system_lib_prefix_;
- /*! \brief Cache ops that need to be frequently used later to reduce lookup
overhead. */
- const Op& alloc_storage_op_ = Op::Get("relax.vm.alloc_storage");
- const Op& alloc_tensor_op_ = Op::Get("relax.vm.alloc_tensor");
- const Op& kill_object_op_ = Op::Get("relax.vm.kill_object");
- const Op& call_builtin_with_ctx_op_ = Op::Get("relax.call_builtin_with_ctx");
- const Op& null_value_op_ = Op::Get("relax.null_value");
-};
-
-/*!
- * \brief Create the Relax VM executable from all relax.Function in mod.
- * and add them to exec_builder. Create extra TIR functions.
- *
- * \param exec_builder Builder to collect executables.
- * \param mod Input module.
- * \return Extra TIR module created.
- */
-IRModule VMTIRCodeGen(ExecBuilder exec_builder, IRModule mod) {
- return CodeGenVMTIR::Run(exec_builder, mod);
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("relax.VMTIRCodeGen", VMTIRCodeGen);
-}
-
-} // namespace codegen_vm
-} // namespace relax
-} // namespace tvm
diff --git a/src/relax/backend/vm/exec_builder.cc
b/src/relax/backend/vm/exec_builder.cc
index 2186ceb252..a4b73e453a 100644
--- a/src/relax/backend/vm/exec_builder.cc
+++ b/src/relax/backend/vm/exec_builder.cc
@@ -181,11 +181,6 @@ void ExecBuilderNode::EmitIf(vm::Instruction::Arg cond,
vm::Index false_offset)
void ExecBuilderNode::CheckExecutable() {
for (auto it = exec_->func_table.cbegin(); it != exec_->func_table.cend();
++it) {
if (it->kind == VMFuncInfo::FuncKind::kPackedFunc) continue;
- if (it->kind == VMFuncInfo::FuncKind::kVMTIRFunc) {
- TVM_FFI_ICHECK_GE(it->register_file_size, it->num_args + 1)
- << "Function " << it->name << " do not meet register file
constraint.";
- continue;
- }
Index num_inputs = it->num_args;
std::unordered_set<RegName> dst_registers;
std::unordered_set<RegName> arg_registers;
@@ -267,7 +262,6 @@ void ExecBuilderNode::Formalize() {
// and decide the number of registers to allocate for each VMFunction in the
VMExecutable
for (auto it = this->exec_->func_table.begin(); it !=
this->exec_->func_table.end(); ++it) {
if (it->kind == VMFuncInfo::FuncKind::kPackedFunc) continue;
- if (it->kind == VMFuncInfo::FuncKind::kVMTIRFunc) continue;
Index num_inputs = it->num_args;
RegName register_idx = num_inputs;
diff --git a/src/relax/transform/kill_after_last_use.cc
b/src/relax/transform/kill_after_last_use.cc
index 098b8a05d6..8785428c73 100644
--- a/src/relax/transform/kill_after_last_use.cc
+++ b/src/relax/transform/kill_after_last_use.cc
@@ -118,12 +118,9 @@ class CollectLastUsage : public ExprVisitor {
// this may be handled more easily at the CodeGenVM level.
//
// Variables bound to `relax.null_value` are excluded for the
- // same reason as constants: both CodeGenVM and CodeGenVMTIR
- // special-case `null_value` to bypass register/anylist-slot
- // allocation, so such a variable is never a valid target for
- // R.vm.kill_object. It is currently the only operator either
- // codegen special-cases this way; a new special case added to
- // either codegen should be reflected here as well.
+ // same reason as constants: CodeGenVM special-cases `null_value` to
bypass register
+ // allocation, so such a variable is never a valid target for
R.vm.kill_object. A new
+ // codegen special case should be reflected here as well.
bool stored_in_vm_register =
!(visitor.constant_tensors_.count(var) ||
visitor.null_value_objects_.count(var) ||
var->ty.as<FuncTypeNode>() || var->ty.as<ShapeTypeNode>() ||
@@ -216,7 +213,7 @@ class CollectLastUsage : public ExprVisitor {
std::unordered_set<const VarNode*> constant_tensors_;
// Variables bound to `relax.null_value`, which do not occupy a VM
- // register in either CodeGenVM or CodeGenVMTIR, and therefore must
+ // register in CodeGenVM, and therefore must
// never be passed to R.vm.kill_object.
std::unordered_set<const VarNode*> null_value_objects_;
diff --git a/src/runtime/vm/builtin.cc b/src/runtime/vm/builtin.cc
index 9c2651e7c8..e39f187335 100644
--- a/src/runtime/vm/builtin.cc
+++ b/src/runtime/vm/builtin.cc
@@ -746,65 +746,3 @@ TVM_FFI_STATIC_INIT_BLOCK() {
} // namespace vm
} // namespace runtime
} // namespace tvm
-
-//-------------------------------------------------
-// AnyList C runtime API: keep in relax for now.
-//--------------------------------------------------
-extern "C" {
-/*!
- * \brief Backend function to get anylist item and set into Packed Func call
arg stack.
- *
- * \param anylist The handle to the anylist, backed by ffi::Any*
- * \param int The index.
- * \param args The args stack.
- * \param arg_offset The offset of argument.
- * \return 0 when no error is thrown, -1 when failure happens
- */
-TVM_RUNTIME_DLL int TVMBackendAnyListSetPackedArg(void* anylist, int index,
TVMFFIAny* args,
- int arg_offset);
-/*!
- * \brief Backend function to get anylist item and set into Packed Func call
arg stack.
- *
- * \param anylist The handle to the anylist, backed by ffi::Any*
- * \param int The index.
- */
-TVM_RUNTIME_DLL int TVMBackendAnyListResetItem(void* anylist, int index);
-
-/*!
- * \brief Backend function to set anylist item by moving from packed func
return.
- *
- * \param anylist The handle to the anylist, backed by ffi::Any*
- * \param int The index.
- * \param args The args stack.
- * \param type_codes The type codes stack.
- * \param arg_offset The offset of argument.
- * \return 0 when no error is thrown, -1 when failure happens.
- */
-TVM_RUNTIME_DLL int TVMBackendAnyListMoveFromPackedReturn(void* anylist, int
index, TVMFFIAny* args,
- int ret_offset);
-
-int TVMBackendAnyListSetPackedArg(void* anylist, int index, TVMFFIAny* args,
int arg_offset) {
- using namespace tvm::runtime;
- TVM_FFI_SAFE_CALL_BEGIN();
- auto* list = static_cast<TVMFFIAny*>(anylist);
- args[arg_offset] = list[index];
- TVM_FFI_SAFE_CALL_END();
-}
-
-int TVMBackendAnyListResetItem(void* anylist, int index) {
- using namespace tvm::runtime;
- TVM_FFI_SAFE_CALL_BEGIN();
- auto* list = static_cast<tvm::ffi::Any*>(anylist);
- list[index] = nullptr;
- TVM_FFI_SAFE_CALL_END();
-}
-
-int TVMBackendAnyListMoveFromPackedReturn(void* anylist, int index, TVMFFIAny*
args,
- int ret_offset) {
- using namespace tvm::runtime;
- TVM_FFI_SAFE_CALL_BEGIN();
- auto* list = static_cast<tvm::ffi::Any*>(anylist);
- list[index] =
tvm::ffi::details::AnyUnsafe::MoveTVMFFIAnyToAny(&args[ret_offset]);
- TVM_FFI_SAFE_CALL_END();
-}
-} // extern "C"
diff --git a/src/runtime/vm/executable.cc b/src/runtime/vm/executable.cc
index 9e3a5f9323..ace1c871a8 100644
--- a/src/runtime/vm/executable.cc
+++ b/src/runtime/vm/executable.cc
@@ -472,10 +472,6 @@ ffi::String VMExecutable::AsText() const {
os << "@" << gfunc.name << " packed_func;\n\n";
continue;
}
- if (gfunc.kind == VMFuncInfo::FuncKind::kVMTIRFunc) {
- os << "@" << gfunc.name << " num_inputs=" << gfunc.num_args << "
vm_tir_func;\n\n";
- continue;
- }
TVM_FFI_ICHECK(gfunc.kind == VMFuncInfo::FuncKind::kVMFunc);
os << "@" << gfunc.name << ":\n";
size_t start_instr = gfunc.start_instr;
@@ -553,9 +549,6 @@ ffi::String VMExecutable::AsPython() const {
if (gfunc.kind == VMFuncInfo::FuncKind::kPackedFunc) {
continue;
}
- if (gfunc.kind == VMFuncInfo::FuncKind::kVMTIRFunc) {
- continue;
- }
TVM_FFI_ICHECK(gfunc.kind == VMFuncInfo::FuncKind::kVMFunc);
os << "with ib.function(\"" << gfunc.name << "\", num_inputs=" <<
gfunc.num_args << "):\n";
diff --git a/src/runtime/vm/vm.cc b/src/runtime/vm/vm.cc
index b5419625d1..89632e2a98 100644
--- a/src/runtime/vm/vm.cc
+++ b/src/runtime/vm/vm.cc
@@ -612,46 +612,20 @@ ffi::Optional<VMClosure>
VirtualMachineImpl::GetClosureInternal(const ffi::Strin
Index gf_idx = it->second;
const VMFuncInfo& finfo = exec_->func_table[gf_idx];
- if (finfo.kind == VMFuncInfo::FuncKind::kVMFunc) {
- // NOTE: should not capture strong ref to self and avoid cyclic ref.
- auto impl = ffi::Function([gf_idx](ffi::PackedArgs args, ffi::Any* rv) {
- // Per convention, ctx ptr is a VirtualMachine*
- VirtualMachine* ctx_ptr =
static_cast<VirtualMachine*>(args[0].cast<void*>());
-
- std::vector<RegType> inputs(args.size() - 1);
- for (size_t i = 0; i < inputs.size(); ++i) {
- inputs[i] = args[i + 1];
- }
- *rv = static_cast<VirtualMachineImpl*>(ctx_ptr)->InvokeBytecode(gf_idx,
inputs);
- });
- return VMClosure(func_name, impl);
- } else {
- TVM_FFI_ICHECK(finfo.kind == VMFuncInfo::FuncKind::kVMTIRFunc)
- << "Cannot support closure with function kind " <<
static_cast<int>(finfo.kind);
- ffi::Optional<ffi::Function> tir_func = GetFuncFromImports("__vmtir__" +
finfo.name);
- TVM_FFI_ICHECK(tir_func.has_value())
- << "Cannot find underlying compiled tirx function of VMTIRFunc " <<
finfo.name;
- auto impl = ffi::Function([this, finfo, tir_func](ffi::PackedArgs args,
ffi::Any* rv) {
- // Per convention, ctx ptr is a VirtualMachine*
- VirtualMachine* ctx_ptr =
static_cast<VirtualMachine*>(args[0].cast<void*>());
- TVM_FFI_ICHECK(ctx_ptr == this);
- TVM_FFI_ICHECK_EQ(args.size() - 1, finfo.num_args)
- << "Function " << finfo.name << " expects " << finfo.num_args << "
arguments";
- TVM_FFI_ICHECK_GE(finfo.register_file_size, finfo.num_args + 1);
- std::vector<ffi::Any> reg_file(finfo.register_file_size);
- for (int64_t i = 0; i < finfo.num_args; ++i) {
- reg_file[i] = args[i + 1];
- }
- void* reg_anylist_handle = reg_file.data();
- void* const_anylist_handle = this->const_pool_.data();
- void* func_anylist_handle = this->func_pool_.data();
- (*tir_func)(static_cast<void*>(ctx_ptr), reg_anylist_handle,
const_anylist_handle,
- func_anylist_handle);
- // Return value always stored after inputs.
- *rv = reg_file[finfo.num_args];
- });
- return VMClosure(func_name, impl);
- }
+ TVM_FFI_ICHECK(finfo.kind == VMFuncInfo::FuncKind::kVMFunc)
+ << "Cannot support closure with function kind " <<
static_cast<int>(finfo.kind);
+ // NOTE: should not capture strong ref to self and avoid cyclic ref.
+ auto impl = ffi::Function([gf_idx](ffi::PackedArgs args, ffi::Any* rv) {
+ // Per convention, ctx ptr is a VirtualMachine*
+ VirtualMachine* ctx_ptr =
static_cast<VirtualMachine*>(args[0].cast<void*>());
+
+ std::vector<RegType> inputs(args.size() - 1);
+ for (size_t i = 0; i < inputs.size(); ++i) {
+ inputs[i] = args[i + 1];
+ }
+ *rv = static_cast<VirtualMachineImpl*>(ctx_ptr)->InvokeBytecode(gf_idx,
inputs);
+ });
+ return VMClosure(func_name, impl);
}
//--------------------------------------------------------------------
@@ -716,8 +690,7 @@ void VirtualMachineImpl::InitFuncPool() {
func_pool_[func_index] = *func;
} else {
- TVM_FFI_ICHECK(info.kind == VMFuncInfo::FuncKind::kVMFunc ||
- info.kind == VMFuncInfo::FuncKind::kVMTIRFunc);
+ TVM_FFI_ICHECK(info.kind == VMFuncInfo::FuncKind::kVMFunc);
auto clo = this->GetClosure(info.name);
func_pool_[func_index] = clo;
}
diff --git a/src/tirx/op/builtin.cc b/src/tirx/op/builtin.cc
index 24ef6af404..3edb8f68e3 100644
--- a/src/tirx/op/builtin.cc
+++ b/src/tirx/op/builtin.cc
@@ -314,19 +314,6 @@ TIR_DEFINE_BUILTIN_FUNC(undef)
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kReadState))
.set_num_inputs(0);
-TIR_DEFINE_BUILTIN_FUNC(anylist_getitem)
- .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kReadState));
-
-TIR_DEFINE_BUILTIN_FUNC(anylist_resetitem)
- .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque))
- .set_attr<TGlobalSymbol>("TGlobalSymbol", "TVMBackendAnyListResetItem");
-
-TIR_DEFINE_BUILTIN_FUNC(anylist_setitem_call_packed)
- .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
-
-TIR_DEFINE_BUILTIN_FUNC(anylist_setitem_call_cpacked)
- .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
-
TIR_DEFINE_BUILTIN_FUNC(get_active_lane_mask)
.set_num_inputs(2)
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kPure))
diff --git a/src/tirx/op/runtime.cc b/src/tirx/op/runtime.cc
deleted file mode 100644
index fb24a82ae6..0000000000
--- a/src/tirx/op/runtime.cc
+++ /dev/null
@@ -1,43 +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 tirx/op/runtime.cc
- * \brief TIR ops for runtime functions.
- */
-#include <tvm/ir/op.h>
-#include <tvm/tirx/op_attr_types.h>
-
-namespace tvm {
-namespace tirx {
-
-TVM_REGISTER_OP("tirx.TVMBackendAnyListSetPackedArg")
- .set_num_inputs(5)
- .set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"), 1)
- .set_attr<TGlobalSymbol>("TGlobalSymbol", "TVMBackendAnyListSetPackedArg")
- .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
-
-TVM_REGISTER_OP("tirx.TVMBackendAnyListMoveFromPackedReturn")
- .set_num_inputs(3)
- .set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"), 1)
- .set_attr<TGlobalSymbol>("TGlobalSymbol",
"TVMBackendAnyListMoveFromPackedReturn")
- .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
-
-} // namespace tirx
-} // namespace tvm
diff --git a/src/tirx/transform/lower_tvm_builtin.cc
b/src/tirx/transform/lower_tvm_builtin.cc
index 1d0e240b7f..73b8cd779d 100644
--- a/src/tirx/transform/lower_tvm_builtin.cc
+++ b/src/tirx/transform/lower_tvm_builtin.cc
@@ -434,10 +434,6 @@ class BuiltinLower : public StmtExprMutator {
return MakeCallPackedGeneric(op, 0, builtin::tvm_call_packed_lowered());
} else if (op->op.same_as(builtin::tvm_call_cpacked())) {
return MakeCallPackedGeneric(op, 0, builtin::tvm_call_cpacked_lowered());
- } else if (op->op.same_as(builtin::anylist_setitem_call_packed())) {
- return MakeAnyListSetItemCallPacked(op,
builtin::tvm_call_packed_lowered());
- } else if (op->op.same_as(builtin::anylist_setitem_call_cpacked())) {
- return MakeAnyListSetItemCallPacked(op,
builtin::tvm_call_cpacked_lowered());
} else if (op->op.same_as(builtin::tvm_stack_make_shape())) {
return MakeShape(op);
} else if (op->op.same_as(builtin::tvm_stack_make_array())) {
@@ -538,78 +534,53 @@ class BuiltinLower : public StmtExprMutator {
void SetPackedArg(Expr arg, const Var& args_stack, size_t stack_offset,
std::vector<tirx::Stmt>* prep_seq) {
- auto* call_pattern = arg.as<CallNode>();
- if (call_pattern && call_pattern->op.same_as(builtin::anylist_getitem())) {
- // call runtime function to set anylist
- static const Op& anylist_set_packed_arg_op =
Op::Get("tirx.TVMBackendAnyListSetPackedArg");
- prep_seq->emplace_back(Evaluate(Call(
- PrimType::Int(32), anylist_set_packed_arg_op,
- {call_pattern->args[0], call_pattern->args[1], args_stack,
ConstInt32(stack_offset)})));
+ int arg_type_index;
+ if (arg.as<StringImmNode>()) {
+ arg_type_index = ffi::TypeIndex::kTVMFFIRawStr;
+ arg = reinterpret(PointerType::VoidPointerTy(), std::move(arg));
+ } else if (arg->ty.as<PointerTypeNode>()) {
+ arg_type_index = IsArrayHandle(arg) ? ffi::TypeIndex::kTVMFFIDLTensorPtr
+ : ffi::TypeIndex::kTVMFFIOpaquePtr;
} else {
- int arg_type_index;
- if (arg.as<StringImmNode>()) {
- arg_type_index = ffi::TypeIndex::kTVMFFIRawStr;
- arg = reinterpret(PointerType::VoidPointerTy(), std::move(arg));
- } else if (arg->ty.as<PointerTypeNode>()) {
- arg_type_index = IsArrayHandle(arg) ?
ffi::TypeIndex::kTVMFFIDLTensorPtr
- : ffi::TypeIndex::kTVMFFIOpaquePtr;
- } else {
- PrimExpr prim_arg = arg.as_or_throw<PrimExpr>();
- PrimType arg_ty = prim_arg.ty();
- PrimType api_ty = APIType(arg_ty);
- if (arg_ty != api_ty) {
- arg = prim::Cast(api_ty, prim_arg);
- }
- if (api_ty.MatchesCode(DLDataTypeCode::kDLBool)) {
- arg_type_index = ffi::TypeIndex::kTVMFFIBool;
- } else if (api_ty.MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt)) {
- arg_type_index = ffi::TypeIndex::kTVMFFIInt;
- } else if (api_ty.code() == DLDataTypeCode::kDLFloat) {
- arg_type_index = ffi::TypeIndex::kTVMFFIFloat;
- } else {
- TVM_FFI_THROW(InternalError) << "Unsupported type: " << api_ty;
- }
+ PrimExpr prim_arg = arg.as_or_throw<PrimExpr>();
+ PrimType arg_ty = prim_arg.ty();
+ PrimType api_ty = APIType(arg_ty);
+ if (arg_ty != api_ty) {
+ arg = prim::Cast(api_ty, prim_arg);
}
-
- // opaque handle need to set the kind properly
- if (arg_type_index == ffi::TypeIndex::kTVMFFIOpaquePtr) {
- prep_seq->emplace_back(
- IfThenElse(Call(PrimType::Bool(), builtin::isnullptr(),
{arg}).as_or_throw<PrimExpr>(),
- TVMStructSet(args_stack, stack_offset,
builtin::kTVMFFIAnyTypeIndex,
- ConstInt32(ffi::TypeIndex::kTVMFFINone)),
- TVMStructSet(args_stack, stack_offset,
builtin::kTVMFFIAnyTypeIndex,
-
ConstInt32(ffi::TypeIndex::kTVMFFIOpaquePtr))));
+ if (api_ty.MatchesCode(DLDataTypeCode::kDLBool)) {
+ arg_type_index = ffi::TypeIndex::kTVMFFIBool;
+ } else if (api_ty.MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt)) {
+ arg_type_index = ffi::TypeIndex::kTVMFFIInt;
+ } else if (api_ty.code() == DLDataTypeCode::kDLFloat) {
+ arg_type_index = ffi::TypeIndex::kTVMFFIFloat;
} else {
- prep_seq->emplace_back(TVMStructSet(args_stack, stack_offset,
builtin::kTVMFFIAnyTypeIndex,
- ConstInt32(arg_type_index)));
+ TVM_FFI_THROW(InternalError) << "Unsupported type: " << api_ty;
}
- // set zero padding to ensure compatibility with FFI convention
- prep_seq->emplace_back(
- TVMStructSet(args_stack, stack_offset,
builtin::kTVMFFIAnyZeroPadding, ConstInt32(0)));
- // handle arg value
- // NOTE: the intrinsic codegen will handle padding value clear for 32bit
- // types or types that are smaller than 64 bits.
+ }
+
+ // opaque handle need to set the kind properly
+ if (arg_type_index == ffi::TypeIndex::kTVMFFIOpaquePtr) {
prep_seq->emplace_back(
- TVMStructSet(args_stack, stack_offset,
builtin::kTVMFFIAnyUnionValue, arg));
+ IfThenElse(Call(PrimType::Bool(), builtin::isnullptr(),
{arg}).as_or_throw<PrimExpr>(),
+ TVMStructSet(args_stack, stack_offset,
builtin::kTVMFFIAnyTypeIndex,
+ ConstInt32(ffi::TypeIndex::kTVMFFINone)),
+ TVMStructSet(args_stack, stack_offset,
builtin::kTVMFFIAnyTypeIndex,
+
ConstInt32(ffi::TypeIndex::kTVMFFIOpaquePtr))));
+ } else {
+ prep_seq->emplace_back(TVMStructSet(args_stack, stack_offset,
builtin::kTVMFFIAnyTypeIndex,
+ ConstInt32(arg_type_index)));
}
+ // set zero padding to ensure compatibility with FFI convention
+ prep_seq->emplace_back(
+ TVMStructSet(args_stack, stack_offset, builtin::kTVMFFIAnyZeroPadding,
ConstInt32(0)));
+ // handle arg value
+ // NOTE: the intrinsic codegen will handle padding value clear for 32bit
+ // types or types that are smaller than 64 bits.
+ prep_seq->emplace_back(
+ TVMStructSet(args_stack, stack_offset, builtin::kTVMFFIAnyUnionValue,
arg));
}
- PrimExpr MakeAnyListSetItemCallPacked(const CallNode* op, const Op&
lowered_op) {
- Expr list_handle = op->args[0];
- PrimExpr list_index = op->args[1].as_or_throw<PrimExpr>();
-
- Call call = MakeCallPackedGeneric(op, 2, lowered_op);
- Expr args_stack = call->args[1];
- // The stack offset of return value stack_end
- PrimExpr ret_offset = call->args[3].as_or_throw<PrimExpr>();
- auto& prep_seq = prep_seq_stack_.back();
- prep_seq.emplace_back(Evaluate(call.as_or_throw<PrimExpr>()));
- static const Op& anylist_move_from_packed_return_op =
- Op::Get("tirx.TVMBackendAnyListMoveFromPackedReturn");
- return Call(PrimType::Int(32), anylist_move_from_packed_return_op,
- {list_handle, list_index, args_stack, ret_offset})
- .as_or_throw<PrimExpr>();
- }
/*!
* \brief Generic tool to make low-level
* packed_call(other_args..., func_name, packed_arg0, packed_arg1...)
diff --git a/tests/python/relax/test_kill_after_last_use.py
b/tests/python/relax/test_kill_after_last_use.py
index 6247782c0d..41d67d2a83 100644
--- a/tests/python/relax/test_kill_after_last_use.py
+++ b/tests/python/relax/test_kill_after_last_use.py
@@ -109,8 +109,7 @@ def test_no_kill_for_null_value():
"""R.null_value() must never be targeted by R.vm.kill_object
A variable bound to `R.null_value()` is never assigned a real VM
- register/anylist slot by either CodeGenVM or CodeGenVMTIR (both
- special-case `null_value` to a sentinel value instead).
+ register by CodeGenVM, which special-cases `null_value` to a sentinel
value.
KillAfterLastUse must therefore never insert `R.vm.kill_object`
for such a variable, even though its type is `R.Any` (the same
type used by legitimate killable objectws such as VM storage).
@@ -148,7 +147,7 @@ def _assert_no_kill_of_null_value(func: tvm.relax.Function):
"""Assert no R.vm.kill_object call in `func` targets a null_value()-bound
var
An `R.null_value()`-bound variable never occupies a VM register in
- either CodeGenVM or CodeGenVMTIR, so passing one to
+ CodeGenVM, so passing one to
R.vm.kill_object is always invalid. Checking this structurally
(rather than only checking that relax.build succeeds) ensures the
test fails if a future change merely makes codegen tolerant of the
@@ -189,7 +188,7 @@ def
test_reapply_after_default_pipeline_builds_successfully():
KillAfterLastUse application, after VMShapeLower has introduced a
`shape_heap: R.Any = R.null_value()` binding), must not insert an
invalid `R.vm.kill_object(shape_heap)`, and the resulting module
- must still build successfully under every exec_mode.
+ must still build successfully using bytecode.
"""
@I.ir_module
@@ -208,8 +207,7 @@ def
test_reapply_after_default_pipeline_builds_successfully():
# Prove the invalid kill is gone, not merely that codegen tolerates it.
_assert_no_kill_of_null_value(second_kill["main"])
- for exec_mode in ["bytecode", "compiled"]:
- tvm.relax.build(second_kill, target=target, relax_pipeline="zero",
exec_mode=exec_mode)
+ tvm.relax.build(second_kill, target=target, relax_pipeline="zero")
if __name__ == "__main__":
diff --git a/tests/python/relax/test_relax_operators.py
b/tests/python/relax/test_relax_operators.py
index 9541e7b06e..d6502b3dbb 100644
--- a/tests/python/relax/test_relax_operators.py
+++ b/tests/python/relax/test_relax_operators.py
@@ -33,9 +33,6 @@ from tvm.script import ir as I
from tvm.script import relax as R
from tvm.script import tirx as T
-exec_mode = tvm.testing.parameter("bytecode", "compiled")
-
-
m = T.dynamic("m")
n = T.dynamic("n")
@@ -49,7 +46,7 @@ class InputModule:
return y, y_sorted
-def run_cpu(mod, func_name, *args, exec_mode):
+def run_cpu(mod, func_name, *args):
if isinstance(mod, relax.Function):
func = mod
args = [func_name, *args]
@@ -57,17 +54,17 @@ def run_cpu(mod, func_name, *args, exec_mode):
mod = tvm.IRModule.from_expr(func)
target = tvm.target.Target("llvm")
- ex = relax.build(mod, target, exec_mode=exec_mode)
+ ex = relax.build(mod, target)
vm = relax.VirtualMachine(ex, tvm.cpu())
return vm[func_name](*args)
-def test_unique(exec_mode):
+def test_unique():
# TODO(prakalp): also add test for compiling and running on CUDA device.
data_numpy = np.random.randint(0, 16, (16, 16))
data = tvm.runtime.tensor(data_numpy)
- result, result_sorted = run_cpu(InputModule, "foo", data,
exec_mode=exec_mode)
+ result, result_sorted = run_cpu(InputModule, "foo", data)
expected_output_sorted, indices = np.unique(data_numpy, return_index=True)
expected_output = [data_numpy.flatten()[index] for index in
sorted(indices)]
@@ -93,7 +90,7 @@ class PrintTest:
return x
-def test_print(exec_mode):
+def test_print():
try:
stdout = sys.stdout
with tempfile.TemporaryFile(mode="w+") as test_out:
@@ -102,7 +99,6 @@ def test_print(exec_mode):
PrintTest,
"foo",
tvm.runtime.tensor(np.array(1).astype("int32")),
- exec_mode=exec_mode,
)
test_out.seek(0)
printed_text = str(test_out.read())
@@ -112,65 +108,65 @@ def test_print(exec_mode):
sys.stdout = stdout
-def test_assert_passes(exec_mode):
+def test_assert_passes():
@R.function(pure=False)
def func(x: R.Tensor((), "int32")):
_ = R.assert_op(relax.const(True))
return x
- run_cpu(func, tvm.runtime.tensor(np.array(1).astype("int32")),
exec_mode=exec_mode)
+ run_cpu(func, tvm.runtime.tensor(np.array(1).astype("int32")))
-def test_assert_passes_with_format_args(exec_mode):
+def test_assert_passes_with_format_args():
@R.function(pure=False)
def func(x: R.Tensor((), "int32")):
_ = R.assert_op(relax.const(True), x, format="You won't see me")
return x
- run_cpu(func, tvm.runtime.tensor(np.array(1).astype("int32")),
exec_mode=exec_mode)
+ run_cpu(func, tvm.runtime.tensor(np.array(1).astype("int32")))
-def test_assert_fails(exec_mode):
+def test_assert_fails():
@R.function(pure=False)
def func(x: R.Tensor((), "int32")):
_ = R.assert_op(relax.const(False))
return x
with pytest.raises(AssertionError, match="Assertion Failed"):
- run_cpu(func, tvm.runtime.tensor(np.array(1).astype("int32")),
exec_mode=exec_mode)
+ run_cpu(func, tvm.runtime.tensor(np.array(1).astype("int32")))
-def test_assert_fails_with_message(exec_mode):
+def test_assert_fails_with_message():
@R.function(pure=False)
def func(x: R.Tensor((), "int32")):
_ = R.assert_op(relax.const(False), format="I failed...")
return x
with pytest.raises(AssertionError, match="I failed..."):
- run_cpu(func, tvm.runtime.tensor(np.array(1).astype("int32")),
exec_mode=exec_mode)
+ run_cpu(func, tvm.runtime.tensor(np.array(1).astype("int32")))
-def test_assert_fails_with_args(exec_mode):
+def test_assert_fails_with_args():
@R.function(pure=False)
def func(x: R.Tensor((), "int32")):
_ = R.assert_op(relax.const(False), [x, x])
return x
with pytest.raises(AssertionError, match="5, 5"):
- run_cpu(func, tvm.runtime.tensor(np.array(5).astype("int32")),
exec_mode=exec_mode)
+ run_cpu(func, tvm.runtime.tensor(np.array(5).astype("int32")))
-def test_assert_fails_with_formatted_args(exec_mode):
+def test_assert_fails_with_formatted_args():
@R.function(pure=False)
def func(x: R.Tensor((), "int32")):
_ = R.assert_op(relax.const(False), x, format="Number: {}")
return x
with pytest.raises(AssertionError, match="Number: 6"):
- run_cpu(func, tvm.runtime.tensor(np.array(6).astype("int32")),
exec_mode=exec_mode)
+ run_cpu(func, tvm.runtime.tensor(np.array(6).astype("int32")))
-def test_assert_on_argument_passes(exec_mode):
+def test_assert_on_argument_passes():
@R.function(pure=False)
def func(condition: R.Tensor((), "bool"), x: R.Tensor((), "int32")):
_ = R.assert_op(condition)
@@ -178,10 +174,10 @@ def test_assert_on_argument_passes(exec_mode):
condition = tvm.runtime.tensor(np.array(True))
x = tvm.runtime.tensor(np.array(5).astype("int32"))
- run_cpu(func, condition, x, exec_mode=exec_mode)
+ run_cpu(func, condition, x)
-def test_assert_on_argument_fails(exec_mode):
+def test_assert_on_argument_fails():
@R.function(pure=False)
def func(condition: R.Tensor((), "bool"), x: R.Tensor((), "int32")):
_ = R.assert_op(condition)
@@ -190,10 +186,10 @@ def test_assert_on_argument_fails(exec_mode):
condition = tvm.runtime.tensor(np.array(False))
x = tvm.runtime.tensor(np.array(5).astype("int32"))
with pytest.raises(AssertionError):
- run_cpu(func, condition, x, exec_mode=exec_mode)
+ run_cpu(func, condition, x)
-def test_assert_on_symbolic_var_passes(exec_mode):
+def test_assert_on_symbolic_var_passes():
N = T.dynamic("N")
@R.function(pure=False)
@@ -202,10 +198,10 @@ def test_assert_on_symbolic_var_passes(exec_mode):
return x
x = tvm.runtime.tensor(np.arange(8, dtype="int32"))
- run_cpu(func, x, exec_mode=exec_mode)
+ run_cpu(func, x)
-def test_assert_on_symbolic_var_fails(exec_mode):
+def test_assert_on_symbolic_var_fails():
N = T.dynamic("N")
@R.function(pure=False)
@@ -215,7 +211,7 @@ def test_assert_on_symbolic_var_fails(exec_mode):
x = tvm.runtime.tensor(np.arange(10, dtype="int32"))
with pytest.raises(AssertionError):
- run_cpu(func, x, exec_mode=exec_mode)
+ run_cpu(func, x)
@tvm.script.ir_module
@@ -242,18 +238,17 @@ class ShapeOfTest:
return R.shape_of(x)
-def test_op_shape_of(exec_mode):
- unit_shape = run_cpu(ShapeOfTest, "get_scalar_shape", exec_mode=exec_mode)
+def test_op_shape_of():
+ unit_shape = run_cpu(ShapeOfTest, "get_scalar_shape")
assert unit_shape == tvm_ffi.Shape([])
- const_shape = run_cpu(ShapeOfTest, "get_constant_shape",
exec_mode=exec_mode)
+ const_shape = run_cpu(ShapeOfTest, "get_constant_shape")
assert const_shape == tvm_ffi.Shape([2, 2])
scalar_shape = run_cpu(
ShapeOfTest,
"get_shape",
tvm.runtime.tensor(np.array(1, dtype="int32")),
- exec_mode=exec_mode,
)
assert scalar_shape == tvm_ffi.Shape([])
@@ -261,7 +256,6 @@ def test_op_shape_of(exec_mode):
ShapeOfTest,
"get_shape",
tvm.runtime.tensor(np.zeros((1, 2, 3)).astype("int32")),
- exec_mode=exec_mode,
)
assert tensor_shape == tvm_ffi.Shape([1, 2, 3])
@@ -269,7 +263,6 @@ def test_op_shape_of(exec_mode):
ShapeOfTest,
"get_constrained_shape",
tvm.runtime.tensor(np.zeros((1,)).astype("int32")),
- exec_mode=exec_mode,
)
assert constrained_shape == tvm_ffi.Shape([1])
@@ -289,7 +282,7 @@ class ShapeToTensorTest:
return R.shape_to_tensor(shape)
-def test_op_shape_to_tensor(exec_mode):
+def test_op_shape_to_tensor():
# Check type
isinstance(ShapeToTensorTest["const_shape"].body.ty, tvm.relax.TensorType)
assert ShapeToTensorTest["const_shape"].body.ty.ndim == 1
@@ -297,26 +290,24 @@ def test_op_shape_to_tensor(exec_mode):
assert ShapeToTensorTest["symbolic_shape"].body.ty.ndim == 1
# Check its functionality
- out2d = run_cpu(ShapeToTensorTest, "const_shape", tvm_ffi.Shape([3, 2]),
exec_mode=exec_mode)
+ out2d = run_cpu(ShapeToTensorTest, "const_shape", tvm_ffi.Shape([3, 2]))
assert isinstance(out2d, tvm.runtime.Tensor)
assert np.array_equal(out2d.numpy(), np.array([3, 2]))
- out3d = run_cpu(ShapeToTensorTest, "const_shape", tvm_ffi.Shape([3, 3,
2]), exec_mode=exec_mode)
+ out3d = run_cpu(ShapeToTensorTest, "const_shape", tvm_ffi.Shape([3, 3, 2]))
assert isinstance(out3d, tvm.runtime.Tensor)
assert np.array_equal(out3d.numpy(), np.array([3, 3, 2]))
- out4d = run_cpu(
- ShapeToTensorTest, "const_shape", tvm_ffi.Shape([3, 3, 2, 2]),
exec_mode=exec_mode
- )
+ out4d = run_cpu(ShapeToTensorTest, "const_shape", tvm_ffi.Shape([3, 3, 2,
2]))
assert isinstance(out4d, tvm.runtime.Tensor)
assert np.array_equal(out4d.numpy(), np.array([3, 3, 2, 2]))
- outs = run_cpu(ShapeToTensorTest, "symbolic_shape", tvm_ffi.Shape([3, 2]),
exec_mode=exec_mode)
+ outs = run_cpu(ShapeToTensorTest, "symbolic_shape", tvm_ffi.Shape([3, 2]))
assert isinstance(outs, tvm.runtime.Tensor)
assert np.array_equal(outs.numpy(), np.array([3, 2]))
-def test_op_call_pure_packed(exec_mode):
+def test_op_call_pure_packed():
@tvm.script.ir_module
class CallPureTest:
@R.function
@@ -328,11 +319,11 @@ def test_op_call_pure_packed(exec_mode):
np.random.seed(0) # to avoid flakiness
arr = np.random.rand(3, 4).astype("float32")
- copy_found = run_cpu(CallPureTest, "pure_copy", tvm.runtime.tensor(arr),
exec_mode=exec_mode)
+ copy_found = run_cpu(CallPureTest, "pure_copy", tvm.runtime.tensor(arr))
assert (copy_found.numpy() == arr).all()
-def test_op_call_inplace_packed(exec_mode):
+def test_op_call_inplace_packed():
# in this case we can use the same test as above
@tvm.script.ir_module
class CallInplaceTest:
@@ -374,9 +365,7 @@ def test_op_call_inplace_packed(exec_mode):
arr_b = np.random.rand(3, 4).astype("float32")
sum = arr_a + arr_b
tvm_arr_a = tvm.runtime.tensor(arr_a)
- result = run_cpu(
- CallInplaceAddTest, "inplace_add", tvm_arr_a,
tvm.runtime.tensor(arr_b), exec_mode=exec_mode
- )
+ result = run_cpu(CallInplaceAddTest, "inplace_add", tvm_arr_a,
tvm.runtime.tensor(arr_b))
assert result == tvm_arr_a
assert (result.numpy() == sum).all()
@@ -410,14 +399,14 @@ def test_op_call_inplace_packed(exec_mode):
sum = arr_a + arr_b
tvm_arr_a = tvm.runtime.tensor(arr_a)
tvm_arr_b = tvm.runtime.tensor(arr_b)
- result = run_cpu(CallInplaceTuple, "inplace_tuple", tvm_arr_a, tvm_arr_b,
exec_mode=exec_mode)
+ result = run_cpu(CallInplaceTuple, "inplace_tuple", tvm_arr_a, tvm_arr_b)
assert result[0] == tvm_arr_a
assert (result[0].numpy() == sum).all()
assert result[1] != tvm_arr_a and result[1] != tvm_arr_b
assert (result[1].numpy() == sum).all()
-def test_op_call_py_func(exec_mode):
+def test_op_call_py_func():
"""Test R.call_py_func operator functionality."""
import torch
@@ -478,14 +467,14 @@ def test_op_call_py_func(exec_mode):
x_data = np.array([-1.0, 0.0, 1.0], dtype=np.float32)
x_tvm = tvm.runtime.tensor(x_data)
- result = run_cpu(CallPyFuncTest, "simple_call", x_tvm, exec_mode=exec_mode)
+ result = run_cpu(CallPyFuncTest, "simple_call", x_tvm)
expected = np.maximum(x_data, 0.0)
assert (result.numpy() == expected).all()
y_data = np.array([-0.5, 0.5], dtype=np.float32)
y_tvm = tvm.runtime.tensor(y_data)
- result2 = run_cpu(CallPyFuncTest, "multiple_calls", y_tvm,
exec_mode=exec_mode)
+ result2 = run_cpu(CallPyFuncTest, "multiple_calls", y_tvm)
expected2 = 1.0 / (1.0 + np.exp(-np.maximum(y_data, 0.0)))
assert (result2.numpy() == expected2).all()
@@ -518,7 +507,7 @@ def test_py_func_registry_is_scoped_to_its_module():
tvm.get_global_func("vm.builtin.unregister_py_func")("registry_probe")
-def test_op_to_device(exec_mode):
+def test_op_to_device():
@tvm.script.ir_module
class CallToDevice:
@R.function
@@ -534,11 +523,11 @@ def test_op_to_device(exec_mode):
np.random.seed(0) # to avoid flakiness
arr = np.random.rand(3, 4).astype("float32")
- copy_found = run_cpu(CallToDevice, "to_dev", tvm.runtime.tensor(arr),
exec_mode=exec_mode)
+ copy_found = run_cpu(CallToDevice, "to_dev", tvm.runtime.tensor(arr))
assert (copy_found.numpy() == arr).all()
-def test_op_to_vdevice(exec_mode):
+def test_op_to_vdevice():
@tvm.script.ir_module
class ToVDevice:
I.module_global_infos({"vdevice": [R.vdevice("llvm")]})
@@ -551,11 +540,11 @@ def test_op_to_vdevice(exec_mode):
np.random.seed(0)
arr = np.random.rand(3, 4).astype("float32")
- copy_found = run_cpu(ToVDevice, "to_vdev", tvm.runtime.tensor(arr),
exec_mode=exec_mode)
+ copy_found = run_cpu(ToVDevice, "to_vdev", tvm.runtime.tensor(arr))
assert (copy_found.numpy() == arr).all()
-def test_scalar_tensor_as_branch_condition(exec_mode):
+def test_scalar_tensor_as_branch_condition():
"""The condition of a branch may be a scalar tensor"""
@R.function
@@ -566,14 +555,14 @@ def test_scalar_tensor_as_branch_condition(exec_mode):
out = R.prim_value(10)
return out
- res = run_cpu(func, tvm.runtime.tensor(np.array(True)),
exec_mode=exec_mode)
+ res = run_cpu(func, tvm.runtime.tensor(np.array(True)))
assert res == 5
- res = run_cpu(func, tvm.runtime.tensor(np.array(False)),
exec_mode=exec_mode)
+ res = run_cpu(func, tvm.runtime.tensor(np.array(False)))
assert res == 10
-def test_prim_value_as_branch_condition(exec_mode):
+def test_prim_value_as_branch_condition():
"""The condition may be a Expr"""
@R.function
@@ -584,14 +573,14 @@ def test_prim_value_as_branch_condition(exec_mode):
out = R.prim_value(10)
return out
- res = run_cpu(func, True, exec_mode=exec_mode)
+ res = run_cpu(func, True)
assert res == 5
- res = run_cpu(func, False, exec_mode=exec_mode)
+ res = run_cpu(func, False)
assert res == 10
-def test_computed_prim_value_as_branch_condition(exec_mode):
+def test_computed_prim_value_as_branch_condition():
"""The primitive scalar condition may be computed within the function"""
N = T.dynamic("N")
@@ -604,10 +593,10 @@ def
test_computed_prim_value_as_branch_condition(exec_mode):
out = R.prim_value(10)
return out
- res = run_cpu(func, tvm.runtime.tensor(np.arange(16)), exec_mode=exec_mode)
+ res = run_cpu(func, tvm.runtime.tensor(np.arange(16)))
assert res == 5
- res = run_cpu(func, tvm.runtime.tensor(np.arange(20)), exec_mode=exec_mode)
+ res = run_cpu(func, tvm.runtime.tensor(np.arange(20)))
assert res == 10
diff --git a/tests/python/relax/test_vm_alloc_storage_with_scope.py
b/tests/python/relax/test_vm_alloc_storage_with_scope.py
index 3cce1d93f2..1d7ccb08e6 100644
--- a/tests/python/relax/test_vm_alloc_storage_with_scope.py
+++ b/tests/python/relax/test_vm_alloc_storage_with_scope.py
@@ -63,7 +63,7 @@ def test_alloc_storage_with_scope_global():
mod = Module
target = "llvm"
with tvm.transform.PassContext(opt_level=3):
- lib = tvm.relax.build(mod, target=target, exec_mode="compiled")
+ lib = tvm.relax.build(mod, target=target)
dev = tvm.cpu()
# This is the important line which tests nd allocator
diff --git a/tests/python/relax/test_vm_build.py
b/tests/python/relax/test_vm_build.py
index 37c8980828..602fc27e4b 100644
--- a/tests/python/relax/test_vm_build.py
+++ b/tests/python/relax/test_vm_build.py
@@ -39,15 +39,8 @@ from tvm.script import tirx as T
from tvm.support import cc, popen_pool, utils
from tvm.testing import env
-EXEC_MODE = ["bytecode", "compiled"]
-
[email protected](params=EXEC_MODE)
-def exec_mode(request):
- return request.param
-
-
-def test_vm_compile_simple(exec_mode):
+def test_vm_compile_simple():
@tvm.script.ir_module
class TestVMCompileStage0:
@R.function
@@ -59,7 +52,7 @@ def test_vm_compile_simple(exec_mode):
mod = TestVMCompileStage0
target = tvm.target.Target("llvm", host="llvm")
- ex = relax.build(mod, target, exec_mode=exec_mode)
+ ex = relax.build(mod, target)
inp1 = tvm.runtime.tensor(np.random.rand(3, 4).astype(np.float32))
inp2 = tvm.runtime.tensor(np.random.rand(3, 4).astype(np.float32))
vm = relax.VirtualMachine(ex, tvm.cpu())
@@ -67,7 +60,7 @@ def test_vm_compile_simple(exec_mode):
tvm.testing.assert_allclose(inp2.numpy(), inp1.numpy(), rtol=1e-7,
atol=1e-7)
-def test_vm_compile_without_target_arg(exec_mode):
+def test_vm_compile_without_target_arg():
"""Like test_vm_compile_simple, but with a default target"""
@tvm.script.ir_module
@@ -79,7 +72,7 @@ def test_vm_compile_without_target_arg(exec_mode):
)
return y
- ex = relax.build(mod, exec_mode=exec_mode)
+ ex = relax.build(mod)
inp1 = tvm.runtime.tensor(np.random.rand(3, 4).astype(np.float32))
inp2 = tvm.runtime.tensor(np.random.rand(3, 4).astype(np.float32))
vm = relax.VirtualMachine(ex, tvm.cpu())
@@ -87,7 +80,7 @@ def test_vm_compile_without_target_arg(exec_mode):
tvm.testing.assert_allclose(inp2.numpy(), inp1.numpy(), rtol=1e-7,
atol=1e-7)
-def test_vm_compile_unlowered_operator_error(exec_mode):
+def test_vm_compile_unlowered_operator_error():
@tvm.script.ir_module
class Conv2dTranspose:
@R.function
@@ -108,10 +101,10 @@ def test_vm_compile_unlowered_operator_error(exec_mode):
tvm.error.InternalError,
match=r"(?s)Offending
call:.*(?:R|relax)\.nn\.conv2d_transpose.*data_layout.*NHWC",
):
- relax.build(mod, target="llvm", exec_mode=exec_mode)
+ relax.build(mod, target="llvm")
-def test_match_check(exec_mode):
+def test_match_check():
n = T.dynamic("n")
m = T.dynamic("m")
@@ -123,7 +116,7 @@ def test_match_check(exec_mode):
mod = TestMatchCheck
target = tvm.target.Target("llvm", host="llvm")
- ex = relax.build(mod, target, exec_mode=exec_mode)
+ ex = relax.build(mod, target)
vm = relax.VirtualMachine(ex, tvm.cpu())
x0 = tvm.runtime.tensor(np.zeros((1, 2)).astype("int32"))
y0 = tvm.runtime.tensor(np.zeros((2, 1)).astype("float32"))
@@ -139,7 +132,7 @@ def test_match_check(exec_mode):
vm["foo"](x0, y2)
-def test_vm_compile_stage2(exec_mode):
+def test_vm_compile_stage2():
n = T.dynamic("n")
m = T.dynamic("m")
@@ -152,7 +145,7 @@ def test_vm_compile_stage2(exec_mode):
mod = TestVMCompileStage2
target = tvm.target.Target("llvm", host="llvm")
- ex = relax.build(mod, target, exec_mode=exec_mode)
+ ex = relax.build(mod, target)
vm = relax.VirtualMachine(ex, tvm.cpu())
shape = (32, 16)
@@ -174,7 +167,7 @@ def test_vm_compile_stage2(exec_mode):
vm["foo"]([])
-def test_vm_compile_stage3(exec_mode):
+def test_vm_compile_stage3():
@tvm.script.ir_module
class TestVMCompileStage3:
@R.function
@@ -186,7 +179,7 @@ def test_vm_compile_stage3(exec_mode):
mod = TestVMCompileStage3
target = tvm.target.Target("llvm", host="llvm")
- ex = relax.build(mod, target, exec_mode=exec_mode)
+ ex = relax.build(mod, target)
vm = relax.VirtualMachine(ex, tvm.cpu())
shape = (32, 16)
@@ -195,7 +188,7 @@ def test_vm_compile_stage3(exec_mode):
tvm.testing.assert_allclose(res.numpy(), inp.numpy(), rtol=1e-7, atol=1e-7)
-def test_vm_compile_e2e(exec_mode):
+def test_vm_compile_e2e():
n = T.dynamic("n")
m = T.dynamic("m")
@@ -212,7 +205,7 @@ def test_vm_compile_e2e(exec_mode):
mod = TestVMCompileE2E
target = tvm.target.Target("llvm", host="llvm")
- ex = relax.build(mod, target, exec_mode=exec_mode)
+ ex = relax.build(mod, target)
vm = relax.VirtualMachine(ex, tvm.cpu())
shape = (32, 16)
@@ -221,7 +214,7 @@ def test_vm_compile_e2e(exec_mode):
tvm.testing.assert_allclose(res.numpy(), np.tile(inp.numpy(), (1, 2)),
rtol=1e-7, atol=1e-7)
-def test_vm_compile_e2e_func_param_with_shape(exec_mode):
+def test_vm_compile_e2e_func_param_with_shape():
m_tir_matmul = T.dynamic("m", "int32")
n_tir_matmul = T.dynamic("n", "int32")
k_tir_matmul = T.dynamic("k", "int32")
@@ -257,7 +250,7 @@ def test_vm_compile_e2e_func_param_with_shape(exec_mode):
mod = TestVMCompileE2E2
target = tvm.target.Target("llvm", host="llvm")
- ex = relax.build(mod, target, exec_mode=exec_mode)
+ ex = relax.build(mod, target)
vm = relax.VirtualMachine(ex, tvm.cpu())
data = tvm.runtime.tensor(np.random.rand(32, 16).astype(np.float32))
@@ -267,7 +260,7 @@ def test_vm_compile_e2e_func_param_with_shape(exec_mode):
tvm.testing.assert_allclose(res.numpy(), expected, rtol=1e-6, atol=1e-6)
-def test_call_tir_inplace_e2e_simple(exec_mode):
+def test_call_tir_inplace_e2e_simple():
@tvm.script.ir_module
class TestCallTIRInplaceE2ESimple:
@Ts.prim_func
@@ -305,7 +298,7 @@ def test_call_tir_inplace_e2e_simple(exec_mode):
mod = TestCallTIRInplaceE2ESimple
target = tvm.target.Target("llvm", host="llvm")
- ex = relax.build(mod, target, exec_mode=exec_mode)
+ ex = relax.build(mod, target)
vm = relax.VirtualMachine(ex, tvm.cpu())
x = tvm.runtime.tensor(np.zeros((2, 3)).astype(np.int32))
@@ -325,7 +318,7 @@ def test_call_tir_inplace_e2e_simple(exec_mode):
tvm.testing.assert_allclose(outs[2].numpy(), z.numpy(), rtol=1e-7,
atol=1e-7)
-def test_call_tir_inplace_e2e_rw(exec_mode):
+def test_call_tir_inplace_e2e_rw():
# read and write from the same tensor
@tvm.script.ir_module
class TestCallTIRInplaceE2ERW:
@@ -352,7 +345,7 @@ def test_call_tir_inplace_e2e_rw(exec_mode):
mod = TestCallTIRInplaceE2ERW
target = tvm.target.Target("llvm", host="llvm")
- ex = relax.build(mod, target, exec_mode=exec_mode)
+ ex = relax.build(mod, target)
vm = relax.VirtualMachine(ex, tvm.cpu())
x = tvm.runtime.tensor(np.ones((2, 3)).astype(np.int32))
@@ -366,7 +359,7 @@ def test_call_tir_inplace_e2e_rw(exec_mode):
tvm.testing.assert_allclose(out.numpy(), expected.numpy(), rtol=1e-7,
atol=1e-7)
-def test_vm_emit_te_extern(exec_mode):
+def test_vm_emit_te_extern():
if not tvm.get_global_func("tvm.contrib.cblas.matmul", True):
print("skip because extern function is not available")
return
@@ -382,7 +375,7 @@ def test_vm_emit_te_extern(exec_mode):
mod = bb.get()
target = tvm.target.Target("llvm", host="llvm")
- ex = relax.build(mod, target, exec_mode=exec_mode)
+ ex = relax.build(mod, target)
vm = relax.VirtualMachine(ex, tvm.cpu())
data = tvm.runtime.tensor(np.random.rand(16, 32).astype(np.float32))
@@ -392,7 +385,7 @@ def test_vm_emit_te_extern(exec_mode):
tvm.testing.assert_allclose(res.numpy(), expected, rtol=1e-6, atol=1e-6)
-def test_vm_emit_te_concat(exec_mode):
+def test_vm_emit_te_concat():
# concatenate of two vectors of size (n,) and (m,)
bb = relax.BlockBuilder()
n, m = tirx.Var("n", "int64"), tirx.Var("m", "int64")
@@ -410,7 +403,7 @@ def test_vm_emit_te_concat(exec_mode):
mod = bb.get()
target = tvm.target.Target("llvm", host="llvm")
- ex = relax.build(mod, target, exec_mode=exec_mode)
+ ex = relax.build(mod, target)
vm = relax.VirtualMachine(ex, tvm.cpu())
inp = tvm.runtime.tensor(
@@ -429,7 +422,7 @@ def test_vm_emit_te_concat(exec_mode):
)
-def test_vm_emit_te_dtype_change(exec_mode):
+def test_vm_emit_te_dtype_change():
bb = relax.BlockBuilder()
n = T.dynamic("n", "int64")
x = relax.Var("x", R.Tensor([n], "float32"))
@@ -446,7 +439,7 @@ def test_vm_emit_te_dtype_change(exec_mode):
mod = bb.get()
target = tvm.target.Target("llvm", host="llvm")
- ex = relax.build(mod, target, exec_mode=exec_mode)
+ ex = relax.build(mod, target)
vm = relax.VirtualMachine(ex, tvm.cpu())
inp = tvm.runtime.tensor(
@@ -458,7 +451,7 @@ def test_vm_emit_te_dtype_change(exec_mode):
tvm.testing.assert_allclose(res.numpy(), inp.numpy().astype("int16"))
-def test_vm_emit_te_floor_symbolic_shape(exec_mode):
+def test_vm_emit_te_floor_symbolic_shape():
bb = relax.BlockBuilder()
n = T.dynamic("n", "int64")
x = relax.Var("x", R.Tensor([n], "float32"))
@@ -474,7 +467,7 @@ def test_vm_emit_te_floor_symbolic_shape(exec_mode):
mod = bb.get()
target = tvm.target.Target("llvm", host="llvm")
- ex = relax.build(mod, target, exec_mode=exec_mode)
+ ex = relax.build(mod, target)
vm = relax.VirtualMachine(ex, tvm.cpu())
shape = (9,)
@@ -488,7 +481,7 @@ def test_vm_emit_te_floor_symbolic_shape(exec_mode):
tvm.testing.assert_allclose(res.numpy(), expected_output(), rtol=1e-7,
atol=1e-7)
-def test_vm_emit_te_constant_param_cpu(exec_mode):
+def test_vm_emit_te_constant_param_cpu():
x_np = np.random.rand(2, 2).astype("float32")
c_np = np.random.rand(2, 2).astype("float32")
@@ -502,7 +495,7 @@ def test_vm_emit_te_constant_param_cpu(exec_mode):
bb.emit_func_output(gv)
mod = bb.get()
- exec = relax.build(mod, "llvm", exec_mode=exec_mode)
+ exec = relax.build(mod, "llvm")
dev = tvm.cpu()
vm = relax.VirtualMachine(exec, dev)
@@ -512,7 +505,7 @@ def test_vm_emit_te_constant_param_cpu(exec_mode):
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_gpu(), reason="need gpu")
-def test_vm_emit_te_constant_param_gpu(exec_mode):
+def test_vm_emit_te_constant_param_gpu():
x_np = np.random.rand(2, 2).astype("float32")
c_np = np.random.rand(2, 2).astype("float32")
@@ -530,7 +523,7 @@ def test_vm_emit_te_constant_param_gpu(exec_mode):
loops = sch.get_loops(sch.get_sblock(name="T_add", func_name="add"))
sch.bind(loops[0], "threadIdx.x")
- exec = relax.build(sch.mod, "cuda", exec_mode=exec_mode)
+ exec = relax.build(sch.mod, "cuda")
def run_and_check():
dev = tvm.cuda()
@@ -541,7 +534,7 @@ def test_vm_emit_te_constant_param_gpu(exec_mode):
tvm.testing.run_with_gpu_lock(run_and_check)
-def test_vm_relax_symbolic_shape(exec_mode):
+def test_vm_relax_symbolic_shape():
bb = relax.BlockBuilder()
n = T.dynamic("n", "int64")
x = relax.Var("x", R.Tensor([n], "float32"))
@@ -558,7 +551,7 @@ def test_vm_relax_symbolic_shape(exec_mode):
mod = bb.get()
target = tvm.target.Target("llvm", host="llvm")
- ex = relax.build(mod, target, exec_mode=exec_mode)
+ ex = relax.build(mod, target)
vm = relax.VirtualMachine(ex, tvm.cpu())
shape1 = (5,)
@@ -573,7 +566,7 @@ def test_vm_relax_symbolic_shape(exec_mode):
tvm.testing.assert_allclose(res.numpy(), expected_output(), rtol=1e-7,
atol=1e-7)
-def test_vm_relax_symbolic_shape_tuple(exec_mode):
+def test_vm_relax_symbolic_shape_tuple():
m = T.dynamic("m")
n = T.dynamic("n")
@@ -584,7 +577,7 @@ def test_vm_relax_symbolic_shape_tuple(exec_mode):
return R.shape([2 * m, 3 * n])
target = tvm.target.Target("llvm", host="llvm")
- ex = relax.build(mod, target, exec_mode=exec_mode)
+ ex = relax.build(mod, target)
vm = relax.VirtualMachine(ex, tvm.cpu())
func = vm["main"]
@@ -598,7 +591,7 @@ def test_vm_relax_symbolic_shape_tuple(exec_mode):
func(R.prim_value(2))
-def test_vm_relax_dyn_tir_shape(exec_mode):
+def test_vm_relax_dyn_tir_shape():
# case where TIR variables are unbound in generated PrimFunc
bb = relax.BlockBuilder()
n = T.dynamic("n", "int64")
@@ -617,7 +610,7 @@ def test_vm_relax_dyn_tir_shape(exec_mode):
mod = bb.get()
target = tvm.target.Target("llvm", host="llvm")
- ex = relax.build(mod, target, exec_mode=exec_mode)
+ ex = relax.build(mod, target)
with utils.tempdir() as temp:
ex.export_library(temp.relpath("exec.so"))
@@ -631,7 +624,7 @@ def test_vm_relax_dyn_tir_shape(exec_mode):
tvm.testing.assert_allclose(res.numpy(), inp2.numpy(), rtol=1e-7,
atol=1e-7)
-def test_vm_tuple(exec_mode):
+def test_vm_tuple():
bb = relax.BlockBuilder()
n = T.dynamic("n", "int64")
@@ -645,7 +638,7 @@ def test_vm_tuple(exec_mode):
mod = bb.get()
target = tvm.target.Target("llvm", host="llvm")
- ex = relax.build(mod, target, exec_mode=exec_mode)
+ ex = relax.build(mod, target)
vm = relax.VirtualMachine(ex, tvm.cpu())
shape = (5,)
@@ -658,7 +651,7 @@ def test_vm_tuple(exec_mode):
tvm.testing.assert_allclose(res3.numpy(), inp.numpy(), rtol=1e-7,
atol=1e-7)
-def test_vm_tuplegetitem(exec_mode):
+def test_vm_tuplegetitem():
@tvm.script.ir_module
class TestVMTupleGetItem:
@R.function
@@ -674,7 +667,7 @@ def test_vm_tuplegetitem(exec_mode):
mod = TestVMTupleGetItem
target = tvm.target.Target("llvm", host="llvm")
- ex = relax.build(mod, target, exec_mode=exec_mode)
+ ex = relax.build(mod, target)
vm = relax.VirtualMachine(ex, tvm.cpu())
x_inp = tvm.runtime.tensor(np.random.rand(2, 3).astype("float32"))
y_inp = tvm.runtime.tensor(np.random.rand(2, 3).astype("float32"))
@@ -682,7 +675,7 @@ def test_vm_tuplegetitem(exec_mode):
tvm.testing.assert_allclose(res.numpy(), x_inp.numpy() + y_inp.numpy(),
rtol=1e-7, atol=1e-7)
-def test_lower_memory_alloc_storage_tensor(exec_mode):
+def test_lower_memory_alloc_storage_tensor():
@tvm.script.ir_module
class TestMemoryAllocStorageTensor:
@R.function
@@ -706,14 +699,14 @@ def test_lower_memory_alloc_storage_tensor(exec_mode):
mod = TestMemoryAllocStorageTensor
target = tvm.target.Target("llvm", host="llvm")
- ex = relax.build(mod, target, exec_mode=exec_mode)
+ ex = relax.build(mod, target)
vm = relax.VirtualMachine(ex, tvm.cpu())
x = tvm.runtime.tensor(np.random.rand(2, 3).astype("float32"))
y = vm["main"](x)
tvm.testing.assert_allclose(y.numpy(), x.numpy(), rtol=1e-7, atol=1e-7)
-def test_sub_func_call(exec_mode):
+def test_sub_func_call():
m = T.dynamic("m", "int32")
n = T.dynamic("n", "int32")
k = T.dynamic("k", "int32")
@@ -758,7 +751,7 @@ def test_sub_func_call(exec_mode):
return gv1
target = tvm.target.Target("llvm", host="llvm")
- ex = relax.build(TestVMSubFunction, target, exec_mode=exec_mode)
+ ex = relax.build(TestVMSubFunction, target)
vm = relax.VirtualMachine(ex, tvm.cpu())
x_inp = tvm.runtime.tensor(np.random.rand(32, 32).astype(np.float32))
y_inp = tvm.runtime.tensor(np.random.rand(32, 32).astype(np.float32))
@@ -768,7 +761,7 @@ def test_sub_func_call(exec_mode):
tvm.testing.assert_allclose(res.numpy(), expected, rtol=1e-6, atol=1e-6)
-def test_recursion(exec_mode):
+def test_recursion():
@tvm.script.ir_module
class TestVMRecursion:
@R.function
@@ -789,7 +782,7 @@ def test_recursion(exec_mode):
return res
target = tvm.target.Target("llvm", host="llvm")
- ex = relax.build(TestVMRecursion, target, exec_mode=exec_mode)
+ ex = relax.build(TestVMRecursion, target)
vm = relax.VirtualMachine(ex, tvm.cpu())
inp = np.empty(1).astype("float32")
@@ -802,7 +795,7 @@ def test_recursion(exec_mode):
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_gpu(), reason="need gpu")
-def test_vm_to_device(exec_mode):
+def test_vm_to_device():
@tvm.script.ir_module
class TestToVDevice:
@R.function
@@ -821,7 +814,7 @@ def test_vm_to_device(exec_mode):
mod = TestToVDevice
target = tvm.target.Target("llvm", host="llvm")
- ex = relax.build(mod, target, exec_mode=exec_mode)
+ ex = relax.build(mod, target)
def run_and_check():
vm = relax.VirtualMachine(ex, tvm.cpu())
@@ -839,7 +832,7 @@ def test_vm_to_device(exec_mode):
tvm.testing.run_with_gpu_lock(run_and_check)
-def test_vm_closure(exec_mode):
+def test_vm_closure():
@tvm.script.ir_module
class TestClosure:
@R.function
@@ -858,7 +851,7 @@ def test_vm_closure(exec_mode):
mod = TestClosure
target = tvm.target.Target("llvm", host="llvm")
- ex = relax.build(mod, target, exec_mode=exec_mode)
+ ex = relax.build(mod, target)
vm = relax.VirtualMachine(ex, tvm.cpu())
x_inp = tvm.runtime.tensor(np.random.rand(2, 3).astype("float32"))
y_inp = tvm.runtime.tensor(np.array([[3.1, 4.0, 5.0], [6.0, 7.1, 9.0]],
dtype="float32"))
@@ -866,7 +859,7 @@ def test_vm_closure(exec_mode):
tvm.testing.assert_allclose(res.numpy(), x_inp.numpy() + y_inp.numpy())
-def test_time_evaluator(exec_mode):
+def test_time_evaluator():
@tvm.script.ir_module
class TestTimeEvaluator:
@R.function
@@ -876,7 +869,7 @@ def test_time_evaluator(exec_mode):
)
target = tvm.target.Target("llvm", host="llvm")
- ex = relax.build(TestTimeEvaluator, target, exec_mode=exec_mode)
+ ex = relax.build(TestTimeEvaluator, target)
vm = relax.VirtualMachine(ex, tvm.cpu())
x = tvm.runtime.tensor(np.random.rand(1).astype("float32"))
y = tvm.runtime.tensor(np.random.rand(1).astype("float32"))
@@ -938,7 +931,7 @@ class TestVMSetInput:
return gv0
-def test_multi_systemlib(exec_mode):
+def test_multi_systemlib():
pytest.importorskip("cloudpickle") # needed by popen_pool.PopenWorker
N = T.dynamic("N")
@@ -976,8 +969,8 @@ def test_multi_systemlib(exec_mode):
return gv0
target = tvm.target.Target("llvm", host="llvm")
- libA = relax.build(ModA, target, exec_mode=exec_mode)
- libB = relax.build(ModB, target, exec_mode=exec_mode)
+ libA = relax.build(ModA, target)
+ libB = relax.build(ModB, target)
temp = utils.tempdir()
pathA = temp.relpath("libA.a")
@@ -1064,10 +1057,10 @@ def set_input_attempt_get(vm: relax.VirtualMachine,
device: tvm.runtime.Device)
_ = vm.get_outputs("main")
-def make_vm(mod, exec_mode, temp) -> tuple[relax.VirtualMachine,
tvm.runtime.Device]:
+def make_vm(mod, temp) -> tuple[relax.VirtualMachine, tvm.runtime.Device]:
"""Returns a local VM for the given mod and the device"""
target = tvm.target.Target("llvm", host="llvm")
- exec = relax.build(mod, target, exec_mode=exec_mode)
+ exec = relax.build(mod, target)
libname = temp.relpath("exec.so")
exec.export_library(libname)
exec_loaded = tvm.runtime.load_module(libname)
@@ -1078,14 +1071,13 @@ def make_vm(mod, exec_mode, temp) ->
tuple[relax.VirtualMachine, tvm.runtime.Dev
def run_on_rpc(
mod: tvm.IRModule,
trial_func: Callable[[relax.VirtualMachine, tvm.runtime.Device], None],
- exec_mode: str,
):
"""
Sets up a VM over localhost using the given mod and runs the given trial
function.
The trial function should take a VM and a device
"""
target = tvm.target.Target("llvm", host="llvm")
- exec = relax.build(mod, target, exec_mode=exec_mode)
+ exec = relax.build(mod, target)
temp = utils.tempdir()
path = temp.relpath("vm_library.so")
exec.export_library(path)
@@ -1109,12 +1101,12 @@ def run_on_rpc(
check_remote(rpc.Server("127.0.0.1"))
-def test_set_input(exec_mode):
+def test_set_input():
temp = utils.tempdir()
- set_input_trial(*make_vm(TestVMSetInput, exec_mode, temp))
+ set_input_trial(*make_vm(TestVMSetInput, temp))
-def test_set_input_tuple(exec_mode):
+def test_set_input_tuple():
@tvm.script.ir_module
class MyMod:
@R.function
@@ -1123,7 +1115,7 @@ def test_set_input_tuple(exec_mode):
return y
temp = utils.tempdir()
- vm, device = make_vm(MyMod, exec_mode, temp)
+ vm, device = make_vm(MyMod, temp)
device = tvm.cpu(0)
a = tvm.runtime.empty((32,), "float32", device=device)
b = tvm.runtime.empty((32,), "float32", device=device)
@@ -1140,14 +1132,14 @@ def save_function_kwargs_trial(vm:
relax.VirtualMachine, device: tvm.runtime.Dev
tvm.testing.assert_allclose(res0.numpy(), a.numpy() * b.numpy(),
rtol=1e-7, atol=1e-7)
-def test_save_function_kwargs(exec_mode):
+def test_save_function_kwargs():
temp = utils.tempdir()
- save_function_kwargs_trial(*make_vm(TestVMSetInput, exec_mode, temp))
+ save_function_kwargs_trial(*make_vm(TestVMSetInput, temp))
-def test_save_function_kwargs_rpc(exec_mode):
+def test_save_function_kwargs_rpc():
pytest.importorskip("cloudpickle") # needed by the popen RPC server
- run_on_rpc(TestVMSetInput, save_function_kwargs_trial, exec_mode)
+ run_on_rpc(TestVMSetInput, save_function_kwargs_trial)
def save_function_time_evaluator_trial(
@@ -1160,61 +1152,61 @@ def save_function_time_evaluator_trial(
vm.time_evaluator("saved_main", device)()
-def test_save_function_time_evaluator(exec_mode):
+def test_save_function_time_evaluator():
temp = utils.tempdir()
- save_function_time_evaluator_trial(*make_vm(TestVMSetInput, exec_mode,
temp))
+ save_function_time_evaluator_trial(*make_vm(TestVMSetInput, temp))
-def test_save_function_time_evaluator_rpc(exec_mode):
+def test_save_function_time_evaluator_rpc():
pytest.importorskip("cloudpickle") # needed by the popen RPC server
- run_on_rpc(TestVMSetInput, save_function_time_evaluator_trial, exec_mode)
+ run_on_rpc(TestVMSetInput, save_function_time_evaluator_trial)
# if you set an input, you should not be able to call statelessly
-def test_set_input_stateless_failure(exec_mode):
+def test_set_input_stateless_failure():
temp = utils.tempdir()
- args = make_vm(TestVMSetInput, exec_mode, temp)
+ args = make_vm(TestVMSetInput, temp)
with pytest.raises(RuntimeError):
set_input_attempt_stateless(*args)
-def test_set_input_stateless_failure_rpc(exec_mode):
+def test_set_input_stateless_failure_rpc():
pytest.importorskip("cloudpickle") # needed by the popen RPC server
with pytest.raises(RuntimeError):
- run_on_rpc(TestVMSetInput, set_input_attempt_stateless, exec_mode)
+ run_on_rpc(TestVMSetInput, set_input_attempt_stateless)
-def test_set_input_invoke_failure(exec_mode):
+def test_set_input_invoke_failure():
temp = utils.tempdir()
- args = make_vm(TestVMSetInput, exec_mode, temp)
+ args = make_vm(TestVMSetInput, temp)
with pytest.raises(ValueError):
set_input_attempt_invoke(*args)
-def test_set_input_invoke_failure_rpc(exec_mode):
+def test_set_input_invoke_failure_rpc():
pytest.importorskip("cloudpickle") # needed by the popen RPC server
with pytest.raises(RuntimeError):
- run_on_rpc(TestVMSetInput, set_input_attempt_invoke, exec_mode)
+ run_on_rpc(TestVMSetInput, set_input_attempt_invoke)
-def test_set_input_get_failure(exec_mode):
+def test_set_input_get_failure():
temp = utils.tempdir()
- args = make_vm(TestVMSetInput, exec_mode, temp)
+ args = make_vm(TestVMSetInput, temp)
with pytest.raises(ValueError):
set_input_attempt_get(*args)
-def test_set_input_get_failure_rpc(exec_mode):
+def test_set_input_get_failure_rpc():
pytest.importorskip("cloudpickle") # needed by the popen RPC server
with pytest.raises(RuntimeError):
- run_on_rpc(TestVMSetInput, set_input_attempt_get, exec_mode)
+ run_on_rpc(TestVMSetInput, set_input_attempt_get)
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_gpu(), reason="need gpu")
-def test_relax_module_with_multiple_targets(exec_mode):
+def test_relax_module_with_multiple_targets():
"""Relax functions may contain kernels for multiple targets
In this example, the module contains one function to execute on
diff --git a/tests/python/relax/test_vm_callback_function.py
b/tests/python/relax/test_vm_callback_function.py
index ef16d8fe58..cc3d0d8be5 100644
--- a/tests/python/relax/test_vm_callback_function.py
+++ b/tests/python/relax/test_vm_callback_function.py
@@ -23,12 +23,10 @@ import tvm
import tvm.testing
from tvm.script import relax as R
-exec_mode = tvm.testing.parameter("bytecode", "compiled")
-
pytestmark = pytest.mark.skipif(not tvm.testing.device_enabled("llvm"),
reason="llvm not enabled")
-def test_pass_tensor_to_function(exec_mode):
+def test_pass_tensor_to_function():
target = "llvm"
dev = tvm.cpu()
@@ -44,7 +42,6 @@ def test_pass_tensor_to_function(exec_mode):
ex = tvm.relax.build(
tvm.IRModule.from_expr(relax_func),
target=target,
- exec_mode=exec_mode,
)
vm = tvm.relax.VirtualMachine(ex, dev)
@@ -63,7 +60,7 @@ def test_pass_tensor_to_function(exec_mode):
np.testing.assert_array_equal(np_A * 2, from_callback.numpy())
-def test_generate_tensor_in_function(exec_mode):
+def test_generate_tensor_in_function():
target = "llvm"
dev = tvm.cpu()
@@ -78,7 +75,6 @@ def test_generate_tensor_in_function(exec_mode):
ex = tvm.relax.build(
tvm.IRModule.from_expr(relax_func),
target=target,
- exec_mode=exec_mode,
)
vm = tvm.relax.VirtualMachine(ex, dev)
@@ -92,7 +88,7 @@ def test_generate_tensor_in_function(exec_mode):
np.testing.assert_array_equal(np_A * 2, output.numpy())
-def test_catch_exception_with_full_stack_trace(exec_mode):
+def test_catch_exception_with_full_stack_trace():
target = "llvm"
dev = tvm.cpu()
@@ -106,7 +102,6 @@ def test_catch_exception_with_full_stack_trace(exec_mode):
ex = tvm.relax.build(
tvm.IRModule.from_expr(relax_func),
target=target,
- exec_mode=exec_mode,
)
vm = tvm.relax.VirtualMachine(ex, dev)
diff --git a/tests/python/relax/test_vm_codegen_only.py
b/tests/python/relax/test_vm_codegen_only.py
index 609e980418..52c6b45d9e 100644
--- a/tests/python/relax/test_vm_codegen_only.py
+++ b/tests/python/relax/test_vm_codegen_only.py
@@ -34,17 +34,14 @@ from tvm.script import relax as R
from tvm.script import s_tir as Ts
from tvm.script import tirx as T
-EXEC_MODE = ["bytecode", "compiled"]
-
-def codegen(mod, target, exec_mode="bytecode"):
+def codegen(mod, target):
builder = relax.ExecBuilder()
- tir_mod = relax.vm_build._vmcodegen(builder, mod, exec_mode=exec_mode)
+ tir_mod = relax.vm_build._vmcodegen(builder, mod)
return relax.vm_build._vmlink(builder, target, tir_mod)
[email protected]("exec_mode", EXEC_MODE)
-def test_vm_copy(exec_mode):
+def test_vm_copy():
@tvm.script.ir_module
class TestVMMove:
@R.function(pure=False)
@@ -55,15 +52,14 @@ def test_vm_copy(exec_mode):
mod = TestVMMove
target = tvm.target.Target("llvm", host="llvm")
- ex = codegen(mod, target, exec_mode)
+ ex = codegen(mod, target)
inp = tvm.runtime.tensor(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)
[email protected]("exec_mode", EXEC_MODE)
-def test_vm_to_device(exec_mode):
+def test_vm_to_device():
@tvm.script.ir_module
class TestVMToDevice:
@R.function(pure=False)
@@ -77,7 +73,7 @@ def test_vm_to_device(exec_mode):
mod = TestVMToDevice
target = tvm.target.Target("llvm", host="llvm")
- ex = codegen(mod, target, exec_mode)
+ ex = codegen(mod, target)
inp = tvm.runtime.tensor(np.random.rand(3, 4).astype(np.float32))
vm = relax.VirtualMachine(ex, tvm.cpu())
res = check_saved_func(vm, "foo", inp)
@@ -88,8 +84,7 @@ def test_vm_to_device(exec_mode):
assert res.device.index == 0
[email protected]("exec_mode", EXEC_MODE)
-def test_if_cond_const(exec_mode):
+def test_if_cond_const():
@tvm.script.ir_module
class TestVMIfCondConst:
@R.function
@@ -103,15 +98,14 @@ def test_if_cond_const(exec_mode):
mod = TestVMIfCondConst
target = tvm.target.Target("llvm", host="llvm")
- ex = codegen(mod, target, exec_mode)
+ ex = codegen(mod, target)
vm = relax.VirtualMachine(ex, tvm.cpu())
inp = tvm.runtime.tensor(np.random.rand(3, 4))
res = vm["main"](inp)
tvm.testing.assert_allclose(res.numpy(), inp.numpy())
[email protected]("exec_mode", EXEC_MODE)
-def test_vm_exec_serialize_export_library(exec_mode):
+def test_vm_exec_serialize_export_library():
@tvm.script.ir_module
class TestVMMove:
@R.function(pure=False)
@@ -133,8 +127,7 @@ def test_vm_exec_serialize_export_library(exec_mode):
assert ex.as_text() == loaded_exec["as_text"]()
[email protected]("exec_mode", EXEC_MODE)
-def test_if_cond(exec_mode):
+def test_if_cond():
@tvm.script.ir_module
class TestVMCompileIf:
@R.function(pure=False)
@@ -148,7 +141,7 @@ def test_if_cond(exec_mode):
mod = TestVMCompileIf
target = tvm.target.Target("llvm", host="llvm")
- ex = codegen(mod, target, exec_mode)
+ ex = codegen(mod, target)
vm = relax.VirtualMachine(ex, tvm.cpu())
inp = tvm.runtime.tensor(np.random.rand(3, 4))
res = vm["ife"](tvm.runtime.tensor(1), inp)
@@ -161,8 +154,7 @@ def test_if_cond(exec_mode):
tvm.testing.assert_allclose(res.numpy(), inp.numpy() * inp.numpy(),
rtol=1e-7, atol=1e-7)
[email protected]("exec_mode", EXEC_MODE)
-def test_vm_return_const_tuple(exec_mode):
+def test_vm_return_const_tuple():
@tvm.script.ir_module
class ReturnConstTuple:
@R.function
@@ -174,7 +166,7 @@ def test_vm_return_const_tuple(exec_mode):
mod = ReturnConstTuple
target = tvm.target.Target("llvm", host="llvm")
- ex = codegen(mod, target, exec_mode)
+ ex = codegen(mod, target)
vm = relax.VirtualMachine(ex, tvm.cpu())
inp = tvm.runtime.tensor(np.random.rand(2, 3))
res0, res1, res2 = vm["main"](inp)
@@ -183,8 +175,7 @@ def test_vm_return_const_tuple(exec_mode):
tvm.testing.assert_allclose(res2.numpy(), inp.numpy())
[email protected]("exec_mode", EXEC_MODE)
-def test_vm_const_as_call_arg(exec_mode):
+def test_vm_const_as_call_arg():
@tvm.script.ir_module
class TestVMConstAsCallArg:
@R.function(pure=False)
@@ -206,15 +197,14 @@ def test_vm_const_as_call_arg(exec_mode):
mod = TestVMConstAsCallArg
target = tvm.target.Target("llvm", host="llvm")
- ex = codegen(mod, target, exec_mode)
+ ex = codegen(mod, target)
vm = relax.VirtualMachine(ex, tvm.cpu())
inp = tvm.runtime.tensor(np.random.rand(1, 2))
res = vm["main"](inp)
tvm.testing.assert_allclose(res.numpy(), np.array([4, 6]) + inp.numpy())
[email protected]("exec_mode", EXEC_MODE)
-def test_shape_check_builtin(exec_mode):
+def test_shape_check_builtin():
MS = MatchShapeCode
MK = MakeShapeCode
# slot assignment:
@@ -267,7 +257,7 @@ def test_shape_check_builtin(exec_mode):
mod = TestVMShapeCheck
target = tvm.target.Target("llvm", host="llvm")
- ex = codegen(mod, target, exec_mode)
+ ex = codegen(mod, target)
vm = relax.VirtualMachine(ex, tvm.cpu())
x = tvm.runtime.tensor(np.zeros((1, 2)).astype("float32"))
res = vm["main"](x)
@@ -286,8 +276,7 @@ def test_shape_check_builtin(exec_mode):
vm["main"](tvm.runtime.tensor(np.zeros((1, 2)).astype("int32")))
[email protected]("exec_mode", EXEC_MODE)
-def test_prim_value(exec_mode):
+def test_prim_value():
@tvm.script.ir_module
class TestVMPrimExpr:
@R.function
@@ -298,14 +287,13 @@ def test_prim_value(exec_mode):
mod = TestVMPrimExpr
target = tvm.target.Target("llvm", host="llvm")
- ex = codegen(mod, target, exec_mode)
+ ex = codegen(mod, target)
vm = relax.VirtualMachine(ex, tvm.cpu())
res = vm["main"]()
assert res == 1
[email protected]("exec_mode", EXEC_MODE)
-def test_string_imm(exec_mode):
+def test_string_imm():
@tvm.script.ir_module
class TestVMStringImm:
@R.function
@@ -316,14 +304,13 @@ def test_string_imm(exec_mode):
mod = TestVMStringImm
target = tvm.target.Target("llvm", host="llvm")
- ex = codegen(mod, target, exec_mode)
+ ex = codegen(mod, target)
vm = relax.VirtualMachine(ex, tvm.cpu())
res = vm["main"]()
assert res == "hello"
[email protected]("exec_mode", EXEC_MODE)
-def test_datatype_imm(exec_mode):
+def test_datatype_imm():
@tvm.script.ir_module
class TestDataTypeImm:
@R.function
@@ -334,14 +321,13 @@ def test_datatype_imm(exec_mode):
mod = TestDataTypeImm
target = tvm.target.Target("llvm", host="llvm")
- ex = codegen(mod, target, exec_mode)
+ ex = codegen(mod, target)
vm = relax.VirtualMachine(ex, tvm.cpu())
res = vm["main"]()
assert res == "float32"
[email protected]("exec_mode", EXEC_MODE)
-def test_vm_builtin_reshape(exec_mode):
+def test_vm_builtin_reshape():
@tvm.script.ir_module
class TestVMBuiltinReshape:
@R.function(pure=False)
@@ -354,7 +340,7 @@ def test_vm_builtin_reshape(exec_mode):
mod = TestVMBuiltinReshape
target = tvm.target.Target("llvm", host="llvm")
- ex = codegen(mod, target, exec_mode)
+ ex = codegen(mod, target)
dev = tvm.cpu()
vm = relax.VirtualMachine(ex, dev)
@@ -365,8 +351,7 @@ def test_vm_builtin_reshape(exec_mode):
tvm.testing.assert_allclose(res.numpy(), expected, rtol=1e-7, atol=1e-7)
[email protected]("exec_mode", EXEC_MODE)
-def test_vm_kill_object(exec_mode):
+def test_vm_kill_object():
@I.ir_module
class TestKillObject:
@Ts.prim_func
@@ -418,7 +403,7 @@ def test_vm_kill_object(exec_mode):
mod = TestKillObject
target = tvm.target.Target("llvm", host="llvm")
- ex = codegen(mod, target, exec_mode)
+ ex = codegen(mod, target)
dev = tvm.cpu()
vm = relax.VirtualMachine(ex, dev)
@@ -426,8 +411,7 @@ def test_vm_kill_object(exec_mode):
tvm.testing.assert_allclose(res.numpy(), np.ones((4,), "float32"))
[email protected]("exec_mode", EXEC_MODE)
-def test_preserve_trivial_bindings(exec_mode):
+def test_preserve_trivial_bindings():
@I.ir_module
class mod:
@R.function(pure=False)
@@ -464,7 +448,7 @@ def test_preserve_trivial_bindings(exec_mode):
)
target = tvm.target.Target("llvm", host="llvm")
- ex = codegen(mod, target, exec_mode)
+ ex = codegen(mod, target)
dev = tvm.cpu()
vm = relax.VirtualMachine(ex, dev)
diff --git a/tests/python/relax/test_vm_codegen_tir.py
b/tests/python/relax/test_vm_codegen_tir.py
deleted file mode 100644
index 5c4bf1d219..0000000000
--- a/tests/python/relax/test_vm_codegen_tir.py
+++ /dev/null
@@ -1,225 +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.
-"""Test the TIR codegen path of VM compiled mode.
-
-Restrictions: all shape lowered, explicit allocation.
-"""
-
-import tvm
-import tvm.testing
-from tvm import relax
-from tvm.ir import assert_structural_equal
-from tvm.script import relax as R
-from tvm.script import s_tir as Ts
-from tvm.script import tirx as T
-
-
-def get_tir_mod(mod):
- builder = relax.ExecBuilder()
- return relax.vm_build._vmcodegen(builder, mod, exec_mode="compiled")
-
-
-def test_add():
- @tvm.script.ir_module
- class Before:
- @R.function(pure=False)
- def foo(x: R.Tensor):
- R.func_attr({"global_symbol": "foo"})
- z = R.call_packed("test.vm.add", x, x, ty_args=(R.Tensor))
- return z
-
- @tvm.script.ir_module
- class Expected:
- @Ts.prim_func
- def __vmtir__foo(ctx_ptr: T.handle, r: T.handle, c: T.handle, f:
T.handle):
- T.func_attr({"global_symbol": "__vmtir__foo"})
- T.anylist_setitem_call_packed(
- r,
- T.int32(2),
- "test.vm.add",
- T.anylist_getitem(r, T.int32(0)),
- T.anylist_getitem(r, T.int32(0)),
- )
- T.anylist_setitem_call_packed(
- r, T.int32(1), "vm.builtin.copy", T.anylist_getitem(r,
T.int32(2))
- )
-
- before = Before
- expected = Expected
- after = get_tir_mod(before)
- assert_structural_equal(expected, after)
-
-
-def test_tir_call():
- @tvm.script.ir_module
- class Before:
- @Ts.prim_func
- def shape_func(H: T.Buffer(T.int64(4), "int64")):
- T.func_attr({"global_symbol": "shape_func"})
- # generated compute function
- H[T.int64(0)] = H[T.int64(0)] + T.int64(1)
-
- @R.function(pure=False)
- def foo(x: R.Tensor([4], "int64")):
- R.func_attr({"global_symbol": "foo"})
- _ = Before.shape_func(x)
- return x
-
- @tvm.script.ir_module
- class Expected:
- @Ts.prim_func
- def shape_func(H: T.Buffer(T.int64(4), "int64")):
- T.func_attr({"global_symbol": "shape_func"})
- # generated compute function
- H[T.int64(0)] = H[T.int64(0)] + T.int64(1)
-
- @Ts.prim_func
- def __vmtir__foo(ctx_ptr: T.handle, r: T.handle, c: T.handle, f:
T.handle):
- T.func_attr({"global_symbol": "__vmtir__foo"})
- T.call_cpacked("shape_func", T.anylist_getitem(r, T.int32(0)))
- T.anylist_setitem_call_packed(
- r, T.int32(1), "vm.builtin.copy", T.anylist_getitem(r,
T.int32(0))
- )
-
- before = Before
- expected = Expected
- after = get_tir_mod(before)
- assert_structural_equal(expected, after)
-
-
-def test_if_cond():
- @tvm.script.ir_module
- class Before:
- @R.function(pure=False)
- def ife(cond: R.Tensor((), "bool"), x: R.Tensor) -> R.Tensor:
- R.func_attr({"global_symbol": "ife"})
- if cond:
- w = R.call_packed("test.vm.add", x, x, ty_args=(R.Tensor))
- else:
- w = R.call_packed("test.vm.mul", x, x, ty_args=(R.Tensor))
- return w
-
- @tvm.script.ir_module
- class Expected:
- @Ts.prim_func
- def __vmtir__ife(ctx_ptr: T.handle, r: T.handle, c: T.handle, f:
T.handle):
- T.func_attr({"global_symbol": "__vmtir__ife"})
- if T.Call(
- tvm.ir.Op.get("tirx.tvm_call_packed"),
- ["vm.builtin.read_if_cond", T.anylist_getitem(r, T.int32(0))],
- ret_ty="bool",
- ):
- T.anylist_setitem_call_packed(
- r,
- T.int32(4),
- "test.vm.add",
- T.anylist_getitem(r, T.int32(1)),
- T.anylist_getitem(r, T.int32(1)),
- )
- T.anylist_setitem_call_packed(
- r, T.int32(3), "vm.builtin.copy", T.anylist_getitem(r,
T.int32(4))
- )
- else:
- T.anylist_setitem_call_packed(
- r,
- T.int32(5),
- "test.vm.mul",
- T.anylist_getitem(r, T.int32(1)),
- T.anylist_getitem(r, T.int32(1)),
- )
- T.anylist_setitem_call_packed(
- r, T.int32(3), "vm.builtin.copy", T.anylist_getitem(r,
T.int32(5))
- )
- T.anylist_setitem_call_packed(
- r, T.int32(2), "vm.builtin.copy", T.anylist_getitem(r,
T.int32(3))
- )
-
- before = Before
- expected = Expected
- after = get_tir_mod(before)
- assert_structural_equal(expected, after)
-
-
-def test_const():
- @tvm.script.ir_module
- class Before:
- @R.function
- def main(x: R.Tensor):
- R.func_attr({"global_symbol": "main"})
- y = R.const([1, 2])
- z = (y, R.const([3, 4]), x)
- return z
-
- @tvm.script.ir_module
- class Expected:
- @Ts.prim_func
- def __vmtir__main(ctx_ptr: T.handle, r: T.handle, c: T.handle, f:
T.handle):
- # function attr dict
- T.func_attr({"global_symbol": "__vmtir__main"})
- # body
- T.anylist_setitem_call_packed(
- r,
- T.int32(2),
- "vm.builtin.make_tuple",
- T.anylist_getitem(c, T.int32(0)),
- T.anylist_getitem(c, T.int32(1)),
- T.anylist_getitem(r, T.int32(0)),
- )
- T.anylist_setitem_call_packed(
- r, T.int32(1), "vm.builtin.copy", T.anylist_getitem(r,
T.int32(2))
- )
-
- before = Before
- expected = Expected
- after = get_tir_mod(before)
- assert_structural_equal(expected, after)
-
-
-def test_const_call():
- @tvm.script.ir_module
- class Before:
- @R.function(pure=False)
- def main(x: R.Tensor):
- R.func_attr({"global_symbol": "main"})
- y = R.const([1, 2])
- z = R.call_packed("test.vm.add", x, y, ty_args=(R.Tensor))
- return z
-
- @tvm.script.ir_module
- class Expected:
- @Ts.prim_func
- def __vmtir__main(ctx_ptr: T.handle, r: T.handle, c: T.handle, f:
T.handle):
- # function attr dict
- T.func_attr({"global_symbol": "__vmtir__main"})
- # body
- T.anylist_setitem_call_packed(
- r,
- 2,
- "test.vm.add",
- T.anylist_getitem(r, 0),
- T.anylist_getitem(c, 0),
- )
- T.anylist_setitem_call_packed(r, 1, "vm.builtin.copy",
T.anylist_getitem(r, 2))
-
- before = Before
- expected = Expected
- after = get_tir_mod(before)
- assert_structural_equal(expected, after)
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/relax/test_vm_cuda_graph.py
b/tests/python/relax/test_vm_cuda_graph.py
index 54c7bdd5b9..3a1e665863 100644
--- a/tests/python/relax/test_vm_cuda_graph.py
+++ b/tests/python/relax/test_vm_cuda_graph.py
@@ -92,9 +92,9 @@ class Module:
# fmt: on
-def codegen(mod, target, exec_mode="bytecode"):
+def codegen(mod, target):
builder = relax.ExecBuilder()
- leftover_mod = relax.vm_build._vmcodegen(builder, mod, exec_mode=exec_mode)
+ leftover_mod = relax.vm_build._vmcodegen(builder, mod)
tir_mod = relax.vm_build._filter_tir(leftover_mod)
return relax.vm_build._vmlink(builder, target, tir_mod)
