You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2025/08/23 16:28:38 UTC
(tvm) 01/01: [FFI][REFACTOR] Establish tvm_ffi as a standalone python module
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch ffi-package
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 9e88a762563d578eb606ca97327775c5ee4ef7c6
Author: tqchen <ti...@gmail.com>
AuthorDate: Wed Aug 20 09:05:55 2025 -0400
[FFI][REFACTOR] Establish tvm_ffi as a standalone python module
This PR establishes tvm_ffi as a standalone python module.
The ffi is structured as a minimal pip module that can be
directly install by path or url.
---
.gitmodules | 5 +-
3rdparty/dlpack | 1 -
CMakeLists.txt | 28 +---
Makefile | 10 --
ci/jenkins/data.py | 7 +-
ci/jenkins/generated/arm_jenkinsfile.groovy | 4 +-
ci/jenkins/generated/cpu_jenkinsfile.groovy | 4 +-
ci/jenkins/generated/gpu_jenkinsfile.groovy | 6 +-
ci/jenkins/generated/hexagon_jenkinsfile.groovy | 4 +-
ci/jenkins/generated/i386_jenkinsfile.groovy | 4 +-
docs/arch/pass_infra.rst | 2 +-
ffi/.clang-format | 8 +
{3rdparty => ffi/3rdparty}/libbacktrace | 0
ffi/CMakeLists.txt | 179 +++++++++++++++------
ffi/README.md | 18 +++
ffi/cmake/Utils/AddGoogleTest.cmake | 2 +-
ffi/cmake/Utils/AddLibbacktrace.cmake | 3 +-
ffi/cmake/Utils/CxxWarning.cmake | 2 +-
ffi/cmake/Utils/Library.cmake | 12 +-
ffi/cmake/tvm_ffi-config.cmake | 41 +++++
ffi/examples/get_started/CMakeLists.txt | 50 ++++++
ffi/examples/get_started/README.md | 4 +
ffi/examples/get_started/run_example.py | 74 +++++++++
.../examples/get_started/run_example.sh | 11 +-
ffi/examples/get_started/src/add_one_cpu.cc | 43 +++++
ffi/examples/get_started/src/add_one_cuda.cu | 59 +++++++
ffi/examples/get_started/src/run_example.cc | 55 +++++++
{licenses => ffi/licenses}/LICENSE.dlpack.txt | 0
.../licenses}/LICENSE.libbacktrace.txt | 0
ffi/pyproject.toml | 152 +++++++++++++++++
{python/tvm/ffi => ffi/python/tvm_ffi}/.gitignore | 0
{python/tvm/ffi => ffi/python/tvm_ffi}/__init__.py | 9 +-
{python/tvm/ffi => ffi/python/tvm_ffi}/_ffi_api.py | 0
.../tvm/ffi => ffi/python/tvm_ffi}/access_path.py | 0
{python/tvm => ffi/python/tvm_ffi}/base.py | 61 +++----
ffi/python/tvm_ffi/config.py | 64 ++++++++
.../tvm/ffi => ffi/python/tvm_ffi}/container.py | 0
{python/tvm/ffi => ffi/python/tvm_ffi}/convert.py | 0
.../tvm/ffi => ffi/python/tvm_ffi}/cython/base.pxi | 0
.../tvm/ffi => ffi/python/tvm_ffi}/cython/core.pyx | 0
.../ffi => ffi/python/tvm_ffi}/cython/device.pxi | 0
.../ffi => ffi/python/tvm_ffi}/cython/dtype.pxi | 0
.../ffi => ffi/python/tvm_ffi}/cython/error.pxi | 0
.../ffi => ffi/python/tvm_ffi}/cython/function.pxi | 4 +-
.../ffi => ffi/python/tvm_ffi}/cython/ndarray.pxi | 0
.../ffi => ffi/python/tvm_ffi}/cython/object.pxi | 0
.../ffi => ffi/python/tvm_ffi}/cython/string.pxi | 0
{python/tvm/ffi => ffi/python/tvm_ffi}/dtype.py | 0
{python/tvm/ffi => ffi/python/tvm_ffi}/error.py | 0
ffi/python/tvm_ffi/libinfo.py | 145 +++++++++++++++++
{python/tvm/ffi => ffi/python/tvm_ffi}/module.py | 13 +-
{python/tvm/ffi => ffi/python/tvm_ffi}/ndarray.py | 6 +-
{python/tvm/ffi => ffi/python/tvm_ffi}/registry.py | 2 +-
.../ffi => ffi/python/tvm_ffi}/serialization.py | 0
{python/tvm/ffi => ffi/python/tvm_ffi}/testing.py | 0
ffi/scripts/benchmark_dlpack.py | 28 ++--
ffi/tests/cpp/CMakeLists.txt | 8 +-
.../ffi => ffi/tests/python}/test_access_path.py | 2 +-
.../ffi => ffi/tests/python}/test_container.py | 2 +-
.../python/ffi => ffi/tests/python}/test_device.py | 4 +-
.../python/ffi => ffi/tests/python}/test_dtype.py | 4 +-
.../python/ffi => ffi/tests/python}/test_error.py | 2 +-
.../ffi => ffi/tests/python}/test_function.py | 2 +-
.../ffi => ffi/tests/python}/test_ndarray.py | 2 +-
.../python/ffi => ffi/tests/python}/test_object.py | 2 +-
.../python/ffi => ffi/tests/python}/test_string.py | 2 +-
include/tvm/runtime/c_backend_api.h | 9 --
python/setup.py | 70 +-------
python/tvm/__init__.py | 6 +-
python/tvm/arith/_ffi_api.py | 4 +-
python/tvm/arith/analyzer.py | 6 +-
python/tvm/arith/int_set.py | 8 +-
python/tvm/arith/int_solver.py | 8 +-
python/tvm/arith/iter_affine_map.py | 12 +-
python/tvm/base.py | 2 +-
python/tvm/contrib/coreml_runtime.py | 4 +-
python/tvm/contrib/cudnn.py | 4 +-
python/tvm/contrib/cutlass/_ffi_api.py | 4 +-
python/tvm/contrib/cutlass/build.py | 2 +-
python/tvm/contrib/cutlass/gen_tensor_op.py | 4 +-
python/tvm/contrib/hexagon/build.py | 2 +-
python/tvm/contrib/hexagon/tools.py | 2 +-
python/tvm/contrib/miopen.py | 4 +-
python/tvm/contrib/mrvl.py | 22 +--
python/tvm/contrib/msc/core/_ffi_api.py | 4 +-
python/tvm/contrib/msc/core/ir/graph.py | 16 +-
.../contrib/msc/framework/tensorflow/_ffi_api.py | 4 +-
.../tvm/contrib/msc/framework/tensorrt/_ffi_api.py | 4 +-
python/tvm/contrib/msc/framework/torch/_ffi_api.py | 4 +-
python/tvm/contrib/msc/framework/tvm/_ffi_api.py | 4 +-
python/tvm/contrib/msc/plugin/_ffi_api.py | 4 +-
python/tvm/contrib/msc/plugin/op/_ffi_api.py | 4 +-
python/tvm/contrib/ndk.py | 2 +-
python/tvm/contrib/nnpack.py | 4 +-
python/tvm/contrib/nvcc.py | 17 +-
python/tvm/contrib/random.py | 4 +-
python/tvm/contrib/rocm.py | 8 +-
python/tvm/contrib/tflite_runtime.py | 4 +-
python/tvm/contrib/thrust.py | 2 +-
python/tvm/dlight/analysis/common_analysis.py | 2 +-
python/tvm/driver/_ffi_api.py | 4 +-
python/tvm/error.py | 2 +-
python/tvm/exec/disco_worker.py | 2 +-
python/tvm/{relax/_ffi_api.py => ffi.py} | 7 +-
python/tvm/ir/_ffi_analysis_api.py | 4 +-
python/tvm/ir/_ffi_api.py | 4 +-
python/tvm/ir/_ffi_instrument_api.py | 4 +-
python/tvm/ir/_ffi_transform_api.py | 4 +-
python/tvm/ir/attrs.py | 6 +-
python/tvm/ir/base.py | 4 +-
python/tvm/ir/container.py | 2 +-
python/tvm/ir/diagnostics/__init__.py | 10 +-
python/tvm/ir/diagnostics/_ffi_api.py | 4 +-
python/tvm/ir/expr.py | 13 +-
python/tvm/ir/function.py | 4 +-
python/tvm/ir/global_info.py | 7 +-
python/tvm/ir/instrument.py | 6 +-
python/tvm/ir/module.py | 5 +-
python/tvm/ir/op.py | 4 +-
python/tvm/ir/supply.py | 5 +-
python/tvm/ir/transform.py | 12 +-
python/tvm/ir/type.py | 14 +-
python/tvm/ir/type_relation.py | 6 +-
python/tvm/libinfo.py | 10 +-
python/tvm/meta_schedule/_ffi_api.py | 2 +-
python/tvm/meta_schedule/arg_info.py | 2 +-
python/tvm/meta_schedule/builder/builder.py | 2 +-
python/tvm/meta_schedule/builder/local_builder.py | 2 +-
python/tvm/meta_schedule/cost_model/cost_model.py | 2 +-
python/tvm/meta_schedule/database/database.py | 2 +-
python/tvm/meta_schedule/database/json_database.py | 2 +-
.../tvm/meta_schedule/database/memory_database.py | 2 +-
.../database/ordered_union_database.py | 2 +-
.../meta_schedule/database/schedule_fn_database.py | 2 +-
.../tvm/meta_schedule/database/union_database.py | 2 +-
python/tvm/meta_schedule/extracted_task.py | 2 +-
.../feature_extractor/feature_extractor.py | 2 +-
.../feature_extractor/per_store_feature.py | 2 +-
.../measure_callback/add_to_database.py | 2 +-
.../measure_callback/measure_callback.py | 2 +-
.../measure_callback/remove_build_artifact.py | 2 +-
.../measure_callback/update_cost_model.py | 2 +-
.../mutator/mutate_compute_location.py | 2 +-
.../tvm/meta_schedule/mutator/mutate_parallel.py | 2 +-
.../meta_schedule/mutator/mutate_thread_binding.py | 2 +-
.../tvm/meta_schedule/mutator/mutate_tile_size.py | 2 +-
python/tvm/meta_schedule/mutator/mutate_unroll.py | 2 +-
python/tvm/meta_schedule/mutator/mutator.py | 2 +-
.../postproc/disallow_async_strided_mem_copy.py | 2 +-
.../postproc/disallow_dynamic_loop.py | 2 +-
python/tvm/meta_schedule/postproc/postproc.py | 2 +-
.../postproc/rewrite_cooperative_fetch.py | 2 +-
.../tvm/meta_schedule/postproc/rewrite_layout.py | 2 +-
.../postproc/rewrite_parallel_vectorize_unroll.py | 2 +-
.../postproc/rewrite_reduction_block.py | 2 +-
.../meta_schedule/postproc/rewrite_tensorize.py | 2 +-
.../postproc/rewrite_unbound_block.py | 2 +-
.../tvm/meta_schedule/postproc/verify_gpu_code.py | 2 +-
.../meta_schedule/postproc/verify_vtcm_limit.py | 2 +-
python/tvm/meta_schedule/profiler.py | 2 +-
python/tvm/meta_schedule/relax_integration.py | 2 +-
python/tvm/meta_schedule/runner/runner.py | 2 +-
.../tvm/meta_schedule/schedule_rule/add_rfactor.py | 2 +-
.../schedule_rule/apply_custom_rule.py | 2 +-
.../tvm/meta_schedule/schedule_rule/auto_bind.py | 2 +-
.../tvm/meta_schedule/schedule_rule/auto_inline.py | 2 +-
.../schedule_rule/cross_thread_reduction.py | 2 +-
.../schedule_rule/multi_level_tiling.py | 2 +-
.../schedule_rule/parallel_vectorize_unroll.py | 2 +-
.../schedule_rule/random_compute_location.py | 2 +-
.../meta_schedule/schedule_rule/schedule_rule.py | 2 +-
.../search_strategy/evolutionary_search.py | 2 +-
.../meta_schedule/search_strategy/replay_func.py | 2 +-
.../meta_schedule/search_strategy/replay_trace.py | 2 +-
.../search_strategy/search_strategy.py | 2 +-
.../space_generator/post_order_apply.py | 2 +-
.../meta_schedule/space_generator/schedule_fn.py | 2 +-
.../space_generator/space_generator.py | 2 +-
.../space_generator/space_generator_union.py | 2 +-
.../meta_schedule/task_scheduler/gradient_based.py | 2 +-
.../meta_schedule/task_scheduler/round_robin.py | 2 +-
.../meta_schedule/task_scheduler/task_scheduler.py | 2 +-
.../tvm/meta_schedule/testing/validate_database.py | 2 +-
python/tvm/meta_schedule/tir_integration.py | 2 +-
python/tvm/meta_schedule/tune_context.py | 2 +-
python/tvm/meta_schedule/utils.py | 2 +-
python/tvm/relax/_ffi_api.py | 4 +-
python/tvm/relax/analysis/_ffi_api.py | 4 +-
python/tvm/relax/backend/_ffi_api.py | 4 +-
python/tvm/relax/backend/cuda/flashinfer.py | 2 +-
python/tvm/relax/backend/metal/coreml.py | 4 +-
python/tvm/relax/binding_rewrite.py | 4 +-
python/tvm/relax/block_builder.py | 3 +-
python/tvm/relax/distributed/_ffi_api.py | 4 +-
python/tvm/relax/distributed/global_info.py | 4 +-
python/tvm/relax/distributed/struct_info.py | 8 +-
python/tvm/relax/distributed/transform/_ffi_api.py | 4 +-
python/tvm/relax/dpl/_ffi.py | 4 +-
python/tvm/relax/dpl/pattern.py | 2 +-
python/tvm/relax/dpl/rewrite.py | 2 +-
python/tvm/relax/exec_builder.py | 6 +-
python/tvm/relax/expr.py | 46 +++---
python/tvm/relax/expr_functor.py | 8 +-
python/tvm/relax/frontend/onnx/onnx_frontend.py | 1 +
python/tvm/relax/op/_ffi_api.py | 4 +-
python/tvm/relax/op/builtin/_ffi_api.py | 4 +-
python/tvm/relax/op/ccl/_ffi_api.py | 4 +-
python/tvm/relax/op/distributed/_ffi_api.py | 4 +-
python/tvm/relax/op/grad/_ffi_api.py | 4 +-
python/tvm/relax/op/image/_ffi_api.py | 4 +-
python/tvm/relax/op/memory/_ffi_api.py | 4 +-
python/tvm/relax/op/nn/_ffi_api.py | 4 +-
python/tvm/relax/op/op_attrs.py | 138 ++++++++--------
python/tvm/relax/op/vm/_ffi_api.py | 4 +-
python/tvm/relax/struct_info.py | 14 +-
python/tvm/relax/testing/transform.py | 3 +-
python/tvm/relax/training/_ffi_api.py | 4 +-
python/tvm/relax/training/utils.py | 2 +-
python/tvm/relax/transform/_ffi_api.py | 4 +-
python/tvm/relax/transform/transform.py | 10 +-
python/tvm/relax/ty.py | 10 +-
python/tvm/relax/utils.py | 3 +-
python/tvm/rpc/_ffi_api.py | 4 +-
python/tvm/rpc/client.py | 4 +-
python/tvm/rpc/server.py | 8 +-
python/tvm/runtime/__init__.py | 2 +-
python/tvm/runtime/_ffi_api.py | 4 +-
python/tvm/runtime/_ffi_node_api.py | 6 +-
python/tvm/runtime/container.py | 2 +-
python/tvm/runtime/device.py | 6 +-
python/tvm/runtime/disco/_ffi_api.py | 2 +-
python/tvm/runtime/disco/process_pool.py | 2 +-
python/tvm/runtime/disco/session.py | 2 +-
python/tvm/runtime/module.py | 2 +-
python/tvm/runtime/ndarray.py | 23 +--
python/tvm/runtime/object.py | 10 +-
python/tvm/runtime/object_generic.py | 2 +-
python/tvm/runtime/packed_func.py | 2 +-
python/tvm/runtime/profiling/_ffi_api.py | 2 +-
python/tvm/runtime/script_printer.py | 4 +-
python/tvm/runtime/support.py | 4 +-
python/tvm/runtime/vm.py | 2 +-
python/tvm/script/_ffi_api.py | 4 +-
python/tvm/script/ir_builder/_ffi_api.py | 4 +-
python/tvm/script/ir_builder/base.py | 2 +-
python/tvm/script/ir_builder/ir/_ffi_api.py | 4 +-
python/tvm/script/ir_builder/ir/frame.py | 2 +-
python/tvm/script/ir_builder/relax/_ffi_api.py | 4 +-
.../ir_builder/relax/distributed/_ffi_api.py | 4 +-
python/tvm/script/ir_builder/relax/frame.py | 2 +-
python/tvm/script/ir_builder/tir/_ffi_api.py | 4 +-
python/tvm/script/ir_builder/tir/frame.py | 2 +-
python/tvm/script/printer/_ffi_api.py | 4 +-
python/tvm/script/printer/doc.py | 4 +-
python/tvm/script/printer/doc_printer.py | 2 +-
python/tvm/support.py | 4 +-
python/tvm/target/_ffi_api.py | 4 +-
python/tvm/target/datatype.py | 6 +-
python/tvm/target/detect_target.py | 2 +-
python/tvm/target/target.py | 8 +-
python/tvm/target/virtual_device.py | 7 +-
python/tvm/target/x86.py | 2 +-
python/tvm/te/_ffi_api.py | 4 +-
python/tvm/te/operation.py | 2 +-
python/tvm/te/tensor.py | 16 +-
python/tvm/testing/_ffi_api.py | 4 +-
python/tvm/testing/attrs.py | 2 +-
python/tvm/testing/popen_pool.py | 8 +-
python/tvm/testing/utils.py | 2 +-
python/tvm/tir/_ffi_api.py | 4 +-
python/tvm/tir/analysis/_ffi_api.py | 4 +-
python/tvm/tir/block_dependence_info.py | 2 +-
python/tvm/tir/block_scope.py | 2 +-
python/tvm/tir/buffer.py | 7 +-
python/tvm/tir/data_layout.py | 6 +-
python/tvm/tir/expr.py | 72 ++++-----
python/tvm/tir/function.py | 8 +-
python/tvm/tir/functor.py | 11 +-
python/tvm/tir/op.py | 9 +-
python/tvm/tir/schedule/_ffi_api.py | 4 +-
python/tvm/tir/schedule/analysis.py | 6 +-
python/tvm/tir/schedule/instruction.py | 2 +-
python/tvm/tir/schedule/schedule.py | 2 +-
python/tvm/tir/schedule/state.py | 2 +-
python/tvm/tir/schedule/trace.py | 2 +-
python/tvm/tir/stmt.py | 36 ++---
python/tvm/tir/tensor_intrin/cuda.py | 2 +-
python/tvm/tir/transform/_ffi_api.py | 4 +-
python/tvm/tir/transform/function_pass.py | 4 +-
python/tvm/topi/cpp/cuda.py | 4 +-
python/tvm/topi/cpp/generic.py | 4 +-
python/tvm/topi/cpp/impl.py | 4 +-
python/tvm/topi/cpp/nn.py | 4 +-
python/tvm/topi/cpp/rocm.py | 4 +-
python/tvm/topi/cpp/utils.py | 4 +-
python/tvm/topi/cpp/vision/__init__.py | 4 +-
python/tvm/topi/cpp/vision/yolo.py | 4 +-
python/tvm/topi/cpp/x86.py | 4 +-
tests/python/codegen/test_gpu_codegen_allreduce.py | 2 +-
tests/python/disco/test_loader.py | 2 +-
tests/python/ir/test_container_structural_equal.py | 2 +-
tests/python/ir/test_ir_container.py | 2 +-
.../meta_schedule/test_meta_schedule_builder.py | 2 +-
.../test_meta_schedule_post_order_apply.py | 2 +-
.../meta_schedule/test_meta_schedule_runner.py | 2 +-
tests/python/relax/test_op_inspect.py | 2 +-
.../tir-base/test_tir_structural_equal_hash.py | 2 +-
.../test_tir_transform_inject_ptx_async_copy.py | 3 +-
.../python/tvmscript/test_tvmscript_parser_tir.py | 5 +-
.../tvmscript/test_tvmscript_printer_annotation.py | 2 +-
.../python/tvmscript/test_tvmscript_printer_doc.py | 2 +-
.../test_tvmscript_printer_structural_equal.py | 2 +-
.../test_tvmscript_printer_underlining.py | 2 +-
tests/scripts/task_python_adreno.sh | 4 +-
tests/scripts/task_python_arm_compute_library.sh | 4 +-
tests/scripts/task_python_docs.sh | 4 +-
tests/scripts/task_python_hexagon.sh | 4 +-
tests/scripts/task_python_integration.sh | 4 +-
tests/scripts/task_python_nightly.sh | 4 +-
tests/scripts/task_python_unittest.sh | 4 +-
tests/scripts/unity/task_python_relax.sh | 4 +-
321 files changed, 1622 insertions(+), 890 deletions(-)
diff --git a/.gitmodules b/.gitmodules
index a481df2438..f984d66a0d 100644
--- a/.gitmodules
+++ b/.gitmodules
@@ -1,14 +1,11 @@
[submodule "dmlc-core"]
path = 3rdparty/dmlc-core
url = https://github.com/dmlc/dmlc-core.git
-[submodule "dlpack"]
- path = 3rdparty/dlpack
- url = https://github.com/dmlc/dlpack.git
[submodule "3rdparty/rang"]
path = 3rdparty/rang
url = https://github.com/agauniyal/rang.git
[submodule "3rdparty/libbacktrace"]
- path = 3rdparty/libbacktrace
+ path = ffi/3rdparty/libbacktrace
url = https://github.com/tlc-pack/libbacktrace.git
[submodule "3rdparty/cutlass"]
path = 3rdparty/cutlass
diff --git a/3rdparty/dlpack b/3rdparty/dlpack
deleted file mode 160000
index 3ea601bb41..0000000000
--- a/3rdparty/dlpack
+++ /dev/null
@@ -1 +0,0 @@
-Subproject commit 3ea601bb413074c49a77c4ce3218bc08f8c4703c
diff --git a/CMakeLists.txt b/CMakeLists.txt
index d8d23f9035..a26016c830 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -501,7 +501,7 @@ if(NOT BUILD_DUMMY_LIBTVM)
$<TARGET_OBJECTS:tvm_libinfo_objs>
${TVM_RUNTIME_EXT_OBJS}
)
-
+ target_link_libraries(tvm PUBLIC tvm_ffi_shared)
else()
# dummy version of libtvm that can be used by downstream to specify dependencies
# the real runner still need a full version of libtvm
@@ -510,6 +510,7 @@ else()
$<TARGET_OBJECTS:tvm_libinfo_objs>
${TVM_RUNTIME_EXT_OBJS}
)
+ target_link_libraries(tvm PUBLIC tvm_ffi_shared)
endif()
target_include_directories(tvm PUBLIC "$<INSTALL_INTERFACE:${CMAKE_INSTALL_INCLUDEDIR}>")
@@ -519,7 +520,6 @@ if(BUILD_STATIC_RUNTIME)
add_library(tvm_runtime STATIC
$<TARGET_OBJECTS:tvm_runtime_objs>
$<TARGET_OBJECTS:tvm_libinfo_objs>
- $<TARGET_OBJECTS:tvm_ffi_objs>
${TVM_RUNTIME_EXT_OBJS}
)
set(NOTICE_MULTILINE
@@ -528,6 +528,7 @@ if(BUILD_STATIC_RUNTIME)
string(CONCAT NOTICE ${NOTICE_MULTILINE})
add_custom_command(TARGET tvm_runtime POST_BUILD
COMMAND ${CMAKE_COMMAND} -E cmake_echo_color --yellow --bold ${NOTICE})
+ target_link_libraries(tvm_runtime PUBLIC tvm_ffi_static)
else()
add_library(tvm_runtime SHARED
$<TARGET_OBJECTS:tvm_runtime_objs>
@@ -535,6 +536,7 @@ else()
${TVM_RUNTIME_EXT_OBJS}
)
set_property(TARGET tvm_runtime APPEND PROPERTY LINK_OPTIONS "${TVM_NO_UNDEFINED_SYMBOLS}")
+ target_link_libraries(tvm_runtime PUBLIC tvm_ffi_shared)
endif()
@@ -602,10 +604,6 @@ endif()
target_link_libraries(tvm PRIVATE ${TVM_RUNTIME_LINKER_LIBS})
target_link_libraries(tvm_runtime PRIVATE ${TVM_RUNTIME_LINKER_LIBS})
-target_link_libraries(tvm PUBLIC tvm_ffi_objs)
-target_link_libraries(tvm_runtime PUBLIC tvm_ffi_objs)
-
-
if(BUILD_FOR_HEXAGON AND DEFINED USE_HEXAGON_GTEST AND EXISTS ${USE_HEXAGON_GTEST})
include(FetchContent)
FetchContent_Declare(googletest SOURCE_DIR "${USE_HEXAGON_GTEST}")
@@ -643,7 +641,6 @@ if (HIDE_PRIVATE_SYMBOLS AND NOT ${CMAKE_SYSTEM_NAME} MATCHES "Darwin")
target_link_libraries(tvm_runtime PRIVATE ${HIDE_SYMBOLS_LINKER_FLAGS})
target_compile_definitions(tvm_allvisible PUBLIC $<TARGET_PROPERTY:tvm,INTERFACE_COMPILE_DEFINITONS>)
target_compile_definitions(tvm_allvisible PRIVATE $<TARGET_PROPERTY:tvm,COMPILE_DEFINITONS>)
- target_link_libraries(tvm_allvisible PUBLIC tvm_ffi_objs)
endif()
# Create the `cpptest` target if we can find GTest. If not, we create dummy
@@ -687,19 +684,6 @@ endif()
# Custom targets
add_custom_target(runtime DEPENDS tvm_runtime)
-# By default add cython to all build
-find_package(Python)
-if(NOT DEFINED ENV{CONDA_BUILD})
- message(STATUS ${CMAKE_CURRENT_BINARY_DIR})
- add_custom_target(
- tvm_cython ALL
- ${Python_EXECUTABLE} setup.py build_ext --inplace
- WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/python
- )
- add_dependencies(tvm_cython tvm)
- message("Add Cython build into the default build step")
-endif()
-
# Installation rules
install(TARGETS tvm DESTINATION lib${LIB_SUFFIX})
install(TARGETS tvm_runtime DESTINATION lib${LIB_SUFFIX})
@@ -779,8 +763,8 @@ if(TVM_IS_DEBUG_BUILD)
endif()
endif()
-add_dsymutil(tvm)
-add_dsymutil(tvm_runtime)
+tvm_ffi_add_dsymutil(tvm)
+tvm_ffi_add_dsymutil(tvm_runtime)
if(BUILD_FOR_HEXAGON)
# Wrap pthread_create to allow setting custom stack size.
diff --git a/Makefile b/Makefile
index ecc891ab76..8b9eddfbfe 100644
--- a/Makefile
+++ b/Makefile
@@ -107,16 +107,6 @@ mypy:
cppdoc:
doxygen docs/Doxyfile
-
-# Cython build
-cython cython3:
- cd python; python3 setup.py build_ext --inplace
-
-cyclean:
- rm -rf python/tvm/*/*/*.so python/tvm/*/*/*.dylib python/tvm/*/*/*.cpp
-
-
-
# EMCC; Web related scripts
web:
$(MAKE) -C $(ROOTDIR)/web
diff --git a/ci/jenkins/data.py b/ci/jenkins/data.py
index e52aaf32a4..3577a0ad00 100644
--- a/ci/jenkins/data.py
+++ b/ci/jenkins/data.py
@@ -30,7 +30,12 @@ files_to_stash = {
# runtime files
"tvm_runtime": ["build/libtvm_runtime.so", "build/config.cmake"],
# compiler files
- "tvm_lib": ["build/libtvm.so", "build/libtvm_runtime.so", "build/config.cmake"],
+ "tvm_lib": [
+ "build/libtvm.so",
+ "build/libtvm_runtime.so",
+ "build/lib/libtvm_ffi.so",
+ "build/config.cmake",
+ ],
# gpu related compiler files
"tvm_lib_gpu_extra": [
"build/3rdparty/libflash_attn/src/libflash_attn.so",
diff --git a/ci/jenkins/generated/arm_jenkinsfile.groovy b/ci/jenkins/generated/arm_jenkinsfile.groovy
index 9e4afc8f13..451c0ec139 100644
--- a/ci/jenkins/generated/arm_jenkinsfile.groovy
+++ b/ci/jenkins/generated/arm_jenkinsfile.groovy
@@ -60,7 +60,7 @@
// 'python3 jenkins/generate.py'
// Note: This timestamp is here to ensure that updates to the Jenkinsfile are
// always rebased on main before merging:
-// Generated at 2025-06-03T18:16:35.851073
+// Generated at 2025-08-23T12:27:21.419416
import org.jenkinsci.plugins.pipeline.modeldefinition.Utils
// These are set at runtime from data in ci/jenkins/docker-images.yml, update
@@ -516,7 +516,7 @@ def run_build(node_type) {
cmake_build(ci_arm, 'build')
make_cpp_tests(ci_arm, 'build')
sh(
- script: "./${jenkins_scripts_root}/s3.py --action upload --bucket ${s3_bucket} --prefix ${s3_prefix}/arm --items build/libtvm.so build/libtvm_runtime.so build/config.cmake build/cpptest build/build.ninja build/CMakeFiles/rules.ninja",
+ script: "./${jenkins_scripts_root}/s3.py --action upload --bucket ${s3_bucket} --prefix ${s3_prefix}/arm --items build/libtvm.so build/libtvm_runtime.so build/lib/libtvm_ffi.so build/config.cmake build/cpptest build/build.ninja build/CMakeFiles/rules.ninja",
label: 'Upload artifacts to S3',
)
})
diff --git a/ci/jenkins/generated/cpu_jenkinsfile.groovy b/ci/jenkins/generated/cpu_jenkinsfile.groovy
index daadc16c76..fa5775b786 100644
--- a/ci/jenkins/generated/cpu_jenkinsfile.groovy
+++ b/ci/jenkins/generated/cpu_jenkinsfile.groovy
@@ -60,7 +60,7 @@
// 'python3 jenkins/generate.py'
// Note: This timestamp is here to ensure that updates to the Jenkinsfile are
// always rebased on main before merging:
-// Generated at 2025-06-03T18:16:35.861918
+// Generated at 2025-08-23T12:27:21.435503
import org.jenkinsci.plugins.pipeline.modeldefinition.Utils
// These are set at runtime from data in ci/jenkins/docker-images.yml, update
@@ -516,7 +516,7 @@ def run_build(node_type) {
cmake_build(ci_cpu, 'build')
make_cpp_tests(ci_cpu, 'build')
sh(
- script: "./${jenkins_scripts_root}/s3.py --action upload --bucket ${s3_bucket} --prefix ${s3_prefix}/cpu --items build/libtvm.so build/libtvm_runtime.so build/config.cmake build/libtvm_allvisible.so build/cpptest build/build.ninja build/CMakeFiles/rules.ninja",
+ script: "./${jenkins_scripts_root}/s3.py --action upload --bucket ${s3_bucket} --prefix ${s3_prefix}/cpu --items build/libtvm.so build/libtvm_runtime.so build/lib/libtvm_ffi.so build/config.cmake build/libtvm_allvisible.so build/cpptest build/build.ninja build/CMakeFiles/rules.ninja",
label: 'Upload artifacts to S3',
)
})
diff --git a/ci/jenkins/generated/gpu_jenkinsfile.groovy b/ci/jenkins/generated/gpu_jenkinsfile.groovy
index 1fc4348c6f..fe4cb9c46a 100644
--- a/ci/jenkins/generated/gpu_jenkinsfile.groovy
+++ b/ci/jenkins/generated/gpu_jenkinsfile.groovy
@@ -60,7 +60,7 @@
// 'python3 jenkins/generate.py'
// Note: This timestamp is here to ensure that updates to the Jenkinsfile are
// always rebased on main before merging:
-// Generated at 2025-06-03T18:16:35.885417
+// Generated at 2025-08-23T12:27:21.381047
import org.jenkinsci.plugins.pipeline.modeldefinition.Utils
// These are set at runtime from data in ci/jenkins/docker-images.yml, update
@@ -512,7 +512,7 @@ def run_build(node_type) {
sh "${docker_run} --no-gpu ${ci_gpu} ./tests/scripts/task_config_build_gpu.sh build"
cmake_build("${ci_gpu} --no-gpu", 'build')
sh(
- script: "./${jenkins_scripts_root}/s3.py --action upload --bucket ${s3_bucket} --prefix ${s3_prefix}/gpu --items build/libtvm.so build/libtvm_runtime.so build/config.cmake build/libtvm_allvisible.so build/3rdparty/libflash_attn/src/libflash_attn.so build/3rdparty/cutlass_fpA_intB_gemm/cutlass_kernels/libfpA_intB_gemm.so",
+ script: "./${jenkins_scripts_root}/s3.py --action upload --bucket ${s3_bucket} --prefix ${s3_prefix}/gpu --items build/libtvm.so build/libtvm_runtime.so build/lib/libtvm_ffi.so build/config.cmake build/libtvm_allvisible.so build/3rdparty/libflash_attn/src/libflash_attn.so build/3rdparty/cutlass_fpA_intB_gemm/cutlass_kernels/libfpA_intB_gemm.so",
label: 'Upload artifacts to S3',
)
@@ -522,7 +522,7 @@ def run_build(node_type) {
sh "${docker_run} --no-gpu ${ci_gpu} ./tests/scripts/task_config_build_gpu_other.sh build"
cmake_build("${ci_gpu} --no-gpu", 'build')
sh(
- script: "./${jenkins_scripts_root}/s3.py --action upload --bucket ${s3_bucket} --prefix ${s3_prefix}/gpu2 --items build/libtvm.so build/libtvm_runtime.so build/config.cmake",
+ script: "./${jenkins_scripts_root}/s3.py --action upload --bucket ${s3_bucket} --prefix ${s3_prefix}/gpu2 --items build/libtvm.so build/libtvm_runtime.so build/lib/libtvm_ffi.so build/config.cmake",
label: 'Upload artifacts to S3',
)
})
diff --git a/ci/jenkins/generated/hexagon_jenkinsfile.groovy b/ci/jenkins/generated/hexagon_jenkinsfile.groovy
index 173506fcce..233e8fe508 100644
--- a/ci/jenkins/generated/hexagon_jenkinsfile.groovy
+++ b/ci/jenkins/generated/hexagon_jenkinsfile.groovy
@@ -60,7 +60,7 @@
// 'python3 jenkins/generate.py'
// Note: This timestamp is here to ensure that updates to the Jenkinsfile are
// always rebased on main before merging:
-// Generated at 2025-06-03T18:16:35.839798
+// Generated at 2025-08-23T12:27:21.325574
import org.jenkinsci.plugins.pipeline.modeldefinition.Utils
// These are set at runtime from data in ci/jenkins/docker-images.yml, update
@@ -520,7 +520,7 @@ def run_build(node_type) {
label: 'Build Hexagon API',
)
sh(
- script: "./${jenkins_scripts_root}/s3.py --action upload --bucket ${s3_bucket} --prefix ${s3_prefix}/hexagon --items build/libtvm.so build/libtvm_runtime.so build/config.cmake build/cpptest build/build.ninja build/CMakeFiles/rules.ninja build/hexagon_api_output",
+ script: "./${jenkins_scripts_root}/s3.py --action upload --bucket ${s3_bucket} --prefix ${s3_prefix}/hexagon --items build/libtvm.so build/libtvm_runtime.so build/lib/libtvm_ffi.so build/config.cmake build/cpptest build/build.ninja build/CMakeFiles/rules.ninja build/hexagon_api_output",
label: 'Upload artifacts to S3',
)
})
diff --git a/ci/jenkins/generated/i386_jenkinsfile.groovy b/ci/jenkins/generated/i386_jenkinsfile.groovy
index 3ef2b532ba..9bcf9588f1 100644
--- a/ci/jenkins/generated/i386_jenkinsfile.groovy
+++ b/ci/jenkins/generated/i386_jenkinsfile.groovy
@@ -60,7 +60,7 @@
// 'python3 jenkins/generate.py'
// Note: This timestamp is here to ensure that updates to the Jenkinsfile are
// always rebased on main before merging:
-// Generated at 2025-06-03T18:16:35.814567
+// Generated at 2025-08-23T12:27:21.401370
import org.jenkinsci.plugins.pipeline.modeldefinition.Utils
// These are set at runtime from data in ci/jenkins/docker-images.yml, update
@@ -516,7 +516,7 @@ def run_build(node_type) {
cmake_build(ci_i386, 'build')
make_cpp_tests(ci_i386, 'build')
sh(
- script: "./${jenkins_scripts_root}/s3.py --action upload --bucket ${s3_bucket} --prefix ${s3_prefix}/i386 --items build/libtvm.so build/libtvm_runtime.so build/config.cmake build/cpptest build/build.ninja build/CMakeFiles/rules.ninja",
+ script: "./${jenkins_scripts_root}/s3.py --action upload --bucket ${s3_bucket} --prefix ${s3_prefix}/i386 --items build/libtvm.so build/libtvm_runtime.so build/lib/libtvm_ffi.so build/config.cmake build/cpptest build/build.ninja build/CMakeFiles/rules.ninja",
label: 'Upload artifacts to S3',
)
})
diff --git a/docs/arch/pass_infra.rst b/docs/arch/pass_infra.rst
index 4bf3abceb0..30e28d20db 100644
--- a/docs/arch/pass_infra.rst
+++ b/docs/arch/pass_infra.rst
@@ -552,7 +552,7 @@ a certain scope.
.. code:: python
- @tvm.ffi.register_object("transform.PassContext")
+ @tvm_ffi.register_object("transform.PassContext")
class PassContext(tvm.runtime.Object):
def __enter__(self):
_transform.EnterPassContext(self)
diff --git a/ffi/.clang-format b/ffi/.clang-format
new file mode 100644
index 0000000000..9d622b98ba
--- /dev/null
+++ b/ffi/.clang-format
@@ -0,0 +1,8 @@
+# Run the following command to reformat a file:
+# clang-format -i -style=Google <file>
+# Or use clang-format-diff to only reformat the changed lines:
+# https://clang.llvm.org/docs/ClangFormat.html
+BasedOnStyle: Google
+DerivePointerAlignment: false
+ColumnLimit: 100
+PointerAlignment: Left
diff --git a/3rdparty/libbacktrace b/ffi/3rdparty/libbacktrace
similarity index 100%
rename from 3rdparty/libbacktrace
rename to ffi/3rdparty/libbacktrace
diff --git a/ffi/CMakeLists.txt b/ffi/CMakeLists.txt
index 466571c288..d4f600a27d 100644
--- a/ffi/CMakeLists.txt
+++ b/ffi/CMakeLists.txt
@@ -14,40 +14,41 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-
-cmake_minimum_required(VERSION 3.14)
+cmake_minimum_required(VERSION 3.18)
project(
tvm_ffi
- VERSION 1.0
- DESCRIPTION "TVM's FFI system"
LANGUAGES CXX C
)
-option(TVM_FFI_BUILD_TESTS "Adding test targets." OFF)
option(TVM_FFI_USE_LIBBACKTRACE "Enable libbacktrace" ON)
option(TVM_FFI_USE_EXTRA_CXX_API "Enable extra CXX API in shared lib" ON)
option(TVM_FFI_BACKTRACE_ON_SEGFAULT "Set signal handler to print traceback on segfault" ON)
-include(cmake/Utils/CxxWarning.cmake)
-include(cmake/Utils/Sanitizer.cmake)
-include(cmake/Utils/Library.cmake)
if (TVM_FFI_USE_LIBBACKTRACE)
- include(cmake/Utils/AddLibbacktrace.cmake)
+ include(${CMAKE_CURRENT_LIST_DIR}/cmake/Utils/AddLibbacktrace.cmake)
endif()
-########## Target: `dlpack_header` ##########
+include(${CMAKE_CURRENT_LIST_DIR}/cmake/Utils/Library.cmake)
-add_library(dlpack_header INTERFACE)
-target_include_directories(dlpack_header INTERFACE "${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/dlpack/include")
########## Target: `tvm_ffi_header` ##########
+# they can be used in cases where user do not want to link into the library
+# in cases like deferred linking
add_library(tvm_ffi_header INTERFACE)
-target_include_directories(tvm_ffi_header INTERFACE "${CMAKE_CURRENT_SOURCE_DIR}/include")
-target_link_libraries(tvm_ffi_header INTERFACE dlpack_header)
+target_include_directories(
+ tvm_ffi_header INTERFACE
+ $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>
+ $<INSTALL_INTERFACE:include>
+)
+target_include_directories(
+ tvm_ffi_header INTERFACE
+ $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/dlpack/include>
+ $<INSTALL_INTERFACE:include>
+)
-########## Target: `tvm_ffi` ##########
+########## Target: `tvm_ffi_objs` ##########
set(tvm_ffi_objs_sources
"${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/traceback.cc"
@@ -60,22 +61,23 @@ set(tvm_ffi_objs_sources
"${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/container.cc"
)
+set(tvm_ffi_extra_objs_sources
+ "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/structural_equal.cc"
+ "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/structural_hash.cc"
+ "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/json_parser.cc"
+ "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/json_writer.cc"
+ "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/serialization.cc"
+ "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/reflection_extra.cc"
+ "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/module.cc"
+ "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/library_module.cc"
+ "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/library_module_system_lib.cc"
+ "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/library_module_dynamic_lib.cc"
+ "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/stream_context.cc"
+ "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/env_c_api.cc"
+ "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/testing.cc"
+)
if (TVM_FFI_USE_EXTRA_CXX_API)
- list(APPEND tvm_ffi_objs_sources
- "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/structural_equal.cc"
- "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/structural_hash.cc"
- "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/json_parser.cc"
- "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/json_writer.cc"
- "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/serialization.cc"
- "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/reflection_extra.cc"
- "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/module.cc"
- "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/library_module.cc"
- "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/library_module_system_lib.cc"
- "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/library_module_dynamic_lib.cc"
- "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/stream_context.cc"
- "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/env_c_api.cc"
- "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/testing.cc"
- )
+ list(APPEND tvm_ffi_objs_sources ${tvm_ffi_extra_objs_sources})
endif()
add_library(tvm_ffi_objs OBJECT ${tvm_ffi_objs_sources})
@@ -90,9 +92,9 @@ set_target_properties(
VISIBILITY_INLINES_HIDDEN ON
PREFIX "lib"
)
-add_cxx_warning(tvm_ffi_objs)
-target_link_libraries(tvm_ffi_objs PRIVATE dlpack_header)
-target_include_directories(tvm_ffi_objs PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/include")
+
+# add the include path as public so they are visible to downstreams
+target_link_libraries(tvm_ffi_objs PUBLIC tvm_ffi_header)
if (TVM_FFI_USE_LIBBACKTRACE)
message(STATUS "Setting C++ macro TVM_FFI_USE_LIBBACKTRACE - 1")
@@ -110,7 +112,9 @@ else()
target_compile_definitions(tvm_ffi_objs PRIVATE TVM_FFI_BACKTRACE_ON_SEGFAULT=0)
endif()
-add_target_from_obj(tvm_ffi tvm_ffi_objs)
+tvm_ffi_add_msvc_flags(tvm_ffi_objs)
+
+tvm_ffi_add_target_from_obj(tvm_ffi tvm_ffi_objs)
if (TARGET libbacktrace)
target_link_libraries(tvm_ffi_objs PRIVATE libbacktrace)
@@ -124,22 +128,107 @@ if (MSVC)
target_link_libraries(tvm_ffi_static PRIVATE DbgHelp.lib)
endif ()
+# expose the headers as public dependencies
target_link_libraries(tvm_ffi_objs PUBLIC tvm_ffi_header)
target_link_libraries(tvm_ffi_shared PUBLIC tvm_ffi_header)
target_link_libraries(tvm_ffi_static PUBLIC tvm_ffi_header)
-install(TARGETS tvm_ffi_static DESTINATION lib${LIB_SUFFIX})
-install(TARGETS tvm_ffi_shared DESTINATION lib${LIB_SUFFIX})
+#----------------------------------------------------------------------------
+# The following code section only is triggered when the project is the root
+# and will be skipped when the project is a subproject.
+#----------------------------------------------------------------------------
+if (NOT ${PROJECT_NAME} STREQUAL ${CMAKE_PROJECT_NAME})
+ return()
+endif()
+
+option(TVM_FFI_USE_RELEASE_FLAGS "Set O3 with RelWithDebInfo, useful to get lineno" OFF)
+option(TVM_FFI_BUILD_TESTS "Adding test targets." OFF)
+
+# related options
+if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang")
+ if (TVM_FFI_USE_RELEASE_FLAGS)
+ message(STATUS "Targeting global CXX flags to -O3 for optimized build")
+ target_compile_options(tvm_ffi_objs PRIVATE -O3)
+ target_compile_options(tvm_ffi_shared PRIVATE -O3)
+ target_compile_options(tvm_ffi_static PRIVATE -O3)
+ endif()
+endif()
+
+
+include(cmake/Utils/CxxWarning.cmake)
+include(cmake/Utils/Sanitizer.cmake)
-add_msvc_flags(tvm_ffi_objs)
+tvm_ffi_add_cxx_warning(tvm_ffi_objs)
-########## Adding tests ##########
+########## Adding cpp tests ##########
+
+# logics below are only executed when the project is the root project.
+# but not when the project is a subproject.
+if (TVM_FFI_BUILD_TESTS)
+ enable_testing()
+ message(STATUS "Enable Testing")
+ include(cmake/Utils/AddGoogleTest.cmake)
+ add_subdirectory(tests/cpp/)
+endif()
-if (${PROJECT_NAME} STREQUAL ${CMAKE_PROJECT_NAME})
- if (TVM_FFI_BUILD_TESTS)
- enable_testing()
- message(STATUS "Enable Testing")
- include(cmake/Utils/AddGoogleTest.cmake)
- add_subdirectory(tests/cpp/)
+########## Install the related for normal cmake library ##########
+
+install(DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/include/tvm/ffi/ DESTINATION include/tvm/ffi/)
+install(DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/dlpack/include/ DESTINATION include/)
+install(TARGETS tvm_ffi_shared tvm_ffi_static DESTINATION lib)
+
+########## Adding python module ##########
+option(TVM_FFI_BUILD_PYTHON_MODULE "Adding python module." OFF)
+
+if (TVM_FFI_BUILD_PYTHON_MODULE)
+ message(STATUS "Building python module")
+ find_package(
+ Python COMPONENTS Interpreter Development.Module Development.SABIModule
+ REQUIRED)
+ set(core_cpp ${CMAKE_CURRENT_BINARY_DIR}/core.cpp)
+ set(core_pyx ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/core.pyx)
+ set(cython_sources
+ ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/core.pyx
+ ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/base.pxi
+ ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/device.pxi
+ ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/dtype.pxi
+ ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/error.pxi
+ ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/function.pxi
+ ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/ndarray.pxi
+ ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/object.pxi
+ ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/string.pxi
+ )
+ add_custom_command(
+ OUTPUT ${core_cpp}
+ COMMAND ${Python_EXECUTABLE} -m cython --cplus ${core_pyx} -o ${core_cpp}
+ COMMENT "Transpiling cython to cpp"
+ DEPENDS ${cython_sources}
+ VERBATIM
+ )
+ message(STATUS "Transling cython to cpp: ${CMAKE_CURRENT_BINARY_DIR}/core.cpp")
+ if(Python_VERSION VERSION_GREATER_EQUAL "3.12")
+ # >= Python3.12, use Use_SABI version
+ Python_add_library(tvm_ffi_cython MODULE "${core_cpp}" USE_SABI 3.12)
+ set_target_properties(tvm_ffi_cython PROPERTIES OUTPUT_NAME "core")
+ set_target_properties(tvm_ffi_cython PROPERTIES SUFFIX ".abi3.so")
+ else()
+ # before Python3.12, use WITH_SOABI version
+ Python_add_library(tvm_ffi_cython MODULE "${core_cpp}" WITH_SOABI)
+ set_target_properties(tvm_ffi_cython PROPERTIES OUTPUT_NAME "core")
endif()
-endif ()
+
+ target_link_libraries(tvm_ffi_cython PRIVATE tvm_ffi_header)
+ target_link_libraries(tvm_ffi_cython PRIVATE tvm_ffi_shared)
+
+ ########## Installing the source ##########
+ install(
+ DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/ DESTINATION 3rdparty/
+ PATTERN ".git" EXCLUDE
+ PATTERN ".git*" EXCLUDE
+ PATTERN "*.tmp" EXCLUDE
+ )
+ install(DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/ DESTINATION src/ffi/)
+ install(DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/cmake/Utils/ DESTINATION cmake/Utils/)
+ install(FILES ${CMAKE_CURRENT_LIST_DIR}/CMakeLists.txt DESTINATION .)
+ install(FILES ${CMAKE_CURRENT_SOURCE_DIR}/cmake/tvm_ffi-config.cmake DESTINATION lib/cmake/tvm_ffi/)
+endif()
diff --git a/ffi/README.md b/ffi/README.md
new file mode 100644
index 0000000000..3b1b1199c2
--- /dev/null
+++ b/ffi/README.md
@@ -0,0 +1,18 @@
+<!--- 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. -->
+
+# tvm ffi
diff --git a/ffi/cmake/Utils/AddGoogleTest.cmake b/ffi/cmake/Utils/AddGoogleTest.cmake
index 10e5938612..85e21ced1b 100644
--- a/ffi/cmake/Utils/AddGoogleTest.cmake
+++ b/ffi/cmake/Utils/AddGoogleTest.cmake
@@ -42,7 +42,7 @@ if (NOT googletest_POPULATED)
)
endif()
-macro(add_googletest target_name)
+macro(tvm_ffi_add_googletest target_name)
add_test(
NAME ${target_name}
COMMAND ${target_name}
diff --git a/ffi/cmake/Utils/AddLibbacktrace.cmake b/ffi/cmake/Utils/AddLibbacktrace.cmake
index 844a8816a6..e920a1f199 100644
--- a/ffi/cmake/Utils/AddLibbacktrace.cmake
+++ b/ffi/cmake/Utils/AddLibbacktrace.cmake
@@ -18,7 +18,7 @@
include(ExternalProject)
function(_libbacktrace_compile)
- set(_libbacktrace_source ${CMAKE_CURRENT_LIST_DIR}/../../../3rdparty/libbacktrace)
+ set(_libbacktrace_source ${CMAKE_CURRENT_LIST_DIR}/../../3rdparty/libbacktrace)
set(_libbacktrace_prefix ${CMAKE_CURRENT_BINARY_DIR}/libbacktrace)
if(CMAKE_SYSTEM_NAME MATCHES "Darwin" AND (CMAKE_C_COMPILER MATCHES "^/Library" OR CMAKE_C_COMPILER MATCHES "^/Applications"))
set(_cmake_c_compiler "/usr/bin/cc")
@@ -36,6 +36,7 @@ function(_libbacktrace_compile)
SOURCE_DIR ${_libbacktrace_source}
BINARY_DIR ${_libbacktrace_prefix}
CONFIGURE_COMMAND
+ "sh"
"${_libbacktrace_source}/configure"
"--prefix=${_libbacktrace_prefix}"
--with-pic
diff --git a/ffi/cmake/Utils/CxxWarning.cmake b/ffi/cmake/Utils/CxxWarning.cmake
index c272bfdf7b..a85e58825b 100644
--- a/ffi/cmake/Utils/CxxWarning.cmake
+++ b/ffi/cmake/Utils/CxxWarning.cmake
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
-function(add_cxx_warning target_name)
+function(tvm_ffi_add_cxx_warning target_name)
# GNU, Clang, or AppleClang
if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang|AppleClang")
target_compile_options(${target_name} PRIVATE "-Werror" "-Wall" "-Wextra" "-Wpedantic" "-Wno-unused-parameter")
diff --git a/ffi/cmake/Utils/Library.cmake b/ffi/cmake/Utils/Library.cmake
index cff7ca35a2..d9ca19ec1c 100644
--- a/ffi/cmake/Utils/Library.cmake
+++ b/ffi/cmake/Utils/Library.cmake
@@ -14,7 +14,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-function(add_dsymutil target_name)
+function(tvm_ffi_add_dsymutil target_name)
# running dsymutil on macos to generate debugging symbols for backtraces
if(APPLE AND TVM_FFI_USE_LIBBACKTRACE)
find_program(DSYMUTIL dsymutil)
@@ -28,7 +28,7 @@ function(add_dsymutil target_name)
endif()
endfunction()
-function(add_msvc_flags target_name)
+function(tvm_ffi_add_msvc_flags target_name)
# running if we are under msvc
if(MSVC)
target_compile_definitions(${target_name} PUBLIC -DWIN32_LEAN_AND_MEAN)
@@ -40,7 +40,7 @@ function(add_msvc_flags target_name)
endif()
endfunction()
-function(add_target_from_obj target_name obj_target_name)
+function(tvm_ffi_add_target_from_obj target_name obj_target_name)
add_library(${target_name}_static STATIC $<TARGET_OBJECTS:${obj_target_name}>)
set_target_properties(
${target_name}_static PROPERTIES
@@ -57,11 +57,9 @@ function(add_target_from_obj target_name obj_target_name)
ARCHIVE_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib"
LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib"
)
- add_custom_target(${target_name})
- add_dependencies(${target_name} ${target_name}_static ${target_name}_shared)
if (MSVC)
target_compile_definitions(${obj_target_name} PRIVATE TVM_FFI_EXPORTS)
endif()
- add_dsymutil(${target_name}_shared)
- add_msvc_flags(${target_name}_shared)
+ tvm_ffi_add_dsymutil(${target_name}_shared)
+ tvm_ffi_add_msvc_flags(${target_name}_shared)
endfunction()
diff --git a/ffi/cmake/tvm_ffi-config.cmake b/ffi/cmake/tvm_ffi-config.cmake
new file mode 100644
index 0000000000..8c947e6c51
--- /dev/null
+++ b/ffi/cmake/tvm_ffi-config.cmake
@@ -0,0 +1,41 @@
+# 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.
+
+find_package(Python COMPONENTS Interpreter REQUIRED)
+
+# call tvm_ffi.config to get the cmake directory and set it to tvm_ffi_ROOT
+execute_process(
+ COMMAND "${Python_EXECUTABLE}" -m tvm_ffi.config --includedir
+ OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE tvm_ffi_INCLUDE_DIR)
+
+execute_process(
+ COMMAND "${Python_EXECUTABLE}" -m tvm_ffi.config --dlpack-includedir
+ OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE tvm_ffi_DLPACK_INCLUDE_DIR)
+
+execute_process(
+ COMMAND "${Python_EXECUTABLE}" -m tvm_ffi.config --libfiles
+ OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE tvm_ffi_LIB_FILES)
+
+add_library(tvm_ffi_header INTERFACE)
+target_include_directories(tvm_ffi_header INTERFACE "${tvm_ffi_INCLUDE_DIR}")
+target_include_directories(tvm_ffi_header INTERFACE "${tvm_ffi_DLPACK_INCLUDE_DIR}")
+
+add_library(tvm_ffi_shared SHARED IMPORTED)
+set_target_properties(
+ tvm_ffi_shared PROPERTIES IMPORTED_LOCATION "${tvm_ffi_LIB_FILES}"
+ INTERFACE_INCLUDE_DIRECTORIES "${tvm_ffi_INCLUDE_DIR};${tvm_ffi_DLPACK_INCLUDE_DIR}"
+)
diff --git a/ffi/examples/get_started/CMakeLists.txt b/ffi/examples/get_started/CMakeLists.txt
new file mode 100644
index 0000000000..8bd3a65e3f
--- /dev/null
+++ b/ffi/examples/get_started/CMakeLists.txt
@@ -0,0 +1,50 @@
+# 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.
+
+cmake_minimum_required(VERSION 3.18)
+project(tvm_ffi_example)
+
+
+# first find python related components
+find_package(Python COMPONENTS Interpreter REQUIRED)
+
+# call tvm_ffi.config to get the cmake directory and set it to tvm_ffi_ROOT
+execute_process(
+ COMMAND "${Python_EXECUTABLE}" -m tvm_ffi.config --cmakedir
+ OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE tvm_ffi_ROOT)
+# find package will automatically include the related projects
+find_package(tvm_ffi CONFIG REQUIRED)
+
+# use the projects as usual
+add_library(add_one_cpu SHARED src/add_one_cpu.cc)
+target_link_libraries(add_one_cpu tvm_ffi_header)
+target_link_libraries(add_one_cpu ${CMAKE_BINARY_DIR}/lib/libtvm_ffi.so)
+# show as add_one_cpu.so
+set_target_properties(add_one_cpu PROPERTIES PREFIX "")
+
+# Check if CUDA is available
+find_package(CUDA QUIET)
+if(CUDA_FOUND)
+ enable_language(CUDA)
+ add_library(add_one_cuda SHARED src/add_one_cuda.cu)
+ target_link_libraries(add_one_cuda tvm_ffi_shared)
+ # show as add_one_cuda.so
+ set_target_properties(add_one_cuda PROPERTIES PREFIX "")
+endif()
+
+add_executable(run_example src/run_example.cc)
+target_link_libraries(run_example tvm_ffi_shared)
diff --git a/ffi/examples/get_started/README.md b/ffi/examples/get_started/README.md
new file mode 100644
index 0000000000..934b7ca552
--- /dev/null
+++ b/ffi/examples/get_started/README.md
@@ -0,0 +1,4 @@
+# Example python binding
+
+Example project that uses the tvm ffi to expose global
+function definitions
diff --git a/ffi/examples/get_started/run_example.py b/ffi/examples/get_started/run_example.py
new file mode 100644
index 0000000000..08debfd68e
--- /dev/null
+++ b/ffi/examples/get_started/run_example.py
@@ -0,0 +1,74 @@
+# 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.
+import tvm_ffi
+import torch
+import numpy
+import ctypes
+
+
+def run_add_one_cpu():
+ """Load the add_one_cpu module and call the add_one_cpu function."""
+ mod = tvm_ffi.load_module("build/add_one_cpu.so")
+ lib = ctypes.CDLL("build/add_one_cpu.so", ctypes.RTLD_GLOBAL)
+ x = torch.tensor([1, 2, 3, 4, 5], dtype=torch.float32)
+ y = torch.empty_like(x)
+ # tvm-ffi automatically handles DLPack compatible tensors
+ # torch tensors can be viewed as ffi::NDArray or DLTensor*
+ # in the background
+ mod.add_one_cpu(x, y)
+ print("torch.result after add_one(x, y)")
+ print(y)
+
+ x = numpy.array([1, 2, 3, 4, 5], dtype=numpy.float32)
+ y = numpy.empty_like(x)
+ # tvm-ffi automatically handles DLPack compatible tensors
+ # torch tensors can be viewed as ffi::NDArray or DLTensor*
+ # in the background
+ mod.add_one_cpu(x, y)
+ print("numpy.result after add_one(x, y)")
+ print(x)
+
+
+def run_add_one_cuda():
+ """Load the add_one_cuda module and call the add_one_cuda function."""
+ if not torch.cuda.is_available():
+ return
+
+ mod = tvm_ffi.load_module("build/add_one_cuda.so")
+ x = torch.tensor([1, 2, 3, 4, 5], dtype=torch.float32, device="cuda")
+ y = torch.empty_like(x)
+
+ stream = torch.cuda.Stream()
+ with torch.cuda.stream(stream):
+ # tvm-ffi automatically handles DLPack compatible tensors
+ # it also handles interactions with torch runtime
+ # torch.cuda.current_stream() will be set and available via TVMFFIEnvGetCurrentStream
+ # when calling the function
+ mod.add_one_cuda(x, y)
+ stream.synchronize()
+ print("torch.result after mod.add_one_cuda(x, y)")
+ print(y)
+
+
+def main():
+ """Main function to run the example."""
+ run_add_one_cpu()
+ run_add_one_cuda()
+
+
+if __name__ == "__main__":
+ main()
diff --git a/python/tvm/ir/_ffi_analysis_api.py b/ffi/examples/get_started/run_example.sh
old mode 100644
new mode 100755
similarity index 84%
copy from python/tvm/ir/_ffi_analysis_api.py
copy to ffi/examples/get_started/run_example.sh
index ca38c2309f..0602b85f37
--- a/python/tvm/ir/_ffi_analysis_api.py
+++ b/ffi/examples/get_started/run_example.sh
@@ -14,9 +14,14 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""FFI APIs for tvm.ir.analysis"""
+#!/bin/bash
+set -ex
-import tvm.ffi
+cmake -B build -S .
+cmake --build build
+# running python example
+python run_example.py
-tvm.ffi._init_api("ir.analysis", __name__)
+# running c++ example
+./build/run_example
diff --git a/ffi/examples/get_started/src/add_one_cpu.cc b/ffi/examples/get_started/src/add_one_cpu.cc
new file mode 100644
index 0000000000..3b7050b5c0
--- /dev/null
+++ b/ffi/examples/get_started/src/add_one_cpu.cc
@@ -0,0 +1,43 @@
+
+
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+#include <tvm/ffi/dtype.h>
+#include <tvm/ffi/error.h>
+#include <tvm/ffi/function.h>
+
+namespace tvm_ffi_example {
+
+void AddOne(DLTensor* x, DLTensor* y) {
+ // implementation of a library function
+ TVM_FFI_ICHECK(x->ndim == 1) << "x must be a 1D tensor";
+ DLDataType f32_dtype{kDLFloat, 32, 1};
+ TVM_FFI_ICHECK(x->dtype == f32_dtype) << "x must be a float tensor";
+ TVM_FFI_ICHECK(y->ndim == 1) << "y must be a 1D tensor";
+ TVM_FFI_ICHECK(y->dtype == f32_dtype) << "y must be a float tensor";
+ TVM_FFI_ICHECK(x->shape[0] == y->shape[0]) << "x and y must have the same shape";
+ for (int i = 0; i < x->shape[0]; ++i) {
+ static_cast<float*>(y->data)[i] = static_cast<float*>(x->data)[i] + 1;
+ }
+}
+
+// Expose global symbol `add_one_cpu` that follows tvm-ffi abi
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(add_one_cpu, tvm_ffi_example::AddOne);
+} // namespace tvm_ffi_example
diff --git a/ffi/examples/get_started/src/add_one_cuda.cu b/ffi/examples/get_started/src/add_one_cuda.cu
new file mode 100644
index 0000000000..4a904b4812
--- /dev/null
+++ b/ffi/examples/get_started/src/add_one_cuda.cu
@@ -0,0 +1,59 @@
+
+
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+#include <tvm/ffi/dtype.h>
+#include <tvm/ffi/error.h>
+#include <tvm/ffi/extra/c_env_api.h>
+#include <tvm/ffi/function.h>
+
+namespace tvm_ffi_example {
+
+__global__ void AddOneKernel(float* x, float* y, int n) {
+ int idx = blockIdx.x * blockDim.x + threadIdx.x;
+ if (idx < n) {
+ y[idx] = x[idx] + 1;
+ }
+}
+
+void AddOneCUDA(DLTensor* x, DLTensor* y) {
+ // implementation of a library function
+ TVM_FFI_ICHECK(x->ndim == 1) << "x must be a 1D tensor";
+ DLDataType f32_dtype{kDLFloat, 32, 1};
+ TVM_FFI_ICHECK(x->dtype == f32_dtype) << "x must be a float tensor";
+ TVM_FFI_ICHECK(y->ndim == 1) << "y must be a 1D tensor";
+ TVM_FFI_ICHECK(y->dtype == f32_dtype) << "y must be a float tensor";
+ TVM_FFI_ICHECK(x->shape[0] == y->shape[0]) << "x and y must have the same shape";
+
+ int64_t n = x->shape[0];
+ int64_t nthread_per_block = 256;
+ int64_t nblock = (n + nthread_per_block - 1) / nthread_per_block;
+ // Obtain the current stream from the environment
+ // it will be set to torch.cuda.current_stream() when calling the function
+ // with torch.Tensors
+ cudaStream_t stream = static_cast<cudaStream_t>(
+ TVMFFIEnvGetCurrentStream(x->device.device_type, x->device.device_id));
+ // launch the kernel
+ AddOneKernel<<<nblock, nthread_per_block, 0, stream>>>(static_cast<float*>(x->data),
+ static_cast<float*>(y->data), n);
+}
+
+// Expose global symbol `add_one_cpu` that follows tvm-ffi abi
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(add_one_cuda, tvm_ffi_example::AddOneCUDA);
+} // namespace tvm_ffi_example
diff --git a/ffi/examples/get_started/src/run_example.cc b/ffi/examples/get_started/src/run_example.cc
new file mode 100644
index 0000000000..3472b99c7e
--- /dev/null
+++ b/ffi/examples/get_started/src/run_example.cc
@@ -0,0 +1,55 @@
+
+
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+#include <tvm/ffi/container/ndarray.h>
+#include <tvm/ffi/extra/module.h>
+
+// This file shows how to load the same compiled module and interact with it in C++
+namespace ffi = tvm::ffi;
+
+struct CPUNDAlloc {
+ void AllocData(DLTensor* tensor) { tensor->data = malloc(ffi::GetDataSize(*tensor)); }
+ void FreeData(DLTensor* tensor) { free(tensor->data); }
+};
+
+inline ffi::NDArray Empty(ffi::Shape shape, DLDataType dtype, DLDevice device) {
+ return ffi::NDArray::FromNDAlloc(CPUNDAlloc(), shape, dtype, device);
+}
+
+int main() {
+ // load the module
+ ffi::Module mod = ffi::Module::LoadFromFile("build/add_one_cpu.so");
+
+ // create an NDArray, alternatively, one can directly pass in a DLTensor*
+ ffi::NDArray x = Empty({5}, DLDataType({kDLFloat, 32, 1}), DLDevice({kDLCPU, 0}));
+ for (int i = 0; i < 5; ++i) {
+ reinterpret_cast<float*>(x->data)[i] = static_cast<float>(i);
+ }
+
+ ffi::Function add_one_cpu = mod->GetFunction("add_one_cpu").value();
+ add_one_cpu(x, x);
+
+ std::cout << "x after add_one_cpu(x, x)" << std::endl;
+ for (int i = 0; i < 5; ++i) {
+ std::cout << reinterpret_cast<float*>(x->data)[i] << " ";
+ }
+ std::cout << std::endl;
+ return 0;
+}
diff --git a/licenses/LICENSE.dlpack.txt b/ffi/licenses/LICENSE.dlpack.txt
similarity index 100%
rename from licenses/LICENSE.dlpack.txt
rename to ffi/licenses/LICENSE.dlpack.txt
diff --git a/licenses/LICENSE.libbacktrace.txt b/ffi/licenses/LICENSE.libbacktrace.txt
similarity index 100%
rename from licenses/LICENSE.libbacktrace.txt
rename to ffi/licenses/LICENSE.libbacktrace.txt
diff --git a/ffi/pyproject.toml b/ffi/pyproject.toml
new file mode 100644
index 0000000000..d9c712a448
--- /dev/null
+++ b/ffi/pyproject.toml
@@ -0,0 +1,152 @@
+# 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.
+
+[project]
+name = "apache-tvm-ffi"
+version = "0.1.0a0"
+description = "tvm ffi"
+
+authors = [{ name = "TVM FFI contributors" }]
+readme = "README.md"
+license = { text = "Apache 2.0" }
+classifiers = [
+ "License :: OSI Approved :: Apache Software License",
+ "Development Status :: 4 - Beta",
+ "Intended Audience :: Developers",
+ "Intended Audience :: Education",
+ "Intended Audience :: Science/Research",
+]
+keywords = ["machine learning", "inference"]
+requires-python = ">=3.9"
+dependencies = []
+
+
+[project.urls]
+Homepage = "https://github.com/apache/tvm/ffi"
+GitHub = "https://github.com/apache/tvm/ffi"
+
+[project.optional-dependencies]
+torch = ["torch"]
+test = ["pytest"]
+
+[project.scripts]
+tvm-ffi-config = "tvm_ffi.config:__main__"
+
+[build-system]
+requires = ["scikit-build-core>=0.10.0", "cython", "cython-cmake"]
+build-backend = "scikit_build_core.build"
+
+[tool.scikit-build]
+wheel.py-api = "cp312"
+minimum-version = "build-system.requires"
+
+# Build configuration
+build-dir = "build"
+build.verbose = true
+
+# CMake configuration
+cmake.version = "CMakeLists.txt"
+cmake.build-type = "RelWithDebugInfo"
+cmake.args = [
+ "-DTVM_FFI_USE_RELEASE_FLAGS=ON",
+ "-DTVM_FFI_BUILD_TESTS=OFF",
+ "-DTVM_FFI_BUILD_PYTHON_MODULE=ON"
+]
+
+# Logging
+logging.level = "INFO"
+
+# Wheel configuration
+wheel.packages = ["python/tvm_ffi"]
+wheel.install-dir = "tvm_ffi"
+
+# Source distribution configuration
+sdist.include = [
+ # Build files
+ "/CMakeLists.txt",
+ "/pyproject.toml",
+ "/cmake/**/*",
+ # Source code
+ "/src/**/*.cc",
+ "/include/**/*",
+
+ # python and cython
+ "/python/tvm_ffi/**/*.py",
+ "/python/tvm_ffi/**/*.pyx",
+ "/python/tvm_ffi/**/*.pyi",
+
+ # Third party files
+ "/3rdparty/**/*",
+
+ # Documentation and metadata
+ "/docs/**/*",
+ "/LICENSE",
+ "/README.md",
+ "/NOTICE",
+
+ # Tests
+ "/tests/**/*",
+]
+
+sdist.exclude = ["**/.git", "**/.github", "**/__pycache__", "**/*.pyc", "build", "dist"]
+
+# Editable install settings
+editable.rebuild = true
+editable.verbose = true
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+
+[tool.black]
+exclude = "3rdparty/*"
+line-length = 100
+skip-magic-trailing-comma = true
+
+[tool.isort]
+profile = "black"
+src_paths = ["python", "tests"]
+extend_skip = ["3rdparty"]
+line_length = 100
+skip_gitignore = true
+
+[tool.cibuildwheel]
+build-verbosity = 1
+# skip pp and low python version
+# sdist should be sufficient
+skip = [
+ "cp36-*",
+ "cp37-*",
+ "cp38-*",
+ "cp39-*",
+ "cp310-*",
+ "cp311-*",
+ "pp*",
+ "*musllinux*",
+] # pypy doesn't play nice with pybind11
+build-frontend = "build[uv]"
+test-command = "pytest {project}/tests -m "
+test-extras = ["test"]
+
+[tool.cibuildwheel.linux]
+archs = ["x86_64", "aarch64"]
+
+[tool.cibuildwheel.macos]
+archs = ["x86_64", "arm64"]
+environment = { MACOSX_DEPLOYMENT_TARGET = "10.14" }
+
+[tool.cibuildwheel.windows]
+archs = ["AMD64"]
diff --git a/python/tvm/ffi/.gitignore b/ffi/python/tvm_ffi/.gitignore
similarity index 100%
rename from python/tvm/ffi/.gitignore
rename to ffi/python/tvm_ffi/.gitignore
diff --git a/python/tvm/ffi/__init__.py b/ffi/python/tvm_ffi/__init__.py
similarity index 94%
rename from python/tvm/ffi/__init__.py
rename to ffi/python/tvm_ffi/__init__.py
index 801a8d2989..6929884609 100644
--- a/python/tvm/ffi/__init__.py
+++ b/ffi/python/tvm_ffi/__init__.py
@@ -14,12 +14,11 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""TVM FFI binding module.
-
-This module binds the TVM FFI C API to python.
-This is a standalone module that can be
-"""
+"""TVM FFI Python package."""
+# base always go first to load the libtvm_ffi
+from . import base
+# package init part
from .registry import register_object, register_func, get_global_func, _init_api
from .dtype import dtype, DataTypeCode
from .core import String, Bytes
diff --git a/python/tvm/ffi/_ffi_api.py b/ffi/python/tvm_ffi/_ffi_api.py
similarity index 100%
rename from python/tvm/ffi/_ffi_api.py
rename to ffi/python/tvm_ffi/_ffi_api.py
diff --git a/python/tvm/ffi/access_path.py b/ffi/python/tvm_ffi/access_path.py
similarity index 100%
rename from python/tvm/ffi/access_path.py
rename to ffi/python/tvm_ffi/access_path.py
diff --git a/python/tvm/base.py b/ffi/python/tvm_ffi/base.py
similarity index 52%
copy from python/tvm/base.py
copy to ffi/python/tvm_ffi/base.py
index 63e097999c..81e06a06e7 100644
--- a/python/tvm/base.py
+++ b/ffi/python/tvm_ffi/base.py
@@ -15,14 +15,16 @@
# specific language governing permissions and limitations
# under the License.
# coding: utf-8
-# pylint: disable=invalid-name, import-outside-toplevel
-"""Base library for TVM."""
+"""Base library for TVM FFI."""
import ctypes
import os
import sys
-
+import subprocess
+import logging
from . import libinfo
+logger = logging.getLogger(__name__)
+
# ----------------------------
# Python3 version.
# ----------------------------
@@ -37,44 +39,45 @@ if not (sys.version_info[0] >= 3 and sys.version_info[1] >= 9):
def _load_lib():
"""Load libary by searching possible path."""
- lib_path = libinfo.find_lib_path()
+ lib_path = libinfo.find_libtvm_ffi()
# The dll search path need to be added explicitly in windows
if sys.platform.startswith("win32"):
for path in libinfo.get_dll_directories():
os.add_dll_directory(path)
- lib = ctypes.CDLL(lib_path[0], ctypes.RTLD_GLOBAL)
- return lib, os.path.basename(lib_path[0])
+ lib = ctypes.CDLL(lib_path, ctypes.RTLD_GLOBAL)
+ return lib
-try:
- # The following import is needed for TVM to work with pdb
- import readline # pylint: disable=unused-import
-except ImportError:
- pass
-# version number
-__version__ = libinfo.__version__
# library instance
-_LIB, _LIB_NAME = _load_lib()
-
-# Whether we are runtime only
-_RUNTIME_ONLY = "runtime" in _LIB_NAME
-
-
-if _RUNTIME_ONLY:
- from .ffi import registry as _tvm_ffi_registry
+_LIB = _load_lib()
- _tvm_ffi_registry._SKIP_UNKNOWN_OBJECTS = True
-# The FFI mode of TVM
-_FFI_MODE = os.environ.get("TVM_FFI", "auto")
+def _setup_cython_redirector():
+ """Setup cython library."""
+ cython_lib_path = libinfo.find_cython_lib()
+ # cython file exists in the same directory
+ # no need to setup redirector
+ if os.path.dirname(os.path.realpath(__file__)) == os.path.dirname(cython_lib_path):
+ return
+ # the install is done via pip install -e, we need to setup redirector
+ logger.info(f"Setting up cython lib redirector from {cython_lib_path}")
+ # pylint: disable=import-outside-toplevel
+ import importlib
+ import importlib.util
+ import importlib.machinery
-if _FFI_MODE == "ctypes":
- raise ImportError("We have phased out ctypes support in favor of cython on wards")
+ class CythonLibRedirector(importlib.abc.MetaPathFinder):
+ def __init__(self, lib_path):
+ self.lib_path = lib_path
+ def find_spec(self, fullname, path, target=None):
+ if fullname != "tvm_ffi.core":
+ return None
+ loader = importlib.machinery.ExtensionFileLoader(fullname, self.lib_path)
+ return importlib.util.spec_from_file_location(fullname, self.lib_path, loader=loader)
-def py_str(x):
- return x.decode("utf-8")
+ sys.meta_path.insert(0, CythonLibRedirector(cython_lib_path))
-TVMError = Exception
+_setup_cython_redirector()
diff --git a/ffi/python/tvm_ffi/config.py b/ffi/python/tvm_ffi/config.py
new file mode 100644
index 0000000000..583ed409ba
--- /dev/null
+++ b/ffi/python/tvm_ffi/config.py
@@ -0,0 +1,64 @@
+# 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.
+"""Config utilities for finding paths to lib and headers"""
+
+import argparse
+import sys
+from . import libinfo
+
+
+def __main__():
+ """Main function"""
+ parser = argparse.ArgumentParser(
+ description="Get various configuration information needed to compile with tvm-ffi",
+ )
+
+ parser.add_argument("--includedir", action="store_true", help="Print include directory")
+ parser.add_argument(
+ "--dlpack-includedir", action="store_true", help="Print dlpack include directory"
+ )
+ parser.add_argument("--cmakedir", action="store_true", help="Print library directory")
+ parser.add_argument("--sourcedir", action="store_true", help="Print source directory")
+ parser.add_argument("--libfiles", action="store_true", help="Fully qualified library filenames")
+ parser.add_argument("--libdir", action="store_true", help="Print library directory")
+ parser.add_argument("--cython-lib-path", action="store_true", help="Print cython path")
+
+ args = parser.parse_args()
+
+ # print help when no arguments are provided
+ if len(sys.argv) == 1:
+ parser.print_help()
+ return
+
+ if args.includedir:
+ print(libinfo.find_include_path())
+ if args.dlpack_includedir:
+ print(libinfo.find_dlpack_include_path())
+ if args.cmakedir:
+ print(libinfo.find_cmake_path())
+ if args.libdir:
+ print(os.path.dirname(libinfo.find_libtvm_ffi()))
+ if args.libfiles:
+ print(libinfo.find_libtvm_ffi())
+ if args.sourcedir:
+ print(libinfo.find_source_path())
+ if args.cython_lib_path:
+ print(libinfo.find_cython_lib())
+
+
+if __name__ == "__main__":
+ __main__()
diff --git a/python/tvm/ffi/container.py b/ffi/python/tvm_ffi/container.py
similarity index 100%
rename from python/tvm/ffi/container.py
rename to ffi/python/tvm_ffi/container.py
diff --git a/python/tvm/ffi/convert.py b/ffi/python/tvm_ffi/convert.py
similarity index 100%
rename from python/tvm/ffi/convert.py
rename to ffi/python/tvm_ffi/convert.py
diff --git a/python/tvm/ffi/cython/base.pxi b/ffi/python/tvm_ffi/cython/base.pxi
similarity index 100%
rename from python/tvm/ffi/cython/base.pxi
rename to ffi/python/tvm_ffi/cython/base.pxi
diff --git a/python/tvm/ffi/cython/core.pyx b/ffi/python/tvm_ffi/cython/core.pyx
similarity index 100%
rename from python/tvm/ffi/cython/core.pyx
rename to ffi/python/tvm_ffi/cython/core.pyx
diff --git a/python/tvm/ffi/cython/device.pxi b/ffi/python/tvm_ffi/cython/device.pxi
similarity index 100%
rename from python/tvm/ffi/cython/device.pxi
rename to ffi/python/tvm_ffi/cython/device.pxi
diff --git a/python/tvm/ffi/cython/dtype.pxi b/ffi/python/tvm_ffi/cython/dtype.pxi
similarity index 100%
rename from python/tvm/ffi/cython/dtype.pxi
rename to ffi/python/tvm_ffi/cython/dtype.pxi
diff --git a/python/tvm/ffi/cython/error.pxi b/ffi/python/tvm_ffi/cython/error.pxi
similarity index 100%
rename from python/tvm/ffi/cython/error.pxi
rename to ffi/python/tvm_ffi/cython/error.pxi
diff --git a/python/tvm/ffi/cython/function.pxi b/ffi/python/tvm_ffi/cython/function.pxi
similarity index 99%
rename from python/tvm/ffi/cython/function.pxi
rename to ffi/python/tvm_ffi/cython/function.pxi
index 4148cc6c88..2a2ee855f5 100644
--- a/python/tvm/ffi/cython/function.pxi
+++ b/ffi/python/tvm_ffi/cython/function.pxi
@@ -307,8 +307,8 @@ class Function(Object):
See Also
--------
- tvm.ffi.register_func: How to register global function.
- tvm.ffi.get_global_func: How to get global function.
+ tvm_ffi.register_func: How to register global function.
+ tvm_ffi.get_global_func: How to get global function.
"""
def __call__(self, *args):
cdef TVMFFIAny result
diff --git a/python/tvm/ffi/cython/ndarray.pxi b/ffi/python/tvm_ffi/cython/ndarray.pxi
similarity index 100%
rename from python/tvm/ffi/cython/ndarray.pxi
rename to ffi/python/tvm_ffi/cython/ndarray.pxi
diff --git a/python/tvm/ffi/cython/object.pxi b/ffi/python/tvm_ffi/cython/object.pxi
similarity index 100%
rename from python/tvm/ffi/cython/object.pxi
rename to ffi/python/tvm_ffi/cython/object.pxi
diff --git a/python/tvm/ffi/cython/string.pxi b/ffi/python/tvm_ffi/cython/string.pxi
similarity index 100%
rename from python/tvm/ffi/cython/string.pxi
rename to ffi/python/tvm_ffi/cython/string.pxi
diff --git a/python/tvm/ffi/dtype.py b/ffi/python/tvm_ffi/dtype.py
similarity index 100%
rename from python/tvm/ffi/dtype.py
rename to ffi/python/tvm_ffi/dtype.py
diff --git a/python/tvm/ffi/error.py b/ffi/python/tvm_ffi/error.py
similarity index 100%
rename from python/tvm/ffi/error.py
rename to ffi/python/tvm_ffi/error.py
diff --git a/ffi/python/tvm_ffi/libinfo.py b/ffi/python/tvm_ffi/libinfo.py
new file mode 100644
index 0000000000..19cb90a3b1
--- /dev/null
+++ b/ffi/python/tvm_ffi/libinfo.py
@@ -0,0 +1,145 @@
+# 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.
+
+import sys
+import os
+import glob
+
+
+def split_env_var(env_var, split):
+ """Splits environment variable string.
+
+ Parameters
+ ----------
+ env_var : str
+ Name of environment variable.
+
+ split : str
+ String to split env_var on.
+
+ Returns
+ -------
+ splits : list(string)
+ If env_var exists, split env_var. Otherwise, empty list.
+ """
+ if os.environ.get(env_var, None):
+ return [p.strip() for p in os.environ[env_var].split(split)]
+ return []
+
+
+def get_dll_directories():
+ """Get the possible dll directories"""
+ ffi_dir = os.path.dirname(os.path.realpath(os.path.expanduser(__file__)))
+ dll_path = [os.path.join(ffi_dir, "lib")]
+ dll_path += [os.path.join(ffi_dir, "..", "..", "build", "lib")]
+ # in source build from parent if needed
+ dll_path += [os.path.join(ffi_dir, "..", "..", "..", "build", "lib")]
+
+ if sys.platform.startswith("linux") or sys.platform.startswith("freebsd"):
+ dll_path.extend(split_env_var("LD_LIBRARY_PATH", ":"))
+ dll_path.extend(split_env_var("PATH", ":"))
+ elif sys.platform.startswith("darwin"):
+ dll_path.extend(split_env_var("DYLD_LIBRARY_PATH", ":"))
+ dll_path.extend(split_env_var("PATH", ":"))
+ elif sys.platform.startswith("win32"):
+ dll_path.extend(split_env_var("PATH", ";"))
+ return [os.path.abspath(x) for x in dll_path if os.path.isdir(x)]
+
+
+def find_libtvm_ffi():
+ """Find libtvm_ffi."""
+ dll_path = get_dll_directories()
+ if sys.platform.startswith("win32"):
+ lib_dll_names = ["libtvm_ffi.dll", "tvm_ffi.dll"]
+ elif sys.platform.startswith("darwin"):
+ lib_dll_names = ["libtvm_ffi.dylib", "libtvm_ffi.so"]
+ else:
+ lib_dll_names = ["libtvm_ffi.so"]
+
+ name = lib_dll_names
+ lib_dll_path = [os.path.join(p, name) for name in lib_dll_names for p in dll_path]
+ lib_found = [p for p in lib_dll_path if os.path.exists(p) and os.path.isfile(p)]
+
+ if not lib_found:
+ raise RuntimeError(f"Cannot find library: {name}\nList of candidates:\n{lib_dll_path}")
+
+ return lib_found[0]
+
+
+def find_source_path():
+ """Find packaged source home path."""
+ candidates = [
+ os.path.join(os.path.dirname(os.path.realpath(__file__))),
+ os.path.join(os.path.dirname(os.path.realpath(__file__)), "..", ".."),
+ ]
+ for candidate in candidates:
+ if os.path.isdir(os.path.join(candidate, "cmake")):
+ return candidate
+ raise RuntimeError("Cannot find home path.")
+
+
+def find_cmake_path():
+ """Find the preferred cmake path."""
+ candidates = [
+ os.path.join(os.path.dirname(os.path.realpath(__file__)), "lib", "cmake"),
+ os.path.join(os.path.dirname(os.path.realpath(__file__)), "..", "..", "cmake"),
+ ]
+ for candidate in candidates:
+ if os.path.isdir(candidate):
+ return candidate
+ raise RuntimeError("Cannot find cmake path.")
+
+
+def find_include_path():
+ """Find header files for C compilation."""
+ candidates = [
+ os.path.join(os.path.dirname(os.path.realpath(__file__)), "include"),
+ os.path.join(os.path.dirname(os.path.realpath(__file__)), "..", "..", "include"),
+ ]
+ for candidate in candidates:
+ if os.path.isdir(candidate):
+ return candidate
+ raise RuntimeError("Cannot find include path.")
+
+
+def find_dlpack_include_path():
+ """Find dlpack header files for C compilation."""
+ install_include_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "include")
+ if os.path.isdir(os.path.join(install_include_path, "dlpack")):
+ return install_include_path
+
+ source_include_path = os.path.join(
+ os.path.dirname(os.path.realpath(__file__)), "..", "..", "3rdparty", "dlpack", "include"
+ )
+ if os.path.isdir(source_include_path):
+ return source_include_path
+
+ raise RuntimeError("Cannot find include path.")
+
+
+def find_cython_lib():
+ """Find the path to tvm cython."""
+ path_candidates = [
+ os.path.dirname(os.path.realpath(__file__)),
+ os.path.join(os.path.dirname(os.path.realpath(__file__)), "..", "..", "build"),
+ os.path.join(os.path.dirname(os.path.realpath(__file__)), "..", "..", "..", "build"),
+ ]
+ suffixes = "so" if sys.platform.startswith("linux") else "pyd"
+ for candidate in path_candidates:
+ for path in glob.glob(os.path.join(candidate, f"core*.{suffixes}")):
+ return os.path.abspath(path)
+ raise RuntimeError("Cannot find tvm cython path.")
diff --git a/python/tvm/ffi/module.py b/ffi/python/tvm_ffi/module.py
similarity index 95%
rename from python/tvm/ffi/module.py
rename to ffi/python/tvm_ffi/module.py
index 0895b317c1..56aa15348e 100644
--- a/python/tvm/ffi/module.py
+++ b/ffi/python/tvm_ffi/module.py
@@ -50,7 +50,7 @@ class Module(core.Object):
Returns
-------
- f : tvm.ffi.Function
+ f : tvm_ffi.Function
The entry function if exist
"""
if self._entry:
@@ -96,6 +96,15 @@ class Module(core.Object):
"""
return _ffi_api.ModuleImplementsFunction(self, name, query_imports)
+ def __getattr__(self, name):
+ """Accessor to allow getting functions as attributes."""
+ try:
+ func = self.get_function(name)
+ self.__dict__[name] = func
+ return func
+ except AttributeError:
+ raise AttributeError(f"Module has no function '{name}'")
+
def get_function(self, name, query_imports=False):
"""Get function from the module.
@@ -109,7 +118,7 @@ class Module(core.Object):
Returns
-------
- f : tvm.ffi.Function
+ f : tvm_ffi.Function
The result function.
"""
func = _ffi_api.ModuleGetFunction(self, name, query_imports)
diff --git a/python/tvm/ffi/ndarray.py b/ffi/python/tvm_ffi/ndarray.py
similarity index 97%
rename from python/tvm/ffi/ndarray.py
rename to ffi/python/tvm_ffi/ndarray.py
index 05856bdae7..d65b8fb361 100644
--- a/python/tvm/ffi/ndarray.py
+++ b/ffi/python/tvm_ffi/ndarray.py
@@ -56,7 +56,7 @@ def device(dev_type, dev_id=0):
Returns
-------
- dev: tvm.ffi.Device
+ dev: tvm_ffi.Device
Examples
--------
@@ -65,8 +65,8 @@ def device(dev_type, dev_id=0):
.. code-block:: python
- assert tvm.ffi.device("cuda:0") == tvm.ffi.cuda(1)
- assert tvm.ffi.device("cpu", 0) == tvm.ffi.cpu(0)
+ assert tvm_ffi.device("cuda:0") == tvm_ffi.cuda(1)
+ assert tvm_ffi.device("cpu", 0) == tvm_ffi.cpu(0)
"""
if isinstance(dev_type, str):
dev_type = dev_type.split(" ")[0]
diff --git a/python/tvm/ffi/registry.py b/ffi/python/tvm_ffi/registry.py
similarity index 98%
rename from python/tvm/ffi/registry.py
rename to ffi/python/tvm_ffi/registry.py
index 9302b25173..e2455c3d33 100644
--- a/python/tvm/ffi/registry.py
+++ b/ffi/python/tvm_ffi/registry.py
@@ -37,7 +37,7 @@ def register_object(type_key=None):
.. code-block:: python
- @tvm.ffi.register_object("test.MyObject")
+ @tvm_ffi.register_object("test.MyObject")
class MyObject(Object):
pass
"""
diff --git a/python/tvm/ffi/serialization.py b/ffi/python/tvm_ffi/serialization.py
similarity index 100%
rename from python/tvm/ffi/serialization.py
rename to ffi/python/tvm_ffi/serialization.py
diff --git a/python/tvm/ffi/testing.py b/ffi/python/tvm_ffi/testing.py
similarity index 100%
rename from python/tvm/ffi/testing.py
rename to ffi/python/tvm_ffi/testing.py
diff --git a/ffi/scripts/benchmark_dlpack.py b/ffi/scripts/benchmark_dlpack.py
index 1453aa95a6..73fbe0f6ac 100644
--- a/ffi/scripts/benchmark_dlpack.py
+++ b/ffi/scripts/benchmark_dlpack.py
@@ -39,7 +39,7 @@ Summary of some takeaways:
import os
import torch
import numpy as np
-from tvm import ffi as tvm_ffi
+import tvm_ffi
import time
@@ -124,11 +124,11 @@ def tvm_ffi_nop(repeat):
for i in range(repeat):
y = tvm_ffi.from_dlpack(x)
end = time.time()
- print_speed("tvm.ffi.nop", (end - start) / repeat)
+ print_speed("tvm_ffi.nop", (end - start) / repeat)
def bench_ffi_nop_from_dlpack(name, x, y, z, repeat):
- """run dlpack conversion + tvm.ffi.nop
+ """run dlpack conversion + tvm_ffi.nop
Measures overhead of running dlpack for each args then invoke
"""
@@ -149,40 +149,40 @@ def bench_ffi_nop_from_dlpack(name, x, y, z, repeat):
def tvm_ffi_nop_from_torch_dlpack(repeat):
- """run dlpack conversion + tvm.ffi.nop
+ """run dlpack conversion + tvm_ffi.nop
Measures overhead of running dlpack for each args then invoke
"""
x = torch.arange(1)
y = torch.arange(1)
z = torch.arange(1)
- bench_ffi_nop_from_dlpack("tvm.ffi.nop+from_dlpack(torch)", x, y, z, repeat)
+ bench_ffi_nop_from_dlpack("tvm_ffi.nop+from_dlpack(torch)", x, y, z, repeat)
def tvm_ffi_nop_from_numpy_dlpack(repeat):
- """run dlpack conversion + tvm.ffi.nop
+ """run dlpack conversion + tvm_ffi.nop
Measures overhead of running dlpack for each args then invoke
"""
x = np.arange(1)
y = np.arange(1)
z = np.arange(1)
- bench_ffi_nop_from_dlpack("tvm.ffi.nop+from_dlpack(numpy)", x, y, z, repeat)
+ bench_ffi_nop_from_dlpack("tvm_ffi.nop+from_dlpack(numpy)", x, y, z, repeat)
def tvm_ffi_self_dlpack_nop(repeat):
- """run dlpack conversion + tvm.ffi.nop
+ """run dlpack conversion + tvm_ffi.nop
Measures overhead of running dlpack for each args then invoke
"""
x = tvm_ffi.from_dlpack(torch.arange(1))
y = tvm_ffi.from_dlpack(torch.arange(1))
z = tvm_ffi.from_dlpack(torch.arange(1))
- bench_ffi_nop_from_dlpack("tvm.ffi.nop+from_dlpack(tvm)", x, y, z, repeat)
+ bench_ffi_nop_from_dlpack("tvm_ffi.nop+from_dlpack(tvm)", x, y, z, repeat)
def bench_ffi_nop_from_dlpack(name, x, y, z, repeat):
- """run dlpack conversion + tvm.ffi.nop
+ """run dlpack conversion + tvm_ffi.nop
Measures overhead of running dlpack for each args then invoke
"""
@@ -227,7 +227,7 @@ def tvm_ffi_nop_from_torch_utils_to_dlpack(repeat):
nop(tx, ty, tz)
end = time.time()
speed = (end - start) / repeat
- print_speed("tvm.ffi.nop+from_dlpack(torch.utils)", speed)
+ print_speed("tvm_ffi.nop+from_dlpack(torch.utils)", speed)
def bench_tvm_ffi_nop_autodlpack(name, x, y, z, repeat):
@@ -257,10 +257,10 @@ def tvm_ffi_nop_autodlpack_from_torch(repeat, device="cpu", stream=False):
if stream:
with torch.cuda.stream(torch.cuda.Stream()):
bench_tvm_ffi_nop_autodlpack(
- f"tvm.ffi.nop.autodlpack(torch[{device}][stream])", x, y, z, repeat
+ f"tvm_ffi.nop.autodlpack(torch[{device}][stream])", x, y, z, repeat
)
else:
- bench_tvm_ffi_nop_autodlpack(f"tvm.ffi.nop.autodlpack(torch[{device}])", x, y, z, repeat)
+ bench_tvm_ffi_nop_autodlpack(f"tvm_ffi.nop.autodlpack(torch[{device}])", x, y, z, repeat)
def tvm_ffi_nop_autodlpack_from_numpy(repeat):
@@ -272,7 +272,7 @@ def tvm_ffi_nop_autodlpack_from_numpy(repeat):
x = np.arange(256)
y = np.arange(256)
z = np.arange(256)
- bench_tvm_ffi_nop_autodlpack("tvm.ffi.nop.autodlpack(numpy)", x, y, z, repeat)
+ bench_tvm_ffi_nop_autodlpack("tvm_ffi.nop.autodlpack(numpy)", x, y, z, repeat)
def bench_to_dlpack(x, name, repeat):
diff --git a/ffi/tests/cpp/CMakeLists.txt b/ffi/tests/cpp/CMakeLists.txt
index 0c820fc80e..ab1e6e9cfd 100644
--- a/ffi/tests/cpp/CMakeLists.txt
+++ b/ffi/tests/cpp/CMakeLists.txt
@@ -20,12 +20,12 @@ set_target_properties(
LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib"
RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/bin"
)
-add_cxx_warning(tvm_ffi_tests)
+tvm_ffi_add_cxx_warning(tvm_ffi_tests)
add_sanitizer_address(tvm_ffi_tests)
-add_dsymutil(tvm_ffi_tests)
-add_msvc_flags(tvm_ffi_tests)
+tvm_ffi_add_dsymutil(tvm_ffi_tests)
+tvm_ffi_add_msvc_flags(tvm_ffi_tests)
target_link_libraries(tvm_ffi_tests PRIVATE tvm_ffi_shared)
-add_googletest(tvm_ffi_tests)
+tvm_ffi_add_googletest(tvm_ffi_tests)
if (MSVC)
target_link_options(tvm_ffi_tests PRIVATE /DEBUG)
diff --git a/tests/python/ffi/test_access_path.py b/ffi/tests/python/test_access_path.py
similarity index 98%
rename from tests/python/ffi/test_access_path.py
rename to ffi/tests/python/test_access_path.py
index 06fbb64ff2..7d9e7af55f 100644
--- a/tests/python/ffi/test_access_path.py
+++ b/ffi/tests/python/test_access_path.py
@@ -16,7 +16,7 @@
# under the License.
import pytest
-from tvm.ffi.access_path import AccessPath, AccessKind
+from tvm_ffi.access_path import AccessPath, AccessKind
def test_root_path():
diff --git a/tests/python/ffi/test_container.py b/ffi/tests/python/test_container.py
similarity index 99%
rename from tests/python/ffi/test_container.py
rename to ffi/tests/python/test_container.py
index 25468f452a..657adbef66 100644
--- a/tests/python/ffi/test_container.py
+++ b/ffi/tests/python/test_container.py
@@ -16,7 +16,7 @@
# under the License.
import pytest
import pickle
-import tvm.ffi as tvm_ffi
+import tvm_ffi
def test_array():
diff --git a/tests/python/ffi/test_device.py b/ffi/tests/python/test_device.py
similarity index 98%
rename from tests/python/ffi/test_device.py
rename to ffi/tests/python/test_device.py
index 5800a0c441..645738710f 100644
--- a/tests/python/ffi/test_device.py
+++ b/ffi/tests/python/test_device.py
@@ -17,8 +17,8 @@
import pytest
import pickle
-from tvm.ffi import Device
-from tvm import ffi as tvm_ffi
+from tvm_ffi import Device
+import tvm_ffi
def test_device():
diff --git a/tests/python/ffi/test_dtype.py b/ffi/tests/python/test_dtype.py
similarity index 97%
rename from tests/python/ffi/test_dtype.py
rename to ffi/tests/python/test_dtype.py
index 332d0e1827..7d09d3def9 100644
--- a/tests/python/ffi/test_dtype.py
+++ b/ffi/tests/python/test_dtype.py
@@ -18,9 +18,7 @@
import pytest
import pickle
import numpy as np
-import tvm
-import tvm.testing
-from tvm import ffi as tvm_ffi
+import tvm_ffi
def test_dtype():
diff --git a/tests/python/ffi/test_error.py b/ffi/tests/python/test_error.py
similarity index 99%
rename from tests/python/ffi/test_error.py
rename to ffi/tests/python/test_error.py
index e3d02234b5..7aa313628c 100644
--- a/tests/python/ffi/test_error.py
+++ b/ffi/tests/python/test_error.py
@@ -17,7 +17,7 @@
import pytest
import platform
-from tvm import ffi as tvm_ffi
+import tvm_ffi
def test_parse_traceback():
diff --git a/tests/python/ffi/test_function.py b/ffi/tests/python/test_function.py
similarity index 99%
rename from tests/python/ffi/test_function.py
rename to ffi/tests/python/test_function.py
index 5a8b4acb1f..cb81f47c7d 100644
--- a/tests/python/ffi/test_function.py
+++ b/ffi/tests/python/test_function.py
@@ -18,7 +18,7 @@
import gc
import ctypes
import numpy as np
-from tvm import ffi as tvm_ffi
+import tvm_ffi
def test_echo():
diff --git a/tests/python/ffi/test_ndarray.py b/ffi/tests/python/test_ndarray.py
similarity index 98%
rename from tests/python/ffi/test_ndarray.py
rename to ffi/tests/python/test_ndarray.py
index 5b75171b55..f0ce0d193c 100644
--- a/tests/python/ffi/test_ndarray.py
+++ b/ffi/tests/python/test_ndarray.py
@@ -21,7 +21,7 @@ try:
except ImportError:
torch = None
-from tvm import ffi as tvm_ffi
+import tvm_ffi
import numpy as np
diff --git a/tests/python/ffi/test_object.py b/ffi/tests/python/test_object.py
similarity index 98%
rename from tests/python/ffi/test_object.py
rename to ffi/tests/python/test_object.py
index d333cbca08..63867b9de1 100644
--- a/tests/python/ffi/test_object.py
+++ b/ffi/tests/python/test_object.py
@@ -16,7 +16,7 @@
# under the License.
import pytest
-from tvm import ffi as tvm_ffi
+import tvm_ffi
def test_make_object():
diff --git a/tests/python/ffi/test_string.py b/ffi/tests/python/test_string.py
similarity index 98%
rename from tests/python/ffi/test_string.py
rename to ffi/tests/python/test_string.py
index 85fed5670c..f334bc4fad 100644
--- a/tests/python/ffi/test_string.py
+++ b/ffi/tests/python/test_string.py
@@ -16,7 +16,7 @@
# under the License.
import pickle
-from tvm import ffi as tvm_ffi
+import tvm_ffi
def test_string():
diff --git a/include/tvm/runtime/c_backend_api.h b/include/tvm/runtime/c_backend_api.h
index 4e6c2f5364..a2eefb5b7d 100644
--- a/include/tvm/runtime/c_backend_api.h
+++ b/include/tvm/runtime/c_backend_api.h
@@ -47,15 +47,6 @@ extern "C" {
TVM_DLL int TVMBackendGetFuncFromEnv(void* mod_node, const char* func_name,
TVMFFIObjectHandle* out);
-/*!
- * \brief Backend function to register system-wide library symbol.
- *
- * \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 TVMFFIEnvModRegisterSystemLibSymbol(const char* name, void* ptr);
-
/*!
* \brief Backend function to allocate temporal workspace.
*
diff --git a/python/setup.py b/python/setup.py
index cf2eff2a3a..a83ad81856 100644
--- a/python/setup.py
+++ b/python/setup.py
@@ -23,14 +23,8 @@ import sys
from setuptools import find_packages
from setuptools.dist import Distribution
-
-# need to use distutils.core for correct placement of cython dll
-if "--inplace" in sys.argv:
- from distutils.core import setup
- from distutils.extension import Extension
-else:
- from setuptools import setup
- from setuptools.extension import Extension
+from setuptools import setup
+from setuptools.extension import Extension
CURRENT_DIR = os.path.dirname(__file__)
FFI_MODE = os.environ.get("TVM_FFI", "auto")
@@ -136,65 +130,6 @@ if not CONDA_BUILD and not INPLACE_BUILD:
_remove_path(f"tvm/{libname}")
-def config_cython():
- """Try to configure cython and return cython configuration"""
- # Enforce cython unless FFI_MODE is explicitly set to ctypes
- # we might consider fully converge to cython later
- if FFI_MODE == "ctypes":
- return []
- try:
- from Cython.Build import cythonize
-
- # for python 3.12+, use limited API for future compact
- limited_api_kwargs = {}
- if sys.version_info >= (3, 12):
- limited_api_kwargs = {
- "define_macros": [
- ("Py_LIMITED_API", 0x030C0000),
- ],
- "py_limited_api": True,
- }
-
- ret = []
- extra_compile_args = ["-std=c++17", "-DDMLC_USE_LOGGING_LIBRARY=<tvm/runtime/logging.h>"]
- if os.name == "nt":
- library_dirs = ["tvm", "../build/Release", "../build"]
- libraries = ["tvm"]
- extra_compile_args = [
- "/std:c++17",
- "/D DMLC_USE_LOGGING_LIBRARY=<tvm/runtime/logging.h>",
- ]
- # library is available via conda env.
- if CONDA_BUILD:
- library_dirs = [os.environ["LIBRARY_LIB"]]
- else:
- library_dirs = None
- libraries = None
-
- # the latest ffi source
- for fn in os.listdir("tvm/ffi/cython"):
- if not fn.endswith(".pyx"):
- continue
- ret.append(
- Extension(
- f"tvm.ffi.{fn[:-4]}",
- ["tvm/ffi/cython/%s" % fn],
- include_dirs=[
- "../ffi/include/",
- "../ffi/3rdparty/dlpack/include",
- ],
- extra_compile_args=extra_compile_args,
- library_dirs=library_dirs,
- libraries=libraries,
- language="c++",
- **limited_api_kwargs,
- )
- )
- return cythonize(ret, compiler_directives={"language_level": 3})
- except ImportError as error:
- raise RuntimeError("Cython is not installed, please pip install cython")
-
-
class BinaryDistribution(Distribution):
def has_ext_modules(self):
return True
@@ -263,7 +198,6 @@ setup(
packages=find_packages(),
package_dir={"tvm": "tvm"},
distclass=BinaryDistribution,
- ext_modules=config_cython(),
**setup_kwargs,
)
diff --git a/python/tvm/__init__.py b/python/tvm/__init__.py
index 150e5d4b1d..59d8e05666 100644
--- a/python/tvm/__init__.py
+++ b/python/tvm/__init__.py
@@ -20,12 +20,12 @@ import multiprocessing
import sys
import os
+# ffi module must load first
+from tvm_ffi import register_object, register_func, get_global_func
+
# top-level alias
-# tvm._ffi
from .base import TVMError, __version__, _RUNTIME_ONLY
-from .ffi import register_object, register_func, get_global_func
-
# top-level alias
# tvm.runtime
from .runtime.object import Object
diff --git a/python/tvm/arith/_ffi_api.py b/python/tvm/arith/_ffi_api.py
index e05405b0fc..aa98839349 100644
--- a/python/tvm/arith/_ffi_api.py
+++ b/python/tvm/arith/_ffi_api.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""FFI APIs for tvm.arith"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("arith", __name__)
+tvm_ffi._init_api("arith", __name__)
diff --git a/python/tvm/arith/analyzer.py b/python/tvm/arith/analyzer.py
index 434e2a3e65..c5c8fc067c 100644
--- a/python/tvm/arith/analyzer.py
+++ b/python/tvm/arith/analyzer.py
@@ -19,7 +19,7 @@
import enum
from typing import Union
-import tvm.ffi
+import tvm_ffi
from tvm import ir, tir
from tvm.arith import IntSet
from tvm.runtime import Object
@@ -47,7 +47,7 @@ class Extension(enum.Flag):
ComparisonOfProductAndSum = 1 << 3
-@tvm.ffi.register_object("arith.ModularSet")
+@tvm_ffi.register_object("arith.ModularSet")
class ModularSet(Object):
"""Represent range of (coeff * x + base) for x in Z"""
@@ -55,7 +55,7 @@ class ModularSet(Object):
self.__init_handle_by_constructor__(_ffi_api.ModularSet, coeff, base)
-@tvm.ffi.register_object("arith.ConstIntBound")
+@tvm_ffi.register_object("arith.ConstIntBound")
class ConstIntBound(Object):
"""Represent constant integer bound
diff --git a/python/tvm/arith/int_set.py b/python/tvm/arith/int_set.py
index 7a0aae5fda..fc6c20dec1 100644
--- a/python/tvm/arith/int_set.py
+++ b/python/tvm/arith/int_set.py
@@ -15,12 +15,12 @@
# specific language governing permissions and limitations
# under the License.
"""Integer set."""
-import tvm.ffi
+import tvm_ffi
from tvm.runtime import Object
from . import _ffi_api
-@tvm.ffi.register_object("ir.IntSet")
+@tvm_ffi.register_object("ir.IntSet")
class IntSet(Object):
"""Represent a set of integer in one dimension."""
@@ -65,7 +65,7 @@ class IntSet(Object):
return _ffi_api.intset_single_point(point)
-@tvm.ffi.register_object("arith.IntervalSet")
+@tvm_ffi.register_object("arith.IntervalSet")
class IntervalSet(IntSet):
"""Represent set of continuous interval [min_value, max_value]
@@ -82,7 +82,7 @@ class IntervalSet(IntSet):
self.__init_handle_by_constructor__(_ffi_api.IntervalSet, min_value, max_value)
-@tvm.ffi.register_object("arith.PresburgerSet")
+@tvm_ffi.register_object("arith.PresburgerSet")
class PresburgerSet(IntSet):
"""Represent of Presburger Set"""
diff --git a/python/tvm/arith/int_solver.py b/python/tvm/arith/int_solver.py
index a97cda10f8..72e4c46896 100644
--- a/python/tvm/arith/int_solver.py
+++ b/python/tvm/arith/int_solver.py
@@ -15,12 +15,12 @@
# specific language governing permissions and limitations
# under the License.
"""integer constraints data structures and solvers"""
-import tvm.ffi
+import tvm_ffi
from tvm.runtime import Object
from . import _ffi_api
-@tvm.ffi.register_object("arith.IntGroupBounds")
+@tvm_ffi.register_object("arith.IntGroupBounds")
class IntGroupBounds(Object):
"""Represent integer grouped bounds which are classified into
lower bounds (include), upper bounds (include) and equalities.
@@ -66,7 +66,7 @@ class IntGroupBounds(Object):
return _ffi_api.IntGroupBounds_FindBestRange(self)
-@tvm.ffi.register_object("arith.IntConstraints")
+@tvm_ffi.register_object("arith.IntConstraints")
class IntConstraints(Object):
"""Represent a set of integer constraints including variables, their ranges and
the relations between them (either equations or inequalities)
@@ -85,7 +85,7 @@ class IntConstraints(Object):
self.__init_handle_by_constructor__(_ffi_api.IntConstraints, variables, ranges, relations)
-@tvm.ffi.register_object("arith.IntConstraintsTransform")
+@tvm_ffi.register_object("arith.IntConstraintsTransform")
class IntConstraintsTransform(Object):
"""We can have different set of variables to represent the same integer constraints.
For example, the following two constrains are equivalent,
diff --git a/python/tvm/arith/iter_affine_map.py b/python/tvm/arith/iter_affine_map.py
index 328bb052b8..69ad3022fb 100644
--- a/python/tvm/arith/iter_affine_map.py
+++ b/python/tvm/arith/iter_affine_map.py
@@ -16,18 +16,18 @@
# under the License.
"""Iterator (quasi)affine mapping patterns."""
from enum import IntEnum
-import tvm.ffi
+import tvm_ffi
from tvm.runtime import Object
from tvm.ir import PrimExpr
from . import _ffi_api
-@tvm.ffi.register_object("arith.IterMapExpr")
+@tvm_ffi.register_object("arith.IterMapExpr")
class IterMapExpr(PrimExpr):
"""Base class of all IterMap expressions."""
-@tvm.ffi.register_object("arith.IterMark")
+@tvm_ffi.register_object("arith.IterMark")
class IterMark(Object):
"""Mark the source as an iterator in [0, extent).
@@ -44,7 +44,7 @@ class IterMark(Object):
self.__init_handle_by_constructor__(_ffi_api.IterMark, source, extent)
-@tvm.ffi.register_object("arith.IterSplitExpr")
+@tvm_ffi.register_object("arith.IterSplitExpr")
class IterSplitExpr(IterMapExpr):
"""Split of an iterator.
@@ -71,7 +71,7 @@ class IterSplitExpr(IterMapExpr):
)
-@tvm.ffi.register_object("arith.IterSumExpr")
+@tvm_ffi.register_object("arith.IterSumExpr")
class IterSumExpr(IterMapExpr):
"""Fuse multiple iterators by summing them with scaling.
@@ -90,7 +90,7 @@ class IterSumExpr(IterMapExpr):
self.__init_handle_by_constructor__(_ffi_api.IterSumExpr, args, base)
-@tvm.ffi.register_object("arith.IterMapResult")
+@tvm_ffi.register_object("arith.IterMapResult")
class IterMapResult(Object):
"""Result of iter map detection."""
diff --git a/python/tvm/base.py b/python/tvm/base.py
index 63e097999c..8e88364e26 100644
--- a/python/tvm/base.py
+++ b/python/tvm/base.py
@@ -62,7 +62,7 @@ _RUNTIME_ONLY = "runtime" in _LIB_NAME
if _RUNTIME_ONLY:
- from .ffi import registry as _tvm_ffi_registry
+ from tvm_ffi import registry as _tvm_ffi_registry
_tvm_ffi_registry._SKIP_UNKNOWN_OBJECTS = True
diff --git a/python/tvm/contrib/coreml_runtime.py b/python/tvm/contrib/coreml_runtime.py
index def5d3c2e0..34e0681d31 100644
--- a/python/tvm/contrib/coreml_runtime.py
+++ b/python/tvm/contrib/coreml_runtime.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""CoreML runtime that load and run coreml models."""
-import tvm.ffi
+import tvm_ffi
from ..rpc import base as rpc_base
@@ -41,7 +41,7 @@ def create(symbol, compiled_model_path, device):
if device_type >= rpc_base.RPC_SESS_MASK:
fcreate = device._rpc_sess.get_function(runtime_func)
else:
- fcreate = tvm.ffi.get_global_func(runtime_func)
+ fcreate = tvm_ffi.get_global_func(runtime_func)
assert fcreate, "Cannot find `tvm.coreml_runtime.create` function."
return CoreMLModule(fcreate(symbol, compiled_model_path))
diff --git a/python/tvm/contrib/cudnn.py b/python/tvm/contrib/cudnn.py
index 1c80d4a3b9..4d39dfd1c6 100644
--- a/python/tvm/contrib/cudnn.py
+++ b/python/tvm/contrib/cudnn.py
@@ -20,7 +20,7 @@ import ctypes
import numpy as np
import tvm
-import tvm.ffi
+import tvm_ffi
from tvm import te
# algos can be read from cudnn.h
@@ -349,7 +349,7 @@ def _conv_find_algo(
dims - 2, pad, stride, dilation, x_shape, w_shape
)
yshape = np.array(y_shape, dtype=np.int32)
- func = tvm.ffi.get_global_func(func_name)
+ func = tvm_ffi.get_global_func(func_name)
return func(
tensor_format,
dims - 2,
diff --git a/python/tvm/contrib/cutlass/_ffi_api.py b/python/tvm/contrib/cutlass/_ffi_api.py
index be71b0d48f..25393a8f99 100644
--- a/python/tvm/contrib/cutlass/_ffi_api.py
+++ b/python/tvm/contrib/cutlass/_ffi_api.py
@@ -15,6 +15,6 @@
# specific language governing permissions and limitations
# under the License.
"""FFI API for CUTLASS BYOC."""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("contrib.cutlass", __name__)
+tvm_ffi._init_api("contrib.cutlass", __name__)
diff --git a/python/tvm/contrib/cutlass/build.py b/python/tvm/contrib/cutlass/build.py
index 0aea5bf141..41f25dd8a5 100644
--- a/python/tvm/contrib/cutlass/build.py
+++ b/python/tvm/contrib/cutlass/build.py
@@ -26,7 +26,7 @@ from typing import Optional, Sequence
import tvm
from tvm import relax, runtime
-from tvm.ffi.registry import register_func
+from tvm_ffi.registry import register_func
from tvm.contrib.nvcc import get_cuda_version
from tvm.topi.utils import get_const_tuple
diff --git a/python/tvm/contrib/cutlass/gen_tensor_op.py b/python/tvm/contrib/cutlass/gen_tensor_op.py
index c594b3897a..e10abf113e 100644
--- a/python/tvm/contrib/cutlass/gen_tensor_op.py
+++ b/python/tvm/contrib/cutlass/gen_tensor_op.py
@@ -24,7 +24,7 @@ import re
import subprocess
import tempfile
-import tvm.ffi
+import tvm_ffi
from tvm.runtime import Object
from tvm.tir import IntImm
@@ -461,7 +461,7 @@ def _get_optional_int_annotation(annotations, key, default=None):
return int(value)
-@tvm.ffi.register_func("contrib.cutlass.instantiate_template")
+@tvm_ffi.register_func("contrib.cutlass.instantiate_template")
def instantiate_template(func_name, annotations, func_args):
"""Return CUTLASS host code based on a template and the provided annotations.
diff --git a/python/tvm/contrib/hexagon/build.py b/python/tvm/contrib/hexagon/build.py
index f4b02ff80f..eb1a0342c7 100644
--- a/python/tvm/contrib/hexagon/build.py
+++ b/python/tvm/contrib/hexagon/build.py
@@ -35,7 +35,7 @@ import tempfile
from typing import Union
from tvm.contrib.hexagon.hexagon_profiler import HexagonProfiler
-from ...ffi import libinfo
+from tvm_ffi import libinfo
from .session import Session
from .tools import HEXAGON_SIMULATOR_NAME
diff --git a/python/tvm/contrib/hexagon/tools.py b/python/tvm/contrib/hexagon/tools.py
index f7f22db721..a26822bc5f 100644
--- a/python/tvm/contrib/hexagon/tools.py
+++ b/python/tvm/contrib/hexagon/tools.py
@@ -29,7 +29,7 @@ import numpy
import tvm
import tvm.contrib.cc as cc
-from ...ffi.registry import register_func
+from tvm_ffi import register_func
# Linking Hexagon shared libraries.
diff --git a/python/tvm/contrib/miopen.py b/python/tvm/contrib/miopen.py
index 22b08f38ca..3aa885f545 100644
--- a/python/tvm/contrib/miopen.py
+++ b/python/tvm/contrib/miopen.py
@@ -19,7 +19,7 @@
import ctypes
import numpy as np
import tvm
-import tvm.ffi
+import tvm_ffi
from tvm import te
@@ -94,7 +94,7 @@ def conv2d_forward(
oshape = np.zeros((len(x.shape)), dtype=np.int32)
xshape = x.shape
wshape = w.shape
- setup_func = tvm.ffi.get_global_func("tvm.contrib.miopen.conv2d.setup")
+ setup_func = tvm_ffi.get_global_func("tvm.contrib.miopen.conv2d.setup")
algo = setup_func(
conv_mode,
data_type,
diff --git a/python/tvm/contrib/mrvl.py b/python/tvm/contrib/mrvl.py
index 36c932cd1a..87e796a515 100644
--- a/python/tvm/contrib/mrvl.py
+++ b/python/tvm/contrib/mrvl.py
@@ -24,10 +24,10 @@ import tempfile
import base64
import numpy as np
import tvm
-import tvm.ffi
+import tvm_ffi
-@tvm.ffi.register_func("tvm.mrvl.find_value_in_KV_pair")
+@tvm_ffi.register_func("tvm.mrvl.find_value_in_KV_pair")
def find_value_in_KV_pair(json_input: str, key_to_find: str) -> str:
"""This function takes the graph_json string and key to be searched in
the json string, using json parser routine it loads the json string
@@ -54,7 +54,7 @@ def find_value_in_KV_pair(json_input: str, key_to_find: str) -> str:
return value
-@tvm.ffi.register_func("tvm.mrvl.GetNodesJSONString")
+@tvm_ffi.register_func("tvm.mrvl.GetNodesJSONString")
def get_nodes_json_string(graph_json):
"""This takes the graph_json string from MrvlJSONSerializer and adds / modifies
the json string to a form suitable for the Marvell Backend.
@@ -206,7 +206,7 @@ def get_nodes_json_string(graph_json):
return nodes_json_string
-@tvm.ffi.register_func("tvm.mrvl.ModifyConstNames")
+@tvm_ffi.register_func("tvm.mrvl.ModifyConstNames")
def modify_const_names(nodes_json_str, consts_json_str):
"""This takes the graph module returned by build an generates nodes and constant
meta data suitable for compilation by the back end.
@@ -329,7 +329,7 @@ def get_working_dir():
return os.getcwd()
-@tvm.ffi.register_func("tvm.mrvl.WriteJsonFile")
+@tvm_ffi.register_func("tvm.mrvl.WriteJsonFile")
def write_json_file(json_string, json_filename):
"""Generate json file under working directory"""
working_dir = get_working_dir()
@@ -351,7 +351,7 @@ def delete_temp_files(symbol_name):
shutil.rmtree(bin_folder)
-@tvm.ffi.register_func("tvm.mrvl.CompileModel")
+@tvm_ffi.register_func("tvm.mrvl.CompileModel")
def compile_model(
symbol_name,
nodes_json_string,
@@ -414,7 +414,7 @@ def compile_model(
raise RuntimeError(error_msg)
-@tvm.ffi.register_func("tvm.mrvl.CleanUpSim")
+@tvm_ffi.register_func("tvm.mrvl.CleanUpSim")
def clean_up_sim(bin_file, input_json, input_bin, out_bin_prefix, num_outputs):
os.remove(bin_file)
os.remove(input_json)
@@ -424,7 +424,7 @@ def clean_up_sim(bin_file, input_json, input_bin, out_bin_prefix, num_outputs):
os.remove(out_bin)
-@tvm.ffi.register_func("tvm.mrvl.SearchPath")
+@tvm_ffi.register_func("tvm.mrvl.SearchPath")
def search_path(file_name):
path = shutil.which(file_name)
if path is None:
@@ -432,7 +432,7 @@ def search_path(file_name):
return os.path.dirname(path)
-@tvm.ffi.register_func("tvm.mrvl.JsonToBin")
+@tvm_ffi.register_func("tvm.mrvl.JsonToBin")
def convert_json_to_bin(json_file, input_bin_file):
with open(json_file) as input_json:
data = json.load(input_json)
@@ -442,7 +442,7 @@ def convert_json_to_bin(json_file, input_bin_file):
f.write(data_b)
-@tvm.ffi.register_func("tvm.mrvl.RunSim")
+@tvm_ffi.register_func("tvm.mrvl.RunSim")
def run_simulation(run_command, sim_directory):
cwd_path = get_working_dir()
os.mkdir(sim_directory)
@@ -452,6 +452,6 @@ def run_simulation(run_command, sim_directory):
shutil.rmtree(sim_directory)
-@tvm.ffi.register_func("tvm.mrvl.TempDir")
+@tvm_ffi.register_func("tvm.mrvl.TempDir")
def get_temp_dir():
return tempfile.gettempdir()
diff --git a/python/tvm/contrib/msc/core/_ffi_api.py b/python/tvm/contrib/msc/core/_ffi_api.py
index f7c975aff9..a8f3614639 100644
--- a/python/tvm/contrib/msc/core/_ffi_api.py
+++ b/python/tvm/contrib/msc/core/_ffi_api.py
@@ -16,6 +16,6 @@
# under the License.
"""tvm.contrib.msc.core._ffi_api"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("msc.core", __name__)
+tvm_ffi._init_api("msc.core", __name__)
diff --git a/python/tvm/contrib/msc/core/ir/graph.py b/python/tvm/contrib/msc/core/ir/graph.py
index 7bd88df5f6..0f81fd4ab9 100644
--- a/python/tvm/contrib/msc/core/ir/graph.py
+++ b/python/tvm/contrib/msc/core/ir/graph.py
@@ -25,7 +25,7 @@ from tvm.contrib.msc.core import _ffi_api
from tvm.contrib.msc.core import utils as msc_utils
-@tvm.ffi.register_object("msc.core.MSCTensor")
+@tvm_ffi.register_object("msc.core.MSCTensor")
class MSCTensor(Object):
"""Tensor in MSCGraph
@@ -194,12 +194,12 @@ class MSCTensor(Object):
return len(self.shape)
-@tvm.ffi.register_object("msc.core.BaseJoint")
+@tvm_ffi.register_object("msc.core.BaseJoint")
class BaseJoint(Object):
"""Base class of all MSC Nodes."""
-@tvm.ffi.register_object("msc.core.MSCJoint")
+@tvm_ffi.register_object("msc.core.MSCJoint")
class MSCJoint(BaseJoint):
"""Node in MSCGraph
@@ -424,7 +424,7 @@ class MSCJoint(BaseJoint):
return msc_utils.dict_equal(self.get_attrs(), other.get_attrs())
-@tvm.ffi.register_object("msc.core.MSCPrim")
+@tvm_ffi.register_object("msc.core.MSCPrim")
class MSCPrim(BaseJoint):
"""Prim in MSCGraph
@@ -448,7 +448,7 @@ class MSCPrim(BaseJoint):
self.__init_handle_by_constructor__(_ffi_api.MSCPrim, index, name, optype, attrs, parents)
-@tvm.ffi.register_object("msc.core.WeightJoint")
+@tvm_ffi.register_object("msc.core.WeightJoint")
class WeightJoint(BaseJoint):
"""Node in WeightGraph
@@ -562,12 +562,12 @@ class WeightJoint(BaseJoint):
return bool(_ffi_api.WeightJointHasAttr(self, key))
-@tvm.ffi.register_object("msc.core.BaseGraph")
+@tvm_ffi.register_object("msc.core.BaseGraph")
class BaseGraph(Object):
"""Base class of all MSC Graphs."""
-@tvm.ffi.register_object("msc.core.MSCGraph")
+@tvm_ffi.register_object("msc.core.MSCGraph")
class MSCGraph(BaseGraph):
"""The MSCGraph
@@ -956,7 +956,7 @@ class MSCGraph(BaseGraph):
return graph_proto
-@tvm.ffi.register_object("msc.core.WeightGraph")
+@tvm_ffi.register_object("msc.core.WeightGraph")
class WeightGraph(BaseGraph):
"""The WeightGraph
diff --git a/python/tvm/contrib/msc/framework/tensorflow/_ffi_api.py b/python/tvm/contrib/msc/framework/tensorflow/_ffi_api.py
index 5b85e16a53..fef10823de 100644
--- a/python/tvm/contrib/msc/framework/tensorflow/_ffi_api.py
+++ b/python/tvm/contrib/msc/framework/tensorflow/_ffi_api.py
@@ -16,6 +16,6 @@
# under the License.
"""tvm.contrib.msc.framework.tensorflow._ffi_api"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("msc.framework.tensorflow", __name__)
+tvm_ffi._init_api("msc.framework.tensorflow", __name__)
diff --git a/python/tvm/contrib/msc/framework/tensorrt/_ffi_api.py b/python/tvm/contrib/msc/framework/tensorrt/_ffi_api.py
index 4db71f3a19..4dc13bd24b 100644
--- a/python/tvm/contrib/msc/framework/tensorrt/_ffi_api.py
+++ b/python/tvm/contrib/msc/framework/tensorrt/_ffi_api.py
@@ -16,6 +16,6 @@
# under the License.
"""tvm.contrib.msc.framework.tensorrt._ffi_api"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("msc.framework.tensorrt", __name__)
+tvm_ffi._init_api("msc.framework.tensorrt", __name__)
diff --git a/python/tvm/contrib/msc/framework/torch/_ffi_api.py b/python/tvm/contrib/msc/framework/torch/_ffi_api.py
index d12fcf2e2f..9ea5136048 100644
--- a/python/tvm/contrib/msc/framework/torch/_ffi_api.py
+++ b/python/tvm/contrib/msc/framework/torch/_ffi_api.py
@@ -16,6 +16,6 @@
# under the License.
"""tvm.contrib.msc.framework.torch._ffi_api"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("msc.framework.torch", __name__)
+tvm_ffi._init_api("msc.framework.torch", __name__)
diff --git a/python/tvm/contrib/msc/framework/tvm/_ffi_api.py b/python/tvm/contrib/msc/framework/tvm/_ffi_api.py
index a3683181b0..dc75eed418 100644
--- a/python/tvm/contrib/msc/framework/tvm/_ffi_api.py
+++ b/python/tvm/contrib/msc/framework/tvm/_ffi_api.py
@@ -16,6 +16,6 @@
# under the License.
"""tvm.contrib.msc.framework.tvm._ffi_api"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("msc.framework.tvm", __name__)
+tvm_ffi._init_api("msc.framework.tvm", __name__)
diff --git a/python/tvm/contrib/msc/plugin/_ffi_api.py b/python/tvm/contrib/msc/plugin/_ffi_api.py
index c566d3b0d3..8bb42c8c02 100644
--- a/python/tvm/contrib/msc/plugin/_ffi_api.py
+++ b/python/tvm/contrib/msc/plugin/_ffi_api.py
@@ -16,6 +16,6 @@
# under the License.
"""tvm.contrib.msc.plugin._ffi_api"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("msc.plugin", __name__)
+tvm_ffi._init_api("msc.plugin", __name__)
diff --git a/python/tvm/contrib/msc/plugin/op/_ffi_api.py b/python/tvm/contrib/msc/plugin/op/_ffi_api.py
index 0d8ad3c5e4..68704bb178 100644
--- a/python/tvm/contrib/msc/plugin/op/_ffi_api.py
+++ b/python/tvm/contrib/msc/plugin/op/_ffi_api.py
@@ -16,6 +16,6 @@
# under the License.
"""tvm.contrib.msc.plugin.op._ffi_api"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("msc.plugin.op", __name__)
+tvm_ffi._init_api("msc.plugin.op", __name__)
diff --git a/python/tvm/contrib/ndk.py b/python/tvm/contrib/ndk.py
index c1441c496a..743f911b48 100644
--- a/python/tvm/contrib/ndk.py
+++ b/python/tvm/contrib/ndk.py
@@ -25,7 +25,7 @@ from typing import Dict
import tempfile
from pathlib import Path
-from ..ffi import register_func
+from tvm_ffi import register_func
from ..base import py_str
from . import utils as _utils, tar as _tar, cc as _cc
from .cc import get_target_by_dump_machine
diff --git a/python/tvm/contrib/nnpack.py b/python/tvm/contrib/nnpack.py
index 1b4f518508..1f1077bf41 100644
--- a/python/tvm/contrib/nnpack.py
+++ b/python/tvm/contrib/nnpack.py
@@ -17,7 +17,7 @@
"""External function interface to NNPACK libraries."""
import tvm
from tvm import te
-import tvm.ffi
+import tvm_ffi
def is_available():
@@ -232,4 +232,4 @@ def convolution_inference_weight_transform(
)
-tvm.ffi._init_api("tvm.contrib.nnpack")
+tvm_ffi._init_api("tvm.contrib.nnpack")
diff --git a/python/tvm/contrib/nvcc.py b/python/tvm/contrib/nvcc.py
index e9d8fac761..cbc88f0ab4 100644
--- a/python/tvm/contrib/nvcc.py
+++ b/python/tvm/contrib/nvcc.py
@@ -23,7 +23,8 @@ import subprocess
import warnings
from typing import Tuple
-import tvm.ffi
+import tvm_ffi
+import tvm
from tvm.target import Target
from ..base import py_str
@@ -88,7 +89,7 @@ def compile_cuda(code, target_format=None, arch=None, options=None, path_target=
temp_code = temp.relpath(f"{file_name}.cu")
temp_target = temp.relpath(f"{file_name}.{target_format}")
- pass_context = tvm.get_global_func("transform.GetCurrentPassContext")()
+ pass_context = tvm_ffi.get_global_func("transform.GetCurrentPassContext")()
kernels_output_dir = (
pass_context.config["cuda.kernels_output_dir"]
if "cuda.kernels_output_dir" in pass_context.config
@@ -311,14 +312,14 @@ def find_nvshmem_paths() -> Tuple[str, str]:
raise RuntimeError("\n".join(error_message))
-@tvm.ffi.register_func
+@tvm_ffi.register_func
def tvm_callback_cuda_compile(code, target): # pylint: disable=unused-argument
"""use nvcc to generate fatbin code for better optimization"""
ptx = compile_cuda(code, target_format="fatbin")
return ptx
-@tvm.ffi.register_func("tvm_callback_libdevice_path")
+@tvm_ffi.register_func("tvm_callback_libdevice_path")
def find_libdevice_path(arch):
"""Utility function to find libdevice
@@ -383,7 +384,7 @@ def callback_libdevice_path(arch):
return ""
-@tvm.ffi.register_func("tvm.contrib.nvcc.get_compute_version")
+@tvm_ffi.register_func("tvm.contrib.nvcc.get_compute_version")
def get_target_compute_version(target=None):
"""Utility function to get compute capability of compilation target.
@@ -528,7 +529,7 @@ def have_cudagraph():
return False
-@tvm.ffi.register_func("tvm.contrib.nvcc.supports_bf16")
+@tvm_ffi.register_func("tvm.contrib.nvcc.supports_bf16")
def have_bf16(compute_version):
"""Either bf16 support is provided in the compute capability or not
@@ -544,7 +545,7 @@ def have_bf16(compute_version):
return False
-@tvm.ffi.register_func("tvm.contrib.nvcc.supports_fp8")
+@tvm_ffi.register_func("tvm.contrib.nvcc.supports_fp8")
def have_fp8(compute_version):
"""Whether fp8 support is provided in the specified compute capability or not
@@ -562,7 +563,7 @@ def have_fp8(compute_version):
return False
-@tvm.ffi.register_func("tvm.contrib.nvcc.supports_fp4")
+@tvm_ffi.register_func("tvm.contrib.nvcc.supports_fp4")
def have_fp4(compute_version):
"""Whether fp4 support is provided in the specified compute capability or not
diff --git a/python/tvm/contrib/random.py b/python/tvm/contrib/random.py
index 6a17693b91..4826399251 100644
--- a/python/tvm/contrib/random.py
+++ b/python/tvm/contrib/random.py
@@ -17,7 +17,7 @@
"""External function interface to random library."""
import tvm
from tvm import te
-import tvm.ffi
+import tvm_ffi
def randint(low, high, size, dtype="int32"):
@@ -112,4 +112,4 @@ def normal(loc, scale, size):
)
-tvm.ffi._init_api("tvm.contrib.random")
+tvm_ffi._init_api("tvm.contrib.random")
diff --git a/python/tvm/contrib/rocm.py b/python/tvm/contrib/rocm.py
index 6e6a985c27..ee9f9e9b79 100644
--- a/python/tvm/contrib/rocm.py
+++ b/python/tvm/contrib/rocm.py
@@ -20,7 +20,7 @@ import subprocess
import os
from os.path import join, exists
-import tvm.ffi
+import tvm_ffi
from tvm.base import py_str
import tvm.runtime
import tvm.target
@@ -99,7 +99,7 @@ def rocm_link(in_file, out_file, lld=None):
raise RuntimeError(msg)
-@tvm.ffi.register_func("tvm_callback_rocm_link")
+@tvm_ffi.register_func("tvm_callback_rocm_link")
def callback_rocm_link(obj_bin):
"""Links object file generated from LLVM to HSA Code Object
@@ -123,7 +123,7 @@ def callback_rocm_link(obj_bin):
return cobj_bin
-@tvm.ffi.register_func("tvm_callback_rocm_bitcode_path")
+@tvm_ffi.register_func("tvm_callback_rocm_bitcode_path")
def callback_rocm_bitcode_path(rocdl_dir=None):
"""Utility function to find ROCm device library bitcodes
@@ -227,7 +227,7 @@ def have_matrixcore(compute_version=None):
return False
-@tvm.ffi.register_func("tvm_callback_rocm_get_arch")
+@tvm_ffi.register_func("tvm_callback_rocm_get_arch")
def get_rocm_arch(rocm_path=None):
"""Utility function to get the AMD GPU architecture
diff --git a/python/tvm/contrib/tflite_runtime.py b/python/tvm/contrib/tflite_runtime.py
index aceeefd248..81c43861c4 100644
--- a/python/tvm/contrib/tflite_runtime.py
+++ b/python/tvm/contrib/tflite_runtime.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""TFLite runtime that load and run tflite models."""
-import tvm.ffi
+import tvm_ffi
from ..rpc import base as rpc_base
@@ -45,7 +45,7 @@ def create(tflite_model_bytes, device, runtime_target="cpu"):
if device_type >= rpc_base.RPC_SESS_MASK:
fcreate = device._rpc_sess.get_function(runtime_func)
else:
- fcreate = tvm.ffi.get_global_func(runtime_func)
+ fcreate = tvm_ffi.get_global_func(runtime_func)
return TFLiteModule(fcreate(bytearray(tflite_model_bytes), device))
diff --git a/python/tvm/contrib/thrust.py b/python/tvm/contrib/thrust.py
index 9a05cfafba..8cf7c59fad 100644
--- a/python/tvm/contrib/thrust.py
+++ b/python/tvm/contrib/thrust.py
@@ -17,7 +17,7 @@
"""Utilities for thrust"""
import logging
-from tvm.ffi import get_global_func
+from tvm_ffi import get_global_func
def maybe_warn(target, func_name):
diff --git a/python/tvm/dlight/analysis/common_analysis.py b/python/tvm/dlight/analysis/common_analysis.py
index a3499274e5..b7f7079c38 100644
--- a/python/tvm/dlight/analysis/common_analysis.py
+++ b/python/tvm/dlight/analysis/common_analysis.py
@@ -20,7 +20,7 @@ from typing import List, Optional, Set, Union
from typing_extensions import Literal
from tvm import ir, tir
-from tvm.ffi import get_global_func
+from tvm_ffi import get_global_func
from tvm.target.target import Target
from tvm.tir import Schedule
from tvm.tir.schedule import BlockRV
diff --git a/python/tvm/driver/_ffi_api.py b/python/tvm/driver/_ffi_api.py
index 1ceecc9c94..b3853345f0 100644
--- a/python/tvm/driver/_ffi_api.py
+++ b/python/tvm/driver/_ffi_api.py
@@ -15,6 +15,6 @@
# specific language governing permissions and limitations
# under the License.
"""FFI APIs for tvm.driver"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("driver", __name__)
+tvm_ffi._init_api("driver", __name__)
diff --git a/python/tvm/error.py b/python/tvm/error.py
index 671f329238..edabbb3a45 100644
--- a/python/tvm/error.py
+++ b/python/tvm/error.py
@@ -25,7 +25,7 @@ copy the examples and raise errors with the same message convention.
Please also refer to :ref:`error-handling-guide`.
"""
-from tvm.ffi import register_error
+from tvm_ffi import register_error
class TVMError(RuntimeError):
diff --git a/python/tvm/exec/disco_worker.py b/python/tvm/exec/disco_worker.py
index 6d1d4b7f33..fc22a50d9b 100644
--- a/python/tvm/exec/disco_worker.py
+++ b/python/tvm/exec/disco_worker.py
@@ -22,7 +22,7 @@ import sys
from typing import Callable
import tvm
-from tvm.ffi import get_global_func, register_func
+from tvm_ffi import get_global_func, register_func
from tvm.runtime import NDArray, ShapeTuple, String
from tvm.runtime.ndarray import array
diff --git a/python/tvm/relax/_ffi_api.py b/python/tvm/ffi.py
similarity index 90%
copy from python/tvm/relax/_ffi_api.py
copy to python/tvm/ffi.py
index db1ca05586..88fa903a92 100644
--- a/python/tvm/relax/_ffi_api.py
+++ b/python/tvm/ffi.py
@@ -14,7 +14,6 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""FFI API for Relax."""
-import tvm.ffi
-
-tvm.ffi._init_api("relax", __name__)
+# pylint: disable=wildcard-import
+"""Redirects to tvm_ffi"""
+from tvm_ffi import *
diff --git a/python/tvm/ir/_ffi_analysis_api.py b/python/tvm/ir/_ffi_analysis_api.py
index ca38c2309f..6ba65fe264 100644
--- a/python/tvm/ir/_ffi_analysis_api.py
+++ b/python/tvm/ir/_ffi_analysis_api.py
@@ -16,7 +16,7 @@
# under the License.
"""FFI APIs for tvm.ir.analysis"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("ir.analysis", __name__)
+tvm_ffi._init_api("ir.analysis", __name__)
diff --git a/python/tvm/ir/_ffi_api.py b/python/tvm/ir/_ffi_api.py
index 6434a3925e..6165d5ea0b 100644
--- a/python/tvm/ir/_ffi_api.py
+++ b/python/tvm/ir/_ffi_api.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""FFI APIs for tvm.ir"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("ir", __name__)
+tvm_ffi._init_api("ir", __name__)
diff --git a/python/tvm/ir/_ffi_instrument_api.py b/python/tvm/ir/_ffi_instrument_api.py
index d88faf7fdd..af0a0ea3eb 100644
--- a/python/tvm/ir/_ffi_instrument_api.py
+++ b/python/tvm/ir/_ffi_instrument_api.py
@@ -15,6 +15,6 @@
# specific language governing permissions and limitations
# under the License.
"""FFI APIs for tvm.instrument"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("instrument", __name__)
+tvm_ffi._init_api("instrument", __name__)
diff --git a/python/tvm/ir/_ffi_transform_api.py b/python/tvm/ir/_ffi_transform_api.py
index 1a27fc5877..eda8d5354b 100644
--- a/python/tvm/ir/_ffi_transform_api.py
+++ b/python/tvm/ir/_ffi_transform_api.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""FFI APIs for tvm.transform"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("transform", __name__)
+tvm_ffi._init_api("transform", __name__)
diff --git a/python/tvm/ir/attrs.py b/python/tvm/ir/attrs.py
index cab982f4e7..fb408cdb8c 100644
--- a/python/tvm/ir/attrs.py
+++ b/python/tvm/ir/attrs.py
@@ -15,14 +15,14 @@
# specific language governing permissions and limitations
# under the License.
"""TVM Attribute module, which is mainly used for defining attributes of operators."""
-import tvm.ffi
+import tvm_ffi
from tvm.runtime import Object
import tvm.runtime._ffi_node_api
from . import _ffi_api
-@tvm.ffi.register_object("ir.Attrs")
+@tvm_ffi.register_object("ir.Attrs")
class Attrs(Object):
"""Attribute node, which is mainly use for defining attributes of operators.
@@ -73,7 +73,7 @@ class Attrs(Object):
return getattr(self, item)
-@tvm.ffi.register_object("ir.DictAttrs")
+@tvm_ffi.register_object("ir.DictAttrs")
class DictAttrs(Attrs):
"""Dictionary attributes."""
diff --git a/python/tvm/ir/base.py b/python/tvm/ir/base.py
index 088ca6b965..a26a80af52 100644
--- a/python/tvm/ir/base.py
+++ b/python/tvm/ir/base.py
@@ -15,9 +15,9 @@
# specific language governing permissions and limitations
# under the License.
"""Common base structures."""
-import tvm.ffi
+import tvm_ffi
import tvm.error
-from tvm.ffi import get_global_func, register_object
+from tvm_ffi import get_global_func, register_object
from tvm.runtime import Object, _ffi_node_api
from . import _ffi_api, json_compact
diff --git a/python/tvm/ir/container.py b/python/tvm/ir/container.py
index 4bc6fcae21..eecc78cba6 100644
--- a/python/tvm/ir/container.py
+++ b/python/tvm/ir/container.py
@@ -15,6 +15,6 @@
# specific language governing permissions and limitations
# under the License.
"""Additional container data structures used across IR variants."""
-from tvm.ffi import Array, Map
+from tvm_ffi import Array, Map
__all__ = ["Array", "Map"]
diff --git a/python/tvm/ir/diagnostics/__init__.py b/python/tvm/ir/diagnostics/__init__.py
index ac4adc3306..3a131b2a14 100644
--- a/python/tvm/ir/diagnostics/__init__.py
+++ b/python/tvm/ir/diagnostics/__init__.py
@@ -22,7 +22,7 @@ three key abstractions: a Diagnostic, the DiagnosticContext,
and the DiagnosticRenderer.
"""
import enum
-import tvm.ffi
+import tvm_ffi
from . import _ffi_api
from ... import get_global_func, register_func, Object
@@ -38,7 +38,7 @@ def get_renderer():
return _ffi_api.GetRenderer()
-@tvm.register_func("diagnostics.override_renderer")
+@tvm_ffi.register_func("diagnostics.override_renderer")
def override_renderer(render_func):
"""
Sets a custom renderer for diagnostics.
@@ -69,7 +69,7 @@ class DiagnosticLevel(enum.IntEnum):
HELP = 50
-@tvm.ffi.register_object("Diagnostic")
+@tvm_ffi.register_object("Diagnostic")
class Diagnostic(Object):
"""A single diagnostic object from TVM."""
@@ -77,7 +77,7 @@ class Diagnostic(Object):
self.__init_handle_by_constructor__(_ffi_api.Diagnostic, level, span, message)
-@tvm.ffi.register_object("DiagnosticRenderer")
+@tvm_ffi.register_object("DiagnosticRenderer")
class DiagnosticRenderer(Object):
"""
A diagnostic renderer, which given a diagnostic context produces a "rendered"
@@ -100,7 +100,7 @@ class DiagnosticRenderer(Object):
# Register the diagnostic context.
-@tvm.ffi.register_object("DiagnosticContext")
+@tvm_ffi.register_object("DiagnosticContext")
class DiagnosticContext(Object):
"""
A diagnostic context which records active errors
diff --git a/python/tvm/ir/diagnostics/_ffi_api.py b/python/tvm/ir/diagnostics/_ffi_api.py
index fb157c9775..0232cac914 100644
--- a/python/tvm/ir/diagnostics/_ffi_api.py
+++ b/python/tvm/ir/diagnostics/_ffi_api.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""FFI for TVM diagnostics."""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("diagnostics", __name__)
+tvm_ffi._init_api("diagnostics", __name__)
diff --git a/python/tvm/ir/expr.py b/python/tvm/ir/expr.py
index 008924c227..19abb6bd1e 100644
--- a/python/tvm/ir/expr.py
+++ b/python/tvm/ir/expr.py
@@ -18,21 +18,22 @@
from numbers import Number
from typing import Optional
-import tvm.ffi
+import tvm
+import tvm_ffi
from ..runtime import Object, Scriptable
from . import _ffi_api
from .base import Node, Span
-@tvm.ffi.register_object("ir.BaseExpr")
+@tvm_ffi.register_object("ir.BaseExpr")
class BaseExpr(Node):
"""Base class of all the expressions."""
span: Optional[Span]
-@tvm.ffi.register_object("ir.PrimExpr")
+@tvm_ffi.register_object("ir.PrimExpr")
class PrimExpr(BaseExpr):
"""Base class of all primitive expressions.
@@ -43,7 +44,7 @@ class PrimExpr(BaseExpr):
dtype: str
-@tvm.ffi.register_object("ir.RelaxExpr")
+@tvm_ffi.register_object("ir.RelaxExpr")
class RelaxExpr(BaseExpr):
"""Base class of all non-primitive expressions."""
@@ -59,7 +60,7 @@ class RelaxExpr(BaseExpr):
return _ffi_api.ExprStructInfo(self)
-@tvm.ffi.register_object("ir.GlobalVar")
+@tvm_ffi.register_object("ir.GlobalVar")
class GlobalVar(RelaxExpr):
"""A global variable in the IR.
@@ -105,7 +106,7 @@ class GlobalVar(RelaxExpr):
raise RuntimeError(f"Do not know how to handle GlobalVar.__call__ for types {arg_types}")
-@tvm.ffi.register_object("ir.Range")
+@tvm_ffi.register_object("ir.Range")
class Range(Node, Scriptable):
"""Represent a range in TVM.
diff --git a/python/tvm/ir/function.py b/python/tvm/ir/function.py
index f6fc42ccbc..75718503aa 100644
--- a/python/tvm/ir/function.py
+++ b/python/tvm/ir/function.py
@@ -19,6 +19,8 @@
from typing import Union, Dict
from enum import IntEnum
+import tvm_ffi
+
import tvm.runtime
from tvm.runtime.object import Object
from .expr import RelaxExpr
@@ -34,7 +36,7 @@ class CallingConv(IntEnum):
DEVICE_KERNEL_LAUNCH = 2
-@tvm.ffi.register_object("ir.BaseFunc")
+@tvm_ffi.register_object("ir.BaseFunc")
class BaseFunc(RelaxExpr):
"""Base class of all functions."""
diff --git a/python/tvm/ir/global_info.py b/python/tvm/ir/global_info.py
index d4b4fdca16..185e10b88c 100644
--- a/python/tvm/ir/global_info.py
+++ b/python/tvm/ir/global_info.py
@@ -16,11 +16,12 @@
# under the License.
"""Global Info."""
import tvm
+import tvm_ffi
from tvm.runtime.object import Object
from . import _ffi_api
-@tvm.ffi.register_object("ir.GlobalInfo")
+@tvm_ffi.register_object("ir.GlobalInfo")
class GlobalInfo(Object):
"""Base node for all global info that can appear in the IR"""
@@ -36,7 +37,7 @@ class GlobalInfo(Object):
return super().__eq__(other)
-@tvm.ffi.register_object("ir.DummyGlobalInfo")
+@tvm_ffi.register_object("ir.DummyGlobalInfo")
class DummyGlobalInfo(GlobalInfo):
def __init__(self) -> None:
self.__init_handle_by_constructor__(
@@ -44,7 +45,7 @@ class DummyGlobalInfo(GlobalInfo):
)
-@tvm.ffi.register_object("ir.VDevice")
+@tvm_ffi.register_object("ir.VDevice")
class VDevice(GlobalInfo):
def __init__(
self,
diff --git a/python/tvm/ir/instrument.py b/python/tvm/ir/instrument.py
index 1e1505858f..7b6749f113 100644
--- a/python/tvm/ir/instrument.py
+++ b/python/tvm/ir/instrument.py
@@ -19,13 +19,13 @@
import inspect
import functools
-import tvm.ffi
+import tvm_ffi
import tvm.runtime
from . import _ffi_instrument_api
-@tvm.ffi.register_object("instrument.PassInstrument")
+@tvm_ffi.register_object("instrument.PassInstrument")
class PassInstrument(tvm.runtime.Object):
"""A pass instrument implementation.
@@ -225,7 +225,7 @@ def pass_instrument(pi_cls=None):
return create_pass_instrument
-@tvm.ffi.register_object("instrument.PassInstrument")
+@tvm_ffi.register_object("instrument.PassInstrument")
class PassTimingInstrument(tvm.runtime.Object):
"""A wrapper to create a passes time instrument that implemented in C++"""
diff --git a/python/tvm/ir/module.py b/python/tvm/ir/module.py
index 3b99db8598..6163528003 100644
--- a/python/tvm/ir/module.py
+++ b/python/tvm/ir/module.py
@@ -20,7 +20,8 @@ from __future__ import annotations
from typing import Dict, Union
-import tvm.ffi
+import tvm
+import tvm_ffi
from tvm.runtime import Scriptable
from tvm.runtime.object import Object
@@ -30,7 +31,7 @@ from .attrs import DictAttrs
from .base import Node
-@tvm.ffi.register_object("ir.IRModule")
+@tvm_ffi.register_object("ir.IRModule")
class IRModule(Node, Scriptable):
"""IRModule that holds functions and type definitions.
diff --git a/python/tvm/ir/op.py b/python/tvm/ir/op.py
index e5111ccc82..5b62d3fe8d 100644
--- a/python/tvm/ir/op.py
+++ b/python/tvm/ir/op.py
@@ -16,13 +16,13 @@
# under the License.
# pylint: disable=invalid-name
"""Primitive operators in the TVM IR."""
-import tvm.ffi
+import tvm_ffi
from . import _ffi_api
from .expr import RelaxExpr
-@tvm.ffi.register_object("ir.Op")
+@tvm_ffi.register_object("ir.Op")
class Op(RelaxExpr):
"""Primitive operator in the IR."""
diff --git a/python/tvm/ir/supply.py b/python/tvm/ir/supply.py
index 2038df4b31..bc38089b22 100644
--- a/python/tvm/ir/supply.py
+++ b/python/tvm/ir/supply.py
@@ -16,11 +16,12 @@
# under the License.
"""Suppliers that are used to guarantee uniqueness of names and GlobalVars."""
import tvm
+import tvm_ffi
from tvm import Object, IRModule
from . import _ffi_api
-@tvm.ffi.register_object("ir.NameSupply")
+@tvm_ffi.register_object("ir.NameSupply")
class NameSupply(Object):
"""NameSupply that can be used to generate unique names.
@@ -77,7 +78,7 @@ class NameSupply(Object):
return _ffi_api.NameSupply_ContainsName(self, name, add_prefix)
-@tvm.ffi.register_object("ir.GlobalVarSupply")
+@tvm_ffi.register_object("ir.GlobalVarSupply")
class GlobalVarSupply(Object):
"""GlobalVarSupply that holds a mapping between names and GlobalVars.
diff --git a/python/tvm/ir/transform.py b/python/tvm/ir/transform.py
index b8f4c36c30..3b9b620081 100644
--- a/python/tvm/ir/transform.py
+++ b/python/tvm/ir/transform.py
@@ -19,13 +19,13 @@
import inspect
import functools
-import tvm.ffi
+import tvm_ffi
import tvm.runtime
from . import _ffi_transform_api
-@tvm.ffi.register_object("transform.PassInfo")
+@tvm_ffi.register_object("transform.PassInfo")
class PassInfo(tvm.runtime.Object):
"""The class contains the meta data required by a pass. It is the
container of information needed by running an optimization or analysis.
@@ -50,7 +50,7 @@ class PassInfo(tvm.runtime.Object):
)
-@tvm.ffi.register_object("transform.PassContext")
+@tvm_ffi.register_object("transform.PassContext")
class PassContext(tvm.runtime.Object):
"""The basis where a TVM optimization/analysis runs on.
Each pass context contains a number of auxiliary information that is used
@@ -138,7 +138,7 @@ class PassContext(tvm.runtime.Object):
return _ffi_transform_api.ListConfigs()
-@tvm.ffi.register_object("transform.Pass")
+@tvm_ffi.register_object("transform.Pass")
class Pass(tvm.runtime.Object):
"""The base class of all passes. All methods here are just simple wrappers
that are implemented in the backend. They are defined for users to
@@ -167,7 +167,7 @@ class Pass(tvm.runtime.Object):
return _ffi_transform_api.RunPass(self, mod)
-@tvm.ffi.register_object("transform.ModulePass")
+@tvm_ffi.register_object("transform.ModulePass")
class ModulePass(Pass):
"""A pass that works on tvm.IRModule. Users don't need to interact with
this class directly. Instead, a module pass should be created through
@@ -178,7 +178,7 @@ class ModulePass(Pass):
"""
-@tvm.ffi.register_object("transform.Sequential")
+@tvm_ffi.register_object("transform.Sequential")
class Sequential(Pass):
"""A pass that works on a sequence of pass objects. Multiple passes can be
executed sequentially using this class.
diff --git a/python/tvm/ir/type.py b/python/tvm/ir/type.py
index 0f287be961..68bed8fb69 100644
--- a/python/tvm/ir/type.py
+++ b/python/tvm/ir/type.py
@@ -16,14 +16,14 @@
# under the License.
"""Unified type system in the project."""
import tvm
-import tvm.ffi
+import tvm_ffi
from tvm.runtime import Scriptable
from . import _ffi_api
from .base import Node
-@tvm.ffi.register_object("ir.Type")
+@tvm_ffi.register_object("ir.Type")
class Type(Node, Scriptable):
"""The base class of all types."""
@@ -39,7 +39,7 @@ class Type(Node, Scriptable):
return super().__eq__(other)
-@tvm.ffi.register_object("ir.PrimType")
+@tvm_ffi.register_object("ir.PrimType")
class PrimType(Type):
"""Primitive data type in the low level IR
@@ -53,7 +53,7 @@ class PrimType(Type):
self.__init_handle_by_constructor__(_ffi_api.PrimType, dtype)
-@tvm.ffi.register_object("ir.PointerType")
+@tvm_ffi.register_object("ir.PointerType")
class PointerType(Type):
"""PointerType used in the low-level TIR.
@@ -70,7 +70,7 @@ class PointerType(Type):
self.__init_handle_by_constructor__(_ffi_api.PointerType, element_type, storage_scope)
-@tvm.ffi.register_object("ir.TupleType")
+@tvm_ffi.register_object("ir.TupleType")
class TupleType(Type):
"""The type of tuple values.
@@ -84,7 +84,7 @@ class TupleType(Type):
self.__init_handle_by_constructor__(_ffi_api.TupleType, fields)
-@tvm.ffi.register_object("ir.FuncType")
+@tvm_ffi.register_object("ir.FuncType")
class FuncType(Type):
"""Function type.
@@ -110,7 +110,7 @@ class FuncType(Type):
)
-@tvm.ffi.register_object("ir.TensorMapType")
+@tvm_ffi.register_object("ir.TensorMapType")
class TensorMapType(Type):
"""TensorMapType used in the low-level TIR.
diff --git a/python/tvm/ir/type_relation.py b/python/tvm/ir/type_relation.py
index d0175fda57..7095095802 100644
--- a/python/tvm/ir/type_relation.py
+++ b/python/tvm/ir/type_relation.py
@@ -15,13 +15,13 @@
# specific language governing permissions and limitations
# under the License.
"""Type relation and function for type checking."""
-import tvm.ffi
+import tvm_ffi
from .type import Type, TypeConstraint
from . import _ffi_api
-@tvm.ffi.register_object("TypeCall")
+@tvm_ffi.register_object("TypeCall")
class TypeCall(Type):
"""Type function application.
@@ -43,7 +43,7 @@ class TypeCall(Type):
self.__init_handle_by_constructor__(_ffi_api.TypeCall, func, args)
-@tvm.ffi.register_object("TypeRelation")
+@tvm_ffi.register_object("TypeRelation")
class TypeRelation(TypeConstraint):
"""User defined type relation, it is an input-output relation on types.
diff --git a/python/tvm/libinfo.py b/python/tvm/libinfo.py
index 2b6d11e0b2..f9f28b6853 100644
--- a/python/tvm/libinfo.py
+++ b/python/tvm/libinfo.py
@@ -232,9 +232,13 @@ def find_include_path(name=None, search_path=None, optional=False):
dmlc_include_path = []
else:
tvm_include_path = [os.path.join(p, "include") for p in header_path]
- tvm_ffi_include_path = [os.path.join(p, "ffi/include") for p in header_path]
- dlpack_include_path = [os.path.join(p, "dlpack/include") for p in header_path]
- dmlc_include_path = [os.path.join(p, "dmlc-core/include") for p in header_path]
+ tvm_ffi_include_path = [os.path.join(p, "ffi", "include") for p in header_path]
+ dlpack_include_path = [
+ os.path.join(p, "ffi", "3rdparty", "dlpack", "include") for p in header_path
+ ]
+ dmlc_include_path = [
+ os.path.join(p, "3rdparty", "dmlc-core", "include") for p in header_path
+ ]
# try to find include path
include_found = [p for p in tvm_include_path if os.path.exists(p) and os.path.isdir(p)]
diff --git a/python/tvm/meta_schedule/_ffi_api.py b/python/tvm/meta_schedule/_ffi_api.py
index 89b8df0860..bb07a22573 100644
--- a/python/tvm/meta_schedule/_ffi_api.py
+++ b/python/tvm/meta_schedule/_ffi_api.py
@@ -15,6 +15,6 @@
# specific language governing permissions and limitations
# under the License.
"""FFI APIs for tvm.meta_schedule"""
-from ..ffi import _init_api
+from tvm_ffi import _init_api
_init_api("meta_schedule", __name__) # pylint: disable=protected-access
diff --git a/python/tvm/meta_schedule/arg_info.py b/python/tvm/meta_schedule/arg_info.py
index 69c8d6d4c5..3f8d721ed1 100644
--- a/python/tvm/meta_schedule/arg_info.py
+++ b/python/tvm/meta_schedule/arg_info.py
@@ -17,7 +17,7 @@
"""The argument information"""
from typing import Any, List, Union
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from tvm.ir import IRModule
from tvm.runtime import DataType, Object, ShapeTuple
from tvm.tir import PrimFunc
diff --git a/python/tvm/meta_schedule/builder/builder.py b/python/tvm/meta_schedule/builder/builder.py
index f323e15bd5..3383ef55ad 100644
--- a/python/tvm/meta_schedule/builder/builder.py
+++ b/python/tvm/meta_schedule/builder/builder.py
@@ -21,7 +21,7 @@ from typing import Callable, Dict, List, Optional, Union
from typing_extensions import Literal
# isort: on
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from tvm.ir import IRModule
from tvm.runtime import NDArray, Object
from tvm.target import Target
diff --git a/python/tvm/meta_schedule/builder/local_builder.py b/python/tvm/meta_schedule/builder/local_builder.py
index 0f68ef7afb..297d6cb610 100644
--- a/python/tvm/meta_schedule/builder/local_builder.py
+++ b/python/tvm/meta_schedule/builder/local_builder.py
@@ -19,7 +19,7 @@ import os
import tempfile
from typing import Callable, Dict, List, Optional, Union
-from tvm.ffi import register_func
+from tvm_ffi import register_func
from tvm.ir import IRModule
from tvm.runtime import Module, NDArray, load_param_dict, save_param_dict
from tvm.target import Target
diff --git a/python/tvm/meta_schedule/cost_model/cost_model.py b/python/tvm/meta_schedule/cost_model/cost_model.py
index 9abd50b94c..f51d2f2ac8 100644
--- a/python/tvm/meta_schedule/cost_model/cost_model.py
+++ b/python/tvm/meta_schedule/cost_model/cost_model.py
@@ -24,7 +24,7 @@ from typing_extensions import Literal
# isort: on
import numpy as np # type: ignore
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from tvm.runtime import Object
from .. import _ffi_api
diff --git a/python/tvm/meta_schedule/database/database.py b/python/tvm/meta_schedule/database/database.py
index 7abaead680..08bcbd33c7 100644
--- a/python/tvm/meta_schedule/database/database.py
+++ b/python/tvm/meta_schedule/database/database.py
@@ -22,7 +22,7 @@ from typing_extensions import Literal
# isort: on
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from tvm.ir.module import IRModule
from tvm.runtime import Object
from tvm.target import Target
diff --git a/python/tvm/meta_schedule/database/json_database.py b/python/tvm/meta_schedule/database/json_database.py
index f3b1884937..cdf08c6e03 100644
--- a/python/tvm/meta_schedule/database/json_database.py
+++ b/python/tvm/meta_schedule/database/json_database.py
@@ -18,7 +18,7 @@
import os.path as osp
from typing import Optional
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from .. import _ffi_api
from .database import Database
diff --git a/python/tvm/meta_schedule/database/memory_database.py b/python/tvm/meta_schedule/database/memory_database.py
index 5375533383..69b129ec21 100644
--- a/python/tvm/meta_schedule/database/memory_database.py
+++ b/python/tvm/meta_schedule/database/memory_database.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""A database that stores TuningRecords in memory"""
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from .. import _ffi_api
from .database import Database
diff --git a/python/tvm/meta_schedule/database/ordered_union_database.py b/python/tvm/meta_schedule/database/ordered_union_database.py
index a451d8ee2f..717d2f3001 100644
--- a/python/tvm/meta_schedule/database/ordered_union_database.py
+++ b/python/tvm/meta_schedule/database/ordered_union_database.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""A database consists of multiple databases."""
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from .. import _ffi_api
from .database import Database
diff --git a/python/tvm/meta_schedule/database/schedule_fn_database.py b/python/tvm/meta_schedule/database/schedule_fn_database.py
index 3b7dfa79f6..477c5664fd 100644
--- a/python/tvm/meta_schedule/database/schedule_fn_database.py
+++ b/python/tvm/meta_schedule/database/schedule_fn_database.py
@@ -17,7 +17,7 @@
"""A database for injecting handcrafted schedule functions."""
from typing import Callable
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from tvm.tir import Schedule
from .. import _ffi_api
diff --git a/python/tvm/meta_schedule/database/union_database.py b/python/tvm/meta_schedule/database/union_database.py
index 7f896c1da6..3a1afbe32a 100644
--- a/python/tvm/meta_schedule/database/union_database.py
+++ b/python/tvm/meta_schedule/database/union_database.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""A database consists of multiple databases."""
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from .. import _ffi_api
from .database import Database
diff --git a/python/tvm/meta_schedule/extracted_task.py b/python/tvm/meta_schedule/extracted_task.py
index 0cdede120b..df66e774e5 100644
--- a/python/tvm/meta_schedule/extracted_task.py
+++ b/python/tvm/meta_schedule/extracted_task.py
@@ -17,7 +17,7 @@
"""Extracted tasks from high-level IR."""
from typing import List
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from tvm.ir import IRModule
from tvm.runtime import Object
from tvm.target import Target
diff --git a/python/tvm/meta_schedule/feature_extractor/feature_extractor.py b/python/tvm/meta_schedule/feature_extractor/feature_extractor.py
index bd37214db9..d4c68fcb93 100644
--- a/python/tvm/meta_schedule/feature_extractor/feature_extractor.py
+++ b/python/tvm/meta_schedule/feature_extractor/feature_extractor.py
@@ -22,7 +22,7 @@ from typing_extensions import Literal
# isort: on
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from tvm.runtime import Object
from tvm.runtime.ndarray import NDArray
diff --git a/python/tvm/meta_schedule/feature_extractor/per_store_feature.py b/python/tvm/meta_schedule/feature_extractor/per_store_feature.py
index b1098bd4ea..673a722955 100644
--- a/python/tvm/meta_schedule/feature_extractor/per_store_feature.py
+++ b/python/tvm/meta_schedule/feature_extractor/per_store_feature.py
@@ -18,7 +18,7 @@
"""We extract one feature vector per BufferStoreNode statement in a TIR Stmt,
so we call this feature as "per-store" feature.
"""
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from .. import _ffi_api
from .feature_extractor import FeatureExtractor
diff --git a/python/tvm/meta_schedule/measure_callback/add_to_database.py b/python/tvm/meta_schedule/measure_callback/add_to_database.py
index f40dffeaad..e0a6f5a273 100644
--- a/python/tvm/meta_schedule/measure_callback/add_to_database.py
+++ b/python/tvm/meta_schedule/measure_callback/add_to_database.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""A callback that adds the measurement results into the database"""
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from .. import _ffi_api
from .measure_callback import MeasureCallback
diff --git a/python/tvm/meta_schedule/measure_callback/measure_callback.py b/python/tvm/meta_schedule/measure_callback/measure_callback.py
index 17a7f45460..885f70e88d 100644
--- a/python/tvm/meta_schedule/measure_callback/measure_callback.py
+++ b/python/tvm/meta_schedule/measure_callback/measure_callback.py
@@ -23,7 +23,7 @@ from typing_extensions import Literal
# isort: on
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from tvm.runtime import Object
from .. import _ffi_api
diff --git a/python/tvm/meta_schedule/measure_callback/remove_build_artifact.py b/python/tvm/meta_schedule/measure_callback/remove_build_artifact.py
index 82c18f8f90..23808b7e99 100644
--- a/python/tvm/meta_schedule/measure_callback/remove_build_artifact.py
+++ b/python/tvm/meta_schedule/measure_callback/remove_build_artifact.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""A callback that removes the build artifacts from the disk"""
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from .. import _ffi_api
from .measure_callback import MeasureCallback
diff --git a/python/tvm/meta_schedule/measure_callback/update_cost_model.py b/python/tvm/meta_schedule/measure_callback/update_cost_model.py
index 5b8b0306d4..7cf60c095b 100644
--- a/python/tvm/meta_schedule/measure_callback/update_cost_model.py
+++ b/python/tvm/meta_schedule/measure_callback/update_cost_model.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""A measure callback that updates the cost model"""
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from .. import _ffi_api
from .measure_callback import MeasureCallback
diff --git a/python/tvm/meta_schedule/mutator/mutate_compute_location.py b/python/tvm/meta_schedule/mutator/mutate_compute_location.py
index 5ebe04a6b1..620e0062cb 100644
--- a/python/tvm/meta_schedule/mutator/mutate_compute_location.py
+++ b/python/tvm/meta_schedule/mutator/mutate_compute_location.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""A mutator that mutates the compute-at location decision of SampleComputeLocation"""
-from tvm.ffi.registry import register_object
+from tvm_ffi.registry import register_object
from .. import _ffi_api
from .mutator import Mutator
diff --git a/python/tvm/meta_schedule/mutator/mutate_parallel.py b/python/tvm/meta_schedule/mutator/mutate_parallel.py
index c7736fdcf7..fc077cd0d4 100644
--- a/python/tvm/meta_schedule/mutator/mutate_parallel.py
+++ b/python/tvm/meta_schedule/mutator/mutate_parallel.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""Mutator that mutates the parallel extent"""
-from tvm.ffi.registry import register_object
+from tvm_ffi.registry import register_object
from .. import _ffi_api
from .mutator import Mutator
diff --git a/python/tvm/meta_schedule/mutator/mutate_thread_binding.py b/python/tvm/meta_schedule/mutator/mutate_thread_binding.py
index 2225ca76c7..4c9fa44c50 100644
--- a/python/tvm/meta_schedule/mutator/mutate_thread_binding.py
+++ b/python/tvm/meta_schedule/mutator/mutate_thread_binding.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""Mutator that mutates the thread binding extent"""
-from tvm.ffi.registry import register_object
+from tvm_ffi.registry import register_object
from .. import _ffi_api
from .mutator import Mutator
diff --git a/python/tvm/meta_schedule/mutator/mutate_tile_size.py b/python/tvm/meta_schedule/mutator/mutate_tile_size.py
index 90cccdc3f5..f40894f5ba 100644
--- a/python/tvm/meta_schedule/mutator/mutate_tile_size.py
+++ b/python/tvm/meta_schedule/mutator/mutate_tile_size.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""Mutator that mutates the decision of instruction Sample-Perfect-Tile"""
-from tvm.ffi.registry import register_object
+from tvm_ffi.registry import register_object
from .. import _ffi_api
from .mutator import Mutator
diff --git a/python/tvm/meta_schedule/mutator/mutate_unroll.py b/python/tvm/meta_schedule/mutator/mutate_unroll.py
index 9575c3fc22..97999c2888 100644
--- a/python/tvm/meta_schedule/mutator/mutate_unroll.py
+++ b/python/tvm/meta_schedule/mutator/mutate_unroll.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""Mutator that mutates auto unroll step"""
-from tvm.ffi.registry import register_object
+from tvm_ffi.registry import register_object
from .. import _ffi_api
from .mutator import Mutator
diff --git a/python/tvm/meta_schedule/mutator/mutator.py b/python/tvm/meta_schedule/mutator/mutator.py
index 6991c72bec..211e2c2b50 100644
--- a/python/tvm/meta_schedule/mutator/mutator.py
+++ b/python/tvm/meta_schedule/mutator/mutator.py
@@ -22,7 +22,7 @@ from typing_extensions import Literal
# isort: on
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from tvm.runtime import Object
from tvm.tir.schedule import Trace
diff --git a/python/tvm/meta_schedule/postproc/disallow_async_strided_mem_copy.py b/python/tvm/meta_schedule/postproc/disallow_async_strided_mem_copy.py
index 5c50b20644..5c18475ea0 100644
--- a/python/tvm/meta_schedule/postproc/disallow_async_strided_mem_copy.py
+++ b/python/tvm/meta_schedule/postproc/disallow_async_strided_mem_copy.py
@@ -16,7 +16,7 @@
# under the License.
"""A postprocessor that checks if the IRModule has any strided memory copies"""
-from tvm.ffi.registry import register_object
+from tvm_ffi.registry import register_object
from .. import _ffi_api
from .postproc import Postproc
diff --git a/python/tvm/meta_schedule/postproc/disallow_dynamic_loop.py b/python/tvm/meta_schedule/postproc/disallow_dynamic_loop.py
index 34c13aded9..da604e42cc 100644
--- a/python/tvm/meta_schedule/postproc/disallow_dynamic_loop.py
+++ b/python/tvm/meta_schedule/postproc/disallow_dynamic_loop.py
@@ -16,7 +16,7 @@
# under the License.
"""A postprocessor that checks if the IRModule has any loop with non-constant extent"""
-from tvm.ffi.registry import register_object
+from tvm_ffi.registry import register_object
from .. import _ffi_api
from .postproc import Postproc
diff --git a/python/tvm/meta_schedule/postproc/postproc.py b/python/tvm/meta_schedule/postproc/postproc.py
index 33daabc395..8e89ad2fe1 100644
--- a/python/tvm/meta_schedule/postproc/postproc.py
+++ b/python/tvm/meta_schedule/postproc/postproc.py
@@ -22,7 +22,7 @@ from typing_extensions import Literal
# isort: on
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from tvm.runtime import Object
from tvm.tir.schedule import Schedule
diff --git a/python/tvm/meta_schedule/postproc/rewrite_cooperative_fetch.py b/python/tvm/meta_schedule/postproc/rewrite_cooperative_fetch.py
index 20c354ce60..d20c22d0f6 100644
--- a/python/tvm/meta_schedule/postproc/rewrite_cooperative_fetch.py
+++ b/python/tvm/meta_schedule/postproc/rewrite_cooperative_fetch.py
@@ -17,7 +17,7 @@
"""A postprocessor that rewrites the cooperative fetch annotation to actual
vectorized cooperative fetching in loop bindings."""
-from tvm.ffi.registry import register_object
+from tvm_ffi.registry import register_object
from .. import _ffi_api
from .postproc import Postproc
diff --git a/python/tvm/meta_schedule/postproc/rewrite_layout.py b/python/tvm/meta_schedule/postproc/rewrite_layout.py
index 13556f1909..73b6dde9f7 100644
--- a/python/tvm/meta_schedule/postproc/rewrite_layout.py
+++ b/python/tvm/meta_schedule/postproc/rewrite_layout.py
@@ -16,7 +16,7 @@
# under the License.
"""A postprocessor that rewrites the layout of input tensor"""
-from tvm.ffi.registry import register_object
+from tvm_ffi.registry import register_object
from .. import _ffi_api
from .postproc import Postproc
diff --git a/python/tvm/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.py b/python/tvm/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.py
index 0be7cdbe11..30235517f9 100644
--- a/python/tvm/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.py
+++ b/python/tvm/meta_schedule/postproc/rewrite_parallel_vectorize_unroll.py
@@ -17,7 +17,7 @@
"""A postprocessor that applies parallelization, vectorization and auto unrolling
according to the annotation of each block"""
-from tvm.ffi.registry import register_object
+from tvm_ffi.registry import register_object
from .. import _ffi_api
from .postproc import Postproc
diff --git a/python/tvm/meta_schedule/postproc/rewrite_reduction_block.py b/python/tvm/meta_schedule/postproc/rewrite_reduction_block.py
index 30c8cf9b06..5bbe2b8838 100644
--- a/python/tvm/meta_schedule/postproc/rewrite_reduction_block.py
+++ b/python/tvm/meta_schedule/postproc/rewrite_reduction_block.py
@@ -16,7 +16,7 @@
# under the License.
"""A postprocessor that rewrites reduction block by moving the init block out."""
-from tvm.ffi.registry import register_object
+from tvm_ffi.registry import register_object
from .. import _ffi_api
from .postproc import Postproc
diff --git a/python/tvm/meta_schedule/postproc/rewrite_tensorize.py b/python/tvm/meta_schedule/postproc/rewrite_tensorize.py
index e04ddcbdf2..8f0edb8695 100644
--- a/python/tvm/meta_schedule/postproc/rewrite_tensorize.py
+++ b/python/tvm/meta_schedule/postproc/rewrite_tensorize.py
@@ -16,7 +16,7 @@
# under the License.
"""A postprocessor that tensorize related components."""
-from tvm.ffi.registry import register_object
+from tvm_ffi.registry import register_object
from .. import _ffi_api
from .postproc import Postproc
diff --git a/python/tvm/meta_schedule/postproc/rewrite_unbound_block.py b/python/tvm/meta_schedule/postproc/rewrite_unbound_block.py
index ca4c9cdcd6..b274c2f55c 100644
--- a/python/tvm/meta_schedule/postproc/rewrite_unbound_block.py
+++ b/python/tvm/meta_schedule/postproc/rewrite_unbound_block.py
@@ -16,7 +16,7 @@
# under the License.
"""A postprocessor that adds thread binding to unbound blocks"""
-from tvm.ffi.registry import register_object
+from tvm_ffi.registry import register_object
from .. import _ffi_api
from .postproc import Postproc
diff --git a/python/tvm/meta_schedule/postproc/verify_gpu_code.py b/python/tvm/meta_schedule/postproc/verify_gpu_code.py
index 1a74eadaa9..48fbe8f4b1 100644
--- a/python/tvm/meta_schedule/postproc/verify_gpu_code.py
+++ b/python/tvm/meta_schedule/postproc/verify_gpu_code.py
@@ -16,7 +16,7 @@
# under the License.
"""A postprocessor that verifies if the GPU code is correct"""
-from tvm.ffi.registry import register_object
+from tvm_ffi.registry import register_object
from .. import _ffi_api
from .postproc import Postproc
diff --git a/python/tvm/meta_schedule/postproc/verify_vtcm_limit.py b/python/tvm/meta_schedule/postproc/verify_vtcm_limit.py
index 51a38624d2..96ece2270b 100644
--- a/python/tvm/meta_schedule/postproc/verify_vtcm_limit.py
+++ b/python/tvm/meta_schedule/postproc/verify_vtcm_limit.py
@@ -16,7 +16,7 @@
# under the License.
"""A postprocessor that verifies the VTCM usage of a given schedule."""
-from tvm.ffi.registry import register_object
+from tvm_ffi.registry import register_object
from .. import _ffi_api
from .postproc import Postproc
diff --git a/python/tvm/meta_schedule/profiler.py b/python/tvm/meta_schedule/profiler.py
index 65c1079d65..1a41f589de 100644
--- a/python/tvm/meta_schedule/profiler.py
+++ b/python/tvm/meta_schedule/profiler.py
@@ -19,7 +19,7 @@
from contextlib import contextmanager
from typing import Dict, Optional
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from tvm.runtime import Object
from . import _ffi_api
diff --git a/python/tvm/meta_schedule/relax_integration.py b/python/tvm/meta_schedule/relax_integration.py
index 613405c8ad..8d041b6caa 100644
--- a/python/tvm/meta_schedule/relax_integration.py
+++ b/python/tvm/meta_schedule/relax_integration.py
@@ -23,7 +23,7 @@ from typing_extensions import Literal
# isort: on
-from tvm.ffi import get_global_func, register_func
+from tvm_ffi import get_global_func, register_func
from tvm.ir import IRModule
from tvm.ir.transform import PassContext
from tvm.runtime import NDArray
diff --git a/python/tvm/meta_schedule/runner/runner.py b/python/tvm/meta_schedule/runner/runner.py
index 0c2609469a..0d7cd32bd7 100644
--- a/python/tvm/meta_schedule/runner/runner.py
+++ b/python/tvm/meta_schedule/runner/runner.py
@@ -22,7 +22,7 @@ from typing_extensions import Literal
# isort: on
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from tvm.runtime import Object
from .. import _ffi_api
diff --git a/python/tvm/meta_schedule/schedule_rule/add_rfactor.py b/python/tvm/meta_schedule/schedule_rule/add_rfactor.py
index ceb18a6c3a..2bef40fffe 100644
--- a/python/tvm/meta_schedule/schedule_rule/add_rfactor.py
+++ b/python/tvm/meta_schedule/schedule_rule/add_rfactor.py
@@ -17,7 +17,7 @@
"""Add-rfactor Rule that add-rfactor to some blocks if needed"""
from typing import Optional
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from .. import _ffi_api
from .schedule_rule import ScheduleRule
diff --git a/python/tvm/meta_schedule/schedule_rule/apply_custom_rule.py b/python/tvm/meta_schedule/schedule_rule/apply_custom_rule.py
index 26f61aa8ce..2e383c75eb 100644
--- a/python/tvm/meta_schedule/schedule_rule/apply_custom_rule.py
+++ b/python/tvm/meta_schedule/schedule_rule/apply_custom_rule.py
@@ -16,7 +16,7 @@
# under the License.
"""Create a rule that applies customized rules registered using block attribute `schedule_rule`.
The rule will be dispatched according to target keys."""
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from .. import _ffi_api
from .schedule_rule import ScheduleRule
diff --git a/python/tvm/meta_schedule/schedule_rule/auto_bind.py b/python/tvm/meta_schedule/schedule_rule/auto_bind.py
index ef34e45061..0704b03f74 100644
--- a/python/tvm/meta_schedule/schedule_rule/auto_bind.py
+++ b/python/tvm/meta_schedule/schedule_rule/auto_bind.py
@@ -17,7 +17,7 @@
"""Auto-bind Rule that binds blocks to threads if needed"""
from typing import List, Optional
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from .. import _ffi_api
from .schedule_rule import ScheduleRule
diff --git a/python/tvm/meta_schedule/schedule_rule/auto_inline.py b/python/tvm/meta_schedule/schedule_rule/auto_inline.py
index 8cd122ec93..b789dd7507 100644
--- a/python/tvm/meta_schedule/schedule_rule/auto_inline.py
+++ b/python/tvm/meta_schedule/schedule_rule/auto_inline.py
@@ -17,7 +17,7 @@
"""Auto-Inline. Rule that inlines spatial blocks if it satisfies some conditions"""
from typing import List, Optional
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from .. import _ffi_api
from .schedule_rule import ScheduleRule
diff --git a/python/tvm/meta_schedule/schedule_rule/cross_thread_reduction.py b/python/tvm/meta_schedule/schedule_rule/cross_thread_reduction.py
index d2c780b728..0c79d4f08b 100644
--- a/python/tvm/meta_schedule/schedule_rule/cross_thread_reduction.py
+++ b/python/tvm/meta_schedule/schedule_rule/cross_thread_reduction.py
@@ -17,7 +17,7 @@
"""Rules which apply cross-thread reduction to some reduction blocks correspondingly when needed"""
from typing import List
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from .. import _ffi_api
from .schedule_rule import ScheduleRule
diff --git a/python/tvm/meta_schedule/schedule_rule/multi_level_tiling.py b/python/tvm/meta_schedule/schedule_rule/multi_level_tiling.py
index 2f389190d6..41c97a7862 100644
--- a/python/tvm/meta_schedule/schedule_rule/multi_level_tiling.py
+++ b/python/tvm/meta_schedule/schedule_rule/multi_level_tiling.py
@@ -18,7 +18,7 @@
from typing import Any, Dict, List, Mapping, NamedTuple, Optional, Callable
from tvm.tir.schedule import Schedule, BlockRV
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from .. import _ffi_api
from .schedule_rule import ScheduleRule
diff --git a/python/tvm/meta_schedule/schedule_rule/parallel_vectorize_unroll.py b/python/tvm/meta_schedule/schedule_rule/parallel_vectorize_unroll.py
index e9626c40e3..259620b3f7 100644
--- a/python/tvm/meta_schedule/schedule_rule/parallel_vectorize_unroll.py
+++ b/python/tvm/meta_schedule/schedule_rule/parallel_vectorize_unroll.py
@@ -18,7 +18,7 @@
each block in a follow-up post processor"""
from typing import List, Optional
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from .. import _ffi_api
from .schedule_rule import ScheduleRule
diff --git a/python/tvm/meta_schedule/schedule_rule/random_compute_location.py b/python/tvm/meta_schedule/schedule_rule/random_compute_location.py
index 81de07afbb..8f1c96f6eb 100644
--- a/python/tvm/meta_schedule/schedule_rule/random_compute_location.py
+++ b/python/tvm/meta_schedule/schedule_rule/random_compute_location.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""Rule that randomly select a compute-at location for a free block"""
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from .. import _ffi_api
from .schedule_rule import ScheduleRule
diff --git a/python/tvm/meta_schedule/schedule_rule/schedule_rule.py b/python/tvm/meta_schedule/schedule_rule/schedule_rule.py
index 5684e68c71..98c81e5b8f 100644
--- a/python/tvm/meta_schedule/schedule_rule/schedule_rule.py
+++ b/python/tvm/meta_schedule/schedule_rule/schedule_rule.py
@@ -25,7 +25,7 @@ from typing_extensions import Literal
# isort: on
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from tvm.runtime import Object
from tvm.tir.schedule import BlockRV, Schedule
diff --git a/python/tvm/meta_schedule/search_strategy/evolutionary_search.py b/python/tvm/meta_schedule/search_strategy/evolutionary_search.py
index 1833ef23bd..04f9310e6e 100644
--- a/python/tvm/meta_schedule/search_strategy/evolutionary_search.py
+++ b/python/tvm/meta_schedule/search_strategy/evolutionary_search.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""Evolutionary Search Strategy"""
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from .. import _ffi_api
from .search_strategy import SearchStrategy
diff --git a/python/tvm/meta_schedule/search_strategy/replay_func.py b/python/tvm/meta_schedule/search_strategy/replay_func.py
index 09e5c58d07..682c9638c5 100644
--- a/python/tvm/meta_schedule/search_strategy/replay_func.py
+++ b/python/tvm/meta_schedule/search_strategy/replay_func.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""Replay Trace Search Strategy"""
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from .. import _ffi_api
from .search_strategy import SearchStrategy
diff --git a/python/tvm/meta_schedule/search_strategy/replay_trace.py b/python/tvm/meta_schedule/search_strategy/replay_trace.py
index a255965244..e04a440da6 100644
--- a/python/tvm/meta_schedule/search_strategy/replay_trace.py
+++ b/python/tvm/meta_schedule/search_strategy/replay_trace.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""Replay Trace Search Strategy"""
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from .. import _ffi_api
from .search_strategy import SearchStrategy
diff --git a/python/tvm/meta_schedule/search_strategy/search_strategy.py b/python/tvm/meta_schedule/search_strategy/search_strategy.py
index ab4a6fb7b6..75b45cf424 100644
--- a/python/tvm/meta_schedule/search_strategy/search_strategy.py
+++ b/python/tvm/meta_schedule/search_strategy/search_strategy.py
@@ -24,7 +24,7 @@ from typing import TYPE_CHECKING, Callable, List, Optional, Union
from typing_extensions import Literal
# isort: on
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from tvm.runtime import Object
from tvm.tir.schedule import Schedule
diff --git a/python/tvm/meta_schedule/space_generator/post_order_apply.py b/python/tvm/meta_schedule/space_generator/post_order_apply.py
index eee9ea0d0e..45b81bdf3e 100644
--- a/python/tvm/meta_schedule/space_generator/post_order_apply.py
+++ b/python/tvm/meta_schedule/space_generator/post_order_apply.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""Post Order Apply Space Generator."""
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from .. import _ffi_api
from .space_generator import (
diff --git a/python/tvm/meta_schedule/space_generator/schedule_fn.py b/python/tvm/meta_schedule/space_generator/schedule_fn.py
index 2cb1538a5a..d01cd7fdcb 100644
--- a/python/tvm/meta_schedule/space_generator/schedule_fn.py
+++ b/python/tvm/meta_schedule/space_generator/schedule_fn.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""Union of meta Schedule design space generators."""
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from .. import _ffi_api
from .space_generator import (
diff --git a/python/tvm/meta_schedule/space_generator/space_generator.py b/python/tvm/meta_schedule/space_generator/space_generator.py
index 8c9effa6e6..35f9e22367 100644
--- a/python/tvm/meta_schedule/space_generator/space_generator.py
+++ b/python/tvm/meta_schedule/space_generator/space_generator.py
@@ -24,7 +24,7 @@ from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple, Union
from typing_extensions import Literal
# isort: on
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from tvm.ir import IRModule
from tvm.runtime import Object
from tvm.tir.schedule import Schedule
diff --git a/python/tvm/meta_schedule/space_generator/space_generator_union.py b/python/tvm/meta_schedule/space_generator/space_generator_union.py
index f512f65355..0b8ceb4531 100644
--- a/python/tvm/meta_schedule/space_generator/space_generator_union.py
+++ b/python/tvm/meta_schedule/space_generator/space_generator_union.py
@@ -17,7 +17,7 @@
"""Union of meta Schedule design space generators."""
from typing import List
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from .. import _ffi_api
from .space_generator import (
diff --git a/python/tvm/meta_schedule/task_scheduler/gradient_based.py b/python/tvm/meta_schedule/task_scheduler/gradient_based.py
index 7bac23bb3f..18d7e2be61 100644
--- a/python/tvm/meta_schedule/task_scheduler/gradient_based.py
+++ b/python/tvm/meta_schedule/task_scheduler/gradient_based.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""Gradient Based Task Scheduler"""
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from .. import _ffi_api
from ..logging import get_logger, get_logging_func
diff --git a/python/tvm/meta_schedule/task_scheduler/round_robin.py b/python/tvm/meta_schedule/task_scheduler/round_robin.py
index 6475b4102a..78504608f9 100644
--- a/python/tvm/meta_schedule/task_scheduler/round_robin.py
+++ b/python/tvm/meta_schedule/task_scheduler/round_robin.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""Round Robin Task Scheduler"""
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from .. import _ffi_api
from ..logging import get_logger, get_logging_func
diff --git a/python/tvm/meta_schedule/task_scheduler/task_scheduler.py b/python/tvm/meta_schedule/task_scheduler/task_scheduler.py
index 9d6fec88b6..4513f60815 100644
--- a/python/tvm/meta_schedule/task_scheduler/task_scheduler.py
+++ b/python/tvm/meta_schedule/task_scheduler/task_scheduler.py
@@ -22,7 +22,7 @@ from typing_extensions import Literal
# isort: on
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from tvm.runtime import Object
from .. import _ffi_api
diff --git a/python/tvm/meta_schedule/testing/validate_database.py b/python/tvm/meta_schedule/testing/validate_database.py
index 14ff32c017..a207f32711 100644
--- a/python/tvm/meta_schedule/testing/validate_database.py
+++ b/python/tvm/meta_schedule/testing/validate_database.py
@@ -25,7 +25,7 @@ import numpy as np # type: ignore
import tvm
from tvm import meta_schedule as ms
-from tvm.ffi import get_global_func, register_func
+from tvm_ffi import get_global_func, register_func
from tvm.ir import IRModule
from tvm.support import describe
from tvm.target import Target
diff --git a/python/tvm/meta_schedule/tir_integration.py b/python/tvm/meta_schedule/tir_integration.py
index b171c97118..d14a130160 100644
--- a/python/tvm/meta_schedule/tir_integration.py
+++ b/python/tvm/meta_schedule/tir_integration.py
@@ -22,7 +22,7 @@ from typing_extensions import Literal
# isort: on
from tvm import ir, tir
-from tvm.ffi import register_func
+from tvm_ffi import register_func
from tvm.target import Target
from tvm.tir.expr import IntImm
diff --git a/python/tvm/meta_schedule/tune_context.py b/python/tvm/meta_schedule/tune_context.py
index 5512b7a268..e0b0e8c6e7 100644
--- a/python/tvm/meta_schedule/tune_context.py
+++ b/python/tvm/meta_schedule/tune_context.py
@@ -24,7 +24,7 @@ from typing_extensions import Literal
# isort: on
from tvm import IRModule
-from tvm.ffi import register_object, register_func
+from tvm_ffi import register_object, register_func
from tvm.runtime import Object
from tvm.target import Target
from tvm.tir import PrimFunc, Schedule
diff --git a/python/tvm/meta_schedule/utils.py b/python/tvm/meta_schedule/utils.py
index 2f18f54a81..76bac88983 100644
--- a/python/tvm/meta_schedule/utils.py
+++ b/python/tvm/meta_schedule/utils.py
@@ -22,7 +22,7 @@ from typing import Any, Callable, List, Optional, Union
import numpy as np # type: ignore
import psutil # type: ignore
-from tvm.ffi import get_global_func, register_func
+from tvm_ffi import get_global_func, register_func
from tvm.error import TVMError
from tvm.ir import Array, IRModule, Map
from tvm.rpc import RPCSession
diff --git a/python/tvm/relax/_ffi_api.py b/python/tvm/relax/_ffi_api.py
index db1ca05586..947ddb089a 100644
--- a/python/tvm/relax/_ffi_api.py
+++ b/python/tvm/relax/_ffi_api.py
@@ -15,6 +15,6 @@
# specific language governing permissions and limitations
# under the License.
"""FFI API for Relax."""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("relax", __name__)
+tvm_ffi._init_api("relax", __name__)
diff --git a/python/tvm/relax/analysis/_ffi_api.py b/python/tvm/relax/analysis/_ffi_api.py
index fb44606f11..d6adf95805 100644
--- a/python/tvm/relax/analysis/_ffi_api.py
+++ b/python/tvm/relax/analysis/_ffi_api.py
@@ -14,6 +14,6 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
"""FFI APIs"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("relax.analysis", __name__)
+tvm_ffi._init_api("relax.analysis", __name__)
diff --git a/python/tvm/relax/backend/_ffi_api.py b/python/tvm/relax/backend/_ffi_api.py
index 17d7a18a33..fbab394294 100644
--- a/python/tvm/relax/backend/_ffi_api.py
+++ b/python/tvm/relax/backend/_ffi_api.py
@@ -16,6 +16,6 @@
# under the License.
"""FFI API for Relax backend."""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("relax.backend", __name__)
+tvm_ffi._init_api("relax.backend", __name__)
diff --git a/python/tvm/relax/backend/cuda/flashinfer.py b/python/tvm/relax/backend/cuda/flashinfer.py
index 47a4946ca9..0f81675a8f 100644
--- a/python/tvm/relax/backend/cuda/flashinfer.py
+++ b/python/tvm/relax/backend/cuda/flashinfer.py
@@ -131,7 +131,7 @@ def _compile_flashinfer_kernels(
FLASHINFER_TVM_BINDING_DIR,
Path(tvm_home).resolve() / "include",
Path(tvm_home).resolve() / "ffi" / "include",
- Path(tvm_home).resolve() / "3rdparty" / "dlpack" / "include",
+ Path(tvm_home).resolve() / "ffi" / "3rdparty" / "dlpack" / "include",
Path(tvm_home).resolve() / "3rdparty" / "dmlc-core" / "include",
] + CUTLASS_INCLUDE_DIRS
diff --git a/python/tvm/relax/backend/metal/coreml.py b/python/tvm/relax/backend/metal/coreml.py
index 139e5cc2b9..56b0eb3a6c 100644
--- a/python/tvm/relax/backend/metal/coreml.py
+++ b/python/tvm/relax/backend/metal/coreml.py
@@ -19,7 +19,7 @@
import os
import shutil
-import tvm.ffi
+import tvm_ffi
from tvm.contrib import coreml_runtime
from tvm.contrib.xcode import compile_coreml
@@ -463,7 +463,7 @@ class CodegenCoreML(PyExprVisitor):
compile_coreml(model, self.model_name, out_dir)
-@tvm.ffi.register_func("relax.ext.coreml")
+@tvm_ffi.register_func("relax.ext.coreml")
def coreml_compiler(funcs, options, constant_names):
"""
Create a CoreML runtime from a Relax module.
diff --git a/python/tvm/relax/binding_rewrite.py b/python/tvm/relax/binding_rewrite.py
index 22215206ac..077f8feebb 100644
--- a/python/tvm/relax/binding_rewrite.py
+++ b/python/tvm/relax/binding_rewrite.py
@@ -20,13 +20,13 @@
from typing import Optional
import tvm
-import tvm.ffi
+import tvm_ffi
from tvm.runtime import Object
from . import Binding, DataflowBlock, Expr, Function, Var
from . import _ffi_api
-@tvm.ffi.register_object("relax.DataflowBlockRewrite")
+@tvm_ffi.register_object("relax.DataflowBlockRewrite")
class DataflowBlockRewrite(Object):
"""
A binding/statement-level dataflow block rewriter.
diff --git a/python/tvm/relax/block_builder.py b/python/tvm/relax/block_builder.py
index e09a9fab26..26a8346b0a 100644
--- a/python/tvm/relax/block_builder.py
+++ b/python/tvm/relax/block_builder.py
@@ -20,6 +20,7 @@
from typing import Any, Callable, Dict, List, Optional, Sequence, Union
import tvm
+import tvm_ffi
from tvm import relax as rx
from tvm import tir
from tvm.ir.module import IRModule
@@ -100,7 +101,7 @@ class TestingScope(object):
self._bb.end_scope()
-@tvm.ffi.register_object("relax.BlockBuilder")
+@tvm_ffi.register_object("relax.BlockBuilder")
class BlockBuilder(Object):
"""A builder to build Relax IR for testing and dev.
diff --git a/python/tvm/relax/distributed/_ffi_api.py b/python/tvm/relax/distributed/_ffi_api.py
index 6544a8d355..89a15a2bc3 100644
--- a/python/tvm/relax/distributed/_ffi_api.py
+++ b/python/tvm/relax/distributed/_ffi_api.py
@@ -15,6 +15,6 @@
# specific language governing permissions and limitations
# under the License.
"""FFI APIs for tvm.relax.distributed"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("relax.distributed", __name__)
+tvm_ffi._init_api("relax.distributed", __name__)
diff --git a/python/tvm/relax/distributed/global_info.py b/python/tvm/relax/distributed/global_info.py
index 3f549ecfa3..34d3f2da47 100644
--- a/python/tvm/relax/distributed/global_info.py
+++ b/python/tvm/relax/distributed/global_info.py
@@ -18,7 +18,7 @@
"""Global Info Data structures for distributed tensor."""
from typing import List, Union, Tuple
-import tvm
+import tvm_ffi
from tvm.ir import Range
from tvm.ir.global_info import GlobalInfo
from tvm.runtime import ShapeTuple
@@ -26,7 +26,7 @@ from tvm.runtime import ShapeTuple
from . import _ffi_api as ffi
-@tvm.ffi.register_object("relax.distributed.DeviceMesh")
+@tvm_ffi.register_object("relax.distributed.DeviceMesh")
class DeviceMesh(GlobalInfo):
"""Device mesh express a view of topology of devices,
represented by an n-d matrix of device ids.
diff --git a/python/tvm/relax/distributed/struct_info.py b/python/tvm/relax/distributed/struct_info.py
index 50087b9884..554c83e474 100644
--- a/python/tvm/relax/distributed/struct_info.py
+++ b/python/tvm/relax/distributed/struct_info.py
@@ -18,7 +18,7 @@
"""Struct Info for distributed tensor."""
import enum
from typing import List
-import tvm
+import tvm_ffi
from tvm.relax.struct_info import StructInfo, TensorStructInfo
from tvm.ir import Span
from tvm.runtime.object import Object
@@ -33,7 +33,7 @@ class PlacementSpecKind(enum.IntEnum):
kReplica = 1
-@tvm.ffi.register_object("relax.distributed.PlacementSpec")
+@tvm_ffi.register_object("relax.distributed.PlacementSpec")
class PlacementSpec(Object):
"""Describes how data is distributed in one dimension of the device mesh
@@ -80,7 +80,7 @@ class PlacementSpec(Object):
return _ffi_api.Replica()
-@tvm.ffi.register_object("relax.distributed.Placement")
+@tvm_ffi.register_object("relax.distributed.Placement")
class Placement(Object):
"""Describes how data is distributed in each dimension of the device mesh
@@ -110,7 +110,7 @@ class Placement(Object):
return _ffi_api.PlacementFromText(text)
-@tvm.ffi.register_object("relax.DTensorStructInfo")
+@tvm_ffi.register_object("relax.DTensorStructInfo")
class DTensorStructInfo(StructInfo):
"""StructInfo of a Distributed Tensor value.
diff --git a/python/tvm/relax/distributed/transform/_ffi_api.py b/python/tvm/relax/distributed/transform/_ffi_api.py
index b694a67116..ffdb09715f 100644
--- a/python/tvm/relax/distributed/transform/_ffi_api.py
+++ b/python/tvm/relax/distributed/transform/_ffi_api.py
@@ -14,6 +14,6 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
"""FFI APIs for tvm.relax.distributed.transform"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("relax.distributed.transform", __name__)
+tvm_ffi._init_api("relax.distributed.transform", __name__)
diff --git a/python/tvm/relax/dpl/_ffi.py b/python/tvm/relax/dpl/_ffi.py
index 72bf073bed..7097ec8c52 100644
--- a/python/tvm/relax/dpl/_ffi.py
+++ b/python/tvm/relax/dpl/_ffi.py
@@ -15,6 +15,6 @@
# specific language governing permissions and limitations
# under the License.
"""DataFlow Pattern Language FFI bindings."""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("relax.dpl", __name__)
+tvm_ffi._init_api("relax.dpl", __name__)
diff --git a/python/tvm/relax/dpl/pattern.py b/python/tvm/relax/dpl/pattern.py
index eca885e03a..490362a1e5 100644
--- a/python/tvm/relax/dpl/pattern.py
+++ b/python/tvm/relax/dpl/pattern.py
@@ -23,7 +23,7 @@ import typing
from typing import Dict, List, Optional, Tuple, Union
import tvm
-import tvm.ffi as tvm_ffi
+import tvm_ffi as tvm_ffi
from tvm.ir.container import Array
from tvm.ir.expr import PrimExpr
from tvm.ir.op import Op
diff --git a/python/tvm/relax/dpl/rewrite.py b/python/tvm/relax/dpl/rewrite.py
index a9782057c8..6dd730e831 100644
--- a/python/tvm/relax/dpl/rewrite.py
+++ b/python/tvm/relax/dpl/rewrite.py
@@ -20,7 +20,7 @@ from typing import Dict, Callable, Union
from tvm.ir import IRModule
from tvm.runtime import Object
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from .pattern import DFPattern
from .context import PatternContext
diff --git a/python/tvm/relax/exec_builder.py b/python/tvm/relax/exec_builder.py
index 4c5647daf7..43f9a2e693 100644
--- a/python/tvm/relax/exec_builder.py
+++ b/python/tvm/relax/exec_builder.py
@@ -19,7 +19,7 @@
from enum import IntEnum
from typing import Optional, Union, List
import tvm
-from tvm.runtime import Object
+import tvm_ffi
from tvm.runtime.container import ShapeTuple
from .vm_build import VMExecutable
from . import _ffi_api
@@ -56,8 +56,8 @@ class VMFuncScope(object):
self.exit_callback()
-@tvm.ffi.register_object("relax.ExecBuilder")
-class ExecBuilder(Object):
+@tvm_ffi.register_object("relax.ExecBuilder")
+class ExecBuilder(tvm_ffi.core.Object):
"""A builder to emit instructions and build executable for the virtual machine."""
def __init__(self) -> None:
diff --git a/python/tvm/relax/expr.py b/python/tvm/relax/expr.py
index ee9caf3a83..9ea043d6da 100644
--- a/python/tvm/relax/expr.py
+++ b/python/tvm/relax/expr.py
@@ -22,7 +22,7 @@ from typing import Any, Callable, Dict, List, Optional, Union, Mapping
import numpy as _np # type: ignore
import tvm
-import tvm.ffi
+import tvm_ffi
import tvm.ir
import tvm.relax
from tvm import DataType
@@ -42,7 +42,7 @@ Type = Union[tvm.ir.Type] # pylint: disable=invalid-name
GlobalVar = Union[tvm.ir.GlobalVar]
-@tvm.ffi.register_object("relax.Id")
+@tvm_ffi.register_object("relax.Id")
class Id(Object):
"""Unique identifier(name) used in Var.
Guaranteed to be stable across all passes.
@@ -56,7 +56,7 @@ class Id(Object):
# NOTE: place base struct info in expr to avoid cyclic dep
# from expr to struct info.
-@tvm.ffi.register_object("ir.StructInfo")
+@tvm_ffi.register_object("ir.StructInfo")
class StructInfo(Node, Scriptable):
"""The base class of all StructInfo.
@@ -528,7 +528,7 @@ class _DLTensorStrideProxy(tvm.runtime.ObjectGeneric):
return tvm.relax.Call(op, [self.tensor, axis])
-@tvm.ffi.register_object("relax.expr.Call")
+@tvm_ffi.register_object("relax.expr.Call")
class Call(ExprWithOp):
"""Function call node in Relax.
@@ -577,7 +577,7 @@ class Call(ExprWithOp):
)
-@tvm.ffi.register_object("relax.expr.If")
+@tvm_ffi.register_object("relax.expr.If")
class If(ExprWithOp):
"""A conditional expression in Relax.
@@ -609,7 +609,7 @@ class If(ExprWithOp):
)
-@tvm.ffi.register_object("relax.expr.Tuple")
+@tvm_ffi.register_object("relax.expr.Tuple")
class Tuple(ExprWithOp):
"""Tuple expression that groups several fields together.
@@ -644,7 +644,7 @@ class Tuple(ExprWithOp):
return len(self.fields)
-@tvm.ffi.register_object("relax.expr.TupleGetItem")
+@tvm_ffi.register_object("relax.expr.TupleGetItem")
class TupleGetItem(ExprWithOp):
"""Get index-th item from a tuple.
@@ -670,7 +670,7 @@ class TupleGetItem(ExprWithOp):
)
-@tvm.ffi.register_object("relax.expr.ShapeExpr")
+@tvm_ffi.register_object("relax.expr.ShapeExpr")
class ShapeExpr(ExprWithOp):
"""A shape expression which allows users to construct a shape containing PrimExpr.
@@ -708,7 +708,7 @@ def make_shape(shape: Union[List[Any], typing.Tuple[Any, ...]]) -> ShapeExpr:
raise ValueError("Wrong type")
-@tvm.ffi.register_object("relax.expr.Constant")
+@tvm_ffi.register_object("relax.expr.Constant")
class Constant(ExprWithOp):
"""Constant Tensor
@@ -742,7 +742,7 @@ class Constant(ExprWithOp):
)
-@tvm.ffi.register_object("relax.expr.Var")
+@tvm_ffi.register_object("relax.expr.Var")
class Var(ExprWithOp):
"""The variable class for all Relax bindings.
@@ -789,7 +789,7 @@ class Var(ExprWithOp):
return name
-@tvm.ffi.register_object("relax.expr.DataflowVar")
+@tvm_ffi.register_object("relax.expr.DataflowVar")
class DataflowVar(Var):
"""A sub-type of the variable node used to mark dataflow variables from
normal visible "function local" bindings.
@@ -838,7 +838,7 @@ class DataflowVar(Var):
)
-@tvm.ffi.register_object("relax.expr.PrimValue")
+@tvm_ffi.register_object("relax.expr.PrimValue")
class PrimValue(Expr, Scriptable):
"""The prim expr representing the value."""
@@ -850,7 +850,7 @@ class PrimValue(Expr, Scriptable):
self.__init_handle_by_constructor__(_ffi_api.PrimValue, value, span) # type: ignore
-@tvm.ffi.register_object("relax.expr.StringImm")
+@tvm_ffi.register_object("relax.expr.StringImm")
class StringImm(Expr, Scriptable):
"""Represent a string literal constant."""
@@ -861,7 +861,7 @@ class StringImm(Expr, Scriptable):
self.__init_handle_by_constructor__(_ffi_api.StringImm, value, span) # type: ignore
-@tvm.ffi.register_object("relax.expr.DataTypeImm")
+@tvm_ffi.register_object("relax.expr.DataTypeImm")
class DataTypeImm(Expr, Scriptable):
"""Represent a data type constant."""
@@ -872,7 +872,7 @@ class DataTypeImm(Expr, Scriptable):
self.__init_handle_by_constructor__(_ffi_api.DataTypeImm, value, span) # type: ignore
-@tvm.ffi.register_object("relax.expr.Binding")
+@tvm_ffi.register_object("relax.expr.Binding")
class Binding(Node, Scriptable):
"""The base class of a binding in Relax."""
@@ -880,7 +880,7 @@ class Binding(Node, Scriptable):
span: Optional[Span]
-@tvm.ffi.register_object("relax.expr.MatchCast")
+@tvm_ffi.register_object("relax.expr.MatchCast")
class MatchCast(Binding):
"""Runtime-match the value to the struct info.
@@ -912,7 +912,7 @@ class MatchCast(Binding):
)
-@tvm.ffi.register_object("relax.expr.VarBinding")
+@tvm_ffi.register_object("relax.expr.VarBinding")
class VarBinding(Binding):
"""Variable binding, bind he variable of the lhs with the rhs.
@@ -934,7 +934,7 @@ class VarBinding(Binding):
self.__init_handle_by_constructor__(_ffi_api.VarBinding, var, value, span) # type: ignore
-@tvm.ffi.register_object("relax.expr.BindingBlock")
+@tvm_ffi.register_object("relax.expr.BindingBlock")
class BindingBlock(Node, Scriptable):
"""base class of binding block, bindings inside can be impure
(with side effect or control flow)"""
@@ -946,7 +946,7 @@ class BindingBlock(Node, Scriptable):
self.__init_handle_by_constructor__(_ffi_api.BindingBlock, bindings, span) # type: ignore
-@tvm.ffi.register_object("relax.expr.DataflowBlock")
+@tvm_ffi.register_object("relax.expr.DataflowBlock")
class DataflowBlock(BindingBlock):
"""dataflow block, bindings inside are pure (no side effect and no control flow)"""
@@ -958,7 +958,7 @@ class DataflowBlock(BindingBlock):
self.__init_handle_by_constructor__(_ffi_api.DataflowBlock, bindings, span) # type: ignore
-@tvm.ffi.register_object("relax.expr.SeqExpr")
+@tvm_ffi.register_object("relax.expr.SeqExpr")
class SeqExpr(ExprWithOp):
"""A sequence of binding blocks followed by an expression."""
@@ -970,7 +970,7 @@ class SeqExpr(ExprWithOp):
self.__init_handle_by_constructor__(_ffi_api.SeqExpr, blocks, body, span) # type: ignore
-@tvm.ffi.register_object("relax.expr.Function")
+@tvm_ffi.register_object("relax.expr.Function")
class Function(BaseFunc, Scriptable):
"""A Relax function."""
@@ -1109,7 +1109,7 @@ class Function(BaseFunc, Scriptable):
return _ffi_api.FunctionInlineFunctions(self, function_map) # type: ignore
-@tvm.ffi.register_object("relax.expr.ExternFunc")
+@tvm_ffi.register_object("relax.expr.ExternFunc")
class ExternFunc(BaseFunc, ExprWithOp):
"""extern function, which represents a PackedFunc."""
@@ -1177,7 +1177,7 @@ def const(
return Constant(value)
-@tvm.ffi.register_object("relax.TEPlaceholderOp")
+@tvm_ffi.register_object("relax.TEPlaceholderOp")
class TEPlaceholderOp(tvm.te.tensor.Operation):
"""The placeholder op that represents a relax expression."""
diff --git a/python/tvm/relax/expr_functor.py b/python/tvm/relax/expr_functor.py
index a40b81c233..e5e77251c6 100644
--- a/python/tvm/relax/expr_functor.py
+++ b/python/tvm/relax/expr_functor.py
@@ -18,7 +18,7 @@
"""The expression functor of Relax."""
from typing import Callable, Optional
-import tvm
+import tvm_ffi
from tvm.ir import Op
from tvm.runtime import Object
from tvm.runtime.support import derived_object
@@ -261,8 +261,8 @@ class ExprFunctor:
raise TypeError("Invalid type: {0}".format(type(var)))
-@tvm.ffi.register_object("expr_functor.PyExprVisitor")
-class _PyExprVisitor(Object):
+@tvm_ffi.register_object("expr_functor.PyExprVisitor")
+class _PyExprVisitor(tvm_ffi.core.Object):
"""
A TVM object to support customization of ExprVisitor on the python side.
This is the decorated result returned from visitor decorator.
@@ -781,7 +781,7 @@ class PyExprVisitor:
return _ffi_api.ExprVisitorVisitSpan(self._outer(), span) # type: ignore
-@tvm.ffi.register_object("expr_functor.PyExprMutator")
+@tvm_ffi.register_object("expr_functor.PyExprMutator")
class _PyExprMutator(Object):
"""
A TVM object to support customization of ExprMutator on the python side.
diff --git a/python/tvm/relax/frontend/onnx/onnx_frontend.py b/python/tvm/relax/frontend/onnx/onnx_frontend.py
index b91106e64a..76c7374436 100644
--- a/python/tvm/relax/frontend/onnx/onnx_frontend.py
+++ b/python/tvm/relax/frontend/onnx/onnx_frontend.py
@@ -44,6 +44,7 @@ import numpy as _np
import onnx.onnx_ml_pb2
import tvm
+import tvm_ffi
from tvm import TVMError, relax, tir, topi
from tvm.ir import IRModule
from tvm.ir.supply import NameSupply
diff --git a/python/tvm/relax/op/_ffi_api.py b/python/tvm/relax/op/_ffi_api.py
index 1d16a024d1..693c9564d5 100644
--- a/python/tvm/relax/op/_ffi_api.py
+++ b/python/tvm/relax/op/_ffi_api.py
@@ -14,6 +14,6 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
"""FFI APIs for tvm.relax.op"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("relax.op", __name__)
+tvm_ffi._init_api("relax.op", __name__)
diff --git a/python/tvm/relax/op/builtin/_ffi_api.py b/python/tvm/relax/op/builtin/_ffi_api.py
index a7f48af576..4ad011b447 100644
--- a/python/tvm/relax/op/builtin/_ffi_api.py
+++ b/python/tvm/relax/op/builtin/_ffi_api.py
@@ -14,6 +14,6 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
"""FFI APIs for tvm.relax.op.builtin"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("relax.op.builtin", __name__)
+tvm_ffi._init_api("relax.op.builtin", __name__)
diff --git a/python/tvm/relax/op/ccl/_ffi_api.py b/python/tvm/relax/op/ccl/_ffi_api.py
index bf605aae6a..eab31a6463 100644
--- a/python/tvm/relax/op/ccl/_ffi_api.py
+++ b/python/tvm/relax/op/ccl/_ffi_api.py
@@ -15,6 +15,6 @@
# specific language governing permissions and limitations
# under the License.
"""Operators serving for Collective Communications Library (CCL) operators"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("relax.op.ccl", __name__)
+tvm_ffi._init_api("relax.op.ccl", __name__)
diff --git a/python/tvm/relax/op/distributed/_ffi_api.py b/python/tvm/relax/op/distributed/_ffi_api.py
index 394cb8c262..03c4bcc988 100644
--- a/python/tvm/relax/op/distributed/_ffi_api.py
+++ b/python/tvm/relax/op/distributed/_ffi_api.py
@@ -15,6 +15,6 @@
# specific language governing permissions and limitations
# under the License.
"""FFI APIs for tvm.relax.op.distributed"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("relax.op.dist", __name__)
+tvm_ffi._init_api("relax.op.dist", __name__)
diff --git a/python/tvm/relax/op/grad/_ffi_api.py b/python/tvm/relax/op/grad/_ffi_api.py
index 415d590f01..d1f96a1d02 100644
--- a/python/tvm/relax/op/grad/_ffi_api.py
+++ b/python/tvm/relax/op/grad/_ffi_api.py
@@ -15,6 +15,6 @@
# specific language governing permissions and limitations
# under the License.
"""FFI APIs for tvm.relax.op.grad"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("relax.op.grad", __name__)
+tvm_ffi._init_api("relax.op.grad", __name__)
diff --git a/python/tvm/relax/op/image/_ffi_api.py b/python/tvm/relax/op/image/_ffi_api.py
index 8c813231f9..b00b26744b 100644
--- a/python/tvm/relax/op/image/_ffi_api.py
+++ b/python/tvm/relax/op/image/_ffi_api.py
@@ -15,6 +15,6 @@
# specific language governing permissions and limitations
# under the License.
"""Constructor APIs"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("relax.op.image", __name__)
+tvm_ffi._init_api("relax.op.image", __name__)
diff --git a/python/tvm/relax/op/memory/_ffi_api.py b/python/tvm/relax/op/memory/_ffi_api.py
index fb829b7db9..f876c2c1e6 100644
--- a/python/tvm/relax/op/memory/_ffi_api.py
+++ b/python/tvm/relax/op/memory/_ffi_api.py
@@ -14,6 +14,6 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
"""FFI APIs for tvm.relax.op.memory"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("relax.op.memory", __name__)
+tvm_ffi._init_api("relax.op.memory", __name__)
diff --git a/python/tvm/relax/op/nn/_ffi_api.py b/python/tvm/relax/op/nn/_ffi_api.py
index b5f735127e..fa8bf8f6d8 100644
--- a/python/tvm/relax/op/nn/_ffi_api.py
+++ b/python/tvm/relax/op/nn/_ffi_api.py
@@ -15,6 +15,6 @@
# specific language governing permissions and limitations
# under the License.
"""Constructor APIs"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("relax.op.nn", __name__)
+tvm_ffi._init_api("relax.op.nn", __name__)
diff --git a/python/tvm/relax/op/op_attrs.py b/python/tvm/relax/op/op_attrs.py
index 9c15cdd966..4062aae0c7 100644
--- a/python/tvm/relax/op/op_attrs.py
+++ b/python/tvm/relax/op/op_attrs.py
@@ -16,344 +16,344 @@
# under the License.
"""The attributes node used for Relax operators"""
from tvm.ir import Attrs
-import tvm.ffi
+import tvm_ffi
-@tvm.ffi.register_object("relax.attrs.CallTIRWithGradAttrs")
+@tvm_ffi.register_object("relax.attrs.CallTIRWithGradAttrs")
class CallTIRWithGradAttrs(Attrs):
"""Attributes used in call_tir_with_grad operator"""
-@tvm.ffi.register_object("relax.attrs.InitAttrs")
+@tvm_ffi.register_object("relax.attrs.InitAttrs")
class InitAttrs(Attrs):
"""Attributes used in full/full_like, ones/ones_like, and zeros/zeros_like operator"""
-@tvm.ffi.register_object("relax.attrs.TriluAttrs")
+@tvm_ffi.register_object("relax.attrs.TriluAttrs")
class TriluAttrs(Attrs):
"""Attributes used in tril and triu operator"""
-@tvm.ffi.register_object("relax.attrs.AstypeAttrs")
+@tvm_ffi.register_object("relax.attrs.AstypeAttrs")
class AstypeAttrs(Attrs):
"""Attributes used in astype operator"""
-@tvm.ffi.register_object("relax.attrs.TakeAttrs")
+@tvm_ffi.register_object("relax.attrs.TakeAttrs")
class TakeAttrs(Attrs):
"""Attributes used in take operator"""
-@tvm.ffi.register_object("relax.attrs.StridedSliceAttrs")
+@tvm_ffi.register_object("relax.attrs.StridedSliceAttrs")
class StridedSliceAttrs(Attrs):
"""Attributes used in strided_slice operator"""
-@tvm.ffi.register_object("relax.attrs.MatmulAttrs")
+@tvm_ffi.register_object("relax.attrs.MatmulAttrs")
class MatmulAttrs(Attrs):
"""Attributes for matmul operator"""
-@tvm.ffi.register_object("relax.attrs.Conv2DAttrs")
+@tvm_ffi.register_object("relax.attrs.Conv2DAttrs")
class Conv2DAttrs(Attrs):
"""Attributes for nn.conv2d"""
-@tvm.ffi.register_object("relax.attrs.Conv3DAttrs")
+@tvm_ffi.register_object("relax.attrs.Conv3DAttrs")
class Conv3DAttrs(Attrs):
"""Attributes for nn.conv3d"""
-@tvm.ffi.register_object("relax.attrs.Conv2DTransposeAttrs")
+@tvm_ffi.register_object("relax.attrs.Conv2DTransposeAttrs")
class Conv2DTransposeAttrs(Attrs):
"""Attributes for nn.conv2d_transpose"""
-@tvm.ffi.register_object("relax.attrs.Pool2DAttrs")
+@tvm_ffi.register_object("relax.attrs.Pool2DAttrs")
class Pool2DAttrs(Attrs):
"""Attributes for nn.max_pool2d"""
-@tvm.ffi.register_object("relax.attrs.AdaptivePool2DAttrs")
+@tvm_ffi.register_object("relax.attrs.AdaptivePool2DAttrs")
class AdaptivePool2DAttrs(Attrs):
"""Attributes for 2d adaptive pool operator"""
-@tvm.ffi.register_object("relax.attrs.SoftmaxAttrs")
+@tvm_ffi.register_object("relax.attrs.SoftmaxAttrs")
class SoftmaxAttrs(Attrs):
"""Attributes for nn.softmax"""
-@tvm.ffi.register_object("relax.attrs.BatchNormAttrs")
+@tvm_ffi.register_object("relax.attrs.BatchNormAttrs")
class BatchNormAttrs(Attrs):
"""Attributes used in batch_norm operator"""
-@tvm.ffi.register_object("relax.attrs.LayerNormAttrs")
+@tvm_ffi.register_object("relax.attrs.LayerNormAttrs")
class LayerNormAttrs(Attrs):
"""Attributes used in layer_norm operator"""
-@tvm.ffi.register_object("relax.attrs.InstanceNormAttrs")
+@tvm_ffi.register_object("relax.attrs.InstanceNormAttrs")
class InstanceNormAttrs(Attrs):
"""Attributes used in instance_norm operator"""
-@tvm.ffi.register_object("relax.attrs.DropoutAttrs")
+@tvm_ffi.register_object("relax.attrs.DropoutAttrs")
class DropoutAttrs(Attrs):
"""Attributes for dropout operator"""
-@tvm.ffi.register_object("relax.attrs.StatisticalAttrs")
+@tvm_ffi.register_object("relax.attrs.StatisticalAttrs")
class StatisticalAttrs(Attrs):
"""Attributes used in statistical operator"""
-@tvm.ffi.register_object("relax.attrs.ConcatAttrs")
+@tvm_ffi.register_object("relax.attrs.ConcatAttrs")
class ConcatAttrs(Attrs):
"""Attributes for concat operator"""
-@tvm.ffi.register_object("relax.attrs.ExpandDimsAttrs")
+@tvm_ffi.register_object("relax.attrs.ExpandDimsAttrs")
class ExpandDimsAttrs(Attrs):
"""Attributes for expand_dims operator"""
-@tvm.ffi.register_object("relax.attrs.PermuteDimsAttrs")
+@tvm_ffi.register_object("relax.attrs.PermuteDimsAttrs")
class PermuteDimsAttrs(Attrs):
"""Attributes for permute_dims operator"""
-@tvm.ffi.register_object("relax.attrs.SortAttrs")
+@tvm_ffi.register_object("relax.attrs.SortAttrs")
class SortAttrs(Attrs):
"""Attributes for sort operator"""
-@tvm.ffi.register_object("relax.attrs.ArgsortAttrs")
+@tvm_ffi.register_object("relax.attrs.ArgsortAttrs")
class ArgsortAttrs(Attrs):
"""Attributes for argsort operator"""
-@tvm.ffi.register_object("relax.attrs.SplitAttrs")
+@tvm_ffi.register_object("relax.attrs.SplitAttrs")
class SplitAttrs(Attrs):
"""Attributes used in split operator"""
-@tvm.ffi.register_object("relax.attrs.SqueezeAttrs")
+@tvm_ffi.register_object("relax.attrs.SqueezeAttrs")
class SqueezeAttrs(Attrs):
"""Attributes for squeeze operator"""
-@tvm.ffi.register_object("relax.attrs.StackAttrs")
+@tvm_ffi.register_object("relax.attrs.StackAttrs")
class StackAttrs(Attrs):
"""Attributes for concat operator"""
-@tvm.ffi.register_object("relax.attrs.IndexPutAttrs")
+@tvm_ffi.register_object("relax.attrs.IndexPutAttrs")
class IndexPutAttrs(Attrs):
"""Attributes for index_put operator"""
-@tvm.ffi.register_object("relax.attrs.LayoutTransformAttrs")
+@tvm_ffi.register_object("relax.attrs.LayoutTransformAttrs")
class LayoutTransformAttrs(Attrs):
"""Attributes used in layout_transform operator"""
-@tvm.ffi.register_object("relax.attrs.Resize2DAttrs")
+@tvm_ffi.register_object("relax.attrs.Resize2DAttrs")
class Resize2DAttrs(Attrs):
"""Attributes used in image resize2d operator"""
-@tvm.ffi.register_object("relax.attrs.ArgmaxArgminAttrs")
+@tvm_ffi.register_object("relax.attrs.ArgmaxArgminAttrs")
class ArgmaxArgminAttrs(Attrs):
"""Attributes for argmax/argmin operator"""
-@tvm.ffi.register_object("relax.attrs.RepeatAttrs")
+@tvm_ffi.register_object("relax.attrs.RepeatAttrs")
class RepeatAttrs(Attrs):
"""Attributes for repeat operator"""
-@tvm.ffi.register_object("relax.attrs.TileAttrs")
+@tvm_ffi.register_object("relax.attrs.TileAttrs")
class TileAttrs(Attrs):
"""Attributes for tile operator"""
-@tvm.ffi.register_object("relax.attrs.ScanopAttrs")
+@tvm_ffi.register_object("relax.attrs.ScanopAttrs")
class ScanopAttrs(Attrs):
"""Attributes for scan operators"""
-@tvm.ffi.register_object("relax.attrs.TopKAttrs")
+@tvm_ffi.register_object("relax.attrs.TopKAttrs")
class TopKAttrs(Attrs):
"""Attributes for topk operators"""
-@tvm.ffi.register_object("relax.attrs.EinsumAttrs")
+@tvm_ffi.register_object("relax.attrs.EinsumAttrs")
class EinsumAttrs(Attrs):
"""Attributes for einsum operator"""
-@tvm.ffi.register_object("relax.attrs.FlipAttrs")
+@tvm_ffi.register_object("relax.attrs.FlipAttrs")
class FlipAttrs(Attrs):
"""Attributes for flip operator"""
-@tvm.ffi.register_object("relax.attrs.PadAttrs")
+@tvm_ffi.register_object("relax.attrs.PadAttrs")
class PadAttrs(Attrs):
"""Attributes used in pad operator"""
-@tvm.ffi.register_object("relax.attrs.MultinomialFromUniformAttrs")
+@tvm_ffi.register_object("relax.attrs.MultinomialFromUniformAttrs")
class MultinomialFromUniformAttrs(Attrs):
"""Attributes for multinomial_from_uniform operator"""
-@tvm.ffi.register_object("relax.attrs.CallInplacePackedAttrs")
+@tvm_ffi.register_object("relax.attrs.CallInplacePackedAttrs")
class CallInplacePackedAttrs(Attrs):
"""Attributes used in call_inplace_packed operator"""
-@tvm.ffi.register_object("relax.attrs.CallTIRInplaceAttrs")
+@tvm_ffi.register_object("relax.attrs.CallTIRInplaceAttrs")
class CallTIRInplaceAttrs(Attrs):
"""Attributes used in call_tir_inplace operator"""
-@tvm.ffi.register_object("relax.attrs.ToVDeviceAttrs")
+@tvm_ffi.register_object("relax.attrs.ToVDeviceAttrs")
class ToVDeviceAttrs(Attrs):
"""Attributes used in to_vdevice operator"""
-@tvm.ffi.register_object("relax.attrs.HintOnDeviceAttrs")
+@tvm_ffi.register_object("relax.attrs.HintOnDeviceAttrs")
class HintOnDeviceAttrs(Attrs):
"""Attributes used in hint_on_device operator"""
-@tvm.ffi.register_object("relax.attrs.ScatterCollectiveAttrs")
+@tvm_ffi.register_object("relax.attrs.ScatterCollectiveAttrs")
class ScatterCollectiveAttrs(Attrs):
"""Attributes used in scatter collective operators"""
-@tvm.ffi.register_object("relax.attrs.AttentionAttrs")
+@tvm_ffi.register_object("relax.attrs.AttentionAttrs")
class AttentionAttrs(Attrs):
"""Attributes used in attention operator"""
-@tvm.ffi.register_object("relax.attrs.Conv1DAttrs")
+@tvm_ffi.register_object("relax.attrs.Conv1DAttrs")
class Conv1DAttrs(Attrs):
"""Attributes for nn.conv1d"""
-@tvm.ffi.register_object("relax.attrs.Conv1DTransposeAttrs")
+@tvm_ffi.register_object("relax.attrs.Conv1DTransposeAttrs")
class Conv1DTransposeAttrs(Attrs):
"""Attributes for nn.conv1d_transpose"""
-@tvm.ffi.register_object("relax.attrs.Pool1DAttrs")
+@tvm_ffi.register_object("relax.attrs.Pool1DAttrs")
class Pool1DAttrs(Attrs):
"""Attributes for nn.max_pool1d and nn.avg_pool1d"""
-@tvm.ffi.register_object("relax.attrs.Pool3DAttrs")
+@tvm_ffi.register_object("relax.attrs.Pool3DAttrs")
class Pool3DAttrs(Attrs):
"""Attributes for nn.max_pool3d and nn.avg_pool3d"""
-@tvm.ffi.register_object("relax.attrs.AdaptivePool1DAttrs")
+@tvm_ffi.register_object("relax.attrs.AdaptivePool1DAttrs")
class AdaptivePool1DAttrs(Attrs):
"""Attributes for 1d adaptive pool operator"""
-@tvm.ffi.register_object("relax.attrs.AdaptivePool3DAttrs")
+@tvm_ffi.register_object("relax.attrs.AdaptivePool3DAttrs")
class AdaptivePool3DAttrs(Attrs):
"""Attributes for 3d adaptive pool operator"""
-@tvm.ffi.register_object("relax.attrs.LeakyReluAttrs")
+@tvm_ffi.register_object("relax.attrs.LeakyReluAttrs")
class LeakyReluAttrs(Attrs):
"""Attributes used in leaky_relu operator"""
-@tvm.ffi.register_object("relax.attrs.SoftplusAttrs")
+@tvm_ffi.register_object("relax.attrs.SoftplusAttrs")
class SoftplusAttrs(Attrs):
"""Attributes used in softplus operator"""
-@tvm.ffi.register_object("relax.attrs.PReluAttrs")
+@tvm_ffi.register_object("relax.attrs.PReluAttrs")
class PReluAttrs(Attrs):
"""Attributes used in prelu operator"""
-@tvm.ffi.register_object("relax.attrs.PixelShuffleAttrs")
+@tvm_ffi.register_object("relax.attrs.PixelShuffleAttrs")
class PixelShuffleAttrs(Attrs):
"""Attributes used in pixel_shuffle operator"""
-@tvm.ffi.register_object("relax.attrs.GroupNormAttrs")
+@tvm_ffi.register_object("relax.attrs.GroupNormAttrs")
class GroupNormAttrs(Attrs):
"""Attributes used in group_norm operator"""
-@tvm.ffi.register_object("relax.attrs.RMSNormAttrs")
+@tvm_ffi.register_object("relax.attrs.RMSNormAttrs")
class RMSNormAttrs(Attrs):
"""Attributes used in rms_norm operator"""
-@tvm.ffi.register_object("relax.attrs.NLLLossAttrs")
+@tvm_ffi.register_object("relax.attrs.NLLLossAttrs")
class NLLLossAttrs(Attrs):
"""Attributes used in nll_loss operator"""
-@tvm.ffi.register_object("relax.attrs.AllReduceAttrs")
+@tvm_ffi.register_object("relax.attrs.AllReduceAttrs")
class AllReduceAttrs(Attrs):
"""Attributes used in allreduce operator"""
-@tvm.ffi.register_object("relax.attrs.AllGatherAttrs")
+@tvm_ffi.register_object("relax.attrs.AllGatherAttrs")
class AllGatherAttrs(Attrs):
"""Attributes used in allgather operator"""
-@tvm.ffi.register_object("relax.attrs.WrapParamAttrs")
+@tvm_ffi.register_object("relax.attrs.WrapParamAttrs")
class WrapParamAttrs(Attrs):
"""Attributes used in wrap_param operator"""
-@tvm.ffi.register_object("relax.attrs.QuantizeAttrs")
+@tvm_ffi.register_object("relax.attrs.QuantizeAttrs")
class QuantizeAttrs(Attrs):
"""Attributes used in quantize/dequantize operators"""
-@tvm.ffi.register_object("relax.attrs.GatherElementsAttrs")
+@tvm_ffi.register_object("relax.attrs.GatherElementsAttrs")
class GatherElementsAttrs(Attrs):
"""Attributes for gather_elements operator"""
-@tvm.ffi.register_object("relax.attrs.GatherNDAttrs")
+@tvm_ffi.register_object("relax.attrs.GatherNDAttrs")
class GatherNDAttrs(Attrs):
"""Attributes for gather_nd operator"""
-@tvm.ffi.register_object("relax.attrs.MeshgridAttrs")
+@tvm_ffi.register_object("relax.attrs.MeshgridAttrs")
class MeshgridAttrs(Attrs):
"""Attributes for meshgrid operator"""
-@tvm.ffi.register_object("relax.attrs.ScatterElementsAttrs")
+@tvm_ffi.register_object("relax.attrs.ScatterElementsAttrs")
class ScatterElementsAttrs(Attrs):
"""Attributes for scatter_elements operator"""
-@tvm.ffi.register_object("relax.attrs.ScatterNDAttrs")
+@tvm_ffi.register_object("relax.attrs.ScatterNDAttrs")
class ScatterNDAttrs(Attrs):
"""Attributes for scatter_nd operator"""
-@tvm.ffi.register_object("relax.attrs.SliceScatterAttrs")
+@tvm_ffi.register_object("relax.attrs.SliceScatterAttrs")
class SliceScatterAttrs(Attrs):
"""Attributes for slice_scatter operator"""
-@tvm.ffi.register_object("relax.attrs.OneHotAttrs")
+@tvm_ffi.register_object("relax.attrs.OneHotAttrs")
class OneHotAttrs(Attrs):
"""Attributes for one_hot operator"""
diff --git a/python/tvm/relax/op/vm/_ffi_api.py b/python/tvm/relax/op/vm/_ffi_api.py
index f3b6cea13b..bd543ad1c9 100644
--- a/python/tvm/relax/op/vm/_ffi_api.py
+++ b/python/tvm/relax/op/vm/_ffi_api.py
@@ -14,6 +14,6 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
"""FFI APIs for tvm.relax.op.vm"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("relax.op.vm", __name__)
+tvm_ffi._init_api("relax.op.vm", __name__)
diff --git a/python/tvm/relax/struct_info.py b/python/tvm/relax/struct_info.py
index c143f09832..e8f6c42435 100644
--- a/python/tvm/relax/struct_info.py
+++ b/python/tvm/relax/struct_info.py
@@ -18,7 +18,7 @@
"""The struct info nodes of the Relax language."""
from typing import List, Optional, Union
-import tvm.ffi
+import tvm_ffi
import tvm
from tvm.ir import Span, EnvFunc, Array, VDevice
@@ -29,7 +29,7 @@ from .expr import StructInfo, Expr, ShapeExpr
from . import _ffi_api, ty, expr
-@tvm.ffi.register_object("relax.ObjectStructInfo")
+@tvm_ffi.register_object("relax.ObjectStructInfo")
class ObjectStructInfo(StructInfo):
"""StructInfo of an Object."""
@@ -37,7 +37,7 @@ class ObjectStructInfo(StructInfo):
self.__init_handle_by_constructor__(_ffi_api.ObjectStructInfo, span) # type: ignore
-@tvm.ffi.register_object("relax.PrimStructInfo")
+@tvm_ffi.register_object("relax.PrimStructInfo")
class PrimStructInfo(StructInfo):
"""StructInfo of a primitive POD value.
@@ -107,7 +107,7 @@ class PrimStructInfo(StructInfo):
) # type: ignore
-@tvm.ffi.register_object("relax.ShapeStructInfo")
+@tvm_ffi.register_object("relax.ShapeStructInfo")
class ShapeStructInfo(StructInfo):
"""StructInfo of a shape value.
@@ -136,7 +136,7 @@ class ShapeStructInfo(StructInfo):
)
-@tvm.ffi.register_object("relax.TensorStructInfo")
+@tvm_ffi.register_object("relax.TensorStructInfo")
class TensorStructInfo(StructInfo):
"""StructInfo of a Tensor value.
@@ -180,7 +180,7 @@ class TensorStructInfo(StructInfo):
)
-@tvm.ffi.register_object("relax.TupleStructInfo")
+@tvm_ffi.register_object("relax.TupleStructInfo")
class TupleStructInfo(StructInfo):
"""StructInfo of a Tuple value.
@@ -197,7 +197,7 @@ class TupleStructInfo(StructInfo):
self.__init_handle_by_constructor__(_ffi_api.TupleStructInfo, fields, span) # type: ignore
-@tvm.ffi.register_object("relax.FuncStructInfo")
+@tvm_ffi.register_object("relax.FuncStructInfo")
class FuncStructInfo(StructInfo):
"""StructInfo of a function value.
diff --git a/python/tvm/relax/testing/transform.py b/python/tvm/relax/testing/transform.py
index 198b07e51e..617ba73f09 100644
--- a/python/tvm/relax/testing/transform.py
+++ b/python/tvm/relax/testing/transform.py
@@ -21,6 +21,7 @@ import logging
import os
from typing import Dict, List, Set, Tuple
import tvm
+import tvm_ffi
from tvm.ir.module import IRModule
from tvm.relax.expr import Call, DataflowBlock, Var
from tvm.runtime.object import Object
@@ -70,7 +71,7 @@ def dataflow_alias_analysis(
return res_alias_sets, res_tuple_map # type: ignore
-@tvm.ffi.register_object("relax.transform.InplaceOpportunity")
+@tvm_ffi.register_object("relax.transform.InplaceOpportunity")
class InplaceOpportunity(Object):
"""
Represents an opportunity to make a binding in-place. Exposed only for testing;
diff --git a/python/tvm/relax/training/_ffi_api.py b/python/tvm/relax/training/_ffi_api.py
index 9b7dbcdee7..84c117f9cb 100644
--- a/python/tvm/relax/training/_ffi_api.py
+++ b/python/tvm/relax/training/_ffi_api.py
@@ -15,6 +15,6 @@
# specific language governing permissions and limitations
# under the License.
"""FFI APIs for tvm.relax.training"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("relax.training", __name__)
+tvm_ffi._init_api("relax.training", __name__)
diff --git a/python/tvm/relax/training/utils.py b/python/tvm/relax/training/utils.py
index dd433435e2..d20c8bcbe9 100644
--- a/python/tvm/relax/training/utils.py
+++ b/python/tvm/relax/training/utils.py
@@ -21,7 +21,7 @@ from typing import Optional, Callable
import tvm
from tvm import relax
-from tvm.ffi.registry import register_func
+from tvm_ffi.registry import register_func
from tvm.relax.block_builder import BlockBuilder
from ..expr import Function, Var, Call
diff --git a/python/tvm/relax/transform/_ffi_api.py b/python/tvm/relax/transform/_ffi_api.py
index 3c4387a3cb..6ae33aef83 100644
--- a/python/tvm/relax/transform/_ffi_api.py
+++ b/python/tvm/relax/transform/_ffi_api.py
@@ -14,6 +14,6 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
"""FFI APIs for tvm.transform"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("relax.transform", __name__)
+tvm_ffi._init_api("relax.transform", __name__)
diff --git a/python/tvm/relax/transform/transform.py b/python/tvm/relax/transform/transform.py
index 57627ceebe..bf813b3dd6 100644
--- a/python/tvm/relax/transform/transform.py
+++ b/python/tvm/relax/transform/transform.py
@@ -23,7 +23,7 @@ import warnings
from typing import Callable, Dict, List, Mapping, Optional, Sequence, Tuple, Union
import numpy as np # type: ignore
-
+import tvm_ffi
import tvm.ir
from tvm.ir.container import Array
from tvm.relax import Expr, Var, StructInfo
@@ -36,14 +36,14 @@ from .legalize_ops.common import LegalizeFunc
from ..expr import Var
-@tvm.ffi.register_object("relax.FunctionPass")
+@tvm_ffi.register_object("relax.FunctionPass")
class FunctionPass(tvm.ir.transform.Pass):
"""A pass that works on each tvm.relax.Function in a module. A function
pass class should be created through `function_pass`.
"""
-@tvm.ffi.register_object("relax.DataflowBlockPass")
+@tvm_ffi.register_object("relax.DataflowBlockPass")
class DataflowBlockPass(tvm.ir.transform.Pass):
"""A pass that works on each tvm.relax.DataflowBlock in a module."""
@@ -820,7 +820,7 @@ def FuseTIR() -> tvm.ir.transform.Pass:
return _ffi_api.FuseTIR() # type: ignore
-@tvm.ffi.register_object("relax.transform.PatternCheckContext")
+@tvm_ffi.register_object("relax.transform.PatternCheckContext")
class PatternCheckContext(Object):
"""
The input of check function `FusionPattern.check`.
@@ -854,7 +854,7 @@ class PatternCheckContext(Object):
value_to_bound_var: Mapping[Expr, Var]
-@tvm.ffi.register_object("relax.transform.FusionPattern")
+@tvm_ffi.register_object("relax.transform.FusionPattern")
class FusionPattern(Object):
"""
The pattern used by `FuseOpsByPattern`. It's mainly DFPattern but with other
diff --git a/python/tvm/relax/ty.py b/python/tvm/relax/ty.py
index 426695c9f1..ebf757f381 100644
--- a/python/tvm/relax/ty.py
+++ b/python/tvm/relax/ty.py
@@ -16,13 +16,13 @@
# under the License.
# pylint: disable=invalid-name, unused-import
"""The type nodes of the Relax language."""
-import tvm.ffi
+import tvm_ffi
from tvm.ir import Type, TupleType, FuncType, Span
from . import _ffi_api
-@tvm.ffi.register_object("relax.ShapeType")
+@tvm_ffi.register_object("relax.ShapeType")
class ShapeType(Type):
"""The type of shape in Relax.
@@ -37,7 +37,7 @@ class ShapeType(Type):
self.__init_handle_by_constructor__(_ffi_api.ShapeType, ndim, span) # type: ignore
-@tvm.ffi.register_object("relax.ObjectType")
+@tvm_ffi.register_object("relax.ObjectType")
class ObjectType(Type):
"""A type that corresponds to tvm::runtime::Object, is base of all possible object
values in TVM."""
@@ -46,7 +46,7 @@ class ObjectType(Type):
self.__init_handle_by_constructor__(_ffi_api.ObjectType, span) # type: ignore
-@tvm.ffi.register_object("relax.DynTensorType")
+@tvm_ffi.register_object("relax.DynTensorType")
class TensorType(Type):
"""A dynamic tensor type in Relax.
@@ -65,7 +65,7 @@ class TensorType(Type):
self.__init_handle_by_constructor__(_ffi_api.TensorType, ndim, dtype, span) # type: ignore
-@tvm.ffi.register_object("relax.PackedFuncType")
+@tvm_ffi.register_object("relax.PackedFuncType")
class PackedFuncType(Type):
"""The type of ExternFunc in Relax."""
diff --git a/python/tvm/relax/utils.py b/python/tvm/relax/utils.py
index 192235d595..7ce188f780 100644
--- a/python/tvm/relax/utils.py
+++ b/python/tvm/relax/utils.py
@@ -26,6 +26,7 @@ from typing import Tuple as typing_Tuple
from typing import Any, Callable, List, Dict, Optional
import tvm
+import tvm_ffi
from .. import tir
from ..tir import PrimExpr
from . import _ffi_api
@@ -99,7 +100,7 @@ def convert_to_expr(value: Any) -> Expr:
if isinstance(value, float):
return PrimValue(tir.FloatImm("float64", value))
- tvm_value = tvm.ffi.convert(value)
+ tvm_value = tvm_ffi.convert(value)
# Case 1
if isinstance(tvm_value, Expr): # type: ignore
return tvm_value
diff --git a/python/tvm/rpc/_ffi_api.py b/python/tvm/rpc/_ffi_api.py
index 3b77e7a552..b1bc8af974 100644
--- a/python/tvm/rpc/_ffi_api.py
+++ b/python/tvm/rpc/_ffi_api.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""FFI APIs for tvm.rpc"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("rpc", __name__)
+tvm_ffi._init_api("rpc", __name__)
diff --git a/python/tvm/rpc/client.py b/python/tvm/rpc/client.py
index 0bb4e8cb7d..90267c0526 100644
--- a/python/tvm/rpc/client.py
+++ b/python/tvm/rpc/client.py
@@ -22,7 +22,7 @@ import stat
import struct
import time
-import tvm.ffi
+import tvm_ffi
from tvm.base import TVMError
from tvm.contrib import utils
from tvm.runtime import ndarray as nd
@@ -263,7 +263,7 @@ class LocalSession(RPCSession):
RPCSession.__init__(self, _ffi_api.LocalSession())
-@tvm.ffi.register_func("rpc.PopenSession")
+@tvm_ffi.register_func("rpc.PopenSession")
def _popen_session(binary):
temp = utils.tempdir()
diff --git a/python/tvm/rpc/server.py b/python/tvm/rpc/server.py
index eb345260e3..17b3f3652e 100644
--- a/python/tvm/rpc/server.py
+++ b/python/tvm/rpc/server.py
@@ -36,7 +36,7 @@ import multiprocessing
import time
import errno
import sys
-import tvm.ffi
+import tvm_ffi
from tvm.base import py_str
from tvm.libinfo import find_lib_path
@@ -70,11 +70,11 @@ def _server_env(load_library, work_path=None):
temp = utils.tempdir()
# pylint: disable=unused-variable
- @tvm.ffi.register_func("tvm.rpc.server.workpath", override=True)
+ @tvm_ffi.register_func("tvm.rpc.server.workpath", override=True)
def get_workpath(path):
return temp.relpath(path)
- @tvm.ffi.register_func("tvm.rpc.server.load_module", override=True)
+ @tvm_ffi.register_func("tvm.rpc.server.load_module", override=True)
def load_module(file_name):
"""Load module from remote side."""
path = temp.relpath(file_name)
@@ -82,7 +82,7 @@ def _server_env(load_library, work_path=None):
logger.info("load_module %s", path)
return m
- @tvm.ffi.register_func("tvm.rpc.server.download_linked_module", override=True)
+ @tvm_ffi.register_func("tvm.rpc.server.download_linked_module", override=True)
def download_linked_module(file_name):
"""Load module from remote side."""
# pylint: disable=import-outside-toplevel
diff --git a/python/tvm/runtime/__init__.py b/python/tvm/runtime/__init__.py
index ca70cf0f45..7cf98787f3 100644
--- a/python/tvm/runtime/__init__.py
+++ b/python/tvm/runtime/__init__.py
@@ -43,4 +43,4 @@ from .params import (
from . import disco
from .support import _regex_match
-from ..ffi import convert, dtype as DataType, DataTypeCode
+from tvm_ffi import convert, dtype as DataType, DataTypeCode
diff --git a/python/tvm/runtime/_ffi_api.py b/python/tvm/runtime/_ffi_api.py
index 88a49f3a63..0357b280bd 100644
--- a/python/tvm/runtime/_ffi_api.py
+++ b/python/tvm/runtime/_ffi_api.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""FFI APIs for tvm.runtime"""
-import tvm.ffi
+import tvm_ffi
# Exports functions registered in runtime namespace.
-tvm.ffi._init_api("runtime", __name__)
+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 4a0edd449c..2e47f6aa32 100644
--- a/python/tvm/runtime/_ffi_node_api.py
+++ b/python/tvm/runtime/_ffi_node_api.py
@@ -17,8 +17,8 @@
# pylint: disable=invalid-name, unused-argument
"""FFI for tvm.node"""
-import tvm.ffi
-import tvm.ffi.core
+import tvm_ffi
+import tvm_ffi.core
# The implementations below are default ones when the corresponding
@@ -37,4 +37,4 @@ def LoadJSON(json_str):
# Exports functions registered in node namespace.
-tvm.ffi._init_api("node", __name__)
+tvm_ffi._init_api("node", __name__)
diff --git a/python/tvm/runtime/container.py b/python/tvm/runtime/container.py
index 3bf149d6b2..37d0d2116c 100644
--- a/python/tvm/runtime/container.py
+++ b/python/tvm/runtime/container.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""Runtime container structures."""
-from tvm.ffi import String, Shape as ShapeTuple
+from tvm_ffi import String, Shape as ShapeTuple
__all__ = ["ShapeTuple", "String"]
diff --git a/python/tvm/runtime/device.py b/python/tvm/runtime/device.py
index d9d6abce50..d86e30605f 100644
--- a/python/tvm/runtime/device.py
+++ b/python/tvm/runtime/device.py
@@ -18,7 +18,7 @@
# pylint: disable=invalid-name
import json
-import tvm.ffi
+import tvm_ffi
from . import _ffi_api
@@ -26,7 +26,7 @@ from . import _ffi_api
RPC_SESS_MASK = 128
-class Device(tvm.ffi.core.Device):
+class Device(tvm_ffi.core.Device):
"""TVM device strucure."""
def _GetDeviceAttr(self, device_type, device_id, attr_id):
@@ -334,4 +334,4 @@ class Device(tvm.ffi.core.Device):
return Device.DEVICE_TYPE_TO_NAME[self.device_type]
-tvm.ffi.core._set_class_device(Device)
+tvm_ffi.core._set_class_device(Device)
diff --git a/python/tvm/runtime/disco/_ffi_api.py b/python/tvm/runtime/disco/_ffi_api.py
index 79e1a52ad4..63a53d8b85 100644
--- a/python/tvm/runtime/disco/_ffi_api.py
+++ b/python/tvm/runtime/disco/_ffi_api.py
@@ -15,6 +15,6 @@
# specific language governing permissions and limitations
# under the License.
"""FFI APIs from C++"""
-from ...ffi import _init_api
+from tvm_ffi import _init_api
_init_api("runtime.disco", __name__)
diff --git a/python/tvm/runtime/disco/process_pool.py b/python/tvm/runtime/disco/process_pool.py
index 8f05f28e91..ba9b512f04 100644
--- a/python/tvm/runtime/disco/process_pool.py
+++ b/python/tvm/runtime/disco/process_pool.py
@@ -20,7 +20,7 @@ import os
import subprocess
import sys
-from tvm.ffi import register_func
+from tvm_ffi import register_func
from tvm.runtime import ShapeTuple
diff --git a/python/tvm/runtime/disco/session.py b/python/tvm/runtime/disco/session.py
index 49449a451a..4e4d030a62 100644
--- a/python/tvm/runtime/disco/session.py
+++ b/python/tvm/runtime/disco/session.py
@@ -25,7 +25,7 @@ from typing import Any, Callable, Optional, Sequence, Union
import numpy as np
-from ...ffi import get_global_func, register_func, register_object
+from tvm_ffi import get_global_func, register_func, register_object
from ..device import Device
from ..container import ShapeTuple
from ..ndarray import NDArray
diff --git a/python/tvm/runtime/module.py b/python/tvm/runtime/module.py
index 9cbc06708b..06d8c81d98 100644
--- a/python/tvm/runtime/module.py
+++ b/python/tvm/runtime/module.py
@@ -26,7 +26,7 @@ from tvm.base import _RUNTIME_ONLY
from tvm.libinfo import find_include_path
from . import _ffi_api
-from ..ffi import (
+from tvm_ffi import (
Module as _Module,
load_module as _load_module,
register_object as _register_object,
diff --git a/python/tvm/runtime/ndarray.py b/python/tvm/runtime/ndarray.py
index 1d960d5dda..f81c678294 100644
--- a/python/tvm/runtime/ndarray.py
+++ b/python/tvm/runtime/ndarray.py
@@ -29,11 +29,12 @@ except ImportError:
from tvm.runtime import Device
-import tvm.ffi
+import tvm_ffi
+import tvm
from . import _ffi_api
-from ..ffi import (
+from tvm_ffi import (
device,
cpu,
cuda,
@@ -63,15 +64,15 @@ def from_dlpack(ext_tensor):
required_contiguous : bool
Whether to check for contiguous memory.
"""
- return tvm.ffi.from_dlpack(
+ return tvm_ffi.from_dlpack(
ext_tensor,
required_alignment=64,
required_contiguous=True,
)
-@tvm.ffi.register_object("ffi.NDArray")
-class NDArray(tvm.ffi.core.NDArray):
+@tvm_ffi.register_object("ffi.NDArray")
+class NDArray(tvm_ffi.core.NDArray):
"""Lightweight NDArray class of TVM runtime.
Strictly this is only an Array Container (a buffer object)
@@ -124,7 +125,7 @@ class NDArray(tvm.ffi.core.NDArray):
f"array must be an array_like data, type {type(source_array)} is not supported"
)
- t = tvm.ffi.dtype(self.dtype)
+ t = tvm_ffi.dtype(self.dtype)
shape, dtype = self.shape, self.dtype
if t.lanes > 1:
shape = shape + (t.lanes,)
@@ -135,7 +136,7 @@ class NDArray(tvm.ffi.core.NDArray):
raise ValueError(
f"array shape do not match the shape of NDArray {source_array.shape} vs {shape}"
)
- numpy_str_map = tvm.ffi.dtype.NUMPY_DTYPE_TO_STR
+ numpy_str_map = tvm_ffi.dtype.NUMPY_DTYPE_TO_STR
np_dtype_str = (
numpy_str_map[source_array.dtype]
if source_array.dtype in numpy_str_map
@@ -182,7 +183,7 @@ class NDArray(tvm.ffi.core.NDArray):
np_arr : numpy.ndarray
The corresponding numpy array.
"""
- t = tvm.ffi.dtype(self.dtype)
+ t = tvm_ffi.dtype(self.dtype)
shape, dtype = self.shape, self.dtype
old_dtype = dtype
if t.lanes > 1:
@@ -247,7 +248,7 @@ class NDArray(tvm.ffi.core.NDArray):
"""
if isinstance(target, NDArray):
return self._copyto(target)
- if isinstance(target, tvm.ffi.core.Device):
+ if isinstance(target, tvm_ffi.core.Device):
res = empty(self.shape, self.dtype, target, mem_scope)
return self._copyto(res)
raise ValueError(f"Unsupported target type {type(target)}")
@@ -330,7 +331,7 @@ def empty(shape, dtype="float32", device=None, mem_scope=None):
device = device or cpu()
if not isinstance(shape, tvm.runtime.ShapeTuple):
shape = tvm.runtime.ShapeTuple([int(dim) for dim in shape])
- dtype = tvm.ffi.dtype(dtype)
+ dtype = tvm_ffi.dtype(dtype)
arr = _ffi_api.TVMArrayAllocWithScope(shape, dtype, device, mem_scope)
return arr
@@ -362,4 +363,4 @@ def array(arr, device=None, mem_scope=None):
# Register back to FFI
-tvm.ffi.core._set_class_ndarray(NDArray)
+tvm_ffi.core._set_class_ndarray(NDArray)
diff --git a/python/tvm/runtime/object.py b/python/tvm/runtime/object.py
index b2fcddc40a..c9dcf2d1a8 100644
--- a/python/tvm/runtime/object.py
+++ b/python/tvm/runtime/object.py
@@ -17,11 +17,11 @@
# pylint: disable=invalid-name, unused-import
"""Runtime Object API"""
-from tvm.ffi.core import Object
-import tvm.ffi.core
+from tvm_ffi.core import Object
+import tvm_ffi.core
from . import _ffi_node_api
-tvm.ffi.core._set_class_object(Object)
-# override the default repr function for tvm.ffi.core.Object
-tvm.ffi.core.__object_repr__ = _ffi_node_api.AsRepr
+tvm_ffi.core._set_class_object(Object)
+# override the default repr function for tvm_ffi.core.Object
+tvm_ffi.core.__object_repr__ = _ffi_node_api.AsRepr
diff --git a/python/tvm/runtime/object_generic.py b/python/tvm/runtime/object_generic.py
index 4ffea01a3c..f5574e4802 100644
--- a/python/tvm/runtime/object_generic.py
+++ b/python/tvm/runtime/object_generic.py
@@ -16,7 +16,7 @@
# under the License.
"""Common implementation of object generic related logic"""
# pylint: disable=unused-import, invalid-name
-from tvm.ffi import ObjectGeneric
+from tvm_ffi import ObjectGeneric
from . import _ffi_node_api
diff --git a/python/tvm/runtime/packed_func.py b/python/tvm/runtime/packed_func.py
index 71a0ba0816..68940103f3 100644
--- a/python/tvm/runtime/packed_func.py
+++ b/python/tvm/runtime/packed_func.py
@@ -17,6 +17,6 @@
# pylint: disable=invalid-name, unused-import
"""Packed Function namespace."""
-from tvm.ffi import Function as PackedFunc
+from tvm_ffi import Function as PackedFunc
__all__ = ["PackedFunc"]
diff --git a/python/tvm/runtime/profiling/_ffi_api.py b/python/tvm/runtime/profiling/_ffi_api.py
index 85e5d4ca02..104aac90a5 100644
--- a/python/tvm/runtime/profiling/_ffi_api.py
+++ b/python/tvm/runtime/profiling/_ffi_api.py
@@ -15,6 +15,6 @@
# specific language governing permissions and limitations
# under the License.
"""FFI for profiling"""
-from ...ffi import _init_api
+from tvm_ffi import _init_api
_init_api("runtime.profiling", __name__)
diff --git a/python/tvm/runtime/script_printer.py b/python/tvm/runtime/script_printer.py
index a00281b435..7442cd9917 100644
--- a/python/tvm/runtime/script_printer.py
+++ b/python/tvm/runtime/script_printer.py
@@ -18,8 +18,8 @@
import os
from typing import Dict, List, Optional, Sequence
-from tvm.ffi import get_global_func, register_object
-from tvm.ffi.access_path import AccessPath
+from tvm_ffi import get_global_func, register_object
+from tvm_ffi.access_path import AccessPath
from tvm.runtime import Object
from . import _ffi_node_api
diff --git a/python/tvm/runtime/support.py b/python/tvm/runtime/support.py
index 2669459d71..99856b8d3b 100644
--- a/python/tvm/runtime/support.py
+++ b/python/tvm/runtime/support.py
@@ -20,10 +20,10 @@
import re
from typing import TypeVar
-import tvm.ffi
+import tvm_ffi
-@tvm.ffi.register_func("tvm.runtime.regex_match")
+@tvm_ffi.register_func("tvm.runtime.regex_match")
def _regex_match(regex_pattern: str, match_against: str) -> bool:
"""Check if a pattern matches a regular expression
diff --git a/python/tvm/runtime/vm.py b/python/tvm/runtime/vm.py
index d69c3308fa..a955835573 100644
--- a/python/tvm/runtime/vm.py
+++ b/python/tvm/runtime/vm.py
@@ -23,7 +23,7 @@ from numbers import Number, Integral
import numpy as np # type: ignore
import tvm
-from tvm.ffi import register_func
+from tvm_ffi import register_func
from tvm.runtime import Device, Object, PackedFunc
from tvm.runtime.profiling import Report
diff --git a/python/tvm/script/_ffi_api.py b/python/tvm/script/_ffi_api.py
index 8ae8f7b7f9..28dcec06bb 100644
--- a/python/tvm/script/_ffi_api.py
+++ b/python/tvm/script/_ffi_api.py
@@ -14,7 +14,7 @@
# specific language governing permissions and limitations
# under the License.
"""FFI APIs for tvm.script"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("script", __name__)
+tvm_ffi._init_api("script", __name__)
diff --git a/python/tvm/script/ir_builder/_ffi_api.py b/python/tvm/script/ir_builder/_ffi_api.py
index 8ee2230519..fdca5f75dc 100644
--- a/python/tvm/script/ir_builder/_ffi_api.py
+++ b/python/tvm/script/ir_builder/_ffi_api.py
@@ -15,6 +15,6 @@
# specific language governing permissions and limitations
# under the License.
"""FFI APIs for tvm.script.ir_builder"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("script.ir_builder", __name__) # pylint: disable=protected-access
+tvm_ffi._init_api("script.ir_builder", __name__) # pylint: disable=protected-access
diff --git a/python/tvm/script/ir_builder/base.py b/python/tvm/script/ir_builder/base.py
index 95b5c50025..a6bb68e250 100644
--- a/python/tvm/script/ir_builder/base.py
+++ b/python/tvm/script/ir_builder/base.py
@@ -17,7 +17,7 @@
"""A generic IRBuilder across the TVM stack"""
from typing import Any, Callable, List
-from tvm.ffi import register_object as _register_object
+from tvm_ffi import register_object as _register_object
from tvm.runtime import Object as _Object
from . import _ffi_api
diff --git a/python/tvm/script/ir_builder/ir/_ffi_api.py b/python/tvm/script/ir_builder/ir/_ffi_api.py
index 5b9d801a6e..23b92904cb 100644
--- a/python/tvm/script/ir_builder/ir/_ffi_api.py
+++ b/python/tvm/script/ir_builder/ir/_ffi_api.py
@@ -15,6 +15,6 @@
# specific language governing permissions and limitations
# under the License.
"""FFI APIs"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("script.ir_builder.ir", __name__) # pylint: disable=protected-access
+tvm_ffi._init_api("script.ir_builder.ir", __name__) # pylint: disable=protected-access
diff --git a/python/tvm/script/ir_builder/ir/frame.py b/python/tvm/script/ir_builder/ir/frame.py
index d2737fde59..45b49221e3 100644
--- a/python/tvm/script/ir_builder/ir/frame.py
+++ b/python/tvm/script/ir_builder/ir/frame.py
@@ -16,7 +16,7 @@
# under the License.
"""Package tvm.script.ir_builder.ir.frame"""
-from tvm.ffi import register_object as _register_object
+from tvm_ffi import register_object as _register_object
from ..base import IRBuilderFrame
diff --git a/python/tvm/script/ir_builder/relax/_ffi_api.py b/python/tvm/script/ir_builder/relax/_ffi_api.py
index 1c767bacc4..251a24d4fa 100644
--- a/python/tvm/script/ir_builder/relax/_ffi_api.py
+++ b/python/tvm/script/ir_builder/relax/_ffi_api.py
@@ -15,6 +15,6 @@
# specific language governing permissions and limitations
# under the License.
"""FFI APIs for tvm.script.ir_builder.relax"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("script.ir_builder.relax", __name__) # pylint: disable=protected-access
+tvm_ffi._init_api("script.ir_builder.relax", __name__) # pylint: disable=protected-access
diff --git a/python/tvm/script/ir_builder/relax/distributed/_ffi_api.py b/python/tvm/script/ir_builder/relax/distributed/_ffi_api.py
index 4d2ba60c20..a69d7f3e38 100644
--- a/python/tvm/script/ir_builder/relax/distributed/_ffi_api.py
+++ b/python/tvm/script/ir_builder/relax/distributed/_ffi_api.py
@@ -15,8 +15,8 @@
# specific language governing permissions and limitations
# under the License.
"""FFI APIs for tvm.script.ir_builder.relax.distributed"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api(
+tvm_ffi._init_api(
"script.ir_builder.relax.distributed", __name__
) # pylint: disable=protected-access
diff --git a/python/tvm/script/ir_builder/relax/frame.py b/python/tvm/script/ir_builder/relax/frame.py
index 181f62ec4f..ed4d948ff9 100644
--- a/python/tvm/script/ir_builder/relax/frame.py
+++ b/python/tvm/script/ir_builder/relax/frame.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""IR Builder Frame for Relax dialect"""
-from tvm.ffi import register_object as _register_object
+from tvm_ffi import register_object as _register_object
from ..base import IRBuilderFrame
diff --git a/python/tvm/script/ir_builder/tir/_ffi_api.py b/python/tvm/script/ir_builder/tir/_ffi_api.py
index 69797f986a..42893a0047 100644
--- a/python/tvm/script/ir_builder/tir/_ffi_api.py
+++ b/python/tvm/script/ir_builder/tir/_ffi_api.py
@@ -15,6 +15,6 @@
# specific language governing permissions and limitations
# under the License.
"""FFI APIs"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("script.ir_builder.tir", __name__) # pylint: disable=protected-access
+tvm_ffi._init_api("script.ir_builder.tir", __name__) # pylint: disable=protected-access
diff --git a/python/tvm/script/ir_builder/tir/frame.py b/python/tvm/script/ir_builder/tir/frame.py
index e3ce2e6e2e..f43b4cf6ed 100644
--- a/python/tvm/script/ir_builder/tir/frame.py
+++ b/python/tvm/script/ir_builder/tir/frame.py
@@ -17,7 +17,7 @@
"""IRBuilder for TIR"""
from typing import List, Union
-from tvm.ffi import register_object as _register_object
+from tvm_ffi import register_object as _register_object
from tvm.tir import Buffer, Var
from ..base import IRBuilderFrame
diff --git a/python/tvm/script/printer/_ffi_api.py b/python/tvm/script/printer/_ffi_api.py
index 9cbf6cfdca..e219c9dbf8 100644
--- a/python/tvm/script/printer/_ffi_api.py
+++ b/python/tvm/script/printer/_ffi_api.py
@@ -15,6 +15,6 @@
# specific language governing permissions and limitations
# under the License.
"""FFI APIs for tvm.script.printer"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("script.printer", __name__) # pylint: disable=protected-access
+tvm_ffi._init_api("script.printer", __name__) # pylint: disable=protected-access
diff --git a/python/tvm/script/printer/doc.py b/python/tvm/script/printer/doc.py
index 382128ef33..62d8c563dd 100644
--- a/python/tvm/script/printer/doc.py
+++ b/python/tvm/script/printer/doc.py
@@ -19,8 +19,8 @@
from enum import IntEnum, unique
from typing import Dict, List, Optional, Sequence, Tuple, Union
-from tvm.ffi import register_object
-from tvm.ffi.access_path import AccessPath
+from tvm_ffi import register_object
+from tvm_ffi.access_path import AccessPath
from tvm.runtime import Object
from tvm.tir import FloatImm, IntImm
diff --git a/python/tvm/script/printer/doc_printer.py b/python/tvm/script/printer/doc_printer.py
index 5f1f980084..0cfc436b6a 100644
--- a/python/tvm/script/printer/doc_printer.py
+++ b/python/tvm/script/printer/doc_printer.py
@@ -18,7 +18,7 @@
from typing import List, Optional
-from tvm.ffi.access_path import AccessPath
+from tvm_ffi.access_path import AccessPath
from tvm.runtime.script_printer import PrinterConfig
from . import _ffi_api
diff --git a/python/tvm/support.py b/python/tvm/support.py
index 7e0ad5875f..5266602fd1 100644
--- a/python/tvm/support.py
+++ b/python/tvm/support.py
@@ -22,11 +22,11 @@ import os
import sys
import tvm
-import tvm.ffi
+import tvm_ffi
from .runtime.module import Module
from . import get_global_func
-tvm.ffi._init_api("support", __name__)
+tvm_ffi._init_api("support", __name__)
def libinfo():
diff --git a/python/tvm/target/_ffi_api.py b/python/tvm/target/_ffi_api.py
index 489b59b4c6..7520482388 100644
--- a/python/tvm/target/_ffi_api.py
+++ b/python/tvm/target/_ffi_api.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""FFI APIs for tvm.target"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("target", __name__)
+tvm_ffi._init_api("target", __name__)
diff --git a/python/tvm/target/datatype.py b/python/tvm/target/datatype.py
index c9be5531c7..0c282906ad 100644
--- a/python/tvm/target/datatype.py
+++ b/python/tvm/target/datatype.py
@@ -26,8 +26,8 @@ from tvm.tir.expr import (
BinaryOpExpr as _BinaryOpExpr,
)
from tvm.tir.op import call_pure_extern
-from tvm.ffi import get_global_func
-from tvm.ffi import register_func as _register_func
+from tvm_ffi import get_global_func
+from tvm_ffi import register_func as _register_func
from tvm.tir import call_intrin
@@ -215,7 +215,7 @@ def register_op(
)
else:
lower_func_name = "tvm.datatype.lower." + target + "." + op_name + "." + src_type_name
- tvm.ffi.register_func(lower_func_name, lower_func)
+ tvm_ffi.register_func(lower_func_name, lower_func)
def register_min_func(func, type_name):
diff --git a/python/tvm/target/detect_target.py b/python/tvm/target/detect_target.py
index ec1875eb90..689825cbe1 100644
--- a/python/tvm/target/detect_target.py
+++ b/python/tvm/target/detect_target.py
@@ -17,7 +17,7 @@
"""Detect target."""
from typing import Union
-from ..ffi import get_global_func
+from tvm_ffi import get_global_func
from ..runtime import Device
from ..runtime.ndarray import device
from . import Target
diff --git a/python/tvm/target/target.py b/python/tvm/target/target.py
index 6c83ef6e5b..64a7a893d8 100644
--- a/python/tvm/target/target.py
+++ b/python/tvm/target/target.py
@@ -20,8 +20,8 @@ import re
import warnings
from typing import Union
-import tvm.ffi
-from tvm.ffi import register_func as _register_func
+import tvm_ffi
+from tvm_ffi import register_func as _register_func
from tvm.runtime import Device
from tvm.runtime import Object, convert
from tvm.runtime.container import String
@@ -30,7 +30,7 @@ from tvm.ir.container import Map, Array
from . import _ffi_api
-@tvm.ffi.register_object("target.TargetKind")
+@tvm_ffi.register_object("target.TargetKind")
class TargetKind(Object):
"""Kind of a compilation target"""
@@ -53,7 +53,7 @@ class TargetFeatures:
return _ffi_api.TargetGetFeature(self.target, name)
-@tvm.ffi.register_object("target.Target")
+@tvm_ffi.register_object("target.Target")
class Target(Object):
"""Target device information, use through TVM API.
diff --git a/python/tvm/target/virtual_device.py b/python/tvm/target/virtual_device.py
index b062feb27a..e73de85cd3 100644
--- a/python/tvm/target/virtual_device.py
+++ b/python/tvm/target/virtual_device.py
@@ -16,14 +16,13 @@
# under the License.
"""Python bindings for creating VirtualDevices."""
-import tvm
-from tvm.runtime import Object
+import tvm_ffi
from . import _ffi_api
-@tvm.ffi.register_object("target.VirtualDevice")
-class VirtualDevice(Object):
+@tvm_ffi.register_object("target.VirtualDevice")
+class VirtualDevice(tvm_ffi.core.Object):
"""A compile time representation for where data is to be stored at runtime,
and how to compile code to compute it."""
diff --git a/python/tvm/target/x86.py b/python/tvm/target/x86.py
index 177021f143..874975383e 100644
--- a/python/tvm/target/x86.py
+++ b/python/tvm/target/x86.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""Common x86 related utilities"""
-from ..ffi import register_func
+from tvm_ffi import register_func
from .codegen import target_has_features
diff --git a/python/tvm/te/_ffi_api.py b/python/tvm/te/_ffi_api.py
index 98e466e9e8..8df8d5ff47 100644
--- a/python/tvm/te/_ffi_api.py
+++ b/python/tvm/te/_ffi_api.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""FFI APIs for tvm.te"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("te", __name__)
+tvm_ffi._init_api("te", __name__)
diff --git a/python/tvm/te/operation.py b/python/tvm/te/operation.py
index 4a5d2425e6..4367209fa6 100644
--- a/python/tvm/te/operation.py
+++ b/python/tvm/te/operation.py
@@ -21,7 +21,7 @@ import inspect
from numbers import Integral as _Integral
from typing import List, Optional, Union
-import tvm.ffi
+import tvm_ffi
import tvm.arith._ffi_api
import tvm.tir
import tvm.tir._ffi_api
diff --git a/python/tvm/te/tensor.py b/python/tvm/te/tensor.py
index 73b995a45e..61102085ef 100644
--- a/python/tvm/te/tensor.py
+++ b/python/tvm/te/tensor.py
@@ -16,7 +16,7 @@
# under the License.
"""Tensor class for computation declaration."""
# pylint: disable=invalid-name
-import tvm.ffi
+import tvm_ffi
from tvm.runtime import Object, ObjectGeneric
from tvm.tir import expr as _expr, DataProducer
@@ -48,7 +48,7 @@ class TensorSlice(ObjectGeneric, _expr.ExprOp):
return self.tensor.dtype
-@tvm.ffi.register_object("te.Tensor")
+@tvm_ffi.register_object("te.Tensor")
class Tensor(DataProducer, _expr.ExprOp):
"""Tensor object, to construct, see function.Tensor"""
@@ -92,7 +92,7 @@ class Tensor(DataProducer, _expr.ExprOp):
return f"{op.name}.v{self.value_index}"
-@tvm.ffi.register_object("te.Operation")
+@tvm_ffi.register_object("te.Operation")
class Operation(Object):
"""Represent an operation that generates a tensor"""
@@ -122,12 +122,12 @@ class Operation(Object):
return _ffi_api.OpInputTensors(self)
-@tvm.ffi.register_object("te.PlaceholderOp")
+@tvm_ffi.register_object("te.PlaceholderOp")
class PlaceholderOp(Operation):
"""Placeholder operation."""
-@tvm.ffi.register_object("te.BaseComputeOp")
+@tvm_ffi.register_object("te.BaseComputeOp")
class BaseComputeOp(Operation):
"""Compute operation."""
@@ -142,12 +142,12 @@ class BaseComputeOp(Operation):
return self.__getattr__("reduce_axis")
-@tvm.ffi.register_object("te.ComputeOp")
+@tvm_ffi.register_object("te.ComputeOp")
class ComputeOp(BaseComputeOp):
"""Scalar operation."""
-@tvm.ffi.register_object("te.ScanOp")
+@tvm_ffi.register_object("te.ScanOp")
class ScanOp(Operation):
"""Scan operation."""
@@ -157,6 +157,6 @@ class ScanOp(Operation):
return self.__getattr__("scan_axis")
-@tvm.ffi.register_object("te.ExternOp")
+@tvm_ffi.register_object("te.ExternOp")
class ExternOp(Operation):
"""External operation."""
diff --git a/python/tvm/testing/_ffi_api.py b/python/tvm/testing/_ffi_api.py
index e3c30d1299..4e57f4feaf 100644
--- a/python/tvm/testing/_ffi_api.py
+++ b/python/tvm/testing/_ffi_api.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""FFI APIs for tvm.testing"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("testing", __name__)
+tvm_ffi._init_api("testing", __name__)
diff --git a/python/tvm/testing/attrs.py b/python/tvm/testing/attrs.py
index ea6f1b1af6..7690af37d3 100644
--- a/python/tvm/testing/attrs.py
+++ b/python/tvm/testing/attrs.py
@@ -17,7 +17,7 @@
# pylint: disable=invalid-name, import-outside-toplevel, unused-variable
"""Testing utilities for attrs"""
from ..ir import Attrs
-from ..ffi import register_object
+from tvm_ffi import register_object
@register_object("attrs.TestAttrs")
diff --git a/python/tvm/testing/popen_pool.py b/python/tvm/testing/popen_pool.py
index 0fc3ce2190..c74829202b 100644
--- a/python/tvm/testing/popen_pool.py
+++ b/python/tvm/testing/popen_pool.py
@@ -16,7 +16,7 @@
# under the License.
# pylint: disable=invalid-name, missing-function-docstring
"""Common functions for popen_pool test cases"""
-import tvm
+import tvm_ffi
from . import _ffi_api
TEST_GLOBAL_STATE_1 = 0
@@ -36,19 +36,19 @@ def after_initializer():
return TEST_GLOBAL_STATE_1, TEST_GLOBAL_STATE_2, TEST_GLOBAL_STATE_3
-@tvm.ffi.register_func("testing.identity_py")
+@tvm_ffi.register_func("testing.identity_py")
def identity_py(arg):
return arg
def register_ffi():
- @tvm.ffi.register_func("testing.nested_identity_py")
+ @tvm_ffi.register_func("testing.nested_identity_py")
def _identity_py(arg): # pylint: disable=unused-variable
return arg
def call_py_ffi(arg):
- _identity_py = tvm.ffi.get_global_func("testing.nested_identity_py")
+ _identity_py = tvm_ffi.get_global_func("testing.nested_identity_py")
return _identity_py(arg)
diff --git a/python/tvm/testing/utils.py b/python/tvm/testing/utils.py
index 6b047de446..7c6f202085 100644
--- a/python/tvm/testing/utils.py
+++ b/python/tvm/testing/utils.py
@@ -87,7 +87,7 @@ import tvm
import tvm.arith
import tvm.tir
import tvm.te
-import tvm.ffi
+import tvm_ffi
from tvm.target import codegen
from tvm.contrib import nvcc, cudnn, rocm
diff --git a/python/tvm/tir/_ffi_api.py b/python/tvm/tir/_ffi_api.py
index 8c438557c8..2a004c9a83 100644
--- a/python/tvm/tir/_ffi_api.py
+++ b/python/tvm/tir/_ffi_api.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""FFI APIs for tvm.tir"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("tir", __name__)
+tvm_ffi._init_api("tir", __name__)
diff --git a/python/tvm/tir/analysis/_ffi_api.py b/python/tvm/tir/analysis/_ffi_api.py
index 40a7b4caf3..f228e8b30c 100644
--- a/python/tvm/tir/analysis/_ffi_api.py
+++ b/python/tvm/tir/analysis/_ffi_api.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""FFI APIs for tvm.tir.analysis"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("tir.analysis", __name__)
+tvm_ffi._init_api("tir.analysis", __name__)
diff --git a/python/tvm/tir/block_dependence_info.py b/python/tvm/tir/block_dependence_info.py
index 67a644967e..7bd6b418fc 100644
--- a/python/tvm/tir/block_dependence_info.py
+++ b/python/tvm/tir/block_dependence_info.py
@@ -18,7 +18,7 @@
to store the block level dependences"""
from typing import Union, Optional
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from tvm.ir.module import IRModule
from tvm.runtime import Object
from tvm.tir import Block, PrimFunc
diff --git a/python/tvm/tir/block_scope.py b/python/tvm/tir/block_scope.py
index b24cca0707..d63771fae9 100644
--- a/python/tvm/tir/block_scope.py
+++ b/python/tvm/tir/block_scope.py
@@ -18,7 +18,7 @@
from enum import IntEnum
from typing import List, Optional, Union
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from tvm.runtime import Object
from tvm.tir import Block, For
diff --git a/python/tvm/tir/buffer.py b/python/tvm/tir/buffer.py
index 1f40520e55..2590176082 100644
--- a/python/tvm/tir/buffer.py
+++ b/python/tvm/tir/buffer.py
@@ -17,14 +17,15 @@
"""Abstraction for array data structures."""
from numbers import Integral
-import tvm.ffi
+import tvm_ffi
+import tvm
from tvm.ir import PointerType, PrimExpr, PrimType, Range
from tvm.runtime import Object, Scriptable, convert
from . import _ffi_api
-@tvm.ffi.register_object("tir.Buffer")
+@tvm_ffi.register_object("tir.Buffer")
class Buffer(Object, Scriptable):
"""Symbolic data buffer in TVM.
@@ -349,6 +350,6 @@ def decl_buffer(
)
-@tvm.ffi.register_object("tir.DataProducer")
+@tvm_ffi.register_object("tir.DataProducer")
class DataProducer(Object):
pass
diff --git a/python/tvm/tir/data_layout.py b/python/tvm/tir/data_layout.py
index 39874640ff..f9c0e0cdc7 100644
--- a/python/tvm/tir/data_layout.py
+++ b/python/tvm/tir/data_layout.py
@@ -17,13 +17,13 @@
"""Data layout."""
from typing import Union
-import tvm.ffi
+import tvm_ffi
from tvm.runtime import Object
from . import _ffi_api
-@tvm.ffi.register_object("tir.Layout")
+@tvm_ffi.register_object("tir.Layout")
class Layout(Object):
"""Layout is composed of upper cases, lower cases and numbers,
where upper case indicates a primal axis and
@@ -81,7 +81,7 @@ class Layout(Object):
return _ffi_api.LayoutFactorOf(self, axis) # type: ignore
-@tvm.ffi.register_object("tir.BijectiveLayout")
+@tvm_ffi.register_object("tir.BijectiveLayout")
class BijectiveLayout(Object):
"""Bijective mapping for two layouts (src-layout and dst-layout).
It provides shape and index conversion between each other.
diff --git a/python/tvm/tir/expr.py b/python/tvm/tir/expr.py
index 2e07cef9a3..4fdee96a93 100644
--- a/python/tvm/tir/expr.py
+++ b/python/tvm/tir/expr.py
@@ -29,7 +29,7 @@ For example, you can use addexp.a to get the left operand of an Add node.
"""
from typing import List, Optional, Union
-import tvm.ffi
+import tvm_ffi
import tvm.ir._ffi_api
from tvm import ir
from tvm.ir import Op, PrimExpr
@@ -349,7 +349,7 @@ class LogicalExpr(PrimExprWithOp):
pass
-@tvm.ffi.register_object("tir.Var")
+@tvm_ffi.register_object("tir.Var")
class Var(PrimExprWithOp):
"""Symbolic variable.
@@ -372,7 +372,7 @@ class Var(PrimExprWithOp):
self.__init_handle_by_constructor__(_ffi_api.Var, name, dtype, span) # type: ignore
-@tvm.ffi.register_object("tir.SizeVar")
+@tvm_ffi.register_object("tir.SizeVar")
class SizeVar(Var):
"""Symbolic variable to represent a tensor index size
which is greater or equal to zero.
@@ -394,7 +394,7 @@ class SizeVar(Var):
self.__init_handle_by_constructor__(_ffi_api.SizeVar, name, dtype, span) # type: ignore
-@tvm.ffi.register_object("tir.IterVar")
+@tvm_ffi.register_object("tir.IterVar")
class IterVar(ExprOp, Object, Scriptable):
"""Represent iteration variable.
@@ -467,7 +467,7 @@ class IterVar(ExprOp, Object, Scriptable):
)
-@tvm.ffi.register_object("tir.CommReducer")
+@tvm_ffi.register_object("tir.CommReducer")
class CommReducer(Object, Scriptable):
"""Commutative reduce operator
@@ -507,7 +507,7 @@ class CommReducer(Object, Scriptable):
)
-@tvm.ffi.register_object("tir.Reduce")
+@tvm_ffi.register_object("tir.Reduce")
class Reduce(PrimExprWithOp):
"""Reduce node.
@@ -558,7 +558,7 @@ class Reduce(PrimExprWithOp):
)
-@tvm.ffi.register_object("ir.FloatImm")
+@tvm_ffi.register_object("ir.FloatImm")
class FloatImm(ConstExpr):
"""Float constant.
@@ -585,7 +585,7 @@ class FloatImm(ConstExpr):
return self.value
-@tvm.ffi.register_object("ir.IntImm")
+@tvm_ffi.register_object("ir.IntImm")
class IntImm(ConstExpr):
"""Int constant.
@@ -627,7 +627,7 @@ class IntImm(ConstExpr):
return self.__nonzero__()
-@tvm.ffi.register_object("tir.StringImm") # type: ignore
+@tvm_ffi.register_object("tir.StringImm") # type: ignore
class StringImm(ConstExpr):
"""String constant.
@@ -659,7 +659,7 @@ class StringImm(ConstExpr):
return PrimExpr.__hash__(self)
-@tvm.ffi.register_object("tir.Cast")
+@tvm_ffi.register_object("tir.Cast")
class Cast(PrimExprWithOp):
"""Cast expression.
@@ -681,7 +681,7 @@ class Cast(PrimExprWithOp):
self.__init_handle_by_constructor__(_ffi_api.Cast, dtype, value, span) # type: ignore
-@tvm.ffi.register_object("tir.Add")
+@tvm_ffi.register_object("tir.Add")
class Add(BinaryOpExpr):
"""Add node.
@@ -701,7 +701,7 @@ class Add(BinaryOpExpr):
self.__init_handle_by_constructor__(_ffi_api.Add, a, b, span) # type: ignore
-@tvm.ffi.register_object("tir.Sub")
+@tvm_ffi.register_object("tir.Sub")
class Sub(BinaryOpExpr):
"""Sub node.
@@ -721,7 +721,7 @@ class Sub(BinaryOpExpr):
self.__init_handle_by_constructor__(_ffi_api.Sub, a, b, span) # type: ignore
-@tvm.ffi.register_object("tir.Mul")
+@tvm_ffi.register_object("tir.Mul")
class Mul(BinaryOpExpr):
"""Mul node.
@@ -741,7 +741,7 @@ class Mul(BinaryOpExpr):
self.__init_handle_by_constructor__(_ffi_api.Mul, a, b, span) # type: ignore
-@tvm.ffi.register_object("tir.Div")
+@tvm_ffi.register_object("tir.Div")
class Div(BinaryOpExpr):
"""Div node.
@@ -761,7 +761,7 @@ class Div(BinaryOpExpr):
self.__init_handle_by_constructor__(_ffi_api.Div, a, b, span) # type: ignore
-@tvm.ffi.register_object("tir.Mod")
+@tvm_ffi.register_object("tir.Mod")
class Mod(BinaryOpExpr):
"""Mod node.
@@ -781,7 +781,7 @@ class Mod(BinaryOpExpr):
self.__init_handle_by_constructor__(_ffi_api.Mod, a, b, span) # type: ignore
-@tvm.ffi.register_object("tir.FloorDiv")
+@tvm_ffi.register_object("tir.FloorDiv")
class FloorDiv(BinaryOpExpr):
"""FloorDiv node.
@@ -801,7 +801,7 @@ class FloorDiv(BinaryOpExpr):
self.__init_handle_by_constructor__(_ffi_api.FloorDiv, a, b, span) # type: ignore
-@tvm.ffi.register_object("tir.FloorMod")
+@tvm_ffi.register_object("tir.FloorMod")
class FloorMod(BinaryOpExpr):
"""FloorMod node.
@@ -821,7 +821,7 @@ class FloorMod(BinaryOpExpr):
self.__init_handle_by_constructor__(_ffi_api.FloorMod, a, b, span) # type: ignore
-@tvm.ffi.register_object("tir.Min")
+@tvm_ffi.register_object("tir.Min")
class Min(BinaryOpExpr):
"""Min node.
@@ -841,7 +841,7 @@ class Min(BinaryOpExpr):
self.__init_handle_by_constructor__(_ffi_api.Min, a, b, span) # type: ignore
-@tvm.ffi.register_object("tir.Max")
+@tvm_ffi.register_object("tir.Max")
class Max(BinaryOpExpr):
"""Max node.
@@ -861,7 +861,7 @@ class Max(BinaryOpExpr):
self.__init_handle_by_constructor__(_ffi_api.Max, a, b, span) # type: ignore
-@tvm.ffi.register_object("tir.EQ")
+@tvm_ffi.register_object("tir.EQ")
class EQ(CmpExpr):
"""EQ node.
@@ -881,7 +881,7 @@ class EQ(CmpExpr):
self.__init_handle_by_constructor__(_ffi_api.EQ, a, b, span) # type: ignore
-@tvm.ffi.register_object("tir.NE")
+@tvm_ffi.register_object("tir.NE")
class NE(CmpExpr):
"""NE node.
@@ -901,7 +901,7 @@ class NE(CmpExpr):
self.__init_handle_by_constructor__(_ffi_api.NE, a, b, span) # type: ignore
-@tvm.ffi.register_object("tir.LT")
+@tvm_ffi.register_object("tir.LT")
class LT(CmpExpr):
"""LT node.
@@ -921,7 +921,7 @@ class LT(CmpExpr):
self.__init_handle_by_constructor__(_ffi_api.LT, a, b, span) # type: ignore
-@tvm.ffi.register_object("tir.LE")
+@tvm_ffi.register_object("tir.LE")
class LE(CmpExpr):
"""LE node.
@@ -941,7 +941,7 @@ class LE(CmpExpr):
self.__init_handle_by_constructor__(_ffi_api.LE, a, b, span) # type: ignore
-@tvm.ffi.register_object("tir.GT")
+@tvm_ffi.register_object("tir.GT")
class GT(CmpExpr):
"""GT node.
@@ -961,7 +961,7 @@ class GT(CmpExpr):
self.__init_handle_by_constructor__(_ffi_api.GT, a, b, span) # type: ignore
-@tvm.ffi.register_object("tir.GE")
+@tvm_ffi.register_object("tir.GE")
class GE(CmpExpr):
"""GE node.
@@ -981,7 +981,7 @@ class GE(CmpExpr):
self.__init_handle_by_constructor__(_ffi_api.GE, a, b, span) # type: ignore
-@tvm.ffi.register_object("tir.And")
+@tvm_ffi.register_object("tir.And")
class And(LogicalExpr):
"""And node.
@@ -1001,7 +1001,7 @@ class And(LogicalExpr):
self.__init_handle_by_constructor__(_ffi_api.And, a, b, span) # type: ignore
-@tvm.ffi.register_object("tir.Or")
+@tvm_ffi.register_object("tir.Or")
class Or(LogicalExpr):
"""Or node.
@@ -1024,7 +1024,7 @@ class Or(LogicalExpr):
self.__init_handle_by_constructor__(_ffi_api.Or, a, b, span) # type: ignore
-@tvm.ffi.register_object("tir.Not")
+@tvm_ffi.register_object("tir.Not")
class Not(LogicalExpr):
"""Not node.
@@ -1043,7 +1043,7 @@ class Not(LogicalExpr):
self.__init_handle_by_constructor__(_ffi_api.Not, a, span) # type: ignore
-@tvm.ffi.register_object("tir.Select")
+@tvm_ffi.register_object("tir.Select")
class Select(PrimExprWithOp):
"""Select node.
@@ -1087,7 +1087,7 @@ class Select(PrimExprWithOp):
)
-@tvm.ffi.register_object("tir.BufferLoad")
+@tvm_ffi.register_object("tir.BufferLoad")
class BufferLoad(PrimExprWithOp):
"""Buffer load node.
@@ -1122,7 +1122,7 @@ class BufferLoad(PrimExprWithOp):
)
-@tvm.ffi.register_object("tir.ProducerLoad")
+@tvm_ffi.register_object("tir.ProducerLoad")
class ProducerLoad(PrimExprWithOp):
"""Producer load node.
@@ -1149,7 +1149,7 @@ class ProducerLoad(PrimExprWithOp):
)
-@tvm.ffi.register_object("tir.Ramp")
+@tvm_ffi.register_object("tir.Ramp")
class Ramp(PrimExprWithOp):
"""Ramp node.
@@ -1180,7 +1180,7 @@ class Ramp(PrimExprWithOp):
)
-@tvm.ffi.register_object("tir.Broadcast")
+@tvm_ffi.register_object("tir.Broadcast")
class Broadcast(PrimExprWithOp):
"""Broadcast node.
@@ -1203,7 +1203,7 @@ class Broadcast(PrimExprWithOp):
self.__init_handle_by_constructor__(_ffi_api.Broadcast, value, lanes, span) # type: ignore
-@tvm.ffi.register_object("tir.Shuffle")
+@tvm_ffi.register_object("tir.Shuffle")
class Shuffle(PrimExprWithOp):
"""Shuffle node.
@@ -1241,7 +1241,7 @@ class CallEffectKind:
Opaque = UpdateState
-@tvm.ffi.register_object("tir.Call")
+@tvm_ffi.register_object("tir.Call")
class Call(PrimExprWithOp):
"""Call node.
@@ -1281,7 +1281,7 @@ class Call(PrimExprWithOp):
self.__init_handle_by_constructor__(_ffi_api.Call, dtype, op, args, span) # type: ignore
-@tvm.ffi.register_object("tir.Let")
+@tvm_ffi.register_object("tir.Let")
class Let(PrimExprWithOp):
"""Let node.
diff --git a/python/tvm/tir/function.py b/python/tvm/tir/function.py
index b85fb39522..e5786eb1a9 100644
--- a/python/tvm/tir/function.py
+++ b/python/tvm/tir/function.py
@@ -22,7 +22,7 @@ import inspect
from typing import Callable, List, Mapping, Optional, Tuple, Union
import tvm
-import tvm.ffi
+import tvm_ffi
import tvm.runtime
from tvm.ir import BaseFunc, Range
from tvm.runtime import Object, Scriptable
@@ -33,7 +33,7 @@ from .buffer import Buffer
from .expr import PrimExpr, Var
-@tvm.ffi.register_object("tir.PrimFunc")
+@tvm_ffi.register_object("tir.PrimFunc")
class PrimFunc(BaseFunc, Scriptable):
"""A function declaration expression.
@@ -174,7 +174,7 @@ class PrimFunc(BaseFunc, Scriptable):
return _ffi_api.Specialize(self, param_map) # type: ignore
-@tvm.ffi.register_object("tir.TensorIntrin")
+@tvm_ffi.register_object("tir.TensorIntrin")
class TensorIntrin(Object):
"""A tensor intrinsic.
@@ -230,7 +230,7 @@ class TensorIntrin(Object):
return _ffi_api.TensorIntrinGet(name, allow_missing) # pylint: type: ignore
-@tvm.ffi.register_object("tir.IndexMap")
+@tvm_ffi.register_object("tir.IndexMap")
class IndexMap(Object):
"""A mapping from multi-dimensional indices to another set of multi-dimensional indices
diff --git a/python/tvm/tir/functor.py b/python/tvm/tir/functor.py
index 06985f6645..c2594835fe 100644
--- a/python/tvm/tir/functor.py
+++ b/python/tvm/tir/functor.py
@@ -18,9 +18,8 @@
"""The expression and statement functor of TIR."""
from typing import Callable
-import tvm
+import tvm_ffi
from tvm.ir import PrimExpr
-from tvm.runtime import Object
from tvm.runtime.support import derived_object
from . import _ffi_api
@@ -144,8 +143,8 @@ Example
"""
-@tvm.ffi.register_object("tir.PyStmtExprVisitor")
-class _PyStmtExprVisitor(Object):
+@tvm_ffi.register_object("tir.PyStmtExprVisitor")
+class _PyStmtExprVisitor(tvm_ffi.core.Object):
"""
An internal wrapper to interface between C++ and Python StmtExprVisitor.
This is the TVM object that wraps PyStmtExprVisitor.
@@ -978,8 +977,8 @@ class PyStmtExprVisitor:
_ffi_api.PyStmtExprVisitorDefaultVisitExpr(self._outer(), op) # type: ignore
-@tvm.ffi.register_object("tir.PyStmtExprMutator")
-class _PyStmtExprMutator(Object):
+@tvm_ffi.register_object("tir.PyStmtExprMutator")
+class _PyStmtExprMutator(tvm_ffi.core.Object):
"""
A TVM object to support customization of StmtExprMutator on the python side.
This is the decorated result returned from stmt_expr_mutator decorator.
diff --git a/python/tvm/tir/op.py b/python/tvm/tir/op.py
index 54c70ede7a..ffd9aeff88 100644
--- a/python/tvm/tir/op.py
+++ b/python/tvm/tir/op.py
@@ -18,7 +18,8 @@
"""Operators used in TIR expression."""
from typing import Any, Optional, Union
-import tvm.ffi
+import tvm_ffi
+import tvm
from tvm import tir
from tvm.ir import Array, Op, PrimExpr
from tvm.ir.base import Span
@@ -1952,7 +1953,7 @@ def all(*args, span=None):
return val
-@tvm.ffi.register_func("tvm.default_trace_action")
+@tvm_ffi.register_func("tvm.default_trace_action")
def _tvm_default_trace_action(*args):
print(list(args))
@@ -3634,7 +3635,7 @@ def get_active_lane_mask(dtype, base, limit):
return call_intrin(dtype, "tir.get_active_lane_mask", base, limit)
-def get_vscale_expr(dtype: Union[str, tvm.ffi.dtype], min_size: int = 128) -> PrimExpr:
+def get_vscale_expr(dtype: Union[str, tvm_ffi.dtype], min_size: int = 128) -> PrimExpr:
"""
Create a datatype dependent scalable expression.
@@ -3646,7 +3647,7 @@ def get_vscale_expr(dtype: Union[str, tvm.ffi.dtype], min_size: int = 128) -> Pr
The minimum size of the scalable vector in bits.
"""
if isinstance(dtype, str):
- dtype = tvm.ffi.dtype(dtype)
+ dtype = tvm_ffi.dtype(dtype)
return min_size // dtype.bits * vscale()
diff --git a/python/tvm/tir/schedule/_ffi_api.py b/python/tvm/tir/schedule/_ffi_api.py
index b854145beb..99b831cdcd 100644
--- a/python/tvm/tir/schedule/_ffi_api.py
+++ b/python/tvm/tir/schedule/_ffi_api.py
@@ -15,6 +15,6 @@
# specific language governing permissions and limitations
# under the License.
"""FFI APIs for tvm.tir.schedule"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("tir.schedule", __name__) # pylint: disable=protected-access
+tvm_ffi._init_api("tir.schedule", __name__) # pylint: disable=protected-access
diff --git a/python/tvm/tir/schedule/analysis.py b/python/tvm/tir/schedule/analysis.py
index 491a689c93..66eab497eb 100644
--- a/python/tvm/tir/schedule/analysis.py
+++ b/python/tvm/tir/schedule/analysis.py
@@ -17,7 +17,7 @@
"""Analysis used in TensorIR scheduling"""
from typing import List, Optional
-import tvm.ffi
+import tvm_ffi
from tvm.runtime import Object
from ..buffer import Buffer
@@ -62,7 +62,7 @@ def suggest_index_map(
)
-@tvm.ffi.register_object("tir.schedule.TensorizeInfo")
+@tvm_ffi.register_object("tir.schedule.TensorizeInfo")
class TensorizeInfo(Object):
"""Necessary information used for tensorization."""
@@ -90,7 +90,7 @@ def get_tensorize_loop_mapping(
return _ffi_api.GetTensorizeLoopMapping(sch, block, desc_func, allow_padding) # type: ignore
-@tvm.ffi.register_object("tir.schedule.AutoTensorizeMappingInfo")
+@tvm_ffi.register_object("tir.schedule.AutoTensorizeMappingInfo")
class AutoTensorizeMappingInfo(Object):
"""Necessary information used to perform transformations for tensorization."""
diff --git a/python/tvm/tir/schedule/instruction.py b/python/tvm/tir/schedule/instruction.py
index 5a8563e652..918292a7bb 100644
--- a/python/tvm/tir/schedule/instruction.py
+++ b/python/tvm/tir/schedule/instruction.py
@@ -17,7 +17,7 @@
"""Schedule instructions each corresponds to a schedule primitive"""
from typing import TYPE_CHECKING, Any, List, Union
-from tvm.ffi import register_object as _register_object
+from tvm_ffi import register_object as _register_object
from tvm.runtime import Object
from . import _ffi_api
diff --git a/python/tvm/tir/schedule/schedule.py b/python/tvm/tir/schedule/schedule.py
index 5325ecdc16..ffa7e7174f 100644
--- a/python/tvm/tir/schedule/schedule.py
+++ b/python/tvm/tir/schedule/schedule.py
@@ -18,7 +18,7 @@
import inspect
from typing import Callable, Dict, List, Literal, Optional, Tuple, Union
-from tvm.ffi import register_object as _register_object
+from tvm_ffi import register_object as _register_object
from tvm.error import TVMError, register_error
from tvm.ir import GlobalVar, IRModule, PrimExpr
from tvm.runtime import Object
diff --git a/python/tvm/tir/schedule/state.py b/python/tvm/tir/schedule/state.py
index f082a9e92e..36436fe957 100644
--- a/python/tvm/tir/schedule/state.py
+++ b/python/tvm/tir/schedule/state.py
@@ -20,7 +20,7 @@ from collections import namedtuple
from enum import IntEnum
from typing import Dict, Optional, Union
-from tvm.ffi import register_object
+from tvm_ffi import register_object
from tvm.ir import IRModule
from tvm.runtime import Object
from tvm.tir import Block, BlockRealize, For, PrimFunc
diff --git a/python/tvm/tir/schedule/trace.py b/python/tvm/tir/schedule/trace.py
index da3508a42e..edc537f3a2 100644
--- a/python/tvm/tir/schedule/trace.py
+++ b/python/tvm/tir/schedule/trace.py
@@ -18,7 +18,7 @@
import os
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional
-from tvm.ffi import register_object as _register_object
+from tvm_ffi import register_object as _register_object
from tvm.runtime import Object
from ...ir import Array, Map, save_json
diff --git a/python/tvm/tir/stmt.py b/python/tvm/tir/stmt.py
index ffb6fd6a70..ed934183a5 100644
--- a/python/tvm/tir/stmt.py
+++ b/python/tvm/tir/stmt.py
@@ -29,7 +29,7 @@ Each statement node have subfields that can be visited from python side.
from enum import IntEnum
from typing import List, Mapping, Optional, Union
-import tvm.ffi
+import tvm_ffi
from tvm.ir import PrimExpr, Range, Span
from tvm.runtime import Object, Scriptable, const, NDArray
@@ -42,7 +42,7 @@ class Stmt(Object, Scriptable):
"""Base class of all the statements."""
-@tvm.ffi.register_object("tir.LetStmt")
+@tvm_ffi.register_object("tir.LetStmt")
class LetStmt(Stmt):
"""LetStmt node.
@@ -72,7 +72,7 @@ class LetStmt(Stmt):
)
-@tvm.ffi.register_object("tir.AssertStmt")
+@tvm_ffi.register_object("tir.AssertStmt")
class AssertStmt(Stmt):
"""AssertStmt node.
@@ -120,7 +120,7 @@ class ForKind(IntEnum):
THREAD_BINDING = 4 # pylint: disable=invalid-name
-@tvm.ffi.register_object("tir.For")
+@tvm_ffi.register_object("tir.For")
class For(Stmt):
"""For node.
@@ -185,7 +185,7 @@ class For(Stmt):
)
-@tvm.ffi.register_object("tir.While")
+@tvm_ffi.register_object("tir.While")
class While(Stmt):
"""While node.
@@ -209,7 +209,7 @@ class While(Stmt):
self.__init_handle_by_constructor__(_ffi_api.While, condition, body, span) # type: ignore
-@tvm.ffi.register_object("tir.BufferStore")
+@tvm_ffi.register_object("tir.BufferStore")
class BufferStore(Stmt):
"""Buffer store node.
@@ -252,7 +252,7 @@ class BufferStore(Stmt):
)
-@tvm.ffi.register_object("tir.BufferRealize")
+@tvm_ffi.register_object("tir.BufferRealize")
class BufferRealize(Stmt):
"""Buffer realize node.
@@ -293,7 +293,7 @@ class BufferRealize(Stmt):
)
-@tvm.ffi.register_object("tir.Allocate")
+@tvm_ffi.register_object("tir.Allocate")
class Allocate(Stmt):
"""Allocate node.
@@ -353,7 +353,7 @@ class Allocate(Stmt):
)
-@tvm.ffi.register_object("tir.AllocateConst")
+@tvm_ffi.register_object("tir.AllocateConst")
class AllocateConst(Stmt):
"""Allocate constant node.
@@ -415,7 +415,7 @@ class AllocateConst(Stmt):
)
-@tvm.ffi.register_object("tir.DeclBuffer")
+@tvm_ffi.register_object("tir.DeclBuffer")
class DeclBuffer(Stmt):
"""DeclBuffer node.
@@ -439,7 +439,7 @@ class DeclBuffer(Stmt):
self.__init_handle_by_constructor__(_ffi_api.DeclBuffer, buffer, body, span)
-@tvm.ffi.register_object("tir.AttrStmt")
+@tvm_ffi.register_object("tir.AttrStmt")
class AttrStmt(Stmt):
"""AttrStmt node.
@@ -475,7 +475,7 @@ class AttrStmt(Stmt):
)
-@tvm.ffi.register_object("tir.SeqStmt")
+@tvm_ffi.register_object("tir.SeqStmt")
class SeqStmt(Stmt):
"""Sequence of statements.
@@ -501,7 +501,7 @@ class SeqStmt(Stmt):
return len(self.seq)
-@tvm.ffi.register_object("tir.IfThenElse")
+@tvm_ffi.register_object("tir.IfThenElse")
class IfThenElse(Stmt):
"""IfThenElse node.
@@ -536,7 +536,7 @@ class IfThenElse(Stmt):
)
-@tvm.ffi.register_object("tir.Evaluate")
+@tvm_ffi.register_object("tir.Evaluate")
class Evaluate(Stmt):
"""Evaluate node.
@@ -556,7 +556,7 @@ class Evaluate(Stmt):
self.__init_handle_by_constructor__(_ffi_api.Evaluate, value, span) # type: ignore
-@tvm.ffi.register_object("tir.BufferRegion")
+@tvm_ffi.register_object("tir.BufferRegion")
class BufferRegion(Object, Scriptable):
"""BufferRegion node.
@@ -576,7 +576,7 @@ class BufferRegion(Object, Scriptable):
self.__init_handle_by_constructor__(_ffi_api.BufferRegion, buffer, region) # type: ignore
-@tvm.ffi.register_object("tir.MatchBufferRegion")
+@tvm_ffi.register_object("tir.MatchBufferRegion")
class MatchBufferRegion(Object, Scriptable):
"""MatchBufferRegion node.
@@ -598,7 +598,7 @@ class MatchBufferRegion(Object, Scriptable):
)
-@tvm.ffi.register_object("tir.Block")
+@tvm_ffi.register_object("tir.Block")
class Block(Stmt):
"""Block node.
@@ -680,7 +680,7 @@ class Block(Stmt):
) # type: ignore
-@tvm.ffi.register_object("tir.BlockRealize")
+@tvm_ffi.register_object("tir.BlockRealize")
class BlockRealize(Stmt):
"""BlockRealize node.
diff --git a/python/tvm/tir/tensor_intrin/cuda.py b/python/tvm/tir/tensor_intrin/cuda.py
index 6f964c9437..104acf2f44 100644
--- a/python/tvm/tir/tensor_intrin/cuda.py
+++ b/python/tvm/tir/tensor_intrin/cuda.py
@@ -18,7 +18,7 @@
"""Intrinsics for tensorization on NVIDIA GPU."""
from typing import Dict, Literal, Optional, Tuple
-from tvm.ffi import register_func
+from tvm_ffi import register_func
from tvm.runtime import convert
from tvm.script import tir as T
from tvm.tir import Cast, IntImm, TensorIntrin
diff --git a/python/tvm/tir/transform/_ffi_api.py b/python/tvm/tir/transform/_ffi_api.py
index 8a6607c11a..6a059ff0cf 100644
--- a/python/tvm/tir/transform/_ffi_api.py
+++ b/python/tvm/tir/transform/_ffi_api.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
"""FFI APIs for tvm.tir.transform"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("tir.transform", __name__)
+tvm_ffi._init_api("tir.transform", __name__)
diff --git a/python/tvm/tir/transform/function_pass.py b/python/tvm/tir/transform/function_pass.py
index b679d4ab16..a85eabd970 100644
--- a/python/tvm/tir/transform/function_pass.py
+++ b/python/tvm/tir/transform/function_pass.py
@@ -19,13 +19,13 @@ import inspect
import functools
from typing import Callable, List, Optional, Union
-import tvm.ffi
+import tvm_ffi
from tvm.ir.transform import Pass, PassInfo
from . import _ffi_api
-@tvm.ffi.register_object("tir.PrimFuncPass")
+@tvm_ffi.register_object("tir.PrimFuncPass")
class PrimFuncPass(Pass):
"""A pass that works on each :py:func:`tvm.tir.PrimFunc` in a module. A function
pass class should be created through py:func:`tvm.tir.transform.function_pass`.
diff --git a/python/tvm/topi/cpp/cuda.py b/python/tvm/topi/cpp/cuda.py
index 22f97293d3..d7d413fcf5 100644
--- a/python/tvm/topi/cpp/cuda.py
+++ b/python/tvm/topi/cpp/cuda.py
@@ -15,6 +15,6 @@
# specific language governing permissions and limitations
# under the License.
"""FFI for CUDA TOPI ops and schedules"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("topi.cuda", "tvm.topi.cpp.cuda")
+tvm_ffi._init_api("topi.cuda", "tvm.topi.cpp.cuda")
diff --git a/python/tvm/topi/cpp/generic.py b/python/tvm/topi/cpp/generic.py
index 3230d5428b..cafcdbcada 100644
--- a/python/tvm/topi/cpp/generic.py
+++ b/python/tvm/topi/cpp/generic.py
@@ -15,6 +15,6 @@
# specific language governing permissions and limitations
# under the License.
"""FFI for generic TOPI ops and schedules"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("topi.generic", "tvm.topi.cpp.generic")
+tvm_ffi._init_api("topi.generic", "tvm.topi.cpp.generic")
diff --git a/python/tvm/topi/cpp/impl.py b/python/tvm/topi/cpp/impl.py
index e5473a7e66..f906fc16d2 100644
--- a/python/tvm/topi/cpp/impl.py
+++ b/python/tvm/topi/cpp/impl.py
@@ -15,6 +15,6 @@
# specific language governing permissions and limitations
# under the License.
"""Load Lib for C++ TOPI ops and schedules"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("topi", "tvm.topi.cpp")
+tvm_ffi._init_api("topi", "tvm.topi.cpp")
diff --git a/python/tvm/topi/cpp/nn.py b/python/tvm/topi/cpp/nn.py
index 2ea1fc3714..b40bf834e0 100644
--- a/python/tvm/topi/cpp/nn.py
+++ b/python/tvm/topi/cpp/nn.py
@@ -15,6 +15,6 @@
# specific language governing permissions and limitations
# under the License.
"""FFI for NN TOPI ops and schedules"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("topi.nn", "tvm.topi.cpp.nn")
+tvm_ffi._init_api("topi.nn", "tvm.topi.cpp.nn")
diff --git a/python/tvm/topi/cpp/rocm.py b/python/tvm/topi/cpp/rocm.py
index 771fc3c3f0..eb14b0c7dc 100644
--- a/python/tvm/topi/cpp/rocm.py
+++ b/python/tvm/topi/cpp/rocm.py
@@ -15,6 +15,6 @@
# specific language governing permissions and limitations
# under the License.
"""FFI for Rocm TOPI ops and schedules"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("topi.rocm", "tvm.topi.cpp.rocm")
+tvm_ffi._init_api("topi.rocm", "tvm.topi.cpp.rocm")
diff --git a/python/tvm/topi/cpp/utils.py b/python/tvm/topi/cpp/utils.py
index b78a6baa0f..3e73ce7a9b 100644
--- a/python/tvm/topi/cpp/utils.py
+++ b/python/tvm/topi/cpp/utils.py
@@ -15,6 +15,6 @@
# specific language governing permissions and limitations
# under the License.
"""FFI for TOPI utility functions"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("topi.utils", "tvm.topi.cpp.utils")
+tvm_ffi._init_api("topi.utils", "tvm.topi.cpp.utils")
diff --git a/python/tvm/topi/cpp/vision/__init__.py b/python/tvm/topi/cpp/vision/__init__.py
index 5fdf1ac4e3..f47a21db78 100644
--- a/python/tvm/topi/cpp/vision/__init__.py
+++ b/python/tvm/topi/cpp/vision/__init__.py
@@ -16,8 +16,8 @@
# under the License.
"""FFI for vision TOPI ops and schedules"""
-import tvm.ffi
+import tvm_ffi
from . import yolo
-tvm.ffi._init_api("topi.vision", "tvm.topi.cpp.vision")
+tvm_ffi._init_api("topi.vision", "tvm.topi.cpp.vision")
diff --git a/python/tvm/topi/cpp/vision/yolo.py b/python/tvm/topi/cpp/vision/yolo.py
index 5d8bdd99d2..a2eb47dadb 100644
--- a/python/tvm/topi/cpp/vision/yolo.py
+++ b/python/tvm/topi/cpp/vision/yolo.py
@@ -15,6 +15,6 @@
# specific language governing permissions and limitations
# under the License.
"""FFI for Yolo TOPI ops and schedules"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("topi.vision.yolo", "tvm.topi.cpp.vision.yolo")
+tvm_ffi._init_api("topi.vision.yolo", "tvm.topi.cpp.vision.yolo")
diff --git a/python/tvm/topi/cpp/x86.py b/python/tvm/topi/cpp/x86.py
index 18de30c668..3432546075 100644
--- a/python/tvm/topi/cpp/x86.py
+++ b/python/tvm/topi/cpp/x86.py
@@ -15,6 +15,6 @@
# specific language governing permissions and limitations
# under the License.
"""FFI for x86 TOPI ops and schedules"""
-import tvm.ffi
+import tvm_ffi
-tvm.ffi._init_api("topi.x86", "tvm.topi.cpp.x86")
+tvm_ffi._init_api("topi.x86", "tvm.topi.cpp.x86")
diff --git a/tests/python/codegen/test_gpu_codegen_allreduce.py b/tests/python/codegen/test_gpu_codegen_allreduce.py
index aa56411cc9..5e8c3a05db 100644
--- a/tests/python/codegen/test_gpu_codegen_allreduce.py
+++ b/tests/python/codegen/test_gpu_codegen_allreduce.py
@@ -102,7 +102,7 @@ def optional_metal_compile_callback(define_metal_compile_callback):
if define_metal_compile_callback:
if cached is None:
- tvm.ffi.registry.remove_global_func(name)
+ tvm_ffi.registry.remove_global_func(name)
else:
tvm.register_func(name, cached, override=True)
diff --git a/tests/python/disco/test_loader.py b/tests/python/disco/test_loader.py
index 5089336f09..cf5955b10d 100644
--- a/tests/python/disco/test_loader.py
+++ b/tests/python/disco/test_loader.py
@@ -25,7 +25,7 @@ import tvm
import tvm.testing
from tvm import dlight as dl
from tvm import relax as rx
-from tvm.ffi import register_func
+from tvm_ffi import register_func
from tvm.contrib import tvmjs
from tvm.runtime import ShapeTuple
from tvm.runtime import disco as di
diff --git a/tests/python/ir/test_container_structural_equal.py b/tests/python/ir/test_container_structural_equal.py
index 251b33f910..957d0946ed 100644
--- a/tests/python/ir/test_container_structural_equal.py
+++ b/tests/python/ir/test_container_structural_equal.py
@@ -18,7 +18,7 @@ import pytest
import tvm
import tvm.testing
-from tvm.ffi.access_path import AccessPath
+from tvm_ffi.access_path import AccessPath
from tvm.ir.base import get_first_structural_mismatch
diff --git a/tests/python/ir/test_ir_container.py b/tests/python/ir/test_ir_container.py
index 1004bad702..7e7fc79960 100644
--- a/tests/python/ir/test_ir_container.py
+++ b/tests/python/ir/test_ir_container.py
@@ -90,7 +90,7 @@ def test_getattr_map():
a = te.var("a")
b = te.var("b")
amap = tvm.runtime.convert({a: 2, b: 3})
- assert isinstance(amap, tvm.ffi.Map)
+ assert isinstance(amap, tvm_ffi.Map)
def test_in_container():
diff --git a/tests/python/meta_schedule/test_meta_schedule_builder.py b/tests/python/meta_schedule/test_meta_schedule_builder.py
index 090a393fbe..a21d5a9195 100644
--- a/tests/python/meta_schedule/test_meta_schedule_builder.py
+++ b/tests/python/meta_schedule/test_meta_schedule_builder.py
@@ -25,7 +25,7 @@ import pytest
import tvm.testing
from tvm import script
-from tvm.ffi import register_func
+from tvm_ffi import register_func
from tvm.meta_schedule.builder import (
BuilderInput,
BuilderResult,
diff --git a/tests/python/meta_schedule/test_meta_schedule_post_order_apply.py b/tests/python/meta_schedule/test_meta_schedule_post_order_apply.py
index 57d9d09610..cbf2530eef 100644
--- a/tests/python/meta_schedule/test_meta_schedule_post_order_apply.py
+++ b/tests/python/meta_schedule/test_meta_schedule_post_order_apply.py
@@ -25,7 +25,7 @@ import tvm
import tvm.testing
from tvm import te
from tvm.ir.module import IRModule
-from tvm.ffi import register_func
+from tvm_ffi import register_func
from tvm.error import TVMError
from tvm.meta_schedule import TuneContext
from tvm.meta_schedule.schedule_rule import PyScheduleRule
diff --git a/tests/python/meta_schedule/test_meta_schedule_runner.py b/tests/python/meta_schedule/test_meta_schedule_runner.py
index e5deefe750..0d6a1e1e7f 100644
--- a/tests/python/meta_schedule/test_meta_schedule_runner.py
+++ b/tests/python/meta_schedule/test_meta_schedule_runner.py
@@ -25,7 +25,7 @@ import numpy as np
import pytest
import tvm
import tvm.testing
-from tvm.ffi import register_func
+from tvm_ffi import register_func
from tvm.meta_schedule.arg_info import TensorInfo
from tvm.meta_schedule.builder import BuilderInput, LocalBuilder
from tvm.meta_schedule.runner import (
diff --git a/tests/python/relax/test_op_inspect.py b/tests/python/relax/test_op_inspect.py
index 2ba9f9a709..130a98b1b3 100644
--- a/tests/python/relax/test_op_inspect.py
+++ b/tests/python/relax/test_op_inspect.py
@@ -170,7 +170,7 @@ def test_strides_of_non_compact_tensor():
expected_strides = [1, 4]
# use transpose to make strides non-compact
x = np.zeros([4, 4], "int32").T
- y = tvm.ffi.from_dlpack(x, required_alignment=4, required_contiguous=False)
+ y = tvm_ffi.from_dlpack(x, required_alignment=4, required_contiguous=False)
res = [vm["main"](y, i) for i, _ in enumerate(view_shape)]
tvm.ir.assert_structural_equal(res, expected_strides)
diff --git a/tests/python/tir-base/test_tir_structural_equal_hash.py b/tests/python/tir-base/test_tir_structural_equal_hash.py
index 601afd8f16..5e7c49ac14 100644
--- a/tests/python/tir-base/test_tir_structural_equal_hash.py
+++ b/tests/python/tir-base/test_tir_structural_equal_hash.py
@@ -18,7 +18,7 @@ import tvm
import numpy as np
import pytest
from tvm import te
-from tvm.ffi.access_path import AccessPath
+from tvm_ffi.access_path import AccessPath
from tvm.script import tir as T, ir as I
diff --git a/tests/python/tir-transform/test_tir_transform_inject_ptx_async_copy.py b/tests/python/tir-transform/test_tir_transform_inject_ptx_async_copy.py
index d4c93bb24a..0fac2177f7 100644
--- a/tests/python/tir-transform/test_tir_transform_inject_ptx_async_copy.py
+++ b/tests/python/tir-transform/test_tir_transform_inject_ptx_async_copy.py
@@ -16,6 +16,7 @@
# under the License.
import tvm
+import tvm_ffi
import tvm.testing
from tvm.script import tir as T
@@ -421,7 +422,7 @@ def postproc_if_missing_async_support():
# Restore previous postproc func to avoid impacting other tests
if prev_postproc is None:
- tvm.ffi.registry.remove_global_func(func_name)
+ tvm_ffi.registry.remove_global_func(func_name)
else:
tvm.register_func(func_name, prev_postproc, override=True)
diff --git a/tests/python/tvmscript/test_tvmscript_parser_tir.py b/tests/python/tvmscript/test_tvmscript_parser_tir.py
index 68e9adeff2..fd196be72a 100644
--- a/tests/python/tvmscript/test_tvmscript_parser_tir.py
+++ b/tests/python/tvmscript/test_tvmscript_parser_tir.py
@@ -18,6 +18,7 @@
import pytest
import tvm.testing
+import tvm_ffi
from tvm.script.parser import tir as T
from tvm import ir, tir
@@ -545,10 +546,10 @@ def test_deterministic_branch():
def test_block_annotation_merge():
- def _to_dict(anno: tvm.ffi.container.Map):
+ def _to_dict(anno: tvm_ffi.container.Map):
result = {}
for k, v in anno.items():
- result[k] = _to_dict(v) if isinstance(v, tvm.ffi.container.Map) else v
+ result[k] = _to_dict(v) if isinstance(v, tvm_ffi.container.Map) else v
return result
@T.prim_func
diff --git a/tests/python/tvmscript/test_tvmscript_printer_annotation.py b/tests/python/tvmscript/test_tvmscript_printer_annotation.py
index c45c0a91c5..74c66fb94c 100644
--- a/tests/python/tvmscript/test_tvmscript_printer_annotation.py
+++ b/tests/python/tvmscript/test_tvmscript_printer_annotation.py
@@ -18,7 +18,7 @@
from typing import Optional
import pytest
-from tvm.ffi.access_path import AccessPath
+from tvm_ffi.access_path import AccessPath
from tvm.script import tir as T
diff --git a/tests/python/tvmscript/test_tvmscript_printer_doc.py b/tests/python/tvmscript/test_tvmscript_printer_doc.py
index 20a705f9ff..f8e20915fa 100644
--- a/tests/python/tvmscript/test_tvmscript_printer_doc.py
+++ b/tests/python/tvmscript/test_tvmscript_printer_doc.py
@@ -21,7 +21,7 @@ Doc objects, then access and modify their attributes correctly.
import pytest
import tvm
-from tvm.ffi.access_path import AccessPath
+from tvm_ffi.access_path import AccessPath
from tvm.script.printer.doc import (
AssertDoc,
AssignDoc,
diff --git a/tests/python/tvmscript/test_tvmscript_printer_structural_equal.py b/tests/python/tvmscript/test_tvmscript_printer_structural_equal.py
index f3a385ca09..70473954eb 100644
--- a/tests/python/tvmscript/test_tvmscript_printer_structural_equal.py
+++ b/tests/python/tvmscript/test_tvmscript_printer_structural_equal.py
@@ -19,7 +19,7 @@ import pytest
import tvm
from tvm.ir import assert_structural_equal
-from tvm.ffi.access_path import AccessPath
+from tvm_ffi.access_path import AccessPath
from tvm.script import ir as I, tir as T
diff --git a/tests/python/tvmscript/test_tvmscript_printer_underlining.py b/tests/python/tvmscript/test_tvmscript_printer_underlining.py
index e36e96c77d..130bb7f237 100644
--- a/tests/python/tvmscript/test_tvmscript_printer_underlining.py
+++ b/tests/python/tvmscript/test_tvmscript_printer_underlining.py
@@ -18,7 +18,7 @@
from typing import Optional
import pytest
-from tvm.ffi.access_path import AccessPath
+from tvm_ffi.access_path import AccessPath
from tvm.script.printer.doc import (
ExprStmtDoc,
IdDoc,
diff --git a/tests/scripts/task_python_adreno.sh b/tests/scripts/task_python_adreno.sh
index f019cd1ecc..f4a5b8c402 100755
--- a/tests/scripts/task_python_adreno.sh
+++ b/tests/scripts/task_python_adreno.sh
@@ -57,8 +57,8 @@ trap "{ kill ${TRACKER_PID}; kill ${DEVICE_PID}; cleanup; }" 0
# cleanup pycache
find . -type f -path "*.pyc" | xargs rm -f
-# setup cython
-cd python; python3 setup.py build_ext --inplace; cd ..
+# setup tvm-ffi
+cd ffi; python3 -m pip install --user -v .; cd -
exit 0
diff --git a/tests/scripts/task_python_arm_compute_library.sh b/tests/scripts/task_python_arm_compute_library.sh
index 1423fb1985..8cf21fbd6d 100755
--- a/tests/scripts/task_python_arm_compute_library.sh
+++ b/tests/scripts/task_python_arm_compute_library.sh
@@ -23,5 +23,5 @@ source tests/scripts/setup-pytest-env.sh
find . -type f -path "*.pyc" | xargs rm -f
-# setup cython
-cd python; python3 setup.py build_ext --inplace; cd ..
+# setup tvm-ffi
+cd ffi; python3 -m pip install --user -v .; cd -
diff --git a/tests/scripts/task_python_docs.sh b/tests/scripts/task_python_docs.sh
index 7b58658bd7..ca86b77d50 100755
--- a/tests/scripts/task_python_docs.sh
+++ b/tests/scripts/task_python_docs.sh
@@ -126,8 +126,8 @@ clean_files
find . -type f -path "*.log" | xargs rm -f
find . -type f -path "*.pyc" | xargs rm -f
-# setup cython
-cd python; python3 setup.py build_ext --inplace; cd ..
+# setup tvm-ffi
+cd ffi; python3 -m pip install --user -v .; cd -
cd docs
diff --git a/tests/scripts/task_python_hexagon.sh b/tests/scripts/task_python_hexagon.sh
index fd53007a37..a146fcdd4f 100755
--- a/tests/scripts/task_python_hexagon.sh
+++ b/tests/scripts/task_python_hexagon.sh
@@ -27,8 +27,8 @@ fi
source tests/scripts/setup-pytest-env.sh
-# setup cython
-cd python; python3 setup.py build_ext --inplace; cd ..
+# setup tvm-ffi
+cd ffi; python3 -m pip install --user -v .; cd -
# disable hexagon tests for now
exit 0
diff --git a/tests/scripts/task_python_integration.sh b/tests/scripts/task_python_integration.sh
index 326743394d..25d34f09a6 100755
--- a/tests/scripts/task_python_integration.sh
+++ b/tests/scripts/task_python_integration.sh
@@ -33,5 +33,5 @@ fi
# cleanup pycache
find . -type f -path "*.pyc" | xargs rm -f
-# setup cython
-cd python; python3 setup.py build_ext --inplace; cd ..
+# setup tvm-ffi
+cd ffi; python3 -m pip install --user -v .; cd -
diff --git a/tests/scripts/task_python_nightly.sh b/tests/scripts/task_python_nightly.sh
index 42cf343e71..57c2300d54 100755
--- a/tests/scripts/task_python_nightly.sh
+++ b/tests/scripts/task_python_nightly.sh
@@ -20,8 +20,8 @@ set -euxo pipefail
source tests/scripts/setup-pytest-env.sh
-# setup cython
-cd python; python3 setup.py build_ext --inplace; cd ..
+# setup tvm-ffi
+cd ffi; python3 -m pip install --user -v .; cd -
# cleanup pycache
find . -type f -path "*.pyc" | xargs rm -f
diff --git a/tests/scripts/task_python_unittest.sh b/tests/scripts/task_python_unittest.sh
index 5417013353..4342aef498 100755
--- a/tests/scripts/task_python_unittest.sh
+++ b/tests/scripts/task_python_unittest.sh
@@ -23,8 +23,8 @@ source tests/scripts/setup-pytest-env.sh
# cleanup pycache
find . -type f -path "*.pyc" | xargs rm -f
-# setup cython
-cd python; python3 setup.py build_ext --inplace; cd ..
+# setup tvm-ffi
+cd ffi; python3 -m pip install --user -v .; cd -
# NOTE: also set by task_python_unittest_gpuonly.sh.
if [ -z "${TVM_UNITTEST_TESTSUITE_NAME:-}" ]; then
diff --git a/tests/scripts/unity/task_python_relax.sh b/tests/scripts/unity/task_python_relax.sh
index 5a72254924..532bf4987f 100755
--- a/tests/scripts/unity/task_python_relax.sh
+++ b/tests/scripts/unity/task_python_relax.sh
@@ -25,8 +25,8 @@ export LD_LIBRARY_PATH="build:${LD_LIBRARY_PATH:-}"
export TVM_BIND_THREADS=0
export TVM_NUM_THREADS=2
-# setup cython
-cd python; python3 setup.py build_ext --inplace; cd ..
+# setup tvm-ffi
+cd ffi; python3 -m pip install --user -v .; cd -
# Run Relax tests
TVM_TEST_TARGETS="${TVM_RELAY_TEST_TARGETS:-llvm}" pytest tests/python/relax