You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by ru...@apache.org on 2025/05/24 15:58:25 UTC
(tvm) branch main updated: [REFACTOR][FFI] Phase out legacy C API (#18010)
This is an automated email from the ASF dual-hosted git repository.
ruihangl pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 2d964b4133 [REFACTOR][FFI] Phase out legacy C API (#18010)
2d964b4133 is described below
commit 2d964b4133aac2f92e4185b3f095df4eb3bf3a90
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Sat May 24 11:58:20 2025 -0400
[REFACTOR][FFI] Phase out legacy C API (#18010)
This PR phases out the legacy c api in favor of the new FFI C API.
Also removes the redirection sccafolding for registry.h
- include <tvm/runtime/registry.h> => include <tvm/ffi/function.h>
- include <tvm/runtime/c_runtime_api.h> => include <tvm/runtime/base.h>
- TVM_REGISTER_GLOBAL => TVM_FFI_REGISTER_GLOBAL
The cleanup will greatly simplify the overall FFI surface of the project
and allows us to move towards an unified clean API based on tvm ffi.
---
3rdparty/cutlass_fpA_intB_gemm | 2 +-
apps/android_rpc/app/src/main/jni/tvm_runtime.h | 4 +-
apps/cpp_rpc/rpc_env.cc | 4 +-
apps/cpp_rpc/rpc_env.h | 2 +-
apps/cpp_rpc/rpc_server.cc | 4 +-
apps/cpp_rpc/rpc_server.h | 2 +-
apps/hexagon_launcher/launcher_core.cc | 2 +-
apps/ios_rpc/tvmrpc/RPCServer.mm | 2 +-
apps/ios_rpc/tvmrpc/TVMRuntime.mm | 8 +-
docs/arch/device_target_interactions.rst | 10 +-
docs/arch/pass_infra.rst | 2 +-
docs/arch/runtime.rst | 8 +-
ffi/include/tvm/ffi/c_api.h | 7 +
ffi/include/tvm/ffi/function.h | 2 +-
include/tvm/ir/op.h | 2 +-
include/tvm/ir/source_map.h | 2 +-
include/tvm/node/node.h | 2 +-
include/tvm/node/reflection.h | 2 +-
include/tvm/node/serialization.h | 2 +-
include/tvm/relax/exec_builder.h | 2 +-
include/tvm/relax/type.h | 2 +-
.../tvm/runtime/base.h | 58 +-
include/tvm/runtime/builtin_fp16.h | 2 +-
include/tvm/runtime/c_backend_api.h | 34 +-
include/tvm/runtime/c_runtime_api.h | 732 -------------------
include/tvm/runtime/data_type.h | 5 +-
include/tvm/runtime/device_api.h | 42 +-
include/tvm/runtime/disco/cuda_ipc_memory.h | 2 +-
include/tvm/runtime/logging.h | 2 +-
include/tvm/runtime/memory/memory_manager.h | 2 +-
include/tvm/runtime/module.h | 2 +-
include/tvm/runtime/ndarray.h | 48 +-
include/tvm/runtime/nvtx.h | 2 +-
include/tvm/runtime/object.h | 2 +-
include/tvm/runtime/packed_func.h | 282 +------
include/tvm/runtime/profiling.h | 6 +-
include/tvm/runtime/registry.h | 102 ---
include/tvm/runtime/relax_vm/executable.h | 2 +-
.../tvm/runtime/relax_vm/ndarray_cache_support.h | 2 +-
include/tvm/runtime/serializer.h | 2 +-
include/tvm/support/parallel_for.h | 2 +-
include/tvm/tir/builtin.h | 20 +-
include/tvm/tir/expr.h | 2 +-
include/tvm/tir/transform.h | 4 +-
python/tvm/contrib/cutlass/gen_tensor_op.py | 2 +-
python/tvm/relax/frontend/nn/op.py | 2 +-
python/tvm/runtime/_ffi_api.py | 4 +-
python/tvm/runtime/_ffi_node_api.py | 6 +-
src/arith/analyzer.cc | 4 +-
src/arith/bound_deducer.cc | 4 +-
src/arith/const_int_bound.cc | 4 +-
src/arith/detect_common_subexpr.cc | 2 +-
src/arith/detect_linear_equation.cc | 6 +-
src/arith/domain_touched.cc | 6 +-
src/arith/int_constraints.cc | 13 +-
src/arith/int_set.cc | 30 +-
src/arith/iter_affine_map.cc | 18 +-
src/arith/modular_set.cc | 4 +-
src/arith/narrow_predicate_expression.cc | 5 +-
src/arith/presburger_set.cc | 4 +-
src/arith/solve_linear_equation.cc | 4 +-
src/arith/solve_linear_inequality.cc | 8 +-
src/contrib/msc/core/ir/graph.cc | 101 +--
src/contrib/msc/core/ir/graph_builder.cc | 4 +-
src/contrib/msc/core/ir/plugin.cc | 8 +-
.../msc/core/transform/bind_named_params.cc | 2 +-
src/contrib/msc/core/transform/bind_shape.cc | 2 +-
src/contrib/msc/core/transform/fuse_tuple.cc | 2 +-
src/contrib/msc/core/transform/inline_params.cc | 2 +-
src/contrib/msc/core/transform/set_byoc_attrs.cc | 2 +-
src/contrib/msc/core/transform/set_expr_layout.cc | 2 +-
src/contrib/msc/core/transform/set_expr_name.cc | 2 +-
src/contrib/msc/core/utils.cc | 12 +-
src/contrib/msc/framework/tensorflow/codegen.cc | 2 +-
src/contrib/msc/framework/tensorrt/codegen.cc | 6 +-
.../msc/framework/tensorrt/transform_tensorrt.cc | 2 +-
src/contrib/msc/framework/torch/codegen.cc | 2 +-
src/contrib/msc/framework/tvm/codegen.cc | 2 +-
src/contrib/msc/plugin/tensorrt_codegen.cc | 2 +-
src/contrib/msc/plugin/torch_codegen.cc | 2 +-
src/contrib/msc/plugin/tvm_codegen.cc | 8 +-
src/ir/analysis.cc | 2 +-
src/ir/apply_pass_to_function.cc | 4 +-
src/ir/attrs.cc | 6 +-
src/ir/diagnostic.cc | 22 +-
src/ir/env_func.cc | 8 +-
src/ir/expr.cc | 14 +-
src/ir/function.cc | 14 +-
src/ir/global_info.cc | 9 +-
src/ir/global_var_supply.cc | 17 +-
src/ir/instrument.cc | 8 +-
src/ir/module.cc | 46 +-
src/ir/name_supply.cc | 11 +-
src/ir/op.cc | 26 +-
src/ir/replace_global_vars.cc | 4 +-
src/ir/source_map.cc | 23 +-
src/ir/transform.cc | 46 +-
src/ir/type.cc | 15 +-
src/meta_schedule/arg_info.cc | 11 +-
src/meta_schedule/builder/builder.cc | 8 +-
src/meta_schedule/cost_model/cost_model.cc | 11 +-
src/meta_schedule/database/database.cc | 45 +-
src/meta_schedule/database/json_database.cc | 3 +-
src/meta_schedule/database/memory_database.cc | 2 +-
.../database/ordered_union_database.cc | 2 +-
src/meta_schedule/database/schedule_fn_database.cc | 2 +-
src/meta_schedule/database/union_database.cc | 3 +-
src/meta_schedule/extracted_task.cc | 2 +-
.../feature_extractor/feature_extractor.cc | 4 +-
.../feature_extractor/per_store_feature.cc | 2 +-
.../measure_callback/add_to_database.cc | 2 +-
.../measure_callback/measure_callback.cc | 6 +-
.../measure_callback/remove_build_artifact.cc | 2 +-
.../measure_callback/update_cost_model.cc | 2 +-
.../mutator/mutate_compute_location.cc | 2 +-
src/meta_schedule/mutator/mutate_parallel.cc | 3 +-
src/meta_schedule/mutator/mutate_thread_binding.cc | 2 +-
src/meta_schedule/mutator/mutate_tile_size.cc | 3 +-
src/meta_schedule/mutator/mutate_unroll.cc | 2 +-
src/meta_schedule/mutator/mutator.cc | 17 +-
.../postproc/disallow_async_strided_mem_copy.cc | 2 +-
.../postproc/disallow_dynamic_loop.cc | 2 +-
src/meta_schedule/postproc/postproc.cc | 16 +-
.../postproc/rewrite_cooperative_fetch.cc | 2 +-
src/meta_schedule/postproc/rewrite_layout.cc | 3 +-
.../postproc/rewrite_parallel_vectorize_unroll.cc | 2 +-
.../postproc/rewrite_reduction_block.cc | 2 +-
src/meta_schedule/postproc/rewrite_tensorize.cc | 2 +-
.../postproc/rewrite_unbound_block.cc | 2 +-
src/meta_schedule/postproc/verify_gpu_code.cc | 3 +-
src/meta_schedule/postproc/verify_vtcm_limit.cc | 2 +-
src/meta_schedule/profiler.cc | 14 +-
src/meta_schedule/runner/runner.cc | 15 +-
src/meta_schedule/schedule/cpu/winograd.cc | 8 +-
src/meta_schedule/schedule/cuda/winograd.cc | 8 +-
src/meta_schedule/schedule_rule/add_rfactor.cc | 2 +-
.../schedule_rule/apply_custom_rule.cc | 2 +-
src/meta_schedule/schedule_rule/auto_bind.cc | 3 +-
src/meta_schedule/schedule_rule/auto_inline.cc | 4 +-
.../schedule_rule/cross_thread_reduction.cc | 2 +-
.../schedule_rule/multi_level_tiling.cc | 2 +-
.../multi_level_tiling_tensor_core.cc | 2 +-
.../multi_level_tiling_wide_vector.cc | 2 +-
.../multi_level_tiling_with_intrin.cc | 2 +-
.../schedule_rule/parallel_vectorize_unroll.cc | 2 +-
.../schedule_rule/random_compute_location.cc | 2 +-
src/meta_schedule/schedule_rule/schedule_rule.cc | 20 +-
.../search_strategy/evolutionary_search.cc | 6 +-
src/meta_schedule/search_strategy/replay_func.cc | 2 +-
src/meta_schedule/search_strategy/replay_trace.cc | 2 +-
.../search_strategy/search_strategy.cc | 16 +-
.../space_generator/post_order_apply.cc | 2 +-
src/meta_schedule/space_generator/schedule_fn.cc | 2 +-
.../space_generator/space_generator.cc | 8 +-
.../space_generator/space_generator_union.cc | 2 +-
src/meta_schedule/task_scheduler/gradient_based.cc | 2 +-
src/meta_schedule/task_scheduler/round_robin.cc | 2 +-
src/meta_schedule/task_scheduler/task_scheduler.cc | 15 +-
src/meta_schedule/trace_apply.cc | 2 +-
src/meta_schedule/tune_context.cc | 8 +-
src/node/container_printing.cc | 2 +-
src/node/object_path.cc | 24 +-
src/node/reflection.cc | 8 +-
src/node/repr_printer.cc | 6 +-
src/node/script_printer.cc | 6 +-
src/node/serialization.cc | 7 +-
src/node/structural_equal.cc | 10 +-
src/node/structural_hash.cc | 4 +-
src/relax/analysis/analysis.cc | 10 +-
src/relax/analysis/computable_at_compile_time.cc | 2 +-
src/relax/analysis/detect_recursion.cc | 2 +-
src/relax/analysis/layout_transformation.cc | 2 +-
src/relax/analysis/struct_info_analysis.cc | 22 +-
src/relax/analysis/tir_op_pattern_kind.cc | 2 +-
src/relax/analysis/udchain.cc | 2 +-
src/relax/analysis/var2value.cc | 4 +-
src/relax/analysis/well_formed.cc | 2 +-
src/relax/backend/contrib/clml/codegen.cc | 7 +-
src/relax/backend/contrib/cublas/codegen.cc | 2 +-
src/relax/backend/contrib/cudnn/codegen.cc | 2 +-
src/relax/backend/contrib/cutlass/codegen.cc | 4 +-
src/relax/backend/contrib/dnnl/codegen.cc | 2 +-
src/relax/backend/contrib/hipblas/codegen.cc | 2 +-
src/relax/backend/contrib/nnapi/codegen.cc | 4 +-
src/relax/backend/contrib/tensorrt/codegen.cc | 7 +-
src/relax/backend/contrib/utils.cc | 2 +-
src/relax/backend/pattern_registry.cc | 9 +-
src/relax/backend/task_extraction.cc | 2 +-
src/relax/backend/vm/codegen_vm.cc | 4 +-
src/relax/backend/vm/codegen_vm_tir.cc | 2 +-
src/relax/backend/vm/exec_builder.cc | 44 +-
src/relax/backend/vm/lower_runtime_builtin.cc | 2 +-
src/relax/backend/vm/vm_shape_lower.cc | 2 +-
src/relax/distributed/global_info.cc | 2 +-
src/relax/distributed/struct_info.cc | 10 +-
.../distributed/transform/legalize_redistribute.cc | 2 +-
src/relax/distributed/transform/lower_distir.cc | 2 +-
.../transform/lower_global_view_to_local_view.cc | 2 +-
.../distributed/transform/propagate_sharding.cc | 2 +-
src/relax/ir/binding_rewrite.cc | 16 +-
src/relax/ir/block_builder.cc | 40 +-
src/relax/ir/dataflow_block_rewriter.cc | 4 +-
src/relax/ir/dataflow_expr_rewriter.cc | 21 +-
src/relax/ir/dataflow_pattern.cc | 74 +-
src/relax/ir/emit_te.cc | 2 +-
src/relax/ir/expr.cc | 62 +-
src/relax/ir/expr_functor.cc | 2 +-
src/relax/ir/py_expr_functor.cc | 50 +-
src/relax/ir/struct_info.cc | 31 +-
src/relax/ir/transform.cc | 6 +-
src/relax/ir/type.cc | 12 +-
src/relax/op/ccl/ccl.cc | 9 +-
src/relax/op/distributed/distributed.cc | 8 +-
src/relax/op/image/resize.cc | 2 +-
src/relax/op/memory/view.cc | 7 +-
src/relax/op/nn/attention.cc | 4 +-
src/relax/op/nn/convolution.cc | 10 +-
src/relax/op/nn/nn.cc | 28 +-
src/relax/op/nn/pooling.cc | 18 +-
src/relax/op/op.cc | 56 +-
src/relax/op/op_common.h | 2 +-
src/relax/op/tensor/binary.h | 2 +-
src/relax/op/tensor/create.cc | 22 +-
src/relax/op/tensor/datatype.cc | 4 +-
src/relax/op/tensor/grad.cc | 14 +-
src/relax/op/tensor/index.cc | 6 +-
src/relax/op/tensor/linear_algebra.cc | 6 +-
src/relax/op/tensor/manipulate.cc | 46 +-
src/relax/op/tensor/qdq.cc | 4 +-
src/relax/op/tensor/sampling.cc | 3 +-
src/relax/op/tensor/search.cc | 4 +-
src/relax/op/tensor/set.cc | 4 +-
src/relax/op/tensor/sorting.cc | 6 +-
src/relax/op/tensor/statistical.cc | 4 +-
src/relax/op/tensor/statistical.h | 2 +-
src/relax/op/tensor/ternary.cc | 2 +-
src/relax/op/tensor/unary.cc | 2 +-
src/relax/testing/transform.cc | 2 +-
src/relax/training/utils.cc | 2 +-
src/relax/transform/adjust_matmul_order.cc | 2 +-
src/relax/transform/allocate_workspace.cc | 2 +-
src/relax/transform/alter_op_impl.cc | 2 +-
src/relax/transform/annotate_tir_op_pattern.cc | 3 +-
.../transform/attach_attr_layout_free_buffers.cc | 2 +-
src/relax/transform/attach_global_symbol.cc | 2 +-
src/relax/transform/bind_params.cc | 4 +-
src/relax/transform/bind_symbolic_vars.cc | 4 +-
src/relax/transform/bundle_model_params.cc | 2 +-
src/relax/transform/call_tir_rewrite.cc | 2 +-
src/relax/transform/canonicalize_bindings.cc | 3 +-
src/relax/transform/combine_parallel_matmul.cc | 3 +-
src/relax/transform/compute_prim_value.cc | 2 +-
src/relax/transform/convert_dataflow.cc | 2 +-
src/relax/transform/convert_layout.cc | 2 +-
src/relax/transform/dataflow_inplace.cc | 10 +-
src/relax/transform/dead_code_elimination.cc | 2 +-
src/relax/transform/decompose_ops.cc | 4 +-
src/relax/transform/eliminate_common_subexpr.cc | 2 +-
src/relax/transform/expand_matmul_of_sum.cc | 2 +-
src/relax/transform/expand_tuple_arguments.cc | 3 +-
src/relax/transform/few_shot_tuning.cc | 2 +-
src/relax/transform/fold_constant.cc | 2 +-
src/relax/transform/fuse_ops.cc | 6 +-
src/relax/transform/fuse_tir.cc | 2 +-
src/relax/transform/gradient.cc | 2 +-
src/relax/transform/inline_functions.cc | 4 +-
src/relax/transform/kill_after_last_use.cc | 2 +-
src/relax/transform/lambda_lift.cc | 2 +-
src/relax/transform/lazy_transform_params.cc | 4 +-
src/relax/transform/legalize_ops.cc | 2 +-
src/relax/transform/lift_transform_params.cc | 2 +-
src/relax/transform/lower_alloc_tensor.cc | 2 +-
src/relax/transform/merge_composite_functions.cc | 2 +-
src/relax/transform/meta_schedule.cc | 7 +-
src/relax/transform/normalize.cc | 4 +-
src/relax/transform/realize_vdevice.cc | 2 +-
src/relax/transform/remove_purity_checking.cc | 3 +-
src/relax/transform/remove_unused_outputs.cc | 2 +-
src/relax/transform/remove_unused_parameters.cc | 2 +-
.../transform/reorder_permute_dims_after_concat.cc | 2 +-
src/relax/transform/reorder_take_after_matmul.cc | 2 +-
src/relax/transform/rewrite_cuda_graph.cc | 2 +-
src/relax/transform/rewrite_dataflow_reshape.cc | 2 +-
src/relax/transform/run_codegen.cc | 2 +-
src/relax/transform/split_call_tir_by_pattern.cc | 3 +-
.../transform/split_layout_rewrite_preproc.cc | 2 +-
src/relax/transform/static_plan_block_memory.cc | 3 +-
src/relax/transform/to_mixed_precision.cc | 2 +-
src/relax/transform/to_non_dataflow.cc | 2 +-
src/relax/transform/topological_sort.cc | 2 +-
src/relax/transform/tuning_api/database.cc | 26 +-
src/relax/transform/tuning_api/primitives.cc | 39 +-
src/relax/transform/update_param_struct_info.cc | 3 +-
src/relax/transform/update_vdevice.cc | 2 +-
src/relax/utils.cc | 2 +-
src/runtime/builtin_fp16.cc | 8 +-
src/runtime/c_runtime_api.cc | 807 ---------------------
src/runtime/const_loader_module.cc | 6 +-
src/runtime/container.cc | 101 ---
src/runtime/contrib/amx/amx_config.cc | 6 +-
.../contrib/arm_compute_lib/acl_allocator.h | 2 +-
src/runtime/contrib/arm_compute_lib/acl_runtime.cc | 6 +-
src/runtime/contrib/bnns/bnns_json_runtime.cc | 6 +-
src/runtime/contrib/cblas/cblas.cc | 8 +-
src/runtime/contrib/cblas/dnnl_blas.cc | 4 +-
src/runtime/contrib/cblas/gemm_common.h | 2 +-
src/runtime/contrib/cblas/mkl.cc | 10 +-
src/runtime/contrib/clml/clml_runtime.cc | 4 +-
src/runtime/contrib/clml/clml_runtime.h | 2 +-
src/runtime/contrib/coreml/coreml_runtime.mm | 7 +-
src/runtime/contrib/cublas/cublas.cc | 8 +-
src/runtime/contrib/cublas/cublas_json_runtime.cc | 6 +-
src/runtime/contrib/cublas/cublas_utils.cc | 2 +-
src/runtime/contrib/cudnn/conv_backward.cc | 10 +-
src/runtime/contrib/cudnn/conv_forward.cc | 10 +-
.../contrib/cudnn/cudnn_frontend/attention.cc | 2 +-
.../contrib/cudnn/cudnn_frontend/attention.h | 2 +-
src/runtime/contrib/cudnn/cudnn_json_runtime.cc | 6 +-
src/runtime/contrib/cudnn/cudnn_utils.cc | 4 +-
src/runtime/contrib/cudnn/softmax.cc | 6 +-
src/runtime/contrib/curand/curand.cc | 6 +-
src/runtime/contrib/curand/helper_cuda_kernels.h | 2 +-
src/runtime/contrib/cutlass/fp16_group_gemm.cu | 4 +-
.../contrib/cutlass/fp8_blockwise_scaled_gemm.cu | 6 +-
src/runtime/contrib/cutlass/fp8_gemm.cu | 8 +-
src/runtime/contrib/cutlass/fp8_group_gemm.cu | 8 +-
src/runtime/contrib/cutlass/weight_preprocess.cc | 4 +-
src/runtime/contrib/dnnl/dnnl.cc | 2 +-
src/runtime/contrib/dnnl/dnnl_json_runtime.cc | 6 +-
src/runtime/contrib/dnnl/dnnl_kernel.h | 4 +-
src/runtime/contrib/edgetpu/edgetpu_runtime.cc | 4 +-
src/runtime/contrib/hipblas/hipblas.cc | 6 +-
.../contrib/hipblas/hipblas_json_runtime.cc | 7 +-
src/runtime/contrib/hipblas/hipblas_utils.cc | 2 +-
src/runtime/contrib/miopen/conv_forward.cc | 6 +-
src/runtime/contrib/miopen/miopen_utils.cc | 2 +-
src/runtime/contrib/miopen/softmax.cc | 6 +-
src/runtime/contrib/mps/conv.mm | 6 +-
src/runtime/contrib/mps/gemm.mm | 2 +-
src/runtime/contrib/mps/mps_utils.h | 2 +-
src/runtime/contrib/mrvl/mrvl_hw_runtime.cc | 6 +-
src/runtime/contrib/mrvl/mrvl_runtime.cc | 6 +-
src/runtime/contrib/mrvl/mrvl_sw_runtime_lib.cc | 2 +-
src/runtime/contrib/msc/tensorrt_runtime.cc | 7 +-
src/runtime/contrib/mscclpp/allreduce.cu | 2 +-
src/runtime/contrib/nnapi/nnapi_runtime.cc | 6 +-
src/runtime/contrib/nvshmem/init.cc | 8 +-
src/runtime/contrib/nvshmem/kv_transfer.cu | 6 +-
src/runtime/contrib/nvshmem/memory_allocator.cc | 6 +-
src/runtime/contrib/papi/papi.cc | 2 +-
src/runtime/contrib/random/random.cc | 12 +-
src/runtime/contrib/rocblas/rocblas.cc | 6 +-
src/runtime/contrib/sort/sort.cc | 21 +-
src/runtime/contrib/tensorrt/tensorrt_runtime.cc | 6 +-
src/runtime/contrib/tflite/tflite_runtime.cc | 6 +-
src/runtime/contrib/thrust/thrust.cu | 8 +-
src/runtime/contrib/vllm/attention_kernels.cu | 8 +-
src/runtime/contrib/vllm/cache_alloc.cc | 4 +-
src/runtime/contrib/vllm/cache_kernels.cu | 8 +-
src/runtime/cpu_device_api.cc | 4 +-
src/runtime/cuda/cuda_device_api.cc | 21 +-
src/runtime/cuda/cuda_module.cc | 8 +-
src/runtime/cuda/l2_cache_flush.cc | 13 +-
src/runtime/debug_compile.cc | 6 +-
src/runtime/device_api.cc | 271 +++++++
src/runtime/disco/bcast_session.cc | 2 +-
src/runtime/disco/builtin.cc | 48 +-
src/runtime/disco/cuda_ipc/cuda_ipc_memory.cc | 9 +-
src/runtime/disco/cuda_ipc/custom_allreduce.cc | 4 +-
src/runtime/disco/disco_worker.cc | 2 +-
src/runtime/disco/distributed/socket_session.cc | 8 +-
src/runtime/disco/loader.cc | 25 +-
src/runtime/disco/nccl/nccl.cc | 35 +-
src/runtime/disco/nccl/nccl_context.h | 4 +-
src/runtime/disco/process_session.cc | 8 +-
src/runtime/disco/protocol.h | 4 +-
src/runtime/disco/session.cc | 24 +-
src/runtime/disco/threaded_session.cc | 2 +-
src/runtime/dso_library.cc | 4 +-
src/runtime/file_utils.cc | 17 +-
src/runtime/hexagon/hexagon_buffer.h | 2 +-
src/runtime/hexagon/hexagon_common.cc | 6 +-
src/runtime/hexagon/hexagon_device_api.cc | 31 +-
src/runtime/hexagon/hexagon_module.cc | 2 +-
src/runtime/hexagon/hexagon_thread_manager.h | 2 +-
src/runtime/hexagon/hexagon_vtcm_pool.h | 2 +-
src/runtime/hexagon/ops/conv2d.h | 2 +-
src/runtime/hexagon/ops/conv2d_fp16_hvx.cc | 30 +-
src/runtime/hexagon/ops/conv2d_quant_hvx.cc | 49 +-
src/runtime/hexagon/rpc/android/session.cc | 4 +-
src/runtime/hexagon/rpc/hexagon/rpc_server.cc | 8 +-
src/runtime/hexagon/rpc/simulator/rpc_server.cc | 8 +-
src/runtime/hexagon/rpc/simulator/session.cc | 4 +-
src/runtime/library_module.cc | 2 +-
src/runtime/library_module.h | 2 +-
src/runtime/logging.cc | 2 +-
src/runtime/memory/memory_manager.cc | 4 +-
src/runtime/meta_data.h | 2 -
src/runtime/metal/metal_common.h | 2 +-
src/runtime/metal/metal_device_api.mm | 8 +-
src/runtime/metal/metal_module.mm | 6 +-
src/runtime/minrpc/minrpc_server.h | 2 +-
src/runtime/module.cc | 28 +-
src/runtime/ndarray.cc | 101 +--
src/runtime/object.cc | 92 ---
src/runtime/object_internal.h | 96 ---
src/runtime/opencl/opencl_common.h | 2 +-
src/runtime/opencl/opencl_device_api.cc | 19 +-
src/runtime/opencl/opencl_module.cc | 10 +-
src/runtime/opencl/opencl_module_spirv.cc | 2 +-
src/runtime/pack_args.h | 9 +-
src/runtime/packed_func.cc | 2 +-
src/runtime/profiling.cc | 32 +-
src/runtime/regex.cc | 11 +-
src/runtime/registry.cc | 266 -------
src/runtime/relax_vm/builtin.cc | 129 ++--
src/runtime/relax_vm/cuda/cuda_graph_builtin.cc | 6 +-
src/runtime/relax_vm/executable.cc | 8 +-
src/runtime/relax_vm/hexagon/builtin.cc | 6 +-
src/runtime/relax_vm/kv_state.cc | 52 +-
src/runtime/relax_vm/kv_state.h | 2 +-
src/runtime/relax_vm/lm_support.cc | 32 +-
src/runtime/relax_vm/ndarray_cache_support.cc | 28 +-
src/runtime/relax_vm/paged_kv_cache.cc | 4 +-
src/runtime/relax_vm/rnn_state.cc | 2 +-
src/runtime/relax_vm/vm.cc | 7 -
src/runtime/rocm/rocm_device_api.cc | 17 +-
src/runtime/rocm/rocm_module.cc | 10 +-
src/runtime/rpc/rpc_device_api.cc | 4 +-
src/runtime/rpc/rpc_endpoint.cc | 5 +-
src/runtime/rpc/rpc_event_impl.cc | 4 +-
src/runtime/rpc/rpc_local_session.cc | 4 +-
src/runtime/rpc/rpc_module.cc | 32 +-
src/runtime/rpc/rpc_pipe_impl.cc | 17 +-
src/runtime/rpc/rpc_server_env.cc | 8 +-
src/runtime/rpc/rpc_socket_impl.cc | 8 +-
src/runtime/runtime_base.h | 67 --
src/runtime/spirv/spirv_shader.h | 2 +-
src/runtime/static_library.cc | 6 +-
src/runtime/system_library.cc | 17 +-
src/runtime/thread_pool.cc | 8 +-
src/runtime/threading_backend.cc | 5 +-
src/runtime/vulkan/vulkan_common.h | 4 +-
src/runtime/vulkan/vulkan_device_api.cc | 11 +-
src/runtime/vulkan/vulkan_module.cc | 6 +-
src/runtime/vulkan/vulkan_wrapped_func.cc | 2 +-
src/script/ir_builder/base.cc | 26 +-
src/script/ir_builder/ir/frame.cc | 2 +-
src/script/ir_builder/ir/ir.cc | 18 +-
src/script/ir_builder/relax/distributed.cc | 2 +-
src/script/ir_builder/relax/ir.cc | 33 +-
src/script/ir_builder/tir/ir.cc | 180 ++---
src/script/printer/doc.cc | 76 +-
.../printer/doc_printer/python_doc_printer.cc | 4 +-
src/script/printer/ir_docsifier.cc | 2 +-
src/script/printer/relax/type.cc | 2 +-
src/support/ffi_testing.cc | 101 +--
src/support/libinfo.cc | 4 +-
src/support/socket.h | 3 +-
src/target/build_common.h | 2 +-
src/target/codegen.cc | 14 +-
src/target/datatype/myfloat/myfloat.cc | 2 +-
src/target/datatype/posit/posit-wrapper.cc | 2 +-
src/target/datatype/registry.cc | 10 +-
src/target/datatype/registry.h | 4 +-
src/target/intrin_rule.h | 2 +-
src/target/llvm/codegen_aarch64.cc | 4 +-
src/target/llvm/codegen_amdgpu.cc | 8 +-
src/target/llvm/codegen_arm.cc | 4 +-
src/target/llvm/codegen_cpu.cc | 8 +-
src/target/llvm/codegen_hexagon.cc | 4 +-
src/target/llvm/codegen_llvm.cc | 13 +-
src/target/llvm/codegen_nvptx.cc | 4 +-
src/target/llvm/codegen_x86_64.cc | 4 +-
src/target/llvm/intrin_rule_llvm.h | 2 +-
src/target/llvm/intrin_rule_nvptx.cc | 2 +-
src/target/llvm/intrin_rule_rocm.cc | 2 +-
src/target/llvm/llvm_module.cc | 59 +-
src/target/opt/build_cuda_on.cc | 2 +-
src/target/source/codegen_c.cc | 17 +-
src/target/source/codegen_c_host.cc | 16 +-
src/target/source/codegen_cuda.cc | 2 +-
src/target/source/codegen_metal.cc | 2 +-
src/target/source/codegen_opencl.cc | 4 +-
src/target/source/codegen_webgpu.cc | 2 +-
src/target/source/source_module.cc | 8 +-
src/target/spirv/build_vulkan.cc | 2 +-
src/target/spirv/intrin_rule_spirv.cc | 2 +-
src/target/tag.cc | 6 +-
src/target/target.cc | 20 +-
src/target/target_info.cc | 2 +-
src/target/target_kind.cc | 11 +-
src/target/virtual_device.cc | 2 +-
src/te/operation/compute_op.cc | 4 +-
src/te/operation/create_primfunc.cc | 21 +-
src/te/operation/extern_op.cc | 4 +-
src/te/operation/graph.cc | 6 +-
src/te/operation/placeholder_op.cc | 4 +-
src/te/operation/scan_op.cc | 4 +-
src/te/tensor.cc | 14 +-
src/tir/analysis/block_access_region_detector.cc | 5 +-
src/tir/analysis/buffer_access_lca_detector.cc | 3 +-
src/tir/analysis/calculate_allocated_memory.cc | 6 +-
src/tir/analysis/control_flow_graph.cc | 2 +-
src/tir/analysis/deep_equal.cc | 4 +-
src/tir/analysis/estimate_flops.cc | 23 +-
src/tir/analysis/identify_memcpy.cc | 2 +-
src/tir/analysis/is_pure_function.cc | 2 +-
src/tir/analysis/oob_checker.cc | 2 +-
src/tir/analysis/stmt_finding.cc | 2 +-
src/tir/analysis/var_use_def_analysis.cc | 2 +-
src/tir/analysis/verify_gpu_code.cc | 6 +-
src/tir/analysis/verify_memory.cc | 6 +-
src/tir/analysis/verify_ssa.cc | 6 +-
src/tir/analysis/verify_well_formed.cc | 4 +-
src/tir/ir/block_dependence_info.cc | 6 +-
src/tir/ir/block_scope.cc | 19 +-
src/tir/ir/buffer.cc | 17 +-
src/tir/ir/data_layout.cc | 34 +-
src/tir/ir/expr.cc | 83 ++-
src/tir/ir/function.cc | 10 +-
src/tir/ir/index_map.cc | 19 +-
src/tir/ir/script/script_complete.cc | 2 +-
src/tir/ir/script/script_complete.h | 2 +-
src/tir/ir/specialize.cc | 4 +-
src/tir/ir/stmt.cc | 47 +-
src/tir/ir/stmt_functor.cc | 10 +-
src/tir/ir/transform.cc | 4 +-
src/tir/op/builtin.cc | 2 +-
src/tir/op/op.cc | 72 +-
src/tir/schedule/analysis/analysis.cc | 27 +-
src/tir/schedule/analysis/layout.cc | 2 +-
src/tir/schedule/instruction.cc | 4 +-
src/tir/schedule/primitive/decompose_padding.cc | 2 +-
src/tir/schedule/primitive/reduction.cc | 2 +-
src/tir/schedule/schedule.cc | 136 ++--
src/tir/schedule/state.cc | 10 +-
src/tir/schedule/trace.cc | 20 +-
src/tir/schedule/transform.cc | 4 +-
src/tir/transforms/annotate_device_regions.cc | 5 +-
src/tir/transforms/bind_params.cc | 2 +-
src/tir/transforms/bound_checker.cc | 4 +-
src/tir/transforms/combine_context_call.cc | 4 +-
src/tir/transforms/common_subexpr_elim.cc | 2 +-
src/tir/transforms/compact_buffer_region.cc | 2 +-
src/tir/transforms/convert_blocks_to_opaque.cc | 3 +-
src/tir/transforms/convert_for_loops_serial.cc | 2 +-
src/tir/transforms/decorate_device_scope.cc | 4 +-
src/tir/transforms/default_gpu_schedule.cc | 2 +-
src/tir/transforms/extract_constants.cc | 4 +-
src/tir/transforms/flatten_buffer.cc | 2 +-
src/tir/transforms/force_narrow_index_to_i32.cc | 2 +-
src/tir/transforms/hoist_expression.cc | 8 +-
src/tir/transforms/inject_double_buffer.cc | 4 +-
src/tir/transforms/inject_permuted_layout.cc | 2 +-
src/tir/transforms/inject_ptx_async_copy.cc | 2 +-
src/tir/transforms/inject_ptx_ldg32.cc | 4 +-
src/tir/transforms/inject_rolling_buffer.cc | 4 +-
src/tir/transforms/inject_software_pipeline.cc | 3 +-
src/tir/transforms/inject_virtual_thread.cc | 4 +-
src/tir/transforms/inline_private_functions.cc | 5 +-
src/tir/transforms/ir_utils.cc | 2 +-
src/tir/transforms/lift_thread_binding.cc | 2 +-
src/tir/transforms/loop_partition.cc | 4 +-
src/tir/transforms/lower_async_dma.cc | 2 +-
src/tir/transforms/lower_cross_thread_reduction.cc | 2 +-
src/tir/transforms/lower_custom_datatypes.cc | 4 +-
src/tir/transforms/lower_device_kernel_launch.cc | 4 +-
.../transforms/lower_device_storage_access_info.cc | 4 +-
src/tir/transforms/lower_init_block.cc | 2 +-
src/tir/transforms/lower_intrin.cc | 4 +-
src/tir/transforms/lower_match_buffer.cc | 2 +-
src/tir/transforms/lower_opaque_block.cc | 2 +-
src/tir/transforms/lower_thread_allreduce.cc | 4 +-
src/tir/transforms/lower_tvm_builtin.cc | 4 +-
src/tir/transforms/lower_vtcm_alloc.cc | 2 +-
src/tir/transforms/lower_warp_memory.cc | 4 +-
src/tir/transforms/make_packed_api.cc | 6 +-
src/tir/transforms/make_unpacked_api.cc | 4 +-
.../manifest_shared_memory_local_stage.cc | 2 +-
src/tir/transforms/memhammer_lower_auto_copy.cc | 4 +-
src/tir/transforms/memhammer_rewrite_rule.h | 2 +-
.../transforms/merge_shared_memory_allocations.cc | 4 +-
src/tir/transforms/narrow_datatype.cc | 4 +-
.../plan_update_buffer_allocation_location.cc | 2 +-
src/tir/transforms/primfunc_utils.cc | 6 +-
src/tir/transforms/profile_instrumentation.cc | 2 +-
.../reduce_branching_through_overcompute.cc | 2 +-
src/tir/transforms/remap_thread_axis.cc | 4 +-
src/tir/transforms/remove_assume.cc | 4 +-
src/tir/transforms/remove_no_op.cc | 4 +-
src/tir/transforms/remove_store_undef.cc | 4 +-
.../remove_weight_layout_rewrite_block.cc | 2 +-
src/tir/transforms/renew_defs.cc | 2 +-
src/tir/transforms/renormalize_split_pattern.cc | 4 +-
src/tir/transforms/rewrite_unsafe_select.cc | 4 +-
src/tir/transforms/simplify.cc | 4 +-
src/tir/transforms/skip_assert.cc | 4 +-
src/tir/transforms/split_host_device.cc | 4 +-
src/tir/transforms/storage_rewrite.cc | 6 +-
src/tir/transforms/tensorcore_infer_fragment.cc | 4 +-
src/tir/transforms/thread_storage_sync.cc | 4 +-
src/tir/transforms/transform_mma_buffer_layout.cc | 2 +-
src/tir/transforms/unify_thread_binding.cc | 2 +-
src/tir/transforms/unroll_loop.cc | 4 +-
src/tir/transforms/unsupported_dtype_legalize.cc | 10 +-
.../transforms/using_assume_to_reduce_branches.cc | 2 +-
src/tir/transforms/vectorize_loop.cc | 4 +-
src/topi/broadcast.cc | 35 +-
src/topi/einsum.cc | 2 +-
src/topi/elemwise.cc | 75 +-
src/topi/nn.cc | 116 +--
src/topi/reduction.cc | 25 +-
src/topi/transform.cc | 169 ++---
src/topi/utils.cc | 8 +-
src/topi/vision.cc | 9 +-
tests/cpp-runtime/hexagon/run_all_tests.cc | 4 +-
tests/cpp-runtime/hexagon/run_unit_tests.cc | 4 +-
tests/cpp-runtime/opencl/texture_copy_test.cc | 2 +-
tests/cpp/llvm_codegen_registry_test.cc | 2 +-
tests/python/contrib/test_hexagon/README_RPC.md | 4 +-
tests/python/runtime/test_runtime_rpc.py | 6 +-
version.py | 4 +-
web/emcc/tvmjs_support.cc | 4 +-
web/emcc/wasm_runtime.cc | 5 +-
625 files changed, 3185 insertions(+), 5408 deletions(-)
diff --git a/3rdparty/cutlass_fpA_intB_gemm b/3rdparty/cutlass_fpA_intB_gemm
index bbccc75af1..3e07e778d7 160000
--- a/3rdparty/cutlass_fpA_intB_gemm
+++ b/3rdparty/cutlass_fpA_intB_gemm
@@ -1 +1 @@
-Subproject commit bbccc75af117473f6de81905bd3314775f41636e
+Subproject commit 3e07e778d78f0fcd047533c1fdaed571a68a396f
diff --git a/apps/android_rpc/app/src/main/jni/tvm_runtime.h b/apps/android_rpc/app/src/main/jni/tvm_runtime.h
index 26085bc366..5255d3f4b1 100644
--- a/apps/android_rpc/app/src/main/jni/tvm_runtime.h
+++ b/apps/android_rpc/app/src/main/jni/tvm_runtime.h
@@ -42,9 +42,8 @@
#include "../ffi/src/ffi/object.cc"
#include "../ffi/src/ffi/testing.cc"
#include "../ffi/src/ffi/traceback.cc"
-#include "../src/runtime/c_runtime_api.cc"
-#include "../src/runtime/container.cc"
#include "../src/runtime/cpu_device_api.cc"
+#include "../src/runtime/device_api.cc"
#include "../src/runtime/dso_library.cc"
#include "../src/runtime/file_utils.cc"
#include "../src/runtime/library_module.cc"
@@ -53,7 +52,6 @@
#include "../src/runtime/minrpc/minrpc_logger.cc"
#include "../src/runtime/module.cc"
#include "../src/runtime/ndarray.cc"
-#include "../src/runtime/object.cc"
#include "../src/runtime/profiling.cc"
#include "../src/runtime/registry.cc"
#include "../src/runtime/rpc/rpc_channel.cc"
diff --git a/apps/cpp_rpc/rpc_env.cc b/apps/cpp_rpc/rpc_env.cc
index 88ad99e47a..e5a5154acb 100644
--- a/apps/cpp_rpc/rpc_env.cc
+++ b/apps/cpp_rpc/rpc_env.cc
@@ -20,7 +20,9 @@
* \file rpc_env.cc
* \brief Server environment of the RPC.
*/
-#include <tvm/runtime/registry.h>
+#include <tvm/ffi/function.h>
+#include <tvm/runtime/logging.h>
+#include <tvm/runtime/module.h>
#include <cerrno>
#ifndef _WIN32
diff --git a/apps/cpp_rpc/rpc_env.h b/apps/cpp_rpc/rpc_env.h
index dbb0a62d2c..a5d3f6957c 100644
--- a/apps/cpp_rpc/rpc_env.h
+++ b/apps/cpp_rpc/rpc_env.h
@@ -24,7 +24,7 @@
#ifndef TVM_APPS_CPP_RPC_ENV_H_
#define TVM_APPS_CPP_RPC_ENV_H_
-#include <tvm/runtime/registry.h>
+#include <tvm/ffi/function.h>
#include <string>
diff --git a/apps/cpp_rpc/rpc_server.cc b/apps/cpp_rpc/rpc_server.cc
index c4ee4d3545..2f74dd309f 100644
--- a/apps/cpp_rpc/rpc_server.cc
+++ b/apps/cpp_rpc/rpc_server.cc
@@ -21,7 +21,7 @@
* \file rpc_server.cc
* \brief RPC Server implementation.
*/
-#include <tvm/runtime/registry.h>
+#include <tvm/ffi/function.h>
#if defined(__linux__) || defined(__ANDROID__) || defined(__APPLE__)
#include <signal.h>
#include <sys/select.h>
@@ -398,6 +398,6 @@ void RPCServerCreate(std::string host, int port, int port_end, std::string track
rpc.Start();
}
-TVM_REGISTER_GLOBAL("rpc.ServerCreate").set_body_typed(RPCServerCreate);
+TVM_FFI_REGISTER_GLOBAL("rpc.ServerCreate").set_body_typed(RPCServerCreate);
} // namespace runtime
} // namespace tvm
diff --git a/apps/cpp_rpc/rpc_server.h b/apps/cpp_rpc/rpc_server.h
index e4565d095b..9bb61065c5 100644
--- a/apps/cpp_rpc/rpc_server.h
+++ b/apps/cpp_rpc/rpc_server.h
@@ -26,7 +26,7 @@
#include <string>
-#include "tvm/runtime/c_runtime_api.h"
+#include "tvm/runtime/base.h"
namespace tvm {
namespace runtime {
diff --git a/apps/hexagon_launcher/launcher_core.cc b/apps/hexagon_launcher/launcher_core.cc
index f4fc9fb365..aebde97a51 100644
--- a/apps/hexagon_launcher/launcher_core.cc
+++ b/apps/hexagon_launcher/launcher_core.cc
@@ -19,9 +19,9 @@
#include "launcher_core.h"
+#include <tvm/ffi/function.h>
#include <tvm/runtime/c_backend_api.h>
#include <tvm/runtime/packed_func.h>
-#include <tvm/runtime/registry.h>
#include <fstream>
#include <ios>
diff --git a/apps/ios_rpc/tvmrpc/RPCServer.mm b/apps/ios_rpc/tvmrpc/RPCServer.mm
index 3dc2fb0c19..4717d71032 100644
--- a/apps/ios_rpc/tvmrpc/RPCServer.mm
+++ b/apps/ios_rpc/tvmrpc/RPCServer.mm
@@ -23,8 +23,8 @@
#import "RPCServer.h"
+#include <tvm/ffi/function.h>
#include <tvm/runtime/packed_func.h>
-#include <tvm/runtime/registry.h>
#include <random>
#include <string>
diff --git a/apps/ios_rpc/tvmrpc/TVMRuntime.mm b/apps/ios_rpc/tvmrpc/TVMRuntime.mm
index 243e4819d0..8d0ae7368d 100644
--- a/apps/ios_rpc/tvmrpc/TVMRuntime.mm
+++ b/apps/ios_rpc/tvmrpc/TVMRuntime.mm
@@ -23,7 +23,7 @@
#import <Foundation/Foundation.h>
-#include <tvm/runtime/registry.h>
+#include <tvm/ffi/function.h>
#include "RPCArgs.h"
@@ -51,14 +51,14 @@ void LogMessageImpl(const std::string& file, int lineno, int level, const std::s
} // namespace detail
-TVM_REGISTER_GLOBAL("tvm.rpc.server.workpath")
+TVM_FFI_REGISTER_GLOBAL("tvm.rpc.server.workpath")
.set_body_packed([](ffi::PackedArgs args, ffi::Any* rv) {
static const std::string base_ = NSTemporaryDirectory().UTF8String;
const auto path = args[0].cast<std::string>();
*rv = base_ + "/" + path;
});
-TVM_REGISTER_GLOBAL("tvm.rpc.server.load_module")
+TVM_FFI_REGISTER_GLOBAL("tvm.rpc.server.load_module")
.set_body_packed([](ffi::PackedArgs args, ffi::Any* rv) {
auto name = args[0].cast<std::string>();
std::string fmt = GetFileFormat(name, "");
@@ -109,7 +109,7 @@ class UnsignedDSOLoader final : public Library {
};
// Add UnsignedDSOLoader plugin in global registry
-TVM_REGISTER_GLOBAL("runtime.module.loadfile_dylib_custom")
+TVM_FFI_REGISTER_GLOBAL("runtime.module.loadfile_dylib_custom")
.set_body_packed([](ffi::PackedArgs args, ffi::Any* rv) {
auto n = make_object<UnsignedDSOLoader>();
n->Init(args[0]);
diff --git a/docs/arch/device_target_interactions.rst b/docs/arch/device_target_interactions.rst
index e39468f0bf..09867f88fa 100644
--- a/docs/arch/device_target_interactions.rst
+++ b/docs/arch/device_target_interactions.rst
@@ -153,18 +153,18 @@ then be registered with the following steps.
#. Register the function to the tvm registry::
- TVM_REGISTER_GLOBAL("device_api.foo").set_body_typed(FooDeviceAPI::Global);
+ TVM_FFI_REGISTER_GLOBAL("device_api.foo").set_body_typed(FooDeviceAPI::Global);
-.. _c_runtime_api.h: https://github.com/apache/tvm/blob/main/include/tvm/runtime/c_runtime_api.h
+.. _base.h: https://github.com/apache/tvm/blob/main/include/tvm/runtime/base.h
#. Add an entry for the new DeviceAPI to the ``TVMDeviceExtType`` enum
- in `c_runtime_api.h`_. The value should be an unused value greater
+ in `base.h`_. The value should be an unused value greater
than ``DLDeviceType::kDLExtDev``, but less than
``DeviceAPIManager::kMaxDeviceAPI``.
#. Add a case in ``DeviceName`` in `device_api.h`_ to convert from the
enum value to a string representation. This string representation
- should match the name given to ``TVM_REGISTER_GLOBAL``.
+ should match the name given to ``TVM_FFI_REGISTER_GLOBAL``.
#. Add entries to the ``DEVICE_TYPE_TO_NAME`` and ``DEVICE_NAME_TO_TYPE`` dictionaries of
:py:class:`tvm.runtime.Device` for the new enum value.
@@ -225,7 +225,7 @@ the same name as was used in the ``TVM_REGISTER_TARGET_KIND``
definition above. ::
tvm::runtime::Module GeneratorFooCode(IRModule mod, Target target);
- TVM_REGISTER_GLOBAL("target.build.foo").set_body_typed(GeneratorFooCode);
+ TVM_FFI_REGISTER_GLOBAL("target.build.foo").set_body_typed(GeneratorFooCode);
The code generator takes two arguments. The first is the ``IRModule``
to compile, and the second is the ``Target`` that describes the device
diff --git a/docs/arch/pass_infra.rst b/docs/arch/pass_infra.rst
index 85e9f45a5f..bf7b52229d 100644
--- a/docs/arch/pass_infra.rst
+++ b/docs/arch/pass_infra.rst
@@ -376,7 +376,7 @@ Python when needed.
return CreateFunctionPass(pass_func, 0, "FoldConstant", {});
}
- TVM_REGISTER_GLOBAL("relax.transform.FoldConstant")
+ TVM_FFI_REGISTER_GLOBAL("relax.transform.FoldConstant")
.set_body_typed(FoldConstant);
} // namespace transform
diff --git a/docs/arch/runtime.rst b/docs/arch/runtime.rst
index f797039ee3..55c523cb4c 100644
--- a/docs/arch/runtime.rst
+++ b/docs/arch/runtime.rst
@@ -80,7 +80,7 @@ The following example registers PackedFunc in C++ and calls from python.
.. code:: c
// register a global packed function in c++
- TVM_REGISTER_GLOBAL("myadd")
+ TVM_FFI_REGISTER_GLOBAL("myadd")
.set_body_packed(MyAdd);
.. code:: python
@@ -110,7 +110,7 @@ we can pass functions from python (as PackedFunc) to C++.
.. code:: c
- TVM_REGISTER_GLOBAL("callhello")
+ TVM_FFI_REGISTER_GLOBAL("callhello")
.set_body_packed([](ffi::PackedArgs args, ffi::Any* rv) {
PackedFunc f = args[0];
f("hello world");
@@ -134,7 +134,7 @@ which allows us to embed the PackedFunc into any languages. Besides python, so f
`java`_ and `javascript`_.
This philosophy of embedded API is very like Lua, except that we don't have a new language but use C++.
-.. _minimum C API: https://github.com/apache/tvm/blob/main/include/tvm/runtime/c_runtime_api.h
+.. _minimum C API: https://github.com/apache/tvm/blob/main/include/tvm/runtime/base.h
.. _java: https://github.com/apache/tvm/tree/main/jvm
.. _javascript: https://github.com/apache/tvm/tree/main/web
@@ -282,7 +282,7 @@ Each argument in PackedFunc contains a union value `TVMValue`_
and a type code. This design allows the dynamically typed language to convert to the corresponding type directly, and statically typed language to
do runtime type checking during conversion.
-.. _TVMValue: https://github.com/apache/tvm/blob/main/include/tvm/runtime/c_runtime_api.h#L135
+.. _TVMValue: https://github.com/apache/tvm/blob/main/include/tvm/runtime/base.h#L135
The relevant files are
diff --git a/ffi/include/tvm/ffi/c_api.h b/ffi/include/tvm/ffi/c_api.h
index 131f2e73e0..df8265d0b9 100644
--- a/ffi/include/tvm/ffi/c_api.h
+++ b/ffi/include/tvm/ffi/c_api.h
@@ -27,6 +27,13 @@
#include <dlpack/dlpack.h>
#include <stdint.h>
+// Macros to do weak linking
+#ifdef _MSC_VER
+#define TVM_FFI_WEAK __declspec(selectany)
+#else
+#define TVM_FFI_WEAK __attribute__((weak))
+#endif
+
#if !defined(TVM_FFI_DLL) && defined(__EMSCRIPTEN__)
#include <emscripten/emscripten.h>
#define TVM_FFI_DLL EMSCRIPTEN_KEEPALIVE
diff --git a/ffi/include/tvm/ffi/function.h b/ffi/include/tvm/ffi/function.h
index 5162df6d83..753d4f50f1 100644
--- a/ffi/include/tvm/ffi/function.h
+++ b/ffi/include/tvm/ffi/function.h
@@ -787,7 +787,7 @@ class Function::Registry {
* .set_body_typed(multiply); // will have type int(int, int)
*
* // will have type int(int, int)
- * TVM_REGISTER_GLOBAL("sub")
+ * TVM_FFI_REGISTER_GLOBAL("sub")
* .set_body_typed([](int a, int b) -> int { return a - b; });
*
* \endcode
diff --git a/include/tvm/ir/op.h b/include/tvm/ir/op.h
index 8eaa62a981..9c758a52b3 100644
--- a/include/tvm/ir/op.h
+++ b/include/tvm/ir/op.h
@@ -25,13 +25,13 @@
#ifndef TVM_IR_OP_H_
#define TVM_IR_OP_H_
+#include <tvm/ffi/function.h>
#include <tvm/ir/attrs.h>
#include <tvm/ir/env_func.h>
#include <tvm/ir/expr.h>
#include <tvm/ir/type.h>
#include <tvm/node/attr_registry_map.h>
#include <tvm/runtime/logging.h>
-#include <tvm/runtime/registry.h>
#include <string>
#include <utility>
diff --git a/include/tvm/ir/source_map.h b/include/tvm/ir/source_map.h
index 7b79a2c894..2752d9951a 100644
--- a/include/tvm/ir/source_map.h
+++ b/include/tvm/ir/source_map.h
@@ -23,10 +23,10 @@
#ifndef TVM_IR_SOURCE_MAP_H_
#define TVM_IR_SOURCE_MAP_H_
+#include <tvm/ffi/function.h>
#include <tvm/node/node.h>
#include <tvm/runtime/object.h>
#include <tvm/runtime/packed_func.h>
-#include <tvm/runtime/registry.h>
#include <fstream>
#include <string>
diff --git a/include/tvm/node/node.h b/include/tvm/node/node.h
index aebff220ab..8a9e763fec 100644
--- a/include/tvm/node/node.h
+++ b/include/tvm/node/node.h
@@ -39,7 +39,7 @@
#include <tvm/node/repr_printer.h>
#include <tvm/node/structural_equal.h>
#include <tvm/node/structural_hash.h>
-#include <tvm/runtime/c_runtime_api.h>
+#include <tvm/runtime/base.h>
#include <tvm/runtime/object.h>
#include <string>
diff --git a/include/tvm/node/reflection.h b/include/tvm/node/reflection.h
index fb15ff4f81..ab197078f3 100644
--- a/include/tvm/node/reflection.h
+++ b/include/tvm/node/reflection.h
@@ -27,7 +27,7 @@
#include <tvm/ffi/memory.h>
#include <tvm/node/structural_equal.h>
#include <tvm/node/structural_hash.h>
-#include <tvm/runtime/c_runtime_api.h>
+#include <tvm/runtime/base.h>
#include <tvm/runtime/data_type.h>
#include <tvm/runtime/ndarray.h>
#include <tvm/runtime/object.h>
diff --git a/include/tvm/node/serialization.h b/include/tvm/node/serialization.h
index c99d0f7f73..5a8e098cfd 100644
--- a/include/tvm/node/serialization.h
+++ b/include/tvm/node/serialization.h
@@ -24,7 +24,7 @@
#ifndef TVM_NODE_SERIALIZATION_H_
#define TVM_NODE_SERIALIZATION_H_
-#include <tvm/runtime/c_runtime_api.h>
+#include <tvm/runtime/base.h>
#include <tvm/runtime/object.h>
#include <string>
diff --git a/include/tvm/relax/exec_builder.h b/include/tvm/relax/exec_builder.h
index 2cee3bca63..81d6d4eb37 100644
--- a/include/tvm/relax/exec_builder.h
+++ b/include/tvm/relax/exec_builder.h
@@ -23,11 +23,11 @@
#ifndef TVM_RELAX_EXEC_BUILDER_H_
#define TVM_RELAX_EXEC_BUILDER_H_
+#include <tvm/ffi/function.h>
#include <tvm/ir/expr.h>
#include <tvm/node/reflection.h>
#include <tvm/node/repr_printer.h>
#include <tvm/runtime/object.h>
-#include <tvm/runtime/registry.h>
#include <tvm/runtime/relax_vm/bytecode.h>
#include <tvm/runtime/relax_vm/executable.h>
diff --git a/include/tvm/relax/type.h b/include/tvm/relax/type.h
index 7e4149fe55..bd75197bfe 100644
--- a/include/tvm/relax/type.h
+++ b/include/tvm/relax/type.h
@@ -24,10 +24,10 @@
#ifndef TVM_RELAX_TYPE_H_
#define TVM_RELAX_TYPE_H_
+#include <tvm/ffi/function.h>
#include <tvm/ir/attrs.h>
#include <tvm/ir/env_func.h>
#include <tvm/ir/type.h>
-#include <tvm/runtime/registry.h>
#include <tvm/tir/expr.h>
#include <string>
diff --git a/src/tir/transforms/decorate_device_scope.cc b/include/tvm/runtime/base.h
similarity index 50%
copy from src/tir/transforms/decorate_device_scope.cc
copy to include/tvm/runtime/base.h
index 5034a85813..c704decb63 100644
--- a/src/tir/transforms/decorate_device_scope.cc
+++ b/include/tvm/runtime/base.h
@@ -17,35 +17,43 @@
* under the License.
*/
-/*!
- * \file decorate_device_scope.cc
+/*
+ * \file tvm/runtime/base.h
+ * \brief base macros
*/
-#include <tvm/runtime/registry.h>
-#include <tvm/tir/op.h>
-#include <tvm/tir/stmt.h>
-#include <tvm/tir/transform.h>
+#ifndef TVM_RUNTIME_BASE_H_
+#define TVM_RUNTIME_BASE_H_
-namespace tvm {
-namespace tir {
+// TVM runtime fully relies on TVM FFI C API
+// we will avoid defining extra C APIs here
+#include <tvm/ffi/c_api.h>
-Stmt DecorateDeviceScope(Stmt&& stmt) {
- Stmt body = AttrStmt(make_zero(DataType::Int(32)), tir::attr::device_scope, 0, stmt);
- return body;
-}
+// TVM version
+#define TVM_VERSION "0.21.dev0"
-namespace transform {
+// define extra macros for TVM DLL exprt
+#ifdef __EMSCRIPTEN__
+#include <emscripten/emscripten.h>
+#define TVM_DLL EMSCRIPTEN_KEEPALIVE
+#endif
-Pass DecorateDeviceScope() {
- auto pass_func = [](PrimFunc f, IRModule m, PassContext ctx) {
- auto* n = f.CopyOnWrite();
- n->body = DecorateDeviceScope(std::move(n->body));
- return f;
- };
- return CreatePrimFuncPass(pass_func, 0, "tir.DecorateDeviceScope", {});
-}
+// helper macro to suppress unused warning
+#if defined(__GNUC__)
+#define TVM_ATTRIBUTE_UNUSED __attribute__((unused))
+#else
+#define TVM_ATTRIBUTE_UNUSED
+#endif
-TVM_REGISTER_GLOBAL("tir.transform.DecorateDeviceScope").set_body_typed(DecorateDeviceScope);
+#ifndef TVM_DLL
+#ifdef _WIN32
+#ifdef TVM_EXPORTS
+#define TVM_DLL __declspec(dllexport)
+#else
+#define TVM_DLL __declspec(dllimport)
+#endif
+#else
+#define TVM_DLL __attribute__((visibility("default")))
+#endif
+#endif
-} // namespace transform
-} // namespace tir
-} // namespace tvm
+#endif // TVM_RUNTIME_BASE_H_
diff --git a/include/tvm/runtime/builtin_fp16.h b/include/tvm/runtime/builtin_fp16.h
index 5b54583da4..3ea670017d 100644
--- a/include/tvm/runtime/builtin_fp16.h
+++ b/include/tvm/runtime/builtin_fp16.h
@@ -24,7 +24,7 @@
#ifndef TVM_RUNTIME_BUILTIN_FP16_H_
#define TVM_RUNTIME_BUILTIN_FP16_H_
-#include <tvm/runtime/c_runtime_api.h>
+#include <tvm/runtime/base.h>
#include <cstdint>
diff --git a/include/tvm/runtime/c_backend_api.h b/include/tvm/runtime/c_backend_api.h
index eb8d7270b1..0d84b55fe3 100644
--- a/include/tvm/runtime/c_backend_api.h
+++ b/include/tvm/runtime/c_backend_api.h
@@ -28,28 +28,12 @@
#ifndef TVM_RUNTIME_C_BACKEND_API_H_
#define TVM_RUNTIME_C_BACKEND_API_H_
-#include <tvm/runtime/c_runtime_api.h>
+#include <tvm/runtime/base.h>
#ifdef __cplusplus
extern "C" {
#endif
-/*!
- * \brief Signature for backend functions exported as DLL.
- *
- * \param args The arguments
- * \param type_codes The type codes of the arguments
- * \param num_args Number of arguments.
- * \param out_ret_value The output value of the return value.
- * \param out_ret_tcode The output type code of the return value.
- * \param resource_handle Pointer to associated resource.
- *
- * \return 0 if success, -1 if failure happens, set error via TVMAPISetLastError.
- */
-typedef int (*TVMBackendPackedCFunc)(TVMValue* args, int* type_codes, int num_args,
- TVMValue* out_ret_value, int* out_ret_tcode,
- void* resource_handle);
-
/*!
* \brief Backend function for modules to get function
* from its environment mod_node (its imports and global function).
@@ -60,7 +44,8 @@ typedef int (*TVMBackendPackedCFunc)(TVMValue* args, int* type_codes, int num_ar
* \param out The result function.
* \return 0 when no error is thrown, -1 when failure happens
*/
-TVM_DLL int TVMBackendGetFuncFromEnv(void* mod_node, const char* func_name, TVMFunctionHandle* out);
+TVM_DLL int TVMBackendGetFuncFromEnv(void* mod_node, const char* func_name,
+ TVMFFIObjectHandle* out);
/*!
* \brief Backend function to register system-wide library symbol.
@@ -100,19 +85,6 @@ TVM_DLL void* TVMBackendAllocWorkspace(int device_type, int device_id, uint64_t
*/
TVM_DLL int TVMBackendFreeWorkspace(int device_type, int device_id, void* ptr);
-/*!
- * \brief Backend function to register execution environment(e.g. python)
- * specific C APIs.
- *
- * \note We only register the C API function when absolutely necessary (e.g. when signal handler
- * cannot trap back into python). In most cases we should use the ffi::Function FFI.
- *
- * \param name The name of the symbol
- * \param ptr The symbol address.
- * \return 0 when no error is thrown, -1 when failure happens
- */
-TVM_DLL int TVMBackendRegisterEnvCAPI(const char* name, void* ptr);
-
/*!
* \brief Environment for TVM parallel task.
*/
diff --git a/include/tvm/runtime/c_runtime_api.h b/include/tvm/runtime/c_runtime_api.h
deleted file mode 100644
index b802dbc228..0000000000
--- a/include/tvm/runtime/c_runtime_api.h
+++ /dev/null
@@ -1,732 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-/*
- * \file tvm/runtime/c_runtime_api.h
- * \brief TVM runtime library.
- *
- * The philosophy of TVM project is to customize the compilation
- * stage to generate code that can used by other projects transparently.
- * So this is a minimum runtime code gluing, and some limited
- * memory management code to enable quick testing.
- *
- * The runtime API is independent from TVM compilation stack and can
- * be linked via libtvm_runtime.
- *
- * The common flow is:
- * - Use TVMFuncListGlobalNames to get global function name
- * - Use TVMFuncCall to call these functions.
- *
- * Possible return values of the API functions:
- * * 0: success
- * * -1: the error can be retrieved through TVMGetLastError.
- * * -2: a frontend error occurred and recorded in the frontend.
- */
-#ifndef TVM_RUNTIME_C_RUNTIME_API_H_
-#define TVM_RUNTIME_C_RUNTIME_API_H_
-
-// Macros to do weak linking
-#ifdef _MSC_VER
-#define TVM_WEAK __declspec(selectany)
-#else
-#define TVM_WEAK __attribute__((weak))
-#endif
-
-#ifdef __EMSCRIPTEN__
-#include <emscripten/emscripten.h>
-#define TVM_DLL EMSCRIPTEN_KEEPALIVE
-#endif
-
-// helper macro to suppress unused warning
-#if defined(__GNUC__)
-#define TVM_ATTRIBUTE_UNUSED __attribute__((unused))
-#else
-#define TVM_ATTRIBUTE_UNUSED
-#endif
-
-#ifndef TVM_DLL
-#ifdef _WIN32
-#ifdef TVM_EXPORTS
-#define TVM_DLL __declspec(dllexport)
-#else
-#define TVM_DLL __declspec(dllimport)
-#endif
-#else
-#define TVM_DLL __attribute__((visibility("default")))
-#endif
-#endif
-
-// TVM version
-#define TVM_VERSION "0.21.dev0"
-
-// TVM Runtime is DLPack compatible.
-#include <dlpack/dlpack.h>
-
-#ifdef __cplusplus
-extern "C" {
-#endif
-#include <stdbool.h>
-#include <stddef.h>
-#include <stdint.h>
-
-/*! \brief type of array index. */
-typedef int64_t tvm_index_t;
-
-/*! \brief Extension device types in TVM
- *
- * Additional enumerators to supplement those provided by
- * DLPack's `DLDeviceType` enumeration.
- *
- * MAINTAINERS NOTE #1: We need to ensure that the two devices
- * are identified by the same integer.
- * Currently this requires manual verification.
- * Discussed here: https://github.com/dmlc/dlpack/issues/111
- * As of DLPack v0.7, the highest-valued enumerator in
- * `DLDeviceType` is kDLHexagon = 16.
- *
- * MAINTAINERS NOTE #2: As of DLPack v0.7, the definition for
- * `DLDeviceType` specifies an underlying storage type of
- * `int32_t`. That guarantees a variable of type
- * `DLDeviceType` is capable of holding any integers provided
- * by *either* of these enumerations.
- *
- * However, the `int32_t` specification only applies when the
- * header file is compiled as C++, and this header file is also
- * meant to work as C code. So the unspecified storage type
- * could be a latent bug when compiled as C.
- */
-#ifdef __cplusplus
-typedef enum : int32_t {
-#else
-typedef enum {
-#endif
- // To help avoid accidental conflicts between `DLDeviceType`
- // and this enumeration, start numbering the new enumerators
- // a little higher than (currently) seems necessary.
- TVMDeviceExtType_End = 36, // sentinel value
-} TVMDeviceExtType;
-
-#ifdef __cplusplus
-// Some other parts of TVM hardcode the integer identifier for
-// some DLPack / TVM devices, rather then using the symbolic
-// enumerator. E.g., `2` rather than `kDLCUDA`.
-// These asserts should alert us when that mapping breaks.
-#define TVM_HARCODED_INTEGER_CHANGED_MSG \
- "Change in compile-time integer. Make sure hardcoded uses of this integer throughout TVM are " \
- "updated."
-static_assert(kDLCPU == 1, TVM_HARCODED_INTEGER_CHANGED_MSG);
-static_assert(kDLCUDA == 2, TVM_HARCODED_INTEGER_CHANGED_MSG);
-static_assert(kDLCUDAHost == 3, TVM_HARCODED_INTEGER_CHANGED_MSG);
-static_assert(kDLOpenCL == 4, TVM_HARCODED_INTEGER_CHANGED_MSG);
-static_assert(kDLVulkan == 7, TVM_HARCODED_INTEGER_CHANGED_MSG);
-static_assert(kDLMetal == 8, TVM_HARCODED_INTEGER_CHANGED_MSG);
-static_assert(kDLVPI == 9, TVM_HARCODED_INTEGER_CHANGED_MSG);
-static_assert(kDLROCM == 10, TVM_HARCODED_INTEGER_CHANGED_MSG);
-static_assert(kDLROCMHost == 11, TVM_HARCODED_INTEGER_CHANGED_MSG);
-static_assert(kDLExtDev == 12, TVM_HARCODED_INTEGER_CHANGED_MSG);
-static_assert(kDLCUDAManaged == 13, TVM_HARCODED_INTEGER_CHANGED_MSG);
-static_assert(kDLOneAPI == 14, TVM_HARCODED_INTEGER_CHANGED_MSG);
-static_assert(kDLWebGPU == 15, TVM_HARCODED_INTEGER_CHANGED_MSG);
-static_assert(kDLHexagon == 16, TVM_HARCODED_INTEGER_CHANGED_MSG);
-
-#undef TVM_HARCODED_INTEGER_CHANGED_MSG
-#endif
-
-/*!
- * \brief The type code in used and only used in TVM FFI for argument passing.
- *
- * DLPack consistency:
- * 1) kTVMArgInt is compatible with kDLInt
- * 2) kTVMArgFloat is compatible with kDLFloat
- * 3) kDLUInt is not in ArgTypeCode, but has a spared slot
- *
- * Downstream consistency:
- * The kDLInt, kDLUInt, kDLFloat are kept consistent with the original ArgType code
- *
- * It is only used in argument passing, and should not be confused with
- * DataType::TypeCode, which is DLPack-compatible.
- *
- * \sa tvm::runtime::DataType::TypeCode
- */
-typedef enum {
- kTVMArgInt = kDLInt,
- kTVMArgFloat = kDLFloat,
- kTVMOpaqueHandle = 3U,
- kTVMNullptr = 4U,
- kTVMDataType = 5U,
- kDLDevice = 6U,
- kTVMDLTensorHandle = 7U,
- kTVMObjectHandle = 8U,
- kTVMModuleHandle = 9U,
- kTVMPackedFuncHandle = 10U,
- kTVMStr = 11U,
- kTVMBytes = 12U,
- kTVMNDArrayHandle = 13U,
- kTVMObjectRValueRefArg = 14U,
- kTVMArgBool = 15U,
- // Extension codes for other frameworks to integrate TVM ffi::Function.
- // To make sure each framework's id do not conflict, use first and
- // last sections to mark ranges.
- // Open an issue at the repo if you need a section of code.
- kTVMExtBegin = 16U,
- kTVMNNVMFirst = 16U,
- kTVMNNVMLast = 20U,
- // The following section of code is used for non-reserved types.
- kTVMExtReserveEnd = 64U,
- kTVMExtEnd = 128U,
-} TVMArgTypeCode;
-
-/*! \brief the array handle */
-typedef DLTensor* TVMArrayHandle;
-
-/*!
- * \brief Union type of values
- * being passed through API and function calls.
- */
-typedef union {
- int64_t v_int64;
- double v_float64;
- void* v_handle;
- const char* v_str;
- DLDataType v_type;
- DLDevice v_device;
-} TVMValue;
-
-/*!
- * \brief Byte array type used to pass in byte array
- * When kTVMBytes is used as data type.
- */
-typedef struct {
- const char* data;
- size_t size;
-} TVMByteArray;
-
-/*! \brief Handle to TVM runtime modules. */
-typedef void* TVMModuleHandle;
-/*! \brief Handle to packed function handle. */
-typedef void* TVMFunctionHandle;
-/*! \brief Handle to hold return value. */
-typedef void* TVMRetValueHandle;
-/*!
- * \brief The stream that is specific to device
- * can be NULL, which indicates the default one.
- */
-typedef void* TVMStreamHandle;
-/*! \brief Handle to Object. */
-typedef void* TVMObjectHandle;
-
-/*!
- * \brief Used for implementing C API function.
- * Set last error message before return.
- * \param msg The error message to be set.
- */
-TVM_DLL void TVMAPISetLastError(const char* msg);
-
-/*!
- * \brief Used for implementing C API function.
- * Set last exception before return.
- * \param py_object The python exception to be set
- */
-TVM_DLL void TVMAPISetLastPythonError(void* py_object);
-
-/*! \brief Return the previous python error, if any.
- *
- * Used to propagate the original Python exception to a python
- * try/except, when there are C++ stack frames between the location thro
- *
- * \return The previous argument passed during the most recent call to
- * TVMAPISetLastPythonError. If TVMAPISetLastPythonError has not
- * been called, or if TVMDropLastPythonError has been called since
- * the most recent to TVMAPISetLastPythonError, returns nullptr.
- */
-TVM_DLL void* TVMGetLastPythonError();
-
-/*!
- * \brief return str message of the last error
- * all function in this file will return 0 when success
- * and nonzero when an error occurred,
- * TVMGetLastError can be called to retrieve the error
- *
- * this function is threadsafe and can be called by different thread
- * \return error info
- */
-TVM_DLL const char* TVMGetLastError(void);
-
-/*!
- * \brief Return the backtrace of the most recent error
- *
- * Returns the backtrace of the most recent error, if an error exists,
- * and the error contains a backtrace. If no error exists or the
- * error does not contain a backtrace, returns nullptr.
- *
- * \return The backtrace of the most recent error
- */
-TVM_DLL const char* TVMGetLastBacktrace();
-
-/*!
- * \brief Remove the propagated python error, if any
- *
- * Removes the TVM-held reference to a thrown python exception object.
- * Because these objects contain references to the stack frames from
- * which the exception was thrown, maintaining a reference to an
- * exception object prevents any local python variables from being
- * garbage-collected. After retrieving the object using
- * TVMGetLastPythonError, the Python FFI interface uses this method to
- * clear the TVM-held reference to the exception, to allow garbage
- * collection to continue.
- */
-TVM_DLL void TVMDropLastPythonError();
-
-/*! \brief Re-throw the most recent error.
- *
- * If an error was previously set using TVMAPISetLastError or
- * TVMAPISetLastPythonError, re-throw the error. This is similar to
- * `LOG(FATAL) << TVMGetLastError()`, but includes handling to
- * propagate a python exception across C++ stack frames, or to append
- * a stack trace to an error message.
- */
-TVM_DLL void TVMThrowLastError();
-
-/*!
- * \brief Load module from file.
- * \param file_name The file name to load the module from.
- * \param format The format of the module.
- * \param out The result module
- *
- * \return 0 when success, nonzero when failure happens
- * \note The resulting module do not contain import relation.
- * It can be reconstructed by TVMModImport.
- */
-TVM_DLL int TVMModLoadFromFile(const char* file_name, const char* format, TVMModuleHandle* out);
-
-/*!
- * \brief Add dep to mod's dependency.
- * This allows functions in this module to use modules.
- *
- * \param mod The module handle.
- * \param dep The dependent module to be imported.
- * \return 0 when success, nonzero when failure happens
- */
-TVM_DLL int TVMModImport(TVMModuleHandle mod, TVMModuleHandle dep);
-
-/*!
- * \brief Get function from the module.
- * \param mod The module handle.
- * \param func_name The name of the function.
- * \param query_imports Whether to query imported modules
- * \param out The result function, can be NULL if it is not available.
- * \return 0 when no error is thrown, nonzero when failure happens
- */
-TVM_DLL int TVMModGetFunction(TVMModuleHandle mod, const char* func_name, int query_imports,
- TVMFunctionHandle* out);
-
-/*!
- * \brief Free the Module
- * \param mod The module to be freed.
- *
- * \note This may not free up the module's resources.
- * If there is active TVMFunctionHandle uses the module
- * Or if this module is imported by another active module.
- *
- * The all functions remains valid until TVMFuncFree is called.
- * \return 0 when success, nonzero when failure happens
- */
-TVM_DLL int TVMModFree(TVMModuleHandle mod);
-
-/*!
- * \brief Free the function when it is no longer needed.
- * \param func The function handle
- * \return 0 when success, nonzero when failure happens
- */
-TVM_DLL int TVMFuncFree(TVMFunctionHandle func);
-
-/*!
- * \brief Call a Packed TVM Function.
- *
- * \param func node handle of the function.
- * \param arg_values The arguments
- * \param type_codes The type codes of the arguments
- * \param num_args Number of arguments.
- *
- * \param ret_val The return value.
- * \param ret_type_code the type code of return value.
- *
- * \return 0 when success, nonzero when failure happens
- * \note TVM calls always exchanges with type bits=64, lanes=1
- *
- * \note API calls always exchanges with type bits=64, lanes=1
- * If API call returns container handles (e.g. FunctionHandle)
- * these handles should be managed by the front-end.
- * The front-end need to call free function (e.g. TVMFuncFree)
- * to free these handles.
- */
-TVM_DLL int TVMFuncCall(TVMFunctionHandle func, TVMValue* arg_values, int* type_codes, int num_args,
- TVMValue* ret_val, int* ret_type_code);
-
-/*!
- * \brief Set the return value of TVMPackedCFunc.
- *
- * This function is called by TVMPackedCFunc to set the return value.
- * When this function is not called, the function returns null by default.
- *
- * \param ret The return value handle, pass by ret in TVMPackedCFunc
- * \param value The value to be returned.
- * \param type_code The type of the value to be returned.
- * \param num_ret Number of return values, for now only 1 is supported.
- */
-TVM_DLL int TVMCFuncSetReturn(TVMRetValueHandle ret, TVMValue* value, int* type_code, int num_ret);
-
-/*!
- * \brief Inplace translate callback argument value to return value.
- * This is only needed for non-POD arguments.
- *
- * \param value The value to be translated.
- * \param code The type code to be translated.
- * \note This function will do a shallow copy when necessary.
- *
- * \return 0 when success, nonzero when failure happens.
- */
-TVM_DLL int TVMCbArgToReturn(TVMValue* value, int* code);
-
-/*!
- * \brief C type of packed function.
- *
- * \param args The arguments
- * \param type_codes The type codes of the arguments
- * \param num_args Number of arguments.
- * \param ret The return value handle.
- * \param resource_handle The handle additional resouce handle from front-end.
- * \return 0 if success, -1 if failure happens, set error via TVMAPISetLastError.
- * \sa TVMCFuncSetReturn
- */
-typedef int (*TVMPackedCFunc)(TVMValue* args, int* type_codes, int num_args, TVMRetValueHandle ret,
- void* resource_handle);
-
-/*!
- * \brief C callback to free the resource handle in C packed function.
- * \param resource_handle The handle additional resouce handle from front-end.
- */
-typedef void (*TVMPackedCFuncFinalizer)(void* resource_handle);
-
-/*!
- * \brief Signature for extension function declarer.
- *
- * TVM call this function to get the extension functions
- * The declarer will call register_func to register function and their name.
- *
- * \param register_func_handle The register function
- * \return 0 if success, -1 if failure happens
- */
-typedef int (*TVMExtensionFuncDeclarer)(TVMFunctionHandle register_func_handle);
-
-/*!
- * \brief Wrap a TVMPackedCFunc to become a FunctionHandle.
- *
- * The resource_handle will be managed by TVM API, until the function is no longer used.
- *
- * \param func The packed C function.
- * \param resource_handle The resource handle from front-end, can be NULL.
- * \param fin The finalizer on resource handle when the FunctionHandle get freed, can be NULL
- * \param out the result function handle.
- * \return 0 when success, nonzero when failure happens
- */
-TVM_DLL int TVMFuncCreateFromCFunc(TVMPackedCFunc func, void* resource_handle,
- TVMPackedCFuncFinalizer fin, TVMFunctionHandle* out);
-
-/*!
- * \brief Register the function to runtime's global table.
- *
- * The registered function then can be pulled by the backend by the name.
- *
- * \param name The name of the function.
- * \param f The function to be registered.
- * \param override Whether allow override already registered function.
- */
-TVM_DLL int TVMFuncRegisterGlobal(const char* name, TVMFunctionHandle f, int override);
-
-/*!
- * \brief Get a global function.
- *
- * \param name The name of the function.
- * \param out the result function pointer, NULL if it does not exist.
- *
- * \note The function handle of global function is managed by TVM runtime,
- * So TVMFuncFree is should not be called when it get deleted.
- */
-TVM_DLL int TVMFuncGetGlobal(const char* name, TVMFunctionHandle* out);
-
-/*!
- * \brief List all the globally registered function name
- * \param out_size The number of functions
- * \param out_array The array of function names.
- * \return 0 when success, nonzero when failure happens
- */
-TVM_DLL int TVMFuncListGlobalNames(int* out_size, const char*** out_array);
-
-/*!
- * \brief Remove a global function.
- * \param name The name of the function.
- */
-TVM_DLL int TVMFuncRemoveGlobal(const char* name);
-
-// Array related apis for quick proptyping
-/*!
- * \brief Allocate a nd-array's memory,
- * including space of shape, of given spec.
- *
- * \param shape The shape of the array, the data content will be copied to out
- * \param ndim The number of dimension of the array.
- * \param dtype_code The type code of the dtype
- * \param dtype_bits The number of bits of dtype
- * \param dtype_lanes The number of lanes in the dtype.
- * \param device_type The device type.
- * \param device_id The device id.
- * \param out The output handle.
- * \return 0 when success, nonzero when failure happens
- */
-TVM_DLL int TVMArrayAlloc(const tvm_index_t* shape, int ndim, int dtype_code, int dtype_bits,
- int dtype_lanes, int device_type, int device_id, TVMArrayHandle* out);
-
-/*!
- * \brief Free the TVM Array.
- * \param handle The array handle to be freed.
- * \return 0 when success, nonzero when failure happens
- */
-TVM_DLL int TVMArrayFree(TVMArrayHandle handle);
-
-/*!
- * \brief Copy array data from CPU byte array.
- * \param handle The array handle.
- * \param data the data pointer
- * \param nbytes The number of bytes to copy.
- * \return 0 when success, nonzero when failure happens
- */
-TVM_DLL int TVMArrayCopyFromBytes(TVMArrayHandle handle, void* data, size_t nbytes);
-
-/*!
- * \brief Copy array data to CPU byte array.
- * \param handle The array handle.
- * \param data the data pointer
- * \param nbytes The number of bytes to copy.
- * \return 0 when success, nonzero when failure happens
- */
-TVM_DLL int TVMArrayCopyToBytes(TVMArrayHandle handle, void* data, size_t nbytes);
-
-/*!
- * \brief Copy the array, both from and to must be valid during the copy.
- * \param from The array to be copied from.
- * \param to The target space.
- * \param stream The stream where the copy happens, can be NULL.
- * \return 0 when success, nonzero when failure happens
- */
-TVM_DLL int TVMArrayCopyFromTo(TVMArrayHandle from, TVMArrayHandle to, TVMStreamHandle stream);
-
-/*!
- * \brief Produce an array from the DLManagedTensor that shares data memory
- * with the DLManagedTensor.
- * \param from The source DLManagedTensor.
- * \param out The output array handle.
- * \return 0 when success, nonzero when failure happens
- */
-TVM_DLL int TVMArrayFromDLPack(DLManagedTensor* from, TVMArrayHandle* out);
-
-/*!
- * \brief Produce a DLMangedTensor from the array that shares data memory with
- * the array.
- * \param from The source array.
- * \param out The DLManagedTensor handle.
- * \return 0 when success, nonzero when failure happens
- */
-TVM_DLL int TVMArrayToDLPack(TVMArrayHandle from, DLManagedTensor** out);
-
-/*!
- * \brief Delete (free) a DLManagedTensor's data.
- * \param dltensor Pointer to the DLManagedTensor.
- */
-TVM_DLL void TVMDLManagedTensorCallDeleter(DLManagedTensor* dltensor);
-
-/*!
- * \brief Create a new runtime stream.
- *
- * \param device_type The device type.
- * \param device_id The device id.
- * \param out The new stream handle.
- * \return 0 when success, nonzero when failure happens
- */
-TVM_DLL int TVMStreamCreate(int device_type, int device_id, TVMStreamHandle* out);
-
-/*!
- * \brief Free a created stream handle.
- *
- * \param device_type The device type.
- * \param device_id The device id.
- * \param stream The stream to be freed.
- * \return 0 when success, nonzero when failure happens
- */
-TVM_DLL int TVMStreamFree(int device_type, int device_id, TVMStreamHandle stream);
-
-/*!
- * \brief Set the runtime stream of current thread to be stream.
- * The subsequent calls to the same device_type
- * will use the setted stream handle.
- * The specific type of stream is runtime device dependent.
- *
- * \param device_type The device type.
- * \param device_id The device id.
- * \param handle The stream handle.
- * \return 0 when success, nonzero when failure happens
- */
-TVM_DLL int TVMSetStream(int device_type, int device_id, TVMStreamHandle handle);
-
-/*!
- * \brief Wait until all computations on stream completes.
- *
- * \param device_type The device type.
- * \param device_id The device id.
- * \param stream The stream to be synchronized.
- * \return 0 when success, nonzero when failure happens
- */
-TVM_DLL int TVMSynchronize(int device_type, int device_id, TVMStreamHandle stream);
-
-/*!
- * \brief Synchronize two streams of execution.
- *
- * \param device_type The device type.
- * \param device_id The device id.
- * \param src The source stream to synchronize.
- * \param dst The destination stream to synchronize.
- * \return 0 when success, nonzero when failure happens
- */
-TVM_DLL int TVMStreamStreamSynchronize(int device_type, int device_id, TVMStreamHandle src,
- TVMStreamHandle dst);
-
-/*!
- * \brief Get the type_index from an object.
- *
- * \param obj The object handle.
- * \param out_tindex the output type index.
- * \return 0 when success, nonzero when failure happens
- */
-TVM_DLL int TVMObjectGetTypeIndex(TVMObjectHandle obj, unsigned* out_tindex);
-
-/*!
- * \brief Convert type key to type index.
- * \param type_key The key of the type.
- * \param out_tindex the corresponding type index.
- * \return 0 when success, nonzero when failure happens
- */
-TVM_DLL int TVMObjectTypeKey2Index(const char* type_key, unsigned* out_tindex);
-
-/*!
- * \brief Convert type index to type key.
- * \param tindex The type index.
- * \param out_type_key The output type key.
- * \return 0 when success, nonzero when failure happens
- */
-TVM_DLL int TVMObjectTypeIndex2Key(unsigned tindex, char** out_type_key);
-
-/*!
- * \brief Increase the reference count of an object.
- *
- * \param obj The object handle.
- * \note Internally we increase the reference counter of the object.
- * \return 0 when success, nonzero when failure happens
- */
-TVM_DLL int TVMObjectRetain(TVMObjectHandle obj);
-
-/*!
- * \brief Free the object.
- *
- * \param obj The object handle.
- * \note Internally we decrease the reference counter of the object.
- * The object will be freed when every reference to the object are removed.
- * \return 0 when success, nonzero when failure happens
- */
-TVM_DLL int TVMObjectFree(TVMObjectHandle obj);
-
-/*!
- * \brief Free a TVMByteArray returned from TVMFuncCall, and associated memory.
- * \param arr The TVMByteArray instance.
- * \return 0 on success, -1 on failure.
- */
-TVM_DLL int TVMByteArrayFree(TVMByteArray* arr);
-
-/*!
- * \brief Allocate a data space on device.
- * \param dev The device to perform operation.
- * \param nbytes The number of bytes in memory.
- * \param alignment The alignment of the memory.
- * \param type_hint The type of elements. Only needed by certain backends such
- * as nbytes & alignment are sufficient for most backends.
- * \param out_data The allocated device pointer.
- * \return 0 when success, nonzero when failure happens
- */
-TVM_DLL int TVMDeviceAllocDataSpace(DLDevice dev, size_t nbytes, size_t alignment,
- DLDataType type_hint, void** out_data);
-
-/*!
- * \brief Allocate a data space on device with special memory scope.
- * \note The memory could use a special multi-dimensional memory layout.
- * That is why we pass shape and dtype instead of raw number of bytes.
- * \param dev The device to perform operation.
- * \param ndim The number of dimension of the tensor.
- * \param shape The shape of the tensor.
- * \param dtype The type of elements.
- * \param mem_scope The memory scope of the tensor,
- * can be nullptr, which indicate the default global DRAM
- * \param out_data The allocated device pointer.
- * \return 0 when success, nonzero when failure happens
- */
-TVM_DLL int TVMDeviceAllocDataSpaceWithScope(DLDevice dev, int ndim, const int64_t* shape,
- DLDataType dtype, const char* mem_scope,
- void** out_data);
-
-/*!
- * \brief Free a data space on device.
- * \param dev The device to perform operation.
- * \param ptr The data space.
- * \return 0 when success, nonzero when failure happens
- */
-TVM_DLL int TVMDeviceFreeDataSpace(DLDevice dev, void* ptr);
-
-/*!
- * \brief Copy data from one place to another.
- * \note This API is designed to support special memory with shape dependent layout.
- * We pass in DLTensor* with shape information to support these cases.
- * \param from The source tensor.
- * \param to The target tensor.
- * \param stream Optional stream object.
- * \return 0 when success, nonzero when failure happens.
- */
-TVM_DLL int TVMDeviceCopyDataFromTo(DLTensor* from, DLTensor* to, TVMStreamHandle stream);
-
-/*!
- * \brief Check that an object is derived from another.
- * \param child_type_index The type index of the derived type.
- * \param parent_type_index The type index of the parent type.
- * \param is_derived A boolean representing whether this predicate holds.
- * \return 0 when success, nonzero when failure happens.
- */
-TVM_DLL int TVMObjectDerivedFrom(uint32_t child_type_index, uint32_t parent_type_index,
- int* is_derived);
-
-#ifdef __cplusplus
-} // TVM_EXTERN_C
-#endif
-#endif // TVM_RUNTIME_C_RUNTIME_API_H_
diff --git a/include/tvm/runtime/data_type.h b/include/tvm/runtime/data_type.h
index b5b6ad19d2..d5f3c6ee3d 100644
--- a/include/tvm/runtime/data_type.h
+++ b/include/tvm/runtime/data_type.h
@@ -24,8 +24,9 @@
#ifndef TVM_RUNTIME_DATA_TYPE_H_
#define TVM_RUNTIME_DATA_TYPE_H_
+#include <tvm/ffi/container/shape.h>
#include <tvm/ffi/dtype.h>
-#include <tvm/runtime/c_runtime_api.h>
+#include <tvm/runtime/base.h>
#include <tvm/runtime/logging.h>
#include <cstring>
@@ -35,6 +36,8 @@
namespace tvm {
namespace runtime {
+using tvm_index_t = ffi::Shape::index_type;
+
/*!
* \brief Runtime primitive data type.
*
diff --git a/include/tvm/runtime/device_api.h b/include/tvm/runtime/device_api.h
index a4b53eb797..7366b9895d 100644
--- a/include/tvm/runtime/device_api.h
+++ b/include/tvm/runtime/device_api.h
@@ -26,10 +26,15 @@
#include <tvm/ffi/any.h>
#include <tvm/ffi/optional.h>
-#include <tvm/runtime/c_runtime_api.h>
+#include <tvm/runtime/base.h>
#include <tvm/runtime/logging.h>
#include <string>
+/*!
+ * \brief The stream that is specific to device
+ * can be NULL, which indicates the default one.
+ */
+typedef void* TVMStreamHandle;
namespace tvm {
@@ -37,6 +42,41 @@ namespace tvm {
using Device = DLDevice;
namespace runtime {
+
+/*! \brief Extension device types in TVM
+ *
+ * Additional enumerators to supplement those provided by
+ * DLPack's `DLDeviceType` enumeration.
+ *
+ * MAINTAINERS NOTE #1: We need to ensure that the two devices
+ * are identified by the same integer.
+ * Currently this requires manual verification.
+ * Discussed here: https://github.com/dmlc/dlpack/issues/111
+ * As of DLPack v0.7, the highest-valued enumerator in
+ * `DLDeviceType` is kDLHexagon = 16.
+ *
+ * MAINTAINERS NOTE #2: As of DLPack v0.7, the definition for
+ * `DLDeviceType` specifies an underlying storage type of
+ * `int32_t`. That guarantees a variable of type
+ * `DLDeviceType` is capable of holding any integers provided
+ * by *either* of these enumerations.
+ *
+ * However, the `int32_t` specification only applies when the
+ * header file is compiled as C++, and this header file is also
+ * meant to work as C code. So the unspecified storage type
+ * could be a latent bug when compiled as C.
+ */
+#ifdef __cplusplus
+typedef enum : int32_t {
+#else
+typedef enum {
+#endif
+ // To help avoid accidental conflicts between `DLDeviceType`
+ // and this enumeration, start numbering the new enumerators
+ // a little higher than (currently) seems necessary.
+ TVMDeviceExtType_End = 36, // sentinel value
+} TVMDeviceExtType;
+
/*!
* \brief the query type into GetAttr
*/
diff --git a/include/tvm/runtime/disco/cuda_ipc_memory.h b/include/tvm/runtime/disco/cuda_ipc_memory.h
index 120e6a5431..ea27205262 100644
--- a/include/tvm/runtime/disco/cuda_ipc_memory.h
+++ b/include/tvm/runtime/disco/cuda_ipc_memory.h
@@ -20,7 +20,7 @@
#ifndef TVM_RUNTIME_DISCO_CUDA_IPC_MEMORY_H_
#define TVM_RUNTIME_DISCO_CUDA_IPC_MEMORY_H_
-#include <tvm/runtime/c_runtime_api.h>
+#include <tvm/runtime/base.h>
#include <tvm/runtime/memory/memory_manager.h>
#include <tvm/runtime/object.h>
diff --git a/include/tvm/runtime/logging.h b/include/tvm/runtime/logging.h
index 807c9dbf30..da715848e0 100644
--- a/include/tvm/runtime/logging.h
+++ b/include/tvm/runtime/logging.h
@@ -32,7 +32,7 @@
#include <dmlc/common.h>
#include <dmlc/thread_local.h>
#include <tvm/ffi/error.h>
-#include <tvm/runtime/c_runtime_api.h>
+#include <tvm/runtime/base.h>
#include <ctime>
#include <iomanip>
diff --git a/include/tvm/runtime/memory/memory_manager.h b/include/tvm/runtime/memory/memory_manager.h
index 7b19cdb1ce..f103c6f30a 100644
--- a/include/tvm/runtime/memory/memory_manager.h
+++ b/include/tvm/runtime/memory/memory_manager.h
@@ -24,7 +24,7 @@
#ifndef TVM_RUNTIME_MEMORY_MEMORY_MANAGER_H_
#define TVM_RUNTIME_MEMORY_MEMORY_MANAGER_H_
-#include <tvm/runtime/c_runtime_api.h>
+#include <tvm/runtime/base.h>
#include <tvm/runtime/ndarray.h>
#include <tvm/runtime/object.h>
diff --git a/include/tvm/runtime/module.h b/include/tvm/runtime/module.h
index a4b192db8c..c02e312b71 100644
--- a/include/tvm/runtime/module.h
+++ b/include/tvm/runtime/module.h
@@ -30,7 +30,7 @@
#include <tvm/ffi/function.h>
#include <tvm/ffi/memory.h>
#include <tvm/ffi/string.h>
-#include <tvm/runtime/c_runtime_api.h>
+#include <tvm/runtime/base.h>
#include <tvm/runtime/object.h>
#include <memory>
diff --git a/include/tvm/runtime/ndarray.h b/include/tvm/runtime/ndarray.h
index 2b36306e59..6eebe49ff1 100644
--- a/include/tvm/runtime/ndarray.h
+++ b/include/tvm/runtime/ndarray.h
@@ -28,7 +28,7 @@
#include <tvm/ffi/container/shape.h>
#include <tvm/ffi/optional.h>
#include <tvm/ffi/string.h>
-#include <tvm/runtime/c_runtime_api.h>
+#include <tvm/runtime/base.h>
#include <tvm/runtime/data_type.h>
#include <tvm/runtime/device_api.h>
#include <tvm/runtime/object.h>
@@ -166,20 +166,15 @@ class NDArray : public tvm::ffi::NDArray {
TVM_DLL static void CopyFromTo(const DLTensor* from, DLTensor* to,
TVMStreamHandle stream = nullptr);
- struct Internal;
-
- protected:
- /*!
- * \brief DecRef resource managed by an FFI array handle.
- * \param handle The array handle.
- */
- inline static void FFIDecRef(TVMArrayHandle handle);
/*!
- * \brief Get FFI Array handle from ndarray.
- * \param nd The object with ndarray type.
- * \return The result array handle.
+ * \brief Function to copy data from one array to a byte buffer.
+ * \param from The source array.
+ * \param to The target byte buffer.
+ * \param nbytes The size of the data buffer.
+ * \param stream The stream used in copy.
*/
- inline static TVMArrayHandle FFIGetHandle(const ObjectRef& nd);
+ TVM_DLL static void CopyToBytes(const DLTensor* from, void* to, size_t nbytes,
+ TVMStreamHandle stream = nullptr);
};
/*!
@@ -211,28 +206,6 @@ inline void NDArray::CopyTo(const NDArray& other) const {
CopyFromTo(get_mutable(), other.get_mutable());
}
-inline TVMArrayHandle NDArray::FFIGetHandle(const ObjectRef& nd) {
- // NOTE: it is necessary to cast to container then to base
- // so that the FFI handle uses the ContainerBase address.
- auto ptr = reinterpret_cast<TVMArrayHandle>(
- TVMFFINDArrayGetDLTensorPtr(static_cast<ffi::NDArrayObj*>(const_cast<Object*>(nd.get()))));
- return ptr;
-}
-
-inline TVMArrayHandle ObjectHandleToTVMArrayHandle(Object* handle) {
- return reinterpret_cast<TVMArrayHandle>(
- TVMFFINDArrayGetDLTensorPtr(static_cast<ffi::NDArrayObj*>(handle)));
-}
-
-inline Object* TVMArrayHandleToObjectHandle(void* handle) {
- // NOTE: legacy patch here for TFM FFI
- return reinterpret_cast<ffi::NDArrayObj*>(reinterpret_cast<char*>(handle) - sizeof(TVMFFIObject));
-}
-
-inline void NDArray::FFIDecRef(TVMArrayHandle handle) {
- ffi::details::ObjectUnsafe::DecRefObjectHandle(TVMArrayHandleToObjectHandle(handle));
-}
-
/*! \brief Magic number for NDArray file */
constexpr uint64_t kTVMNDArrayMagic = 0xDD5E40F096B4A13F;
@@ -271,10 +244,7 @@ inline bool SaveDLTensor(dmlc::Stream* strm, const DLTensor* tensor) {
strm->Write(tensor->data, data_byte_size);
} else {
std::vector<uint8_t> bytes(data_byte_size);
- ICHECK_EQ(
- TVMArrayCopyToBytes(const_cast<DLTensor*>(tensor), dmlc::BeginPtr(bytes), data_byte_size),
- 0)
- << TVMGetLastError();
+ NDArray::CopyToBytes(const_cast<DLTensor*>(tensor), dmlc::BeginPtr(bytes), data_byte_size);
if (!DMLC_IO_NO_ENDIAN_SWAP) {
dmlc::ByteSwap(dmlc::BeginPtr(bytes), type_bytes, num_elems);
}
diff --git a/include/tvm/runtime/nvtx.h b/include/tvm/runtime/nvtx.h
index db99154b0b..289837c1fd 100644
--- a/include/tvm/runtime/nvtx.h
+++ b/include/tvm/runtime/nvtx.h
@@ -19,7 +19,7 @@
#ifndef TVM_RUNTIME_NVTX_H_
#define TVM_RUNTIME_NVTX_H_
-#include <tvm/runtime/c_runtime_api.h>
+#include <tvm/runtime/base.h>
#include <string>
namespace tvm {
diff --git a/include/tvm/runtime/object.h b/include/tvm/runtime/object.h
index c686b54d09..6ce95eea1e 100644
--- a/include/tvm/runtime/object.h
+++ b/include/tvm/runtime/object.h
@@ -26,7 +26,7 @@
#include <tvm/ffi/cast.h>
#include <tvm/ffi/object.h>
#include <tvm/ffi/optional.h>
-#include <tvm/runtime/c_runtime_api.h>
+#include <tvm/runtime/base.h>
#include <utility>
diff --git a/include/tvm/runtime/packed_func.h b/include/tvm/runtime/packed_func.h
index 235bdcf3e3..6da06c1191 100644
--- a/include/tvm/runtime/packed_func.h
+++ b/include/tvm/runtime/packed_func.h
@@ -26,7 +26,7 @@
#include <tvm/ffi/any.h>
#include <tvm/ffi/function.h>
-#include <tvm/runtime/c_runtime_api.h>
+#include <tvm/runtime/base.h>
#include <tvm/runtime/logging.h>
#include <tvm/runtime/module.h>
#include <tvm/runtime/ndarray.h>
@@ -41,286 +41,6 @@ namespace runtime {
using ffi::Any;
using ffi::AnyView;
-/*!
- * \brief Utility function to convert legacy ffi::AnyView to AnyView
- * \note This routine is not fastest, but serves purpose to do transition of ABI.
- */
-inline TVMFFIAny LegacyTVMArgValueToFFIAny(TVMValue value, int type_code) {
- TVMFFIAny res;
- // clear first to ensure consistent hash
- res.v_uint64 = 0;
- switch (type_code) {
- case kTVMArgInt: {
- res.type_index = ffi::TypeIndex::kTVMFFIInt;
- res.v_int64 = value.v_int64;
- return res;
- }
- case kTVMArgFloat: {
- res.type_index = ffi::TypeIndex::kTVMFFIFloat;
- res.v_float64 = value.v_float64;
- return res;
- }
- case kTVMOpaqueHandle: {
- res.type_index = ffi::TypeIndex::kTVMFFIOpaquePtr;
- res.v_ptr = value.v_handle;
- return res;
- }
- case kTVMNullptr: {
- res.type_index = ffi::TypeIndex::kTVMFFINone;
- return res;
- }
- case kTVMDataType: {
- res.type_index = ffi::TypeIndex::kTVMFFIDataType;
- res.v_dtype = value.v_type;
- return res;
- }
- case kDLDevice: {
- res.type_index = ffi::TypeIndex::kTVMFFIDevice;
- res.v_device = value.v_device;
- return res;
- }
- case kTVMDLTensorHandle: {
- res.type_index = ffi::TypeIndex::kTVMFFIDLTensorPtr;
- res.v_ptr = value.v_handle;
- return res;
- }
- case kTVMObjectHandle: {
- res.v_obj = static_cast<TVMFFIObject*>(value.v_handle);
- res.type_index = res.v_obj->type_index;
- return res;
- }
- case kTVMModuleHandle: {
- res.type_index = ffi::TypeIndex::kTVMFFIModule;
- res.v_obj = static_cast<TVMFFIObject*>(value.v_handle);
- return res;
- }
- case kTVMPackedFuncHandle: {
- res.type_index = ffi::TypeIndex::kTVMFFIFunction;
- res.v_obj = static_cast<TVMFFIObject*>(value.v_handle);
- return res;
- }
- case kTVMStr: {
- res.type_index = ffi::TypeIndex::kTVMFFIRawStr;
- res.v_c_str = value.v_str;
- return res;
- }
- case kTVMBytes: {
- res.type_index = ffi::TypeIndex::kTVMFFIByteArrayPtr;
- res.v_ptr = value.v_handle;
- return res;
- }
- case kTVMNDArrayHandle: {
- res.type_index = ffi::TypeIndex::kTVMFFINDArray;
- res.v_obj = reinterpret_cast<TVMFFIObject*>(TVMArrayHandleToObjectHandle(value.v_handle));
- return res;
- }
- case kTVMArgBool: {
- res.type_index = ffi::TypeIndex::kTVMFFIBool;
- res.v_int64 = value.v_int64;
- return res;
- }
- case kTVMObjectRValueRefArg: {
- res.type_index = ffi::TypeIndex::kTVMFFIObjectRValueRef;
- res.v_ptr = value.v_handle;
- return res;
- }
- default: {
- LOG(FATAL) << "Unsupported type code: " << type_code;
- TVM_FFI_UNREACHABLE();
- }
- }
-}
-
-/*!
- * \brief Utility function to convert legacy ffi::AnyView to AnyView
- * \note This routine is not fastest, but serves purpose to do transition of ABI.
- */
-inline AnyView LegacyTVMArgValueToAnyView(TVMValue value, int type_code) {
- return AnyView::CopyFromTVMFFIAny(LegacyTVMArgValueToFFIAny(value, type_code));
-}
-
-/*!
- * \brief Utility function to convert legacy ffi::AnyView to Any
- * \note This routine is not fastest, but serves purpose to do transition of ABI.
- */
-inline Any MoveLegacyTVMArgValueToAny(TVMValue value, int type_code) {
- return ffi::details::AnyUnsafe::MoveTVMFFIAnyToAny(LegacyTVMArgValueToFFIAny(value, type_code));
-}
-
-/*
- * \brief Convert AnyView to legacy TVMValue and type_code
- * \param src The AnyView to convert
- * \param value The TVMValue to store the result
- * \param type_code The type code to store the result
- * \note This routine is not fastest, but serves purpose to do transition of ABI.
- */
-inline void AnyViewToLegacyTVMArgValue(TVMFFIAny src, TVMValue* value, int* type_code) {
- switch (src.type_index) {
- case ffi::TypeIndex::kTVMFFIBool: {
- type_code[0] = kTVMArgBool;
- value[0].v_int64 = src.v_int64;
- break;
- }
- case ffi::TypeIndex::kTVMFFIInt: {
- type_code[0] = kDLInt;
- value[0].v_int64 = src.v_int64;
- break;
- }
- case ffi::TypeIndex::kTVMFFIFloat: {
- type_code[0] = kDLFloat;
- value[0].v_float64 = src.v_float64;
- break;
- }
- case ffi::TypeIndex::kTVMFFIOpaquePtr: {
- type_code[0] = kTVMOpaqueHandle;
- value[0].v_handle = src.v_ptr;
- break;
- }
- case ffi::TypeIndex::kTVMFFINone: {
- type_code[0] = kTVMNullptr;
- break;
- }
- case ffi::TypeIndex::kTVMFFIDataType: {
- type_code[0] = kTVMDataType;
- value[0].v_type = src.v_dtype;
- break;
- }
- case ffi::TypeIndex::kTVMFFIDevice: {
- type_code[0] = kDLDevice;
- value[0].v_device = src.v_device;
- break;
- }
- case ffi::TypeIndex::kTVMFFIDLTensorPtr: {
- type_code[0] = kTVMDLTensorHandle;
- value[0].v_handle = src.v_ptr;
- break;
- }
- case ffi::TypeIndex::kTVMFFIRawStr: {
- type_code[0] = kTVMStr;
- value[0].v_str = src.v_c_str;
- break;
- }
- case ffi::TypeIndex::kTVMFFIByteArrayPtr: {
- type_code[0] = kTVMBytes;
- value[0].v_handle = src.v_ptr;
- break;
- }
- case ffi::TypeIndex::kTVMFFINDArray: {
- type_code[0] = kTVMNDArrayHandle;
- value[0].v_handle = ObjectHandleToTVMArrayHandle(reinterpret_cast<Object*>(src.v_obj));
- break;
- }
- case ffi::TypeIndex::kTVMFFIModule: {
- type_code[0] = kTVMModuleHandle;
- value[0].v_handle = src.v_obj;
- break;
- }
- case ffi::TypeIndex::kTVMFFIFunction: {
- type_code[0] = kTVMPackedFuncHandle;
- value[0].v_handle = src.v_obj;
- break;
- }
- case ffi::TypeIndex::kTVMFFIObjectRValueRef: {
- type_code[0] = kTVMObjectRValueRefArg;
- value[0].v_handle = src.v_ptr;
- break;
- }
- default: {
- if (src.type_index >= ffi::TypeIndex::kTVMFFIStaticObjectBegin) {
- type_code[0] = kTVMObjectHandle;
- value[0].v_handle = src.v_obj;
- break;
- }
- LOG(FATAL) << "Unsupported type index: " << src.type_index;
- }
- }
-}
-
-/*
- * \brief Move Any to legacy TVMValue and type_code
- * \param src The Any to move
- * \param value The TVMValue to store the result
- * \param type_code The type code to store the result
- */
-inline void MoveAnyToLegacyTVMValue(Any&& src, TVMValue* value, int* type_code) {
- TVMFFIAny val = ffi::details::AnyUnsafe::MoveAnyToTVMFFIAny(std::move(src));
- // NOTE: conversion rule is the same as AnyViewToLegacyTVMArgValue
- AnyViewToLegacyTVMArgValue(val, value, type_code);
-}
-
-/*!
- * \brief Translate legacy ffi::PackedArgs to PackedArgs
- * \param value The TVMValue array
- * \param type_code The type code array
- * \param num_args The number of arguments
- * \param dst The destination AnyView array
- */
-inline void LegacyTVMArgsToPackedArgs(const TVMValue* value, const int* type_code, int num_args,
- AnyView* dst) {
- for (int i = 0; i < num_args; ++i) {
- dst[i] = LegacyTVMArgValueToAnyView(value[i], type_code[i]);
- }
-}
-
-/*!
- * \brief Translate legacy ffi::PackedArgs to PackedArgs
- * \param args The AnyView array
- * \param num_args The number of arguments
- * \param value The TVMValue array
- * \param type_code The type code array
- */
-inline void PackedArgsToLegacyTVMArgs(const AnyView* args, int num_args, TVMValue* value,
- int* type_code) {
- for (int i = 0; i < num_args; ++i) {
- AnyViewToLegacyTVMArgValue(args[i].CopyToTVMFFIAny(), value + i, type_code + i);
- }
-}
-
-/*!
- * \brief Convert argument type code to string.
- * \param type_code The input type code.
- * \return The corresponding string repr.
- */
-inline const char* ArgTypeCode2Str(int type_code) {
- switch (type_code) {
- case kDLInt:
- return "int";
- case kTVMArgBool:
- return "bool";
- case kDLUInt:
- return "uint";
- case kDLFloat:
- return "float";
- case kTVMStr:
- return "str";
- case kTVMBytes:
- return "bytes";
- case kTVMOpaqueHandle:
- return "handle";
- case kTVMNullptr:
- return "NULL";
- case kTVMDLTensorHandle:
- return "ArrayHandle";
- case kTVMDataType:
- return "DLDataType";
- case kDLDevice:
- return "DLDevice";
- case kTVMPackedFuncHandle:
- return "FunctionHandle";
- case kTVMModuleHandle:
- return "ModuleHandle";
- case kTVMNDArrayHandle:
- return "NDArrayContainer";
- case kTVMObjectHandle:
- return "Object";
- case kTVMObjectRValueRefArg:
- return "ObjectRValueRefArg";
- default:
- LOG(FATAL) << "unknown type_code=" << static_cast<int>(type_code);
- }
- throw;
-}
-
namespace details {
template <typename T>
diff --git a/include/tvm/runtime/profiling.h b/include/tvm/runtime/profiling.h
index 63814a6375..1e950f9f7a 100644
--- a/include/tvm/runtime/profiling.h
+++ b/include/tvm/runtime/profiling.h
@@ -26,11 +26,11 @@
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/container/map.h>
-#include <tvm/runtime/c_runtime_api.h>
+#include <tvm/ffi/function.h>
+#include <tvm/runtime/base.h>
#include <tvm/runtime/device_api.h>
#include <tvm/runtime/object.h>
#include <tvm/runtime/packed_func.h>
-#include <tvm/runtime/registry.h>
#include <stack>
#include <string>
@@ -133,7 +133,7 @@ class Timer : public ObjectRef {
* };
* TVM_REGISTER_OBJECT_TYPE(CPUTimerNode);
*
- * TVM_REGISTER_GLOBAL("profiling.timer.cpu").set_body_typed([](Device dev) {
+ * TVM_FFI_REGISTER_GLOBAL("profiling.timer.cpu").set_body_typed([](Device dev) {
* return Timer(make_object<CPUTimerNode>());
* });
* \endcode
diff --git a/include/tvm/runtime/registry.h b/include/tvm/runtime/registry.h
deleted file mode 100644
index 94463dc725..0000000000
--- a/include/tvm/runtime/registry.h
+++ /dev/null
@@ -1,102 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-/*!
- * \file tvm/runtime/registry.h
- * \brief This file defines the TVM global function registry.
- */
-#ifndef TVM_RUNTIME_REGISTRY_H_
-#define TVM_RUNTIME_REGISTRY_H_
-
-#include <tvm/ffi/string.h>
-#include <tvm/runtime/packed_func.h>
-
-#include <type_traits>
-#include <utility>
-#include <vector>
-
-namespace tvm {
-namespace runtime {
-
-/*! \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.
- * Allows for the creation of strong references to Python objects, which prevent them from being
- * garbage-collected as long as the wrapper object exists.
- */
-class WrappedPythonObject {
- public:
- /*! \brief Construct a wrapper that doesn't own anything */
- WrappedPythonObject() : python_obj_(nullptr) {}
-
- /*! \brief Conversion constructor from nullptr */
- explicit WrappedPythonObject(std::nullptr_t) : python_obj_(nullptr) {}
-
- /*! \brief Take ownership of a python object
- *
- * A new strong reference is created for the underlying python
- * object.
- *
- * \param python_obj A PyObject* from the Python.h API. A new
- * strong reference is created using Py_IncRef.
- */
- explicit WrappedPythonObject(void* python_obj);
-
- /*! \brief Drop ownership of a python object
- *
- * Removes the strong reference held by the wrapper.
- */
- ~WrappedPythonObject();
-
- WrappedPythonObject(WrappedPythonObject&&);
- WrappedPythonObject& operator=(WrappedPythonObject&&);
-
- WrappedPythonObject(const WrappedPythonObject&);
- WrappedPythonObject& operator=(const WrappedPythonObject&);
- WrappedPythonObject& operator=(std::nullptr_t);
-
- operator bool() { return python_obj_; }
-
- void* raw_pointer() { return python_obj_; }
-
- private:
- void* python_obj_ = nullptr;
-};
-
-/*!
- * \brief Register a function globally.
- * \code
- * TVM_REGISTER_GLOBAL("MyPrint")
- * .set_body_packed([](ffi::PackedArgs args, ffi::Any* rv) {
- * });
- * \endcode
- */
-#define TVM_REGISTER_GLOBAL TVM_FFI_REGISTER_GLOBAL
-
-#define TVM_STRINGIZE_DETAIL(x) #x
-#define TVM_STRINGIZE(x) TVM_STRINGIZE_DETAIL(x)
-#define TVM_DESCRIBE(...) describe(__VA_ARGS__ "\n\nFrom:" __FILE__ ":" TVM_STRINGIZE(__LINE__))
-/*!
- * \brief Macro to include current line as string
- */
-#define TVM_ADD_FILELINE "\n\nDefined in " __FILE__ ":L" TVM_STRINGIZE(__LINE__)
-
-} // namespace runtime
-} // namespace tvm
-#endif // TVM_RUNTIME_REGISTRY_H_
diff --git a/include/tvm/runtime/relax_vm/executable.h b/include/tvm/runtime/relax_vm/executable.h
index dc9d870253..afaaea9e41 100644
--- a/include/tvm/runtime/relax_vm/executable.h
+++ b/include/tvm/runtime/relax_vm/executable.h
@@ -23,9 +23,9 @@
#ifndef TVM_RUNTIME_RELAX_VM_EXECUTABLE_H_
#define TVM_RUNTIME_RELAX_VM_EXECUTABLE_H_
+#include <tvm/ffi/function.h>
#include <tvm/runtime/object.h>
#include <tvm/runtime/packed_func.h>
-#include <tvm/runtime/registry.h>
#include <string>
#include <unordered_map>
diff --git a/include/tvm/runtime/relax_vm/ndarray_cache_support.h b/include/tvm/runtime/relax_vm/ndarray_cache_support.h
index f595d81ffe..579fbf306f 100644
--- a/include/tvm/runtime/relax_vm/ndarray_cache_support.h
+++ b/include/tvm/runtime/relax_vm/ndarray_cache_support.h
@@ -20,8 +20,8 @@
#define TVM_RUNTIME_RELAX_VM_NDARRAY_CACHE_SUPPORT_H_
#include <tvm/ffi/container/array.h>
+#include <tvm/ffi/function.h>
#include <tvm/runtime/ndarray.h>
-#include <tvm/runtime/registry.h>
#include <string>
#include <unordered_map>
diff --git a/include/tvm/runtime/serializer.h b/include/tvm/runtime/serializer.h
index b35cad3688..2cfd1de44d 100644
--- a/include/tvm/runtime/serializer.h
+++ b/include/tvm/runtime/serializer.h
@@ -27,7 +27,7 @@
#include <dmlc/io.h>
#include <dmlc/serializer.h>
-#include <tvm/runtime/c_runtime_api.h>
+#include <tvm/runtime/base.h>
#include <tvm/runtime/ndarray.h>
namespace dmlc {
diff --git a/include/tvm/support/parallel_for.h b/include/tvm/support/parallel_for.h
index 8bd2e6b825..aa9da30d8f 100644
--- a/include/tvm/support/parallel_for.h
+++ b/include/tvm/support/parallel_for.h
@@ -24,7 +24,7 @@
#ifndef TVM_SUPPORT_PARALLEL_FOR_H_
#define TVM_SUPPORT_PARALLEL_FOR_H_
-#include <tvm/runtime/c_runtime_api.h>
+#include <tvm/runtime/base.h>
#include <functional>
#include <vector>
diff --git a/include/tvm/tir/builtin.h b/include/tvm/tir/builtin.h
index 2035a511c1..c057422a02 100644
--- a/include/tvm/tir/builtin.h
+++ b/include/tvm/tir/builtin.h
@@ -415,17 +415,15 @@ TVM_DLL const Op& tvm_thread_invariant();
* type codes are explicitly allocated.
*
* return_type tvm_call_packed_lowered(name,
- * TVMValue* value_stack,
- * int* tcode_stack,
+ * TVMFFIAny* args_stack,
* int begin,
* int end) {
* ModuleNode* env = GetCurrentEnv();
* const ffi::Function* f = env->GetFuncFromEnv(name);
- * f->CallPacked(ffi::PackedArgs(value_stack[begin:end],
- * tcode_stack[begin:end]),
- * ffi::Any(value_stack + end, tcode_stack + end));
+ * f->CallPacked(ffi::PackedArgs(args_stack[begin:end]),
+ * ffi::Any(args_stack + end));
* // return type can be int, float, handle.
- * return cast(return_type, load_return_from(tcode_stack + end))
+ * return cast(return_type, load_return_from(args_stack + end))
* }
*/
TVM_DLL const Op& tvm_call_packed_lowered();
@@ -451,17 +449,15 @@ TVM_DLL const Op& tvm_call_cpacked_lowered();
* (end - 1) value on the stack.
*
* return_type tvm_call_trace_packed_lowered(name,
- * TVMValue* value_stack,
- * int* tcode_stack,
+ * TVMFFIAny* args_stack,
* int begin,
* int end) {
* ModuleNode* env = GetCurrentEnv();
* const ffi::Function* f = env->GetFuncFromEnv(name);
- * f->CallPacked(ffi::PackedArgs(value_stack[begin:end],
- * tcode_stack[begin:end]),
- * ffi::Any(value_stack + end, tcode_stack + end));
+ * f->CallPacked(ffi::PackedArgs(args_stack[begin:end]),
+ * ffi::Any(args_stack + end));
* // return type can be int, float, handle.
- * return cast(return_type, load_return_from(tcode_stack + end))
+ * return cast(return_type, load_return_from(args_stack + end))
* }
*/
TVM_DLL const Op& tvm_call_trace_packed_lowered();
diff --git a/include/tvm/tir/expr.h b/include/tvm/tir/expr.h
index 078b41e327..5f058f7d5e 100644
--- a/include/tvm/tir/expr.h
+++ b/include/tvm/tir/expr.h
@@ -31,7 +31,7 @@
#include <tvm/ir/expr.h>
#include <tvm/node/functor.h>
#include <tvm/node/node.h>
-#include <tvm/runtime/c_runtime_api.h>
+#include <tvm/runtime/base.h>
#include <tvm/runtime/data_type.h>
#include <tvm/tir/buffer.h>
#include <tvm/tir/var.h>
diff --git a/include/tvm/tir/transform.h b/include/tvm/tir/transform.h
index b80d4456c0..eb64d87f95 100644
--- a/include/tvm/tir/transform.h
+++ b/include/tvm/tir/transform.h
@@ -166,9 +166,9 @@ TVM_DLL Pass InstrumentBoundCheckers();
* f()
*
* if num_packed_args is not zero:
- * f(TVMArg* packed_args, int* packed_arg_type_ids, int num_packed_args,
+ * f(void *, TVMFFIAny* packed_args, int num_packed_args,
* api_arg_k, api_arg_k+1, ... api_arg_n,
- * TVMValue* out_ret_val, int* out_ret_tcode)
+ * TVMFFIAny* out_ret_val)
*
* where n == len(api_args), k == num_packed_args
*
diff --git a/python/tvm/contrib/cutlass/gen_tensor_op.py b/python/tvm/contrib/cutlass/gen_tensor_op.py
index 5d04cf13e6..88581b1cb4 100644
--- a/python/tvm/contrib/cutlass/gen_tensor_op.py
+++ b/python/tvm/contrib/cutlass/gen_tensor_op.py
@@ -487,7 +487,7 @@ def instantiate_template(func_name, annotations, func_args):
if k in annotations:
attrs[k] = annotations[k]
- headers = ["tvm/runtime/registry.h"]
+ headers = ["tvm/ffi/function.h"]
if "relu" in func_name:
headers.append("cutlass/epilogue/thread/linear_combination_bias_relu.h")
diff --git a/python/tvm/relax/frontend/nn/op.py b/python/tvm/relax/frontend/nn/op.py
index f7df8d4417..ab416ef141 100644
--- a/python/tvm/relax/frontend/nn/op.py
+++ b/python/tvm/relax/frontend/nn/op.py
@@ -2087,7 +2087,7 @@ def extern(
out: OutType,
) -> OutType:
"""Invoke an extern function during runtime. The extern function must be registered with the "
- TVM runtime using `TVM_REGISTER_GLOBAL` (C++), or `tvm.register_func` (Python).
+ TVM runtime using `TVM_FFI_REGISTER_GLOBAL` (C++), or `tvm.register_func` (Python).
Parameters
----------
diff --git a/python/tvm/runtime/_ffi_api.py b/python/tvm/runtime/_ffi_api.py
index a07193ea98..f0d1bdeb0a 100644
--- a/python/tvm/runtime/_ffi_api.py
+++ b/python/tvm/runtime/_ffi_api.py
@@ -17,6 +17,6 @@
"""FFI APIs for tvm.runtime"""
import tvm._ffi
-# Exports functions registered via TVM_REGISTER_GLOBAL with the "runtime" prefix.
-# e.g. TVM_REGISTER_GLOBAL("runtime.ModuleLoadFromFile")
+# Exports functions registered via TVM_FFI_REGISTER_GLOBAL with the "runtime" prefix.
+# e.g. TVM_FFI_REGISTER_GLOBAL("runtime.ModuleLoadFromFile")
tvm._ffi._init_api("runtime", __name__)
diff --git a/python/tvm/runtime/_ffi_node_api.py b/python/tvm/runtime/_ffi_node_api.py
index 395496d16b..600206d425 100644
--- a/python/tvm/runtime/_ffi_node_api.py
+++ b/python/tvm/runtime/_ffi_node_api.py
@@ -24,7 +24,7 @@ import tvm.ffi.core
# The implementations below are default ones when the corresponding
# functions are not available in the runtime only mode.
# They will be overriden via _init_api to the ones registered
-# via TVM_REGISTER_GLOBAL in the compiler mode.
+# via TVM_FFI_REGISTER_GLOBAL in the compiler mode.
def AsRepr(obj):
return type(obj).__name__ + "(" + obj.__ctypes_handle__().value + ")"
@@ -45,6 +45,6 @@ def LoadJSON(json_str):
raise RuntimeError("Do not support object serialization in runtime only mode")
-# Exports functions registered via TVM_REGISTER_GLOBAL with the "node" prefix.
-# e.g. TVM_REGISTER_GLOBAL("node.AsRepr")
+# Exports functions registered via TVM_FFI_REGISTER_GLOBAL with the "node" prefix.
+# e.g. TVM_FFI_REGISTER_GLOBAL("node.AsRepr")
tvm._ffi._init_api("node", __name__)
diff --git a/src/arith/analyzer.cc b/src/arith/analyzer.cc
index 8dcb67190a..f0a317659d 100644
--- a/src/arith/analyzer.cc
+++ b/src/arith/analyzer.cc
@@ -21,7 +21,7 @@
* \file tvm/arith/analyzer.cc
*/
#include <tvm/arith/analyzer.h>
-#include <tvm/runtime/registry.h>
+#include <tvm/ffi/function.h>
#include <tvm/tir/expr.h>
#include <tvm/tir/op.h>
@@ -269,7 +269,7 @@ PrimExpr Analyzer::Simplify(const PrimExpr& expr, int steps) {
return res;
}
-TVM_REGISTER_GLOBAL("arith.CreateAnalyzer")
+TVM_FFI_REGISTER_GLOBAL("arith.CreateAnalyzer")
.set_body_packed([](ffi::PackedArgs args, ffi::Any* ret) {
using ffi::Function;
using ffi::TypedFunction;
diff --git a/src/arith/bound_deducer.cc b/src/arith/bound_deducer.cc
index d52ae7e6fd..b8b5d64824 100644
--- a/src/arith/bound_deducer.cc
+++ b/src/arith/bound_deducer.cc
@@ -22,7 +22,7 @@
* \brief Utility to deduce bound of expression
*/
#include <tvm/arith/analyzer.h>
-#include <tvm/runtime/registry.h>
+#include <tvm/ffi/function.h>
#include <tvm/tir/expr.h>
#include <tvm/tir/expr_functor.h>
@@ -402,7 +402,7 @@ IntSet DeduceBound(PrimExpr v, PrimExpr e, const Map<Var, IntSet>& hint_map,
return DeduceBound(v, e, hmap, rmap);
}
-TVM_REGISTER_GLOBAL("arith.DeduceBound")
+TVM_FFI_REGISTER_GLOBAL("arith.DeduceBound")
.set_body_typed([](PrimExpr v, PrimExpr cond, const Map<Var, IntSet> hint_map,
const Map<Var, IntSet> relax_map) {
return DeduceBound(v, cond, hint_map, relax_map);
diff --git a/src/arith/const_int_bound.cc b/src/arith/const_int_bound.cc
index 7409ecc6f3..a440b52074 100644
--- a/src/arith/const_int_bound.cc
+++ b/src/arith/const_int_bound.cc
@@ -21,7 +21,7 @@
* \file tvm/arith/const_int_bound.cc
*/
#include <tvm/arith/analyzer.h>
-#include <tvm/runtime/registry.h>
+#include <tvm/ffi/function.h>
#include <tvm/tir/builtin.h>
#include <tvm/tir/expr_functor.h>
@@ -51,7 +51,7 @@ ConstIntBound MakeConstIntBound(int64_t min_value, int64_t max_value) {
return ConstIntBound(min_value, max_value);
}
-TVM_REGISTER_GLOBAL("arith.ConstIntBound").set_body_typed(MakeConstIntBound);
+TVM_FFI_REGISTER_GLOBAL("arith.ConstIntBound").set_body_typed(MakeConstIntBound);
inline void PrintBoundValue(std::ostream& os, int64_t val) {
if (val == ConstIntBound::kPosInf) {
diff --git a/src/arith/detect_common_subexpr.cc b/src/arith/detect_common_subexpr.cc
index b496e7fefc..303360002e 100644
--- a/src/arith/detect_common_subexpr.cc
+++ b/src/arith/detect_common_subexpr.cc
@@ -69,6 +69,6 @@ Map<PrimExpr, Integer> DetectCommonSubExpr(const PrimExpr& e, int thresh) {
return results;
}
-TVM_REGISTER_GLOBAL("arith.DetectCommonSubExpr").set_body_typed(DetectCommonSubExpr);
+TVM_FFI_REGISTER_GLOBAL("arith.DetectCommonSubExpr").set_body_typed(DetectCommonSubExpr);
} // namespace arith
} // namespace tvm
diff --git a/src/arith/detect_linear_equation.cc b/src/arith/detect_linear_equation.cc
index 4d3164cbd3..0dcbc76235 100644
--- a/src/arith/detect_linear_equation.cc
+++ b/src/arith/detect_linear_equation.cc
@@ -22,7 +22,7 @@
* \brief Utility to detect patterns in the expression.
*/
#include <tvm/arith/analyzer.h>
-#include <tvm/runtime/registry.h>
+#include <tvm/ffi/function.h>
#include <tvm/tir/analysis.h>
#include <tvm/tir/expr.h>
#include <tvm/tir/expr_functor.h>
@@ -290,9 +290,9 @@ Array<PrimExpr> DetectClipBound(const PrimExpr& e, const Array<Var>& vars) {
return ret;
}
-TVM_REGISTER_GLOBAL("arith.DetectLinearEquation").set_body_typed(DetectLinearEquation);
+TVM_FFI_REGISTER_GLOBAL("arith.DetectLinearEquation").set_body_typed(DetectLinearEquation);
-TVM_REGISTER_GLOBAL("arith.DetectClipBound")
+TVM_FFI_REGISTER_GLOBAL("arith.DetectClipBound")
.set_body_typed([](const PrimExpr& e, const Array<Var>& vars) {
return DetectClipBound(e, vars);
});
diff --git a/src/arith/domain_touched.cc b/src/arith/domain_touched.cc
index 13243ddb5e..5f9d780030 100644
--- a/src/arith/domain_touched.cc
+++ b/src/arith/domain_touched.cc
@@ -21,7 +21,7 @@
* \file bound_deducer.cc
* \brief Utility to deduce bound of expression
*/
-#include <tvm/runtime/registry.h>
+#include <tvm/ffi/function.h>
#include <tvm/te/tensor.h>
#include <tvm/tir/expr.h>
#include <tvm/tir/stmt_functor.h>
@@ -162,8 +162,8 @@ Map<Buffer, Array<ObjectRef>> DomainTouchedAccessMap(const PrimFunc& func) {
return ret;
}
-TVM_REGISTER_GLOBAL("arith.DomainTouched").set_body_typed(DomainTouched);
-TVM_REGISTER_GLOBAL("arith.DomainTouchedAccessMap").set_body_typed(DomainTouchedAccessMap);
+TVM_FFI_REGISTER_GLOBAL("arith.DomainTouched").set_body_typed(DomainTouched);
+TVM_FFI_REGISTER_GLOBAL("arith.DomainTouchedAccessMap").set_body_typed(DomainTouchedAccessMap);
} // namespace arith
} // namespace tvm
diff --git a/src/arith/int_constraints.cc b/src/arith/int_constraints.cc
index 8c314992ab..01e7a30969 100644
--- a/src/arith/int_constraints.cc
+++ b/src/arith/int_constraints.cc
@@ -23,7 +23,7 @@
*/
#include <tvm/arith/analyzer.h>
#include <tvm/arith/int_solver.h>
-#include <tvm/runtime/registry.h>
+#include <tvm/ffi/function.h>
#include <tvm/tir/expr.h>
#include <tvm/tir/expr_functor.h>
#include <tvm/tir/op.h>
@@ -195,15 +195,16 @@ Range IntGroupBounds::FindBestRange(const Map<Var, Range>& vranges_addl) const {
TVM_REGISTER_NODE_TYPE(IntGroupBoundsNode);
-TVM_REGISTER_GLOBAL("arith.IntGroupBounds")
+TVM_FFI_REGISTER_GLOBAL("arith.IntGroupBounds")
.set_body_typed([](PrimExpr coef, Array<PrimExpr> lower, Array<PrimExpr> equal,
Array<PrimExpr> upper) {
return IntGroupBounds(coef, lower, equal, upper);
});
-TVM_REGISTER_GLOBAL("arith.IntGroupBounds_from_range").set_body_typed(IntGroupBounds::FromRange);
+TVM_FFI_REGISTER_GLOBAL("arith.IntGroupBounds_from_range")
+ .set_body_typed(IntGroupBounds::FromRange);
-TVM_REGISTER_GLOBAL("arith.IntGroupBounds_FindBestRange")
+TVM_FFI_REGISTER_GLOBAL("arith.IntGroupBounds_FindBestRange")
.set_body_packed([](ffi::PackedArgs args, ffi::Any* ret) {
ICHECK(args.size() == 1 || args.size() == 2);
auto bounds = args[0].cast<IntGroupBounds>();
@@ -243,7 +244,7 @@ IntConstraints::IntConstraints(Array<Var> variables, Map<Var, Range> ranges,
TVM_REGISTER_NODE_TYPE(IntConstraintsNode);
-TVM_REGISTER_GLOBAL("arith.IntConstraints")
+TVM_FFI_REGISTER_GLOBAL("arith.IntConstraints")
.set_body_typed([](Array<Var> variables, Map<Var, Range> ranges, Array<PrimExpr> relations) {
return IntConstraints(variables, ranges, relations);
});
@@ -288,7 +289,7 @@ IntConstraintsTransform IntConstraintsTransform::operator+(
TVM_REGISTER_NODE_TYPE(IntConstraintsTransformNode);
-TVM_REGISTER_GLOBAL("arith.IntConstraintsTransform")
+TVM_FFI_REGISTER_GLOBAL("arith.IntConstraintsTransform")
.set_body_typed([](IntConstraints src, IntConstraints dst, Map<Var, PrimExpr> src_to_dst,
Map<Var, PrimExpr> dst_to_src) {
return IntConstraintsTransform(src, dst, src_to_dst, dst_to_src);
diff --git a/src/arith/int_set.cc b/src/arith/int_set.cc
index 94a2a369a6..d3b7b30628 100644
--- a/src/arith/int_set.cc
+++ b/src/arith/int_set.cc
@@ -23,7 +23,7 @@
*/
#include <tvm/arith/int_set.h>
#include <tvm/arith/iter_affine_map.h>
-#include <tvm/runtime/registry.h>
+#include <tvm/ffi/function.h>
#include <tvm/tir/expr.h>
#include <tvm/tir/expr_functor.h>
@@ -57,7 +57,7 @@ IntervalSet MakeIntervalSet(PrimExpr min_value, PrimExpr max_value) {
return IntervalSet(min_value, max_value);
}
-TVM_REGISTER_GLOBAL("arith.IntervalSet").set_body_typed(MakeIntervalSet);
+TVM_FFI_REGISTER_GLOBAL("arith.IntervalSet").set_body_typed(MakeIntervalSet);
IntervalSet Intersect(Analyzer* analyzer, IntervalSet a, IntervalSet b) {
PrimExpr max_value = min(a->max_value, b->max_value);
@@ -1192,42 +1192,42 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
<< "[" << op->min_value << ", " << op->max_value << ']';
});
-TVM_REGISTER_GLOBAL("arith.intset_single_point").set_body_typed(IntSet::SinglePoint);
+TVM_FFI_REGISTER_GLOBAL("arith.intset_single_point").set_body_typed(IntSet::SinglePoint);
-TVM_REGISTER_GLOBAL("arith.intset_vector").set_body_typed(IntSet::Vector);
+TVM_FFI_REGISTER_GLOBAL("arith.intset_vector").set_body_typed(IntSet::Vector);
-TVM_REGISTER_GLOBAL("arith.intset_interval").set_body_typed(IntSet::Interval);
+TVM_FFI_REGISTER_GLOBAL("arith.intset_interval").set_body_typed(IntSet::Interval);
-TVM_REGISTER_GLOBAL("arith.IntervalSetGetMin").set_body_method(&IntSet::min);
+TVM_FFI_REGISTER_GLOBAL("arith.IntervalSetGetMin").set_body_method(&IntSet::min);
-TVM_REGISTER_GLOBAL("arith.IntervalSetGetMax").set_body_method(&IntSet::max);
+TVM_FFI_REGISTER_GLOBAL("arith.IntervalSetGetMax").set_body_method(&IntSet::max);
-TVM_REGISTER_GLOBAL("arith.IntSetIsNothing").set_body_method(&IntSet::IsNothing);
+TVM_FFI_REGISTER_GLOBAL("arith.IntSetIsNothing").set_body_method(&IntSet::IsNothing);
-TVM_REGISTER_GLOBAL("arith.IntSetIsEverything").set_body_method(&IntSet::IsEverything);
+TVM_FFI_REGISTER_GLOBAL("arith.IntSetIsEverything").set_body_method(&IntSet::IsEverything);
-TVM_REGISTER_GLOBAL("arith.EstimateRegionLowerBound")
+TVM_FFI_REGISTER_GLOBAL("arith.EstimateRegionLowerBound")
.set_body_typed([](Array<Range> region, Map<Var, Range> var_dom,
PrimExpr predicate) -> Optional<Array<IntSet>> {
Analyzer analyzer;
return EstimateRegionLowerBound(region, var_dom, predicate, &analyzer);
});
-TVM_REGISTER_GLOBAL("arith.EstimateRegionStrictBound")
+TVM_FFI_REGISTER_GLOBAL("arith.EstimateRegionStrictBound")
.set_body_typed([](Array<Range> region, Map<Var, Range> var_dom,
PrimExpr predicate) -> Optional<Array<IntSet>> {
Analyzer analyzer;
return EstimateRegionStrictBound(region, var_dom, predicate, &analyzer);
});
-TVM_REGISTER_GLOBAL("arith.EstimateRegionUpperBound")
+TVM_FFI_REGISTER_GLOBAL("arith.EstimateRegionUpperBound")
.set_body_typed([](Array<Range> region, Map<Var, Range> var_dom,
PrimExpr predicate) -> Optional<Array<IntSet>> {
Analyzer analyzer;
return EstimateRegionUpperBound(region, var_dom, predicate, &analyzer);
});
-TVM_REGISTER_GLOBAL("arith.PosInf").set_body_typed([]() { return SymbolicLimits::pos_inf_; });
-TVM_REGISTER_GLOBAL("arith.NegInf").set_body_typed([]() { return SymbolicLimits::neg_inf_; });
-TVM_REGISTER_GLOBAL("arith.UnionLowerBound").set_body_typed(UnionLowerBound);
+TVM_FFI_REGISTER_GLOBAL("arith.PosInf").set_body_typed([]() { return SymbolicLimits::pos_inf_; });
+TVM_FFI_REGISTER_GLOBAL("arith.NegInf").set_body_typed([]() { return SymbolicLimits::neg_inf_; });
+TVM_FFI_REGISTER_GLOBAL("arith.UnionLowerBound").set_body_typed(UnionLowerBound);
} // namespace arith
} // namespace tvm
diff --git a/src/arith/iter_affine_map.cc b/src/arith/iter_affine_map.cc
index 5b1feef4e6..2aa0ca6b64 100644
--- a/src/arith/iter_affine_map.cc
+++ b/src/arith/iter_affine_map.cc
@@ -47,7 +47,7 @@ IterMark::IterMark(PrimExpr source, PrimExpr extent) {
data_ = std::move(n);
}
-TVM_REGISTER_GLOBAL("arith.IterMark").set_body_typed([](PrimExpr source, PrimExpr extent) {
+TVM_FFI_REGISTER_GLOBAL("arith.IterMark").set_body_typed([](PrimExpr source, PrimExpr extent) {
return IterMark(source, extent);
});
@@ -92,7 +92,7 @@ IterSplitExpr::IterSplitExpr(IterMark source, PrimExpr lower_factor, PrimExpr ex
data_ = std::move(n);
}
-TVM_REGISTER_GLOBAL("arith.IterSplitExpr")
+TVM_FFI_REGISTER_GLOBAL("arith.IterSplitExpr")
.set_body_typed([](IterMark source, PrimExpr lower_factor, PrimExpr extent, PrimExpr scale) {
return IterSplitExpr(source, lower_factor, extent, scale);
});
@@ -114,7 +114,7 @@ IterSumExpr::IterSumExpr(Array<IterSplitExpr> args, PrimExpr base) {
data_ = std::move(n);
}
-TVM_REGISTER_GLOBAL("arith.IterSumExpr")
+TVM_FFI_REGISTER_GLOBAL("arith.IterSumExpr")
.set_body_typed([](Array<IterSplitExpr> args, PrimExpr base) {
return IterSumExpr(args, base);
});
@@ -1513,7 +1513,7 @@ IterMapResult DetectIterMap(const Array<PrimExpr>& indices, const Map<Var, Range
return result;
}
-TVM_REGISTER_GLOBAL("arith.DetectIterMap")
+TVM_FFI_REGISTER_GLOBAL("arith.DetectIterMap")
.set_body_typed([](const Array<PrimExpr>& indices, const Map<Var, Range>& input_iters,
const PrimExpr& input_pred, int check_level,
bool simplify_trivial_iterators) {
@@ -1538,7 +1538,7 @@ IterSumExpr NormalizeToIterSum(PrimExpr index, const Map<Var, Range>& input_iter
return rewriter.RewriteToNormalizedIterSum(index);
}
-TVM_REGISTER_GLOBAL("arith.NormalizeToIterSum")
+TVM_FFI_REGISTER_GLOBAL("arith.NormalizeToIterSum")
.set_body_typed([](PrimExpr index, const Map<Var, Range>& input_iters) {
arith::Analyzer ana;
return NormalizeToIterSum(index, input_iters, &ana);
@@ -2133,7 +2133,7 @@ PrimExpr NormalizeIterMapToExpr(const PrimExpr& expr) {
return normalizer.Convert(expr);
}
-TVM_REGISTER_GLOBAL("arith.NormalizeIterMapToExpr").set_body_typed(NormalizeIterMapToExpr);
+TVM_FFI_REGISTER_GLOBAL("arith.NormalizeIterMapToExpr").set_body_typed(NormalizeIterMapToExpr);
Array<PrimExpr> IterMapSimplify(const Array<PrimExpr>& indices, const Map<Var, Range>& input_iters,
const PrimExpr& input_pred, IterMapLevel check_level,
@@ -2162,7 +2162,7 @@ Array<PrimExpr> IterMapSimplify(const Array<PrimExpr>& indices, const Map<Var, R
return simplified;
}
-TVM_REGISTER_GLOBAL("arith.IterMapSimplify")
+TVM_FFI_REGISTER_GLOBAL("arith.IterMapSimplify")
.set_body_typed([](const Array<PrimExpr>& indices, const Map<Var, Range>& input_iters,
const PrimExpr& input_pred, int check_level,
bool simplify_trivial_iterators) {
@@ -2495,7 +2495,7 @@ Array<Array<IterMark>> SubspaceDivide(const Array<PrimExpr>& bindings,
return results;
}
-TVM_REGISTER_GLOBAL("arith.SubspaceDivide")
+TVM_FFI_REGISTER_GLOBAL("arith.SubspaceDivide")
.set_body_typed([](const Array<PrimExpr>& bindings, const Map<Var, Range>& root_iters,
const Array<Var>& sub_iters, const PrimExpr& predicate, int check_level,
bool simplify_trivial_iterators) {
@@ -2634,7 +2634,7 @@ Map<Var, PrimExpr> InverseAffineIterMap(const Array<IterSumExpr>& iter_map,
return InverseAffineIterMapTransformer(&analyzer)(iter_map, outputs);
}
-TVM_REGISTER_GLOBAL("arith.InverseAffineIterMap").set_body_typed(InverseAffineIterMap);
+TVM_FFI_REGISTER_GLOBAL("arith.InverseAffineIterMap").set_body_typed(InverseAffineIterMap);
TVM_REGISTER_NODE_TYPE(IterMapResultNode);
diff --git a/src/arith/modular_set.cc b/src/arith/modular_set.cc
index 197e5ec8b8..fa4891d5a0 100644
--- a/src/arith/modular_set.cc
+++ b/src/arith/modular_set.cc
@@ -22,7 +22,7 @@
* \brief Modular set analysis
*/
#include <tvm/arith/analyzer.h>
-#include <tvm/runtime/registry.h>
+#include <tvm/ffi/function.h>
#include <tvm/tir/builtin.h>
#include <tvm/tir/expr_functor.h>
#include <tvm/tir/op.h>
@@ -57,7 +57,7 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
ModularSet MakeModularSet(int64_t coeff, int64_t base) { return ModularSet(coeff, base); }
-TVM_REGISTER_GLOBAL("arith.ModularSet").set_body_typed(MakeModularSet);
+TVM_FFI_REGISTER_GLOBAL("arith.ModularSet").set_body_typed(MakeModularSet);
// internal entry for const int bound
struct ModularSetAnalyzer::Entry {
diff --git a/src/arith/narrow_predicate_expression.cc b/src/arith/narrow_predicate_expression.cc
index 40c7ab3c54..a1a9768110 100644
--- a/src/arith/narrow_predicate_expression.cc
+++ b/src/arith/narrow_predicate_expression.cc
@@ -22,7 +22,7 @@
* \brief Utility to deduce bound of expression
*/
#include <tvm/arith/int_solver.h>
-#include <tvm/runtime/registry.h>
+#include <tvm/ffi/function.h>
#include <tvm/tir/analysis.h>
#include <tvm/tir/expr.h>
#include <tvm/tir/op.h>
@@ -212,7 +212,8 @@ PrimExpr NarrowPredicateExpression(PrimExpr expr, Map<Var, Range> free_parameter
return ExpressionNarrower::Apply(std::move(expr), std::move(free_parameters));
}
-TVM_REGISTER_GLOBAL("arith.NarrowPredicateExpression").set_body_typed(NarrowPredicateExpression);
+TVM_FFI_REGISTER_GLOBAL("arith.NarrowPredicateExpression")
+ .set_body_typed(NarrowPredicateExpression);
} // namespace arith
} // namespace tvm
diff --git a/src/arith/presburger_set.cc b/src/arith/presburger_set.cc
index 4f4d7e1857..e514ad1b1a 100644
--- a/src/arith/presburger_set.cc
+++ b/src/arith/presburger_set.cc
@@ -26,7 +26,7 @@
#include <tvm/arith/int_set.h>
#include <tvm/arith/int_solver.h>
#include <tvm/arith/pattern.h>
-#include <tvm/runtime/registry.h>
+#include <tvm/ffi/function.h>
#include <tvm/tir/expr.h>
#include <tvm/tir/expr_functor.h>
#include <tvm/tir/stmt_functor.h>
@@ -272,7 +272,7 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
PresburgerSet MakePresburgerSet(const PrimExpr& constraint) { return PresburgerSet(constraint); }
-TVM_REGISTER_GLOBAL("arith.PresburgerSet").set_body_typed(MakePresburgerSet);
+TVM_FFI_REGISTER_GLOBAL("arith.PresburgerSet").set_body_typed(MakePresburgerSet);
TVM_REGISTER_NODE_TYPE(PresburgerSetNode);
diff --git a/src/arith/solve_linear_equation.cc b/src/arith/solve_linear_equation.cc
index fb6250a778..4d90c61ea3 100644
--- a/src/arith/solve_linear_equation.cc
+++ b/src/arith/solve_linear_equation.cc
@@ -24,8 +24,8 @@
#include <tvm/arith/analyzer.h>
#include <tvm/arith/int_solver.h>
#include <tvm/arith/pattern.h>
+#include <tvm/ffi/function.h>
#include <tvm/runtime/data_type.h>
-#include <tvm/runtime/registry.h>
#include <tvm/tir/expr.h>
#include <tvm/tir/op.h>
#include <tvm/tir/stmt_functor.h>
@@ -454,7 +454,7 @@ IntConstraintsTransform SolveLinearEquations(const IntConstraints& system_to_sol
return transform;
}
-TVM_REGISTER_GLOBAL("arith.SolveLinearEquations")
+TVM_FFI_REGISTER_GLOBAL("arith.SolveLinearEquations")
.set_body_packed([](ffi::PackedArgs args, ffi::Any* ret) {
if (args.size() == 1) {
*ret = SolveLinearEquations(args[0].cast<IntConstraints>());
diff --git a/src/arith/solve_linear_inequality.cc b/src/arith/solve_linear_inequality.cc
index 0e5e6d485e..62f314d190 100644
--- a/src/arith/solve_linear_inequality.cc
+++ b/src/arith/solve_linear_inequality.cc
@@ -24,8 +24,8 @@
#include <tvm/arith/analyzer.h>
#include <tvm/arith/int_solver.h>
#include <tvm/arith/pattern.h>
+#include <tvm/ffi/function.h>
#include <tvm/runtime/data_type.h>
-#include <tvm/runtime/registry.h>
#include <tvm/tir/analysis.h>
#include <tvm/tir/expr.h>
#include <tvm/tir/op.h>
@@ -535,7 +535,7 @@ IntConstraintsTransform SolveInequalitiesDeskewRange(const IntConstraints& inequ
return transform;
}
-TVM_REGISTER_GLOBAL("arith.SolveInequalitiesAsCondition")
+TVM_FFI_REGISTER_GLOBAL("arith.SolveInequalitiesAsCondition")
.set_body_packed([](ffi::PackedArgs args, ffi::Any* ret) {
IntConstraints problem;
PartialSolvedInequalities ret_ineq;
@@ -553,7 +553,7 @@ TVM_REGISTER_GLOBAL("arith.SolveInequalitiesAsCondition")
*ret = AsConditions(problem->variables, ret_ineq.first, ret_ineq.second);
});
-TVM_REGISTER_GLOBAL("arith.SolveInequalitiesToRange")
+TVM_FFI_REGISTER_GLOBAL("arith.SolveInequalitiesToRange")
.set_body_packed([](ffi::PackedArgs args, ffi::Any* ret) {
if (args.size() == 1) {
*ret = SolveInequalitiesToRange(args[0].cast<IntConstraints>());
@@ -568,7 +568,7 @@ TVM_REGISTER_GLOBAL("arith.SolveInequalitiesToRange")
}
});
-TVM_REGISTER_GLOBAL("arith.SolveInequalitiesDeskewRange")
+TVM_FFI_REGISTER_GLOBAL("arith.SolveInequalitiesDeskewRange")
.set_body_packed([](ffi::PackedArgs args, ffi::Any* ret) {
if (args.size() == 1) {
*ret = SolveInequalitiesDeskewRange(args[0].cast<IntConstraints>());
diff --git a/src/contrib/msc/core/ir/graph.cc b/src/contrib/msc/core/ir/graph.cc
index 26f46bd032..d38beab5b4 100644
--- a/src/contrib/msc/core/ir/graph.cc
+++ b/src/contrib/msc/core/ir/graph.cc
@@ -1431,14 +1431,14 @@ TVM_REGISTER_NODE_TYPE(MSCGraphNode);
TVM_REGISTER_NODE_TYPE(WeightGraphNode);
-TVM_REGISTER_GLOBAL("msc.core.MSCTensor")
+TVM_FFI_REGISTER_GLOBAL("msc.core.MSCTensor")
.set_body_typed([](const String& name, const DataType& dtype, const String& layout,
const Array<Integer>& shape, const String& alias,
const Array<String>& prims) -> MSCTensor {
return MSCTensor(name, dtype, layout, shape, alias, prims);
});
-TVM_REGISTER_GLOBAL("msc.core.MSCTensorToJson")
+TVM_FFI_REGISTER_GLOBAL("msc.core.MSCTensorToJson")
.set_body_typed([](const MSCTensor& tensor) -> String {
const auto& tensor_json = tensor->ToJson();
std::ostringstream os;
@@ -1447,10 +1447,10 @@ TVM_REGISTER_GLOBAL("msc.core.MSCTensorToJson")
return os.str();
});
-TVM_REGISTER_GLOBAL("msc.core.MSCTensorFromJson")
+TVM_FFI_REGISTER_GLOBAL("msc.core.MSCTensorFromJson")
.set_body_typed([](const String& tensor_json) -> MSCTensor { return MSCTensor(tensor_json); });
-TVM_REGISTER_GLOBAL("msc.core.MSCJoint")
+TVM_FFI_REGISTER_GLOBAL("msc.core.MSCJoint")
.set_body_typed([](Integer index, const String& name, const String& shared_ref,
const String& optype, const Map<String, String>& attrs,
const Array<String>& scope, const Array<MSCJoint>& parents,
@@ -1464,7 +1464,7 @@ TVM_REGISTER_GLOBAL("msc.core.MSCJoint")
weights);
});
-TVM_REGISTER_GLOBAL("msc.core.MSCPrim")
+TVM_FFI_REGISTER_GLOBAL("msc.core.MSCPrim")
.set_body_typed([](Integer index, const String& name, const String& optype,
const Map<String, String>& attrs, const Array<MSCPrim>& parents) -> MSCPrim {
Array<BaseJoint> b_parents;
@@ -1474,7 +1474,7 @@ TVM_REGISTER_GLOBAL("msc.core.MSCPrim")
return MSCPrim(index->value, name, optype, b_parents, attrs);
});
-TVM_REGISTER_GLOBAL("msc.core.WeightJoint")
+TVM_FFI_REGISTER_GLOBAL("msc.core.WeightJoint")
.set_body_typed([](Integer index, const String& name, const String& shared_ref,
const String& weight_type, const MSCTensor& weight,
const Array<BaseJoint> parents, const Map<String, String>& attrs,
@@ -1490,108 +1490,109 @@ TVM_REGISTER_GLOBAL("msc.core.WeightJoint")
b_friends);
});
-TVM_REGISTER_GLOBAL("msc.core.WeightJointSetAttr")
+TVM_FFI_REGISTER_GLOBAL("msc.core.WeightJointSetAttr")
.set_body_typed([](const WeightJoint& node, const String& key, const String& value) {
node->attrs.Set(key, value);
});
-TVM_REGISTER_GLOBAL("msc.core.MSCGraph")
+TVM_FFI_REGISTER_GLOBAL("msc.core.MSCGraph")
.set_body_typed([](const String& name, const Array<MSCJoint>& nodes,
const Array<String>& input_names, const Array<String>& output_names,
const Array<MSCPrim>& prims) -> MSCGraph {
return MSCGraph(name, nodes, input_names, output_names, prims);
});
-TVM_REGISTER_GLOBAL("msc.core.WeightGraph")
+TVM_FFI_REGISTER_GLOBAL("msc.core.WeightGraph")
.set_body_typed([](const MSCGraph& graph, const Map<String, Array<String>>& main_wtypes,
const Map<String, String>& relation_wtypes) -> WeightGraph {
return WeightGraph(graph, main_wtypes, relation_wtypes);
});
// MSC Graph APIS
-TVM_REGISTER_GLOBAL("msc.core.MSCGraphHasNode")
+TVM_FFI_REGISTER_GLOBAL("msc.core.MSCGraphHasNode")
.set_body_typed([](const MSCGraph& graph, const String& name) -> Bool {
return Bool(graph->HasNode(name));
});
-TVM_REGISTER_GLOBAL("msc.core.MSCGraphFindNode")
+TVM_FFI_REGISTER_GLOBAL("msc.core.MSCGraphFindNode")
.set_body_typed([](const MSCGraph& graph, const String& name) -> MSCJoint {
return graph->FindNode(name);
});
-TVM_REGISTER_GLOBAL("msc.core.MSCGraphFindPrim")
+TVM_FFI_REGISTER_GLOBAL("msc.core.MSCGraphFindPrim")
.set_body_typed([](const MSCGraph& graph, const String& name) -> MSCPrim {
return graph->FindPrim(name);
});
-TVM_REGISTER_GLOBAL("msc.core.MSCGraphHasTensor")
+TVM_FFI_REGISTER_GLOBAL("msc.core.MSCGraphHasTensor")
.set_body_typed([](const MSCGraph& graph, const String& name) -> Bool {
return Bool(graph->HasTensor(name));
});
-TVM_REGISTER_GLOBAL("msc.core.MSCGraphFindTensor")
+TVM_FFI_REGISTER_GLOBAL("msc.core.MSCGraphFindTensor")
.set_body_typed([](const MSCGraph& graph, const String& name) -> MSCTensor {
return graph->FindTensor(name);
});
-TVM_REGISTER_GLOBAL("msc.core.MSCGraphSetTensorAlias")
+TVM_FFI_REGISTER_GLOBAL("msc.core.MSCGraphSetTensorAlias")
.set_body_typed([](const MSCGraph& graph, const MSCTensor& tensor, const String& alias) {
tensor->alias = alias;
graph->tensor_alias.Set(alias, tensor->name);
});
-TVM_REGISTER_GLOBAL("msc.core.MSCGraphFindProducer")
+TVM_FFI_REGISTER_GLOBAL("msc.core.MSCGraphFindProducer")
.set_body_typed([](const MSCGraph& graph, const String& name) -> MSCJoint {
return graph->FindProducer(name);
});
-TVM_REGISTER_GLOBAL("msc.core.MSCGraphFindConsumers")
+TVM_FFI_REGISTER_GLOBAL("msc.core.MSCGraphFindConsumers")
.set_body_typed([](const MSCGraph& graph, const String& name) -> Array<MSCJoint> {
return graph->FindConsumers(name);
});
-TVM_REGISTER_GLOBAL("msc.core.MSCGraphInputAt")
+TVM_FFI_REGISTER_GLOBAL("msc.core.MSCGraphInputAt")
.set_body_typed([](const MSCGraph& graph, int index) -> MSCTensor {
return graph->InputAt(index);
});
-TVM_REGISTER_GLOBAL("msc.core.MSCGraphOutputAt")
+TVM_FFI_REGISTER_GLOBAL("msc.core.MSCGraphOutputAt")
.set_body_typed([](const MSCGraph& graph, int index) -> MSCTensor {
return graph->OutputAt(index);
});
-TVM_REGISTER_GLOBAL("msc.core.MSCGraphGetInputs")
+TVM_FFI_REGISTER_GLOBAL("msc.core.MSCGraphGetInputs")
.set_body_typed([](const MSCGraph& graph) -> Array<MSCTensor> { return graph->GetInputs(); });
-TVM_REGISTER_GLOBAL("msc.core.MSCGraphGetOutputs")
+TVM_FFI_REGISTER_GLOBAL("msc.core.MSCGraphGetOutputs")
.set_body_typed([](const MSCGraph& graph) -> Array<MSCTensor> { return graph->GetOutputs(); });
-TVM_REGISTER_GLOBAL("msc.core.MSCGraphToJson").set_body_typed([](const MSCGraph& graph) -> String {
- const auto& graph_json = graph->ToJson();
- std::ostringstream os;
- dmlc::JSONWriter writer(&os);
- graph_json.Save(&writer);
- return os.str();
-});
+TVM_FFI_REGISTER_GLOBAL("msc.core.MSCGraphToJson")
+ .set_body_typed([](const MSCGraph& graph) -> String {
+ const auto& graph_json = graph->ToJson();
+ std::ostringstream os;
+ dmlc::JSONWriter writer(&os);
+ graph_json.Save(&writer);
+ return os.str();
+ });
-TVM_REGISTER_GLOBAL("msc.core.MSCGraphFromJson")
+TVM_FFI_REGISTER_GLOBAL("msc.core.MSCGraphFromJson")
.set_body_typed([](const String& graph_json) -> MSCGraph { return MSCGraph(graph_json); });
-TVM_REGISTER_GLOBAL("msc.core.MSCGraphToPrototxt")
+TVM_FFI_REGISTER_GLOBAL("msc.core.MSCGraphToPrototxt")
.set_body_typed([](const MSCGraph& graph) -> String { return graph->ToPrototxt(); });
// Weight Graph APIS
-TVM_REGISTER_GLOBAL("msc.core.WeightGraphHasNode")
+TVM_FFI_REGISTER_GLOBAL("msc.core.WeightGraphHasNode")
.set_body_typed([](const WeightGraph& graph, const String& name) -> Bool {
return Bool(graph->HasNode(name));
});
-TVM_REGISTER_GLOBAL("msc.core.WeightGraphFindNode")
+TVM_FFI_REGISTER_GLOBAL("msc.core.WeightGraphFindNode")
.set_body_typed([](const WeightGraph& graph, const String& name) -> WeightJoint {
return graph->FindNode(name);
});
-TVM_REGISTER_GLOBAL("msc.core.WeightGraphToJson")
+TVM_FFI_REGISTER_GLOBAL("msc.core.WeightGraphToJson")
.set_body_typed([](const WeightGraph& graph) -> String {
const auto& graph_json = graph->ToJson();
std::ostringstream os;
@@ -1600,69 +1601,69 @@ TVM_REGISTER_GLOBAL("msc.core.WeightGraphToJson")
return os.str();
});
-TVM_REGISTER_GLOBAL("msc.core.WeightGraphFromJson")
+TVM_FFI_REGISTER_GLOBAL("msc.core.WeightGraphFromJson")
.set_body_typed([](const String& graph_json) -> WeightGraph {
return WeightGraph(graph_json);
});
-TVM_REGISTER_GLOBAL("msc.core.WeightGraphToPrototxt")
+TVM_FFI_REGISTER_GLOBAL("msc.core.WeightGraphToPrototxt")
.set_body_typed([](const WeightGraph& graph) -> String { return graph->ToPrototxt(); });
-TVM_REGISTER_GLOBAL("msc.core.MSCJointInputAt")
+TVM_FFI_REGISTER_GLOBAL("msc.core.MSCJointInputAt")
.set_body_typed([](const MSCJoint& node, int index) -> MSCTensor {
return node->InputAt(index);
});
-TVM_REGISTER_GLOBAL("msc.core.MSCJointOutputAt")
+TVM_FFI_REGISTER_GLOBAL("msc.core.MSCJointOutputAt")
.set_body_typed([](const MSCJoint& node, int index) -> MSCTensor {
return node->OutputAt(index);
});
-TVM_REGISTER_GLOBAL("msc.core.MSCJointWeightAt")
+TVM_FFI_REGISTER_GLOBAL("msc.core.MSCJointWeightAt")
.set_body_typed([](const MSCJoint& node, const String& wtype) -> MSCTensor {
return node->WeightAt(wtype);
});
-TVM_REGISTER_GLOBAL("msc.core.MSCJointGetInputs")
+TVM_FFI_REGISTER_GLOBAL("msc.core.MSCJointGetInputs")
.set_body_typed([](const MSCJoint& node) -> Array<MSCTensor> { return node->GetInputs(); });
-TVM_REGISTER_GLOBAL("msc.core.MSCJointGetOutputs")
+TVM_FFI_REGISTER_GLOBAL("msc.core.MSCJointGetOutputs")
.set_body_typed([](const MSCJoint& node) -> Array<MSCTensor> { return node->GetOutputs(); });
-TVM_REGISTER_GLOBAL("msc.core.MSCJointGetWeights")
+TVM_FFI_REGISTER_GLOBAL("msc.core.MSCJointGetWeights")
.set_body_typed([](const MSCJoint& node) -> Map<String, MSCTensor> { return node->weights; });
-TVM_REGISTER_GLOBAL("msc.core.MSCJointHasAttr")
+TVM_FFI_REGISTER_GLOBAL("msc.core.MSCJointHasAttr")
.set_body_typed([](const MSCJoint& node, const String& key) -> Bool {
return Bool(node->HasAttr(key));
});
-TVM_REGISTER_GLOBAL("msc.core.MSCJointGetAttrs")
+TVM_FFI_REGISTER_GLOBAL("msc.core.MSCJointGetAttrs")
.set_body_typed([](const MSCJoint& node) -> Map<String, String> { return node->attrs; });
-TVM_REGISTER_GLOBAL("msc.core.WeightJointHasAttr")
+TVM_FFI_REGISTER_GLOBAL("msc.core.WeightJointHasAttr")
.set_body_typed([](const WeightJoint& node, const String& key) -> Bool {
return Bool(node->HasAttr(key));
});
-TVM_REGISTER_GLOBAL("msc.core.WeightJointGetAttrs")
+TVM_FFI_REGISTER_GLOBAL("msc.core.WeightJointGetAttrs")
.set_body_typed([](const WeightJoint& node) -> Map<String, String> { return node->attrs; });
-TVM_REGISTER_GLOBAL("msc.core.MSCTensorDTypeName")
+TVM_FFI_REGISTER_GLOBAL("msc.core.MSCTensorDTypeName")
.set_body_typed([](const MSCTensor& tensor) -> String { return tensor->DTypeName(); });
-TVM_REGISTER_GLOBAL("msc.core.MSCTensorDimAt")
+TVM_FFI_REGISTER_GLOBAL("msc.core.MSCTensorDimAt")
.set_body_typed([](const MSCTensor& tensor, const String& axis) -> Integer {
return tensor->DimAt(axis);
});
-TVM_REGISTER_GLOBAL("msc.core.MSCTensorGetSize")
+TVM_FFI_REGISTER_GLOBAL("msc.core.MSCTensorGetSize")
.set_body_typed([](const MSCTensor& tensor) -> Integer { return tensor->GetSize(); });
-TVM_REGISTER_GLOBAL("msc.core.MSCTensorSetAlias")
+TVM_FFI_REGISTER_GLOBAL("msc.core.MSCTensorSetAlias")
.set_body_typed([](const MSCTensor& tensor, const String& alias) { tensor->alias = alias; });
-TVM_REGISTER_GLOBAL("msc.core.PruneWeights")
+TVM_FFI_REGISTER_GLOBAL("msc.core.PruneWeights")
.set_body_typed([](const MSCGraph& graph,
const Map<String, MSCTensor>& pruned_tensors) -> MSCGraph {
return PruneWeights(graph, pruned_tensors);
diff --git a/src/contrib/msc/core/ir/graph_builder.cc b/src/contrib/msc/core/ir/graph_builder.cc
index 030df82dc6..2550f5652f 100644
--- a/src/contrib/msc/core/ir/graph_builder.cc
+++ b/src/contrib/msc/core/ir/graph_builder.cc
@@ -834,7 +834,7 @@ void WeightsExtractor::VisitExpr_(const CallNode* op) {
}
}
-TVM_REGISTER_GLOBAL("msc.core.BuildFromRelax")
+TVM_FFI_REGISTER_GLOBAL("msc.core.BuildFromRelax")
.set_body_typed([](const IRModule& module, const String& entry_name,
const String& options) -> MSCGraph {
auto builder = GraphBuilder(module, entry_name, options);
@@ -844,7 +844,7 @@ TVM_REGISTER_GLOBAL("msc.core.BuildFromRelax")
return builder.Build(func);
});
-TVM_REGISTER_GLOBAL("msc.core.GetRelaxWeights")
+TVM_FFI_REGISTER_GLOBAL("msc.core.GetRelaxWeights")
.set_body_typed([](const IRModule& module,
const String& entry_name) -> Map<MSCTensor, NDArray> {
const auto& func = Downcast<Function>(module->Lookup(entry_name));
diff --git a/src/contrib/msc/core/ir/plugin.cc b/src/contrib/msc/core/ir/plugin.cc
index d34972639a..fc6000a20f 100644
--- a/src/contrib/msc/core/ir/plugin.cc
+++ b/src/contrib/msc/core/ir/plugin.cc
@@ -305,20 +305,20 @@ const Plugin GetPlugin(const String& name) { return PluginRegistry::Global()->Ge
bool IsPlugin(const String& name) { return PluginRegistry::Global()->Registered(name); }
-TVM_REGISTER_GLOBAL("msc.core.RegisterPlugin")
+TVM_FFI_REGISTER_GLOBAL("msc.core.RegisterPlugin")
.set_body_typed([](const String& name, const String& json_str) {
PluginRegistry::Global()->Register(name, json_str);
});
-TVM_REGISTER_GLOBAL("msc.core.ListPluginNames").set_body_typed([]() -> Array<String> {
+TVM_FFI_REGISTER_GLOBAL("msc.core.ListPluginNames").set_body_typed([]() -> Array<String> {
return ListPluginNames();
});
-TVM_REGISTER_GLOBAL("msc.core.GetPlugin").set_body_typed([](const String& name) -> Plugin {
+TVM_FFI_REGISTER_GLOBAL("msc.core.GetPlugin").set_body_typed([](const String& name) -> Plugin {
return GetPlugin(name);
});
-TVM_REGISTER_GLOBAL("msc.core.IsPlugin").set_body_typed([](const String& name) -> Bool {
+TVM_FFI_REGISTER_GLOBAL("msc.core.IsPlugin").set_body_typed([](const String& name) -> Bool {
return Bool(IsPlugin(name));
});
diff --git a/src/contrib/msc/core/transform/bind_named_params.cc b/src/contrib/msc/core/transform/bind_named_params.cc
index 523db32b3a..0225ff3190 100644
--- a/src/contrib/msc/core/transform/bind_named_params.cc
+++ b/src/contrib/msc/core/transform/bind_named_params.cc
@@ -154,7 +154,7 @@ Pass BindNamedParams(String func_name, Map<ObjectRef, ObjectRef> params) {
return CreateModulePass(pass_func, 0, "BindNamedParams", {});
}
-TVM_REGISTER_GLOBAL("relax.transform.BindNamedParams").set_body_typed(BindNamedParams);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.BindNamedParams").set_body_typed(BindNamedParams);
} // namespace transform
diff --git a/src/contrib/msc/core/transform/bind_shape.cc b/src/contrib/msc/core/transform/bind_shape.cc
index ca91b3424e..b554e08ab8 100644
--- a/src/contrib/msc/core/transform/bind_shape.cc
+++ b/src/contrib/msc/core/transform/bind_shape.cc
@@ -132,7 +132,7 @@ Pass BindShape(const String& entry_name) {
return CreateModulePass(pass_func, 0, "BindShape", {});
}
-TVM_REGISTER_GLOBAL("relax.transform.BindShape").set_body_typed(BindShape);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.BindShape").set_body_typed(BindShape);
} // namespace transform
} // namespace relax
diff --git a/src/contrib/msc/core/transform/fuse_tuple.cc b/src/contrib/msc/core/transform/fuse_tuple.cc
index 1bcbe076fc..1eabf3306f 100644
--- a/src/contrib/msc/core/transform/fuse_tuple.cc
+++ b/src/contrib/msc/core/transform/fuse_tuple.cc
@@ -231,7 +231,7 @@ Pass FuseTuple(const String& target, const String& entry_name) {
return CreateModulePass(pass_func, 0, "FuseTuple", {});
}
-TVM_REGISTER_GLOBAL("relax.transform.FuseTuple").set_body_typed(FuseTuple);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.FuseTuple").set_body_typed(FuseTuple);
} // namespace transform
} // namespace relax
diff --git a/src/contrib/msc/core/transform/inline_params.cc b/src/contrib/msc/core/transform/inline_params.cc
index e892d3b36a..a91eb590af 100644
--- a/src/contrib/msc/core/transform/inline_params.cc
+++ b/src/contrib/msc/core/transform/inline_params.cc
@@ -184,7 +184,7 @@ Pass InlineParams(const String& entry_name) {
return CreateModulePass(pass_func, 0, "InlineParams", {});
}
-TVM_REGISTER_GLOBAL("relax.transform.InlineParams").set_body_typed(InlineParams);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.InlineParams").set_body_typed(InlineParams);
} // namespace transform
} // namespace relax
diff --git a/src/contrib/msc/core/transform/set_byoc_attrs.cc b/src/contrib/msc/core/transform/set_byoc_attrs.cc
index 6f3a29346b..4755ebf389 100644
--- a/src/contrib/msc/core/transform/set_byoc_attrs.cc
+++ b/src/contrib/msc/core/transform/set_byoc_attrs.cc
@@ -101,7 +101,7 @@ Pass SetBYOCAttrs(const String& target, const String& entry_name) {
return CreateModulePass(pass_func, 0, "SetBYOCAttrs", {});
}
-TVM_REGISTER_GLOBAL("relax.transform.SetBYOCAttrs").set_body_typed(SetBYOCAttrs);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.SetBYOCAttrs").set_body_typed(SetBYOCAttrs);
} // namespace transform
} // namespace relax
diff --git a/src/contrib/msc/core/transform/set_expr_layout.cc b/src/contrib/msc/core/transform/set_expr_layout.cc
index 80416bafd0..dd87e60e7b 100644
--- a/src/contrib/msc/core/transform/set_expr_layout.cc
+++ b/src/contrib/msc/core/transform/set_expr_layout.cc
@@ -1359,7 +1359,7 @@ Pass SetExprLayout(bool allow_missing, const String& entry_name) {
return CreateModulePass(pass_func, 0, "SetExprLayout", {});
}
-TVM_REGISTER_GLOBAL("relax.transform.SetExprLayout").set_body_typed(SetExprLayout);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.SetExprLayout").set_body_typed(SetExprLayout);
} // namespace transform
} // namespace relax
diff --git a/src/contrib/msc/core/transform/set_expr_name.cc b/src/contrib/msc/core/transform/set_expr_name.cc
index 4c36e2ba27..4d0cc0314e 100644
--- a/src/contrib/msc/core/transform/set_expr_name.cc
+++ b/src/contrib/msc/core/transform/set_expr_name.cc
@@ -324,7 +324,7 @@ Pass SetRelaxExprName(const String& entry_name, const String& target,
return CreateModulePass(pass_func, 0, "SetRelaxExprName", {});
}
-TVM_REGISTER_GLOBAL("relax.transform.SetRelaxExprName").set_body_typed(SetRelaxExprName);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.SetRelaxExprName").set_body_typed(SetRelaxExprName);
} // namespace transform
} // namespace relax
diff --git a/src/contrib/msc/core/utils.cc b/src/contrib/msc/core/utils.cc
index 3c53e32170..d03f3ba82b 100644
--- a/src/contrib/msc/core/utils.cc
+++ b/src/contrib/msc/core/utils.cc
@@ -523,27 +523,27 @@ const DataType ExprUtils::GetDataType(const Expr& expr) {
return Downcast<TensorStructInfo>(GetStructInfo(expr))->dtype;
}
-TVM_REGISTER_GLOBAL("msc.core.SpanGetAttr").set_body_typed(SpanUtils::GetAttr);
+TVM_FFI_REGISTER_GLOBAL("msc.core.SpanGetAttr").set_body_typed(SpanUtils::GetAttr);
-TVM_REGISTER_GLOBAL("msc.core.SpanGetAttrs").set_body_typed(SpanUtils::GetAttrs);
+TVM_FFI_REGISTER_GLOBAL("msc.core.SpanGetAttrs").set_body_typed(SpanUtils::GetAttrs);
-TVM_REGISTER_GLOBAL("msc.core.SpanCreateWithAttr")
+TVM_FFI_REGISTER_GLOBAL("msc.core.SpanCreateWithAttr")
.set_body_typed([](const String& key, const String& value) -> Span {
return SpanUtils::CreateWithAttr(key, value);
});
-TVM_REGISTER_GLOBAL("msc.core.SpanSetAttr")
+TVM_FFI_REGISTER_GLOBAL("msc.core.SpanSetAttr")
.set_body_typed([](const Span& span, const String& key, const String& value) -> Span {
return SpanUtils::SetAttr(span, key, value);
});
-TVM_REGISTER_GLOBAL("msc.core.CompareVersion")
+TVM_FFI_REGISTER_GLOBAL("msc.core.CompareVersion")
.set_body_typed([](const Array<Integer>& given_version,
const Array<Integer>& target_version) -> Integer {
return Integer(CommonUtils::CompareVersion(given_version, target_version));
});
-TVM_REGISTER_GLOBAL("msc.core.ToAttrKey").set_body_typed([](const String& key) -> String {
+TVM_FFI_REGISTER_GLOBAL("msc.core.ToAttrKey").set_body_typed([](const String& key) -> String {
return CommonUtils::ToAttrKey(key);
});
diff --git a/src/contrib/msc/framework/tensorflow/codegen.cc b/src/contrib/msc/framework/tensorflow/codegen.cc
index 9506d4eac8..4bceb76d46 100644
--- a/src/contrib/msc/framework/tensorflow/codegen.cc
+++ b/src/contrib/msc/framework/tensorflow/codegen.cc
@@ -150,7 +150,7 @@ const Array<Doc> TensorflowCodeGen::GetOpCodes(const MSCJoint& node) {
}
}
-TVM_REGISTER_GLOBAL("msc.framework.tensorflow.GetTensorflowSources")
+TVM_FFI_REGISTER_GLOBAL("msc.framework.tensorflow.GetTensorflowSources")
.set_body_typed([](const MSCGraph& graph, const String& codegen_config,
const String& print_config) -> Map<String, String> {
TensorflowCodeGen codegen = TensorflowCodeGen(graph, codegen_config);
diff --git a/src/contrib/msc/framework/tensorrt/codegen.cc b/src/contrib/msc/framework/tensorrt/codegen.cc
index 0271b63112..8b85f2e88f 100644
--- a/src/contrib/msc/framework/tensorrt/codegen.cc
+++ b/src/contrib/msc/framework/tensorrt/codegen.cc
@@ -574,7 +574,7 @@ const Map<String, String> TensorRTCodeGen::GetStepCtx() {
return step_ctx;
}
-TVM_REGISTER_GLOBAL("msc.framework.tensorrt.GetTensorRTSources")
+TVM_FFI_REGISTER_GLOBAL("msc.framework.tensorrt.GetTensorRTSources")
.set_body_typed([](const MSCGraph& graph, const String& codegen_config,
const String& print_config) -> Map<String, String> {
TensorRTCodeGen codegen = TensorRTCodeGen(graph, codegen_config);
@@ -582,7 +582,7 @@ TVM_REGISTER_GLOBAL("msc.framework.tensorrt.GetTensorRTSources")
return codegen.GetSources(print_config);
});
-TVM_REGISTER_GLOBAL("msc.framework.tensorrt.GetTensorRTRoot").set_body_typed([]() -> String {
+TVM_FFI_REGISTER_GLOBAL("msc.framework.tensorrt.GetTensorRTRoot").set_body_typed([]() -> String {
#ifdef TENSORRT_ROOT_DIR
return TENSORRT_ROOT_DIR;
#else
@@ -618,7 +618,7 @@ Array<runtime::Module> MSCTensorRTCompiler(Array<Function> functions,
return compiled_functions;
}
-TVM_REGISTER_GLOBAL("relax.ext.msc_tensorrt").set_body_typed(MSCTensorRTCompiler);
+TVM_FFI_REGISTER_GLOBAL("relax.ext.msc_tensorrt").set_body_typed(MSCTensorRTCompiler);
} // namespace msc
} // namespace contrib
diff --git a/src/contrib/msc/framework/tensorrt/transform_tensorrt.cc b/src/contrib/msc/framework/tensorrt/transform_tensorrt.cc
index fc0e538252..67f453268e 100644
--- a/src/contrib/msc/framework/tensorrt/transform_tensorrt.cc
+++ b/src/contrib/msc/framework/tensorrt/transform_tensorrt.cc
@@ -913,7 +913,7 @@ Pass TransformTensorRT(const String& config) {
return CreateFunctionPass(pass_func, 0, "TransformTensorRT", {});
}
-TVM_REGISTER_GLOBAL("relax.transform.TransformTensorRT").set_body_typed(TransformTensorRT);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.TransformTensorRT").set_body_typed(TransformTensorRT);
} // namespace transform
} // namespace relax
diff --git a/src/contrib/msc/framework/torch/codegen.cc b/src/contrib/msc/framework/torch/codegen.cc
index 547c1c22ba..228efa4381 100644
--- a/src/contrib/msc/framework/torch/codegen.cc
+++ b/src/contrib/msc/framework/torch/codegen.cc
@@ -151,7 +151,7 @@ const Array<Doc> TorchCodeGen::GetOpCodes(const MSCJoint& node) {
}
}
-TVM_REGISTER_GLOBAL("msc.framework.torch.GetTorchSources")
+TVM_FFI_REGISTER_GLOBAL("msc.framework.torch.GetTorchSources")
.set_body_typed([](const MSCGraph& graph, const String& codegen_config,
const String& print_config) -> Map<String, String> {
TorchCodeGen codegen = TorchCodeGen(graph, codegen_config);
diff --git a/src/contrib/msc/framework/tvm/codegen.cc b/src/contrib/msc/framework/tvm/codegen.cc
index 1d6d74d7e4..53d1bc0562 100644
--- a/src/contrib/msc/framework/tvm/codegen.cc
+++ b/src/contrib/msc/framework/tvm/codegen.cc
@@ -210,7 +210,7 @@ const Array<Doc> RelaxCodeGen::GetOpCodes(const MSCJoint& node) {
}
}
-TVM_REGISTER_GLOBAL("msc.framework.tvm.GetRelaxSources")
+TVM_FFI_REGISTER_GLOBAL("msc.framework.tvm.GetRelaxSources")
.set_body_typed([](const MSCGraph& graph, const String& codegen_config,
const String& print_config) -> Map<String, String> {
RelaxCodeGen codegen = RelaxCodeGen(graph, codegen_config);
diff --git a/src/contrib/msc/plugin/tensorrt_codegen.cc b/src/contrib/msc/plugin/tensorrt_codegen.cc
index d16bb5d924..02904c3bd9 100644
--- a/src/contrib/msc/plugin/tensorrt_codegen.cc
+++ b/src/contrib/msc/plugin/tensorrt_codegen.cc
@@ -883,7 +883,7 @@ void TensorRTPluginCodeGen::CodegenEnqueue(const Plugin& plugin, bool dynamic) {
}
}
-TVM_REGISTER_GLOBAL("msc.plugin.GetTensorRTPluginSources")
+TVM_FFI_REGISTER_GLOBAL("msc.plugin.GetTensorRTPluginSources")
.set_body_typed([](const String& codegen_config, const String& print_config,
const String& codegen_type) -> Map<String, String> {
TensorRTPluginCodeGen codegen = TensorRTPluginCodeGen(codegen_config);
diff --git a/src/contrib/msc/plugin/torch_codegen.cc b/src/contrib/msc/plugin/torch_codegen.cc
index 215341af01..59b99f22c7 100644
--- a/src/contrib/msc/plugin/torch_codegen.cc
+++ b/src/contrib/msc/plugin/torch_codegen.cc
@@ -492,7 +492,7 @@ void TorchPluginCodeGen::CodeGenCompute(const Plugin& plugin, const String& devi
}
}
-TVM_REGISTER_GLOBAL("msc.plugin.GetTorchPluginSources")
+TVM_FFI_REGISTER_GLOBAL("msc.plugin.GetTorchPluginSources")
.set_body_typed([](const String& codegen_config, const String& print_config,
const String& codegen_type) -> Map<String, String> {
TorchPluginCodeGen codegen = TorchPluginCodeGen(codegen_config);
diff --git a/src/contrib/msc/plugin/tvm_codegen.cc b/src/contrib/msc/plugin/tvm_codegen.cc
index e1d3c9960f..610fbc4c32 100644
--- a/src/contrib/msc/plugin/tvm_codegen.cc
+++ b/src/contrib/msc/plugin/tvm_codegen.cc
@@ -213,12 +213,12 @@ void TVMPluginCodeGen::CodeGenOpDefine(const Plugin& plugin) {
stack_.func_end("infer_output");
// register funcs
- stack_.func_call("TVM_REGISTER_GLOBAL")
+ stack_.func_call("TVM_FFI_REGISTER_GLOBAL")
.call_arg(DocUtils::ToStr("msc.plugin.op.InferStructInfo" + plugin->name))
.method_call("set_body_typed")
.call_arg("InferStructInfo" + plugin->name)
.line()
- .func_call("TVM_REGISTER_GLOBAL")
+ .func_call("TVM_FFI_REGISTER_GLOBAL")
.call_arg(DocUtils::ToStr("msc.plugin.op.InferLayout" + plugin->name))
.method_call("set_body_typed")
.call_arg("InferLayout" + plugin->name)
@@ -260,7 +260,7 @@ void TVMPluginCodeGen::CodeGenOpRuntime(const Plugin& plugin) {
CodeGenCompute(plugin, "cpu");
stack_.cond_end().func_end();
// register the compute
- stack_.func_call("TVM_REGISTER_GLOBAL")
+ stack_.func_call("TVM_FFI_REGISTER_GLOBAL")
.call_arg(DocUtils::ToStr(plugin->name))
.method_call("set_body")
.call_arg(func_name)
@@ -393,7 +393,7 @@ void TVMPluginCodeGen::CodeGenCompute(const Plugin& plugin, const String& device
}
}
-TVM_REGISTER_GLOBAL("msc.plugin.GetTVMPluginSources")
+TVM_FFI_REGISTER_GLOBAL("msc.plugin.GetTVMPluginSources")
.set_body_typed([](const String& codegen_config, const String& print_config,
const String& codegen_type) -> Map<String, String> {
TVMPluginCodeGen codegen = TVMPluginCodeGen(codegen_config);
diff --git a/src/ir/analysis.cc b/src/ir/analysis.cc
index 9de36b0a28..3a54085c22 100644
--- a/src/ir/analysis.cc
+++ b/src/ir/analysis.cc
@@ -43,7 +43,7 @@ Map<GlobalVar, Array<GlobalVar>> CollectCallMap(const IRModule& mod) {
return call_map;
}
-TVM_REGISTER_GLOBAL("ir.analysis.CollectCallMap").set_body_typed(CollectCallMap);
+TVM_FFI_REGISTER_GLOBAL("ir.analysis.CollectCallMap").set_body_typed(CollectCallMap);
} // namespace ir
} // namespace tvm
diff --git a/src/ir/apply_pass_to_function.cc b/src/ir/apply_pass_to_function.cc
index 9e43e33a6c..877530f4c3 100644
--- a/src/ir/apply_pass_to_function.cc
+++ b/src/ir/apply_pass_to_function.cc
@@ -21,9 +21,9 @@
* \file src/ir/apply_pass_to_function.cc
* \brief Utility transformation that applies an inner pass to a subset of an IRModule
*/
+#include <tvm/ffi/function.h>
#include <tvm/ir/transform.h>
#include <tvm/relax/expr.h>
-#include <tvm/runtime/registry.h>
#include <tvm/tir/function.h>
#include <unordered_set>
@@ -130,7 +130,7 @@ Pass ApplyPassToFunction(Pass pass, String func_name_regex,
return CreateModulePass(pass_func, 0, pass_name, {});
}
-TVM_REGISTER_GLOBAL("transform.ApplyPassToFunction").set_body_typed(ApplyPassToFunction);
+TVM_FFI_REGISTER_GLOBAL("transform.ApplyPassToFunction").set_body_typed(ApplyPassToFunction);
} // namespace transform
} // namespace tvm
diff --git a/src/ir/attrs.cc b/src/ir/attrs.cc
index fd87c2bc8e..52a2cceeaf 100644
--- a/src/ir/attrs.cc
+++ b/src/ir/attrs.cc
@@ -20,8 +20,8 @@
/*!
* \file attrs.cc
*/
+#include <tvm/ffi/function.h>
#include <tvm/ir/attrs.h>
-#include <tvm/runtime/registry.h>
#include "attr_functor.h"
@@ -73,11 +73,11 @@ TVM_REGISTER_NODE_TYPE(DictAttrsNode);
TVM_REGISTER_NODE_TYPE(AttrFieldInfoNode);
-TVM_REGISTER_GLOBAL("ir.DictAttrsGetDict").set_body_typed([](DictAttrs attrs) {
+TVM_FFI_REGISTER_GLOBAL("ir.DictAttrsGetDict").set_body_typed([](DictAttrs attrs) {
return attrs->dict;
});
-TVM_REGISTER_GLOBAL("ir.AttrsListFieldInfo").set_body_typed([](Attrs attrs) {
+TVM_FFI_REGISTER_GLOBAL("ir.AttrsListFieldInfo").set_body_typed([](Attrs attrs) {
return attrs->ListFieldInfo();
});
diff --git a/src/ir/diagnostic.cc b/src/ir/diagnostic.cc
index ec11f2c04f..7019707431 100644
--- a/src/ir/diagnostic.cc
+++ b/src/ir/diagnostic.cc
@@ -33,7 +33,7 @@ namespace tvm {
/* Diagnostic */
TVM_REGISTER_NODE_TYPE(DiagnosticNode);
-TVM_REGISTER_GLOBAL("diagnostics.Diagnostic")
+TVM_FFI_REGISTER_GLOBAL("diagnostics.Diagnostic")
.set_body_typed([](int level, Span span, String message) {
return Diagnostic(static_cast<DiagnosticLevel>(level), span, message);
});
@@ -106,7 +106,7 @@ TVM_DLL DiagnosticRenderer::DiagnosticRenderer(
data_ = std::move(n);
}
-TVM_REGISTER_GLOBAL("diagnostics.DiagnosticRenderer")
+TVM_FFI_REGISTER_GLOBAL("diagnostics.DiagnosticRenderer")
.set_body_typed([](ffi::TypedFunction<void(DiagnosticContext ctx)> renderer) {
return DiagnosticRenderer(renderer);
});
@@ -134,7 +134,7 @@ void DiagnosticContext::Render() {
}
}
-TVM_REGISTER_GLOBAL("diagnostics.DiagnosticRendererRender")
+TVM_FFI_REGISTER_GLOBAL("diagnostics.DiagnosticRendererRender")
.set_body_typed([](DiagnosticRenderer renderer, DiagnosticContext ctx) {
renderer.Render(ctx);
});
@@ -147,7 +147,7 @@ DiagnosticContext::DiagnosticContext(const IRModule& module, const DiagnosticRen
data_ = std::move(n);
}
-TVM_REGISTER_GLOBAL("diagnostics.DiagnosticContext")
+TVM_FFI_REGISTER_GLOBAL("diagnostics.DiagnosticContext")
.set_body_typed([](const IRModule& module, const DiagnosticRenderer& renderer) {
return DiagnosticContext(module, renderer);
});
@@ -157,12 +157,12 @@ void DiagnosticContext::Emit(const Diagnostic& diagnostic) {
(*this)->diagnostics.push_back(diagnostic);
}
-TVM_REGISTER_GLOBAL("diagnostics.Emit")
+TVM_FFI_REGISTER_GLOBAL("diagnostics.Emit")
.set_body_typed([](DiagnosticContext ctx, const Diagnostic& diagnostic) {
return ctx.Emit(diagnostic);
});
-TVM_REGISTER_GLOBAL("diagnostics.DiagnosticContextRender")
+TVM_FFI_REGISTER_GLOBAL("diagnostics.DiagnosticContextRender")
.set_body_typed([](DiagnosticContext context) { return context.Render(); });
/*! \brief Emit a diagnostic. */
@@ -195,7 +195,7 @@ DiagnosticContext DiagnosticContext::Default(const IRModule& module) {
return DiagnosticContext(module, renderer);
}
-TVM_REGISTER_GLOBAL("diagnostics.Default").set_body_typed([](const IRModule& module) {
+TVM_FFI_REGISTER_GLOBAL("diagnostics.Default").set_body_typed([](const IRModule& module) {
return DiagnosticContext::Default(module);
});
@@ -311,11 +311,13 @@ DiagnosticRenderer TerminalRenderer(std::ostream& out) {
});
}
-TVM_REGISTER_GLOBAL(DEFAULT_RENDERER).set_body_typed([]() { return TerminalRenderer(std::cerr); });
+TVM_FFI_REGISTER_GLOBAL(DEFAULT_RENDERER).set_body_typed([]() {
+ return TerminalRenderer(std::cerr);
+});
-TVM_REGISTER_GLOBAL("diagnostics.GetRenderer").set_body_typed([]() { return GetRenderer(); });
+TVM_FFI_REGISTER_GLOBAL("diagnostics.GetRenderer").set_body_typed([]() { return GetRenderer(); });
-TVM_REGISTER_GLOBAL("diagnostics.ClearRenderer").set_body_typed([]() {
+TVM_FFI_REGISTER_GLOBAL("diagnostics.ClearRenderer").set_body_typed([]() {
tvm::ffi::Function::RemoveGlobal(OVERRIDE_RENDERER);
});
diff --git a/src/ir/env_func.cc b/src/ir/env_func.cc
index 9713f88f7d..ce40df21eb 100644
--- a/src/ir/env_func.cc
+++ b/src/ir/env_func.cc
@@ -20,8 +20,8 @@
/*!
* \file env_func.cc
*/
+#include <tvm/ffi/function.h>
#include <tvm/ir/env_func.h>
-#include <tvm/runtime/registry.h>
#include <tvm/tir/expr.h>
namespace tvm {
@@ -47,15 +47,15 @@ ObjectPtr<Object> CreateEnvNode(const std::string& name) {
EnvFunc EnvFunc::Get(const String& name) { return EnvFunc(CreateEnvNode(name)); }
-TVM_REGISTER_GLOBAL("ir.EnvFuncGet").set_body_typed(EnvFunc::Get);
+TVM_FFI_REGISTER_GLOBAL("ir.EnvFuncGet").set_body_typed(EnvFunc::Get);
-TVM_REGISTER_GLOBAL("ir.EnvFuncCall").set_body_packed([](ffi::PackedArgs args, ffi::Any* rv) {
+TVM_FFI_REGISTER_GLOBAL("ir.EnvFuncCall").set_body_packed([](ffi::PackedArgs args, ffi::Any* rv) {
EnvFunc env = args[0].cast<EnvFunc>();
ICHECK_GE(args.size(), 1);
env->func.CallPacked(args.Slice(1), rv);
});
-TVM_REGISTER_GLOBAL("ir.EnvFuncGetFunction").set_body_typed([](const EnvFunc& n) {
+TVM_FFI_REGISTER_GLOBAL("ir.EnvFuncGetFunction").set_body_typed([](const EnvFunc& n) {
return n->func;
});
diff --git a/src/ir/expr.cc b/src/ir/expr.cc
index de665dcd22..387572f642 100644
--- a/src/ir/expr.cc
+++ b/src/ir/expr.cc
@@ -22,9 +22,9 @@
* \brief The expression AST nodes for the common IR infra.
*/
#include <tvm/arith/analyzer.h>
+#include <tvm/ffi/function.h>
#include <tvm/ir/expr.h>
#include <tvm/ir/function.h>
-#include <tvm/runtime/registry.h>
#include <tvm/te/tensor.h>
#include <tvm/tir/expr.h>
@@ -64,7 +64,7 @@ IntImm::IntImm(DataType dtype, int64_t value, Span span) {
data_ = std::move(node);
}
-TVM_REGISTER_GLOBAL("ir.IntImm").set_body_typed([](DataType dtype, int64_t value, Span span) {
+TVM_FFI_REGISTER_GLOBAL("ir.IntImm").set_body_typed([](DataType dtype, int64_t value, Span span) {
return IntImm(dtype, value, span);
});
@@ -115,7 +115,7 @@ FloatImm::FloatImm(DataType dtype, double value, Span span) {
data_ = std::move(node);
}
-TVM_REGISTER_GLOBAL("ir.FloatImm").set_body_typed([](DataType dtype, double value, Span span) {
+TVM_FFI_REGISTER_GLOBAL("ir.FloatImm").set_body_typed([](DataType dtype, double value, Span span) {
return FloatImm(dtype, value, span);
});
@@ -128,9 +128,9 @@ Range Range::FromMinExtent(PrimExpr min, PrimExpr extent, Span span) {
return Range(make_object<RangeNode>(min, extent, span));
}
-TVM_REGISTER_GLOBAL("ir.Range_from_min_extent").set_body_typed(Range::FromMinExtent);
+TVM_FFI_REGISTER_GLOBAL("ir.Range_from_min_extent").set_body_typed(Range::FromMinExtent);
-TVM_REGISTER_GLOBAL("ir.Range")
+TVM_FFI_REGISTER_GLOBAL("ir.Range")
.set_body_typed([](PrimExpr begin, Optional<PrimExpr> end, Span span) -> Range {
if (end.defined()) {
return Range(begin, end.value(), span);
@@ -151,11 +151,11 @@ GlobalVar::GlobalVar(String name_hint, Type type, Span span) {
TVM_REGISTER_NODE_TYPE(GlobalVarNode);
-TVM_REGISTER_GLOBAL("ir.GlobalVar").set_body_typed([](String name, Type type) {
+TVM_FFI_REGISTER_GLOBAL("ir.GlobalVar").set_body_typed([](String name, Type type) {
return GlobalVar(name, type);
});
-TVM_REGISTER_GLOBAL("ir.DebugPrint").set_body_typed([](ObjectRef ref) {
+TVM_FFI_REGISTER_GLOBAL("ir.DebugPrint").set_body_typed([](ObjectRef ref) {
std::stringstream ss;
ss << ref;
return ss.str();
diff --git a/src/ir/function.cc b/src/ir/function.cc
index 8f543b0326..66d66e3c81 100644
--- a/src/ir/function.cc
+++ b/src/ir/function.cc
@@ -21,19 +21,21 @@
* \file src/ir/function.cc
* \brief The function data structure.
*/
+#include <tvm/ffi/function.h>
#include <tvm/ffi/rvalue_ref.h>
#include <tvm/ir/function.h>
#include <tvm/relax/expr.h>
-#include <tvm/runtime/registry.h>
#include <tvm/tir/function.h>
namespace tvm {
-TVM_REGISTER_GLOBAL("ir.BaseFunc_Attrs").set_body_typed([](BaseFunc func) { return func->attrs; });
+TVM_FFI_REGISTER_GLOBAL("ir.BaseFunc_Attrs").set_body_typed([](BaseFunc func) {
+ return func->attrs;
+});
-TVM_REGISTER_GLOBAL("ir.BaseFuncCopy").set_body_typed([](BaseFunc func) { return func; });
+TVM_FFI_REGISTER_GLOBAL("ir.BaseFuncCopy").set_body_typed([](BaseFunc func) { return func; });
-TVM_REGISTER_GLOBAL("ir.BaseFuncWithAttr")
+TVM_FFI_REGISTER_GLOBAL("ir.BaseFuncWithAttr")
.set_body_typed([](ffi::RValueRef<BaseFunc> func_ref, String key, Any value) -> BaseFunc {
BaseFunc func = *std::move(func_ref);
if (func->IsInstance<tir::PrimFuncNode>()) {
@@ -45,7 +47,7 @@ TVM_REGISTER_GLOBAL("ir.BaseFuncWithAttr")
}
});
-TVM_REGISTER_GLOBAL("ir.BaseFuncWithAttrs")
+TVM_FFI_REGISTER_GLOBAL("ir.BaseFuncWithAttrs")
.set_body_typed([](ffi::RValueRef<BaseFunc> func_ref,
Map<String, ffi::Any> attr_map) -> BaseFunc {
BaseFunc func = *std::move(func_ref);
@@ -61,7 +63,7 @@ TVM_REGISTER_GLOBAL("ir.BaseFuncWithAttrs")
TVM_FFI_UNREACHABLE();
});
-TVM_REGISTER_GLOBAL("ir.BaseFuncWithoutAttr")
+TVM_FFI_REGISTER_GLOBAL("ir.BaseFuncWithoutAttr")
.set_body_typed([](ffi::RValueRef<BaseFunc> func_ref, String key) -> BaseFunc {
BaseFunc func = *std::move(func_ref);
if (func->IsInstance<tir::PrimFuncNode>()) {
diff --git a/src/ir/global_info.cc b/src/ir/global_info.cc
index 6abac574e1..3df9ae00fb 100644
--- a/src/ir/global_info.cc
+++ b/src/ir/global_info.cc
@@ -25,7 +25,7 @@
#include <tvm/ir/global_info.h>
namespace tvm {
TVM_REGISTER_NODE_TYPE(DummyGlobalInfoNode);
-TVM_REGISTER_GLOBAL("ir.DummyGlobalInfo").set_body_typed([]() {
+TVM_FFI_REGISTER_GLOBAL("ir.DummyGlobalInfo").set_body_typed([]() {
auto n = DummyGlobalInfo(make_object<DummyGlobalInfoNode>());
return n;
});
@@ -39,7 +39,8 @@ VDevice::VDevice(Target tgt, int dev_id, MemoryScope mem_scope) {
}
TVM_REGISTER_NODE_TYPE(VDeviceNode);
-TVM_REGISTER_GLOBAL("ir.VDevice").set_body_typed([](Target tgt, int dev_id, MemoryScope mem_scope) {
- return VDevice(tgt, dev_id, mem_scope);
-});
+TVM_FFI_REGISTER_GLOBAL("ir.VDevice")
+ .set_body_typed([](Target tgt, int dev_id, MemoryScope mem_scope) {
+ return VDevice(tgt, dev_id, mem_scope);
+ });
} // namespace tvm
diff --git a/src/ir/global_var_supply.cc b/src/ir/global_var_supply.cc
index 3d3b891991..1b47c1a896 100644
--- a/src/ir/global_var_supply.cc
+++ b/src/ir/global_var_supply.cc
@@ -23,7 +23,7 @@
*/
#include "tvm/ir/global_var_supply.h"
-#include <tvm/runtime/registry.h>
+#include <tvm/ffi/function.h>
#include <utility>
@@ -92,24 +92,23 @@ GlobalVar GlobalVarSupplyNode::FreshGlobal(String name, bool add_prefix) {
TVM_REGISTER_NODE_TYPE(GlobalVarSupplyNode);
-TVM_REGISTER_GLOBAL("ir.GlobalVarSupply_NameSupply")
+TVM_FFI_REGISTER_GLOBAL("ir.GlobalVarSupply_NameSupply")
.set_body_typed([](const NameSupply& name_supply) { return GlobalVarSupply(name_supply); });
-TVM_REGISTER_GLOBAL("ir.GlobalVarSupply_IRModule").set_body_typed([](IRModule mod) {
+TVM_FFI_REGISTER_GLOBAL("ir.GlobalVarSupply_IRModule").set_body_typed([](IRModule mod) {
return GlobalVarSupply(std::move(mod));
});
-TVM_REGISTER_GLOBAL("ir.GlobalVarSupply_IRModules").set_body_typed([](const Array<IRModule>& mods) {
- return GlobalVarSupply(mods);
-});
+TVM_FFI_REGISTER_GLOBAL("ir.GlobalVarSupply_IRModules")
+ .set_body_typed([](const Array<IRModule>& mods) { return GlobalVarSupply(mods); });
-TVM_REGISTER_GLOBAL("ir.GlobalVarSupply_FreshGlobal")
+TVM_FFI_REGISTER_GLOBAL("ir.GlobalVarSupply_FreshGlobal")
.set_body_method(&GlobalVarSupplyNode::FreshGlobal);
-TVM_REGISTER_GLOBAL("ir.GlobalVarSupply_UniqueGlobalFor")
+TVM_FFI_REGISTER_GLOBAL("ir.GlobalVarSupply_UniqueGlobalFor")
.set_body_method(&GlobalVarSupplyNode::UniqueGlobalFor);
-TVM_REGISTER_GLOBAL("ir.GlobalVarSupply_ReserveGlobalVar")
+TVM_FFI_REGISTER_GLOBAL("ir.GlobalVarSupply_ReserveGlobalVar")
.set_body_method(&GlobalVarSupplyNode::ReserveGlobalVar);
} // namespace tvm
diff --git a/src/ir/instrument.cc b/src/ir/instrument.cc
index ad66f29448..a273245c1b 100644
--- a/src/ir/instrument.cc
+++ b/src/ir/instrument.cc
@@ -22,10 +22,10 @@
* \brief Infrastructure for instrumentation.
*/
#include <dmlc/thread_local.h>
+#include <tvm/ffi/function.h>
#include <tvm/ir/instrument.h>
#include <tvm/ir/transform.h>
#include <tvm/node/repr_printer.h>
-#include <tvm/runtime/registry.h>
#include <stack>
@@ -175,7 +175,7 @@ void BasePassInstrumentNode::RunAfterPass(const IRModule& ir_module,
TVM_REGISTER_NODE_TYPE(BasePassInstrumentNode);
-TVM_REGISTER_GLOBAL("instrument.PassInstrument")
+TVM_FFI_REGISTER_GLOBAL("instrument.PassInstrument")
.set_body_typed(
[](String name, ffi::TypedFunction<void()> enter_pass_ctx,
ffi::TypedFunction<void()> exit_pass_ctx,
@@ -308,9 +308,9 @@ String RenderPassProfiles() {
return os.str();
}
-TVM_REGISTER_GLOBAL("instrument.RenderTimePassProfiles").set_body_typed(RenderPassProfiles);
+TVM_FFI_REGISTER_GLOBAL("instrument.RenderTimePassProfiles").set_body_typed(RenderPassProfiles);
-TVM_REGISTER_GLOBAL("instrument.MakePassTimingInstrument").set_body_typed([]() {
+TVM_FFI_REGISTER_GLOBAL("instrument.MakePassTimingInstrument").set_body_typed([]() {
auto run_before_pass = [](const IRModule&, const transform::PassInfo& pass_info) {
PassProfile::EnterPass(pass_info->name);
return true;
diff --git a/src/ir/module.cc b/src/ir/module.cc
index 0223aef4d1..3166ffba97 100644
--- a/src/ir/module.cc
+++ b/src/ir/module.cc
@@ -21,12 +21,12 @@
* \brief The global module in TVM.
*/
#include <tvm/ffi/container/variant.h>
+#include <tvm/ffi/function.h>
#include <tvm/ffi/rvalue_ref.h>
#include <tvm/ir/global_var_supply.h>
#include <tvm/ir/module.h>
#include <tvm/ir/type_functor.h>
#include <tvm/node/structural_equal.h>
-#include <tvm/runtime/registry.h>
#include <algorithm>
#include <fstream>
@@ -242,7 +242,7 @@ IRModule IRModule::FromExpr(const RelaxExpr& expr,
TVM_REGISTER_NODE_TYPE(IRModuleNode);
-TVM_REGISTER_GLOBAL("ir.IRModule")
+TVM_FFI_REGISTER_GLOBAL("ir.IRModule")
.set_body_typed([](tvm::Map<GlobalVar, BaseFunc> funcs, tvm::ObjectRef attrs,
Map<String, Array<GlobalInfo>> global_infos) {
auto dict_attrs = [&attrs]() {
@@ -260,20 +260,20 @@ TVM_REGISTER_GLOBAL("ir.IRModule")
return IRModule(funcs, {}, dict_attrs, global_infos);
});
-TVM_REGISTER_GLOBAL("ir.Module_Clone").set_body_typed([](IRModule mod) -> IRModule {
+TVM_FFI_REGISTER_GLOBAL("ir.Module_Clone").set_body_typed([](IRModule mod) -> IRModule {
IRModule clone = mod;
clone.CopyOnWrite();
return clone;
});
-TVM_REGISTER_GLOBAL("ir.Module_Add")
+TVM_FFI_REGISTER_GLOBAL("ir.Module_Add")
.set_body_typed([](IRModule mod, GlobalVar var, ObjectRef val, bool update) -> IRModule {
ICHECK(val->IsInstance<RelaxExprNode>());
mod->Add(var, Downcast<BaseFunc>(val), update);
return mod;
});
-TVM_REGISTER_GLOBAL("ir.Module_Remove")
+TVM_FFI_REGISTER_GLOBAL("ir.Module_Remove")
.set_body_typed([](IRModule mod, Variant<String, GlobalVar> var) -> IRModule {
GlobalVar gvar = [&]() {
if (auto opt = var.as<GlobalVar>()) {
@@ -289,7 +289,7 @@ TVM_REGISTER_GLOBAL("ir.Module_Remove")
return mod;
});
-TVM_REGISTER_GLOBAL("ir.Module_Contains")
+TVM_FFI_REGISTER_GLOBAL("ir.Module_Contains")
.set_body_typed([](IRModule mod, Variant<String, GlobalVar> var) -> bool {
if (auto opt = var.as<GlobalVar>()) {
return mod->functions.count(opt.value());
@@ -301,55 +301,57 @@ TVM_REGISTER_GLOBAL("ir.Module_Contains")
}
});
-TVM_REGISTER_GLOBAL("ir.Module_GetGlobalVar").set_body_method(&IRModuleNode::GetGlobalVar);
+TVM_FFI_REGISTER_GLOBAL("ir.Module_GetGlobalVar").set_body_method(&IRModuleNode::GetGlobalVar);
-TVM_REGISTER_GLOBAL("ir.Module_GetGlobalVars").set_body_method(&IRModuleNode::GetGlobalVars);
+TVM_FFI_REGISTER_GLOBAL("ir.Module_GetGlobalVars").set_body_method(&IRModuleNode::GetGlobalVars);
-TVM_REGISTER_GLOBAL("ir.Module_ContainGlobalVar").set_body_method(&IRModuleNode::ContainGlobalVar);
+TVM_FFI_REGISTER_GLOBAL("ir.Module_ContainGlobalVar")
+ .set_body_method(&IRModuleNode::ContainGlobalVar);
-TVM_REGISTER_GLOBAL("ir.Module_Lookup").set_body_typed([](IRModule mod, GlobalVar var) {
+TVM_FFI_REGISTER_GLOBAL("ir.Module_Lookup").set_body_typed([](IRModule mod, GlobalVar var) {
return mod->Lookup(var);
});
-TVM_REGISTER_GLOBAL("ir.Module_Lookup_str").set_body_typed([](IRModule mod, String var) {
+TVM_FFI_REGISTER_GLOBAL("ir.Module_Lookup_str").set_body_typed([](IRModule mod, String var) {
return mod->Lookup(var);
});
-TVM_REGISTER_GLOBAL("ir.Module_FromExpr").set_body_typed(&IRModule::FromExpr);
+TVM_FFI_REGISTER_GLOBAL("ir.Module_FromExpr").set_body_typed(&IRModule::FromExpr);
-TVM_REGISTER_GLOBAL("ir.Module_Update").set_body_typed([](IRModule mod, IRModule from) {
+TVM_FFI_REGISTER_GLOBAL("ir.Module_Update").set_body_typed([](IRModule mod, IRModule from) {
mod->Update(from);
});
-TVM_REGISTER_GLOBAL("ir.Module_UpdateFunction")
+TVM_FFI_REGISTER_GLOBAL("ir.Module_UpdateFunction")
.set_body_typed([](IRModule mod, GlobalVar gv, BaseFunc func) { mod->Update(gv, func); });
-TVM_REGISTER_GLOBAL("ir.Module_UpdateGlobalInfo")
+TVM_FFI_REGISTER_GLOBAL("ir.Module_UpdateGlobalInfo")
.set_body_typed([](IRModule mod, String name, Array<GlobalInfo> global_info) {
mod->UpdateGlobalInfo(name, global_info);
});
-TVM_REGISTER_GLOBAL("ir.Module_GetAttrs").set_body_typed([](IRModule mod) -> ObjectRef {
+TVM_FFI_REGISTER_GLOBAL("ir.Module_GetAttrs").set_body_typed([](IRModule mod) -> ObjectRef {
return mod->GetAttrs();
});
-TVM_REGISTER_GLOBAL("ir.Module_WithAttr")
+TVM_FFI_REGISTER_GLOBAL("ir.Module_WithAttr")
.set_body_typed([](ffi::RValueRef<IRModule> mod, String key, ffi::Any value) -> IRModule {
return WithAttr(*std::move(mod), key, value);
});
-TVM_REGISTER_GLOBAL("ir.Module_WithoutAttr")
+TVM_FFI_REGISTER_GLOBAL("ir.Module_WithoutAttr")
.set_body_typed([](ffi::RValueRef<IRModule> mod, String key) -> IRModule {
return WithoutAttr(*std::move(mod), key);
});
-TVM_REGISTER_GLOBAL("ir.Module_WithAttrs")
+TVM_FFI_REGISTER_GLOBAL("ir.Module_WithAttrs")
.set_body_typed([](ffi::RValueRef<IRModule> mod, Map<String, ffi::Any> attr_map) -> IRModule {
return WithAttrs(*std::move(mod), attr_map);
});
-TVM_REGISTER_GLOBAL("ir.Module_GetAttr").set_body_typed([](IRModule mod, String key) -> ObjectRef {
- return mod->GetAttr<ObjectRef>(key);
-});
+TVM_FFI_REGISTER_GLOBAL("ir.Module_GetAttr")
+ .set_body_typed([](IRModule mod, String key) -> ObjectRef {
+ return mod->GetAttr<ObjectRef>(key);
+ });
} // namespace tvm
diff --git a/src/ir/name_supply.cc b/src/ir/name_supply.cc
index 087fc82a50..e73b0e63e3 100644
--- a/src/ir/name_supply.cc
+++ b/src/ir/name_supply.cc
@@ -23,7 +23,7 @@
*/
#include "tvm/ir/name_supply.h"
-#include <tvm/runtime/registry.h>
+#include <tvm/ffi/function.h>
#include <utility>
@@ -92,14 +92,15 @@ std::string NameSupplyNode::GetUniqueName(std::string name, bool add_underscore)
TVM_REGISTER_NODE_TYPE(NameSupplyNode);
-TVM_REGISTER_GLOBAL("ir.NameSupply").set_body_typed([](String prefix) {
+TVM_FFI_REGISTER_GLOBAL("ir.NameSupply").set_body_typed([](String prefix) {
return NameSupply(prefix);
});
-TVM_REGISTER_GLOBAL("ir.NameSupply_FreshName").set_body_method(&NameSupplyNode::FreshName);
+TVM_FFI_REGISTER_GLOBAL("ir.NameSupply_FreshName").set_body_method(&NameSupplyNode::FreshName);
-TVM_REGISTER_GLOBAL("ir.NameSupply_ReserveName").set_body_method(&NameSupplyNode::ReserveName);
+TVM_FFI_REGISTER_GLOBAL("ir.NameSupply_ReserveName").set_body_method(&NameSupplyNode::ReserveName);
-TVM_REGISTER_GLOBAL("ir.NameSupply_ContainsName").set_body_method(&NameSupplyNode::ContainsName);
+TVM_FFI_REGISTER_GLOBAL("ir.NameSupply_ContainsName")
+ .set_body_method(&NameSupplyNode::ContainsName);
} // namespace tvm
diff --git a/src/ir/op.cc b/src/ir/op.cc
index 70f7528e5e..4917f8336b 100644
--- a/src/ir/op.cc
+++ b/src/ir/op.cc
@@ -75,13 +75,13 @@ void OpRegEntry::UpdateAttr(const String& key, ffi::Any value, int plevel) {
}
// Frontend APIs
-TVM_REGISTER_GLOBAL("ir.ListOpNames").set_body_typed([]() {
+TVM_FFI_REGISTER_GLOBAL("ir.ListOpNames").set_body_typed([]() {
return OpRegistry::Global()->ListAllNames();
});
-TVM_REGISTER_GLOBAL("ir.GetOp").set_body_typed([](String name) -> Op { return Op::Get(name); });
+TVM_FFI_REGISTER_GLOBAL("ir.GetOp").set_body_typed([](String name) -> Op { return Op::Get(name); });
-TVM_REGISTER_GLOBAL("ir.OpGetAttr").set_body_typed([](Op op, String attr_name) -> ffi::Any {
+TVM_FFI_REGISTER_GLOBAL("ir.OpGetAttr").set_body_typed([](Op op, String attr_name) -> ffi::Any {
auto op_map = Op::GetAttrMap<ffi::Any>(attr_name);
ffi::Any rv;
if (op_map.count(op)) {
@@ -90,50 +90,50 @@ TVM_REGISTER_GLOBAL("ir.OpGetAttr").set_body_typed([](Op op, String attr_name) -
return rv;
});
-TVM_REGISTER_GLOBAL("ir.OpHasAttr").set_body_typed([](Op op, String attr_name) -> bool {
+TVM_FFI_REGISTER_GLOBAL("ir.OpHasAttr").set_body_typed([](Op op, String attr_name) -> bool {
return Op::HasAttrMap(attr_name);
});
-TVM_REGISTER_GLOBAL("ir.OpSetAttr")
+TVM_FFI_REGISTER_GLOBAL("ir.OpSetAttr")
.set_body_typed([](Op op, String attr_name, ffi::AnyView value, int plevel) {
auto& reg = OpRegistry::Global()->RegisterOrGet(op->name).set_name();
reg.set_attr(attr_name, value, plevel);
});
-TVM_REGISTER_GLOBAL("ir.OpResetAttr").set_body_typed([](Op op, String attr_name) {
+TVM_FFI_REGISTER_GLOBAL("ir.OpResetAttr").set_body_typed([](Op op, String attr_name) {
auto& reg = OpRegistry::Global()->RegisterOrGet(op->name);
reg.reset_attr(attr_name);
});
-TVM_REGISTER_GLOBAL("ir.RegisterOp").set_body_typed([](String op_name, String descr) {
+TVM_FFI_REGISTER_GLOBAL("ir.RegisterOp").set_body_typed([](String op_name, String descr) {
const OpRegEntry* reg = OpRegistry::Global()->Get(op_name);
ICHECK(reg == nullptr) << "AttributeError: Operator " << op_name << " is registered before";
auto& op = OpRegistry::Global()->RegisterOrGet(op_name).set_name();
op.describe(descr);
});
-TVM_REGISTER_GLOBAL("ir.OpAddArgument")
+TVM_FFI_REGISTER_GLOBAL("ir.OpAddArgument")
.set_body_typed([](Op op, String name, String type, String description) {
auto& reg = OpRegistry::Global()->RegisterOrGet(op->name).set_name();
reg.add_argument(name, type, description);
});
-TVM_REGISTER_GLOBAL("ir.OpSetSupportLevel").set_body_typed([](Op op, int level) {
+TVM_FFI_REGISTER_GLOBAL("ir.OpSetSupportLevel").set_body_typed([](Op op, int level) {
auto& reg = OpRegistry::Global()->RegisterOrGet(op->name).set_name();
reg.set_support_level(level);
});
-TVM_REGISTER_GLOBAL("ir.OpSetNumInputs").set_body_typed([](Op op, int n) {
+TVM_FFI_REGISTER_GLOBAL("ir.OpSetNumInputs").set_body_typed([](Op op, int n) {
auto& reg = OpRegistry::Global()->RegisterOrGet(op->name).set_name();
reg.set_num_inputs(n);
});
-TVM_REGISTER_GLOBAL("ir.OpSetAttrsTypeKey").set_body_typed([](Op op, String key) {
+TVM_FFI_REGISTER_GLOBAL("ir.OpSetAttrsTypeKey").set_body_typed([](Op op, String key) {
auto& reg = OpRegistry::Global()->RegisterOrGet(op->name).set_name();
reg.set_attrs_type_key(key);
});
-TVM_REGISTER_GLOBAL("ir.RegisterOpAttr")
+TVM_FFI_REGISTER_GLOBAL("ir.RegisterOpAttr")
.set_body_typed([](String op_name, String attr_key, ffi::AnyView value, int plevel) {
auto& reg = OpRegistry::Global()->RegisterOrGet(op_name).set_name();
// enable resgiteration and override of certain properties
@@ -146,7 +146,7 @@ TVM_REGISTER_GLOBAL("ir.RegisterOpAttr")
}
});
-TVM_REGISTER_GLOBAL("ir.RegisterOpLowerIntrinsic")
+TVM_FFI_REGISTER_GLOBAL("ir.RegisterOpLowerIntrinsic")
.set_body_typed([](String name, ffi::Function f, String target, int plevel) {
tvm::OpRegEntry::RegisterOrGet(name).set_attr<FLowerIntrinsic>(target + ".FLowerIntrinsic", f,
plevel);
diff --git a/src/ir/replace_global_vars.cc b/src/ir/replace_global_vars.cc
index 44a6a22db7..0dca973024 100644
--- a/src/ir/replace_global_vars.cc
+++ b/src/ir/replace_global_vars.cc
@@ -62,7 +62,7 @@ IRModule ReplaceGlobalVars(IRModule mod, Map<GlobalVar, GlobalVar> replacements)
return mod;
}
-TVM_REGISTER_GLOBAL("transform.ReplaceGlobalVars").set_body_typed(ReplaceGlobalVars);
+TVM_FFI_REGISTER_GLOBAL("transform.ReplaceGlobalVars").set_body_typed(ReplaceGlobalVars);
IRModule ModuleReplaceGlobalVars(
IRModule mod, Map<Variant<String, GlobalVar>, Variant<String, GlobalVar>> replacements) {
@@ -93,7 +93,7 @@ IRModule ModuleReplaceGlobalVars(
return ReplaceGlobalVars(mod, gvar_replacements);
}
-TVM_REGISTER_GLOBAL("ir.Module_ReplaceGlobalVars").set_body_typed(ModuleReplaceGlobalVars);
+TVM_FFI_REGISTER_GLOBAL("ir.Module_ReplaceGlobalVars").set_body_typed(ModuleReplaceGlobalVars);
} // namespace transform
} // namespace tvm
diff --git a/src/ir/source_map.cc b/src/ir/source_map.cc
index 8e25b25a4c..482e1dfa10 100644
--- a/src/ir/source_map.cc
+++ b/src/ir/source_map.cc
@@ -20,9 +20,9 @@
* \file source_map.cc
* \brief The implementation of the source map data structure.
*/
+#include <tvm/ffi/function.h>
#include <tvm/ir/source_map.h>
#include <tvm/ir/transform.h>
-#include <tvm/runtime/registry.h>
#include <algorithm>
@@ -50,7 +50,7 @@ ObjectPtr<Object> GetSourceNameNodeByStr(const std::string& name) {
SourceName SourceName::Get(const String& name) { return SourceName(GetSourceNameNode(name)); }
-TVM_REGISTER_GLOBAL("ir.SourceName").set_body_typed(SourceName::Get);
+TVM_FFI_REGISTER_GLOBAL("ir.SourceName").set_body_typed(SourceName::Get);
TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
.set_dispatch<SourceNameNode>([](const ObjectRef& ref, ReprPrinter* p) {
@@ -129,12 +129,12 @@ SequentialSpan::SequentialSpan(std::initializer_list<Span> init) {
TVM_REGISTER_NODE_TYPE(SequentialSpanNode);
-TVM_REGISTER_GLOBAL("ir.Span").set_body_typed([](SourceName source_name, int line, int end_line,
- int column, int end_column) {
+TVM_FFI_REGISTER_GLOBAL("ir.Span").set_body_typed([](SourceName source_name, int line, int end_line,
+ int column, int end_column) {
return Span(source_name, line, end_line, column, end_column);
});
-TVM_REGISTER_GLOBAL("ir.SequentialSpan").set_body_typed([](tvm::Array<Span> spans) {
+TVM_FFI_REGISTER_GLOBAL("ir.SequentialSpan").set_body_typed([](tvm::Array<Span> spans) {
return SequentialSpan(spans);
});
@@ -218,11 +218,12 @@ SourceMap::SourceMap(Map<SourceName, Source> source_map) {
void SourceMap::Add(const Source& source) { (*this)->source_map.Set(source->source_name, source); }
-TVM_REGISTER_GLOBAL("SourceMapAdd").set_body_typed([](SourceMap map, String name, String content) {
- auto src_name = SourceName::Get(name);
- Source source(src_name, content);
- map.Add(source);
- return src_name;
-});
+TVM_FFI_REGISTER_GLOBAL("SourceMapAdd")
+ .set_body_typed([](SourceMap map, String name, String content) {
+ auto src_name = SourceName::Get(name);
+ Source source(src_name, content);
+ map.Add(source);
+ return src_name;
+ });
} // namespace tvm
diff --git a/src/ir/transform.cc b/src/ir/transform.cc
index 44dafcbf1d..db4e47ca0d 100644
--- a/src/ir/transform.cc
+++ b/src/ir/transform.cc
@@ -22,6 +22,7 @@
* \brief Infrastructure for transformation passes.
*/
#include <dmlc/thread_local.h>
+#include <tvm/ffi/function.h>
#include <tvm/ffi/rvalue_ref.h>
#include <tvm/ir/transform.h>
#include <tvm/node/repr_printer.h>
@@ -29,14 +30,12 @@
#include <tvm/relax/expr.h>
#include <tvm/relax/tuning_api.h>
#include <tvm/runtime/device_api.h>
-#include <tvm/runtime/registry.h>
#include <chrono>
#include <iomanip>
#include <stack>
#include <unordered_set>
-#include "../runtime/object_internal.h"
#include "../runtime/regex.h"
namespace tvm {
@@ -532,12 +531,12 @@ Pass CreateModulePass(std::function<IRModule(IRModule, PassContext)> pass_func,
TVM_REGISTER_NODE_TYPE(PassInfoNode);
-TVM_REGISTER_GLOBAL("transform.PassInfo")
+TVM_FFI_REGISTER_GLOBAL("transform.PassInfo")
.set_body_typed([](int opt_level, String name, tvm::Array<String> required, bool traceable) {
return PassInfo(opt_level, name, required, traceable);
});
-TVM_REGISTER_GLOBAL("transform.Info").set_body_packed([](ffi::PackedArgs args, ffi::Any* ret) {
+TVM_FFI_REGISTER_GLOBAL("transform.Info").set_body_packed([](ffi::PackedArgs args, ffi::Any* ret) {
Pass pass = args[0].cast<Pass>();
*ret = pass->Info();
});
@@ -562,7 +561,7 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
TVM_REGISTER_NODE_TYPE(ModulePassNode);
-TVM_REGISTER_GLOBAL("transform.MakeModulePass")
+TVM_FFI_REGISTER_GLOBAL("transform.MakeModulePass")
.set_body_typed(
[](ffi::TypedFunction<IRModule(ffi::RValueRef<IRModule>, PassContext)> pass_func,
PassInfo pass_info) {
@@ -572,7 +571,7 @@ TVM_REGISTER_GLOBAL("transform.MakeModulePass")
return ModulePass(wrapped_pass_func, pass_info);
});
-TVM_REGISTER_GLOBAL("transform.RunPass")
+TVM_FFI_REGISTER_GLOBAL("transform.RunPass")
.set_body_typed([](Pass pass, ffi::RValueRef<IRModule> mod) { return pass(*std::move(mod)); });
TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
@@ -585,7 +584,7 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
TVM_REGISTER_NODE_TYPE(SequentialNode);
-TVM_REGISTER_GLOBAL("transform.Sequential")
+TVM_FFI_REGISTER_GLOBAL("transform.Sequential")
.set_body_packed([](ffi::PackedArgs args, ffi::Any* ret) {
auto passes = args[0].cast<tvm::Array<Pass>>();
int opt_level = args[1].cast<int>();
@@ -612,7 +611,7 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
TVM_REGISTER_NODE_TYPE(PassContextNode);
-TVM_REGISTER_GLOBAL("transform.PassContext")
+TVM_FFI_REGISTER_GLOBAL("transform.PassContext")
.set_body_typed([](int opt_level, Array<String> required, Array<String> disabled,
Array<instrument::PassInstrument> instruments,
Optional<Map<String, ffi::Any>> config, Array<ObjectRef> trace_stack,
@@ -658,24 +657,27 @@ class PassContext::Internal {
static void ExitScope(PassContext pass_ctx) { pass_ctx.ExitWithScope(); }
};
-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")
+TVM_FFI_REGISTER_GLOBAL("transform.GetTraceStack").set_body_method(&PassContextNode::GetTraceStack);
+TVM_FFI_REGISTER_GLOBAL("transform.PushTrace").set_body_method(&PassContextNode::PushTrace);
+TVM_FFI_REGISTER_GLOBAL("transform.PopTrace").set_body_method(&PassContextNode::PopTrace);
+TVM_FFI_REGISTER_GLOBAL("transform.GetTraceStackSize")
.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")
+TVM_FFI_REGISTER_GLOBAL("transform.GetCurrentTrace")
+ .set_body_method(&PassContextNode::GetCurrentTrace);
+TVM_FFI_REGISTER_GLOBAL("transform.SetNumEvals").set_body_method(&PassContextNode::SetNumEvals);
+TVM_FFI_REGISTER_GLOBAL("transform.IncNumEvals").set_body_method(&PassContextNode::IncNumEvals);
+TVM_FFI_REGISTER_GLOBAL("transform.GetTuningAPIDatabase")
.set_body_method(&PassContextNode::GetTuningAPIDatabase);
-TVM_REGISTER_GLOBAL("transform.GetCurrentPassContext").set_body_typed(PassContext::Current);
+TVM_FFI_REGISTER_GLOBAL("transform.GetCurrentPassContext").set_body_typed(PassContext::Current);
-TVM_REGISTER_GLOBAL("transform.EnterPassContext").set_body_typed(PassContext::Internal::EnterScope);
+TVM_FFI_REGISTER_GLOBAL("transform.EnterPassContext")
+ .set_body_typed(PassContext::Internal::EnterScope);
-TVM_REGISTER_GLOBAL("transform.ExitPassContext").set_body_typed(PassContext::Internal::ExitScope);
+TVM_FFI_REGISTER_GLOBAL("transform.ExitPassContext")
+ .set_body_typed(PassContext::Internal::ExitScope);
-TVM_REGISTER_GLOBAL("transform.OverrideInstruments")
+TVM_FFI_REGISTER_GLOBAL("transform.OverrideInstruments")
.set_body_typed([](PassContext pass_ctx, Array<instrument::PassInstrument> instruments) {
pass_ctx.InstrumentExitPassContext();
pass_ctx->instruments = instruments;
@@ -690,9 +692,9 @@ Pass PrintIR(String header, bool show_meta_data) {
return CreateModulePass(pass_func, 0, "PrintIR", {}, /* traceable */ false);
}
-TVM_REGISTER_GLOBAL("transform.PrintIR").set_body_typed(PrintIR);
+TVM_FFI_REGISTER_GLOBAL("transform.PrintIR").set_body_typed(PrintIR);
-TVM_REGISTER_GLOBAL("transform.ListConfigs").set_body_typed(PassContext::ListConfigs);
+TVM_FFI_REGISTER_GLOBAL("transform.ListConfigs").set_body_typed(PassContext::ListConfigs);
} // namespace transform
} // namespace tvm
diff --git a/src/ir/type.cc b/src/ir/type.cc
index 3c648418c6..8bc48a1114 100644
--- a/src/ir/type.cc
+++ b/src/ir/type.cc
@@ -21,8 +21,8 @@
* \file src/ir/type.cc
* \brief Common type system AST nodes throughout the IR.
*/
+#include <tvm/ffi/function.h>
#include <tvm/ir/type.h>
-#include <tvm/runtime/registry.h>
namespace tvm {
PrimType::PrimType(runtime::DataType dtype, Span span) {
@@ -34,7 +34,7 @@ PrimType::PrimType(runtime::DataType dtype, Span span) {
TVM_REGISTER_NODE_TYPE(PrimTypeNode);
-TVM_REGISTER_GLOBAL("ir.PrimType").set_body_typed([](runtime::DataType dtype) {
+TVM_FFI_REGISTER_GLOBAL("ir.PrimType").set_body_typed([](runtime::DataType dtype) {
return PrimType(dtype);
});
@@ -47,7 +47,7 @@ PointerType::PointerType(Type element_type, String storage_scope) {
TVM_REGISTER_NODE_TYPE(PointerTypeNode);
-TVM_REGISTER_GLOBAL("ir.PointerType")
+TVM_FFI_REGISTER_GLOBAL("ir.PointerType")
.set_body_typed([](Type element_type, String storage_scope = "") {
return PointerType(element_type, storage_scope);
});
@@ -62,9 +62,10 @@ FuncType::FuncType(tvm::Array<Type> arg_types, Type ret_type, Span span) {
TVM_REGISTER_NODE_TYPE(FuncTypeNode);
-TVM_REGISTER_GLOBAL("ir.FuncType").set_body_typed([](tvm::Array<Type> arg_types, Type ret_type) {
- return FuncType(arg_types, ret_type);
-});
+TVM_FFI_REGISTER_GLOBAL("ir.FuncType")
+ .set_body_typed([](tvm::Array<Type> arg_types, Type ret_type) {
+ return FuncType(arg_types, ret_type);
+ });
TupleType::TupleType(Array<Type> fields, Span span) {
ObjectPtr<TupleTypeNode> n = make_object<TupleTypeNode>();
@@ -77,7 +78,7 @@ TupleType TupleType::Empty() { return TupleType(Array<Type>()); }
TVM_REGISTER_NODE_TYPE(TupleTypeNode);
-TVM_REGISTER_GLOBAL("ir.TupleType").set_body_typed([](Array<Type> fields) {
+TVM_FFI_REGISTER_GLOBAL("ir.TupleType").set_body_typed([](Array<Type> fields) {
return TupleType(fields);
});
diff --git a/src/meta_schedule/arg_info.cc b/src/meta_schedule/arg_info.cc
index bc0fef1ff4..58c4a7b33c 100644
--- a/src/meta_schedule/arg_info.cc
+++ b/src/meta_schedule/arg_info.cc
@@ -162,11 +162,12 @@ 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(&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);
-TVM_REGISTER_GLOBAL("meta_schedule.TensorInfo")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.ArgInfoAsJSON").set_body_method(&ArgInfoNode::AsJSON);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.ArgInfoFromPrimFunc").set_body_typed(ArgInfo::FromPrimFunc);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.ArgInfoFromEntryFunc")
+ .set_body_typed(ArgInfo::FromEntryFunc);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.ArgInfoFromJSON").set_body_typed(ArgInfo::FromJSON);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.TensorInfo")
.set_body_typed([](runtime::DataType dtype, ffi::Shape shape) -> TensorInfo {
return TensorInfo(dtype, shape);
});
diff --git a/src/meta_schedule/builder/builder.cc b/src/meta_schedule/builder/builder.cc
index 9d725e91e2..85e189e732 100644
--- a/src/meta_schedule/builder/builder.cc
+++ b/src/meta_schedule/builder/builder.cc
@@ -52,21 +52,21 @@ TVM_REGISTER_NODE_TYPE(BuilderResultNode);
TVM_REGISTER_OBJECT_TYPE(BuilderNode);
TVM_REGISTER_NODE_TYPE(PyBuilderNode);
-TVM_REGISTER_GLOBAL("meta_schedule.BuilderInput")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.BuilderInput")
.set_body_typed([](IRModule mod, Target target,
Optional<Map<String, runtime::NDArray>> params) -> BuilderInput {
return BuilderInput(mod, target, params);
});
-TVM_REGISTER_GLOBAL("meta_schedule.BuilderResult")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.BuilderResult")
.set_body_typed([](Optional<String> artifact_path,
Optional<String> error_msg) -> BuilderResult {
return BuilderResult(artifact_path, error_msg);
});
-TVM_REGISTER_GLOBAL("meta_schedule.BuilderBuild").set_body_method(&BuilderNode::Build);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.BuilderBuild").set_body_method(&BuilderNode::Build);
-TVM_REGISTER_GLOBAL("meta_schedule.BuilderPyBuilder").set_body_typed(Builder::PyBuilder);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.BuilderPyBuilder").set_body_typed(Builder::PyBuilder);
} // namespace meta_schedule
} // namespace tvm
diff --git a/src/meta_schedule/cost_model/cost_model.cc b/src/meta_schedule/cost_model/cost_model.cc
index 1d28eb19d7..5c1c7a5685 100644
--- a/src/meta_schedule/cost_model/cost_model.cc
+++ b/src/meta_schedule/cost_model/cost_model.cc
@@ -71,10 +71,10 @@ 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(&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")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.CostModelLoad").set_body_method(&CostModelNode::Load);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.CostModelSave").set_body_method(&CostModelNode::Save);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.CostModelUpdate").set_body_method(&CostModelNode::Update);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.CostModelPredict")
.set_body_typed([](CostModel model, //
const TuneContext& context, //
Array<MeasureCandidate> candidates, //
@@ -82,7 +82,8 @@ TVM_REGISTER_GLOBAL("meta_schedule.CostModelPredict")
std::vector<double> result = model->Predict(context, candidates);
std::copy(result.begin(), result.end(), static_cast<double*>(p_addr));
});
-TVM_REGISTER_GLOBAL("meta_schedule.CostModelPyCostModel").set_body_typed(CostModel::PyCostModel);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.CostModelPyCostModel")
+ .set_body_typed(CostModel::PyCostModel);
} // namespace meta_schedule
} // namespace tvm
diff --git a/src/meta_schedule/database/database.cc b/src/meta_schedule/database/database.cc
index 9c23d9e845..034294eedc 100644
--- a/src/meta_schedule/database/database.cc
+++ b/src/meta_schedule/database/database.cc
@@ -282,43 +282,46 @@ TVM_REGISTER_NODE_TYPE(WorkloadNode);
TVM_REGISTER_NODE_TYPE(TuningRecordNode);
TVM_REGISTER_OBJECT_TYPE(DatabaseNode);
TVM_REGISTER_NODE_TYPE(PyDatabaseNode);
-TVM_REGISTER_GLOBAL("meta_schedule.Workload").set_body_typed([](IRModule mod) {
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.Workload").set_body_typed([](IRModule mod) {
return Workload(mod);
});
-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")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.WorkloadAsJSON").set_body_method(&WorkloadNode::AsJSON);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.WorkloadFromJSON").set_body_typed(&Workload::FromJSON);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.TuningRecord")
.set_body_typed([](tir::Trace trace, Workload workload, Optional<Array<FloatImm>> run_secs,
Optional<Target> target, Optional<Array<ArgInfo>> args_info) {
return TuningRecord(trace, workload, run_secs, target, args_info);
});
-TVM_REGISTER_GLOBAL("meta_schedule.TuningRecordAsMeasureCandidate")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.TuningRecordAsMeasureCandidate")
.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")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.TuningRecordAsJSON")
+ .set_body_method(&TuningRecordNode::AsJSON);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.TuningRecordFromJSON")
+ .set_body_typed(TuningRecord::FromJSON);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.DatabaseEnterWithScope")
.set_body_method(&Database::EnterWithScope);
-TVM_REGISTER_GLOBAL("meta_schedule.DatabaseExitWithScope")
+TVM_FFI_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")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.DatabaseCurrent").set_body_typed(Database::Current);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.DatabaseHasWorkload")
.set_body_method(&DatabaseNode::HasWorkload);
-TVM_REGISTER_GLOBAL("meta_schedule.DatabaseCommitWorkload")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.DatabaseCommitWorkload")
.set_body_method(&DatabaseNode::CommitWorkload);
-TVM_REGISTER_GLOBAL("meta_schedule.DatabaseCommitTuningRecord")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.DatabaseCommitTuningRecord")
.set_body_method(&DatabaseNode::CommitTuningRecord);
-TVM_REGISTER_GLOBAL("meta_schedule.DatabaseGetTopK").set_body_method(&DatabaseNode::GetTopK);
-TVM_REGISTER_GLOBAL("meta_schedule.DatabaseGetAllTuningRecords")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.DatabaseGetTopK").set_body_method(&DatabaseNode::GetTopK);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.DatabaseGetAllTuningRecords")
.set_body_method(&DatabaseNode::GetAllTuningRecords);
-TVM_REGISTER_GLOBAL("meta_schedule.DatabaseSize").set_body_method(&DatabaseNode::Size);
-TVM_REGISTER_GLOBAL("meta_schedule.DatabaseQueryTuningRecord")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.DatabaseSize").set_body_method(&DatabaseNode::Size);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.DatabaseQueryTuningRecord")
.set_body_method(&DatabaseNode::QueryTuningRecord);
-TVM_REGISTER_GLOBAL("meta_schedule.DatabaseQuerySchedule")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.DatabaseQuerySchedule")
.set_body_method(&DatabaseNode::QuerySchedule);
-TVM_REGISTER_GLOBAL("meta_schedule.DatabaseQueryIRModule")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.DatabaseQueryIRModule")
.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);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.DatabaseDumpPruned")
+ .set_body_method(&DatabaseNode::DumpPruned);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.DatabasePyDatabase").set_body_typed(Database::PyDatabase);
} // namespace meta_schedule
} // namespace tvm
diff --git a/src/meta_schedule/database/json_database.cc b/src/meta_schedule/database/json_database.cc
index a59b334b22..2a6b93f8cb 100644
--- a/src/meta_schedule/database/json_database.cc
+++ b/src/meta_schedule/database/json_database.cc
@@ -214,7 +214,8 @@ Database Database::JSONDatabase(String path_workload, String path_tuning_record,
}
TVM_REGISTER_NODE_TYPE(JSONDatabaseNode);
-TVM_REGISTER_GLOBAL("meta_schedule.DatabaseJSONDatabase").set_body_typed(Database::JSONDatabase);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.DatabaseJSONDatabase")
+ .set_body_typed(Database::JSONDatabase);
} // namespace meta_schedule
} // namespace tvm
diff --git a/src/meta_schedule/database/memory_database.cc b/src/meta_schedule/database/memory_database.cc
index 3d418206b0..cbc811752c 100644
--- a/src/meta_schedule/database/memory_database.cc
+++ b/src/meta_schedule/database/memory_database.cc
@@ -97,7 +97,7 @@ Database Database::MemoryDatabase(String mod_eq_name) {
}
TVM_REGISTER_NODE_TYPE(MemoryDatabaseNode);
-TVM_REGISTER_GLOBAL("meta_schedule.DatabaseMemoryDatabase")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.DatabaseMemoryDatabase")
.set_body_typed(Database::MemoryDatabase);
} // namespace meta_schedule
diff --git a/src/meta_schedule/database/ordered_union_database.cc b/src/meta_schedule/database/ordered_union_database.cc
index df151764a6..87f5c03a71 100644
--- a/src/meta_schedule/database/ordered_union_database.cc
+++ b/src/meta_schedule/database/ordered_union_database.cc
@@ -79,7 +79,7 @@ Database Database::OrderedUnionDatabase(Array<Database> databases) {
}
TVM_REGISTER_NODE_TYPE(OrderedUnionDatabaseNode);
-TVM_REGISTER_GLOBAL("meta_schedule.DatabaseOrderedUnionDatabase")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.DatabaseOrderedUnionDatabase")
.set_body_typed(Database::OrderedUnionDatabase);
} // namespace meta_schedule
diff --git a/src/meta_schedule/database/schedule_fn_database.cc b/src/meta_schedule/database/schedule_fn_database.cc
index 0dbac9616c..c66ec5f4f0 100644
--- a/src/meta_schedule/database/schedule_fn_database.cc
+++ b/src/meta_schedule/database/schedule_fn_database.cc
@@ -99,7 +99,7 @@ Database Database::ScheduleFnDatabase(ffi::TypedFunction<bool(tir::Schedule)> sc
}
TVM_REGISTER_NODE_TYPE(ScheduleFnDatabaseNode);
-TVM_REGISTER_GLOBAL("meta_schedule.DatabaseScheduleFnDatabase")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.DatabaseScheduleFnDatabase")
.set_body_typed(Database::ScheduleFnDatabase);
} // namespace meta_schedule
diff --git a/src/meta_schedule/database/union_database.cc b/src/meta_schedule/database/union_database.cc
index 876affbdb2..2bc82b459c 100644
--- a/src/meta_schedule/database/union_database.cc
+++ b/src/meta_schedule/database/union_database.cc
@@ -82,7 +82,8 @@ Database Database::UnionDatabase(Array<Database> databases) {
}
TVM_REGISTER_NODE_TYPE(UnionDatabaseNode);
-TVM_REGISTER_GLOBAL("meta_schedule.DatabaseUnionDatabase").set_body_typed(Database::UnionDatabase);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.DatabaseUnionDatabase")
+ .set_body_typed(Database::UnionDatabase);
} // namespace meta_schedule
} // namespace tvm
diff --git a/src/meta_schedule/extracted_task.cc b/src/meta_schedule/extracted_task.cc
index ec04361f51..fb26e6eb69 100644
--- a/src/meta_schedule/extracted_task.cc
+++ b/src/meta_schedule/extracted_task.cc
@@ -39,7 +39,7 @@ ExtractedTask::ExtractedTask(String task_name, IRModule mod, Target target,
}
TVM_REGISTER_NODE_TYPE(ExtractedTaskNode);
-TVM_REGISTER_GLOBAL("meta_schedule.ExtractedTask")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.ExtractedTask")
.set_body_typed([](String task_name, IRModule mod, Target target, Array<IRModule> dispatched,
int weight) -> ExtractedTask {
return ExtractedTask(task_name, mod, target, dispatched, weight);
diff --git a/src/meta_schedule/feature_extractor/feature_extractor.cc b/src/meta_schedule/feature_extractor/feature_extractor.cc
index 093558d228..9a3cecf4ce 100644
--- a/src/meta_schedule/feature_extractor/feature_extractor.cc
+++ b/src/meta_schedule/feature_extractor/feature_extractor.cc
@@ -48,9 +48,9 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
TVM_REGISTER_OBJECT_TYPE(FeatureExtractorNode);
TVM_REGISTER_NODE_TYPE(PyFeatureExtractorNode);
-TVM_REGISTER_GLOBAL("meta_schedule.FeatureExtractorExtractFrom")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.FeatureExtractorExtractFrom")
.set_body_method(&FeatureExtractorNode::ExtractFrom);
-TVM_REGISTER_GLOBAL("meta_schedule.FeatureExtractorPyFeatureExtractor")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.FeatureExtractorPyFeatureExtractor")
.set_body_typed(FeatureExtractor::PyFeatureExtractor);
} // namespace meta_schedule
diff --git a/src/meta_schedule/feature_extractor/per_store_feature.cc b/src/meta_schedule/feature_extractor/per_store_feature.cc
index 5d91bc34f9..2fc8878546 100644
--- a/src/meta_schedule/feature_extractor/per_store_feature.cc
+++ b/src/meta_schedule/feature_extractor/per_store_feature.cc
@@ -1442,7 +1442,7 @@ FeatureExtractor FeatureExtractor::PerStoreFeature(int buffers_per_store,
}
TVM_REGISTER_NODE_TYPE(PerStoreFeatureNode);
-TVM_REGISTER_GLOBAL("meta_schedule.FeatureExtractorPerStoreFeature")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.FeatureExtractorPerStoreFeature")
.set_body_typed(FeatureExtractor::PerStoreFeature);
} // namespace meta_schedule
diff --git a/src/meta_schedule/measure_callback/add_to_database.cc b/src/meta_schedule/measure_callback/add_to_database.cc
index 68a4b93ea9..becd9d2110 100644
--- a/src/meta_schedule/measure_callback/add_to_database.cc
+++ b/src/meta_schedule/measure_callback/add_to_database.cc
@@ -65,7 +65,7 @@ MeasureCallback MeasureCallback::AddToDatabase() {
}
TVM_REGISTER_NODE_TYPE(AddToDatabaseNode);
-TVM_REGISTER_GLOBAL("meta_schedule.MeasureCallbackAddToDatabase")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.MeasureCallbackAddToDatabase")
.set_body_typed(MeasureCallback::AddToDatabase);
} // namespace meta_schedule
diff --git a/src/meta_schedule/measure_callback/measure_callback.cc b/src/meta_schedule/measure_callback/measure_callback.cc
index 8f94e29846..0ee49f2ab4 100644
--- a/src/meta_schedule/measure_callback/measure_callback.cc
+++ b/src/meta_schedule/measure_callback/measure_callback.cc
@@ -59,11 +59,11 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
TVM_REGISTER_OBJECT_TYPE(MeasureCallbackNode);
TVM_REGISTER_NODE_TYPE(PyMeasureCallbackNode);
-TVM_REGISTER_GLOBAL("meta_schedule.MeasureCallbackApply")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.MeasureCallbackApply")
.set_body_method(&MeasureCallbackNode::Apply);
-TVM_REGISTER_GLOBAL("meta_schedule.MeasureCallbackPyMeasureCallback")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.MeasureCallbackPyMeasureCallback")
.set_body_typed(MeasureCallback::PyMeasureCallback);
-TVM_REGISTER_GLOBAL("meta_schedule.MeasureCallbackDefault")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.MeasureCallbackDefault")
.set_body_typed(MeasureCallback::Default);
} // namespace meta_schedule
diff --git a/src/meta_schedule/measure_callback/remove_build_artifact.cc b/src/meta_schedule/measure_callback/remove_build_artifact.cc
index 9242e79912..da74e85cac 100644
--- a/src/meta_schedule/measure_callback/remove_build_artifact.cc
+++ b/src/meta_schedule/measure_callback/remove_build_artifact.cc
@@ -46,7 +46,7 @@ MeasureCallback MeasureCallback::RemoveBuildArtifact() {
}
TVM_REGISTER_NODE_TYPE(RemoveBuildArtifactNode);
-TVM_REGISTER_GLOBAL("meta_schedule.MeasureCallbackRemoveBuildArtifact")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.MeasureCallbackRemoveBuildArtifact")
.set_body_typed(MeasureCallback::RemoveBuildArtifact);
} // namespace meta_schedule
diff --git a/src/meta_schedule/measure_callback/update_cost_model.cc b/src/meta_schedule/measure_callback/update_cost_model.cc
index 63c32b189e..1969d7fc83 100644
--- a/src/meta_schedule/measure_callback/update_cost_model.cc
+++ b/src/meta_schedule/measure_callback/update_cost_model.cc
@@ -63,7 +63,7 @@ MeasureCallback MeasureCallback::UpdateCostModel() {
}
TVM_REGISTER_NODE_TYPE(UpdateCostModelNode);
-TVM_REGISTER_GLOBAL("meta_schedule.MeasureCallbackUpdateCostModel")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.MeasureCallbackUpdateCostModel")
.set_body_typed(MeasureCallback::UpdateCostModel);
} // namespace meta_schedule
diff --git a/src/meta_schedule/mutator/mutate_compute_location.cc b/src/meta_schedule/mutator/mutate_compute_location.cc
index f74d1640e4..8f8c077aa8 100644
--- a/src/meta_schedule/mutator/mutate_compute_location.cc
+++ b/src/meta_schedule/mutator/mutate_compute_location.cc
@@ -127,7 +127,7 @@ Mutator Mutator::MutateComputeLocation() {
}
TVM_REGISTER_NODE_TYPE(MutateComputeLocationNode);
-TVM_REGISTER_GLOBAL("meta_schedule.MutatorMutateComputeLocation")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.MutatorMutateComputeLocation")
.set_body_typed(Mutator::MutateComputeLocation);
} // namespace meta_schedule
diff --git a/src/meta_schedule/mutator/mutate_parallel.cc b/src/meta_schedule/mutator/mutate_parallel.cc
index 6ddb6ea96b..a6a34e47a9 100644
--- a/src/meta_schedule/mutator/mutate_parallel.cc
+++ b/src/meta_schedule/mutator/mutate_parallel.cc
@@ -312,7 +312,8 @@ Mutator Mutator::MutateParallel(int64_t max_jobs_per_core) {
}
TVM_REGISTER_NODE_TYPE(MutateParallelNode);
-TVM_REGISTER_GLOBAL("meta_schedule.MutatorMutateParallel").set_body_typed(Mutator::MutateParallel);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.MutatorMutateParallel")
+ .set_body_typed(Mutator::MutateParallel);
} // namespace meta_schedule
} // namespace tvm
diff --git a/src/meta_schedule/mutator/mutate_thread_binding.cc b/src/meta_schedule/mutator/mutate_thread_binding.cc
index 3666e84bd5..269b052404 100644
--- a/src/meta_schedule/mutator/mutate_thread_binding.cc
+++ b/src/meta_schedule/mutator/mutate_thread_binding.cc
@@ -165,7 +165,7 @@ Optional<Trace> MutateThreadBindingNode::Apply(const Trace& trace, TRandState* r
Mutator Mutator::MutateThreadBinding() { return Mutator(make_object<MutateThreadBindingNode>()); }
TVM_REGISTER_NODE_TYPE(MutateThreadBindingNode);
-TVM_REGISTER_GLOBAL("meta_schedule.MutateThreadBinding")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.MutateThreadBinding")
.set_body_typed(Mutator::MutateThreadBinding);
} // namespace meta_schedule
diff --git a/src/meta_schedule/mutator/mutate_tile_size.cc b/src/meta_schedule/mutator/mutate_tile_size.cc
index 35dc23ceba..e8a728d050 100644
--- a/src/meta_schedule/mutator/mutate_tile_size.cc
+++ b/src/meta_schedule/mutator/mutate_tile_size.cc
@@ -269,7 +269,8 @@ Optional<Trace> MutateTileSizeNode::Apply(const Trace& trace, TRandState* rand_s
Mutator Mutator::MutateTileSize() { return Mutator(make_object<MutateTileSizeNode>()); }
TVM_REGISTER_NODE_TYPE(MutateTileSizeNode);
-TVM_REGISTER_GLOBAL("meta_schedule.MutatorMutateTileSize").set_body_typed(Mutator::MutateTileSize);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.MutatorMutateTileSize")
+ .set_body_typed(Mutator::MutateTileSize);
} // namespace meta_schedule
} // namespace tvm
diff --git a/src/meta_schedule/mutator/mutate_unroll.cc b/src/meta_schedule/mutator/mutate_unroll.cc
index 782da45a35..28fcf3668f 100644
--- a/src/meta_schedule/mutator/mutate_unroll.cc
+++ b/src/meta_schedule/mutator/mutate_unroll.cc
@@ -138,7 +138,7 @@ Optional<Trace> MutateUnrollNode::Apply(const Trace& trace, TRandState* rand_sta
Mutator Mutator::MutateUnroll() { return Mutator(make_object<MutateUnrollNode>()); }
TVM_REGISTER_NODE_TYPE(MutateUnrollNode);
-TVM_REGISTER_GLOBAL("meta_schedule.MutatorMutateUnroll").set_body_typed(Mutator::MutateUnroll);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.MutatorMutateUnroll").set_body_typed(Mutator::MutateUnroll);
} // namespace meta_schedule
} // namespace tvm
diff --git a/src/meta_schedule/mutator/mutator.cc b/src/meta_schedule/mutator/mutator.cc
index e1831d213d..e415b3909f 100644
--- a/src/meta_schedule/mutator/mutator.cc
+++ b/src/meta_schedule/mutator/mutator.cc
@@ -88,20 +88,21 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
TVM_REGISTER_OBJECT_TYPE(MutatorNode);
TVM_REGISTER_NODE_TYPE(PyMutatorNode);
-TVM_REGISTER_GLOBAL("meta_schedule.MutatorInitializeWithTuneContext")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.MutatorInitializeWithTuneContext")
.set_body_method(&MutatorNode::InitializeWithTuneContext);
-TVM_REGISTER_GLOBAL("meta_schedule.MutatorApply")
+TVM_FFI_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(&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);
-TVM_REGISTER_GLOBAL("meta_schedule.MutatorDefaultCUDATensorCore")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.MutatorClone").set_body_method(&MutatorNode::Clone);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.MutatorPyMutator").set_body_typed(Mutator::PyMutator);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.MutatorDefaultLLVM").set_body_typed(Mutator::DefaultLLVM);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.MutatorDefaultCUDA").set_body_typed(Mutator::DefaultCUDA);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.MutatorDefaultCUDATensorCore")
.set_body_typed(Mutator::DefaultCUDATensorCore);
-TVM_REGISTER_GLOBAL("meta_schedule.MutatorDefaultHexagon").set_body_typed(Mutator::DefaultHexagon);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.MutatorDefaultHexagon")
+ .set_body_typed(Mutator::DefaultHexagon);
} // namespace meta_schedule
} // namespace tvm
diff --git a/src/meta_schedule/postproc/disallow_async_strided_mem_copy.cc b/src/meta_schedule/postproc/disallow_async_strided_mem_copy.cc
index c32c7515fa..01a75a5bfb 100644
--- a/src/meta_schedule/postproc/disallow_async_strided_mem_copy.cc
+++ b/src/meta_schedule/postproc/disallow_async_strided_mem_copy.cc
@@ -184,7 +184,7 @@ Postproc Postproc::DisallowAsyncStridedMemCopy() {
}
TVM_REGISTER_NODE_TYPE(DisallowAsyncStridedMemCopyNode);
-TVM_REGISTER_GLOBAL("meta_schedule.PostprocDisallowAsyncStridedMemCopy")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.PostprocDisallowAsyncStridedMemCopy")
.set_body_typed(Postproc::DisallowAsyncStridedMemCopy);
} // namespace meta_schedule
diff --git a/src/meta_schedule/postproc/disallow_dynamic_loop.cc b/src/meta_schedule/postproc/disallow_dynamic_loop.cc
index 8362da552e..fd099ac5dd 100644
--- a/src/meta_schedule/postproc/disallow_dynamic_loop.cc
+++ b/src/meta_schedule/postproc/disallow_dynamic_loop.cc
@@ -83,7 +83,7 @@ Postproc Postproc::DisallowDynamicLoop() {
}
TVM_REGISTER_NODE_TYPE(DisallowDynamicLoopNode);
-TVM_REGISTER_GLOBAL("meta_schedule.PostprocDisallowDynamicLoop")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.PostprocDisallowDynamicLoop")
.set_body_typed(Postproc::DisallowDynamicLoop);
} // namespace meta_schedule
diff --git a/src/meta_schedule/postproc/postproc.cc b/src/meta_schedule/postproc/postproc.cc
index b5d62634c2..e29f9dd54c 100644
--- a/src/meta_schedule/postproc/postproc.cc
+++ b/src/meta_schedule/postproc/postproc.cc
@@ -112,16 +112,16 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
TVM_REGISTER_OBJECT_TYPE(PostprocNode);
TVM_REGISTER_NODE_TYPE(PyPostprocNode);
-TVM_REGISTER_GLOBAL("meta_schedule.PostprocInitializeWithTuneContext")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.PostprocInitializeWithTuneContext")
.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);
-TVM_REGISTER_GLOBAL("meta_schedule.PostprocDefaultCUDATensorCore")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.PostprocApply").set_body_method(&PostprocNode::Apply);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.PostprocClone").set_body_method(&PostprocNode::Clone);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.PostprocPyPostproc").set_body_typed(Postproc::PyPostproc);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.PostprocDefaultLLVM").set_body_typed(Postproc::DefaultLLVM);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.PostprocDefaultCUDA").set_body_typed(Postproc::DefaultCUDA);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.PostprocDefaultCUDATensorCore")
.set_body_typed(Postproc::DefaultCUDATensorCore);
-TVM_REGISTER_GLOBAL("meta_schedule.PostprocDefaultHexagon")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.PostprocDefaultHexagon")
.set_body_typed(Postproc::DefaultHexagon);
} // namespace meta_schedule
diff --git a/src/meta_schedule/postproc/rewrite_cooperative_fetch.cc b/src/meta_schedule/postproc/rewrite_cooperative_fetch.cc
index aa206764cc..d23e07795c 100644
--- a/src/meta_schedule/postproc/rewrite_cooperative_fetch.cc
+++ b/src/meta_schedule/postproc/rewrite_cooperative_fetch.cc
@@ -227,7 +227,7 @@ Postproc Postproc::RewriteCooperativeFetch() {
}
TVM_REGISTER_NODE_TYPE(RewriteCooperativeFetchNode);
-TVM_REGISTER_GLOBAL("meta_schedule.PostprocRewriteCooperativeFetch")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.PostprocRewriteCooperativeFetch")
.set_body_typed(Postproc::RewriteCooperativeFetch);
} // namespace meta_schedule
diff --git a/src/meta_schedule/postproc/rewrite_layout.cc b/src/meta_schedule/postproc/rewrite_layout.cc
index cdfc6c56e5..84dc33ec98 100644
--- a/src/meta_schedule/postproc/rewrite_layout.cc
+++ b/src/meta_schedule/postproc/rewrite_layout.cc
@@ -273,7 +273,8 @@ Postproc Postproc::RewriteLayout() {
}
TVM_REGISTER_NODE_TYPE(RewriteLayoutNode);
-TVM_REGISTER_GLOBAL("meta_schedule.PostprocRewriteLayout").set_body_typed(Postproc::RewriteLayout);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.PostprocRewriteLayout")
+ .set_body_typed(Postproc::RewriteLayout);
} // namespace meta_schedule
} // namespace tvm
diff --git a/src/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.cc b/src/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.cc
index c43c8aac12..3f665cd8d8 100644
--- a/src/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.cc
+++ b/src/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.cc
@@ -464,7 +464,7 @@ Postproc Postproc::RewriteParallelVectorizeUnroll() {
}
TVM_REGISTER_NODE_TYPE(RewriteParallelVectorizeUnrollNode);
-TVM_REGISTER_GLOBAL("meta_schedule.PostprocRewriteParallelVectorizeUnroll")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.PostprocRewriteParallelVectorizeUnroll")
.set_body_typed(Postproc::RewriteParallelVectorizeUnroll);
} // namespace meta_schedule
diff --git a/src/meta_schedule/postproc/rewrite_reduction_block.cc b/src/meta_schedule/postproc/rewrite_reduction_block.cc
index 295044596f..3ffe0f9234 100644
--- a/src/meta_schedule/postproc/rewrite_reduction_block.cc
+++ b/src/meta_schedule/postproc/rewrite_reduction_block.cc
@@ -172,7 +172,7 @@ Postproc Postproc::RewriteReductionBlock() {
}
TVM_REGISTER_NODE_TYPE(RewriteReductionBlockNode);
-TVM_REGISTER_GLOBAL("meta_schedule.PostprocRewriteReductionBlock")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.PostprocRewriteReductionBlock")
.set_body_typed(Postproc::RewriteReductionBlock);
} // namespace meta_schedule
diff --git a/src/meta_schedule/postproc/rewrite_tensorize.cc b/src/meta_schedule/postproc/rewrite_tensorize.cc
index 4f8e0fb213..0f98484dd4 100644
--- a/src/meta_schedule/postproc/rewrite_tensorize.cc
+++ b/src/meta_schedule/postproc/rewrite_tensorize.cc
@@ -107,7 +107,7 @@ Postproc Postproc::RewriteTensorize(bool vectorize_init_loop) {
}
TVM_REGISTER_NODE_TYPE(RewriteTensorizeNode);
-TVM_REGISTER_GLOBAL("meta_schedule.PostprocRewriteTensorize")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.PostprocRewriteTensorize")
.set_body_typed(Postproc::RewriteTensorize);
} // namespace meta_schedule
diff --git a/src/meta_schedule/postproc/rewrite_unbound_block.cc b/src/meta_schedule/postproc/rewrite_unbound_block.cc
index 27ce34a8cb..a2c9d1364a 100644
--- a/src/meta_schedule/postproc/rewrite_unbound_block.cc
+++ b/src/meta_schedule/postproc/rewrite_unbound_block.cc
@@ -146,7 +146,7 @@ Postproc Postproc::RewriteUnboundBlock(int max_threadblocks) {
}
TVM_REGISTER_NODE_TYPE(RewriteUnboundBlockNode);
-TVM_REGISTER_GLOBAL("meta_schedule.PostprocRewriteUnboundBlock")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.PostprocRewriteUnboundBlock")
.set_body_typed(Postproc::RewriteUnboundBlock);
} // namespace meta_schedule
diff --git a/src/meta_schedule/postproc/verify_gpu_code.cc b/src/meta_schedule/postproc/verify_gpu_code.cc
index f87e1ed9d5..8ffc424e44 100644
--- a/src/meta_schedule/postproc/verify_gpu_code.cc
+++ b/src/meta_schedule/postproc/verify_gpu_code.cc
@@ -215,7 +215,8 @@ Postproc Postproc::VerifyGPUCode() {
}
TVM_REGISTER_NODE_TYPE(VerifyGPUCodeNode);
-TVM_REGISTER_GLOBAL("meta_schedule.PostprocVerifyGPUCode").set_body_typed(Postproc::VerifyGPUCode);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.PostprocVerifyGPUCode")
+ .set_body_typed(Postproc::VerifyGPUCode);
} // namespace meta_schedule
} // namespace tvm
diff --git a/src/meta_schedule/postproc/verify_vtcm_limit.cc b/src/meta_schedule/postproc/verify_vtcm_limit.cc
index 4de9750896..7da2f8546b 100644
--- a/src/meta_schedule/postproc/verify_vtcm_limit.cc
+++ b/src/meta_schedule/postproc/verify_vtcm_limit.cc
@@ -69,7 +69,7 @@ Postproc Postproc::VerifyVTCMLimit() {
}
TVM_REGISTER_NODE_TYPE(VerifyVTCMLimitNode);
-TVM_REGISTER_GLOBAL("meta_schedule.PostprocVerifyVTCMLimit")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.PostprocVerifyVTCMLimit")
.set_body_typed(Postproc::VerifyVTCMLimit);
} // namespace meta_schedule
diff --git a/src/meta_schedule/profiler.cc b/src/meta_schedule/profiler.cc
index b1386710b2..2a034a7be2 100644
--- a/src/meta_schedule/profiler.cc
+++ b/src/meta_schedule/profiler.cc
@@ -121,17 +121,17 @@ Optional<Profiler> Profiler::Current() {
}
TVM_REGISTER_NODE_TYPE(ProfilerNode);
-TVM_REGISTER_GLOBAL("meta_schedule.Profiler").set_body_typed([]() -> Profiler {
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.Profiler").set_body_typed([]() -> Profiler {
return Profiler();
});
-TVM_REGISTER_GLOBAL("meta_schedule.ProfilerEnterWithScope")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.ProfilerEnterWithScope")
.set_body_method(&Profiler::EnterWithScope);
-TVM_REGISTER_GLOBAL("meta_schedule.ProfilerExitWithScope")
+TVM_FFI_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(&ProfilerNode::Get);
-TVM_REGISTER_GLOBAL("meta_schedule.ProfilerTable").set_body_method(&ProfilerNode::Table);
-TVM_REGISTER_GLOBAL("meta_schedule.ProfilerTimedScope").set_body_typed(ProfilerTimedScope);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.ProfilerCurrent").set_body_typed(Profiler::Current);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.ProfilerGet").set_body_method(&ProfilerNode::Get);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.ProfilerTable").set_body_method(&ProfilerNode::Table);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.ProfilerTimedScope").set_body_typed(ProfilerTimedScope);
} // namespace meta_schedule
} // namespace tvm
diff --git a/src/meta_schedule/runner/runner.cc b/src/meta_schedule/runner/runner.cc
index 0d5edf299e..38d4225f0f 100644
--- a/src/meta_schedule/runner/runner.cc
+++ b/src/meta_schedule/runner/runner.cc
@@ -56,24 +56,25 @@ TVM_REGISTER_NODE_TYPE(RunnerResultNode);
TVM_REGISTER_NODE_TYPE(RunnerFutureNode);
TVM_REGISTER_OBJECT_TYPE(RunnerNode);
TVM_REGISTER_NODE_TYPE(PyRunnerNode);
-TVM_REGISTER_GLOBAL("meta_schedule.RunnerInput")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.RunnerInput")
.set_body_typed([](String artifact_path, String device_type,
Array<ArgInfo> args_info) -> RunnerInput {
return RunnerInput(artifact_path, device_type, args_info);
});
-TVM_REGISTER_GLOBAL("meta_schedule.RunnerResult")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.RunnerResult")
.set_body_typed([](Optional<Array<FloatImm>> run_secs,
Optional<String> error_msg) -> RunnerResult {
return RunnerResult(run_secs, error_msg);
});
-TVM_REGISTER_GLOBAL("meta_schedule.RunnerFuture")
+TVM_FFI_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(&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);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.RunnerFutureDone").set_body_method(&RunnerFutureNode::Done);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.RunnerFutureResult")
+ .set_body_method(&RunnerFutureNode::Result);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.RunnerRun").set_body_method(&RunnerNode::Run);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.RunnerPyRunner").set_body_typed(Runner::PyRunner);
} // namespace meta_schedule
} // namespace tvm
diff --git a/src/meta_schedule/schedule/cpu/winograd.cc b/src/meta_schedule/schedule/cpu/winograd.cc
index 16e53b5692..4e09fa729b 100644
--- a/src/meta_schedule/schedule/cpu/winograd.cc
+++ b/src/meta_schedule/schedule/cpu/winograd.cc
@@ -59,7 +59,7 @@ static Array<tir::LoopRV> ScheduleDataPack(tir::Schedule sch, tir::BlockRV block
return {t0[0], t1[0], t0[1], t1[1]};
}
-TVM_REGISTER_GLOBAL("meta_schedule.cpu.conv2d_nhwc_winograd_data_pack")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.cpu.conv2d_nhwc_winograd_data_pack")
.set_body_typed([](Schedule sch, BlockRV data_pack) -> Array<Schedule> {
BlockRV input_tile = GetWinogradProducerAndInlineConst(sch, data_pack);
BlockRV data_pad = GetWinogradProducerAndInlineConst(sch, input_tile);
@@ -71,14 +71,14 @@ TVM_REGISTER_GLOBAL("meta_schedule.cpu.conv2d_nhwc_winograd_data_pack")
return {sch};
});
-TVM_REGISTER_GLOBAL("meta_schedule.cpu.conv2d_nhwc_winograd_inverse")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.cpu.conv2d_nhwc_winograd_inverse")
.set_body_typed([](Schedule sch, BlockRV block) -> Array<Schedule> {
GetWinogradProducerAndInlineConst(sch, block);
ScheduleDataPack(sch, block, {2, 3}, {0, 1, 4, 5});
return {sch};
});
-TVM_REGISTER_GLOBAL("meta_schedule.cpu.conv2d_nchw_winograd_data_pack")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.cpu.conv2d_nchw_winograd_data_pack")
.set_body_typed([](Schedule sch, BlockRV data_pack) -> Array<Schedule> {
BlockRV input_tile = GetWinogradProducerAndInlineConst(sch, data_pack);
BlockRV data_pad = GetWinogradProducerAndInlineConst(sch, input_tile);
@@ -90,7 +90,7 @@ TVM_REGISTER_GLOBAL("meta_schedule.cpu.conv2d_nchw_winograd_data_pack")
return {sch};
});
-TVM_REGISTER_GLOBAL("meta_schedule.cpu.conv2d_nchw_winograd_inverse")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.cpu.conv2d_nchw_winograd_inverse")
.set_body_typed([](Schedule sch, BlockRV block) -> Array<Schedule> {
GetWinogradProducerAndInlineConst(sch, block);
ScheduleDataPack(sch, block, {0, 1}, {2, 3, 4, 5});
diff --git a/src/meta_schedule/schedule/cuda/winograd.cc b/src/meta_schedule/schedule/cuda/winograd.cc
index 46ef6f3663..c80141f528 100644
--- a/src/meta_schedule/schedule/cuda/winograd.cc
+++ b/src/meta_schedule/schedule/cuda/winograd.cc
@@ -63,7 +63,7 @@ static Array<tir::LoopRV> ScheduleDataPack(tir::Schedule sch, tir::BlockRV block
return {t0[0], t1[0], t0[1], t1[1]};
}
-TVM_REGISTER_GLOBAL("meta_schedule.cuda.conv2d_nhwc_winograd_data_pack")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.cuda.conv2d_nhwc_winograd_data_pack")
.set_body_typed([](Schedule sch, BlockRV data_pack) -> Array<Schedule> {
BlockRV input_tile = GetWinogradProducerAndInlineConst(sch, data_pack);
BlockRV data_pad = GetWinogradProducerAndInlineConst(sch, input_tile);
@@ -88,7 +88,7 @@ TVM_REGISTER_GLOBAL("meta_schedule.cuda.conv2d_nhwc_winograd_data_pack")
return {sch};
});
-TVM_REGISTER_GLOBAL("meta_schedule.cuda.conv2d_nhwc_winograd_inverse")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.cuda.conv2d_nhwc_winograd_inverse")
.set_body_typed([](Schedule sch, BlockRV inverse) -> Array<Schedule> {
GetWinogradProducerAndInlineConst(sch, inverse);
ScheduleDataPack(sch, inverse, /*tiled=*/{2, 3}, /*unrolled=*/{0, 1, 4, 5});
@@ -101,7 +101,7 @@ TVM_REGISTER_GLOBAL("meta_schedule.cuda.conv2d_nhwc_winograd_inverse")
return {sch};
});
-TVM_REGISTER_GLOBAL("meta_schedule.cuda.conv2d_nchw_winograd_data_pack")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.cuda.conv2d_nchw_winograd_data_pack")
.set_body_typed([](Schedule sch, BlockRV data_pack) -> Array<Schedule> {
int64_t max_threadblocks = 256;
int64_t max_threads_per_block = 1024;
@@ -132,7 +132,7 @@ TVM_REGISTER_GLOBAL("meta_schedule.cuda.conv2d_nchw_winograd_data_pack")
return {sch};
});
-TVM_REGISTER_GLOBAL("meta_schedule.cuda.conv2d_nchw_winograd_inverse")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.cuda.conv2d_nchw_winograd_inverse")
.set_body_typed([](Schedule sch, BlockRV inverse) -> Array<Schedule> {
GetWinogradProducerAndInlineConst(sch, inverse);
// loops on top of the inverse block: [CO, P, tile_size, tile_size, alpha, alpha]
diff --git a/src/meta_schedule/schedule_rule/add_rfactor.cc b/src/meta_schedule/schedule_rule/add_rfactor.cc
index 2fc1352677..48149ed871 100644
--- a/src/meta_schedule/schedule_rule/add_rfactor.cc
+++ b/src/meta_schedule/schedule_rule/add_rfactor.cc
@@ -120,7 +120,7 @@ Array<tir::Schedule> AddRFactorNode::Apply(const tir::Schedule& sch, const tir::
}
TVM_REGISTER_NODE_TYPE(AddRFactorNode);
-TVM_REGISTER_GLOBAL("meta_schedule.ScheduleRuleAddRFactor")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.ScheduleRuleAddRFactor")
.set_body_typed(ScheduleRule::AddRFactor);
} // 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 9fdcaa4ef5..92de19163a 100644
--- a/src/meta_schedule/schedule_rule/apply_custom_rule.cc
+++ b/src/meta_schedule/schedule_rule/apply_custom_rule.cc
@@ -87,7 +87,7 @@ bool ScheduleRule::IsApplyCustomRule(const ScheduleRule& rule) {
}
TVM_REGISTER_NODE_TYPE(ApplyCustomRuleNode);
-TVM_REGISTER_GLOBAL("meta_schedule.ScheduleRuleApplyCustomRule")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.ScheduleRuleApplyCustomRule")
.set_body_typed(ScheduleRule::ApplyCustomRule);
} // namespace meta_schedule
diff --git a/src/meta_schedule/schedule_rule/auto_bind.cc b/src/meta_schedule/schedule_rule/auto_bind.cc
index fa47d1edb8..892a79ea92 100644
--- a/src/meta_schedule/schedule_rule/auto_bind.cc
+++ b/src/meta_schedule/schedule_rule/auto_bind.cc
@@ -82,7 +82,8 @@ ScheduleRule ScheduleRule::AutoBind(int max_threadblocks, Array<Integer> thread_
}
TVM_REGISTER_NODE_TYPE(AutoBindNode);
-TVM_REGISTER_GLOBAL("meta_schedule.ScheduleRuleAutoBind").set_body_typed(ScheduleRule::AutoBind);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.ScheduleRuleAutoBind")
+ .set_body_typed(ScheduleRule::AutoBind);
} // namespace meta_schedule
} // namespace tvm
diff --git a/src/meta_schedule/schedule_rule/auto_inline.cc b/src/meta_schedule/schedule_rule/auto_inline.cc
index d9e033eff8..948632e580 100644
--- a/src/meta_schedule/schedule_rule/auto_inline.cc
+++ b/src/meta_schedule/schedule_rule/auto_inline.cc
@@ -191,7 +191,7 @@ ScheduleRule ScheduleRule::AutoInline(bool into_producer, //
}
TVM_REGISTER_NODE_TYPE(AutoInlineNode);
-TVM_REGISTER_GLOBAL("meta_schedule.ScheduleRuleAutoInline")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.ScheduleRuleAutoInline")
.set_body_typed(ScheduleRule::AutoInline);
/*! \brief Inline blocks that produce a constant scalar. */
@@ -232,7 +232,7 @@ ScheduleRule ScheduleRule::InlineConstantScalars() {
}
TVM_REGISTER_NODE_TYPE(InlineConstantScalarsNode);
-TVM_REGISTER_GLOBAL("meta_schedule.ScheduleRuleInlineConstantScalars")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.ScheduleRuleInlineConstantScalars")
.set_body_typed(ScheduleRule::InlineConstantScalars);
} // namespace meta_schedule
} // namespace tvm
diff --git a/src/meta_schedule/schedule_rule/cross_thread_reduction.cc b/src/meta_schedule/schedule_rule/cross_thread_reduction.cc
index b71a14cc04..e06817e37c 100644
--- a/src/meta_schedule/schedule_rule/cross_thread_reduction.cc
+++ b/src/meta_schedule/schedule_rule/cross_thread_reduction.cc
@@ -291,7 +291,7 @@ ScheduleRule ScheduleRule::CrossThreadReduction(Array<Integer> thread_extents) {
}
TVM_REGISTER_NODE_TYPE(CrossThreadReductionNode);
-TVM_REGISTER_GLOBAL("meta_schedule.ScheduleRuleCrossThreadReduction")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.ScheduleRuleCrossThreadReduction")
.set_body_typed(ScheduleRule::CrossThreadReduction);
} // namespace meta_schedule
diff --git a/src/meta_schedule/schedule_rule/multi_level_tiling.cc b/src/meta_schedule/schedule_rule/multi_level_tiling.cc
index 79cff3bad7..f020c8efd0 100644
--- a/src/meta_schedule/schedule_rule/multi_level_tiling.cc
+++ b/src/meta_schedule/schedule_rule/multi_level_tiling.cc
@@ -405,7 +405,7 @@ ScheduleRule ScheduleRule::MultiLevelTiling(String structure, Optional<Array<Str
}
TVM_REGISTER_NODE_TYPE(MultiLevelTilingNode);
-TVM_REGISTER_GLOBAL("meta_schedule.ScheduleRuleMultiLevelTiling")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.ScheduleRuleMultiLevelTiling")
.set_body_typed(ScheduleRule::MultiLevelTiling);
} // namespace meta_schedule
diff --git a/src/meta_schedule/schedule_rule/multi_level_tiling_tensor_core.cc b/src/meta_schedule/schedule_rule/multi_level_tiling_tensor_core.cc
index 2edb18ca85..0728789cd2 100644
--- a/src/meta_schedule/schedule_rule/multi_level_tiling_tensor_core.cc
+++ b/src/meta_schedule/schedule_rule/multi_level_tiling_tensor_core.cc
@@ -923,7 +923,7 @@ ScheduleRule ScheduleRule::MultiLevelTilingTensorCore(
}
TVM_REGISTER_NODE_TYPE(MultiLevelTilingTensorCoreNode);
-TVM_REGISTER_GLOBAL("meta_schedule.ScheduleRuleMultiLevelTilingTensorCore")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.ScheduleRuleMultiLevelTilingTensorCore")
.set_body_typed(ScheduleRule::MultiLevelTilingTensorCore);
} // namespace meta_schedule
diff --git a/src/meta_schedule/schedule_rule/multi_level_tiling_wide_vector.cc b/src/meta_schedule/schedule_rule/multi_level_tiling_wide_vector.cc
index c7bbf82bf4..0da8ee35cf 100644
--- a/src/meta_schedule/schedule_rule/multi_level_tiling_wide_vector.cc
+++ b/src/meta_schedule/schedule_rule/multi_level_tiling_wide_vector.cc
@@ -124,7 +124,7 @@ ScheduleRule ScheduleRule::MultiLevelTilingWideVector(String structure,
}
TVM_REGISTER_NODE_TYPE(MultiLevelTilingWideVectorNode);
-TVM_REGISTER_GLOBAL("meta_schedule.ScheduleRuleMultiLevelTilingWideVector")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.ScheduleRuleMultiLevelTilingWideVector")
.set_body_typed(ScheduleRule::MultiLevelTilingWideVector);
} // namespace meta_schedule
diff --git a/src/meta_schedule/schedule_rule/multi_level_tiling_with_intrin.cc b/src/meta_schedule/schedule_rule/multi_level_tiling_with_intrin.cc
index 554bb9dfb6..731860e8d6 100644
--- a/src/meta_schedule/schedule_rule/multi_level_tiling_with_intrin.cc
+++ b/src/meta_schedule/schedule_rule/multi_level_tiling_with_intrin.cc
@@ -106,7 +106,7 @@ ScheduleRule ScheduleRule::MultiLevelTilingWithIntrin(String intrin_name, String
}
TVM_REGISTER_NODE_TYPE(MultiLevelTilingWithIntrinNode);
-TVM_REGISTER_GLOBAL("meta_schedule.ScheduleRuleMultiLevelTilingWithIntrin")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.ScheduleRuleMultiLevelTilingWithIntrin")
.set_body_typed(ScheduleRule::MultiLevelTilingWithIntrin);
} // namespace meta_schedule
diff --git a/src/meta_schedule/schedule_rule/parallel_vectorize_unroll.cc b/src/meta_schedule/schedule_rule/parallel_vectorize_unroll.cc
index 54091cfd73..905f8d8ce6 100644
--- a/src/meta_schedule/schedule_rule/parallel_vectorize_unroll.cc
+++ b/src/meta_schedule/schedule_rule/parallel_vectorize_unroll.cc
@@ -134,7 +134,7 @@ ScheduleRule ScheduleRule::ParallelizeVectorizeUnroll(int max_jobs_per_core,
}
TVM_REGISTER_NODE_TYPE(ParallelizeVectorizeUnrollNode);
-TVM_REGISTER_GLOBAL("meta_schedule.ScheduleRuleParallelizeVectorizeUnroll")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.ScheduleRuleParallelizeVectorizeUnroll")
.set_body_typed(ScheduleRule::ParallelizeVectorizeUnroll);
} // namespace meta_schedule
diff --git a/src/meta_schedule/schedule_rule/random_compute_location.cc b/src/meta_schedule/schedule_rule/random_compute_location.cc
index 7796eddd44..ed71baade0 100644
--- a/src/meta_schedule/schedule_rule/random_compute_location.cc
+++ b/src/meta_schedule/schedule_rule/random_compute_location.cc
@@ -122,7 +122,7 @@ ScheduleRule ScheduleRule::RandomComputeLocation() {
}
TVM_REGISTER_NODE_TYPE(RandomComputeLocationNode);
-TVM_REGISTER_GLOBAL("meta_schedule.ScheduleRuleRandomComputeLocation")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.ScheduleRuleRandomComputeLocation")
.set_body_typed(ScheduleRule::RandomComputeLocation);
} // namespace meta_schedule
} // namespace tvm
diff --git a/src/meta_schedule/schedule_rule/schedule_rule.cc b/src/meta_schedule/schedule_rule/schedule_rule.cc
index 5c5e2583dd..3640694b4e 100644
--- a/src/meta_schedule/schedule_rule/schedule_rule.cc
+++ b/src/meta_schedule/schedule_rule/schedule_rule.cc
@@ -402,21 +402,23 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
TVM_REGISTER_OBJECT_TYPE(ScheduleRuleNode);
TVM_REGISTER_NODE_TYPE(PyScheduleRuleNode);
-TVM_REGISTER_GLOBAL("meta_schedule.ScheduleRuleInitializeWithTuneContext")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.ScheduleRuleInitializeWithTuneContext")
.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")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.ScheduleRuleApply")
+ .set_body_method(&ScheduleRuleNode::Apply);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.ScheduleRuleClone")
+ .set_body_method(&ScheduleRuleNode::Clone);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.ScheduleRulePyScheduleRule")
.set_body_typed(ScheduleRule::PyScheduleRule);
-TVM_REGISTER_GLOBAL("meta_schedule.ScheduleRuleDefaultLLVM")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.ScheduleRuleDefaultLLVM")
.set_body_typed(ScheduleRule::DefaultLLVM);
-TVM_REGISTER_GLOBAL("meta_schedule.ScheduleRuleDefaultCUDA")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.ScheduleRuleDefaultCUDA")
.set_body_typed(ScheduleRule::DefaultCUDA);
-TVM_REGISTER_GLOBAL("meta_schedule.ScheduleRuleDefaultCUDATensorCore")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.ScheduleRuleDefaultCUDATensorCore")
.set_body_typed(ScheduleRule::DefaultCUDATensorCore);
-TVM_REGISTER_GLOBAL("meta_schedule.ScheduleRuleDefaultHexagon")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.ScheduleRuleDefaultHexagon")
.set_body_typed(ScheduleRule::DefaultHexagon);
-TVM_REGISTER_GLOBAL("meta_schedule.ScheduleRuleDefaultARM")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.ScheduleRuleDefaultARM")
.set_body_typed(ScheduleRule::DefaultARM);
} // namespace meta_schedule
diff --git a/src/meta_schedule/search_strategy/evolutionary_search.cc b/src/meta_schedule/search_strategy/evolutionary_search.cc
index f4812d56f1..e43e15c999 100644
--- a/src/meta_schedule/search_strategy/evolutionary_search.cc
+++ b/src/meta_schedule/search_strategy/evolutionary_search.cc
@@ -801,11 +801,11 @@ Array<Schedule> EvolutionarySearchEvolveWithCostModel(EvolutionarySearch self,
}
TVM_REGISTER_NODE_TYPE(EvolutionarySearchNode);
-TVM_REGISTER_GLOBAL("meta_schedule.SearchStrategyEvolutionarySearch")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.SearchStrategyEvolutionarySearch")
.set_body_typed(SearchStrategy::EvolutionarySearch);
-TVM_REGISTER_GLOBAL("meta_schedule.SearchStrategyEvolutionarySearchSampleInitPopulation")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.SearchStrategyEvolutionarySearchSampleInitPopulation")
.set_body_typed(EvolutionarySearchSampleInitPopulation);
-TVM_REGISTER_GLOBAL("meta_schedule.SearchStrategyEvolutionarySearchEvolveWithCostModel")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.SearchStrategyEvolutionarySearchEvolveWithCostModel")
.set_body_typed(EvolutionarySearchEvolveWithCostModel);
} // namespace meta_schedule
diff --git a/src/meta_schedule/search_strategy/replay_func.cc b/src/meta_schedule/search_strategy/replay_func.cc
index 60552b2e14..51cc408391 100644
--- a/src/meta_schedule/search_strategy/replay_func.cc
+++ b/src/meta_schedule/search_strategy/replay_func.cc
@@ -157,7 +157,7 @@ SearchStrategy SearchStrategy::ReplayFunc() {
}
TVM_REGISTER_NODE_TYPE(ReplayFuncNode);
-TVM_REGISTER_GLOBAL("meta_schedule.SearchStrategyReplayFunc")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.SearchStrategyReplayFunc")
.set_body_typed(SearchStrategy::ReplayFunc);
} // namespace meta_schedule
diff --git a/src/meta_schedule/search_strategy/replay_trace.cc b/src/meta_schedule/search_strategy/replay_trace.cc
index f7960688cd..c9a7459fdf 100644
--- a/src/meta_schedule/search_strategy/replay_trace.cc
+++ b/src/meta_schedule/search_strategy/replay_trace.cc
@@ -191,7 +191,7 @@ SearchStrategy SearchStrategy::ReplayTrace(int max_fail_count) {
}
TVM_REGISTER_NODE_TYPE(ReplayTraceNode);
-TVM_REGISTER_GLOBAL("meta_schedule.SearchStrategyReplayTrace")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.SearchStrategyReplayTrace")
.set_body_typed(SearchStrategy::ReplayTrace);
} // namespace meta_schedule
diff --git a/src/meta_schedule/search_strategy/search_strategy.cc b/src/meta_schedule/search_strategy/search_strategy.cc
index 1bc71502ad..8fc6538b59 100644
--- a/src/meta_schedule/search_strategy/search_strategy.cc
+++ b/src/meta_schedule/search_strategy/search_strategy.cc
@@ -86,23 +86,23 @@ TVM_REGISTER_NODE_TYPE(MeasureCandidateNode);
TVM_REGISTER_OBJECT_TYPE(SearchStrategyNode);
TVM_REGISTER_NODE_TYPE(PySearchStrategyNode);
-TVM_REGISTER_GLOBAL("meta_schedule.MeasureCandidate")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.MeasureCandidate")
.set_body_typed([](tir::Schedule sch, Optional<Array<ArgInfo>> args_info) -> MeasureCandidate {
return MeasureCandidate(sch, args_info.value_or({}));
});
-TVM_REGISTER_GLOBAL("meta_schedule.SearchStrategyPySearchStrategy")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.SearchStrategyPySearchStrategy")
.set_body_typed(SearchStrategy::PySearchStrategy);
-TVM_REGISTER_GLOBAL("meta_schedule.SearchStrategyInitializeWithTuneContext")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.SearchStrategyInitializeWithTuneContext")
.set_body_method(&SearchStrategyNode::InitializeWithTuneContext);
-TVM_REGISTER_GLOBAL("meta_schedule.SearchStrategyPreTuning")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.SearchStrategyPreTuning")
.set_body_method(&SearchStrategyNode::PreTuning);
-TVM_REGISTER_GLOBAL("meta_schedule.SearchStrategyPostTuning")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.SearchStrategyPostTuning")
.set_body_method(&SearchStrategyNode::PostTuning);
-TVM_REGISTER_GLOBAL("meta_schedule.SearchStrategyGenerateMeasureCandidates")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.SearchStrategyGenerateMeasureCandidates")
.set_body_method(&SearchStrategyNode::GenerateMeasureCandidates);
-TVM_REGISTER_GLOBAL("meta_schedule.SearchStrategyNotifyRunnerResults")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.SearchStrategyNotifyRunnerResults")
.set_body_method(&SearchStrategyNode::NotifyRunnerResults);
-TVM_REGISTER_GLOBAL("meta_schedule.SearchStrategyClone")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.SearchStrategyClone")
.set_body_method(&SearchStrategyNode::Clone);
} // namespace meta_schedule
diff --git a/src/meta_schedule/space_generator/post_order_apply.cc b/src/meta_schedule/space_generator/post_order_apply.cc
index da2178c736..91d5ba53d5 100644
--- a/src/meta_schedule/space_generator/post_order_apply.cc
+++ b/src/meta_schedule/space_generator/post_order_apply.cc
@@ -116,7 +116,7 @@ SpaceGenerator SpaceGenerator::PostOrderApply(ffi::Function f_block_filter,
}
TVM_REGISTER_NODE_TYPE(PostOrderApplyNode);
-TVM_REGISTER_GLOBAL("meta_schedule.SpaceGeneratorPostOrderApply")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.SpaceGeneratorPostOrderApply")
.set_body_typed(SpaceGenerator::PostOrderApply);
} // namespace meta_schedule
diff --git a/src/meta_schedule/space_generator/schedule_fn.cc b/src/meta_schedule/space_generator/schedule_fn.cc
index 89a02876f3..f7f2a3ba19 100644
--- a/src/meta_schedule/space_generator/schedule_fn.cc
+++ b/src/meta_schedule/space_generator/schedule_fn.cc
@@ -97,7 +97,7 @@ SpaceGenerator SpaceGenerator::ScheduleFn(ffi::Function schedule_fn,
}
TVM_REGISTER_NODE_TYPE(ScheduleFnNode);
-TVM_REGISTER_GLOBAL("meta_schedule.SpaceGeneratorScheduleFn")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.SpaceGeneratorScheduleFn")
.set_body_typed(SpaceGenerator::ScheduleFn);
} // namespace meta_schedule
diff --git a/src/meta_schedule/space_generator/space_generator.cc b/src/meta_schedule/space_generator/space_generator.cc
index 8712f5ad48..7306fffcb1 100644
--- a/src/meta_schedule/space_generator/space_generator.cc
+++ b/src/meta_schedule/space_generator/space_generator.cc
@@ -190,13 +190,13 @@ SpaceGenerator SpaceGenerator::PySpaceGenerator(
TVM_REGISTER_OBJECT_TYPE(SpaceGeneratorNode);
TVM_REGISTER_NODE_TYPE(PySpaceGeneratorNode);
-TVM_REGISTER_GLOBAL("meta_schedule.SpaceGeneratorInitializeWithTuneContext")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.SpaceGeneratorInitializeWithTuneContext")
.set_body_method(&SpaceGeneratorNode::InitializeWithTuneContext);
-TVM_REGISTER_GLOBAL("meta_schedule.SpaceGeneratorGenerateDesignSpace")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.SpaceGeneratorGenerateDesignSpace")
.set_body_method(&SpaceGeneratorNode::GenerateDesignSpace);
-TVM_REGISTER_GLOBAL("meta_schedule.SpaceGeneratorPySpaceGenerator")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.SpaceGeneratorPySpaceGenerator")
.set_body_typed(SpaceGenerator::PySpaceGenerator);
-TVM_REGISTER_GLOBAL("meta_schedule.SpaceGeneratorClone")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.SpaceGeneratorClone")
.set_body_method(&SpaceGeneratorNode::Clone);
} // namespace meta_schedule
diff --git a/src/meta_schedule/space_generator/space_generator_union.cc b/src/meta_schedule/space_generator/space_generator_union.cc
index 819a4ee5f7..12bf753494 100644
--- a/src/meta_schedule/space_generator/space_generator_union.cc
+++ b/src/meta_schedule/space_generator/space_generator_union.cc
@@ -82,7 +82,7 @@ SpaceGenerator SpaceGenerator::SpaceGeneratorUnion(Array<SpaceGenerator> space_g
}
TVM_REGISTER_NODE_TYPE(SpaceGeneratorUnionNode);
-TVM_REGISTER_GLOBAL("meta_schedule.SpaceGeneratorSpaceGeneratorUnion")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.SpaceGeneratorSpaceGeneratorUnion")
.set_body_typed(SpaceGenerator::SpaceGeneratorUnion);
} // namespace meta_schedule
diff --git a/src/meta_schedule/task_scheduler/gradient_based.cc b/src/meta_schedule/task_scheduler/gradient_based.cc
index c750067ace..23d23e6243 100644
--- a/src/meta_schedule/task_scheduler/gradient_based.cc
+++ b/src/meta_schedule/task_scheduler/gradient_based.cc
@@ -145,7 +145,7 @@ TaskScheduler TaskScheduler::GradientBased(ffi::Function logger, double alpha, i
}
TVM_REGISTER_NODE_TYPE(GradientBasedNode);
-TVM_REGISTER_GLOBAL("meta_schedule.TaskSchedulerGradientBased")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.TaskSchedulerGradientBased")
.set_body_typed(TaskScheduler::GradientBased);
} // namespace meta_schedule
diff --git a/src/meta_schedule/task_scheduler/round_robin.cc b/src/meta_schedule/task_scheduler/round_robin.cc
index d7c6f37e12..9792fa7e7c 100644
--- a/src/meta_schedule/task_scheduler/round_robin.cc
+++ b/src/meta_schedule/task_scheduler/round_robin.cc
@@ -63,7 +63,7 @@ TaskScheduler TaskScheduler::RoundRobin(ffi::Function logger) {
}
TVM_REGISTER_NODE_TYPE(RoundRobinNode);
-TVM_REGISTER_GLOBAL("meta_schedule.TaskSchedulerRoundRobin")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.TaskSchedulerRoundRobin")
.set_body_typed(TaskScheduler::RoundRobin);
} // namespace meta_schedule
diff --git a/src/meta_schedule/task_scheduler/task_scheduler.cc b/src/meta_schedule/task_scheduler/task_scheduler.cc
index 4411bd0792..85a4063653 100644
--- a/src/meta_schedule/task_scheduler/task_scheduler.cc
+++ b/src/meta_schedule/task_scheduler/task_scheduler.cc
@@ -364,18 +364,19 @@ void PyTaskSchedulerNode::Tune(Array<TuneContext> tasks, Array<FloatImm> task_we
TVM_REGISTER_NODE_TYPE(TaskRecordNode);
TVM_REGISTER_OBJECT_TYPE(TaskSchedulerNode);
TVM_REGISTER_NODE_TYPE(PyTaskSchedulerNode);
-TVM_REGISTER_GLOBAL("meta_schedule.TaskSchedulerPyTaskScheduler")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.TaskSchedulerPyTaskScheduler")
.set_body_typed(TaskScheduler::PyTaskScheduler);
-TVM_REGISTER_GLOBAL("meta_schedule.TaskSchedulerTune").set_body_method(&TaskSchedulerNode::Tune);
-TVM_REGISTER_GLOBAL("meta_schedule.TaskSchedulerJoinRunningTask")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.TaskSchedulerTune")
+ .set_body_method(&TaskSchedulerNode::Tune);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.TaskSchedulerJoinRunningTask")
.set_body_method(&TaskSchedulerNode::JoinRunningTask);
-TVM_REGISTER_GLOBAL("meta_schedule.TaskSchedulerNextTaskId")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.TaskSchedulerNextTaskId")
.set_body_method(&TaskSchedulerNode::NextTaskId);
-TVM_REGISTER_GLOBAL("meta_schedule.TaskSchedulerTerminateTask")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.TaskSchedulerTerminateTask")
.set_body_method(&TaskSchedulerNode::TerminateTask);
-TVM_REGISTER_GLOBAL("meta_schedule.TaskSchedulerTouchTask")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.TaskSchedulerTouchTask")
.set_body_method(&TaskSchedulerNode::TouchTask);
-TVM_REGISTER_GLOBAL("meta_schedule.TaskSchedulerPrintTuningStatistics")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.TaskSchedulerPrintTuningStatistics")
.set_body_method(&TaskSchedulerNode::PrintTuningStatistics);
} // namespace meta_schedule
diff --git a/src/meta_schedule/trace_apply.cc b/src/meta_schedule/trace_apply.cc
index 5ba3f3123c..9d22554d91 100644
--- a/src/meta_schedule/trace_apply.cc
+++ b/src/meta_schedule/trace_apply.cc
@@ -254,7 +254,7 @@ void ScheduleUsingAnchorTrace(Schedule sch, const Trace& anchor_trace, const tvm
}
}
-TVM_REGISTER_GLOBAL("meta_schedule.ScheduleUsingAnchorTrace")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.ScheduleUsingAnchorTrace")
.set_body_typed(ScheduleUsingAnchorTrace);
} // namespace meta_schedule
diff --git a/src/meta_schedule/tune_context.cc b/src/meta_schedule/tune_context.cc
index 275f8d124c..31120ce45d 100644
--- a/src/meta_schedule/tune_context.cc
+++ b/src/meta_schedule/tune_context.cc
@@ -63,7 +63,7 @@ void TuneContextNode::Initialize() {
}
TVM_REGISTER_NODE_TYPE(TuneContextNode);
-TVM_REGISTER_GLOBAL("meta_schedule.TuneContext")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.TuneContext")
.set_body_typed([](Optional<IRModule> mod, Optional<Target> target,
Optional<SpaceGenerator> space_generator,
Optional<SearchStrategy> search_strategy, Optional<String> task_name,
@@ -72,10 +72,10 @@ TVM_REGISTER_GLOBAL("meta_schedule.TuneContext")
return TuneContext(mod, target, space_generator, search_strategy, task_name, num_threads,
rand_state, logger);
});
-TVM_REGISTER_GLOBAL("meta_schedule._SHash2Hex").set_body_typed(SHash2Hex);
-TVM_REGISTER_GLOBAL("meta_schedule.TuneContextInitialize")
+TVM_FFI_REGISTER_GLOBAL("meta_schedule._SHash2Hex").set_body_typed(SHash2Hex);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.TuneContextInitialize")
.set_body_method(&TuneContextNode::Initialize);
-TVM_REGISTER_GLOBAL("meta_schedule.TuneContextClone").set_body_method(&TuneContextNode::Clone);
+TVM_FFI_REGISTER_GLOBAL("meta_schedule.TuneContextClone").set_body_method(&TuneContextNode::Clone);
} // namespace meta_schedule
} // namespace tvm
diff --git a/src/node/container_printing.cc b/src/node/container_printing.cc
index 5fb3d503f5..7441db7832 100644
--- a/src/node/container_printing.cc
+++ b/src/node/container_printing.cc
@@ -21,9 +21,9 @@
* Printer implementation for containers
* \file node/container_printint.cc
*/
+#include <tvm/ffi/function.h>
#include <tvm/node/functor.h>
#include <tvm/node/repr_printer.h>
-#include <tvm/runtime/registry.h>
namespace tvm {
diff --git a/src/node/object_path.cc b/src/node/object_path.cc
index bfb05aa4bc..a99835ea17 100644
--- a/src/node/object_path.cc
+++ b/src/node/object_path.cc
@@ -17,10 +17,10 @@
* under the License.
*/
+#include <tvm/ffi/function.h>
#include <tvm/ffi/memory.h>
#include <tvm/node/object_path.h>
#include <tvm/node/repr_printer.h>
-#include <tvm/runtime/registry.h>
#include <algorithm>
#include <cstring>
@@ -40,13 +40,13 @@ Optional<ObjectPath> ObjectPathNode::GetParent() const {
return Downcast<Optional<ObjectPath>>(parent_);
}
-TVM_REGISTER_GLOBAL("node.ObjectPathGetParent").set_body_method(&ObjectPathNode::GetParent);
+TVM_FFI_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(&ObjectPathNode::Length);
+TVM_FFI_REGISTER_GLOBAL("node.ObjectPathLength").set_body_method(&ObjectPathNode::Length);
// --- GetPrefix ---
@@ -63,7 +63,7 @@ ObjectPath ObjectPathNode::GetPrefix(int32_t length) const {
return GetRef<ObjectPath>(node);
}
-TVM_REGISTER_GLOBAL("node.ObjectPathGetPrefix").set_body_method(&ObjectPathNode::GetPrefix);
+TVM_FFI_REGISTER_GLOBAL("node.ObjectPathGetPrefix").set_body_method(&ObjectPathNode::GetPrefix);
// --- IsPrefixOf ---
@@ -75,7 +75,7 @@ bool ObjectPathNode::IsPrefixOf(const ObjectPath& other) const {
return this->PathsEqual(other->GetPrefix(this_len));
}
-TVM_REGISTER_GLOBAL("node.ObjectPathIsPrefixOf").set_body_method(&ObjectPathNode::IsPrefixOf);
+TVM_FFI_REGISTER_GLOBAL("node.ObjectPathIsPrefixOf").set_body_method(&ObjectPathNode::IsPrefixOf);
// --- Attr ---
@@ -95,7 +95,7 @@ ObjectPath ObjectPathNode::Attr(Optional<String> attr_key) const {
}
}
-TVM_REGISTER_GLOBAL("node.ObjectPathAttr")
+TVM_FFI_REGISTER_GLOBAL("node.ObjectPathAttr")
.set_body_typed([](const ObjectPath& object_path, Optional<String> attr_key) {
return object_path->Attr(attr_key);
});
@@ -106,7 +106,7 @@ ObjectPath ObjectPathNode::ArrayIndex(int32_t index) const {
return ObjectPath(make_object<ArrayIndexPathNode>(this, index));
}
-TVM_REGISTER_GLOBAL("node.ObjectPathArrayIndex").set_body_method(&ObjectPathNode::ArrayIndex);
+TVM_FFI_REGISTER_GLOBAL("node.ObjectPathArrayIndex").set_body_method(&ObjectPathNode::ArrayIndex);
// --- MissingArrayElement ---
@@ -114,7 +114,7 @@ ObjectPath ObjectPathNode::MissingArrayElement(int32_t index) const {
return ObjectPath(make_object<MissingArrayElementPathNode>(this, index));
}
-TVM_REGISTER_GLOBAL("node.ObjectPathMissingArrayElement")
+TVM_FFI_REGISTER_GLOBAL("node.ObjectPathMissingArrayElement")
.set_body_method(&ObjectPathNode::MissingArrayElement);
// --- MapValue ---
@@ -123,7 +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(&ObjectPathNode::MapValue);
+TVM_FFI_REGISTER_GLOBAL("node.ObjectPathMapValue").set_body_method(&ObjectPathNode::MapValue);
// --- MissingMapEntry ---
@@ -131,7 +131,7 @@ ObjectPath ObjectPathNode::MissingMapEntry() const {
return ObjectPath(make_object<MissingMapEntryPathNode>(this));
}
-TVM_REGISTER_GLOBAL("node.ObjectPathMissingMapEntry")
+TVM_FFI_REGISTER_GLOBAL("node.ObjectPathMissingMapEntry")
.set_body_method(&ObjectPathNode::MissingMapEntry);
// --- PathsEqual ----
@@ -158,7 +158,7 @@ bool ObjectPathNode::PathsEqual(const ObjectPath& other) const {
return lhs == nullptr && rhs == nullptr;
}
-TVM_REGISTER_GLOBAL("node.ObjectPathEqual").set_body_method(&ObjectPathNode::PathsEqual);
+TVM_FFI_REGISTER_GLOBAL("node.ObjectPathEqual").set_body_method(&ObjectPathNode::PathsEqual);
// --- Repr ---
@@ -191,7 +191,7 @@ const ObjectPathNode* ObjectPathNode::ParentNode() const {
return ObjectPath(make_object<RootPathNode>(name));
}
-TVM_REGISTER_GLOBAL("node.ObjectPathRoot").set_body_typed(ObjectPath::Root);
+TVM_FFI_REGISTER_GLOBAL("node.ObjectPathRoot").set_body_typed(ObjectPath::Root);
// ============== Individual path classes ==============
diff --git a/src/node/reflection.cc b/src/node/reflection.cc
index 6ba6a4ef1f..2290403d37 100644
--- a/src/node/reflection.cc
+++ b/src/node/reflection.cc
@@ -21,10 +21,10 @@
* Reflection utilities.
* \file node/reflection.cc
*/
+#include <tvm/ffi/function.h>
#include <tvm/ir/attrs.h>
#include <tvm/node/node.h>
#include <tvm/node/reflection.h>
-#include <tvm/runtime/registry.h>
namespace tvm {
@@ -292,11 +292,11 @@ void MakeNode(const ffi::PackedArgs& args, ffi::Any* rv) {
*rv = ReflectionVTable::Global()->CreateObject(type_key, args.Slice(1));
}
-TVM_REGISTER_GLOBAL("node.NodeGetAttr").set_body_packed(NodeGetAttr);
+TVM_FFI_REGISTER_GLOBAL("node.NodeGetAttr").set_body_packed(NodeGetAttr);
-TVM_REGISTER_GLOBAL("node.NodeListAttrNames").set_body_packed(NodeListAttrNames);
+TVM_FFI_REGISTER_GLOBAL("node.NodeListAttrNames").set_body_packed(NodeListAttrNames);
-TVM_REGISTER_GLOBAL("node.MakeNode").set_body_packed(MakeNode);
+TVM_FFI_REGISTER_GLOBAL("node.MakeNode").set_body_packed(MakeNode);
namespace {
// Attribute visitor class for finding the attribute key by its address
diff --git a/src/node/repr_printer.cc b/src/node/repr_printer.cc
index 3e80751e66..aa999655c0 100644
--- a/src/node/repr_printer.cc
+++ b/src/node/repr_printer.cc
@@ -21,9 +21,9 @@
* Printer utilities
* \file node/repr_printer.cc
*/
+#include <tvm/ffi/function.h>
#include <tvm/node/repr_printer.h>
#include <tvm/runtime/device_api.h>
-#include <tvm/runtime/registry.h>
namespace tvm {
@@ -133,12 +133,12 @@ void Dump(const runtime::ObjectRef& n) { std::cerr << n << "\n"; }
void Dump(const runtime::Object* n) { Dump(runtime::GetRef<runtime::ObjectRef>(n)); }
-TVM_REGISTER_GLOBAL("node.AsRepr").set_body_typed([](ffi::Any obj) {
+TVM_FFI_REGISTER_GLOBAL("node.AsRepr").set_body_typed([](ffi::Any obj) {
std::ostringstream os;
os << obj;
return os.str();
});
-TVM_REGISTER_GLOBAL("node.AsLegacyRepr").set_body_typed(ffi::AsLegacyRepr);
+TVM_FFI_REGISTER_GLOBAL("node.AsLegacyRepr").set_body_typed(ffi::AsLegacyRepr);
} // namespace tvm
diff --git a/src/node/script_printer.cc b/src/node/script_printer.cc
index dd249b77db..ee7880f448 100644
--- a/src/node/script_printer.cc
+++ b/src/node/script_printer.cc
@@ -16,10 +16,10 @@
* specific language governing permissions and limitations
* under the License.
*/
+#include <tvm/ffi/function.h>
#include <tvm/ir/expr.h>
#include <tvm/node/repr_printer.h>
#include <tvm/node/script_printer.h>
-#include <tvm/runtime/registry.h>
#include <algorithm>
@@ -135,9 +135,9 @@ Array<String> PrinterConfigNode::GetBuiltinKeywords() {
}
TVM_REGISTER_NODE_TYPE(PrinterConfigNode);
-TVM_REGISTER_GLOBAL("node.PrinterConfig").set_body_typed([](Map<String, Any> config_dict) {
+TVM_FFI_REGISTER_GLOBAL("node.PrinterConfig").set_body_typed([](Map<String, Any> config_dict) {
return PrinterConfig(config_dict);
});
-TVM_REGISTER_GLOBAL("node.TVMScriptPrinterScript").set_body_typed(TVMScriptPrinter::Script);
+TVM_FFI_REGISTER_GLOBAL("node.TVMScriptPrinterScript").set_body_typed(TVMScriptPrinter::Script);
} // namespace tvm
diff --git a/src/node/serialization.cc b/src/node/serialization.cc
index 7f74012c0d..631d70e235 100644
--- a/src/node/serialization.cc
+++ b/src/node/serialization.cc
@@ -23,18 +23,17 @@
*/
#include <dmlc/json.h>
#include <dmlc/memory_io.h>
+#include <tvm/ffi/function.h>
#include <tvm/ir/attrs.h>
#include <tvm/node/reflection.h>
#include <tvm/node/serialization.h>
#include <tvm/runtime/ndarray.h>
#include <tvm/runtime/packed_func.h>
-#include <tvm/runtime/registry.h>
#include <cctype>
#include <map>
#include <string>
-#include "../runtime/object_internal.h"
#include "../support/base64.h"
namespace tvm {
@@ -701,7 +700,7 @@ Any LoadJSON(std::string json_str) {
return nodes.at(jgraph.root);
}
-TVM_REGISTER_GLOBAL("node.SaveJSON").set_body_typed(SaveJSON);
+TVM_FFI_REGISTER_GLOBAL("node.SaveJSON").set_body_typed(SaveJSON);
-TVM_REGISTER_GLOBAL("node.LoadJSON").set_body_typed(LoadJSON);
+TVM_FFI_REGISTER_GLOBAL("node.LoadJSON").set_body_typed(LoadJSON);
} // namespace tvm
diff --git a/src/node/structural_equal.cc b/src/node/structural_equal.cc
index 17a96b5e92..6b19fb5355 100644
--- a/src/node/structural_equal.cc
+++ b/src/node/structural_equal.cc
@@ -19,13 +19,13 @@
/*!
* \file src/node/structural_equal.cc
*/
+#include <tvm/ffi/function.h>
#include <tvm/ir/module.h>
#include <tvm/node/functor.h>
#include <tvm/node/node.h>
#include <tvm/node/object_path.h>
#include <tvm/node/reflection.h>
#include <tvm/node/structural_equal.h>
-#include <tvm/runtime/registry.h>
#include <optional>
#include <unordered_map>
@@ -36,12 +36,12 @@ namespace tvm {
TVM_REGISTER_OBJECT_TYPE(ObjectPathPairNode);
-TVM_REGISTER_GLOBAL("node.ObjectPathPairLhsPath")
+TVM_FFI_REGISTER_GLOBAL("node.ObjectPathPairLhsPath")
.set_body_typed([](const ObjectPathPair& object_path_pair) {
return object_path_pair->lhs_path;
});
-TVM_REGISTER_GLOBAL("node.ObjectPathPairRhsPath")
+TVM_FFI_REGISTER_GLOBAL("node.ObjectPathPairRhsPath")
.set_body_typed([](const ObjectPathPair& object_path_pair) {
return object_path_pair->rhs_path;
});
@@ -595,7 +595,7 @@ bool SEqualHandlerDefault::DispatchSEqualReduce(const ObjectRef& lhs, const Obje
return impl->DispatchSEqualReduce(lhs, rhs, map_free_vars, current_paths);
}
-TVM_REGISTER_GLOBAL("node.StructuralEqual")
+TVM_FFI_REGISTER_GLOBAL("node.StructuralEqual")
.set_body_typed([](const Any& lhs, const Any& rhs, bool assert_mode, bool map_free_vars) {
// If we are asserting on failure, then the `defer_fails` option
// should be enabled, to provide better error messages. For
@@ -608,7 +608,7 @@ TVM_REGISTER_GLOBAL("node.StructuralEqual")
.Equal(lhs, rhs, map_free_vars);
});
-TVM_REGISTER_GLOBAL("node.GetFirstStructuralMismatch")
+TVM_FFI_REGISTER_GLOBAL("node.GetFirstStructuralMismatch")
.set_body_typed([](const Any& lhs, const Any& rhs, bool map_free_vars) {
Optional<ObjectPathPair> first_mismatch;
bool equal =
diff --git a/src/node/structural_hash.cc b/src/node/structural_hash.cc
index 14c7b8a39a..efaa7037b0 100644
--- a/src/node/structural_hash.cc
+++ b/src/node/structural_hash.cc
@@ -20,13 +20,13 @@
* \file src/node/structural_hash.cc
*/
#include <dmlc/memory_io.h>
+#include <tvm/ffi/function.h>
#include <tvm/node/functor.h>
#include <tvm/node/node.h>
#include <tvm/node/object_path.h>
#include <tvm/node/reflection.h>
#include <tvm/node/structural_hash.h>
#include <tvm/runtime/profiling.h>
-#include <tvm/runtime/registry.h>
#include <tvm/target/codegen.h>
#include <algorithm>
@@ -291,7 +291,7 @@ void SHashHandlerDefault::DispatchSHash(const ObjectRef& key, bool map_free_vars
impl->DispatchSHash(key, map_free_vars);
}
-TVM_REGISTER_GLOBAL("node.StructuralHash")
+TVM_FFI_REGISTER_GLOBAL("node.StructuralHash")
.set_body_typed([](const Any& object, bool map_free_vars) -> int64_t {
uint64_t hashed_value = SHashHandlerDefault().Hash(object, map_free_vars);
return static_cast<int64_t>(hashed_value);
diff --git a/src/relax/analysis/analysis.cc b/src/relax/analysis/analysis.cc
index b4a5a039bb..98122d1e1e 100644
--- a/src/relax/analysis/analysis.cc
+++ b/src/relax/analysis/analysis.cc
@@ -197,15 +197,15 @@ bool ContainsImpureCall(const Expr& expr, const Optional<Expr>& own_name) {
return FindImpureCall(expr, own_name).defined();
}
-TVM_REGISTER_GLOBAL("relax.analysis.free_vars").set_body_typed(FreeVars);
+TVM_FFI_REGISTER_GLOBAL("relax.analysis.free_vars").set_body_typed(FreeVars);
-TVM_REGISTER_GLOBAL("relax.analysis.bound_vars").set_body_typed(BoundVars);
+TVM_FFI_REGISTER_GLOBAL("relax.analysis.bound_vars").set_body_typed(BoundVars);
-TVM_REGISTER_GLOBAL("relax.analysis.all_vars").set_body_typed(AllVars);
+TVM_FFI_REGISTER_GLOBAL("relax.analysis.all_vars").set_body_typed(AllVars);
-TVM_REGISTER_GLOBAL("relax.analysis.all_global_vars").set_body_typed(AllGlobalVars);
+TVM_FFI_REGISTER_GLOBAL("relax.analysis.all_global_vars").set_body_typed(AllGlobalVars);
-TVM_REGISTER_GLOBAL("relax.analysis.contains_impure_call").set_body_typed(ContainsImpureCall);
+TVM_FFI_REGISTER_GLOBAL("relax.analysis.contains_impure_call").set_body_typed(ContainsImpureCall);
} // namespace relax
} // namespace tvm
diff --git a/src/relax/analysis/computable_at_compile_time.cc b/src/relax/analysis/computable_at_compile_time.cc
index 37bbf3a977..ba163b51d6 100644
--- a/src/relax/analysis/computable_at_compile_time.cc
+++ b/src/relax/analysis/computable_at_compile_time.cc
@@ -92,7 +92,7 @@ Array<Var> ComputableAtCompileTime(const Function& func) {
return CompileTimeCollector::Collect(func);
}
-TVM_REGISTER_GLOBAL("relax.analysis.computable_at_compile_time")
+TVM_FFI_REGISTER_GLOBAL("relax.analysis.computable_at_compile_time")
.set_body_typed(ComputableAtCompileTime);
} // namespace relax
diff --git a/src/relax/analysis/detect_recursion.cc b/src/relax/analysis/detect_recursion.cc
index 9c150fed8b..48ec7880b1 100644
--- a/src/relax/analysis/detect_recursion.cc
+++ b/src/relax/analysis/detect_recursion.cc
@@ -392,7 +392,7 @@ tvm::Array<tvm::Array<GlobalVar>> DetectRecursion(const IRModule& m) {
return ret;
}
-TVM_REGISTER_GLOBAL("relax.analysis.detect_recursion").set_body_typed(DetectRecursion);
+TVM_FFI_REGISTER_GLOBAL("relax.analysis.detect_recursion").set_body_typed(DetectRecursion);
} // namespace relax
} // namespace tvm
diff --git a/src/relax/analysis/layout_transformation.cc b/src/relax/analysis/layout_transformation.cc
index f0658dabb3..ab32abab5b 100644
--- a/src/relax/analysis/layout_transformation.cc
+++ b/src/relax/analysis/layout_transformation.cc
@@ -614,7 +614,7 @@ Map<tir::Block, Map<ObjectRef, tir::IndexMap>> SuggestLayoutTransforms(
return analyzer.GetSuggestedTransforms();
}
-TVM_REGISTER_GLOBAL(("relax.analysis.suggest_layout_transforms"))
+TVM_FFI_REGISTER_GLOBAL(("relax.analysis.suggest_layout_transforms"))
.set_body_typed([](PrimFunc fn, Array<tir::IndexMap> write_buffer_transformations) {
return SuggestLayoutTransforms(fn, write_buffer_transformations);
});
diff --git a/src/relax/analysis/struct_info_analysis.cc b/src/relax/analysis/struct_info_analysis.cc
index 5ad8a173c0..e09f061001 100644
--- a/src/relax/analysis/struct_info_analysis.cc
+++ b/src/relax/analysis/struct_info_analysis.cc
@@ -72,7 +72,7 @@ class StaticTypeDeriver : public StructInfoFunctor<Type(const StructInfo&)> {
Type GetStaticType(const StructInfo& info) { return StaticTypeDeriver()(info); }
-TVM_REGISTER_GLOBAL("relax.analysis.GetStaticType").set_body_typed([](const StructInfo& info) {
+TVM_FFI_REGISTER_GLOBAL("relax.analysis.GetStaticType").set_body_typed([](const StructInfo& info) {
return GetStaticType(info);
});
@@ -285,7 +285,7 @@ StructInfo EraseToWellDefined(const StructInfo& info, Map<tir::Var, PrimExpr> sh
return EraseToWellDefined(info, f_shape_var_map, f_var_map, ana);
}
-TVM_REGISTER_GLOBAL("relax.analysis.EraseToWellDefined")
+TVM_FFI_REGISTER_GLOBAL("relax.analysis.EraseToWellDefined")
.set_body_typed([](const StructInfo& info, Map<tir::Var, PrimExpr> shape_var_map,
Map<Var, Expr> var_map) {
return EraseToWellDefined(info, shape_var_map, var_map);
@@ -595,7 +595,7 @@ BaseCheckResult StructInfoBaseCheck(const StructInfo& base, const StructInfo& de
}
}
-TVM_REGISTER_GLOBAL("relax.analysis.StructInfoBaseCheck")
+TVM_FFI_REGISTER_GLOBAL("relax.analysis.StructInfoBaseCheck")
.set_body_typed([](const StructInfo& base, const StructInfo& derived) -> int {
return static_cast<int>(StructInfoBaseCheck(base, derived));
});
@@ -604,7 +604,7 @@ bool IsBaseOf(const StructInfo& base, const StructInfo& derived, arith::Analyzer
return StructInfoBaseCheck(base, derived, ana) == BaseCheckResult::kPass;
}
-TVM_REGISTER_GLOBAL("relax.StructInfoIsBaseOf")
+TVM_FFI_REGISTER_GLOBAL("relax.StructInfoIsBaseOf")
.set_body_typed([](const StructInfo& base, const StructInfo& derived) {
return IsBaseOf(base, derived);
});
@@ -955,7 +955,7 @@ StructInfo DeriveCallRetStructInfo(const FuncStructInfo& finfo, const Call& call
}
}
-TVM_REGISTER_GLOBAL("relax.analysis.DeriveCallRetStructInfo")
+TVM_FFI_REGISTER_GLOBAL("relax.analysis.DeriveCallRetStructInfo")
.set_body_typed([](const FuncStructInfo& finfo, const Call& call, const BlockBuilder& ctx) {
return DeriveCallRetStructInfo(finfo, call, ctx);
});
@@ -1158,7 +1158,7 @@ StructInfo StructInfoLCA(const StructInfo& lhs, const StructInfo& rhs, arith::An
}
}
-TVM_REGISTER_GLOBAL("relax.analysis.StructInfoLCA")
+TVM_FFI_REGISTER_GLOBAL("relax.analysis.StructInfoLCA")
.set_body_typed([](const StructInfo& lhs, const StructInfo& rhs) {
return StructInfoLCA(lhs, rhs);
});
@@ -1241,9 +1241,9 @@ Array<tir::Var> DefinableTIRVarsInStructInfo(const StructInfo& sinfo) {
return detector.GetTIRVars();
}
-TVM_REGISTER_GLOBAL("relax.analysis.TIRVarsInStructInfo").set_body_typed(TIRVarsInStructInfo);
+TVM_FFI_REGISTER_GLOBAL("relax.analysis.TIRVarsInStructInfo").set_body_typed(TIRVarsInStructInfo);
-TVM_REGISTER_GLOBAL("relax.analysis.DefinableTIRVarsInStructInfo")
+TVM_FFI_REGISTER_GLOBAL("relax.analysis.DefinableTIRVarsInStructInfo")
.set_body_typed(DefinableTIRVarsInStructInfo);
class NonNegativeExpressionCollector : relax::StructInfoVisitor {
@@ -1288,7 +1288,7 @@ Array<PrimExpr> CollectNonNegativeExpressions(const StructInfo& sinfo) {
return NonNegativeExpressionCollector::Collect(sinfo);
}
-TVM_REGISTER_GLOBAL("relax.analysis.CollectNonNegativeExpressions")
+TVM_FFI_REGISTER_GLOBAL("relax.analysis.CollectNonNegativeExpressions")
.set_body_typed(CollectNonNegativeExpressions);
class SymbolicVarCollector : public relax::ExprVisitor,
@@ -1436,9 +1436,9 @@ Array<tir::Var> DefinedSymbolicVars(const Expr& expr) {
}
Array<tir::Var> FreeSymbolicVars(const Expr& expr) { return SymbolicVarCollector::Free(expr); }
-TVM_REGISTER_GLOBAL("relax.analysis.DefinedSymbolicVars").set_body_typed(DefinedSymbolicVars);
+TVM_FFI_REGISTER_GLOBAL("relax.analysis.DefinedSymbolicVars").set_body_typed(DefinedSymbolicVars);
-TVM_REGISTER_GLOBAL("relax.analysis.FreeSymbolicVars").set_body_typed(FreeSymbolicVars);
+TVM_FFI_REGISTER_GLOBAL("relax.analysis.FreeSymbolicVars").set_body_typed(FreeSymbolicVars);
} // namespace relax
} // namespace tvm
diff --git a/src/relax/analysis/tir_op_pattern_kind.cc b/src/relax/analysis/tir_op_pattern_kind.cc
index 177fdeff58..0845ec092f 100644
--- a/src/relax/analysis/tir_op_pattern_kind.cc
+++ b/src/relax/analysis/tir_op_pattern_kind.cc
@@ -537,7 +537,7 @@ bool HasReshapePattern(const PrimFunc& func) {
return ReshapeDetector::Detect(src_buffer, dst_buffer, func->body);
}
-TVM_REGISTER_GLOBAL("relax.analysis.has_reshape_pattern").set_body_typed(HasReshapePattern);
+TVM_FFI_REGISTER_GLOBAL("relax.analysis.has_reshape_pattern").set_body_typed(HasReshapePattern);
} // namespace relax
} // namespace tvm
diff --git a/src/relax/analysis/udchain.cc b/src/relax/analysis/udchain.cc
index 1c16f25c99..f62254b695 100644
--- a/src/relax/analysis/udchain.cc
+++ b/src/relax/analysis/udchain.cc
@@ -118,7 +118,7 @@ Map<Var, Array<Var>> DataflowBlockUseDef(const DataflowBlock& dfb) {
return usage.downstream_usage;
}
-TVM_REGISTER_GLOBAL("relax.analysis.udchain").set_body_typed(DataflowBlockUseDef);
+TVM_FFI_REGISTER_GLOBAL("relax.analysis.udchain").set_body_typed(DataflowBlockUseDef);
VarUsageInfo CollectVarUsage(const Expr& expr) { return UDChain::Collect(expr); }
diff --git a/src/relax/analysis/var2value.cc b/src/relax/analysis/var2value.cc
index 7625841eb6..a367d33ca4 100644
--- a/src/relax/analysis/var2value.cc
+++ b/src/relax/analysis/var2value.cc
@@ -58,7 +58,7 @@ Map<Var, Expr> AnalyzeVar2Value(const IRModule& m) {
return std::move(var2val_analysis.var2value_);
}
-TVM_REGISTER_GLOBAL(("relax.analysis.get_var2val")).set_body_typed([](const Function& f) {
+TVM_FFI_REGISTER_GLOBAL(("relax.analysis.get_var2val")).set_body_typed([](const Function& f) {
return AnalyzeVar2Value(f);
});
@@ -85,7 +85,7 @@ Map<String, Array<Binding>> NameToBinding(const Function& fn) {
std::make_move_iterator(analysis.name2bindings_.end()));
}
-TVM_REGISTER_GLOBAL(("relax.analysis.name_to_binding")).set_body_typed(NameToBinding);
+TVM_FFI_REGISTER_GLOBAL(("relax.analysis.name_to_binding")).set_body_typed(NameToBinding);
} // namespace relax
} // namespace tvm
diff --git a/src/relax/analysis/well_formed.cc b/src/relax/analysis/well_formed.cc
index 316e87f351..243033e945 100644
--- a/src/relax/analysis/well_formed.cc
+++ b/src/relax/analysis/well_formed.cc
@@ -646,7 +646,7 @@ bool WellFormed(Variant<IRModule, Function> obj, bool check_struct_info) {
return WellFormedChecker::Check(obj, check_struct_info);
}
-TVM_REGISTER_GLOBAL(("relax.analysis.well_formed")).set_body_typed(WellFormed);
+TVM_FFI_REGISTER_GLOBAL(("relax.analysis.well_formed")).set_body_typed(WellFormed);
} // namespace relax
} // namespace tvm
diff --git a/src/relax/backend/contrib/clml/codegen.cc b/src/relax/backend/contrib/clml/codegen.cc
index 3c87079f99..ff64504f61 100644
--- a/src/relax/backend/contrib/clml/codegen.cc
+++ b/src/relax/backend/contrib/clml/codegen.cc
@@ -322,7 +322,7 @@ Array<runtime::Module> OpenCLMLCompiler(Array<Function> functions, Map<String, A
return compiled_functions;
}
-TVM_REGISTER_GLOBAL("relax.ext.openclml").set_body_typed(OpenCLMLCompiler);
+TVM_FFI_REGISTER_GLOBAL("relax.ext.openclml").set_body_typed(OpenCLMLCompiler);
/*!
* \brief Check whether OpenCLML graph executor is enabled.
@@ -348,8 +348,9 @@ Integer GetOpenCLMLVersion() {
#endif // TVM_GRAPH_EXECUTOR_CLML
}
-TVM_REGISTER_GLOBAL("relax.is_openclml_runtime_enabled").set_body_typed(IsOpenCLMLRuntimeEnabled);
-TVM_REGISTER_GLOBAL("relax.get_openclml_version").set_body_typed(GetOpenCLMLVersion);
+TVM_FFI_REGISTER_GLOBAL("relax.is_openclml_runtime_enabled")
+ .set_body_typed(IsOpenCLMLRuntimeEnabled);
+TVM_FFI_REGISTER_GLOBAL("relax.get_openclml_version").set_body_typed(GetOpenCLMLVersion);
} // namespace contrib
} // namespace relax
diff --git a/src/relax/backend/contrib/cublas/codegen.cc b/src/relax/backend/contrib/cublas/codegen.cc
index 431265d1a7..085535b87f 100644
--- a/src/relax/backend/contrib/cublas/codegen.cc
+++ b/src/relax/backend/contrib/cublas/codegen.cc
@@ -125,7 +125,7 @@ Array<runtime::Module> CublasCompiler(Array<Function> functions, Map<String, ffi
return compiled_functions;
}
-TVM_REGISTER_GLOBAL("relax.ext.cublas").set_body_typed(CublasCompiler);
+TVM_FFI_REGISTER_GLOBAL("relax.ext.cublas").set_body_typed(CublasCompiler);
} // namespace contrib
} // namespace relax
diff --git a/src/relax/backend/contrib/cudnn/codegen.cc b/src/relax/backend/contrib/cudnn/codegen.cc
index 43933c7d2a..7026396b51 100644
--- a/src/relax/backend/contrib/cudnn/codegen.cc
+++ b/src/relax/backend/contrib/cudnn/codegen.cc
@@ -149,7 +149,7 @@ Array<runtime::Module> cuDNNCompiler(Array<Function> functions, Map<String, ffi:
return compiled_functions;
}
-TVM_REGISTER_GLOBAL("relax.ext.cudnn").set_body_typed(cuDNNCompiler);
+TVM_FFI_REGISTER_GLOBAL("relax.ext.cudnn").set_body_typed(cuDNNCompiler);
} // namespace contrib
} // namespace relax
diff --git a/src/relax/backend/contrib/cutlass/codegen.cc b/src/relax/backend/contrib/cutlass/codegen.cc
index 175be21834..a05907a231 100644
--- a/src/relax/backend/contrib/cutlass/codegen.cc
+++ b/src/relax/backend/contrib/cutlass/codegen.cc
@@ -101,7 +101,7 @@ class CodegenResult : public ObjectRef {
TVM_REGISTER_NODE_TYPE(CodegenResultNode);
-TVM_REGISTER_GLOBAL("contrib.cutlass.CodegenResult")
+TVM_FFI_REGISTER_GLOBAL("contrib.cutlass.CodegenResult")
.set_body_typed([](String code, Array<String> headers) {
return CodegenResult(code, headers);
});
@@ -385,7 +385,7 @@ Array<runtime::Module> CUTLASSCompiler(Array<Function> functions, Map<String, ff
return {cutlass_mod};
}
-TVM_REGISTER_GLOBAL("relax.ext.cutlass").set_body_typed(CUTLASSCompiler);
+TVM_FFI_REGISTER_GLOBAL("relax.ext.cutlass").set_body_typed(CUTLASSCompiler);
} // namespace contrib
} // namespace relax
diff --git a/src/relax/backend/contrib/dnnl/codegen.cc b/src/relax/backend/contrib/dnnl/codegen.cc
index ba219be4c6..906cf941d4 100644
--- a/src/relax/backend/contrib/dnnl/codegen.cc
+++ b/src/relax/backend/contrib/dnnl/codegen.cc
@@ -97,7 +97,7 @@ Array<runtime::Module> DNNLCompiler(Array<Function> functions, Map<String, ffi::
return compiled_functions;
}
-TVM_REGISTER_GLOBAL("relax.ext.dnnl").set_body_typed(DNNLCompiler);
+TVM_FFI_REGISTER_GLOBAL("relax.ext.dnnl").set_body_typed(DNNLCompiler);
} // namespace contrib
} // namespace relax
diff --git a/src/relax/backend/contrib/hipblas/codegen.cc b/src/relax/backend/contrib/hipblas/codegen.cc
index 54260693f9..4a4666a92d 100644
--- a/src/relax/backend/contrib/hipblas/codegen.cc
+++ b/src/relax/backend/contrib/hipblas/codegen.cc
@@ -102,7 +102,7 @@ Array<runtime::Module> HipblasCompiler(Array<Function> functions, Map<String, ff
return compiled_functions;
}
-TVM_REGISTER_GLOBAL("relax.ext.hipblas").set_body_typed(HipblasCompiler);
+TVM_FFI_REGISTER_GLOBAL("relax.ext.hipblas").set_body_typed(HipblasCompiler);
} // namespace contrib
} // namespace relax
diff --git a/src/relax/backend/contrib/nnapi/codegen.cc b/src/relax/backend/contrib/nnapi/codegen.cc
index f0517fa359..d9f26c4172 100644
--- a/src/relax/backend/contrib/nnapi/codegen.cc
+++ b/src/relax/backend/contrib/nnapi/codegen.cc
@@ -17,11 +17,11 @@
* under the License.
*/
+#include <tvm/ffi/function.h>
#include <tvm/ir/module.h>
#include <tvm/relax/analysis.h>
#include <tvm/relax/type.h>
#include <tvm/runtime/logging.h>
-#include <tvm/runtime/registry.h>
#include <memory>
#include <string>
@@ -264,7 +264,7 @@ Array<runtime::Module> NNAPICompiler(Array<Function> functions, Map<String, ffi:
return compiled_functions;
}
-TVM_REGISTER_GLOBAL("relax.ext.nnapi").set_body_typed(NNAPICompiler);
+TVM_FFI_REGISTER_GLOBAL("relax.ext.nnapi").set_body_typed(NNAPICompiler);
} // namespace contrib
} // namespace relax
diff --git a/src/relax/backend/contrib/tensorrt/codegen.cc b/src/relax/backend/contrib/tensorrt/codegen.cc
index c9cf80695f..8f9b6b5f51 100644
--- a/src/relax/backend/contrib/tensorrt/codegen.cc
+++ b/src/relax/backend/contrib/tensorrt/codegen.cc
@@ -231,7 +231,7 @@ Array<runtime::Module> TensorRTCompiler(Array<Function> functions, Map<String, f
return compiled_functions;
}
-TVM_REGISTER_GLOBAL("relax.ext.tensorrt").set_body_typed(TensorRTCompiler);
+TVM_FFI_REGISTER_GLOBAL("relax.ext.tensorrt").set_body_typed(TensorRTCompiler);
/*!
* \brief Check whether TensorRT graph executor is enabled.
@@ -258,8 +258,9 @@ Array<Integer> GetTensorRTVersion() {
#endif // TVM_GRAPH_EXECUTOR_TENSORRT
}
-TVM_REGISTER_GLOBAL("relax.is_tensorrt_runtime_enabled").set_body_typed(IsTensorRTRuntimeEnabled);
-TVM_REGISTER_GLOBAL("relax.get_tensorrt_version").set_body_typed(GetTensorRTVersion);
+TVM_FFI_REGISTER_GLOBAL("relax.is_tensorrt_runtime_enabled")
+ .set_body_typed(IsTensorRTRuntimeEnabled);
+TVM_FFI_REGISTER_GLOBAL("relax.get_tensorrt_version").set_body_typed(GetTensorRTVersion);
} // namespace contrib
} // namespace relax
diff --git a/src/relax/backend/contrib/utils.cc b/src/relax/backend/contrib/utils.cc
index 8e214809dd..6574ccc37a 100644
--- a/src/relax/backend/contrib/utils.cc
+++ b/src/relax/backend/contrib/utils.cc
@@ -75,7 +75,7 @@ bool EndsWithPattern(const std::string& str, const std::string& pattern) {
return str.compare(str.length() - pattern.length(), pattern.length(), pattern) == 0;
}
-TVM_REGISTER_GLOBAL("relax.contrib.extract_arg_idx").set_body_typed(ExtractArgIdx);
+TVM_FFI_REGISTER_GLOBAL("relax.contrib.extract_arg_idx").set_body_typed(ExtractArgIdx);
} // namespace backend
} // namespace relax
diff --git a/src/relax/backend/pattern_registry.cc b/src/relax/backend/pattern_registry.cc
index 9feeca6662..840b44c128 100644
--- a/src/relax/backend/pattern_registry.cc
+++ b/src/relax/backend/pattern_registry.cc
@@ -67,10 +67,11 @@ Optional<FusionPattern> GetPattern(const String& pattern_name) {
return std::nullopt;
}
-TVM_REGISTER_GLOBAL("relax.backend.RegisterPatterns").set_body_typed(RegisterPatterns);
-TVM_REGISTER_GLOBAL("relax.backend.RemovePatterns").set_body_typed(RemovePatterns);
-TVM_REGISTER_GLOBAL("relax.backend.GetPatternsWithPrefix").set_body_typed(GetPatternsWithPrefix);
-TVM_REGISTER_GLOBAL("relax.backend.GetPattern").set_body_typed(GetPattern);
+TVM_FFI_REGISTER_GLOBAL("relax.backend.RegisterPatterns").set_body_typed(RegisterPatterns);
+TVM_FFI_REGISTER_GLOBAL("relax.backend.RemovePatterns").set_body_typed(RemovePatterns);
+TVM_FFI_REGISTER_GLOBAL("relax.backend.GetPatternsWithPrefix")
+ .set_body_typed(GetPatternsWithPrefix);
+TVM_FFI_REGISTER_GLOBAL("relax.backend.GetPattern").set_body_typed(GetPattern);
} // namespace backend
} // namespace relax
diff --git a/src/relax/backend/task_extraction.cc b/src/relax/backend/task_extraction.cc
index af6e83cf5f..686d24de62 100644
--- a/src/relax/backend/task_extraction.cc
+++ b/src/relax/backend/task_extraction.cc
@@ -139,7 +139,7 @@ class TaskExtractor : public ExprVisitor {
std::optional<tvm::ffi::Function> normalize_mod_func_;
};
-TVM_REGISTER_GLOBAL("relax.backend.MetaScheduleExtractTask")
+TVM_FFI_REGISTER_GLOBAL("relax.backend.MetaScheduleExtractTask")
.set_body_typed([](IRModule mod, Target target, String mod_eq_name) {
return TaskExtractor::ExtractTask(std::move(mod), std::move(target), std::move(mod_eq_name));
});
diff --git a/src/relax/backend/vm/codegen_vm.cc b/src/relax/backend/vm/codegen_vm.cc
index 44acd6ea81..f61579e25e 100644
--- a/src/relax/backend/vm/codegen_vm.cc
+++ b/src/relax/backend/vm/codegen_vm.cc
@@ -425,7 +425,7 @@ IRModule VMCodeGen(ExecBuilder exec_builder, IRModule mod) {
return CodeGenVM::Run(exec_builder, mod);
}
-TVM_REGISTER_GLOBAL("relax.VMCodeGen").set_body_typed(VMCodeGen);
+TVM_FFI_REGISTER_GLOBAL("relax.VMCodeGen").set_body_typed(VMCodeGen);
/*!
* \brief Link the modules together, possibly create a constant module.
@@ -490,7 +490,7 @@ Module VMLink(ExecBuilder builder, Target target, Optional<Module> lib, Array<Mo
return Module(executable);
}
-TVM_REGISTER_GLOBAL("relax.VMLink").set_body_typed(VMLink);
+TVM_FFI_REGISTER_GLOBAL("relax.VMLink").set_body_typed(VMLink);
} // namespace relax_vm
} // namespace relax
diff --git a/src/relax/backend/vm/codegen_vm_tir.cc b/src/relax/backend/vm/codegen_vm_tir.cc
index cfcdf00efb..aa228822de 100644
--- a/src/relax/backend/vm/codegen_vm_tir.cc
+++ b/src/relax/backend/vm/codegen_vm_tir.cc
@@ -530,7 +530,7 @@ IRModule VMTIRCodeGen(ExecBuilder exec_builder, IRModule mod) {
return CodeGenVMTIR::Run(exec_builder, mod);
}
-TVM_REGISTER_GLOBAL("relax.VMTIRCodeGen").set_body_typed(VMTIRCodeGen);
+TVM_FFI_REGISTER_GLOBAL("relax.VMTIRCodeGen").set_body_typed(VMTIRCodeGen);
} // namespace relax_vm
} // namespace relax
diff --git a/src/relax/backend/vm/exec_builder.cc b/src/relax/backend/vm/exec_builder.cc
index 7bc5cafc0a..56b035212e 100644
--- a/src/relax/backend/vm/exec_builder.cc
+++ b/src/relax/backend/vm/exec_builder.cc
@@ -327,9 +327,9 @@ void ExecBuilderNode::Formalize() {
}
}
-TVM_REGISTER_GLOBAL("relax.ExecBuilderCreate").set_body_typed(ExecBuilderNode::Create);
+TVM_FFI_REGISTER_GLOBAL("relax.ExecBuilderCreate").set_body_typed(ExecBuilderNode::Create);
-TVM_REGISTER_GLOBAL("relax.ExecBuilderConvertConstant")
+TVM_FFI_REGISTER_GLOBAL("relax.ExecBuilderConvertConstant")
.set_body_packed([](ffi::PackedArgs args, ffi::Any* ret) {
ExecBuilder builder = args[0].cast<ExecBuilder>();
ffi::Any rt;
@@ -337,20 +337,21 @@ TVM_REGISTER_GLOBAL("relax.ExecBuilderConvertConstant")
*ret = builder->ConvertConstant(rt).data();
});
-TVM_REGISTER_GLOBAL("relax.ExecBuilderEmitFunction")
+TVM_FFI_REGISTER_GLOBAL("relax.ExecBuilderEmitFunction")
.set_body_typed([](ExecBuilder builder, String func, int64_t num_inputs,
Optional<Array<String>> param_names) {
builder->EmitFunction(func, num_inputs, param_names);
});
-TVM_REGISTER_GLOBAL("relax.ExecBuilderEndFunction").set_body_method(&ExecBuilderNode::EndFunction);
+TVM_FFI_REGISTER_GLOBAL("relax.ExecBuilderEndFunction")
+ .set_body_method(&ExecBuilderNode::EndFunction);
-TVM_REGISTER_GLOBAL("relax.ExecBuilderDeclareFunction")
+TVM_FFI_REGISTER_GLOBAL("relax.ExecBuilderDeclareFunction")
.set_body_typed([](ExecBuilder builder, String name, int32_t kind) {
builder->DeclareFunction(name, static_cast<VMFuncInfo::FuncKind>(kind));
});
-TVM_REGISTER_GLOBAL("relax.ExecBuilderEmitCall")
+TVM_FFI_REGISTER_GLOBAL("relax.ExecBuilderEmitCall")
.set_body_typed([](ExecBuilder builder, String name, Array<IntImm> args, int64_t dst) {
std::vector<Instruction::Arg> args_;
for (size_t i = 0; i < args.size(); ++i) {
@@ -360,35 +361,38 @@ TVM_REGISTER_GLOBAL("relax.ExecBuilderEmitCall")
builder->EmitCall(name, args_, dst_.value());
});
-TVM_REGISTER_GLOBAL("relax.ExecBuilderEmitRet")
+TVM_FFI_REGISTER_GLOBAL("relax.ExecBuilderEmitRet")
.set_body_typed([](ExecBuilder builder, int64_t data) {
builder->EmitRet(Instruction::Arg::FromData(data));
});
-TVM_REGISTER_GLOBAL("relax.ExecBuilderEmitGoto").set_body_method(&ExecBuilderNode::EmitGoto);
+TVM_FFI_REGISTER_GLOBAL("relax.ExecBuilderEmitGoto").set_body_method(&ExecBuilderNode::EmitGoto);
-TVM_REGISTER_GLOBAL("relax.ExecBuilderEmitIf")
+TVM_FFI_REGISTER_GLOBAL("relax.ExecBuilderEmitIf")
.set_body_typed([](ExecBuilder builder, int64_t data, vm::Index false_offset) {
builder->EmitIf(Instruction::Arg::FromData(data), false_offset);
});
-TVM_REGISTER_GLOBAL("relax.ExecBuilderR").set_body_typed([](ExecBuilder builder, int64_t value) {
- return Instruction::Arg::Register(value).data();
-});
+TVM_FFI_REGISTER_GLOBAL("relax.ExecBuilderR")
+ .set_body_typed([](ExecBuilder builder, int64_t value) {
+ return Instruction::Arg::Register(value).data();
+ });
-TVM_REGISTER_GLOBAL("relax.ExecBuilderImm").set_body_typed([](ExecBuilder builder, int64_t value) {
- return Instruction::Arg::Immediate(value).data();
-});
+TVM_FFI_REGISTER_GLOBAL("relax.ExecBuilderImm")
+ .set_body_typed([](ExecBuilder builder, int64_t value) {
+ return Instruction::Arg::Immediate(value).data();
+ });
-TVM_REGISTER_GLOBAL("relax.ExecBuilderC").set_body_typed([](ExecBuilder builder, int64_t value) {
- return Instruction::Arg::ConstIdx(value).data();
-});
+TVM_FFI_REGISTER_GLOBAL("relax.ExecBuilderC")
+ .set_body_typed([](ExecBuilder builder, int64_t value) {
+ return Instruction::Arg::ConstIdx(value).data();
+ });
-TVM_REGISTER_GLOBAL("relax.ExecBuilderF").set_body_typed([](ExecBuilder builder, String value) {
+TVM_FFI_REGISTER_GLOBAL("relax.ExecBuilderF").set_body_typed([](ExecBuilder builder, String value) {
return builder->GetFunction(value).data();
});
-TVM_REGISTER_GLOBAL("relax.ExecBuilderGet").set_body_typed([](ExecBuilder builder) {
+TVM_FFI_REGISTER_GLOBAL("relax.ExecBuilderGet").set_body_typed([](ExecBuilder builder) {
ObjectPtr<VMExecutable> p_exec = builder->Get();
return runtime::Module(p_exec);
});
diff --git a/src/relax/backend/vm/lower_runtime_builtin.cc b/src/relax/backend/vm/lower_runtime_builtin.cc
index f746e4a5af..7757195bcb 100644
--- a/src/relax/backend/vm/lower_runtime_builtin.cc
+++ b/src/relax/backend/vm/lower_runtime_builtin.cc
@@ -231,7 +231,7 @@ Pass LowerRuntimeBuiltin() {
return CreateFunctionPass(pass_func, 0, "LowerRuntimeBuiltin", {});
}
-TVM_REGISTER_GLOBAL("relax.transform.LowerRuntimeBuiltin").set_body_typed(LowerRuntimeBuiltin);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.LowerRuntimeBuiltin").set_body_typed(LowerRuntimeBuiltin);
} // namespace transform
} // namespace relax
diff --git a/src/relax/backend/vm/vm_shape_lower.cc b/src/relax/backend/vm/vm_shape_lower.cc
index 6400a33321..0b60553034 100644
--- a/src/relax/backend/vm/vm_shape_lower.cc
+++ b/src/relax/backend/vm/vm_shape_lower.cc
@@ -813,7 +813,7 @@ Pass VMShapeLower(bool emit_err_ctx) {
return CreateModulePass(pass_func, 0, "VMShapeLower", {});
}
-TVM_REGISTER_GLOBAL("relax.transform.VMShapeLower").set_body_typed([](bool emit_err_ctx) {
+TVM_FFI_REGISTER_GLOBAL("relax.transform.VMShapeLower").set_body_typed([](bool emit_err_ctx) {
return VMShapeLower(emit_err_ctx);
});
diff --git a/src/relax/distributed/global_info.cc b/src/relax/distributed/global_info.cc
index c0d6204c87..e1cc32fc82 100644
--- a/src/relax/distributed/global_info.cc
+++ b/src/relax/distributed/global_info.cc
@@ -57,7 +57,7 @@ DeviceMesh::DeviceMesh(ffi::Shape shape, Range device_range) {
}
TVM_REGISTER_NODE_TYPE(DeviceMeshNode);
-TVM_REGISTER_GLOBAL("relax.distributed.DeviceMesh")
+TVM_FFI_REGISTER_GLOBAL("relax.distributed.DeviceMesh")
.set_body_typed([](ffi::Shape shape, Array<Integer> device_ids, Optional<Range> device_range) {
if (device_range.defined())
return DeviceMesh(shape, device_range.value());
diff --git a/src/relax/distributed/struct_info.cc b/src/relax/distributed/struct_info.cc
index 3569b15385..0ff9d4d6fa 100644
--- a/src/relax/distributed/struct_info.cc
+++ b/src/relax/distributed/struct_info.cc
@@ -43,11 +43,11 @@ PlacementSpec PlacementSpec::Replica() {
TVM_REGISTER_NODE_TYPE(PlacementSpecNode);
-TVM_REGISTER_GLOBAL("relax.distributed.Sharding").set_body_typed([](int axis) {
+TVM_FFI_REGISTER_GLOBAL("relax.distributed.Sharding").set_body_typed([](int axis) {
return PlacementSpec::Sharding(axis);
});
-TVM_REGISTER_GLOBAL("relax.distributed.Replica").set_body_typed([]() {
+TVM_FFI_REGISTER_GLOBAL("relax.distributed.Replica").set_body_typed([]() {
return PlacementSpec::Replica();
});
@@ -106,8 +106,8 @@ Placement Placement::FromText(String text_repr) {
}
TVM_REGISTER_NODE_TYPE(PlacementNode);
-TVM_REGISTER_GLOBAL("relax.distributed.PlacementFromText").set_body_typed(Placement::FromText);
-TVM_REGISTER_GLOBAL("relax.distributed.Placement")
+TVM_FFI_REGISTER_GLOBAL("relax.distributed.PlacementFromText").set_body_typed(Placement::FromText);
+TVM_FFI_REGISTER_GLOBAL("relax.distributed.Placement")
.set_body_typed([](Array<PlacementSpec> dim_specs) { return Placement(dim_specs); });
// DTensor
@@ -130,7 +130,7 @@ DTensorStructInfo::DTensorStructInfo(TensorStructInfo tensor_sinfo, DeviceMesh d
TVM_REGISTER_NODE_TYPE(DTensorStructInfoNode);
-TVM_REGISTER_GLOBAL("relax.distributed.DTensorStructInfo")
+TVM_FFI_REGISTER_GLOBAL("relax.distributed.DTensorStructInfo")
.set_body_typed([](TensorStructInfo tensor_sinfo, DeviceMesh device_mesh, Placement placement,
Span span) {
return DTensorStructInfo(tensor_sinfo, device_mesh, placement, span);
diff --git a/src/relax/distributed/transform/legalize_redistribute.cc b/src/relax/distributed/transform/legalize_redistribute.cc
index 5ed9478587..1df1d2110b 100644
--- a/src/relax/distributed/transform/legalize_redistribute.cc
+++ b/src/relax/distributed/transform/legalize_redistribute.cc
@@ -115,7 +115,7 @@ Pass LegalizeRedistribute() {
};
return CreateModulePass(pass_func, 1, "LegalizeRedistribute", {});
}
-TVM_REGISTER_GLOBAL("relax.distributed.transform.LegalizeRedistribute")
+TVM_FFI_REGISTER_GLOBAL("relax.distributed.transform.LegalizeRedistribute")
.set_body_typed(LegalizeRedistribute);
} // namespace transform
diff --git a/src/relax/distributed/transform/lower_distir.cc b/src/relax/distributed/transform/lower_distir.cc
index 59de65450e..e4f811b83d 100644
--- a/src/relax/distributed/transform/lower_distir.cc
+++ b/src/relax/distributed/transform/lower_distir.cc
@@ -262,7 +262,7 @@ Pass LowerDistIR() {
auto pass_func = [=](IRModule m, PassContext pc) { return DistIRSharder::LowerDistIR(m); };
return CreateModulePass(pass_func, 1, "LowerDistIR", {});
}
-TVM_REGISTER_GLOBAL("relax.distributed.transform.LowerDistIR").set_body_typed(LowerDistIR);
+TVM_FFI_REGISTER_GLOBAL("relax.distributed.transform.LowerDistIR").set_body_typed(LowerDistIR);
} // namespace transform
} // namespace distributed
diff --git a/src/relax/distributed/transform/lower_global_view_to_local_view.cc b/src/relax/distributed/transform/lower_global_view_to_local_view.cc
index 514a98ef44..c8abe2b1d1 100644
--- a/src/relax/distributed/transform/lower_global_view_to_local_view.cc
+++ b/src/relax/distributed/transform/lower_global_view_to_local_view.cc
@@ -432,7 +432,7 @@ Pass LowerGlobalViewToLocalView() {
auto pass_func = [=](IRModule m, PassContext pc) { return LowerTIRToLocalView(m).Lower(); };
return CreateModulePass(pass_func, 1, "LowerGlobalViewToLocalView", {});
}
-TVM_REGISTER_GLOBAL("relax.distributed.transform.LowerGlobalViewToLocalView")
+TVM_FFI_REGISTER_GLOBAL("relax.distributed.transform.LowerGlobalViewToLocalView")
.set_body_typed(LowerGlobalViewToLocalView);
} // namespace transform
diff --git a/src/relax/distributed/transform/propagate_sharding.cc b/src/relax/distributed/transform/propagate_sharding.cc
index 15b372fb83..f5f276c2b8 100644
--- a/src/relax/distributed/transform/propagate_sharding.cc
+++ b/src/relax/distributed/transform/propagate_sharding.cc
@@ -615,7 +615,7 @@ Pass PropagateSharding() {
};
return CreateModulePass(pass_func, 1, "PropagateSharding", {});
}
-TVM_REGISTER_GLOBAL("relax.distributed.transform.PropagateSharding")
+TVM_FFI_REGISTER_GLOBAL("relax.distributed.transform.PropagateSharding")
.set_body_typed(PropagateSharding);
} // namespace transform
diff --git a/src/relax/ir/binding_rewrite.cc b/src/relax/ir/binding_rewrite.cc
index 8a70f18324..f35b443b5b 100644
--- a/src/relax/ir/binding_rewrite.cc
+++ b/src/relax/ir/binding_rewrite.cc
@@ -51,7 +51,7 @@ DataflowBlockRewrite::DataflowBlockRewrite(DataflowBlock dfb, Function root_fn)
data_ = std::move(n);
}
-TVM_REGISTER_GLOBAL("relax.DataflowBlockRewrite")
+TVM_FFI_REGISTER_GLOBAL("relax.DataflowBlockRewrite")
.set_body_typed([](DataflowBlock dfb, Function root_fn) {
return DataflowBlockRewrite(dfb, root_fn);
});
@@ -110,7 +110,7 @@ void DataflowBlockRewriteNode::ReplaceAllUses(Var old_var, Var new_var) {
}
}
-TVM_REGISTER_GLOBAL("relax.dfb_rewrite_replace_all_uses")
+TVM_FFI_REGISTER_GLOBAL("relax.dfb_rewrite_replace_all_uses")
.set_body_typed([](DataflowBlockRewrite rwt, Var old_var, Var new_var) {
rwt->ReplaceAllUses(old_var, new_var);
});
@@ -178,10 +178,10 @@ void DataflowBlockRewriteNode::Add(Binding binding) {
}
}
-TVM_REGISTER_GLOBAL("relax.dfb_rewrite_add_binding")
+TVM_FFI_REGISTER_GLOBAL("relax.dfb_rewrite_add_binding")
.set_body_typed([](DataflowBlockRewrite rwt, Binding vb) { rwt->Add(vb); });
-TVM_REGISTER_GLOBAL("relax.dfb_rewrite_add")
+TVM_FFI_REGISTER_GLOBAL("relax.dfb_rewrite_add")
.set_body_typed([](DataflowBlockRewrite rwt, Expr expr, Optional<String> name, bool is_dfvar) {
if (name.has_value()) {
rwt->Add(name.value(), expr, is_dfvar);
@@ -292,7 +292,7 @@ void DataflowBlockRewriteNode::RemoveUnused(Var unused, bool allow_undef) {
to_users_.erase(unused); // update use-def chain.
}
-TVM_REGISTER_GLOBAL("relax.dfb_rewrite_remove_unused")
+TVM_FFI_REGISTER_GLOBAL("relax.dfb_rewrite_remove_unused")
.set_body_typed([](DataflowBlockRewrite rwt, Var unused, bool allow_undef) {
rwt->RemoveUnused(unused, allow_undef);
});
@@ -314,7 +314,7 @@ void DataflowBlockRewriteNode::RemoveAllUnused() {
for (const auto& unused : remover.unused_vars) to_users_.erase(unused);
}
-TVM_REGISTER_GLOBAL("relax.dfb_rewrite_remove_all_unused")
+TVM_FFI_REGISTER_GLOBAL("relax.dfb_rewrite_remove_all_unused")
.set_body_typed([](DataflowBlockRewrite rwt) { rwt->RemoveAllUnused(); });
Expr RemoveAllUnused(Expr expr) {
@@ -333,7 +333,7 @@ Expr RemoveAllUnused(Expr expr) {
return remover.VisitExpr(std::move(expr));
}
-TVM_REGISTER_GLOBAL("relax.analysis.remove_all_unused").set_body_typed(RemoveAllUnused);
+TVM_FFI_REGISTER_GLOBAL("relax.analysis.remove_all_unused").set_body_typed(RemoveAllUnused);
IRModule DataflowBlockRewriteNode::MutateIRModule(IRModule irmod) {
BlockBuilder builder = BlockBuilder::Create(irmod);
@@ -348,7 +348,7 @@ IRModule DataflowBlockRewriteNode::MutateIRModule(IRModule irmod) {
return builder->GetContextIRModule();
}
-TVM_REGISTER_GLOBAL("relax.dfb_rewrite_mutate_irmodule")
+TVM_FFI_REGISTER_GLOBAL("relax.dfb_rewrite_mutate_irmodule")
.set_body_typed([](DataflowBlockRewrite rwt, IRModule irmod) {
return rwt->MutateIRModule(irmod);
});
diff --git a/src/relax/ir/block_builder.cc b/src/relax/ir/block_builder.cc
index db173fe9d0..63288201e7 100644
--- a/src/relax/ir/block_builder.cc
+++ b/src/relax/ir/block_builder.cc
@@ -21,6 +21,7 @@
* \file src/relax/block_builder.cc
*/
#include <tvm/arith/analyzer.h>
+#include <tvm/ffi/function.h>
#include <tvm/relax/analysis.h>
#include <tvm/relax/block_builder.h>
#include <tvm/relax/expr_functor.h>
@@ -29,7 +30,6 @@
#include <tvm/relax/struct_info_functor.h>
#include <tvm/relax/transform.h>
#include <tvm/relax/type.h>
-#include <tvm/runtime/registry.h>
#include <tvm/tir/function.h>
#include <memory>
@@ -1054,65 +1054,67 @@ BlockBuilder BlockBuilder::Create(Optional<IRModule> mod,
//---------------------------------------
TVM_REGISTER_OBJECT_TYPE(BlockBuilderNode);
-TVM_REGISTER_GLOBAL("relax.BlockBuilderCreate").set_body_typed([](Optional<IRModule> mod) {
+TVM_FFI_REGISTER_GLOBAL("relax.BlockBuilderCreate").set_body_typed([](Optional<IRModule> mod) {
return BlockBuilder::Create(mod);
});
-TVM_REGISTER_GLOBAL("relax.BlockBuilderBeginDataflowBlock")
+TVM_FFI_REGISTER_GLOBAL("relax.BlockBuilderBeginDataflowBlock")
.set_body_method(&BlockBuilderNode::BeginDataflowBlock);
-TVM_REGISTER_GLOBAL("relax.BlockBuilderBeginBindingBlock")
+TVM_FFI_REGISTER_GLOBAL("relax.BlockBuilderBeginBindingBlock")
.set_body_method(&BlockBuilderNode::BeginBindingBlock);
-TVM_REGISTER_GLOBAL("relax.BlockBuilderEndBlock").set_body_method(&BlockBuilderNode::EndBlock);
+TVM_FFI_REGISTER_GLOBAL("relax.BlockBuilderEndBlock").set_body_method(&BlockBuilderNode::EndBlock);
-TVM_REGISTER_GLOBAL("relax.BlockBuilderNormalize").set_body_method(&BlockBuilderNode::Normalize);
+TVM_FFI_REGISTER_GLOBAL("relax.BlockBuilderNormalize")
+ .set_body_method(&BlockBuilderNode::Normalize);
-TVM_REGISTER_GLOBAL("relax.BlockBuilderEmit")
+TVM_FFI_REGISTER_GLOBAL("relax.BlockBuilderEmit")
.set_body_typed([](BlockBuilder builder, Expr expr, String name_hint) {
return builder->Emit(expr, name_hint);
});
-TVM_REGISTER_GLOBAL("relax.BlockBuilderEmitMatchCast")
+TVM_FFI_REGISTER_GLOBAL("relax.BlockBuilderEmitMatchCast")
.set_body_typed([](BlockBuilder builder, Expr value, StructInfo struct_info, String name_hint) {
return builder->EmitMatchCast(value, struct_info, name_hint);
});
-TVM_REGISTER_GLOBAL("relax.BlockBuilderEmitOutput")
+TVM_FFI_REGISTER_GLOBAL("relax.BlockBuilderEmitOutput")
.set_body_typed([](BlockBuilder builder, const Expr& output, String name_hint) {
return builder->EmitOutput(output, name_hint);
});
-TVM_REGISTER_GLOBAL("relax.BlockBuilderEmitNormalized")
+TVM_FFI_REGISTER_GLOBAL("relax.BlockBuilderEmitNormalized")
.set_body_typed([](BlockBuilder builder, Binding binding) {
return builder->EmitNormalized(binding);
});
-TVM_REGISTER_GLOBAL("relax.BlockBuilderGetUniqueName")
+TVM_FFI_REGISTER_GLOBAL("relax.BlockBuilderGetUniqueName")
.set_body_typed([](BlockBuilder builder, String name_hint) {
return builder->name_supply()->FreshName(name_hint, /*add_prefix*/ false,
/*add_underscore*/ false);
});
-TVM_REGISTER_GLOBAL("relax.BlockBuilderAddFunction")
+TVM_FFI_REGISTER_GLOBAL("relax.BlockBuilderAddFunction")
.set_body_method(&BlockBuilderNode::AddFunction);
-TVM_REGISTER_GLOBAL("relax.BlockBuilderUpdateFunction")
+TVM_FFI_REGISTER_GLOBAL("relax.BlockBuilderUpdateFunction")
.set_body_method(&BlockBuilderNode::UpdateFunction);
-TVM_REGISTER_GLOBAL("relax.BlockBuilderGetContextIRModule")
+TVM_FFI_REGISTER_GLOBAL("relax.BlockBuilderGetContextIRModule")
.set_body_method(&BlockBuilderNode::GetContextIRModule);
-TVM_REGISTER_GLOBAL("relax.BlockBuilderFinalize").set_body_method(&BlockBuilderNode::Finalize);
+TVM_FFI_REGISTER_GLOBAL("relax.BlockBuilderFinalize").set_body_method(&BlockBuilderNode::Finalize);
-TVM_REGISTER_GLOBAL("relax.BlockBuilderCurrentBlockIsDataFlow")
+TVM_FFI_REGISTER_GLOBAL("relax.BlockBuilderCurrentBlockIsDataFlow")
.set_body_method(&BlockBuilderNode::CurrentBlockIsDataFlow);
-TVM_REGISTER_GLOBAL("relax.BlockBuilderLookupBinding")
+TVM_FFI_REGISTER_GLOBAL("relax.BlockBuilderLookupBinding")
.set_body_method(&BlockBuilderNode::LookupBinding);
-TVM_REGISTER_GLOBAL("relax.BlockBuilderBeginScope").set_body_method(&BlockBuilderNode::BeginScope);
+TVM_FFI_REGISTER_GLOBAL("relax.BlockBuilderBeginScope")
+ .set_body_method(&BlockBuilderNode::BeginScope);
-TVM_REGISTER_GLOBAL("relax.BlockBuilderEndScope").set_body_method(&BlockBuilderNode::EndScope);
+TVM_FFI_REGISTER_GLOBAL("relax.BlockBuilderEndScope").set_body_method(&BlockBuilderNode::EndScope);
} // namespace relax
} // namespace tvm
diff --git a/src/relax/ir/dataflow_block_rewriter.cc b/src/relax/ir/dataflow_block_rewriter.cc
index b82269f228..172f4d7bcb 100644
--- a/src/relax/ir/dataflow_block_rewriter.cc
+++ b/src/relax/ir/dataflow_block_rewriter.cc
@@ -362,7 +362,7 @@ Optional<Map<DFPattern, Var>> MatchGraph(const PatternContext& ctx, const Datafl
return MatchGraph(ctx, dfb->bindings, AnalyzeVar2Value(dfb));
}
-TVM_REGISTER_GLOBAL("relax.dpl.match_dfb")
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.match_dfb")
.set_body_typed([](const PatternContext& ctx, const DataflowBlock& dfb) {
return MatchGraph(ctx, dfb);
});
@@ -447,7 +447,7 @@ Function RewriteBindings(
return Downcast<Function>(PatternContextRewriter(ctx, rewriter)(func));
}
-TVM_REGISTER_GLOBAL("relax.dpl.rewrite_bindings").set_body_typed(RewriteBindings);
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.rewrite_bindings").set_body_typed(RewriteBindings);
} // namespace relax
} // namespace tvm
diff --git a/src/relax/ir/dataflow_expr_rewriter.cc b/src/relax/ir/dataflow_expr_rewriter.cc
index 1c5aa4a0cc..c398305d93 100644
--- a/src/relax/ir/dataflow_expr_rewriter.cc
+++ b/src/relax/ir/dataflow_expr_rewriter.cc
@@ -193,17 +193,16 @@ void RewriteSpec::Append(RewriteSpec other) {
TVM_REGISTER_NODE_TYPE(PatternMatchingRewriterNode);
-TVM_REGISTER_GLOBAL("relax.dpl.PatternMatchingRewriterFromPattern")
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.PatternMatchingRewriterFromPattern")
.set_body_typed([](DFPattern pattern,
ffi::TypedFunction<Optional<Expr>(Expr, Map<DFPattern, Expr>)> func) {
return PatternMatchingRewriter::FromPattern(pattern, func);
});
-TVM_REGISTER_GLOBAL("relax.dpl.PatternMatchingRewriterFromModule").set_body_typed([](IRModule mod) {
- return PatternMatchingRewriter::FromModule(mod);
-});
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.PatternMatchingRewriterFromModule")
+ .set_body_typed([](IRModule mod) { return PatternMatchingRewriter::FromModule(mod); });
-TVM_REGISTER_GLOBAL("relax.dpl.PatternMatchingRewriterApply")
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.PatternMatchingRewriterApply")
.set_body_typed([](PatternMatchingRewriter rewriter,
Variant<Expr, IRModule> obj) -> Variant<Expr, IRModule> {
if (auto expr = obj.as<Expr>()) {
@@ -259,7 +258,7 @@ Optional<Expr> ExprPatternRewriterNode::RewriteExpr(const Expr& expr,
return std::nullopt;
}
-TVM_REGISTER_GLOBAL("relax.dpl.PatternRewriter")
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.PatternRewriter")
.set_body_typed([](DFPattern pattern,
ffi::TypedFunction<Optional<Expr>(Expr, Map<DFPattern, Expr>)> func) {
return ExprPatternRewriter(pattern, func);
@@ -308,7 +307,7 @@ RewriteSpec OrRewriterNode::RewriteBindings(const Array<Binding>& bindings) cons
return lhs_match;
}
-TVM_REGISTER_GLOBAL("relax.dpl.OrRewriter")
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.OrRewriter")
.set_body_typed([](PatternMatchingRewriter lhs, PatternMatchingRewriter rhs) {
return OrRewriter(lhs, rhs);
});
@@ -603,7 +602,7 @@ std::optional<std::vector<Expr>> TupleRewriterNode::TryMatchByBindingIndex(
return rewrites;
}
-TVM_REGISTER_GLOBAL("relax.dpl.TupleRewriter")
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.TupleRewriter")
.set_body_typed([](Array<DFPattern> patterns,
ffi::TypedFunction<Optional<Expr>(Expr, Map<DFPattern, Expr>)> func) {
return TupleRewriter(patterns, func);
@@ -796,13 +795,13 @@ Optional<Map<DFPattern, Expr>> ExtractMatchedExpr(DFPattern pattern, Expr expr,
return matcher.GetMemo();
}
-TVM_REGISTER_GLOBAL("relax.dpl.extract_matched_expr").set_body_typed(ExtractMatchedExpr);
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.extract_matched_expr").set_body_typed(ExtractMatchedExpr);
bool MatchExpr(DFPattern pattern, Expr expr, Optional<Map<Var, Expr>> bindings_opt) {
return static_cast<bool>(ExtractMatchedExpr(pattern, expr, bindings_opt));
}
-TVM_REGISTER_GLOBAL("relax.dpl.match_expr").set_body_typed(MatchExpr);
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.match_expr").set_body_typed(MatchExpr);
/*!
* \brief Apply pattern matching to each expression, replacing
@@ -1074,7 +1073,7 @@ Function RewriteCall(const DFPattern& pat,
return Downcast<Function>(PatternMatchingRewriter::FromPattern(pat, rewriter)(func));
}
-TVM_REGISTER_GLOBAL("relax.dpl.rewrite_call").set_body_typed(RewriteCall);
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.rewrite_call").set_body_typed(RewriteCall);
} // namespace relax
} // namespace tvm
diff --git a/src/relax/ir/dataflow_pattern.cc b/src/relax/ir/dataflow_pattern.cc
index 7f045ebe48..db242b773b 100644
--- a/src/relax/ir/dataflow_pattern.cc
+++ b/src/relax/ir/dataflow_pattern.cc
@@ -44,7 +44,7 @@ ExternFuncPattern::ExternFuncPattern(String global_symbol) {
n->global_symbol_ = std::move(global_symbol);
data_ = std::move(n);
}
-TVM_REGISTER_GLOBAL("relax.dpl.ExternFuncPattern").set_body_typed([](String global_symbol) {
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.ExternFuncPattern").set_body_typed([](String global_symbol) {
return ExternFuncPattern(global_symbol);
});
RELAX_PATTERN_PRINTER_DEF(ExternFuncPatternNode, [](auto p, auto node) {
@@ -57,7 +57,7 @@ VarPattern::VarPattern(String name_hint) {
n->name = std::move(name_hint);
data_ = std::move(n);
}
-TVM_REGISTER_GLOBAL("relax.dpl.VarPattern").set_body_typed([](String name_hint) {
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.VarPattern").set_body_typed([](String name_hint) {
return VarPattern(name_hint);
});
RELAX_PATTERN_PRINTER_DEF(VarPatternNode, [](auto p, auto node) {
@@ -65,7 +65,7 @@ RELAX_PATTERN_PRINTER_DEF(VarPatternNode, [](auto p, auto node) {
});
TVM_REGISTER_NODE_TYPE(DataflowVarPatternNode);
-TVM_REGISTER_GLOBAL("relax.dpl.DataflowVarPattern").set_body_typed([](String name_hint) {
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.DataflowVarPattern").set_body_typed([](String name_hint) {
return DataflowVarPattern(name_hint);
});
DataflowVarPattern::DataflowVarPattern(String name_hint) {
@@ -83,7 +83,7 @@ GlobalVarPattern::GlobalVarPattern(String name_hint) {
n->name = std::move(name_hint);
data_ = std::move(n);
}
-TVM_REGISTER_GLOBAL("relax.dpl.GlobalVarPattern").set_body_typed([](String name_hint) {
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.GlobalVarPattern").set_body_typed([](String name_hint) {
return GlobalVarPattern(name_hint);
});
RELAX_PATTERN_PRINTER_DEF(GlobalVarPatternNode, [](auto p, auto node) {
@@ -96,11 +96,13 @@ ExprPattern::ExprPattern(Expr expr) {
n->expr = std::move(expr);
data_ = std::move(n);
}
-TVM_REGISTER_GLOBAL("relax.dpl.ExprPattern").set_body_typed([](Expr e) { return ExprPattern(e); });
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.ExprPattern").set_body_typed([](Expr e) {
+ return ExprPattern(e);
+});
RELAX_PATTERN_PRINTER_DEF(ExprPatternNode, [](auto p, auto node) { p->Print(node->expr); });
TVM_REGISTER_NODE_TYPE(ConstantPatternNode);
-TVM_REGISTER_GLOBAL("relax.dpl.ConstantPattern").set_body_typed([]() {
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.ConstantPattern").set_body_typed([]() {
auto c = ConstantPattern(make_object<ConstantPatternNode>());
return c;
});
@@ -115,7 +117,7 @@ CallPattern::CallPattern(DFPattern op, Array<DFPattern> args, bool varg_default_
n->varg_default_wildcard = varg_default_wildcard;
data_ = std::move(n);
}
-TVM_REGISTER_GLOBAL("relax.dpl.CallPattern")
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.CallPattern")
.set_body_typed([](DFPattern op, Array<DFPattern> args, bool varg_default_wildcard) {
return CallPattern(op, args, varg_default_wildcard);
});
@@ -138,7 +140,7 @@ PrimArrPattern::PrimArrPattern(Array<PrimExpr> arr) {
n->fields = std::move(arr);
data_ = std::move(n);
}
-TVM_REGISTER_GLOBAL("relax.dpl.PrimArrPattern").set_body_typed([](Array<PrimExpr> arr) {
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.PrimArrPattern").set_body_typed([](Array<PrimExpr> arr) {
return PrimArrPattern(std::move(arr));
});
RELAX_PATTERN_PRINTER_DEF(PrimArrPatternNode, [](auto p, auto node) {
@@ -152,7 +154,7 @@ FunctionPattern::FunctionPattern(Array<DFPattern> params, DFPattern body) {
n->body = std::move(body);
data_ = std::move(n);
}
-TVM_REGISTER_GLOBAL("relax.dpl.FunctionPattern")
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.FunctionPattern")
.set_body_typed([](Array<DFPattern> params, DFPattern body) {
return FunctionPattern(params, body);
});
@@ -166,7 +168,7 @@ TuplePattern::TuplePattern(tvm::Array<DFPattern> fields) {
n->fields = std::move(fields);
data_ = std::move(n);
}
-TVM_REGISTER_GLOBAL("relax.dpl.TuplePattern").set_body_typed([](tvm::Array<DFPattern> fields) {
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.TuplePattern").set_body_typed([](tvm::Array<DFPattern> fields) {
return TuplePattern(fields);
});
RELAX_PATTERN_PRINTER_DEF(TuplePatternNode, [](auto p, auto node) {
@@ -179,7 +181,7 @@ UnorderedTuplePattern::UnorderedTuplePattern(tvm::Array<DFPattern> fields) {
n->fields = std::move(fields);
data_ = std::move(n);
}
-TVM_REGISTER_GLOBAL("relax.dpl.UnorderedTuplePattern")
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.UnorderedTuplePattern")
.set_body_typed([](tvm::Array<DFPattern> fields) { return UnorderedTuplePattern(fields); });
RELAX_PATTERN_PRINTER_DEF(UnorderedTuplePatternNode, [](auto p, auto node) {
p->stream << "UnorderedTuplePattern(" << node->fields << ")";
@@ -192,9 +194,8 @@ TupleGetItemPattern::TupleGetItemPattern(DFPattern tuple, int index) {
n->index = index;
data_ = std::move(n);
}
-TVM_REGISTER_GLOBAL("relax.dpl.TupleGetItemPattern").set_body_typed([](DFPattern tuple, int index) {
- return TupleGetItemPattern(tuple, index);
-});
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.TupleGetItemPattern")
+ .set_body_typed([](DFPattern tuple, int index) { return TupleGetItemPattern(tuple, index); });
RELAX_PATTERN_PRINTER_DEF(TupleGetItemPatternNode, [](auto p, auto node) {
p->stream << "TupleGetItemPattern(" << node->tuple << ", " << node->index << ")";
});
@@ -206,7 +207,7 @@ AndPattern::AndPattern(DFPattern left, DFPattern right) {
n->right = std::move(right);
data_ = std::move(n);
}
-TVM_REGISTER_GLOBAL("relax.dpl.AndPattern").set_body_typed([](DFPattern left, DFPattern right) {
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.AndPattern").set_body_typed([](DFPattern left, DFPattern right) {
return AndPattern(left, right);
});
RELAX_PATTERN_PRINTER_DEF(AndPatternNode, [](auto p, auto node) {
@@ -220,7 +221,7 @@ OrPattern::OrPattern(DFPattern left, DFPattern right) {
n->right = std::move(right);
data_ = std::move(n);
}
-TVM_REGISTER_GLOBAL("relax.dpl.OrPattern").set_body_typed([](DFPattern left, DFPattern right) {
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.OrPattern").set_body_typed([](DFPattern left, DFPattern right) {
return OrPattern(left, right);
});
RELAX_PATTERN_PRINTER_DEF(OrPatternNode, [](auto p, auto node) {
@@ -233,7 +234,7 @@ NotPattern::NotPattern(DFPattern reject) {
n->reject = std::move(reject);
data_ = std::move(n);
}
-TVM_REGISTER_GLOBAL("relax.dpl.NotPattern").set_body_typed([](DFPattern reject) {
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.NotPattern").set_body_typed([](DFPattern reject) {
return NotPattern(reject);
});
RELAX_PATTERN_PRINTER_DEF(NotPatternNode,
@@ -241,7 +242,9 @@ RELAX_PATTERN_PRINTER_DEF(NotPatternNode,
TVM_REGISTER_NODE_TYPE(WildcardPatternNode);
WildcardPattern::WildcardPattern() { data_ = make_object<WildcardPatternNode>(); }
-TVM_REGISTER_GLOBAL("relax.dpl.WildcardPattern").set_body_typed([]() { return WildcardPattern(); });
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.WildcardPattern").set_body_typed([]() {
+ return WildcardPattern();
+});
RELAX_PATTERN_PRINTER_DEF(WildcardPatternNode, [](auto p, auto node) { p->stream << "*"; });
TVM_REGISTER_NODE_TYPE(TypePatternNode);
@@ -251,7 +254,7 @@ TypePattern::TypePattern(DFPattern pattern, Type type) {
n->type = std::move(type);
data_ = std::move(n);
}
-TVM_REGISTER_GLOBAL("relax.dpl.TypePattern").set_body_typed([](DFPattern pattern, Type type) {
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.TypePattern").set_body_typed([](DFPattern pattern, Type type) {
return TypePattern(pattern, type);
});
RELAX_PATTERN_PRINTER_DEF(TypePatternNode, [](auto p, auto node) {
@@ -265,7 +268,7 @@ StructInfoPattern::StructInfoPattern(DFPattern pattern, StructInfo struct_info)
n->struct_info = std::move(struct_info);
data_ = std::move(n);
}
-TVM_REGISTER_GLOBAL("relax.dpl.StructInfoPattern")
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.StructInfoPattern")
.set_body_typed([](DFPattern pattern, StructInfo struct_info) {
return StructInfoPattern(pattern, struct_info);
});
@@ -281,7 +284,7 @@ ShapePattern::ShapePattern(DFPattern pattern, Array<PrimExpr> shape) {
n->shape = std::move(shape);
data_ = std::move(n);
}
-TVM_REGISTER_GLOBAL("relax.dpl.ShapePattern")
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.ShapePattern")
.set_body_typed([](DFPattern pattern, Array<PrimExpr> shape) {
return ShapePattern(pattern, shape);
});
@@ -299,7 +302,7 @@ SameShapeConstraint::SameShapeConstraint(Array<DFPattern> args) {
ctx.value().add_constraint(*this);
}
}
-TVM_REGISTER_GLOBAL("relax.dpl.SameShapeConstraint").set_body_typed([](Array<DFPattern> args) {
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.SameShapeConstraint").set_body_typed([](Array<DFPattern> args) {
return SameShapeConstraint(args);
});
RELAX_PATTERN_PRINTER_DEF(SameShapeConstraintNode, [](auto p, auto node) {
@@ -320,7 +323,7 @@ DataTypePattern::DataTypePattern(DFPattern pattern, DataType dtype) {
n->dtype = std::move(dtype);
data_ = std::move(n);
}
-TVM_REGISTER_GLOBAL("relax.dpl.DataTypePattern")
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.DataTypePattern")
.set_body_typed([](DFPattern pattern, DataType dtype) {
return DataTypePattern(pattern, dtype);
});
@@ -335,9 +338,8 @@ AttrPattern::AttrPattern(DFPattern pattern, DictAttrs attrs) {
n->attrs = std::move(attrs);
data_ = std::move(n);
}
-TVM_REGISTER_GLOBAL("relax.dpl.AttrPattern").set_body_typed([](DFPattern pattern, DictAttrs attrs) {
- return AttrPattern(pattern, attrs);
-});
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.AttrPattern")
+ .set_body_typed([](DFPattern pattern, DictAttrs attrs) { return AttrPattern(pattern, attrs); });
RELAX_PATTERN_PRINTER_DEF(AttrPatternNode, [](auto p, auto node) {
p->stream << "AttrPattern(" << node->pattern << " has attributes " << node->attrs << ")";
});
@@ -511,7 +513,7 @@ PatternSeq PatternSeq::dup() const {
return ret;
}
-TVM_REGISTER_GLOBAL("relax.dpl.PatternSeq")
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.PatternSeq")
.set_body_typed([](Array<DFPattern> patterns, bool only_used_by) {
return PatternSeq(std::move(patterns), only_used_by);
});
@@ -525,12 +527,12 @@ RELAX_PATTERN_PRINTER_DEF(PatternSeqNode, [](auto p, auto node) {
p->stream << "]";
});
-TVM_REGISTER_GLOBAL("relax.dpl.used_by")
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.used_by")
.set_body_typed([](PatternSeq lhs, PatternSeq rhs, int index) {
return lhs.UsedBy(rhs, index);
});
-TVM_REGISTER_GLOBAL("relax.dpl.only_used_by")
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.only_used_by")
.set_body_typed([](PatternSeq lhs, PatternSeq rhs, int index) {
return lhs.OnlyUsedBy(rhs, index);
});
@@ -643,25 +645,27 @@ DFPattern DFPattern::dup() const {
return pattern;
}
-TVM_REGISTER_GLOBAL("relax.dpl.dup_pattern").set_body_typed([](DFPattern pattern) {
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.dup_pattern").set_body_typed([](DFPattern pattern) {
return pattern.dup();
});
-TVM_REGISTER_GLOBAL("relax.dpl.dup_seq").set_body_typed([](PatternSeq seq) { return seq.dup(); });
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.dup_seq").set_body_typed([](PatternSeq seq) {
+ return seq.dup();
+});
-TVM_REGISTER_GLOBAL("relax.dpl.PatternContext").set_body_typed([](bool incre) {
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.PatternContext").set_body_typed([](bool incre) {
return PatternContext(incre);
});
-TVM_REGISTER_GLOBAL("relax.dpl.current_context").set_body_typed([] {
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.current_context").set_body_typed([] {
return PatternContext::Current();
});
-TVM_REGISTER_GLOBAL("relax.dpl.enter_context").set_body_typed([](const PatternContext& ctx) {
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.enter_context").set_body_typed([](const PatternContext& ctx) {
ctx.EnterWithScope();
});
-TVM_REGISTER_GLOBAL("relax.dpl.exit_context").set_body_typed([](const PatternContext& ctx) {
+TVM_FFI_REGISTER_GLOBAL("relax.dpl.exit_context").set_body_typed([](const PatternContext& ctx) {
ctx.ExitWithScope();
});
diff --git a/src/relax/ir/emit_te.cc b/src/relax/ir/emit_te.cc
index 86c38d08a1..e75dc3c2d7 100644
--- a/src/relax/ir/emit_te.cc
+++ b/src/relax/ir/emit_te.cc
@@ -72,7 +72,7 @@ te::Tensor TETensor(Expr value, Map<tir::Var, PrimExpr> tir_var_map, std::string
return te::PlaceholderOp(n).output(0);
}
-TVM_REGISTER_GLOBAL("relax.TETensor").set_body_typed(TETensor);
+TVM_FFI_REGISTER_GLOBAL("relax.TETensor").set_body_typed(TETensor);
} // namespace relax
} // namespace tvm
diff --git a/src/relax/ir/expr.cc b/src/relax/ir/expr.cc
index 4981e9b9f8..238cece41f 100644
--- a/src/relax/ir/expr.cc
+++ b/src/relax/ir/expr.cc
@@ -97,7 +97,7 @@ Call WithFields(Call call, Optional<Expr> opt_op, Optional<Array<Expr>> opt_args
TVM_REGISTER_NODE_TYPE(CallNode);
-TVM_REGISTER_GLOBAL("relax.Call")
+TVM_FFI_REGISTER_GLOBAL("relax.Call")
.set_body_typed([](Expr op, Array<Expr> args, Attrs attrs, Array<StructInfo> sinfo_args,
Span span) { return Call(op, args, attrs, sinfo_args, span); });
@@ -132,7 +132,7 @@ If WithFields(If if_expr, Optional<Expr> opt_cond, Optional<Expr> opt_true_branc
TVM_REGISTER_NODE_TYPE(IfNode);
-TVM_REGISTER_GLOBAL("relax.If")
+TVM_FFI_REGISTER_GLOBAL("relax.If")
.set_body_typed([](Expr cond, Expr true_branch, Expr false_branch, Span span) {
return If(cond, true_branch, false_branch, span);
});
@@ -162,7 +162,7 @@ Tuple::Tuple(tvm::Array<Expr> fields, Span span) {
TVM_REGISTER_NODE_TYPE(TupleNode);
-TVM_REGISTER_GLOBAL("relax.Tuple").set_body_typed([](tvm::Array<Expr> fields, Span span) {
+TVM_FFI_REGISTER_GLOBAL("relax.Tuple").set_body_typed([](tvm::Array<Expr> fields, Span span) {
return Tuple(fields, span);
});
@@ -226,7 +226,7 @@ TupleGetItem WithFields(TupleGetItem tuple_get_item, Optional<Expr> opt_tuple,
TVM_REGISTER_NODE_TYPE(TupleGetItemNode);
-TVM_REGISTER_GLOBAL("relax.TupleGetItem").set_body_typed([](Expr tuple, int index, Span span) {
+TVM_FFI_REGISTER_GLOBAL("relax.TupleGetItem").set_body_typed([](Expr tuple, int index, Span span) {
return TupleGetItem(tuple, index, span);
});
@@ -249,7 +249,7 @@ ShapeExpr::ShapeExpr(Array<PrimExpr> values, Span span) {
data_ = std::move(n);
}
-TVM_REGISTER_GLOBAL("relax.ShapeExpr").set_body_typed([](Array<PrimExpr> values, Span span) {
+TVM_FFI_REGISTER_GLOBAL("relax.ShapeExpr").set_body_typed([](Array<PrimExpr> values, Span span) {
return ShapeExpr(values, span);
});
@@ -285,12 +285,12 @@ VarNode* Var::CopyOnWrite() {
return static_cast<VarNode*>(data_.get());
}
-TVM_REGISTER_GLOBAL("relax.Var")
+TVM_FFI_REGISTER_GLOBAL("relax.Var")
.set_body_typed([](String name_hint, Optional<StructInfo> struct_info_annotation, Span span) {
return Var(name_hint, struct_info_annotation, span);
});
-TVM_REGISTER_GLOBAL("relax.VarFromId")
+TVM_FFI_REGISTER_GLOBAL("relax.VarFromId")
.set_body_typed([](Id vid, Optional<StructInfo> struct_info_annotation, Span span) {
return Var(vid, struct_info_annotation, span);
});
@@ -309,12 +309,12 @@ DataflowVar::DataflowVar(Id vid, Optional<StructInfo> struct_info_annotation, Sp
data_ = std::move(n);
}
-TVM_REGISTER_GLOBAL("relax.DataflowVar")
+TVM_FFI_REGISTER_GLOBAL("relax.DataflowVar")
.set_body_typed([](String name_hint, Optional<StructInfo> struct_info_annotation, Span span) {
return DataflowVar(name_hint, struct_info_annotation, span);
});
-TVM_REGISTER_GLOBAL("relax.DataflowVarFromId")
+TVM_FFI_REGISTER_GLOBAL("relax.DataflowVarFromId")
.set_body_typed([](Id vid, Optional<StructInfo> struct_info_annotation, Span span) {
return DataflowVar(vid, struct_info_annotation, span);
});
@@ -344,7 +344,7 @@ Constant::Constant(runtime::NDArray data, Optional<StructInfo> struct_info_annot
TVM_REGISTER_NODE_TYPE(ConstantNode);
-TVM_REGISTER_GLOBAL("relax.Constant")
+TVM_FFI_REGISTER_GLOBAL("relax.Constant")
.set_body_typed([](runtime::NDArray data,
Optional<StructInfo> struct_info_annotation = std::nullopt,
Span span = Span()) {
@@ -366,7 +366,7 @@ PrimValue PrimValue::Int64(int64_t value, Span span) {
TVM_REGISTER_NODE_TYPE(PrimValueNode);
-TVM_REGISTER_GLOBAL("relax.PrimValue").set_body_typed([](PrimExpr value, Span span) {
+TVM_FFI_REGISTER_GLOBAL("relax.PrimValue").set_body_typed([](PrimExpr value, Span span) {
return PrimValue(value, span);
});
@@ -383,7 +383,7 @@ StringImm::StringImm(String value, Span span) {
TVM_REGISTER_NODE_TYPE(StringImmNode);
-TVM_REGISTER_GLOBAL("relax.StringImm").set_body_typed([](String value, Span span) {
+TVM_FFI_REGISTER_GLOBAL("relax.StringImm").set_body_typed([](String value, Span span) {
return StringImm(value, span);
});
@@ -400,7 +400,7 @@ DataTypeImm::DataTypeImm(DataType value, Span span) {
TVM_REGISTER_NODE_TYPE(DataTypeImmNode);
-TVM_REGISTER_GLOBAL("relax.DataTypeImm").set_body_typed([](DataType value, Span span) {
+TVM_FFI_REGISTER_GLOBAL("relax.DataTypeImm").set_body_typed([](DataType value, Span span) {
return DataTypeImm(value, span);
});
@@ -416,7 +416,7 @@ MatchCast::MatchCast(Var var, Expr value, StructInfo struct_info, Span span) {
data_ = std::move(n);
}
-TVM_REGISTER_GLOBAL("relax.MatchCast")
+TVM_FFI_REGISTER_GLOBAL("relax.MatchCast")
.set_body_typed([](Var var, Expr value, StructInfo struct_info, Span span) {
return MatchCast(var, value, struct_info, span);
});
@@ -458,7 +458,7 @@ VarBinding::VarBinding(Var var, Expr value, Span span) {
data_ = std::move(n);
}
-TVM_REGISTER_GLOBAL("relax.VarBinding").set_body_typed([](Var var, Expr value, Span span) {
+TVM_FFI_REGISTER_GLOBAL("relax.VarBinding").set_body_typed([](Var var, Expr value, Span span) {
return VarBinding(var, value, span);
});
@@ -513,9 +513,10 @@ BindingBlockNode* BindingBlock::CopyOnWrite() {
return static_cast<BindingBlockNode*>(data_.get());
}
-TVM_REGISTER_GLOBAL("relax.BindingBlock").set_body_typed([](Array<Binding> bindings, Span span) {
- return BindingBlock(bindings, span);
-});
+TVM_FFI_REGISTER_GLOBAL("relax.BindingBlock")
+ .set_body_typed([](Array<Binding> bindings, Span span) {
+ return BindingBlock(bindings, span);
+ });
TVM_REGISTER_NODE_TYPE(DataflowBlockNode);
@@ -526,9 +527,10 @@ DataflowBlock::DataflowBlock(Array<Binding> bindings, Span span) {
data_ = std::move(n);
}
-TVM_REGISTER_GLOBAL("relax.DataflowBlock").set_body_typed([](Array<Binding> bindings, Span span) {
- return DataflowBlock(bindings, span);
-});
+TVM_FFI_REGISTER_GLOBAL("relax.DataflowBlock")
+ .set_body_typed([](Array<Binding> bindings, Span span) {
+ return DataflowBlock(bindings, span);
+ });
TVM_REGISTER_NODE_TYPE(SeqExprNode);
@@ -548,7 +550,7 @@ SeqExpr::SeqExpr(Array<BindingBlock> blocks, Expr body, Span span) {
data_ = std::move(n);
}
-TVM_REGISTER_GLOBAL("relax.SeqExpr")
+TVM_FFI_REGISTER_GLOBAL("relax.SeqExpr")
.set_body_typed([](Array<BindingBlock> blocks, Expr body, Span span) {
return SeqExpr(blocks, body, span);
});
@@ -624,7 +626,7 @@ Function::Function(Array<Var> params, Expr body, Optional<StructInfo> ret_struct
data_ = std::move(n);
}
-TVM_REGISTER_GLOBAL("relax.Function")
+TVM_FFI_REGISTER_GLOBAL("relax.Function")
.set_body_typed([](Array<Var> params, Expr body, Optional<StructInfo> ret_struct_info,
bool is_pure, DictAttrs attrs, Span span) {
return Function(params, body, ret_struct_info, is_pure, attrs, span);
@@ -662,7 +664,7 @@ Function Function::CreateEmpty(Array<Var> params, StructInfo ret_struct_info, bo
return Function(std::move(n));
}
-TVM_REGISTER_GLOBAL("relax.FunctionCreateEmpty")
+TVM_FFI_REGISTER_GLOBAL("relax.FunctionCreateEmpty")
.set_body_typed([](Array<Var> params, StructInfo ret_struct_info, bool is_pure, DictAttrs attrs,
Span span) {
return Function::CreateEmpty(params, ret_struct_info, is_pure, attrs, span);
@@ -670,7 +672,7 @@ TVM_REGISTER_GLOBAL("relax.FunctionCreateEmpty")
// Special opaque derivation function for ExternFunc
// Take look at sinfo_args to figure out the return StructInfo.
-TVM_REGISTER_GLOBAL("tvm.relax.struct_info.infer_by_sinfo_args")
+TVM_FFI_REGISTER_GLOBAL("tvm.relax.struct_info.infer_by_sinfo_args")
.set_body_typed([](const Call& call, const BlockBuilder& ctx) -> StructInfo {
ICHECK(call->sinfo_args.defined()) << "sinfo_args field of CallNode should always be defined";
if (call->sinfo_args.empty()) {
@@ -708,7 +710,7 @@ ExternFunc::ExternFunc(String global_symbol, StructInfo struct_info, Span span)
data_ = std::move(n);
}
-TVM_REGISTER_GLOBAL("relax.ExternFunc")
+TVM_FFI_REGISTER_GLOBAL("relax.ExternFunc")
.set_body_typed([](String global_symbol, Optional<StructInfo> struct_info, Span span) {
if (struct_info.defined()) {
return ExternFunc(global_symbol, struct_info.value(), span);
@@ -732,11 +734,11 @@ Expr GetShapeOf(const Expr& expr) {
return call_shape_of;
}
-TVM_REGISTER_GLOBAL("relax.GetShapeOf").set_body_typed([](const Expr& expr) {
+TVM_FFI_REGISTER_GLOBAL("relax.GetShapeOf").set_body_typed([](const Expr& expr) {
return GetShapeOf(expr);
});
-TVM_REGISTER_GLOBAL("relax.FuncWithAttr")
+TVM_FFI_REGISTER_GLOBAL("relax.FuncWithAttr")
.set_body_typed([](BaseFunc func, String key, ObjectRef value) -> Optional<Function> {
if (func->IsInstance<relax::FunctionNode>()) {
return WithAttr(Downcast<relax::Function>(std::move(func)), key, value);
@@ -744,7 +746,7 @@ TVM_REGISTER_GLOBAL("relax.FuncWithAttr")
return std::nullopt;
});
-TVM_REGISTER_GLOBAL("relax.FuncWithAttrs")
+TVM_FFI_REGISTER_GLOBAL("relax.FuncWithAttrs")
.set_body_typed([](BaseFunc func, Map<String, ffi::Any> attr_map) -> Optional<Function> {
if (func->IsInstance<relax::FunctionNode>()) {
return WithAttrs(Downcast<relax::Function>(std::move(func)), attr_map);
@@ -752,7 +754,7 @@ TVM_REGISTER_GLOBAL("relax.FuncWithAttrs")
return std::nullopt;
});
-TVM_REGISTER_GLOBAL("relax.FuncWithoutAttr")
+TVM_FFI_REGISTER_GLOBAL("relax.FuncWithoutAttr")
.set_body_typed([](BaseFunc func, String key) -> Optional<Function> {
if (func->IsInstance<relax::FunctionNode>()) {
return WithoutAttr(Downcast<relax::Function>(std::move(func)), key);
diff --git a/src/relax/ir/expr_functor.cc b/src/relax/ir/expr_functor.cc
index 0b497a97de..5e04453a12 100644
--- a/src/relax/ir/expr_functor.cc
+++ b/src/relax/ir/expr_functor.cc
@@ -326,7 +326,7 @@ void PostOrderVisit(const Expr& e, std::function<void(const Expr&)> fvisit) {
ExprApplyVisit(fvisit).VisitExpr(e);
}
-TVM_REGISTER_GLOBAL("relax.analysis.post_order_visit")
+TVM_FFI_REGISTER_GLOBAL("relax.analysis.post_order_visit")
.set_body_typed([](Expr expr, ffi::Function f) {
PostOrderVisit(expr, [f](const Expr& n) { f(n); });
});
diff --git a/src/relax/ir/py_expr_functor.cc b/src/relax/ir/py_expr_functor.cc
index 4a36bf2148..dc355cef90 100644
--- a/src/relax/ir/py_expr_functor.cc
+++ b/src/relax/ir/py_expr_functor.cc
@@ -540,30 +540,30 @@ class PyExprMutator : public ObjectRef {
TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(PyExprMutator, ObjectRef, PyExprMutatorNode);
};
-TVM_REGISTER_GLOBAL("relax.MakePyExprVisitor").set_body_typed(PyExprVisitor::MakePyExprVisitor);
+TVM_FFI_REGISTER_GLOBAL("relax.MakePyExprVisitor").set_body_typed(PyExprVisitor::MakePyExprVisitor);
-TVM_REGISTER_GLOBAL("relax.PyExprVisitorVisitExpr")
+TVM_FFI_REGISTER_GLOBAL("relax.PyExprVisitorVisitExpr")
.set_body_typed([](PyExprVisitor visitor, const Expr& expr) { visitor->VisitExpr(expr); });
-TVM_REGISTER_GLOBAL("relax.PyExprVisitorVisitBinding")
+TVM_FFI_REGISTER_GLOBAL("relax.PyExprVisitorVisitBinding")
.set_body_typed([](PyExprVisitor visitor, const Binding& binding) {
visitor->VisitBinding(binding);
});
-TVM_REGISTER_GLOBAL("relax.PyExprVisitorVisitBindingBlock")
+TVM_FFI_REGISTER_GLOBAL("relax.PyExprVisitorVisitBindingBlock")
.set_body_typed([](PyExprVisitor visitor, const BindingBlock& block) {
visitor->VisitBindingBlock(block);
});
-TVM_REGISTER_GLOBAL("relax.PyExprVisitorVisitVarDef")
+TVM_FFI_REGISTER_GLOBAL("relax.PyExprVisitorVisitVarDef")
.set_body_typed([](PyExprVisitor visitor, const Var& var) { visitor->VisitVarDef(var); });
-TVM_REGISTER_GLOBAL("relax.ExprVisitorVisitExpr")
+TVM_FFI_REGISTER_GLOBAL("relax.ExprVisitorVisitExpr")
.set_body_typed([](PyExprVisitor visitor, const Expr& expr) {
visitor->ExprVisitor::VisitExpr(expr);
});
-TVM_REGISTER_GLOBAL("relax.ExprVisitorVisitBinding")
+TVM_FFI_REGISTER_GLOBAL("relax.ExprVisitorVisitBinding")
.set_body_typed([](PyExprVisitor visitor, const Binding& binding) {
if (const auto* ptr = binding.as<VarBindingNode>()) {
visitor->ExprVisitor::VisitBinding_(ptr);
@@ -574,7 +574,7 @@ TVM_REGISTER_GLOBAL("relax.ExprVisitorVisitBinding")
}
});
-TVM_REGISTER_GLOBAL("relax.ExprVisitorVisitBindingBlock")
+TVM_FFI_REGISTER_GLOBAL("relax.ExprVisitorVisitBindingBlock")
.set_body_typed([](PyExprVisitor visitor, const BindingBlock& block) {
if (const auto* ptr = block.as<DataflowBlockNode>()) {
visitor->ExprVisitor::VisitBindingBlock_(ptr);
@@ -585,7 +585,7 @@ TVM_REGISTER_GLOBAL("relax.ExprVisitorVisitBindingBlock")
}
});
-TVM_REGISTER_GLOBAL("relax.ExprVisitorVisitVarDef")
+TVM_FFI_REGISTER_GLOBAL("relax.ExprVisitorVisitVarDef")
.set_body_typed([](PyExprVisitor visitor, const Var& var) {
if (const auto* node = var.as<DataflowVarNode>()) {
visitor->ExprVisitor::VisitVarDef_(node);
@@ -596,39 +596,39 @@ TVM_REGISTER_GLOBAL("relax.ExprVisitorVisitVarDef")
}
});
-TVM_REGISTER_GLOBAL("relax.ExprVisitorVisitSpan")
+TVM_FFI_REGISTER_GLOBAL("relax.ExprVisitorVisitSpan")
.set_body_typed([](PyExprVisitor visitor, const Span& span) {
visitor->ExprVisitor::VisitSpan(span);
});
-TVM_REGISTER_GLOBAL("relax.MakePyExprMutator").set_body_typed(PyExprMutator::MakePyExprMutator);
+TVM_FFI_REGISTER_GLOBAL("relax.MakePyExprMutator").set_body_typed(PyExprMutator::MakePyExprMutator);
-TVM_REGISTER_GLOBAL("relax.PyExprMutatorVisitExpr")
+TVM_FFI_REGISTER_GLOBAL("relax.PyExprMutatorVisitExpr")
.set_body_typed([](PyExprMutator mutator, const Expr& expr) {
return mutator->VisitExpr(expr);
});
-TVM_REGISTER_GLOBAL("relax.PyExprMutatorVisitBinding")
+TVM_FFI_REGISTER_GLOBAL("relax.PyExprMutatorVisitBinding")
.set_body_typed([](PyExprMutator mutator, const Binding& binding) {
mutator->VisitBinding(binding);
});
-TVM_REGISTER_GLOBAL("relax.PyExprMutatorVisitBindingBlock")
+TVM_FFI_REGISTER_GLOBAL("relax.PyExprMutatorVisitBindingBlock")
.set_body_typed([](PyExprMutator mutator, const BindingBlock& block) {
return mutator->VisitBindingBlock(block);
});
-TVM_REGISTER_GLOBAL("relax.PyExprMutatorVisitVarDef")
+TVM_FFI_REGISTER_GLOBAL("relax.PyExprMutatorVisitVarDef")
.set_body_typed([](PyExprMutator mutator, const Var& var) {
return mutator->VisitVarDef(var);
});
-TVM_REGISTER_GLOBAL("relax.ExprMutatorVisitExpr")
+TVM_FFI_REGISTER_GLOBAL("relax.ExprMutatorVisitExpr")
.set_body_typed([](PyExprMutator mutator, const Expr& expr) {
return mutator->ExprMutator::VisitExpr(expr);
});
-TVM_REGISTER_GLOBAL("relax.ExprMutatorVisitBinding")
+TVM_FFI_REGISTER_GLOBAL("relax.ExprMutatorVisitBinding")
.set_body_typed([](PyExprMutator mutator, const Binding& binding) {
if (const auto* ptr = binding.as<VarBindingNode>()) {
return mutator->ExprMutator::VisitBinding_(ptr);
@@ -639,7 +639,7 @@ TVM_REGISTER_GLOBAL("relax.ExprMutatorVisitBinding")
}
});
-TVM_REGISTER_GLOBAL("relax.ExprMutatorVisitBindingBlock")
+TVM_FFI_REGISTER_GLOBAL("relax.ExprMutatorVisitBindingBlock")
.set_body_typed([](PyExprMutator mutator, const BindingBlock& block) {
if (const auto* node = block.as<DataflowBlockNode>()) {
return mutator->ExprMutator::VisitBindingBlock_(node);
@@ -650,7 +650,7 @@ TVM_REGISTER_GLOBAL("relax.ExprMutatorVisitBindingBlock")
}
});
-TVM_REGISTER_GLOBAL("relax.ExprMutatorVisitVarDef")
+TVM_FFI_REGISTER_GLOBAL("relax.ExprMutatorVisitVarDef")
.set_body_typed([](PyExprMutator mutator, const Var& var) {
if (const auto* node = var.as<DataflowVarNode>()) {
return mutator->ExprMutator::VisitVarDef_(node);
@@ -661,32 +661,32 @@ TVM_REGISTER_GLOBAL("relax.ExprMutatorVisitVarDef")
}
});
-TVM_REGISTER_GLOBAL("relax.PyExprMutatorVisitExprPostOrder")
+TVM_FFI_REGISTER_GLOBAL("relax.PyExprMutatorVisitExprPostOrder")
.set_body_typed([](PyExprMutator mutator, const Expr& expr) {
return mutator->VisitExprPostOrder(expr);
});
-TVM_REGISTER_GLOBAL("relax.PyExprMutatorVisitWithNewScope")
+TVM_FFI_REGISTER_GLOBAL("relax.PyExprMutatorVisitWithNewScope")
.set_body_typed([](PyExprMutator mutator, const Expr& expr) {
return mutator->VisitWithNewScope(expr);
});
-TVM_REGISTER_GLOBAL("relax.PyExprMutatorLookupBinding")
+TVM_FFI_REGISTER_GLOBAL("relax.PyExprMutatorLookupBinding")
.set_body_typed([](PyExprMutator mutator, const Var& var) {
return mutator->LookupBinding(var);
});
-TVM_REGISTER_GLOBAL("relax.PyExprMutatorWithStructInfo")
+TVM_FFI_REGISTER_GLOBAL("relax.PyExprMutatorWithStructInfo")
.set_body_typed([](PyExprMutator mutator, Var var, StructInfo sinfo) {
return mutator->WithStructInfo(var, sinfo);
});
-TVM_REGISTER_GLOBAL("relax.PyExprMutatorSetVarRemap")
+TVM_FFI_REGISTER_GLOBAL("relax.PyExprMutatorSetVarRemap")
.set_body_typed([](PyExprMutator mutator, Id id, Var var) {
return mutator->var_remap_[id] = var;
});
-TVM_REGISTER_GLOBAL("relax.PyExprMutatorGetVarRemap")
+TVM_FFI_REGISTER_GLOBAL("relax.PyExprMutatorGetVarRemap")
.set_body_typed([](PyExprMutator mutator, Id id) { return mutator->var_remap_[id]; });
} // namespace relax
diff --git a/src/relax/ir/struct_info.cc b/src/relax/ir/struct_info.cc
index 38e54d4794..feb1f910a4 100644
--- a/src/relax/ir/struct_info.cc
+++ b/src/relax/ir/struct_info.cc
@@ -21,10 +21,10 @@
* \file src/relax/ir/struct_info.cc
* \brief Relax struct info.
*/
+#include <tvm/ffi/function.h>
#include <tvm/relax/analysis.h>
#include <tvm/relax/struct_info.h>
#include <tvm/relax/struct_info_functor.h>
-#include <tvm/runtime/registry.h>
namespace tvm {
namespace relax {
@@ -37,7 +37,7 @@ ObjectStructInfo::ObjectStructInfo(Span span) {
TVM_REGISTER_NODE_TYPE(ObjectStructInfoNode);
-TVM_REGISTER_GLOBAL("relax.ObjectStructInfo").set_body_typed([](Span span) {
+TVM_FFI_REGISTER_GLOBAL("relax.ObjectStructInfo").set_body_typed([](Span span) {
return ObjectStructInfo(span);
});
@@ -60,13 +60,11 @@ PrimStructInfo::PrimStructInfo(DataType dtype, Span span) {
TVM_REGISTER_NODE_TYPE(PrimStructInfoNode);
-TVM_REGISTER_GLOBAL("relax.PrimStructInfoFromDtype").set_body_typed([](DataType dtype, Span span) {
- return PrimStructInfo(dtype, span);
-});
+TVM_FFI_REGISTER_GLOBAL("relax.PrimStructInfoFromDtype")
+ .set_body_typed([](DataType dtype, Span span) { return PrimStructInfo(dtype, span); });
-TVM_REGISTER_GLOBAL("relax.PrimStructInfoFromValue").set_body_typed([](PrimExpr value, Span span) {
- return PrimStructInfo(value, span);
-});
+TVM_FFI_REGISTER_GLOBAL("relax.PrimStructInfoFromValue")
+ .set_body_typed([](PrimExpr value, Span span) { return PrimStructInfo(value, span); });
// Shape
ShapeStructInfo::ShapeStructInfo(Array<PrimExpr> values, Span span) {
@@ -94,7 +92,7 @@ ShapeStructInfo::ShapeStructInfo(int ndim, Span span) {
TVM_REGISTER_NODE_TYPE(ShapeStructInfoNode);
-TVM_REGISTER_GLOBAL("relax.ShapeStructInfo")
+TVM_FFI_REGISTER_GLOBAL("relax.ShapeStructInfo")
.set_body_typed([](Optional<Array<PrimExpr>> values, int ndim, Span span) {
if (values.defined()) {
CHECK_EQ(ndim, kUnknownNDim) << "ValueError: Cannot both specify values and ndim";
@@ -135,7 +133,7 @@ TensorStructInfo::TensorStructInfo(DataType dtype, int ndim, Optional<VDevice> v
TVM_REGISTER_NODE_TYPE(TensorStructInfoNode);
-TVM_REGISTER_GLOBAL("relax.TensorStructInfo")
+TVM_FFI_REGISTER_GLOBAL("relax.TensorStructInfo")
.set_body_typed([](Optional<Expr> shape, Optional<DataType> dtype, int ndim, VDevice vdevice,
Span span) {
if (shape.defined()) {
@@ -156,7 +154,7 @@ TupleStructInfo::TupleStructInfo(Array<StructInfo> fields, Span span) {
TVM_REGISTER_NODE_TYPE(TupleStructInfoNode);
-TVM_REGISTER_GLOBAL("relax.TupleStructInfo")
+TVM_FFI_REGISTER_GLOBAL("relax.TupleStructInfo")
.set_body_typed([](Array<StructInfo> fields, Span span) {
return TupleStructInfo(fields, span);
});
@@ -191,12 +189,12 @@ FuncStructInfo FuncStructInfo::OpaqueFunc(StructInfo ret, bool purity, Span span
TVM_REGISTER_NODE_TYPE(FuncStructInfoNode);
-TVM_REGISTER_GLOBAL("relax.FuncStructInfo")
+TVM_FFI_REGISTER_GLOBAL("relax.FuncStructInfo")
.set_body_typed([](Array<StructInfo> params, StructInfo ret, bool purity, Span span) {
return FuncStructInfo(params, ret, purity, span);
});
-TVM_REGISTER_GLOBAL("relax.FuncStructInfoOpaqueFunc")
+TVM_FFI_REGISTER_GLOBAL("relax.FuncStructInfoOpaqueFunc")
.set_body_typed([](Optional<StructInfo> ret, Optional<StructInfoDeriveFunc> derive_func,
bool purity, Span span) {
if (derive_func.defined()) {
@@ -220,11 +218,10 @@ void UpdateStructInfo(Expr expr, StructInfo struct_info) {
expr->checked_type_ = GetStaticType(struct_info);
}
-TVM_REGISTER_GLOBAL("relax.UpdateStructInfo").set_body_typed([](Expr expr, StructInfo struct_info) {
- UpdateStructInfo(expr, struct_info);
-});
+TVM_FFI_REGISTER_GLOBAL("relax.UpdateStructInfo")
+ .set_body_typed([](Expr expr, StructInfo struct_info) { UpdateStructInfo(expr, struct_info); });
-TVM_REGISTER_GLOBAL("ir.ExprStructInfo").set_body_typed([](Expr expr) {
+TVM_FFI_REGISTER_GLOBAL("ir.ExprStructInfo").set_body_typed([](Expr expr) {
return GetStructInfo(expr);
});
diff --git a/src/relax/ir/transform.cc b/src/relax/ir/transform.cc
index d79d8b3fd5..a44deba0fe 100644
--- a/src/relax/ir/transform.cc
+++ b/src/relax/ir/transform.cc
@@ -22,13 +22,13 @@
* \brief Relax specific transformation passes.
*/
#include <dmlc/thread_local.h>
+#include <tvm/ffi/function.h>
#include <tvm/ffi/rvalue_ref.h>
#include <tvm/node/repr_printer.h>
#include <tvm/relax/analysis.h>
#include <tvm/relax/expr_functor.h>
#include <tvm/relax/struct_info_functor.h>
#include <tvm/relax/transform.h>
-#include <tvm/runtime/registry.h>
namespace tvm {
namespace relax {
@@ -163,7 +163,7 @@ Pass CreateFunctionPass(std::function<Function(Function, IRModule, PassContext)>
TVM_REGISTER_NODE_TYPE(FunctionPassNode);
-TVM_REGISTER_GLOBAL("relax.transform.MakeFunctionPass")
+TVM_FFI_REGISTER_GLOBAL("relax.transform.MakeFunctionPass")
.set_body_typed(
[](ffi::TypedFunction<Function(ffi::RValueRef<Function>, IRModule, PassContext)> pass_func,
PassInfo pass_info) {
@@ -383,7 +383,7 @@ Pass CreateDataflowBlockPass(
TVM_REGISTER_NODE_TYPE(DataflowBlockPassNode);
-TVM_REGISTER_GLOBAL("relax.transform.MakeDataflowBlockPass")
+TVM_FFI_REGISTER_GLOBAL("relax.transform.MakeDataflowBlockPass")
.set_body_typed(
[](ffi::TypedFunction<DataflowBlock(ffi::RValueRef<DataflowBlock>, IRModule, PassContext)>
pass_func,
diff --git a/src/relax/ir/type.cc b/src/relax/ir/type.cc
index 82b95b556b..8b70bcf2c7 100644
--- a/src/relax/ir/type.cc
+++ b/src/relax/ir/type.cc
@@ -21,8 +21,8 @@
* \file src/relax/ir/type.cc
* \brief Relax type system.
*/
+#include <tvm/ffi/function.h>
#include <tvm/relax/type.h>
-#include <tvm/runtime/registry.h>
namespace tvm {
namespace relax {
@@ -36,7 +36,7 @@ ShapeType::ShapeType(int ndim, Span span) {
data_ = std::move(n);
}
-TVM_REGISTER_GLOBAL("relax.ShapeType").set_body_typed([](int ndim, Span span) {
+TVM_FFI_REGISTER_GLOBAL("relax.ShapeType").set_body_typed([](int ndim, Span span) {
return ShapeType(ndim, span);
});
@@ -48,7 +48,9 @@ ObjectType::ObjectType(Span span) {
TVM_REGISTER_NODE_TYPE(ObjectTypeNode);
-TVM_REGISTER_GLOBAL("relax.ObjectType").set_body_typed([](Span span) { return ObjectType(span); });
+TVM_FFI_REGISTER_GLOBAL("relax.ObjectType").set_body_typed([](Span span) {
+ return ObjectType(span);
+});
TensorType::TensorType(int ndim, DataType dtype, Span span) {
ObjectPtr<TensorTypeNode> n = make_object<TensorTypeNode>();
@@ -68,7 +70,7 @@ TensorType TensorType::CreateUnknownNDim(DataType dtype, Span span) {
TVM_REGISTER_NODE_TYPE(TensorTypeNode);
-TVM_REGISTER_GLOBAL("relax.TensorType").set_body_typed([](int ndim, DataType dtype, Span span) {
+TVM_FFI_REGISTER_GLOBAL("relax.TensorType").set_body_typed([](int ndim, DataType dtype, Span span) {
return TensorType(ndim, dtype, span);
});
@@ -80,7 +82,7 @@ PackedFuncType::PackedFuncType(Span span) {
TVM_REGISTER_NODE_TYPE(PackedFuncTypeNode);
-TVM_REGISTER_GLOBAL("relax.PackedFuncType").set_body_typed([](Span span) {
+TVM_FFI_REGISTER_GLOBAL("relax.PackedFuncType").set_body_typed([](Span span) {
return PackedFuncType(span);
});
diff --git a/src/relax/op/ccl/ccl.cc b/src/relax/op/ccl/ccl.cc
index c32cdc3aac..2f6314221e 100644
--- a/src/relax/op/ccl/ccl.cc
+++ b/src/relax/op/ccl/ccl.cc
@@ -36,7 +36,7 @@ Expr allreduce(Expr x, String op_type, bool in_group) {
return Call(op, {std::move(x)}, Attrs{attrs}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.ccl.allreduce").set_body_typed(allreduce);
+TVM_FFI_REGISTER_GLOBAL("relax.op.ccl.allreduce").set_body_typed(allreduce);
StructInfo InferStructInfoAllReduce(const Call& call, const BlockBuilder& ctx) {
TensorStructInfo input_sinfo = GetUnaryInputTensorStructInfo(call, ctx);
@@ -63,7 +63,7 @@ Expr allgather(Expr x, int num_workers, bool in_group) {
return Call(op, {std::move(x)}, Attrs{attrs}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.ccl.allgather").set_body_typed(allgather);
+TVM_FFI_REGISTER_GLOBAL("relax.op.ccl.allgather").set_body_typed(allgather);
StructInfo InferStructInfoAllGather(const Call& call, const BlockBuilder& ctx) {
TensorStructInfo input_sinfo = GetUnaryInputTensorStructInfo(call, ctx);
@@ -94,7 +94,8 @@ Expr broadcast_from_worker0(Expr x) {
return Call(op, {std::move(x)}, {}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.ccl.broadcast_from_worker0").set_body_typed(broadcast_from_worker0);
+TVM_FFI_REGISTER_GLOBAL("relax.op.ccl.broadcast_from_worker0")
+ .set_body_typed(broadcast_from_worker0);
StructInfo InferStructInfoBroadcastFromZero(const Call& call, const BlockBuilder& ctx) {
TensorStructInfo input_sinfo = GetUnaryInputTensorStructInfo(call, ctx);
@@ -120,7 +121,7 @@ Expr scatter_from_worker0(Expr data, int num_workers, int axis) {
return Call(op, {std::move(data)}, Attrs{attrs}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.ccl.scatter_from_worker0").set_body_typed(scatter_from_worker0);
+TVM_FFI_REGISTER_GLOBAL("relax.op.ccl.scatter_from_worker0").set_body_typed(scatter_from_worker0);
StructInfo InferStructInfoScatter(const Call& call, const BlockBuilder& ctx) {
TensorStructInfo input_sinfo = GetUnaryInputTensorStructInfo(call, ctx);
diff --git a/src/relax/op/distributed/distributed.cc b/src/relax/op/distributed/distributed.cc
index cdeb537c3d..84750c0c9c 100644
--- a/src/relax/op/distributed/distributed.cc
+++ b/src/relax/op/distributed/distributed.cc
@@ -48,7 +48,7 @@ Expr annotate_sharding(Expr input, distributed::DeviceMesh device_mesh,
return Call(op, {std::move(input)}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.dist.annotate_sharding").set_body_typed(annotate_sharding);
+TVM_FFI_REGISTER_GLOBAL("relax.op.dist.annotate_sharding").set_body_typed(annotate_sharding);
StructInfo InferStructInfoAnnotateSharding(const Call& call, const BlockBuilder& ctx) {
return GetStructInfo(call->args[0]);
@@ -73,7 +73,7 @@ Expr redistribute(Expr input, distributed::DeviceMesh device_mesh,
return Call(op, {std::move(input)}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.dist.redistribute").set_body_typed(redistribute);
+TVM_FFI_REGISTER_GLOBAL("relax.op.dist.redistribute").set_body_typed(redistribute);
StructInfo InferDistStructInfoRedistribute(const Call& call, const BlockBuilder& ctx) {
const auto* attrs = call->attrs.as<DistributionAttrs>();
@@ -139,7 +139,7 @@ Expr MakeCallTIRLocalView(Expr func, Tuple args,
return call;
}
-TVM_REGISTER_GLOBAL("relax.op.dist.call_tir_local_view").set_body_typed(MakeCallTIRLocalView);
+TVM_FFI_REGISTER_GLOBAL("relax.op.dist.call_tir_local_view").set_body_typed(MakeCallTIRLocalView);
StructInfo InferStructInfoRtoS(const Call& call, const BlockBuilder& ctx) {
TensorStructInfo input_sinfo = GetUnaryInputTensorStructInfo(call, ctx);
@@ -208,7 +208,7 @@ Expr redistribute_replica_to_shard(Expr input, int num_workers, int axis) {
return Call(op, {std::move(input)}, Attrs{attrs}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.dist.redistribute_replica_to_shard")
+TVM_FFI_REGISTER_GLOBAL("relax.op.dist.redistribute_replica_to_shard")
.set_body_typed(redistribute_replica_to_shard);
TVM_REGISTER_OP("relax.dist.redistribute_replica_to_shard")
diff --git a/src/relax/op/image/resize.cc b/src/relax/op/image/resize.cc
index 5b6550c729..c45e24df2b 100644
--- a/src/relax/op/image/resize.cc
+++ b/src/relax/op/image/resize.cc
@@ -50,7 +50,7 @@ Expr resize2d(Expr data, Expr size, Array<FloatImm> roi, String layout, String m
return Call(op, {std::move(data), std::move(size)}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.image.resize2d").set_body_typed(resize2d);
+TVM_FFI_REGISTER_GLOBAL("relax.op.image.resize2d").set_body_typed(resize2d);
StructInfo InferStructInfoResize2D(const Call& call, const BlockBuilder& ctx) {
if (call->args.size() != 1 && call->args.size() != 2) {
diff --git a/src/relax/op/memory/view.cc b/src/relax/op/memory/view.cc
index 0b8243a889..a7465db868 100644
--- a/src/relax/op/memory/view.cc
+++ b/src/relax/op/memory/view.cc
@@ -40,7 +40,7 @@ Expr view(Expr x, Optional<Expr> shape, Optional<Expr> dtype, Optional<Expr> rel
});
}
-TVM_REGISTER_GLOBAL("relax.op.memory.view").set_body_typed(view);
+TVM_FFI_REGISTER_GLOBAL("relax.op.memory.view").set_body_typed(view);
StructInfo InferStructInfoView(const Call& call, const BlockBuilder& ctx) {
if (call->args.size() != 4) {
@@ -289,7 +289,8 @@ StructInfo InferStructInfoView(const Call& call, const BlockBuilder& ctx) {
}
}
-TVM_REGISTER_GLOBAL("tvm.relax.struct_info.infer_view_sinfo").set_body_typed(InferStructInfoView);
+TVM_FFI_REGISTER_GLOBAL("tvm.relax.struct_info.infer_view_sinfo")
+ .set_body_typed(InferStructInfoView);
Expr LowerBuiltinView(const BlockBuilder& bb, const Call& call) {
Expr data = call->args[0];
@@ -360,7 +361,7 @@ Expr ensure_zero_offset(const Expr& x) {
return Call(op, {x});
}
-TVM_REGISTER_GLOBAL("relax.op.memory.ensure_zero_offset").set_body_typed(ensure_zero_offset);
+TVM_FFI_REGISTER_GLOBAL("relax.op.memory.ensure_zero_offset").set_body_typed(ensure_zero_offset);
StructInfo InferStructInfoEnsureZeroOffset(const Call& call, const BlockBuilder& ctx) {
if (call->args.size() != 1) {
diff --git a/src/relax/op/nn/attention.cc b/src/relax/op/nn/attention.cc
index d6e7fa707d..a084747a5c 100644
--- a/src/relax/op/nn/attention.cc
+++ b/src/relax/op/nn/attention.cc
@@ -57,8 +57,8 @@ Expr attention_var_len(Expr query, Expr key, Expr value, Expr seqstart_q, Expr s
{});
}
-TVM_REGISTER_GLOBAL("relax.op.nn.attention").set_body_typed(attention);
-TVM_REGISTER_GLOBAL("relax.op.nn.attention_var_len").set_body_typed(attention_var_len);
+TVM_FFI_REGISTER_GLOBAL("relax.op.nn.attention").set_body_typed(attention);
+TVM_FFI_REGISTER_GLOBAL("relax.op.nn.attention_var_len").set_body_typed(attention_var_len);
StructInfo InferStructInfoAttention(const Call& call, const BlockBuilder& ctx) {
Array<TensorStructInfo> input_sinfo = GetInputTensorStructInfo(call, ctx);
diff --git a/src/relax/op/nn/convolution.cc b/src/relax/op/nn/convolution.cc
index cca50689cb..f335bb9e7c 100644
--- a/src/relax/op/nn/convolution.cc
+++ b/src/relax/op/nn/convolution.cc
@@ -51,7 +51,7 @@ Expr conv1d(Expr data, Expr weight, Array<IntImm> strides, Array<IntImm> padding
out_dtype.value_or(DataType::Void()), /*op_name=*/"relax.nn.conv1d");
}
-TVM_REGISTER_GLOBAL("relax.op.nn.conv1d").set_body_typed(conv1d);
+TVM_FFI_REGISTER_GLOBAL("relax.op.nn.conv1d").set_body_typed(conv1d);
StructInfo InferStructInfoConv1d(const Call& call, const BlockBuilder& ctx) {
Array<TensorStructInfo> input_sinfo = GetInputTensorStructInfo(call, ctx);
@@ -214,7 +214,7 @@ Expr conv2d(Expr data, Expr weight, Array<IntImm> strides, Array<IntImm> padding
out_dtype.value_or(DataType::Void()), /*op_name=*/"relax.nn.conv2d");
}
-TVM_REGISTER_GLOBAL("relax.op.nn.conv2d").set_body_typed(conv2d);
+TVM_FFI_REGISTER_GLOBAL("relax.op.nn.conv2d").set_body_typed(conv2d);
StructInfo InferStructInfoConv2d(const Call& call, const BlockBuilder& ctx) {
Array<TensorStructInfo> input_sinfo = GetInputTensorStructInfo(call, ctx);
@@ -413,7 +413,7 @@ Expr conv3d(Expr data, Expr weight, Array<IntImm> strides, Array<IntImm> padding
out_dtype.value_or(DataType::Void()), /*op_name=*/"relax.nn.conv3d");
}
-TVM_REGISTER_GLOBAL("relax.op.nn.conv3d").set_body_typed(conv3d);
+TVM_FFI_REGISTER_GLOBAL("relax.op.nn.conv3d").set_body_typed(conv3d);
StructInfo InferStructInfoConv3d(const Call& call, const BlockBuilder& ctx) {
Array<TensorStructInfo> input_sinfo = GetInputTensorStructInfo(call, ctx);
@@ -593,7 +593,7 @@ Expr conv1d_transpose(Expr data, Expr weight, Array<IntImm> strides, Array<IntIm
return Call(op, {data, weight}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.nn.conv1d_transpose").set_body_typed(conv1d_transpose);
+TVM_FFI_REGISTER_GLOBAL("relax.op.nn.conv1d_transpose").set_body_typed(conv1d_transpose);
StructInfo InferStructInfoConv1dTranspose(const Call& call, const BlockBuilder& ctx) {
Array<TensorStructInfo> input_sinfo = GetInputTensorStructInfo(call, ctx);
@@ -730,7 +730,7 @@ Expr conv2d_transpose(Expr data, Expr weight, Array<IntImm> strides, Array<IntIm
return Call(op, {data, weight}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.nn.conv2d_transpose").set_body_typed(conv2d_transpose);
+TVM_FFI_REGISTER_GLOBAL("relax.op.nn.conv2d_transpose").set_body_typed(conv2d_transpose);
StructInfo InferStructInfoConv2dTranspose(const Call& call, const BlockBuilder& ctx) {
Array<TensorStructInfo> input_sinfo = GetInputTensorStructInfo(call, ctx);
diff --git a/src/relax/op/nn/nn.cc b/src/relax/op/nn/nn.cc
index 2296c364fc..b79690d3a9 100644
--- a/src/relax/op/nn/nn.cc
+++ b/src/relax/op/nn/nn.cc
@@ -50,7 +50,7 @@ Expr leakyrelu(Expr data, double alpha) {
return Call(op, {data}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.nn.leakyrelu").set_body_typed(leakyrelu);
+TVM_FFI_REGISTER_GLOBAL("relax.op.nn.leakyrelu").set_body_typed(leakyrelu);
TVM_REGISTER_OP("relax.nn.leakyrelu")
.set_num_inputs(1)
@@ -71,7 +71,7 @@ Expr softplus(Expr data, double beta, double threshold) {
return Call(op, {data}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.nn.softplus").set_body_typed(softplus);
+TVM_FFI_REGISTER_GLOBAL("relax.op.nn.softplus").set_body_typed(softplus);
TVM_REGISTER_OP("relax.nn.softplus")
.set_num_inputs(1)
@@ -91,7 +91,7 @@ Expr prelu(Expr data, Expr alpha, int axis = 1) {
return Call(op, {data, alpha}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.nn.prelu").set_body_typed(prelu);
+TVM_FFI_REGISTER_GLOBAL("relax.op.nn.prelu").set_body_typed(prelu);
TVM_REGISTER_OP("relax.nn.prelu")
.set_num_inputs(2)
@@ -112,7 +112,7 @@ Expr softmax(Expr data, int axis) {
return Call(op, {data}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.nn.softmax").set_body_typed(softmax);
+TVM_FFI_REGISTER_GLOBAL("relax.op.nn.softmax").set_body_typed(softmax);
StructInfo InferStructInfoSoftmax(const Call& call, const BlockBuilder& ctx) {
TensorStructInfo data_sinfo = GetUnaryInputTensorStructInfo(call, ctx);
@@ -170,7 +170,7 @@ Expr log_softmax(Expr data, int axis) {
return Call(op, {data}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.nn.log_softmax").set_body_typed(log_softmax);
+TVM_FFI_REGISTER_GLOBAL("relax.op.nn.log_softmax").set_body_typed(log_softmax);
TVM_REGISTER_OP("relax.nn.log_softmax")
.set_num_inputs(1)
@@ -191,7 +191,7 @@ Expr pad(Expr data, Array<Integer> pad_width, String pad_mode, double pad_value)
return Call(op, {data}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.nn.pad").set_body_typed(pad);
+TVM_FFI_REGISTER_GLOBAL("relax.op.nn.pad").set_body_typed(pad);
StructInfo InferStructInfoPad(const Call& call, const BlockBuilder& ctx) {
Array<TensorStructInfo> input_sinfo = GetInputTensorStructInfo(call, ctx);
@@ -234,7 +234,7 @@ Expr pixel_shuffle(Expr data, int upscale_factor) {
return Call(op, {data}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.nn.pixel_shuffle").set_body_typed(pixel_shuffle);
+TVM_FFI_REGISTER_GLOBAL("relax.op.nn.pixel_shuffle").set_body_typed(pixel_shuffle);
StructInfo InferStructInfoPixelShuffle(const Call& call, const BlockBuilder& ctx) {
Array<TensorStructInfo> input_sinfo = GetInputTensorStructInfo(call, ctx);
@@ -383,7 +383,7 @@ Expr batch_norm(Expr data, Expr gamma, Expr beta, Expr moving_mean, Expr moving_
std::move(moving_var)},
Attrs{attrs}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.nn.batch_norm").set_body_typed(batch_norm);
+TVM_FFI_REGISTER_GLOBAL("relax.op.nn.batch_norm").set_body_typed(batch_norm);
StructInfo InferStructInfoBatchNorm(const Call& call, const BlockBuilder& ctx) {
Array<TensorStructInfo> input_sinfo = GetInputTensorStructInfo(call, ctx);
@@ -461,7 +461,7 @@ Expr layer_norm(Expr data, Expr gamma, Expr beta, Array<Integer> axes, double ep
return Call(op, {std::move(data), std::move(gamma), std::move(beta)}, Attrs{attrs}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.nn.layer_norm").set_body_typed(layer_norm);
+TVM_FFI_REGISTER_GLOBAL("relax.op.nn.layer_norm").set_body_typed(layer_norm);
StructInfo InferStructInfoLayerNorm(const Call& call, const BlockBuilder& ctx) {
Array<TensorStructInfo> input_sinfo = GetInputTensorStructInfo(call, ctx);
@@ -529,7 +529,7 @@ Expr group_norm(Expr data, Expr gamma, Expr beta, int num_groups, int channel_ax
return Call(op, {std::move(data), std::move(gamma), std::move(beta)}, Attrs{attrs}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.nn.group_norm").set_body_typed(group_norm);
+TVM_FFI_REGISTER_GLOBAL("relax.op.nn.group_norm").set_body_typed(group_norm);
StructInfo InferStructInfoGroupNorm(const Call& call, const BlockBuilder& ctx) {
Op op = Downcast<Op>(call->op);
@@ -636,7 +636,7 @@ Expr rms_norm(Expr data, Expr weight, Array<Integer> axes, double epsilon) {
return Call(op, {std::move(data), std::move(weight)}, Attrs{attrs}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.nn.rms_norm").set_body_typed(rms_norm);
+TVM_FFI_REGISTER_GLOBAL("relax.op.nn.rms_norm").set_body_typed(rms_norm);
StructInfo InferStructInfoRMSNorm(const Call& call, const BlockBuilder& ctx) {
Array<TensorStructInfo> input_sinfo = GetInputTensorStructInfo(call, ctx);
@@ -695,7 +695,7 @@ Expr dropout(Expr data, double rate) {
return Call(op, {std::move(data)}, Attrs{attrs}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.nn.dropout").set_body_typed(dropout);
+TVM_FFI_REGISTER_GLOBAL("relax.op.nn.dropout").set_body_typed(dropout);
StructInfo InferStructInfoDropout(const Call& call, const BlockBuilder& ctx) {
TensorStructInfo data_sinfo = GetUnaryInputTensorStructInfo(call, ctx);
@@ -763,7 +763,7 @@ Expr cross_entropy_with_logits(Expr predictions, Expr labels) {
return Call(op, {std::move(predictions), std::move(labels)}, {}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.nn.cross_entropy_with_logits")
+TVM_FFI_REGISTER_GLOBAL("relax.op.nn.cross_entropy_with_logits")
.set_body_typed(cross_entropy_with_logits);
TVM_REGISTER_OP("relax.nn.cross_entropy_with_logits")
@@ -797,7 +797,7 @@ Expr nll_loss(Expr predictions, Expr targets, Optional<Expr> weights, String red
}
}
-TVM_REGISTER_GLOBAL("relax.op.nn.nll_loss").set_body_typed(nll_loss);
+TVM_FFI_REGISTER_GLOBAL("relax.op.nn.nll_loss").set_body_typed(nll_loss);
StructInfo InferStructInfoNLLLoss(const Call& call, const BlockBuilder& ctx) {
if (call->args.size() < 2 || call->args.size() > 3) {
diff --git a/src/relax/op/nn/pooling.cc b/src/relax/op/nn/pooling.cc
index 391edda9ef..0161a4d419 100644
--- a/src/relax/op/nn/pooling.cc
+++ b/src/relax/op/nn/pooling.cc
@@ -62,7 +62,7 @@ Expr max_pool1d(Expr data, Array<IntImm> pool_size, Array<IntImm> strides, Array
count_include_pad, layout, out_layout);
}
-TVM_REGISTER_GLOBAL("relax.op.nn.max_pool1d").set_body_typed(max_pool1d);
+TVM_FFI_REGISTER_GLOBAL("relax.op.nn.max_pool1d").set_body_typed(max_pool1d);
StructInfo InferStructInfoPool1D(const Call& call, const BlockBuilder& ctx) {
TensorStructInfo data_sinfo = GetUnaryInputTensorStructInfo(call, ctx);
@@ -175,7 +175,7 @@ Expr max_pool2d(Expr data, Array<IntImm> pool_size, Array<IntImm> strides, Array
count_include_pad, layout, out_layout);
}
-TVM_REGISTER_GLOBAL("relax.op.nn.max_pool2d").set_body_typed(max_pool2d);
+TVM_FFI_REGISTER_GLOBAL("relax.op.nn.max_pool2d").set_body_typed(max_pool2d);
StructInfo InferStructInfoPool2D(const Call& call, const BlockBuilder& ctx) {
TensorStructInfo data_sinfo = GetUnaryInputTensorStructInfo(call, ctx);
@@ -314,7 +314,7 @@ Expr max_pool3d(Expr data, Array<IntImm> pool_size, Array<IntImm> strides, Array
count_include_pad, layout, out_layout);
}
-TVM_REGISTER_GLOBAL("relax.op.nn.max_pool3d").set_body_typed(max_pool3d);
+TVM_FFI_REGISTER_GLOBAL("relax.op.nn.max_pool3d").set_body_typed(max_pool3d);
StructInfo InferStructInfoPool3D(const Call& call, const BlockBuilder& ctx) {
TensorStructInfo data_sinfo = GetUnaryInputTensorStructInfo(call, ctx);
@@ -401,7 +401,7 @@ Expr avg_pool1d(Expr data, Array<IntImm> pool_size, Array<IntImm> strides, Array
count_include_pad, layout, out_layout);
}
-TVM_REGISTER_GLOBAL("relax.op.nn.avg_pool1d").set_body_typed(avg_pool1d);
+TVM_FFI_REGISTER_GLOBAL("relax.op.nn.avg_pool1d").set_body_typed(avg_pool1d);
TVM_REGISTER_OP("relax.nn.avg_pool1d")
.set_num_inputs(1)
@@ -420,7 +420,7 @@ Expr avg_pool2d(Expr data, Array<IntImm> pool_size, Array<IntImm> strides, Array
count_include_pad, layout, out_layout);
}
-TVM_REGISTER_GLOBAL("relax.op.nn.avg_pool2d").set_body_typed(avg_pool2d);
+TVM_FFI_REGISTER_GLOBAL("relax.op.nn.avg_pool2d").set_body_typed(avg_pool2d);
TVM_REGISTER_OP("relax.nn.avg_pool2d")
.set_num_inputs(1)
@@ -439,7 +439,7 @@ Expr avg_pool3d(Expr data, Array<IntImm> pool_size, Array<IntImm> strides, Array
count_include_pad, layout, out_layout);
}
-TVM_REGISTER_GLOBAL("relax.op.nn.avg_pool3d").set_body_typed(avg_pool3d);
+TVM_FFI_REGISTER_GLOBAL("relax.op.nn.avg_pool3d").set_body_typed(avg_pool3d);
TVM_REGISTER_OP("relax.nn.avg_pool3d")
.set_num_inputs(1)
@@ -470,7 +470,7 @@ Expr adaptive_avg_pool1d(Expr data, Optional<Array<IntImm>> output_size, String
return Call(op, {std::move(data)}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.nn.adaptive_avg_pool1d").set_body_typed(adaptive_avg_pool1d);
+TVM_FFI_REGISTER_GLOBAL("relax.op.nn.adaptive_avg_pool1d").set_body_typed(adaptive_avg_pool1d);
StructInfo InferStructInfoAdaptiveAvgPool1D(const Call& call, const BlockBuilder& ctx) {
TensorStructInfo data_sinfo = GetUnaryInputTensorStructInfo(call, ctx);
@@ -553,7 +553,7 @@ Expr adaptive_avg_pool2d(Expr data, Optional<Array<IntImm>> output_size, String
return Call(op, {std::move(data)}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.nn.adaptive_avg_pool2d").set_body_typed(adaptive_avg_pool2d);
+TVM_FFI_REGISTER_GLOBAL("relax.op.nn.adaptive_avg_pool2d").set_body_typed(adaptive_avg_pool2d);
StructInfo InferStructInfoAdaptiveAvgPool2D(const Call& call, const BlockBuilder& ctx) {
TensorStructInfo data_sinfo = GetUnaryInputTensorStructInfo(call, ctx);
@@ -652,7 +652,7 @@ Expr adaptive_avg_pool3d(Expr data, Optional<Array<IntImm>> output_size, String
return Call(op, {std::move(data)}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.nn.adaptive_avg_pool3d").set_body_typed(adaptive_avg_pool3d);
+TVM_FFI_REGISTER_GLOBAL("relax.op.nn.adaptive_avg_pool3d").set_body_typed(adaptive_avg_pool3d);
StructInfo InferStructInfoAdaptiveAvgPool3D(const Call& call, const BlockBuilder& ctx) {
TensorStructInfo data_sinfo = GetUnaryInputTensorStructInfo(call, ctx);
diff --git a/src/relax/op/op.cc b/src/relax/op/op.cc
index 94fabd2d0e..c581b20835 100644
--- a/src/relax/op/op.cc
+++ b/src/relax/op/op.cc
@@ -119,7 +119,7 @@ Expr MakeCallPurePacked(const Expr& callee, Array<Expr> args, const Attrs& attrs
return Call(op, call_args, attrs, sinfo_args);
}
-TVM_REGISTER_GLOBAL("relax.op.call_pure_packed").set_body_typed(MakeCallPurePacked);
+TVM_FFI_REGISTER_GLOBAL("relax.op.call_pure_packed").set_body_typed(MakeCallPurePacked);
// call_inplace_packed
@@ -238,7 +238,7 @@ Expr MakeCallInplacePacked(Expr func, Array<Expr> args, Array<Integer> inplace_i
return Call(op, call_args, Attrs(attrs), sinfo_args);
}
-TVM_REGISTER_GLOBAL("relax.op.call_inplace_packed").set_body_typed(MakeCallInplacePacked);
+TVM_FFI_REGISTER_GLOBAL("relax.op.call_inplace_packed").set_body_typed(MakeCallInplacePacked);
// call_tir
@@ -600,7 +600,7 @@ Expr MakeCallTIR(Expr func, Tuple args, Array<TensorStructInfo> out_sinfo_list,
return call;
}
-TVM_REGISTER_GLOBAL("relax.op.call_tir").set_body_typed(MakeCallTIR);
+TVM_FFI_REGISTER_GLOBAL("relax.op.call_tir").set_body_typed(MakeCallTIR);
// call_tir_with_grad
@@ -652,7 +652,7 @@ Expr MakeCallTIRWithGrad(Expr func, Tuple args, Array<TensorStructInfo> out_sinf
return call;
}
-TVM_REGISTER_GLOBAL("relax.op.call_tir_with_grad").set_body_typed(MakeCallTIRWithGrad);
+TVM_FFI_REGISTER_GLOBAL("relax.op.call_tir_with_grad").set_body_typed(MakeCallTIRWithGrad);
// call_tir_inplace
@@ -793,7 +793,7 @@ Expr MakeCallTIRInplace(Expr func, Tuple args, Array<Integer> inplace_indices,
return call;
}
-TVM_REGISTER_GLOBAL("relax.op.call_tir_inplace").set_body_typed(MakeCallTIRInplace);
+TVM_FFI_REGISTER_GLOBAL("relax.op.call_tir_inplace").set_body_typed(MakeCallTIRInplace);
// call_dps_packed
@@ -834,7 +834,7 @@ Expr MakeCallDPSPacked(Expr func, Tuple args, Array<TensorStructInfo> out_sinfo_
return Call(op, {func, args}, {}, {out_sinfo});
}
-TVM_REGISTER_GLOBAL("relax.op.call_dps_packed").set_body_typed(MakeCallDPSPacked);
+TVM_FFI_REGISTER_GLOBAL("relax.op.call_dps_packed").set_body_typed(MakeCallDPSPacked);
// call builtin
StructInfo InferStructInfoCallBuiltinWithCtx(const Call& call, const BlockBuilder& ctx) {
@@ -860,7 +860,7 @@ Expr MakeCallBuiltinWithCtx(Expr func, Tuple args, Array<StructInfo> sinfo_args)
return Call(op, {func, args}, Attrs(), sinfo_args);
}
-TVM_REGISTER_GLOBAL("relax.op.call_builtin_with_ctx").set_body_typed(MakeCallBuiltinWithCtx);
+TVM_FFI_REGISTER_GLOBAL("relax.op.call_builtin_with_ctx").set_body_typed(MakeCallBuiltinWithCtx);
TVM_REGISTER_OP("relax.null_value")
.set_num_inputs(0)
@@ -872,7 +872,7 @@ Expr MakeCallNullValue() {
return Call(op, {}, {}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.null_value").set_body_typed(MakeCallNullValue);
+TVM_FFI_REGISTER_GLOBAL("relax.op.null_value").set_body_typed(MakeCallNullValue);
// print
@@ -895,7 +895,7 @@ Expr MakePrint(Array<Expr> vals, StringImm format) {
return Call(op, params);
}
-TVM_REGISTER_GLOBAL("relax.op.print").set_body_typed(MakePrint);
+TVM_FFI_REGISTER_GLOBAL("relax.op.print").set_body_typed(MakePrint);
// assert_op
@@ -938,7 +938,7 @@ Expr MakeAssertOp(Expr condition, Array<Expr> vals, StringImm format) {
return Call(op, args);
}
-TVM_REGISTER_GLOBAL("relax.op.assert_op").set_body_typed(MakeAssertOp);
+TVM_FFI_REGISTER_GLOBAL("relax.op.assert_op").set_body_typed(MakeAssertOp);
// make_closure
@@ -954,7 +954,7 @@ Expr MakeClosure(Expr func, Tuple args) {
return Call(op, {func, args}, {}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.make_closure").set_body_typed(MakeClosure);
+TVM_FFI_REGISTER_GLOBAL("relax.op.make_closure").set_body_typed(MakeClosure);
// invoke_closure
@@ -981,7 +981,7 @@ Expr InvokeClosure(Expr closure, Tuple args, Array<StructInfo> sinfo_args) {
return Call(op, {closure, args}, {}, sinfo_args);
}
-TVM_REGISTER_GLOBAL("relax.op.invoke_closure").set_body_typed(InvokeClosure);
+TVM_FFI_REGISTER_GLOBAL("relax.op.invoke_closure").set_body_typed(InvokeClosure);
// invoke_pure_closure
@@ -997,7 +997,7 @@ Expr InvokePureClosure(Expr closure, Tuple args, Array<StructInfo> sinfo_args) {
return Call(op, {closure, args}, {}, sinfo_args);
}
-TVM_REGISTER_GLOBAL("relax.op.invoke_pure_closure").set_body_typed(InvokePureClosure);
+TVM_FFI_REGISTER_GLOBAL("relax.op.invoke_pure_closure").set_body_typed(InvokePureClosure);
// shape_of
@@ -1012,7 +1012,7 @@ Expr MakeShapeOf(Expr expr) {
return Call(op, {expr}, {}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.shape_of").set_body_typed(MakeShapeOf);
+TVM_FFI_REGISTER_GLOBAL("relax.op.shape_of").set_body_typed(MakeShapeOf);
// tensor_to_shape
@@ -1046,7 +1046,7 @@ Expr MakeTensorToShape(Expr expr) {
return Call(op, {expr}, {}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.tensor_to_shape").set_body_typed(MakeTensorToShape);
+TVM_FFI_REGISTER_GLOBAL("relax.op.tensor_to_shape").set_body_typed(MakeTensorToShape);
// shape_to_tensor
StructInfo ReturnShapeToTensorStructInfo(const Call& call, const BlockBuilder& ctx) {
@@ -1070,7 +1070,7 @@ Expr MakeShapeToTensor(Expr expr) {
return Call(op, {expr}, {}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.shape_to_tensor").set_body_typed(MakeShapeToTensor);
+TVM_FFI_REGISTER_GLOBAL("relax.op.shape_to_tensor").set_body_typed(MakeShapeToTensor);
// alloc_tensor
@@ -1107,7 +1107,7 @@ Expr MakeAllocTensor(Expr shape, DataTypeImm dtype, PrimValue runtime_device_ind
return Call(op, {shape, dtype, runtime_device_index, storage_scope}, Attrs(), {});
}
-TVM_REGISTER_GLOBAL("relax.op.builtin.alloc_tensor").set_body_typed(MakeAllocTensor);
+TVM_FFI_REGISTER_GLOBAL("relax.op.builtin.alloc_tensor").set_body_typed(MakeAllocTensor);
// memory planning alloc_storage
@@ -1132,7 +1132,7 @@ Expr MakeAllocStorage(Expr size, PrimValue virtual_device_index, StringImm stora
return Call(op, {size, virtual_device_index, storage_scope, dtype}, Attrs(), {});
}
-TVM_REGISTER_GLOBAL("relax.op.memory.alloc_storage").set_body_typed(MakeAllocStorage);
+TVM_FFI_REGISTER_GLOBAL("relax.op.memory.alloc_storage").set_body_typed(MakeAllocStorage);
// memory planning alloc_tensor
@@ -1163,7 +1163,7 @@ Expr MakeMemAllocTensor(Expr storage, PrimValue offset, Expr shape, DataTypeImm
return Call(op, {storage, offset, shape, dtype}, Attrs(), {});
}
-TVM_REGISTER_GLOBAL("relax.op.memory.alloc_tensor").set_body_typed(MakeMemAllocTensor);
+TVM_FFI_REGISTER_GLOBAL("relax.op.memory.alloc_tensor").set_body_typed(MakeMemAllocTensor);
// memory planning kill_storage
@@ -1179,7 +1179,7 @@ Expr MakeMemKillStorage(Expr storage) {
return Call(op, {storage}, {}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.memory.kill_storage").set_body_typed(MakeMemKillStorage);
+TVM_FFI_REGISTER_GLOBAL("relax.op.memory.kill_storage").set_body_typed(MakeMemKillStorage);
// memory planning kill_tensor
@@ -1195,7 +1195,7 @@ Expr MakeMemKillTensor(Expr tensor) {
return Call(op, {tensor}, {}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.memory.kill_tensor").set_body_typed(MakeMemKillTensor);
+TVM_FFI_REGISTER_GLOBAL("relax.op.memory.kill_tensor").set_body_typed(MakeMemKillTensor);
// vm alloc_storage
@@ -1219,7 +1219,7 @@ Expr MakeVMAllocStorage(Expr size, PrimValue runtime_device_index, DataTypeImm d
return Call(op, {size, runtime_device_index, dtype, storage_scope}, Attrs(), {});
}
-TVM_REGISTER_GLOBAL("relax.op.vm.alloc_storage").set_body_typed(MakeVMAllocStorage);
+TVM_FFI_REGISTER_GLOBAL("relax.op.vm.alloc_storage").set_body_typed(MakeVMAllocStorage);
// vm alloc_tensor
@@ -1257,7 +1257,7 @@ Expr MakeVMAllocTensor(Expr storage, PrimValue offset, Expr shape, DataTypeImm d
return Call(op, {storage, offset, shape, dtype}, Attrs(), {});
}
-TVM_REGISTER_GLOBAL("relax.op.vm.alloc_tensor").set_body_typed(MakeVMAllocTensor);
+TVM_FFI_REGISTER_GLOBAL("relax.op.vm.alloc_tensor").set_body_typed(MakeVMAllocTensor);
// vm kill_object
@@ -1273,7 +1273,7 @@ Expr MakeVMKillObject(Expr obj) {
return Call(op, {std::move(obj)}, Attrs(), {});
}
-TVM_REGISTER_GLOBAL("relax.op.vm.kill_object").set_body_typed(MakeVMKillObject);
+TVM_FFI_REGISTER_GLOBAL("relax.op.vm.kill_object").set_body_typed(MakeVMKillObject);
// vm call_tir_dyn
@@ -1291,7 +1291,7 @@ Expr MakeCallTIRDyn(Expr func, Tuple args) {
return Call(op, {func, args}, Attrs(), {});
}
-TVM_REGISTER_GLOBAL("relax.op.vm.call_tir_dyn").set_body_typed(MakeCallTIRDyn);
+TVM_FFI_REGISTER_GLOBAL("relax.op.vm.call_tir_dyn").set_body_typed(MakeCallTIRDyn);
// builtin stop_lift_params
StructInfo InferStructInfoStopLiftParams(const Call& call, const BlockBuilder& ctx) {
@@ -1309,7 +1309,7 @@ Expr MakeStopLiftParams(Expr x) {
return Call(op, {x}, Attrs(), {});
}
-TVM_REGISTER_GLOBAL("relax.op.builtin.stop_lift_params").set_body_typed(MakeStopLiftParams);
+TVM_FFI_REGISTER_GLOBAL("relax.op.builtin.stop_lift_params").set_body_typed(MakeStopLiftParams);
// to_vdevice
TVM_REGISTER_NODE_TYPE(ToVDeviceAttrs);
@@ -1340,7 +1340,7 @@ Expr MakeToVDevice(Expr data, VDevice dst_vdev) {
return Call(op, {data}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.to_vdevice").set_body_typed(MakeToVDevice);
+TVM_FFI_REGISTER_GLOBAL("relax.op.to_vdevice").set_body_typed(MakeToVDevice);
// hint_on_device
TVM_REGISTER_NODE_TYPE(HintOnDeviceAttrs);
@@ -1367,7 +1367,7 @@ Expr MakeHintOnDevice(Expr data, Device device) {
return Call(op, {data}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.hint_on_device").set_body_typed(MakeHintOnDevice);
+TVM_FFI_REGISTER_GLOBAL("relax.op.hint_on_device").set_body_typed(MakeHintOnDevice);
} // namespace relax
} // namespace tvm
diff --git a/src/relax/op/op_common.h b/src/relax/op/op_common.h
index df29d9c885..d7d50f8fa7 100644
--- a/src/relax/op/op_common.h
+++ b/src/relax/op/op_common.h
@@ -181,7 +181,7 @@ std::tuple<ArgTypes...> GetArgStructInfo(const Call& call, const BlockBuilder& c
static const Op& op = Op::Get("relax." OpRegName); \
return Call(op, {std::move(x)}, Attrs(), {}); \
} \
- TVM_REGISTER_GLOBAL("relax.op." OpRegName).set_body_typed(OpName)
+ TVM_FFI_REGISTER_GLOBAL("relax.op." OpRegName).set_body_typed(OpName)
/************ Utilities ************/
diff --git a/src/relax/op/tensor/binary.h b/src/relax/op/tensor/binary.h
index 6b106f760d..ae36d45b36 100644
--- a/src/relax/op/tensor/binary.h
+++ b/src/relax/op/tensor/binary.h
@@ -42,7 +42,7 @@ namespace relax {
static const Op& op = Op::Get("relax." #OpName); \
return Call(op, {x1, x2}, Attrs(), {}); \
} \
- TVM_REGISTER_GLOBAL("relax.op." #OpName).set_body_typed(OpName); \
+ TVM_FFI_REGISTER_GLOBAL("relax.op." #OpName).set_body_typed(OpName); \
TVM_REGISTER_OP("relax." #OpName) \
.set_num_inputs(2) \
.add_argument("x1", "Tensor", "The first input tensor.") \
diff --git a/src/relax/op/tensor/create.cc b/src/relax/op/tensor/create.cc
index c06b51ae2e..b2355b1af7 100644
--- a/src/relax/op/tensor/create.cc
+++ b/src/relax/op/tensor/create.cc
@@ -54,7 +54,7 @@ Expr full(Variant<Expr, Array<PrimExpr>> shape, Expr fill_value, Optional<DataTy
return Call(op, {std::move(shape_in_expr), std::move(fill_value)}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.full").set_body_typed(full);
+TVM_FFI_REGISTER_GLOBAL("relax.op.full").set_body_typed(full);
StructInfo InferStructInfoFull(const Call& call, const BlockBuilder& ctx) {
if (call->args.size() != 2) {
@@ -96,7 +96,7 @@ Expr full_like(Expr x, Expr fill_value, Optional<DataType> dtype) {
return Call(op, {std::move(x), std::move(fill_value)}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.full_like").set_body_typed(full_like);
+TVM_FFI_REGISTER_GLOBAL("relax.op.full_like").set_body_typed(full_like);
StructInfo InferStructInfoFullLike(const Call& call, const BlockBuilder& ctx) {
Array<TensorStructInfo> input_sinfo = GetInputTensorStructInfo(call, ctx);
@@ -174,8 +174,8 @@ Expr ones_like(Expr x, Optional<DataType> dtype) {
return Call(op, {std::move(x)}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.ones").set_body_typed(ones);
-TVM_REGISTER_GLOBAL("relax.op.ones_like").set_body_typed(ones_like);
+TVM_FFI_REGISTER_GLOBAL("relax.op.ones").set_body_typed(ones);
+TVM_FFI_REGISTER_GLOBAL("relax.op.ones_like").set_body_typed(ones_like);
TVM_REGISTER_OP("relax.ones")
.set_attrs_type<InitAttrs>()
@@ -209,8 +209,8 @@ Expr zeros_like(Expr x, Optional<DataType> dtype) {
return Call(op, {std::move(x)}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.zeros").set_body_typed(zeros);
-TVM_REGISTER_GLOBAL("relax.op.zeros_like").set_body_typed(zeros_like);
+TVM_FFI_REGISTER_GLOBAL("relax.op.zeros").set_body_typed(zeros);
+TVM_FFI_REGISTER_GLOBAL("relax.op.zeros_like").set_body_typed(zeros_like);
TVM_REGISTER_OP("relax.zeros")
.set_attrs_type<InitAttrs>()
@@ -242,8 +242,8 @@ Expr eye_like(Expr x, PrimValue k, Optional<DataType> dtype) {
return Call(op, {std::move(x), std::move(k)}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.eye").set_body_typed(eye);
-TVM_REGISTER_GLOBAL("relax.op.eye_like").set_body_typed(eye_like);
+TVM_FFI_REGISTER_GLOBAL("relax.op.eye").set_body_typed(eye);
+TVM_FFI_REGISTER_GLOBAL("relax.op.eye_like").set_body_typed(eye_like);
StructInfo InferStructInfoEye(const Call& call, const BlockBuilder& ctx) {
if (call->args.size() != 3) {
@@ -319,7 +319,7 @@ Expr arange(PrimValue start, PrimValue stop, PrimValue step, DataType dtype) {
return Call(op, {std::move(start), std::move(stop), std::move(step)}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.arange").set_body_typed(arange);
+TVM_FFI_REGISTER_GLOBAL("relax.op.arange").set_body_typed(arange);
StructInfo InferStructInfoArange(const Call& call, const BlockBuilder& ctx) {
if (call->args.size() != 3) {
@@ -380,8 +380,8 @@ Expr triu(Expr x, Expr k) {
Expr triu(Expr x, int k) { return triu(x, relax::PrimValue::Int64(k)); }
-TVM_REGISTER_GLOBAL("relax.op.tril").set_body_typed(static_cast<Expr (*)(Expr, Expr)>(tril));
-TVM_REGISTER_GLOBAL("relax.op.triu").set_body_typed(static_cast<Expr (*)(Expr, Expr)>(triu));
+TVM_FFI_REGISTER_GLOBAL("relax.op.tril").set_body_typed(static_cast<Expr (*)(Expr, Expr)>(tril));
+TVM_FFI_REGISTER_GLOBAL("relax.op.triu").set_body_typed(static_cast<Expr (*)(Expr, Expr)>(triu));
StructInfo InferStructInfoTrilTriu(const Call& call, const BlockBuilder& ctx) {
auto [data_sinfo, offset] = GetArgStructInfo<TensorStructInfo, PrimStructInfo>(call, ctx);
diff --git a/src/relax/op/tensor/datatype.cc b/src/relax/op/tensor/datatype.cc
index bc24285cf9..d1d5bbccbc 100644
--- a/src/relax/op/tensor/datatype.cc
+++ b/src/relax/op/tensor/datatype.cc
@@ -40,7 +40,7 @@ Expr astype(Expr x, DataType dtype) {
return Call(op, {std::move(x)}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.astype").set_body_typed(astype);
+TVM_FFI_REGISTER_GLOBAL("relax.op.astype").set_body_typed(astype);
StructInfo InferStructInfoAstype(const Call& call, const BlockBuilder& ctx) {
TensorStructInfo sinfo = GetUnaryInputTensorStructInfo(call, ctx);
@@ -70,7 +70,7 @@ Expr MakeWrapParam(Expr data, DataType dtype) {
return Call(op, {std::move(data)}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.wrap_param").set_body_typed(MakeWrapParam);
+TVM_FFI_REGISTER_GLOBAL("relax.op.wrap_param").set_body_typed(MakeWrapParam);
StructInfo InferStructInfoWrapParam(const Call& call, const BlockBuilder& ctx) {
TensorStructInfo sinfo = GetUnaryInputTensorStructInfo(call, ctx);
diff --git a/src/relax/op/tensor/grad.cc b/src/relax/op/tensor/grad.cc
index d8aecb3461..c25d587052 100644
--- a/src/relax/op/tensor/grad.cc
+++ b/src/relax/op/tensor/grad.cc
@@ -35,7 +35,7 @@ Expr no_grad(Expr input) {
return Call(op, {std::move(input)}, {}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.grad.no_grad").set_body_typed(no_grad);
+TVM_FFI_REGISTER_GLOBAL("relax.op.grad.no_grad").set_body_typed(no_grad);
StructInfo InferStructInfoNoGrad(const Call& call, const BlockBuilder& ctx) {
return GetStructInfo(call->args[0]);
@@ -53,7 +53,7 @@ Expr start_checkpoint(Expr input) {
return Call(op, {std::move(input)}, {}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.grad.start_checkpoint").set_body_typed(start_checkpoint);
+TVM_FFI_REGISTER_GLOBAL("relax.op.grad.start_checkpoint").set_body_typed(start_checkpoint);
StructInfo InferStructInfoStartCheckpoint(const Call& call, const BlockBuilder& ctx) {
if (!call->args[0].as<VarNode>()) {
@@ -75,7 +75,7 @@ Expr end_checkpoint(Expr input) {
return Call(op, {std::move(input)}, {}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.grad.end_checkpoint").set_body_typed(end_checkpoint);
+TVM_FFI_REGISTER_GLOBAL("relax.op.grad.end_checkpoint").set_body_typed(end_checkpoint);
StructInfo InferStructInfoEndCheckpoint(const Call& call, const BlockBuilder& ctx) {
if (!call->args[0].as<VarNode>()) {
@@ -111,7 +111,7 @@ Expr nll_loss_backward(Expr output_grad, Expr predictions, Expr targets, Optiona
}
}
-TVM_REGISTER_GLOBAL("relax.op.grad.nll_loss_backward").set_body_typed(nll_loss_backward);
+TVM_FFI_REGISTER_GLOBAL("relax.op.grad.nll_loss_backward").set_body_typed(nll_loss_backward);
StructInfo InferStructInfoNLLLossBackward(const Call& call, const BlockBuilder& ctx) {
return GetStructInfo(call->args[1]);
@@ -145,7 +145,7 @@ Expr max_pool2d_backward(Expr output_grad, Expr data, Array<IntImm> pool_size,
return Call(op, {std::move(output_grad), std::move(data)}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.grad.max_pool2d_backward").set_body_typed(max_pool2d_backward);
+TVM_FFI_REGISTER_GLOBAL("relax.op.grad.max_pool2d_backward").set_body_typed(max_pool2d_backward);
StructInfo InferStructInfoMaxPool2DBackward(const Call& call, const BlockBuilder& ctx) {
return GetStructInfo(call->args[1]);
@@ -177,7 +177,7 @@ Expr avg_pool2d_backward(Expr output_grad, Expr data, Array<IntImm> pool_size,
return Call(op, {std::move(output_grad), std::move(data)}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.grad.avg_pool2d_backward").set_body_typed(avg_pool2d_backward);
+TVM_FFI_REGISTER_GLOBAL("relax.op.grad.avg_pool2d_backward").set_body_typed(avg_pool2d_backward);
StructInfo InferStructInfoAvgPool2DBackward(const Call& call, const BlockBuilder& ctx) {
return GetStructInfo(call->args[1]);
@@ -202,7 +202,7 @@ Expr take_backward(Expr output_grad, Expr x, Expr indices, Optional<int64_t> axi
return Call(op, {std::move(output_grad), std::move(x), std::move(indices)}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.grad.take_backward").set_body_typed(take_backward);
+TVM_FFI_REGISTER_GLOBAL("relax.op.grad.take_backward").set_body_typed(take_backward);
StructInfo InferStructInfoTakeBackward(const Call& call, const BlockBuilder& ctx) {
return GetStructInfo(call->args[1]);
diff --git a/src/relax/op/tensor/index.cc b/src/relax/op/tensor/index.cc
index 8f262ce38d..26978f2fad 100644
--- a/src/relax/op/tensor/index.cc
+++ b/src/relax/op/tensor/index.cc
@@ -47,7 +47,7 @@ Expr take(Expr x, Expr indices, Optional<int64_t> axis) {
return Call(op, {std::move(x), std::move(indices)}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.take").set_body_typed(take);
+TVM_FFI_REGISTER_GLOBAL("relax.op.take").set_body_typed(take);
StructInfo InferStructInfoTake(const Call& call, const BlockBuilder& ctx) {
CheckNumArguments(call, ctx);
@@ -169,7 +169,7 @@ Expr strided_slice(Expr x, Expr axes, Expr begin, Expr end, Optional<Expr> strid
return call;
}
-TVM_REGISTER_GLOBAL("relax.op.strided_slice").set_body_typed(strided_slice);
+TVM_FFI_REGISTER_GLOBAL("relax.op.strided_slice").set_body_typed(strided_slice);
/* \brief Helper function to unpack a relax::Tuple
*
@@ -477,7 +477,7 @@ Expr dynamic_strided_slice(Expr x, //
return Call(op, {std::move(x), std::move(begin), std::move(end), std::move(strides)}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.dynamic_strided_slice").set_body_typed(dynamic_strided_slice);
+TVM_FFI_REGISTER_GLOBAL("relax.op.dynamic_strided_slice").set_body_typed(dynamic_strided_slice);
StructInfo InferStructInfoDynStridedSlice(const Call& call, const BlockBuilder& ctx) {
const auto* data_sinfo = GetStructInfoAs<TensorStructInfoNode>(call->args[0]);
diff --git a/src/relax/op/tensor/linear_algebra.cc b/src/relax/op/tensor/linear_algebra.cc
index 4ca42bffec..28af375f00 100644
--- a/src/relax/op/tensor/linear_algebra.cc
+++ b/src/relax/op/tensor/linear_algebra.cc
@@ -44,7 +44,7 @@ Expr matmul(Expr x1, Expr x2, Optional<DataType> out_dtype) {
return Call(op, {std::move(x1), std::move(x2)}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.matmul").set_body_typed(matmul);
+TVM_FFI_REGISTER_GLOBAL("relax.op.matmul").set_body_typed(matmul);
StructInfo InferStructInfoMatmul(const Call& call, const BlockBuilder& ctx) {
Array<TensorStructInfo> input_sinfo = GetInputTensorStructInfo(call, ctx);
@@ -176,7 +176,7 @@ Expr einsum(Expr operands, String subscripts) {
return Call(op, {std::move(operands)}, Attrs{attrs}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.einsum").set_body_typed(einsum);
+TVM_FFI_REGISTER_GLOBAL("relax.op.einsum").set_body_typed(einsum);
StructInfo InferStructInfoEinsum(const Call& call, const BlockBuilder& ctx) {
if (call->args.size() != 1) {
@@ -258,7 +258,7 @@ Expr outer(Expr x1, Expr x2) {
return Call(op, {std::move(x1), std::move(x2)}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.outer").set_body_typed(outer);
+TVM_FFI_REGISTER_GLOBAL("relax.op.outer").set_body_typed(outer);
StructInfo InferStructInfoOuter(const Call& call, const BlockBuilder& ctx) {
auto input_sinfo = GetInputTensorStructInfo(call, ctx);
diff --git a/src/relax/op/tensor/manipulate.cc b/src/relax/op/tensor/manipulate.cc
index 9b653046f7..e98ba946c5 100644
--- a/src/relax/op/tensor/manipulate.cc
+++ b/src/relax/op/tensor/manipulate.cc
@@ -41,7 +41,7 @@ Expr broadcast_to(Expr x, Expr shape) {
return Call(op, {std::move(x), std::move(shape)}, Attrs(), {});
}
-TVM_REGISTER_GLOBAL("relax.op.broadcast_to").set_body_typed(broadcast_to);
+TVM_FFI_REGISTER_GLOBAL("relax.op.broadcast_to").set_body_typed(broadcast_to);
StructInfo InferStructInfoBroadcastTo(const Call& call, const BlockBuilder& ctx) {
if (call->args.size() != 2) {
@@ -124,7 +124,7 @@ Expr concat(Expr tensors, Optional<int64_t> axis) {
return Call(op, {std::move(tensors)}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.concat").set_body_typed(concat);
+TVM_FFI_REGISTER_GLOBAL("relax.op.concat").set_body_typed(concat);
Optional<Array<PrimExpr>> CheckConcatOutputShape(const Call& call, const BlockBuilder& ctx,
const std::vector<Array<PrimExpr>>& shape_values,
@@ -340,7 +340,7 @@ Expr expand_dims(Expr x, Array<Integer> axis) {
return Call(op, {std::move(x)}, Attrs{attrs}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.expand_dims").set_body_typed(expand_dims);
+TVM_FFI_REGISTER_GLOBAL("relax.op.expand_dims").set_body_typed(expand_dims);
StructInfo InferStructInfoExpandDims(const Call& call, const BlockBuilder& ctx) {
TensorStructInfo data_sinfo = GetUnaryInputTensorStructInfo(call, ctx);
@@ -446,7 +446,7 @@ Expr flatten(Expr x) {
return Call(op, {std::move(x)}, {}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.flatten").set_body_typed(flatten);
+TVM_FFI_REGISTER_GLOBAL("relax.op.flatten").set_body_typed(flatten);
StructInfo InferStructInfoFlatten(const Call& call, const BlockBuilder& ctx) {
TensorStructInfo data_sinfo = GetUnaryInputTensorStructInfo(call, ctx);
@@ -481,7 +481,7 @@ Expr index_tensor(Expr first, Expr tensors) {
return Call(op, {std::move(first), std::move(tensors)}, Attrs(), {});
}
-TVM_REGISTER_GLOBAL("relax.op.index_tensor").set_body_typed(index_tensor);
+TVM_FFI_REGISTER_GLOBAL("relax.op.index_tensor").set_body_typed(index_tensor);
StructInfo InferStructInfoIndexTensor(const Call& call, const BlockBuilder& ctx) {
if (call->args.size() != 2) {
@@ -635,7 +635,7 @@ Expr layout_transform(Expr x, tir::IndexMap index_map, Optional<PrimValue> pad_v
return Call(op, {std::move(x)}, Attrs{attrs}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.layout_transform").set_body_typed(layout_transform);
+TVM_FFI_REGISTER_GLOBAL("relax.op.layout_transform").set_body_typed(layout_transform);
StructInfo InferStructInfoLayoutTransform(const Call& call, const BlockBuilder& ctx) {
TensorStructInfo data_sinfo = GetUnaryInputTensorStructInfo(call, ctx);
@@ -702,7 +702,7 @@ Expr permute_dims(Expr x, Optional<Array<Integer>> axes) {
return Call(op, {std::move(x)}, Attrs{attrs}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.permute_dims").set_body_typed(permute_dims);
+TVM_FFI_REGISTER_GLOBAL("relax.op.permute_dims").set_body_typed(permute_dims);
bool IsIdentityPermutation(const std::vector<int>& permutation) {
for (int i = 0; i < static_cast<int>(permutation.size()); ++i) {
@@ -910,7 +910,7 @@ Expr reshape(Expr x, Variant<Expr, Array<PrimExpr>> shape) {
return Call(op, {std::move(x), std::move(shape_in_expr)}, Attrs(), {});
}
-TVM_REGISTER_GLOBAL("relax.op.reshape").set_body_typed(reshape);
+TVM_FFI_REGISTER_GLOBAL("relax.op.reshape").set_body_typed(reshape);
StructInfo InferStructInfoReshape(const Call& call, const BlockBuilder& ctx) {
if (call->args.size() != 2) {
@@ -997,7 +997,7 @@ Expr split(Expr x, Variant<IntImm, Array<IntImm>> indices_or_sections, int axis)
return Call(op, {std::move(x)}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.split").set_body_typed(split);
+TVM_FFI_REGISTER_GLOBAL("relax.op.split").set_body_typed(split);
StructInfo InferStructInfoSplit(const Call& call, const BlockBuilder& ctx) {
TensorStructInfo data_sinfo = GetUnaryInputTensorStructInfo(call, ctx);
@@ -1150,7 +1150,7 @@ Expr squeeze(Expr x, Optional<Array<Integer>> axis) {
return Call(op, {std::move(x)}, Attrs{attrs}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.squeeze").set_body_typed(squeeze);
+TVM_FFI_REGISTER_GLOBAL("relax.op.squeeze").set_body_typed(squeeze);
StructInfo InferStructInfoSqueeze(const Call& call, const BlockBuilder& ctx) {
TensorStructInfo data_sinfo = GetUnaryInputTensorStructInfo(call, ctx);
@@ -1350,7 +1350,7 @@ Expr stack(Expr tensors, Optional<Integer> axis) {
return Call(op, {std::move(tensors)}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.stack").set_body_typed(stack);
+TVM_FFI_REGISTER_GLOBAL("relax.op.stack").set_body_typed(stack);
Optional<Array<PrimExpr>> CheckStackOutputShape(const Call& call, const BlockBuilder& ctx,
const std::vector<Array<PrimExpr>>& shape_values,
@@ -1554,7 +1554,7 @@ Expr collapse_sum_like(Expr data, Expr collapse_target) {
return Call(op, {std::move(data), std::move(collapse_target)}, Attrs(), {});
}
-TVM_REGISTER_GLOBAL("relax.op.collapse_sum_like").set_body_typed(collapse_sum_like);
+TVM_FFI_REGISTER_GLOBAL("relax.op.collapse_sum_like").set_body_typed(collapse_sum_like);
StructInfo InferStructInfoCollapseSumLike(const Call& call, const BlockBuilder& ctx) {
Array<TensorStructInfo> input_sinfo = GetInputTensorStructInfo(call, ctx);
@@ -1600,7 +1600,7 @@ Expr collapse_sum_to(Expr data, Expr shape) {
return Call(op, {std::move(data), std::move(shape)}, Attrs(), {});
}
-TVM_REGISTER_GLOBAL("relax.op.collapse_sum_to").set_body_typed(collapse_sum_to);
+TVM_FFI_REGISTER_GLOBAL("relax.op.collapse_sum_to").set_body_typed(collapse_sum_to);
StructInfo InferStructInfoCollapseSumTo(const Call& call, const BlockBuilder& ctx) {
if (call->args.size() != 2) {
@@ -1655,7 +1655,7 @@ Expr repeat(Expr data, int repeats, Optional<int64_t> axis) {
return Call(op, {std::move(data)}, Attrs{attrs}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.repeat").set_body_typed(repeat);
+TVM_FFI_REGISTER_GLOBAL("relax.op.repeat").set_body_typed(repeat);
StructInfo InferStructInfoRepeat(const Call& call, const BlockBuilder& ctx) {
arith::Analyzer* analyzer = ctx->GetAnalyzer();
@@ -1720,7 +1720,7 @@ Expr tile(Expr data, Array<Integer> repeats) {
return Call(op, {std::move(data)}, Attrs{attrs}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.tile").set_body_typed(tile);
+TVM_FFI_REGISTER_GLOBAL("relax.op.tile").set_body_typed(tile);
StructInfo InferStructInfoTile(const Call& call, const BlockBuilder& ctx) {
arith::Analyzer* analyzer = ctx->GetAnalyzer();
@@ -1783,7 +1783,7 @@ Expr flip(Expr data, Integer axis) {
return Call(op, {std::move(data)}, Attrs{attrs}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.flip").set_body_typed(flip);
+TVM_FFI_REGISTER_GLOBAL("relax.op.flip").set_body_typed(flip);
StructInfo InferStructInfoFlip(const Call& call, const BlockBuilder& ctx) {
if (call->args.size() != 1) {
@@ -1820,7 +1820,7 @@ Expr gather_elements(Expr data, Expr indices, int axis) {
return Call(op, {data, indices}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.gather_elements").set_body_typed(gather_elements);
+TVM_FFI_REGISTER_GLOBAL("relax.op.gather_elements").set_body_typed(gather_elements);
StructInfo InferStructInfoGatherElements(const Call& call, const BlockBuilder& ctx) {
const auto* data_sinfo = GetStructInfoAs<TensorStructInfoNode>(call->args[0]);
@@ -1889,7 +1889,7 @@ Expr gather_nd(Expr data, Expr indices, int batch_dims) {
return Call(op, {data, indices}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.gather_nd").set_body_typed(gather_nd);
+TVM_FFI_REGISTER_GLOBAL("relax.op.gather_nd").set_body_typed(gather_nd);
StructInfo InferStructInfoGatherND(const Call& call, const BlockBuilder& ctx) {
const auto* data_sinfo = GetStructInfoAs<TensorStructInfoNode>(call->args[0]);
@@ -1983,7 +1983,7 @@ Expr index_put(Expr data, Expr indices, Expr values, bool accumulate) {
return Call(op, {data, indices, values}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.index_put").set_body_typed(index_put);
+TVM_FFI_REGISTER_GLOBAL("relax.op.index_put").set_body_typed(index_put);
StructInfo InferStructInfoIndexPut(const Call& call, const BlockBuilder& ctx) {
const auto* data_sinfo = GetStructInfoAs<TensorStructInfoNode>(call->args[0]);
@@ -2106,7 +2106,7 @@ Expr meshgrid(Expr tensors, Optional<String> indexing) {
return Call(op, {std::move(tensors)}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.meshgrid").set_body_typed(meshgrid);
+TVM_FFI_REGISTER_GLOBAL("relax.op.meshgrid").set_body_typed(meshgrid);
StructInfo InferStructInfoMeshgrid(const Call& call, const BlockBuilder& ctx) {
if (call->args.size() != 1) {
@@ -2210,7 +2210,7 @@ Expr scatter_elements(Expr data, Expr indices, Expr updates, int axis, String re
return Call(op, {data, indices, updates}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.scatter_elements").set_body_typed(scatter_elements);
+TVM_FFI_REGISTER_GLOBAL("relax.op.scatter_elements").set_body_typed(scatter_elements);
StructInfo InferStructInfoScatterElements(const Call& call, const BlockBuilder& ctx) {
arith::Analyzer* analyzer = ctx->GetAnalyzer();
@@ -2324,7 +2324,7 @@ Expr scatter_nd(Expr data, Expr indices, Expr updates, String reduction) {
return Call(op, {data, indices, updates}, Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.scatter_nd").set_body_typed(scatter_nd);
+TVM_FFI_REGISTER_GLOBAL("relax.op.scatter_nd").set_body_typed(scatter_nd);
StructInfo InferStructInfoScatterND(const Call& call, const BlockBuilder& ctx) {
// `call->args` contains: [data, indices, updates]
@@ -2467,7 +2467,7 @@ Expr one_hot(Expr indices, PrimValue on_value, PrimValue off_value, int depth, i
return Call(op, {indices, on_value, off_value}, Attrs(attrs), {});
} // namespace relax
-TVM_REGISTER_GLOBAL("relax.op.one_hot").set_body_typed(one_hot);
+TVM_FFI_REGISTER_GLOBAL("relax.op.one_hot").set_body_typed(one_hot);
StructInfo InferStructInfoOneHot(const Call& call, const BlockBuilder& ctx) {
TensorStructInfo indices_sinfo = GetInputTensorStructInfo(call, 0, ctx);
diff --git a/src/relax/op/tensor/qdq.cc b/src/relax/op/tensor/qdq.cc
index 0189ef9678..78ba6fec34 100644
--- a/src/relax/op/tensor/qdq.cc
+++ b/src/relax/op/tensor/qdq.cc
@@ -44,7 +44,7 @@ Expr quantize(Expr data, Expr scale, Expr zero_point, int axis, DataType out_dty
return Call(op, {std::move(data), std::move(scale), std::move(zero_point)}, Attrs(attrs));
}
-TVM_REGISTER_GLOBAL("relax.op.quantize").set_body_typed(quantize);
+TVM_FFI_REGISTER_GLOBAL("relax.op.quantize").set_body_typed(quantize);
StructInfo InferStructInfoQuantize(const Call& call, const BlockBuilder& ctx) {
const auto* attrs = call->attrs.as<QuantizeAttrs>();
@@ -128,7 +128,7 @@ Expr dequantize(Expr data, Expr scale, Expr zero_point, int axis, DataType out_d
return Call(op, {std::move(data), std::move(scale), std::move(zero_point)}, Attrs(attrs));
}
-TVM_REGISTER_GLOBAL("relax.op.dequantize").set_body_typed(dequantize);
+TVM_FFI_REGISTER_GLOBAL("relax.op.dequantize").set_body_typed(dequantize);
StructInfo InferStructInfoDequantize(const Call& call, const BlockBuilder& ctx) {
const auto* attrs = call->attrs.as<QuantizeAttrs>();
diff --git a/src/relax/op/tensor/sampling.cc b/src/relax/op/tensor/sampling.cc
index 35ee4c486b..80bbf48fd4 100644
--- a/src/relax/op/tensor/sampling.cc
+++ b/src/relax/op/tensor/sampling.cc
@@ -43,7 +43,8 @@ Expr multinomial_from_uniform(Expr prob, Expr uniform_sample, Expr sample_indice
Attrs(attrs), {});
}
-TVM_REGISTER_GLOBAL("relax.op.multinomial_from_uniform").set_body_typed(multinomial_from_uniform);
+TVM_FFI_REGISTER_GLOBAL("relax.op.multinomial_from_uniform")
+ .set_body_typed(multinomial_from_uniform);
StructInfo InferStructInfoMultinomialFromUniform(const Call& call, const BlockBuilder& ctx) {
CheckNumArguments(call, ctx);
diff --git a/src/relax/op/tensor/search.cc b/src/relax/op/tensor/search.cc
index 4df1662154..83e0e246b1 100644
--- a/src/relax/op/tensor/search.cc
+++ b/src/relax/op/tensor/search.cc
@@ -36,7 +36,7 @@ Expr where(Expr condition, Expr x1, Expr x2) {
return Call(op, {std::move(condition), std::move(x1), std::move(x2)}, Attrs(), {});
}
-TVM_REGISTER_GLOBAL("relax.op.where").set_body_typed(where);
+TVM_FFI_REGISTER_GLOBAL("relax.op.where").set_body_typed(where);
StructInfo InferStructInfoWhere(const Call& call, const BlockBuilder& ctx) {
Array<TensorStructInfo> input_sinfo = GetInputTensorStructInfo(call, ctx);
@@ -195,7 +195,7 @@ StructInfo InferStructInfoArgmaxArgmin(const Call& call, const BlockBuilder& ctx
static const Op& op = Op::Get("relax." #OpName); \
return Call(op, {std::move(x)}, Attrs(attrs)); \
} \
- TVM_REGISTER_GLOBAL("relax.op." #OpName).set_body_typed(OpName); \
+ TVM_FFI_REGISTER_GLOBAL("relax.op." #OpName).set_body_typed(OpName); \
TVM_REGISTER_OP("relax." #OpName) \
.set_num_inputs(1) \
.add_argument("x", "Tensor", "The input data tensor") \
diff --git a/src/relax/op/tensor/set.cc b/src/relax/op/tensor/set.cc
index f2f840ba31..e321b326d2 100644
--- a/src/relax/op/tensor/set.cc
+++ b/src/relax/op/tensor/set.cc
@@ -46,7 +46,7 @@ Expr unique(Expr x, PrimValue sorted, PrimValue return_index, PrimValue return_i
return call;
}
-TVM_REGISTER_GLOBAL("relax.op.unique").set_body_typed(unique);
+TVM_FFI_REGISTER_GLOBAL("relax.op.unique").set_body_typed(unique);
StructInfo InferStructInfoUnique(const Call& call, const BlockBuilder& ctx) {
TensorStructInfo data_sinfo = Downcast<TensorStructInfo>(call->args[0]->struct_info_);
@@ -144,7 +144,7 @@ Expr nonzero(Expr x) {
return Call(op, {std::move(x)});
}
-TVM_REGISTER_GLOBAL("relax.op.nonzero").set_body_typed(nonzero);
+TVM_FFI_REGISTER_GLOBAL("relax.op.nonzero").set_body_typed(nonzero);
StructInfo InferStructInfoNonzero(const Call& call, const BlockBuilder& ctx) {
TensorStructInfo data_sinfo = GetInputTensorStructInfo(call, 0, ctx);
diff --git a/src/relax/op/tensor/sorting.cc b/src/relax/op/tensor/sorting.cc
index 9f8545e9b3..1cd061084e 100644
--- a/src/relax/op/tensor/sorting.cc
+++ b/src/relax/op/tensor/sorting.cc
@@ -41,7 +41,7 @@ Expr sort(Expr data, int axis, bool descending) {
return Call(op, {std::move(data)}, Attrs{attrs}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.sort").set_body_typed(sort);
+TVM_FFI_REGISTER_GLOBAL("relax.op.sort").set_body_typed(sort);
StructInfo InferStructInfoSort(const Call& call, const BlockBuilder& ctx) {
return GetUnaryInputTensorStructInfo(call, ctx);
@@ -67,7 +67,7 @@ Expr argsort(Expr data, int axis, bool descending, DataType dtype) {
return Call(op, {std::move(data)}, Attrs{attrs}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.argsort").set_body_typed(argsort);
+TVM_FFI_REGISTER_GLOBAL("relax.op.argsort").set_body_typed(argsort);
StructInfo InferStructInfoArgsort(const Call& call, const BlockBuilder& ctx) {
TensorStructInfo data_sinfo = GetUnaryInputTensorStructInfo(call, ctx);
@@ -101,7 +101,7 @@ Expr topk(Expr data, int k, int axis, String ret_type, bool largest, DataType dt
return Call(op, {std::move(data)}, Attrs{attrs}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.topk").set_body_typed(topk);
+TVM_FFI_REGISTER_GLOBAL("relax.op.topk").set_body_typed(topk);
StructInfo InferStructInfoTopK(const Call& call, const BlockBuilder& ctx) {
TensorStructInfo data_sinfo = GetUnaryInputTensorStructInfo(call, ctx);
diff --git a/src/relax/op/tensor/statistical.cc b/src/relax/op/tensor/statistical.cc
index e4765c8ddb..73e74578fc 100644
--- a/src/relax/op/tensor/statistical.cc
+++ b/src/relax/op/tensor/statistical.cc
@@ -186,7 +186,7 @@ Expr cumprod(Expr data, Optional<int64_t> axis, Optional<DataType> dtype, Bool e
return Call(op, {std::move(data)}, Attrs{attrs}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.cumprod").set_body_typed(cumprod);
+TVM_FFI_REGISTER_GLOBAL("relax.op.cumprod").set_body_typed(cumprod);
TVM_REGISTER_OP("relax.cumprod")
.set_attrs_type<ScanopAttrs>()
@@ -206,7 +206,7 @@ Expr cumsum(Expr data, Optional<int64_t> axis, Optional<DataType> dtype, Bool ex
return Call(op, {std::move(data)}, Attrs{attrs}, {});
}
-TVM_REGISTER_GLOBAL("relax.op.cumsum").set_body_typed(cumsum);
+TVM_FFI_REGISTER_GLOBAL("relax.op.cumsum").set_body_typed(cumsum);
TVM_REGISTER_OP("relax.cumsum")
.set_attrs_type<ScanopAttrs>()
diff --git a/src/relax/op/tensor/statistical.h b/src/relax/op/tensor/statistical.h
index aa312f5df7..331562454e 100644
--- a/src/relax/op/tensor/statistical.h
+++ b/src/relax/op/tensor/statistical.h
@@ -50,7 +50,7 @@ namespace relax {
static const Op& op = Op::Get("relax." #OpName); \
return Call(op, {std::move(x)}, Attrs{attrs}, {}); \
} \
- TVM_REGISTER_GLOBAL("relax.op." #OpName).set_body_typed(OpName); \
+ TVM_FFI_REGISTER_GLOBAL("relax.op." #OpName).set_body_typed(OpName); \
TVM_REGISTER_OP("relax." #OpName) \
.set_num_inputs(1) \
.add_argument("x", "Tensor", "The input data tensor") \
diff --git a/src/relax/op/tensor/ternary.cc b/src/relax/op/tensor/ternary.cc
index cc265ad9a1..91a6e8d0ae 100644
--- a/src/relax/op/tensor/ternary.cc
+++ b/src/relax/op/tensor/ternary.cc
@@ -143,7 +143,7 @@ Expr ewise_fma(Expr x1, Expr x2, Expr x3) {
return Call(op, {x1, x2, x3}, Attrs(), {});
}
-TVM_REGISTER_GLOBAL("relax.op.ewise_fma").set_body_typed(ewise_fma);
+TVM_FFI_REGISTER_GLOBAL("relax.op.ewise_fma").set_body_typed(ewise_fma);
} // namespace relax
} // namespace tvm
diff --git a/src/relax/op/tensor/unary.cc b/src/relax/op/tensor/unary.cc
index f95eb721fc..828a91dde2 100644
--- a/src/relax/op/tensor/unary.cc
+++ b/src/relax/op/tensor/unary.cc
@@ -85,7 +85,7 @@ Expr clip(Expr x, Expr min, Expr max) {
return Call(op, {std::move(x), std::move(min), std::move(max)});
}
-TVM_REGISTER_GLOBAL("relax.op.clip").set_body_typed(clip);
+TVM_FFI_REGISTER_GLOBAL("relax.op.clip").set_body_typed(clip);
/***************** Check operators *****************/
diff --git a/src/relax/testing/transform.cc b/src/relax/testing/transform.cc
index eed2329e3d..c4e41d5afc 100644
--- a/src/relax/testing/transform.cc
+++ b/src/relax/testing/transform.cc
@@ -35,7 +35,7 @@ tvm::transform::Pass ApplyEmptyCppMutator() {
"relax.testing.ApplyEmptyCppMutator", {});
}
-TVM_REGISTER_GLOBAL("relax.testing.transform.ApplyEmptyCppMutator")
+TVM_FFI_REGISTER_GLOBAL("relax.testing.transform.ApplyEmptyCppMutator")
.set_body_typed(ApplyEmptyCppMutator);
} // namespace testing
diff --git a/src/relax/training/utils.cc b/src/relax/training/utils.cc
index 43fa2357f0..cb44339f19 100644
--- a/src/relax/training/utils.cc
+++ b/src/relax/training/utils.cc
@@ -215,7 +215,7 @@ Pass AppendLoss(String func_name, Function loss_function, int num_backbone_outpu
/*required=*/{});
}
-TVM_REGISTER_GLOBAL("relax.training.AppendLoss").set_body_typed(AppendLoss);
+TVM_FFI_REGISTER_GLOBAL("relax.training.AppendLoss").set_body_typed(AppendLoss);
} // namespace transform
diff --git a/src/relax/transform/adjust_matmul_order.cc b/src/relax/transform/adjust_matmul_order.cc
index 8d268cfef7..46dc803018 100644
--- a/src/relax/transform/adjust_matmul_order.cc
+++ b/src/relax/transform/adjust_matmul_order.cc
@@ -213,7 +213,7 @@ Pass AdjustMatmulOrder() {
return CreateFunctionPass(pass_func, 1, "AdjustMatmulOrder", {});
}
-TVM_REGISTER_GLOBAL("relax.transform.AdjustMatmulOrder").set_body_typed(AdjustMatmulOrder);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.AdjustMatmulOrder").set_body_typed(AdjustMatmulOrder);
} // namespace transform
} // namespace relax
diff --git a/src/relax/transform/allocate_workspace.cc b/src/relax/transform/allocate_workspace.cc
index c8a2cef400..763a009a24 100644
--- a/src/relax/transform/allocate_workspace.cc
+++ b/src/relax/transform/allocate_workspace.cc
@@ -201,7 +201,7 @@ Pass AllocateWorkspace() {
return CreateModulePass(pass_func, 0, "AllocateWorkspace", {});
}
-TVM_REGISTER_GLOBAL("relax.transform.AllocateWorkspace").set_body_typed(AllocateWorkspace);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.AllocateWorkspace").set_body_typed(AllocateWorkspace);
} // namespace transform
} // namespace tvm
diff --git a/src/relax/transform/alter_op_impl.cc b/src/relax/transform/alter_op_impl.cc
index 61d7725e6c..63521e4e8f 100644
--- a/src/relax/transform/alter_op_impl.cc
+++ b/src/relax/transform/alter_op_impl.cc
@@ -438,7 +438,7 @@ Pass AlterOpImpl(const Map<String, tir::PrimFunc>& op_impl_map,
/*required=*/{});
}
-TVM_REGISTER_GLOBAL("relax.transform.AlterOpImpl").set_body_typed(AlterOpImpl);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.AlterOpImpl").set_body_typed(AlterOpImpl);
} // namespace transform
} // namespace relax
diff --git a/src/relax/transform/annotate_tir_op_pattern.cc b/src/relax/transform/annotate_tir_op_pattern.cc
index 5127b3df0c..e2b0fc2c28 100644
--- a/src/relax/transform/annotate_tir_op_pattern.cc
+++ b/src/relax/transform/annotate_tir_op_pattern.cc
@@ -47,7 +47,8 @@ Pass AnnotateTIROpPattern() {
return tir::transform::CreatePrimFuncPass(pass_func, 0, "AnnotateTIROpPattern", {});
}
-TVM_REGISTER_GLOBAL("relax.transform.AnnotateTIROpPattern").set_body_typed(AnnotateTIROpPattern);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.AnnotateTIROpPattern")
+ .set_body_typed(AnnotateTIROpPattern);
} // namespace transform
diff --git a/src/relax/transform/attach_attr_layout_free_buffers.cc b/src/relax/transform/attach_attr_layout_free_buffers.cc
index 3593b22c10..cef7489080 100644
--- a/src/relax/transform/attach_attr_layout_free_buffers.cc
+++ b/src/relax/transform/attach_attr_layout_free_buffers.cc
@@ -105,7 +105,7 @@ Pass AttachAttrLayoutFreeBuffers() {
return tvm::transform::Sequential({pass, DeadCodeElimination()}, "AttachAttrLayoutFreeBuffers");
}
-TVM_REGISTER_GLOBAL("relax.transform.AttachAttrLayoutFreeBuffers")
+TVM_FFI_REGISTER_GLOBAL("relax.transform.AttachAttrLayoutFreeBuffers")
.set_body_typed(AttachAttrLayoutFreeBuffers);
} // namespace transform
} // namespace relax
diff --git a/src/relax/transform/attach_global_symbol.cc b/src/relax/transform/attach_global_symbol.cc
index 97226491a8..905d2bcd83 100644
--- a/src/relax/transform/attach_global_symbol.cc
+++ b/src/relax/transform/attach_global_symbol.cc
@@ -79,7 +79,7 @@ Pass AttachGlobalSymbol() {
return CreateModulePass(pass_func, 0, "AttachGlobalSymbol", {});
}
-TVM_REGISTER_GLOBAL("relax.transform.AttachGlobalSymbol").set_body_typed(AttachGlobalSymbol);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.AttachGlobalSymbol").set_body_typed(AttachGlobalSymbol);
} // namespace transform
} // namespace relax
} // namespace tvm
diff --git a/src/relax/transform/bind_params.cc b/src/relax/transform/bind_params.cc
index 6a871edeaa..2a5c6f525d 100644
--- a/src/relax/transform/bind_params.cc
+++ b/src/relax/transform/bind_params.cc
@@ -196,7 +196,7 @@ IRModule BindParam(IRModule m, String func_name, Map<ObjectRef, ObjectRef> bind_
return GetRef<IRModule>(new_module);
}
-TVM_REGISTER_GLOBAL("relax.FunctionBindParams").set_body_typed(FunctionBindParams);
+TVM_FFI_REGISTER_GLOBAL("relax.FunctionBindParams").set_body_typed(FunctionBindParams);
namespace transform {
@@ -207,7 +207,7 @@ Pass BindParams(String func_name, Map<ObjectRef, ObjectRef> params) {
return CreateModulePass(pass_func, 0, "BindParams", {});
}
-TVM_REGISTER_GLOBAL("relax.transform.BindParams").set_body_typed(BindParams);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.BindParams").set_body_typed(BindParams);
} // namespace transform
diff --git a/src/relax/transform/bind_symbolic_vars.cc b/src/relax/transform/bind_symbolic_vars.cc
index 2df9ed1f01..49af21c107 100644
--- a/src/relax/transform/bind_symbolic_vars.cc
+++ b/src/relax/transform/bind_symbolic_vars.cc
@@ -148,7 +148,7 @@ IRModule ModuleBindSymbolicVars(IRModule mod, Map<ObjectRef, PrimExpr> binding_m
}
} // namespace
-TVM_REGISTER_GLOBAL("relax.FunctionBindSymbolicVars").set_body_typed(FunctionBindSymbolicVars);
+TVM_FFI_REGISTER_GLOBAL("relax.FunctionBindSymbolicVars").set_body_typed(FunctionBindSymbolicVars);
namespace transform {
@@ -170,7 +170,7 @@ Pass BindSymbolicVars(Map<ObjectRef, PrimExpr> binding_map, Optional<String> fun
return tvm::transform::CreateModulePass(pass_func, 1, "relax.BindSymbolicVars", {});
}
-TVM_REGISTER_GLOBAL("relax.transform.BindSymbolicVars").set_body_typed(BindSymbolicVars);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.BindSymbolicVars").set_body_typed(BindSymbolicVars);
} // namespace transform
} // namespace relax
diff --git a/src/relax/transform/bundle_model_params.cc b/src/relax/transform/bundle_model_params.cc
index a011841c13..982e1ac0c3 100644
--- a/src/relax/transform/bundle_model_params.cc
+++ b/src/relax/transform/bundle_model_params.cc
@@ -115,7 +115,7 @@ Pass BundleModelParams(Optional<String> param_tuple_name) {
return CreateModulePass(pass_func, 1, "BundleModelParams", {});
}
-TVM_REGISTER_GLOBAL("relax.transform.BundleModelParams").set_body_typed(BundleModelParams);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.BundleModelParams").set_body_typed(BundleModelParams);
} // namespace transform
} // namespace relax
diff --git a/src/relax/transform/call_tir_rewrite.cc b/src/relax/transform/call_tir_rewrite.cc
index e3bb2bcbae..25b4abadc7 100644
--- a/src/relax/transform/call_tir_rewrite.cc
+++ b/src/relax/transform/call_tir_rewrite.cc
@@ -183,7 +183,7 @@ Pass CallTIRRewrite() {
/*required=*/{});
}
-TVM_REGISTER_GLOBAL("relax.transform.CallTIRRewrite").set_body_typed(CallTIRRewrite);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.CallTIRRewrite").set_body_typed(CallTIRRewrite);
} // namespace transform
diff --git a/src/relax/transform/canonicalize_bindings.cc b/src/relax/transform/canonicalize_bindings.cc
index 3d8c65de95..ecbb9e7751 100644
--- a/src/relax/transform/canonicalize_bindings.cc
+++ b/src/relax/transform/canonicalize_bindings.cc
@@ -591,7 +591,8 @@ Pass CanonicalizeBindings() {
"CanonicalizeBindings");
}
-TVM_REGISTER_GLOBAL("relax.transform.CanonicalizeBindings").set_body_typed(CanonicalizeBindings);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.CanonicalizeBindings")
+ .set_body_typed(CanonicalizeBindings);
} // namespace transform
diff --git a/src/relax/transform/combine_parallel_matmul.cc b/src/relax/transform/combine_parallel_matmul.cc
index 916061aa57..6201863203 100644
--- a/src/relax/transform/combine_parallel_matmul.cc
+++ b/src/relax/transform/combine_parallel_matmul.cc
@@ -387,7 +387,8 @@ Pass CombineParallelMatmul(FCheck check) {
/*required=*/{});
}
-TVM_REGISTER_GLOBAL("relax.transform.CombineParallelMatmul").set_body_typed(CombineParallelMatmul);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.CombineParallelMatmul")
+ .set_body_typed(CombineParallelMatmul);
} // namespace transform
diff --git a/src/relax/transform/compute_prim_value.cc b/src/relax/transform/compute_prim_value.cc
index 25a1d1b0ed..e6db2eb73f 100644
--- a/src/relax/transform/compute_prim_value.cc
+++ b/src/relax/transform/compute_prim_value.cc
@@ -86,7 +86,7 @@ Pass ComputePrimValue() {
return CreateModulePass(pass_func, 0, "ComputePrimValue", {});
}
-TVM_REGISTER_GLOBAL("relax.transform.ComputePrimValue").set_body_typed(ComputePrimValue);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.ComputePrimValue").set_body_typed(ComputePrimValue);
} // namespace transform
diff --git a/src/relax/transform/convert_dataflow.cc b/src/relax/transform/convert_dataflow.cc
index 1c4cef892a..c359afdebc 100644
--- a/src/relax/transform/convert_dataflow.cc
+++ b/src/relax/transform/convert_dataflow.cc
@@ -159,7 +159,7 @@ Pass ConvertToDataflow(int min_size) {
return tvm::transform::Sequential({pass, CanonicalizeBindings()});
}
-TVM_REGISTER_GLOBAL("relax.transform.ConvertToDataflow").set_body_typed(ConvertToDataflow);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.ConvertToDataflow").set_body_typed(ConvertToDataflow);
} // namespace transform
diff --git a/src/relax/transform/convert_layout.cc b/src/relax/transform/convert_layout.cc
index 78670bd51a..0c06cac75d 100644
--- a/src/relax/transform/convert_layout.cc
+++ b/src/relax/transform/convert_layout.cc
@@ -350,7 +350,7 @@ Pass ConvertLayout(Map<String, Array<String>> desired_layouts) {
return CreateDataflowBlockPass(pass_func, 0, "ConvertLayout", {});
}
-TVM_REGISTER_GLOBAL("relax.transform.ConvertLayout").set_body_typed(ConvertLayout);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.ConvertLayout").set_body_typed(ConvertLayout);
} // namespace transform
} // namespace relax
diff --git a/src/relax/transform/dataflow_inplace.cc b/src/relax/transform/dataflow_inplace.cc
index aee2c015fc..51ab6bb230 100644
--- a/src/relax/transform/dataflow_inplace.cc
+++ b/src/relax/transform/dataflow_inplace.cc
@@ -1016,13 +1016,13 @@ Array<Array<InplaceOpportunity>> DataflowInplaceAnalysis(const DataflowBlock& bl
}
// these are exposed only for testing
-TVM_REGISTER_GLOBAL("relax.testing.transform.DataflowLivenessAnalysis")
+TVM_FFI_REGISTER_GLOBAL("relax.testing.transform.DataflowLivenessAnalysis")
.set_body_typed(DataflowLivenessAnalysis);
-TVM_REGISTER_GLOBAL("relax.testing.transform.DataflowAliasAnalysis")
+TVM_FFI_REGISTER_GLOBAL("relax.testing.transform.DataflowAliasAnalysis")
.set_body_typed(DataflowAliasAnalysis);
-TVM_REGISTER_GLOBAL("relax.testing.transform.DataflowInplaceAnalysis")
+TVM_FFI_REGISTER_GLOBAL("relax.testing.transform.DataflowInplaceAnalysis")
.set_body_typed(DataflowInplaceAnalysis);
-TVM_REGISTER_GLOBAL("relax.testing.transform.SingleInplaceCall")
+TVM_FFI_REGISTER_GLOBAL("relax.testing.transform.SingleInplaceCall")
.set_body_typed([](const IRModule& mod, const Call& call,
const Array<Integer>& inplace_indices) -> Array<ObjectRef> {
ModuleInplaceTransformer transformer(mod);
@@ -1031,7 +1031,7 @@ TVM_REGISTER_GLOBAL("relax.testing.transform.SingleInplaceCall")
});
// actually exposed
-TVM_REGISTER_GLOBAL("relax.transform.DataflowUseInplaceCalls")
+TVM_FFI_REGISTER_GLOBAL("relax.transform.DataflowUseInplaceCalls")
.set_body_typed(DataflowUseInplaceCalls);
} // namespace transform
diff --git a/src/relax/transform/dead_code_elimination.cc b/src/relax/transform/dead_code_elimination.cc
index abceb517c1..7de1da329f 100644
--- a/src/relax/transform/dead_code_elimination.cc
+++ b/src/relax/transform/dead_code_elimination.cc
@@ -140,7 +140,7 @@ Pass DeadCodeElimination(Array<String> entry_functions) {
return CreateModulePass(pass_func, 1, "DeadCodeElimination", {});
}
-TVM_REGISTER_GLOBAL("relax.transform.DeadCodeElimination").set_body_typed(DeadCodeElimination);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.DeadCodeElimination").set_body_typed(DeadCodeElimination);
} // namespace transform
} // namespace relax
diff --git a/src/relax/transform/decompose_ops.cc b/src/relax/transform/decompose_ops.cc
index 1a4cd21625..eec27f3b78 100644
--- a/src/relax/transform/decompose_ops.cc
+++ b/src/relax/transform/decompose_ops.cc
@@ -250,10 +250,10 @@ Pass DecomposeOpsForTraining(Optional<String> func_name) {
}
}
-TVM_REGISTER_GLOBAL("relax.transform.DecomposeOpsForInference")
+TVM_FFI_REGISTER_GLOBAL("relax.transform.DecomposeOpsForInference")
.set_body_typed(DecomposeOpsForInference);
-TVM_REGISTER_GLOBAL("relax.transform.DecomposeOpsForTraining")
+TVM_FFI_REGISTER_GLOBAL("relax.transform.DecomposeOpsForTraining")
.set_body_typed(DecomposeOpsForTraining);
} // namespace transform
diff --git a/src/relax/transform/eliminate_common_subexpr.cc b/src/relax/transform/eliminate_common_subexpr.cc
index 5173df7ef3..8a5ce1db04 100644
--- a/src/relax/transform/eliminate_common_subexpr.cc
+++ b/src/relax/transform/eliminate_common_subexpr.cc
@@ -221,7 +221,7 @@ Pass EliminateCommonSubexpr(bool call_only) {
return CreateFunctionPass(pass_func, 1, "EliminateCommonSubexpr", {});
}
-TVM_REGISTER_GLOBAL("relax.transform.EliminateCommonSubexpr")
+TVM_FFI_REGISTER_GLOBAL("relax.transform.EliminateCommonSubexpr")
.set_body_typed(EliminateCommonSubexpr);
} // namespace transform
diff --git a/src/relax/transform/expand_matmul_of_sum.cc b/src/relax/transform/expand_matmul_of_sum.cc
index 134eca5572..d7bf2dd95f 100644
--- a/src/relax/transform/expand_matmul_of_sum.cc
+++ b/src/relax/transform/expand_matmul_of_sum.cc
@@ -104,7 +104,7 @@ Pass ExpandMatmulOfSum() {
return CreateFunctionPass(pass_func, 1, "ExpandMatmulOfSum", {});
}
-TVM_REGISTER_GLOBAL("relax.transform.ExpandMatmulOfSum").set_body_typed(ExpandMatmulOfSum);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.ExpandMatmulOfSum").set_body_typed(ExpandMatmulOfSum);
} // namespace transform
} // namespace relax
diff --git a/src/relax/transform/expand_tuple_arguments.cc b/src/relax/transform/expand_tuple_arguments.cc
index 8d5c833d43..1a9afadf7e 100644
--- a/src/relax/transform/expand_tuple_arguments.cc
+++ b/src/relax/transform/expand_tuple_arguments.cc
@@ -178,7 +178,8 @@ Pass ExpandTupleArguments() {
"ExpandTupleArguments");
}
-TVM_REGISTER_GLOBAL("relax.transform.ExpandTupleArguments").set_body_typed(ExpandTupleArguments);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.ExpandTupleArguments")
+ .set_body_typed(ExpandTupleArguments);
} // namespace transform
diff --git a/src/relax/transform/few_shot_tuning.cc b/src/relax/transform/few_shot_tuning.cc
index 084791bde1..4ccf6c25ab 100644
--- a/src/relax/transform/few_shot_tuning.cc
+++ b/src/relax/transform/few_shot_tuning.cc
@@ -172,7 +172,7 @@ Pass FewShotTuning(int valid_count, bool benchmark) {
/*required=*/{});
}
-TVM_REGISTER_GLOBAL("relax.transform.FewShotTuning").set_body_typed(FewShotTuning);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.FewShotTuning").set_body_typed(FewShotTuning);
} // namespace transform
} // namespace relax
diff --git a/src/relax/transform/fold_constant.cc b/src/relax/transform/fold_constant.cc
index 2d916d0391..c2bac5daa7 100644
--- a/src/relax/transform/fold_constant.cc
+++ b/src/relax/transform/fold_constant.cc
@@ -326,7 +326,7 @@ Pass FoldConstant() {
return CreateFunctionPass(pass_func, 0, "FoldConstant", {});
}
-TVM_REGISTER_GLOBAL("relax.transform.FoldConstant").set_body_typed(FoldConstant);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.FoldConstant").set_body_typed(FoldConstant);
} // namespace transform
diff --git a/src/relax/transform/fuse_ops.cc b/src/relax/transform/fuse_ops.cc
index 74141724ee..f9ffcd9302 100644
--- a/src/relax/transform/fuse_ops.cc
+++ b/src/relax/transform/fuse_ops.cc
@@ -1395,7 +1395,7 @@ FusionPattern::FusionPattern(String name, DFPattern pattern,
}
TVM_REGISTER_NODE_TYPE(FusionPatternNode);
-TVM_REGISTER_GLOBAL("relax.transform.FusionPattern")
+TVM_FFI_REGISTER_GLOBAL("relax.transform.FusionPattern")
.set_body_typed([](String name, DFPattern pattern, Map<String, DFPattern> annotation_patterns,
Optional<ffi::Function> check, Optional<ffi::Function> attrs_getter) {
return FusionPattern(name, pattern, annotation_patterns, check, attrs_getter);
@@ -1429,7 +1429,7 @@ Pass FuseOps(int fuse_opt_level) {
/*required=*/{});
}
-TVM_REGISTER_GLOBAL("relax.transform.FuseOps").set_body_typed(FuseOps);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.FuseOps").set_body_typed(FuseOps);
Pass FuseOpsByPattern(const tvm::Array<FusionPattern>& patterns, bool bind_constants,
bool annotate_codegen, const Array<String>& entry_function_names) {
@@ -1444,7 +1444,7 @@ Pass FuseOpsByPattern(const tvm::Array<FusionPattern>& patterns, bool bind_const
/*required=*/{});
}
-TVM_REGISTER_GLOBAL("relax.transform.FuseOpsByPattern").set_body_typed(FuseOpsByPattern);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.FuseOpsByPattern").set_body_typed(FuseOpsByPattern);
} // namespace transform
diff --git a/src/relax/transform/fuse_tir.cc b/src/relax/transform/fuse_tir.cc
index 0dc807a5a6..05b7bf4218 100644
--- a/src/relax/transform/fuse_tir.cc
+++ b/src/relax/transform/fuse_tir.cc
@@ -1262,7 +1262,7 @@ Pass FuseTIR() {
"FuseTIR");
}
-TVM_REGISTER_GLOBAL("relax.transform.FuseTIR").set_body_typed(FuseTIR);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.FuseTIR").set_body_typed(FuseTIR);
} // namespace transform
diff --git a/src/relax/transform/gradient.cc b/src/relax/transform/gradient.cc
index 35c6d9af5e..9998b6da93 100644
--- a/src/relax/transform/gradient.cc
+++ b/src/relax/transform/gradient.cc
@@ -787,7 +787,7 @@ Pass Gradient(String func_name, Optional<Array<Var>> require_grads, int target_i
/*required=*/{});
}
-TVM_REGISTER_GLOBAL("relax.transform.Gradient").set_body_typed(Gradient);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.Gradient").set_body_typed(Gradient);
} // namespace transform
diff --git a/src/relax/transform/inline_functions.cc b/src/relax/transform/inline_functions.cc
index 981a5e654c..26b106373f 100644
--- a/src/relax/transform/inline_functions.cc
+++ b/src/relax/transform/inline_functions.cc
@@ -164,7 +164,7 @@ Function FunctionInlineFunctions(Function func,
return Downcast<Function>(mutator(std::move(func)));
}
-TVM_REGISTER_GLOBAL("relax.FunctionInlineFunctions").set_body_typed(FunctionInlineFunctions);
+TVM_FFI_REGISTER_GLOBAL("relax.FunctionInlineFunctions").set_body_typed(FunctionInlineFunctions);
namespace transform {
@@ -219,7 +219,7 @@ Pass InlinePrivateFunctions() {
return tvm::transform::CreateModulePass(pass_func, 0, "InlinePrivateFunctions", {});
}
-TVM_REGISTER_GLOBAL("relax.transform.InlinePrivateFunctions")
+TVM_FFI_REGISTER_GLOBAL("relax.transform.InlinePrivateFunctions")
.set_body_typed(InlinePrivateFunctions);
} // namespace transform
diff --git a/src/relax/transform/kill_after_last_use.cc b/src/relax/transform/kill_after_last_use.cc
index 20ec5eb434..730f65f701 100644
--- a/src/relax/transform/kill_after_last_use.cc
+++ b/src/relax/transform/kill_after_last_use.cc
@@ -265,7 +265,7 @@ Pass KillAfterLastUse() {
return CreateFunctionPass(pass_func, /*opt_level=*/0, "KillAfterLastUse", {});
}
-TVM_REGISTER_GLOBAL("relax.transform.KillAfterLastUse").set_body_typed(KillAfterLastUse);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.KillAfterLastUse").set_body_typed(KillAfterLastUse);
} // namespace transform
} // namespace relax
diff --git a/src/relax/transform/lambda_lift.cc b/src/relax/transform/lambda_lift.cc
index 41a970abfe..e5e28cb553 100644
--- a/src/relax/transform/lambda_lift.cc
+++ b/src/relax/transform/lambda_lift.cc
@@ -503,7 +503,7 @@ Pass LambdaLift() {
return tvm::transform::CreateModulePass(pass_func, 1, "LambdaLift", {});
}
-TVM_REGISTER_GLOBAL("relax.transform.LambdaLift").set_body_typed(LambdaLift);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.LambdaLift").set_body_typed(LambdaLift);
} // namespace transform
} // namespace relax
diff --git a/src/relax/transform/lazy_transform_params.cc b/src/relax/transform/lazy_transform_params.cc
index d1cf4ff714..32f63e1e14 100644
--- a/src/relax/transform/lazy_transform_params.cc
+++ b/src/relax/transform/lazy_transform_params.cc
@@ -259,7 +259,7 @@ Pass LazyGetInput() {
/*required=*/{});
}
-TVM_REGISTER_GLOBAL("relax.transform.LazyGetInput").set_body_typed(LazyGetInput);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.LazyGetInput").set_body_typed(LazyGetInput);
Pass LazySetOutput() {
auto pass_func = [](Function func, IRModule, PassContext) -> Function {
@@ -274,7 +274,7 @@ Pass LazySetOutput() {
/*required=*/{});
}
-TVM_REGISTER_GLOBAL("relax.transform.LazySetOutput").set_body_typed(LazySetOutput);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.LazySetOutput").set_body_typed(LazySetOutput);
} // namespace transform
} // namespace relax
diff --git a/src/relax/transform/legalize_ops.cc b/src/relax/transform/legalize_ops.cc
index 71a55a3d84..a0ac6fffb6 100644
--- a/src/relax/transform/legalize_ops.cc
+++ b/src/relax/transform/legalize_ops.cc
@@ -404,7 +404,7 @@ Pass LegalizeOps(Optional<Map<String, ffi::Function>> cmap, bool enable_warning)
/*required=*/{});
}
-TVM_REGISTER_GLOBAL("relax.transform.LegalizeOps").set_body_typed(LegalizeOps);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.LegalizeOps").set_body_typed(LegalizeOps);
} // namespace transform
diff --git a/src/relax/transform/lift_transform_params.cc b/src/relax/transform/lift_transform_params.cc
index df7242efd4..9013737df5 100644
--- a/src/relax/transform/lift_transform_params.cc
+++ b/src/relax/transform/lift_transform_params.cc
@@ -867,7 +867,7 @@ Pass LiftTransformParams(Variant<Bool, Array<String>> shared_transform) {
"LiftTransformParams");
}
-TVM_REGISTER_GLOBAL("relax.transform.LiftTransformParams").set_body_typed(LiftTransformParams);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.LiftTransformParams").set_body_typed(LiftTransformParams);
} // namespace transform
} // namespace relax
diff --git a/src/relax/transform/lower_alloc_tensor.cc b/src/relax/transform/lower_alloc_tensor.cc
index 13705c0908..3bdbfd0b94 100644
--- a/src/relax/transform/lower_alloc_tensor.cc
+++ b/src/relax/transform/lower_alloc_tensor.cc
@@ -99,7 +99,7 @@ Pass LowerAllocTensor() {
return CreateFunctionPass(pass_func, /*opt_level=*/0, "LowerAllocTensor", {});
}
-TVM_REGISTER_GLOBAL("relax.transform.LowerAllocTensor").set_body_typed(LowerAllocTensor);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.LowerAllocTensor").set_body_typed(LowerAllocTensor);
} // namespace transform
} // namespace relax
diff --git a/src/relax/transform/merge_composite_functions.cc b/src/relax/transform/merge_composite_functions.cc
index 9538249b77..ffeddd08c4 100644
--- a/src/relax/transform/merge_composite_functions.cc
+++ b/src/relax/transform/merge_composite_functions.cc
@@ -421,7 +421,7 @@ Pass MergeCompositeFunctions() {
/*required=*/{});
}
-TVM_REGISTER_GLOBAL("relax.transform.MergeCompositeFunctions")
+TVM_FFI_REGISTER_GLOBAL("relax.transform.MergeCompositeFunctions")
.set_body_typed(MergeCompositeFunctions);
} // namespace transform
diff --git a/src/relax/transform/meta_schedule.cc b/src/relax/transform/meta_schedule.cc
index 57012a40dd..cf7b9fc03a 100644
--- a/src/relax/transform/meta_schedule.cc
+++ b/src/relax/transform/meta_schedule.cc
@@ -201,10 +201,11 @@ Pass MetaScheduleTuneTIR(String work_dir, Integer max_trials_global) {
/*traceable*/ true);
}
-TVM_REGISTER_GLOBAL("relax.transform.MetaScheduleApplyDatabase")
+TVM_FFI_REGISTER_GLOBAL("relax.transform.MetaScheduleApplyDatabase")
.set_body_typed(MetaScheduleApplyDatabase);
-TVM_REGISTER_GLOBAL("relax.transform.MetaScheduleTuneIRMod").set_body_typed(MetaScheduleTuneIRMod);
-TVM_REGISTER_GLOBAL("relax.transform.MetaScheduleTuneTIR").set_body_typed(MetaScheduleTuneTIR);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.MetaScheduleTuneIRMod")
+ .set_body_typed(MetaScheduleTuneIRMod);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.MetaScheduleTuneTIR").set_body_typed(MetaScheduleTuneTIR);
} // namespace transform
} // namespace relax
} // namespace tvm
diff --git a/src/relax/transform/normalize.cc b/src/relax/transform/normalize.cc
index 2c57c89a26..07ca6a1133 100644
--- a/src/relax/transform/normalize.cc
+++ b/src/relax/transform/normalize.cc
@@ -279,7 +279,7 @@ Pass Normalize() {
return CreateFunctionPass(pass_func, 1, "Normalize", {});
}
-TVM_REGISTER_GLOBAL("relax.transform.Normalize").set_body_typed(Normalize);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.Normalize").set_body_typed(Normalize);
Pass NormalizeGlobalVar() {
auto pass_func = [=](IRModule mod, PassContext pc) {
@@ -290,7 +290,7 @@ Pass NormalizeGlobalVar() {
/*pass_name=*/"NormalizeGlobalVar",
/*required=*/{});
}
-TVM_REGISTER_GLOBAL("relax.transform.NormalizeGlobalVar").set_body_typed(NormalizeGlobalVar);
+TVM_FFI_REGISTER_GLOBAL("relax.transform.NormalizeGlobalVar").set_body_typed(NormalizeGlobalVar);
} // namespace transform
diff --git a/src/relax/transform/realize_vdevice.cc b/src/relax/transform/realize_vdevice.cc
index 15474eb8f8..ee4773fb3a 100644
--- a/src/relax/transform/realize_vdevice.cc
+++ b/src/relax/transform/realize_vdevice.cc
@@ -415,7 +415,7 @@ Pass RealizeVDevice() {
/*required=*/{});
... 13474 lines suppressed ...