You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2025/04/28 23:34:23 UTC
(tvm) branch refactor-s2 updated (889ae5444b -> 3d8cc1bfd2)
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a change to branch refactor-s2
in repository https://gitbox.apache.org/repos/asf/tvm.git
discard 889ae5444b Swap global registry to use ffi::Function
new 3d8cc1bfd2 Swap global registry to use ffi::Function
This update added new revisions after undoing existing revisions.
That is to say, some revisions that were in the old version of the
branch are not in the new version. This situation occurs
when a user --force pushes a change and generates a repository
containing something like this:
* -- * -- B -- O -- O -- O (889ae5444b)
\
N -- N -- N refs/heads/refactor-s2 (3d8cc1bfd2)
You should already have received notification emails for all of the O
revisions, and so the following emails describe only the N revisions
from the common base, B.
Any revisions marked "omit" are not gone; other references still
refer to them. Any revisions marked "discard" are gone forever.
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
Summary of changes:
python/tvm/ffi/cython/base.pxi | 4 ++--
python/tvm/ffi/cython/device.pxi | 1 -
python/tvm/ffi/cython/function.pxi | 4 ++--
src/relax/backend/contrib/cutlass/codegen.cc | 2 +-
src/relax/backend/contrib/dnnl/codegen.cc | 2 +-
src/runtime/disco/builtin.cc | 2 +-
tests/python/relax/test_vm_codegen_tir.py | 4 +---
7 files changed, 8 insertions(+), 11 deletions(-)
(tvm) 01/01: Swap global registry to use ffi::Function
Posted by tq...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch refactor-s2
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 3d8cc1bfd261365016ef1716ddee185d3d15defa
Author: tqchen <ti...@gmail.com>
AuthorDate: Mon Apr 28 18:54:04 2025 -0400
Swap global registry to use ffi::Function
---
apps/cpp_rpc/rpc_env.cc | 66 ++---
apps/cpp_rpc/rpc_server.cc | 2 +-
apps/hexagon_launcher/launcher_core.cc | 4 +-
apps/ios_rpc/tvmrpc/RPCServer.mm | 4 +-
apps/ios_rpc/tvmrpc/TVMRuntime.mm | 53 ++--
docs/arch/pass_infra.rst | 6 +-
docs/arch/runtime.rst | 4 +-
include/tvm/relax/tuning_api.h | 12 +-
include/tvm/runtime/registry.h | 307 +--------------------
python/tvm/contrib/cutlass/attention_operation.py | 20 +-
python/tvm/contrib/cutlass/conv2d_operation.py | 5 +-
python/tvm/contrib/cutlass/gemm_operation.py | 10 +-
python/tvm/contrib/cutlass/layer_norm_operation.py | 5 +-
python/tvm/contrib/cutlass/rms_norm_operation.py | 5 +-
python/tvm/ffi/cython/base.pxi | 4 +-
python/tvm/ffi/cython/device.pxi | 1 -
python/tvm/ffi/cython/function.pxi | 4 +-
src/arith/analyzer.cc | 2 +-
src/arith/int_constraints.cc | 2 +-
src/arith/solve_linear_equation.cc | 28 +-
src/arith/solve_linear_inequality.cc | 28 +-
src/contrib/msc/core/codegen/cpp_codegen.h | 19 +-
src/contrib/msc/core/transform/set_expr_layout.cc | 4 +-
src/contrib/msc/framework/tensorrt/codegen.cc | 15 +-
src/ir/diagnostic.cc | 8 +-
src/ir/env_func.cc | 6 +-
src/ir/function.cc | 2 +-
src/ir/global_var_supply.cc | 6 +-
src/ir/module.cc | 9 +-
src/ir/name_supply.cc | 9 +-
src/ir/transform.cc | 35 +--
src/meta_schedule/arg_info.cc | 2 +-
src/meta_schedule/builder/builder.cc | 2 +-
src/meta_schedule/cost_model/cost_model.cc | 7 +-
src/meta_schedule/database/database.cc | 30 +-
.../feature_extractor/feature_extractor.cc | 2 +-
.../measure_callback/measure_callback.cc | 2 +-
.../measure_callback/remove_build_artifact.cc | 5 +-
src/meta_schedule/mutator/mutator.cc | 4 +-
src/meta_schedule/postproc/postproc.cc | 6 +-
src/meta_schedule/profiler.cc | 4 +-
src/meta_schedule/runner/runner.cc | 8 +-
.../schedule_rule/apply_custom_rule.cc | 6 +-
src/meta_schedule/schedule_rule/schedule_rule.cc | 8 +-
.../search_strategy/search_strategy.cc | 12 +-
.../space_generator/space_generator.cc | 20 +-
src/meta_schedule/task_scheduler/task_scheduler.cc | 13 +-
src/meta_schedule/tune_context.cc | 5 +-
src/meta_schedule/utils.h | 12 +-
src/node/object_path.cc | 24 +-
src/node/reflection.cc | 6 +-
src/relax/backend/contrib/clml/codegen.cc | 5 +-
src/relax/backend/contrib/cublas/codegen.cc | 5 +-
src/relax/backend/contrib/cudnn/codegen.cc | 5 +-
src/relax/backend/contrib/cutlass/codegen.cc | 23 +-
src/relax/backend/contrib/dnnl/codegen.cc | 5 +-
src/relax/backend/contrib/hipblas/codegen.cc | 3 +-
src/relax/backend/contrib/nnapi/codegen.cc | 5 +-
src/relax/backend/contrib/tensorrt/codegen.cc | 5 +-
src/relax/backend/contrib/utils.h | 35 +--
src/relax/backend/task_extraction.cc | 6 +-
src/relax/backend/vm/exec_builder.cc | 8 +-
src/relax/ir/block_builder.cc | 29 +-
src/relax/transform/few_shot_tuning.cc | 21 +-
src/relax/transform/fold_constant.cc | 10 +-
src/relax/transform/gradient.cc | 6 +-
src/relax/transform/meta_schedule.cc | 15 +-
src/relax/transform/run_codegen.cc | 9 +-
src/relax/transform/tuning_api/database.cc | 19 +-
src/relax/transform/tuning_api/primitives.cc | 34 ++-
src/runtime/c_runtime_api.cc | 23 +-
src/runtime/contrib/amx/amx_config.cc | 4 +-
src/runtime/contrib/cblas/cblas.cc | 23 +-
src/runtime/contrib/cblas/dnnl_blas.cc | 2 +-
src/runtime/contrib/cblas/mkl.cc | 40 +--
src/runtime/contrib/clml/clml_runtime.cc | 8 +-
src/runtime/contrib/coreml/coreml_runtime.mm | 2 +-
src/runtime/contrib/cublas/cublas.cc | 120 ++++----
src/runtime/contrib/cublas/cublas_json_runtime.cc | 5 +-
src/runtime/contrib/cudnn/conv_backward.cc | 8 +-
src/runtime/contrib/cudnn/conv_forward.cc | 8 +-
src/runtime/contrib/cudnn/cudnn_utils.cc | 4 +-
src/runtime/contrib/cudnn/softmax.cc | 6 +-
src/runtime/contrib/curand/curand.cc | 5 +-
src/runtime/contrib/cutlass/fp16_group_gemm.cu | 5 +-
.../contrib/cutlass/fp8_blockwise_scaled_gemm.cu | 8 +-
src/runtime/contrib/cutlass/fp8_gemm.cu | 6 +-
src/runtime/contrib/cutlass/fp8_group_gemm.cu | 5 +-
src/runtime/contrib/dnnl/dnnl.cc | 2 +-
src/runtime/contrib/edgetpu/edgetpu_runtime.cc | 7 +-
src/runtime/contrib/hipblas/hipblas.cc | 39 +--
.../contrib/hipblas/hipblas_json_runtime.cc | 6 +-
src/runtime/contrib/miopen/conv_forward.cc | 227 +++++++--------
src/runtime/contrib/miopen/miopen_utils.cc | 4 +-
src/runtime/contrib/miopen/softmax.cc | 6 +-
src/runtime/contrib/mps/conv.mm | 108 ++++----
src/runtime/contrib/mps/gemm.mm | 2 +-
src/runtime/contrib/mps/mps_utils.mm | 4 +-
src/runtime/contrib/mrvl/mrvl_runtime.cc | 2 +-
src/runtime/contrib/mrvl/mrvl_sw_runtime_lib.cc | 10 +-
src/runtime/contrib/msc/tensorrt_runtime.cc | 8 +-
src/runtime/contrib/random/random.cc | 116 ++++----
src/runtime/contrib/rocblas/rocblas.cc | 74 ++---
src/runtime/contrib/sort/sort.cc | 129 ++++-----
src/runtime/contrib/tflite/tflite_runtime.cc | 2 +-
src/runtime/contrib/thrust/thrust.cu | 6 +-
src/runtime/cpu_device_api.cc | 2 +-
src/runtime/cuda/cuda_device_api.cc | 4 +-
src/runtime/cuda/l2_cache_flush.cc | 2 +-
src/runtime/disco/bcast_session.cc | 6 +-
src/runtime/disco/builtin.cc | 15 +-
src/runtime/disco/disco_worker.cc | 4 +-
src/runtime/disco/distributed/socket_session.cc | 10 +-
src/runtime/disco/loader.cc | 2 +-
src/runtime/disco/nccl/nccl_context.h | 4 +-
src/runtime/disco/process_session.cc | 2 +-
src/runtime/disco/protocol.h | 9 +-
src/runtime/disco/session.cc | 21 +-
src/runtime/dso_library.cc | 9 +-
src/runtime/hexagon/hexagon_common.cc | 9 +-
src/runtime/hexagon/hexagon_device_api.cc | 143 +++++-----
src/runtime/hexagon/rpc/android/session.cc | 2 +-
src/runtime/hexagon/rpc/hexagon/rpc_server.cc | 10 +-
src/runtime/hexagon/rpc/simulator/rpc_server.cc | 12 +-
src/runtime/hexagon/rpc/simulator/session.cc | 2 +-
src/runtime/library_module.cc | 18 +-
src/runtime/memory/memory_manager.cc | 2 +-
src/runtime/metal/metal_device_api.mm | 2 +-
src/runtime/module.cc | 39 ++-
src/runtime/opencl/opencl_device_api.cc | 61 ++--
src/runtime/profiling.cc | 7 +-
src/runtime/regex.cc | 9 +-
src/runtime/registry.cc | 79 +-----
src/runtime/relax_vm/builtin.cc | 10 +-
src/runtime/relax_vm/cuda/cuda_graph_builtin.cc | 4 +-
src/runtime/relax_vm/kv_state.cc | 42 ++-
src/runtime/relax_vm/lm_support.cc | 2 +-
src/runtime/relax_vm/paged_kv_cache.cc | 18 +-
src/runtime/relax_vm/vm.cc | 8 +-
src/runtime/rocm/rocm_device_api.cc | 4 +-
src/runtime/rpc/rpc_device_api.cc | 2 +-
src/runtime/rpc/rpc_endpoint.cc | 8 +-
src/runtime/rpc/rpc_local_session.cc | 2 +-
src/runtime/rpc/rpc_module.cc | 19 +-
src/runtime/rpc/rpc_pipe_impl.cc | 2 +-
src/runtime/rpc/rpc_server_env.cc | 10 +-
src/runtime/rpc/rpc_socket_impl.cc | 4 +-
src/runtime/system_library.cc | 2 +-
src/runtime/thread_pool.cc | 2 +-
src/runtime/vulkan/vulkan_device_api.cc | 2 +-
src/script/ir_builder/base.cc | 8 +-
src/script/printer/doc.cc | 7 +-
src/script/printer/tir/expr.cc | 2 +-
src/support/errno_handling.h | 4 +-
src/support/ffi_testing.cc | 58 ++--
src/target/codegen.cc | 8 +-
src/target/datatype/registry.cc | 46 +--
src/target/datatype/registry.h | 23 +-
src/target/llvm/codegen_aarch64.cc | 2 +-
src/target/llvm/codegen_amdgpu.cc | 13 +-
src/target/llvm/codegen_arm.cc | 2 +-
src/target/llvm/codegen_cpu.cc | 2 +-
src/target/llvm/codegen_hexagon.cc | 8 +-
src/target/llvm/codegen_llvm.cc | 4 +-
src/target/llvm/codegen_nvptx.cc | 6 +-
src/target/llvm/codegen_x86_64.cc | 2 +-
src/target/llvm/intrin_rule_hexagon.cc | 16 +-
src/target/opt/build_cuda_on.cc | 9 +-
src/target/parsers/cpu.cc | 6 +-
src/target/source/codegen_c_host.cc | 1 -
src/target/source/codegen_cuda.cc | 6 +-
src/target/source/codegen_metal.cc | 3 +-
src/target/source/codegen_opencl.cc | 3 +-
src/target/spirv/spirv_utils.cc | 5 +-
src/target/target.cc | 9 +-
src/target/target_info.cc | 4 +-
src/target/target_kind.cc | 3 +-
src/te/operation/create_primfunc.cc | 15 +-
src/te/tensor.cc | 4 +-
src/tir/analysis/var_use_def_analysis.cc | 19 +-
src/tir/ir/block_dependence_info.cc | 2 +-
src/tir/ir/block_scope.cc | 6 +-
src/tir/ir/buffer.cc | 2 +-
src/tir/ir/expr.cc | 3 +-
src/tir/op/op.cc | 22 +-
src/tir/schedule/schedule.cc | 113 ++++----
src/tir/schedule/state.cc | 4 +-
src/tir/schedule/trace.cc | 16 +-
src/tir/transforms/transform_mma_buffer_layout.cc | 8 +-
src/tir/transforms/unsupported_dtype_legalize.cc | 5 +-
src/topi/broadcast.cc | 4 +-
src/topi/einsum.cc | 2 +-
src/topi/elemwise.cc | 80 +++---
src/topi/nn.cc | 60 ++--
src/topi/reduction.cc | 18 +-
src/topi/transform.cc | 92 +++---
src/topi/utils.cc | 16 +-
src/topi/vision.cc | 2 +-
tests/cpp-runtime/hexagon/run_all_tests.cc | 2 +-
tests/cpp-runtime/hexagon/run_unit_tests.cc | 2 +-
tests/cpp/llvm_codegen_registry_test.cc | 18 +-
tests/cpp/target_test.cc | 4 +-
tests/python/contrib/test_hexagon/README_RPC.md | 4 +-
tests/python/relax/test_vm_codegen_tir.py | 4 +-
web/emcc/tvmjs_support.cc | 20 +-
web/emcc/wasm_runtime.cc | 58 ++--
web/emcc/webgpu_runtime.cc | 14 +-
207 files changed, 1580 insertions(+), 2070 deletions(-)
diff --git a/apps/cpp_rpc/rpc_env.cc b/apps/cpp_rpc/rpc_env.cc
index 21e366076f..146b9cf050 100644
--- a/apps/cpp_rpc/rpc_env.cc
+++ b/apps/cpp_rpc/rpc_env.cc
@@ -126,42 +126,44 @@ RPCEnv::RPCEnv(const std::string& wd) {
mkdir(base_.c_str(), 0777);
}
- TVM_REGISTER_GLOBAL("tvm.rpc.server.workpath").set_body([this](TVMArgs args, TVMRetValue* rv) {
- *rv = this->GetPath(args[0]);
- });
+ ffi::Function::SetGlobal(
+ "tvm.rpc.server.workpath",
+ ffi::Function::FromUnpacked([this](const std::string& path) { return this->GetPath(path); }));
- TVM_REGISTER_GLOBAL("tvm.rpc.server.listdir").set_body([this](TVMArgs args, TVMRetValue* rv) {
- std::string dir = this->GetPath(args[0]);
- std::ostringstream os;
- for (auto d : ListDir(dir)) {
- os << d << ",";
- }
- *rv = os.str();
- });
+ ffi::Function::SetGlobal("tvm.rpc.server.listdir",
+ ffi::Function::FromUnpacked([this](const std::string& path) {
+ std::string dir = this->GetPath(path);
+ std::ostringstream os;
+ for (auto d : ListDir(dir)) {
+ os << d << ",";
+ }
+ return os.str();
+ }));
- TVM_REGISTER_GLOBAL("tvm.rpc.server.load_module").set_body([this](TVMArgs args, TVMRetValue* rv) {
- std::string file_name = this->GetPath(args[0]);
- file_name = BuildSharedLibrary(file_name);
- *rv = Module::LoadFromFile(file_name, "");
- LOG(INFO) << "Load module from " << file_name << " ...";
- });
+ ffi::Function::SetGlobal("tvm.rpc.server.load_module",
+ ffi::Function::FromUnpacked([this](const std::string& path) {
+ std::string file_name = this->GetPath(path);
+ file_name = BuildSharedLibrary(file_name);
+ LOG(INFO) << "Load module from " << file_name << " ...";
+ return Module::LoadFromFile(file_name, "");
+ }));
- TVM_REGISTER_GLOBAL("tvm.rpc.server.download_linked_module")
- .set_body([this](TVMArgs args, TVMRetValue* rv) {
- std::string file_name = this->GetPath(args[0]);
- file_name = BuildSharedLibrary(file_name);
- std::string bin;
+ ffi::Function::SetGlobal("tvm.rpc.server.download_linked_module",
+ ffi::Function::FromUnpacked([this](const std::string& path) {
+ std::string file_name = this->GetPath(path);
+ file_name = BuildSharedLibrary(file_name);
+ std::string bin;
- std::ifstream fs(file_name, std::ios::in | std::ios::binary);
- ICHECK(!fs.fail()) << "Cannot open " << file_name;
- fs.seekg(0, std::ios::end);
- size_t size = static_cast<size_t>(fs.tellg());
- fs.seekg(0, std::ios::beg);
- bin.resize(size);
- fs.read(dmlc::BeginPtr(bin), size);
- *rv = ffi::Bytes(bin);
- LOG(INFO) << "Send linked module " << file_name << " to client";
- });
+ std::ifstream fs(file_name, std::ios::in | std::ios::binary);
+ ICHECK(!fs.fail()) << "Cannot open " << file_name;
+ fs.seekg(0, std::ios::end);
+ size_t size = static_cast<size_t>(fs.tellg());
+ fs.seekg(0, std::ios::beg);
+ bin.resize(size);
+ fs.read(dmlc::BeginPtr(bin), size);
+ LOG(INFO) << "Send linked module " << file_name << " to client";
+ return ffi::Bytes(bin);
+ }));
}
/*!
* \brief GetPath To get the work path from packed function
diff --git a/apps/cpp_rpc/rpc_server.cc b/apps/cpp_rpc/rpc_server.cc
index 217baf133b..00fa489d01 100644
--- a/apps/cpp_rpc/rpc_server.cc
+++ b/apps/cpp_rpc/rpc_server.cc
@@ -398,7 +398,7 @@ void RPCServerCreate(std::string host, int port, int port_end, std::string track
rpc.Start();
}
-TVM_REGISTER_GLOBAL("rpc.ServerCreate").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("rpc.ServerCreate").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
RPCServerCreate(args[0], args[1], args[2], args[3], args[4], args[5], args[6], args[7]);
});
} // namespace runtime
diff --git a/apps/hexagon_launcher/launcher_core.cc b/apps/hexagon_launcher/launcher_core.cc
index 737404f213..267440bcd1 100644
--- a/apps/hexagon_launcher/launcher_core.cc
+++ b/apps/hexagon_launcher/launcher_core.cc
@@ -138,7 +138,7 @@ Model::Model(tvm::runtime::Module executor, tvm::runtime::Module module, std::st
}
const tvm::runtime::PackedFunc get_runtime_func(const std::string& name) {
- if (const tvm::runtime::PackedFunc* pf = tvm::runtime::Registry::Get(name)) {
+ if (auto pf = tvm::ffi::Function::GetGlobal(name)) {
return *pf;
}
return tvm::runtime::PackedFunc();
@@ -151,7 +151,7 @@ const tvm::runtime::PackedFunc get_module_func(tvm::runtime::Module module,
void reset_device_api() {
const tvm::runtime::PackedFunc api = get_runtime_func("device_api.hexagon");
- tvm::runtime::Registry::Register("device_api.cpu", true).set_body(api);
+ tvm::runtime::Registry::Register("device_api.cpu", true).set_body_packed(api);
}
tvm::runtime::Module load_module(const std::string& file_name) {
diff --git a/apps/ios_rpc/tvmrpc/RPCServer.mm b/apps/ios_rpc/tvmrpc/RPCServer.mm
index 5dcf8c1eb9..284a2cfcd9 100644
--- a/apps/ios_rpc/tvmrpc/RPCServer.mm
+++ b/apps/ios_rpc/tvmrpc/RPCServer.mm
@@ -63,8 +63,8 @@ using FEventHandler = PackedFunc;
*/
FEventHandler CreateServerEventHandler(NSOutputStream* outputStream, std::string name,
std::string remote_key) {
- const PackedFunc* event_handler_factory = Registry::Get("rpc.CreateEventDrivenServer");
- ICHECK(event_handler_factory != nullptr)
+ auto event_handler_factory = tvm::ffi::Function::GetGlobal("rpc.CreateEventDrivenServer");
+ ICHECK(event_handler_factory.has_value())
<< "You are using tvm_runtime module built without RPC support. "
<< "Please rebuild it with USE_RPC flag.";
diff --git a/apps/ios_rpc/tvmrpc/TVMRuntime.mm b/apps/ios_rpc/tvmrpc/TVMRuntime.mm
index 4e1d66ed64..f33ae12d4b 100644
--- a/apps/ios_rpc/tvmrpc/TVMRuntime.mm
+++ b/apps/ios_rpc/tvmrpc/TVMRuntime.mm
@@ -51,36 +51,37 @@ void LogMessageImpl(const std::string& file, int lineno, int level, const std::s
} // namespace detail
-TVM_REGISTER_GLOBAL("tvm.rpc.server.workpath").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("tvm.rpc.server.workpath").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
static const std::string base_ = NSTemporaryDirectory().UTF8String;
const std::string path = args[0];
*rv = base_ + "/" + path;
});
-TVM_REGISTER_GLOBAL("tvm.rpc.server.load_module").set_body([](TVMArgs args, TVMRetValue* rv) {
- std::string name = args[0];
- std::string fmt = GetFileFormat(name, "");
- NSString* base;
- if (fmt == "dylib") {
- // only load dylib from frameworks.
- NSBundle* bundle = [NSBundle mainBundle];
- base = [[bundle privateFrameworksPath] stringByAppendingPathComponent:@"tvm"];
-
- if (Registry::Get("runtime.module.loadfile_dylib_custom")) {
- // Custom dso loader is present. Will use it.
- base = NSTemporaryDirectory();
- fmt = "dylib_custom";
- }
- } else {
- // Load other modules in tempdir.
- base = NSTemporaryDirectory();
- }
- NSString* path =
- [base stringByAppendingPathComponent:[NSString stringWithUTF8String:name.c_str()]];
- name = [path UTF8String];
- *rv = Module::LoadFromFile(name, fmt);
- LOG(INFO) << "Load module from " << name << " ...";
-});
+TVM_REGISTER_GLOBAL("tvm.rpc.server.load_module")
+ .set_body_packed([](TVMArgs args, TVMRetValue* rv) {
+ std::string name = args[0];
+ std::string fmt = GetFileFormat(name, "");
+ NSString* base;
+ if (fmt == "dylib") {
+ // only load dylib from frameworks.
+ NSBundle* bundle = [NSBundle mainBundle];
+ base = [[bundle privateFrameworksPath] stringByAppendingPathComponent:@"tvm"];
+
+ if (tvm::ffi::Function::GetGlobal("runtime.module.loadfile_dylib_custom")) {
+ // Custom dso loader is present. Will use it.
+ base = NSTemporaryDirectory();
+ fmt = "dylib_custom";
+ }
+ } else {
+ // Load other modules in tempdir.
+ base = NSTemporaryDirectory();
+ }
+ NSString* path =
+ [base stringByAppendingPathComponent:[NSString stringWithUTF8String:name.c_str()]];
+ name = [path UTF8String];
+ *rv = Module::LoadFromFile(name, fmt);
+ LOG(INFO) << "Load module from " << name << " ...";
+ });
#if defined(USE_CUSTOM_DSO_LOADER) && USE_CUSTOM_DSO_LOADER == 1
@@ -108,7 +109,7 @@ class UnsignedDSOLoader final : public Library {
// Add UnsignedDSOLoader plugin in global registry
TVM_REGISTER_GLOBAL("runtime.module.loadfile_dylib_custom")
- .set_body([](TVMArgs args, TVMRetValue* rv) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* rv) {
auto n = make_object<UnsignedDSOLoader>();
n->Init(args[0]);
*rv = CreateModuleFromLibrary(n);
diff --git a/docs/arch/pass_infra.rst b/docs/arch/pass_infra.rst
index 13c9c37122..85e9f45a5f 100644
--- a/docs/arch/pass_infra.rst
+++ b/docs/arch/pass_infra.rst
@@ -306,9 +306,9 @@ pass is registered with an API endpoint as we will show later.
Pass GetPass(const std::string& pass_name) {
using tvm::runtime::Registry;
std::string fpass_name = "relax.transform." + pass_name;
- const auto* f = Registry::Get(fpass_name);
- ICHECK(f != nullptr) << "Cannot find " << fpass_name
- << "to create the pass " << pass_name;
+ const std::optional<tvm::ffi::Function> f = tvm::ffi::Function::GetGlobal(fpass_name);
+ ICHECK(f.has_value()) << "Cannot find " << fpass_name
+ << "to create the pass " << pass_name;
return (*f)();
}
diff --git a/docs/arch/runtime.rst b/docs/arch/runtime.rst
index dfda00c1d6..9f640654ae 100644
--- a/docs/arch/runtime.rst
+++ b/docs/arch/runtime.rst
@@ -81,7 +81,7 @@ The following example registers PackedFunc in C++ and calls from python.
// register a global packed function in c++
TVM_REGISTER_GLOBAL("myadd")
- .set_body(MyAdd);
+ .set_body_packed(MyAdd);
.. code:: python
@@ -111,7 +111,7 @@ we can pass functions from python (as PackedFunc) to C++.
.. code:: c
TVM_REGISTER_GLOBAL("callhello")
- .set_body([](TVMArgs args, TVMRetValue* rv) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* rv) {
PackedFunc f = args[0];
f("hello world");
});
diff --git a/include/tvm/relax/tuning_api.h b/include/tvm/relax/tuning_api.h
index 637c65201c..1e9a74cbfb 100644
--- a/include/tvm/relax/tuning_api.h
+++ b/include/tvm/relax/tuning_api.h
@@ -67,17 +67,17 @@ class ChoiceNode : public runtime::Object {
/*! \brief Getter for constr_func. */
const runtime::PackedFunc GetConstrFunc() {
- const auto* constr_func = tvm::runtime::Registry::Get(constr_func_key);
- ICHECK(constr_func != nullptr) << "constr_func_key is not registered: " << constr_func_key;
- return *constr_func;
+ const auto constr_func = tvm::ffi::Function::GetGlobal(constr_func_key);
+ ICHECK(constr_func.has_value()) << "constr_func_key is not registered: " << constr_func_key;
+ return *std::move(constr_func);
}
/*! \brief Getter for transform_func. */
const runtime::PackedFunc GetTransformFunc() {
- auto* transform_func = tvm::runtime::Registry::Get(transform_func_key);
- ICHECK(transform_func != nullptr)
+ auto transform_func = tvm::ffi::Function::GetGlobal(transform_func_key);
+ ICHECK(transform_func.has_value())
<< "transform_func_key is not registered: " << transform_func_key;
- return *transform_func;
+ return *std::move(transform_func);
}
/*! \brief Perform constr_func. */
diff --git a/include/tvm/runtime/registry.h b/include/tvm/runtime/registry.h
index 038225216a..e6ee664b5a 100644
--- a/include/tvm/runtime/registry.h
+++ b/include/tvm/runtime/registry.h
@@ -20,25 +20,6 @@
/*!
* \file tvm/runtime/registry.h
* \brief This file defines the TVM global function registry.
- *
- * The registered functions will be made available to front-end
- * as well as backend users.
- *
- * The registry stores type-erased functions.
- * Each registered function is automatically exposed
- * to front-end language(e.g. python).
- *
- * Front-end can also pass callbacks as PackedFunc, or register
- * then into the same global registry in C++.
- * The goal is to mix the front-end language and the TVM back-end.
- *
- * \code
- * // register the function as MyAPIFuncName
- * TVM_REGISTER_GLOBAL(MyAPIFuncName)
- * .set_body([](TVMArgs args, TVMRetValue* rv) {
- * // my code.
- * });
- * \endcode
*/
#ifndef TVM_RUNTIME_REGISTRY_H_
#define TVM_RUNTIME_REGISTRY_H_
@@ -53,50 +34,6 @@
namespace tvm {
namespace runtime {
-/*!
- * \brief Check if signals have been sent to the process and if so
- * invoke the registered signal handler in the frontend environment.
- *
- * When running TVM in another language (Python), the signal handler
- * may not be immediately executed, but instead the signal is marked
- * in the interpreter state (to ensure non-blocking of the signal handler).
- *
- * This function can be explicitly invoked to check the cached signal
- * and run the related processing if a signal is marked.
- *
- * On Linux, when siginterrupt() is set, invoke this function whenever a syscall returns EINTR.
- * When it is not set, invoke it between long-running syscalls when you will not immediately
- * return to the frontend. On Windows, the same rules apply, but due to differences in signal
- * processing, these are likely to only make a difference when used with Ctrl+C and socket calls.
- *
- * Not inserting this function will not cause any correctness
- * issue, but will delay invoking the Python-side signal handler until the function returns to
- * the Python side. This means that the effect of e.g. pressing Ctrl+C or sending signals the
- * process will be delayed until function return. When a C function is blocked on a syscall
- * such as accept(), it needs to be called when EINTR is received.
- * So this function is not needed in most API functions, which can finish quickly in a
- * reasonable, deterministic amount of time.
- *
- * \code
- *
- * int check_signal_every_k_iter = 10;
- *
- * for (int iter = 0; iter < very_large_number; ++iter) {
- * if (iter % check_signal_every_k_iter == 0) {
- * tvm::runtime::EnvCheckSignals();
- * }
- * // do work here
- * }
- *
- * \endcode
- *
- * \note This function is a nop when no PyErr_CheckSignals is registered.
- *
- * \throws This function throws an exception when the frontend signal handler
- * indicate an error happens, otherwise it returns normally.
- */
-TVM_DLL void EnvCheckSignals();
-
/*! \brief A class that wraps a Python object and preserves its ownership.
* This class is used to wrap a PyObject* from the Python API and preserve its ownership.
@@ -142,255 +79,15 @@ class WrappedPythonObject {
void* python_obj_ = nullptr;
};
-/*! \brief Registry for global function */
-class Registry {
- public:
- /*!
- * \brief set the body of the function to be f
- * \param f The body of the function.
- */
- TVM_DLL Registry& set_body(PackedFunc f); // NOLINT(*)
- /*!
- * \brief set the body of the function to be f
- * \param f The body of the function.
- */
- template <typename TCallable,
- typename = typename std::enable_if_t<
- std::is_convertible<TCallable, std::function<void(TVMArgs, TVMRetValue*)>>::value &&
- !std::is_base_of<PackedFunc, TCallable>::value>>
- Registry& set_body(TCallable f) { // NOLINT(*)
- return set_body_packed(f);
- }
- /*!
- * \brief Set body to be to use the new Packed convention.
- *
- * \param f The body of the function.
- */
- template <typename FLambda>
- Registry& set_body_packed(FLambda f) {
- return set_body(ffi::Function::FromPacked(f));
- }
- /*!
- * \brief set the body of the function to the given function.
- * Note that this will ignore default arg values and always require all arguments to be
- * provided.
- *
- * \code
- *
- * int multiply(int x, int y) {
- * return x * y;
- * }
- *
- * TVM_REGISTER_GLOBAL("multiply")
- * .set_body_typed(multiply); // will have type int(int, int)
- *
- * // will have type int(int, int)
- * TVM_REGISTER_GLOBAL("sub")
- * .set_body_typed([](int a, int b) -> int { return a - b; });
- *
- * \endcode
- *
- * \param f The function to forward to.
- * \tparam FLambda The signature of the function.
- */
- template <typename FLambda>
- Registry& set_body_typed(FLambda f) {
- return set_body(ffi::Function::FromUnpacked(f, name_));
- }
- /*!
- * \brief set the body of the function to be the passed method pointer.
- * Note that this will ignore default arg values and always require all arguments to be
- * provided.
- *
- * \code
- *
- * // node subclass:
- * struct Example {
- * int doThing(int x);
- * }
- * TVM_REGISTER_GLOBAL("Example_doThing")
- * .set_body_method(&Example::doThing); // will have type int(Example, int)
- *
- * \endcode
- *
- * \param f the method pointer to forward to.
- * \tparam T the type containing the method (inferred).
- * \tparam R the return type of the function (inferred).
- * \tparam Args the argument types of the function (inferred).
- */
- template <typename T, typename R, typename... Args>
- Registry& set_body_method(R (T::*f)(Args...)) {
- using R_ = typename std::remove_reference<R>::type;
- auto fwrap = [f](T target, Args... params) -> R_ {
- // call method pointer
- return (target.*f)(params...);
- };
- return set_body(ffi::Function::FromUnpacked(fwrap, name_));
- }
-
- /*!
- * \brief set the body of the function to be the passed method pointer.
- * Note that this will ignore default arg values and always require all arguments to be
- * provided.
- *
- * \code
- *
- * // node subclass:
- * struct Example {
- * int doThing(int x);
- * }
- * TVM_REGISTER_GLOBAL("Example_doThing")
- * .set_body_method(&Example::doThing); // will have type int(Example, int)
- *
- * \endcode
- *
- * \param f the method pointer to forward to.
- * \tparam T the type containing the method (inferred).
- * \tparam R the return type of the function (inferred).
- * \tparam Args the argument types of the function (inferred).
- */
- template <typename T, typename R, typename... Args>
- Registry& set_body_method(R (T::*f)(Args...) const) {
- auto fwrap = [f](const T target, Args... params) -> R {
- // call method pointer
- return (target.*f)(params...);
- };
- return set_body(ffi::Function::FromUnpacked(fwrap, name_));
- }
-
- /*!
- * \brief set the body of the function to be the passed method pointer.
- * Used when calling a method on a Node subclass through a ObjectRef subclass.
- * Note that this will ignore default arg values and always require all arguments to be
- * provided.
- *
- * \code
- *
- * // node subclass:
- * struct ExampleNode: BaseNode {
- * int doThing(int x);
- * }
- *
- * // noderef subclass
- * struct Example;
- *
- * TVM_REGISTER_GLOBAL("Example_doThing")
- * .set_body_method<Example>(&ExampleNode::doThing); // will have type int(Example, int)
- *
- * // note that just doing:
- * // .set_body_method(&ExampleNode::doThing);
- * // wouldn't work, because ExampleNode can't be taken from a TVMArgValue.
- *
- * \endcode
- *
- * \param f the method pointer to forward to.
- * \tparam TObjectRef the node reference type to call the method on
- * \tparam TNode the node type containing the method (inferred).
- * \tparam R the return type of the function (inferred).
- * \tparam Args the argument types of the function (inferred).
- */
- template <typename TObjectRef, typename TNode, typename R, typename... Args,
- typename = typename std::enable_if<std::is_base_of<ObjectRef, TObjectRef>::value>::type>
- Registry& set_body_method(R (TNode::*f)(Args...)) {
- auto fwrap = [f](TObjectRef ref, Args... params) {
- TNode* target = ref.operator->();
- // call method pointer
- return (target->*f)(params...);
- };
- return set_body(ffi::Function::FromUnpacked(fwrap, name_));
- }
-
- /*!
- * \brief set the body of the function to be the passed method pointer.
- * Used when calling a method on a Node subclass through a ObjectRef subclass.
- * Note that this will ignore default arg values and always require all arguments to be
- * provided.
- *
- * \code
- *
- * // node subclass:
- * struct ExampleNode: BaseNode {
- * int doThing(int x);
- * }
- *
- * // noderef subclass
- * struct Example;
- *
- * TVM_REGISTER_GLOBAL("Example_doThing")
- * .set_body_method<Example>(&ExampleNode::doThing); // will have type int(Example, int)
- *
- * // note that just doing:
- * // .set_body_method(&ExampleNode::doThing);
- * // wouldn't work, because ExampleNode can't be taken from a TVMArgValue.
- *
- * \endcode
- *
- * \param f the method pointer to forward to.
- * \tparam TObjectRef the node reference type to call the method on
- * \tparam TNode the node type containing the method (inferred).
- * \tparam R the return type of the function (inferred).
- * \tparam Args the argument types of the function (inferred).
- */
- template <typename TObjectRef, typename TNode, typename R, typename... Args,
- typename = typename std::enable_if<std::is_base_of<ObjectRef, TObjectRef>::value>::type>
- Registry& set_body_method(R (TNode::*f)(Args...) const) {
- auto fwrap = [f](TObjectRef ref, Args... params) {
- const TNode* target = ref.operator->();
- // call method pointer
- return (target->*f)(params...);
- };
- return set_body(ffi::Function::FromUnpacked(fwrap, name_));
- }
-
- /*!
- * \brief Register a function with given name
- * \param name The name of the function.
- * \param override Whether allow override existing function.
- * \return Reference to the registry.
- */
- TVM_DLL static Registry& Register(const String& name, bool override = false); // NOLINT(*)
- /*!
- * \brief Erase global function from registry, if exist.
- * \param name The name of the function.
- * \return Whether function exist.
- */
- TVM_DLL static bool Remove(const String& name);
- /*!
- * \brief Get the global function by name.
- * \param name The name of the function.
- * \return pointer to the registered function,
- * nullptr if it does not exist.
- */
- TVM_DLL static const PackedFunc* Get(const String& name); // NOLINT(*)
- /*!
- * \brief Get the names of currently registered global function.
- * \return The names
- */
- TVM_DLL static std::vector<String> ListNames();
-
- // Internal class.
- struct Manager;
-
- protected:
- /*! \brief name of the function */
- String name_;
- /*! \brief internal packed function */
- PackedFunc func_;
- friend struct Manager;
-};
-
-#define TVM_FUNC_REG_VAR_DEF static TVM_ATTRIBUTE_UNUSED ::tvm::runtime::Registry& __mk_##TVM
-
/*!
* \brief Register a function globally.
* \code
* TVM_REGISTER_GLOBAL("MyPrint")
- * .set_body([](TVMArgs args, TVMRetValue* rv) {
+ * .set_body_packed([](TVMArgs args, TVMRetValue* rv) {
* });
* \endcode
*/
-#define TVM_REGISTER_GLOBAL(OpName) \
- TVM_STR_CONCAT(TVM_FUNC_REG_VAR_DEF, __COUNTER__) = ::tvm::runtime::Registry::Register(OpName)
+#define TVM_REGISTER_GLOBAL TVM_FFI_REGISTER_GLOBAL
#define TVM_STRINGIZE_DETAIL(x) #x
#define TVM_STRINGIZE(x) TVM_STRINGIZE_DETAIL(x)
diff --git a/python/tvm/contrib/cutlass/attention_operation.py b/python/tvm/contrib/cutlass/attention_operation.py
index 69298453cb..e89e2f747e 100644
--- a/python/tvm/contrib/cutlass/attention_operation.py
+++ b/python/tvm/contrib/cutlass/attention_operation.py
@@ -147,9 +147,8 @@ def instantiate_attention_template(attrs):
}
CHECK(Attention::check_supported(p));
- auto func = tvm::runtime::Registry::Get("runtime.get_cuda_stream");
- ICHECK(func != nullptr);
- cudaStream_t stream = static_cast<cudaStream_t>((*func)().operator void*());
+ auto func = tvm::ffi::Function::GetGlobalRequired("runtime.get_cuda_stream");
+ cudaStream_t stream = static_cast<cudaStream_t>(func().operator void*());
kernel_fn<<<p.getBlocksGrid(), p.getThreadsGrid(), smem_bytes, stream>>>(p);
@@ -187,9 +186,8 @@ def instantiate_flash_attention_template(attrs):
int v_batch_stride = v_row_stride * ${num_keys};
int o_batch_stride = o_row_stride * ${num_queries};
- auto func = tvm::runtime::Registry::Get("runtime.get_cuda_stream");
- ICHECK(func != nullptr);
- cudaStream_t stream = static_cast<cudaStream_t>((*func)().operator void*());
+ auto func = tvm::ffi::Function::GetGlobalRequired("runtime.get_cuda_stream");
+ cudaStream_t stream = static_cast<cudaStream_t>(func().operator void*());
flash_attn::flash_attention_forward(
static_cast<const cutlass::half_t*>(${query}->data),
@@ -239,9 +237,8 @@ def instantiate_flash_attention_template(attrs):
int v_batch_stride = v_row_stride * ${num_keys};
int o_batch_stride = o_row_stride * ${num_queries};
- auto func = tvm::runtime::Registry::Get("runtime.get_cuda_stream");
- ICHECK(func != nullptr);
- cudaStream_t stream = static_cast<cudaStream_t>((*func)().operator void*());
+ auto func = tvm::ffi::Function::GetGlobalRequired("runtime.get_cuda_stream");
+ cudaStream_t stream = static_cast<cudaStream_t>(func().operator void*());
flash_attn::flash_attention_forward(
static_cast<const cutlass::half_t*>(${qkv}->data),
@@ -297,9 +294,8 @@ def instantiate_flash_attention_var_len_template(attrs):
int v_row_stride = v_head_stride * ${num_kv_heads};
int o_row_stride = o_head_stride * ${num_q_heads};
- auto func = tvm::runtime::Registry::Get("runtime.get_cuda_stream");
- ICHECK(func != nullptr);
- cudaStream_t stream = static_cast<cudaStream_t>((*func)().operator void*());
+ auto func = tvm::ffi::Function::GetGlobalRequired("runtime.get_cuda_stream");
+ cudaStream_t stream = static_cast<cudaStream_t>(func().operator void*());
flash_attn::flash_attention_var_len_forward(
static_cast<const cutlass::half_t*>(${query}->data),
diff --git a/python/tvm/contrib/cutlass/conv2d_operation.py b/python/tvm/contrib/cutlass/conv2d_operation.py
index 77f4449db2..4585933c04 100644
--- a/python/tvm/contrib/cutlass/conv2d_operation.py
+++ b/python/tvm/contrib/cutlass/conv2d_operation.py
@@ -424,9 +424,8 @@ def instantiate_conv2d_template(attrs):
CHECK(status == cutlass::Status::kSuccess);
${split_k_update}
- auto func = tvm::runtime::Registry::Get("runtime.get_cuda_stream");
- ICHECK(func != nullptr);
- cudaStream_t stream = static_cast<cudaStream_t>((*func)().operator void*());
+ auto func = tvm::ffi::Function::GetGlobalRequired("runtime.get_cuda_stream");
+ cudaStream_t stream = static_cast<cudaStream_t>(func().operator void*());
status = conv2d_op(stream);
CHECK(status == cutlass::Status::kSuccess);
diff --git a/python/tvm/contrib/cutlass/gemm_operation.py b/python/tvm/contrib/cutlass/gemm_operation.py
index 2639a0359a..dce354b711 100644
--- a/python/tvm/contrib/cutlass/gemm_operation.py
+++ b/python/tvm/contrib/cutlass/gemm_operation.py
@@ -345,9 +345,8 @@ def instantiate_gemm_template(attrs):
status = gemm_op.initialize(arguments, workspace.get());
CHECK(status == cutlass::Status::kSuccess);
- auto func = tvm::runtime::Registry::Get("runtime.get_cuda_stream");
- ICHECK(func != nullptr);
- cudaStream_t stream = static_cast<cudaStream_t>((*func)().operator void*());
+ auto func = tvm::ffi::Function::GetGlobalRequired("runtime.get_cuda_stream");
+ cudaStream_t stream = static_cast<cudaStream_t>(func().operator void*());
status = gemm_op(stream);
CHECK(status == cutlass::Status::kSuccess);
@@ -429,9 +428,8 @@ def emit_fp16A_intB_matmul(attrs):
int n = ${B_arg}->shape[1] * ${float_per_int};
int k = ${B_arg}->shape[0];
- auto func = tvm::runtime::Registry::Get("runtime.get_cuda_stream");
- ICHECK(func != nullptr);
- cudaStream_t stream = static_cast<cudaStream_t>((*func)().operator void*());
+ auto func = tvm::ffi::Function::GetGlobalRequired("runtime.get_cuda_stream");
+ cudaStream_t stream = static_cast<cudaStream_t>(func().operator void*());
""",
attrs,
)
diff --git a/python/tvm/contrib/cutlass/layer_norm_operation.py b/python/tvm/contrib/cutlass/layer_norm_operation.py
index 760ddec3bf..0afe4cba19 100644
--- a/python/tvm/contrib/cutlass/layer_norm_operation.py
+++ b/python/tvm/contrib/cutlass/layer_norm_operation.py
@@ -39,9 +39,8 @@ def instantiate_layer_norm_template(attrs):
cutlass::TensorRef<data_type, RowMajor> _beta((data_type*)${beta}->data, layout_channels);
cutlass::TensorRef<data_type, RowMajor> _output((data_type*)out0->data, layout_2D);
- auto func = tvm::runtime::Registry::Get("runtime.get_cuda_stream");
- ICHECK(func != nullptr);
- cudaStream_t stream = static_cast<cudaStream_t>((*func)().operator void*());
+ auto func = tvm::ffi::Function::GetGlobalRequired("runtime.get_cuda_stream");
+ cudaStream_t stream = static_cast<cudaStream_t>(func().operator void*());
cutlass::layernorm(size, _output, _input, _gamma, _beta, stream);
"""
diff --git a/python/tvm/contrib/cutlass/rms_norm_operation.py b/python/tvm/contrib/cutlass/rms_norm_operation.py
index 9b10620fd4..ce0843c88f 100644
--- a/python/tvm/contrib/cutlass/rms_norm_operation.py
+++ b/python/tvm/contrib/cutlass/rms_norm_operation.py
@@ -38,9 +38,8 @@ def instantiate_rms_norm_template(attrs):
cutlass::TensorRef<data_type, RowMajor> _weight((data_type*)${weight}->data, layout_channels);
cutlass::TensorRef<data_type, RowMajor> _output((data_type*)out0->data, layout_2D);
- auto func = tvm::runtime::Registry::Get("runtime.get_cuda_stream");
- ICHECK(func != nullptr);
- cudaStream_t stream = static_cast<cudaStream_t>((*func)().operator void*());
+ auto func = tvm::ffi::Function::GetGlobalRequired("runtime.get_cuda_stream");
+ cudaStream_t stream = static_cast<cudaStream_t>(func().operator void*());
cutlass::rmsnorm(size, _output, _input, _weight, stream, ${rms_eps});
"""
diff --git a/python/tvm/ffi/cython/base.pxi b/python/tvm/ffi/cython/base.pxi
index a2f90d1ff5..5e6d53af3d 100644
--- a/python/tvm/ffi/cython/base.pxi
+++ b/python/tvm/ffi/cython/base.pxi
@@ -117,8 +117,8 @@ cdef extern from "tvm/ffi/c_api.h":
const char* traceback
ctypedef int (*TVMFFISafeCallType)(
- void* self, const TVMFFIAny* args, int32_t num_args,
- TVMFFIAny* result) nogil
+ void* ctx, const TVMFFIAny* args, int32_t num_args,
+ TVMFFIAny* result) noexcept
int TVMFFIObjectFree(TVMFFIObjectHandle obj) nogil
int TVMFFIObjectGetTypeIndex(TVMFFIObjectHandle obj) nogil
diff --git a/python/tvm/ffi/cython/device.pxi b/python/tvm/ffi/cython/device.pxi
index 8a165be52a..16b5510ff2 100644
--- a/python/tvm/ffi/cython/device.pxi
+++ b/python/tvm/ffi/cython/device.pxi
@@ -149,4 +149,3 @@ cdef inline object make_ret_device(TVMFFIAny result):
_set_class_device(Device)
-
diff --git a/python/tvm/ffi/cython/function.pxi b/python/tvm/ffi/cython/function.pxi
index 0287e474bf..4a9a26a637 100644
--- a/python/tvm/ffi/cython/function.pxi
+++ b/python/tvm/ffi/cython/function.pxi
@@ -236,7 +236,7 @@ def _get_global_func(name, allow_missing):
# handle callbacks
-cdef void tvm_ffi_callback_deleter(void* fhandle) with gil:
+cdef void tvm_ffi_callback_deleter(void* fhandle) noexcept with gil:
local_pyfunc = <object>(fhandle)
Py_DECREF(local_pyfunc)
@@ -244,7 +244,7 @@ cdef void tvm_ffi_callback_deleter(void* fhandle) with gil:
cdef int tvm_ffi_callback(void* context,
const TVMFFIAny* packed_args,
int32_t num_args,
- TVMFFIAny* result) with gil:
+ TVMFFIAny* result) noexcept with gil:
cdef list pyargs
cdef TVMFFIAny temp_result
local_pyfunc = <object>(context)
diff --git a/src/arith/analyzer.cc b/src/arith/analyzer.cc
index 6d651d5058..cab660096f 100644
--- a/src/arith/analyzer.cc
+++ b/src/arith/analyzer.cc
@@ -268,7 +268,7 @@ PrimExpr Analyzer::Simplify(const PrimExpr& expr, int steps) {
return res;
}
-TVM_REGISTER_GLOBAL("arith.CreateAnalyzer").set_body([](TVMArgs args, TVMRetValue* ret) {
+TVM_REGISTER_GLOBAL("arith.CreateAnalyzer").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
using runtime::PackedFunc;
using runtime::TypedPackedFunc;
auto self = std::make_shared<Analyzer>();
diff --git a/src/arith/int_constraints.cc b/src/arith/int_constraints.cc
index 4382a2ef89..923f6f2241 100644
--- a/src/arith/int_constraints.cc
+++ b/src/arith/int_constraints.cc
@@ -204,7 +204,7 @@ TVM_REGISTER_GLOBAL("arith.IntGroupBounds")
TVM_REGISTER_GLOBAL("arith.IntGroupBounds_from_range").set_body_typed(IntGroupBounds::FromRange);
TVM_REGISTER_GLOBAL("arith.IntGroupBounds_FindBestRange")
- .set_body([](TVMArgs args, TVMRetValue* ret) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
ICHECK(args.size() == 1 || args.size() == 2);
IntGroupBounds bounds = args[0];
if (args.size() == 1) {
diff --git a/src/arith/solve_linear_equation.cc b/src/arith/solve_linear_equation.cc
index 3d18437887..bff1ddea71 100644
--- a/src/arith/solve_linear_equation.cc
+++ b/src/arith/solve_linear_equation.cc
@@ -454,19 +454,21 @@ IntConstraintsTransform SolveLinearEquations(const IntConstraints& system_to_sol
return transform;
}
-TVM_REGISTER_GLOBAL("arith.SolveLinearEquations").set_body([](TVMArgs args, TVMRetValue* ret) {
- if (args.size() == 1) {
- *ret = SolveLinearEquations(args[0]);
- } else if (args.size() == 3) {
- Optional<Array<Var>> opt_vars = args[0];
- Optional<Map<Var, Range>> opt_map = args[1];
- Optional<Array<PrimExpr>> opt_relations = args[2];
- IntConstraints problem(opt_vars.value_or({}), opt_map.value_or({}), opt_relations.value_or({}));
- *ret = SolveLinearEquations(problem);
- } else {
- LOG(FATAL) << "arith.SolveLinearEquations expects 1 or 3 arguments, gets " << args.size();
- }
-});
+TVM_REGISTER_GLOBAL("arith.SolveLinearEquations")
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
+ if (args.size() == 1) {
+ *ret = SolveLinearEquations(args[0]);
+ } else if (args.size() == 3) {
+ Optional<Array<Var>> opt_vars = args[0];
+ Optional<Map<Var, Range>> opt_map = args[1];
+ Optional<Array<PrimExpr>> opt_relations = args[2];
+ IntConstraints problem(opt_vars.value_or({}), opt_map.value_or({}),
+ opt_relations.value_or({}));
+ *ret = SolveLinearEquations(problem);
+ } else {
+ LOG(FATAL) << "arith.SolveLinearEquations expects 1 or 3 arguments, gets " << args.size();
+ }
+ });
} // namespace arith
} // namespace tvm
diff --git a/src/arith/solve_linear_inequality.cc b/src/arith/solve_linear_inequality.cc
index c6887c4541..cd0bb3409d 100644
--- a/src/arith/solve_linear_inequality.cc
+++ b/src/arith/solve_linear_inequality.cc
@@ -536,7 +536,7 @@ IntConstraintsTransform SolveInequalitiesDeskewRange(const IntConstraints& inequ
}
TVM_REGISTER_GLOBAL("arith.SolveInequalitiesAsCondition")
- .set_body([](TVMArgs args, TVMRetValue* ret) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
IntConstraints problem;
PartialSolvedInequalities ret_ineq;
if (args.size() == 1) {
@@ -552,20 +552,22 @@ TVM_REGISTER_GLOBAL("arith.SolveInequalitiesAsCondition")
*ret = AsConditions(problem->variables, ret_ineq.first, ret_ineq.second);
});
-TVM_REGISTER_GLOBAL("arith.SolveInequalitiesToRange").set_body([](TVMArgs args, TVMRetValue* ret) {
- if (args.size() == 1) {
- *ret = SolveInequalitiesToRange(args[0]);
- } else if (args.size() == 3) {
- Optional<Map<Var, Range>> opt_map = args[1];
- IntConstraints problem(args[0], opt_map.value_or({}), args[2]);
- *ret = SolveInequalitiesToRange(problem);
- } else {
- LOG(FATAL) << "arith.SolveInequalitiesToRange expects 1 or 3 arguments, gets " << args.size();
- }
-});
+TVM_REGISTER_GLOBAL("arith.SolveInequalitiesToRange")
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
+ if (args.size() == 1) {
+ *ret = SolveInequalitiesToRange(args[0]);
+ } else if (args.size() == 3) {
+ Optional<Map<Var, Range>> opt_map = args[1];
+ IntConstraints problem(args[0], opt_map.value_or({}), args[2]);
+ *ret = SolveInequalitiesToRange(problem);
+ } else {
+ LOG(FATAL) << "arith.SolveInequalitiesToRange expects 1 or 3 arguments, gets "
+ << args.size();
+ }
+ });
TVM_REGISTER_GLOBAL("arith.SolveInequalitiesDeskewRange")
- .set_body([](TVMArgs args, TVMRetValue* ret) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
if (args.size() == 1) {
*ret = SolveInequalitiesDeskewRange(args[0]);
} else if (args.size() == 3) {
diff --git a/src/contrib/msc/core/codegen/cpp_codegen.h b/src/contrib/msc/core/codegen/cpp_codegen.h
index 81b7d1e871..397bf50bc7 100644
--- a/src/contrib/msc/core/codegen/cpp_codegen.h
+++ b/src/contrib/msc/core/codegen/cpp_codegen.h
@@ -114,19 +114,18 @@ class CppCodeGen : public BaseCodeGen<ConfigType, HelperType> {
this->stack_.comment(this->Comment(node));
// process inputs and weights by tools
if (use_tools) {
- const auto* pf = runtime::Registry::Get("msc_tool.codegen_tensor");
- ICHECK(pf != nullptr) << "Cannot find codegen_tensor func.";
+ const auto pf = tvm::ffi::Function::GetGlobalRequired("msc_tool.codegen_tensor");
for (size_t i = 0; i < node->inputs.size(); i++) {
const auto& input = node->InputAt(i);
- const Array<String>& lines = (*pf)(GetTensorCtx(input), input->name, node->name,
- this->config()->tools_scope, this->config()->tools_tag);
+ const Array<String>& lines = pf(GetTensorCtx(input), input->name, node->name,
+ this->config()->tools_scope, this->config()->tools_tag);
for (const auto& l : lines) {
this->stack_.line(l);
}
}
for (const auto& pair : node->weights) {
- const Array<String>& lines = (*pf)(GetTensorCtx(pair.second), pair.second->name, node->name,
- this->config()->tools_scope, this->config()->tools_tag);
+ const Array<String>& lines = pf(GetTensorCtx(pair.second), pair.second->name, node->name,
+ this->config()->tools_scope, this->config()->tools_tag);
for (const auto& l : lines) {
this->stack_.line(l);
}
@@ -137,15 +136,13 @@ class CppCodeGen : public BaseCodeGen<ConfigType, HelperType> {
}
// process graph outputs by tools
if (use_tools) {
- const auto* pf = runtime::Registry::Get("msc_tool.codegen_tensor");
- ICHECK(pf != nullptr) << "Cannot find codegen_tensor func.";
+ const auto pf = tvm::ffi::Function::GetGlobalRequired("msc_tool.codegen_tensor");
for (size_t i = 0; i < node->outputs.size(); i++) {
int index = static_cast<int>(i);
if (graph_outputs_.count(node->OutputAt(index))) {
const auto& output = node->OutputAt(index);
- const Array<String>& lines =
- (*pf)(GetTensorCtx(output), output->name, node->name, this->config()->tools_scope,
- this->config()->tools_tag);
+ const Array<String>& lines = pf(GetTensorCtx(output), output->name, node->name,
+ this->config()->tools_scope, this->config()->tools_tag);
for (const auto& l : lines) {
this->stack_.line(l);
}
diff --git a/src/contrib/msc/core/transform/set_expr_layout.cc b/src/contrib/msc/core/transform/set_expr_layout.cc
index 0b274a2ab6..22da70c7f0 100644
--- a/src/contrib/msc/core/transform/set_expr_layout.cc
+++ b/src/contrib/msc/core/transform/set_expr_layout.cc
@@ -514,8 +514,8 @@ InferLayoutOutput ForwardInferLayoutPlugin(const Call& call,
return InferLayoutOutput();
}
const auto& name = Downcast<ExternFunc>(call->args[0])->global_symbol;
- const auto* pf = runtime::Registry::Get("msc.plugin.op.InferLayout" + name);
- if (pf == nullptr) {
+ const auto pf = tvm::ffi::Function::GetGlobal("msc.plugin.op.InferLayout" + name);
+ if (!pf.has_value()) {
return InferLayoutOutput();
}
const auto& args = Downcast<Tuple>(call->args[1]);
diff --git a/src/contrib/msc/framework/tensorrt/codegen.cc b/src/contrib/msc/framework/tensorrt/codegen.cc
index f3def93f52..32eabb771c 100644
--- a/src/contrib/msc/framework/tensorrt/codegen.cc
+++ b/src/contrib/msc/framework/tensorrt/codegen.cc
@@ -119,9 +119,8 @@ void TensorRTCodeGen::CodeGenClassDefine() {
stack_.func_arg("logger", "TRTLogger&").func_start();
// save codegen before build
if (config()->use_tools) {
- const auto* pf = runtime::Registry::Get("msc_tool.codegen_step");
- ICHECK(pf != nullptr) << "Cannot find msc_tool.codegen_step func.";
- before_build_codes_ = (*pf)(GetStepCtx(), "before_build", graph()->name, config()->tools_tag);
+ const auto pf = tvm::ffi::Function::GetGlobalRequired("msc_tool.codegen_step");
+ before_build_codes_ = pf(GetStepCtx(), "before_build", graph()->name, config()->tools_tag);
}
if (graph()->weight_holders.size() > 0) {
stack_.func_call("TRTUtils::LoadWeights", "mWeights")
@@ -205,9 +204,8 @@ void TensorRTCodeGen::CodeGenClassDefine() {
}
// save codegen after build
if (config()->use_tools) {
- const auto* pf = runtime::Registry::Get("msc_tool.codegen_step");
- ICHECK(pf != nullptr) << "Cannot find msc_tool.codegen_step func.";
- after_build_codes_ = (*pf)(GetStepCtx(), "after_build", graph()->name, config()->tools_tag);
+ const auto pf = tvm::ffi::Function::GetGlobalRequired("msc_tool.codegen_step");
+ after_build_codes_ = pf(GetStepCtx(), "after_build", graph()->name, config()->tools_tag);
}
// end define build method
stack_.func_end("true");
@@ -610,10 +608,9 @@ Array<runtime::Module> MSCTensorRTCompiler(Array<Function> functions,
MSCJSONSerializer serializer(constant_names, options);
serializer.serialize(func);
std::string graph_json = serializer.GetJSON();
- const auto* pf = runtime::Registry::Get("runtime.msc_tensorrt_runtime_create");
- ICHECK(pf != nullptr) << "Cannot find TensorRT runtime module create function.";
+ const auto pf = tvm::ffi::Function::GetGlobalRequired("runtime.msc_tensorrt_runtime_create");
VLOG(1) << "Creating msc_tensorrt runtime::Module for '" << func_name << "'";
- compiled_functions.push_back((*pf)(func_name, graph_json, serializer.GetConstantNames()));
+ compiled_functions.push_back(pf(func_name, graph_json, serializer.GetConstantNames()));
}
return compiled_functions;
}
diff --git a/src/ir/diagnostic.cc b/src/ir/diagnostic.cc
index 2c9d7e8e93..0e58a6847c 100644
--- a/src/ir/diagnostic.cc
+++ b/src/ir/diagnostic.cc
@@ -176,13 +176,13 @@ static const char* DEFAULT_RENDERER = "diagnostics.DefaultRenderer";
static const char* OVERRIDE_RENDERER = "diagnostics.OverrideRenderer";
DiagnosticRenderer GetRenderer() {
- auto override_pf = tvm::runtime::Registry::Get(OVERRIDE_RENDERER);
+ auto override_pf = tvm::ffi::Function::GetGlobal(OVERRIDE_RENDERER);
tvm::runtime::TypedPackedFunc<ObjectRef()> pf;
if (override_pf) {
pf = tvm::runtime::TypedPackedFunc<ObjectRef()>(*override_pf);
} else {
- auto default_pf = tvm::runtime::Registry::Get(DEFAULT_RENDERER);
- ICHECK(default_pf != nullptr)
+ auto default_pf = tvm::ffi::Function::GetGlobal(DEFAULT_RENDERER);
+ ICHECK(default_pf.has_value())
<< "Can not find registered function for " << DEFAULT_RENDERER << "." << std::endl
<< "Either this is an internal error or the default function was overloaded incorrectly.";
pf = tvm::runtime::TypedPackedFunc<ObjectRef()>(*default_pf);
@@ -316,7 +316,7 @@ TVM_REGISTER_GLOBAL(DEFAULT_RENDERER).set_body_typed([]() { return TerminalRende
TVM_REGISTER_GLOBAL("diagnostics.GetRenderer").set_body_typed([]() { return GetRenderer(); });
TVM_REGISTER_GLOBAL("diagnostics.ClearRenderer").set_body_typed([]() {
- tvm::runtime::Registry::Remove(OVERRIDE_RENDERER);
+ tvm::ffi::Function::RemoveGlobal(OVERRIDE_RENDERER);
});
} // namespace tvm
diff --git a/src/ir/env_func.cc b/src/ir/env_func.cc
index f272281d5b..6cd245bc99 100644
--- a/src/ir/env_func.cc
+++ b/src/ir/env_func.cc
@@ -37,8 +37,8 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
});
ObjectPtr<Object> CreateEnvNode(const std::string& name) {
- auto* f = runtime::Registry::Get(name);
- ICHECK(f != nullptr) << "Cannot find global function \'" << name << '\'';
+ auto f = tvm::ffi::Function::GetGlobal(name);
+ ICHECK(f.has_value()) << "Cannot find global function \'" << name << '\'';
ObjectPtr<EnvFuncNode> n = make_object<EnvFuncNode>();
n->func = *f;
n->name = name;
@@ -49,7 +49,7 @@ EnvFunc EnvFunc::Get(const String& name) { return EnvFunc(CreateEnvNode(name));
TVM_REGISTER_GLOBAL("ir.EnvFuncGet").set_body_typed(EnvFunc::Get);
-TVM_REGISTER_GLOBAL("ir.EnvFuncCall").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("ir.EnvFuncCall").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
EnvFunc env = args[0];
ICHECK_GE(args.size(), 1);
env->func.CallPacked(args.Slice(1), rv);
diff --git a/src/ir/function.cc b/src/ir/function.cc
index 16b28e515d..bb5a01551e 100644
--- a/src/ir/function.cc
+++ b/src/ir/function.cc
@@ -52,7 +52,7 @@ TVM_REGISTER_GLOBAL("ir.BaseFuncWithAttrs")
if (func->IsInstance<tir::PrimFuncNode>()) {
return WithAttrs(Downcast<tir::PrimFunc>(std::move(func)), attr_map);
}
- if (const auto* f = runtime::Registry::Get("relax.FuncWithAttrs")) {
+ if (const auto f = tvm::ffi::Function::GetGlobal("relax.FuncWithAttrs")) {
if (Optional<BaseFunc> ret = (*f)(func, attr_map)) {
return ret.value();
}
diff --git a/src/ir/global_var_supply.cc b/src/ir/global_var_supply.cc
index 571a7f304c..3d3b891991 100644
--- a/src/ir/global_var_supply.cc
+++ b/src/ir/global_var_supply.cc
@@ -104,12 +104,12 @@ TVM_REGISTER_GLOBAL("ir.GlobalVarSupply_IRModules").set_body_typed([](const Arra
});
TVM_REGISTER_GLOBAL("ir.GlobalVarSupply_FreshGlobal")
- .set_body_method<GlobalVarSupply>(&GlobalVarSupplyNode::FreshGlobal);
+ .set_body_method(&GlobalVarSupplyNode::FreshGlobal);
TVM_REGISTER_GLOBAL("ir.GlobalVarSupply_UniqueGlobalFor")
- .set_body_method<GlobalVarSupply>(&GlobalVarSupplyNode::UniqueGlobalFor);
+ .set_body_method(&GlobalVarSupplyNode::UniqueGlobalFor);
TVM_REGISTER_GLOBAL("ir.GlobalVarSupply_ReserveGlobalVar")
- .set_body_method<GlobalVarSupply>(&GlobalVarSupplyNode::ReserveGlobalVar);
+ .set_body_method(&GlobalVarSupplyNode::ReserveGlobalVar);
} // namespace tvm
diff --git a/src/ir/module.cc b/src/ir/module.cc
index 336a90fa98..02e353c8a6 100644
--- a/src/ir/module.cc
+++ b/src/ir/module.cc
@@ -301,14 +301,11 @@ TVM_REGISTER_GLOBAL("ir.Module_Contains")
}
});
-TVM_REGISTER_GLOBAL("ir.Module_GetGlobalVar")
- .set_body_method<IRModule>(&IRModuleNode::GetGlobalVar);
+TVM_REGISTER_GLOBAL("ir.Module_GetGlobalVar").set_body_method(&IRModuleNode::GetGlobalVar);
-TVM_REGISTER_GLOBAL("ir.Module_GetGlobalVars")
- .set_body_method<IRModule>(&IRModuleNode::GetGlobalVars);
+TVM_REGISTER_GLOBAL("ir.Module_GetGlobalVars").set_body_method(&IRModuleNode::GetGlobalVars);
-TVM_REGISTER_GLOBAL("ir.Module_ContainGlobalVar")
- .set_body_method<IRModule>(&IRModuleNode::ContainGlobalVar);
+TVM_REGISTER_GLOBAL("ir.Module_ContainGlobalVar").set_body_method(&IRModuleNode::ContainGlobalVar);
TVM_REGISTER_GLOBAL("ir.Module_Lookup").set_body_typed([](IRModule mod, GlobalVar var) {
return mod->Lookup(var);
diff --git a/src/ir/name_supply.cc b/src/ir/name_supply.cc
index 65cb134f11..087fc82a50 100644
--- a/src/ir/name_supply.cc
+++ b/src/ir/name_supply.cc
@@ -96,13 +96,10 @@ TVM_REGISTER_GLOBAL("ir.NameSupply").set_body_typed([](String prefix) {
return NameSupply(prefix);
});
-TVM_REGISTER_GLOBAL("ir.NameSupply_FreshName")
- .set_body_method<NameSupply>(&NameSupplyNode::FreshName);
+TVM_REGISTER_GLOBAL("ir.NameSupply_FreshName").set_body_method(&NameSupplyNode::FreshName);
-TVM_REGISTER_GLOBAL("ir.NameSupply_ReserveName")
- .set_body_method<NameSupply>(&NameSupplyNode::ReserveName);
+TVM_REGISTER_GLOBAL("ir.NameSupply_ReserveName").set_body_method(&NameSupplyNode::ReserveName);
-TVM_REGISTER_GLOBAL("ir.NameSupply_ContainsName")
- .set_body_method<NameSupply>(&NameSupplyNode::ContainsName);
+TVM_REGISTER_GLOBAL("ir.NameSupply_ContainsName").set_body_method(&NameSupplyNode::ContainsName);
} // namespace tvm
diff --git a/src/ir/transform.cc b/src/ir/transform.cc
index 9f56ef883b..a7383bf2e2 100644
--- a/src/ir/transform.cc
+++ b/src/ir/transform.cc
@@ -459,13 +459,13 @@ void SequentialNode::ResolveDependency(const IRModule& mod) {
}
Pass GetPass(const String& pass_name) {
- using tvm::runtime::Registry;
- const runtime::PackedFunc* f = nullptr;
+ std::optional<tvm::ffi::Function> f;
if (pass_name.operator std::string().find("transform.") != std::string::npos) {
- f = Registry::Get(pass_name);
- } else if ((f = Registry::Get("transform." + pass_name))) {
+ f = tvm::ffi::Function::GetGlobal(pass_name);
+ } else {
+ f = tvm::ffi::Function::GetGlobal("transform." + pass_name);
}
- ICHECK(f != nullptr) << "Cannot use " << pass_name << " to create the pass";
+ ICHECK(f.has_value()) << "Cannot use " << pass_name << " to create the pass";
return (*f)();
}
@@ -538,7 +538,7 @@ TVM_REGISTER_GLOBAL("transform.PassInfo")
return PassInfo(opt_level, name, required, traceable);
});
-TVM_REGISTER_GLOBAL("transform.Info").set_body([](TVMArgs args, TVMRetValue* ret) {
+TVM_REGISTER_GLOBAL("transform.Info").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
Pass pass = args[0];
*ret = pass->Info();
});
@@ -586,7 +586,7 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
TVM_REGISTER_NODE_TYPE(SequentialNode);
-TVM_REGISTER_GLOBAL("transform.Sequential").set_body([](TVMArgs args, TVMRetValue* ret) {
+TVM_REGISTER_GLOBAL("transform.Sequential").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
tvm::Array<Pass> passes = args[0];
int opt_level = args[1];
std::string name = args[2];
@@ -658,21 +658,16 @@ class PassContext::Internal {
static void ExitScope(PassContext pass_ctx) { pass_ctx.ExitWithScope(); }
};
-TVM_REGISTER_GLOBAL("transform.GetTraceStack")
- .set_body_method<PassContext>(&PassContextNode::GetTraceStack);
-TVM_REGISTER_GLOBAL("transform.PushTrace")
- .set_body_method<PassContext>(&PassContextNode::PushTrace);
-TVM_REGISTER_GLOBAL("transform.PopTrace").set_body_method<PassContext>(&PassContextNode::PopTrace);
+TVM_REGISTER_GLOBAL("transform.GetTraceStack").set_body_method(&PassContextNode::GetTraceStack);
+TVM_REGISTER_GLOBAL("transform.PushTrace").set_body_method(&PassContextNode::PushTrace);
+TVM_REGISTER_GLOBAL("transform.PopTrace").set_body_method(&PassContextNode::PopTrace);
TVM_REGISTER_GLOBAL("transform.GetTraceStackSize")
- .set_body_method<PassContext>(&PassContextNode::GetTraceStackSize);
-TVM_REGISTER_GLOBAL("transform.GetCurrentTrace")
- .set_body_method<PassContext>(&PassContextNode::GetCurrentTrace);
-TVM_REGISTER_GLOBAL("transform.SetNumEvals")
- .set_body_method<PassContext>(&PassContextNode::SetNumEvals);
-TVM_REGISTER_GLOBAL("transform.IncNumEvals")
- .set_body_method<PassContext>(&PassContextNode::IncNumEvals);
+ .set_body_method(&PassContextNode::GetTraceStackSize);
+TVM_REGISTER_GLOBAL("transform.GetCurrentTrace").set_body_method(&PassContextNode::GetCurrentTrace);
+TVM_REGISTER_GLOBAL("transform.SetNumEvals").set_body_method(&PassContextNode::SetNumEvals);
+TVM_REGISTER_GLOBAL("transform.IncNumEvals").set_body_method(&PassContextNode::IncNumEvals);
TVM_REGISTER_GLOBAL("transform.GetTuningAPIDatabase")
- .set_body_method<PassContext>(&PassContextNode::GetTuningAPIDatabase);
+ .set_body_method(&PassContextNode::GetTuningAPIDatabase);
TVM_REGISTER_GLOBAL("transform.GetCurrentPassContext").set_body_typed(PassContext::Current);
diff --git a/src/meta_schedule/arg_info.cc b/src/meta_schedule/arg_info.cc
index 40c8005d5b..f3b8345141 100644
--- a/src/meta_schedule/arg_info.cc
+++ b/src/meta_schedule/arg_info.cc
@@ -162,7 +162,7 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
TVM_REGISTER_OBJECT_TYPE(ArgInfoNode);
TVM_REGISTER_NODE_TYPE(TensorInfoNode);
-TVM_REGISTER_GLOBAL("meta_schedule.ArgInfoAsJSON").set_body_method<ArgInfo>(&ArgInfoNode::AsJSON);
+TVM_REGISTER_GLOBAL("meta_schedule.ArgInfoAsJSON").set_body_method(&ArgInfoNode::AsJSON);
TVM_REGISTER_GLOBAL("meta_schedule.ArgInfoFromPrimFunc").set_body_typed(ArgInfo::FromPrimFunc);
TVM_REGISTER_GLOBAL("meta_schedule.ArgInfoFromEntryFunc").set_body_typed(ArgInfo::FromEntryFunc);
TVM_REGISTER_GLOBAL("meta_schedule.ArgInfoFromJSON").set_body_typed(ArgInfo::FromJSON);
diff --git a/src/meta_schedule/builder/builder.cc b/src/meta_schedule/builder/builder.cc
index 25adfbfdfa..9d725e91e2 100644
--- a/src/meta_schedule/builder/builder.cc
+++ b/src/meta_schedule/builder/builder.cc
@@ -64,7 +64,7 @@ TVM_REGISTER_GLOBAL("meta_schedule.BuilderResult")
return BuilderResult(artifact_path, error_msg);
});
-TVM_REGISTER_GLOBAL("meta_schedule.BuilderBuild").set_body_method<Builder>(&BuilderNode::Build);
+TVM_REGISTER_GLOBAL("meta_schedule.BuilderBuild").set_body_method(&BuilderNode::Build);
TVM_REGISTER_GLOBAL("meta_schedule.BuilderPyBuilder").set_body_typed(Builder::PyBuilder);
diff --git a/src/meta_schedule/cost_model/cost_model.cc b/src/meta_schedule/cost_model/cost_model.cc
index aabab5d83a..1d28eb19d7 100644
--- a/src/meta_schedule/cost_model/cost_model.cc
+++ b/src/meta_schedule/cost_model/cost_model.cc
@@ -71,10 +71,9 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
TVM_REGISTER_OBJECT_TYPE(CostModelNode);
TVM_REGISTER_NODE_TYPE(PyCostModelNode);
-TVM_REGISTER_GLOBAL("meta_schedule.CostModelLoad").set_body_method<CostModel>(&CostModelNode::Load);
-TVM_REGISTER_GLOBAL("meta_schedule.CostModelSave").set_body_method<CostModel>(&CostModelNode::Save);
-TVM_REGISTER_GLOBAL("meta_schedule.CostModelUpdate")
- .set_body_method<CostModel>(&CostModelNode::Update);
+TVM_REGISTER_GLOBAL("meta_schedule.CostModelLoad").set_body_method(&CostModelNode::Load);
+TVM_REGISTER_GLOBAL("meta_schedule.CostModelSave").set_body_method(&CostModelNode::Save);
+TVM_REGISTER_GLOBAL("meta_schedule.CostModelUpdate").set_body_method(&CostModelNode::Update);
TVM_REGISTER_GLOBAL("meta_schedule.CostModelPredict")
.set_body_typed([](CostModel model, //
const TuneContext& context, //
diff --git a/src/meta_schedule/database/database.cc b/src/meta_schedule/database/database.cc
index 153cb3e9ee..8e800bdb9b 100644
--- a/src/meta_schedule/database/database.cc
+++ b/src/meta_schedule/database/database.cc
@@ -285,8 +285,7 @@ TVM_REGISTER_NODE_TYPE(PyDatabaseNode);
TVM_REGISTER_GLOBAL("meta_schedule.Workload").set_body_typed([](IRModule mod) {
return Workload(mod);
});
-TVM_REGISTER_GLOBAL("meta_schedule.WorkloadAsJSON")
- .set_body_method<Workload>(&WorkloadNode::AsJSON);
+TVM_REGISTER_GLOBAL("meta_schedule.WorkloadAsJSON").set_body_method(&WorkloadNode::AsJSON);
TVM_REGISTER_GLOBAL("meta_schedule.WorkloadFromJSON").set_body_typed(&Workload::FromJSON);
TVM_REGISTER_GLOBAL("meta_schedule.TuningRecord")
.set_body_typed([](tir::Trace trace, Workload workload, Optional<Array<FloatImm>> run_secs,
@@ -294,9 +293,8 @@ TVM_REGISTER_GLOBAL("meta_schedule.TuningRecord")
return TuningRecord(trace, workload, run_secs, target, args_info);
});
TVM_REGISTER_GLOBAL("meta_schedule.TuningRecordAsMeasureCandidate")
- .set_body_method<TuningRecord>(&TuningRecordNode::AsMeasureCandidate);
-TVM_REGISTER_GLOBAL("meta_schedule.TuningRecordAsJSON")
- .set_body_method<TuningRecord>(&TuningRecordNode::AsJSON);
+ .set_body_method(&TuningRecordNode::AsMeasureCandidate);
+TVM_REGISTER_GLOBAL("meta_schedule.TuningRecordAsJSON").set_body_method(&TuningRecordNode::AsJSON);
TVM_REGISTER_GLOBAL("meta_schedule.TuningRecordFromJSON").set_body_typed(TuningRecord::FromJSON);
TVM_REGISTER_GLOBAL("meta_schedule.DatabaseEnterWithScope")
.set_body_method(&Database::EnterWithScope);
@@ -304,24 +302,22 @@ TVM_REGISTER_GLOBAL("meta_schedule.DatabaseExitWithScope")
.set_body_method(&Database::ExitWithScope);
TVM_REGISTER_GLOBAL("meta_schedule.DatabaseCurrent").set_body_typed(Database::Current);
TVM_REGISTER_GLOBAL("meta_schedule.DatabaseHasWorkload")
- .set_body_method<Database>(&DatabaseNode::HasWorkload);
+ .set_body_method(&DatabaseNode::HasWorkload);
TVM_REGISTER_GLOBAL("meta_schedule.DatabaseCommitWorkload")
- .set_body_method<Database>(&DatabaseNode::CommitWorkload);
+ .set_body_method(&DatabaseNode::CommitWorkload);
TVM_REGISTER_GLOBAL("meta_schedule.DatabaseCommitTuningRecord")
- .set_body_method<Database>(&DatabaseNode::CommitTuningRecord);
-TVM_REGISTER_GLOBAL("meta_schedule.DatabaseGetTopK")
- .set_body_method<Database>(&DatabaseNode::GetTopK);
+ .set_body_method(&DatabaseNode::CommitTuningRecord);
+TVM_REGISTER_GLOBAL("meta_schedule.DatabaseGetTopK").set_body_method(&DatabaseNode::GetTopK);
TVM_REGISTER_GLOBAL("meta_schedule.DatabaseGetAllTuningRecords")
- .set_body_method<Database>(&DatabaseNode::GetAllTuningRecords);
-TVM_REGISTER_GLOBAL("meta_schedule.DatabaseSize").set_body_method<Database>(&DatabaseNode::Size);
+ .set_body_method(&DatabaseNode::GetAllTuningRecords);
+TVM_REGISTER_GLOBAL("meta_schedule.DatabaseSize").set_body_method(&DatabaseNode::Size);
TVM_REGISTER_GLOBAL("meta_schedule.DatabaseQueryTuningRecord")
- .set_body_method<Database>(&DatabaseNode::QueryTuningRecord);
+ .set_body_method(&DatabaseNode::QueryTuningRecord);
TVM_REGISTER_GLOBAL("meta_schedule.DatabaseQuerySchedule")
- .set_body_method<Database>(&DatabaseNode::QuerySchedule);
+ .set_body_method(&DatabaseNode::QuerySchedule);
TVM_REGISTER_GLOBAL("meta_schedule.DatabaseQueryIRModule")
- .set_body_method<Database>(&DatabaseNode::QueryIRModule);
-TVM_REGISTER_GLOBAL("meta_schedule.DatabaseDumpPruned")
- .set_body_method<Database>(&DatabaseNode::DumpPruned);
+ .set_body_method(&DatabaseNode::QueryIRModule);
+TVM_REGISTER_GLOBAL("meta_schedule.DatabaseDumpPruned").set_body_method(&DatabaseNode::DumpPruned);
TVM_REGISTER_GLOBAL("meta_schedule.DatabasePyDatabase").set_body_typed(Database::PyDatabase);
} // namespace meta_schedule
diff --git a/src/meta_schedule/feature_extractor/feature_extractor.cc b/src/meta_schedule/feature_extractor/feature_extractor.cc
index 1ebbb6e2e2..093558d228 100644
--- a/src/meta_schedule/feature_extractor/feature_extractor.cc
+++ b/src/meta_schedule/feature_extractor/feature_extractor.cc
@@ -49,7 +49,7 @@ TVM_REGISTER_OBJECT_TYPE(FeatureExtractorNode);
TVM_REGISTER_NODE_TYPE(PyFeatureExtractorNode);
TVM_REGISTER_GLOBAL("meta_schedule.FeatureExtractorExtractFrom")
- .set_body_method<FeatureExtractor>(&FeatureExtractorNode::ExtractFrom);
+ .set_body_method(&FeatureExtractorNode::ExtractFrom);
TVM_REGISTER_GLOBAL("meta_schedule.FeatureExtractorPyFeatureExtractor")
.set_body_typed(FeatureExtractor::PyFeatureExtractor);
diff --git a/src/meta_schedule/measure_callback/measure_callback.cc b/src/meta_schedule/measure_callback/measure_callback.cc
index f16fb73c52..8f94e29846 100644
--- a/src/meta_schedule/measure_callback/measure_callback.cc
+++ b/src/meta_schedule/measure_callback/measure_callback.cc
@@ -60,7 +60,7 @@ TVM_REGISTER_OBJECT_TYPE(MeasureCallbackNode);
TVM_REGISTER_NODE_TYPE(PyMeasureCallbackNode);
TVM_REGISTER_GLOBAL("meta_schedule.MeasureCallbackApply")
- .set_body_method<MeasureCallback>(&MeasureCallbackNode::Apply);
+ .set_body_method(&MeasureCallbackNode::Apply);
TVM_REGISTER_GLOBAL("meta_schedule.MeasureCallbackPyMeasureCallback")
.set_body_typed(MeasureCallback::PyMeasureCallback);
TVM_REGISTER_GLOBAL("meta_schedule.MeasureCallbackDefault")
diff --git a/src/meta_schedule/measure_callback/remove_build_artifact.cc b/src/meta_schedule/measure_callback/remove_build_artifact.cc
index 41e52adbae..9242e79912 100644
--- a/src/meta_schedule/measure_callback/remove_build_artifact.cc
+++ b/src/meta_schedule/measure_callback/remove_build_artifact.cc
@@ -27,12 +27,11 @@ class RemoveBuildArtifactNode : public MeasureCallbackNode {
const Array<MeasureCandidate>& measure_candidates,
const Array<BuilderResult>& builder_results,
const Array<RunnerResult>& runner_results) final {
- static const PackedFunc* f_rm = runtime::Registry::Get("meta_schedule.remove_build_dir");
- ICHECK(f_rm != nullptr) << "The `remove_build_dir` func is not in tvm registry.";
+ static auto f_rm = tvm::ffi::Function::GetGlobalRequired("meta_schedule.remove_build_dir");
auto _ = Profiler::TimedScope("MeasureCallback/RemoveBuildArtifact");
for (const BuilderResult& build_result : builder_results) {
if (Optional<String> path = build_result->artifact_path) {
- (*f_rm)(path.value());
+ f_rm(path.value());
}
}
}
diff --git a/src/meta_schedule/mutator/mutator.cc b/src/meta_schedule/mutator/mutator.cc
index 8e9bfc8bde..e1831d213d 100644
--- a/src/meta_schedule/mutator/mutator.cc
+++ b/src/meta_schedule/mutator/mutator.cc
@@ -89,13 +89,13 @@ TVM_REGISTER_OBJECT_TYPE(MutatorNode);
TVM_REGISTER_NODE_TYPE(PyMutatorNode);
TVM_REGISTER_GLOBAL("meta_schedule.MutatorInitializeWithTuneContext")
- .set_body_method<Mutator>(&MutatorNode::InitializeWithTuneContext);
+ .set_body_method(&MutatorNode::InitializeWithTuneContext);
TVM_REGISTER_GLOBAL("meta_schedule.MutatorApply")
.set_body_typed([](Mutator self, tir::Trace trace, TRandState seed) -> Optional<tir::Trace> {
TRandState seed_ = (seed != -1) ? seed : support::LinearCongruentialEngine::DeviceRandom();
return self->Apply(trace, &seed_);
});
-TVM_REGISTER_GLOBAL("meta_schedule.MutatorClone").set_body_method<Mutator>(&MutatorNode::Clone);
+TVM_REGISTER_GLOBAL("meta_schedule.MutatorClone").set_body_method(&MutatorNode::Clone);
TVM_REGISTER_GLOBAL("meta_schedule.MutatorPyMutator").set_body_typed(Mutator::PyMutator);
TVM_REGISTER_GLOBAL("meta_schedule.MutatorDefaultLLVM").set_body_typed(Mutator::DefaultLLVM);
TVM_REGISTER_GLOBAL("meta_schedule.MutatorDefaultCUDA").set_body_typed(Mutator::DefaultCUDA);
diff --git a/src/meta_schedule/postproc/postproc.cc b/src/meta_schedule/postproc/postproc.cc
index a5f833044c..b5d62634c2 100644
--- a/src/meta_schedule/postproc/postproc.cc
+++ b/src/meta_schedule/postproc/postproc.cc
@@ -113,9 +113,9 @@ TVM_REGISTER_OBJECT_TYPE(PostprocNode);
TVM_REGISTER_NODE_TYPE(PyPostprocNode);
TVM_REGISTER_GLOBAL("meta_schedule.PostprocInitializeWithTuneContext")
- .set_body_method<Postproc>(&PostprocNode::InitializeWithTuneContext);
-TVM_REGISTER_GLOBAL("meta_schedule.PostprocApply").set_body_method<Postproc>(&PostprocNode::Apply);
-TVM_REGISTER_GLOBAL("meta_schedule.PostprocClone").set_body_method<Postproc>(&PostprocNode::Clone);
+ .set_body_method(&PostprocNode::InitializeWithTuneContext);
+TVM_REGISTER_GLOBAL("meta_schedule.PostprocApply").set_body_method(&PostprocNode::Apply);
+TVM_REGISTER_GLOBAL("meta_schedule.PostprocClone").set_body_method(&PostprocNode::Clone);
TVM_REGISTER_GLOBAL("meta_schedule.PostprocPyPostproc").set_body_typed(Postproc::PyPostproc);
TVM_REGISTER_GLOBAL("meta_schedule.PostprocDefaultLLVM").set_body_typed(Postproc::DefaultLLVM);
TVM_REGISTER_GLOBAL("meta_schedule.PostprocDefaultCUDA").set_body_typed(Postproc::DefaultCUDA);
diff --git a/src/meta_schedule/profiler.cc b/src/meta_schedule/profiler.cc
index 98aaef59a5..ecbb1c50ce 100644
--- a/src/meta_schedule/profiler.cc
+++ b/src/meta_schedule/profiler.cc
@@ -129,8 +129,8 @@ TVM_REGISTER_GLOBAL("meta_schedule.ProfilerEnterWithScope")
TVM_REGISTER_GLOBAL("meta_schedule.ProfilerExitWithScope")
.set_body_method(&Profiler::ExitWithScope);
TVM_REGISTER_GLOBAL("meta_schedule.ProfilerCurrent").set_body_typed(Profiler::Current);
-TVM_REGISTER_GLOBAL("meta_schedule.ProfilerGet").set_body_method<Profiler>(&ProfilerNode::Get);
-TVM_REGISTER_GLOBAL("meta_schedule.ProfilerTable").set_body_method<Profiler>(&ProfilerNode::Table);
+TVM_REGISTER_GLOBAL("meta_schedule.ProfilerGet").set_body_method(&ProfilerNode::Get);
+TVM_REGISTER_GLOBAL("meta_schedule.ProfilerTable").set_body_method(&ProfilerNode::Table);
TVM_REGISTER_GLOBAL("meta_schedule.ProfilerTimedScope").set_body_typed(ProfilerTimedScope);
} // namespace meta_schedule
diff --git a/src/meta_schedule/runner/runner.cc b/src/meta_schedule/runner/runner.cc
index a85c4d3f38..0d5edf299e 100644
--- a/src/meta_schedule/runner/runner.cc
+++ b/src/meta_schedule/runner/runner.cc
@@ -70,11 +70,9 @@ TVM_REGISTER_GLOBAL("meta_schedule.RunnerFuture")
.set_body_typed([](RunnerFuture::FDone f_done, RunnerFuture::FResult f_result) -> RunnerFuture {
return RunnerFuture(f_done, f_result);
});
-TVM_REGISTER_GLOBAL("meta_schedule.RunnerFutureDone")
- .set_body_method<RunnerFuture>(&RunnerFutureNode::Done);
-TVM_REGISTER_GLOBAL("meta_schedule.RunnerFutureResult")
- .set_body_method<RunnerFuture>(&RunnerFutureNode::Result);
-TVM_REGISTER_GLOBAL("meta_schedule.RunnerRun").set_body_method<Runner>(&RunnerNode::Run);
+TVM_REGISTER_GLOBAL("meta_schedule.RunnerFutureDone").set_body_method(&RunnerFutureNode::Done);
+TVM_REGISTER_GLOBAL("meta_schedule.RunnerFutureResult").set_body_method(&RunnerFutureNode::Result);
+TVM_REGISTER_GLOBAL("meta_schedule.RunnerRun").set_body_method(&RunnerNode::Run);
TVM_REGISTER_GLOBAL("meta_schedule.RunnerPyRunner").set_body_typed(Runner::PyRunner);
} // namespace meta_schedule
diff --git a/src/meta_schedule/schedule_rule/apply_custom_rule.cc b/src/meta_schedule/schedule_rule/apply_custom_rule.cc
index 011e96e895..e972be32e1 100644
--- a/src/meta_schedule/schedule_rule/apply_custom_rule.cc
+++ b/src/meta_schedule/schedule_rule/apply_custom_rule.cc
@@ -41,9 +41,9 @@ class ApplyCustomRuleNode : public ScheduleRuleNode {
if (Optional<String> ann = tir::GetAnn<String>(sch->GetSRef(block_rv), "schedule_rule")) {
if (ann.value() != "None") {
for (const String& key : keys) {
- if (const runtime::PackedFunc* custom_schedule_fn =
- runtime::Registry::Get(GetCustomRuleName(ann.value(), key))) {
- Array<tir::Schedule> result = ((*custom_schedule_fn)(sch, block_rv));
+ if (const auto custom_schedule_fn =
+ tvm::ffi::Function::GetGlobal(GetCustomRuleName(ann.value(), key))) {
+ Array<tir::Schedule> result = (*custom_schedule_fn)(sch, block_rv);
return result;
}
}
diff --git a/src/meta_schedule/schedule_rule/schedule_rule.cc b/src/meta_schedule/schedule_rule/schedule_rule.cc
index d5bbe8207a..ec3c6ac480 100644
--- a/src/meta_schedule/schedule_rule/schedule_rule.cc
+++ b/src/meta_schedule/schedule_rule/schedule_rule.cc
@@ -403,11 +403,9 @@ TVM_REGISTER_OBJECT_TYPE(ScheduleRuleNode);
TVM_REGISTER_NODE_TYPE(PyScheduleRuleNode);
TVM_REGISTER_GLOBAL("meta_schedule.ScheduleRuleInitializeWithTuneContext")
- .set_body_method<ScheduleRule>(&ScheduleRuleNode::InitializeWithTuneContext);
-TVM_REGISTER_GLOBAL("meta_schedule.ScheduleRuleApply")
- .set_body_method<ScheduleRule>(&ScheduleRuleNode::Apply);
-TVM_REGISTER_GLOBAL("meta_schedule.ScheduleRuleClone")
- .set_body_method<ScheduleRule>(&ScheduleRuleNode::Clone);
+ .set_body_method(&ScheduleRuleNode::InitializeWithTuneContext);
+TVM_REGISTER_GLOBAL("meta_schedule.ScheduleRuleApply").set_body_method(&ScheduleRuleNode::Apply);
+TVM_REGISTER_GLOBAL("meta_schedule.ScheduleRuleClone").set_body_method(&ScheduleRuleNode::Clone);
TVM_REGISTER_GLOBAL("meta_schedule.ScheduleRulePyScheduleRule")
.set_body_typed(ScheduleRule::PyScheduleRule);
TVM_REGISTER_GLOBAL("meta_schedule.ScheduleRuleDefaultLLVM")
diff --git a/src/meta_schedule/search_strategy/search_strategy.cc b/src/meta_schedule/search_strategy/search_strategy.cc
index e1846dc5cb..1bc71502ad 100644
--- a/src/meta_schedule/search_strategy/search_strategy.cc
+++ b/src/meta_schedule/search_strategy/search_strategy.cc
@@ -93,17 +93,17 @@ TVM_REGISTER_GLOBAL("meta_schedule.MeasureCandidate")
TVM_REGISTER_GLOBAL("meta_schedule.SearchStrategyPySearchStrategy")
.set_body_typed(SearchStrategy::PySearchStrategy);
TVM_REGISTER_GLOBAL("meta_schedule.SearchStrategyInitializeWithTuneContext")
- .set_body_method<SearchStrategy>(&SearchStrategyNode::InitializeWithTuneContext);
+ .set_body_method(&SearchStrategyNode::InitializeWithTuneContext);
TVM_REGISTER_GLOBAL("meta_schedule.SearchStrategyPreTuning")
- .set_body_method<SearchStrategy>(&SearchStrategyNode::PreTuning);
+ .set_body_method(&SearchStrategyNode::PreTuning);
TVM_REGISTER_GLOBAL("meta_schedule.SearchStrategyPostTuning")
- .set_body_method<SearchStrategy>(&SearchStrategyNode::PostTuning);
+ .set_body_method(&SearchStrategyNode::PostTuning);
TVM_REGISTER_GLOBAL("meta_schedule.SearchStrategyGenerateMeasureCandidates")
- .set_body_method<SearchStrategy>(&SearchStrategyNode::GenerateMeasureCandidates);
+ .set_body_method(&SearchStrategyNode::GenerateMeasureCandidates);
TVM_REGISTER_GLOBAL("meta_schedule.SearchStrategyNotifyRunnerResults")
- .set_body_method<SearchStrategy>(&SearchStrategyNode::NotifyRunnerResults);
+ .set_body_method(&SearchStrategyNode::NotifyRunnerResults);
TVM_REGISTER_GLOBAL("meta_schedule.SearchStrategyClone")
- .set_body_method<SearchStrategy>(&SearchStrategyNode::Clone);
+ .set_body_method(&SearchStrategyNode::Clone);
} // namespace meta_schedule
} // namespace tvm
diff --git a/src/meta_schedule/space_generator/space_generator.cc b/src/meta_schedule/space_generator/space_generator.cc
index 400532750c..036705e43b 100644
--- a/src/meta_schedule/space_generator/space_generator.cc
+++ b/src/meta_schedule/space_generator/space_generator.cc
@@ -24,17 +24,15 @@ namespace meta_schedule {
String GetRuleKindFromTarget(const Target& target) {
if (target->kind->name == "llvm") {
- static const PackedFunc* target_has_feature_fn_ptr =
- runtime::Registry::Get("target.target_has_feature");
- ICHECK(target_has_feature_fn_ptr != nullptr)
- << "The `target.target_has_feature` func is not in tvm registry.";
- bool have_avx512vnni = (*target_has_feature_fn_ptr)("avx512vnni", target);
- bool have_avxvnni = (*target_has_feature_fn_ptr)("avxvnni", target);
+ static auto target_has_feature_fn_ptr =
+ tvm::ffi::Function::GetGlobalRequired("target.target_has_feature");
+ bool have_avx512vnni = target_has_feature_fn_ptr("avx512vnni", target);
+ bool have_avxvnni = target_has_feature_fn_ptr("avxvnni", target);
if (have_avx512vnni || have_avxvnni) {
return "vnni";
} else {
- bool have_avx512f = (*target_has_feature_fn_ptr)("avx512f", target);
- bool have_avx512bw = (*target_has_feature_fn_ptr)("avx512bw", target);
+ bool have_avx512f = target_has_feature_fn_ptr("avx512f", target);
+ bool have_avx512bw = target_has_feature_fn_ptr("avx512bw", target);
if (have_avx512bw && have_avx512f) {
return "avx512";
}
@@ -193,13 +191,13 @@ TVM_REGISTER_OBJECT_TYPE(SpaceGeneratorNode);
TVM_REGISTER_NODE_TYPE(PySpaceGeneratorNode);
TVM_REGISTER_GLOBAL("meta_schedule.SpaceGeneratorInitializeWithTuneContext")
- .set_body_method<SpaceGenerator>(&SpaceGeneratorNode::InitializeWithTuneContext);
+ .set_body_method(&SpaceGeneratorNode::InitializeWithTuneContext);
TVM_REGISTER_GLOBAL("meta_schedule.SpaceGeneratorGenerateDesignSpace")
- .set_body_method<SpaceGenerator>(&SpaceGeneratorNode::GenerateDesignSpace);
+ .set_body_method(&SpaceGeneratorNode::GenerateDesignSpace);
TVM_REGISTER_GLOBAL("meta_schedule.SpaceGeneratorPySpaceGenerator")
.set_body_typed(SpaceGenerator::PySpaceGenerator);
TVM_REGISTER_GLOBAL("meta_schedule.SpaceGeneratorClone")
- .set_body_method<SpaceGenerator>(&SpaceGeneratorNode::Clone);
+ .set_body_method(&SpaceGeneratorNode::Clone);
} // namespace meta_schedule
} // namespace tvm
diff --git a/src/meta_schedule/task_scheduler/task_scheduler.cc b/src/meta_schedule/task_scheduler/task_scheduler.cc
index 404ee01983..f17cfd55cf 100644
--- a/src/meta_schedule/task_scheduler/task_scheduler.cc
+++ b/src/meta_schedule/task_scheduler/task_scheduler.cc
@@ -366,18 +366,17 @@ TVM_REGISTER_OBJECT_TYPE(TaskSchedulerNode);
TVM_REGISTER_NODE_TYPE(PyTaskSchedulerNode);
TVM_REGISTER_GLOBAL("meta_schedule.TaskSchedulerPyTaskScheduler")
.set_body_typed(TaskScheduler::PyTaskScheduler);
-TVM_REGISTER_GLOBAL("meta_schedule.TaskSchedulerTune")
- .set_body_method<TaskScheduler>(&TaskSchedulerNode::Tune);
+TVM_REGISTER_GLOBAL("meta_schedule.TaskSchedulerTune").set_body_method(&TaskSchedulerNode::Tune);
TVM_REGISTER_GLOBAL("meta_schedule.TaskSchedulerJoinRunningTask")
- .set_body_method<TaskScheduler>(&TaskSchedulerNode::JoinRunningTask);
+ .set_body_method(&TaskSchedulerNode::JoinRunningTask);
TVM_REGISTER_GLOBAL("meta_schedule.TaskSchedulerNextTaskId")
- .set_body_method<TaskScheduler>(&TaskSchedulerNode::NextTaskId);
+ .set_body_method(&TaskSchedulerNode::NextTaskId);
TVM_REGISTER_GLOBAL("meta_schedule.TaskSchedulerTerminateTask")
- .set_body_method<TaskScheduler>(&TaskSchedulerNode::TerminateTask);
+ .set_body_method(&TaskSchedulerNode::TerminateTask);
TVM_REGISTER_GLOBAL("meta_schedule.TaskSchedulerTouchTask")
- .set_body_method<TaskScheduler>(&TaskSchedulerNode::TouchTask);
+ .set_body_method(&TaskSchedulerNode::TouchTask);
TVM_REGISTER_GLOBAL("meta_schedule.TaskSchedulerPrintTuningStatistics")
- .set_body_method<TaskScheduler>(&TaskSchedulerNode::PrintTuningStatistics);
+ .set_body_method(&TaskSchedulerNode::PrintTuningStatistics);
} // namespace meta_schedule
} // namespace tvm
diff --git a/src/meta_schedule/tune_context.cc b/src/meta_schedule/tune_context.cc
index 768c958571..e9cb17eeee 100644
--- a/src/meta_schedule/tune_context.cc
+++ b/src/meta_schedule/tune_context.cc
@@ -73,9 +73,8 @@ TVM_REGISTER_GLOBAL("meta_schedule.TuneContext")
});
TVM_REGISTER_GLOBAL("meta_schedule._SHash2Hex").set_body_typed(SHash2Hex);
TVM_REGISTER_GLOBAL("meta_schedule.TuneContextInitialize")
- .set_body_method<TuneContext>(&TuneContextNode::Initialize);
-TVM_REGISTER_GLOBAL("meta_schedule.TuneContextClone")
- .set_body_method<TuneContext>(&TuneContextNode::Clone);
+ .set_body_method(&TuneContextNode::Initialize);
+TVM_REGISTER_GLOBAL("meta_schedule.TuneContextClone").set_body_method(&TuneContextNode::Clone);
} // namespace meta_schedule
} // namespace tvm
diff --git a/src/meta_schedule/utils.h b/src/meta_schedule/utils.h
index f8812dab37..decd8731b8 100644
--- a/src/meta_schedule/utils.h
+++ b/src/meta_schedule/utils.h
@@ -124,7 +124,7 @@ class PyLogMessage {
*/
inline bool using_ipython() {
bool flag = false;
- const auto* f_using_ipython = runtime::Registry::Get("meta_schedule.using_ipython");
+ const auto f_using_ipython = tvm::ffi::Function::GetGlobal("meta_schedule.using_ipython");
if (f_using_ipython) {
flag = (*f_using_ipython)();
}
@@ -136,9 +136,9 @@ inline bool using_ipython() {
* \param str The serialized performance table.
*/
inline void print_interactive_table(const String& data) {
- const auto* f_print_interactive_table =
- runtime::Registry::Get("meta_schedule.print_interactive_table");
- ICHECK(f_print_interactive_table->defined())
+ const auto f_print_interactive_table =
+ tvm::ffi::Function::GetGlobal("meta_schedule.print_interactive_table");
+ ICHECK(f_print_interactive_table.has_value())
<< "Cannot find print_interactive_table function in registry.";
(*f_print_interactive_table)(data);
}
@@ -380,8 +380,8 @@ struct ThreadedTraceApply {
inline int GetTargetNumCores(const Target& target) {
int num_cores = target->GetAttr<Integer>("num-cores").value_or(-1).IntValue();
if (num_cores == -1) {
- static const auto* f_cpu_count = runtime::Registry::Get("meta_schedule.cpu_count");
- ICHECK(f_cpu_count)
+ static const auto f_cpu_count = tvm::ffi::Function::GetGlobal("meta_schedule.cpu_count");
+ ICHECK(f_cpu_count.has_value())
<< "ValueError: Cannot find the packed function \"meta_schedule._cpu_count\"";
num_cores = (*f_cpu_count)(false);
LOG(FATAL)
diff --git a/src/node/object_path.cc b/src/node/object_path.cc
index d7660e6153..fda5d3d9d1 100644
--- a/src/node/object_path.cc
+++ b/src/node/object_path.cc
@@ -40,14 +40,13 @@ Optional<ObjectPath> ObjectPathNode::GetParent() const {
return Downcast<Optional<ObjectPath>>(parent_);
}
-TVM_REGISTER_GLOBAL("node.ObjectPathGetParent")
- .set_body_method<ObjectPath>(&ObjectPathNode::GetParent);
+TVM_REGISTER_GLOBAL("node.ObjectPathGetParent").set_body_method(&ObjectPathNode::GetParent);
// --- Length ---
int32_t ObjectPathNode::Length() const { return length_; }
-TVM_REGISTER_GLOBAL("node.ObjectPathLength").set_body_method<ObjectPath>(&ObjectPathNode::Length);
+TVM_REGISTER_GLOBAL("node.ObjectPathLength").set_body_method(&ObjectPathNode::Length);
// --- GetPrefix ---
@@ -64,8 +63,7 @@ ObjectPath ObjectPathNode::GetPrefix(int32_t length) const {
return GetRef<ObjectPath>(node);
}
-TVM_REGISTER_GLOBAL("node.ObjectPathGetPrefix")
- .set_body_method<ObjectPath>(&ObjectPathNode::GetPrefix);
+TVM_REGISTER_GLOBAL("node.ObjectPathGetPrefix").set_body_method(&ObjectPathNode::GetPrefix);
// --- IsPrefixOf ---
@@ -77,8 +75,7 @@ bool ObjectPathNode::IsPrefixOf(const ObjectPath& other) const {
return this->PathsEqual(other->GetPrefix(this_len));
}
-TVM_REGISTER_GLOBAL("node.ObjectPathIsPrefixOf")
- .set_body_method<ObjectPath>(&ObjectPathNode::IsPrefixOf);
+TVM_REGISTER_GLOBAL("node.ObjectPathIsPrefixOf").set_body_method(&ObjectPathNode::IsPrefixOf);
// --- Attr ---
@@ -109,8 +106,7 @@ ObjectPath ObjectPathNode::ArrayIndex(int32_t index) const {
return ObjectPath(make_object<ArrayIndexPathNode>(this, index));
}
-TVM_REGISTER_GLOBAL("node.ObjectPathArrayIndex")
- .set_body_method<ObjectPath>(&ObjectPathNode::ArrayIndex);
+TVM_REGISTER_GLOBAL("node.ObjectPathArrayIndex").set_body_method(&ObjectPathNode::ArrayIndex);
// --- MissingArrayElement ---
@@ -119,7 +115,7 @@ ObjectPath ObjectPathNode::MissingArrayElement(int32_t index) const {
}
TVM_REGISTER_GLOBAL("node.ObjectPathMissingArrayElement")
- .set_body_method<ObjectPath>(&ObjectPathNode::MissingArrayElement);
+ .set_body_method(&ObjectPathNode::MissingArrayElement);
// --- MapValue ---
@@ -127,8 +123,7 @@ ObjectPath ObjectPathNode::MapValue(Any key) const {
return ObjectPath(make_object<MapValuePathNode>(this, std::move(key)));
}
-TVM_REGISTER_GLOBAL("node.ObjectPathMapValue")
- .set_body_method<ObjectPath>(&ObjectPathNode::MapValue);
+TVM_REGISTER_GLOBAL("node.ObjectPathMapValue").set_body_method(&ObjectPathNode::MapValue);
// --- MissingMapEntry ---
@@ -137,7 +132,7 @@ ObjectPath ObjectPathNode::MissingMapEntry() const {
}
TVM_REGISTER_GLOBAL("node.ObjectPathMissingMapEntry")
- .set_body_method<ObjectPath>(&ObjectPathNode::MissingMapEntry);
+ .set_body_method(&ObjectPathNode::MissingMapEntry);
// --- PathsEqual ----
@@ -163,8 +158,7 @@ bool ObjectPathNode::PathsEqual(const ObjectPath& other) const {
return lhs == nullptr && rhs == nullptr;
}
-TVM_REGISTER_GLOBAL("node.ObjectPathEqual")
- .set_body_method<ObjectPath>(&ObjectPathNode::PathsEqual);
+TVM_REGISTER_GLOBAL("node.ObjectPathEqual").set_body_method(&ObjectPathNode::PathsEqual);
// --- Repr ---
diff --git a/src/node/reflection.cc b/src/node/reflection.cc
index 7d55059cca..92687a3bb6 100644
--- a/src/node/reflection.cc
+++ b/src/node/reflection.cc
@@ -292,11 +292,11 @@ void MakeNode(const TVMArgs& args, TVMRetValue* rv) {
*rv = ReflectionVTable::Global()->CreateObject(type_key, args.Slice(1));
}
-TVM_REGISTER_GLOBAL("node.NodeGetAttr").set_body(NodeGetAttr);
+TVM_REGISTER_GLOBAL("node.NodeGetAttr").set_body_packed(NodeGetAttr);
-TVM_REGISTER_GLOBAL("node.NodeListAttrNames").set_body(NodeListAttrNames);
+TVM_REGISTER_GLOBAL("node.NodeListAttrNames").set_body_packed(NodeListAttrNames);
-TVM_REGISTER_GLOBAL("node.MakeNode").set_body(MakeNode);
+TVM_REGISTER_GLOBAL("node.MakeNode").set_body_packed(MakeNode);
namespace {
// Attribute visitor class for finding the attribute key by its address
diff --git a/src/relax/backend/contrib/clml/codegen.cc b/src/relax/backend/contrib/clml/codegen.cc
index 5fd8a923a1..7d3d243fc9 100644
--- a/src/relax/backend/contrib/clml/codegen.cc
+++ b/src/relax/backend/contrib/clml/codegen.cc
@@ -314,11 +314,10 @@ Array<runtime::Module> OpenCLMLCompiler(Array<Function> functions, Map<String, A
serializer.serialize(func);
std::string graph_json = serializer.GetJSON();
auto constant_names = serializer.GetConstantNames();
- const auto* pf = runtime::Registry::Get("runtime.clml_runtime_create");
- ICHECK(pf != nullptr) << "Cannot find OpenCLML runtime module create function.";
+ const auto pf = tvm::ffi::Function::GetGlobalRequired("runtime.clml_runtime_create");
std::string func_name = GetExtSymbol(func);
VLOG(1) << "Creating clml runtime::Module for '" << func_name << "'";
- compiled_functions.push_back((*pf)(func_name, graph_json, constant_names));
+ compiled_functions.push_back(pf(func_name, graph_json, constant_names));
}
return compiled_functions;
}
diff --git a/src/relax/backend/contrib/cublas/codegen.cc b/src/relax/backend/contrib/cublas/codegen.cc
index 0e581c9d56..bea9ed7baa 100644
--- a/src/relax/backend/contrib/cublas/codegen.cc
+++ b/src/relax/backend/contrib/cublas/codegen.cc
@@ -117,10 +117,9 @@ Array<runtime::Module> CublasCompiler(Array<Function> functions, Map<String, ffi
serializer.serialize(func);
auto graph_json = serializer.GetJSON();
auto constant_names = serializer.GetConstantNames();
- const auto* pf = runtime::Registry::Get("runtime.CublasJSONRuntimeCreate");
- ICHECK(pf != nullptr) << "Cannot find CUBLAS runtime module create function.";
+ const auto pf = tvm::ffi::Function::GetGlobalRequired("runtime.CublasJSONRuntimeCreate");
auto func_name = GetExtSymbol(func);
- compiled_functions.push_back((*pf)(func_name, graph_json, constant_names));
+ compiled_functions.push_back(pf(func_name, graph_json, constant_names));
}
return compiled_functions;
diff --git a/src/relax/backend/contrib/cudnn/codegen.cc b/src/relax/backend/contrib/cudnn/codegen.cc
index 55471fb80d..e183dadc26 100644
--- a/src/relax/backend/contrib/cudnn/codegen.cc
+++ b/src/relax/backend/contrib/cudnn/codegen.cc
@@ -141,10 +141,9 @@ Array<runtime::Module> cuDNNCompiler(Array<Function> functions, Map<String, ffi:
serializer.serialize(func);
auto graph_json = serializer.GetJSON();
auto constant_names = serializer.GetConstantNames();
- const auto* pf = runtime::Registry::Get("runtime.cuDNNJSONRuntimeCreate");
- ICHECK(pf != nullptr) << "Cannot find cuDNN runtime module create function.";
+ const auto pf = tvm::ffi::Function::GetGlobalRequired("runtime.cuDNNJSONRuntimeCreate");
auto func_name = GetExtSymbol(func);
- compiled_functions.push_back((*pf)(func_name, graph_json, constant_names));
+ compiled_functions.push_back(pf(func_name, graph_json, constant_names));
}
return compiled_functions;
diff --git a/src/relax/backend/contrib/cutlass/codegen.cc b/src/relax/backend/contrib/cutlass/codegen.cc
index 7b852ecbc8..4aa5000bbb 100644
--- a/src/relax/backend/contrib/cutlass/codegen.cc
+++ b/src/relax/backend/contrib/cutlass/codegen.cc
@@ -67,10 +67,9 @@ runtime::Module Finalize(const std::string& code, const Array<String>& func_name
default_headers << "#include <cutlass/tensor_ref.h>\n";
default_headers << "#include <cutlass/util/host_tensor.h>\n";
- const auto* pf = runtime::Registry::Get("runtime.CSourceModuleCreate");
- ICHECK(pf != nullptr) << "Cannot find CSource module to create the external runtime module";
+ const auto pf = tvm::ffi::Function::GetGlobalRequired("runtime.CSourceModuleCreate");
VLOG(1) << "Generated CUTLASS code:" << std::endl << code;
- return (*pf)(default_headers.str() + code, "cu", func_names, /*const_vars=*/Array<String>());
+ return pf(default_headers.str() + code, "cu", func_names, /*const_vars=*/Array<String>());
}
class CodegenResultNode : public Object {
@@ -128,11 +127,9 @@ GenerateBodyOutput GenerateBody(const std::string& func_name, const std::string&
}
decl_stream << ");";
- const auto* instantiate_template_func =
- runtime::Registry::Get("contrib.cutlass.instantiate_template");
- ICHECK(instantiate_template_func);
-
- CodegenResult codegen_res = (*instantiate_template_func)(func_name, attrs, func_args);
+ const auto instantiate_template_func =
+ tvm::ffi::Function::GetGlobalRequired("contrib.cutlass.instantiate_template");
+ CodegenResult codegen_res = instantiate_template_func(func_name, attrs, func_args);
ret.decl = codegen_res->code;
ret.headers = codegen_res->headers;
@@ -369,17 +366,17 @@ class CutlassModuleCodegen {
Array<runtime::Module> CUTLASSCompiler(Array<Function> functions, Map<String, ffi::Any> options,
Map<Constant, String> /*unused*/) {
- const auto* tune_func = runtime::Registry::Get("contrib.cutlass.tune_relax_function");
- ICHECK(tune_func != nullptr)
+ const auto tune_func = tvm::ffi::Function::GetGlobal("contrib.cutlass.tune_relax_function");
+ ICHECK(tune_func.has_value())
<< "The packed function contrib.cutlass.tune_relax_function not found, "
"please import tvm.contrib.cutlass.build";
Array<Function> annotated_functions = (*tune_func)(functions, options);
auto source_mod = CutlassModuleCodegen().CreateCSourceModule(annotated_functions, options);
- const auto* pf = runtime::Registry::Get("contrib.cutlass.compile");
- ICHECK(pf != nullptr) << "The packed function contrib.cutlass.compile not found, please import "
- "tvm.contrib.cutlass.build";
+ const auto pf = tvm::ffi::Function::GetGlobal("contrib.cutlass.compile");
+ ICHECK(pf.has_value()) << "The packed function contrib.cutlass.compile not found, please import "
+ "tvm.contrib.cutlass.build";
runtime::Module cutlass_mod = (*pf)(source_mod, options);
return {cutlass_mod};
diff --git a/src/relax/backend/contrib/dnnl/codegen.cc b/src/relax/backend/contrib/dnnl/codegen.cc
index 3e7cbdf079..d52823498c 100644
--- a/src/relax/backend/contrib/dnnl/codegen.cc
+++ b/src/relax/backend/contrib/dnnl/codegen.cc
@@ -89,10 +89,9 @@ Array<runtime::Module> DNNLCompiler(Array<Function> functions, Map<String, ffi::
serializer.serialize(func);
auto graph_json = serializer.GetJSON();
auto constant_names = serializer.GetConstantNames();
- const auto* pf = runtime::Registry::Get("runtime.DNNLJSONRuntimeCreate");
- ICHECK(pf != nullptr) << "Cannot find DNNL runtime module create function.";
+ const auto pf = tvm::ffi::Function::GetGlobalRequired("runtime.DNNLJSONRuntimeCreate");
auto func_name = GetExtSymbol(func);
- compiled_functions.push_back((*pf)(func_name, graph_json, constant_names));
+ compiled_functions.push_back(pf(func_name, graph_json, constant_names));
}
return compiled_functions;
diff --git a/src/relax/backend/contrib/hipblas/codegen.cc b/src/relax/backend/contrib/hipblas/codegen.cc
index d229336cc6..54260693f9 100644
--- a/src/relax/backend/contrib/hipblas/codegen.cc
+++ b/src/relax/backend/contrib/hipblas/codegen.cc
@@ -94,8 +94,7 @@ Array<runtime::Module> HipblasCompiler(Array<Function> functions, Map<String, ff
serializer.serialize(func);
auto graph_json = serializer.GetJSON();
auto constant_names = serializer.GetConstantNames();
- const auto* pf = runtime::Registry::Get("runtime.HipblasJSONRuntimeCreate");
- ICHECK(pf != nullptr) << "Cannot find HIPBLAS runtime module create function.";
+ const auto pf = tvm::ffi::Function::GetGlobalRequired("runtime.HipblasJSONRuntimeCreate");
auto func_name = GetExtSymbol(func);
compiled_functions.push_back((*pf)(func_name, graph_json, constant_names));
}
diff --git a/src/relax/backend/contrib/nnapi/codegen.cc b/src/relax/backend/contrib/nnapi/codegen.cc
index b120bad2f2..f0517fa359 100644
--- a/src/relax/backend/contrib/nnapi/codegen.cc
+++ b/src/relax/backend/contrib/nnapi/codegen.cc
@@ -256,10 +256,9 @@ Array<runtime::Module> NNAPICompiler(Array<Function> functions, Map<String, ffi:
serializer.serialize(func);
auto graph_json = serializer.GetJSON();
auto constant_names = serializer.GetConstantNames();
- const auto* pf = runtime::Registry::Get("runtime.nnapi_runtime_create");
- ICHECK(pf != nullptr) << "Cannot find NNAPI runtime module create function.";
+ const auto pf = tvm::ffi::Function::GetGlobalRequired("runtime.nnapi_runtime_create");
auto func_name = GetExtSymbol(func);
- compiled_functions.push_back((*pf)(func_name, graph_json, constant_names));
+ compiled_functions.push_back(pf(func_name, graph_json, constant_names));
}
return compiled_functions;
diff --git a/src/relax/backend/contrib/tensorrt/codegen.cc b/src/relax/backend/contrib/tensorrt/codegen.cc
index 12b99b8252..f458ff85ef 100644
--- a/src/relax/backend/contrib/tensorrt/codegen.cc
+++ b/src/relax/backend/contrib/tensorrt/codegen.cc
@@ -223,11 +223,10 @@ Array<runtime::Module> TensorRTCompiler(Array<Function> functions, Map<String, f
std::string graph_json = serializer.GetJSON();
VLOG(1) << "TensorRT JSON:" << std::endl << graph_json;
auto constant_names = serializer.GetConstantNames();
- const auto* pf = runtime::Registry::Get("runtime.tensorrt_runtime_create");
- ICHECK(pf != nullptr) << "Cannot find TensorRT runtime module create function.";
+ const auto pf = tvm::ffi::Function::GetGlobalRequired("runtime.tensorrt_runtime_create");
std::string func_name = GetExtSymbol(func);
VLOG(1) << "Creating tensorrt runtime::Module for '" << func_name << "'";
- compiled_functions.push_back((*pf)(func_name, graph_json, constant_names));
+ compiled_functions.push_back(pf(func_name, graph_json, constant_names));
}
return compiled_functions;
}
diff --git a/src/relax/backend/contrib/utils.h b/src/relax/backend/contrib/utils.h
index e63e99548c..bbff798b86 100644
--- a/src/relax/backend/contrib/utils.h
+++ b/src/relax/backend/contrib/utils.h
@@ -36,16 +36,6 @@ namespace tvm {
namespace relax {
namespace backend {
-/*!
- * \brief Get the Packed Func
- *
- * \param func_name
- * \return const PackedFunc*
- */
-inline const PackedFunc* GetPackedFunc(const std::string& func_name) {
- return tvm::runtime::Registry::Get(func_name);
-}
-
/*!
* \brief Extract shape from an IndexExpr array to std::vector<int64_t>
*
@@ -69,30 +59,7 @@ inline std::vector<int64_t> GetIntShape(const Array<PrimExpr>& shape) {
* \return std::string string format of type
*/
inline std::string DType2String(const tvm::DataType dtype) {
- std::ostringstream os;
- if (dtype.is_float()) {
- os << "float";
- } else if (dtype.is_float8_e4m3fn()) {
- return "float8_e4m3fn";
- } else if (dtype.is_float8_e5m2()) {
- return "float8_e5m2";
- } else if (dtype.is_float4_e2m1fn()) {
- return "float4_e2m1fn";
- } else if (dtype.is_int()) {
- os << "int";
- } else if (dtype.is_uint()) {
- os << "uint";
- } else if (dtype.is_bfloat16()) {
- os << "bfloat";
- } else if ((*GetPackedFunc("runtime._datatype_get_type_registered"))(dtype.code())) {
- os << "custom["
- << (*GetPackedFunc("runtime._datatype_get_type_name"))(dtype.code()).operator std::string()
- << "]";
- } else {
- LOG(FATAL) << "Unknown type with code " << static_cast<unsigned>(dtype.code());
- }
- os << dtype.bits();
- return os.str();
+ return tvm::ffi::DLDataTypeToString(dtype);
}
/*!
diff --git a/src/relax/backend/task_extraction.cc b/src/relax/backend/task_extraction.cc
index ebaee74f47..aeee9d2285 100644
--- a/src/relax/backend/task_extraction.cc
+++ b/src/relax/backend/task_extraction.cc
@@ -87,8 +87,8 @@ class TaskExtractor : public ExprVisitor {
target_(std::move(target)),
mod_eq_(ModuleEquality::Create(mod_eq_name)),
func2task_(/*bucket_count*/ 0, ModuleHash(*mod_eq_), ModuleEqual(*mod_eq_)) {
- normalize_mod_func_ = runtime::Registry::Get("tvm.meta_schedule.normalize_mod");
- ICHECK(normalize_mod_func_) << "Normalization function is not found.";
+ normalize_mod_func_ = tvm::ffi::Function::GetGlobal("tvm.meta_schedule.normalize_mod");
+ ICHECK(normalize_mod_func_.has_value()) << "Normalization function is not found.";
}
void VisitExpr_(const CallNode* call) final {
@@ -136,7 +136,7 @@ class TaskExtractor : public ExprVisitor {
Target target_;
std::unique_ptr<ModuleEquality> mod_eq_;
std::unordered_map<IRModule, ExtractedTask, ModuleHash, ModuleEqual> func2task_;
- const runtime::PackedFunc* normalize_mod_func_;
+ std::optional<tvm::ffi::Function> normalize_mod_func_;
};
TVM_REGISTER_GLOBAL("relax.backend.MetaScheduleExtractTask")
diff --git a/src/relax/backend/vm/exec_builder.cc b/src/relax/backend/vm/exec_builder.cc
index 573c27daee..e32c043b06 100644
--- a/src/relax/backend/vm/exec_builder.cc
+++ b/src/relax/backend/vm/exec_builder.cc
@@ -330,7 +330,7 @@ void ExecBuilderNode::Formalize() {
TVM_REGISTER_GLOBAL("relax.ExecBuilderCreate").set_body_typed(ExecBuilderNode::Create);
TVM_REGISTER_GLOBAL("relax.ExecBuilderConvertConstant")
- .set_body([](TVMArgs args, TVMRetValue* ret) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
ExecBuilder builder = args[0];
TVMRetValue rt;
rt = args[1];
@@ -343,8 +343,7 @@ TVM_REGISTER_GLOBAL("relax.ExecBuilderEmitFunction")
builder->EmitFunction(func, num_inputs, param_names);
});
-TVM_REGISTER_GLOBAL("relax.ExecBuilderEndFunction")
- .set_body_method<ExecBuilder>(&ExecBuilderNode::EndFunction);
+TVM_REGISTER_GLOBAL("relax.ExecBuilderEndFunction").set_body_method(&ExecBuilderNode::EndFunction);
TVM_REGISTER_GLOBAL("relax.ExecBuilderDeclareFunction")
.set_body_typed([](ExecBuilder builder, String name, int32_t kind) {
@@ -366,8 +365,7 @@ TVM_REGISTER_GLOBAL("relax.ExecBuilderEmitRet")
builder->EmitRet(Instruction::Arg::FromData(data));
});
-TVM_REGISTER_GLOBAL("relax.ExecBuilderEmitGoto")
- .set_body_method<ExecBuilder>(&ExecBuilderNode::EmitGoto);
+TVM_REGISTER_GLOBAL("relax.ExecBuilderEmitGoto").set_body_method(&ExecBuilderNode::EmitGoto);
TVM_REGISTER_GLOBAL("relax.ExecBuilderEmitIf")
.set_body_typed([](ExecBuilder builder, int64_t data, vm::Index false_offset) {
diff --git a/src/relax/ir/block_builder.cc b/src/relax/ir/block_builder.cc
index 8df9a67a26..5f56e7cf45 100644
--- a/src/relax/ir/block_builder.cc
+++ b/src/relax/ir/block_builder.cc
@@ -1059,16 +1059,14 @@ TVM_REGISTER_GLOBAL("relax.BlockBuilderCreate").set_body_typed([](Optional<IRMod
});
TVM_REGISTER_GLOBAL("relax.BlockBuilderBeginDataflowBlock")
- .set_body_method<BlockBuilder>(&BlockBuilderNode::BeginDataflowBlock);
+ .set_body_method(&BlockBuilderNode::BeginDataflowBlock);
TVM_REGISTER_GLOBAL("relax.BlockBuilderBeginBindingBlock")
- .set_body_method<BlockBuilder>(&BlockBuilderNode::BeginBindingBlock);
+ .set_body_method(&BlockBuilderNode::BeginBindingBlock);
-TVM_REGISTER_GLOBAL("relax.BlockBuilderEndBlock")
- .set_body_method<BlockBuilder>(&BlockBuilderNode::EndBlock);
+TVM_REGISTER_GLOBAL("relax.BlockBuilderEndBlock").set_body_method(&BlockBuilderNode::EndBlock);
-TVM_REGISTER_GLOBAL("relax.BlockBuilderNormalize")
- .set_body_method<BlockBuilder>(&BlockBuilderNode::Normalize);
+TVM_REGISTER_GLOBAL("relax.BlockBuilderNormalize").set_body_method(&BlockBuilderNode::Normalize);
TVM_REGISTER_GLOBAL("relax.BlockBuilderEmit")
.set_body_typed([](BlockBuilder builder, Expr expr, String name_hint) {
@@ -1097,27 +1095,24 @@ TVM_REGISTER_GLOBAL("relax.BlockBuilderGetUniqueName")
});
TVM_REGISTER_GLOBAL("relax.BlockBuilderAddFunction")
- .set_body_method<BlockBuilder>(&BlockBuilderNode::AddFunction);
+ .set_body_method(&BlockBuilderNode::AddFunction);
TVM_REGISTER_GLOBAL("relax.BlockBuilderUpdateFunction")
- .set_body_method<BlockBuilder>(&BlockBuilderNode::UpdateFunction);
+ .set_body_method(&BlockBuilderNode::UpdateFunction);
TVM_REGISTER_GLOBAL("relax.BlockBuilderGetContextIRModule")
- .set_body_method<BlockBuilder>(&BlockBuilderNode::GetContextIRModule);
+ .set_body_method(&BlockBuilderNode::GetContextIRModule);
-TVM_REGISTER_GLOBAL("relax.BlockBuilderFinalize")
- .set_body_method<BlockBuilder>(&BlockBuilderNode::Finalize);
+TVM_REGISTER_GLOBAL("relax.BlockBuilderFinalize").set_body_method(&BlockBuilderNode::Finalize);
TVM_REGISTER_GLOBAL("relax.BlockBuilderCurrentBlockIsDataFlow")
- .set_body_method<BlockBuilder>(&BlockBuilderNode::CurrentBlockIsDataFlow);
+ .set_body_method(&BlockBuilderNode::CurrentBlockIsDataFlow);
TVM_REGISTER_GLOBAL("relax.BlockBuilderLookupBinding")
- .set_body_method<BlockBuilder>(&BlockBuilderNode::LookupBinding);
+ .set_body_method(&BlockBuilderNode::LookupBinding);
-TVM_REGISTER_GLOBAL("relax.BlockBuilderBeginScope")
- .set_body_method<BlockBuilder>(&BlockBuilderNode::BeginScope);
+TVM_REGISTER_GLOBAL("relax.BlockBuilderBeginScope").set_body_method(&BlockBuilderNode::BeginScope);
-TVM_REGISTER_GLOBAL("relax.BlockBuilderEndScope")
- .set_body_method<BlockBuilder>(&BlockBuilderNode::EndScope);
+TVM_REGISTER_GLOBAL("relax.BlockBuilderEndScope").set_body_method(&BlockBuilderNode::EndScope);
} // namespace relax
} // namespace tvm
diff --git a/src/relax/transform/few_shot_tuning.cc b/src/relax/transform/few_shot_tuning.cc
index a8281f81c2..d7e0ead986 100644
--- a/src/relax/transform/few_shot_tuning.cc
+++ b/src/relax/transform/few_shot_tuning.cc
@@ -28,29 +28,24 @@ namespace transform {
tir::PrimFunc FewShotTunePrimFunc(const tir::PrimFunc& prim_func, const Target& target,
int64_t valid_count, bool benchmark) {
// fetch a local builder
- static const auto* f_get_local_builder =
- runtime::Registry::Get("meta_schedule.builder.get_local_builder");
- ICHECK(f_get_local_builder)
- << "ValueError: Cannot find the packed function \"meta_schedule.builder.get_local_builder\"";
- meta_schedule::Builder builder = (*f_get_local_builder)();
+ static const auto f_get_local_builder =
+ tvm::ffi::Function::GetGlobalRequired("meta_schedule.builder.get_local_builder");
+ meta_schedule::Builder builder = f_get_local_builder();
ICHECK(builder.defined()) << "ValueError: The local builder is not defined!";
// fetch a local runner
meta_schedule::Runner runner{nullptr};
if (benchmark) {
- static const auto* f_get_local_runner =
- runtime::Registry::Get("meta_schedule.runner.get_local_runner");
- ICHECK(f_get_local_runner) << "ValueError: Cannot find the packed function "
- "\"meta_schedule.builder.get_local_runner\"";
- runner = (*f_get_local_runner)();
+ static const auto f_get_local_runner =
+ tvm::ffi::Function::GetGlobalRequired("meta_schedule.runner.get_local_runner");
+ runner = f_get_local_runner();
ICHECK(runner.defined()) << "ValueError: The local runner is not defined!";
}
// create an IRModule
IRModule mod = IRModule(Map<GlobalVar, BaseFunc>(
{{GlobalVar("main"), WithAttr(prim_func, tvm::attr::kGlobalSymbol, String("main"))}}));
// fetch the number of physical cores
- static const auto* f_cpu_count = runtime::Registry::Get("meta_schedule.cpu_count");
- ICHECK(f_cpu_count) << "ValueError: Cannot find the packed function \"meta_schedule._cpu_count\"";
- int num_threads = (*f_cpu_count)(false);
+ static const auto f_cpu_count = tvm::ffi::Function::GetGlobalRequired("meta_schedule.cpu_count");
+ int num_threads = f_cpu_count(false);
// store the results
Array<IRModule> results;
std::vector<double> costs;
diff --git a/src/relax/transform/fold_constant.cc b/src/relax/transform/fold_constant.cc
index 18c1eadd8e..352092c433 100644
--- a/src/relax/transform/fold_constant.cc
+++ b/src/relax/transform/fold_constant.cc
@@ -115,10 +115,9 @@ class ConstantFolder : public ExprMutator {
// already scheduled to only work on GPU, we will need to skip this in the const folder for
// now
// TODO(Hongyi): further check and narrow the scope of foldable function
- auto* pf = runtime::Registry::Get("tir.build");
- ICHECK(pf != nullptr) << "Cannot find tir.build in registry";
+ const auto pf = tvm::ffi::Function::GetGlobalRequired("tir.build");
func = WithAttr(func, tvm::attr::kGlobalSymbol, String("tir_function"));
- runtime::Module rt_module = (*pf)(func, eval_cpu_target);
+ runtime::Module rt_module = pf(func, eval_cpu_target);
build_func = rt_module.GetFunction("tir_function");
} catch (const tvm::Error& err) {
// build failure may happen in which case we skip
@@ -294,9 +293,8 @@ class ConstantFolder : public ExprMutator {
is_known &= (val.dtype() == DataType::Int(64));
}
if (is_known) {
- const auto* func = tvm::runtime::Registry::Get("relax.run.shape_to_tensor");
- ICHECK(func != nullptr);
- runtime::NDArray vals = (*func)(arr);
+ const auto func = tvm::ffi::Function::GetGlobalRequired("relax.run.shape_to_tensor");
+ runtime::NDArray vals = func(arr);
return Constant(vals);
}
}
diff --git a/src/relax/transform/gradient.cc b/src/relax/transform/gradient.cc
index e0dd2d1a1f..1100e2d457 100644
--- a/src/relax/transform/gradient.cc
+++ b/src/relax/transform/gradient.cc
@@ -365,10 +365,10 @@ class BackwardBindingGenerator : private ExprVisitor {
} else if (call_op == Op::Get("relax.call_tir_with_grad")) {
// tir gradient registering
auto te_grad_name = call->attrs.as<CallTIRWithGradAttrs>()->te_grad_name;
- auto* grad_func = tvm::runtime::Registry::Get(te_grad_func_prefix + te_grad_name);
- CHECK(grad_func) << "TIR gradient function " << te_grad_name << " is not registered";
+ const auto grad_func =
+ tvm::ffi::Function::GetGlobalRequired(te_grad_func_prefix + te_grad_name);
Var partials =
- (*grad_func)(checkpoint_var, Downcast<Call>(checkpoint_call), adjoint_var, builder_);
+ grad_func(checkpoint_var, Downcast<Call>(checkpoint_call), adjoint_var, builder_);
Tuple args = Downcast<Tuple>(call->args[1]);
auto* tuple_sinfo = GetStructInfoAs<TupleStructInfoNode>(partials);
if (!tuple_sinfo) {
diff --git a/src/relax/transform/meta_schedule.cc b/src/relax/transform/meta_schedule.cc
index eb1451147a..a1fc299ec3 100644
--- a/src/relax/transform/meta_schedule.cc
+++ b/src/relax/transform/meta_schedule.cc
@@ -44,10 +44,8 @@ class MetaScheduleTuner {
max_trials_per_task_(max_trials_per_task),
op_names_(op_names),
params_(params) {
- // candgen_func_ = runtime::Registry::Get("relax.tuning_api.default_generate_candidate");
- // ICHECK(candgen_func_) << "Default candidate generation function is not found.";
- normalize_mod_func_ = runtime::Registry::Get("tvm.meta_schedule.normalize_mod");
- ICHECK(normalize_mod_func_) << "Normalization function is not found.";
+ normalize_mod_func_ = tvm::ffi::Function::GetGlobal("tvm.meta_schedule.normalize_mod");
+ ICHECK(normalize_mod_func_.has_value()) << "Normalization function is not found.";
}
// TODO(@sunggg): Currently, only supports basic arguments.
@@ -96,16 +94,15 @@ class MetaScheduleTuner {
Integer max_trials_per_task_;
Optional<Array<String>> op_names_;
Map<String, runtime::NDArray> params_;
- // const runtime::PackedFunc* candgen_func_;
- const runtime::PackedFunc* normalize_mod_func_;
+ std::optional<tvm::ffi::Function> normalize_mod_func_;
};
Pass MetaScheduleApplyDatabase(Optional<String> work_dir, bool enable_warning = false) {
using tvm::meta_schedule::Database;
Target target = Target::Current(false);
- const runtime::PackedFunc* normalize_mod_func_ =
- runtime::Registry::Get("tvm.meta_schedule.normalize_mod");
- ICHECK(normalize_mod_func_) << "Normalization function is not found.";
+ const std::optional<tvm::ffi::Function> normalize_mod_func_ =
+ tvm::ffi::Function::GetGlobalRequired("tvm.meta_schedule.normalize_mod");
+ ICHECK(normalize_mod_func_.has_value()) << "Normalization function is not found.";
auto pass_func = [=](IRModule mod, PassContext ctx) {
Database database{nullptr};
diff --git a/src/relax/transform/run_codegen.cc b/src/relax/transform/run_codegen.cc
index b385297007..73949a19f8 100644
--- a/src/relax/transform/run_codegen.cc
+++ b/src/relax/transform/run_codegen.cc
@@ -128,9 +128,8 @@ class CodeGenRunner : ExprMutator {
extern_funcs_[gvar_node] = new_func;
// Remove the global symbol and codegen attributes from the function so that it can be
// removed the module.
- static const runtime::PackedFunc* RemoveFuncAttrFunc =
- runtime::Registry::Get("ir.BaseFuncWithoutAttr");
- ICHECK(RemoveFuncAttrFunc);
+ const auto RemoveFuncAttrFunc = tvm::ffi::Function::GetGlobal("ir.BaseFuncWithoutAttr");
+ ICHECK(RemoveFuncAttrFunc.has_value());
func = (*RemoveFuncAttrFunc)(func, tvm::attr::kGlobalSymbol);
func = (*RemoveFuncAttrFunc)(func, attr::kCodegen);
builder_->UpdateFunction(gvar, func);
@@ -192,8 +191,8 @@ class CodeGenRunner : ExprMutator {
// Start the codegen process.
// Get the codegen with its ffi key.
String codegen_name = "relax.ext." + target;
- auto codegen = runtime::Registry::Get(codegen_name);
- ICHECK(codegen) << "Codegen is not found: " << codegen_name << "\n";
+ const auto codegen = tvm::ffi::Function::GetGlobal(codegen_name);
+ ICHECK(codegen.has_value()) << "Codegen is not found: " << codegen_name << "\n";
Array<runtime::Module> compiled_functions = (*codegen)(functions, options, constant_names);
ext_mods.insert(ext_mods.end(), compiled_functions.begin(), compiled_functions.end());
diff --git a/src/relax/transform/tuning_api/database.cc b/src/relax/transform/tuning_api/database.cc
index 3bcb1b6da6..cccd8e7f87 100644
--- a/src/relax/transform/tuning_api/database.cc
+++ b/src/relax/transform/tuning_api/database.cc
@@ -324,26 +324,25 @@ TVM_REGISTER_GLOBAL("relax.tuning_api.TuningRecord")
return TuningRecord(trace, run_secs);
});
TVM_REGISTER_GLOBAL("relax.tuning_api.TuningRecordAsJSON")
- .set_body_method<TuningRecord>(&TuningRecordNode::AsJSON);
+ .set_body_method(&TuningRecordNode::AsJSON);
TVM_REGISTER_GLOBAL("relax.tuning_api.TuningRecordFromJSON").set_body_typed(TuningRecord::FromJSON);
TVM_REGISTER_OBJECT_TYPE(DatabaseNode);
TVM_REGISTER_GLOBAL("relax.tuning_api.DatabaseHasWorkload")
- .set_body_method<Database>(&DatabaseNode::HasWorkload);
+ .set_body_method(&DatabaseNode::HasWorkload);
TVM_REGISTER_GLOBAL("relax.tuning_api.DatabaseHasMeasurementRecord")
- .set_body_method<Database>(&DatabaseNode::HasMeasurementRecord);
+ .set_body_method(&DatabaseNode::HasMeasurementRecord);
TVM_REGISTER_GLOBAL("relax.tuning_api.DatabaseHasTuningRecord")
- .set_body_method<Database>(&DatabaseNode::HasTuningRecord);
+ .set_body_method(&DatabaseNode::HasTuningRecord);
TVM_REGISTER_GLOBAL("relax.tuning_api.DatabaseCommitMeasurementRecord")
- .set_body_method<Database>(&DatabaseNode::CommitMeasurementRecord);
+ .set_body_method(&DatabaseNode::CommitMeasurementRecord);
TVM_REGISTER_GLOBAL("relax.tuning_api.DatabaseCommitWorkload")
- .set_body_method<Database>(&DatabaseNode::CommitWorkload);
+ .set_body_method(&DatabaseNode::CommitWorkload);
TVM_REGISTER_GLOBAL("relax.tuning_api.DatabaseCommitTuningRecord")
- .set_body_method<Database>(&DatabaseNode::CommitTuningRecord);
-TVM_REGISTER_GLOBAL("relax.tuning_api.DatabaseGetTopK")
- .set_body_method<Database>(&DatabaseNode::GetTopK);
+ .set_body_method(&DatabaseNode::CommitTuningRecord);
+TVM_REGISTER_GLOBAL("relax.tuning_api.DatabaseGetTopK").set_body_method(&DatabaseNode::GetTopK);
TVM_REGISTER_GLOBAL("relax.tuning_api.DatabaseGetMeasurementRecord")
- .set_body_method<Database>(&DatabaseNode::GetMeasurementRecord);
+ .set_body_method(&DatabaseNode::GetMeasurementRecord);
TVM_REGISTER_NODE_TYPE(JSONDatabaseNode);
TVM_REGISTER_GLOBAL("relax.tuning_api.DatabaseJSONDatabase").set_body_typed(Database::JSONDatabase);
diff --git a/src/relax/transform/tuning_api/primitives.cc b/src/relax/transform/tuning_api/primitives.cc
index a19022344b..9d60a3e185 100644
--- a/src/relax/transform/tuning_api/primitives.cc
+++ b/src/relax/transform/tuning_api/primitives.cc
@@ -148,9 +148,9 @@ Trace::Trace() { data_ = make_object<TraceNode>(); }
Trace::Trace(IRModule in_mod, Array<Knob> knobs, Array<String> decisions) {
ICHECK(knobs.size() == decisions.size()) << "Size of knobs and decisions should match";
// Deep-copy IRModule
- auto func_deepcopy = runtime::Registry::Get("relax.tuning_api.deepcopy_irmodule");
- ICHECK(func_deepcopy);
- IRModule out_mod = (*func_deepcopy)(in_mod);
+ const auto func_deepcopy =
+ tvm::ffi::Function::GetGlobalRequired("relax.tuning_api.deepcopy_irmodule");
+ IRModule out_mod = func_deepcopy(in_mod);
// Apply the decision history if provided
int size = knobs.size();
for (int i = 0; i < size; i++) {
@@ -236,38 +236,36 @@ TVM_REGISTER_GLOBAL("relax.tuning_api.Choice")
String constr_func_key, Array<Any> constr_func_args) {
return Choice(transform_func_key, transform_func_args, constr_func_key, constr_func_args);
});
-TVM_REGISTER_GLOBAL("relax.tuning_api.ChoiceAsJSON").set_body_method<Choice>(&ChoiceNode::AsJSON);
+TVM_REGISTER_GLOBAL("relax.tuning_api.ChoiceAsJSON").set_body_method(&ChoiceNode::AsJSON);
TVM_REGISTER_GLOBAL("relax.tuning_api.ChoiceFromJSON").set_body_typed(Choice::FromJSON);
TVM_REGISTER_GLOBAL("relax.tuning_api.ChoiceGetTransformFunc")
- .set_body_method<Choice>(&ChoiceNode::GetTransformFunc);
+ .set_body_method(&ChoiceNode::GetTransformFunc);
TVM_REGISTER_GLOBAL("relax.tuning_api.ChoiceGetConstrFunc")
- .set_body_method<Choice>(&ChoiceNode::GetConstrFunc);
+ .set_body_method(&ChoiceNode::GetConstrFunc);
TVM_REGISTER_GLOBAL("relax.tuning_api.ChoiceApplyTransformFunc")
- .set_body_method<Choice>(&ChoiceNode::ApplyTransformFunc);
-TVM_REGISTER_GLOBAL("relax.tuning_api.ChoiceCheckConstr")
- .set_body_method<Choice>(&ChoiceNode::CheckConstr);
+ .set_body_method(&ChoiceNode::ApplyTransformFunc);
+TVM_REGISTER_GLOBAL("relax.tuning_api.ChoiceCheckConstr").set_body_method(&ChoiceNode::CheckConstr);
TVM_REGISTER_NODE_TYPE(KnobNode);
TVM_REGISTER_GLOBAL("relax.tuning_api.Knob")
.set_body_typed([](String name, Map<String, Choice> choices) { return Knob(name, choices); });
-TVM_REGISTER_GLOBAL("relax.tuning_api.KnobAsJSON").set_body_method<Knob>(&KnobNode::AsJSON);
+TVM_REGISTER_GLOBAL("relax.tuning_api.KnobAsJSON").set_body_method(&KnobNode::AsJSON);
TVM_REGISTER_GLOBAL("relax.tuning_api.KnobFromJSON").set_body_typed(Knob::FromJSON);
TVM_REGISTER_GLOBAL("relax.tuning_api.KnobIsValidDecision")
- .set_body_method<Knob>(&KnobNode::IsValidDecision);
-TVM_REGISTER_GLOBAL("relax.tuning_api.KnobApply").set_body_method<Knob>(&KnobNode::Apply);
+ .set_body_method(&KnobNode::IsValidDecision);
+TVM_REGISTER_GLOBAL("relax.tuning_api.KnobApply").set_body_method(&KnobNode::Apply);
TVM_REGISTER_NODE_TYPE(TraceNode);
TVM_REGISTER_GLOBAL("relax.tuning_api.Trace")
.set_body_typed([](IRModule in_mod, Array<Knob> knobs, Array<String> decisions) {
return Trace(in_mod, knobs, decisions);
});
-TVM_REGISTER_GLOBAL("relax.tuning_api.TraceVerify").set_body_method<Trace>(&TraceNode::Verify);
-TVM_REGISTER_GLOBAL("relax.tuning_api.TraceAdd").set_body_method<Trace>(&TraceNode::Add);
-TVM_REGISTER_GLOBAL("relax.tuning_api.TraceSetPerf").set_body_method<Trace>(&TraceNode::SetPerf);
-TVM_REGISTER_GLOBAL("relax.tuning_api.TraceSetOutMod")
- .set_body_method<Trace>(&TraceNode::SetOutMod);
+TVM_REGISTER_GLOBAL("relax.tuning_api.TraceVerify").set_body_method(&TraceNode::Verify);
+TVM_REGISTER_GLOBAL("relax.tuning_api.TraceAdd").set_body_method(&TraceNode::Add);
+TVM_REGISTER_GLOBAL("relax.tuning_api.TraceSetPerf").set_body_method(&TraceNode::SetPerf);
+TVM_REGISTER_GLOBAL("relax.tuning_api.TraceSetOutMod").set_body_method(&TraceNode::SetOutMod);
-TVM_REGISTER_GLOBAL("relax.tuning_api.TraceAsJSON").set_body_method<Trace>(&TraceNode::AsJSON);
+TVM_REGISTER_GLOBAL("relax.tuning_api.TraceAsJSON").set_body_method(&TraceNode::AsJSON);
TVM_REGISTER_GLOBAL("relax.tuning_api.TraceFromJSON").set_body_typed(Trace::FromJSON);
} // namespace relax
} // namespace tvm
diff --git a/src/runtime/c_runtime_api.cc b/src/runtime/c_runtime_api.cc
index c1a3662eaf..9f17a03216 100644
--- a/src/runtime/c_runtime_api.cc
+++ b/src/runtime/c_runtime_api.cc
@@ -46,21 +46,18 @@ namespace tvm {
namespace runtime {
std::string GetCustomTypeName(uint8_t type_code) {
- auto f = tvm::runtime::Registry::Get("runtime._datatype_get_type_name");
- ICHECK(f) << "Function runtime._datatype_get_type_name not found";
- return (*f)(type_code).operator String();
+ const auto f = tvm::ffi::Function::GetGlobalRequired("runtime._datatype_get_type_name");
+ return f(type_code).operator String();
}
uint8_t GetCustomTypeCode(const std::string& type_name) {
- auto f = tvm::runtime::Registry::Get("runtime._datatype_get_type_code");
- ICHECK(f) << "Function runtime._datatype_get_type_code not found";
- return (*f)(type_name).operator int();
+ const auto f = tvm::ffi::Function::GetGlobalRequired("runtime._datatype_get_type_code");
+ return f(type_name).operator int();
}
bool GetCustomTypeRegistered(uint8_t type_code) {
- auto f = tvm::runtime::Registry::Get("runtime._datatype_get_type_registered");
- ICHECK(f) << "Function runtime._datatype_get_type_registered not found";
- return (*f)(type_code).operator bool();
+ const auto f = tvm::ffi::Function::GetGlobalRequired("runtime._datatype_get_type_registered");
+ return f(type_code).operator bool();
}
uint8_t ParseCustomDatatype(const std::string& s, const char** scan) {
@@ -129,8 +126,8 @@ class DeviceAPIManager {
}
DeviceAPI* GetAPI(const std::string name, bool allow_missing) {
std::string factory = "device_api." + name;
- auto* f = Registry::Get(factory);
- if (f == nullptr) {
+ const auto f = tvm::ffi::Function::GetGlobal(factory);
+ if (!f.has_value()) {
ICHECK(allow_missing) << "Device API " << name << " is not enabled.";
return nullptr;
}
@@ -749,7 +746,7 @@ int TVMDeviceCopyDataFromTo(DLTensor* from, DLTensor* to, TVMStreamHandle stream
// set device api
TVM_REGISTER_GLOBAL(tvm::runtime::symbol::tvm_set_device)
- .set_body([](TVMArgs args, TVMRetValue* ret) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
DLDevice dev;
dev.device_type = static_cast<DLDeviceType>(args[0].operator int());
dev.device_id = args[1];
@@ -757,7 +754,7 @@ TVM_REGISTER_GLOBAL(tvm::runtime::symbol::tvm_set_device)
});
// set device api
-TVM_REGISTER_GLOBAL("runtime.GetDeviceAttr").set_body([](TVMArgs args, TVMRetValue* ret) {
+TVM_REGISTER_GLOBAL("runtime.GetDeviceAttr").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
DLDevice dev;
dev.device_type = static_cast<DLDeviceType>(args[0].operator int());
dev.device_id = args[1];
diff --git a/src/runtime/contrib/amx/amx_config.cc b/src/runtime/contrib/amx/amx_config.cc
index 8aae82fbdf..a3a216b993 100644
--- a/src/runtime/contrib/amx/amx_config.cc
+++ b/src/runtime/contrib/amx/amx_config.cc
@@ -76,7 +76,7 @@ void init_tile_config(__tilecfg_u* dst, uint16_t cols, uint8_t rows) {
_tile_loadconfig(dst->a);
}
-TVM_REGISTER_GLOBAL("runtime.amx_tileconfig").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("runtime.amx_tileconfig").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
int rows = args[0];
int cols = args[1];
LOG(INFO) << "rows: " << rows << ", cols:" << cols;
@@ -89,7 +89,7 @@ TVM_REGISTER_GLOBAL("runtime.amx_tileconfig").set_body([](TVMArgs args, TVMRetVa
});
// register a global packed function in c++，to init the system for AMX config
-TVM_REGISTER_GLOBAL("runtime.amx_init").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("runtime.amx_init").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
// -----------Detect and request for AMX control----------------------
uint64_t bitmask = 0;
int64_t status = syscall(SYS_arch_prctl, ARCH_GET_XCOMP_PERM, &bitmask);
diff --git a/src/runtime/contrib/cblas/cblas.cc b/src/runtime/contrib/cblas/cblas.cc
index fbac622248..9d7a86e09d 100644
--- a/src/runtime/contrib/cblas/cblas.cc
+++ b/src/runtime/contrib/cblas/cblas.cc
@@ -123,7 +123,7 @@ struct CblasDgemmBatchIterativeOp {
};
// matrix multiplication for row major
-TVM_REGISTER_GLOBAL("tvm.contrib.cblas.matmul").set_body([](TVMArgs args, TVMRetValue* ret) {
+TVM_REGISTER_GLOBAL("tvm.contrib.cblas.matmul").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
DLTensor* A = args[0];
ICHECK(TypeMatch(A->dtype, kDLFloat, 32) || TypeMatch(A->dtype, kDLFloat, 64));
@@ -133,18 +133,19 @@ TVM_REGISTER_GLOBAL("tvm.contrib.cblas.matmul").set_body([](TVMArgs args, TVMRet
CallGemm(args, ret, CblasDgemmOp());
});
-TVM_REGISTER_GLOBAL("tvm.contrib.cblas.batch_matmul").set_body([](TVMArgs args, TVMRetValue* ret) {
- DLTensor* A = args[0];
- ICHECK(TypeMatch(A->dtype, kDLFloat, 32) || TypeMatch(A->dtype, kDLFloat, 64));
- if (TypeMatch(A->dtype, kDLFloat, 32)) {
- CallBatchGemm(args, ret, CblasSgemmBatchOp());
- } else {
- CallBatchGemm(args, ret, CblasDgemmBatchOp());
- }
-});
+TVM_REGISTER_GLOBAL("tvm.contrib.cblas.batch_matmul")
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
+ DLTensor* A = args[0];
+ ICHECK(TypeMatch(A->dtype, kDLFloat, 32) || TypeMatch(A->dtype, kDLFloat, 64));
+ if (TypeMatch(A->dtype, kDLFloat, 32)) {
+ CallBatchGemm(args, ret, CblasSgemmBatchOp());
+ } else {
+ CallBatchGemm(args, ret, CblasDgemmBatchOp());
+ }
+ });
TVM_REGISTER_GLOBAL("tvm.contrib.cblas.batch_matmul_iterative")
- .set_body([](TVMArgs args, TVMRetValue* ret) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
DLTensor* A = args[0];
ICHECK(TypeMatch(A->dtype, kDLFloat, 32) || TypeMatch(A->dtype, kDLFloat, 64));
if (TypeMatch(A->dtype, kDLFloat, 32)) {
diff --git a/src/runtime/contrib/cblas/dnnl_blas.cc b/src/runtime/contrib/cblas/dnnl_blas.cc
index 43d5ab35b4..6b9391f9c4 100644
--- a/src/runtime/contrib/cblas/dnnl_blas.cc
+++ b/src/runtime/contrib/cblas/dnnl_blas.cc
@@ -46,7 +46,7 @@ struct DNNLSgemmOp {
};
// matrix multiplication for row major
-TVM_REGISTER_GLOBAL("tvm.contrib.dnnl.matmul").set_body([](TVMArgs args, TVMRetValue* ret) {
+TVM_REGISTER_GLOBAL("tvm.contrib.dnnl.matmul").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
DLTensor* A = args[0];
ICHECK(TypeMatch(A->dtype, kDLFloat, 32));
CallGemm(args, ret, DNNLSgemmOp());
diff --git a/src/runtime/contrib/cblas/mkl.cc b/src/runtime/contrib/cblas/mkl.cc
index 4323878db2..6eac35418d 100644
--- a/src/runtime/contrib/cblas/mkl.cc
+++ b/src/runtime/contrib/cblas/mkl.cc
@@ -154,7 +154,7 @@ struct MKLDgemmBatchIterativeOp {
};
// matrix multiplication for row major
-TVM_REGISTER_GLOBAL("tvm.contrib.mkl.matmul").set_body([](TVMArgs args, TVMRetValue* ret) {
+TVM_REGISTER_GLOBAL("tvm.contrib.mkl.matmul").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
DLTensor* A = args[0];
ICHECK(TypeMatch(A->dtype, kDLFloat, 32) || TypeMatch(A->dtype, kDLFloat, 64));
@@ -165,28 +165,30 @@ TVM_REGISTER_GLOBAL("tvm.contrib.mkl.matmul").set_body([](TVMArgs args, TVMRetVa
});
// integer matrix multiplication for row major
-TVM_REGISTER_GLOBAL("tvm.contrib.mkl.matmul_u8s8s32").set_body([](TVMArgs args, TVMRetValue* ret) {
- DLTensor* A = args[0];
- DLTensor* B = args[1];
- DLTensor* C = args[2];
- ICHECK(TypeMatch(A->dtype, kDLUInt, 8) && TypeMatch(B->dtype, kDLInt, 8) &&
- TypeMatch(C->dtype, kDLInt, 32));
+TVM_REGISTER_GLOBAL("tvm.contrib.mkl.matmul_u8s8s32")
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
+ DLTensor* A = args[0];
+ DLTensor* B = args[1];
+ DLTensor* C = args[2];
+ ICHECK(TypeMatch(A->dtype, kDLUInt, 8) && TypeMatch(B->dtype, kDLInt, 8) &&
+ TypeMatch(C->dtype, kDLInt, 32));
- CallU8S8S32Gemm(args, ret, MKLGemmU8S8S32Op());
-});
+ CallU8S8S32Gemm(args, ret, MKLGemmU8S8S32Op());
+ });
-TVM_REGISTER_GLOBAL("tvm.contrib.mkl.batch_matmul").set_body([](TVMArgs args, TVMRetValue* ret) {
- DLTensor* A = args[0];
- ICHECK(TypeMatch(A->dtype, kDLFloat, 32) || TypeMatch(A->dtype, kDLFloat, 64));
- if (TypeMatch(A->dtype, kDLFloat, 32)) {
- CallBatchGemm(args, ret, MKLSgemmBatchOp());
- } else {
- CallBatchGemm(args, ret, MKLDgemmBatchOp());
- }
-});
+TVM_REGISTER_GLOBAL("tvm.contrib.mkl.batch_matmul")
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
+ DLTensor* A = args[0];
+ ICHECK(TypeMatch(A->dtype, kDLFloat, 32) || TypeMatch(A->dtype, kDLFloat, 64));
+ if (TypeMatch(A->dtype, kDLFloat, 32)) {
+ CallBatchGemm(args, ret, MKLSgemmBatchOp());
+ } else {
+ CallBatchGemm(args, ret, MKLDgemmBatchOp());
+ }
+ });
TVM_REGISTER_GLOBAL("tvm.contrib.mkl.batch_matmul_iterative")
- .set_body([](TVMArgs args, TVMRetValue* ret) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
DLTensor* A = args[0];
ICHECK(TypeMatch(A->dtype, kDLFloat, 32) || TypeMatch(A->dtype, kDLFloat, 64));
if (TypeMatch(A->dtype, kDLFloat, 32)) {
diff --git a/src/runtime/contrib/clml/clml_runtime.cc b/src/runtime/contrib/clml/clml_runtime.cc
index 4998f2b476..21584997a7 100644
--- a/src/runtime/contrib/clml/clml_runtime.cc
+++ b/src/runtime/contrib/clml/clml_runtime.cc
@@ -311,8 +311,8 @@ class CLMLRuntime : public JSONRuntimeBase {
dump_tensors.Set(node_name, narr);
}
- const PackedFunc* f = Registry::Get("runtime.SaveParams");
- if (nullptr != f) {
+ const auto f = tvm::ffi::Function::GetGlobal("runtime.SaveParams");
+ if (f.has_value()) {
std::string dump_bytes = (*f)(dump_tensors);
std::ostringstream oss;
/*TODO(Siva) HEX encoding doubles the size, look for better encode that can cross the RPC. */
@@ -478,7 +478,7 @@ class CLMLRuntime : public JSONRuntimeBase {
LOG_CLML << "Execution by Rec Queue";
if (cws->workspace->IsProfiling(cws->tentry->device)) {
Timer t;
- auto f = Registry::Get(std::string("profiling.timer.opencl"));
+ auto f = tvm::ffi::Function::GetGlobal(std::string("profiling.timer.opencl"));
t = f->operator()(cws->tentry->device);
t->Start();
queue = CLML_QUEUE;
@@ -499,7 +499,7 @@ class CLMLRuntime : public JSONRuntimeBase {
LOG_CLML << "Run Layer:" << this->layer_.layer_names[i];
if (cws->workspace->IsProfiling(cws->tentry->device)) {
Timer t;
- auto f = Registry::Get(std::string("profiling.timer.opencl"));
+ auto f = tvm::ffi::Function::GetGlobal(std::string("profiling.timer.opencl"));
t = f->operator()(cws->tentry->device);
t->Start();
queue = CLML_QUEUE;
diff --git a/src/runtime/contrib/coreml/coreml_runtime.mm b/src/runtime/contrib/coreml/coreml_runtime.mm
index ab1372e228..272e871ead 100644
--- a/src/runtime/contrib/coreml/coreml_runtime.mm
+++ b/src/runtime/contrib/coreml/coreml_runtime.mm
@@ -189,7 +189,7 @@ Module CoreMLRuntimeCreate(const std::string& symbol, const std::string& model_p
return Module(exec);
}
-TVM_REGISTER_GLOBAL("tvm.coreml_runtime.create").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("tvm.coreml_runtime.create").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = CoreMLRuntimeCreate(args[0], args[1]);
});
diff --git a/src/runtime/contrib/cublas/cublas.cc b/src/runtime/contrib/cublas/cublas.cc
index 8272531774..f5d72db859 100644
--- a/src/runtime/contrib/cublas/cublas.cc
+++ b/src/runtime/contrib/cublas/cublas.cc
@@ -513,69 +513,71 @@ inline void CallBatchGemmEx(TVMArgs args, TVMRetValue* ret, cublasHandle_t hdl)
}
// matrix multiplication for row major
-TVM_REGISTER_GLOBAL("tvm.contrib.cublas.matmul").set_body([](TVMArgs args, TVMRetValue* ret) {
- DLTensor* A = args[0];
- DLTensor* C = args[2];
-
- CuBlasThreadEntry* entry_ptr = CuBlasThreadEntry::ThreadLocal();
-
- CUBLASTryEnableTensorCore(entry_ptr->handle);
-
- if (TypeEqual(A->dtype, C->dtype)) {
- ICHECK(TypeMatch(A->dtype, kDLFloat, 16) || TypeMatch(A->dtype, kDLFloat, 32) ||
- TypeMatch(A->dtype, kDLFloat, 64));
-
- if (TypeMatch(A->dtype, kDLFloat, 16))
- CallGemm(args, ret, CublasHgemmOp(entry_ptr->handle));
- else if (TypeMatch(A->dtype, kDLFloat, 32))
- CallGemm(args, ret, CublasSgemmOp(entry_ptr->handle));
- else
- CallGemm(args, ret, CublasDgemmOp(entry_ptr->handle));
- } else {
- CallGemmEx(args, ret, entry_ptr->handle);
- }
-});
+TVM_REGISTER_GLOBAL("tvm.contrib.cublas.matmul")
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
+ DLTensor* A = args[0];
+ DLTensor* C = args[2];
+
+ CuBlasThreadEntry* entry_ptr = CuBlasThreadEntry::ThreadLocal();
+
+ CUBLASTryEnableTensorCore(entry_ptr->handle);
+
+ if (TypeEqual(A->dtype, C->dtype)) {
+ ICHECK(TypeMatch(A->dtype, kDLFloat, 16) || TypeMatch(A->dtype, kDLFloat, 32) ||
+ TypeMatch(A->dtype, kDLFloat, 64));
+
+ if (TypeMatch(A->dtype, kDLFloat, 16))
+ CallGemm(args, ret, CublasHgemmOp(entry_ptr->handle));
+ else if (TypeMatch(A->dtype, kDLFloat, 32))
+ CallGemm(args, ret, CublasSgemmOp(entry_ptr->handle));
+ else
+ CallGemm(args, ret, CublasDgemmOp(entry_ptr->handle));
+ } else {
+ CallGemmEx(args, ret, entry_ptr->handle);
+ }
+ });
#if CUDART_VERSION >= 10010
-TVM_REGISTER_GLOBAL("tvm.contrib.cublaslt.matmul").set_body([](TVMArgs args, TVMRetValue* ret) {
- DLTensor* A = args[0];
-
- CuBlasThreadEntry* entry_ptr = CuBlasThreadEntry::ThreadLocal();
-
- CUBLASTryEnableTensorCore(entry_ptr->handle);
-
- ICHECK(TypeMatch(A->dtype, kDLInt, 8)) << "Expects dtype to be int8\n";
- cublasLtHandle_t ltHandle;
- CHECK_CUBLAS_ERROR(cublasLtCreate(<Handle));
- auto func = tvm::runtime::Registry::Get("runtime.get_cuda_stream");
- ICHECK(func != nullptr);
- cudaStream_t stream = static_cast<cudaStream_t>((*func)().operator void*());
- CallLtIgemm(args, ret, ltHandle, stream);
- CHECK_CUBLAS_ERROR(cublasLtDestroy(ltHandle));
-});
+TVM_REGISTER_GLOBAL("tvm.contrib.cublaslt.matmul")
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
+ DLTensor* A = args[0];
+
+ CuBlasThreadEntry* entry_ptr = CuBlasThreadEntry::ThreadLocal();
+
+ CUBLASTryEnableTensorCore(entry_ptr->handle);
+
+ ICHECK(TypeMatch(A->dtype, kDLInt, 8)) << "Expects dtype to be int8\n";
+ cublasLtHandle_t ltHandle;
+ CHECK_CUBLAS_ERROR(cublasLtCreate(<Handle));
+ auto func = tvm::ffi::Function::GetGlobalRequired("runtime.get_cuda_stream");
+ cudaStream_t stream = static_cast<cudaStream_t>(func().operator void*());
+ CallLtIgemm(args, ret, ltHandle, stream);
+ CHECK_CUBLAS_ERROR(cublasLtDestroy(ltHandle));
+ });
#endif // CUDART_VERSION >= 10010
-TVM_REGISTER_GLOBAL("tvm.contrib.cublas.batch_matmul").set_body([](TVMArgs args, TVMRetValue* ret) {
- DLTensor* A = args[0];
- DLTensor* C = args[2];
-
- CuBlasThreadEntry* entry_ptr = CuBlasThreadEntry::ThreadLocal();
-
- CUBLASTryEnableTensorCore(entry_ptr->handle);
- if (TypeEqual(A->dtype, C->dtype)) {
- ICHECK(TypeMatch(A->dtype, kDLFloat, 16) || TypeMatch(A->dtype, kDLFloat, 32) ||
- TypeMatch(A->dtype, kDLFloat, 64));
-
- if (TypeMatch(A->dtype, kDLFloat, 16))
- CallBatchGemm(args, ret, CublasHgemmBatchOp(entry_ptr->handle));
- else if (TypeMatch(A->dtype, kDLFloat, 32))
- CallBatchGemm(args, ret, CublasSgemmBatchOp(entry_ptr->handle));
- else
- CallBatchGemm(args, ret, CublasDgemmBatchOp(entry_ptr->handle));
- } else {
- CallBatchGemmEx(args, ret, entry_ptr->handle);
- }
-});
+TVM_REGISTER_GLOBAL("tvm.contrib.cublas.batch_matmul")
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
+ DLTensor* A = args[0];
+ DLTensor* C = args[2];
+
+ CuBlasThreadEntry* entry_ptr = CuBlasThreadEntry::ThreadLocal();
+
+ CUBLASTryEnableTensorCore(entry_ptr->handle);
+ if (TypeEqual(A->dtype, C->dtype)) {
+ ICHECK(TypeMatch(A->dtype, kDLFloat, 16) || TypeMatch(A->dtype, kDLFloat, 32) ||
+ TypeMatch(A->dtype, kDLFloat, 64));
+
+ if (TypeMatch(A->dtype, kDLFloat, 16))
+ CallBatchGemm(args, ret, CublasHgemmBatchOp(entry_ptr->handle));
+ else if (TypeMatch(A->dtype, kDLFloat, 32))
+ CallBatchGemm(args, ret, CublasSgemmBatchOp(entry_ptr->handle));
+ else
+ CallBatchGemm(args, ret, CublasDgemmBatchOp(entry_ptr->handle));
+ } else {
+ CallBatchGemmEx(args, ret, entry_ptr->handle);
+ }
+ });
} // namespace contrib
} // namespace tvm
diff --git a/src/runtime/contrib/cublas/cublas_json_runtime.cc b/src/runtime/contrib/cublas/cublas_json_runtime.cc
index 295089dd07..1d1593ae97 100644
--- a/src/runtime/contrib/cublas/cublas_json_runtime.cc
+++ b/src/runtime/contrib/cublas/cublas_json_runtime.cc
@@ -67,9 +67,8 @@ class CublasJSONRuntime : public JSONRuntimeBase {
void Run(TVMArgs args) {
auto* entry_ptr = tvm::contrib::CuBlasLtThreadEntry::ThreadLocal();
- auto func = tvm::runtime::Registry::Get("runtime.get_cuda_stream");
- ICHECK(func != nullptr);
- cudaStream_t stream = static_cast<cudaStream_t>((*func)().operator void*());
+ auto func = tvm::ffi::Function::GetGlobalRequired("runtime.get_cuda_stream");
+ cudaStream_t stream = static_cast<cudaStream_t>(func().operator void*());
std::vector<const DLTensor*> dl_tensors(NumEntries());
diff --git a/src/runtime/contrib/cudnn/conv_backward.cc b/src/runtime/contrib/cudnn/conv_backward.cc
index a58448ae67..d27bdccfcb 100644
--- a/src/runtime/contrib/cudnn/conv_backward.cc
+++ b/src/runtime/contrib/cudnn/conv_backward.cc
@@ -186,7 +186,7 @@ void BackwardFilterFindAlgo(int format, int dims, int groups, const int pad[], c
}
TVM_REGISTER_GLOBAL("tvm.contrib.cudnn.conv2d.backward_data")
- .set_body([](TVMArgs args, TVMRetValue* ret) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
int mode = args[0];
int format = args[1];
int algo = args[2];
@@ -207,7 +207,7 @@ TVM_REGISTER_GLOBAL("tvm.contrib.cudnn.conv2d.backward_data")
});
TVM_REGISTER_GLOBAL("tvm.contrib.cudnn.conv.backward_data_find_algo")
- .set_body([](TVMArgs args, TVMRetValue* ret) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
int format = args[0];
int dims = args[1];
int* pad = static_cast<int*>(static_cast<void*>(args[2]));
@@ -226,7 +226,7 @@ TVM_REGISTER_GLOBAL("tvm.contrib.cudnn.conv.backward_data_find_algo")
});
TVM_REGISTER_GLOBAL("tvm.contrib.cudnn.conv2d.backward_filter")
- .set_body([](TVMArgs args, TVMRetValue* ret) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
int mode = args[0];
int format = args[1];
int algo = args[2];
@@ -247,7 +247,7 @@ TVM_REGISTER_GLOBAL("tvm.contrib.cudnn.conv2d.backward_filter")
});
TVM_REGISTER_GLOBAL("tvm.contrib.cudnn.conv.backward_filter_find_algo")
- .set_body([](TVMArgs args, TVMRetValue* ret) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
int format = args[0];
int dims = args[1];
int* pad = static_cast<int*>(static_cast<void*>(args[2]));
diff --git a/src/runtime/contrib/cudnn/conv_forward.cc b/src/runtime/contrib/cudnn/conv_forward.cc
index c5f46d4847..5f21f31d61 100644
--- a/src/runtime/contrib/cudnn/conv_forward.cc
+++ b/src/runtime/contrib/cudnn/conv_forward.cc
@@ -154,7 +154,7 @@ void FindAlgo(int format, int dims, int groups, const int pad[], const int strid
}
TVM_REGISTER_GLOBAL("tvm.contrib.cudnn.conv2d.forward")
- .set_body([](TVMArgs args, TVMRetValue* ret) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
int mode = args[0];
int format = args[1];
int algo = args[2];
@@ -175,7 +175,7 @@ TVM_REGISTER_GLOBAL("tvm.contrib.cudnn.conv2d.forward")
});
TVM_REGISTER_GLOBAL("tvm.contrib.cudnn.conv2d+bias+act.forward")
- .set_body([](TVMArgs args, TVMRetValue* ret) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
int mode = args[0];
int format = args[1];
int algo = args[2];
@@ -199,7 +199,7 @@ TVM_REGISTER_GLOBAL("tvm.contrib.cudnn.conv2d+bias+act.forward")
});
TVM_REGISTER_GLOBAL("tvm.contrib.cudnn.conv3d.forward")
- .set_body([](TVMArgs args, TVMRetValue* ret) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
int mode = args[0];
int format = args[1];
int algo = args[2];
@@ -220,7 +220,7 @@ TVM_REGISTER_GLOBAL("tvm.contrib.cudnn.conv3d.forward")
});
TVM_REGISTER_GLOBAL("tvm.contrib.cudnn.conv.forward_find_algo")
- .set_body([](TVMArgs args, TVMRetValue* ret) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
int format = args[0];
int dims = args[1];
int* pad = static_cast<int*>(static_cast<void*>(args[2]));
diff --git a/src/runtime/contrib/cudnn/cudnn_utils.cc b/src/runtime/contrib/cudnn/cudnn_utils.cc
index 07c6611f7e..fa27c55211 100644
--- a/src/runtime/contrib/cudnn/cudnn_utils.cc
+++ b/src/runtime/contrib/cudnn/cudnn_utils.cc
@@ -101,8 +101,8 @@ const void* CuDNNDataType::GetConst<1>(cudnnDataType_t type) {
CuDNNThreadEntry::CuDNNThreadEntry() {
auto stream = runtime::CUDAThreadEntry::ThreadLocal()->stream;
- auto func = runtime::Registry::Get("device_api.cuda");
- void* ret = (*func)();
+ auto func = tvm::ffi::Function::GetGlobalRequired("device_api.cuda");
+ void* ret = func();
cuda_api = static_cast<runtime::DeviceAPI*>(ret);
// If no CuDNN-capable device is present, allow the CuDNNThreadEntry
diff --git a/src/runtime/contrib/cudnn/softmax.cc b/src/runtime/contrib/cudnn/softmax.cc
index 4b37c428e0..e178df812a 100644
--- a/src/runtime/contrib/cudnn/softmax.cc
+++ b/src/runtime/contrib/cudnn/softmax.cc
@@ -78,12 +78,14 @@ void softmax_impl(cudnnSoftmaxAlgorithm_t alg, TVMArgs args, TVMRetValue* ret) {
}
TVM_REGISTER_GLOBAL("tvm.contrib.cudnn.softmax.forward")
- .set_body([](TVMArgs args, TVMRetValue* ret) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
softmax_impl(CUDNN_SOFTMAX_ACCURATE, args, ret);
});
TVM_REGISTER_GLOBAL("tvm.contrib.cudnn.log_softmax.forward")
- .set_body([](TVMArgs args, TVMRetValue* ret) { softmax_impl(CUDNN_SOFTMAX_LOG, args, ret); });
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
+ softmax_impl(CUDNN_SOFTMAX_LOG, args, ret);
+ });
} // namespace contrib
} // namespace tvm
diff --git a/src/runtime/contrib/curand/curand.cc b/src/runtime/contrib/curand/curand.cc
index 50600d913f..e8aaf31fc5 100644
--- a/src/runtime/contrib/curand/curand.cc
+++ b/src/runtime/contrib/curand/curand.cc
@@ -51,9 +51,8 @@ class CURandGenerator {
};
DeviceAPI* GetCUDADeviceAPI() {
- const PackedFunc* get_cuda_api = runtime::Registry::Get("device_api.cuda");
- ICHECK(get_cuda_api) << "ValueError: TVM is not built with USE_CUDA=ON";
- void* ret = (*get_cuda_api)();
+ const auto get_cuda_api = tvm::ffi::Function::GetGlobalRequired("device_api.cuda");
+ void* ret = get_cuda_api();
runtime::DeviceAPI* cuda_api = static_cast<runtime::DeviceAPI*>(ret);
return cuda_api;
}
diff --git a/src/runtime/contrib/cutlass/fp16_group_gemm.cu b/src/runtime/contrib/cutlass/fp16_group_gemm.cu
index 3c051819b2..123bf1b8d6 100644
--- a/src/runtime/contrib/cutlass/fp16_group_gemm.cu
+++ b/src/runtime/contrib/cutlass/fp16_group_gemm.cu
@@ -42,8 +42,8 @@ void tvm_cutlass_group_gemm_sm90(NDArray x, NDArray weight, NDArray indptr, NDAr
NDArray out) {
// Workspace is used for storing device-side group gemm arguments and cutlass internal workspace.
// Recommened size is 4MB.
- auto func = tvm::runtime::Registry::Get("runtime.get_cuda_stream");
- ICHECK(func != nullptr);
+ static auto func = tvm::ffi::Function::GetGlobalRequired("runtime.get_cuda_stream");
+ cudaStream_t stream = static_cast<cudaStream_t>(func().operator void*());
CHECK_EQ(x->ndim, 2);
CHECK_EQ(weight->ndim, 3);
CHECK_EQ(indptr->ndim, 1);
@@ -54,7 +54,6 @@ void tvm_cutlass_group_gemm_sm90(NDArray x, NDArray weight, NDArray indptr, NDAr
int k = weight->shape[2];
float alpha = 1.0f;
float beta = 0.0f;
- cudaStream_t stream = static_cast<cudaStream_t>((*func)().operator void*());
cutlass_group_gemm(static_cast<ElementA*>(x->data), static_cast<ElementB*>(weight->data),
static_cast<int64_t*>(indptr->data), static_cast<uint8_t*>(workspace->data),
workspace->shape[0], n, k, num_groups, alpha, beta,
diff --git a/src/runtime/contrib/cutlass/fp8_blockwise_scaled_gemm.cu b/src/runtime/contrib/cutlass/fp8_blockwise_scaled_gemm.cu
index 4ac5a621a0..6de7201856 100644
--- a/src/runtime/contrib/cutlass/fp8_blockwise_scaled_gemm.cu
+++ b/src/runtime/contrib/cutlass/fp8_blockwise_scaled_gemm.cu
@@ -39,8 +39,8 @@ void tvm_cutlass_fp8_blockwise_scaled_gemm(NDArray a, NDArray b, NDArray scales_
// Workspace is used for storing device-side gemm arguments and cutlass internal workspace.
// Recommened size is 4MB.
- auto get_stream_func = tvm::runtime::Registry::Get("runtime.get_cuda_stream");
- ICHECK(get_stream_func != nullptr);
+ const auto get_stream_func = tvm::ffi::Function::GetGlobal("runtime.get_cuda_stream");
+ ICHECK(get_stream_func.has_value());
cudaStream_t stream = static_cast<cudaStream_t>((*get_stream_func)().operator void*());
CHECK_GE(a->ndim, 2);
@@ -100,8 +100,8 @@ void tvm_cutlass_fp8_blockwise_scaled_bmm(NDArray a, NDArray b, NDArray scales_a
// Workspace is used for storing device-side gemm arguments and cutlass internal workspace.
// Recommened size is 4MB.
- auto get_stream_func = tvm::runtime::Registry::Get("runtime.get_cuda_stream");
- ICHECK(get_stream_func != nullptr);
+ const auto get_stream_func = tvm::ffi::Function::GetGlobal("runtime.get_cuda_stream");
+ ICHECK(get_stream_func.has_value());
cudaStream_t stream = static_cast<cudaStream_t>((*get_stream_func)().operator void*());
CHECK_EQ(a->ndim, 3);
diff --git a/src/runtime/contrib/cutlass/fp8_gemm.cu b/src/runtime/contrib/cutlass/fp8_gemm.cu
index 67e502a163..dfffc34e92 100644
--- a/src/runtime/contrib/cutlass/fp8_gemm.cu
+++ b/src/runtime/contrib/cutlass/fp8_gemm.cu
@@ -42,8 +42,9 @@ void tvm_cutlass_fp8_gemm(NDArray x, NDArray weight, NDArray workspace, NDArray
NDArray out) {
// Workspace is used for storing device-side gemm arguments and cutlass internal workspace.
// Recommened size is 4MB.
- auto func = tvm::runtime::Registry::Get("runtime.get_cuda_stream");
- ICHECK(func != nullptr);
+ static auto func = tvm::ffi::Function::GetGlobalRequired("runtime.get_cuda_stream");
+ cudaStream_t stream = static_cast<cudaStream_t>(func().operator void*());
+
CHECK_GE(x->ndim, 2);
CHECK_EQ(weight->ndim, 2);
CHECK_EQ(workspace->ndim, 1);
@@ -60,7 +61,6 @@ void tvm_cutlass_fp8_gemm(NDArray x, NDArray weight, NDArray workspace, NDArray
CHECK_EQ(x->shape[x->ndim - 1], weight->shape[1]) << "Only col-major weight is supported now.";
int64_t k = x->shape[x->ndim - 1];
const float* beta = nullptr;
- cudaStream_t stream = static_cast<cudaStream_t>((*func)().operator void*());
if (m <= 64) {
cutlass_gemm<KernelTraitsM64>(
static_cast<ElementA*>(x->data), static_cast<ElementB*>(weight->data),
diff --git a/src/runtime/contrib/cutlass/fp8_group_gemm.cu b/src/runtime/contrib/cutlass/fp8_group_gemm.cu
index 31ad4367af..4c1689ca80 100644
--- a/src/runtime/contrib/cutlass/fp8_group_gemm.cu
+++ b/src/runtime/contrib/cutlass/fp8_group_gemm.cu
@@ -45,8 +45,8 @@ void tvm_cutlass_fp8_group_gemm(NDArray x, NDArray weight, NDArray indptr, NDArr
NDArray alpha, NDArray out) {
// Workspace is used for storing device-side group gemm arguments and cutlass internal workspace.
// Recommened size is 4MB.
- auto func = tvm::runtime::Registry::Get("runtime.get_cuda_stream");
- ICHECK(func != nullptr);
+ static auto func = tvm::ffi::Function::GetGlobalRequired("runtime.get_cuda_stream");
+ cudaStream_t stream = static_cast<cudaStream_t>(func().operator void*());
CHECK_EQ(x->ndim, 2);
CHECK_EQ(weight->ndim, 3);
CHECK_EQ(indptr->ndim, 1);
@@ -60,7 +60,6 @@ void tvm_cutlass_fp8_group_gemm(NDArray x, NDArray weight, NDArray indptr, NDArr
int n = weight->shape[1];
int k = x->shape[1];
const float* beta = nullptr;
- cudaStream_t stream = static_cast<cudaStream_t>((*func)().operator void*());
cutlass_group_gemm(static_cast<ElementA*>(x->data), static_cast<ElementB*>(weight->data),
static_cast<int64_t*>(indptr->data), static_cast<uint8_t*>(workspace->data),
workspace->shape[0], n, k, num_groups, static_cast<float*>(alpha->data), beta,
diff --git a/src/runtime/contrib/dnnl/dnnl.cc b/src/runtime/contrib/dnnl/dnnl.cc
index 54ee635488..c3afebcc47 100644
--- a/src/runtime/contrib/dnnl/dnnl.cc
+++ b/src/runtime/contrib/dnnl/dnnl.cc
@@ -348,7 +348,7 @@ extern "C" void dnnl_binary_op(float* data, float* weight, float* out, int algo_
}
// DNNL Conv2d single OP
-TVM_REGISTER_GLOBAL("tvm.contrib.dnnl.conv2d").set_body([](TVMArgs args, TVMRetValue* ret) {
+TVM_REGISTER_GLOBAL("tvm.contrib.dnnl.conv2d").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
DLTensor* input = args[0];
DLTensor* weights = args[1];
DLTensor* output = args[2];
diff --git a/src/runtime/contrib/edgetpu/edgetpu_runtime.cc b/src/runtime/contrib/edgetpu/edgetpu_runtime.cc
index a2f159b55d..a0c11fea22 100644
--- a/src/runtime/contrib/edgetpu/edgetpu_runtime.cc
+++ b/src/runtime/contrib/edgetpu/edgetpu_runtime.cc
@@ -68,8 +68,9 @@ Module EdgeTPURuntimeCreate(const std::string& tflite_model_bytes, Device dev) {
return Module(exec);
}
-TVM_REGISTER_GLOBAL("tvm.edgetpu_runtime.create").set_body([](TVMArgs args, TVMRetValue* rv) {
- *rv = EdgeTPURuntimeCreate(args[0], args[1]);
-});
+TVM_REGISTER_GLOBAL("tvm.edgetpu_runtime.create")
+ .set_body_packed([](TVMArgs args, TVMRetValue* rv) {
+ *rv = EdgeTPURuntimeCreate(args[0], args[1]);
+ });
} // namespace runtime
} // namespace tvm
diff --git a/src/runtime/contrib/hipblas/hipblas.cc b/src/runtime/contrib/hipblas/hipblas.cc
index c135a2855d..e402121d04 100644
--- a/src/runtime/contrib/hipblas/hipblas.cc
+++ b/src/runtime/contrib/hipblas/hipblas.cc
@@ -407,30 +407,31 @@ inline void CallBatchGemmEx(TVMArgs args, TVMRetValue* ret, hipblasHandle_t hdl)
}
// matrix multiplication for row major
-TVM_REGISTER_GLOBAL("tvm.contrib.hipblas.matmul").set_body([](TVMArgs args, TVMRetValue* ret) {
- DLTensor* A = args[0];
- DLTensor* C = args[2];
+TVM_REGISTER_GLOBAL("tvm.contrib.hipblas.matmul")
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
+ DLTensor* A = args[0];
+ DLTensor* C = args[2];
- HipBlasThreadEntry* entry_ptr = HipBlasThreadEntry::ThreadLocal();
+ HipBlasThreadEntry* entry_ptr = HipBlasThreadEntry::ThreadLocal();
- if (TypeEqual(A->dtype, C->dtype)) {
- ICHECK(TypeMatch(A->dtype, kDLFloat, 16) || TypeMatch(A->dtype, kDLFloat, 32) ||
- TypeMatch(A->dtype, kDLFloat, 64));
+ if (TypeEqual(A->dtype, C->dtype)) {
+ ICHECK(TypeMatch(A->dtype, kDLFloat, 16) || TypeMatch(A->dtype, kDLFloat, 32) ||
+ TypeMatch(A->dtype, kDLFloat, 64));
- if (TypeMatch(A->dtype, kDLFloat, 16)) {
- CallGemm(args, ret, HipblasHgemmOp(entry_ptr->handle));
- } else if (TypeMatch(A->dtype, kDLFloat, 32)) {
- CallGemm(args, ret, HipblasSgemmOp(entry_ptr->handle));
- } else {
- CallGemm(args, ret, HipblasDgemmOp(entry_ptr->handle));
- }
- } else {
- CallGemmEx(args, ret, entry_ptr->handle);
- }
-});
+ if (TypeMatch(A->dtype, kDLFloat, 16)) {
+ CallGemm(args, ret, HipblasHgemmOp(entry_ptr->handle));
+ } else if (TypeMatch(A->dtype, kDLFloat, 32)) {
+ CallGemm(args, ret, HipblasSgemmOp(entry_ptr->handle));
+ } else {
+ CallGemm(args, ret, HipblasDgemmOp(entry_ptr->handle));
+ }
+ } else {
+ CallGemmEx(args, ret, entry_ptr->handle);
+ }
+ });
TVM_REGISTER_GLOBAL("tvm.contrib.hipblas.batch_matmul")
- .set_body([](TVMArgs args, TVMRetValue* ret) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
DLTensor* A = args[0];
DLTensor* C = args[2];
diff --git a/src/runtime/contrib/hipblas/hipblas_json_runtime.cc b/src/runtime/contrib/hipblas/hipblas_json_runtime.cc
index a6e7949e45..78f0673e91 100644
--- a/src/runtime/contrib/hipblas/hipblas_json_runtime.cc
+++ b/src/runtime/contrib/hipblas/hipblas_json_runtime.cc
@@ -64,10 +64,8 @@ class HipblasJSONRuntime : public JSONRuntimeBase {
void Run(TVMArgs args) {
auto* entry_ptr = tvm::contrib::HipBlasLtThreadEntry::ThreadLocal();
-
- auto func = tvm::runtime::Registry::Get("runtime.get_rocm_stream");
- ICHECK(func != nullptr);
- hipStream_t stream = static_cast<hipStream_t>((*func)().operator void*());
+ static auto func = tvm::ffi::Function::GetGlobalRequired("runtime.get_rocm_stream");
+ hipStream_t stream = static_cast<hipStream_t>(func().operator void*());
std::vector<const DLTensor*> dl_tensors(NumEntries());
diff --git a/src/runtime/contrib/miopen/conv_forward.cc b/src/runtime/contrib/miopen/conv_forward.cc
index e1d6eb6197..70772209c3 100644
--- a/src/runtime/contrib/miopen/conv_forward.cc
+++ b/src/runtime/contrib/miopen/conv_forward.cc
@@ -34,121 +34,122 @@ namespace miopen {
using namespace runtime;
-TVM_REGISTER_GLOBAL("tvm.contrib.miopen.conv2d.setup").set_body([](TVMArgs args, TVMRetValue* ret) {
- const int mode = args[0];
- const int dtype = args[1];
- const int pad_h = args[2];
- const int pad_w = args[3];
- const int stride_h = args[4];
- const int stride_w = args[5];
- const int dilation_h = args[6];
- const int dilation_w = args[7];
- const int x_dim0 = args[8];
- const int x_dim1 = args[9];
- const int x_dim2 = args[10];
- const int x_dim3 = args[11];
- const int w_dim0 = args[12];
- const int w_dim1 = args[13];
- const int w_dim2 = args[14];
- const int w_dim3 = args[15];
- const int n_group = args[16];
- void* out_shape = args[17];
-
- MIOpenThreadEntry* entry_ptr = MIOpenThreadEntry::ThreadLocal();
- assert(n_group > 0 && "Group Size > 0 is expected");
- if (n_group > 1) assert(mode > 1 && "Group /Depthwise Conv mode when num of groups > 1");
- // Set Mode
- entry_ptr->conv_entry.mode = static_cast<miopenConvolutionMode_t>(mode);
- // Set Device
- entry_ptr->conv_entry.device = Device{kDLROCM, 0};
- // Set Data Type
- entry_ptr->conv_entry.data_type =
- static_cast<miopenDataType_t>(dtype); // MIOpen supports fp32(miopenFloat), fp16(miopenHalf),
- // int32, int8 at this moment.
- // Set Desc
- MIOPEN_CALL(miopenInitConvolutionDescriptor(entry_ptr->conv_entry.conv_desc,
- entry_ptr->conv_entry.mode, pad_h, pad_w, stride_h,
- stride_w, dilation_h, dilation_w));
- if (n_group > 1)
- MIOPEN_CALL(miopenSetConvolutionGroupCount(entry_ptr->conv_entry.conv_desc, n_group));
- // Set Filter
- MIOPEN_CALL(miopenSet4dTensorDescriptor(entry_ptr->conv_entry.filter_desc,
- entry_ptr->conv_entry.data_type, w_dim0, w_dim1 / n_group,
- w_dim2, w_dim3));
- // Set Input
- MIOPEN_CALL(miopenSet4dTensorDescriptor(entry_ptr->conv_entry.input_desc,
- entry_ptr->conv_entry.data_type, x_dim0, x_dim1, x_dim2,
- x_dim3));
-
- // Set Output shape
- MIOPEN_CALL(miopenGetConvolutionForwardOutputDim(
- entry_ptr->conv_entry.conv_desc, entry_ptr->conv_entry.input_desc,
- entry_ptr->conv_entry.filter_desc, static_cast<int*>(out_shape),
- static_cast<int*>(out_shape) + 1, static_cast<int*>(out_shape) + 2,
- static_cast<int*>(out_shape) + 3));
-
- const int* oshape = static_cast<int*>(out_shape);
- // Set Output
- MIOPEN_CALL(miopenSet4dTensorDescriptor(entry_ptr->conv_entry.output_desc,
- entry_ptr->conv_entry.data_type, oshape[0], oshape[1],
- oshape[2], oshape[3]));
-
- // Set workspace
- size_t workspace_size = 0;
- MIOPEN_CALL(miopenConvolutionForwardGetWorkSpaceSize(
- entry_ptr->handle, entry_ptr->conv_entry.filter_desc, entry_ptr->conv_entry.input_desc,
- entry_ptr->conv_entry.conv_desc, entry_ptr->conv_entry.output_desc, &workspace_size));
- entry_ptr->conv_entry.UpdateWorkspace(workspace_size);
-
- const size_t input_size = x_dim0 * x_dim1 * x_dim2 * x_dim3;
- const size_t filter_size = w_dim0 * w_dim1 * w_dim2 * w_dim3;
- const size_t output_size = oshape[0] * oshape[1] * oshape[2] * oshape[3];
-
- runtime::DeviceAPI* rocm_api = entry_ptr->conv_entry.rocm_api;
- float* input_buf = static_cast<float*>(
- rocm_api->AllocWorkspace(entry_ptr->conv_entry.device, input_size * sizeof(float)));
- float* filter_buf = static_cast<float*>(
- rocm_api->AllocWorkspace(entry_ptr->conv_entry.device, filter_size * sizeof(float)));
- float* output_buf = static_cast<float*>(
- rocm_api->AllocWorkspace(entry_ptr->conv_entry.device, output_size * sizeof(float)));
-
- const int request_algo_count = 4;
- const bool exhaustive_search = false;
- void* workspace = entry_ptr->conv_entry.workspace;
- if (workspace_size == 0) workspace = nullptr;
- int returned_algo_count = 0;
- miopenConvAlgoPerf_t perfs[4];
-
- MIOPEN_CALL(miopenFindConvolutionForwardAlgorithm(
- entry_ptr->handle, entry_ptr->conv_entry.input_desc, input_buf,
- entry_ptr->conv_entry.filter_desc, filter_buf, entry_ptr->conv_entry.conv_desc,
- entry_ptr->conv_entry.output_desc, output_buf, request_algo_count, &returned_algo_count,
- perfs, workspace, workspace_size, exhaustive_search));
-
- rocm_api->FreeWorkspace(entry_ptr->conv_entry.device, input_buf);
- rocm_api->FreeWorkspace(entry_ptr->conv_entry.device, filter_buf);
- rocm_api->FreeWorkspace(entry_ptr->conv_entry.device, output_buf);
-
- const std::vector<std::string> fwd_algo_names{
- "miopenConvolutionFwdAlgoGEMM",
- "miopenConvolutionFwdAlgoDirect",
- "miopenConvolutionFwdAlgoFFT",
- "miopenConvolutionFwdAlgoWinograd",
- };
- const auto best_algo = perfs[0].fwd_algo;
- LOG(INFO) << "\tMIOpen Found " << returned_algo_count << " fwd algorithms, choosing "
- << fwd_algo_names[best_algo];
- for (int i = 0; i < returned_algo_count; ++i) {
- LOG(INFO) << "\t\t" << i << ") " << fwd_algo_names[perfs[i].fwd_algo]
- << " - time: " << perfs[i].time << " ms"
- << ", Memory: " << perfs[i].memory;
- }
- // Set Algo
- ret[0] = static_cast<int>(best_algo);
-});
+TVM_REGISTER_GLOBAL("tvm.contrib.miopen.conv2d.setup")
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
+ const int mode = args[0];
+ const int dtype = args[1];
+ const int pad_h = args[2];
+ const int pad_w = args[3];
+ const int stride_h = args[4];
+ const int stride_w = args[5];
+ const int dilation_h = args[6];
+ const int dilation_w = args[7];
+ const int x_dim0 = args[8];
+ const int x_dim1 = args[9];
+ const int x_dim2 = args[10];
+ const int x_dim3 = args[11];
+ const int w_dim0 = args[12];
+ const int w_dim1 = args[13];
+ const int w_dim2 = args[14];
+ const int w_dim3 = args[15];
+ const int n_group = args[16];
+ void* out_shape = args[17];
+
+ MIOpenThreadEntry* entry_ptr = MIOpenThreadEntry::ThreadLocal();
+ assert(n_group > 0 && "Group Size > 0 is expected");
+ if (n_group > 1) assert(mode > 1 && "Group /Depthwise Conv mode when num of groups > 1");
+ // Set Mode
+ entry_ptr->conv_entry.mode = static_cast<miopenConvolutionMode_t>(mode);
+ // Set Device
+ entry_ptr->conv_entry.device = Device{kDLROCM, 0};
+ // Set Data Type
+ entry_ptr->conv_entry.data_type =
+ static_cast<miopenDataType_t>(dtype); // MIOpen supports fp32(miopenFloat),
+ // fp16(miopenHalf), int32, int8 at this moment.
+ // Set Desc
+ MIOPEN_CALL(miopenInitConvolutionDescriptor(entry_ptr->conv_entry.conv_desc,
+ entry_ptr->conv_entry.mode, pad_h, pad_w,
+ stride_h, stride_w, dilation_h, dilation_w));
+ if (n_group > 1)
+ MIOPEN_CALL(miopenSetConvolutionGroupCount(entry_ptr->conv_entry.conv_desc, n_group));
+ // Set Filter
+ MIOPEN_CALL(miopenSet4dTensorDescriptor(entry_ptr->conv_entry.filter_desc,
+ entry_ptr->conv_entry.data_type, w_dim0,
+ w_dim1 / n_group, w_dim2, w_dim3));
+ // Set Input
+ MIOPEN_CALL(miopenSet4dTensorDescriptor(entry_ptr->conv_entry.input_desc,
+ entry_ptr->conv_entry.data_type, x_dim0, x_dim1,
+ x_dim2, x_dim3));
+
+ // Set Output shape
+ MIOPEN_CALL(miopenGetConvolutionForwardOutputDim(
+ entry_ptr->conv_entry.conv_desc, entry_ptr->conv_entry.input_desc,
+ entry_ptr->conv_entry.filter_desc, static_cast<int*>(out_shape),
+ static_cast<int*>(out_shape) + 1, static_cast<int*>(out_shape) + 2,
+ static_cast<int*>(out_shape) + 3));
+
+ const int* oshape = static_cast<int*>(out_shape);
+ // Set Output
+ MIOPEN_CALL(miopenSet4dTensorDescriptor(entry_ptr->conv_entry.output_desc,
+ entry_ptr->conv_entry.data_type, oshape[0], oshape[1],
+ oshape[2], oshape[3]));
+
+ // Set workspace
+ size_t workspace_size = 0;
+ MIOPEN_CALL(miopenConvolutionForwardGetWorkSpaceSize(
+ entry_ptr->handle, entry_ptr->conv_entry.filter_desc, entry_ptr->conv_entry.input_desc,
+ entry_ptr->conv_entry.conv_desc, entry_ptr->conv_entry.output_desc, &workspace_size));
+ entry_ptr->conv_entry.UpdateWorkspace(workspace_size);
+
+ const size_t input_size = x_dim0 * x_dim1 * x_dim2 * x_dim3;
+ const size_t filter_size = w_dim0 * w_dim1 * w_dim2 * w_dim3;
+ const size_t output_size = oshape[0] * oshape[1] * oshape[2] * oshape[3];
+
+ runtime::DeviceAPI* rocm_api = entry_ptr->conv_entry.rocm_api;
+ float* input_buf = static_cast<float*>(
+ rocm_api->AllocWorkspace(entry_ptr->conv_entry.device, input_size * sizeof(float)));
+ float* filter_buf = static_cast<float*>(
+ rocm_api->AllocWorkspace(entry_ptr->conv_entry.device, filter_size * sizeof(float)));
+ float* output_buf = static_cast<float*>(
+ rocm_api->AllocWorkspace(entry_ptr->conv_entry.device, output_size * sizeof(float)));
+
+ const int request_algo_count = 4;
+ const bool exhaustive_search = false;
+ void* workspace = entry_ptr->conv_entry.workspace;
+ if (workspace_size == 0) workspace = nullptr;
+ int returned_algo_count = 0;
+ miopenConvAlgoPerf_t perfs[4];
+
+ MIOPEN_CALL(miopenFindConvolutionForwardAlgorithm(
+ entry_ptr->handle, entry_ptr->conv_entry.input_desc, input_buf,
+ entry_ptr->conv_entry.filter_desc, filter_buf, entry_ptr->conv_entry.conv_desc,
+ entry_ptr->conv_entry.output_desc, output_buf, request_algo_count, &returned_algo_count,
+ perfs, workspace, workspace_size, exhaustive_search));
+
+ rocm_api->FreeWorkspace(entry_ptr->conv_entry.device, input_buf);
+ rocm_api->FreeWorkspace(entry_ptr->conv_entry.device, filter_buf);
+ rocm_api->FreeWorkspace(entry_ptr->conv_entry.device, output_buf);
+
+ const std::vector<std::string> fwd_algo_names{
+ "miopenConvolutionFwdAlgoGEMM",
+ "miopenConvolutionFwdAlgoDirect",
+ "miopenConvolutionFwdAlgoFFT",
+ "miopenConvolutionFwdAlgoWinograd",
+ };
+ const auto best_algo = perfs[0].fwd_algo;
+ LOG(INFO) << "\tMIOpen Found " << returned_algo_count << " fwd algorithms, choosing "
+ << fwd_algo_names[best_algo];
+ for (int i = 0; i < returned_algo_count; ++i) {
+ LOG(INFO) << "\t\t" << i << ") " << fwd_algo_names[perfs[i].fwd_algo]
+ << " - time: " << perfs[i].time << " ms"
+ << ", Memory: " << perfs[i].memory;
+ }
+ // Set Algo
+ ret[0] = static_cast<int>(best_algo);
+ });
TVM_REGISTER_GLOBAL("tvm.contrib.miopen.conv2d.forward")
- .set_body([](TVMArgs args, TVMRetValue* ret) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
const int mode = args[0];
const int dtype = args[1];
const int pad_h = args[2];
diff --git a/src/runtime/contrib/miopen/miopen_utils.cc b/src/runtime/contrib/miopen/miopen_utils.cc
index ff748b8826..b750e56c7e 100644
--- a/src/runtime/contrib/miopen/miopen_utils.cc
+++ b/src/runtime/contrib/miopen/miopen_utils.cc
@@ -43,8 +43,8 @@ std::string miopenGetErrorString(int error_code) {
// MiopenThreadEntry
MIOpenThreadEntry::MIOpenThreadEntry() {
auto stream = runtime::ROCMThreadEntry::ThreadLocal()->stream;
- auto func = runtime::Registry::Get("device_api.rocm");
- void* ret = (*func)();
+ const auto get_rocm_api = tvm::ffi::Function::GetGlobalRequired("device_api.rocm");
+ void* ret = get_rocm_api();
rocm_api = static_cast<runtime::DeviceAPI*>(ret);
MIOPEN_CALL(miopenCreate(&handle));
MIOPEN_CALL(miopenSetStream(handle, stream));
diff --git a/src/runtime/contrib/miopen/softmax.cc b/src/runtime/contrib/miopen/softmax.cc
index 5a0f24ed7a..32c648edee 100644
--- a/src/runtime/contrib/miopen/softmax.cc
+++ b/src/runtime/contrib/miopen/softmax.cc
@@ -80,12 +80,14 @@ void softmax_impl(TVMArgs args, TVMRetValue* ret, miopenSoftmaxAlgorithm_t alg)
}
TVM_REGISTER_GLOBAL("tvm.contrib.miopen.softmax.forward")
- .set_body([](TVMArgs args, TVMRetValue* ret) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
softmax_impl(args, ret, MIOPEN_SOFTMAX_ACCURATE);
});
TVM_REGISTER_GLOBAL("tvm.contrib.miopen.log_softmax.forward")
- .set_body([](TVMArgs args, TVMRetValue* ret) { softmax_impl(args, ret, MIOPEN_SOFTMAX_LOG); });
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
+ softmax_impl(args, ret, MIOPEN_SOFTMAX_LOG);
+ });
} // namespace miopen
} // namespace contrib
diff --git a/src/runtime/contrib/mps/conv.mm b/src/runtime/contrib/mps/conv.mm
index 84ca5e10f9..d268eae4c3 100644
--- a/src/runtime/contrib/mps/conv.mm
+++ b/src/runtime/contrib/mps/conv.mm
@@ -24,57 +24,59 @@ namespace contrib {
using namespace runtime;
-TVM_REGISTER_GLOBAL("tvm.contrib.mps.buffer2img").set_body([](TVMArgs args, TVMRetValue* ret) {
- DLTensor* buf = args[0];
- DLTensor* img = args[1];
- // copy to temp
- id<MTLBuffer> mtlbuf = (__bridge id<MTLBuffer>)(buf->data);
- MetalThreadEntry* entry_ptr = MetalThreadEntry::ThreadLocal();
- runtime::metal::MetalThreadEntry* rt = runtime::metal::MetalThreadEntry::ThreadLocal();
- id<MTLDevice> dev = entry_ptr->metal_api->GetDevice(buf->device);
- id<MTLBuffer> temp = rt->GetTempBuffer(buf->device, [mtlbuf length]);
- entry_ptr->metal_api->CopyDataFromTo((__bridge void*)mtlbuf, 0, (__bridge void*)temp, 0,
- [mtlbuf length], buf -> device, buf -> device, buf -> dtype,
- nullptr);
-
- MPSImageDescriptor* desc =
- [MPSImageDescriptor imageDescriptorWithChannelFormat:MPSImageFeatureChannelFormatFloat32
- width:buf->shape[2]
- height:buf->shape[1]
- featureChannels:buf->shape[3]];
-
- MPSImage* mpsimg = entry_ptr->AllocMPSImage(dev, desc);
-
- [mpsimg writeBytes:[temp contents]
- dataLayout:MPSDataLayoutHeightxWidthxFeatureChannels
- imageIndex:0];
-
- img->data = (__bridge void*)mpsimg;
-
- [mpsimg readBytes:[temp contents]
- dataLayout:MPSDataLayoutHeightxWidthxFeatureChannels
- imageIndex:0];
-});
-
-TVM_REGISTER_GLOBAL("tvm.contrib.mps.img2buffer").set_body([](TVMArgs args, TVMRetValue* ret) {
- DLTensor* img = args[0];
- DLTensor* buf = args[1];
- id<MTLBuffer> mtlbuf = (__bridge id<MTLBuffer>)(buf->data);
- MPSImage* mpsimg = (__bridge MPSImage*)(img->data);
- MetalThreadEntry* entry_ptr = MetalThreadEntry::ThreadLocal();
- runtime::metal::MetalThreadEntry* rt = runtime::metal::MetalThreadEntry::ThreadLocal();
- id<MTLBuffer> temp = rt->GetTempBuffer(buf->device, [mtlbuf length]);
-
- [mpsimg readBytes:[temp contents]
- dataLayout:MPSDataLayoutHeightxWidthxFeatureChannels
- imageIndex:0];
-
- entry_ptr->metal_api->CopyDataFromTo((__bridge void*)temp, 0, (__bridge void*)mtlbuf, 0,
- [mtlbuf length], buf -> device, buf -> device, buf -> dtype,
- nullptr);
-});
-
-TVM_REGISTER_GLOBAL("tvm.contrib.mps.conv2d").set_body([](TVMArgs args, TVMRetValue* ret) {
+TVM_REGISTER_GLOBAL("tvm.contrib.mps.buffer2img")
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
+ DLTensor* buf = args[0];
+ DLTensor* img = args[1];
+ // copy to temp
+ id<MTLBuffer> mtlbuf = (__bridge id<MTLBuffer>)(buf->data);
+ MetalThreadEntry* entry_ptr = MetalThreadEntry::ThreadLocal();
+ runtime::metal::MetalThreadEntry* rt = runtime::metal::MetalThreadEntry::ThreadLocal();
+ id<MTLDevice> dev = entry_ptr->metal_api->GetDevice(buf->device);
+ id<MTLBuffer> temp = rt->GetTempBuffer(buf->device, [mtlbuf length]);
+ entry_ptr->metal_api->CopyDataFromTo((__bridge void*)mtlbuf, 0, (__bridge void*)temp, 0,
+ [mtlbuf length], buf -> device, buf -> device,
+ buf -> dtype, nullptr);
+
+ MPSImageDescriptor* desc =
+ [MPSImageDescriptor imageDescriptorWithChannelFormat:MPSImageFeatureChannelFormatFloat32
+ width:buf->shape[2]
+ height:buf->shape[1]
+ featureChannels:buf->shape[3]];
+
+ MPSImage* mpsimg = entry_ptr->AllocMPSImage(dev, desc);
+
+ [mpsimg writeBytes:[temp contents]
+ dataLayout:MPSDataLayoutHeightxWidthxFeatureChannels
+ imageIndex:0];
+
+ img->data = (__bridge void*)mpsimg;
+
+ [mpsimg readBytes:[temp contents]
+ dataLayout:MPSDataLayoutHeightxWidthxFeatureChannels
+ imageIndex:0];
+ });
+
+TVM_REGISTER_GLOBAL("tvm.contrib.mps.img2buffer")
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
+ DLTensor* img = args[0];
+ DLTensor* buf = args[1];
+ id<MTLBuffer> mtlbuf = (__bridge id<MTLBuffer>)(buf->data);
+ MPSImage* mpsimg = (__bridge MPSImage*)(img->data);
+ MetalThreadEntry* entry_ptr = MetalThreadEntry::ThreadLocal();
+ runtime::metal::MetalThreadEntry* rt = runtime::metal::MetalThreadEntry::ThreadLocal();
+ id<MTLBuffer> temp = rt->GetTempBuffer(buf->device, [mtlbuf length]);
+
+ [mpsimg readBytes:[temp contents]
+ dataLayout:MPSDataLayoutHeightxWidthxFeatureChannels
+ imageIndex:0];
+
+ entry_ptr->metal_api->CopyDataFromTo((__bridge void*)temp, 0, (__bridge void*)mtlbuf, 0,
+ [mtlbuf length], buf -> device, buf -> device,
+ buf -> dtype, nullptr);
+ });
+
+TVM_REGISTER_GLOBAL("tvm.contrib.mps.conv2d").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
// MPS-NHWC
DLTensor* data = args[0];
DLTensor* weight = args[1];
@@ -97,8 +99,8 @@ TVM_REGISTER_GLOBAL("tvm.contrib.mps.conv2d").set_body([](TVMArgs args, TVMRetVa
int kW = weight->shape[2];
int iCh = weight->shape[3];
- auto f_buf2img = runtime::Registry::Get("tvm.contrib.mps.buffer2img");
- auto f_img2buf = runtime::Registry::Get("tvm.contrib.mps.img2buffer");
+ const auto f_buf2img = tvm::ffi::Function::GetGlobal("tvm.contrib.mps.buffer2img");
+ const auto f_img2buf = tvm::ffi::Function::GetGlobal("tvm.contrib.mps.img2buffer");
// Get Metal device API
MetalThreadEntry* entry_ptr = MetalThreadEntry::ThreadLocal();
runtime::metal::MetalThreadEntry* rt = runtime::metal::MetalThreadEntry::ThreadLocal();
diff --git a/src/runtime/contrib/mps/gemm.mm b/src/runtime/contrib/mps/gemm.mm
index db3a80ba49..288277ab5b 100644
--- a/src/runtime/contrib/mps/gemm.mm
+++ b/src/runtime/contrib/mps/gemm.mm
@@ -24,7 +24,7 @@ namespace contrib {
using namespace runtime;
-TVM_REGISTER_GLOBAL("tvm.contrib.mps.matmul").set_body([](TVMArgs args, TVMRetValue* ret) {
+TVM_REGISTER_GLOBAL("tvm.contrib.mps.matmul").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
DLTensor* A = args[0];
DLTensor* B = args[1];
DLTensor* C = args[2];
diff --git a/src/runtime/contrib/mps/mps_utils.mm b/src/runtime/contrib/mps/mps_utils.mm
index f9f8043116..f5dfbdd530 100644
--- a/src/runtime/contrib/mps/mps_utils.mm
+++ b/src/runtime/contrib/mps/mps_utils.mm
@@ -72,8 +72,8 @@ MPSTemporaryImage* MetalThreadEntry::AllocTempImage(id<MTLCommandBuffer> cb,
}
MetalThreadEntry::MetalThreadEntry() {
- auto func = runtime::Registry::Get("device_api.metal");
- void* ret = (*func)();
+ const auto get_metal_api = tvm::ffi::Function::GetGlobalRequired("device_api.metal");
+ void* ret = get_metal_api();
metal_api = static_cast<runtime::metal::MetalWorkspace*>(ret);
}
diff --git a/src/runtime/contrib/mrvl/mrvl_runtime.cc b/src/runtime/contrib/mrvl/mrvl_runtime.cc
index bd81df3cc2..312cf7ab3c 100644
--- a/src/runtime/contrib/mrvl/mrvl_runtime.cc
+++ b/src/runtime/contrib/mrvl/mrvl_runtime.cc
@@ -132,7 +132,7 @@ class MarvellSimulatorModuleNode : public ModuleNode {
}
void set_num_inputs_outputs() {
- const auto* get_value_from_key = runtime::Registry::Get("tvm.mrvl.find_value_in_KV_pair");
+ const auto get_value_from_key = tvm::ffi::Function::GetGlobal("tvm.mrvl.find_value_in_KV_pair");
std::string value_for_inputs = (*get_value_from_key)(nodes_json_, "num_subgraph_inputs");
num_inputs_ = std::stoi(value_for_inputs);
diff --git a/src/runtime/contrib/mrvl/mrvl_sw_runtime_lib.cc b/src/runtime/contrib/mrvl/mrvl_sw_runtime_lib.cc
index f5e222255c..38b6ba760b 100644
--- a/src/runtime/contrib/mrvl/mrvl_sw_runtime_lib.cc
+++ b/src/runtime/contrib/mrvl/mrvl_sw_runtime_lib.cc
@@ -90,7 +90,7 @@ static void ReadInputsAndGenerateInputBin(TVMArgs args, const std::string& input
file_out << std::endl << "\t]" << std::endl;
file_out << "}" << std::endl;
- const auto* json_to_bin = tvm::runtime::Registry::Get("tvm.mrvl.JsonToBin");
+ const auto json_to_bin = tvm::ffi::Function::GetGlobal("tvm.mrvl.JsonToBin");
(*json_to_bin)(input_json, input_bin);
}
@@ -100,7 +100,7 @@ static void RunInferenceOnMlModel(const std::string& symbol_name, const std::str
auto command = bin_directory + "/mrvl-mlsim " + "-m " + bin_file + " -d " + input_bin + " -o " +
out_bin_prefix;
std::string sim_directory = "mrvl_sw_sim_" + symbol_name;
- const auto* run_sim = tvm::runtime::Registry::Get("tvm.mrvl.RunSim");
+ const auto run_sim = tvm::ffi::Function::GetGlobal("tvm.mrvl.RunSim");
(*run_sim)(command, sim_directory);
}
@@ -144,7 +144,7 @@ static void ReadOutputsAndUpdateRuntime(TVMArgs args, size_t num_inputs,
static void CleanUp(TVMArgs args, const std::string& bin_file, const std::string& input_json,
const std::string& input_bin, const std::string& out_bin_prefix,
size_t num_outputs) {
- const auto* clean_up = tvm::runtime::Registry::Get("tvm.mrvl.CleanUpSim");
+ const auto clean_up = tvm::ffi::Function::GetGlobal("tvm.mrvl.CleanUpSim");
(*clean_up)(bin_file, input_json, input_bin, out_bin_prefix, num_outputs);
}
@@ -153,14 +153,14 @@ void tvm::runtime::contrib::mrvl::RunMarvellSimulator(TVMArgs args, const std::s
size_t num_inputs, size_t num_outputs) {
// check $PATH for the presence of MRVL dependent tools/scripts
std::string file_name("mrvl-mlsim");
- const auto* search_path = tvm::runtime::Registry::Get("tvm.mrvl.SearchPath");
+ const auto search_path = tvm::ffi::Function::GetGlobal("tvm.mrvl.SearchPath");
std::string tools_directory = (*search_path)(file_name);
if (tools_directory.empty()) {
ICHECK(false) << "mrvl-mlsim simulator not found! Please specify the path to Marvell "
"tools by adding it to $PATH.";
}
- const auto* temp_dir = tvm::runtime::Registry::Get("tvm.mrvl.TempDir");
+ const auto temp_dir = tvm::ffi::Function::GetGlobal("tvm.mrvl.TempDir");
std::string working_directory = (*temp_dir)();
auto bin_file = working_directory + "/" + symbol_name + ".bin";
auto input_json = working_directory + "/indata.json";
diff --git a/src/runtime/contrib/msc/tensorrt_runtime.cc b/src/runtime/contrib/msc/tensorrt_runtime.cc
index 8b7ba0c4b5..7ddbcb34ad 100644
--- a/src/runtime/contrib/msc/tensorrt_runtime.cc
+++ b/src/runtime/contrib/msc/tensorrt_runtime.cc
@@ -119,8 +119,8 @@ class MSCTensorRTRuntime : public JSONRuntimeBase {
void Run() override {
SetInputOutputBinds();
if (tool_tag_.size() > 0) {
- const auto* pf = runtime::Registry::Get("msc_tool.callback_step");
- ICHECK(pf != nullptr) << "Cannot find msc_tool.callback_step func.";
+ const auto pf = tvm::ffi::Function::GetGlobal("msc_tool.callback_step");
+ ICHECK(pf.has_value()) << "Cannot find msc_tool.callback_step func.";
Map<String, runtime::NDArray> input_datas;
for (const auto& pair : input_bindings_) {
const auto& tensor_name = engine_->getBindingName(pair.first);
@@ -150,8 +150,8 @@ class MSCTensorRTRuntime : public JSONRuntimeBase {
}
}
if (tool_tag_.size() > 0) {
- const auto* pf = runtime::Registry::Get("msc_tool.callback_step");
- ICHECK(pf != nullptr) << "Cannot find msc_tool.callback_step func.";
+ const auto pf = tvm::ffi::Function::GetGlobal("msc_tool.callback_step");
+ ICHECK(pf.has_value()) << "Cannot find msc_tool.callback_step func.";
Map<String, runtime::NDArray> output_datas;
for (int bid = 0; bid < engine_->getNbBindings(); bid++) {
if (input_bindings_.count(bid)) {
diff --git a/src/runtime/contrib/random/random.cc b/src/runtime/contrib/random/random.cc
index 38c2de6555..70a3e59f15 100644
--- a/src/runtime/contrib/random/random.cc
+++ b/src/runtime/contrib/random/random.cc
@@ -69,66 +69,70 @@ RandomThreadLocalEntry* RandomThreadLocalEntry::ThreadLocal() {
return RandomThreadLocalStore::Get();
}
-TVM_REGISTER_GLOBAL("tvm.contrib.random.randint").set_body([](TVMArgs args, TVMRetValue* ret) {
- RandomThreadLocalEntry* entry = RandomThreadLocalEntry::ThreadLocal();
- int64_t low = args[0];
- int64_t high = args[1];
- DLTensor* out = args[2];
- ICHECK_GT(high, low) << "high must be bigger than low";
- ICHECK(out->strides == nullptr);
-
- DLDataType dtype = out->dtype;
- int64_t size = 1;
- for (int i = 0; i < out->ndim; ++i) {
- size *= out->shape[i];
- }
+TVM_REGISTER_GLOBAL("tvm.contrib.random.randint")
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
+ RandomThreadLocalEntry* entry = RandomThreadLocalEntry::ThreadLocal();
+ int64_t low = args[0];
+ int64_t high = args[1];
+ DLTensor* out = args[2];
+ ICHECK_GT(high, low) << "high must be bigger than low";
+ ICHECK(out->strides == nullptr);
+
+ DLDataType dtype = out->dtype;
+ int64_t size = 1;
+ for (int i = 0; i < out->ndim; ++i) {
+ size *= out->shape[i];
+ }
+
+ DLPACK_INTEGER_TYPE_SWITCH(dtype, DType, {
+ int64_t numeric_low = std::numeric_limits<DType>::min();
+ int64_t numeric_high = std::numeric_limits<DType>::max();
+ numeric_high += 1; // exclusive upper bound
+ low = std::max(low, numeric_low);
+ high = std::min(high, numeric_high);
+
+ if (out->device.device_type == kDLCPU) {
+ // file the data with random byte
+ std::generate_n(static_cast<DType*>(out->data), size, [&]() {
+ unsigned rint = entry->random_engine.GetRandInt();
+ return low + rint % (high - low);
+ });
+ } else {
+ LOG(FATAL) << "Do not support random.randint on this device yet";
+ }
+ })
+ });
- DLPACK_INTEGER_TYPE_SWITCH(dtype, DType, {
- int64_t numeric_low = std::numeric_limits<DType>::min();
- int64_t numeric_high = std::numeric_limits<DType>::max();
- numeric_high += 1; // exclusive upper bound
- low = std::max(low, numeric_low);
- high = std::min(high, numeric_high);
-
- if (out->device.device_type == kDLCPU) {
- // file the data with random byte
- std::generate_n(static_cast<DType*>(out->data), size, [&]() {
- unsigned rint = entry->random_engine.GetRandInt();
- return low + rint % (high - low);
- });
- } else {
- LOG(FATAL) << "Do not support random.randint on this device yet";
- }
- })
-});
-
-TVM_REGISTER_GLOBAL("tvm.contrib.random.uniform").set_body([](TVMArgs args, TVMRetValue* ret) {
- RandomThreadLocalEntry* entry = RandomThreadLocalEntry::ThreadLocal();
- double low = args[0];
- double high = args[1];
- DLTensor* out = args[2];
- entry->random_engine.SampleUniform(out, low, high);
-});
-
-TVM_REGISTER_GLOBAL("tvm.contrib.random.normal").set_body([](TVMArgs args, TVMRetValue* ret) {
- RandomThreadLocalEntry* entry = RandomThreadLocalEntry::ThreadLocal();
- double loc = args[0];
- double scale = args[1];
- DLTensor* out = args[2];
- entry->random_engine.SampleNormal(out, loc, scale);
-});
-
-TVM_REGISTER_GLOBAL("tvm.contrib.random.random_fill").set_body([](TVMArgs args, TVMRetValue* ret) {
- RandomThreadLocalEntry* entry = RandomThreadLocalEntry::ThreadLocal();
- DLTensor* out = args[0];
- entry->random_engine.RandomFill(out);
-});
+TVM_REGISTER_GLOBAL("tvm.contrib.random.uniform")
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
+ RandomThreadLocalEntry* entry = RandomThreadLocalEntry::ThreadLocal();
+ double low = args[0];
+ double high = args[1];
+ DLTensor* out = args[2];
+ entry->random_engine.SampleUniform(out, low, high);
+ });
+
+TVM_REGISTER_GLOBAL("tvm.contrib.random.normal")
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
+ RandomThreadLocalEntry* entry = RandomThreadLocalEntry::ThreadLocal();
+ double loc = args[0];
+ double scale = args[1];
+ DLTensor* out = args[2];
+ entry->random_engine.SampleNormal(out, loc, scale);
+ });
+
+TVM_REGISTER_GLOBAL("tvm.contrib.random.random_fill")
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
+ RandomThreadLocalEntry* entry = RandomThreadLocalEntry::ThreadLocal();
+ DLTensor* out = args[0];
+ entry->random_engine.RandomFill(out);
+ });
TVM_REGISTER_GLOBAL("tvm.contrib.random.random_fill_for_measure")
- .set_body([](TVMArgs args, TVMRetValue* ret) -> void {
- static const PackedFunc* curand = Registry::Get("runtime.contrib.curand.RandomFill");
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) -> void {
+ const auto curand = tvm::ffi::Function::GetGlobal("runtime.contrib.curand.RandomFill");
DLTensor* out = args[0];
- if (curand && out->device.device_type == DLDeviceType::kDLCUDA) {
+ if (curand.has_value() && out->device.device_type == DLDeviceType::kDLCUDA) {
if (out->dtype.code == DLDataTypeCode::kDLFloat) {
(*curand)(out);
return;
diff --git a/src/runtime/contrib/rocblas/rocblas.cc b/src/runtime/contrib/rocblas/rocblas.cc
index 57646b45b9..97951eda6e 100644
--- a/src/runtime/contrib/rocblas/rocblas.cc
+++ b/src/runtime/contrib/rocblas/rocblas.cc
@@ -65,44 +65,46 @@ struct RocBlasThreadEntry {
typedef dmlc::ThreadLocalStore<RocBlasThreadEntry> RocBlasThreadStore;
// matrix multiplication for row major
-TVM_REGISTER_GLOBAL("tvm.contrib.rocblas.matmul").set_body([](TVMArgs args, TVMRetValue* ret) {
- DLTensor* A = args[0];
- DLTensor* B = args[1];
- DLTensor* C = args[2];
- bool transa = args[3];
- bool transb = args[4];
- // call gemm for simple compact code.
- ICHECK_EQ(A->ndim, 2);
- ICHECK_EQ(B->ndim, 2);
- ICHECK_EQ(C->ndim, 2);
- ICHECK(C->strides == nullptr);
- ICHECK(B->strides == nullptr);
- ICHECK(A->strides == nullptr);
- ICHECK(TypeMatch(A->dtype, kDLFloat, 32));
- ICHECK(TypeMatch(B->dtype, kDLFloat, 32));
- ICHECK(TypeMatch(C->dtype, kDLFloat, 32));
-
- float alpha = 1.0;
- float beta = 0.0;
- float* A_ptr = reinterpret_cast<float*>(static_cast<char*>(A->data) + A->byte_offset);
- float* B_ptr = reinterpret_cast<float*>(static_cast<char*>(B->data) + B->byte_offset);
- float* C_ptr = reinterpret_cast<float*>(static_cast<char*>(C->data) + C->byte_offset);
-
- rocblas_operation roc_trans_A = transa ? rocblas_operation_transpose : rocblas_operation_none;
- rocblas_operation roc_trans_B = transb ? rocblas_operation_transpose : rocblas_operation_none;
- size_t N = transb ? B->shape[0] : B->shape[1];
- size_t M = transa ? A->shape[1] : A->shape[0];
- size_t K = transb ? B->shape[1] : B->shape[0];
- size_t lda = transa ? M : K;
- size_t ldb = transb ? K : N;
- size_t ldc = N;
-
- CHECK_ROCBLAS_ERROR(rocblas_sgemm(RocBlasThreadStore::Get()->handle, roc_trans_B, roc_trans_A, N,
- M, K, &alpha, B_ptr, ldb, A_ptr, lda, &beta, C_ptr, ldc));
-});
+TVM_REGISTER_GLOBAL("tvm.contrib.rocblas.matmul")
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
+ DLTensor* A = args[0];
+ DLTensor* B = args[1];
+ DLTensor* C = args[2];
+ bool transa = args[3];
+ bool transb = args[4];
+ // call gemm for simple compact code.
+ ICHECK_EQ(A->ndim, 2);
+ ICHECK_EQ(B->ndim, 2);
+ ICHECK_EQ(C->ndim, 2);
+ ICHECK(C->strides == nullptr);
+ ICHECK(B->strides == nullptr);
+ ICHECK(A->strides == nullptr);
+ ICHECK(TypeMatch(A->dtype, kDLFloat, 32));
+ ICHECK(TypeMatch(B->dtype, kDLFloat, 32));
+ ICHECK(TypeMatch(C->dtype, kDLFloat, 32));
+
+ float alpha = 1.0;
+ float beta = 0.0;
+ float* A_ptr = reinterpret_cast<float*>(static_cast<char*>(A->data) + A->byte_offset);
+ float* B_ptr = reinterpret_cast<float*>(static_cast<char*>(B->data) + B->byte_offset);
+ float* C_ptr = reinterpret_cast<float*>(static_cast<char*>(C->data) + C->byte_offset);
+
+ rocblas_operation roc_trans_A = transa ? rocblas_operation_transpose : rocblas_operation_none;
+ rocblas_operation roc_trans_B = transb ? rocblas_operation_transpose : rocblas_operation_none;
+ size_t N = transb ? B->shape[0] : B->shape[1];
+ size_t M = transa ? A->shape[1] : A->shape[0];
+ size_t K = transb ? B->shape[1] : B->shape[0];
+ size_t lda = transa ? M : K;
+ size_t ldb = transb ? K : N;
+ size_t ldc = N;
+
+ CHECK_ROCBLAS_ERROR(rocblas_sgemm(RocBlasThreadStore::Get()->handle, roc_trans_B, roc_trans_A,
+ N, M, K, &alpha, B_ptr, ldb, A_ptr, lda, &beta, C_ptr,
+ ldc));
+ });
TVM_REGISTER_GLOBAL("tvm.contrib.rocblas.batch_matmul")
- .set_body([](TVMArgs args, TVMRetValue* ret) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
DLTensor* A = args[0];
DLTensor* B = args[1];
DLTensor* C = args[2];
diff --git a/src/runtime/contrib/sort/sort.cc b/src/runtime/contrib/sort/sort.cc
index 381e94db33..81f498fc6b 100644
--- a/src/runtime/contrib/sort/sort.cc
+++ b/src/runtime/contrib/sort/sort.cc
@@ -77,80 +77,81 @@ struct float16 {
// If input tensor has dimension (d0, d1, ..., d(k-1), dk, d(k+1), ..., d(n-1))
// and sort axis is dk. sort_num should have dimension of
// (d1, d2, ..., d(k-1), d(k+1), ..., dn).
-TVM_REGISTER_GLOBAL("tvm.contrib.sort.argsort_nms").set_body([](TVMArgs args, TVMRetValue* ret) {
- DLTensor* input = args[0];
- DLTensor* sort_num = args[1];
- DLTensor* output = args[2];
- int32_t axis = args[3];
- bool is_ascend = args[4];
-
- auto dtype = input->dtype;
- auto data_ptr = static_cast<float*>(input->data);
- auto sort_num_ptr = static_cast<int32_t*>(sort_num->data);
- std::vector<std::pair<int32_t, float>> sorter;
- int64_t axis_mul_before = 1;
- int64_t axis_mul_after = 1;
-
- if (axis < 0) {
- axis = input->ndim + axis;
- }
+TVM_REGISTER_GLOBAL("tvm.contrib.sort.argsort_nms")
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
+ DLTensor* input = args[0];
+ DLTensor* sort_num = args[1];
+ DLTensor* output = args[2];
+ int32_t axis = args[3];
+ bool is_ascend = args[4];
+
+ auto dtype = input->dtype;
+ auto data_ptr = static_cast<float*>(input->data);
+ auto sort_num_ptr = static_cast<int32_t*>(sort_num->data);
+ std::vector<std::pair<int32_t, float>> sorter;
+ int64_t axis_mul_before = 1;
+ int64_t axis_mul_after = 1;
+
+ if (axis < 0) {
+ axis = input->ndim + axis;
+ }
- // Currently only supports input dtype to be float32.
- ICHECK_EQ(dtype.code, 2) << "Currently only supports input dtype "
- "to be float.";
+ // Currently only supports input dtype to be float32.
+ ICHECK_EQ(dtype.code, 2) << "Currently only supports input dtype "
+ "to be float.";
#if (__ARM_FEATURE_FP16_SCALAR_ARITHMETIC != 1)
- ICHECK_EQ(dtype.bits, 32) << "Currently only supports input dtype "
- "to be float32.";
+ ICHECK_EQ(dtype.bits, 32) << "Currently only supports input dtype "
+ "to be float32.";
#endif
- ICHECK_LT(axis, input->ndim) << "Axis out of boundary for "
- "input ndim "
- << input->ndim;
-
- for (int i = 0; i < input->ndim; ++i) {
- if (i < axis) {
- axis_mul_before *= input->shape[i];
- } else if (i > axis) {
- axis_mul_after *= input->shape[i];
- }
- }
-
- for (int64_t i = 0; i < axis_mul_before; ++i) {
- for (int64_t j = 0; j < axis_mul_after; ++j) {
- sorter.clear();
- int32_t current_sort_num = *(sort_num_ptr + i * axis_mul_after + j);
- int64_t base_idx = i * input->shape[axis] * axis_mul_after + j;
- for (int64_t k = 0; k < current_sort_num; ++k) {
- int64_t full_idx = base_idx + k * axis_mul_after;
- sorter.emplace_back(std::make_pair(k, *(data_ptr + full_idx)));
+ ICHECK_LT(axis, input->ndim) << "Axis out of boundary for "
+ "input ndim "
+ << input->ndim;
+
+ for (int i = 0; i < input->ndim; ++i) {
+ if (i < axis) {
+ axis_mul_before *= input->shape[i];
+ } else if (i > axis) {
+ axis_mul_after *= input->shape[i];
+ }
}
- if (is_ascend) {
+
+ for (int64_t i = 0; i < axis_mul_before; ++i) {
+ for (int64_t j = 0; j < axis_mul_after; ++j) {
+ sorter.clear();
+ int32_t current_sort_num = *(sort_num_ptr + i * axis_mul_after + j);
+ int64_t base_idx = i * input->shape[axis] * axis_mul_after + j;
+ for (int64_t k = 0; k < current_sort_num; ++k) {
+ int64_t full_idx = base_idx + k * axis_mul_after;
+ sorter.emplace_back(std::make_pair(k, *(data_ptr + full_idx)));
+ }
+ if (is_ascend) {
#if (__ARM_FEATURE_FP16_SCALAR_ARITHMETIC == 1)
- if (dtype.bits == 16) {
- std::stable_sort(sorter.begin(), sorter.end(), CompareAscend<__fp16>);
- } else {
+ if (dtype.bits == 16) {
+ std::stable_sort(sorter.begin(), sorter.end(), CompareAscend<__fp16>);
+ } else {
#endif
- std::stable_sort(sorter.begin(), sorter.end(), CompareAscend<float>);
+ std::stable_sort(sorter.begin(), sorter.end(), CompareAscend<float>);
#if (__ARM_FEATURE_FP16_SCALAR_ARITHMETIC == 1)
- }
+ }
#endif
- } else {
+ } else {
#if (__ARM_FEATURE_FP16_SCALAR_ARITHMETIC == 1)
- if (dtype.bits == 16) {
- std::stable_sort(sorter.begin(), sorter.end(), CompareDescend<__fp16>);
- } else {
+ if (dtype.bits == 16) {
+ std::stable_sort(sorter.begin(), sorter.end(), CompareDescend<__fp16>);
+ } else {
#endif
- std::stable_sort(sorter.begin(), sorter.end(), CompareDescend<float>);
+ std::stable_sort(sorter.begin(), sorter.end(), CompareDescend<float>);
#if (__ARM_FEATURE_FP16_SCALAR_ARITHMETIC == 1)
- }
+ }
#endif
+ }
+ for (int32_t k = 0; k < input->shape[axis]; ++k) {
+ *(static_cast<int32_t*>(output->data) + base_idx + k * axis_mul_after) =
+ k < static_cast<int32_t>(sorter.size()) ? sorter[k].first : k;
+ }
+ }
}
- for (int32_t k = 0; k < input->shape[axis]; ++k) {
- *(static_cast<int32_t*>(output->data) + base_idx + k * axis_mul_after) =
- k < static_cast<int32_t>(sorter.size()) ? sorter[k].first : k;
- }
- }
- }
-});
+ });
template <typename DataType, typename OutType>
void sort_impl(
@@ -215,7 +216,7 @@ void sort(DLTensor* input, DLTensor* output, int32_t axis, bool is_ascend) {
// If input tensor has dimension (d0, d1, ..., d(k-1), dk, d(k+1), ..., d(n-1))
// and sort axis is dk. sort_num should have dimension of
// (d1, d2, ..., d(k-1), d(k+1), ..., dn).
-TVM_REGISTER_GLOBAL("tvm.contrib.sort.argsort").set_body([](TVMArgs args, TVMRetValue* ret) {
+TVM_REGISTER_GLOBAL("tvm.contrib.sort.argsort").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
DLTensor* input = args[0];
DLTensor* output = args[1];
int32_t axis = args[2];
@@ -310,7 +311,7 @@ TVM_REGISTER_GLOBAL("tvm.contrib.sort.argsort").set_body([](TVMArgs args, TVMRet
// If input tensor has dimension (d0, d1, ..., d(k-1), dk, d(k+1), ..., d(n-1))
// and sort axis is dk. sort_num should have dimension of
// (d1, d2, ..., d(k-1), d(k+1), ..., dn).
-TVM_REGISTER_GLOBAL("tvm.contrib.sort.sort").set_body([](TVMArgs args, TVMRetValue* ret) {
+TVM_REGISTER_GLOBAL("tvm.contrib.sort.sort").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
DLTensor* input = args[0];
DLTensor* output = args[1];
int32_t axis = args[2];
@@ -439,7 +440,7 @@ void topk(DLTensor* input, DLTensor* out_values, DLTensor* out_indices, int k, i
// If input tensor has dimension (d0, d1, ..., d(k-1), dk, d(k+1), ..., d(n-1))
// and sort axis is dk. sort_num should have dimension of
// (d1, d2, ..., d(k-1), d(k+1), ..., dn).
-TVM_REGISTER_GLOBAL("tvm.contrib.sort.topk").set_body([](TVMArgs args, TVMRetValue* ret) {
+TVM_REGISTER_GLOBAL("tvm.contrib.sort.topk").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
DLTensor* input = args[0];
DLTensor* values_out = nullptr;
DLTensor* indices_out = nullptr;
diff --git a/src/runtime/contrib/tflite/tflite_runtime.cc b/src/runtime/contrib/tflite/tflite_runtime.cc
index 17ca44174b..70231eac01 100644
--- a/src/runtime/contrib/tflite/tflite_runtime.cc
+++ b/src/runtime/contrib/tflite/tflite_runtime.cc
@@ -180,7 +180,7 @@ Module TFLiteRuntimeCreate(const std::string& tflite_model_bytes, Device dev) {
return Module(exec);
}
-TVM_REGISTER_GLOBAL("tvm.tflite_runtime.create").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("tvm.tflite_runtime.create").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = TFLiteRuntimeCreate(args[0], args[1]);
});
diff --git a/src/runtime/contrib/thrust/thrust.cu b/src/runtime/contrib/thrust/thrust.cu
index eaaad534ff..754a5942fb 100644
--- a/src/runtime/contrib/thrust/thrust.cu
+++ b/src/runtime/contrib/thrust/thrust.cu
@@ -232,7 +232,7 @@ void thrust_sort_common(DLTensor* input, DLTensor* values_out, DLTensor* indices
}
}
-TVM_REGISTER_GLOBAL("tvm.contrib.thrust.sort").set_body([](TVMArgs args, TVMRetValue* ret) {
+TVM_REGISTER_GLOBAL("tvm.contrib.thrust.sort").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
ICHECK_GE(args.num_args, 4);
DLTensor* input = args[0];
DLTensor* values_out = args[1];
@@ -279,7 +279,7 @@ void thrust_stable_sort_by_key(DLTensor* keys_in, DLTensor* values_in, DLTensor*
}
TVM_REGISTER_GLOBAL("tvm.contrib.thrust.stable_sort_by_key")
- .set_body([](TVMArgs args, TVMRetValue* ret) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
ICHECK_GE(args.num_args, 5);
DLTensor* keys_in = args[0];
DLTensor* values_in = args[1];
@@ -393,7 +393,7 @@ void thrust_scan(DLTensor* data, DLTensor* output, bool exclusive, DLTensor* wor
}
}
-TVM_REGISTER_GLOBAL("tvm.contrib.thrust.sum_scan").set_body([](TVMArgs args, TVMRetValue* ret) {
+TVM_REGISTER_GLOBAL("tvm.contrib.thrust.sum_scan").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
ICHECK(args.num_args == 2 || args.num_args == 3 || args.num_args == 4);
DLTensor* data = args[0];
DLTensor* output = args[1];
diff --git a/src/runtime/cpu_device_api.cc b/src/runtime/cpu_device_api.cc
index 5dd470f00d..867f262264 100644
--- a/src/runtime/cpu_device_api.cc
+++ b/src/runtime/cpu_device_api.cc
@@ -150,7 +150,7 @@ void CPUDeviceAPI::FreeWorkspace(Device dev, void* data) {
dmlc::ThreadLocalStore<CPUWorkspacePool>::Get()->FreeWorkspace(dev, data);
}
-TVM_REGISTER_GLOBAL("device_api.cpu").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("device_api.cpu").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
DeviceAPI* ptr = CPUDeviceAPI::Global();
*rv = static_cast<void*>(ptr);
});
diff --git a/src/runtime/cuda/cuda_device_api.cc b/src/runtime/cuda/cuda_device_api.cc
index 82b8d90626..79f54494ce 100644
--- a/src/runtime/cuda/cuda_device_api.cc
+++ b/src/runtime/cuda/cuda_device_api.cc
@@ -286,12 +286,12 @@ CUDAThreadEntry::CUDAThreadEntry() : pool(kDLCUDA, CUDADeviceAPI::Global()) {}
CUDAThreadEntry* CUDAThreadEntry::ThreadLocal() { return CUDAThreadStore::Get(); }
-TVM_REGISTER_GLOBAL("device_api.cuda").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("device_api.cuda").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
DeviceAPI* ptr = CUDADeviceAPI::Global();
*rv = static_cast<void*>(ptr);
});
-TVM_REGISTER_GLOBAL("device_api.cuda_host").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("device_api.cuda_host").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
DeviceAPI* ptr = CUDADeviceAPI::Global();
*rv = static_cast<void*>(ptr);
});
diff --git a/src/runtime/cuda/l2_cache_flush.cc b/src/runtime/cuda/l2_cache_flush.cc
index 6b2c466530..23eddac2be 100644
--- a/src/runtime/cuda/l2_cache_flush.cc
+++ b/src/runtime/cuda/l2_cache_flush.cc
@@ -32,7 +32,7 @@ typedef dmlc::ThreadLocalStore<L2Flush> L2FlushStore;
L2Flush* L2Flush::ThreadLocal() { return L2FlushStore::Get(); }
-TVM_REGISTER_GLOBAL("l2_cache_flush_cuda").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("l2_cache_flush_cuda").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
ICHECK(L2Flush::ThreadLocal() != nullptr) << "L2Flush::ThreadLocal do not exist.";
cudaStream_t stream = CUDAThreadEntry::ThreadLocal()->stream;
L2Flush::ThreadLocal()->Flush(stream);
diff --git a/src/runtime/disco/bcast_session.cc b/src/runtime/disco/bcast_session.cc
index 5645c372b2..5effbe37a4 100644
--- a/src/runtime/disco/bcast_session.cc
+++ b/src/runtime/disco/bcast_session.cc
@@ -69,9 +69,9 @@ void BcastSessionObj::Shutdown() {
}
void BcastSessionObj::InitCCL(String ccl, IntTuple device_ids) {
- const auto* pf = runtime::Registry::Get("runtime.disco." + ccl + ".init_ccl");
- CHECK(pf) << "ValueError: Cannot initialize CCL `" << ccl
- << "`, because cannot find function: runtime.disco." << ccl << ".init_ccl";
+ const auto pf = tvm::ffi::Function::GetGlobal("runtime.disco." + ccl + ".init_ccl");
+ CHECK(pf.has_value()) << "ValueError: Cannot initialize CCL `" << ccl
+ << "`, because cannot find function: runtime.disco." << ccl << ".init_ccl";
(*pf)(GetRef<Session>(this), device_ids);
}
diff --git a/src/runtime/disco/builtin.cc b/src/runtime/disco/builtin.cc
index 760a330a7a..57707656c3 100644
--- a/src/runtime/disco/builtin.cc
+++ b/src/runtime/disco/builtin.cc
@@ -70,12 +70,12 @@ NDArray DiscoEmptyNDArray(ShapeTuple shape, DataType dtype, Device device) {
return NDArray::Empty(shape, dtype, UseDefaultDeviceIfNone(device));
}
-const PackedFunc& GetCCLFunc(const char* name) {
+PackedFunc GetCCLFunc(const char* name) {
std::string ccl = DiscoWorker::ThreadLocal()->ccl;
std::string pf_name = "runtime.disco." + ccl + "." + name;
- const PackedFunc* pf = tvm::runtime::Registry::Get(pf_name);
- CHECK(pf != nullptr) << "ValueError: Cannot find the `" << name << "` function for `" << ccl
- << "` via `" << pf_name << "`";
+ const auto pf = tvm::ffi::Function::GetGlobal(pf_name);
+ CHECK(pf.has_value()) << "ValueError: Cannot find the `" << name << "` function for `" << ccl
+ << "` via `" << pf_name << "`";
return *pf;
}
@@ -164,10 +164,9 @@ TVM_REGISTER_GLOBAL("runtime.disco.device").set_body_typed([]() -> Device {
TVM_REGISTER_GLOBAL("runtime.disco.bind_worker_to_cpu_core").set_body_typed([](IntTuple cpu_ids) {
int worker_id = WorkerId();
ICHECK_LT(worker_id, static_cast<int>(cpu_ids.size()));
- const PackedFunc* f_set_thread_affinity =
- Registry::Get("tvm.runtime.threading.set_current_thread_affinity");
- ICHECK_NOTNULL(f_set_thread_affinity);
- (*f_set_thread_affinity)(IntTuple{cpu_ids[worker_id]});
+ const auto f_set_thread_affinity =
+ tvm::ffi::Function::GetGlobalRequired("tvm.runtime.threading.set_current_thread_affinity");
+ f_set_thread_affinity(IntTuple{cpu_ids[worker_id]});
});
} // namespace runtime
diff --git a/src/runtime/disco/disco_worker.cc b/src/runtime/disco/disco_worker.cc
index 1dfd98f543..541d293a5a 100644
--- a/src/runtime/disco/disco_worker.cc
+++ b/src/runtime/disco/disco_worker.cc
@@ -106,8 +106,8 @@ struct DiscoWorker::Impl {
static void Shutdown(DiscoWorker* self) {}
static void GetGlobalFunc(DiscoWorker* self, int reg_id, const std::string& name) {
- const PackedFunc* pf = runtime::Registry::Get(name);
- CHECK(pf) << "ValueError: Cannot find global function: " << name;
+ const auto pf = tvm::ffi::Function::GetGlobal(name);
+ CHECK(pf.has_value()) << "ValueError: Cannot find global function: " << name;
if (reg_id != 0) {
GetReg(self, reg_id) = *pf;
}
diff --git a/src/runtime/disco/distributed/socket_session.cc b/src/runtime/disco/distributed/socket_session.cc
index 22b567e8e0..f56c6d35d4 100644
--- a/src/runtime/disco/distributed/socket_session.cc
+++ b/src/runtime/disco/distributed/socket_session.cc
@@ -57,9 +57,9 @@ class SocketSessionObj : public BcastSessionObj {
explicit SocketSessionObj(int num_nodes, int num_workers_per_node, int num_groups,
const String& host, int port)
: num_nodes_(num_nodes), num_workers_per_node_(num_workers_per_node) {
- const PackedFunc* f_create_local_session =
- Registry::Get("runtime.disco.create_socket_session_local_workers");
- ICHECK(f_create_local_session != nullptr)
+ const auto f_create_local_session =
+ tvm::ffi::Function::GetGlobal("runtime.disco.create_socket_session_local_workers");
+ ICHECK(f_create_local_session.has_value())
<< "Cannot find function runtime.disco.create_socket_session_local_workers";
local_session_ = ((*f_create_local_session)(num_workers_per_node)).operator BcastSession();
DRef f_init_workers =
@@ -269,8 +269,8 @@ class RemoteSocketSession {
private:
void InitLocalSession() {
- const PackedFunc* f_create_local_session =
- Registry::Get("runtime.disco.create_socket_session_local_workers");
+ const auto f_create_local_session =
+ tvm::ffi::Function::GetGlobal("runtime.disco.create_socket_session_local_workers");
local_session_ = ((*f_create_local_session)(num_workers_per_node_)).operator BcastSession();
DRef f_init_workers =
diff --git a/src/runtime/disco/loader.cc b/src/runtime/disco/loader.cc
index 7ed1ffb7f7..79618fcd0f 100644
--- a/src/runtime/disco/loader.cc
+++ b/src/runtime/disco/loader.cc
@@ -198,7 +198,7 @@ ObjectRef ShardLoaderObj::Create(const std::string& path_to_metadata, const std:
const std::string& name = shard_func.name;
if (PackedFunc f = mod.defined() ? mod->GetFunction(name, true) : nullptr; f != nullptr) {
n->shard_funcs_[name] = f;
- } else if (const PackedFunc* f = runtime::Registry::Get(name)) {
+ } else if (const auto f = tvm::ffi::Function::GetGlobal(name)) {
n->shard_funcs_[name] = *f;
} else {
LOG(FATAL) << "ValueError: Undefined function: " << name;
diff --git a/src/runtime/disco/nccl/nccl_context.h b/src/runtime/disco/nccl/nccl_context.h
index 6c1eaf749a..aa466c5263 100644
--- a/src/runtime/disco/nccl/nccl_context.h
+++ b/src/runtime/disco/nccl/nccl_context.h
@@ -149,8 +149,8 @@ struct CCLThreadLocalContext {
}
deviceStream_t GetDefaultStream() {
- const auto* func = tvm::runtime::Registry::Get("runtime.get_" TVM_DISCO_DEVICE_NAME "_stream");
- ICHECK(func != nullptr);
+ const auto func = tvm::ffi::Function::GetGlobal("runtime.get_" TVM_DISCO_DEVICE_NAME "_stream");
+ ICHECK(func.has_value());
deviceStream_t stream = static_cast<deviceStream_t>((*func)().operator void*());
return stream == nullptr ? default_stream : stream;
}
diff --git a/src/runtime/disco/process_session.cc b/src/runtime/disco/process_session.cc
index cebe1f898f..eb8c3c7355 100644
--- a/src/runtime/disco/process_session.cc
+++ b/src/runtime/disco/process_session.cc
@@ -180,7 +180,7 @@ Session Session::ProcessSession(int num_workers, int num_group, String process_p
String entrypoint) {
CHECK_EQ(num_workers % num_group, 0)
<< "The number of workers should be divisible by the number of worker group.";
- const PackedFunc* pf = Registry::Get(process_pool_creator);
+ const auto pf = tvm::ffi::Function::GetGlobal(process_pool_creator);
CHECK(pf) << "ValueError: Cannot find function " << process_pool_creator
<< " in the registry. Please check if it is registered.";
PackedFunc process_pool = (*pf)(num_workers, num_group, entrypoint);
diff --git a/src/runtime/disco/protocol.h b/src/runtime/disco/protocol.h
index 6bcba6e470..c5964cc6df 100644
--- a/src/runtime/disco/protocol.h
+++ b/src/runtime/disco/protocol.h
@@ -232,8 +232,9 @@ inline std::string DiscoDebugObject::SaveToStr() const {
return result;
} else if (auto opt_obj = this->data.as<ObjectRef>()) {
ObjectRef obj = opt_obj.value();
- const PackedFunc* f = runtime::Registry::Get("node.SaveJSON");
- CHECK(f) << "ValueError: Cannot serialize object in non-debugging mode: " << obj->GetTypeKey();
+ const auto f = tvm::ffi::Function::GetGlobal("node.SaveJSON");
+ CHECK(f.has_value()) << "ValueError: Cannot serialize object in non-debugging mode: "
+ << obj->GetTypeKey();
std::string result = (*f)(obj);
result.push_back('0');
return result;
@@ -248,8 +249,8 @@ inline ObjectPtr<DiscoDebugObject> DiscoDebugObject::LoadFromStr(std::string jso
json_str.pop_back();
ObjectPtr<DiscoDebugObject> result = make_object<DiscoDebugObject>();
if (control_bit == '0') {
- const PackedFunc* f = runtime::Registry::Get("node.LoadJSON");
- CHECK(f) << "ValueError: Cannot deserialize object in non-debugging mode";
+ const auto f = tvm::ffi::Function::GetGlobal("node.LoadJSON");
+ CHECK(f.has_value()) << "ValueError: Cannot deserialize object in non-debugging mode";
result->data = (*f)(json_str);
} else if (control_bit == '1') {
dmlc::MemoryStringStream mstrm(&json_str);
diff --git a/src/runtime/disco/session.cc b/src/runtime/disco/session.cc
index 7113360eda..bea6f6eef8 100644
--- a/src/runtime/disco/session.cc
+++ b/src/runtime/disco/session.cc
@@ -34,28 +34,25 @@ TVM_REGISTER_OBJECT_TYPE(DRefObj);
TVM_REGISTER_OBJECT_TYPE(SessionObj);
TVM_REGISTER_GLOBAL("runtime.disco.SessionThreaded").set_body_typed(Session::ThreadedSession);
TVM_REGISTER_GLOBAL("runtime.disco.DRefDebugGetFromRemote")
- .set_body_method<DRef>(&DRefObj::DebugGetFromRemote);
-TVM_REGISTER_GLOBAL("runtime.disco.DRefDebugCopyFrom")
- .set_body_method<DRef>(&DRefObj::DebugCopyFrom);
+ .set_body_method(&DRefObj::DebugGetFromRemote);
+TVM_REGISTER_GLOBAL("runtime.disco.DRefDebugCopyFrom").set_body_method(&DRefObj::DebugCopyFrom);
TVM_REGISTER_GLOBAL("runtime.disco.SessionGetNumWorkers")
- .set_body_method<Session>(&SessionObj::GetNumWorkers);
+ .set_body_method(&SessionObj::GetNumWorkers);
TVM_REGISTER_GLOBAL("runtime.disco.SessionGetGlobalFunc")
- .set_body_method<Session>(&SessionObj::GetGlobalFunc);
+ .set_body_method(&SessionObj::GetGlobalFunc);
TVM_REGISTER_GLOBAL("runtime.disco.SessionCopyFromWorker0")
- .set_body_method<Session>(&SessionObj::CopyFromWorker0);
+ .set_body_method(&SessionObj::CopyFromWorker0);
TVM_REGISTER_GLOBAL("runtime.disco.SessionCopyToWorker0")
- .set_body_method<Session>(&SessionObj::CopyToWorker0);
-TVM_REGISTER_GLOBAL("runtime.disco.SessionSyncWorker")
- .set_body_method<Session>(&SessionObj::SyncWorker);
+ .set_body_method(&SessionObj::CopyToWorker0);
+TVM_REGISTER_GLOBAL("runtime.disco.SessionSyncWorker").set_body_method(&SessionObj::SyncWorker);
TVM_REGISTER_GLOBAL("runtime.disco.SessionInitCCL") //
- .set_body_method<Session>(&SessionObj::InitCCL);
+ .set_body_method(&SessionObj::InitCCL);
TVM_REGISTER_GLOBAL("runtime.disco.SessionCallPacked")
.set_body_packed([](TVMArgs args, TVMRetValue* rv) {
Session self = args[0];
*rv = SessionObj::FFI::CallWithPacked(self, args.Slice(1));
});
-TVM_REGISTER_GLOBAL("runtime.disco.SessionShutdown")
- .set_body_method<Session>(&SessionObj::Shutdown);
+TVM_REGISTER_GLOBAL("runtime.disco.SessionShutdown").set_body_method(&SessionObj::Shutdown);
} // namespace runtime
} // namespace tvm
diff --git a/src/runtime/dso_library.cc b/src/runtime/dso_library.cc
index e4f4937a8a..f162499787 100644
--- a/src/runtime/dso_library.cc
+++ b/src/runtime/dso_library.cc
@@ -149,9 +149,10 @@ ObjectPtr<Library> CreateDSOLibraryObject(std::string library_path) {
return n;
}
-TVM_REGISTER_GLOBAL("runtime.module.loadfile_so").set_body([](TVMArgs args, TVMRetValue* rv) {
- ObjectPtr<Library> n = CreateDSOLibraryObject(args[0]);
- *rv = CreateModuleFromLibrary(n);
-});
+TVM_REGISTER_GLOBAL("runtime.module.loadfile_so")
+ .set_body_packed([](TVMArgs args, TVMRetValue* rv) {
+ ObjectPtr<Library> n = CreateDSOLibraryObject(args[0]);
+ *rv = CreateModuleFromLibrary(n);
+ });
} // namespace runtime
} // namespace tvm
diff --git a/src/runtime/hexagon/hexagon_common.cc b/src/runtime/hexagon/hexagon_common.cc
index 79acbf8b99..60a4ab3b4c 100644
--- a/src/runtime/hexagon/hexagon_common.cc
+++ b/src/runtime/hexagon/hexagon_common.cc
@@ -89,10 +89,11 @@ void LogMessageImpl(const std::string& file, int lineno, int level, const std::s
}
} // namespace detail
-TVM_REGISTER_GLOBAL("runtime.module.loadfile_hexagon").set_body([](TVMArgs args, TVMRetValue* rv) {
- ObjectPtr<Library> n = CreateDSOLibraryObject(args[0]);
- *rv = CreateModuleFromLibrary(n);
-});
+TVM_REGISTER_GLOBAL("runtime.module.loadfile_hexagon")
+ .set_body_packed([](TVMArgs args, TVMRetValue* rv) {
+ ObjectPtr<Library> n = CreateDSOLibraryObject(args[0]);
+ *rv = CreateModuleFromLibrary(n);
+ });
} // namespace runtime
} // namespace tvm
diff --git a/src/runtime/hexagon/hexagon_device_api.cc b/src/runtime/hexagon/hexagon_device_api.cc
index 65162e7cc6..9eda329cdd 100644
--- a/src/runtime/hexagon/hexagon_device_api.cc
+++ b/src/runtime/hexagon/hexagon_device_api.cc
@@ -191,7 +191,7 @@ void HexagonDeviceAPI::CopyDataFromTo(const void* from, size_t from_offset, void
}
TVM_REGISTER_GLOBAL("device_api.hexagon.dma_copy_dltensor")
- .set_body([](TVMArgs args, TVMRetValue* rv) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* rv) {
DLTensor* dst = args[0];
DLTensor* src = args[1];
int size = args[2];
@@ -209,102 +209,107 @@ TVM_REGISTER_GLOBAL("device_api.hexagon.dma_copy_dltensor")
*rv = static_cast<int32_t>(0);
});
-TVM_REGISTER_GLOBAL("device_api.hexagon.dma_copy").set_body([](TVMArgs args, TVMRetValue* rv) {
- uint32_t queue_id = static_cast<int>(args[0]);
- void* dst = args[1];
- void* src = args[2];
- uint32_t size = static_cast<int>(args[3]);
- ICHECK(size > 0);
- bool bypass_cache = args[4];
-
- int ret = DMA_RETRY;
- do {
- ret = HexagonDeviceAPI::Global()->UserDMA()->Copy(queue_id, dst, src, size, bypass_cache);
- } while (ret == DMA_RETRY);
- CHECK(ret == DMA_SUCCESS);
- *rv = static_cast<int32_t>(ret);
-});
+TVM_REGISTER_GLOBAL("device_api.hexagon.dma_copy")
+ .set_body_packed([](TVMArgs args, TVMRetValue* rv) {
+ uint32_t queue_id = static_cast<int>(args[0]);
+ void* dst = args[1];
+ void* src = args[2];
+ uint32_t size = static_cast<int>(args[3]);
+ ICHECK(size > 0);
+ bool bypass_cache = args[4];
-TVM_REGISTER_GLOBAL("device_api.hexagon.dma_wait").set_body([](TVMArgs args, TVMRetValue* rv) {
- uint32_t queue_id = static_cast<int>(args[0]);
- int inflight = args[1];
- ICHECK(inflight >= 0);
- HexagonDeviceAPI::Global()->UserDMA()->Wait(queue_id, inflight);
- *rv = static_cast<int32_t>(0);
-});
+ int ret = DMA_RETRY;
+ do {
+ ret = HexagonDeviceAPI::Global()->UserDMA()->Copy(queue_id, dst, src, size, bypass_cache);
+ } while (ret == DMA_RETRY);
+ CHECK(ret == DMA_SUCCESS);
+ *rv = static_cast<int32_t>(ret);
+ });
+
+TVM_REGISTER_GLOBAL("device_api.hexagon.dma_wait")
+ .set_body_packed([](TVMArgs args, TVMRetValue* rv) {
+ uint32_t queue_id = static_cast<int>(args[0]);
+ int inflight = args[1];
+ ICHECK(inflight >= 0);
+ HexagonDeviceAPI::Global()->UserDMA()->Wait(queue_id, inflight);
+ *rv = static_cast<int32_t>(0);
+ });
TVM_REGISTER_GLOBAL("device_api.hexagon.dma_start_group")
- .set_body([](TVMArgs args, TVMRetValue* rv) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* rv) {
uint32_t queue_id = static_cast<int>(args[0]);
HexagonDeviceAPI::Global()->UserDMA()->StartGroup(queue_id);
*rv = static_cast<int32_t>(0);
});
-TVM_REGISTER_GLOBAL("device_api.hexagon.dma_end_group").set_body([](TVMArgs args, TVMRetValue* rv) {
- uint32_t queue_id = static_cast<int>(args[0]);
- HexagonDeviceAPI::Global()->UserDMA()->EndGroup(queue_id);
- *rv = static_cast<int32_t>(0);
-});
+TVM_REGISTER_GLOBAL("device_api.hexagon.dma_end_group")
+ .set_body_packed([](TVMArgs args, TVMRetValue* rv) {
+ uint32_t queue_id = static_cast<int>(args[0]);
+ HexagonDeviceAPI::Global()->UserDMA()->EndGroup(queue_id);
+ *rv = static_cast<int32_t>(0);
+ });
-TVM_REGISTER_GLOBAL("device_api.hexagon.alloc_nd").set_body([](TVMArgs args, TVMRetValue* rv) {
- int32_t device_type = args[0];
- int32_t device_id = args[1];
- int32_t dtype_code_hint = args[2];
- int32_t dtype_bits_hint = args[3];
- std::string scope = args[4];
- CHECK(scope.find("global.vtcm") != std::string::npos);
- int64_t ndim = args[5];
- CHECK((ndim == 1 || ndim == 2) && "Hexagon Device API supports only 1d and 2d allocations");
- int64_t* shape = static_cast<int64_t*>(static_cast<void*>(args[6]));
-
- Device dev;
- dev.device_type = static_cast<DLDeviceType>(device_type);
- dev.device_id = device_id;
-
- DLDataType type_hint;
- type_hint.code = static_cast<decltype(type_hint.code)>(dtype_code_hint);
- type_hint.bits = static_cast<decltype(type_hint.bits)>(dtype_bits_hint);
- type_hint.lanes = 1;
-
- HexagonDeviceAPI* hexapi = HexagonDeviceAPI::Global();
- *rv = hexapi->AllocDataSpace(dev, ndim, shape, type_hint, String(scope));
-});
+TVM_REGISTER_GLOBAL("device_api.hexagon.alloc_nd")
+ .set_body_packed([](TVMArgs args, TVMRetValue* rv) {
+ int32_t device_type = args[0];
+ int32_t device_id = args[1];
+ int32_t dtype_code_hint = args[2];
+ int32_t dtype_bits_hint = args[3];
+ std::string scope = args[4];
+ CHECK(scope.find("global.vtcm") != std::string::npos);
+ int64_t ndim = args[5];
+ CHECK((ndim == 1 || ndim == 2) && "Hexagon Device API supports only 1d and 2d allocations");
+ int64_t* shape = static_cast<int64_t*>(static_cast<void*>(args[6]));
+
+ Device dev;
+ dev.device_type = static_cast<DLDeviceType>(device_type);
+ dev.device_id = device_id;
+
+ DLDataType type_hint;
+ type_hint.code = static_cast<decltype(type_hint.code)>(dtype_code_hint);
+ type_hint.bits = static_cast<decltype(type_hint.bits)>(dtype_bits_hint);
+ type_hint.lanes = 1;
+
+ HexagonDeviceAPI* hexapi = HexagonDeviceAPI::Global();
+ *rv = hexapi->AllocDataSpace(dev, ndim, shape, type_hint, String(scope));
+ });
-TVM_REGISTER_GLOBAL("device_api.hexagon.free_nd").set_body([](TVMArgs args, TVMRetValue* rv) {
- int32_t device_type = args[0];
- int32_t device_id = args[1];
- std::string scope = args[2];
- CHECK(scope.find("global.vtcm") != std::string::npos);
- void* ptr = args[3];
+TVM_REGISTER_GLOBAL("device_api.hexagon.free_nd")
+ .set_body_packed([](TVMArgs args, TVMRetValue* rv) {
+ int32_t device_type = args[0];
+ int32_t device_id = args[1];
+ std::string scope = args[2];
+ CHECK(scope.find("global.vtcm") != std::string::npos);
+ void* ptr = args[3];
- Device dev;
- dev.device_type = static_cast<DLDeviceType>(device_type);
- dev.device_id = device_id;
+ Device dev;
+ dev.device_type = static_cast<DLDeviceType>(device_type);
+ dev.device_id = device_id;
- HexagonDeviceAPI* hexapi = HexagonDeviceAPI::Global();
- hexapi->FreeDataSpace(dev, ptr);
- *rv = static_cast<int32_t>(0);
-});
+ HexagonDeviceAPI* hexapi = HexagonDeviceAPI::Global();
+ hexapi->FreeDataSpace(dev, ptr);
+ *rv = static_cast<int32_t>(0);
+ });
TVM_REGISTER_GLOBAL("device_api.hexagon.acquire_resources")
- .set_body([](TVMArgs args, TVMRetValue* rv) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* rv) {
HexagonDeviceAPI* api = HexagonDeviceAPI::Global();
api->AcquireResources();
});
TVM_REGISTER_GLOBAL("device_api.hexagon.release_resources")
- .set_body([](TVMArgs args, TVMRetValue* rv) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* rv) {
HexagonDeviceAPI* api = HexagonDeviceAPI::Global();
api->ReleaseResources();
});
TVM_REGISTER_GLOBAL("device_api.hexagon.vtcm_device_bytes")
- .set_body([](TVMArgs args, TVMRetValue* rv) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* rv) {
HexagonDeviceAPI* api = HexagonDeviceAPI::Global();
*rv = static_cast<int32_t>(api->VtcmPool()->VtcmDeviceBytes());
});
-TVM_REGISTER_GLOBAL("device_api.hexagon").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("device_api.hexagon").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
DeviceAPI* ptr = HexagonDeviceAPI::Global();
*rv = static_cast<void*>(ptr);
});
diff --git a/src/runtime/hexagon/rpc/android/session.cc b/src/runtime/hexagon/rpc/android/session.cc
index 59f1a2a122..aafd5e82df 100644
--- a/src/runtime/hexagon/rpc/android/session.cc
+++ b/src/runtime/hexagon/rpc/android/session.cc
@@ -110,7 +110,7 @@ class HexagonTransportChannel : public RPCChannel {
};
TVM_REGISTER_GLOBAL("tvm.contrib.hexagon.create_hexagon_session")
- .set_body([](TVMArgs args, TVMRetValue* rv) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* rv) {
ICHECK(args.size() >= 4) << args.size() << " is less than 4";
std::string session_name = args[0];
diff --git a/src/runtime/hexagon/rpc/hexagon/rpc_server.cc b/src/runtime/hexagon/rpc/hexagon/rpc_server.cc
index 7659b685c3..30122d29e8 100644
--- a/src/runtime/hexagon/rpc/hexagon/rpc_server.cc
+++ b/src/runtime/hexagon/rpc/hexagon/rpc_server.cc
@@ -244,7 +244,7 @@ tvm::runtime::hexagon::HexagonRPCServer* get_hexagon_rpc_server(
} // namespace
const tvm::runtime::PackedFunc get_runtime_func(const std::string& name) {
- if (const tvm::runtime::PackedFunc* pf = tvm::runtime::Registry::Get(name)) {
+ if (const auto pf = tvm::ffi::Function::GetGlobal(name)) {
return *pf;
}
return tvm::runtime::PackedFunc();
@@ -253,7 +253,7 @@ const tvm::runtime::PackedFunc get_runtime_func(const std::string& name) {
void reset_device_api() {
const tvm::runtime::PackedFunc api = get_runtime_func("device_api.hexagon");
// Registering device_api.cpu as device_api.hexagon since we use hexagon as sub-target of LLVM.
- tvm::runtime::Registry::Register("device_api.cpu", true).set_body(api);
+ tvm::runtime::Registry::Register("device_api.cpu", true).set_body_packed(api);
}
int __QAIC_HEADER(hexagon_rpc_open)(const char* uri, remote_handle64* handle) {
@@ -330,14 +330,14 @@ __attribute__((weak)) void _Parse_fde_instr() {}
}
TVM_REGISTER_GLOBAL("tvm.hexagon.load_module")
- .set_body([](tvm::runtime::TVMArgs args, tvm::runtime::TVMRetValue* rv) {
+ .set_body_packed([](tvm::runtime::TVMArgs args, tvm::runtime::TVMRetValue* rv) {
std::string soname = args[0];
tvm::ObjectPtr<tvm::runtime::Library> n = tvm::runtime::CreateDSOLibraryObject(soname);
*rv = CreateModuleFromLibrary(n);
});
TVM_REGISTER_GLOBAL("tvm.hexagon.get_profile_output")
- .set_body([](tvm::runtime::TVMArgs args, tvm::runtime::TVMRetValue* rv) {
+ .set_body_packed([](tvm::runtime::TVMArgs args, tvm::runtime::TVMRetValue* rv) {
std::string profiling_mode = args[0];
std::string out_file = args[1];
if (profiling_mode.compare("lwp") == 0) {
@@ -355,7 +355,7 @@ void SaveBinaryToFile(const std::string& file_name, const std::string& data) {
}
TVM_REGISTER_GLOBAL("tvm.rpc.server.upload")
- .set_body([](tvm::runtime::TVMArgs args, tvm::runtime::TVMRetValue* rv) {
+ .set_body_packed([](tvm::runtime::TVMArgs args, tvm::runtime::TVMRetValue* rv) {
std::string file_name = args[0];
std::string data = args[1];
SaveBinaryToFile(file_name, data);
diff --git a/src/runtime/hexagon/rpc/simulator/rpc_server.cc b/src/runtime/hexagon/rpc/simulator/rpc_server.cc
index f4370bd3c8..101e2beb80 100644
--- a/src/runtime/hexagon/rpc/simulator/rpc_server.cc
+++ b/src/runtime/hexagon/rpc/simulator/rpc_server.cc
@@ -301,9 +301,9 @@ int main(int argc, char* argv[]) {
void* self = dlopen(argv[0], RTLD_GLOBAL);
ICHECK(self != nullptr);
- const auto* api = tvm::runtime::Registry::Get("device_api.hexagon");
- ICHECK(api != nullptr);
- tvm::runtime::Registry::Register("device_api.cpu", true).set_body(*api);
+ const auto api = tvm::ffi::Function::GetGlobal("device_api.hexagon");
+ ICHECK(api.has_value());
+ tvm::runtime::Registry::Register("device_api.cpu", true).set_body_packed(*api);
tvm::runtime::hexagon::SimulatorRPCServer server;
@@ -333,14 +333,14 @@ __attribute__((weak)) void _Parse_fde_instr() {}
}
TVM_REGISTER_GLOBAL("tvm.hexagon.load_module")
- .set_body([](tvm::runtime::TVMArgs args, tvm::runtime::TVMRetValue* rv) {
+ .set_body_packed([](tvm::runtime::TVMArgs args, tvm::runtime::TVMRetValue* rv) {
std::string soname = args[0];
tvm::ObjectPtr<tvm::runtime::Library> n = tvm::runtime::CreateDSOLibraryObject(soname);
*rv = CreateModuleFromLibrary(n);
});
TVM_REGISTER_GLOBAL("tvm.hexagon.get_profile_output")
- .set_body([](tvm::runtime::TVMArgs args, tvm::runtime::TVMRetValue* rv) {
+ .set_body_packed([](tvm::runtime::TVMArgs args, tvm::runtime::TVMRetValue* rv) {
std::string profiling_mode = args[0];
std::string out_file = args[1];
if (profiling_mode.compare("lwp") == 0) {
@@ -358,7 +358,7 @@ void SaveBinaryToFile(const std::string& file_name, const std::string& data) {
}
TVM_REGISTER_GLOBAL("tvm.rpc.server.upload")
- .set_body([](tvm::runtime::TVMArgs args, tvm::runtime::TVMRetValue* rv) {
+ .set_body_packed([](tvm::runtime::TVMArgs args, tvm::runtime::TVMRetValue* rv) {
std::string file_name = args[0];
std::string data = args[1];
SaveBinaryToFile(file_name, data);
diff --git a/src/runtime/hexagon/rpc/simulator/session.cc b/src/runtime/hexagon/rpc/simulator/session.cc
index bec400ce60..5d7a550629 100644
--- a/src/runtime/hexagon/rpc/simulator/session.cc
+++ b/src/runtime/hexagon/rpc/simulator/session.cc
@@ -1371,7 +1371,7 @@ std::optional<HEXAPI_Nullptr> SimulatorRPCChannel::to_nullptr(const detail::Mayb
}
TVM_REGISTER_GLOBAL("tvm.contrib.hexagon.create_hexagon_session")
- .set_body([](TVMArgs args, TVMRetValue* rv) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* rv) {
ICHECK(args.size() >= 4) << args.size() << " is less than 4";
std::string session_name = args[0];
diff --git a/src/runtime/library_module.cc b/src/runtime/library_module.cc
index f346c463a0..c4a73f2d29 100644
--- a/src/runtime/library_module.cc
+++ b/src/runtime/library_module.cc
@@ -96,21 +96,11 @@ void InitContextFunctions(std::function<void*(const char*)> fgetsymbol) {
Module LoadModuleFromBinary(const std::string& type_key, dmlc::Stream* stream) {
std::string loadkey = "runtime.module.loadbinary_";
std::string fkey = loadkey + type_key;
- const PackedFunc* f = Registry::Get(fkey);
- if (f == nullptr) {
- std::string loaders = "";
- for (auto reg_name : Registry::ListNames()) {
- std::string name = reg_name;
- if (name.find(loadkey, 0) == 0) {
- if (loaders.size() > 0) {
- loaders += ", ";
- }
- loaders += name.substr(loadkey.size());
- }
- }
+ const auto f = tvm::ffi::Function::GetGlobal(fkey);
+ if (!f.has_value()) {
LOG(FATAL) << "Binary was created using {" << type_key
- << "} but a loader of that name is not registered. Available loaders are " << loaders
- << ". Perhaps you need to recompile with this runtime enabled.";
+ << "} but a loader of that name is not registered."
+ << "Perhaps you need to recompile with this runtime enabled.";
}
return (*f)(static_cast<void*>(stream));
diff --git a/src/runtime/memory/memory_manager.cc b/src/runtime/memory/memory_manager.cc
index bc09f464ab..9ca438d6b0 100644
--- a/src/runtime/memory/memory_manager.cc
+++ b/src/runtime/memory/memory_manager.cc
@@ -145,7 +145,7 @@ std::string DeviceTypeStr(DLDeviceType type) {
Allocator* GetDeviceSpecificAllocator(Device dev, AllocatorType type) {
std::string dev_str = DeviceTypeStr(dev.device_type);
- auto* device_alloc_helper = tvm::runtime::Registry::Get("DeviceAllocator." + dev_str);
+ auto device_alloc_helper = tvm::ffi::Function::GetGlobal("DeviceAllocator." + dev_str);
void* valloc;
Allocator* allocator = nullptr;
if (device_alloc_helper) {
diff --git a/src/runtime/metal/metal_device_api.mm b/src/runtime/metal/metal_device_api.mm
index cbdacb5c09..ead9c1804a 100644
--- a/src/runtime/metal/metal_device_api.mm
+++ b/src/runtime/metal/metal_device_api.mm
@@ -362,7 +362,7 @@ typedef dmlc::ThreadLocalStore<MetalThreadEntry> MetalThreadStore;
MetalThreadEntry* MetalThreadEntry::ThreadLocal() { return MetalThreadStore::Get(); }
-TVM_REGISTER_GLOBAL("device_api.metal").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("device_api.metal").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
DeviceAPI* ptr = MetalWorkspace::Global();
*rv = static_cast<void*>(ptr);
});
diff --git a/src/runtime/module.cc b/src/runtime/module.cc
index 68aa98250d..39a48a8e76 100644
--- a/src/runtime/module.cc
+++ b/src/runtime/module.cc
@@ -36,12 +36,8 @@ namespace runtime {
void ModuleNode::Import(Module other) {
// specially handle rpc
if (!std::strcmp(this->type_key(), "rpc")) {
- static const PackedFunc* fimport_ = nullptr;
- if (fimport_ == nullptr) {
- fimport_ = runtime::Registry::Get("rpc.ImportRemoteModule");
- ICHECK(fimport_ != nullptr);
- }
- (*fimport_)(GetRef<Module>(this), other);
+ static auto fimport_ = tvm::ffi::Function::GetGlobalRequired("rpc.ImportRemoteModule");
+ fimport_(GetRef<Module>(this), other);
return;
}
// cyclic detection.
@@ -84,11 +80,11 @@ Module Module::LoadFromFile(const String& file_name, const String& format) {
}
std::string load_f_name = "runtime.module.loadfile_" + fmt;
VLOG(1) << "Loading module from '" << file_name << "' of format '" << fmt << "'";
- const PackedFunc* f = Registry::Get(load_f_name);
- ICHECK(f != nullptr) << "Loader for `." << format << "` files is not registered,"
- << " resolved to (" << load_f_name << ") in the global registry."
- << "Ensure that you have loaded the correct runtime code, and"
- << "that you are on the correct hardware architecture.";
+ const auto f = tvm::ffi::Function::GetGlobal(load_f_name);
+ ICHECK(f.has_value()) << "Loader for `." << format << "` files is not registered,"
+ << " resolved to (" << load_f_name << ") in the global registry."
+ << "Ensure that you have loaded the correct runtime code, and"
+ << "that you are on the correct hardware architecture.";
Module m = (*f)(file_name, format);
return m;
}
@@ -115,13 +111,14 @@ const PackedFunc* ModuleNode::GetFuncFromEnv(const String& name) {
if (pf != nullptr) break;
}
if (pf == nullptr) {
- const PackedFunc* f = Registry::Get(name);
- ICHECK(f != nullptr) << "Cannot find function " << name
- << " in the imported modules or global registry."
- << " If this involves ops from a contrib library like"
- << " cuDNN, ensure TVM was built with the relevant"
- << " library.";
- return f;
+ const auto f = tvm::ffi::Function::GetGlobal(name);
+ ICHECK(f.has_value()) << "Cannot find function " << name
+ << " in the imported modules or global registry."
+ << " If this involves ops from a contrib library like"
+ << " cuDNN, ensure TVM was built with the relevant"
+ << " library.";
+ import_cache_.insert(std::make_pair(name, std::make_shared<PackedFunc>(*f)));
+ return import_cache_.at(name).get();
} else {
import_cache_.insert(std::make_pair(name, std::make_shared<PackedFunc>(pf)));
return import_cache_.at(name).get();
@@ -160,13 +157,13 @@ bool RuntimeEnabled(const String& target_str) {
} else if (target.length() >= 4 && target.substr(0, 4) == "rocm") {
f_name = "device_api.rocm";
} else if (target.length() >= 4 && target.substr(0, 4) == "llvm") {
- const PackedFunc* pf = runtime::Registry::Get("codegen.llvm_target_enabled");
- if (pf == nullptr) return false;
+ const auto pf = tvm::ffi::Function::GetGlobal("codegen.llvm_target_enabled");
+ if (!pf.has_value()) return false;
return (*pf)(target);
} else {
LOG(FATAL) << "Unknown optional runtime " << target;
}
- return runtime::Registry::Get(f_name) != nullptr;
+ return tvm::ffi::Function::GetGlobal(f_name).has_value();
}
TVM_REGISTER_GLOBAL("runtime.RuntimeEnabled").set_body_typed(RuntimeEnabled);
diff --git a/src/runtime/opencl/opencl_device_api.cc b/src/runtime/opencl/opencl_device_api.cc
index 2e5aff9413..0df96b0b62 100644
--- a/src/runtime/opencl/opencl_device_api.cc
+++ b/src/runtime/opencl/opencl_device_api.cc
@@ -760,34 +760,35 @@ void OpenCLWorkspace::Init(const std::string& type_key, const std::string& devic
initialized_ = true;
}
-TVM_REGISTER_GLOBAL("device_api.opencl.alloc_nd").set_body([](TVMArgs args, TVMRetValue* rv) {
- int32_t device_type = args[0];
- int32_t device_id = args[1];
- int32_t dtype_code_hint = args[2];
- int32_t dtype_bits_hint = args[3];
- std::string scope = args[4];
- CHECK(scope.find("texture") != std::string::npos);
- int64_t ndim = args[5];
- CHECK_EQ(ndim, 2);
- int64_t* shape = static_cast<int64_t*>(static_cast<void*>(args[6]));
- int64_t width = shape[0];
- int64_t height = shape[1];
-
- Device dev;
- dev.device_type = static_cast<DLDeviceType>(device_type);
- dev.device_id = device_id;
-
- DLDataType type_hint;
- type_hint.code = static_cast<decltype(type_hint.code)>(dtype_code_hint);
- type_hint.bits = static_cast<decltype(type_hint.bits)>(dtype_bits_hint);
- type_hint.lanes = 1;
-
- *rv = OpenCLWorkspace::Global()->AllocDataSpace(dev, static_cast<size_t>(width),
- static_cast<size_t>(height), type_hint,
- String("global.texture"));
-});
-
-TVM_REGISTER_GLOBAL("device_api.opencl.free_nd").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("device_api.opencl.alloc_nd")
+ .set_body_packed([](TVMArgs args, TVMRetValue* rv) {
+ int32_t device_type = args[0];
+ int32_t device_id = args[1];
+ int32_t dtype_code_hint = args[2];
+ int32_t dtype_bits_hint = args[3];
+ std::string scope = args[4];
+ CHECK(scope.find("texture") != std::string::npos);
+ int64_t ndim = args[5];
+ CHECK_EQ(ndim, 2);
+ int64_t* shape = static_cast<int64_t*>(static_cast<void*>(args[6]));
+ int64_t width = shape[0];
+ int64_t height = shape[1];
+
+ Device dev;
+ dev.device_type = static_cast<DLDeviceType>(device_type);
+ dev.device_id = device_id;
+
+ DLDataType type_hint;
+ type_hint.code = static_cast<decltype(type_hint.code)>(dtype_code_hint);
+ type_hint.bits = static_cast<decltype(type_hint.bits)>(dtype_bits_hint);
+ type_hint.lanes = 1;
+
+ *rv = OpenCLWorkspace::Global()->AllocDataSpace(dev, static_cast<size_t>(width),
+ static_cast<size_t>(height), type_hint,
+ String("global.texture"));
+ });
+
+TVM_REGISTER_GLOBAL("device_api.opencl.free_nd").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
int32_t device_type = args[0];
int32_t device_id = args[1];
std::string scope = args[2];
@@ -801,7 +802,7 @@ TVM_REGISTER_GLOBAL("device_api.opencl.free_nd").set_body([](TVMArgs args, TVMRe
*rv = static_cast<int32_t>(0);
});
-TVM_REGISTER_GLOBAL("device_api.opencl").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("device_api.opencl").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
DeviceAPI* ptr = OpenCLWorkspace::Global();
*rv = static_cast<void*>(ptr);
});
@@ -891,7 +892,7 @@ class OpenCLPooledAllocator final : public memory::PooledAllocator {
}
};
-TVM_REGISTER_GLOBAL("DeviceAllocator.opencl").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("DeviceAllocator.opencl").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
Allocator* alloc = new OpenCLPooledAllocator();
*rv = static_cast<void*>(alloc);
});
diff --git a/src/runtime/profiling.cc b/src/runtime/profiling.cc
index 267228303a..97a03ba9ae 100644
--- a/src/runtime/profiling.cc
+++ b/src/runtime/profiling.cc
@@ -93,8 +93,9 @@ std::set<DLDeviceType> seen_devices;
std::mutex seen_devices_lock;
Timer Timer::Start(Device dev) {
- auto f = Registry::Get(std::string("profiling.timer.") + DLDeviceType2Str(dev.device_type));
- if (f == nullptr) {
+ auto f = tvm::ffi::Function::GetGlobal(std::string("profiling.timer.") +
+ DLDeviceType2Str(dev.device_type));
+ if (!f.has_value()) {
{
std::lock_guard<std::mutex> lock(seen_devices_lock);
if (seen_devices.find(dev.device_type) == seen_devices.end()) {
@@ -787,7 +788,7 @@ TVM_REGISTER_OBJECT_TYPE(ReportNode);
TVM_REGISTER_OBJECT_TYPE(DeviceWrapperNode);
TVM_REGISTER_OBJECT_TYPE(MetricCollectorNode);
-TVM_REGISTER_GLOBAL("runtime.profiling.AsTable").set_body_method<Report>(&ReportNode::AsTable);
+TVM_REGISTER_GLOBAL("runtime.profiling.AsTable").set_body_method(&ReportNode::AsTable);
TVM_REGISTER_GLOBAL("runtime.profiling.AsCSV").set_body_typed([](Report n) { return n->AsCSV(); });
TVM_REGISTER_GLOBAL("runtime.profiling.AsJSON").set_body_typed([](Report n) {
return n->AsJSON();
diff --git a/src/runtime/regex.cc b/src/runtime/regex.cc
index ef6c068edf..c99ec88a93 100644
--- a/src/runtime/regex.cc
+++ b/src/runtime/regex.cc
@@ -30,10 +30,11 @@ namespace tvm {
namespace runtime {
bool regex_match(const std::string& match_against, const std::string& regex_pattern) {
- const auto* regex_match_func = tvm::runtime::Registry::Get("tvm.runtime.regex_match");
- CHECK(regex_match_func) << "RuntimeError: "
- << "The PackedFunc 'tvm.runtime.regex_match' has not been registered. "
- << "This can occur if the TVM Python library has not yet been imported.";
+ const auto regex_match_func = tvm::ffi::Function::GetGlobal("tvm.runtime.regex_match");
+ CHECK(regex_match_func.has_value())
+ << "RuntimeError: "
+ << "The PackedFunc 'tvm.runtime.regex_match' has not been registered. "
+ << "This can occur if the TVM Python library has not yet been imported.";
return (*regex_match_func)(regex_pattern, match_against);
}
diff --git a/src/runtime/registry.cc b/src/runtime/registry.cc
index d5631d148a..5ea1d7677a 100644
--- a/src/runtime/registry.cc
+++ b/src/runtime/registry.cc
@@ -36,73 +36,6 @@
namespace tvm {
namespace runtime {
-struct Registry::Manager {
- // map storing the functions.
- // We deliberately used raw pointer.
- // This is because PackedFunc can contain callbacks into the host language (Python) and the
- // resource can become invalid because of indeterministic order of destruction and forking.
- // The resources will only be recycled during program exit.
- std::unordered_map<String, Registry*> fmap;
- // mutex
- std::mutex mutex;
-
- Manager() {}
-
- static Manager* Global() {
- // We deliberately leak the Manager instance, to avoid leak sanitizers
- // complaining about the entries in Manager::fmap being leaked at program
- // exit.
- static Manager* inst = new Manager();
- return inst;
- }
-};
-
-Registry& Registry::set_body(PackedFunc f) { // NOLINT(*)
- func_ = f;
- return *this;
-}
-
-Registry& Registry::Register(const String& name, bool can_override) { // NOLINT(*)
- Manager* m = Manager::Global();
- std::lock_guard<std::mutex> lock(m->mutex);
- if (m->fmap.count(name)) {
- ICHECK(can_override) << "Global PackedFunc " << name << " is already registered";
- }
-
- Registry* r = new Registry();
- r->name_ = name;
- m->fmap[name] = r;
- return *r;
-}
-
-bool Registry::Remove(const String& name) {
- Manager* m = Manager::Global();
- std::lock_guard<std::mutex> lock(m->mutex);
- auto it = m->fmap.find(name);
- if (it == m->fmap.end()) return false;
- m->fmap.erase(it);
- return true;
-}
-
-const PackedFunc* Registry::Get(const String& name) {
- Manager* m = Manager::Global();
- std::lock_guard<std::mutex> lock(m->mutex);
- auto it = m->fmap.find(name);
- if (it == m->fmap.end()) return nullptr;
- return &(it->second->func_);
-}
-
-std::vector<String> Registry::ListNames() {
- Manager* m = Manager::Global();
- std::lock_guard<std::mutex> lock(m->mutex);
- std::vector<String> keys;
- keys.reserve(m->fmap.size());
- for (const auto& kv : m->fmap) {
- keys.push_back(kv.first);
- }
- return keys;
-}
-
/*!
* \brief Execution environment specific API registry.
*
@@ -290,15 +223,15 @@ typedef dmlc::ThreadLocalStore<TVMFuncThreadLocalEntry> TVMFuncThreadLocalStore;
int TVMFuncRegisterGlobal(const char* name, TVMFunctionHandle f, int override) {
API_BEGIN();
using tvm::runtime::GetRef;
- tvm::runtime::Registry::Register(name, override != 0)
- .set_body(GetRef<tvm::ffi::Function>(static_cast<tvm::ffi::FunctionObj*>(f)));
+ tvm::ffi::Function::SetGlobal(
+ name, GetRef<tvm::ffi::Function>(static_cast<tvm::ffi::FunctionObj*>(f)), override != 0);
API_END();
}
int TVMFuncGetGlobal(const char* name, TVMFunctionHandle* out) {
API_BEGIN();
- const tvm::runtime::PackedFunc* fp = tvm::runtime::Registry::Get(name);
- if (fp != nullptr) {
+ const auto fp = tvm::ffi::Function::GetGlobal(name);
+ if (fp.has_value()) {
TVMFFIAny val = tvm::ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(tvm::ffi::Any(*fp));
*out = val.v_obj;
} else {
@@ -310,7 +243,7 @@ int TVMFuncGetGlobal(const char* name, TVMFunctionHandle* out) {
int TVMFuncListGlobalNames(int* out_size, const char*** out_array) {
API_BEGIN();
TVMFuncThreadLocalEntry* ret = TVMFuncThreadLocalStore::Get();
- ret->ret_vec_str = tvm::runtime::Registry::ListNames();
+ ret->ret_vec_str = tvm::ffi::Function::ListGlobalNames();
ret->ret_vec_charp.clear();
for (size_t i = 0; i < ret->ret_vec_str.size(); ++i) {
ret->ret_vec_charp.push_back(ret->ret_vec_str[i].c_str());
@@ -322,7 +255,7 @@ int TVMFuncListGlobalNames(int* out_size, const char*** out_array) {
int TVMFuncRemoveGlobal(const char* name) {
API_BEGIN();
- tvm::runtime::Registry::Remove(name);
+ tvm::ffi::Function::RemoveGlobal(name);
API_END();
}
diff --git a/src/runtime/relax_vm/builtin.cc b/src/runtime/relax_vm/builtin.cc
index bfc94d1897..002134c209 100644
--- a/src/runtime/relax_vm/builtin.cc
+++ b/src/runtime/relax_vm/builtin.cc
@@ -358,7 +358,7 @@ Storage VMAllocStorage(void* ctx_ptr, ShapeTuple buffer_shape, Index device_inde
TVM_REGISTER_GLOBAL("vm.builtin.alloc_storage").set_body_typed(VMAllocStorage);
-TVM_REGISTER_GLOBAL("vm.builtin.alloc_tensor").set_body_method<Storage>(&StorageObj::AllocNDArray);
+TVM_REGISTER_GLOBAL("vm.builtin.alloc_tensor").set_body_method(&StorageObj::AllocNDArray);
//-------------------------------------------------
// Closure function handling, calling convention
@@ -472,10 +472,10 @@ TVM_REGISTER_GLOBAL("vm.builtin.invoke_debug_func")
ObjectRef io_effect = args[0];
ICHECK(!io_effect.defined()) << "ValueError: IOEffect is expected to be lowered to None.";
String debug_func_name = args[1];
- const PackedFunc* debug_func = runtime::Registry::Get(debug_func_name);
- CHECK(debug_func) << "ValueError: " << debug_func_name << " is not found. "
- << "Use the decorator `@tvm.register_func(\"" << debug_func_name
- << "\")` to register it.";
+ const auto debug_func = tvm::ffi::Function::GetGlobal(debug_func_name);
+ CHECK(debug_func.has_value()) << "ValueError: " << debug_func_name << " is not found. "
+ << "Use the decorator `@tvm.register_func(\"" << debug_func_name
+ << "\")` to register it.";
String line_info = args[2];
std::vector<AnyView> call_args(num_args + 1);
{
diff --git a/src/runtime/relax_vm/cuda/cuda_graph_builtin.cc b/src/runtime/relax_vm/cuda/cuda_graph_builtin.cc
index 757318dbad..91bf54bdbf 100644
--- a/src/runtime/relax_vm/cuda/cuda_graph_builtin.cc
+++ b/src/runtime/relax_vm/cuda/cuda_graph_builtin.cc
@@ -240,7 +240,7 @@ class CUDAGraphExtension : public VMExtension {
};
TVM_REGISTER_GLOBAL("vm.builtin.cuda_graph.run_or_capture")
- .set_body([](TVMArgs args, TVMRetValue* rv) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* rv) {
ICHECK(args.size() == 5 || args.size() == 4);
VirtualMachine* vm = VirtualMachine::GetContextPtr(args[0]);
auto extension = vm->GetOrCreateExtension<CUDAGraphExtension>();
@@ -255,7 +255,7 @@ TVM_REGISTER_GLOBAL("vm.builtin.cuda_graph.run_or_capture")
});
TVM_REGISTER_GLOBAL("vm.builtin.cuda_graph.get_cached_alloc")
- .set_body([](TVMArgs args, TVMRetValue* rv) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* rv) {
ICHECK_EQ(args.size(), 3);
VirtualMachine* vm = VirtualMachine::GetContextPtr(args[0]);
auto extension = vm->GetOrCreateExtension<CUDAGraphExtension>();
diff --git a/src/runtime/relax_vm/kv_state.cc b/src/runtime/relax_vm/kv_state.cc
index 9a5478411a..7a3d97fe7c 100644
--- a/src/runtime/relax_vm/kv_state.cc
+++ b/src/runtime/relax_vm/kv_state.cc
@@ -31,16 +31,14 @@ TVM_REGISTER_OBJECT_TYPE(AttentionKVCacheObj);
TVM_REGISTER_OBJECT_TYPE(RNNStateObj);
// KV State base methods
-TVM_REGISTER_GLOBAL("vm.builtin.kv_state_clear").set_body_method<KVState>(&KVStateObj::Clear);
-TVM_REGISTER_GLOBAL("vm.builtin.kv_state_add_sequence")
- .set_body_method<KVState>(&KVStateObj::AddSequence);
+TVM_REGISTER_GLOBAL("vm.builtin.kv_state_clear").set_body_method(&KVStateObj::Clear);
+TVM_REGISTER_GLOBAL("vm.builtin.kv_state_add_sequence").set_body_method(&KVStateObj::AddSequence);
TVM_REGISTER_GLOBAL("vm.builtin.kv_state_remove_sequence")
- .set_body_method<KVState>(&KVStateObj::RemoveSequence);
-TVM_REGISTER_GLOBAL("vm.builtin.kv_state_fork_sequence")
- .set_body_method<KVState>(&KVStateObj::ForkSequence);
-TVM_REGISTER_GLOBAL("vm.builtin.kv_state_popn").set_body_method<KVState>(&KVStateObj::PopN);
+ .set_body_method(&KVStateObj::RemoveSequence);
+TVM_REGISTER_GLOBAL("vm.builtin.kv_state_fork_sequence").set_body_method(&KVStateObj::ForkSequence);
+TVM_REGISTER_GLOBAL("vm.builtin.kv_state_popn").set_body_method(&KVStateObj::PopN);
TVM_REGISTER_GLOBAL("vm.builtin.kv_state_begin_forward")
- .set_body([](TVMArgs args, TVMRetValue* rv) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* rv) {
CHECK(args.size() == 3 || args.size() == 4)
<< "KVState BeginForward only accepts 3 or 4 arguments";
KVState kv_state = args[0];
@@ -52,30 +50,29 @@ TVM_REGISTER_GLOBAL("vm.builtin.kv_state_begin_forward")
}
kv_state->BeginForward(seq_ids, append_lengths, token_tree_parent_ptr);
});
-TVM_REGISTER_GLOBAL("vm.builtin.kv_state_end_forward")
- .set_body_method<KVState>(&KVStateObj::EndForward);
+TVM_REGISTER_GLOBAL("vm.builtin.kv_state_end_forward").set_body_method(&KVStateObj::EndForward);
// Attention KV Cache methods
TVM_REGISTER_GLOBAL("vm.builtin.kv_cache_disagg_prepare_recv")
- .set_body_method<AttentionKVCache>(&AttentionKVCacheObj::DisaggPrepareRecv);
+ .set_body_method(&AttentionKVCacheObj::DisaggPrepareRecv);
TVM_REGISTER_GLOBAL("vm.builtin.kv_cache_disagg_mark_send")
- .set_body_method<AttentionKVCache>(&AttentionKVCacheObj::DisaggMarkSend);
+ .set_body_method(&AttentionKVCacheObj::DisaggMarkSend);
TVM_REGISTER_GLOBAL("vm.builtin.attention_kv_cache_enable_sliding_window_for_seq")
- .set_body_method<AttentionKVCache>(&AttentionKVCacheObj::EnableSlidingWindowForSeq);
+ .set_body_method(&AttentionKVCacheObj::EnableSlidingWindowForSeq);
TVM_REGISTER_GLOBAL("vm.builtin.attention_kv_cache_commit_accepted_token_tree_nodes")
- .set_body_method<AttentionKVCache>(&AttentionKVCacheObj::CommitAcceptedTokenTreeNodes);
+ .set_body_method(&AttentionKVCacheObj::CommitAcceptedTokenTreeNodes);
TVM_REGISTER_GLOBAL("vm.builtin.attention_kv_cache_empty")
- .set_body_method<AttentionKVCache>(&AttentionKVCacheObj::Empty);
+ .set_body_method(&AttentionKVCacheObj::Empty);
TVM_REGISTER_GLOBAL("vm.builtin.attention_kv_cache_get_num_available_pages")
- .set_body_method<AttentionKVCache>(&AttentionKVCacheObj::GetNumAvailablePages);
+ .set_body_method(&AttentionKVCacheObj::GetNumAvailablePages);
TVM_REGISTER_GLOBAL("vm.builtin.attention_kv_cache_get_total_sequence_length")
- .set_body_method<AttentionKVCache>(&AttentionKVCacheObj::GetTotalSequenceLength);
+ .set_body_method(&AttentionKVCacheObj::GetTotalSequenceLength);
TVM_REGISTER_GLOBAL("vm.builtin.attention_kv_cache_get_query_positions")
- .set_body_method<AttentionKVCache>(&AttentionKVCacheObj::GetQueryPositions);
+ .set_body_method(&AttentionKVCacheObj::GetQueryPositions);
TVM_REGISTER_GLOBAL("vm.builtin.attention_kv_cache_debug_get_kv")
- .set_body_method<AttentionKVCache>(&AttentionKVCacheObj::DebugGetKV);
+ .set_body_method(&AttentionKVCacheObj::DebugGetKV);
TVM_REGISTER_GLOBAL("vm.builtin.attention_kv_cache_debug_get_kv_mla")
- .set_body_method<AttentionKVCache>(&AttentionKVCacheObj::DebugGetKVMLA);
+ .set_body_method(&AttentionKVCacheObj::DebugGetKVMLA);
TVM_REGISTER_GLOBAL("vm.builtin.attention_kv_cache_attention_with_fused_qkv")
.set_body_typed([](AttentionKVCache kv_cache, int64_t layer_id, double sm_scale,
NDArray qkv_data, NDArray o_data) {
@@ -107,14 +104,13 @@ TVM_REGISTER_GLOBAL("vm.builtin.attention_kv_cache_merge_attn_output_inplace")
});
// RNN State methods
-TVM_REGISTER_GLOBAL("vm.builtin.rnn_state_get").set_body_method<RNNState>(&RNNStateObj::Get);
+TVM_REGISTER_GLOBAL("vm.builtin.rnn_state_get").set_body_method(&RNNStateObj::Get);
TVM_REGISTER_GLOBAL("vm.builtin.rnn_state_set")
.set_body_typed([](RNNState state, int64_t layer_id, int64_t state_id, NDArray data) {
state->Set(layer_id, state_id, data);
return state;
});
-TVM_REGISTER_GLOBAL("vm.builtin.rnn_state_debug_get")
- .set_body_method<RNNState>(&RNNStateObj::DebugGet);
+TVM_REGISTER_GLOBAL("vm.builtin.rnn_state_debug_get").set_body_method(&RNNStateObj::DebugGet);
} // namespace relax_vm
} // namespace runtime
diff --git a/src/runtime/relax_vm/lm_support.cc b/src/runtime/relax_vm/lm_support.cc
index 675ffd472b..e3f5cd3c3d 100644
--- a/src/runtime/relax_vm/lm_support.cc
+++ b/src/runtime/relax_vm/lm_support.cc
@@ -301,7 +301,7 @@ NDArray AttentionKVCacheView(AttentionKVCacheLegacy cache, ShapeTuple shape) {
}
TVM_REGISTER_GLOBAL("vm.builtin.attention_kv_cache_view")
- .set_body([](TVMArgs args, TVMRetValue* rv) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* rv) {
CHECK(args.size() == 1 || args.size() == 2)
<< "ValueError: `vm.builtin.attention_kv_cache_view` expects 1 or 2 arguments, but got "
<< args.size() << ".";
diff --git a/src/runtime/relax_vm/paged_kv_cache.cc b/src/runtime/relax_vm/paged_kv_cache.cc
index d654a6defe..0d12def981 100644
--- a/src/runtime/relax_vm/paged_kv_cache.cc
+++ b/src/runtime/relax_vm/paged_kv_cache.cc
@@ -334,10 +334,12 @@ class PagedAttentionKVCacheObj : public AttentionKVCacheObj {
for (AttnKind attn_kind : attn_kinds_) {
CHECK(attn_kind == AttnKind::kMHA);
}
- CHECK(Registry::Get("runtime.disco.nvshmem.init_nvshmem") != nullptr)
+ const auto f_nvshmem_init =
+ tvm::ffi::Function::GetGlobal("runtime.disco.nvshmem.init_nvshmem");
+ CHECK(f_nvshmem_init.has_value())
<< "NVSHMEM is not enabled. Please make sure NVSHMEM is enabled when compiling TVM.";
- const PackedFunc* f_nvshmem_empty = runtime::Registry::Get("runtime.disco.nvshmem.empty");
- ICHECK_NOTNULL(f_nvshmem_empty);
+ const auto f_nvshmem_empty = tvm::ffi::Function::GetGlobal("runtime.disco.nvshmem.empty");
+ ICHECK(f_nvshmem_empty.has_value());
nvshmem_pages_ = (*f_nvshmem_empty)(
ShapeTuple({num_layers, num_total_pages, 2, num_kv_heads, page_size, qk_head_dim}), dtype,
device);
@@ -348,11 +350,11 @@ class PagedAttentionKVCacheObj : public AttentionKVCacheObj {
nvshmem_pages_.DataType().bytes()));
}
- const PackedFunc* f_transfer_kv_ptr = Registry::Get("nvshmem.KVTransfer");
- const PackedFunc* f_transfer_kv_page_to_page_ptr =
- Registry::Get("nvshmem.KVTransferPageToPage");
- ICHECK_NOTNULL(f_transfer_kv_ptr);
- ICHECK_NOTNULL(f_transfer_kv_page_to_page_ptr);
+ const auto f_transfer_kv_ptr = tvm::ffi::Function::GetGlobal("nvshmem.KVTransfer");
+ const auto f_transfer_kv_page_to_page_ptr =
+ tvm::ffi::Function::GetGlobal("nvshmem.KVTransferPageToPage");
+ ICHECK(f_transfer_kv_ptr.has_value());
+ ICHECK(f_transfer_kv_page_to_page_ptr.has_value());
f_transfer_kv_ = *f_transfer_kv_ptr;
f_transfer_kv_page_to_page_ = *f_transfer_kv_page_to_page_ptr;
} else {
diff --git a/src/runtime/relax_vm/vm.cc b/src/runtime/relax_vm/vm.cc
index 461b64f50e..e53c50aa1f 100644
--- a/src/runtime/relax_vm/vm.cc
+++ b/src/runtime/relax_vm/vm.cc
@@ -705,8 +705,8 @@ void VirtualMachineImpl::InitFuncPool() {
// only look through imports first
PackedFunc func = GetFuncFromImports(info.name);
if (!func.defined()) {
- const PackedFunc* p_func = Registry::Get(info.name);
- if (p_func != nullptr) func = *(p_func);
+ const auto p_func = tvm::ffi::Function::GetGlobal(info.name);
+ if (p_func.has_value()) func = *(p_func);
}
ICHECK(func.defined())
<< "Error: Cannot find PackedFunc " << info.name
@@ -900,8 +900,8 @@ void VirtualMachineImpl::_SetInstrument(TVMArgs args, TVMRetValue* rv) {
this->SetInstrument(args[0]);
} else {
String func_name = args[0];
- const PackedFunc* factory = Registry::Get(func_name);
- CHECK(factory) << "Cannot find factory " << func_name;
+ const auto factory = tvm::ffi::Function::GetGlobal(func_name);
+ CHECK(factory.has_value()) << "Cannot find factory " << func_name;
TVMRetValue rv;
factory->CallPacked(args.Slice(1), &rv);
this->SetInstrument(rv);
diff --git a/src/runtime/rocm/rocm_device_api.cc b/src/runtime/rocm/rocm_device_api.cc
index e7f103daad..afd3a55c22 100644
--- a/src/runtime/rocm/rocm_device_api.cc
+++ b/src/runtime/rocm/rocm_device_api.cc
@@ -251,12 +251,12 @@ ROCMThreadEntry::ROCMThreadEntry() : pool(kDLROCM, ROCMDeviceAPI::Global()) {}
ROCMThreadEntry* ROCMThreadEntry::ThreadLocal() { return ROCMThreadStore::Get(); }
-TVM_REGISTER_GLOBAL("device_api.rocm").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("device_api.rocm").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
DeviceAPI* ptr = ROCMDeviceAPI::Global();
*rv = static_cast<void*>(ptr);
});
-TVM_REGISTER_GLOBAL("device_api.rocm_host").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("device_api.rocm_host").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
DeviceAPI* ptr = ROCMDeviceAPI::Global();
*rv = static_cast<void*>(ptr);
});
diff --git a/src/runtime/rpc/rpc_device_api.cc b/src/runtime/rpc/rpc_device_api.cc
index a5c8541dc0..729cc10f0b 100644
--- a/src/runtime/rpc/rpc_device_api.cc
+++ b/src/runtime/rpc/rpc_device_api.cc
@@ -150,7 +150,7 @@ class RPCDeviceAPI final : public DeviceAPI {
}
};
-TVM_REGISTER_GLOBAL("device_api.rpc").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("device_api.rpc").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
static RPCDeviceAPI inst;
DeviceAPI* ptr = &inst;
*rv = static_cast<void*>(ptr);
diff --git a/src/runtime/rpc/rpc_endpoint.cc b/src/runtime/rpc/rpc_endpoint.cc
index cae9072697..9a7657fce5 100644
--- a/src/runtime/rpc/rpc_endpoint.cc
+++ b/src/runtime/rpc/rpc_endpoint.cc
@@ -592,8 +592,8 @@ class RPCEndpoint::EventHandler : public dmlc::Stream {
constructor_args = args.Slice(1);
}
- auto* fconstructor = Registry::Get(constructor_name);
- ICHECK(fconstructor != nullptr) << " Cannot find session constructor " << constructor_name;
+ auto fconstructor = tvm::ffi::Function::GetGlobal(constructor_name);
+ ICHECK(fconstructor.has_value()) << " Cannot find session constructor " << constructor_name;
TVMRetValue con_ret;
try {
@@ -812,12 +812,12 @@ void RPCEndpoint::Shutdown() {
}
void RPCEndpoint::ServerLoop() {
- if (const auto* f = Registry::Get("tvm.rpc.server.start")) {
+ if (const auto f = tvm::ffi::Function::GetGlobal("tvm.rpc.server.start")) {
(*f)();
}
TVMRetValue rv;
ICHECK(HandleUntilReturnEvent(false, [](TVMArgs) {}) == RPCCode::kShutdown);
- if (const auto* f = Registry::Get("tvm.rpc.server.shutdown")) {
+ if (const auto f = tvm::ffi::Function::GetGlobal("tvm.rpc.server.shutdown")) {
(*f)();
}
channel_.reset(nullptr);
diff --git a/src/runtime/rpc/rpc_local_session.cc b/src/runtime/rpc/rpc_local_session.cc
index 771b8c60bc..45b1ad2a7b 100644
--- a/src/runtime/rpc/rpc_local_session.cc
+++ b/src/runtime/rpc/rpc_local_session.cc
@@ -33,7 +33,7 @@ namespace tvm {
namespace runtime {
RPCSession::PackedFuncHandle LocalSession::GetFunction(const std::string& name) {
- if (auto* fp = tvm::runtime::Registry::Get(name)) {
+ if (auto fp = tvm::ffi::Function::GetGlobal(name)) {
// return raw handle because the remote need to explicitly manage it.
Any ret = *fp;
TVMFFIAny ret_any = ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(std::move(ret));
diff --git a/src/runtime/rpc/rpc_module.cc b/src/runtime/rpc/rpc_module.cc
index 67cc2b23bc..9e60e18172 100644
--- a/src/runtime/rpc/rpc_module.cc
+++ b/src/runtime/rpc/rpc_module.cc
@@ -404,8 +404,8 @@ TVM_REGISTER_GLOBAL("runtime.RPCTimeEvaluator")
} else {
PackedFunc f_preproc;
if (!f_preproc_name.empty()) {
- auto* pf_preproc = runtime::Registry::Get(f_preproc_name);
- ICHECK(pf_preproc != nullptr)
+ auto pf_preproc = tvm::ffi::Function::GetGlobal(f_preproc_name);
+ ICHECK(pf_preproc.has_value())
<< "Cannot find " << f_preproc_name << " in the global function";
f_preproc = *pf_preproc;
}
@@ -416,12 +416,12 @@ TVM_REGISTER_GLOBAL("runtime.RPCTimeEvaluator")
repeats_to_cooldown, cache_flush_bytes, f_preproc);
}
} else {
- auto* pf = runtime::Registry::Get(name);
- ICHECK(pf != nullptr) << "Cannot find " << name << " in the global function";
+ auto pf = tvm::ffi::Function::GetGlobal(name);
+ ICHECK(pf.has_value()) << "Cannot find " << name << " in the global function";
PackedFunc f_preproc;
if (!f_preproc_name.empty()) {
- auto* pf_preproc = runtime::Registry::Get(f_preproc_name);
- ICHECK(pf_preproc != nullptr)
+ auto pf_preproc = tvm::ffi::Function::GetGlobal(f_preproc_name);
+ ICHECK(pf_preproc.has_value())
<< "Cannot find " << f_preproc_name << " in the global function";
f_preproc = *pf_preproc;
}
@@ -431,9 +431,8 @@ TVM_REGISTER_GLOBAL("runtime.RPCTimeEvaluator")
}
});
-TVM_REGISTER_GLOBAL("cache_flush_cpu_non_first_arg").set_body([](TVMArgs args, TVMRetValue* rv) {
- CPUCacheFlush(1, args);
-});
+TVM_REGISTER_GLOBAL("cache_flush_cpu_non_first_arg")
+ .set_body_packed([](TVMArgs args, TVMRetValue* rv) { CPUCacheFlush(1, args); });
// server function registration.
TVM_REGISTER_GLOBAL("tvm.rpc.server.ImportModule").set_body_typed([](Module parent, Module child) {
@@ -458,7 +457,7 @@ TVM_REGISTER_GLOBAL("rpc.ImportRemoteModule").set_body_typed([](Module parent, M
static_cast<RPCModuleNode*>(parent.operator->())->ImportModule(child);
});
-TVM_REGISTER_GLOBAL("rpc.SessTableIndex").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("rpc.SessTableIndex").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
Module m = args[0];
std::string tkey = m->type_key();
ICHECK_EQ(tkey, "rpc");
diff --git a/src/runtime/rpc/rpc_pipe_impl.cc b/src/runtime/rpc/rpc_pipe_impl.cc
index b84dfaaa1c..bfcea856ca 100644
--- a/src/runtime/rpc/rpc_pipe_impl.cc
+++ b/src/runtime/rpc/rpc_pipe_impl.cc
@@ -112,7 +112,7 @@ Module CreatePipeClient(std::vector<std::string> cmd) {
return CreateRPCSessionModule(CreateClientSession(endpt));
}
-TVM_REGISTER_GLOBAL("rpc.CreatePipeClient").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("rpc.CreatePipeClient").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
std::vector<std::string> cmd;
for (int i = 0; i < args.size(); ++i) {
cmd.push_back(args[i].operator std::string());
diff --git a/src/runtime/rpc/rpc_server_env.cc b/src/runtime/rpc/rpc_server_env.cc
index 9908c89167..605b2bc0c6 100644
--- a/src/runtime/rpc/rpc_server_env.cc
+++ b/src/runtime/rpc/rpc_server_env.cc
@@ -30,18 +30,18 @@ namespace runtime {
std::string RPCGetPath(const std::string& name) {
// do live lookup everytime as workpath can change.
- const PackedFunc* f = runtime::Registry::Get("tvm.rpc.server.workpath");
- ICHECK(f != nullptr) << "require tvm.rpc.server.workpath";
+ const auto f = tvm::ffi::Function::GetGlobal("tvm.rpc.server.workpath");
+ ICHECK(f.has_value()) << "require tvm.rpc.server.workpath";
return (*f)(name);
}
-TVM_REGISTER_GLOBAL("tvm.rpc.server.upload").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("tvm.rpc.server.upload").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
std::string file_name = RPCGetPath(args[0]);
std::string data = args[1];
SaveBinaryToFile(file_name, data);
});
-TVM_REGISTER_GLOBAL("tvm.rpc.server.download").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("tvm.rpc.server.download").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
std::string file_name = RPCGetPath(args[0]);
std::string data;
LoadBinaryFromFile(file_name, &data);
@@ -49,7 +49,7 @@ TVM_REGISTER_GLOBAL("tvm.rpc.server.download").set_body([](TVMArgs args, TVMRetV
*rv = ffi::Bytes(data);
});
-TVM_REGISTER_GLOBAL("tvm.rpc.server.remove").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("tvm.rpc.server.remove").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
std::string file_name = RPCGetPath(args[0]);
RemoveFile(file_name);
});
diff --git a/src/runtime/rpc/rpc_socket_impl.cc b/src/runtime/rpc/rpc_socket_impl.cc
index fb356f9100..11d198b227 100644
--- a/src/runtime/rpc/rpc_socket_impl.cc
+++ b/src/runtime/rpc/rpc_socket_impl.cc
@@ -124,7 +124,7 @@ void RPCServerLoop(PackedFunc fsend, PackedFunc frecv) {
->ServerLoop();
}
-TVM_REGISTER_GLOBAL("rpc.Connect").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("rpc.Connect").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
std::string url = args[0];
int port = args[1];
std::string key = args[2];
@@ -132,7 +132,7 @@ TVM_REGISTER_GLOBAL("rpc.Connect").set_body([](TVMArgs args, TVMRetValue* rv) {
*rv = RPCClientConnect(url, port, key, enable_logging, args.Slice(4));
});
-TVM_REGISTER_GLOBAL("rpc.ServerLoop").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("rpc.ServerLoop").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
if (auto opt_int = args[0].as<int64_t>()) {
RPCServerLoop(opt_int.value());
} else {
diff --git a/src/runtime/system_library.cc b/src/runtime/system_library.cc
index 55335649a7..3d8fa86abe 100644
--- a/src/runtime/system_library.cc
+++ b/src/runtime/system_library.cc
@@ -112,7 +112,7 @@ class SystemLibModuleRegistry {
std::unordered_map<std::string, runtime::Module> lib_map_;
};
-TVM_REGISTER_GLOBAL("runtime.SystemLib").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("runtime.SystemLib").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
std::string symbol_prefix = "";
if (args.size() != 0) {
symbol_prefix = args[0].operator std::string();
diff --git a/src/runtime/thread_pool.cc b/src/runtime/thread_pool.cc
index 39b668c7fc..56fa1e7e0b 100644
--- a/src/runtime/thread_pool.cc
+++ b/src/runtime/thread_pool.cc
@@ -379,7 +379,7 @@ class ThreadPool {
* \brief args[0] is the AffinityMode, args[1] is the number of threads.
* args2 is a list of CPUs which is used to set the CPU affinity.
*/
-TVM_REGISTER_GLOBAL("runtime.config_threadpool").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("runtime.config_threadpool").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
threading::ThreadGroup::AffinityMode mode =
static_cast<threading::ThreadGroup::AffinityMode>(static_cast<int>(args[0]));
int nthreads = args[1];
diff --git a/src/runtime/vulkan/vulkan_device_api.cc b/src/runtime/vulkan/vulkan_device_api.cc
index af7b35e85e..8ec8101d83 100644
--- a/src/runtime/vulkan/vulkan_device_api.cc
+++ b/src/runtime/vulkan/vulkan_device_api.cc
@@ -455,7 +455,7 @@ VulkanDevice& VulkanDeviceAPI::device(size_t device_id) {
return const_cast<VulkanDevice&>(const_cast<const VulkanDeviceAPI*>(this)->device(device_id));
}
-TVM_REGISTER_GLOBAL("device_api.vulkan").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("device_api.vulkan").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
DeviceAPI* ptr = VulkanDeviceAPI::Global();
*rv = static_cast<void*>(ptr);
});
diff --git a/src/script/ir_builder/base.cc b/src/script/ir_builder/base.cc
index 879db4f3d7..f295b48f51 100644
--- a/src/script/ir_builder/base.cc
+++ b/src/script/ir_builder/base.cc
@@ -102,18 +102,18 @@ void Namer::Name(ObjectRef node, String name) {
TVM_REGISTER_NODE_TYPE(IRBuilderFrameNode);
TVM_REGISTER_NODE_TYPE(IRBuilderNode);
TVM_REGISTER_GLOBAL("script.ir_builder.IRBuilderFrameEnter")
- .set_body_method<IRBuilderFrame>(&IRBuilderFrameNode::EnterWithScope);
+ .set_body_method(&IRBuilderFrameNode::EnterWithScope);
TVM_REGISTER_GLOBAL("script.ir_builder.IRBuilderFrameExit")
- .set_body_method<IRBuilderFrame>(&IRBuilderFrameNode::ExitWithScope);
+ .set_body_method(&IRBuilderFrameNode::ExitWithScope);
TVM_REGISTER_GLOBAL("script.ir_builder.IRBuilderFrameAddCallback")
- .set_body_method<IRBuilderFrame>(&IRBuilderFrameNode::AddCallback);
+ .set_body_method(&IRBuilderFrameNode::AddCallback);
TVM_REGISTER_GLOBAL("script.ir_builder.IRBuilder").set_body_typed([]() { return IRBuilder(); });
TVM_REGISTER_GLOBAL("script.ir_builder.IRBuilderEnter").set_body_method(&IRBuilder::EnterWithScope);
TVM_REGISTER_GLOBAL("script.ir_builder.IRBuilderExit").set_body_method(&IRBuilder::ExitWithScope);
TVM_REGISTER_GLOBAL("script.ir_builder.IRBuilderCurrent").set_body_typed(IRBuilder::Current);
TVM_REGISTER_GLOBAL("script.ir_builder.IRBuilderIsInScope").set_body_typed(IRBuilder::IsInScope);
TVM_REGISTER_GLOBAL("script.ir_builder.IRBuilderGet")
- .set_body_method<IRBuilder>(&IRBuilderNode::Get<ObjectRef>);
+ .set_body_method(&IRBuilderNode::Get<ObjectRef>);
TVM_REGISTER_GLOBAL("script.ir_builder.IRBuilderName").set_body_typed(IRBuilder::Name<ObjectRef>);
} // namespace ir_builder
diff --git a/src/script/printer/doc.cc b/src/script/printer/doc.cc
index 1db4e090dc..615032ea00 100644
--- a/src/script/printer/doc.cc
+++ b/src/script/printer/doc.cc
@@ -241,11 +241,10 @@ TVM_REGISTER_GLOBAL("script.printer.DocSetSourcePaths")
TVM_REGISTER_NODE_TYPE(ExprDocNode);
TVM_REGISTER_GLOBAL("script.printer.ExprDocAttr")
- .set_body_method<ExprDoc, ExprDocNode, ExprDoc, String>(&ExprDocNode::Attr);
-TVM_REGISTER_GLOBAL("script.printer.ExprDocIndex")
- .set_body_method<ExprDoc>(&ExprDocNode::operator[]);
+ .set_body_method<ExprDocNode, ExprDoc, String>(&ExprDocNode::Attr);
+TVM_REGISTER_GLOBAL("script.printer.ExprDocIndex").set_body_method(&ExprDocNode::operator[]);
TVM_REGISTER_GLOBAL("script.printer.ExprDocCall")
- .set_body_method<ExprDoc, ExprDocNode, ExprDoc, Array<ExprDoc>, Array<String>, Array<ExprDoc>>(
+ .set_body_method<ExprDocNode, ExprDoc, Array<ExprDoc>, Array<String>, Array<ExprDoc>>(
&ExprDocNode::Call);
TVM_REGISTER_NODE_TYPE(StmtDocNode);
diff --git a/src/script/printer/tir/expr.cc b/src/script/printer/tir/expr.cc
index 8ac0931496..63c79af55c 100644
--- a/src/script/printer/tir/expr.cc
+++ b/src/script/printer/tir/expr.cc
@@ -254,7 +254,7 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
int n_args = call->args.size();
int64_t id = call->args[0].as<IntImmNode>()->value;
auto f_llvm_lookup_intrinsic_name =
- tvm::runtime::Registry::Get("target.llvm_get_intrinsic_name");
+ tvm::ffi::Function::GetGlobal("target.llvm_get_intrinsic_name");
Array<ExprDoc> args;
args.reserve(n_args + 1);
diff --git a/src/support/errno_handling.h b/src/support/errno_handling.h
index 0bdfdfdf02..748e390e78 100644
--- a/src/support/errno_handling.h
+++ b/src/support/errno_handling.h
@@ -55,7 +55,9 @@ inline ssize_t RetryCallOnEINTR(FuncType func, GetErrorCodeFuncType fgeterrorcod
// environment specific(e.g. python) signal exceptions.
// This function will throw an exception if there is
// if the process received a signal that requires TVM to return immediately (e.g. SIGINT).
- runtime::EnvCheckSignals();
+ if (TVMFFIEnvCheckSignals() != 0) {
+ throw ffi::EnvErrorAlreadySet();
+ }
} else {
// other errors
return ret;
diff --git a/src/support/ffi_testing.cc b/src/support/ffi_testing.cc
index fbd9e3188f..fc7315eb83 100644
--- a/src/support/ffi_testing.cc
+++ b/src/support/ffi_testing.cc
@@ -53,19 +53,15 @@ struct TestAttrs : public AttrsNode<TestAttrs> {
TVM_REGISTER_NODE_TYPE(TestAttrs);
-TVM_REGISTER_GLOBAL("testing.nop").set_body([](TVMArgs args, TVMRetValue* ret) {});
-TVM_REGISTER_GLOBAL("testing.echo").set_body([](TVMArgs args, TVMRetValue* ret) {
- *ret = args[0];
-});
-
-TVM_REGISTER_GLOBAL("testing.test_wrap_callback").set_body([](TVMArgs args, TVMRetValue* ret) {
- PackedFunc pf = args[0];
- *ret = runtime::TypedPackedFunc<void()>([pf]() { pf(); });
-});
+TVM_REGISTER_GLOBAL("testing.test_wrap_callback")
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
+ PackedFunc pf = args[0];
+ *ret = runtime::TypedPackedFunc<void()>([pf]() { pf(); });
+ });
TVM_REGISTER_GLOBAL("testing.test_wrap_callback_suppress_err")
- .set_body([](TVMArgs args, TVMRetValue* ret) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
PackedFunc pf = args[0];
auto result = runtime::TypedPackedFunc<void()>([pf]() {
try {
@@ -77,18 +73,19 @@ TVM_REGISTER_GLOBAL("testing.test_wrap_callback_suppress_err")
});
TVM_REGISTER_GLOBAL("testing.test_raise_error_callback")
- .set_body([](TVMArgs args, TVMRetValue* ret) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
std::string msg = args[0];
*ret = runtime::TypedPackedFunc<void()>([msg]() { LOG(FATAL) << msg; });
});
-TVM_REGISTER_GLOBAL("testing.test_check_eq_callback").set_body([](TVMArgs args, TVMRetValue* ret) {
- std::string msg = args[0];
- *ret =
- runtime::TypedPackedFunc<void(int x, int y)>([msg](int x, int y) { CHECK_EQ(x, y) << msg; });
-});
+TVM_REGISTER_GLOBAL("testing.test_check_eq_callback")
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
+ std::string msg = args[0];
+ *ret = runtime::TypedPackedFunc<void(int x, int y)>(
+ [msg](int x, int y) { CHECK_EQ(x, y) << msg; });
+ });
-TVM_REGISTER_GLOBAL("testing.device_test").set_body([](TVMArgs args, TVMRetValue* ret) {
+TVM_REGISTER_GLOBAL("testing.device_test").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
Device dev = args[0];
int dtype = args[1];
int did = args[2];
@@ -97,17 +94,9 @@ TVM_REGISTER_GLOBAL("testing.device_test").set_body([](TVMArgs args, TVMRetValue
*ret = dev;
});
-TVM_REGISTER_GLOBAL("testing.run_check_signal").set_body_typed([](int nsec) {
- for (int i = 0; i < nsec; ++i) {
- tvm::runtime::EnvCheckSignals();
- std::this_thread::sleep_for(std::chrono::seconds(1));
- }
- LOG(INFO) << "Function finished without catching signal";
-});
-
-TVM_REGISTER_GLOBAL("testing.identity_cpp").set_body([](TVMArgs args, TVMRetValue* ret) {
- const auto* identity_func = tvm::runtime::Registry::Get("testing.identity_py");
- ICHECK(identity_func != nullptr)
+TVM_REGISTER_GLOBAL("testing.identity_cpp").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
+ const auto identity_func = tvm::ffi::Function::GetGlobal("testing.identity_py");
+ ICHECK(identity_func.has_value())
<< "AttributeError: \"testing.identity_py\" is not registered. Please check "
"if the python module is properly loaded";
*ret = (*identity_func)(args[0]);
@@ -126,7 +115,7 @@ void ErrorTest(int x, int y) {
TVM_REGISTER_GLOBAL("testing.ErrorTest").set_body_typed(ErrorTest);
// internal function used for debug and testing purposes
-TVM_REGISTER_GLOBAL("testing.object_use_count").set_body([](TVMArgs args, TVMRetValue* ret) {
+TVM_REGISTER_GLOBAL("testing.object_use_count").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
runtime::ObjectRef obj = args[0];
// substract the current one because we always copy
// and get another value.
@@ -178,13 +167,6 @@ TVM_REGISTER_GLOBAL("testing.sleep_in_ffi").set_body_typed([](double timeout) {
std::this_thread::sleep_for(duration);
});
-TVM_REGISTER_GLOBAL("testing.check_signals").set_body_typed([](double sleep_period) {
- while (true) {
- std::chrono::duration<int64_t, std::nano> duration(static_cast<int64_t>(sleep_period * 1e9));
- std::this_thread::sleep_for(duration);
- runtime::EnvCheckSignals();
- }
-});
TVM_REGISTER_GLOBAL("testing.ReturnsVariant").set_body_typed([](int x) -> Variant<String, IntImm> {
if (x % 2 == 0) {
@@ -288,7 +270,7 @@ class TestingEventLogger {
std::vector<Entry> entries_;
};
-TVM_REGISTER_GLOBAL("testing.record_event").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("testing.record_event").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
if (args.size() != 0 && args[0].as<String>()) {
TestingEventLogger::ThreadLocal()->Record(args[0]);
} else {
@@ -296,7 +278,7 @@ TVM_REGISTER_GLOBAL("testing.record_event").set_body([](TVMArgs args, TVMRetValu
}
});
-TVM_REGISTER_GLOBAL("testing.reset_events").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("testing.reset_events").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
TestingEventLogger::ThreadLocal()->Reset();
});
diff --git a/src/target/codegen.cc b/src/target/codegen.cc
index 78a89872ce..87c294050f 100644
--- a/src/target/codegen.cc
+++ b/src/target/codegen.cc
@@ -51,8 +51,8 @@ runtime::Module Build(IRModule mod, Target target) {
// the build function.
std::string build_f_name = "target.build." + target->kind->name;
- const PackedFunc* bf = runtime::Registry::Get(build_f_name);
- ICHECK(bf != nullptr) << build_f_name << " is not enabled";
+ const auto bf = tvm::ffi::Function::GetGlobal(build_f_name);
+ ICHECK(bf.has_value()) << build_f_name << " is not enabled";
return (*bf)(mod, target);
}
@@ -355,8 +355,8 @@ runtime::Module PackImportsToLLVM(const runtime::Module& mod, bool system_lib,
// Call codegen_blob to generate LLVM module
std::string codegen_f_name = "codegen.codegen_blob";
// the codegen function.
- const PackedFunc* codegen_f = runtime::Registry::Get(codegen_f_name);
- ICHECK(codegen_f != nullptr) << "codegen.codegen_blob is not presented.";
+ const auto codegen_f = tvm::ffi::Function::GetGlobal(codegen_f_name);
+ ICHECK(codegen_f.has_value()) << "codegen.codegen_blob is not presented.";
return (*codegen_f)(ffi::Bytes(blob), system_lib, llvm_target_string, c_symbol_prefix);
}
diff --git a/src/target/datatype/registry.cc b/src/target/datatype/registry.cc
index e12927e8c7..e308b3cf35 100644
--- a/src/target/datatype/registry.cc
+++ b/src/target/datatype/registry.cc
@@ -26,20 +26,23 @@ namespace datatype {
using runtime::TVMArgs;
using runtime::TVMRetValue;
-TVM_REGISTER_GLOBAL("dtype.register_custom_type").set_body([](TVMArgs args, TVMRetValue* ret) {
- datatype::Registry::Global()->Register(args[0], static_cast<uint8_t>(args[1].operator int()));
-});
+TVM_REGISTER_GLOBAL("dtype.register_custom_type")
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
+ datatype::Registry::Global()->Register(args[0], static_cast<uint8_t>(args[1].operator int()));
+ });
-TVM_REGISTER_GLOBAL("dtype.get_custom_type_code").set_body([](TVMArgs args, TVMRetValue* ret) {
- *ret = datatype::Registry::Global()->GetTypeCode(args[0]);
-});
+TVM_REGISTER_GLOBAL("dtype.get_custom_type_code")
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
+ *ret = datatype::Registry::Global()->GetTypeCode(args[0]);
+ });
-TVM_REGISTER_GLOBAL("dtype.get_custom_type_name").set_body([](TVMArgs args, TVMRetValue* ret) {
- *ret = Registry::Global()->GetTypeName(args[0].operator int());
-});
+TVM_REGISTER_GLOBAL("dtype.get_custom_type_name")
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
+ *ret = Registry::Global()->GetTypeName(args[0].operator int());
+ });
TVM_REGISTER_GLOBAL("runtime._datatype_get_type_registered")
- .set_body([](TVMArgs args, TVMRetValue* ret) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
*ret = Registry::Global()->GetTypeRegistered(args[0].operator int());
});
@@ -67,8 +70,8 @@ std::string Registry::GetTypeName(uint8_t type_code) {
return code_to_name_[type_code];
}
-const runtime::PackedFunc* GetCastLowerFunc(const std::string& target, uint8_t type_code,
- uint8_t src_type_code) {
+std::optional<tvm::ffi::Function> GetCastLowerFunc(const std::string& target, uint8_t type_code,
+ uint8_t src_type_code) {
std::ostringstream ss;
ss << "tvm.datatype.lower.";
ss << target << ".";
@@ -88,27 +91,28 @@ const runtime::PackedFunc* GetCastLowerFunc(const std::string& target, uint8_t t
} else {
ss << ffi::details::DLDataTypeCodeAsCStr(static_cast<DLDataTypeCode>(src_type_code));
}
- return runtime::Registry::Get(ss.str());
+ return tvm::ffi::Function::GetGlobal(ss.str());
}
-const runtime::PackedFunc* GetMinFunc(uint8_t type_code) {
+std::optional<tvm::ffi::Function> GetMinFunc(uint8_t type_code) {
std::ostringstream ss;
ss << "tvm.datatype.min.";
ss << datatype::Registry::Global()->GetTypeName(type_code);
- return runtime::Registry::Get(ss.str());
+ return tvm::ffi::Function::GetGlobal(ss.str());
}
-const runtime::PackedFunc* GetFloatImmLowerFunc(const std::string& target, uint8_t type_code) {
+std::optional<tvm::ffi::Function> GetFloatImmLowerFunc(const std::string& target,
+ uint8_t type_code) {
std::ostringstream ss;
ss << "tvm.datatype.lower.";
ss << target;
ss << ".FloatImm.";
ss << datatype::Registry::Global()->GetTypeName(type_code);
- return runtime::Registry::Get(ss.str());
+ return tvm::ffi::Function::GetGlobal(ss.str());
}
-const runtime::PackedFunc* GetIntrinLowerFunc(const std::string& target, const std::string& name,
- uint8_t type_code) {
+std::optional<tvm::ffi::Function> GetIntrinLowerFunc(const std::string& target,
+ const std::string& name, uint8_t type_code) {
std::ostringstream ss;
ss << "tvm.datatype.lower.";
ss << target;
@@ -116,14 +120,14 @@ const runtime::PackedFunc* GetIntrinLowerFunc(const std::string& target, const s
ss << name;
ss << ".";
ss << datatype::Registry::Global()->GetTypeName(type_code);
- return runtime::Registry::Get(ss.str());
+ return tvm::ffi::Function::GetGlobal(ss.str());
}
uint64_t ConvertConstScalar(uint8_t type_code, double value) {
std::ostringstream ss;
ss << "tvm.datatype.convertconstscalar.float.";
ss << datatype::Registry::Global()->GetTypeName(type_code);
- auto make_const_scalar_func = runtime::Registry::Get(ss.str());
+ auto make_const_scalar_func = tvm::ffi::Function::GetGlobal(ss.str());
return (*make_const_scalar_func)(value).operator uint64_t();
}
diff --git a/src/target/datatype/registry.h b/src/target/datatype/registry.h
index 01e7b82a45..46b189880f 100644
--- a/src/target/datatype/registry.h
+++ b/src/target/datatype/registry.h
@@ -118,7 +118,7 @@ uint64_t ConvertConstScalar(uint8_t type_code, double value);
* \param type_code The datatype
* \return Function which takes the width of the datatype and returns the min value
*/
-const runtime::PackedFunc* GetMinFunc(uint8_t type_code);
+std::optional<tvm::ffi::Function> GetMinFunc(uint8_t type_code);
/*!
* \brief Get lowering function for Cast ops
@@ -127,8 +127,8 @@ const runtime::PackedFunc* GetMinFunc(uint8_t type_code);
* \param src_type_code The datatype being cast from
* \return Lowering function for Cast ops for the provided target, type, and source type
*/
-const runtime::PackedFunc* GetCastLowerFunc(const std::string& target, uint8_t type_code,
- uint8_t src_type_code);
+std::optional<tvm::ffi::Function> GetCastLowerFunc(const std::string& target, uint8_t type_code,
+ uint8_t src_type_code);
/*!
* \brief Get lowering function for FloatImms
@@ -136,7 +136,8 @@ const runtime::PackedFunc* GetCastLowerFunc(const std::string& target, uint8_t t
* \param type_code The datatype of the FloatImm
* \return Lowering function for FloatImms for the provided target and type
*/
-const runtime::PackedFunc* GetFloatImmLowerFunc(const std::string& target, uint8_t type_code);
+std::optional<tvm::ffi::Function> GetFloatImmLowerFunc(const std::string& target,
+ uint8_t type_code);
/*!
* \brief Get lowering function for intrinsic Calls/pure intrinsic Calls
@@ -145,8 +146,8 @@ const runtime::PackedFunc* GetFloatImmLowerFunc(const std::string& target, uint8
* \param name The intrinsic name
* \return Lowering function for intrinsic Calls for the provided target and type
*/
-const runtime::PackedFunc* GetIntrinLowerFunc(const std::string& target, const std::string& name,
- uint8_t type_code);
+std::optional<tvm::ffi::Function> GetIntrinLowerFunc(const std::string& target,
+ const std::string& name, uint8_t type_code);
/*!
* \brief Get lowering function for other ops
@@ -154,11 +155,11 @@ const runtime::PackedFunc* GetIntrinLowerFunc(const std::string& target, const s
* \param type_code The datatype of the op
* \return Lowering function for other ops for the provided target and type
*/
-#define DEFINE_GET_LOWER_FUNC_(OP) \
- inline const runtime::PackedFunc* Get##OP##LowerFunc(const std::string& target, \
- uint8_t type_code) { \
- return runtime::Registry::Get("tvm.datatype.lower." + target + "." #OP "." + \
- datatype::Registry::Global()->GetTypeName(type_code)); \
+#define DEFINE_GET_LOWER_FUNC_(OP) \
+ inline std::optional<tvm::ffi::Function> Get##OP##LowerFunc(const std::string& target, \
+ uint8_t type_code) { \
+ return tvm::ffi::Function::GetGlobal("tvm.datatype.lower." + target + "." #OP "." + \
+ datatype::Registry::Global()->GetTypeName(type_code)); \
}
DEFINE_GET_LOWER_FUNC_(Add)
diff --git a/src/target/llvm/codegen_aarch64.cc b/src/target/llvm/codegen_aarch64.cc
index 2510c8bd77..37a11d8576 100644
--- a/src/target/llvm/codegen_aarch64.cc
+++ b/src/target/llvm/codegen_aarch64.cc
@@ -107,7 +107,7 @@ void CodeGenAArch64::VisitStmt_(const AttrStmtNode* op) {
}
TVM_REGISTER_GLOBAL("tvm.codegen.llvm.target_aarch64")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
+ .set_body_packed([](const TVMArgs& targs, TVMRetValue* rv) {
*rv = static_cast<void*>(new CodeGenAArch64());
});
diff --git a/src/target/llvm/codegen_amdgpu.cc b/src/target/llvm/codegen_amdgpu.cc
index 8c3da713a9..70ecdb5494 100644
--- a/src/target/llvm/codegen_amdgpu.cc
+++ b/src/target/llvm/codegen_amdgpu.cc
@@ -278,8 +278,8 @@ runtime::Module BuildAMDGPU(IRModule mod, Target target) {
cg->AddFunctionsOrdered(mod->functions.begin(), mod->functions.end());
llvm::TargetMachine* tm = llvm_target->GetOrCreateTargetMachine();
- const auto* find_rocm_bitcodes = tvm::runtime::Registry::Get("tvm_callback_rocm_bitcode_path");
- Array<runtime::String> bitcode_files = (*find_rocm_bitcodes)();
+ auto fbitcode = tvm::ffi::Function::GetGlobalRequired("tvm_callback_rocm_bitcode_path");
+ Array<runtime::String> bitcode_files = fbitcode();
for (auto& bitcode_path : bitcode_files) {
std::unique_ptr<llvm::Module> mlib = llvm_instance.LoadIR(bitcode_path);
@@ -343,14 +343,15 @@ runtime::Module BuildAMDGPU(IRModule mod, Target target) {
passAsm.run(*mAsm);
std::string assembly(dataAsm.begin(), dataAsm.end());
- const auto* f = tvm::runtime::Registry::Get("tvm_callback_rocm_link");
- ICHECK(f != nullptr) << "Require tvm_callback_rocm_link to exist, do import tvm.contrib.rocm";
+ auto flink = tvm::ffi::Function::GetGlobal("tvm_callback_rocm_link");
+ ICHECK(flink.has_value())
+ << "Require tvm_callback_rocm_link to exist, do import tvm.contrib.rocm";
TVMFFIByteArray arr;
arr.data = &obj[0];
arr.size = obj.length();
- std::string hsaco = (*f)(&arr);
+ std::string hsaco = (*flink)(&arr);
std::string ll(data_ll.begin(), data_ll.end());
return ROCMModuleCreate(hsaco, "hsaco", ExtractFuncInfo(mod), ll, assembly);
}
@@ -358,7 +359,7 @@ runtime::Module BuildAMDGPU(IRModule mod, Target target) {
TVM_REGISTER_GLOBAL("target.build.rocm").set_body_typed(BuildAMDGPU);
TVM_REGISTER_GLOBAL("tvm.codegen.llvm.target_rocm")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
+ .set_body_packed([](const TVMArgs& targs, TVMRetValue* rv) {
*rv = static_cast<void*>(new CodeGenAMDGPU());
});
diff --git a/src/target/llvm/codegen_arm.cc b/src/target/llvm/codegen_arm.cc
index ba9641bb06..0fd6673f49 100644
--- a/src/target/llvm/codegen_arm.cc
+++ b/src/target/llvm/codegen_arm.cc
@@ -133,7 +133,7 @@ PrimExpr CodeGenARM::ARMPopcount(const CallNode* call) {
}
TVM_REGISTER_GLOBAL("tvm.codegen.llvm.target_arm")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
+ .set_body_packed([](const TVMArgs& targs, TVMRetValue* rv) {
*rv = static_cast<void*>(new CodeGenARM());
});
diff --git a/src/target/llvm/codegen_cpu.cc b/src/target/llvm/codegen_cpu.cc
index 25485b70d2..d16d139f7e 100644
--- a/src/target/llvm/codegen_cpu.cc
+++ b/src/target/llvm/codegen_cpu.cc
@@ -1158,7 +1158,7 @@ void CodeGenCPU::VisitStmt_(const ForNode* op) {
}
TVM_REGISTER_GLOBAL("tvm.codegen.llvm.target_cpu")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
+ .set_body_packed([](const TVMArgs& targs, TVMRetValue* rv) {
*rv = static_cast<void*>(new CodeGenCPU());
});
diff --git a/src/target/llvm/codegen_hexagon.cc b/src/target/llvm/codegen_hexagon.cc
index 7c35b2ead6..d188f7e5a8 100644
--- a/src/target/llvm/codegen_hexagon.cc
+++ b/src/target/llvm/codegen_hexagon.cc
@@ -567,9 +567,9 @@ runtime::Module BuildHexagon(IRModule mod, Target target) {
std::string so_name(o_name, 0, o_name.size() - 1);
so_name += "so";
- const auto* f = tvm::runtime::Registry::Get("tvm.contrib.hexagon.link_shared");
- ICHECK(f != nullptr) << "tvm.contrib.hexagon.link_shared does not to exist, "
- "do import tvm.contrib.hexagon";
+ const auto f = tvm::ffi::Function::GetGlobal("tvm.contrib.hexagon.link_shared");
+ ICHECK(f.has_value()) << "tvm.contrib.hexagon.link_shared does not to exist, "
+ "do import tvm.contrib.hexagon";
Array<PrimExpr> o_names = {StringImm(o_name)};
Map<String, String> extra_args;
@@ -592,7 +592,7 @@ runtime::Module BuildHexagon(IRModule mod, Target target) {
TVM_REGISTER_GLOBAL("target.build.hexagon").set_body_typed(BuildHexagon);
TVM_REGISTER_GLOBAL("tvm.codegen.llvm.target_hexagon")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
+ .set_body_packed([](const TVMArgs& targs, TVMRetValue* rv) {
*rv = static_cast<void*>(new CodeGenHexagon());
});
diff --git a/src/target/llvm/codegen_llvm.cc b/src/target/llvm/codegen_llvm.cc
index 4ab6576ff3..68a1a4ee68 100644
--- a/src/target/llvm/codegen_llvm.cc
+++ b/src/target/llvm/codegen_llvm.cc
@@ -122,9 +122,9 @@ std::unique_ptr<CodeGenLLVM> CodeGenLLVM::Create(LLVMTarget* llvm_target) {
std::string target = llvm_target->GetOrCreateTargetMachine()->getTarget().getName();
std::string factory_template = "tvm.codegen.llvm.target_";
void* handle = nullptr;
- if (const PackedFunc* f = runtime::Registry::Get(factory_template + target)) {
+ if (auto f = tvm::ffi::Function::GetGlobal(factory_template + target)) {
handle = (*f)();
- } else if (const PackedFunc* f = runtime::Registry::Get(factory_template + "cpu")) {
+ } else if (auto f = tvm::ffi::Function::GetGlobal(factory_template + "cpu")) {
handle = (*f)();
} else {
LOG(FATAL) << "no factory function for codegen for target " << target;
diff --git a/src/target/llvm/codegen_nvptx.cc b/src/target/llvm/codegen_nvptx.cc
index 492248ea18..bb903522e1 100644
--- a/src/target/llvm/codegen_nvptx.cc
+++ b/src/target/llvm/codegen_nvptx.cc
@@ -329,8 +329,8 @@ runtime::Module BuildNVPTX(IRModule mod, Target target) {
cg->AddFunctionsOrdered(mod->functions.begin(), mod->functions.end());
llvm::TargetMachine* tm = llvm_target->GetOrCreateTargetMachine();
- const auto* flibdevice_path = tvm::runtime::Registry::Get("tvm_callback_libdevice_path");
- if (flibdevice_path != nullptr) {
+ const auto flibdevice_path = tvm::ffi::Function::GetGlobal("tvm_callback_libdevice_path");
+ if (flibdevice_path.has_value()) {
std::string path = (*flibdevice_path)(compute_ver);
if (path.length() != 0) {
std::unique_ptr<llvm::Module> mlib = llvm_instance.LoadIR(path);
@@ -371,7 +371,7 @@ runtime::Module BuildNVPTX(IRModule mod, Target target) {
TVM_REGISTER_GLOBAL("target.build.nvptx").set_body_typed(BuildNVPTX);
TVM_REGISTER_GLOBAL("tvm.codegen.llvm.target_nvptx")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
+ .set_body_packed([](const TVMArgs& targs, TVMRetValue* rv) {
*rv = static_cast<void*>(new CodeGenNVPTX());
});
diff --git a/src/target/llvm/codegen_x86_64.cc b/src/target/llvm/codegen_x86_64.cc
index 6329a77d7a..563b7be2ca 100644
--- a/src/target/llvm/codegen_x86_64.cc
+++ b/src/target/llvm/codegen_x86_64.cc
@@ -133,7 +133,7 @@ llvm::Value* CodeGenX86_64::CallVectorIntrin(llvm::Intrinsic::ID id, size_t intr
}
TVM_REGISTER_GLOBAL("tvm.codegen.llvm.target_x86-64")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
+ .set_body_packed([](const TVMArgs& targs, TVMRetValue* rv) {
*rv = static_cast<void*>(new CodeGenX86_64());
});
diff --git a/src/target/llvm/intrin_rule_hexagon.cc b/src/target/llvm/intrin_rule_hexagon.cc
index 83c6871d5c..b38ff06749 100644
--- a/src/target/llvm/intrin_rule_hexagon.cc
+++ b/src/target/llvm/intrin_rule_hexagon.cc
@@ -54,8 +54,8 @@ inline PrimExpr DispatchTVMQHLWrapperFp16(const PrimExpr& e) {
Array<PrimExpr> new_args;
#if ENABLE_QHL
// Check target for qfloat enablement
- const auto* f = tvm::runtime::Registry::Get("target.TargetCurrent");
- ICHECK(f != nullptr);
+ const auto f = tvm::ffi::Function::GetGlobal("target.TargetCurrent");
+ ICHECK(f.has_value()) << "target.TargetCurrent is not registered";
const auto ret = (*f)(true);
bool useqhl = true;
if (auto opt_target = ret.as<Target>()) {
@@ -104,8 +104,8 @@ TVM_REGISTER_OP("tir.tanh")
#if ENABLE_QHL
// Check target for qfloat enablement
- const auto* f = tvm::runtime::Registry::Get("target.TargetCurrent");
- ICHECK(f != nullptr);
+ const auto f = tvm::ffi::Function::GetGlobal("target.TargetCurrent");
+ ICHECK(f.has_value()) << "target.TargetCurrent is not registered";
const auto ret = (*f)(true);
bool useqhl = true;
if (auto opt_target = ret.as<Target>()) {
@@ -139,8 +139,8 @@ TVM_REGISTER_OP("tir.tan").set_attr<FLowerIntrinsic>(
const PrimExpr& x = call->args[0];
#if ENABLE_QHL
// Check target for qfloat enablement
- const auto* f = tvm::runtime::Registry::Get("target.TargetCurrent");
- ICHECK(f != nullptr);
+ const auto f = tvm::ffi::Function::GetGlobal("target.TargetCurrent");
+ ICHECK(f.has_value()) << "target.TargetCurrent is not registered";
const auto ret = (*f)(true);
bool useqhl = true;
if (auto opt_target = ret.as<Target>()) {
@@ -169,8 +169,8 @@ TVM_REGISTER_OP("tir.sigmoid")
const PrimExpr& x = call->args[0];
#if ENABLE_QHL
// Check target for qfloat enablement
- const auto* f = tvm::runtime::Registry::Get("target.TargetCurrent");
- ICHECK(f != nullptr);
+ const auto f = tvm::ffi::Function::GetGlobal("target.TargetCurrent");
+ ICHECK(f.has_value()) << "target.TargetCurrent is not registered";
const auto ret = (*f)(true);
bool useqhl = true;
if (auto opt_target = ret.as<Target>()) {
diff --git a/src/target/opt/build_cuda_on.cc b/src/target/opt/build_cuda_on.cc
index 87351051ad..449125ff15 100644
--- a/src/target/opt/build_cuda_on.cc
+++ b/src/target/opt/build_cuda_on.cc
@@ -126,7 +126,6 @@ std::string NVRTCCompile(const std::string& code, bool include_path = false) {
}
runtime::Module BuildCUDA(IRModule mod, Target target) {
- using tvm::runtime::Registry;
bool output_ssa = false;
CodeGenCUDA cg;
cg.Init(output_ssa);
@@ -150,14 +149,14 @@ runtime::Module BuildCUDA(IRModule mod, Target target) {
std::string code = cg.Finish();
- if (const auto* f = Registry::Get("tvm_callback_cuda_postproc")) {
+ if (auto f = ffi::Function::GetGlobal("tvm_callback_cuda_postproc")) {
code = (*f)(code, target).operator std::string();
}
std::string fmt = "ptx";
std::string ptx;
- const auto* f_enter = Registry::Get("target.TargetEnterScope");
+ auto f_enter = ffi::Function::GetGlobal("target.TargetEnterScope");
(*f_enter)(target);
- if (const auto* f = Registry::Get("tvm_callback_cuda_compile")) {
+ if (auto f = ffi::Function::GetGlobal("tvm_callback_cuda_compile")) {
ptx = (*f)(code, target).operator std::string();
// Dirty matching to check PTX vs cubin.
// TODO(tqchen) more reliable checks
@@ -165,7 +164,7 @@ runtime::Module BuildCUDA(IRModule mod, Target target) {
} else {
ptx = NVRTCCompile(code, cg.need_include_path());
}
- const auto* f_exit = Registry::Get("target.TargetExitScope");
+ auto f_exit = ffi::Function::GetGlobal("target.TargetExitScope");
(*f_exit)(target);
return CUDAModuleCreate(ptx, fmt, ExtractFuncInfo(mod), code);
}
diff --git a/src/target/parsers/cpu.cc b/src/target/parsers/cpu.cc
index c72401b5a0..cccc53fd05 100644
--- a/src/target/parsers/cpu.cc
+++ b/src/target/parsers/cpu.cc
@@ -30,9 +30,9 @@ namespace cpu {
Optional<String> DetectSystemTriple() {
#ifdef TVM_LLVM_VERSION
- auto pf = tvm::runtime::Registry::Get("target.llvm_get_system_triple");
- ICHECK(pf != nullptr) << "The target llvm_get_system_triple was not found, "
- "please compile with USE_LLVM = ON";
+ auto pf = tvm::ffi::Function::GetGlobal("target.llvm_get_system_triple");
+ ICHECK(pf.has_value()) << "The target llvm_get_system_triple was not found, "
+ "please compile with USE_LLVM = ON";
return (*pf)();
#endif
return {};
diff --git a/src/target/source/codegen_c_host.cc b/src/target/source/codegen_c_host.cc
index 684b32c642..c34d1206ab 100644
--- a/src/target/source/codegen_c_host.cc
+++ b/src/target/source/codegen_c_host.cc
@@ -353,7 +353,6 @@ inline void CodeGenCHost::PrintTernaryCondExpr(const T* op, const char* compare,
}
runtime::Module BuildCHost(IRModule mod, Target target) {
- using tvm::runtime::Registry;
bool output_ssa = false;
bool emit_asserts = false;
bool emit_fwd_func_decl = true;
diff --git a/src/target/source/codegen_cuda.cc b/src/target/source/codegen_cuda.cc
index a97e66d346..6bbb388a94 100644
--- a/src/target/source/codegen_cuda.cc
+++ b/src/target/source/codegen_cuda.cc
@@ -1025,9 +1025,9 @@ void CodeGenCUDA::VisitExpr_(const CallNode* op, std::ostream& os) {
// To store the 32x8 output back to a 16x16 tile in shared or global memory, we invert this map
// to determine the output location for each 8 element.
- const auto* index_map_func =
- runtime::Registry::Get("tir.index_map.shared_16x16_to_ldmatrix_32x8_layout");
- ICHECK(index_map_func);
+ const auto index_map_func =
+ tvm::ffi::Function::GetGlobal("tir.index_map.shared_16x16_to_ldmatrix_32x8_layout");
+ ICHECK(index_map_func.has_value());
arith::Analyzer analyzer;
auto inverse_index_map =
diff --git a/src/target/source/codegen_metal.cc b/src/target/source/codegen_metal.cc
index 2908514988..ecb5962d95 100644
--- a/src/target/source/codegen_metal.cc
+++ b/src/target/source/codegen_metal.cc
@@ -431,13 +431,12 @@ void CodeGenMetal::VisitExpr_(const FloatImmNode* op, std::ostream& os) { // NO
}
runtime::Module BuildMetal(IRModule mod, Target target) {
- using tvm::runtime::Registry;
bool output_ssa = false;
mod = tir::transform::PointerValueTypeRewrite()(std::move(mod));
std::ostringstream source_maker;
std::unordered_map<std::string, std::string> smap;
- const auto* fmetal_compile = Registry::Get("tvm_callback_metal_compile");
+ const auto fmetal_compile = tvm::ffi::Function::GetGlobal("tvm_callback_metal_compile");
std::string fmt = fmetal_compile ? "metallib" : "metal";
for (auto kv : mod->functions) {
diff --git a/src/target/source/codegen_opencl.cc b/src/target/source/codegen_opencl.cc
index b447c07297..f62a8847f7 100644
--- a/src/target/source/codegen_opencl.cc
+++ b/src/target/source/codegen_opencl.cc
@@ -640,7 +640,6 @@ runtime::Module BuildOpenCL(IRModule mod, Target target) {
}
#endif
- using tvm::runtime::Registry;
bool output_ssa = false;
Map<GlobalVar, PrimFunc> functions;
@@ -654,7 +653,7 @@ runtime::Module BuildOpenCL(IRModule mod, Target target) {
}
std::stringstream code;
- const auto* fpostproc = Registry::Get("tvm_callback_opencl_postproc");
+ const auto fpostproc = tvm::ffi::Function::GetGlobal("tvm_callback_opencl_postproc");
for (auto [gvar, prim_func] : functions) {
code << "// Function: " << gvar->name_hint << std::endl;
CodeGenOpenCL cg;
diff --git a/src/target/spirv/spirv_utils.cc b/src/target/spirv/spirv_utils.cc
index 82c355eda8..180f8f31a6 100644
--- a/src/target/spirv/spirv_utils.cc
+++ b/src/target/spirv/spirv_utils.cc
@@ -111,14 +111,13 @@ class SPIRVTools {
std::pair<std::unordered_map<std::string, runtime::SPIRVShader>, std::string> LowerToSPIRV(
IRModule mod, Target target) {
- using tvm::runtime::Registry;
using tvm::runtime::SPIRVShader;
std::ostringstream code_data;
SPIRVTools spirv_tools(target);
std::unordered_map<std::string, SPIRVShader> smap;
- const auto* postproc = Registry::Get("tvm_callback_vulkan_postproc");
+ auto postproc = tvm::ffi::Function::GetGlobal("tvm_callback_vulkan_postproc");
mod = tir::transform::PointerValueTypeRewrite()(std::move(mod));
@@ -157,7 +156,7 @@ std::pair<std::unordered_map<std::string, runtime::SPIRVShader>, std::string> Lo
spirv_tools.ValidateShader(shader.data);
}
- if (postproc != nullptr) {
+ if (postproc) {
TVMFFIByteArray arr;
arr.data = reinterpret_cast<const char*>(dmlc::BeginPtr(shader.data));
arr.size = shader.data.size() * sizeof(uint32_t);
diff --git a/src/target/target.cc b/src/target/target.cc
index 7dab6d365b..4aa4db3e1e 100644
--- a/src/target/target.cc
+++ b/src/target/target.cc
@@ -794,9 +794,10 @@ ObjectPtr<Object> TargetInternal::FromString(const String& tag_or_config_or_targ
}
ObjectPtr<Object> TargetInternal::FromConfigString(const String& config_str) {
- const auto* loader = tvm::runtime::Registry::Get("target._load_config_dict");
- ICHECK(loader) << "AttributeError: \"target._load_config_dict\" is not registered. Please check "
- "if the python module is properly loaded";
+ const auto loader = tvm::ffi::Function::GetGlobal("target._load_config_dict");
+ ICHECK(loader.has_value())
+ << "AttributeError: \"target._load_config_dict\" is not registered. Please check "
+ "if the python module is properly loaded";
Optional<Map<String, ffi::Any>> config = (*loader)(config_str);
if (!config.defined()) {
TVM_FFI_THROW(ValueError) << "Cannot load config dict with python JSON loader";
@@ -1008,7 +1009,7 @@ std::unordered_map<String, ffi::Any> TargetInternal::QueryDevice(int device_id,
/********** Registry **********/
-TVM_REGISTER_GLOBAL("target.Target").set_body(TargetInternal::ConstructorDispatcher);
+TVM_REGISTER_GLOBAL("target.Target").set_body_packed(TargetInternal::ConstructorDispatcher);
TVM_REGISTER_GLOBAL("target.TargetEnterScope").set_body_typed(TargetInternal::EnterScope);
TVM_REGISTER_GLOBAL("target.TargetExitScope").set_body_typed(TargetInternal::ExitScope);
TVM_REGISTER_GLOBAL("target.TargetCurrent").set_body_typed(Target::Current);
diff --git a/src/target/target_info.cc b/src/target/target_info.cc
index d83c8beac4..b047d3a96b 100644
--- a/src/target/target_info.cc
+++ b/src/target/target_info.cc
@@ -40,8 +40,8 @@ TVM_REGISTER_NODE_TYPE(MemoryInfoNode);
MemoryInfo GetMemoryInfo(const std::string& scope) {
std::string fname = "tvm.info.mem." + scope;
- const runtime::PackedFunc* f = runtime::Registry::Get(fname);
- if (f == nullptr) {
+ const auto f = tvm::ffi::Function::GetGlobal(fname);
+ if (!f.has_value()) {
LOG(WARNING) << "MemoryInfo for scope = " << scope << " is undefined";
return MemoryInfo();
} else {
diff --git a/src/target/target_kind.cc b/src/target/target_kind.cc
index 840f0b2e49..e51a4c21ff 100644
--- a/src/target/target_kind.cc
+++ b/src/target/target_kind.cc
@@ -214,7 +214,6 @@ TargetJSON UpdateNVPTXAttrs(TargetJSON target) {
* \return The updated attributes
*/
TargetJSON UpdateROCmAttrs(TargetJSON target) {
- using tvm::runtime::Registry;
CheckOrSetAttr(&target, "mtriple", "amdgcn-amd-amdhsa-hcc");
// Update -mcpu=gfx
std::string arch = "gfx900";
@@ -224,7 +223,7 @@ TargetJSON UpdateROCmAttrs(TargetJSON target) {
ICHECK(!arch.empty()) << "ValueError: ROCm target gets an invalid GFX version: -mcpu=" << mcpu;
} else {
TVMRetValue val;
- if (const auto* f_get_rocm_arch = Registry::Get("tvm_callback_rocm_get_arch")) {
+ if (const auto f_get_rocm_arch = tvm::ffi::Function::GetGlobal("tvm_callback_rocm_get_arch")) {
arch = (*f_get_rocm_arch)().operator std::string();
}
target.Set("mcpu", String(arch));
diff --git a/src/te/operation/create_primfunc.cc b/src/te/operation/create_primfunc.cc
index b5da43d670..186cd28811 100644
--- a/src/te/operation/create_primfunc.cc
+++ b/src/te/operation/create_primfunc.cc
@@ -742,9 +742,9 @@ PrimFunc GenerateAndCompletePrimFunc(const Array<te::Tensor>& arg_list,
/*ret_type=*/VoidType(),
/*buffer_map=*/std::move(buffer_map)),
{{"global_symbol", String("main")}, {"tir.noalias", true}});
- const auto* complete = runtime::Registry::Get("script.Complete");
- ICHECK(complete);
- func = (*complete)(std::move(func), info->root_alloc);
+ const auto fcomplete = tvm::ffi::Function::GetGlobal("script.Complete");
+ ICHECK(fcomplete.has_value());
+ func = (*fcomplete)(std::move(func), info->root_alloc);
return func;
}
@@ -784,7 +784,7 @@ PrimFunc CreatePrimFunc(const Array<te::Tensor>& arg_list,
return CreatePrimFuncWithConstants(arg_list, {}, index_dtype_override);
}
-TVM_REGISTER_GLOBAL("te.CreatePrimFunc").set_body([](TVMArgs args, TVMRetValue* ret) {
+TVM_REGISTER_GLOBAL("te.CreatePrimFunc").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
Array<ObjectRef> arg_list = args[0];
std::optional<DataType> index_dtype_override{std::nullopt};
// Add conversion to make std::optional compatible with FFI.
@@ -816,10 +816,9 @@ PrimFunc GenerateAndCompletePrimFunc(const Array<ObjectRef>& arg_tir_var_list,
/*ret_type=*/VoidType(),
/*buffer_map=*/std::move(buffer_map)),
{{"global_symbol", String("main")}, {"tir.noalias", true}});
-
- const auto* complete = runtime::Registry::Get("script.Complete");
- ICHECK(complete);
- func = (*complete)(std::move(func), info->root_alloc);
+ const auto fcomplete = tvm::ffi::Function::GetGlobal("script.Complete");
+ ICHECK(fcomplete.has_value());
+ func = (*fcomplete)(std::move(func), info->root_alloc);
return func;
}
diff --git a/src/te/tensor.cc b/src/te/tensor.cc
index f1c621975a..f46c095f3b 100644
--- a/src/te/tensor.cc
+++ b/src/te/tensor.cc
@@ -122,9 +122,9 @@ TVM_REGISTER_GLOBAL("te.OpGetOutput").set_body_typed([](Operation op, int64_t ou
return op.output(static_cast<size_t>(output));
});
-TVM_REGISTER_GLOBAL("te.OpNumOutputs").set_body_method<Operation>(&OperationNode::num_outputs);
+TVM_REGISTER_GLOBAL("te.OpNumOutputs").set_body_method(&OperationNode::num_outputs);
-TVM_REGISTER_GLOBAL("te.OpInputTensors").set_body_method<Operation>(&OperationNode::InputTensors);
+TVM_REGISTER_GLOBAL("te.OpInputTensors").set_body_method(&OperationNode::InputTensors);
} // namespace te
} // namespace tvm
diff --git a/src/tir/analysis/var_use_def_analysis.cc b/src/tir/analysis/var_use_def_analysis.cc
index 9975a872bd..1ca5b670f6 100644
--- a/src/tir/analysis/var_use_def_analysis.cc
+++ b/src/tir/analysis/var_use_def_analysis.cc
@@ -199,14 +199,15 @@ Array<Var> UndefinedVars(const PrimExpr& expr, const Array<Var>& args) {
return m.undefined_;
}
-TVM_REGISTER_GLOBAL("tir.analysis.UndefinedVars").set_body([](TVMArgs args, TVMRetValue* rv) {
- if (auto opt_stmt = args[0].as<Stmt>()) {
- *rv = UndefinedVars(opt_stmt.value(), args[1]);
- } else if (auto opt_expr = args[0].as<PrimExpr>()) {
- *rv = UndefinedVars(opt_expr.value(), args[1]);
- } else {
- LOG(FATAL) << "either UndefinedVars(stmt, args) or UndefinedVars(expr, args) is expected";
- }
-});
+TVM_REGISTER_GLOBAL("tir.analysis.UndefinedVars")
+ .set_body_packed([](TVMArgs args, TVMRetValue* rv) {
+ if (auto opt_stmt = args[0].as<Stmt>()) {
+ *rv = UndefinedVars(opt_stmt.value(), args[1]);
+ } else if (auto opt_expr = args[0].as<PrimExpr>()) {
+ *rv = UndefinedVars(opt_expr.value(), args[1]);
+ } else {
+ LOG(FATAL) << "either UndefinedVars(stmt, args) or UndefinedVars(expr, args) is expected";
+ }
+ });
} // namespace tir
} // namespace tvm
diff --git a/src/tir/ir/block_dependence_info.cc b/src/tir/ir/block_dependence_info.cc
index cd4801d286..13f38c86bf 100644
--- a/src/tir/ir/block_dependence_info.cc
+++ b/src/tir/ir/block_dependence_info.cc
@@ -88,7 +88,7 @@ TVM_REGISTER_NODE_TYPE(BlockDependenceInfoNode);
TVM_REGISTER_GLOBAL("tir.BlockDependenceInfo")
.set_body_typed([](IRModule mod) -> BlockDependenceInfo { return BlockDependenceInfo(mod); });
TVM_REGISTER_GLOBAL("tir.BlockDependenceInfoGetBlockScope")
- .set_body_method<BlockDependenceInfo>(&BlockDependenceInfoNode::GetBlockScope);
+ .set_body_method(&BlockDependenceInfoNode::GetBlockScope);
TVM_REGISTER_GLOBAL("tir.BlockDependenceInfoGetSRef")
.set_body_typed([](BlockDependenceInfo self, Stmt stmt) -> Optional<StmtSRef> {
auto it = self->stmt2ref.find(stmt.get());
diff --git a/src/tir/ir/block_scope.cc b/src/tir/ir/block_scope.cc
index 4bec28619e..5320c5d68a 100644
--- a/src/tir/ir/block_scope.cc
+++ b/src/tir/ir/block_scope.cc
@@ -200,10 +200,8 @@ TVM_REGISTER_GLOBAL("tir.StmtSRefRootMark") //
.set_body_typed(StmtSRef::RootMark);
TVM_REGISTER_GLOBAL("tir.StmtSRefInlineMark") //
.set_body_typed(StmtSRef::InlineMark);
-TVM_REGISTER_GLOBAL("tir.BlockScopeGetDepsBySrc")
- .set_body_method<BlockScope>(&BlockScopeNode::GetDepsBySrc);
-TVM_REGISTER_GLOBAL("tir.BlockScopeGetDepsByDst")
- .set_body_method<BlockScope>(&BlockScopeNode::GetDepsByDst);
+TVM_REGISTER_GLOBAL("tir.BlockScopeGetDepsBySrc").set_body_method(&BlockScopeNode::GetDepsBySrc);
+TVM_REGISTER_GLOBAL("tir.BlockScopeGetDepsByDst").set_body_method(&BlockScopeNode::GetDepsByDst);
} // namespace tir
} // namespace tvm
diff --git a/src/tir/ir/buffer.cc b/src/tir/ir/buffer.cc
index fefcc5b0de..4c16c9805a 100644
--- a/src/tir/ir/buffer.cc
+++ b/src/tir/ir/buffer.cc
@@ -640,7 +640,7 @@ tir::Buffer BufferWithOffsetAlignment(Array<PrimExpr> shape, DataType dtype, std
TVM_REGISTER_NODE_TYPE(BufferNode);
-TVM_REGISTER_GLOBAL("tir.Buffer").set_body([](TVMArgs args, TVMRetValue* ret) {
+TVM_REGISTER_GLOBAL("tir.Buffer").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
ICHECK_EQ(args.size(), 11);
auto buffer_type = args[8].operator String();
BufferType type = (buffer_type == "auto_broadcast") ? kAutoBroadcast : kDefault;
diff --git a/src/tir/ir/expr.cc b/src/tir/ir/expr.cc
index 1590c51f38..7270c014cc 100644
--- a/src/tir/ir/expr.cc
+++ b/src/tir/ir/expr.cc
@@ -696,8 +696,7 @@ TVM_REGISTER_GLOBAL("tir.CommReducer")
return CommReducer(lhs, rhs, result, identity_element, span);
});
-TVM_REGISTER_GLOBAL("tir.CommReducerCombine")
- .set_body_method<tir::CommReducer>(&tir::CommReducerNode::operator());
+TVM_REGISTER_GLOBAL("tir.CommReducerCombine").set_body_method(&tir::CommReducerNode::operator());
TVM_REGISTER_NODE_TYPE(CommReducerNode);
diff --git a/src/tir/op/op.cc b/src/tir/op/op.cc
index d6f90457b0..ecfca1730c 100644
--- a/src/tir/op/op.cc
+++ b/src/tir/op/op.cc
@@ -1120,17 +1120,17 @@ TVM_REGISTER_GLOBAL("tir.reinterpret").set_body_typed(tvm::reinterpret);
return (Func(a, b, span)); \
})
-#define REGISTER_MAKE_BIT_OP(Node, Func) \
- TVM_REGISTER_GLOBAL("tir." #Node).set_body([](TVMArgs args, TVMRetValue* ret) { \
- bool lhs_is_int = args[0].type_index() == ffi::TypeIndex::kTVMFFIInt; \
- bool rhs_is_int = args[1].type_index() == ffi::TypeIndex::kTVMFFIInt; \
- if (lhs_is_int) { \
- *ret = (Func(args[0].operator int(), args[1].operator PrimExpr(), args[2])); \
- } else if (rhs_is_int) { \
- *ret = (Func(args[0].operator PrimExpr(), args[1].operator int(), args[2])); \
- } else { \
- *ret = (Func(args[0].operator PrimExpr(), args[1].operator PrimExpr(), args[2])); \
- } \
+#define REGISTER_MAKE_BIT_OP(Node, Func) \
+ TVM_REGISTER_GLOBAL("tir." #Node).set_body_packed([](TVMArgs args, TVMRetValue* ret) { \
+ bool lhs_is_int = args[0].type_index() == ffi::TypeIndex::kTVMFFIInt; \
+ bool rhs_is_int = args[1].type_index() == ffi::TypeIndex::kTVMFFIInt; \
+ if (lhs_is_int) { \
+ *ret = (Func(args[0].operator int(), args[1].operator PrimExpr(), args[2])); \
+ } else if (rhs_is_int) { \
+ *ret = (Func(args[0].operator PrimExpr(), args[1].operator int(), args[2])); \
+ } else { \
+ *ret = (Func(args[0].operator PrimExpr(), args[1].operator PrimExpr(), args[2])); \
+ } \
})
REGISTER_MAKE_BINARY_OP(_OpAdd, add);
diff --git a/src/tir/schedule/schedule.cc b/src/tir/schedule/schedule.cc
index 6e3a1762e5..29e40d4003 100644
--- a/src/tir/schedule/schedule.cc
+++ b/src/tir/schedule/schedule.cc
@@ -44,21 +44,21 @@ TVM_REGISTER_NODE_TYPE(LoopRVNode);
TVM_REGISTER_OBJECT_TYPE(ScheduleNode);
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleGetMod") //
- .set_body_method<Schedule>(&ScheduleNode::mod);
+ .set_body_method(&ScheduleNode::mod);
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleGetState") //
- .set_body_method<Schedule>(&ScheduleNode::state);
+ .set_body_method(&ScheduleNode::state);
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleGetTrace") //
- .set_body_method<Schedule>(&ScheduleNode::trace);
+ .set_body_method(&ScheduleNode::trace);
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleGetFuncWorkingOn") //
- .set_body_method<Schedule>(&ScheduleNode::func_working_on);
+ .set_body_method(&ScheduleNode::func_working_on);
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleCopy") //
- .set_body_method<Schedule>(&ScheduleNode::Copy);
+ .set_body_method(&ScheduleNode::Copy);
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleSeed") //
- .set_body_method<Schedule>(&ScheduleNode::Seed);
+ .set_body_method(&ScheduleNode::Seed);
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleForkSeed") //
- .set_body_method<Schedule>(&ScheduleNode::ForkSeed);
+ .set_body_method(&ScheduleNode::ForkSeed);
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleWorkOn") //
- .set_body_method<Schedule>(&ScheduleNode::WorkOn);
+ .set_body_method(&ScheduleNode::WorkOn);
/**************** (FFI) Constructor ****************/
@@ -127,18 +127,16 @@ TVM_REGISTER_GLOBAL("tir.schedule.ScheduleRemoveRV")
/******** (FFI) Sampling ********/
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleSampleCategorical")
- .set_body_method<Schedule>(&ScheduleNode::SampleCategorical);
+ .set_body_method(&ScheduleNode::SampleCategorical);
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleSamplePerfectTile")
- .set_body_method<Schedule>(&ScheduleNode::SamplePerfectTile);
+ .set_body_method(&ScheduleNode::SamplePerfectTile);
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleSamplePartitionedTile")
- .set_body_method<Schedule>(&ScheduleNode::SamplePartitionedTile);
+ .set_body_method(&ScheduleNode::SamplePartitionedTile);
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleSampleComputeLocation")
- .set_body_method<Schedule>(&ScheduleNode::SampleComputeLocation);
+ .set_body_method(&ScheduleNode::SampleComputeLocation);
/******** (FFI) Get blocks & loops ********/
-TVM_REGISTER_GLOBAL("tir.schedule.ScheduleGetBlock")
- .set_body_method<Schedule>(&ScheduleNode::GetBlock);
-TVM_REGISTER_GLOBAL("tir.schedule.ScheduleGetLoops")
- .set_body_method<Schedule>(&ScheduleNode::GetLoops);
+TVM_REGISTER_GLOBAL("tir.schedule.ScheduleGetBlock").set_body_method(&ScheduleNode::GetBlock);
+TVM_REGISTER_GLOBAL("tir.schedule.ScheduleGetLoops").set_body_method(&ScheduleNode::GetLoops);
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleGetChildBlocks")
.set_body_typed([](Schedule self, ObjectRef rv) {
if (auto block_rv = rv.as<BlockRV>()) {
@@ -152,21 +150,20 @@ TVM_REGISTER_GLOBAL("tir.schedule.ScheduleGetChildBlocks")
throw;
});
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleGetProducers")
- .set_body_method<Schedule>(&ScheduleNode::GetProducers);
+ .set_body_method(&ScheduleNode::GetProducers);
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleGetConsumers")
- .set_body_method<Schedule>(&ScheduleNode::GetConsumers);
+ .set_body_method(&ScheduleNode::GetConsumers);
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleGetOutputBlocks")
- .set_body_method<Schedule>(&ScheduleNode::GetOutputBlocks);
+ .set_body_method(&ScheduleNode::GetOutputBlocks);
/******** (FFI) Transform loops ********/
-TVM_REGISTER_GLOBAL("tir.schedule.ScheduleMerge").set_body_method<Schedule>(&ScheduleNode::Merge);
-TVM_REGISTER_GLOBAL("tir.schedule.ScheduleFuse").set_body_method<Schedule>(&ScheduleNode::Fuse);
-TVM_REGISTER_GLOBAL("tir.schedule.ScheduleSplit").set_body_method<Schedule>(&ScheduleNode::Split);
+TVM_REGISTER_GLOBAL("tir.schedule.ScheduleMerge").set_body_method(&ScheduleNode::Merge);
+TVM_REGISTER_GLOBAL("tir.schedule.ScheduleFuse").set_body_method(&ScheduleNode::Fuse);
+TVM_REGISTER_GLOBAL("tir.schedule.ScheduleSplit").set_body_method(&ScheduleNode::Split);
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleLoopPartition")
- .set_body_method<Schedule>(&ScheduleNode::LoopPartition);
-TVM_REGISTER_GLOBAL("tir.schedule.ScheduleReorder")
- .set_body_method<Schedule>(&ScheduleNode::Reorder);
+ .set_body_method(&ScheduleNode::LoopPartition);
+TVM_REGISTER_GLOBAL("tir.schedule.ScheduleReorder").set_body_method(&ScheduleNode::Reorder);
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleReorderBlockIterVar")
- .set_body_method<Schedule>(&ScheduleNode::ReorderBlockIterVar);
+ .set_body_method(&ScheduleNode::ReorderBlockIterVar);
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleAddUnitLoop")
.set_body_typed([](Schedule self, ObjectRef rv) -> LoopRV {
if (auto loop_rv = rv.as<LoopRV>()) {
@@ -180,55 +177,46 @@ TVM_REGISTER_GLOBAL("tir.schedule.ScheduleAddUnitLoop")
}
});
/******** (FFI) Manipulate ForKind ********/
-TVM_REGISTER_GLOBAL("tir.schedule.ScheduleParallel")
- .set_body_method<Schedule>(&ScheduleNode::Parallel);
-TVM_REGISTER_GLOBAL("tir.schedule.ScheduleVectorize")
- .set_body_method<Schedule>(&ScheduleNode::Vectorize);
-TVM_REGISTER_GLOBAL("tir.schedule.ScheduleBind").set_body_method<Schedule>(&ScheduleNode::Bind);
-TVM_REGISTER_GLOBAL("tir.schedule.ScheduleUnroll").set_body_method<Schedule>(&ScheduleNode::Unroll);
+TVM_REGISTER_GLOBAL("tir.schedule.ScheduleParallel").set_body_method(&ScheduleNode::Parallel);
+TVM_REGISTER_GLOBAL("tir.schedule.ScheduleVectorize").set_body_method(&ScheduleNode::Vectorize);
+TVM_REGISTER_GLOBAL("tir.schedule.ScheduleBind").set_body_method(&ScheduleNode::Bind);
+TVM_REGISTER_GLOBAL("tir.schedule.ScheduleUnroll").set_body_method(&ScheduleNode::Unroll);
/******** (FFI) Insert cache stages ********/
-TVM_REGISTER_GLOBAL("tir.schedule.ScheduleCacheRead")
- .set_body_method<Schedule>(&ScheduleNode::CacheRead);
-TVM_REGISTER_GLOBAL("tir.schedule.ScheduleCacheWrite")
- .set_body_method<Schedule>(&ScheduleNode::CacheWrite);
+TVM_REGISTER_GLOBAL("tir.schedule.ScheduleCacheRead").set_body_method(&ScheduleNode::CacheRead);
+TVM_REGISTER_GLOBAL("tir.schedule.ScheduleCacheWrite").set_body_method(&ScheduleNode::CacheWrite);
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleReindexCacheRead")
- .set_body_method<Schedule>(&ScheduleNode::ReindexCacheRead);
+ .set_body_method(&ScheduleNode::ReindexCacheRead);
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleReindexCacheWrite")
- .set_body_method<Schedule>(&ScheduleNode::ReindexCacheWrite);
+ .set_body_method(&ScheduleNode::ReindexCacheWrite);
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleCacheInplace")
- .set_body_method<Schedule>(&ScheduleNode::CacheInplace);
-TVM_REGISTER_GLOBAL("tir.schedule.ScheduleCacheIndex")
- .set_body_method<Schedule>(&ScheduleNode::CacheIndex);
+ .set_body_method(&ScheduleNode::CacheInplace);
+TVM_REGISTER_GLOBAL("tir.schedule.ScheduleCacheIndex").set_body_method(&ScheduleNode::CacheIndex);
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleReIndex")
.set_body_typed([](Schedule self, const BlockRV& block_rv, int buffer_index,
int buffer_index_type) {
return self->ReIndex(block_rv, buffer_index, static_cast<BufferIndexType>(buffer_index_type));
});
/******** (FFI) Data movement ********/
-TVM_REGISTER_GLOBAL("tir.schedule.ScheduleReadAt").set_body_method<Schedule>(&ScheduleNode::ReadAt);
-TVM_REGISTER_GLOBAL("tir.schedule.ScheduleWriteAt")
- .set_body_method<Schedule>(&ScheduleNode::WriteAt);
+TVM_REGISTER_GLOBAL("tir.schedule.ScheduleReadAt").set_body_method(&ScheduleNode::ReadAt);
+TVM_REGISTER_GLOBAL("tir.schedule.ScheduleWriteAt").set_body_method(&ScheduleNode::WriteAt);
/******** (FFI) Compute location ********/
-TVM_REGISTER_GLOBAL("tir.schedule.ScheduleComputeAt")
- .set_body_method<Schedule>(&ScheduleNode::ComputeAt);
+TVM_REGISTER_GLOBAL("tir.schedule.ScheduleComputeAt").set_body_method(&ScheduleNode::ComputeAt);
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleReverseComputeAt")
- .set_body_method<Schedule>(&ScheduleNode::ReverseComputeAt);
+ .set_body_method(&ScheduleNode::ReverseComputeAt);
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleComputeInline")
- .set_body_method<Schedule>(&ScheduleNode::ComputeInline);
+ .set_body_method(&ScheduleNode::ComputeInline);
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleReverseComputeInline")
- .set_body_method<Schedule>(&ScheduleNode::ReverseComputeInline);
+ .set_body_method(&ScheduleNode::ReverseComputeInline);
/******** (FFI) Reduction ********/
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleDecomposeReduction")
- .set_body_method<Schedule>(&ScheduleNode::DecomposeReduction);
-TVM_REGISTER_GLOBAL("tir.schedule.ScheduleRFactor")
- .set_body_method<Schedule>(&ScheduleNode::RFactor);
+ .set_body_method(&ScheduleNode::DecomposeReduction);
+TVM_REGISTER_GLOBAL("tir.schedule.ScheduleRFactor").set_body_method(&ScheduleNode::RFactor);
/******** (FFI) Block annotation ********/
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleStorageAlign")
- .set_body_method<Schedule>(&ScheduleNode::StorageAlign);
-TVM_REGISTER_GLOBAL("tir.schedule.ScheduleSetScope")
- .set_body_method<Schedule>(&ScheduleNode::SetScope);
+ .set_body_method(&ScheduleNode::StorageAlign);
+TVM_REGISTER_GLOBAL("tir.schedule.ScheduleSetScope").set_body_method(&ScheduleNode::SetScope);
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleUnsafeSetDType")
- .set_body_method<Schedule>(&ScheduleNode::UnsafeSetDType);
+ .set_body_method(&ScheduleNode::UnsafeSetDType);
/******** (FFI) Blockize & Tensorize ********/
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleBlockize")
.set_body_typed([](Schedule self, ObjectRef target, bool preserve_unit_iters) {
@@ -287,7 +275,7 @@ TVM_REGISTER_GLOBAL("tir.schedule.ScheduleTransformLayout")
pad_value, assume_injective_transform);
});
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleTransformBlockLayout")
- .set_body_method<Schedule>(&ScheduleNode::TransformBlockLayout);
+ .set_body_method(&ScheduleNode::TransformBlockLayout);
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleSetAxisSeparator")
.set_body_typed([](Schedule self, const BlockRV& block_rv, int buffer_index,
int buffer_index_type, const Array<IntImm>& axis_separators) {
@@ -297,17 +285,16 @@ TVM_REGISTER_GLOBAL("tir.schedule.ScheduleSetAxisSeparator")
/******** (FFI) Padding decomposition ********/
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleDecomposePadding")
- .set_body_method<Schedule>(&ScheduleNode::DecomposePadding);
-TVM_REGISTER_GLOBAL("tir.schedule.SchedulePadEinsum")
- .set_body_method<Schedule>(&ScheduleNode::PadEinsum);
+ .set_body_method(&ScheduleNode::DecomposePadding);
+TVM_REGISTER_GLOBAL("tir.schedule.SchedulePadEinsum").set_body_method(&ScheduleNode::PadEinsum);
/******** (FFI) Buffer transformation ********/
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleRollingBuffer")
- .set_body_method<Schedule>(&ScheduleNode::RollingBuffer);
+ .set_body_method(&ScheduleNode::RollingBuffer);
/******** (FFI) Misc ********/
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleEnterPostproc")
- .set_body_method<Schedule>(&ScheduleNode::EnterPostproc);
+ .set_body_method(&ScheduleNode::EnterPostproc);
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleUnsafeHideBufferAccess")
- .set_body_method<Schedule>(&ScheduleNode::UnsafeHideBufferAccess);
+ .set_body_method(&ScheduleNode::UnsafeHideBufferAccess);
/******** (FFI) Annotate buffer access ********/
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleAnnotateBufferAccess")
.set_body_typed([](Schedule self, const BlockRV& block_rv, int buffer_index,
diff --git a/src/tir/schedule/state.cc b/src/tir/schedule/state.cc
index f2301a1a77..43a3c9008f 100644
--- a/src/tir/schedule/state.cc
+++ b/src/tir/schedule/state.cc
@@ -1017,9 +1017,9 @@ TVM_REGISTER_GLOBAL("tir.schedule.ScheduleState")
return ScheduleState(mod, debug_mask, enable_check);
});
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleStateGetBlockScope")
- .set_body_method<ScheduleState>(&ScheduleStateNode::GetBlockScope);
+ .set_body_method(&ScheduleStateNode::GetBlockScope);
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleStateReplace")
- .set_body_method<ScheduleState>(&ScheduleStateNode::Replace);
+ .set_body_method(&ScheduleStateNode::Replace);
TVM_REGISTER_GLOBAL("tir.schedule.ScheduleStateGetSRef")
.set_body_typed([](ScheduleState self, Stmt stmt) -> Optional<StmtSRef> {
auto it = self->stmt2ref.find(stmt.get());
diff --git a/src/tir/schedule/trace.cc b/src/tir/schedule/trace.cc
index 8eb95be8fb..13bb93a798 100644
--- a/src/tir/schedule/trace.cc
+++ b/src/tir/schedule/trace.cc
@@ -567,8 +567,7 @@ TVM_REGISTER_GLOBAL("tir.schedule.Trace")
Optional<Map<Instruction, Any>> decisions) {
return Trace(insts.value_or(Array<Instruction>()), decisions.value_or({}));
});
-TVM_REGISTER_GLOBAL("tir.schedule.TraceGetDecision")
- .set_body_method<Trace>(&TraceNode::GetDecision);
+TVM_REGISTER_GLOBAL("tir.schedule.TraceGetDecision").set_body_method(&TraceNode::GetDecision);
TVM_REGISTER_GLOBAL("tir.schedule.TraceAppend")
.set_body_typed([](Trace self, Instruction inst, Optional<ObjectRef> decision) {
if (decision.defined()) {
@@ -577,14 +576,13 @@ TVM_REGISTER_GLOBAL("tir.schedule.TraceAppend")
return self->Append(inst);
}
});
-TVM_REGISTER_GLOBAL("tir.schedule.TracePop").set_body_method<Trace>(&TraceNode::Pop);
+TVM_REGISTER_GLOBAL("tir.schedule.TracePop").set_body_method(&TraceNode::Pop);
TVM_REGISTER_GLOBAL("tir.schedule.TraceApplyToSchedule")
- .set_body_method<Trace>(&TraceNode::ApplyToSchedule);
-TVM_REGISTER_GLOBAL("tir.schedule.TraceAsJSON").set_body_method<Trace>(&TraceNode::AsJSON);
-TVM_REGISTER_GLOBAL("tir.schedule.TraceAsPython").set_body_method<Trace>(&TraceNode::AsPython);
-TVM_REGISTER_GLOBAL("tir.schedule.TraceWithDecision")
- .set_body_method<Trace>(&TraceNode::WithDecision);
-TVM_REGISTER_GLOBAL("tir.schedule.TraceSimplified").set_body_method<Trace>(&TraceNode::Simplified);
+ .set_body_method(&TraceNode::ApplyToSchedule);
+TVM_REGISTER_GLOBAL("tir.schedule.TraceAsJSON").set_body_method(&TraceNode::AsJSON);
+TVM_REGISTER_GLOBAL("tir.schedule.TraceAsPython").set_body_method(&TraceNode::AsPython);
+TVM_REGISTER_GLOBAL("tir.schedule.TraceWithDecision").set_body_method(&TraceNode::WithDecision);
+TVM_REGISTER_GLOBAL("tir.schedule.TraceSimplified").set_body_method(&TraceNode::Simplified);
TVM_REGISTER_GLOBAL("tir.schedule.TraceApplyJSONToSchedule")
.set_body_typed(Trace::ApplyJSONToSchedule);
diff --git a/src/tir/transforms/transform_mma_buffer_layout.cc b/src/tir/transforms/transform_mma_buffer_layout.cc
index 899f292b8f..5332fcfff1 100644
--- a/src/tir/transforms/transform_mma_buffer_layout.cc
+++ b/src/tir/transforms/transform_mma_buffer_layout.cc
@@ -127,8 +127,8 @@ class MmaBufferLayoutTransformer : public StmtExprMutator {
if (buffer_map_.count(store->buffer)) {
auto* n = store.CopyOnWrite();
if (store->buffer.scope() == "m16n8k8.matrixC") {
- const auto* index_map_func = runtime::Registry::Get("tir.index_map_m16n8k8.matrixC");
- ICHECK(index_map_func);
+ const auto index_map_func = tvm::ffi::Function::GetGlobal("tir.index_map_m16n8k8.matrixC");
+ ICHECK(index_map_func.has_value());
auto index_map = IndexMap::FromFunc(2, *index_map_func);
auto new_indices = index_map->MapIndices(store->indices, &analyzer);
n->buffer = buffer_map_[store->buffer];
@@ -146,8 +146,8 @@ class MmaBufferLayoutTransformer : public StmtExprMutator {
if (buffer_map_.count(load->buffer)) {
auto* n = load.CopyOnWrite();
if (load->buffer.scope() == "m16n8k8.matrixC") {
- const auto* index_map_func = runtime::Registry::Get("tir.index_map_m16n8k8.matrixC");
- ICHECK(index_map_func);
+ const auto index_map_func = tvm::ffi::Function::GetGlobal("tir.index_map_m16n8k8.matrixC");
+ ICHECK(index_map_func.has_value());
auto index_map = IndexMap::FromFunc(2, *index_map_func);
auto new_indices = index_map->MapIndices(load->indices, &analyzer);
n->buffer = buffer_map_[load->buffer];
diff --git a/src/tir/transforms/unsupported_dtype_legalize.cc b/src/tir/transforms/unsupported_dtype_legalize.cc
index 5374c68f17..29ac31617f 100644
--- a/src/tir/transforms/unsupported_dtype_legalize.cc
+++ b/src/tir/transforms/unsupported_dtype_legalize.cc
@@ -737,10 +737,9 @@ namespace transform {
bool CheckDataTypeSupport(const Target& target, const std::string& support_func_name) {
bool has_native_support = false;
if (target->kind->name == "cuda") {
- if (const PackedFunc* get_cv =
- tvm::runtime::Registry::Get("tvm.contrib.nvcc.get_compute_version")) {
+ if (auto get_cv = tvm::ffi::Function::GetGlobal("tvm.contrib.nvcc.get_compute_version")) {
std::string compute_version = (*get_cv)(target);
- if (const PackedFunc* check_support = tvm::runtime::Registry::Get(support_func_name)) {
+ if (auto check_support = tvm::ffi::Function::GetGlobal(support_func_name)) {
has_native_support = (*check_support)(compute_version);
}
}
diff --git a/src/topi/broadcast.cc b/src/topi/broadcast.cc
index 31caf7280a..cb34617f65 100644
--- a/src/topi/broadcast.cc
+++ b/src/topi/broadcast.cc
@@ -33,7 +33,7 @@ using namespace tvm;
using namespace tvm::runtime;
#define TOPI_REGISTER_BCAST_OP(OpName, Op) \
- TVM_REGISTER_GLOBAL(OpName).set_body([](TVMArgs args, TVMRetValue* rv) { \
+ TVM_REGISTER_GLOBAL(OpName).set_body_packed([](TVMArgs args, TVMRetValue* rv) { \
bool lhs_is_tensor = args[0].as<tvm::te::Tensor>().has_value(); \
bool rhs_is_tensor = args[1].as<tvm::te::Tensor>().has_value(); \
if (lhs_is_tensor && rhs_is_tensor) { \
@@ -73,7 +73,7 @@ TOPI_REGISTER_BCAST_OP("topi.not_equal", topi::not_equal);
TOPI_REGISTER_BCAST_OP("topi.greater_equal", topi::greater_equal);
TOPI_REGISTER_BCAST_OP("topi.less_equal", topi::less_equal);
-TVM_REGISTER_GLOBAL("topi.broadcast_to").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.broadcast_to").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = broadcast_to(args[0], args[1]);
});
diff --git a/src/topi/einsum.cc b/src/topi/einsum.cc
index c69baf3ebd..47e5aeb4d2 100644
--- a/src/topi/einsum.cc
+++ b/src/topi/einsum.cc
@@ -355,7 +355,7 @@ Array<PrimExpr> InferEinsumShape(const std::string& subscripts,
return einsum_builder.InferShape();
}
-TVM_REGISTER_GLOBAL("topi.einsum").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.einsum").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = einsum(args[0], args[1]);
});
diff --git a/src/topi/elemwise.cc b/src/topi/elemwise.cc
index 7a39758f6c..a7d554ca09 100644
--- a/src/topi/elemwise.cc
+++ b/src/topi/elemwise.cc
@@ -31,127 +31,139 @@ namespace topi {
using namespace tvm;
using namespace tvm::runtime;
-TVM_REGISTER_GLOBAL("topi.acos").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.acos").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = acos(args[0]);
});
-TVM_REGISTER_GLOBAL("topi.acosh").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.acosh").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = acosh(args[0]);
});
-TVM_REGISTER_GLOBAL("topi.asin").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.asin").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = asin(args[0]);
});
-TVM_REGISTER_GLOBAL("topi.asinh").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.asinh").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = asinh(args[0]);
});
-TVM_REGISTER_GLOBAL("topi.atanh").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.atanh").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = atanh(args[0]);
});
-TVM_REGISTER_GLOBAL("topi.exp").set_body([](TVMArgs args, TVMRetValue* rv) { *rv = exp(args[0]); });
+TVM_REGISTER_GLOBAL("topi.exp").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
+ *rv = exp(args[0]);
+});
-TVM_REGISTER_GLOBAL("topi.fast_exp").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.fast_exp").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = fast_exp(args[0]);
});
-TVM_REGISTER_GLOBAL("topi.erf").set_body([](TVMArgs args, TVMRetValue* rv) { *rv = erf(args[0]); });
+TVM_REGISTER_GLOBAL("topi.erf").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
+ *rv = erf(args[0]);
+});
-TVM_REGISTER_GLOBAL("topi.fast_erf").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.fast_erf").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = fast_erf(args[0]);
});
-TVM_REGISTER_GLOBAL("topi.tan").set_body([](TVMArgs args, TVMRetValue* rv) { *rv = tan(args[0]); });
+TVM_REGISTER_GLOBAL("topi.tan").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
+ *rv = tan(args[0]);
+});
-TVM_REGISTER_GLOBAL("topi.cos").set_body([](TVMArgs args, TVMRetValue* rv) { *rv = cos(args[0]); });
+TVM_REGISTER_GLOBAL("topi.cos").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
+ *rv = cos(args[0]);
+});
-TVM_REGISTER_GLOBAL("topi.cosh").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.cosh").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = cosh(args[0]);
});
-TVM_REGISTER_GLOBAL("topi.sin").set_body([](TVMArgs args, TVMRetValue* rv) { *rv = sin(args[0]); });
+TVM_REGISTER_GLOBAL("topi.sin").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
+ *rv = sin(args[0]);
+});
-TVM_REGISTER_GLOBAL("topi.sinh").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.sinh").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = sinh(args[0]);
});
-TVM_REGISTER_GLOBAL("topi.tanh").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.tanh").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = tanh(args[0]);
});
-TVM_REGISTER_GLOBAL("topi.fast_tanh").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.fast_tanh").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = fast_tanh(args[0]);
});
-TVM_REGISTER_GLOBAL("topi.atan").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.atan").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = atan(args[0]);
});
-TVM_REGISTER_GLOBAL("topi.sigmoid").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.sigmoid").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = sigmoid(args[0]);
});
-TVM_REGISTER_GLOBAL("topi.sqrt").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.sqrt").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = sqrt(args[0]);
});
-TVM_REGISTER_GLOBAL("topi.rsqrt").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.rsqrt").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = rsqrt(args[0]);
});
-TVM_REGISTER_GLOBAL("topi.log").set_body([](TVMArgs args, TVMRetValue* rv) { *rv = log(args[0]); });
+TVM_REGISTER_GLOBAL("topi.log").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
+ *rv = log(args[0]);
+});
-TVM_REGISTER_GLOBAL("topi.log2").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.log2").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = log2(args[0]);
});
-TVM_REGISTER_GLOBAL("topi.log10").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.log10").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = log10(args[0]);
});
-TVM_REGISTER_GLOBAL("topi.identity").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.identity").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = identity(args[0]);
});
-TVM_REGISTER_GLOBAL("topi.negative").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.negative").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = negative(args[0]);
});
-TVM_REGISTER_GLOBAL("topi.clip").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.clip").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = clip(args[0], args[1], args[2]);
});
-TVM_REGISTER_GLOBAL("topi.cast").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.cast").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = cast(args[0], args[1]);
});
-TVM_REGISTER_GLOBAL("topi.reinterpret").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.reinterpret").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = reinterpret(args[0], args[1]);
});
-TVM_REGISTER_GLOBAL("topi.elemwise_sum").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.elemwise_sum").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = elemwise_sum(args[0]);
});
-TVM_REGISTER_GLOBAL("topi.sign").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.sign").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = sign(args[0]);
});
-TVM_REGISTER_GLOBAL("topi.full").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.full").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = full(args[0], args[1], args[2]);
});
-TVM_REGISTER_GLOBAL("topi.full_like").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.full_like").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = full_like(args[0], args[1]);
});
-TVM_REGISTER_GLOBAL("topi.logical_not").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.logical_not").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = logical_not(args[0]);
});
-TVM_REGISTER_GLOBAL("topi.bitwise_not").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.bitwise_not").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = bitwise_not(args[0]);
});
diff --git a/src/topi/nn.cc b/src/topi/nn.cc
index 09859e3318..331e3737ec 100644
--- a/src/topi/nn.cc
+++ b/src/topi/nn.cc
@@ -45,145 +45,145 @@ using namespace tvm;
using namespace tvm::runtime;
/* Ops from nn.h */
-TVM_REGISTER_GLOBAL("topi.nn.relu").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.nn.relu").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = relu<float>(args[0]);
});
-TVM_REGISTER_GLOBAL("topi.nn.leaky_relu").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.nn.leaky_relu").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = leaky_relu(args[0], args[1]);
});
-TVM_REGISTER_GLOBAL("topi.nn.prelu").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.nn.prelu").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = prelu(args[0], args[1], args[2]);
});
-TVM_REGISTER_GLOBAL("topi.nn.pad").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.nn.pad").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = pad(args[0], args[1], args[2], args[3]);
});
-TVM_REGISTER_GLOBAL("topi.nn.space_to_batch_nd").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.nn.space_to_batch_nd").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = space_to_batch_nd(args[0], args[1], args[2], args[3], args[4]);
});
-TVM_REGISTER_GLOBAL("topi.nn.batch_to_space_nd").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.nn.batch_to_space_nd").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = batch_to_space_nd(args[0], args[1], args[2], args[3]);
});
-TVM_REGISTER_GLOBAL("topi.nn.nll_loss").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.nn.nll_loss").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = nll_loss(args[0], args[1], args[2], args[3], args[4]);
});
/* Ops from nn/dense.h */
-TVM_REGISTER_GLOBAL("topi.nn.dense").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.nn.dense").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = nn::dense(args[0], args[1], args[2], args[3]);
});
/* Ops from nn/bias_add.h */
-TVM_REGISTER_GLOBAL("topi.nn.bias_add").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.nn.bias_add").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = nn::bias_add(args[0], args[1], args[2]);
});
/* Ops from nn/dilate.h */
-TVM_REGISTER_GLOBAL("topi.nn.dilate").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.nn.dilate").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = nn::dilate(args[0], args[1], args[2]);
});
/* Ops from nn/flatten.h */
-TVM_REGISTER_GLOBAL("topi.nn.flatten").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.nn.flatten").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = nn::flatten(args[0]);
});
/* Ops from nn/mapping.h */
-TVM_REGISTER_GLOBAL("topi.nn.scale_shift_nchw").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.nn.scale_shift_nchw").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = nn::scale_shift_nchw(args[0], args[1], args[2]);
});
-TVM_REGISTER_GLOBAL("topi.nn.scale_shift_nhwc").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.nn.scale_shift_nhwc").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = nn::scale_shift_nhwc(args[0], args[1], args[2]);
});
/* Ops from nn/pooling.h */
-TVM_REGISTER_GLOBAL("topi.nn.pool_grad").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.nn.pool_grad").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = nn::pool_grad(args[0], args[1], args[2], args[3], args[4],
static_cast<nn::PoolType>(static_cast<int>(args[5])), args[6], args[7],
args[8]);
});
-TVM_REGISTER_GLOBAL("topi.nn.global_pool").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.nn.global_pool").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = nn::global_pool(args[0], static_cast<nn::PoolType>(static_cast<int>(args[1])), args[2]);
});
-TVM_REGISTER_GLOBAL("topi.nn.adaptive_pool1d").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.nn.adaptive_pool1d").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = nn::adaptive_pool1d(args[0], args[1], static_cast<nn::PoolType>(static_cast<int>(args[2])),
args[3]);
});
-TVM_REGISTER_GLOBAL("topi.nn.adaptive_pool").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.nn.adaptive_pool").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = nn::adaptive_pool(args[0], args[1], static_cast<nn::PoolType>(static_cast<int>(args[2])),
args[3]);
});
-TVM_REGISTER_GLOBAL("topi.nn.adaptive_pool3d").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.nn.adaptive_pool3d").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = nn::adaptive_pool3d(args[0], args[1], static_cast<nn::PoolType>(static_cast<int>(args[2])),
args[3]);
});
-TVM_REGISTER_GLOBAL("topi.nn.pool1d").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.nn.pool1d").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = nn::pool1d(args[0], args[1], args[2], args[3], args[4],
static_cast<nn::PoolType>(static_cast<int>(args[5])), args[6], args[7], args[8]);
});
-TVM_REGISTER_GLOBAL("topi.nn.pool2d").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.nn.pool2d").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = nn::pool2d(args[0], args[1], args[2], args[3], args[4],
static_cast<nn::PoolType>(static_cast<int>(args[5])), args[6], args[7], args[8]);
});
-TVM_REGISTER_GLOBAL("topi.nn.pool3d").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.nn.pool3d").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = nn::pool3d(args[0], args[1], args[2], args[3], args[4],
static_cast<nn::PoolType>(static_cast<int>(args[5])), args[6], args[7], args[8]);
});
/* Ops from nn/softmax.h */
-TVM_REGISTER_GLOBAL("topi.nn.softmax").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.nn.softmax").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = nn::softmax(args[0], args[1]);
});
-TVM_REGISTER_GLOBAL("topi.nn.log_softmax").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.nn.log_softmax").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = nn::log_softmax(args[0]);
});
-TVM_REGISTER_GLOBAL("topi.nn.lrn").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.nn.lrn").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = nn::lrn(args[0], args[1], args[2], static_cast<double>(args[3]),
static_cast<double>(args[4]), static_cast<double>(args[5]));
});
/* Ops from nn/bnn.h */
-TVM_REGISTER_GLOBAL("topi.nn.binarize_pack").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.nn.binarize_pack").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = nn::binarize_pack(args[0], args[1]);
});
-TVM_REGISTER_GLOBAL("topi.nn.binary_dense").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.nn.binary_dense").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = nn::binary_dense(args[0], args[1]);
});
/* Ops from nn/layer_norm.h */
-TVM_REGISTER_GLOBAL("topi.nn.layer_norm").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.nn.layer_norm").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = nn::layer_norm(args[0], args[1], args[2], args[3], static_cast<double>(args[4]));
});
/* Ops from nn/group_norm.h */
-TVM_REGISTER_GLOBAL("topi.nn.group_norm").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.nn.group_norm").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = nn::group_norm(args[0], args[1], args[2], static_cast<int>(args[3]),
static_cast<int>(args[4]), args[5], static_cast<double>(args[6]));
});
/* Ops from nn/instance_norm.h */
-TVM_REGISTER_GLOBAL("topi.nn.instance_norm").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.nn.instance_norm").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = nn::instance_norm(args[0], args[1], args[2], args[3], static_cast<double>(args[4]));
});
/* Ops from nn/rms_norm.h */
-TVM_REGISTER_GLOBAL("topi.nn.rms_norm").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.nn.rms_norm").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = nn::rms_norm(args[0], args[1], args[2], static_cast<double>(args[3]));
});
diff --git a/src/topi/reduction.cc b/src/topi/reduction.cc
index a9d692cc07..2400534b4f 100644
--- a/src/topi/reduction.cc
+++ b/src/topi/reduction.cc
@@ -32,39 +32,39 @@ namespace topi {
using namespace tvm;
using namespace tvm::runtime;
-TVM_REGISTER_GLOBAL("topi.sum").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.sum").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = topi::sum(args[0], ArrayOrInt(args[1]), args[2]);
});
-TVM_REGISTER_GLOBAL("topi.min").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.min").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = topi::min(args[0], ArrayOrInt(args[1]), args[2]);
});
-TVM_REGISTER_GLOBAL("topi.max").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.max").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = topi::max(args[0], ArrayOrInt(args[1]), args[2]);
});
-TVM_REGISTER_GLOBAL("topi.argmin").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.argmin").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = topi::argmin(args[0], ArrayOrInt(args[1]), args[2], false, args[3]);
});
-TVM_REGISTER_GLOBAL("topi.argmax").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.argmax").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = topi::argmax(args[0], ArrayOrInt(args[1]), args[2], false, args[3]);
});
-TVM_REGISTER_GLOBAL("topi.prod").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.prod").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = topi::prod(args[0], ArrayOrInt(args[1]), args[2]);
});
-TVM_REGISTER_GLOBAL("topi.all").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.all").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = topi::all(args[0], ArrayOrInt(args[1]), args[2]);
});
-TVM_REGISTER_GLOBAL("topi.any").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.any").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = topi::any(args[0], ArrayOrInt(args[1]), args[2]);
});
-TVM_REGISTER_GLOBAL("topi.collapse_sum").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.collapse_sum").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = topi::collapse_sum(args[0], args[1]);
});
diff --git a/src/topi/transform.cc b/src/topi/transform.cc
index 44c46e8725..f56b2978c8 100644
--- a/src/topi/transform.cc
+++ b/src/topi/transform.cc
@@ -37,52 +37,52 @@ namespace topi {
using namespace tvm;
using namespace tvm::runtime;
-TVM_REGISTER_GLOBAL("topi.expand_dims").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.expand_dims").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = expand_dims(args[0], args[1], args[2]);
});
-TVM_REGISTER_GLOBAL("topi.transpose").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.transpose").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = transpose(args[0], args[1]);
});
-TVM_REGISTER_GLOBAL("topi.flip").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.flip").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
// pass empty seq_lengths tensor to reverse_sequence
*rv = reverse_sequence(args[0], Tensor(), args[1]);
});
-TVM_REGISTER_GLOBAL("topi.reverse_sequence").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.reverse_sequence").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = reverse_sequence(args[0], args[1], args[2], args[3]);
});
-TVM_REGISTER_GLOBAL("topi.reshape").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.reshape").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = reshape(args[0], args[1]);
});
-TVM_REGISTER_GLOBAL("topi.sliding_window").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.sliding_window").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = sliding_window(args[0], args[1], args[2], args[3]);
});
-TVM_REGISTER_GLOBAL("topi.squeeze").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.squeeze").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = squeeze(args[0], ArrayOrInt(args[1]));
});
-TVM_REGISTER_GLOBAL("topi.concatenate").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.concatenate").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = concatenate(args[0], args[1]);
});
-TVM_REGISTER_GLOBAL("topi.stack").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.stack").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = stack(args[0], args[1]);
});
-TVM_REGISTER_GLOBAL("topi.shape").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.shape").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = shape(args[0], args[1]);
});
-TVM_REGISTER_GLOBAL("topi.ndarray_size").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.ndarray_size").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = ndarray_size(args[0], args[1]);
});
-TVM_REGISTER_GLOBAL("topi.split").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.split").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
if (args[1].as<int>()) {
*rv = split_n_sections(args[0], args[1], args[2]);
} else {
@@ -90,11 +90,11 @@ TVM_REGISTER_GLOBAL("topi.split").set_body([](TVMArgs args, TVMRetValue* rv) {
}
});
-TVM_REGISTER_GLOBAL("topi.layout_transform").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.layout_transform").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = layout_transform(args[0], args[1], args[2], args[3]);
});
-TVM_REGISTER_GLOBAL("topi.take").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.take").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
if (args.size() == 4) {
std::string mode = args[3];
int batch_dims = args[2];
@@ -108,50 +108,50 @@ TVM_REGISTER_GLOBAL("topi.take").set_body([](TVMArgs args, TVMRetValue* rv) {
}
});
-TVM_REGISTER_GLOBAL("topi.sequence_mask").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.sequence_mask").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
double pad_val = args[2];
int axis = args[3];
*rv = sequence_mask(args[0], args[1], pad_val, axis);
});
-TVM_REGISTER_GLOBAL("topi.where").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.where").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = where(args[0], args[1], args[2]);
});
-TVM_REGISTER_GLOBAL("topi.arange").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.arange").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = arange(args[0], args[1], args[2], args[3]);
});
-TVM_REGISTER_GLOBAL("topi.meshgrid").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.meshgrid").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = meshgrid(args[0], args[1]);
});
-TVM_REGISTER_GLOBAL("topi.repeat").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.repeat").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = repeat(args[0], args[1], args[2]);
});
-TVM_REGISTER_GLOBAL("topi.tile").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.tile").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = tile(args[0], args[1]);
});
-TVM_REGISTER_GLOBAL("topi.gather").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.gather").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = gather(args[0], args[1], args[2]);
});
-TVM_REGISTER_GLOBAL("topi.gather_nd").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.gather_nd").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
int batch_dims = args[2];
*rv = gather_nd(args[0], args[1], batch_dims);
});
-TVM_REGISTER_GLOBAL("topi.unravel_index").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.unravel_index").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = unravel_index(args[0], args[1]);
});
-TVM_REGISTER_GLOBAL("topi.sparse_to_dense").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.sparse_to_dense").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = sparse_to_dense(args[0], args[1], args[2], args[3]);
});
-TVM_REGISTER_GLOBAL("topi.matmul").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.matmul").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
switch (args.size()) {
case 2:
*rv = matmul(args[0], args[1]);
@@ -167,7 +167,7 @@ TVM_REGISTER_GLOBAL("topi.matmul").set_body([](TVMArgs args, TVMRetValue* rv) {
}
});
-TVM_REGISTER_GLOBAL("topi.tensordot").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.tensordot").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
if (args.size() == 2) {
*rv = tensordot(args[0], args[1]);
} else if (args.size() == 3) {
@@ -178,7 +178,7 @@ TVM_REGISTER_GLOBAL("topi.tensordot").set_body([](TVMArgs args, TVMRetValue* rv)
}
});
-TVM_REGISTER_GLOBAL("topi.strided_slice").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.strided_slice").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
Tensor x = args[0];
Array<PrimExpr> begin = args[1];
Array<PrimExpr> end = args[2];
@@ -205,29 +205,31 @@ TVM_REGISTER_GLOBAL("topi.strided_slice").set_body([](TVMArgs args, TVMRetValue*
}
});
-TVM_REGISTER_GLOBAL("topi.dynamic_strided_slice").set_body([](TVMArgs args, TVMRetValue* rv) {
- te::Tensor begin = args[1];
- te::Tensor end = args[2];
- te::Tensor strides = args[3];
- *rv = dynamic_strided_slice(args[0], begin, end, strides);
-});
-
-TVM_REGISTER_GLOBAL("topi.relax_dynamic_strided_slice").set_body([](TVMArgs args, TVMRetValue* rv) {
- te::Tensor begin = args[1];
- te::Tensor end = args[2];
- te::Tensor strides = args[3];
- Array<PrimExpr> output_shape = args[4];
- *rv = relax::dynamic_strided_slice(args[0], begin, end, strides, output_shape);
-});
-
-TVM_REGISTER_GLOBAL("topi.one_hot").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.dynamic_strided_slice")
+ .set_body_packed([](TVMArgs args, TVMRetValue* rv) {
+ te::Tensor begin = args[1];
+ te::Tensor end = args[2];
+ te::Tensor strides = args[3];
+ *rv = dynamic_strided_slice(args[0], begin, end, strides);
+ });
+
+TVM_REGISTER_GLOBAL("topi.relax_dynamic_strided_slice")
+ .set_body_packed([](TVMArgs args, TVMRetValue* rv) {
+ te::Tensor begin = args[1];
+ te::Tensor end = args[2];
+ te::Tensor strides = args[3];
+ Array<PrimExpr> output_shape = args[4];
+ *rv = relax::dynamic_strided_slice(args[0], begin, end, strides, output_shape);
+ });
+
+TVM_REGISTER_GLOBAL("topi.one_hot").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
int depth = args[3];
int axis = args[4];
DataType dtype = args[5];
*rv = one_hot(args[0], args[1], args[2], depth, axis, dtype);
});
-TVM_REGISTER_GLOBAL("topi.matrix_set_diag").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.matrix_set_diag").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
int k1 = args[2];
int k2 = args[3];
bool super_diag_right_align = args[4];
@@ -235,7 +237,7 @@ TVM_REGISTER_GLOBAL("topi.matrix_set_diag").set_body([](TVMArgs args, TVMRetValu
*rv = matrix_set_diag(args[0], args[1], k1, k2, super_diag_right_align, sub_diag_right_align);
});
-TVM_REGISTER_GLOBAL("topi.adv_index").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.adv_index").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = adv_index(args[0], args[1]);
});
diff --git a/src/topi/utils.cc b/src/topi/utils.cc
index 94a5bf2dac..c3ad01b1a5 100644
--- a/src/topi/utils.cc
+++ b/src/topi/utils.cc
@@ -28,17 +28,19 @@
namespace tvm {
namespace topi {
-TVM_REGISTER_GLOBAL("topi.utils.is_empty_shape").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.utils.is_empty_shape").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = topi::detail::is_empty_shape(args[0]);
});
-TVM_REGISTER_GLOBAL("topi.utils.bilinear_sample_nchw").set_body([](TVMArgs args, TVMRetValue* rv) {
- *rv = detail::bilinear_sample_nchw(args[0], args[1], args[2], args[3]);
-});
+TVM_REGISTER_GLOBAL("topi.utils.bilinear_sample_nchw")
+ .set_body_packed([](TVMArgs args, TVMRetValue* rv) {
+ *rv = detail::bilinear_sample_nchw(args[0], args[1], args[2], args[3]);
+ });
-TVM_REGISTER_GLOBAL("topi.utils.bilinear_sample_nhwc").set_body([](TVMArgs args, TVMRetValue* rv) {
- *rv = detail::bilinear_sample_nhwc(args[0], args[1], args[2], args[3]);
-});
+TVM_REGISTER_GLOBAL("topi.utils.bilinear_sample_nhwc")
+ .set_body_packed([](TVMArgs args, TVMRetValue* rv) {
+ *rv = detail::bilinear_sample_nhwc(args[0], args[1], args[2], args[3]);
+ });
} // namespace topi
} // namespace tvm
diff --git a/src/topi/vision.cc b/src/topi/vision.cc
index e64eec452c..68eab43fc2 100644
--- a/src/topi/vision.cc
+++ b/src/topi/vision.cc
@@ -31,7 +31,7 @@ namespace topi {
using namespace tvm;
using namespace tvm::runtime;
-TVM_REGISTER_GLOBAL("topi.vision.reorg").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("topi.vision.reorg").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
*rv = vision::reorg(args[0], args[1]);
});
diff --git a/tests/cpp-runtime/hexagon/run_all_tests.cc b/tests/cpp-runtime/hexagon/run_all_tests.cc
index 720249d9d9..5187277dcb 100644
--- a/tests/cpp-runtime/hexagon/run_all_tests.cc
+++ b/tests/cpp-runtime/hexagon/run_all_tests.cc
@@ -38,7 +38,7 @@ namespace tvm {
namespace runtime {
namespace hexagon {
-TVM_REGISTER_GLOBAL("hexagon.run_all_tests").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("hexagon.run_all_tests").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
// gtest args are passed into this packed func as a singular string
// split gtest args using <space> delimiter and build argument vector
std::vector<std::string> parsed_args = tvm::support::Split(args[0], ' ');
diff --git a/tests/cpp-runtime/hexagon/run_unit_tests.cc b/tests/cpp-runtime/hexagon/run_unit_tests.cc
index 6ad770da33..37a521457a 100644
--- a/tests/cpp-runtime/hexagon/run_unit_tests.cc
+++ b/tests/cpp-runtime/hexagon/run_unit_tests.cc
@@ -80,7 +80,7 @@ class GtestPrinter : public testing::EmptyTestEventListener {
std::string GetOutput() { return gtest_out_.str(); }
};
-TVM_REGISTER_GLOBAL("hexagon.run_unit_tests").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("hexagon.run_unit_tests").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
// gtest args are passed into this packed func as a singular string
// split gtest args using <space> delimiter and build argument vector
std::vector<std::string> parsed_args = tvm::support::Split(args[0], ' ');
diff --git a/tests/cpp/llvm_codegen_registry_test.cc b/tests/cpp/llvm_codegen_registry_test.cc
index 967f506000..5e2a167214 100644
--- a/tests/cpp/llvm_codegen_registry_test.cc
+++ b/tests/cpp/llvm_codegen_registry_test.cc
@@ -35,27 +35,27 @@
TEST(LLVMCodeGen, CodeGenFactoryPresent) {
std::initializer_list<std::string> cpu_targets = {CPU_TARGETS};
for (const std::string& s : cpu_targets) {
- auto* pf = tvm::runtime::Registry::Get("tvm.codegen.llvm.target_" + s);
- EXPECT_NE(pf, nullptr);
+ auto pf = tvm::ffi::Function::GetGlobal("tvm.codegen.llvm.target_" + s);
+ EXPECT_TRUE(pf.has_value());
}
std::initializer_list<std::string> optional_targets = {OPTIONAL_TARGETS};
for (const std::string& s : optional_targets) {
- if (tvm::runtime::Registry::Get("device_api." + s)) {
- auto* pf = tvm::runtime::Registry::Get("tvm.codegen.llvm.target_" + s);
- EXPECT_NE(pf, nullptr);
+ if (tvm::ffi::Function::GetGlobal("device_api." + s)) {
+ auto pf = tvm::ffi::Function::GetGlobal("tvm.codegen.llvm.target_" + s);
+ EXPECT_TRUE(pf.has_value());
}
}
- auto* pf_bad = tvm::runtime::Registry::Get("tvm.codegen.llvm.target_invalid-target");
- EXPECT_EQ(pf_bad, nullptr);
+ auto pf_bad = tvm::ffi::Function::GetGlobal("tvm.codegen.llvm.target_invalid-target");
+ EXPECT_FALSE(pf_bad.has_value());
}
TEST(LLVMCodeGen, CodeGenFactoryWorks) {
std::initializer_list<std::string> all_targets = {ALL_TARGETS};
for (const std::string& s : all_targets) {
- if (auto* pf = tvm::runtime::Registry::Get("tvm.codegen.llvm.target_" + s)) {
- auto* cg = static_cast<void*>((*pf)());
+ if (auto pf = tvm::ffi::Function::GetGlobal("tvm.codegen.llvm.target_" + s)) {
+ auto cg = static_cast<void*>((*pf)());
EXPECT_NE(cg, nullptr);
delete static_cast<tvm::codegen::CodeGenLLVM*>(cg);
}
diff --git a/tests/cpp/target_test.cc b/tests/cpp/target_test.cc
index a75ba09e82..baa5e24bb5 100644
--- a/tests/cpp/target_test.cc
+++ b/tests/cpp/target_test.cc
@@ -456,8 +456,8 @@ TEST(TargetCreation, DetectSystemTriple) {
Target target = Target(config);
ICHECK_EQ(target->kind, TargetKind::Get("llvm").value());
- auto pf = tvm::runtime::Registry::Get("target.llvm_get_system_triple");
- if (pf == nullptr) {
+ auto pf = tvm::ffi::Function::GetGlobal("target.llvm_get_system_triple");
+ if (!pf.has_value()) {
GTEST_SKIP() << "LLVM is not available, skipping test";
}
diff --git a/tests/python/contrib/test_hexagon/README_RPC.md b/tests/python/contrib/test_hexagon/README_RPC.md
index 228922d3f1..5beffd800b 100644
--- a/tests/python/contrib/test_hexagon/README_RPC.md
+++ b/tests/python/contrib/test_hexagon/README_RPC.md
@@ -80,7 +80,7 @@ Which eventually jumps to the following line in C++, which creates a RPC client
[https://github.com/apache/tvm/blob/2cca934aad1635e3a83b712958ea83ff65704316/src/runtime/rpc/rpc_socket_impl.cc#L123-L129](https://github.com/apache/tvm/blob/2cca934aad1635e3a83b712958ea83ff65704316/src/runtime/rpc/rpc_socket_impl.cc#L123-L129)
```cpp
-TVM_REGISTER_GLOBAL("rpc.Connect").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("rpc.Connect").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
std::string url = args[0];
int port = args[1];
std::string key = args[2];
@@ -95,7 +95,7 @@ TVM_REGISTER_GLOBAL("rpc.Connect").set_body([](TVMArgs args, TVMRetValue* rv) {
```cpp
TVM_REGISTER_GLOBAL("tvm.contrib.hexagon.create_hexagon_session")
- .set_body([](TVMArgs args, TVMRetValue* rv) {
+ .set_body_packed([](TVMArgs args, TVMRetValue* rv) {
std::string session_name = args[0];
int remote_stack_size_bytes = args[1];
HexagonTransportChannel* hexagon_channel =
diff --git a/tests/python/relax/test_vm_codegen_tir.py b/tests/python/relax/test_vm_codegen_tir.py
index 41f8e81735..c60b307d68 100644
--- a/tests/python/relax/test_vm_codegen_tir.py
+++ b/tests/python/relax/test_vm_codegen_tir.py
@@ -88,9 +88,7 @@ def test_tir_call():
@T.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.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))
)
diff --git a/web/emcc/tvmjs_support.cc b/web/emcc/tvmjs_support.cc
index 970a1d9d31..620095391c 100644
--- a/web/emcc/tvmjs_support.cc
+++ b/web/emcc/tvmjs_support.cc
@@ -106,10 +106,10 @@ class AsyncLocalSession : public LocalSession {
PackedFuncHandle GetFunction(const std::string& name) final {
if (name == "runtime.RPCTimeEvaluator") {
return get_time_eval_placeholder_.get();
- } else if (auto* fp = tvm::runtime::Registry::Get(name)) {
+ } else if (auto fp = tvm::ffi::Function::GetGlobal(name)) {
TVMFFIAny val = tvm::ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(tvm::ffi::Any(*fp));
return val.v_obj;
- } else if (auto* fp = tvm::runtime::Registry::Get("__async." + name)) {
+ } else if (auto fp = tvm::ffi::Function::GetGlobal("__async." + name)) {
TVMFFIAny val = tvm::ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(tvm::ffi::Any(*fp));
async_func_set_.insert(val.v_obj);
return val.v_obj;
@@ -219,10 +219,10 @@ class AsyncLocalSession : public LocalSession {
on_complete(RPCCode::kReturn, ffi::PackedArgs(packed_args, 1));
} else {
CHECK(dev.device_type == static_cast<DLDeviceType>(kDLWebGPU));
- if (async_wait_ == nullptr) {
- async_wait_ = tvm::runtime::Registry::Get("__async.wasm.WebGPUWaitForTasks");
+ if (!async_wait_.has_value()) {
+ async_wait_ = tvm::ffi::Function::GetGlobal("__async.wasm.WebGPUWaitForTasks");
}
- CHECK(async_wait_ != nullptr);
+ CHECK(async_wait_.has_value());
PackedFunc packed_callback([on_complete](TVMArgs args, TVMRetValue*) {
int code = args[0];
on_complete(static_cast<RPCCode>(code), args.Slice(1));
@@ -236,7 +236,7 @@ class AsyncLocalSession : public LocalSession {
private:
std::unordered_set<void*> async_func_set_;
std::unique_ptr<PackedFunc> get_time_eval_placeholder_ = std::make_unique<PackedFunc>();
- const PackedFunc* async_wait_{nullptr};
+ std::optional<PackedFunc> async_wait_;
// time evaluator
PackedFunc GetTimeEvaluator(Optional<Module> opt_mod, std::string name, int device_type,
@@ -254,8 +254,8 @@ class AsyncLocalSession : public LocalSession {
limit_zero_time_iterations, cooldown_interval_ms,
repeats_to_cooldown);
} else {
- auto* pf = runtime::Registry::Get(name);
- CHECK(pf != nullptr) << "Cannot find " << name << " in the global function";
+ auto pf = tvm::ffi::Function::GetGlobal(name);
+ CHECK(pf.has_value()) << "Cannot find " << name << " in the global function";
return WrapWasmTimeEvaluator(*pf, dev, number, repeat, min_repeat_ms,
limit_zero_time_iterations, cooldown_interval_ms,
repeats_to_cooldown);
@@ -279,8 +279,8 @@ class AsyncLocalSession : public LocalSession {
pf.CallPacked(invoke_args, &temp);
}
};
- auto* time_exec = runtime::Registry::Get("__async.wasm.TimeExecution");
- CHECK(time_exec != nullptr) << "Cannot find wasm.GetTimer in the global function";
+ auto time_exec = tvm::ffi::Function::GetGlobal("__async.wasm.TimeExecution");
+ CHECK(time_exec.has_value()) << "Cannot find wasm.GetTimer in the global function";
(*time_exec)(TypedPackedFunc<void(int)>(finvoke), dev, number, repeat, min_repeat_ms,
limit_zero_time_iterations, cooldown_interval_ms, repeats_to_cooldown,
/*cache_flush_bytes=*/0, on_complete);
diff --git a/web/emcc/wasm_runtime.cc b/web/emcc/wasm_runtime.cc
index 6e2dfb366a..1939a48be6 100644
--- a/web/emcc/wasm_runtime.cc
+++ b/web/emcc/wasm_runtime.cc
@@ -102,35 +102,35 @@ void LogMessageImpl(const std::string& file, int lineno, int level, const std::s
} // namespace detail
-TVM_REGISTER_GLOBAL("testing.echo").set_body([](TVMArgs args, TVMRetValue* ret) {
+TVM_REGISTER_GLOBAL("testing.echo").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
*ret = args[0];
});
-TVM_REGISTER_GLOBAL("testing.call").set_body([](TVMArgs args, TVMRetValue* ret) {
+TVM_REGISTER_GLOBAL("testing.call").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
(args[0].operator PackedFunc()).CallPacked(args.Slice(1), ret);
});
-TVM_REGISTER_GLOBAL("testing.ret_string").set_body([](TVMArgs args, TVMRetValue* ret) {
+TVM_REGISTER_GLOBAL("testing.ret_string").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
*ret = args[0].operator String();
});
-TVM_REGISTER_GLOBAL("testing.log_info_str").set_body([](TVMArgs args, TVMRetValue* ret) {
+TVM_REGISTER_GLOBAL("testing.log_info_str").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
LOG(INFO) << args[0].operator String();
});
-TVM_REGISTER_GLOBAL("testing.log_fatal_str").set_body([](TVMArgs args, TVMRetValue* ret) {
+TVM_REGISTER_GLOBAL("testing.log_fatal_str").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
LOG(FATAL) << args[0].operator String();
});
TVM_REGISTER_GLOBAL("testing.add_one").set_body_typed([](int x) { return x + 1; });
-TVM_REGISTER_GLOBAL("testing.wrap_callback").set_body([](TVMArgs args, TVMRetValue* ret) {
+TVM_REGISTER_GLOBAL("testing.wrap_callback").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
PackedFunc pf = args[0];
*ret = runtime::TypedPackedFunc<void()>([pf]() { pf(); });
});
// internal function used for debug and testing purposes
-TVM_REGISTER_GLOBAL("testing.object_use_count").set_body([](TVMArgs args, TVMRetValue* ret) {
+TVM_REGISTER_GLOBAL("testing.object_use_count").set_body_packed([](TVMArgs args, TVMRetValue* ret) {
runtime::ObjectRef obj = args[0];
// subtract the current one because we always copy
// and get another value.
@@ -161,19 +161,20 @@ void ArrayDecodeStorage(NDArray cpu_arr, std::string bytes, std::string format,
TVM_REGISTER_GLOBAL("tvmjs.array.decode_storage").set_body_typed(ArrayDecodeStorage);
// Concatenate n TVMArrays
-TVM_REGISTER_GLOBAL("tvmjs.runtime.ArrayConcat").set_body([](TVMArgs args, TVMRetValue* ret) {
- std::vector<Any> data;
- for (int i = 0; i < args.size(); ++i) {
- // Get i-th TVMArray
- auto* arr_i = args[i].as<ArrayObj>();
- ICHECK(arr_i != nullptr);
- for (size_t j = 0; j < arr_i->size(); ++j) {
- // Push back each j-th element of the i-th array
- data.push_back(arr_i->at(j));
- }
- }
- *ret = Array<Any>(data);
-});
+TVM_REGISTER_GLOBAL("tvmjs.runtime.ArrayConcat")
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
+ std::vector<Any> data;
+ for (int i = 0; i < args.size(); ++i) {
+ // Get i-th TVMArray
+ auto* arr_i = args[i].as<ArrayObj>();
+ ICHECK(arr_i != nullptr);
+ for (size_t j = 0; j < arr_i->size(); ++j) {
+ // Push back each j-th element of the i-th array
+ data.push_back(arr_i->at(j));
+ }
+ }
+ *ret = Array<Any>(data);
+ });
NDArray ConcatEmbeddings(const std::vector<NDArray>& embeddings) {
// Get output shape
@@ -210,14 +211,15 @@ NDArray ConcatEmbeddings(const std::vector<NDArray>& embeddings) {
}
// Concatenate n NDArrays
-TVM_REGISTER_GLOBAL("tvmjs.runtime.ConcatEmbeddings").set_body([](TVMArgs args, TVMRetValue* ret) {
- std::vector<NDArray> embeddings;
- for (int i = 0; i < args.size(); ++i) {
- embeddings.push_back(args[i].operator NDArray());
- }
- NDArray result = ConcatEmbeddings(std::move(embeddings));
- *ret = result;
-});
+TVM_REGISTER_GLOBAL("tvmjs.runtime.ConcatEmbeddings")
+ .set_body_packed([](TVMArgs args, TVMRetValue* ret) {
+ std::vector<NDArray> embeddings;
+ for (int i = 0; i < args.size(); ++i) {
+ embeddings.push_back(args[i].operator NDArray());
+ }
+ NDArray result = ConcatEmbeddings(std::move(embeddings));
+ *ret = result;
+ });
} // namespace runtime
} // namespace tvm
diff --git a/web/emcc/webgpu_runtime.cc b/web/emcc/webgpu_runtime.cc
index 1d7dbe0787..0b36f392c5 100644
--- a/web/emcc/webgpu_runtime.cc
+++ b/web/emcc/webgpu_runtime.cc
@@ -58,8 +58,8 @@ class WebGPUThreadEntry {
class WebGPUDeviceAPI : public DeviceAPI {
public:
WebGPUDeviceAPI() {
- auto* fp = tvm::runtime::Registry::Get("wasm.WebGPUDeviceAPI");
- CHECK(fp != nullptr) << "Cannot find wasm.WebGPUContext in the env";
+ auto fp = tvm::ffi::Function::GetGlobal("wasm.WebGPUDeviceAPI");
+ CHECK(fp.has_value()) << "Cannot find wasm.WebGPUContext in the env";
auto getter = TypedPackedFunc<PackedFunc(std::string)>(*fp);
alloc_space_ = getter("deviceAllocDataSpace");
free_space_ = getter("deviceFreeDataSpace");
@@ -113,8 +113,8 @@ class WebGPUDeviceAPI : public DeviceAPI {
}
void StreamSync(Device dev, TVMStreamHandle stream) final {
- static const PackedFunc* func = runtime::Registry::Get("__asyncify.WebGPUWaitForTasks");
- ICHECK(func != nullptr) << "Stream sync inside c++ only supported in asyncify mode";
+ static auto func = tvm::ffi::Function::GetGlobal("__asyncify.WebGPUWaitForTasks");
+ ICHECK(func.has_value()) << "Stream sync inside c++ only supported in asyncify mode";
(*func)();
}
@@ -158,8 +158,8 @@ class WebGPUModuleNode final : public runtime::ModuleNode {
explicit WebGPUModuleNode(std::unordered_map<std::string, std::string> smap,
std::unordered_map<std::string, FunctionInfo> fmap)
: smap_(smap), fmap_(fmap) {
- auto* fp = tvm::runtime::Registry::Get("wasm.WebGPUCreateShader");
- CHECK(fp != nullptr);
+ auto fp = tvm::ffi::Function::GetGlobal("wasm.WebGPUCreateShader");
+ CHECK(fp.has_value());
create_shader_ = *fp;
}
@@ -242,7 +242,7 @@ Module WebGPUModuleLoadBinary(void* strm) {
// for now webgpu is hosted via a vulkan module.
TVM_REGISTER_GLOBAL("runtime.module.loadbinary_webgpu").set_body_typed(WebGPUModuleLoadBinary);
-TVM_REGISTER_GLOBAL("device_api.webgpu").set_body([](TVMArgs args, TVMRetValue* rv) {
+TVM_REGISTER_GLOBAL("device_api.webgpu").set_body_packed([](TVMArgs args, TVMRetValue* rv) {
DeviceAPI* ptr = WebGPUDeviceAPI::Global();
*rv = static_cast<void*>(ptr);
});