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 2026/10/12 19:27:33 UTC
(tvm) branch main updated: [REFACTOR] Clarify core ownership and native compilation interfaces (#20678)
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new ca9f2259be [REFACTOR] Clarify core ownership and native compilation interfaces (#20678)
ca9f2259be is described below
commit ca9f2259be5779f3dd0d06ec9078d29d4fe2eaed
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Mon Oct 12 15:27:26 2026 -0400
[REFACTOR] Clarify core ownership and native compilation interfaces (#20678)
Clarify code ownership and dependency boundaries while preserving the
current library layout and native identities. Core-owned declarations
use `TVM_TIRX_DLL`, and `tvm.tirx.compile` becomes the canonical native
compilation API returning `Executable`.
- Keep scheduled lowering in S-TIR and native device-scope decoration in
TIRx.
- Separate local timing from RPC dispatch and let dialects register
their own Python behavior.
- Keep WASM integration parent-owned through explicit callbacks, place
narrow Hexagon linker support with its backend, and bind optional
runtime implementations to runtime-extra.
- Place VM constant loading and LLVM target canonicalization with their
owning source groups, preserving registrations and compilation behavior.
---
CMakeLists.txt | 6 +-
cmake/modules/contrib/AMX.cmake | 6 +-
cmake/modules/contrib/CLML.cmake | 4 +-
cmake/modules/contrib/ExampleNPU.cmake | 4 +-
cmake/modules/contrib/NNAPI.cmake | 4 +-
cmake/modules/contrib/TensorRT.cmake | 4 +-
docs/arch/codegen.rst | 14 +-
docs/arch/relax_vm.rst | 4 +-
.../tensor_ir/tutorials/tir_transformation.py | 2 +-
docs/reference/api/python/support.rst | 4 +-
docs/tirx/api/compilation.rst | 6 +-
include/tvm/backend/cuda/op/iket.h | 12 +-
include/tvm/backend/cuda/op/tensormap.h | 2 +-
include/tvm/backend/metal/op/cooperative_tensor.h | 8 +-
include/tvm/backend/opencl/op/memory.h | 6 +-
include/tvm/ir/base_expr.h | 30 ++--
include/tvm/ir/env_func.h | 2 +-
include/tvm/ir/expr.h | 94 +++++-----
include/tvm/ir/expr_functor.h | 4 +-
include/tvm/ir/location.h | 16 +-
include/tvm/ir/module.h | 44 ++---
include/tvm/ir/object_functor.h | 4 +-
include/tvm/ir/op.h | 26 +--
include/tvm/ir/prim/expr.h | 58 +++---
include/tvm/ir/prim/op.h | 196 ++++++++++-----------
include/tvm/ir/prim/vector_expr.h | 13 +-
include/tvm/ir/stmt.h | 47 ++---
include/tvm/ir/transform.h | 38 ++--
include/tvm/ir/type.h | 16 +-
include/tvm/ir/unique_name_supply.h | 6 +-
include/tvm/runtime/base.h | 8 +
include/tvm/script/ir_builder/ir.h | 6 +-
include/tvm/script/printer/config.h | 4 +-
include/tvm/script/printer/doc.h | 88 ++++-----
include/tvm/script/printer/doc_translator.h | 5 +-
include/tvm/script/printer/printer.h | 8 +-
include/tvm/sym/analyzer.h | 90 +++++-----
include/tvm/sym/int_set.h | 12 +-
include/tvm/sym/iter_affine_map.h | 12 +-
include/tvm/target/target.h | 22 +--
include/tvm/target/target_kind.h | 12 +-
include/tvm/target/target_tag_registry.h | 10 +-
include/tvm/tirx/analysis.h | 14 +-
include/tvm/tirx/expr.h | 25 +--
include/tvm/tirx/function.h | 10 +-
include/tvm/tirx/layout.h | 12 +-
include/tvm/tirx/op/abi.h | 32 ++--
include/tvm/tirx/op/annotation.h | 4 +-
include/tvm/tirx/op/cpu.h | 2 +-
include/tvm/tirx/op/gpu.h | 26 +--
include/tvm/tirx/op/math.h | 4 +-
include/tvm/tirx/op/memory.h | 32 ++--
include/tvm/tirx/op/region.h | 10 +-
include/tvm/tirx/op/vector.h | 8 +-
include/tvm/tirx/stmt_functor.h | 8 +-
include/tvm/tirx/tile/context.h | 16 +-
include/tvm/tirx/tile/dispatch.h | 2 +-
include/tvm/tirx/tile/exec_scope.h | 6 +-
include/tvm/tirx/tile/op.h | 2 +-
include/tvm/tirx/transform.h | 65 +++----
include/tvm/tirx/type.h | 14 +-
jvm/core/src/test/scripts/prepare_test_libs.py | 4 +-
python/tvm/backend/cuda/iket.py | 2 +-
python/tvm/backend/hexagon/__init__.py | 5 +-
python/tvm/backend/hexagon/linker.py | 112 ++++++++++++
python/tvm/backend/trn/pipeline.py | 8 +-
python/tvm/{support => contrib}/emcc.py | 4 +-
python/tvm/contrib/hexagon/tools.py | 98 ++---------
python/tvm/contrib/tvmjs.py | 4 +-
python/tvm/driver/build_module.py | 3 +-
python/tvm/ir/_tensor_expr_overload.py | 37 ++++
python/tvm/ir/base.py | 14 ++
python/tvm/relax/__init__.py | 2 +
python/tvm/relax/expr.py | 1 -
python/tvm/relax/op/__init__.py | 68 +++----
python/tvm/relax/script/printer.py | 65 +++++++
python/tvm/relax/vm_build.py | 2 +-
python/tvm/runtime/module.py | 56 +++---
python/tvm/s_tir/dlight/benchmark/bench.py | 2 +-
.../s_tir/meta_schedule/builder/local_builder.py | 4 +-
python/tvm/s_tir/pipeline.py | 13 +-
python/tvm/script/printer/__init__.py | 4 +-
python/tvm/script/printer/scriptable.py | 50 +-----
python/tvm/tirx/__init__.py | 2 +-
python/tvm/tirx/{build.py => _compile.py} | 43 ++---
python/tvm/tirx/compilation_pipeline.py | 19 +-
python/tvm/tirx/transform/transform.py | 5 +
src/backend/cuda/tile/exec_context.h | 17 +-
src/backend/cuda/tile/exec_scope.h | 9 +-
src/backend/cuda/tile/filter_canonical.h | 4 +-
.../hexagon/codegen/llvm/codegen_hexagon.cc | 2 +-
src/backend/vulkan/codegen/codegen_spirv.cc | 3 +-
src/ir/instrument.cc | 2 +-
src/ir/prim/op_utils.h | 2 +-
src/relax/backend/vm/codegen_vm.cc | 2 +-
src/relax/transform/fold_constant.cc | 2 +-
src/runtime/extra/contrib/amx/amx_config.cc | 9 +-
src/runtime/rpc/rpc_module.cc | 132 +++-----------
src/runtime/timer.cc | 96 ++++++++++
src/runtime/{ => vm}/const_loader_module.cc | 12 +-
src/runtime/{ => vm}/const_loader_module.h | 12 +-
.../space_generator/space_generator.cc | 2 +-
src/s_tir/transform/decorate_device_scope.cc | 13 +-
src/sym/interval_set.h | 6 +-
src/sym/presburger_set.h | 8 +-
src/sym/z3_prover.cc | 28 +--
.../llvm => llvm/canonicalizer}/arm_aprofile.cc | 4 +-
.../llvm => llvm/canonicalizer}/arm_aprofile.h | 8 +-
.../llvm => llvm/canonicalizer}/arm_mprofile.cc | 2 +-
.../llvm => llvm/canonicalizer}/arm_mprofile.h | 8 +-
.../llvm => llvm/canonicalizer}/canonicalize.cc | 0
.../llvm => llvm/canonicalizer}/canonicalize.h | 8 +-
src/target/target_kind.cc | 2 +-
src/tirx/ir/transform.cc | 2 +-
.../transform/decorate_device_scope.cc | 25 +--
.../transform/decorate_device_scope.h} | 21 ++-
tests/python/codegen/test_codegen_assert.py | 16 +-
.../python/codegen/test_codegen_error_handling.py | 28 +--
tests/python/codegen/test_gpu_codegen_allreduce.py | 6 +-
tests/python/codegen/test_target_codegen.py | 10 +-
.../python/codegen/test_target_codegen_aarch64.py | 30 ++--
tests/python/codegen/test_target_codegen_arm.py | 6 +-
tests/python/codegen/test_target_codegen_blob.py | 4 +-
tests/python/codegen/test_target_codegen_bool.py | 4 +-
tests/python/codegen/test_target_codegen_c_host.py | 20 +--
.../codegen/test_target_codegen_cross_llvm.py | 2 +-
tests/python/codegen/test_target_codegen_cuda.py | 64 +++----
.../codegen/test_target_codegen_cuda_fastmath.py | 2 +-
.../python/codegen/test_target_codegen_cuda_fp4.py | 6 +-
.../python/codegen/test_target_codegen_cuda_fp8.py | 24 +--
tests/python/codegen/test_target_codegen_device.py | 4 +-
tests/python/codegen/test_target_codegen_extern.py | 4 +-
.../codegen/test_target_codegen_gpu_common.py | 2 +-
.../python/codegen/test_target_codegen_hexagon.py | 10 +-
tests/python/codegen/test_target_codegen_llvm.py | 94 +++++-----
.../python/codegen/test_target_codegen_llvm_vla.py | 10 +-
tests/python/codegen/test_target_codegen_metal.py | 22 +--
tests/python/codegen/test_target_codegen_opencl.py | 16 +-
tests/python/codegen/test_target_codegen_riscv.py | 8 +-
tests/python/codegen/test_target_codegen_rocm.py | 10 +-
.../codegen/test_target_codegen_static_init.py | 2 +-
tests/python/codegen/test_target_codegen_vulkan.py | 28 +--
tests/python/codegen/test_target_codegen_webgpu.py | 4 +-
tests/python/codegen/test_target_codegen_x86.py | 2 +-
.../test_runtime_builtin_kv_cache_transfer.py | 2 +-
.../relax/test_frontend_nn_llm_kernel_config.py | 13 +-
...test_frontend_nn_llm_sequence_prefill_masked.py | 2 +-
...runtime_builtin_paged_attention_kv_cache_cpu.py | 2 +-
..._builtin_paged_attention_kv_cache_flashinfer.py | 2 +-
...ltin_paged_attention_kv_cache_mla_flashinfer.py | 2 +-
...ime_builtin_paged_attention_kv_cache_mla_tir.py | 2 +-
...runtime_builtin_paged_attention_kv_cache_tir.py | 4 +-
.../python/relax/test_runtime_builtin_rnn_state.py | 2 +-
tests/python/runtime/test_executable.py | 14 +-
tests/python/runtime/test_runtime_measure.py | 2 +-
tests/python/runtime/test_runtime_module_export.py | 4 +-
tests/python/runtime/test_runtime_module_load.py | 26 +--
.../python/runtime/test_runtime_module_property.py | 2 +-
.../python/s_tir/base/test_compilation_pipeline.py | 2 +-
.../python/s_tir/base/test_s_tir_texture_scope.py | 4 +-
tests/python/s_tir/dlight/test_cpu_reduction.py | 8 +-
...st_s_tir_transform_force_narrow_index_to_i32.py | 2 +-
tests/python/te/test_te_create_function.py | 2 +-
tests/python/tirx-base/test_tir_base.py | 2 +-
tests/python/tirx-base/test_tir_imm_values.py | 48 ++---
tests/python/tirx-base/test_tir_intrin.py | 14 +-
tests/python/tirx-base/test_tir_ptx_cp_async.py | 2 +-
.../tirx-base/test_tir_ptx_griddepcontrol.py | 2 +-
tests/python/tirx-base/test_tir_ptx_ldmatrix.py | 4 +-
tests/python/tirx-base/test_tir_ptx_mma.py | 34 ++--
tests/python/tirx-base/test_tir_ptx_mma_sp.py | 4 +-
.../tirx-base/test_tir_ptx_scalar_f32_math.py | 2 +-
.../test_tir_transform_flatten_tensor.py | 2 +-
.../test_tir_transform_lower_intrin.py | 4 +-
.../test_tir_transform_lower_tvm_builtin.py | 6 +-
.../tirx-transform/test_tir_transform_vectorize.py | 8 +-
tests/python/tirx/codegen/test_codegen_ampere.py | 2 +-
.../python/tirx/codegen/test_codegen_blackwell.py | 2 +-
tests/python/tirx/codegen/test_codegen_cuda.py | 10 +-
tests/python/tirx/codegen/test_codegen_dsmem.py | 2 +-
tests/python/tirx/codegen/test_codegen_hopper.py | 24 +--
tests/python/tirx/codegen/test_codegen_nki.py | 2 +-
tests/python/tirx/codegen/test_codegen_nvshmem.py | 2 +-
tests/python/tirx/codegen/test_cuda_cta_reduce.py | 2 +-
tests/python/tirx/codegen/test_cuda_wait_until.py | 2 +-
tests/python/tirx/codegen/test_cuda_warp_reduce.py | 2 +-
tests/python/tirx/codegen/test_ptx_addr.py | 2 +-
tests/python/tirx/codegen/test_ptx_dialect.py | 6 +-
tests/python/tirx/codegen/test_ptx_ld_st_ops.py | 14 +-
.../codegen/test_tcgen05_descriptor_encoders.py | 2 +-
tests/python/tirx/iket/iket_profile_workload.py | 2 +-
tests/python/tirx/iket/test_iket_profiler.py | 8 +-
.../tirx/script/test_tirx_script_dynamic_shape.py | 4 +-
.../tirx/script/test_tirx_script_ir_builder.py | 2 +-
tests/python/tirx/test_compilation_pipeline.py | 23 +--
tests/python/tirx/test_control_flow.py | 2 +-
tests/python/tirx/test_cuda_backend_config.py | 16 +-
tests/python/tirx/test_cuda_launch.py | 6 +-
tests/python/tirx/test_jit.py | 2 +-
tests/python/tirx/test_tile_namespace.py | 4 +-
tests/python/tirx/tile/cuda/copy/test_gmem_smem.py | 6 +-
.../python/tirx/tile/cuda/copy/test_ld_stmatrix.py | 2 +-
tests/python/tirx/tile/cuda/copy/test_reg.py | 34 ++--
.../tirx/tile/cuda/copy/test_vec_forced_cache.py | 4 +-
.../python/tirx/tile/cuda/copy_async/test_dsmem.py | 2 +-
.../tirx/tile/cuda/copy_async/test_ldgsts.py | 4 +-
.../tirx/tile/cuda/copy_async/test_tcgen05_cp.py | 4 +-
.../tirx/tile/cuda/copy_async/test_tcgen05_ldst.py | 32 ++--
tests/python/tirx/tile/cuda/copy_async/test_tma.py | 4 +-
.../tirx/tile/cuda/elementwise/test_binary.py | 26 +--
.../python/tirx/tile/cuda/elementwise/test_fma.py | 16 +-
.../tirx/tile/cuda/elementwise/test_unary.py | 38 ++--
.../tirx/tile/cuda/gemm/test_gemm_mma_m16n8k_.py | 14 +-
.../tirx/tile/cuda/gemm_async/test_gemm_async.py | 92 +++++-----
web/README.md | 2 +-
web/tests/python/prepare_test_libs.py | 8 +-
web/tests/python/relax_rpc_test.py | 4 +-
web/tests/python/webgpu_rpc_test.py | 6 +-
218 files changed, 1776 insertions(+), 1649 deletions(-)
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 6070108922..c3bff01777 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -342,7 +342,7 @@ tvm_file_glob(GLOB CODEGEN_SRCS
src/target/*.cc
src/target/source/*.cc
src/target/canonicalizer/*.cc
- src/target/canonicalizer/llvm/*.cc
+ src/target/llvm/canonicalizer/*.cc
src/backend/cuda/codegen/*.cc
src/backend/cuda/op/*.cc
src/backend/cuda/transforms/*.cc
@@ -367,9 +367,11 @@ list(APPEND COMPILER_SRCS ${CODEGEN_SRCS})
tvm_file_glob(GLOB RUNTIME_SRCS
src/runtime/*.cc
- src/runtime/vm/*.cc
src/runtime/rpc/minrpc/*.cc
)
+# VM-owned runtime implementations, including external-module constant loading.
+tvm_file_glob(GLOB RUNTIME_VM_SRCS src/runtime/vm/*.cc)
+list(APPEND RUNTIME_SRCS ${RUNTIME_VM_SRCS})
# Note: src/runtime/extra/disco/** moves to libtvm_runtime_extra.
# Note: src/backend/*/runtime sources move to per-backend DSOs.
set(TVM_RUNTIME_EXT_OBJS "")
diff --git a/cmake/modules/contrib/AMX.cmake b/cmake/modules/contrib/AMX.cmake
index f01377417a..553171400d 100644
--- a/cmake/modules/contrib/AMX.cmake
+++ b/cmake/modules/contrib/AMX.cmake
@@ -17,7 +17,9 @@
if(USE_AMX)
file(GLOB AMX_RUNTIME_CONFIG src/runtime/extra/contrib/amx/amx_config.cc)
- list(APPEND COMPILER_SRCS ${AMX_RUNTIME_CONFIG})
- set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -march=sapphirerapids")
+ add_library(tvm_amx_objs OBJECT ${AMX_RUNTIME_CONFIG})
+ target_link_libraries(tvm_amx_objs PRIVATE tvm_runtime_extra_defs)
+ target_compile_options(tvm_amx_objs PRIVATE -march=sapphirerapids)
+ target_link_libraries(tvm_runtime_extra PRIVATE tvm_amx_objs)
message(STATUS "Build with Intel AMX support...")
endif()
diff --git a/cmake/modules/contrib/CLML.cmake b/cmake/modules/contrib/CLML.cmake
index 2f12713b13..d3e80a2e84 100644
--- a/cmake/modules/contrib/CLML.cmake
+++ b/cmake/modules/contrib/CLML.cmake
@@ -21,7 +21,9 @@ if(USE_CLML)
include_directories(SYSTEM "3rdparty/OpenCL-Headers")
list(APPEND COMPILER_SRCS ${CLML_RELAX_CONTRIB_SRC})
if(NOT USE_CLML_GRAPH_EXECUTOR)
- list(APPEND COMPILER_SRCS ${CLML_RUNTIME_MODULE})
+ add_library(tvm_clml_objs OBJECT ${CLML_RUNTIME_MODULE})
+ target_link_libraries(tvm_clml_objs PRIVATE tvm_runtime_extra_defs)
+ target_link_libraries(tvm_runtime_extra PRIVATE tvm_clml_objs)
endif()
message(STATUS "Build with CLML support : " ${USE_CLML})
if(NOT USE_CLML STREQUAL "ON")
diff --git a/cmake/modules/contrib/ExampleNPU.cmake b/cmake/modules/contrib/ExampleNPU.cmake
index b2567a44fd..6ae6257ce7 100644
--- a/cmake/modules/contrib/ExampleNPU.cmake
+++ b/cmake/modules/contrib/ExampleNPU.cmake
@@ -24,7 +24,9 @@ if(USE_EXAMPLE_NPU_CODEGEN)
tvm_file_glob(GLOB RUNTIME_EXAMPLE_NPU_SRCS src/runtime/extra/contrib/example_npu/*.cc)
if(NOT USE_EXAMPLE_NPU_RUNTIME)
- list(APPEND COMPILER_SRCS ${RUNTIME_EXAMPLE_NPU_SRCS})
+ add_library(tvm_example_npu_objs OBJECT ${RUNTIME_EXAMPLE_NPU_SRCS})
+ target_link_libraries(tvm_example_npu_objs PRIVATE tvm_runtime_extra_defs)
+ target_link_libraries(tvm_runtime_extra PRIVATE tvm_example_npu_objs)
endif()
endif()
diff --git a/cmake/modules/contrib/NNAPI.cmake b/cmake/modules/contrib/NNAPI.cmake
index 6d1c44b842..0b3f5e5db5 100644
--- a/cmake/modules/contrib/NNAPI.cmake
+++ b/cmake/modules/contrib/NNAPI.cmake
@@ -23,7 +23,9 @@ if(USE_NNAPI_CODEGEN)
tvm_file_glob(GLOB RUNTIME_NNAPI_SRCS src/runtime/extra/contrib/nnapi/*.cc)
list(APPEND COMPILER_SRCS ${COMPILER_NNAPI_SRCS})
if(NOT USE_NNAPI_RUNTIME)
- list(APPEND COMPILER_SRCS ${RUNTIME_NNAPI_SRCS})
+ add_library(tvm_nnapi_objs OBJECT ${RUNTIME_NNAPI_SRCS})
+ target_link_libraries(tvm_nnapi_objs PRIVATE tvm_runtime_extra_defs)
+ target_link_libraries(tvm_runtime_extra PRIVATE tvm_nnapi_objs)
endif()
endif()
diff --git a/cmake/modules/contrib/TensorRT.cmake b/cmake/modules/contrib/TensorRT.cmake
index 0b70d77572..299140b92b 100644
--- a/cmake/modules/contrib/TensorRT.cmake
+++ b/cmake/modules/contrib/TensorRT.cmake
@@ -29,7 +29,9 @@ if(USE_TENSORRT_CODEGEN)
set_source_files_properties(${RUNTIME_TENSORRT_SRCS} PROPERTIES COMPILE_FLAGS "-Wno-deprecated-declarations")
list(APPEND COMPILER_SRCS ${COMPILER_TENSORRT_SRCS})
if(NOT USE_TENSORRT_RUNTIME)
- list(APPEND COMPILER_SRCS ${RUNTIME_TENSORRT_SRCS})
+ add_library(tvm_tensorrt_objs OBJECT ${RUNTIME_TENSORRT_SRCS})
+ target_link_libraries(tvm_tensorrt_objs PRIVATE tvm_runtime_extra_defs)
+ target_link_libraries(tvm_runtime_extra PRIVATE tvm_tensorrt_objs)
endif()
endif()
diff --git a/docs/arch/codegen.rst b/docs/arch/codegen.rst
index 8190f580a0..5b09d1d0b6 100644
--- a/docs/arch/codegen.rst
+++ b/docs/arch/codegen.rst
@@ -36,7 +36,7 @@ When a user calls ``tvm.compile()``, the compilation proceeds in two phases:
``VMCodeGen`` translates Relax functions into VM bytecode (see :ref:`relax-vm-arch`).
2. **TIR phase**: TIR ``Function``\ s (the actual compute kernels) are compiled to native code.
-The TIR phase is handled internally by ``tirx.build()`` (called from ``relax.build()``).
+The TIR phase is handled internally by ``tirx.compile()`` (called from ``relax.build()``).
It performs these steps:
.. code-block:: text
@@ -46,11 +46,11 @@ It performs these steps:
▼ TIR pipeline ← lowering passes (flatten buffers, lower intrinsics, etc.)
TIR Functions (lowered)
│
- ▼ split_host_device_mods() ← separate host and device functions
+ ▼ _split_host_device_modules() ← separate host and device functions
Host IRModule + Device IRModule(s)
│ │
▼ ▼
- codegen_build() codegen_build() ← target-specific code generation
+ _build_for_target() _build_for_target() ← target-specific code generation
│ │
▼ ▼
Host Module Device Module(s)
@@ -65,7 +65,7 @@ Target Dispatch
---------------
The core dispatch logic lives in ``codegen::Build()`` (``src/target/codegen.cc``), which is
-called from the Python-side ``codegen_build()`` in ``tirx/build.py``. It selects the correct
+called from the Python-side ``_build_for_target()`` in ``tirx/_compile.py``. It selects the correct
backend based on the ``Target`` object:
.. code-block:: cpp
@@ -221,7 +221,7 @@ When compiling for GPU targets, TIR functions are split into two categories:
- **Device functions** — the actual compute kernels that run on the GPU. Compiled with the
target-specific codegen (``target.build.cuda``, etc.).
-``split_host_device_mods()`` (``python/tvm/tirx/build.py``) separates functions by their
+``_split_host_device_modules()`` (``python/tvm/tirx/_compile.py``) separates functions by their
``target`` attribute: functions whose target kind is ``"llvm"`` or ``"c"`` go to the host
module; all others go to device modules grouped by target.
@@ -279,8 +279,8 @@ Source Code Map
* - Path
- Contents
- * - ``python/tvm/tirx/build.py``
- - ``tirx.build()``: TIR compilation entry, host/device split, module linking
+ * - ``python/tvm/tirx/_compile.py``
+ - ``tirx.compile()``: TIR compilation entry, host/device split, module linking
* - ``src/target/codegen.cc``
- ``codegen::Build()``: target dispatch via ``"target.build.<kind>"``
* - ``src/target/llvm/codegen_llvm.h``
diff --git a/docs/arch/relax_vm.rst b/docs/arch/relax_vm.rst
index 26ced6469d..9bf3580786 100644
--- a/docs/arch/relax_vm.rst
+++ b/docs/arch/relax_vm.rst
@@ -47,7 +47,7 @@ The end-to-end flow from model to execution is:
▼ VMCodeGen
ExecBuilder (bytecode) + IRModule (TIR only)
│ │
- │ ▼ tirx.build()
+ │ ▼ tirx.compile()
│ runtime.Module (native kernels)
│ │
▼ VMLink ▼
@@ -87,7 +87,7 @@ Internally, ``relax.build()`` performs these steps:
2. Create an ``ExecBuilder`` and run **VMCodeGen** (``src/relax/backend/vm/codegen_vm.cc``),
which walks each ``relax.Function`` and emits bytecode instructions. The Relax functions are
removed from the IRModule; only TIR functions remain.
-3. Compile the remaining TIR functions to native code via ``tirx.build()``.
+3. Compile the remaining TIR functions to native code via ``tirx.compile()``.
4. **Link** the bytecode executable with the compiled native module using ``VMLink``, producing
a ``VMExecutable``.
diff --git a/docs/deep_dive/tensor_ir/tutorials/tir_transformation.py b/docs/deep_dive/tensor_ir/tutorials/tir_transformation.py
index f221924756..6f53a1029c 100644
--- a/docs/deep_dive/tensor_ir/tutorials/tir_transformation.py
+++ b/docs/deep_dive/tensor_ir/tutorials/tir_transformation.py
@@ -83,7 +83,7 @@ c_nd = tvm.runtime.tensor(np.zeros((128, 128), dtype="float32"))
def evaluate(mod: tvm.IRModule):
- lib = tvm.tirx.build(mod, target="llvm")
+ lib = tvm.tirx.compile(mod, target="llvm").mod
# check correctness
lib(a_nd, b_nd, c_nd)
np.testing.assert_allclose(c_nd.numpy(), c_np, rtol=1e-5)
diff --git a/docs/reference/api/python/support.rst b/docs/reference/api/python/support.rst
index 4663acd2aa..8042ce34c1 100644
--- a/docs/reference/api/python/support.rst
+++ b/docs/reference/api/python/support.rst
@@ -52,9 +52,9 @@ tvm.support.clang
.. automodule:: tvm.support.clang
:members:
-tvm.support.emcc
+tvm.contrib.emcc
~~~~~~~~~~~~~~~~
-.. automodule:: tvm.support.emcc
+.. automodule:: tvm.contrib.emcc
:members:
tvm.support.popen_pool
diff --git a/docs/tirx/api/compilation.rst b/docs/tirx/api/compilation.rst
index 8406c25985..95e180a1c9 100644
--- a/docs/tirx/api/compilation.rst
+++ b/docs/tirx/api/compilation.rst
@@ -18,11 +18,11 @@
Compilation
===========
-tvm.tirx.build
---------------
+tvm.tirx.compile
+----------------
.. currentmodule:: tvm.tirx
-.. autofunction:: build
+.. autofunction:: compile
:no-index:
tvm.tirx.compilation_pipeline
diff --git a/include/tvm/backend/cuda/op/iket.h b/include/tvm/backend/cuda/op/iket.h
index a4d36ac3d8..026f42b5ad 100644
--- a/include/tvm/backend/cuda/op/iket.h
+++ b/include/tvm/backend/cuda/op/iket.h
@@ -48,7 +48,7 @@ namespace cuda {
* Operands: [0] name, the marker name; [1...] args, the optional payload tail.
* Result: void. Effects: opaque marker annotation.
*/
-TVM_DLL const Op& iket_mark_op();
+TVM_TIRX_DLL const Op& iket_mark_op();
/*!
* \brief Get the NVIDIA IKET token-range start operator, tirx.cuda.iket_range_start.
@@ -57,7 +57,7 @@ TVM_DLL const Op& iket_mark_op();
* Operands: [0] name, the range name; [1...] args, the optional payload tail.
* Result: uint32 range token. Effects: opaque range-start annotation.
*/
-TVM_DLL const Op& iket_range_start_op();
+TVM_TIRX_DLL const Op& iket_range_start_op();
/*!
* \brief Get the NVIDIA IKET token-range end operator, tirx.cuda.iket_range_end.
@@ -67,7 +67,7 @@ TVM_DLL const Op& iket_range_start_op();
* [1...] args, the optional payload tail. IKET lowering requires a uint32 token.
* Result: void. Effects: opaque range-end annotation.
*/
-TVM_DLL const Op& iket_range_end_op();
+TVM_TIRX_DLL const Op& iket_range_end_op();
/*!
* \brief Get the NVIDIA IKET stack-range push operator, tirx.cuda.iket_range_push.
@@ -76,7 +76,7 @@ TVM_DLL const Op& iket_range_end_op();
* Operands: [0] name, the range name; [1...] args, the optional payload tail.
* Result: void. Effects: opaque stack-range push annotation.
*/
-TVM_DLL const Op& iket_range_push_op();
+TVM_TIRX_DLL const Op& iket_range_push_op();
/*!
* \brief Get the NVIDIA IKET stack-range pop operator, tirx.cuda.iket_range_pop.
@@ -84,7 +84,7 @@ TVM_DLL const Op& iket_range_push_op();
* Operands: none; IKET lowering enforces the empty operand list.
* Result: void. Effects: opaque stack-range pop annotation.
*/
-TVM_DLL const Op& iket_range_pop_op();
+TVM_TIRX_DLL const Op& iket_range_pop_op();
/*!
* \brief Get the NVIDIA IKET sentinel-token operator, tirx.cuda.iket_sentinel_token.
@@ -93,7 +93,7 @@ TVM_DLL const Op& iket_range_pop_op();
* Result: uint32 sentinel token. Registered effects: opaque. The sentinel carries
* token-flow identity; IKET lowering emits no runtime event for it.
*/
-TVM_DLL const Op& iket_sentinel_token_op();
+TVM_TIRX_DLL const Op& iket_sentinel_token_op();
/*! \} */
diff --git a/include/tvm/backend/cuda/op/tensormap.h b/include/tvm/backend/cuda/op/tensormap.h
index 9da1fa1b9b..cfb06e7014 100644
--- a/include/tvm/backend/cuda/op/tensormap.h
+++ b/include/tvm/backend/cuda/op/tensormap.h
@@ -53,7 +53,7 @@ struct TensorMapEncodeTiledAttr : public AttrsNode {
* Arguments are descriptor and data pointers, global dimensions (rank), byte
* strides (rank - 1), box dimensions (rank), then element strides (rank).
*/
-TVM_DLL const Op& tensormap_encode_tiled_op();
+TVM_TIRX_DLL const Op& tensormap_encode_tiled_op();
} // namespace cuda
} // namespace backend
diff --git a/include/tvm/backend/metal/op/cooperative_tensor.h b/include/tvm/backend/metal/op/cooperative_tensor.h
index 5d83adb328..7ad13a1134 100644
--- a/include/tvm/backend/metal/op/cooperative_tensor.h
+++ b/include/tvm/backend/metal/op/cooperative_tensor.h
@@ -41,7 +41,7 @@ namespace metal {
* - args[3]: rows, The number of rows.
* - args[4]: cols, The number of columns.
*/
-TVM_DLL const Op& cooperative_tensor_fill_op();
+TVM_TIRX_DLL const Op& cooperative_tensor_fill_op();
/*!
* \brief Load data from device or threadgroup memory into a cooperative_tensor.
@@ -59,7 +59,7 @@ TVM_DLL const Op& cooperative_tensor_fill_op();
* - args[9]: mma_K, The K dimension of the matrix operation.
* - args[10]: operand_role, 0 for left (A), 1 for right (B), 2 for destination (C).
*/
-TVM_DLL const Op& cooperative_tensor_load_op();
+TVM_TIRX_DLL const Op& cooperative_tensor_load_op();
/*!
* \brief Store data from a cooperative_tensor to device or threadgroup memory.
@@ -77,7 +77,7 @@ TVM_DLL const Op& cooperative_tensor_load_op();
* - args[9]: mma_K, The K dimension of the matrix operation.
* - args[10]: operand_role, 0 for left (A), 1 for right (B), 2 for destination (C).
*/
-TVM_DLL const Op& cooperative_tensor_store_op();
+TVM_TIRX_DLL const Op& cooperative_tensor_store_op();
/*!
* \brief Multiply and accumulate two matrices using cooperative_tensor
@@ -98,7 +98,7 @@ TVM_DLL const Op& cooperative_tensor_store_op();
* - args[11]: transpose_a, Whether to transpose A.
* - args[12]: transpose_b, Whether to transpose B.
*/
-TVM_DLL const Op& cooperative_tensor_multiply_accumulate_op();
+TVM_TIRX_DLL const Op& cooperative_tensor_multiply_accumulate_op();
} // namespace metal
} // namespace backend
diff --git a/include/tvm/backend/opencl/op/memory.h b/include/tvm/backend/opencl/op/memory.h
index c5d5cf0c33..25ee078f1a 100644
--- a/include/tvm/backend/opencl/op/memory.h
+++ b/include/tvm/backend/opencl/op/memory.h
@@ -41,7 +41,7 @@ namespace opencl {
* - args[4]: channel_size, The number of channels.
* - args[5]: value, The value to use.
*/
-TVM_DLL const Op& texture2d_store_op();
+TVM_TIRX_DLL const Op& texture2d_store_op();
/*!
* \brief Load from texture 2d memory.
@@ -54,7 +54,7 @@ TVM_DLL const Op& texture2d_store_op();
* - args[4]: channel_size, The number of channels.
* - args[5]: element_index, The element index within a texture channel.
*/
-TVM_DLL const Op& texture2d_load_op();
+TVM_TIRX_DLL const Op& texture2d_load_op();
/*!
* \brief Create an Nd memory allocation with storage scope.
@@ -65,7 +65,7 @@ TVM_DLL const Op& texture2d_load_op();
* - args[2]: shape, The shape.
* - args[3...]: args, trailing Expr operands.
*/
-TVM_DLL const Op& nd_mem_alloc_with_scope_op();
+TVM_TIRX_DLL const Op& nd_mem_alloc_with_scope_op();
} // namespace opencl
} // namespace backend
diff --git a/include/tvm/ir/base_expr.h b/include/tvm/ir/base_expr.h
index 421ab3cb46..309eb9abd7 100644
--- a/include/tvm/ir/base_expr.h
+++ b/include/tvm/ir/base_expr.h
@@ -80,7 +80,7 @@ class TypeNode : public ffi::Object {
class Type : public ffi::ObjectRef {
public:
/*! \brief Construct a MissingType for type information not yet populated. */
- TVM_DLL static Type Missing();
+ TVM_TIRX_DLL static Type Missing();
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Type, ffi::ObjectRef, TypeNode);
};
@@ -104,7 +104,7 @@ class MissingTypeNode final : public TypeNode {
/*! \brief Managed reference to the MissingTypeNode singleton. */
class MissingType final : public Type {
public:
- TVM_DLL MissingType();
+ TVM_TIRX_DLL MissingType();
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(MissingType, Type, MissingTypeNode);
};
@@ -128,7 +128,7 @@ class OpaqueTypeNode final : public TypeNode {
/*! \brief Managed reference to OpaqueTypeNode. */
class OpaqueType final : public Type {
public:
- TVM_DLL OpaqueType();
+ TVM_TIRX_DLL OpaqueType();
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(OpaqueType, Type, OpaqueTypeNode);
};
@@ -164,7 +164,7 @@ class PrimType final : public Type {
* \brief Construct from a raw DLPack dtype.
* \param dtype The corresponding DLPack dtype.
*/
- TVM_DLL explicit PrimType(DLDataType dtype);
+ TVM_TIRX_DLL explicit PrimType(DLDataType dtype);
/*!
* \brief Construct from DLPack dtype fields.
@@ -172,27 +172,27 @@ class PrimType final : public Type {
* \param bits The scalar bit width.
* \param lanes The fixed lane count.
*/
- TVM_DLL PrimType(DLDataTypeCode code, int bits, int lanes = 1);
+ TVM_TIRX_DLL PrimType(DLDataTypeCode code, int bits, int lanes = 1);
/*! \brief Construct a signed integer type with fixed lanes. */
- TVM_DLL static PrimType Int(int bits, int lanes = 1);
+ TVM_TIRX_DLL static PrimType Int(int bits, int lanes = 1);
/*! \brief Construct an unsigned integer type with fixed lanes. */
- TVM_DLL static PrimType UInt(int bits, int lanes = 1);
+ TVM_TIRX_DLL static PrimType UInt(int bits, int lanes = 1);
/*! \brief Construct a floating-point type with fixed lanes. */
- TVM_DLL static PrimType Float(int bits, int lanes = 1);
+ TVM_TIRX_DLL static PrimType Float(int bits, int lanes = 1);
/*! \brief Construct a bfloat type with fixed lanes. */
- TVM_DLL static PrimType BFloat(int bits, int lanes = 1);
+ TVM_TIRX_DLL static PrimType BFloat(int bits, int lanes = 1);
/*! \brief Construct a boolean type with fixed lanes. */
- TVM_DLL static PrimType Bool(int lanes = 1);
+ TVM_TIRX_DLL static PrimType Bool(int lanes = 1);
/*! \brief Construct the void sentinel type, encoded as handle(0, 0). */
- TVM_DLL static PrimType Void();
+ TVM_TIRX_DLL static PrimType Void();
/*!
* \brief Construct a scalable vector type.
* \param code The DLPack dtype code.
* \param bits The scalar bit width.
* \param lanes The positive vscale factor to encode in the DLPack lane field.
*/
- TVM_DLL static PrimType ScalableVector(DLDataTypeCode code, int bits, int lanes);
+ TVM_TIRX_DLL static PrimType ScalableVector(DLDataTypeCode code, int bits, int lanes);
/*! \return The DLPack dtype code. */
TVM_FFI_INLINE DLDataTypeCode code() const {
@@ -453,18 +453,18 @@ class PrimExpr : public TypedExpr<PrimType> {
* PrimType.
* \param call The call to view as a primitive expression.
*/
- TVM_DLL PrimExpr(Call call); // NOLINT(*)
+ TVM_TIRX_DLL PrimExpr(Call call); // NOLINT(*)
/*!
* \brief construct from integer.
* \param value The value to be constructed.
*/
- TVM_DLL PrimExpr(int32_t value); // NOLINT(*)
+ TVM_TIRX_DLL PrimExpr(int32_t value); // NOLINT(*)
/*!
* \brief construct from float.
* \param value The value to be constructed.
*/
- TVM_DLL PrimExpr(float value); // NOLINT(*)
+ TVM_TIRX_DLL PrimExpr(float value); // NOLINT(*)
explicit PrimExpr(ffi::ObjectPtr<ExprNode> node) : TypedExpr<PrimType>(std::move(node)) {}
diff --git a/include/tvm/ir/env_func.h b/include/tvm/ir/env_func.h
index 08e4572e9c..861dab0548 100644
--- a/include/tvm/ir/env_func.h
+++ b/include/tvm/ir/env_func.h
@@ -95,7 +95,7 @@ class EnvFunc : public ffi::ObjectRef {
* \return The created global function.
* \note The function can be unique
*/
- TVM_DLL static EnvFunc Get(const ffi::String& name);
+ TVM_TIRX_DLL static EnvFunc Get(const ffi::String& name);
/*! \brief specify container node */
using ContainerType = EnvFuncObj;
};
diff --git a/include/tvm/ir/expr.h b/include/tvm/ir/expr.h
index e9c17fe69c..5d48559bbf 100644
--- a/include/tvm/ir/expr.h
+++ b/include/tvm/ir/expr.h
@@ -76,7 +76,7 @@ class GenericConstNode : public ConstantNode {
/*! \brief Managed reference to a generic constant. */
class GenericConst : public Constant {
public:
- TVM_DLL GenericConst(ffi::Any value, Type ty, Location loc = UnknownLoc());
+ TVM_TIRX_DLL GenericConst(ffi::Any value, Type ty, Location loc = UnknownLoc());
explicit GenericConst(ffi::ObjectPtr<GenericConstNode> node) : Constant(std::move(node)) {}
@@ -98,7 +98,7 @@ class StringImmNode : public ConstantNode {
/*! \brief Managed reference to a string literal. */
class StringImm : public Constant {
public:
- TVM_DLL explicit StringImm(ffi::String value, Location loc = UnknownLoc());
+ TVM_TIRX_DLL explicit StringImm(ffi::String value, Location loc = UnknownLoc());
explicit StringImm(ffi::ObjectPtr<StringImmNode> node) : Constant(std::move(node)) {}
@@ -120,7 +120,7 @@ class DataTypeImmNode : public ConstantNode {
/*! \brief Managed reference to a data type literal. */
class DataTypeImm : public Constant {
public:
- TVM_DLL explicit DataTypeImm(DLDataType value, Location loc = UnknownLoc());
+ TVM_TIRX_DLL explicit DataTypeImm(DLDataType value, Location loc = UnknownLoc());
explicit DataTypeImm(ffi::ObjectPtr<DataTypeImmNode> node) : Constant(std::move(node)) {}
@@ -157,7 +157,7 @@ class IntImm : public PrimExpr {
* \param value The internal value.
* \param loc The location of this object in the source code.
*/
- TVM_DLL IntImm(PrimType value_ty, ffi::BigInt value, Location loc = UnknownLoc());
+ TVM_TIRX_DLL IntImm(PrimType value_ty, ffi::BigInt value, Location loc = UnknownLoc());
template <typename Enum, std::enable_if_t<std::is_enum_v<Enum>, int> = 0>
IntImm(PrimType value_ty, Enum value, Location loc = UnknownLoc())
@@ -237,7 +237,7 @@ class FloatImm : public PrimExpr {
* \param value The internal value.
* \param loc The location in the source code.
*/
- TVM_DLL FloatImm(PrimType value_ty, double value, Location loc = UnknownLoc());
+ TVM_TIRX_DLL FloatImm(PrimType value_ty, double value, Location loc = UnknownLoc());
explicit FloatImm(ffi::ObjectPtr<FloatImmNode> node) : PrimExpr(std::move(node)) {}
@@ -381,7 +381,7 @@ class Tuple : public Expr {
* \param fields The fields of the tuple.
* \param loc The source location of the expression.
*/
- TVM_DLL explicit Tuple(ffi::Array<Expr> fields, Location loc = UnknownLoc());
+ TVM_TIRX_DLL explicit Tuple(ffi::Array<Expr> fields, Location loc = UnknownLoc());
explicit Tuple(ffi::ObjectPtr<TupleNode> node) : Expr(std::move(node)) {}
@@ -431,7 +431,7 @@ class TupleGetItem : public Expr {
* \param index The field index.
* \param loc The source location of the expression.
*/
- TVM_DLL TupleGetItem(Expr tuple, int index, Location loc = UnknownLoc());
+ TVM_TIRX_DLL TupleGetItem(Expr tuple, int index, Location loc = UnknownLoc());
explicit TupleGetItem(ffi::ObjectPtr<TupleGetItemNode> node) : Expr(std::move(node)) {}
@@ -479,7 +479,7 @@ class TensorLoad : public PrimExpr {
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr operator+(PrimExpr a, PrimExpr b);
+TVM_TIRX_DLL PrimExpr operator+(PrimExpr a, PrimExpr b);
/*!
* \brief subtraction operator
@@ -490,7 +490,7 @@ TVM_DLL PrimExpr operator+(PrimExpr a, PrimExpr b);
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr operator-(PrimExpr a, PrimExpr b);
+TVM_TIRX_DLL PrimExpr operator-(PrimExpr a, PrimExpr b);
/*!
* \brief negation.
@@ -500,7 +500,7 @@ TVM_DLL PrimExpr operator-(PrimExpr a, PrimExpr b);
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr operator-(PrimExpr a);
+TVM_TIRX_DLL PrimExpr operator-(PrimExpr a);
/*!
* \brief multiplication operator
@@ -511,7 +511,7 @@ TVM_DLL PrimExpr operator-(PrimExpr a);
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr operator*(PrimExpr a, PrimExpr b);
+TVM_TIRX_DLL PrimExpr operator*(PrimExpr a, PrimExpr b);
/*!
* \brief division operator
@@ -522,7 +522,7 @@ TVM_DLL PrimExpr operator*(PrimExpr a, PrimExpr b);
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr operator/(PrimExpr a, PrimExpr b);
+TVM_TIRX_DLL PrimExpr operator/(PrimExpr a, PrimExpr b);
/*!
* \brief left shift operator
@@ -533,7 +533,7 @@ TVM_DLL PrimExpr operator/(PrimExpr a, PrimExpr b);
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr operator<<(PrimExpr a, PrimExpr b);
+TVM_TIRX_DLL PrimExpr operator<<(PrimExpr a, PrimExpr b);
/*!
* \brief right shift operator
@@ -544,7 +544,7 @@ TVM_DLL PrimExpr operator<<(PrimExpr a, PrimExpr b);
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr operator>>(PrimExpr a, PrimExpr b);
+TVM_TIRX_DLL PrimExpr operator>>(PrimExpr a, PrimExpr b);
/*!
* \brief greater
@@ -555,7 +555,7 @@ TVM_DLL PrimExpr operator>>(PrimExpr a, PrimExpr b);
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr operator>(PrimExpr a, PrimExpr b);
+TVM_TIRX_DLL PrimExpr operator>(PrimExpr a, PrimExpr b);
/*!
* \brief greater_equal
@@ -566,7 +566,7 @@ TVM_DLL PrimExpr operator>(PrimExpr a, PrimExpr b);
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr operator>=(PrimExpr a, PrimExpr b);
+TVM_TIRX_DLL PrimExpr operator>=(PrimExpr a, PrimExpr b);
/*!
* \brief less
@@ -577,7 +577,7 @@ TVM_DLL PrimExpr operator>=(PrimExpr a, PrimExpr b);
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr operator<(PrimExpr a, PrimExpr b);
+TVM_TIRX_DLL PrimExpr operator<(PrimExpr a, PrimExpr b);
/*!
* \brief less_equal
@@ -588,7 +588,7 @@ TVM_DLL PrimExpr operator<(PrimExpr a, PrimExpr b);
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr operator<=(PrimExpr a, PrimExpr b);
+TVM_TIRX_DLL PrimExpr operator<=(PrimExpr a, PrimExpr b);
/*!
* \brief equal
@@ -599,7 +599,7 @@ TVM_DLL PrimExpr operator<=(PrimExpr a, PrimExpr b);
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr operator==(PrimExpr a, PrimExpr b);
+TVM_TIRX_DLL PrimExpr operator==(PrimExpr a, PrimExpr b);
/*!
* \brief not_equal
@@ -610,7 +610,7 @@ TVM_DLL PrimExpr operator==(PrimExpr a, PrimExpr b);
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr operator!=(PrimExpr a, PrimExpr b);
+TVM_TIRX_DLL PrimExpr operator!=(PrimExpr a, PrimExpr b);
/*!
* \brief and
@@ -620,7 +620,7 @@ TVM_DLL PrimExpr operator!=(PrimExpr a, PrimExpr b);
* \return The result expression.
* \note This operator does eager constant folding.
*/
-TVM_DLL PrimExpr operator&&(PrimExpr a, PrimExpr b);
+TVM_TIRX_DLL PrimExpr operator&&(PrimExpr a, PrimExpr b);
/*!
* \brief or
@@ -630,7 +630,7 @@ TVM_DLL PrimExpr operator&&(PrimExpr a, PrimExpr b);
* \return The result expression.
* \note This operator does eager constant folding.
*/
-TVM_DLL PrimExpr operator||(PrimExpr a, PrimExpr b);
+TVM_TIRX_DLL PrimExpr operator||(PrimExpr a, PrimExpr b);
/*!
* \brief not
@@ -639,7 +639,7 @@ TVM_DLL PrimExpr operator||(PrimExpr a, PrimExpr b);
* \return The result expression.
* \note This operator does eager constant folding.
*/
-TVM_DLL PrimExpr operator!(PrimExpr a);
+TVM_TIRX_DLL PrimExpr operator!(PrimExpr a);
/*!
* \brief take bitwise and of two values
@@ -650,7 +650,7 @@ TVM_DLL PrimExpr operator!(PrimExpr a);
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr operator&(PrimExpr a, PrimExpr b);
+TVM_TIRX_DLL PrimExpr operator&(PrimExpr a, PrimExpr b);
/*!
* \brief take bitwise or of two values
@@ -661,7 +661,7 @@ TVM_DLL PrimExpr operator&(PrimExpr a, PrimExpr b);
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr operator|(PrimExpr a, PrimExpr b);
+TVM_TIRX_DLL PrimExpr operator|(PrimExpr a, PrimExpr b);
/*!
* \brief take bitwise xor of two values
@@ -672,7 +672,7 @@ TVM_DLL PrimExpr operator|(PrimExpr a, PrimExpr b);
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr operator^(PrimExpr a, PrimExpr b);
+TVM_TIRX_DLL PrimExpr operator^(PrimExpr a, PrimExpr b);
/*!
* \brief take bitwise negation of two values
@@ -682,7 +682,7 @@ TVM_DLL PrimExpr operator^(PrimExpr a, PrimExpr b);
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr operator~(PrimExpr a);
+TVM_TIRX_DLL PrimExpr operator~(PrimExpr a);
/*!
* \brief A local variable in the IR.
@@ -713,20 +713,20 @@ class VarNode : public ExprNode {
/*! \brief Managed reference to VarNode. */
class Var : public Expr {
public:
- TVM_DLL explicit Var(ffi::String name, ffi::Optional<Type> ty_annotation,
- Location loc = UnknownLoc());
+ TVM_TIRX_DLL explicit Var(ffi::String name, ffi::Optional<Type> ty_annotation,
+ Location loc = UnknownLoc());
/*! \brief Return a fresh ordinary Var with the same type and a new name. */
- TVM_DLL Var CopyWithName(const ffi::String& name) const;
+ TVM_TIRX_DLL Var CopyWithName(const ffi::String& name) const;
/*! \brief Return a fresh ordinary Var with a suffix appended to its name. */
- TVM_DLL Var CopyWithSuffix(const ffi::String& suffix) const;
+ TVM_TIRX_DLL Var CopyWithSuffix(const ffi::String& suffix) const;
/*! \brief Return a fresh ordinary Var with a new primitive type. */
- TVM_DLL Var CopyWithDType(PrimType dtype) const;
+ TVM_TIRX_DLL Var CopyWithDType(PrimType dtype) const;
/*! \brief Return a fresh ordinary Var with a new type, retaining its metadata. */
- TVM_DLL Var CopyWithType(Type type) const;
+ TVM_TIRX_DLL Var CopyWithType(Type type) const;
explicit Var(ffi::ObjectPtr<VarNode> node) : Expr(std::move(node)) {}
@@ -790,7 +790,7 @@ class LambdaExprNode : public StagingExprNode {
explicit LambdaExprNode(ffi::UnsafeInit tag) : body(tag) {}
explicit LambdaExprNode(Expr body) : body(std::move(body)) {}
/*! \brief Simultaneously substitute arguments for the bound parameters. */
- TVM_DLL Expr Apply(const ffi::Array<Expr>& arguments) const;
+ TVM_TIRX_DLL Expr Apply(const ffi::Array<Expr>& arguments) const;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
@@ -805,7 +805,7 @@ class LambdaExprNode : public StagingExprNode {
/*! \brief Managed reference to a typed staging lambda. */
class LambdaExpr : public StagingExpr {
public:
- TVM_DLL explicit LambdaExpr(ffi::Array<Var> vars, Expr body);
+ TVM_TIRX_DLL explicit LambdaExpr(ffi::Array<Var> vars, Expr body);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(LambdaExpr, StagingExpr, LambdaExprNode);
};
@@ -852,7 +852,7 @@ class GlobalVarNode : public ExprNode {
*/
class GlobalVar : public Expr {
public:
- TVM_DLL explicit GlobalVar(ffi::String name_hint, Location loc = UnknownLoc());
+ TVM_TIRX_DLL explicit GlobalVar(ffi::String name_hint, Location loc = UnknownLoc());
explicit GlobalVar(ffi::ObjectPtr<GlobalVarNode> node) : Expr(std::move(node)) {}
@@ -974,19 +974,19 @@ class Call : public Expr {
/*! \brief Construct a Call, inferring only when ret_ty is absent.
* Construction preserves provisional IR; invoke Validate separately when needed.
*/
- TVM_DLL Call(ffi::Optional<Type> ret_ty, Expr op, ffi::Array<Expr> args,
- ffi::Optional<Attrs> attrs = std::nullopt,
- ffi::Array<Type> ty_args = ffi::Array<Type>(), Location loc = UnknownLoc());
+ TVM_TIRX_DLL Call(ffi::Optional<Type> ret_ty, Expr op, ffi::Array<Expr> args,
+ ffi::Optional<Attrs> attrs = std::nullopt,
+ ffi::Array<Type> ty_args = ffi::Array<Type>(), Location loc = UnknownLoc());
/*! \brief Check this Call against its registered operator contract. */
- TVM_DLL void Validate() const;
+ TVM_TIRX_DLL void Validate() const;
/*! \brief Recompute a result type from the Call's current explicit inputs.
*
* This ignores the Call's stored result type and does not mutate the Call.
* Returns Missing when no deduction is available; inference errors propagate.
* This does not invoke the operator's validator.
*/
- TVM_DLL static Type ReinferType(const CallNode* call);
+ TVM_TIRX_DLL static Type ReinferType(const CallNode* call);
explicit Call(ffi::ObjectPtr<CallNode> node) : Expr(std::move(node)) {}
@@ -1030,7 +1030,7 @@ class Range : public ffi::ObjectRef {
* \param end The end of the range.
* \param loc The location of the Range in the source.
*/
- TVM_DLL Range(PrimExpr begin, PrimExpr end, Location loc = UnknownLoc());
+ TVM_TIRX_DLL Range(PrimExpr begin, PrimExpr end, Location loc = UnknownLoc());
/*!
* \brief construct a new range with min and extent
* The corresponding constructor is removed,
@@ -1041,7 +1041,8 @@ class Range : public ffi::ObjectRef {
* \param extent The extent of the range.
* \param loc The location of the Range in the source.
*/
- TVM_DLL static Range FromMinExtent(PrimExpr min, PrimExpr extent, Location loc = UnknownLoc());
+ TVM_TIRX_DLL static Range FromMinExtent(PrimExpr min, PrimExpr extent,
+ Location loc = UnknownLoc());
// declare range.
explicit Range(ffi::UnsafeInit tag, ffi::ObjectPtr<RangeNode> n) : ffi::ObjectRef(tag) {
data_ = std::move(n);
@@ -1077,7 +1078,8 @@ class TensorRegionNode : public ExprNode {
/*! \brief Managed reference to TensorRegionNode. */
class TensorRegion : public Expr {
public:
- TVM_DLL TensorRegion(Expr source, ffi::Array<Range> region, Type ty, Location loc = UnknownLoc());
+ TVM_TIRX_DLL TensorRegion(Expr source, ffi::Array<Range> region, Type ty,
+ Location loc = UnknownLoc());
explicit TensorRegion(ffi::ObjectPtr<TensorRegionNode> node) : Expr(std::move(node)) {}
@@ -1099,7 +1101,7 @@ class TensorRegion : public Expr {
* \param expr The expression to inspect.
* \return The strongest runtime effect of the expression.
*/
-TVM_DLL CallEffectKind SideEffect(const Expr& expr);
+TVM_TIRX_DLL CallEffectKind SideEffect(const Expr& expr);
namespace ffi {
@@ -1108,7 +1110,7 @@ inline constexpr bool use_default_type_traits_v<TensorRegion> = false;
template <>
struct TypeTraits<TensorRegion> : public ObjectRefWithFallbackTraitsBase<TensorRegion, TensorLoad> {
- TVM_DLL static TensorRegion ConvertFallbackValue(TensorLoad value);
+ TVM_TIRX_DLL static TensorRegion ConvertFallbackValue(TensorLoad value);
};
template <>
diff --git a/include/tvm/ir/expr_functor.h b/include/tvm/ir/expr_functor.h
index 64df3d6e76..3fbab04085 100644
--- a/include/tvm/ir/expr_functor.h
+++ b/include/tvm/ir/expr_functor.h
@@ -333,7 +333,7 @@ class ExprFunctor<R(const Expr&, Args...)> {
* \endcode
* Use tvm::ExprVisitor explicitly when dialect visitors are also in scope.
*/
-class TVM_DLL ExprVisitor : public ObjectVisitor {
+class TVM_TIRX_DLL ExprVisitor : public ObjectVisitor {
public:
/*! \brief Construct a visitor with the core expression hooks. */
TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(ExprVisitor, ObjectVisitor)
@@ -428,7 +428,7 @@ class TVM_DLL ExprVisitor : public ObjectVisitor {
* \endcode
* Use tvm::ExprMutator explicitly when dialect mutators are also in scope.
*/
-class TVM_DLL ExprMutator : public ObjectMutator {
+class TVM_TIRX_DLL ExprMutator : public ObjectMutator {
public:
/*! \brief Construct a mutator with the core expression hooks. */
TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(ExprMutator, ObjectMutator)
diff --git a/include/tvm/ir/location.h b/include/tvm/ir/location.h
index c9186286b4..2fb42aeec6 100644
--- a/include/tvm/ir/location.h
+++ b/include/tvm/ir/location.h
@@ -69,7 +69,7 @@ class SourceName : public ffi::ObjectRef {
* \param name Name of the operator.
* \return SourceName valid throughout program lifetime.
*/
- TVM_DLL static SourceName Get(const ffi::String& name);
+ TVM_TIRX_DLL static SourceName Get(const ffi::String& name);
explicit SourceName(ffi::UnsafeInit tag, ffi::ObjectPtr<SourceNameNode> n) : ffi::ObjectRef(tag) {
data_ = std::move(n);
@@ -104,7 +104,7 @@ class UnknownLocNode : public LocationNode {
/*! \brief Reference to the shared immutable unknown location. */
class UnknownLoc : public Location {
public:
- TVM_DLL UnknownLoc();
+ TVM_TIRX_DLL UnknownLoc();
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(UnknownLoc, Location, UnknownLocNode);
};
@@ -141,11 +141,11 @@ class SourceLocNode : public LocationNode {
/*! \brief A concrete source range. */
class SourceLoc : public Location {
public:
- TVM_DLL SourceLoc(SourceName source_name, int start_line, int start_column, int end_line,
- int end_column);
+ TVM_TIRX_DLL SourceLoc(SourceName source_name, int start_line, int start_column, int end_line,
+ int end_column);
/*! \brief Merge two ranges in the same source. */
- TVM_DLL SourceLoc Merge(const SourceLoc& other) const;
+ TVM_TIRX_DLL SourceLoc Merge(const SourceLoc& other) const;
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(SourceLoc, Location, SourceLocNode);
};
@@ -169,7 +169,7 @@ class CallSiteLocNode : public LocationNode {
/*! \brief Explicit inline callee/caller provenance. */
class CallSiteLoc : public Location {
public:
- TVM_DLL CallSiteLoc(Location callee, Location caller);
+ TVM_TIRX_DLL CallSiteLoc(Location callee, Location caller);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(CallSiteLoc, Location, CallSiteLocNode);
};
@@ -205,8 +205,8 @@ class SourceObj : public ffi::Object {
class Source : public ffi::ObjectRef {
public:
- TVM_DLL Source(SourceName src_name, std::string source);
- TVM_DLL tvm::ffi::String GetLine(int line);
+ TVM_TIRX_DLL Source(SourceName src_name, std::string source);
+ TVM_TIRX_DLL tvm::ffi::String GetLine(int line);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Source, ffi::ObjectRef, SourceObj);
};
diff --git a/include/tvm/ir/module.h b/include/tvm/ir/module.h
index a31f5b35c6..afd7a101d4 100644
--- a/include/tvm/ir/module.h
+++ b/include/tvm/ir/module.h
@@ -145,10 +145,10 @@ class IRModuleNode : public ffi::Object {
.def(tvm::ffi::reflection::type_attr::kSHash, &IRModuleNode::SHash);
}
- TVM_DLL bool SEqual(const IRModuleNode* other,
- ffi::TypedFunction<bool(AnyView, AnyView, bool, AnyView)> equal) const;
- TVM_DLL int64_t SHash(int64_t init_hash,
- ffi::TypedFunction<int64_t(AnyView, int64_t, bool)> hash) const;
+ TVM_TIRX_DLL bool SEqual(const IRModuleNode* other,
+ ffi::TypedFunction<bool(AnyView, AnyView, bool, AnyView)> equal) const;
+ TVM_TIRX_DLL int64_t SHash(int64_t init_hash,
+ ffi::TypedFunction<int64_t(AnyView, int64_t, bool)> hash) const;
/*!
* \brief Add a function to the global environment.
@@ -157,7 +157,7 @@ class IRModuleNode : public ffi::Object {
* \param update Controls whether you can replace a definition in the
* environment.
*/
- TVM_DLL void Add(const GlobalVar& var, const BaseFunc& func, bool update = false);
+ TVM_TIRX_DLL void Add(const GlobalVar& var, const BaseFunc& func, bool update = false);
/*!
* \brief Add a function to the global environment.
@@ -166,79 +166,79 @@ class IRModuleNode : public ffi::Object {
*
* It does not do type inference as Add does.
*/
- TVM_DLL void AddUnchecked(const GlobalVar& var, const BaseFunc& func);
+ TVM_TIRX_DLL void AddUnchecked(const GlobalVar& var, const BaseFunc& func);
/*!
* \brief Update a function in the global environment.
* \param var The name of the global function to update.
* \param func The new function.
*/
- TVM_DLL void Update(const GlobalVar& var, const BaseFunc& func);
+ TVM_TIRX_DLL void Update(const GlobalVar& var, const BaseFunc& func);
/*!
* \brief Update an array of global infos in the global environment.
* \param name The name of the global info.
* \param info The new array of global infos.
*/
- TVM_DLL void UpdateGlobalInfo(const ffi::String& name, const ffi::Array<GlobalInfo>& info);
+ TVM_TIRX_DLL void UpdateGlobalInfo(const ffi::String& name, const ffi::Array<GlobalInfo>& info);
/*!
* \brief Remove a function from the global environment.
* \param var The name of the global function to update.
*/
- TVM_DLL void Remove(const GlobalVar& var);
+ TVM_TIRX_DLL void Remove(const GlobalVar& var);
/*!
* \brief Check if the global_var_map_ contains a global variable.
* \param name The variable name.
* \returns true if contains, otherise false.
*/
- TVM_DLL bool ContainGlobalVar(const ffi::String& name) const;
+ TVM_TIRX_DLL bool ContainGlobalVar(const ffi::String& name) const;
/*!
* \brief Lookup a global function by its variable.
* \param str The unique string specifying the global variable.
* \returns The global variable.
*/
- TVM_DLL GlobalVar GetGlobalVar(const ffi::String& str) const;
+ TVM_TIRX_DLL GlobalVar GetGlobalVar(const ffi::String& str) const;
/*!
* \brief Collect all global vars defined in this module, ordered by
* the global variable name.
* \returns An array of global vars
*/
- TVM_DLL ffi::Array<GlobalVar> GetGlobalVars() const;
+ TVM_TIRX_DLL ffi::Array<GlobalVar> GetGlobalVars() const;
/*!
* \brief Look up a global function by its variable.
* \param var The global var to lookup.
* \returns The function named by the variable argument.
*/
- TVM_DLL BaseFunc Lookup(const GlobalVar& var) const;
+ TVM_TIRX_DLL BaseFunc Lookup(const GlobalVar& var) const;
/*!
* \brief Look up a global function by its string name
* \param name The name of the function.
* \returns The function named by the argument.
*/
- TVM_DLL BaseFunc Lookup(const ffi::String& name) const;
+ TVM_TIRX_DLL BaseFunc Lookup(const ffi::String& name) const;
/*!
* \brief Update the functions inside this environment by
* functions in another environment.
* \param other The other environment.
*/
- TVM_DLL void Update(const IRModule& other);
+ TVM_TIRX_DLL void Update(const IRModule& other);
/*!
* \brief Create a shallow copy of this IRModule.
* \returns The shallow copy of the IRModule.
*/
- TVM_DLL IRModule ShallowCopy();
+ TVM_TIRX_DLL IRModule ShallowCopy();
/*!
* \brief The set of imported files.
*/
- TVM_DLL std::unordered_set<ffi::String> Imports() const;
+ TVM_TIRX_DLL std::unordered_set<ffi::String> Imports() const;
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
@@ -262,9 +262,9 @@ class IRModule : public ffi::ObjectRef {
* \param attrs The module meta-data attributes.
* \param global_infos Global infos in the module.
*/
- TVM_DLL explicit IRModule(ffi::Map<GlobalVar, BaseFunc> functions, SourceMap map = {},
- DictAttrs attrs = DictAttrs(),
- ffi::Map<ffi::String, ffi::Array<GlobalInfo>> global_infos = {});
+ TVM_TIRX_DLL explicit IRModule(ffi::Map<GlobalVar, BaseFunc> functions, SourceMap map = {},
+ DictAttrs attrs = DictAttrs(),
+ ffi::Map<ffi::String, ffi::Array<GlobalInfo>> global_infos = {});
/*! \brief default constructor */
IRModule() : IRModule(ffi::Map<GlobalVar, BaseFunc>({})) {}
@@ -287,8 +287,8 @@ class IRModule : public ffi::ObjectRef {
* \brief As for \p FromExprInContext, but assuming \p expr is bound to 'main' and no
* imports.
*/
- TVM_DLL static IRModule FromExpr(const Expr& expr,
- const ffi::Map<GlobalVar, BaseFunc>& global_funcs = {});
+ TVM_TIRX_DLL static IRModule FromExpr(const Expr& expr,
+ const ffi::Map<GlobalVar, BaseFunc>& global_funcs = {});
/*!
* \brief Create a shallow copy of an IRModule.
diff --git a/include/tvm/ir/object_functor.h b/include/tvm/ir/object_functor.h
index 5c2493307d..23ce91b396 100644
--- a/include/tvm/ir/object_functor.h
+++ b/include/tvm/ir/object_functor.h
@@ -225,7 +225,7 @@ using ffi::VisitInterrupt;
* interrupts and errors own their values. A hook controls descent into its node.
* One visitor must not be shared by overlapping traversals.
*/
-class TVM_DLL ObjectVisitor : public ffi::StructuralVisitorObj {
+class TVM_TIRX_DLL ObjectVisitor : public ffi::StructuralVisitorObj {
public:
/*! \brief Construct a visitor using structural fallback for every value. */
ObjectVisitor() : ObjectVisitor(GlobalVTable()) {}
@@ -343,7 +343,7 @@ class TVM_DLL ObjectVisitor : public ffi::StructuralVisitorObj {
* replacements and errors own their values. Overrides forward inplace_mode
* to child calls, following the structural mutation ownership contract.
*/
-class TVM_DLL ObjectMutator : public ffi::StructuralMapEngineBase {
+class TVM_TIRX_DLL ObjectMutator : public ffi::StructuralMapEngineBase {
public:
/*! \brief Construct a mutator using structural fallback for every object. */
ObjectMutator() : ObjectMutator(GlobalVTable()) {}
diff --git a/include/tvm/ir/op.h b/include/tvm/ir/op.h
index 13cf14beae..7d2ca57c45 100644
--- a/include/tvm/ir/op.h
+++ b/include/tvm/ir/op.h
@@ -110,7 +110,7 @@ class OpNode : public ExprNode {
/*! \brief Whether the validator was explicitly registered rather than generated. */
bool validator_is_custom{false};
- TVM_DLL static void RegisterReflection();
+ TVM_TIRX_DLL static void RegisterReflection();
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindUniqueInstance;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.Op", OpNode, ExprNode);
@@ -165,16 +165,16 @@ class Op : public Expr {
* \param attr_name The column name.
* \return Whether the column has been registered.
*/
- TVM_DLL static bool HasAttrMap(const ffi::String& attr_name);
+ TVM_TIRX_DLL static bool HasAttrMap(const ffi::String& attr_name);
/*!
* \brief Look up a registered operator by name.
* \param op_name The canonical name.
* \return A copy of the reference-counted handle to the canonical node.
* \throws AttributeError if the operator is not registered.
*/
- TVM_DLL static Op Get(const ffi::String& op_name);
+ TVM_TIRX_DLL static Op Get(const ffi::String& op_name);
/*! \brief List registered operator names. \return Names in unspecified order. */
- TVM_DLL static ffi::Array<ffi::String> ListNames();
+ TVM_TIRX_DLL static ffi::Array<ffi::String> ListNames();
/*! \brief Validate a Call with this Op's callback, when one is registered. */
TVM_FFI_INLINE void Validate(const CallNode* call) const {
@@ -193,8 +193,8 @@ class Op : public Expr {
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Op, Expr, OpNode);
private:
- TVM_DLL static void ThrowInvalidCall(const OpNode* op);
- TVM_DLL static ffi::List<ffi::Any> GetAttrColumn(const ffi::String& attr_name);
+ TVM_TIRX_DLL static void ThrowInvalidCall(const OpNode* op);
+ TVM_TIRX_DLL static ffi::List<ffi::Any> GetAttrColumn(const ffi::String& attr_name);
};
namespace sig {
@@ -443,13 +443,13 @@ class OpDef {
* \brief Get or create the named Op without changing its documentation.
* \param name The canonical operator name.
*/
- TVM_DLL explicit OpDef(const ffi::String& name);
+ TVM_TIRX_DLL explicit OpDef(const ffi::String& name);
/*!
* \brief Get or create the named Op and set its documentation.
* \param name The canonical operator name.
* \param doc Operator documentation, replacing any existing documentation.
*/
- TVM_DLL OpDef(const ffi::String& name, const ffi::String& doc);
+ TVM_TIRX_DLL OpDef(const ffi::String& name, const ffi::String& doc);
OpDef(const OpDef&) = delete;
OpDef& operator=(const OpDef&) = delete;
OpDef(OpDef&&) = delete;
@@ -462,7 +462,7 @@ class OpDef {
* \param doc Argument documentation.
* \return This builder. Without a variadic descriptor, this is part of the complete list.
*/
- TVM_DLL OpDef& add_arg(const ffi::String& name, const ffi::String& doc);
+ TVM_TIRX_DLL OpDef& add_arg(const ffi::String& name, const ffi::String& doc);
/*!
* \brief Set the attribute object type key and runtime index together.
* \tparam AttrsType The attribute object node type.
@@ -480,14 +480,14 @@ class OpDef {
* \param key A registered attribute object type key.
* \return This builder. An unknown key raises an error before modifying metadata.
*/
- TVM_DLL OpDef& set_attrs_type_key(const ffi::String& key);
+ TVM_TIRX_DLL OpDef& set_attrs_type_key(const ffi::String& key);
/*!
* \brief Append a descriptor for a type argument in Call.ty_args.
* \param name Type-argument name.
* \param doc Type-argument documentation.
* \return This builder. Inference retains responsibility for type-argument counts.
*/
- TVM_DLL OpDef& add_ty_arg(const ffi::String& name, const ffi::String& doc);
+ TVM_TIRX_DLL OpDef& add_ty_arg(const ffi::String& name, const ffi::String& doc);
/*!
* \brief Register a complete typed Call signature and its validator.
*
@@ -573,7 +573,7 @@ class OpDef {
* \param attr_name Attribute column name; missing columns and values are ignored.
* \return This builder. Existing cached maps observe the removal.
*/
- TVM_DLL OpDef& reset_attr(const ffi::String& attr_name);
+ TVM_TIRX_DLL OpDef& reset_attr(const ffi::String& attr_name);
private:
template <typename Spec>
@@ -629,7 +629,7 @@ class OpDef {
}
OpNode* get() { return const_cast<OpNode*>(op_.operator->()); }
- TVM_DLL void UpdateAttr(const ffi::String& attr_name, ffi::Any value, bool override);
+ TVM_TIRX_DLL void UpdateAttr(const ffi::String& attr_name, ffi::Any value, bool override);
Op op_;
};
diff --git a/include/tvm/ir/prim/expr.h b/include/tvm/ir/prim/expr.h
index 206ec438eb..61cd3e0e8a 100644
--- a/include/tvm/ir/prim/expr.h
+++ b/include/tvm/ir/prim/expr.h
@@ -73,7 +73,7 @@ class CastNode : public ExprNode {
*/
class Cast : public PrimExpr {
public:
- TVM_DLL Cast(PrimType value_ty, PrimExpr value, Location loc = UnknownLoc());
+ TVM_TIRX_DLL Cast(PrimType value_ty, PrimExpr value, Location loc = UnknownLoc());
explicit Cast(ffi::ObjectPtr<CastNode> node) : PrimExpr(std::move(node)) {}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Cast, PrimExpr, CastNode);
@@ -117,7 +117,7 @@ class AddNode : public BinaryOpNode<AddNode> {
*/
class Add : public PrimExpr {
public:
- TVM_DLL Add(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+ TVM_TIRX_DLL Add(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
explicit Add(ffi::ObjectPtr<AddNode> node) : PrimExpr(std::move(node)) {}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Add, PrimExpr, AddNode);
@@ -138,7 +138,7 @@ class LShiftNode : public BinaryOpNode<LShiftNode> {
*/
class LShift : public PrimExpr {
public:
- TVM_DLL LShift(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+ TVM_TIRX_DLL LShift(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
explicit LShift(ffi::ObjectPtr<LShiftNode> node) : PrimExpr(std::move(node)) {}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(LShift, PrimExpr, LShiftNode);
@@ -159,7 +159,7 @@ class RShiftNode : public BinaryOpNode<RShiftNode> {
*/
class RShift : public PrimExpr {
public:
- TVM_DLL RShift(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+ TVM_TIRX_DLL RShift(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
explicit RShift(ffi::ObjectPtr<RShiftNode> node) : PrimExpr(std::move(node)) {}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(RShift, PrimExpr, RShiftNode);
@@ -180,7 +180,7 @@ class BitwiseAndNode : public BinaryOpNode<BitwiseAndNode> {
*/
class BitwiseAnd : public PrimExpr {
public:
- TVM_DLL BitwiseAnd(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+ TVM_TIRX_DLL BitwiseAnd(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
explicit BitwiseAnd(ffi::ObjectPtr<BitwiseAndNode> node) : PrimExpr(std::move(node)) {}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(BitwiseAnd, PrimExpr, BitwiseAndNode);
@@ -201,7 +201,7 @@ class BitwiseOrNode : public BinaryOpNode<BitwiseOrNode> {
*/
class BitwiseOr : public PrimExpr {
public:
- TVM_DLL BitwiseOr(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+ TVM_TIRX_DLL BitwiseOr(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
explicit BitwiseOr(ffi::ObjectPtr<BitwiseOrNode> node) : PrimExpr(std::move(node)) {}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(BitwiseOr, PrimExpr, BitwiseOrNode);
@@ -222,7 +222,7 @@ class BitwiseXorNode : public BinaryOpNode<BitwiseXorNode> {
*/
class BitwiseXor : public PrimExpr {
public:
- TVM_DLL BitwiseXor(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+ TVM_TIRX_DLL BitwiseXor(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
explicit BitwiseXor(ffi::ObjectPtr<BitwiseXorNode> node) : PrimExpr(std::move(node)) {}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(BitwiseXor, PrimExpr, BitwiseXorNode);
@@ -243,7 +243,7 @@ class SubNode : public BinaryOpNode<SubNode> {
*/
class Sub : public PrimExpr {
public:
- TVM_DLL Sub(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+ TVM_TIRX_DLL Sub(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
explicit Sub(ffi::ObjectPtr<SubNode> node) : PrimExpr(std::move(node)) {}
@@ -265,7 +265,7 @@ class MulNode : public BinaryOpNode<MulNode> {
*/
class Mul : public PrimExpr {
public:
- TVM_DLL Mul(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+ TVM_TIRX_DLL Mul(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
explicit Mul(ffi::ObjectPtr<MulNode> node) : PrimExpr(std::move(node)) {}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Mul, PrimExpr, MulNode);
@@ -289,7 +289,7 @@ class DivNode : public BinaryOpNode<DivNode> {
*/
class Div : public PrimExpr {
public:
- TVM_DLL Div(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+ TVM_TIRX_DLL Div(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
explicit Div(ffi::ObjectPtr<DivNode> node) : PrimExpr(std::move(node)) {}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Div, PrimExpr, DivNode);
@@ -313,7 +313,7 @@ class ModNode : public BinaryOpNode<ModNode> {
*/
class Mod : public PrimExpr {
public:
- TVM_DLL Mod(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+ TVM_TIRX_DLL Mod(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
explicit Mod(ffi::ObjectPtr<ModNode> node) : PrimExpr(std::move(node)) {}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Mod, PrimExpr, ModNode);
@@ -334,7 +334,7 @@ class FloorDivNode : public BinaryOpNode<FloorDivNode> {
*/
class FloorDiv : public PrimExpr {
public:
- TVM_DLL FloorDiv(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+ TVM_TIRX_DLL FloorDiv(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
explicit FloorDiv(ffi::ObjectPtr<FloorDivNode> node) : PrimExpr(std::move(node)) {}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(FloorDiv, PrimExpr, FloorDivNode);
@@ -355,7 +355,7 @@ class FloorModNode : public BinaryOpNode<FloorModNode> {
*/
class FloorMod : public PrimExpr {
public:
- TVM_DLL FloorMod(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+ TVM_TIRX_DLL FloorMod(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
explicit FloorMod(ffi::ObjectPtr<FloorModNode> node) : PrimExpr(std::move(node)) {}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(FloorMod, PrimExpr, FloorModNode);
@@ -376,7 +376,7 @@ class MinNode : public BinaryOpNode<MinNode> {
*/
class Min : public PrimExpr {
public:
- TVM_DLL Min(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+ TVM_TIRX_DLL Min(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
explicit Min(ffi::ObjectPtr<MinNode> node) : PrimExpr(std::move(node)) {}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Min, PrimExpr, MinNode);
@@ -397,7 +397,7 @@ class MaxNode : public BinaryOpNode<MaxNode> {
*/
class Max : public PrimExpr {
public:
- TVM_DLL Max(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+ TVM_TIRX_DLL Max(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
explicit Max(ffi::ObjectPtr<MaxNode> node) : PrimExpr(std::move(node)) {}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Max, PrimExpr, MaxNode);
@@ -441,7 +441,7 @@ class EQNode : public CmpOpNode<EQNode> {
*/
class EQ : public PrimExpr {
public:
- TVM_DLL EQ(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+ TVM_TIRX_DLL EQ(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
explicit EQ(ffi::ObjectPtr<EQNode> node) : PrimExpr(std::move(node)) {}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(EQ, PrimExpr, EQNode);
@@ -462,7 +462,7 @@ class NENode : public CmpOpNode<NENode> {
*/
class NE : public PrimExpr {
public:
- TVM_DLL NE(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+ TVM_TIRX_DLL NE(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
explicit NE(ffi::ObjectPtr<NENode> node) : PrimExpr(std::move(node)) {}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(NE, PrimExpr, NENode);
@@ -483,7 +483,7 @@ class LTNode : public CmpOpNode<LTNode> {
*/
class LT : public PrimExpr {
public:
- TVM_DLL LT(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+ TVM_TIRX_DLL LT(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
explicit LT(ffi::ObjectPtr<LTNode> node) : PrimExpr(std::move(node)) {}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(LT, PrimExpr, LTNode);
@@ -504,7 +504,7 @@ struct LENode : public CmpOpNode<LENode> {
*/
class LE : public PrimExpr {
public:
- TVM_DLL LE(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+ TVM_TIRX_DLL LE(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
explicit LE(ffi::ObjectPtr<LENode> node) : PrimExpr(std::move(node)) {}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(LE, PrimExpr, LENode);
@@ -525,7 +525,7 @@ class GTNode : public CmpOpNode<GTNode> {
*/
class GT : public PrimExpr {
public:
- TVM_DLL GT(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+ TVM_TIRX_DLL GT(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
explicit GT(ffi::ObjectPtr<GTNode> node) : PrimExpr(std::move(node)) {}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(GT, PrimExpr, GTNode);
@@ -546,7 +546,7 @@ class GENode : public CmpOpNode<GENode> {
*/
class GE : public PrimExpr {
public:
- TVM_DLL GE(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+ TVM_TIRX_DLL GE(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
explicit GE(ffi::ObjectPtr<GENode> node) : PrimExpr(std::move(node)) {}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(GE, PrimExpr, GENode);
@@ -577,7 +577,7 @@ class AndNode : public ExprNode {
*/
class And : public PrimExpr {
public:
- TVM_DLL And(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+ TVM_TIRX_DLL And(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
explicit And(ffi::ObjectPtr<AndNode> node) : PrimExpr(std::move(node)) {}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(And, PrimExpr, AndNode);
@@ -608,7 +608,7 @@ class OrNode : public ExprNode {
*/
class Or : public PrimExpr {
public:
- TVM_DLL Or(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+ TVM_TIRX_DLL Or(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
explicit Or(ffi::ObjectPtr<OrNode> node) : PrimExpr(std::move(node)) {}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Or, PrimExpr, OrNode);
@@ -637,7 +637,7 @@ class NotNode : public ExprNode {
*/
class Not : public PrimExpr {
public:
- TVM_DLL Not(PrimExpr a, Location loc = UnknownLoc());
+ TVM_TIRX_DLL Not(PrimExpr a, Location loc = UnknownLoc());
explicit Not(ffi::ObjectPtr<NotNode> node) : PrimExpr(std::move(node)) {}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Not, PrimExpr, NotNode);
@@ -666,7 +666,7 @@ class BitwiseNotNode : public ExprNode {
*/
class BitwiseNot : public PrimExpr {
public:
- TVM_DLL BitwiseNot(PrimExpr a, Location loc = UnknownLoc());
+ TVM_TIRX_DLL BitwiseNot(PrimExpr a, Location loc = UnknownLoc());
explicit BitwiseNot(ffi::ObjectPtr<BitwiseNotNode> node) : PrimExpr(std::move(node)) {}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(BitwiseNot, PrimExpr, BitwiseNotNode);
@@ -711,8 +711,8 @@ class SelectNode : public ExprNode {
*/
class Select : public PrimExpr {
public:
- TVM_DLL Select(PrimExpr condition, PrimExpr true_value, PrimExpr false_value,
- Location loc = UnknownLoc());
+ TVM_TIRX_DLL Select(PrimExpr condition, PrimExpr true_value, PrimExpr false_value,
+ Location loc = UnknownLoc());
explicit Select(ffi::ObjectPtr<SelectNode> node) : PrimExpr(std::move(node)) {}
@@ -752,7 +752,7 @@ class LetNode : public ExprNode {
*/
class Let : public PrimExpr {
public:
- TVM_DLL Let(Var var, PrimExpr value, PrimExpr body, Location loc = UnknownLoc());
+ TVM_TIRX_DLL Let(Var var, PrimExpr value, PrimExpr body, Location loc = UnknownLoc());
explicit Let(ffi::ObjectPtr<LetNode> node) : PrimExpr(std::move(node)) {}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Let, PrimExpr, LetNode);
@@ -794,7 +794,7 @@ inline std::unordered_map<K, V> as_unordered_map(const ffi::Map<K, V>& dmap) {
*/
struct ExprDeepEqual {
public:
- TVM_DLL bool operator()(const PrimExpr& lhs, const PrimExpr& rhs) const;
+ TVM_TIRX_DLL bool operator()(const PrimExpr& lhs, const PrimExpr& rhs) const;
};
} // namespace prim
diff --git a/include/tvm/ir/prim/op.h b/include/tvm/ir/prim/op.h
index 8fdf390cd6..c79c4334dc 100644
--- a/include/tvm/ir/prim/op.h
+++ b/include/tvm/ir/prim/op.h
@@ -36,16 +36,16 @@ namespace prim {
* \brief Get the target's vscale value. It will be lowered to llvm.vscale intrinsic
* (https://llvm.org/docs/LangRef.html#llvm-vscale-intrinsic)
*/
-TVM_DLL const Op& vscale_op();
+TVM_TIRX_DLL const Op& vscale_op();
/*! \brief Round up to the nearest integral value. */
-TVM_DLL const Op& ceil_op();
+TVM_TIRX_DLL const Op& ceil_op();
/*! \brief Base-two logarithm. */
-TVM_DLL const Op& log2_op();
+TVM_TIRX_DLL const Op& log2_op();
/*! \brief Count leading zero bits. */
-TVM_DLL const Op& clz_op();
+TVM_TIRX_DLL const Op& clz_op();
/*!
* \brief Same as select, used for unsafe memory access.
@@ -54,10 +54,10 @@ TVM_DLL const Op& clz_op();
* return cond ? a : b;
* }
*/
-TVM_DLL const Op& if_then_else_op();
+TVM_TIRX_DLL const Op& if_then_else_op();
/*! \brief Marks a condition is likely going to happen. */
-TVM_DLL const Op& likely_op();
+TVM_TIRX_DLL const Op& likely_op();
} // namespace prim
@@ -71,7 +71,7 @@ TVM_DLL const Op& likely_op();
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr add(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr add(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
/*!
* \brief subtraction operator
@@ -83,7 +83,7 @@ TVM_DLL PrimExpr add(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr sub(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr sub(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
/*!
* \brief negation.
@@ -94,7 +94,7 @@ TVM_DLL PrimExpr sub(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr neg(PrimExpr a, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr neg(PrimExpr a, Location loc = UnknownLoc());
/*!
* \brief multiplication operator
@@ -106,7 +106,7 @@ TVM_DLL PrimExpr neg(PrimExpr a, Location loc = UnknownLoc());
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr mul(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr mul(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
/*!
* \brief left shift operator
@@ -118,7 +118,7 @@ TVM_DLL PrimExpr mul(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr left_shift(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr left_shift(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
/*!
* \brief right shift operator
@@ -130,7 +130,7 @@ TVM_DLL PrimExpr left_shift(PrimExpr a, PrimExpr b, Location loc = UnknownLoc())
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr right_shift(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr right_shift(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
/*!
* \brief greater
@@ -142,7 +142,7 @@ TVM_DLL PrimExpr right_shift(PrimExpr a, PrimExpr b, Location loc = UnknownLoc()
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr greater(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr greater(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
/*!
* \brief greater_equal
@@ -154,7 +154,7 @@ TVM_DLL PrimExpr greater(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr greater_equal(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr greater_equal(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
/*!
* \brief less
@@ -166,7 +166,7 @@ TVM_DLL PrimExpr greater_equal(PrimExpr a, PrimExpr b, Location loc = UnknownLoc
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr less(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr less(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
/*!
* \brief less_equal
@@ -178,7 +178,7 @@ TVM_DLL PrimExpr less(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr less_equal(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr less_equal(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
/*!
* \brief equal
@@ -190,7 +190,7 @@ TVM_DLL PrimExpr less_equal(PrimExpr a, PrimExpr b, Location loc = UnknownLoc())
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr equal(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr equal(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
/*!
* \brief not_equal
@@ -202,7 +202,7 @@ TVM_DLL PrimExpr equal(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr not_equal(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr not_equal(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
/*!
* \brief and
@@ -213,7 +213,7 @@ TVM_DLL PrimExpr not_equal(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
* \return The result expression.
* \note This operator does eager constant folding.
*/
-TVM_DLL PrimExpr logical_and(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr logical_and(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
/*!
* \brief or
@@ -224,7 +224,7 @@ TVM_DLL PrimExpr logical_and(PrimExpr a, PrimExpr b, Location loc = UnknownLoc()
* \return The result expression.
* \note This operator does eager constant folding.
*/
-TVM_DLL PrimExpr logical_or(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr logical_or(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
/*!
* \brief not
@@ -234,7 +234,7 @@ TVM_DLL PrimExpr logical_or(PrimExpr a, PrimExpr b, Location loc = UnknownLoc())
* \return The result expression.
* \note This operator does eager constant folding.
*/
-TVM_DLL PrimExpr logical_not(PrimExpr a, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr logical_not(PrimExpr a, Location loc = UnknownLoc());
/*!
* \brief compute division in C semantics.
@@ -250,7 +250,7 @@ TVM_DLL PrimExpr logical_not(PrimExpr a, Location loc = UnknownLoc());
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr div(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr div(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
/*!
* \brief compute trunc(a / b)
@@ -264,7 +264,7 @@ TVM_DLL PrimExpr div(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr truncdiv(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr truncdiv(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
/*!
* \brief compute the remainder of truncdiv
@@ -278,7 +278,7 @@ TVM_DLL PrimExpr truncdiv(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr truncmod(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr truncmod(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
/*!
* \brief compute floor(a / b) where a and b are non-negative.
@@ -295,7 +295,7 @@ TVM_DLL PrimExpr truncmod(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr indexdiv(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr indexdiv(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
/*!
* \brief compute ceil(a / b) where a and b are non-negative.
@@ -312,7 +312,7 @@ TVM_DLL PrimExpr indexdiv(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
* \note this function does eager constant folding for
* shape types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr shapediv(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr shapediv(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
/*!
* \brief compute the remainder floor(a / b) where a and b are non-negative.
@@ -328,7 +328,7 @@ TVM_DLL PrimExpr shapediv(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr indexmod(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr indexmod(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
/*!
* \brief compute floor(a / b)
@@ -340,7 +340,7 @@ TVM_DLL PrimExpr indexmod(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr floordiv(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr floordiv(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
/*!
* \brief compute ceil(a / b)
@@ -353,7 +353,7 @@ TVM_DLL PrimExpr floordiv(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr ceildiv(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr ceildiv(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
/*!
* \brief compute the remainder of floordiv
@@ -365,7 +365,7 @@ TVM_DLL PrimExpr ceildiv(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr floormod(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr floormod(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
/*!
* \brief take maximum of two values
@@ -377,7 +377,7 @@ TVM_DLL PrimExpr floormod(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr max(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr max(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
/*!
* \brief take minimum of two values
@@ -389,7 +389,7 @@ TVM_DLL PrimExpr max(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr min(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr min(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
/*!
* \brief take bitwise and of two values
@@ -401,7 +401,7 @@ TVM_DLL PrimExpr min(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr bitwise_and(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr bitwise_and(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
/*!
* \brief take bitwise or of two values
@@ -413,7 +413,7 @@ TVM_DLL PrimExpr bitwise_and(PrimExpr a, PrimExpr b, Location loc = UnknownLoc()
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr bitwise_or(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr bitwise_or(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
/*!
* \brief take bitwise xor of two values
@@ -425,7 +425,7 @@ TVM_DLL PrimExpr bitwise_or(PrimExpr a, PrimExpr b, Location loc = UnknownLoc())
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr bitwise_xor(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr bitwise_xor(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
/*!
* \brief take bitwise negation of two values
@@ -436,7 +436,7 @@ TVM_DLL PrimExpr bitwise_xor(PrimExpr a, PrimExpr b, Location loc = UnknownLoc()
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr bitwise_neg(PrimExpr a, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr bitwise_neg(PrimExpr a, Location loc = UnknownLoc());
/*!
* \brief Conditional expression.
@@ -449,8 +449,8 @@ TVM_DLL PrimExpr bitwise_neg(PrimExpr a, Location loc = UnknownLoc());
* \note this function does eager constant folding for
* index types(int32, int64) when possible.
*/
-TVM_DLL PrimExpr if_then_else(PrimExpr cond, PrimExpr true_value, PrimExpr false_value,
- Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr if_then_else(PrimExpr cond, PrimExpr true_value, PrimExpr false_value,
+ Location loc = UnknownLoc());
/*!
* \brief Mark condition as likely.
@@ -458,21 +458,21 @@ TVM_DLL PrimExpr if_then_else(PrimExpr cond, PrimExpr true_value, PrimExpr false
* \param loc The location of this operation in the source.
* \return The marked expression.
*/
-TVM_DLL PrimExpr likely(PrimExpr cond, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr likely(PrimExpr cond, Location loc = UnknownLoc());
/*! \brief Round up to the nearest integral value, preserving the input type. */
-TVM_DLL PrimExpr ceil(PrimExpr x, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr ceil(PrimExpr x, Location loc = UnknownLoc());
/*! \brief Construct a base-two logarithm with the input's primitive type. */
-TVM_DLL PrimExpr log2(PrimExpr x, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr log2(PrimExpr x, Location loc = UnknownLoc());
namespace prim {
/*! \brief Count leading zero bits, preserving the input primitive type. */
-TVM_DLL PrimExpr clz(PrimExpr x, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr clz(PrimExpr x, Location loc = UnknownLoc());
/*!
* Query the minimum possible value of dtype.
* \param dtype The primitive type.
* \param loc The location of this operation in the source.
* \return the minimum possible value in this format.
*/
-TVM_DLL PrimExpr min_value(PrimType dtype, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr min_value(PrimType dtype, Location loc = UnknownLoc());
/*!
* Query the maximum possible value of dtype.
@@ -480,7 +480,7 @@ TVM_DLL PrimExpr min_value(PrimType dtype, Location loc = UnknownLoc());
* \param loc The location of this operation in the source.
* \return the maximum possible value in this format.
*/
-TVM_DLL PrimExpr max_value(PrimType dtype, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr max_value(PrimType dtype, Location loc = UnknownLoc());
/*!
* \brief cast value to type.
@@ -492,10 +492,10 @@ TVM_DLL PrimExpr max_value(PrimType dtype, Location loc = UnknownLoc());
* \note This function may return value if the type is the same. Constant folding
* uses MakeConst for scalar and vector constants.
*/
-TVM_DLL PrimExpr cast(PrimType t, PrimExpr value, Location loc = UnknownLoc());
-TVM_DLL PrimExpr cast(DLDataType t, PrimExpr value, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr cast(PrimType t, PrimExpr value, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr cast(DLDataType t, PrimExpr value, Location loc = UnknownLoc());
/*! \brief Construct integer absolute value; floating absolute value belongs to TIRX. */
-TVM_DLL PrimExpr IntegerAbs(PrimExpr x, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr IntegerAbs(PrimExpr x, Location loc = UnknownLoc());
/*!
* \brief Make a const value with certain data type.
*
@@ -752,7 +752,7 @@ inline bool IsConstNumber(const PrimExpr& x);
* \param shift The output shift if x is power of two.
* \return whether x is constant power of two
*/
-TVM_DLL bool IsPowerOfTwoInt(const PrimExpr& x, int* shift);
+TVM_TIRX_DLL bool IsPowerOfTwoInt(const PrimExpr& x, int* shift);
/*! \brief Check whether a is a positive integer constant. */
inline bool IsPositiveConst(const PrimExpr& a);
@@ -799,7 +799,7 @@ inline bool IsZero(const PrimExpr& x) { return IsConstInt(x, 0); }
* \param loc The location of this operation in the source.
* \return the infinity value in this format.
*/
-TVM_DLL PrimExpr infinity(PrimType dtype, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr infinity(PrimType dtype, Location loc = UnknownLoc());
/*!
* \brief Calculate power(x, y)
@@ -807,7 +807,7 @@ TVM_DLL PrimExpr infinity(PrimType dtype, Location loc = UnknownLoc());
* \param y The right operand.
* \param loc The location of this operation in the source.
*/
-TVM_DLL PrimExpr pow(PrimExpr x, PrimExpr y, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr pow(PrimExpr x, PrimExpr y, Location loc = UnknownLoc());
/*!
* \brief Calculate absolute value of x.
@@ -816,7 +816,7 @@ TVM_DLL PrimExpr pow(PrimExpr x, PrimExpr y, Location loc = UnknownLoc());
*
* \return The absolute value of input data x
*/
-TVM_DLL PrimExpr abs(PrimExpr x, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr abs(PrimExpr x, Location loc = UnknownLoc());
/*!
* \brief Check if x is NaN.
@@ -824,7 +824,7 @@ TVM_DLL PrimExpr abs(PrimExpr x, Location loc = UnknownLoc());
* \param loc The location of this operation in the source.
* \return The result expression.
*/
-TVM_DLL PrimExpr isnan(PrimExpr x, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr isnan(PrimExpr x, Location loc = UnknownLoc());
/*!
* \brief Check if x is finite.
@@ -832,7 +832,7 @@ TVM_DLL PrimExpr isnan(PrimExpr x, Location loc = UnknownLoc());
* \param loc The location of this operation in the source.
* \return The result expression.
*/
-TVM_DLL PrimExpr isfinite(PrimExpr x, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr isfinite(PrimExpr x, Location loc = UnknownLoc());
/*!
* \brief Check if x is infinite.
@@ -840,7 +840,7 @@ TVM_DLL PrimExpr isfinite(PrimExpr x, Location loc = UnknownLoc());
* \param loc The location of this operation in the source.
* \return The result expression.
*/
-TVM_DLL PrimExpr isinf(PrimExpr x, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr isinf(PrimExpr x, Location loc = UnknownLoc());
/*!
* \brief Calculate floor(x)
@@ -848,7 +848,7 @@ TVM_DLL PrimExpr isinf(PrimExpr x, Location loc = UnknownLoc());
* \param loc The location of this operation in the source.
* \return The result expression.
*/
-TVM_DLL PrimExpr floor(PrimExpr x, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr floor(PrimExpr x, Location loc = UnknownLoc());
/*!
* \brief Round x to the nearest integer, ties to even.
@@ -860,7 +860,7 @@ TVM_DLL PrimExpr floor(PrimExpr x, Location loc = UnknownLoc());
* \param loc The location of this operation in the source.
* \return The result expression.
*/
-TVM_DLL PrimExpr round(PrimExpr x, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr round(PrimExpr x, Location loc = UnknownLoc());
/*!
* \brief Round x to the nearest integer, ties to even.
@@ -871,7 +871,7 @@ TVM_DLL PrimExpr round(PrimExpr x, Location loc = UnknownLoc());
* \param loc The location of this operation in the source.
* \return The result expression.
*/
-TVM_DLL PrimExpr nearbyint(PrimExpr x, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr nearbyint(PrimExpr x, Location loc = UnknownLoc());
/*!
* \brief Calculate trunc(x)
@@ -879,88 +879,88 @@ TVM_DLL PrimExpr nearbyint(PrimExpr x, Location loc = UnknownLoc());
* \param loc The location of this operation in the source.
* \return The result expression.
*/
-TVM_DLL PrimExpr trunc(PrimExpr x, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr trunc(PrimExpr x, Location loc = UnknownLoc());
/*! \brief Floating-point remainder of x divided by y. */
-TVM_DLL PrimExpr fmod(PrimExpr x, PrimExpr y, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr fmod(PrimExpr x, PrimExpr y, Location loc = UnknownLoc());
/*! \brief Raise x to the power y. Arguments: x, y. */
-TVM_DLL const Op& pow_op();
+TVM_TIRX_DLL const Op& pow_op();
/*! \brief Absolute value. Argument: x. */
-TVM_DLL const Op& fabs_op();
+TVM_TIRX_DLL const Op& fabs_op();
/*! \brief Floating-point remainder. Arguments: x, y. */
-TVM_DLL const Op& fmod_op();
+TVM_TIRX_DLL const Op& fmod_op();
/*! \brief Round toward negative infinity. Argument: x. */
-TVM_DLL const Op& floor_op();
+TVM_TIRX_DLL const Op& floor_op();
/*! \brief Round to nearest integer, ties to even. Argument: x. */
-TVM_DLL const Op& round_op();
+TVM_TIRX_DLL const Op& round_op();
/*! \brief Round to nearest integer, ties to even. Argument: x. */
-TVM_DLL const Op& nearbyint_op();
+TVM_TIRX_DLL const Op& nearbyint_op();
/*! \brief Round toward zero. Argument: x. */
-TVM_DLL const Op& trunc_op();
+TVM_TIRX_DLL const Op& trunc_op();
/*! \brief Natural exponential. Argument: x. */
-TVM_DLL const Op& exp_op();
+TVM_TIRX_DLL const Op& exp_op();
/*! \brief Base-two exponential. Argument: x. */
-TVM_DLL const Op& exp2_op();
+TVM_TIRX_DLL const Op& exp2_op();
/*! \brief Base-ten exponential. Argument: x. */
-TVM_DLL const Op& exp10_op();
+TVM_TIRX_DLL const Op& exp10_op();
/*! \brief Error function. Argument: x. */
-TVM_DLL const Op& erf_op();
+TVM_TIRX_DLL const Op& erf_op();
/*! \brief Hyperbolic tangent. Argument: x. */
-TVM_DLL const Op& tanh_op();
+TVM_TIRX_DLL const Op& tanh_op();
/*! \brief Logistic sigmoid. Argument: x. */
-TVM_DLL const Op& sigmoid_op();
+TVM_TIRX_DLL const Op& sigmoid_op();
/*! \brief Square root. Argument: x. */
-TVM_DLL const Op& sqrt_op();
+TVM_TIRX_DLL const Op& sqrt_op();
/*! \brief Reciprocal square root. Argument: x. */
-TVM_DLL const Op& rsqrt_op();
+TVM_TIRX_DLL const Op& rsqrt_op();
/*! \brief Natural logarithm. Argument: x. */
-TVM_DLL const Op& log_op();
+TVM_TIRX_DLL const Op& log_op();
/*! \brief Natural logarithm of one plus x. Argument: x. */
-TVM_DLL const Op& log1p_op();
+TVM_TIRX_DLL const Op& log1p_op();
/*! \brief Base-ten logarithm. Argument: x. */
-TVM_DLL const Op& log10_op();
+TVM_TIRX_DLL const Op& log10_op();
/*! \brief Tangent. Argument: x. */
-TVM_DLL const Op& tan_op();
+TVM_TIRX_DLL const Op& tan_op();
/*! \brief Cosine. Argument: x. */
-TVM_DLL const Op& cos_op();
+TVM_TIRX_DLL const Op& cos_op();
/*! \brief Hyperbolic cosine. Argument: x. */
-TVM_DLL const Op& cosh_op();
+TVM_TIRX_DLL const Op& cosh_op();
/*! \brief Sine. Argument: x. */
-TVM_DLL const Op& sin_op();
+TVM_TIRX_DLL const Op& sin_op();
/*! \brief Hyperbolic sine. Argument: x. */
-TVM_DLL const Op& sinh_op();
+TVM_TIRX_DLL const Op& sinh_op();
/*! \brief Inverse sine. Argument: x. */
-TVM_DLL const Op& asin_op();
+TVM_TIRX_DLL const Op& asin_op();
/*! \brief Inverse cosine. Argument: x. */
-TVM_DLL const Op& acos_op();
+TVM_TIRX_DLL const Op& acos_op();
/*! \brief Inverse tangent. Argument: x. */
-TVM_DLL const Op& atan_op();
+TVM_TIRX_DLL const Op& atan_op();
/*! \brief Inverse hyperbolic cosine. Argument: x. */
-TVM_DLL const Op& acosh_op();
+TVM_TIRX_DLL const Op& acosh_op();
/*! \brief Inverse hyperbolic sine. Argument: x. */
-TVM_DLL const Op& asinh_op();
+TVM_TIRX_DLL const Op& asinh_op();
/*! \brief Inverse hyperbolic tangent. Argument: x. */
-TVM_DLL const Op& atanh_op();
+TVM_TIRX_DLL const Op& atanh_op();
/*! \brief Inverse tangent of y/x with quadrant selection. Arguments: y, x. */
-TVM_DLL const Op& atan2_op();
+TVM_TIRX_DLL const Op& atan2_op();
/*! \brief Next representable value from x toward y. Arguments: x, y. */
-TVM_DLL const Op& nextafter_op();
+TVM_TIRX_DLL const Op& nextafter_op();
/*! \brief Euclidean length of x and y. Arguments: x, y. */
-TVM_DLL const Op& hypot_op();
+TVM_TIRX_DLL const Op& hypot_op();
/*! \brief Magnitude of x with the sign of y. Arguments: x, y. */
-TVM_DLL const Op& copysign_op();
+TVM_TIRX_DLL const Op& copysign_op();
/*! \brief Multiply x by two raised to exponent. Arguments: x, exponent. */
-TVM_DLL const Op& ldexp_op();
+TVM_TIRX_DLL const Op& ldexp_op();
/*! \brief Check for NaN, preserving the lane count. Argument: x. */
-TVM_DLL const Op& isnan_op();
+TVM_TIRX_DLL const Op& isnan_op();
/*! \brief Count set bits. Argument: x. */
-TVM_DLL const Op& popcount_op();
+TVM_TIRX_DLL const Op& popcount_op();
/*! \brief Fused multiply-add, x * y + z. Arguments: x, y, z. */
-TVM_DLL const Op& fma_op();
+TVM_TIRX_DLL const Op& fma_op();
/*! \brief Record a known condition for simplification. Argument: condition. */
-TVM_DLL const Op& assume_op();
+TVM_TIRX_DLL const Op& assume_op();
/*! \brief Record a known condition for compile-time simplification. */
-TVM_DLL PrimExpr assume(PrimExpr condition, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr assume(PrimExpr condition, Location loc = UnknownLoc());
/*! \brief Fused multiply-add of x, y and z, in that order. */
inline PrimExpr fma(PrimExpr x, PrimExpr y, PrimExpr z, Location loc = UnknownLoc()) {
diff --git a/include/tvm/ir/prim/vector_expr.h b/include/tvm/ir/prim/vector_expr.h
index a11acac402..940be97f2e 100644
--- a/include/tvm/ir/prim/vector_expr.h
+++ b/include/tvm/ir/prim/vector_expr.h
@@ -67,7 +67,7 @@ class RampNode : public ExprNode {
*/
class Ramp : public PrimExpr {
public:
- TVM_DLL Ramp(PrimExpr base, PrimExpr stride, PrimExpr lanes, Location loc = UnknownLoc());
+ TVM_TIRX_DLL Ramp(PrimExpr base, PrimExpr stride, PrimExpr lanes, Location loc = UnknownLoc());
explicit Ramp(ffi::ObjectPtr<RampNode> node) : PrimExpr(std::move(node)) {}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Ramp, PrimExpr, RampNode);
@@ -101,7 +101,7 @@ class BroadcastNode : public ExprNode {
*/
class Broadcast : public PrimExpr {
public:
- TVM_DLL Broadcast(PrimExpr value, PrimExpr lanes, Location loc = UnknownLoc());
+ TVM_TIRX_DLL Broadcast(PrimExpr value, PrimExpr lanes, Location loc = UnknownLoc());
explicit Broadcast(ffi::ObjectPtr<BroadcastNode> node) : PrimExpr(std::move(node)) {}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Broadcast, PrimExpr, BroadcastNode);
@@ -135,10 +135,11 @@ class ShuffleNode : public ExprNode {
*/
class Shuffle : public PrimExpr {
public:
- TVM_DLL Shuffle(ffi::Array<PrimExpr> vectors, ffi::Array<PrimExpr> indices,
- Location loc = UnknownLoc());
- TVM_DLL static PrimExpr Concat(ffi::Array<PrimExpr> vectors, Location loc = UnknownLoc());
- TVM_DLL static PrimExpr ExtractElement(PrimExpr vector, int index, Location loc = UnknownLoc());
+ TVM_TIRX_DLL Shuffle(ffi::Array<PrimExpr> vectors, ffi::Array<PrimExpr> indices,
+ Location loc = UnknownLoc());
+ TVM_TIRX_DLL static PrimExpr Concat(ffi::Array<PrimExpr> vectors, Location loc = UnknownLoc());
+ TVM_TIRX_DLL static PrimExpr ExtractElement(PrimExpr vector, int index,
+ Location loc = UnknownLoc());
explicit Shuffle(ffi::ObjectPtr<ShuffleNode> node) : PrimExpr(std::move(node)) {}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Shuffle, PrimExpr, ShuffleNode);
diff --git a/include/tvm/ir/stmt.h b/include/tvm/ir/stmt.h
index f1c9e0d1de..7a9e8e1455 100644
--- a/include/tvm/ir/stmt.h
+++ b/include/tvm/ir/stmt.h
@@ -84,7 +84,7 @@ class SeqStmtNode : public StmtNode {
*/
Stmt operator[](size_t index) const { return seq[index]; }
- TVM_DLL static void RegisterReflection();
+ TVM_TIRX_DLL static void RegisterReflection();
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.SeqStmt", SeqStmtNode, StmtNode);
};
@@ -96,9 +96,9 @@ class SeqStmt : public Stmt {
* \param seq The sequence.
* \param loc The location of this object in the source code.
*/
- TVM_DLL explicit SeqStmt(ffi::Array<Stmt> seq, Location loc = UnknownLoc());
+ TVM_TIRX_DLL explicit SeqStmt(ffi::Array<Stmt> seq, Location loc = UnknownLoc());
/*! \brief Wrap a statement, reusing an existing sequence when possible. */
- TVM_DLL SeqStmt(Stmt stmt, ffi::Optional<Location> loc = std::nullopt);
+ TVM_TIRX_DLL SeqStmt(Stmt stmt, ffi::Optional<Location> loc = std::nullopt);
SeqStmt(std::initializer_list<Stmt> seq, Location loc = UnknownLoc())
: SeqStmt(ffi::Array<Stmt>(seq), std::move(loc)) {}
@@ -158,7 +158,7 @@ class BindNode : public StmtNode {
*/
class Bind : public Stmt {
public:
- TVM_DLL Bind(Var var, Expr value, Location loc = UnknownLoc());
+ TVM_TIRX_DLL Bind(Var var, Expr value, Location loc = UnknownLoc());
explicit Bind(ffi::ObjectPtr<BindNode> node) : Stmt(std::move(node)) {}
@@ -169,7 +169,7 @@ class Bind : public Stmt {
/*!
* \brief Check FRegionGetBodyParams registration without invoking it or creating variables.
*/
-TVM_DLL bool IsRegionOp(const Op& op);
+TVM_TIRX_DLL bool IsRegionOp(const Op& op);
/*!
* \brief Validate region operands and construct fresh typed body parameters.
@@ -179,7 +179,7 @@ TVM_DLL bool IsRegionOp(const Op& op);
* \return Parameters from the required FRegionGetBodyParams hook.
* \throws ValueError if the operation has no registered region hook.
*/
-TVM_DLL ffi::Array<Var> GetRegionBodyParams(Op op, ffi::Array<Expr> args, DictAttrs attrs);
+TVM_TIRX_DLL ffi::Array<Var> GetRegionBodyParams(Op op, ffi::Array<Expr> args, DictAttrs attrs);
/*!
* \brief A single-body statement whose semantics are defined by an operator.
@@ -223,8 +223,9 @@ class RegionStmtNode : public StmtNode {
/*! \brief Managed reference to RegionStmtNode. */
class RegionStmt : public Stmt {
public:
- TVM_DLL RegionStmt(Op op, ffi::Array<Expr> args, ffi::Array<Var> body_params, DictAttrs attrs,
- SeqStmt body, ffi::Array<Var> result_vars = {}, Location loc = UnknownLoc());
+ TVM_TIRX_DLL RegionStmt(Op op, ffi::Array<Expr> args, ffi::Array<Var> body_params,
+ DictAttrs attrs, SeqStmt body, ffi::Array<Var> result_vars = {},
+ Location loc = UnknownLoc());
explicit RegionStmt(ffi::ObjectPtr<RegionStmtNode> node) : Stmt(std::move(node)) {}
@@ -286,8 +287,8 @@ class AssertStmtNode : public StmtNode {
*/
class AssertStmt : public Stmt {
public:
- TVM_DLL AssertStmt(PrimExpr condition, StringImm error_kind, ffi::Array<StringImm> message_parts,
- Location loc = UnknownLoc());
+ TVM_TIRX_DLL AssertStmt(PrimExpr condition, StringImm error_kind,
+ ffi::Array<StringImm> message_parts, Location loc = UnknownLoc());
explicit AssertStmt(ffi::ObjectPtr<AssertStmtNode> node) : Stmt(std::move(node)) {}
@@ -333,8 +334,8 @@ class TensorStoreNode : public StmtNode {
*/
class TensorStore : public Stmt {
public:
- TVM_DLL explicit TensorStore(Expr dest, ffi::Array<PrimExpr> indices, PrimExpr value,
- Location loc = UnknownLoc());
+ TVM_TIRX_DLL explicit TensorStore(Expr dest, ffi::Array<PrimExpr> indices, PrimExpr value,
+ Location loc = UnknownLoc());
explicit TensorStore(ffi::ObjectPtr<TensorStoreNode> node) : Stmt(std::move(node)) {}
@@ -370,7 +371,7 @@ class EvaluateNode : public StmtNode {
*/
class Evaluate : public Stmt {
public:
- TVM_DLL explicit Evaluate(Expr value, Location loc = UnknownLoc());
+ TVM_TIRX_DLL explicit Evaluate(Expr value, Location loc = UnknownLoc());
explicit Evaluate(int value, Location loc = UnknownLoc()) : Evaluate(PrimExpr(value), loc) {}
@@ -413,8 +414,8 @@ class IfNode : public StmtNode {
*/
class If : public Stmt {
public:
- TVM_DLL If(PrimExpr condition, SeqStmt then_case, ffi::Optional<SeqStmt> else_case = std::nullopt,
- Location loc = UnknownLoc());
+ TVM_TIRX_DLL If(PrimExpr condition, SeqStmt then_case,
+ ffi::Optional<SeqStmt> else_case = std::nullopt, Location loc = UnknownLoc());
explicit If(ffi::ObjectPtr<IfNode> node) : Stmt(std::move(node)) {}
@@ -509,9 +510,9 @@ class ForNode : public StmtNode {
*/
class For : public Stmt {
public:
- TVM_DLL For(PrimVar loop_var, PrimExpr min, PrimExpr extent, ForKind kind, SeqStmt body,
- ffi::Map<ffi::String, ffi::Any> annotations = {},
- ffi::Optional<PrimExpr> step = std::nullopt, Location loc = UnknownLoc());
+ TVM_TIRX_DLL For(PrimVar loop_var, PrimExpr min, PrimExpr extent, ForKind kind, SeqStmt body,
+ ffi::Map<ffi::String, ffi::Any> annotations = {},
+ ffi::Optional<PrimExpr> step = std::nullopt, Location loc = UnknownLoc());
explicit For(ffi::ObjectPtr<ForNode> node) : Stmt(std::move(node)) {}
@@ -556,7 +557,7 @@ class WhileNode : public StmtNode {
*/
class While : public Stmt {
public:
- TVM_DLL While(PrimExpr condition, SeqStmt body, Location loc = UnknownLoc());
+ TVM_TIRX_DLL While(PrimExpr condition, SeqStmt body, Location loc = UnknownLoc());
explicit While(ffi::ObjectPtr<WhileNode> node) : Stmt(std::move(node)) {}
@@ -590,7 +591,7 @@ class ReturnNode : public StmtNode {
*/
class Return : public Stmt {
public:
- TVM_DLL explicit Return(Expr value, Location loc = UnknownLoc());
+ TVM_TIRX_DLL explicit Return(Expr value, Location loc = UnknownLoc());
explicit Return(ffi::ObjectPtr<ReturnNode> node) : Stmt(std::move(node)) {}
@@ -617,7 +618,7 @@ class BreakNode : public StmtNode {
*/
class Break : public Stmt {
public:
- TVM_DLL explicit Break(Location loc);
+ TVM_TIRX_DLL explicit Break(Location loc);
explicit Break(ffi::ObjectPtr<BreakNode> node) : Stmt(std::move(node)) {}
@@ -644,7 +645,7 @@ class ContinueNode : public StmtNode {
*/
class Continue : public Stmt {
public:
- TVM_DLL explicit Continue(Location loc);
+ TVM_TIRX_DLL explicit Continue(Location loc);
explicit Continue(ffi::ObjectPtr<ContinueNode> node) : Stmt(std::move(node)) {}
@@ -653,7 +654,7 @@ class Continue : public Stmt {
};
// overload printing of for type.
-TVM_DLL std::ostream& operator<<(std::ostream& os, ForKind kind);
+TVM_TIRX_DLL std::ostream& operator<<(std::ostream& os, ForKind kind);
// inline implementations
inline const char* ForKind2String(ForKind t) {
diff --git a/include/tvm/ir/transform.h b/include/tvm/ir/transform.h
index a5c5590814..7c7dfd1d76 100644
--- a/include/tvm/ir/transform.h
+++ b/include/tvm/ir/transform.h
@@ -85,32 +85,32 @@ class PassContext : public ffi::ObjectRef {
* \brief Construct a PassContext containing the default configurations.
* \return The new PassContext.
*/
- TVM_DLL static PassContext Create();
+ TVM_TIRX_DLL static PassContext Create();
/*!
* \brief Get the default pass context in the current scope.
* \return The pass context.
*/
- TVM_DLL static PassContext Current();
+ TVM_TIRX_DLL static PassContext Current();
/*!
* \brief Get all supported configuration names and metadata, registered within the PassContext.
* \return Map indexed by the config name, pointing to the metadata map as key-value
*/
- TVM_DLL static ffi::Map<ffi::String, ffi::Map<ffi::String, ffi::String>> ListConfigs();
+ TVM_TIRX_DLL static ffi::Map<ffi::String, ffi::Map<ffi::String, ffi::String>> ListConfigs();
/*!
* \brief Call instrument implementations' callbacks when entering PassContext.
* The callbacks are called in order, and if one raises an exception, the rest will not be
* called.
*/
- TVM_DLL void InstrumentEnterPassContext();
+ TVM_TIRX_DLL void InstrumentEnterPassContext();
/*!
* \brief Call instrument implementations' callbacks when exiting PassContext.
* The callbacks are called in order, and if one raises an exception, the rest will not be
* called.
*/
- TVM_DLL void InstrumentExitPassContext();
+ TVM_TIRX_DLL void InstrumentExitPassContext();
/*!
* \brief Call instrument implementations' callbacks before a pass run.
@@ -122,7 +122,7 @@ class PassContext : public ffi::ObjectRef {
*
* \return false: the pass is skipped; true: the pass runs.
*/
- TVM_DLL bool InstrumentBeforePass(const IRModule& mod, const PassInfo& info) const;
+ TVM_TIRX_DLL bool InstrumentBeforePass(const IRModule& mod, const PassInfo& info) const;
/*!
* \brief Call instrument implementations callbacks after a pass run.
@@ -132,14 +132,14 @@ class PassContext : public ffi::ObjectRef {
* \param mod The module that an optimization pass runs on.
* \param info The pass information.
*/
- TVM_DLL void InstrumentAfterPass(const IRModule& mod, const PassInfo& info) const;
+ TVM_TIRX_DLL void InstrumentAfterPass(const IRModule& mod, const PassInfo& info) const;
/*!
* \brief Check whether a pass is enabled.
* \param info The pass information.
* \return true if the pass is enabled. Otherwise, false.
*/
- TVM_DLL bool PassEnabled(const PassInfo& info) const;
+ TVM_TIRX_DLL bool PassEnabled(const PassInfo& info) const;
/*!
* \brief Register a valid configuration option and its ValueType for validation.
@@ -190,12 +190,12 @@ class PassContext : public ffi::ObjectRef {
private:
// The entry of a pass context scope.
- TVM_DLL void EnterWithScope();
+ TVM_TIRX_DLL void EnterWithScope();
// The exit of a pass context scope.
- TVM_DLL void ExitWithScope();
+ TVM_TIRX_DLL void ExitWithScope();
// Register configuration key value type.
- TVM_DLL static void RegisterConfigOption(const char* key, ffi::String value_type_str,
- std::function<ffi::Any(ffi::Any)> legalization);
+ TVM_TIRX_DLL static void RegisterConfigOption(const char* key, ffi::String value_type_str,
+ std::function<ffi::Any(ffi::Any)> legalization);
// Classes to get the Python `with` like syntax.
friend class Internal;
@@ -255,7 +255,7 @@ class PassInfo : public ffi::ObjectRef {
* \param opt_level The optimization level
* \param name Name of the pass.
*/
- TVM_DLL PassInfo(int opt_level, ffi::String name);
+ TVM_TIRX_DLL PassInfo(int opt_level, ffi::String name);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PassInfo, ffi::ObjectRef, PassInfoObj);
};
@@ -384,7 +384,7 @@ class Sequential : public Pass {
* \param passes The passes to apply.
* \param pass_info The pass metadata.
*/
- TVM_DLL Sequential(ffi::Array<Pass> passes, PassInfo pass_info);
+ TVM_TIRX_DLL Sequential(ffi::Array<Pass> passes, PassInfo pass_info);
/*!
* \brief The constructor of `Sequential`.
@@ -394,7 +394,7 @@ class Sequential : public Pass {
* This allows users to only provide a list of passes and execute them
* under a given context.
*/
- TVM_DLL Sequential(ffi::Array<Pass> passes, ffi::String name = "sequential");
+ TVM_TIRX_DLL Sequential(ffi::Array<Pass> passes, ffi::String name = "sequential");
explicit Sequential(ffi::UnsafeInit tag) : Pass(tag) {}
@@ -411,15 +411,15 @@ class Sequential : public Pass {
*
* \return The created module pass.
*/
-TVM_DLL Pass CreateModulePass(std::function<IRModule(IRModule, PassContext)> pass_func,
- int opt_level, ffi::String name);
+TVM_TIRX_DLL Pass CreateModulePass(std::function<IRModule(IRModule, PassContext)> pass_func,
+ int opt_level, ffi::String name);
/*!
* \brief A special trace pass that prints the header and IR to LOG(INFO).
* \param header The header to be attached to the output.
* \return The pass.
*/
-TVM_DLL Pass PrintIR(ffi::String header = "");
+TVM_TIRX_DLL Pass PrintIR(ffi::String header = "");
/*!
* \brief Enrich a pass-time error with a TVMScript-rendered, underlined source
@@ -442,7 +442,7 @@ TVM_DLL Pass PrintIR(ffi::String header = "");
* \p mod->functions[func]; otherwise use the whole module.
* \return The enriched (or, on any fallback, the original) error.
*/
-TVM_DLL ffi::Error EnrichPassErrorWithContext(
+TVM_TIRX_DLL ffi::Error EnrichPassErrorWithContext(
const ffi::Error& err, const IRModule& mod, ffi::String pass_name,
ffi::Optional<GlobalVar> func = ffi::Optional<GlobalVar>(std::nullopt));
diff --git a/include/tvm/ir/type.h b/include/tvm/ir/type.h
index fccc044bad..5b4694cb37 100644
--- a/include/tvm/ir/type.h
+++ b/include/tvm/ir/type.h
@@ -52,7 +52,7 @@ class AnyTypeNode : public TypeNode {
*/
class AnyType : public Type {
public:
- TVM_DLL AnyType(Location loc = UnknownLoc());
+ TVM_TIRX_DLL AnyType(Location loc = UnknownLoc());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(AnyType, Type, AnyTypeNode);
};
@@ -71,7 +71,7 @@ class TensorRegionTypeNode : public TypeNode {
/*! \brief Managed reference to TensorRegionTypeNode. */
class TensorRegionType : public Type {
public:
- TVM_DLL TensorRegionType();
+ TVM_TIRX_DLL TensorRegionType();
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TensorRegionType, Type, TensorRegionTypeNode);
};
@@ -86,7 +86,7 @@ class StringTypeNode : public TypeNode {
/*! \brief Managed reference to StringTypeNode. */
class StringType : public Type {
public:
- TVM_DLL StringType();
+ TVM_TIRX_DLL StringType();
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(StringType, Type, StringTypeNode);
};
@@ -131,10 +131,10 @@ class PtrType : public Type {
* \param element_type The type of the element which the pointer points to.
* \param storage_scope The storage scope into which the pointer addresses
*/
- TVM_DLL explicit PtrType(Type element_type, ffi::String storage_scope = "");
+ TVM_TIRX_DLL explicit PtrType(Type element_type, ffi::String storage_scope = "");
/*! \brief Construct an opaque pointer with void element type. */
- TVM_DLL static PtrType VoidPtrType(ffi::String storage_scope = "");
+ TVM_TIRX_DLL static PtrType VoidPtrType(ffi::String storage_scope = "");
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PtrType, Type, PtrTypeNode);
};
@@ -168,13 +168,13 @@ class TupleType : public Type {
* \param fields Fields in the tuple.
* \param loc The loc of the type.
*/
- TVM_DLL explicit TupleType(ffi::Array<Type> fields, Location loc = UnknownLoc());
+ TVM_TIRX_DLL explicit TupleType(ffi::Array<Type> fields, Location loc = UnknownLoc());
/*!
* \brief Create an empty tuple type that constains nothing.
* \return A empty tuple type.
*/
- TVM_DLL TupleType static Empty();
+ TVM_TIRX_DLL TupleType static Empty();
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TupleType, Type, TupleTypeNode);
};
@@ -230,7 +230,7 @@ class FuncType : public Type {
* \param loc The loc information.
* \sa FuncTypeNode for more docs about these fields.
*/
- TVM_DLL FuncType(ffi::Array<Type> arg_types, Type ret_type, Location loc = UnknownLoc());
+ TVM_TIRX_DLL FuncType(ffi::Array<Type> arg_types, Type ret_type, Location loc = UnknownLoc());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(FuncType, Type, FuncTypeNode);
};
diff --git a/include/tvm/ir/unique_name_supply.h b/include/tvm/ir/unique_name_supply.h
index ec760bdc81..9435fc87e4 100644
--- a/include/tvm/ir/unique_name_supply.h
+++ b/include/tvm/ir/unique_name_supply.h
@@ -119,8 +119,8 @@ class UniqueNameSupply : public ffi::ObjectRef {
* \param prefix The prefix to be used with this UniqueNameSupply.
* \param name_map An optional map.
*/
- TVM_DLL explicit UniqueNameSupply(const ffi::String& prefix = "",
- ffi::Map<ffi::String, int64_t> name_map = {});
+ TVM_TIRX_DLL explicit UniqueNameSupply(const ffi::String& prefix = "",
+ ffi::Map<ffi::String, int64_t> name_map = {});
/*!
* \brief Construct UniqueNameSupply by reserving names from the given iterator range.
@@ -128,7 +128,7 @@ class UniqueNameSupply : public ffi::ObjectRef {
* The functor should return the name of the dereferenced object.
*/
template <typename Iter, typename Lambda>
- TVM_DLL UniqueNameSupply(Iter begin, Iter end, Lambda f) : UniqueNameSupply("") {
+ TVM_TIRX_DLL UniqueNameSupply(Iter begin, Iter end, Lambda f) : UniqueNameSupply("") {
for (auto it = begin; it != end; ++it) {
this->operator->()->ReserveName(f(*it), false);
}
diff --git a/include/tvm/runtime/base.h b/include/tvm/runtime/base.h
index 60920d258b..bc31b47f07 100644
--- a/include/tvm/runtime/base.h
+++ b/include/tvm/runtime/base.h
@@ -75,6 +75,14 @@
#define TVM_DLL_EXPORT __attribute__((visibility("default")))
#endif
+// Core IR, TIRx and shared compiler infrastructure are currently linked into
+// libtvm_compiler. Keep their ownership explicit while inheriting the compiler
+// export policy, including downstream TVM_DLL overrides. Embedders may override
+// TVM_TIRX_DLL independently.
+#ifndef TVM_TIRX_DLL
+#define TVM_TIRX_DLL TVM_DLL
+#endif
+
// --- TVM_RUNTIME_DLL family (libtvm_runtime) ---
#if !defined(TVM_RUNTIME_DLL) && defined(__EMSCRIPTEN__)
#define TVM_RUNTIME_DLL EMSCRIPTEN_KEEPALIVE
diff --git a/include/tvm/script/ir_builder/ir.h b/include/tvm/script/ir_builder/ir.h
index eda4b3bceb..725434a73c 100644
--- a/include/tvm/script/ir_builder/ir.h
+++ b/include/tvm/script/ir_builder/ir.h
@@ -36,7 +36,7 @@ namespace ir {
* \brief The IRModule declaration statement.
* \return The IRModuleFrame.
*/
-TVM_DLL IRModuleFrame IRModule();
+TVM_TIRX_DLL IRModuleFrame IRModule();
/*!
* \brief Declare a Function without given the specific function implementation.
@@ -46,14 +46,14 @@ TVM_DLL IRModuleFrame IRModule();
* (i.e. func params and func return type/shape).
* \return The corresponding GlobalVar.
*/
-TVM_DLL GlobalVar DeclFunction(const ffi::String& func_name, const BaseFunc& func_signature);
+TVM_TIRX_DLL GlobalVar DeclFunction(const ffi::String& func_name, const BaseFunc& func_signature);
/*!
* \brief Define the function which is declared before.
* \param func_name The function unique name.
* \param func The given function implementation
*/
-TVM_DLL void DefFunction(const ffi::String& func_name, const BaseFunc& func);
+TVM_TIRX_DLL void DefFunction(const ffi::String& func_name, const BaseFunc& func);
/*!
* \brief Construct a For statement of the given kind.
diff --git a/include/tvm/script/printer/config.h b/include/tvm/script/printer/config.h
index 4d119f05f1..dfb9dbbe40 100644
--- a/include/tvm/script/printer/config.h
+++ b/include/tvm/script/printer/config.h
@@ -116,7 +116,7 @@ class PrinterConfigObj : public ffi::Object {
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.PrinterConfig", PrinterConfigObj, ffi::Object);
};
-class TVM_DLL PrinterConfig : public ffi::ObjectRef {
+class TVM_TIRX_DLL PrinterConfig : public ffi::ObjectRef {
public:
explicit PrinterConfig(
ffi::Map<ffi::String, ffi::Any> config_dict = ffi::Map<ffi::String, ffi::Any>());
@@ -131,7 +131,7 @@ class TVM_DLL PrinterConfig : public ffi::ObjectRef {
* address string. Defined in src/script/printer/script_printer.cc so that
* <tvm/runtime/logging.h> is not pulled into this public header.
*/
-TVM_DLL std::string RedirectedReprPrinterMethod(const ffi::ObjectRef& obj);
+TVM_TIRX_DLL std::string RedirectedReprPrinterMethod(const ffi::ObjectRef& obj);
} // namespace tvm
#endif // TVM_SCRIPT_PRINTER_CONFIG_H_
diff --git a/include/tvm/script/printer/doc.h b/include/tvm/script/printer/doc.h
index c634448959..b8afc30f40 100644
--- a/include/tvm/script/printer/doc.h
+++ b/include/tvm/script/printer/doc.h
@@ -95,19 +95,19 @@ class ExprDocNode : public DocNode {
* \brief Create a doc representing attribute access on the current ExprDoc
* \param attr The attribute to access.
*/
- TVM_DLL ExprDoc Attr(ffi::String attr) const;
+ TVM_TIRX_DLL ExprDoc Attr(ffi::String attr) const;
/*!
* \brief Create a doc representing index access on the current ExprDoc
* \param indices The indices to access.
*/
- TVM_DLL ExprDoc operator[](ffi::Array<Doc> indices) const;
+ TVM_TIRX_DLL ExprDoc operator[](ffi::Array<Doc> indices) const;
/*!
* \brief Create a doc representing calling the current ExprDoc
* \param args The positional arguments of the function call.
*/
- TVM_DLL ExprDoc Call(ffi::Array<ExprDoc, void> args) const;
+ TVM_TIRX_DLL ExprDoc Call(ffi::Array<ExprDoc, void> args) const;
/*!
* \brief Create a doc representing attribute access on the current ExprDoc
@@ -115,9 +115,9 @@ class ExprDocNode : public DocNode {
* \param kwargs_keys Keys of keywords arguments of the function call.
* \param kwargs_values Values of keywords arguments of the function call.
*/
- TVM_DLL ExprDoc Call(ffi::Array<ExprDoc, void> args, //
- ffi::Array<ffi::String> kwargs_keys, //
- ffi::Array<ExprDoc, void> kwargs_values) const;
+ TVM_TIRX_DLL ExprDoc Call(ffi::Array<ExprDoc, void> args, //
+ ffi::Array<ffi::String> kwargs_keys, //
+ ffi::Array<ExprDoc, void> kwargs_values) const;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
@@ -140,7 +140,7 @@ class ExprDoc : public Doc {
* \brief Create a doc representing index access on the current ExprDoc
* \param indices The indices to access.
*/
- TVM_DLL ExprDoc operator[](ffi::Array<Doc> indices) const;
+ TVM_TIRX_DLL ExprDoc operator[](ffi::Array<Doc> indices) const;
explicit ExprDoc(ffi::ObjectPtr<ExprDocNode> data) : Doc(ffi::UnsafeInit{}) {
TVM_FFI_ICHECK(data != nullptr);
@@ -215,7 +215,7 @@ class StmtBlockDoc : public Doc {
* \brief Constructor of StmtBlockDoc.
* \param stmts The list of statements.
*/
- TVM_DLL explicit StmtBlockDoc(ffi::Array<StmtDoc> stmts);
+ TVM_TIRX_DLL explicit StmtBlockDoc(ffi::Array<StmtDoc> stmts);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(StmtBlockDoc, Doc, StmtBlockDocNode);
};
@@ -251,7 +251,7 @@ class LiteralDocNode : public ExprDocNode {
*/
class LiteralDoc : public ExprDoc {
protected:
- TVM_DLL explicit LiteralDoc(ffi::Any value, const ffi::Optional<AccessPath>& object_path);
+ TVM_TIRX_DLL explicit LiteralDoc(ffi::Any value, const ffi::Optional<AccessPath>& object_path);
public:
/*!
@@ -356,7 +356,7 @@ class ExprStringDoc : public ExprDoc {
* \param value The expression to render as a string.
* \param object_path The object path.
*/
- TVM_DLL explicit ExprStringDoc(ExprDoc value, const ffi::Optional<AccessPath>& object_path);
+ TVM_TIRX_DLL explicit ExprStringDoc(ExprDoc value, const ffi::Optional<AccessPath>& object_path);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ExprStringDoc, ExprDoc, ExprStringDocNode);
};
@@ -389,7 +389,7 @@ class IdDoc : public ExprDoc {
* \brief Constructor of IdDoc.
* \param name The name of identifier.
*/
- TVM_DLL explicit IdDoc(ffi::String name);
+ TVM_TIRX_DLL explicit IdDoc(ffi::String name);
explicit IdDoc(std::nullptr_t) : ExprDoc(nullptr) {}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IdDoc, ExprDoc, IdDocNode);
};
@@ -427,7 +427,7 @@ class AttrAccessDoc : public ExprDoc {
* \param value The target expression of attribute access.
* \param name The name of attribute to access.
*/
- TVM_DLL explicit AttrAccessDoc(ExprDoc value, ffi::String name);
+ TVM_TIRX_DLL explicit AttrAccessDoc(ExprDoc value, ffi::String name);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(AttrAccessDoc, ExprDoc, AttrAccessDocNode);
};
@@ -470,7 +470,7 @@ class IndexDoc : public ExprDoc {
* \param value The target expression of index access.
* \param indices The indices to access.
*/
- TVM_DLL explicit IndexDoc(ExprDoc value, ffi::Array<Doc> indices);
+ TVM_TIRX_DLL explicit IndexDoc(ExprDoc value, ffi::Array<Doc> indices);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IndexDoc, ExprDoc, IndexDocNode);
};
@@ -520,8 +520,8 @@ class CallDoc : public ExprDoc {
* \param kwargs_keys Keys of keyword arguments.
* \param kwargs_values Values of keyword arguments, must have the same length as `kwargs_keys.
*/
- TVM_DLL CallDoc(ExprDoc callee, ffi::Array<ExprDoc> args, ffi::Array<ffi::String> kwargs_keys,
- ffi::Array<ExprDoc> kwargs_values);
+ TVM_TIRX_DLL CallDoc(ExprDoc callee, ffi::Array<ExprDoc> args,
+ ffi::Array<ffi::String> kwargs_keys, ffi::Array<ExprDoc> kwargs_values);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(CallDoc, ExprDoc, CallDocNode);
};
@@ -600,7 +600,7 @@ class OperationDoc : public ExprDoc {
* \param kind The kind of operation.
* \param operands Operands of this expression.
*/
- TVM_DLL explicit OperationDoc(OperationDocNode::Kind kind, ffi::Array<ExprDoc> operands);
+ TVM_TIRX_DLL explicit OperationDoc(OperationDocNode::Kind kind, ffi::Array<ExprDoc> operands);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(OperationDoc, ExprDoc, OperationDocNode);
};
@@ -640,7 +640,7 @@ class LambdaDoc : public ExprDoc {
* \param args Arguments of this function.
* \param body Body expression of this function.
*/
- TVM_DLL explicit LambdaDoc(ffi::Array<IdDoc> args, ExprDoc body);
+ TVM_TIRX_DLL explicit LambdaDoc(ffi::Array<IdDoc> args, ExprDoc body);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(LambdaDoc, ExprDoc, LambdaDocNode);
};
@@ -676,7 +676,7 @@ class TupleDoc : public ExprDoc {
* \brief Constructor of TupleDoc
* \param elements Elements of tuple.
*/
- TVM_DLL explicit TupleDoc(ffi::Array<ExprDoc> elements);
+ TVM_TIRX_DLL explicit TupleDoc(ffi::Array<ExprDoc> elements);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TupleDoc, ExprDoc, TupleDocNode);
};
@@ -712,7 +712,7 @@ class ListDoc : public ExprDoc {
* \brief Constructor of ListDoc
* \param elements Elements of list.
*/
- TVM_DLL explicit ListDoc(ffi::Array<ExprDoc> elements);
+ TVM_TIRX_DLL explicit ListDoc(ffi::Array<ExprDoc> elements);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ListDoc, ExprDoc, ListDocNode);
};
@@ -758,7 +758,7 @@ class DictDoc : public ExprDoc {
* \param keys Keys of dictionary.
* \param values Values of dictionary, must have same length as `keys`.
*/
- TVM_DLL explicit DictDoc(ffi::Array<ExprDoc> keys, ffi::Array<ExprDoc> values);
+ TVM_TIRX_DLL explicit DictDoc(ffi::Array<ExprDoc> keys, ffi::Array<ExprDoc> values);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DictDoc, ExprDoc, DictDocNode);
};
@@ -801,8 +801,8 @@ class SliceDoc : public Doc {
* \param stop The exclusive end of slice.
* \param step The step of slice.
*/
- TVM_DLL explicit SliceDoc(ffi::Optional<ExprDoc> start, ffi::Optional<ExprDoc> stop,
- ffi::Optional<ExprDoc> step);
+ TVM_TIRX_DLL explicit SliceDoc(ffi::Optional<ExprDoc> start, ffi::Optional<ExprDoc> stop,
+ ffi::Optional<ExprDoc> step);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(SliceDoc, Doc, SliceDocNode);
};
@@ -847,8 +847,8 @@ class AssignDoc : public StmtDoc {
* \param rhs The right hand side of the assignment.
* \param annotation The type annotation of this assignment.
*/
- TVM_DLL explicit AssignDoc(ExprDoc lhs, ffi::Optional<ExprDoc> rhs,
- ffi::Optional<ExprDoc> annotation);
+ TVM_TIRX_DLL explicit AssignDoc(ExprDoc lhs, ffi::Optional<ExprDoc> rhs,
+ ffi::Optional<ExprDoc> annotation);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(AssignDoc, StmtDoc, AssignDocNode);
};
@@ -889,8 +889,8 @@ class IfDoc : public StmtDoc {
* \param then_branch The then branch of the if-then-else statement.
* \param else_branch The else branch, or nullopt when no else branch is present.
*/
- TVM_DLL explicit IfDoc(ExprDoc predicate, ffi::Array<StmtDoc> then_branch,
- ffi::Optional<ffi::Array<StmtDoc>> else_branch);
+ TVM_TIRX_DLL explicit IfDoc(ExprDoc predicate, ffi::Array<StmtDoc> then_branch,
+ ffi::Optional<ffi::Array<StmtDoc>> else_branch);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IfDoc, StmtDoc, IfDocNode);
};
@@ -927,7 +927,7 @@ class WhileDoc : public StmtDoc {
* \param predicate The predicate of the while statement.
* \param body The body of the while statement.
*/
- TVM_DLL explicit WhileDoc(ExprDoc predicate, ffi::Array<StmtDoc> body);
+ TVM_TIRX_DLL explicit WhileDoc(ExprDoc predicate, ffi::Array<StmtDoc> body);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(WhileDoc, StmtDoc, WhileDocNode);
};
@@ -956,7 +956,7 @@ class BreakDoc : public StmtDoc {
/*!
* \brief Constructor of BreakDoc.
*/
- TVM_DLL explicit BreakDoc();
+ TVM_TIRX_DLL explicit BreakDoc();
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(BreakDoc, StmtDoc, BreakDocNode);
};
@@ -985,7 +985,7 @@ class ContinueDoc : public StmtDoc {
/*!
* \brief Constructor of ContinueDoc.
*/
- TVM_DLL explicit ContinueDoc();
+ TVM_TIRX_DLL explicit ContinueDoc();
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ContinueDoc, StmtDoc, ContinueDocNode);
};
@@ -1030,7 +1030,7 @@ class ForDoc : public StmtDoc {
* \param rhs The right hand side of the assignment of iterating variable.
* \param body The body of the for statement.
*/
- TVM_DLL explicit ForDoc(ExprDoc lhs, ExprDoc rhs, ffi::Array<StmtDoc> body);
+ TVM_TIRX_DLL explicit ForDoc(ExprDoc lhs, ExprDoc rhs, ffi::Array<StmtDoc> body);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ForDoc, StmtDoc, ForDocNode);
};
@@ -1080,15 +1080,15 @@ class ScopeDoc : public StmtDoc {
* \param body The body of the scope doc.
* \param allow_concise_scoping Whether a tail-position renderer may use a concise form.
*/
- TVM_DLL explicit ScopeDoc(ffi::Optional<ExprDoc> lhs, ExprDoc rhs, ffi::Array<StmtDoc> body,
- bool allow_concise_scoping = false);
+ TVM_TIRX_DLL explicit ScopeDoc(ffi::Optional<ExprDoc> lhs, ExprDoc rhs, ffi::Array<StmtDoc> body,
+ bool allow_concise_scoping = false);
/*!
* \brief Constructor of ScopeDoc.
* \param rhs The value of the scoped variable.
* \param body The body of the scope doc.
*/
- TVM_DLL explicit ScopeDoc(ExprDoc rhs, ffi::Array<StmtDoc> body);
+ TVM_TIRX_DLL explicit ScopeDoc(ExprDoc rhs, ffi::Array<StmtDoc> body);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ScopeDoc, StmtDoc, ScopeDocNode);
};
@@ -1121,7 +1121,7 @@ class ExprStmtDoc : public StmtDoc {
* \brief Constructor of ExprStmtDoc.
* \param expr The expression represented by this doc.
*/
- TVM_DLL explicit ExprStmtDoc(ExprDoc expr);
+ TVM_TIRX_DLL explicit ExprStmtDoc(ExprDoc expr);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ExprStmtDoc, StmtDoc, ExprStmtDocNode);
};
@@ -1158,7 +1158,7 @@ class AssertDoc : public StmtDoc {
* \param test The expression to test.
* \param msg The optional error message when assertion failed.
*/
- TVM_DLL explicit AssertDoc(ExprDoc test, ffi::Optional<ExprDoc> msg = std::nullopt);
+ TVM_TIRX_DLL explicit AssertDoc(ExprDoc test, ffi::Optional<ExprDoc> msg = std::nullopt);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(AssertDoc, StmtDoc, AssertDocNode);
};
@@ -1190,7 +1190,7 @@ class ReturnDoc : public StmtDoc {
* \brief Constructor of ReturnDoc.
* \param value The value to return.
*/
- TVM_DLL explicit ReturnDoc(ExprDoc value);
+ TVM_TIRX_DLL explicit ReturnDoc(ExprDoc value);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ReturnDoc, StmtDoc, ReturnDocNode);
};
@@ -1255,9 +1255,10 @@ class FunctionDoc : public StmtDoc {
* \param body The body of function.
* \param type_params The PEP 695 type parameters of the function.
*/
- TVM_DLL explicit FunctionDoc(IdDoc name, ffi::Array<AssignDoc> args,
- ffi::Array<ExprDoc> decorators, ffi::Optional<ExprDoc> return_type,
- ffi::Array<StmtDoc> body, ffi::Array<Doc> type_params = {});
+ TVM_TIRX_DLL explicit FunctionDoc(IdDoc name, ffi::Array<AssignDoc> args,
+ ffi::Array<ExprDoc> decorators,
+ ffi::Optional<ExprDoc> return_type, ffi::Array<StmtDoc> body,
+ ffi::Array<Doc> type_params = {});
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(FunctionDoc, StmtDoc, FunctionDocNode);
};
@@ -1298,7 +1299,8 @@ class ClassDoc : public StmtDoc {
* \param decorators The decorator of class.
* \param body The body of class.
*/
- TVM_DLL explicit ClassDoc(IdDoc name, ffi::Array<ExprDoc> decorators, ffi::Array<StmtDoc> body);
+ TVM_TIRX_DLL explicit ClassDoc(IdDoc name, ffi::Array<ExprDoc> decorators,
+ ffi::Array<StmtDoc> body);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ClassDoc, StmtDoc, ClassDocNode);
};
@@ -1323,7 +1325,7 @@ class CommentDocNode : public StmtDocNode {
*/
class CommentDoc : public StmtDoc {
public:
- TVM_DLL explicit CommentDoc(ffi::String comment);
+ TVM_TIRX_DLL explicit CommentDoc(ffi::String comment);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(CommentDoc, StmtDoc, CommentDocNode);
};
@@ -1348,7 +1350,7 @@ class DocStringDocNode : public StmtDocNode {
*/
class DocStringDoc : public StmtDoc {
public:
- TVM_DLL explicit DocStringDoc(ffi::String docs);
+ TVM_TIRX_DLL explicit DocStringDoc(ffi::String docs);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DocStringDoc, StmtDoc, DocStringDocNode);
};
@@ -1378,7 +1380,7 @@ class NamespaceDoc : public ExprDoc {
* \brief Construct a deferred namespace reference.
* \param canonical_name The namespace name before alias resolution.
*/
- TVM_DLL explicit NamespaceDoc(ffi::String canonical_name);
+ TVM_TIRX_DLL explicit NamespaceDoc(ffi::String canonical_name);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(NamespaceDoc, ExprDoc, NamespaceDocNode);
};
diff --git a/include/tvm/script/printer/doc_translator.h b/include/tvm/script/printer/doc_translator.h
index 49613f2612..7ff75e86ee 100644
--- a/include/tvm/script/printer/doc_translator.h
+++ b/include/tvm/script/printer/doc_translator.h
@@ -374,8 +374,9 @@ class DocTranslator : public ffi::ObjectRef {
* \param extra_config The translation configuration snapshot.
* \return The translated Doc, with text headers and diagnostic recovery left to the text adapter.
*/
-TVM_DLL Doc DocTranslate(ffi::AnyView ir, ffi::Dict<Doc, ffi::ObjectRef>* doc_origins = nullptr,
- ffi::Map<ffi::String, ffi::Any> extra_config = {});
+TVM_TIRX_DLL Doc DocTranslate(ffi::AnyView ir,
+ ffi::Dict<Doc, ffi::ObjectRef>* doc_origins = nullptr,
+ ffi::Map<ffi::String, ffi::Any> extra_config = {});
namespace type_attr {
/*!
diff --git a/include/tvm/script/printer/printer.h b/include/tvm/script/printer/printer.h
index 308d424bfa..b9f76d16e5 100644
--- a/include/tvm/script/printer/printer.h
+++ b/include/tvm/script/printer/printer.h
@@ -38,8 +38,8 @@ namespace tvm {
* \param config Optional translation and rendering configuration.
* \return The rendered script or fallback representation.
*/
-TVM_DLL std::string Script(const ffi::ObjectRef& node,
- const ffi::Optional<PrinterConfig>& config = std::nullopt);
+TVM_TIRX_DLL std::string Script(const ffi::ObjectRef& node,
+ const ffi::Optional<PrinterConfig>& config = std::nullopt);
namespace script {
namespace printer {
@@ -49,10 +49,10 @@ namespace printer {
* \param key The existing prefix configuration key, such as "tirx.prefix".
* \param default_alias The alias reserved before translation assigns variable names.
*/
-TVM_DLL void RegisterNamespaceAlias(const ffi::String& key, const ffi::String& default_alias);
+TVM_TIRX_DLL void RegisterNamespaceAlias(const ffi::String& key, const ffi::String& default_alias);
/*! \brief Read the registered namespace aliases. */
-TVM_DLL const ffi::Map<ffi::String, ffi::String>& GetNamespaceAliases();
+TVM_TIRX_DLL const ffi::Map<ffi::String, ffi::String>& GetNamespaceAliases();
} // namespace printer
} // namespace script
diff --git a/include/tvm/sym/analyzer.h b/include/tvm/sym/analyzer.h
index 05e21c71ed..146277301b 100644
--- a/include/tvm/sym/analyzer.h
+++ b/include/tvm/sym/analyzer.h
@@ -121,7 +121,7 @@ class ConstIntBound : public ffi::ObjectRef {
* \param min_value The mininum value.
* \param max_value The maximum value.
*/
- TVM_DLL ConstIntBound(int64_t min_value, int64_t max_value);
+ TVM_TIRX_DLL ConstIntBound(int64_t min_value, int64_t max_value);
static const constexpr int64_t kPosInf = ConstIntBoundObj::kPosInf;
static const constexpr int64_t kNegInf = ConstIntBoundObj::kNegInf;
@@ -140,7 +140,7 @@ class ConstIntBoundAnalyzer {
* \param expr The expression of interest.
* \return the result of the analysis.
*/
- TVM_DLL ConstIntBound operator()(const PrimExpr& expr) const;
+ TVM_TIRX_DLL ConstIntBound operator()(const PrimExpr& expr) const;
/*!
* \brief analyze the expr with the intermediate memorized to avoid redundant computation
@@ -148,7 +148,7 @@ class ConstIntBoundAnalyzer {
* \param bound The lookup table to store the intermediate results
* \return the result of the analysis.
*/
- TVM_DLL ConstIntBound operator()(const PrimExpr& expr, BoundMapType* bound);
+ TVM_TIRX_DLL ConstIntBound operator()(const PrimExpr& expr, BoundMapType* bound);
/*!
* \brief Update constant int bound information of var.
@@ -157,7 +157,7 @@ class ConstIntBoundAnalyzer {
* \param info The bound information.
* \param allow_override whether we allow override of existing information.
*/
- TVM_DLL void Update(const Var& var, const ConstIntBound& info, bool allow_override = false);
+ TVM_TIRX_DLL void Update(const Var& var, const ConstIntBound& info, bool allow_override = false);
/*!
* \brief Bind variable to a range.
@@ -166,20 +166,20 @@ class ConstIntBoundAnalyzer {
* \param range The range we bind to.
* \param allow_override Whether we allow overriding an existing var's range.
*/
- TVM_DLL void Bind(const Var& var, const Range& range, bool allow_override = false);
+ TVM_TIRX_DLL void Bind(const Var& var, const Range& range, bool allow_override = false);
/*!
* \brief Check if a variable is bound to a range.
* \param var The variable.
* \return Whether the variable is bound to a range.
*/
- TVM_DLL bool IsBound(const Var& var) const;
+ TVM_TIRX_DLL bool IsBound(const Var& var) const;
private:
friend class AnalyzerObj;
friend class ConstraintContext;
explicit ConstIntBoundAnalyzer(AnalyzerObj* parent);
- TVM_DLL ~ConstIntBoundAnalyzer();
+ TVM_TIRX_DLL ~ConstIntBoundAnalyzer();
void CopyFrom(const ConstIntBoundAnalyzer& other);
/*!
* \brief Update the internal state to enter constraint.
@@ -230,7 +230,7 @@ class ModularSetObj : public ffi::Object {
*/
class ModularSet : public ffi::ObjectRef {
public:
- TVM_DLL ModularSet(int64_t coeff, int64_t base);
+ TVM_TIRX_DLL ModularSet(int64_t coeff, int64_t base);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ModularSet, ffi::ObjectRef, ModularSetObj);
};
@@ -245,7 +245,7 @@ class ModularSetAnalyzer {
* \param expr The expression of interest.
* \return the result of the analysis.
*/
- TVM_DLL ModularSet operator()(const PrimExpr& expr);
+ TVM_TIRX_DLL ModularSet operator()(const PrimExpr& expr);
/*!
* \brief Update constant int bound information of var.
*
@@ -253,13 +253,13 @@ class ModularSetAnalyzer {
* \param info The bound information.
* \param allow_override whether we allow override of existing information.
*/
- TVM_DLL void Update(const Var& var, const ModularSet& info, bool allow_override = false);
+ TVM_TIRX_DLL void Update(const Var& var, const ModularSet& info, bool allow_override = false);
private:
friend class AnalyzerObj;
friend class ConstraintContext;
explicit ModularSetAnalyzer(AnalyzerObj* parent);
- TVM_DLL ~ModularSetAnalyzer();
+ TVM_TIRX_DLL ~ModularSetAnalyzer();
void CopyFrom(const ModularSetAnalyzer& other);
/*!
* \brief Update the internal state to enter constraint.
@@ -284,7 +284,7 @@ class RewriteSimplifier {
* \param expr The expression of interest.
* \return the result of the analysis.
*/
- TVM_DLL PrimExpr operator()(const PrimExpr& expr);
+ TVM_TIRX_DLL PrimExpr operator()(const PrimExpr& expr);
/*!
* \brief Update binding of var to a new expression.
@@ -293,7 +293,7 @@ class RewriteSimplifier {
* \param new_expr
* \param allow_override Whether we allow override of existing information.
*/
- TVM_DLL void Update(const Var& var, const PrimExpr& new_expr, bool allow_override = false);
+ TVM_TIRX_DLL void Update(const Var& var, const PrimExpr& new_expr, bool allow_override = false);
/*!
* \brief Update the internal state to enter constraint.
@@ -302,7 +302,8 @@ class RewriteSimplifier {
*
* \return an exit function that must be called to cleanup the constraint can be nullptr.
*/
- TVM_DLL std::function<void()> EnterConstraint(const PrimExpr& constraint, bool is_assume = false);
+ TVM_TIRX_DLL std::function<void()> EnterConstraint(const PrimExpr& constraint,
+ bool is_assume = false);
/*! \brief Flags to enable more computationally-intensive simplifications
*
@@ -383,16 +384,16 @@ class RewriteSimplifier {
* \param flags A bitwise OR of all optional extensions that should
* be enabled.
*/
- TVM_DLL void SetEnabledExtensions(Extension flags);
+ TVM_TIRX_DLL void SetEnabledExtensions(Extension flags);
/*! \brief Return the currently enabled extensions */
- TVM_DLL Extension GetEnabledExtensions() const;
+ TVM_TIRX_DLL Extension GetEnabledExtensions() const;
/*! \brief Return the statistics counters */
- TVM_DLL ffi::ObjectRef GetStatsCounters() const;
+ TVM_TIRX_DLL ffi::ObjectRef GetStatsCounters() const;
/*! \brief Reset the statistics counters */
- TVM_DLL void ResetStatsCounters();
+ TVM_TIRX_DLL void ResetStatsCounters();
/*! \brief Set the maximum allowed number of rewrite steps
*
@@ -408,14 +409,14 @@ class RewriteSimplifier {
* the analyzer temporarily may receive the underlying `AnalyzerObj*`
* from their calling scope.
*/
- TVM_DLL void SetMaximumRewriteSteps(int64_t maximum);
+ TVM_TIRX_DLL void SetMaximumRewriteSteps(int64_t maximum);
private:
friend class AnalyzerObj;
friend class ConstraintContext;
friend class CanonicalSimplifier;
explicit RewriteSimplifier(AnalyzerObj* parent);
- TVM_DLL ~RewriteSimplifier();
+ TVM_TIRX_DLL ~RewriteSimplifier();
void CopyFrom(const RewriteSimplifier& other);
class Impl;
/*! \brief Internal impl */
@@ -432,7 +433,7 @@ class CanonicalSimplifier {
* \param expr The expression of interest.
* \return the result of the analysis.
*/
- TVM_DLL PrimExpr operator()(const PrimExpr& expr);
+ TVM_TIRX_DLL PrimExpr operator()(const PrimExpr& expr);
/*!
* \brief Update binding of var to a new expression.
@@ -441,13 +442,13 @@ class CanonicalSimplifier {
* \param new_expr
* \param allow_override whether we allow override of existing information.
*/
- TVM_DLL void Update(const Var& var, const PrimExpr& new_expr, bool allow_override = false);
+ TVM_TIRX_DLL void Update(const Var& var, const PrimExpr& new_expr, bool allow_override = false);
private:
friend class AnalyzerObj;
friend class ConstraintContext;
explicit CanonicalSimplifier(AnalyzerObj* parent);
- TVM_DLL ~CanonicalSimplifier();
+ TVM_TIRX_DLL ~CanonicalSimplifier();
void CopyFrom(const CanonicalSimplifier& other);
class Impl;
/*! \brief Internal impl */
@@ -502,8 +503,8 @@ class TransitiveComparisonAnalyzer {
* \return The most specific result that can be proven about the
* comparison. If nothing can be proven, returns kUnknown.
*/
- TVM_DLL CompareResult TryCompare(const PrimExpr& lhs, const PrimExpr& rhs,
- bool propagate_inequalities = true);
+ TVM_TIRX_DLL CompareResult TryCompare(const PrimExpr& lhs, const PrimExpr& rhs,
+ bool propagate_inequalities = true);
/*! \brief Bind a variable as being equal to a known expression
*
@@ -511,7 +512,7 @@ class TransitiveComparisonAnalyzer {
* \param expr The bound expression
* \param allow_override Whether to allow override of existing information.
*/
- TVM_DLL void Bind(const Var& var, const PrimExpr& expr, bool allow_override = false);
+ TVM_TIRX_DLL void Bind(const Var& var, const PrimExpr& expr, bool allow_override = false);
/*! \brief Bind a variable as being within a specified range
*
@@ -519,7 +520,7 @@ class TransitiveComparisonAnalyzer {
* \param range The known range
* \param allow_override Whether to allow override of existing information.
*/
- TVM_DLL void Bind(const Var& var, const Range& range, bool allow_override = false);
+ TVM_TIRX_DLL void Bind(const Var& var, const Range& range, bool allow_override = false);
/*!
* \brief Update the internal state to enter constraint.
@@ -527,13 +528,13 @@ class TransitiveComparisonAnalyzer {
*
* \return an exit function that must be called to cleanup the constraint can be nullptr.
*/
- TVM_DLL std::function<void()> EnterConstraint(const PrimExpr& constraint);
+ TVM_TIRX_DLL std::function<void()> EnterConstraint(const PrimExpr& constraint);
private:
friend class AnalyzerObj;
friend class ConstraintContext;
TransitiveComparisonAnalyzer();
- TVM_DLL ~TransitiveComparisonAnalyzer();
+ TVM_TIRX_DLL ~TransitiveComparisonAnalyzer();
void CopyFrom(const TransitiveComparisonAnalyzer& other);
class Impl;
/*! \brief Internal impl */
@@ -553,7 +554,7 @@ class IntSetAnalyzer {
* \param dom_map The domain map to indicate which variable to relax.
* \return the result of the analysis.
*/
- TVM_DLL IntSet operator()(const PrimExpr& expr, const ffi::Map<Var, IntSet>& dom_map);
+ TVM_TIRX_DLL IntSet operator()(const PrimExpr& expr, const ffi::Map<Var, IntSet>& dom_map);
/*!
* \brief Find a symbolic integer set that contains all possible
@@ -563,7 +564,7 @@ class IntSetAnalyzer {
* \param expr The expression of interest.
* \return the result of the analysis.
*/
- TVM_DLL IntSet operator()(const PrimExpr& expr);
+ TVM_TIRX_DLL IntSet operator()(const PrimExpr& expr);
/*!
* \brief Update binding of var to a new expression.
@@ -572,7 +573,8 @@ class IntSetAnalyzer {
* \param new_interval_set The set of allowed values for this var.
* \param allow_override whether we allow override of existing information.
*/
- TVM_DLL void Update(const Var& var, const IntSet& new_interval_set, bool allow_override = false);
+ TVM_TIRX_DLL void Update(const Var& var, const IntSet& new_interval_set,
+ bool allow_override = false);
/*!
* \brief Update binding of var to a new expression.
@@ -581,14 +583,14 @@ class IntSetAnalyzer {
* \param new_range The range of allowed values for this var.
* \param allow_override whether we allow override of existing information.
*/
- TVM_DLL void Bind(const Var& var, const Range& new_range, bool allow_override = false);
+ TVM_TIRX_DLL void Bind(const Var& var, const Range& new_range, bool allow_override = false);
std::function<void()> EnterConstraint(const PrimExpr& constraint);
private:
friend class AnalyzerObj;
explicit IntSetAnalyzer(AnalyzerObj* parent);
- TVM_DLL ~IntSetAnalyzer();
+ TVM_TIRX_DLL ~IntSetAnalyzer();
void CopyFrom(const IntSetAnalyzer& other);
class Impl;
/*! \brief Internal impl */
@@ -604,7 +606,7 @@ class Z3Prover {
* \param new_range The range of allowed values for this var.
* \param allow_override whether we allow override of existing information.
*/
- TVM_DLL void Bind(const Var& var, const Range& new_range, bool allow_override = false);
+ TVM_TIRX_DLL void Bind(const Var& var, const Range& new_range, bool allow_override = false);
/*!
* \brief Update binding of var to a new expression.
@@ -613,14 +615,14 @@ class Z3Prover {
* \param expr The bound expression.
* \param allow_override whether we allow override of existing information.
*/
- TVM_DLL void Bind(const Var& var, const PrimExpr& expr, bool allow_override = false);
+ TVM_TIRX_DLL void Bind(const Var& var, const PrimExpr& expr, bool allow_override = false);
/*!
* \brief Whether the Z3 backend is compiled into this build (USE_Z3=ON).
*
* \return true if the real Z3 prover is available, false for the stub.
*/
- TVM_DLL bool IsEnabled() const;
+ TVM_TIRX_DLL bool IsEnabled() const;
/*!
* \brief Whether can we prove expr is always true.
@@ -628,7 +630,7 @@ class Z3Prover {
* \param expr The expression.
* \return Whether we can prove it.
*/
- TVM_DLL bool CanProve(const PrimExpr& expr);
+ TVM_TIRX_DLL bool CanProve(const PrimExpr& expr);
/*!
* \brief Update the internal state to enter constraint.
@@ -687,14 +689,14 @@ class Z3Prover {
* \param min_consecutive Minimum consecutive count requirement.
* \return The number of distinct values that satisfy the constraints, or a negative error code.
*/
- TVM_DLL int64_t CountSatisfyingValues(const Var& var, int64_t max_count = 2048,
- int64_t min_consecutive = 1);
+ TVM_TIRX_DLL int64_t CountSatisfyingValues(const Var& var, int64_t max_count = 2048,
+ int64_t min_consecutive = 1);
private:
friend class AnalyzerObj;
friend class Analyzer;
explicit Z3Prover(AnalyzerObj* parent);
- TVM_DLL ~Z3Prover();
+ TVM_TIRX_DLL ~Z3Prover();
void CopyFrom(const Z3Prover& other);
class Impl;
std::unique_ptr<Impl> impl_;
@@ -710,7 +712,7 @@ class Z3Prover {
* If the analyzer uses memoization, we need to clear the internal
* cache when information about a Var has been overridden.
*/
-class TVM_DLL AnalyzerObj : public ffi::Object {
+class TVM_TIRX_DLL AnalyzerObj : public ffi::Object {
public:
/*! \brief sub-analyzer: const integer bound */
ConstIntBoundAnalyzer const_int_bound;
@@ -976,9 +978,9 @@ class ConstraintContext {
ConstraintContext(AnalyzerObj* analyzer, PrimExpr constraint, bool is_assume)
: ConstraintContext(ffi::GetRef<Analyzer>(analyzer), std::move(constraint), is_assume) {}
// enter the scope.
- TVM_DLL void EnterWithScope();
+ TVM_TIRX_DLL void EnterWithScope();
// exit the scope.
- TVM_DLL void ExitWithScope();
+ TVM_TIRX_DLL void ExitWithScope();
/*! \brief Analyzer kept alive while the context is active. */
Analyzer analyzer_;
/*! \brief The constraint */
diff --git a/include/tvm/sym/int_set.h b/include/tvm/sym/int_set.h
index a90c58d3e3..d5a9c3f6be 100644
--- a/include/tvm/sym/int_set.h
+++ b/include/tvm/sym/int_set.h
@@ -279,7 +279,7 @@ ffi::Map<Var, sym::IntSet> AsIntSet(const ffi::Map<Var, Range>& var_dom);
* \return std::nullopt if the detection fails, or an array of sym::IntSet as the result of
* analysis
*/
-TVM_DLL ffi::Optional<ffi::Array<IntSet>> EstimateRegionStrictBound(
+TVM_TIRX_DLL ffi::Optional<ffi::Array<IntSet>> EstimateRegionStrictBound(
const ffi::Array<Range>& region, const ffi::Map<Var, Range>& var_dom, const PrimExpr& predicate,
const sym::Analyzer& analyzer);
@@ -293,7 +293,7 @@ TVM_DLL ffi::Optional<ffi::Array<IntSet>> EstimateRegionStrictBound(
* \return std::nullopt if the detection fails, or an array of sym::IntSet as the result of
* analysis
*/
-TVM_DLL ffi::Optional<ffi::Array<IntSet>> EstimateRegionLowerBound(
+TVM_TIRX_DLL ffi::Optional<ffi::Array<IntSet>> EstimateRegionLowerBound(
const ffi::Array<Range>& region, const ffi::Map<Var, Range>& var_dom, const PrimExpr& predicate,
const sym::Analyzer& analyzer);
@@ -307,10 +307,10 @@ TVM_DLL ffi::Optional<ffi::Array<IntSet>> EstimateRegionLowerBound(
* \param analyzer The analyzer used
* \return an array of sym::IntSet as the result of analysis
*/
-TVM_DLL ffi::Array<IntSet> EstimateRegionUpperBound(const ffi::Array<Range>& region,
- const ffi::Map<Var, Range>& var_dom,
- const PrimExpr& predicate,
- const sym::Analyzer& analyzer);
+TVM_TIRX_DLL ffi::Array<IntSet> EstimateRegionUpperBound(const ffi::Array<Range>& region,
+ const ffi::Map<Var, Range>& var_dom,
+ const PrimExpr& predicate,
+ const sym::Analyzer& analyzer);
} // namespace sym
} // namespace tvm
diff --git a/include/tvm/sym/iter_affine_map.h b/include/tvm/sym/iter_affine_map.h
index fa55862006..7b2880dc9e 100644
--- a/include/tvm/sym/iter_affine_map.h
+++ b/include/tvm/sym/iter_affine_map.h
@@ -125,7 +125,7 @@ class IterMark : public ffi::ObjectRef {
* \param source The source expression.
* \param extent The extent of the iterator.
*/
- TVM_DLL IterMark(PrimExpr source, PrimExpr extent);
+ TVM_TIRX_DLL IterMark(PrimExpr source, PrimExpr extent);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IterMark, ffi::ObjectRef, IterMarkNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(IterMarkNode);
@@ -182,13 +182,13 @@ class IterSplitExpr : public IterMapExpr {
* \brief constructor from just source.
* \param source The source expression.
*/
- TVM_DLL explicit IterSplitExpr(IterMark source);
+ TVM_TIRX_DLL explicit IterSplitExpr(IterMark source);
/*!
* \brief constructor from just source.
* \param source The source expression.
* \param scale The additional scaling factor.
*/
- TVM_DLL explicit IterSplitExpr(IterMark source, PrimExpr scale);
+ TVM_TIRX_DLL explicit IterSplitExpr(IterMark source, PrimExpr scale);
/*!
* \brief constructor
* \param source The source expression.
@@ -196,8 +196,8 @@ class IterSplitExpr : public IterMapExpr {
* \param extent The extent of the split.
* \param scale The additional scaling factor.
*/
- TVM_DLL explicit IterSplitExpr(IterMark source, PrimExpr lower_factor, PrimExpr extent,
- PrimExpr scale);
+ TVM_TIRX_DLL explicit IterSplitExpr(IterMark source, PrimExpr lower_factor, PrimExpr extent,
+ PrimExpr scale);
explicit IterSplitExpr(ffi::ObjectPtr<IterSplitExprNode> node) : IterMapExpr(std::move(node)) {}
@@ -242,7 +242,7 @@ class IterSumExpr : public IterMapExpr {
* \param args The args to the sum.
* \param base The base offset.
*/
- TVM_DLL IterSumExpr(ffi::Array<IterSplitExpr> args, PrimExpr base);
+ TVM_TIRX_DLL IterSumExpr(ffi::Array<IterSplitExpr> args, PrimExpr base);
explicit IterSumExpr(ffi::ObjectPtr<IterSumExprNode> node) : IterMapExpr(std::move(node)) {}
diff --git a/include/tvm/target/target.h b/include/tvm/target/target.h
index a1ff2dd45c..b1dce829e9 100644
--- a/include/tvm/target/target.h
+++ b/include/tvm/target/target.h
@@ -61,13 +61,13 @@ class TargetObj : public ffi::Object {
* \brief The JSON string representation of the target
* \return JSON string of the target configuration (e.g. {"kind": "llvm", "mcpu": "cortex-a53"})
*/
- TVM_DLL const std::string& str() const;
+ TVM_TIRX_DLL const std::string& str() const;
/*! \return Export target to JSON-like configuration */
- TVM_DLL ffi::Map<ffi::String, ffi::Any> ToConfig() const;
+ TVM_TIRX_DLL ffi::Map<ffi::String, ffi::Any> ToConfig() const;
/*! \return The ffi::Optional<Target> typed target host of the TargetObj */
- TVM_DLL ffi::Optional<Target> GetHost() const;
+ TVM_TIRX_DLL ffi::Optional<Target> GetHost() const;
/*! \return The device type for this target */
- TVM_DLL int GetTargetDeviceType() const;
+ TVM_TIRX_DLL int GetTargetDeviceType() const;
/*!
* \brief Check if the target contains a key
@@ -77,7 +77,7 @@ class TargetObj : public ffi::Object {
* \return True if the target's `TargetObj::keys` contains the
* specified key, False otherwise.
*/
- TVM_DLL bool HasKey(const std::string& query_key) const;
+ TVM_TIRX_DLL bool HasKey(const std::string& query_key) const;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
@@ -139,12 +139,12 @@ class Target : public ffi::ObjectRef {
* \brief Construct a Target given a string
* \param tag_or_config_or_target_str the string to parse for target
*/
- TVM_DLL explicit Target(const ffi::String& tag_or_config_or_target_str);
+ TVM_TIRX_DLL explicit Target(const ffi::String& tag_or_config_or_target_str);
/*!
* \brief Construct a Target using a JSON-like configuration
* \param config The JSON-like configuration for target
*/
- TVM_DLL explicit Target(const ffi::Map<ffi::String, ffi::Any>& config);
+ TVM_TIRX_DLL explicit Target(const ffi::Map<ffi::String, ffi::Any>& config);
/*!
* \brief Get the current target context from thread local storage.
* \param allow_not_defined If the context stack is empty and this is set to true, std::nullopt
@@ -153,13 +153,13 @@ class Target : public ffi::ObjectRef {
* \return The target that is the current context. The result may be absent if
* allow_not_defined is true.
*/
- TVM_DLL static ffi::Optional<Target> Current(bool allow_not_defined = true);
+ TVM_TIRX_DLL static ffi::Optional<Target> Current(bool allow_not_defined = true);
/*!
* \brief Construct a Target given target and host
* \param target The Target typed object with host field undefined for target
* \param host The Target typed object for target host
*/
- TVM_DLL explicit Target(Target target, ffi::Optional<Target> host);
+ TVM_TIRX_DLL explicit Target(Target target, ffi::Optional<Target> host);
explicit Target(ffi::UnsafeInit tag, ffi::ObjectPtr<TargetObj> n) : ffi::ObjectRef(tag) {
data_ = std::move(n);
}
@@ -183,12 +183,12 @@ class Target : public ffi::ObjectRef {
* The Target on top of the stack is used to determine which
* specialization to use when invoking a GenericFunc.
*/
- TVM_DLL void EnterWithScope();
+ TVM_TIRX_DLL void EnterWithScope();
/*!
* \brief Pop a target off the thread local context stack,
* restoring the previous target as the current context.
*/
- TVM_DLL void ExitWithScope();
+ TVM_TIRX_DLL void ExitWithScope();
};
} // namespace tvm
diff --git a/include/tvm/target/target_kind.h b/include/tvm/target/target_kind.h
index f017104872..3ee45dd619 100644
--- a/include/tvm/target/target_kind.h
+++ b/include/tvm/target/target_kind.h
@@ -95,7 +95,7 @@ class TargetKind : public ffi::ObjectRef {
* \param target_kind_name Name of the target kind
* \return The canonical kind, or nullopt when the name is unknown.
*/
- TVM_DLL static ffi::Optional<TargetKind> Get(const ffi::String& target_kind_name);
+ TVM_TIRX_DLL static ffi::Optional<TargetKind> Get(const ffi::String& target_kind_name);
/*! \brief Mutable access to the container class */
TargetKindObj* operator->() { return static_cast<TargetKindObj*>(data_.get()); }
@@ -120,31 +120,31 @@ class TargetKindRegistry {
* \brief Access the process-wide registry.
* \return The singleton registry.
*/
- TVM_DLL static TargetKindRegistry* Global();
+ TVM_TIRX_DLL static TargetKindRegistry* Global();
/*!
* \brief Register a name or return its existing canonical kind.
* \param name Target kind name.
* \return The canonical target kind for name. Standard target options are
* declared when the name is first registered.
*/
- TVM_DLL TargetKind RegisterOrGet(const ffi::String& name);
+ TVM_TIRX_DLL TargetKind RegisterOrGet(const ffi::String& name);
/*!
* \brief Look up a registered kind.
* \param name Target kind name.
* \return The canonical kind, or nullopt when name is unknown.
*/
- TVM_DLL ffi::Optional<TargetKind> Get(const ffi::String& name);
+ TVM_TIRX_DLL ffi::Optional<TargetKind> Get(const ffi::String& name);
/*!
* \brief List registered kind names.
* \return Names of all registered target kinds.
*/
- TVM_DLL ffi::Array<ffi::String> ListTargetKinds();
+ TVM_TIRX_DLL ffi::Array<ffi::String> ListTargetKinds();
/*!
* \brief List declared option names and types for a kind.
* \param kind Registered target kind.
* \return Map from option name to type string.
*/
- TVM_DLL ffi::Map<ffi::String, ffi::String> ListTargetKindOptions(const TargetKind& kind);
+ TVM_TIRX_DLL ffi::Map<ffi::String, ffi::String> ListTargetKindOptions(const TargetKind& kind);
private:
ffi::Map<ffi::String, TargetKind> kinds_;
diff --git a/include/tvm/target/target_tag_registry.h b/include/tvm/target/target_tag_registry.h
index c6b6ae15cb..f4be0effca 100644
--- a/include/tvm/target/target_tag_registry.h
+++ b/include/tvm/target/target_tag_registry.h
@@ -45,27 +45,27 @@ class TargetTagRegistry {
using Config = ffi::Map<ffi::String, ffi::Any>;
/*! \brief Return the process-wide registry. */
- TVM_DLL static TargetTagRegistry* Global();
+ TVM_TIRX_DLL static TargetTagRegistry* Global();
/*!
* \brief Construct the target registered under a name.
* \param name The tag name to look up.
* \return The constructed target, or nullopt if the name is unknown.
*/
- TVM_DLL ffi::Optional<Target> Get(const ffi::String& name) const;
+ TVM_TIRX_DLL ffi::Optional<Target> Get(const ffi::String& name) const;
/*!
* \brief Return a tag's stored configuration without constructing a target.
* \param name The tag name to look up.
* \return The configuration, or nullopt if the name is unknown.
*/
- TVM_DLL ffi::Optional<Config> GetConfig(const ffi::String& name) const;
+ TVM_TIRX_DLL ffi::Optional<Config> GetConfig(const ffi::String& name) const;
/*!
* \brief List all registered names and their constructed targets.
* \return A map from tag names to targets; empty if no tags are registered.
*/
- TVM_DLL ffi::Map<ffi::String, Target> ListTags() const;
+ TVM_TIRX_DLL ffi::Map<ffi::String, Target> ListTags() const;
/*!
* \brief Register a target configuration under a tag name.
@@ -74,7 +74,7 @@ class TargetTagRegistry {
* \param override Whether to replace an existing registration; otherwise a duplicate is rejected.
* \return The target constructed from config.
*/
- TVM_DLL Target AddTag(ffi::String name, Config config, bool override);
+ TVM_TIRX_DLL Target AddTag(ffi::String name, Config config, bool override);
private:
TargetTagRegistry() = default;
diff --git a/include/tvm/tirx/analysis.h b/include/tvm/tirx/analysis.h
index 3de8282f31..1f97156b30 100644
--- a/include/tvm/tirx/analysis.h
+++ b/include/tvm/tirx/analysis.h
@@ -48,20 +48,20 @@ namespace tirx {
* \param func The function to be verified.
* \return Success of memory verification.
*/
-TVM_DLL bool VerifyMemory(const Function& func);
+TVM_TIRX_DLL bool VerifyMemory(const Function& func);
/*!
* \brief Calculate the expression complexity based on number of symbols it contains.
* \param expr The expr to be calculated.
*/
-TVM_DLL size_t CalculateExprComplexity(const PrimExpr& expr);
+TVM_TIRX_DLL size_t CalculateExprComplexity(const PrimExpr& expr);
/*!
* \brief Calculate the constants size in bytes needed by the TIR allocates inside the TIR Function
* \param func The TIR Function for which the constants size to be calculated
* \param constant_byte_alignment The byte alignment required for each constant allocated
*/
-TVM_DLL size_t CalculateConstantBytes(const Function& func, int64_t constant_byte_alignment);
+TVM_TIRX_DLL size_t CalculateConstantBytes(const Function& func, int64_t constant_byte_alignment);
/*!
* \brief Calculate the workspace size in bytes needed by the TIR allocates inside the TIR Function
@@ -69,7 +69,7 @@ TVM_DLL size_t CalculateConstantBytes(const Function& func, int64_t constant_byt
* \param workspace_byte_alignment The byte alignment required for each tensor allocated in this
* workspace
*/
-TVM_DLL size_t CalculateWorkspaceBytes(const Function& func, int64_t workspace_byte_alignment);
+TVM_TIRX_DLL size_t CalculateWorkspaceBytes(const Function& func, int64_t workspace_byte_alignment);
/*!
* \brief Verify if the given TIR is well-formed. The verification includes:
@@ -90,7 +90,7 @@ TVM_DLL size_t CalculateWorkspaceBytes(const Function& func, int64_t workspace_b
* \param assert_mode The indicator if it raises an error when the function is not well-formed.
* \return Whether it is a well-formed TIR function.
*/
-TVM_DLL bool VerifyWellFormed(const Function& func, bool assert_mode = true);
+TVM_TIRX_DLL bool VerifyWellFormed(const Function& func, bool assert_mode = true);
/*!
* \brief Verify if the TIR in the given IRMOdule is well-formed.
@@ -104,7 +104,7 @@ TVM_DLL bool VerifyWellFormed(const Function& func, bool assert_mode = true);
* \param assert_mode The indicator if it raises an error when the function is not well-formed.
* \return Whether it is a well-formed TIR module.
*/
-TVM_DLL bool VerifyWellFormed(const IRModule& mod, bool assert_mode = true);
+TVM_TIRX_DLL bool VerifyWellFormed(const IRModule& mod, bool assert_mode = true);
/*!
* \brief Find the entry function of the given IRModule, i.e, functions marked by
@@ -128,7 +128,7 @@ using tvm::transform::PassContext;
* \returns The pass.
* \sa tvm::tirx::VerifyMemory
*/
-TVM_DLL Pass VerifyMemory();
+TVM_TIRX_DLL Pass VerifyMemory();
} // namespace transform
} // namespace tirx
diff --git a/include/tvm/tirx/expr.h b/include/tvm/tirx/expr.h
index 8aa8f5bc42..ccf97307f2 100644
--- a/include/tvm/tirx/expr.h
+++ b/include/tvm/tirx/expr.h
@@ -50,7 +50,7 @@ namespace tirx {
class TensorVar : public Var {
public:
/*! \brief Construct a fresh tensor variable from an explicit TensorType. */
- TVM_DLL explicit TensorVar(ffi::String name, TensorType type, Location loc = UnknownLoc());
+ TVM_TIRX_DLL explicit TensorVar(ffi::String name, TensorType type, Location loc = UnknownLoc());
/*! \brief Create a checked tensor view over an existing ordinary Var. */
explicit TensorVar(Var var) : Var(std::move(var)) {
@@ -85,7 +85,7 @@ class TensorVar : public Var {
* \param alloc_or_decl Whether to consider alloc_scalar and decl_scalar as scalar. True for
* alloc_scalar, False for decl_scalar.
*/
- TVM_DLL bool IsScalar(bool alloc_or_decl = true) const;
+ TVM_TIRX_DLL bool IsScalar(bool alloc_or_decl = true) const;
/*! \return primitive element type for compiler-side uses. */
PrimType ElementType() const { return (*this)->ElementType(); }
@@ -138,9 +138,9 @@ inline TensorVar RebuildTensorVar(const TensorVar& var, ffi::ObjectPtr<TensorTyp
* \return The created tensor.
* \sa TensorVar for complete constructor.
*/
-TVM_DLL TensorVar decl_tensor(ffi::Array<PrimExpr> shape, PrimType dtype = PrimType::Float(32),
- ffi::String name = "tensor", ffi::String storage_scope = "",
- Location loc = UnknownLoc());
+TVM_TIRX_DLL TensorVar decl_tensor(ffi::Array<PrimExpr> shape, PrimType dtype = PrimType::Float(32),
+ ffi::String name = "tensor", ffi::String storage_scope = "",
+ Location loc = UnknownLoc());
/*!
* \brief Creates a TIR tensor for the provided parameters.
@@ -154,9 +154,10 @@ TVM_DLL TensorVar decl_tensor(ffi::Array<PrimExpr> shape, PrimType dtype = PrimT
* A default value will be picked.
* \param memory_scope memory scope of the tensor
*/
-TVM_DLL tirx::TensorVar TensorWithOffsetAlignment(ffi::Array<PrimExpr> shape, PrimType dtype,
- std::string name, int data_alignment,
- int offset_factor, std::string memory_scope = "");
+TVM_TIRX_DLL tirx::TensorVar TensorWithOffsetAlignment(ffi::Array<PrimExpr> shape, PrimType dtype,
+ std::string name, int data_alignment,
+ int offset_factor,
+ std::string memory_scope = "");
/*!
* \brief Construct a TensorLoad from a TensorVar.
@@ -165,13 +166,13 @@ TVM_DLL tirx::TensorVar TensorWithOffsetAlignment(ffi::Array<PrimExpr> shape, Pr
* is derived from the tensor element type and index lanes, and every tirx
* TensorLoad is required to have a TensorVar source.
*/
-TVM_DLL TensorLoad MakeTensorLoad(TensorVar tensor, ffi::Array<PrimExpr> indices,
- Location loc = UnknownLoc());
+TVM_TIRX_DLL TensorLoad MakeTensorLoad(TensorVar tensor, ffi::Array<PrimExpr> indices,
+ Location loc = UnknownLoc());
/*! \brief Select the entire tensor. */
-TVM_DLL TensorRegion FullTensorRegion(TensorVar tensor);
+TVM_TIRX_DLL TensorRegion FullTensorRegion(TensorVar tensor);
/*! \brief Construct unit or vector-lane ranges from point indices. */
-TVM_DLL TensorRegion TensorRegionFromPoint(TensorVar tensor, ffi::Array<PrimExpr> indices);
+TVM_TIRX_DLL TensorRegion TensorRegionFromPoint(TensorVar tensor, ffi::Array<PrimExpr> indices);
} // namespace tirx
} // namespace tvm
diff --git a/include/tvm/tirx/function.h b/include/tvm/tirx/function.h
index 8185209a07..9880e02f66 100644
--- a/include/tvm/tirx/function.h
+++ b/include/tvm/tirx/function.h
@@ -93,7 +93,7 @@ class FunctionNode : public BaseFuncNode {
* \note The function type annotation of PrimExpr is
* directly derived from the Vars without the need of type inference.
*/
- TVM_DLL FuncType func_type_annotation() const;
+ TVM_TIRX_DLL FuncType func_type_annotation() const;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.Function", FunctionNode, BaseFuncNode);
};
@@ -117,9 +117,9 @@ class Function : public BaseFunc {
*
* \param loc The location of this object in the source code.
*/
- TVM_DLL Function(ffi::Array<tvm::Var> params, ffi::Optional<SeqStmt> body,
- Type ret_type = VoidType(), DictAttrs attrs = DictAttrs(),
- Location loc = UnknownLoc());
+ TVM_TIRX_DLL Function(ffi::Array<tvm::Var> params, ffi::Optional<SeqStmt> body,
+ Type ret_type = VoidType(), DictAttrs attrs = DictAttrs(),
+ Location loc = UnknownLoc());
explicit Function(ffi::ObjectPtr<FunctionNode> node) : BaseFunc(std::move(node)) {}
@@ -136,7 +136,7 @@ class Function : public BaseFunc {
* \param func The function to renew.
* \return The function with fresh variable definitions.
*/
-TVM_DLL Function RenewDef(Function func);
+TVM_TIRX_DLL Function RenewDef(Function func);
/*!
* \brief Specialize parameters of Function.
diff --git a/include/tvm/tirx/layout.h b/include/tvm/tirx/layout.h
index c7bb6b99e3..e46e222117 100644
--- a/include/tvm/tirx/layout.h
+++ b/include/tvm/tirx/layout.h
@@ -176,7 +176,7 @@ class Axis : public ffi::Enum {
* \param name Name of the axis to look up or register.
* \return The registered axis singleton.
*/
- TVM_DLL static Axis Get(const ffi::String& name);
+ TVM_TIRX_DLL static Axis Get(const ffi::String& name);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Axis, ffi::Enum, AxisNode);
};
@@ -206,7 +206,7 @@ class IterNode : public ffi::Object {
class Iter : public ffi::ObjectRef {
public:
- TVM_DLL explicit Iter(PrimExpr extent, PrimExpr stride, Axis axis);
+ TVM_TIRX_DLL explicit Iter(PrimExpr extent, PrimExpr stride, Axis axis);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Iter, ffi::ObjectRef, IterNode);
};
@@ -301,8 +301,8 @@ class TileLayoutNode : public LayoutNode {
class TileLayout : public Layout {
public:
- TVM_DLL explicit TileLayout(ffi::Array<Iter> shard, ffi::Array<Iter> replica,
- ffi::Map<Axis, PrimExpr> offset);
+ TVM_TIRX_DLL explicit TileLayout(ffi::Array<Iter> shard, ffi::Array<Iter> replica,
+ ffi::Map<Axis, PrimExpr> offset);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TileLayout, Layout, TileLayoutNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(TileLayoutNode);
@@ -391,8 +391,8 @@ class ComposeLayoutNode : public LayoutNode {
class ComposeLayout : public Layout {
public:
- TVM_DLL explicit ComposeLayout(int per_element, int swizzle_len, int atom_len,
- TileLayout tile_layout, bool swizzle_inner = true);
+ TVM_TIRX_DLL explicit ComposeLayout(int per_element, int swizzle_len, int atom_len,
+ TileLayout tile_layout, bool swizzle_inner = true);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ComposeLayout, Layout, ComposeLayoutNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(ComposeLayoutNode);
diff --git a/include/tvm/tirx/op/abi.h b/include/tvm/tirx/op/abi.h
index fe5006bf00..d99e8ef7b8 100644
--- a/include/tvm/tirx/op/abi.h
+++ b/include/tvm/tirx/op/abi.h
@@ -38,7 +38,7 @@ namespace tirx {
* - args[0]: func_name, The function name.
* - args[1...]: args, trailing Expr operands.
*/
-TVM_DLL const Op& call_extern_op();
+TVM_TIRX_DLL const Op& call_extern_op();
/*!
* \brief Call a pure external C function with types supplied by the caller.
@@ -47,7 +47,7 @@ TVM_DLL const Op& call_extern_op();
* - args[0]: func_name, The function name.
* - args[1...]: args, trailing Expr operands.
*/
-TVM_DLL const Op& call_pure_extern_op();
+TVM_TIRX_DLL const Op& call_pure_extern_op();
/*!
* \brief Invoke an LLVM intrinsic.
@@ -56,7 +56,7 @@ TVM_DLL const Op& call_pure_extern_op();
* - args[0]: intrin_id, The intrinsic identifier.
* - args[1...]: args, trailing Expr operands.
*/
-TVM_DLL const Op& call_llvm_intrin_op();
+TVM_TIRX_DLL const Op& call_llvm_intrin_op();
/*!
* \brief Invoke a pure LLVM intrinsic.
@@ -65,7 +65,7 @@ TVM_DLL const Op& call_llvm_intrin_op();
* - args[0]: intrin_id, The intrinsic identifier.
* - args[1...]: args, trailing Expr operands.
*/
-TVM_DLL const Op& call_llvm_pure_intrin_op();
+TVM_TIRX_DLL const Op& call_llvm_pure_intrin_op();
/*!
* \brief Invoke a pure GLSL450 SPIR-V intrinsic.
@@ -74,7 +74,7 @@ TVM_DLL const Op& call_llvm_pure_intrin_op();
* - args[0]: intrin_id, The intrinsic identifier.
* - args[1...]: args, trailing PrimExpr operands.
*/
-TVM_DLL const Op& call_spirv_pure_glsl450_op();
+TVM_TIRX_DLL const Op& call_spirv_pure_glsl450_op();
enum TVMStructFieldKind : int {
// DLTensor fields
@@ -108,7 +108,7 @@ enum TVMStructFieldKind : int {
* - args[1]: index, The index.
* - args[2]: field, The field index.
*/
-TVM_DLL const Op& abi_field_get_op();
+TVM_TIRX_DLL const Op& abi_field_get_op();
/*!
* \brief Write a runtime structure field.
@@ -119,12 +119,12 @@ TVM_DLL const Op& abi_field_get_op();
* - args[2]: field, The field index.
* - args[3]: value, The value to use.
*/
-TVM_DLL const Op& abi_field_set_op();
+TVM_TIRX_DLL const Op& abi_field_set_op();
/*!
* \brief Raise the last runtime error.
*/
-TVM_DLL const Op& throw_last_error_op();
+TVM_TIRX_DLL const Op& throw_last_error_op();
/*!
* \brief Allocate stack storage for runtime values.
@@ -133,7 +133,7 @@ TVM_DLL const Op& throw_last_error_op();
* - args[0]: dtype_str, The data type name.
* - args[1]: num, The number of entries.
*/
-TVM_DLL const Op& stack_alloca_op();
+TVM_TIRX_DLL const Op& stack_alloca_op();
/*!
* \brief Construct a stack-allocated shape tuple.
@@ -141,7 +141,7 @@ TVM_DLL const Op& stack_alloca_op();
* Arguments, in order:
* - args[0...]: args, trailing IntExpr operands.
*/
-TVM_DLL const Op& stack_make_shape_op();
+TVM_TIRX_DLL const Op& stack_make_shape_op();
/*!
* \brief Construct a stack-allocated DLTensor.
@@ -154,7 +154,7 @@ TVM_DLL const Op& stack_make_shape_op();
* - args[4]: arr_dtype, The array data type.
* - args[5]: elem_offset, The element offset.
*/
-TVM_DLL const Op& stack_make_dltensor_op();
+TVM_TIRX_DLL const Op& stack_make_dltensor_op();
/*!
* \brief Invoke a runtime packed function.
@@ -163,7 +163,7 @@ TVM_DLL const Op& stack_make_dltensor_op();
* - args[0]: func_name, The function name.
* - args[1...]: args, trailing Expr operands.
*/
-TVM_DLL const Op& call_packed_op();
+TVM_TIRX_DLL const Op& call_packed_op();
/*! \brief Launch metadata for call_ffi_kernel. */
struct CallFFIKernelAttr : public AttrsNode {
@@ -193,7 +193,7 @@ struct CallFFIKernelAttr : public AttrsNode {
* CallFFIKernelAttr::launch_params describes the launch-value suffix.
* Host backends may consume this call directly or lower it to call_packed.
*/
-TVM_DLL const Op& call_ffi_kernel_op();
+TVM_TIRX_DLL const Op& call_ffi_kernel_op();
/*!
* \brief Invoke a C packed function.
@@ -202,7 +202,7 @@ TVM_DLL const Op& call_ffi_kernel_op();
* - args[0]: func_name, The function name.
* - args[1...]: args, trailing Expr operands.
*/
-TVM_DLL const Op& call_cpacked_op();
+TVM_TIRX_DLL const Op& call_cpacked_op();
/*!
* \brief Invoke a packed function using an explicit argument stack.
@@ -213,7 +213,7 @@ TVM_DLL const Op& call_cpacked_op();
* - args[2]: begin, The start index.
* - args[3]: end, The end index.
*/
-TVM_DLL const Op& call_packed_lowered_op();
+TVM_TIRX_DLL const Op& call_packed_lowered_op();
/*!
* \brief Invoke a C packed function using an explicit argument stack.
@@ -224,7 +224,7 @@ TVM_DLL const Op& call_packed_lowered_op();
* - args[2]: begin, The start index.
* - args[3]: end, The end index.
*/
-TVM_DLL const Op& call_cpacked_lowered_op();
+TVM_TIRX_DLL const Op& call_cpacked_lowered_op();
} // namespace tirx
} // namespace tvm
diff --git a/include/tvm/tirx/op/annotation.h b/include/tvm/tirx/op/annotation.h
index 27a96edf56..3276909c1e 100644
--- a/include/tvm/tirx/op/annotation.h
+++ b/include/tvm/tirx/op/annotation.h
@@ -33,7 +33,7 @@ namespace tirx {
/*!
* \brief Marker where a transform should replace generated kernel initialization.
*/
-TVM_DLL const Op& kernel_replace_point_op();
+TVM_TIRX_DLL const Op& kernel_replace_point_op();
/*!
* \brief Annotate a predicate not be considered as target condition of loop partition.
@@ -41,7 +41,7 @@ TVM_DLL const Op& kernel_replace_point_op();
* Arguments, in order:
* - args[0]: predicate, The predicate.
*/
-TVM_DLL const Op& ignore_loop_partition_op();
+TVM_TIRX_DLL const Op& ignore_loop_partition_op();
} // namespace tirx
} // namespace tvm
diff --git a/include/tvm/tirx/op/cpu.h b/include/tvm/tirx/op/cpu.h
index d5238a0472..28360aa039 100644
--- a/include/tvm/tirx/op/cpu.h
+++ b/include/tvm/tirx/op/cpu.h
@@ -36,7 +36,7 @@ namespace tirx {
* Every worker must reach this operation. It must be outside partitioned
* parallel loops, whose iteration counts can differ between workers.
*/
-TVM_DLL const Op& cpu_parallel_barrier_op();
+TVM_TIRX_DLL const Op& cpu_parallel_barrier_op();
} // namespace tirx
} // namespace tvm
diff --git a/include/tvm/tirx/op/gpu.h b/include/tvm/tirx/op/gpu.h
index 79c9046caf..cedb485d08 100644
--- a/include/tvm/tirx/op/gpu.h
+++ b/include/tvm/tirx/op/gpu.h
@@ -33,7 +33,7 @@ namespace tirx {
/*!
* \brief Return from a GPU thread without returning a function value.
*/
-TVM_DLL const Op& gpu_thread_return_op();
+TVM_TIRX_DLL const Op& gpu_thread_return_op();
/*!
* \brief Return from a GPU thread without returning a function value.
@@ -41,7 +41,7 @@ TVM_DLL const Op& gpu_thread_return_op();
* \param loc The location of this operation in the source.
* \return The thread return expression.
*/
-TVM_DLL PrimExpr gpu_thread_return(Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr gpu_thread_return(Location loc = UnknownLoc());
/*!
* \brief Mark an opaque predicate that selects one thread on an active-set axis.
@@ -51,7 +51,7 @@ TVM_DLL PrimExpr gpu_thread_return(Location loc = UnknownLoc());
* - args[1]: pred, The runtime predicate used as an If condition.
* Canonical thread predicates can be used directly without this wrapper.
*/
-TVM_DLL const Op& gpu_thread_filter_op();
+TVM_TIRX_DLL const Op& gpu_thread_filter_op();
/*!
* \brief Analysis-only active-thread selector.
@@ -61,7 +61,7 @@ TVM_DLL const Op& gpu_thread_filter_op();
* ExecContext/DispatchContext metadata, for predicates such as
* ``ptx.elect_sync()`` whose selected lane cannot be inferred structurally.
*/
-TVM_DLL const Op& gpu_active_thread_selector_op();
+TVM_TIRX_DLL const Op& gpu_active_thread_selector_op();
/*!
* \brief Mark a condition to be thread invariant.
@@ -69,7 +69,7 @@ TVM_DLL const Op& gpu_active_thread_selector_op();
* Arguments, in order:
* - args[0]: cond, The condition.
*/
-TVM_DLL const Op& gpu_thread_invariant_op();
+TVM_TIRX_DLL const Op& gpu_thread_invariant_op();
/*!
* \brief Synchronize accesses in a storage scope.
@@ -77,7 +77,7 @@ TVM_DLL const Op& gpu_thread_invariant_op();
* Arguments, in order:
* - args[0]: storage_scope, The storage scope.
*/
-TVM_DLL const Op& gpu_storage_sync_op();
+TVM_TIRX_DLL const Op& gpu_storage_sync_op();
/*!
* \brief Read the value from a selected lane in the warp.
@@ -89,7 +89,7 @@ TVM_DLL const Op& gpu_storage_sync_op();
* - args[3]: width, The width.
* - args[4]: warp_size, The number of threads per warp.
*/
-TVM_DLL const Op& gpu_warp_shuffle_op();
+TVM_TIRX_DLL const Op& gpu_warp_shuffle_op();
/*!
* \brief Read the value from a lower lane in the warp.
@@ -101,7 +101,7 @@ TVM_DLL const Op& gpu_warp_shuffle_op();
* - args[3]: width, The width.
* - args[4]: warp_size, The number of threads per warp.
*/
-TVM_DLL const Op& gpu_warp_shuffle_up_op();
+TVM_TIRX_DLL const Op& gpu_warp_shuffle_up_op();
/*!
* \brief Read the value from a higher lane in the warp.
@@ -113,7 +113,7 @@ TVM_DLL const Op& gpu_warp_shuffle_up_op();
* - args[3]: width, The width.
* - args[4]: warp_size, The number of threads per warp.
*/
-TVM_DLL const Op& gpu_warp_shuffle_down_op();
+TVM_TIRX_DLL const Op& gpu_warp_shuffle_down_op();
/*!
* \brief Read the value from the lane selected by an XOR mask.
@@ -125,12 +125,12 @@ TVM_DLL const Op& gpu_warp_shuffle_down_op();
* - args[3]: width, The width.
* - args[4]: warp_size, The number of threads per warp.
*/
-TVM_DLL const Op& gpu_warp_shuffle_xor_op();
+TVM_TIRX_DLL const Op& gpu_warp_shuffle_xor_op();
/*!
* \brief Return the bit mask of active lanes in the warp.
*/
-TVM_DLL const Op& gpu_warp_activemask_op();
+TVM_TIRX_DLL const Op& gpu_warp_activemask_op();
/*!
* \brief Cross-thread reduction with an explicit typed combiner and identities.
@@ -150,7 +150,7 @@ TVM_DLL const Op& gpu_warp_activemask_op();
* Other thread indices remain fixed. The operation writes the reduced values
* to the destination tensors and returns void.
*/
-TVM_DLL const Op& gpu_thread_allreduce_op();
+TVM_TIRX_DLL const Op& gpu_thread_allreduce_op();
/*!
* \brief Dot product of two int8x4 vectors and add an optional accumulator.
@@ -160,7 +160,7 @@ TVM_DLL const Op& gpu_thread_allreduce_op();
* - args[1]: vec2, The second input vector.
* - args[2]: acc, The accumulator.
*/
-TVM_DLL const Op& gpu_dp4a_op();
+TVM_TIRX_DLL const Op& gpu_dp4a_op();
} // namespace tirx
} // namespace tvm
diff --git a/include/tvm/tirx/op/math.h b/include/tvm/tirx/op/math.h
index f4a7e4cbaa..12b1f99c21 100644
--- a/include/tvm/tirx/op/math.h
+++ b/include/tvm/tirx/op/math.h
@@ -37,7 +37,7 @@ namespace tirx {
* \param loc The location of this operation in the source.
* \return The result expression.
*/
-TVM_DLL PrimExpr logaddexp(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
+TVM_TIRX_DLL PrimExpr logaddexp(PrimExpr a, PrimExpr b, Location loc = UnknownLoc());
TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(logaddexp);
@@ -48,7 +48,7 @@ TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(logaddexp);
* \param bits The number of bits in the type.
* \return The constructed expression.
*/
-TVM_DLL PrimExpr fast_erf_float_expr(PrimExpr arg, int bits);
+TVM_TIRX_DLL PrimExpr fast_erf_float_expr(PrimExpr arg, int bits);
} // namespace tirx
} // namespace tvm
diff --git a/include/tvm/tirx/op/memory.h b/include/tvm/tirx/op/memory.h
index 92df1fe024..f5f8b8f4b7 100644
--- a/include/tvm/tirx/op/memory.h
+++ b/include/tvm/tirx/op/memory.h
@@ -37,7 +37,7 @@ namespace tirx {
* DictAttrs directly holds allocation annotations. The TensorType owns all
* shape, dtype, storage scope and access/layout metadata, and is the result type.
*/
-TVM_DLL const Op& alloc_tensor_op();
+TVM_TIRX_DLL const Op& alloc_tensor_op();
/*!
* \brief Declare a tensor view with exactly one TensorType in ty_args.
@@ -45,7 +45,7 @@ TVM_DLL const Op& alloc_tensor_op();
* args[0] is the existing physical pointer backing the view. There are no
* attributes. The TensorType is the result type; no memory is allocated.
*/
-TVM_DLL const Op& decl_tensor_op();
+TVM_TIRX_DLL const Op& decl_tensor_op();
/*!
* \brief Reinterpret the bits of a value as an explicit primitive or pointer type.
@@ -53,7 +53,7 @@ TVM_DLL const Op& decl_tensor_op();
* Arguments, in order:
* - args[0]: x, The input value.
*/
-TVM_DLL const Op& reinterpret_op();
+TVM_TIRX_DLL const Op& reinterpret_op();
/*!
* \brief perform reinterpret cast value to type.
@@ -64,7 +64,7 @@ TVM_DLL const Op& reinterpret_op();
* \return The result expression.
* \note This function may return value if the type is the same.
*/
-TVM_DLL Expr reinterpret(Type target_ty, Expr value, Location loc = UnknownLoc());
+TVM_TIRX_DLL Expr reinterpret(Type target_ty, Expr value, Location loc = UnknownLoc());
inline PrimExpr reinterpret(PrimType t, PrimExpr value, Location loc = UnknownLoc()) {
return reinterpret(Type(std::move(t)), Expr(std::move(value)), std::move(loc))
@@ -82,7 +82,7 @@ inline Expr ConstHandle(int64_t value, Location loc = UnknownLoc()) {
* Arguments, in order:
* - args[0]: obj, The referenced object.
*/
-TVM_DLL const Op& address_of_op();
+TVM_TIRX_DLL const Op& address_of_op();
/*!
* \brief Test whether a pointer is null.
@@ -90,7 +90,7 @@ TVM_DLL const Op& address_of_op();
* Arguments, in order:
* - args[0]: x, The input value.
*/
-TVM_DLL const Op& isnullptr_op();
+TVM_TIRX_DLL const Op& isnullptr_op();
/*!
* \brief Prefetch a memory address.
@@ -101,7 +101,7 @@ TVM_DLL const Op& isnullptr_op();
* - args[2]: locality, The locality hint.
* - args[3]: cache_type, The cache policy.
*/
-TVM_DLL const Op& prefetch_op();
+TVM_TIRX_DLL const Op& prefetch_op();
/*!
* \brief Cast a handle to a typed pointer after adding a byte offset.
@@ -111,12 +111,12 @@ TVM_DLL const Op& prefetch_op();
* - args[1]: byte_offset, Offset in bytes.
* The result type specifies the pointed-to element type.
*/
-TVM_DLL const Op& ptr_byte_offset_op();
+TVM_TIRX_DLL const Op& ptr_byte_offset_op();
/*!
* \brief Create a function local static handle that iniitalizes to nullptr.
*/
-TVM_DLL const Op& static_handle_op();
+TVM_TIRX_DLL const Op& static_handle_op();
/*!
* \brief Add a byte offset to an opaque pointer.
@@ -125,7 +125,7 @@ TVM_DLL const Op& static_handle_op();
* - args[0]: handle, The handle.
* - args[1]: offset, The offset.
*/
-TVM_DLL const Op& handle_add_byte_offset_op();
+TVM_TIRX_DLL const Op& handle_add_byte_offset_op();
/*!
* \brief atomic add instruction, corresponding e.g. to atomicAdd in CUDA.
@@ -134,7 +134,7 @@ TVM_DLL const Op& handle_add_byte_offset_op();
* - args[0]: ptr, The pointer.
* - args[1]: value, The value to use.
*/
-TVM_DLL const Op& atomic_add_op();
+TVM_TIRX_DLL const Op& atomic_add_op();
/*!
* \brief Assume a tensor's base address has the given constant byte alignment.
@@ -142,12 +142,12 @@ TVM_DLL const Op& atomic_add_op();
* This leaf operation carries a compiler fact, without checking or changing
* the address. Alignment must be a power of two between 1 and 2^27 bytes.
*/
-TVM_DLL const Op& assume_aligned_op();
+TVM_TIRX_DLL const Op& assume_aligned_op();
/*!
* \brief Returns an initialized but arbitrary value.
*/
-TVM_DLL const Op& undef_op();
+TVM_TIRX_DLL const Op& undef_op();
/*!
* \brief Masked tensor load.
@@ -157,7 +157,7 @@ TVM_DLL const Op& undef_op();
* - args[1]: index, The index.
* - args[2...]: args, trailing PrimExpr operands.
*/
-TVM_DLL const Op& masked_load_op();
+TVM_TIRX_DLL const Op& masked_load_op();
/*!
* \brief Masked tensor store.
@@ -168,7 +168,7 @@ TVM_DLL const Op& masked_load_op();
* - args[2]: index, The index.
* - args[3...]: args, trailing PrimExpr operands.
*/
-TVM_DLL const Op& masked_store_op();
+TVM_TIRX_DLL const Op& masked_store_op();
/*!
* \brief Project the physical pointer associated with a TensorVar definition.
@@ -178,7 +178,7 @@ TVM_DLL const Op& masked_store_op();
* scope of the sole TensorVar argument. This operation is consumed by TIRx
* lowering and code generation.
*/
-TVM_DLL const Op& tensor_data_ptr_op();
+TVM_TIRX_DLL const Op& tensor_data_ptr_op();
} // namespace tirx
} // namespace tvm
diff --git a/include/tvm/tirx/op/region.h b/include/tvm/tirx/op/region.h
index e54f9045ba..f405475235 100644
--- a/include/tvm/tirx/op/region.h
+++ b/include/tvm/tirx/op/region.h
@@ -36,32 +36,32 @@ namespace tirx {
* The sole body parameter is a fresh thread-index PrimVar matching the extent type.
* Tags starting with vthread denote virtual threads. There are no attrs or results.
*/
-TVM_DLL const Op& launch_thread_op();
+TVM_TIRX_DLL const Op& launch_thread_op();
/*!
* \brief Mark a user-facing device entry containing device scope definitions.
* Takes no operands, body parameters, attributes or results.
*/
-TVM_DLL const Op& device_entry_op();
+TVM_TIRX_DLL const Op& device_entry_op();
/*!
* \brief Supply lexical device context for allocation and packed-call lowering.
* Operands are integer device type and device ID; there are no body parameters,
* attributes or results. The region does not change the active runtime device.
*/
-TVM_DLL const Op& device_context_op();
+TVM_TIRX_DLL const Op& device_context_op();
/*!
* \brief Outline the body as a CPU compute helper named by a StringImm operand.
* There are no body parameters, attributes or results.
*/
-TVM_DLL const Op& compute_scope_op();
+TVM_TIRX_DLL const Op& compute_scope_op();
/*!
* \brief Launch a CPU worker team around parallel loops and team barriers.
* Takes no operands, body parameters, attributes or results.
*/
-TVM_DLL const Op& parallel_launch_op();
+TVM_TIRX_DLL const Op& parallel_launch_op();
} // namespace tirx
} // namespace tvm
diff --git a/include/tvm/tirx/op/vector.h b/include/tvm/tirx/op/vector.h
index 4a5e8036c4..f5d7b05acc 100644
--- a/include/tvm/tirx/op/vector.h
+++ b/include/tvm/tirx/op/vector.h
@@ -36,7 +36,7 @@ namespace tirx {
* Arguments, in order:
* - args[0]: vec, The input vector.
*/
-TVM_DLL const Op& vector_high_op();
+TVM_TIRX_DLL const Op& vector_high_op();
/*!
* \brief Get the low-level half of the vector.
@@ -44,7 +44,7 @@ TVM_DLL const Op& vector_high_op();
* Arguments, in order:
* - args[0]: vec, The input vector.
*/
-TVM_DLL const Op& vector_low_op();
+TVM_TIRX_DLL const Op& vector_low_op();
/*!
* \brief Concat two vectors.
@@ -53,7 +53,7 @@ TVM_DLL const Op& vector_low_op();
* - args[0]: vec1, The first input vector.
* - args[1]: vec2, The second input vector.
*/
-TVM_DLL const Op& vector_combine_op();
+TVM_TIRX_DLL const Op& vector_combine_op();
/*!
* \brief Calculate a predicate mask given an upper bound (limit) and a current value (base).
@@ -62,7 +62,7 @@ TVM_DLL const Op& vector_combine_op();
* - args[0]: base, The base value.
* - args[1]: limit, The limit value.
*/
-TVM_DLL const Op& get_active_lane_mask_op();
+TVM_TIRX_DLL const Op& get_active_lane_mask_op();
} // namespace tirx
} // namespace tvm
diff --git a/include/tvm/tirx/stmt_functor.h b/include/tvm/tirx/stmt_functor.h
index 4d59cd5536..e706800802 100644
--- a/include/tvm/tirx/stmt_functor.h
+++ b/include/tvm/tirx/stmt_functor.h
@@ -173,7 +173,7 @@ class StmtFunctor<R(const Stmt&, Args...)> {
* Native hooks match exact types. Unregistered OpaqueExprNode subclasses use
* structural traversal; leaf types register non-descending structural hooks.
*/
-class TVM_DLL StmtExprVisitor : public tvm::ExprVisitor {
+class TVM_TIRX_DLL StmtExprVisitor : public tvm::ExprVisitor {
public:
TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(StmtExprVisitor, tvm::ExprVisitor)
@@ -218,7 +218,7 @@ class TVM_DLL StmtExprVisitor : public tvm::ExprVisitor {
* including on in-place writes. Passes that change dtypes or index lanes must
* provide the corresponding TensorLoad, TensorStore, or TensorRegion inference.
*/
-class TVM_DLL StmtExprMutator : public tvm::ExprMutator {
+class TVM_TIRX_DLL StmtExprMutator : public tvm::ExprMutator {
public:
TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(StmtExprMutator, tvm::ExprMutator)
using tvm::ExprMutator::Mutate;
@@ -293,7 +293,7 @@ class TVM_DLL StmtExprMutator : public tvm::ExprMutator {
*
* \return The result.
*/
-TVM_DLL Stmt SubstituteWithDataTypeLegalization(
+TVM_TIRX_DLL Stmt SubstituteWithDataTypeLegalization(
Stmt stmt, std::function<ffi::Optional<PrimExpr>(const Var&)> vmap);
/*!
@@ -305,7 +305,7 @@ TVM_DLL Stmt SubstituteWithDataTypeLegalization(
*
* \return The result.
*/
-TVM_DLL PrimExpr SubstituteWithDataTypeLegalization(
+TVM_TIRX_DLL PrimExpr SubstituteWithDataTypeLegalization(
PrimExpr expr, std::function<ffi::Optional<PrimExpr>(const Var&)> vmap);
/*!
diff --git a/include/tvm/tirx/tile/context.h b/include/tvm/tirx/tile/context.h
index 5a6d7b5c9a..23591e400a 100644
--- a/include/tvm/tirx/tile/context.h
+++ b/include/tvm/tirx/tile/context.h
@@ -123,14 +123,14 @@ class DispatchContextObj : public ffi::Object {
*/
class DispatchContext : public ffi::ObjectRef {
public:
- TVM_DLL DispatchContext(Target target, ExecScope exec_scope,
- ffi::Map<ffi::String, ffi::Tuple<PrimVar, PrimExpr>> launch_params = {},
- ffi::Map<Var, Range> var_range_map = {}, bool alloc_only = false,
- ffi::Map<ffi::String, ffi::ObjectRef> callbacks = {},
- ffi::Map<ffi::String, ffi::ObjectRef> shared_state = {},
- ffi::Map<ffi::String, ffi::Array<PrimExpr>> inter = {},
- ffi::Map<ffi::String, ffi::Array<PrimExpr>> intra = {},
- ffi::String scope_kind = "");
+ TVM_TIRX_DLL DispatchContext(
+ Target target, ExecScope exec_scope,
+ ffi::Map<ffi::String, ffi::Tuple<PrimVar, PrimExpr>> launch_params = {},
+ ffi::Map<Var, Range> var_range_map = {}, bool alloc_only = false,
+ ffi::Map<ffi::String, ffi::ObjectRef> callbacks = {},
+ ffi::Map<ffi::String, ffi::ObjectRef> shared_state = {},
+ ffi::Map<ffi::String, ffi::Array<PrimExpr>> inter = {},
+ ffi::Map<ffi::String, ffi::Array<PrimExpr>> intra = {}, ffi::String scope_kind = "");
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DispatchContext, ffi::ObjectRef,
DispatchContextObj);
diff --git a/include/tvm/tirx/tile/dispatch.h b/include/tvm/tirx/tile/dispatch.h
index d2672ec3f2..9dc30c48b7 100644
--- a/include/tvm/tirx/tile/dispatch.h
+++ b/include/tvm/tirx/tile/dispatch.h
@@ -45,7 +45,7 @@ inline constexpr const char* kTileOpDispatch = "FTileOpDispatch";
* \param context The operator's execution context and accumulated callbacks.
* \return The instruction implementation, whose body is expanded by TileDispatch.
*/
-TVM_DLL Function Dispatch(const Call& call, const DispatchContext& context);
+TVM_TIRX_DLL Function Dispatch(const Call& call, const DispatchContext& context);
} // namespace tvm::tirx::tile
#endif // TVM_TIRX_TILE_DISPATCH_H_
diff --git a/include/tvm/tirx/tile/exec_scope.h b/include/tvm/tirx/tile/exec_scope.h
index 1cc36672a8..541f6391cc 100644
--- a/include/tvm/tirx/tile/exec_scope.h
+++ b/include/tvm/tirx/tile/exec_scope.h
@@ -25,7 +25,7 @@
namespace tvm::tirx {
/*! \brief Register a backend's execution hierarchy. Smaller rank means wider scope. */
-TVM_DLL void RegisterExecScope(const ffi::String& backend, const ffi::String& name, int rank);
+TVM_TIRX_DLL void RegisterExecScope(const ffi::String& backend, const ffi::String& name, int rank);
class ExecScopeNode : public ffi::Object {
public:
@@ -46,12 +46,12 @@ class ExecScopeNode : public ffi::Object {
class ExecScope : public ffi::ObjectRef {
public:
- TVM_DLL ExecScope(ffi::String backend, ffi::String name);
+ TVM_TIRX_DLL ExecScope(ffi::String backend, ffi::String name);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ExecScope, ffi::ObjectRef, ExecScopeNode);
};
/*! \brief Compare scopes within the same backend's hierarchy. */
-TVM_DLL bool ScopeHigher(const ExecScope& a, const ExecScope& b);
+TVM_TIRX_DLL bool ScopeHigher(const ExecScope& a, const ExecScope& b);
} // namespace tvm::tirx
#endif // TVM_TIRX_TILE_EXEC_SCOPE_H_
diff --git a/include/tvm/tirx/tile/op.h b/include/tvm/tirx/tile/op.h
index 4ceda3fac8..d235d73e00 100644
--- a/include/tvm/tirx/tile/op.h
+++ b/include/tvm/tirx/tile/op.h
@@ -31,7 +31,7 @@ inline constexpr const char* kTensorCallValidate = "FTensorCallValidate";
} // namespace op_attr
/*! \brief Attach the native Call contract and backend validator to a tile instruction. */
-TVM_DLL void ConfigureTensorInstruction(Op op, ffi::Function validator);
+TVM_TIRX_DLL void ConfigureTensorInstruction(Op op, ffi::Function validator);
} // namespace tvm::tirx::tile
#endif // TVM_TIRX_TILE_OP_H_
diff --git a/include/tvm/tirx/transform.h b/include/tvm/tirx/transform.h
index e53a0e051a..aaf1b1a31a 100644
--- a/include/tvm/tirx/transform.h
+++ b/include/tvm/tirx/transform.h
@@ -46,6 +46,9 @@ using tvm::transform::PassInfoObj;
using tvm::transform::PassObj;
using tvm::transform::Sequential;
+/*! \brief Wrap a function body in a device scope region. */
+TVM_TIRX_DLL Pass DecorateDeviceScope();
+
/*
* \brief Create a function pass that optimizes Functions.
*
@@ -55,7 +58,7 @@ using tvm::transform::Sequential;
*
* \return The created function pass.
*/
-TVM_DLL Pass CreateFunctionPass(
+TVM_TIRX_DLL Pass CreateFunctionPass(
std::function<ffi::Optional<Function>(Function, IRModule, PassContext)> pass_func,
int opt_level, ffi::String name);
@@ -71,7 +74,7 @@ TVM_DLL Pass CreateFunctionPass(
*
* \return The pass.
*/
-TVM_DLL Pass VectorizeLoop(bool enable_vectorize = true);
+TVM_TIRX_DLL Pass VectorizeLoop(bool enable_vectorize = true);
/*!
* \brief Rewrite storage allocation pattern.
@@ -81,7 +84,7 @@ TVM_DLL Pass VectorizeLoop(bool enable_vectorize = true);
*
* \return The pass.
*/
-TVM_DLL Pass StorageRewrite();
+TVM_TIRX_DLL Pass StorageRewrite();
/*!
* \brief unroll the constant loop marked by unroll.
@@ -89,33 +92,33 @@ TVM_DLL Pass StorageRewrite();
*
* \return The pass.
*/
-TVM_DLL Pass UnrollLoop();
+TVM_TIRX_DLL Pass UnrollLoop();
/*!
* \brief Remove No Op from the Stmt.
*
* \return The pass.
*/
-TVM_DLL Pass RemoveNoOp();
+TVM_TIRX_DLL Pass RemoveNoOp();
/*!
* \brief Hoist loop-invariant If statements outside eligible loops.
* \return The pass.
*/
-TVM_DLL Pass HoistIf();
+TVM_TIRX_DLL Pass HoistIf();
/*!
* \brief Renormalize the split pattern from floordiv(floormod()) to floormod(floordiv()).
* \return The pass.
*/
-TVM_DLL Pass RenormalizeSplitPattern();
+TVM_TIRX_DLL Pass RenormalizeSplitPattern();
/*!
* \brief Run statement-level arithmetic simplifications on the TIR Function.
*
* \return The pass.
*/
-TVM_DLL Pass StmtSimplify();
+TVM_TIRX_DLL Pass StmtSimplify();
/*!
* \brief Convert an IRModule to be SSA form.
@@ -128,10 +131,10 @@ TVM_DLL Pass StmtSimplify();
*
* \return The pass.
*/
-TVM_DLL Pass ConvertSSA();
+TVM_TIRX_DLL Pass ConvertSSA();
/*! \brief Lower cross-thread reductions to target-specific communication. */
-TVM_DLL Pass LowerThreadAllreduce();
+TVM_TIRX_DLL Pass LowerThreadAllreduce();
/*!
* \brief Transform the high-level Function to a low-level version
@@ -159,7 +162,7 @@ TVM_DLL Pass LowerThreadAllreduce();
*
* \return The pass.
*/
-TVM_DLL Pass MakePackedAPI();
+TVM_TIRX_DLL Pass MakePackedAPI();
/*!
* \brief Annotate, split, and lower host/device functions.
@@ -179,14 +182,14 @@ TVM_DLL Pass MakePackedAPI();
*
* \return The pass.
*/
-TVM_DLL Pass SplitHostDevice();
+TVM_TIRX_DLL Pass SplitHostDevice();
/*!
* \brief skip assert stmt.
*
* \return The pass.
*/
-TVM_DLL Pass SkipAssert();
+TVM_TIRX_DLL Pass SkipAssert();
/*!
* \brief This annotation is for nodes to be disabled for builtin lowering
@@ -197,20 +200,20 @@ static constexpr const char* kDisableLowerTVMBuiltin = "disable_lower_builtin";
* \brief Lower builtin intrinsics.
* \return The pass.
*/
-TVM_DLL Pass LowerTVMBuiltin();
+TVM_TIRX_DLL Pass LowerTVMBuiltin();
/*!
* \brief Lower the target specific function intrinsics in each of the function.
*
* \return The pass.
*/
-TVM_DLL Pass LowerIntrin();
+TVM_TIRX_DLL Pass LowerIntrin();
/*!
* \brief Lower warp memory access to low-level device related function calls.
* \return The pass.
*/
-TVM_DLL Pass LowerWarpMemory();
+TVM_TIRX_DLL Pass LowerWarpMemory();
/*!
* \brief Narrow down PrimExpr datatype in stmt to target_bits.
@@ -220,7 +223,7 @@ TVM_DLL Pass LowerWarpMemory();
* \note Run this pass after storage flatten.
* \return The pass.
*/
-TVM_DLL Pass NarrowDataType(int target_bits);
+TVM_TIRX_DLL Pass NarrowDataType(int target_bits);
/*!
* \brief Force to narrow down indexing expressions and integer tensors to int32 dtype.
@@ -231,14 +234,14 @@ TVM_DLL Pass NarrowDataType(int target_bits);
* \return The pass.
* \note This pass should not be used in default cases.
*/
-TVM_DLL Pass ForceNarrowIndexToInt32();
+TVM_TIRX_DLL Pass ForceNarrowIndexToInt32();
/*!
* \brief Legalize bf16 compute Ops. Add a cast to fp32
* before Ops, then add a cast back to bf16.
* \return The pass.
*/
-TVM_DLL Pass BF16ComputeLegalize();
+TVM_TIRX_DLL Pass BF16ComputeLegalize();
/*!
* \brief Legalize fp8 compute Ops. Add a cast to fp16/fp32
@@ -247,20 +250,20 @@ TVM_DLL Pass BF16ComputeLegalize();
* \note Must be run after BindTarget, as it relies on target attributes for Functions
* \return The pass.
*/
-TVM_DLL Pass FP8ComputeLegalize(ffi::String promote_dtype = "float16");
+TVM_TIRX_DLL Pass FP8ComputeLegalize(ffi::String promote_dtype = "float16");
/*!
* \brief Legalize bf16 storage types to u16.
* \return The pass.
*/
-TVM_DLL Pass BF16StorageLegalize();
+TVM_TIRX_DLL Pass BF16StorageLegalize();
/*!
* \brief Legalize fp8 storage types to u8.
* \note Must be run after BindTarget, as it relies on target attributes for Functions
* \return The pass.
*/
-TVM_DLL Pass FP8StorageLegalize();
+TVM_TIRX_DLL Pass FP8StorageLegalize();
/*!
* \brief Rewrite the pointer content type of arguments,
@@ -270,21 +273,21 @@ TVM_DLL Pass FP8StorageLegalize();
*
* \return The pass.
*/
-TVM_DLL Pass PointerValueTypeRewrite();
+TVM_TIRX_DLL Pass PointerValueTypeRewrite();
/*!
* \brief Flatten the multi-dimensional TensorLoad and TensorStore to single dimensional
* TensorLoad/TensorStore for the TIR not contains opaque block.
* \return The pass.
*/
-TVM_DLL Pass FlattenTensor();
+TVM_TIRX_DLL Pass FlattenTensor();
/*!
* \brief Implements Common Subexpression Elimination (CSE) for TIR
* which introduces Bind statements for duplicated sub-expressions.
* \return The pass.
*/
-TVM_DLL Pass CommonSubexprElim();
+TVM_TIRX_DLL Pass CommonSubexprElim();
/*!
* \brief Annotate a Function with a given target.
@@ -292,19 +295,19 @@ TVM_DLL Pass CommonSubexprElim();
* Lower dialect-specific thread bindings before classifying host/device calls.
* \return The pass.
*/
-TVM_DLL Pass BindTarget(Target target);
+TVM_TIRX_DLL Pass BindTarget(Target target);
/*!
* \brief Set a Function as the entry point if it is only function in IRModule.
* \return The pass.
*/
-TVM_DLL Pass AnnotateEntryFunc();
+TVM_TIRX_DLL Pass AnnotateEntryFunc();
/*!
* \brief Filter Functions with a given condition.
* \return The pass.
*/
-TVM_DLL Pass Filter(ffi::TypedFunction<bool(Function)> fcond);
+TVM_TIRX_DLL Pass Filter(ffi::TypedFunction<bool(Function)> fcond);
/*!
* \brief Lower TIRx op calls using registered op dispatchers for the given target.
@@ -313,19 +316,19 @@ TVM_DLL Pass Filter(ffi::TypedFunction<bool(Function)> fcond);
* launch configuration, preserving execution-context information during dispatch.
* \return The pass.
*/
-TVM_DLL Pass TileDispatch();
+TVM_TIRX_DLL Pass TileDispatch();
/*!
* \brief Finalize TIRx lowering by applying layout rewriters and cleanup passes.
* \return The pass.
*/
-TVM_DLL Pass LowerTIRxCleanup();
+TVM_TIRX_DLL Pass LowerTIRxCleanup();
/*!
* \brief Lower the TIR to a lower level IR for the given target.
* \return The pass.
*/
-TVM_DLL Pass LowerTIRx();
+TVM_TIRX_DLL Pass LowerTIRx();
} // namespace transform
} // namespace tirx
diff --git a/include/tvm/tirx/type.h b/include/tvm/tirx/type.h
index b3c447d70d..5829310523 100644
--- a/include/tvm/tirx/type.h
+++ b/include/tvm/tirx/type.h
@@ -121,10 +121,10 @@ class TensorTypeNode : public TypeNode {
PtrType DataPtrType() const { return PtrType(dtype, storage_scope); }
/*! \brief Whether this type supports scalar tensor syntax. */
- TVM_DLL bool IsScalar(bool alloc_or_decl = true) const;
+ TVM_TIRX_DLL bool IsScalar(bool alloc_or_decl = true) const;
/*! \brief Return the constant element count, or nullopt for symbolic extents or overflow. */
- TVM_DLL std::optional<int64_t> ConstantAllocationSize() const;
+ TVM_TIRX_DLL std::optional<int64_t> ConstantAllocationSize() const;
/*! \brief Determine the offset in the tensor of the given index.
*
@@ -145,10 +145,10 @@ class TensorTypeNode : public TypeNode {
*/
class TensorType : public Type {
public:
- TVM_DLL TensorType(ffi::String storage_scope, PrimType dtype, ffi::Array<PrimExpr> shape,
- ffi::Array<PrimExpr> strides, ffi::Optional<PrimExpr> elem_offset,
- int data_alignment, int offset_factor,
- ffi::Optional<Layout> layout = std::nullopt, Location loc = UnknownLoc());
+ TVM_TIRX_DLL TensorType(ffi::String storage_scope, PrimType dtype, ffi::Array<PrimExpr> shape,
+ ffi::Array<PrimExpr> strides, ffi::Optional<PrimExpr> elem_offset,
+ int data_alignment, int offset_factor,
+ ffi::Optional<Layout> layout = std::nullopt, Location loc = UnknownLoc());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TensorType, Type, TensorTypeNode);
@@ -177,7 +177,7 @@ class TensorMapTypeNode : public TypeNode {
*/
class TensorMapType : public Type {
public:
- TVM_DLL TensorMapType(Location loc = UnknownLoc());
+ TVM_TIRX_DLL TensorMapType(Location loc = UnknownLoc());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TensorMapType, Type, TensorMapTypeNode);
};
diff --git a/jvm/core/src/test/scripts/prepare_test_libs.py b/jvm/core/src/test/scripts/prepare_test_libs.py
index 71362fb6c7..aa14514380 100644
--- a/jvm/core/src/test/scripts/prepare_test_libs.py
+++ b/jvm/core/src/test/scripts/prepare_test_libs.py
@@ -53,7 +53,7 @@ def prepare_cpu_lib(base_path):
B = te.placeholder((n,), name="B")
C = te.compute(A.shape, lambda *i: A(*i) + B(*i), name="C")
mod = tvm.IRModule.from_expr(te.create_function([A, B, C]).with_attr("global_symbol", "myadd"))
- fadd = tvm.tirx.build(mod, target)
+ fadd = tvm.tirx.compile(mod, target).mod
lib_path = os.path.join(base_path, "add_cpu.so")
fadd.export_library(lib_path)
@@ -73,7 +73,7 @@ def prepare_gpu_lib(base_path):
i0, i1 = sch.split(i, [None, 32])
sch.bind(i0, "blockIdx.x")
sch.bind(i1, "threadIdx.x")
- fadd = tvm.tirx.build(sch.mod, "cuda")
+ fadd = tvm.tirx.compile(sch.mod, "cuda").mod
lib_path = os.path.join(base_path, "add_cuda.so")
fadd.export_library(lib_path)
diff --git a/python/tvm/backend/cuda/iket.py b/python/tvm/backend/cuda/iket.py
index 3912adcb0d..9b09cf9be5 100644
--- a/python/tvm/backend/cuda/iket.py
+++ b/python/tvm/backend/cuda/iket.py
@@ -316,7 +316,7 @@ class IketProfiler:
if not isinstance(mod, tvm.IRModule):
raise TypeError("IketProfiler.compile expects a TIRx Function or IRModule")
enabled_mod = mod.with_attr("tirx.iket.enabled", True)
- executable = tvm.compile(enabled_mod, target=target, tir_pipeline=tir_pipeline)
+ executable = tvm.tirx.compile(enabled_mod, target=target, pipeline=tir_pipeline)
return _OfficialIketExecutable(executable)
diff --git a/python/tvm/backend/hexagon/__init__.py b/python/tvm/backend/hexagon/__init__.py
index 5d8aa202e5..2db10464f6 100644
--- a/python/tvm/backend/hexagon/__init__.py
+++ b/python/tvm/backend/hexagon/__init__.py
@@ -23,7 +23,7 @@ from tvm_ffi.libinfo import load_lib_ctypes
from tvm.base import _LOADED_LIBS
-_LAZY_SUBMODULES = {"target_tags"}
+_LAZY_SUBMODULES = {"linker", "target_tags"}
def register_backend():
@@ -40,6 +40,7 @@ def register_backend():
except (OSError, FileNotFoundError, RuntimeError):
pass
import_module(f"{__name__}.target_tags")
+ import_module(f"{__name__}.linker")
def __getattr__(name: str):
@@ -48,4 +49,4 @@ def __getattr__(name: str):
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
-__all__ = ["register_backend", "target_tags"]
+__all__ = ["linker", "register_backend", "target_tags"]
diff --git a/python/tvm/backend/hexagon/linker.py b/python/tvm/backend/hexagon/linker.py
new file mode 100644
index 0000000000..0469b32a67
--- /dev/null
+++ b/python/tvm/backend/hexagon/linker.py
@@ -0,0 +1,112 @@
+# 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.
+"""Hexagon shared-library linking and custom linker registration."""
+
+import os
+from pathlib import Path
+
+import tvm_ffi
+
+from tvm.ir import StringImm
+from tvm.support import cc
+
+HEXAGON_TOOLCHAIN = os.environ.get("HEXAGON_TOOLCHAIN", "")
+HEXAGON_LINK_MAIN = Path(HEXAGON_TOOLCHAIN) / "bin" / "hexagon-link"
+
+
+def register_linker(f):
+ """Register a function that will return the path to the Hexagon linker."""
+ return tvm_ffi.register_global_func("tvm.contrib.hexagon.hexagon_link", f, True)
+
+
+@tvm_ffi.register_global_func("tvm.contrib.hexagon.hexagon_link")
+def hexagon_link() -> str:
+ """Return path to the Hexagon linker."""
+ return str(HEXAGON_LINK_MAIN)
+
+
+@tvm_ffi.register_global_func("tvm.contrib.hexagon.link_shared")
+def link_shared(so_name, objs, extra_args=None):
+ """Link shared library on Hexagon using the registered Hexagon linker.
+
+ Parameters
+ ----------
+ so_name : str
+ Name of the shared library file.
+ objs : list[str, tvm.ir.StringImm]
+ extra_args : dict (str->str) or Map<String,String>
+ Additional arguments:
+ 'hex_arch' - Hexagon architecture, e.g. v68
+ 'verbose' - Print additional information if the key is present
+
+ Returns
+ -------
+ ret_val : int
+ This function returns 0 at the moment.
+ """
+
+ # The list of object files can be passed as built-in Python strings,
+ # or as tvm.ir.StringImm's.
+ def to_str(s):
+ if isinstance(s, StringImm):
+ return s.value
+ assert isinstance(s, str), 'argument "' + str(s) + '" should be a string or StrImm'
+ return s
+
+ objs = [to_str(s) for s in objs]
+
+ if not extra_args:
+ extra_args = {}
+ hex_arch = extra_args.get("hex_arch") or "v68"
+ linker = tvm_ffi.get_global_func("tvm.contrib.hexagon.hexagon_link")()
+ if extra_args.get("verbose"):
+ print("tvm.contrib.hexagon.link_shared:")
+ print(" Using linker:", linker)
+ print(" Library name:", so_name)
+ print(" Object files:", objs)
+ print(" Architecture:", hex_arch)
+ if not os.access(linker, os.X_OK):
+ message = 'The linker "' + linker + '" does not exist or is not executable.'
+ if not os.environ.get("HEXAGON_TOOLCHAIN"):
+ message += (
+ " The environment variable HEXAGON_TOOLCHAIN is unset. Please export "
+ + "HEXAGON_TOOLCHAIN in your environment, so that ${HEXAGON_TOOLCHAIN}/bin/"
+ + "hexagon-link exists."
+ )
+ else:
+ message += (
+ " Please verify the value of the HEXAGON_LINKER environment variable "
+ + '(currently set to "'
+ + HEXAGON_TOOLCHAIN
+ + '").'
+ )
+ raise Exception(message)
+
+ libpath = os.path.join(HEXAGON_TOOLCHAIN, "target", "hexagon", "lib", hex_arch, "G0")
+ cc.create_shared(
+ so_name,
+ objs,
+ # pylint: disable=bad-whitespace
+ options=[
+ "-Bdynamic",
+ "-shared",
+ "-export-dynamic",
+ os.path.join(libpath, "pic", "libgcc.so"),
+ ],
+ cc=linker,
+ )
+ return 0
diff --git a/python/tvm/backend/trn/pipeline.py b/python/tvm/backend/trn/pipeline.py
index 9fdb87b28c..0b8daf57ca 100644
--- a/python/tvm/backend/trn/pipeline.py
+++ b/python/tvm/backend/trn/pipeline.py
@@ -18,26 +18,24 @@
import tvm
from tvm import tirx
-from tvm.tirx.compilation_pipeline import _bind_target, _lower_thread_binding, finalize_host_passes
+from tvm.tirx.compilation_pipeline import _bind_target, finalize_host_passes
from . import transform as trn_transform
def trn_pipeline(*, target=None):
- """The Trainium pipeline used in tvm.tirx.build."""
+ """The Trainium pipeline used in tvm.tirx.compile."""
@tvm.transform.module_pass(opt_level=0)
def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.IRModule:
"""Lower TIRx for the Trainium backend."""
- mod = tvm.s_tir.transform.LoopPartition()(mod)
- mod = _lower_thread_binding()(mod)
mod = _bind_target(mod, target)
passes = [
trn_transform.TrnPrivateTensorAlloc(),
trn_transform.TrnNaiveAllocator(),
tirx.transform.TileDispatch(),
trn_transform.LowerTrainiumLayout(),
- tvm.s_tir.transform.DecorateDeviceScope(),
+ tirx.transform.DecorateDeviceScope(),
tirx.transform.StmtSimplify(),
tirx.transform.HoistIf(),
tirx.transform.StmtSimplify(),
diff --git a/python/tvm/support/emcc.py b/python/tvm/contrib/emcc.py
similarity index 97%
rename from python/tvm/support/emcc.py
rename to python/tvm/contrib/emcc.py
index f1455c580b..296f9ee9a2 100644
--- a/python/tvm/support/emcc.py
+++ b/python/tvm/contrib/emcc.py
@@ -62,7 +62,7 @@ def find_wasm_lib(name, optional=False):
return found
-def create_tvmjs_wasm(output, objects, options=None, cc="emcc", libs=None):
+def create_wasm(output, objects, options=None, cc="emcc", libs=None):
"""Create wasm that is supposed to run with the tvmjs.
Parameters
@@ -142,4 +142,4 @@ def create_tvmjs_wasm(output, objects, options=None, cc="emcc", libs=None):
raise RuntimeError(msg)
-create_tvmjs_wasm.object_format = "bc"
+create_wasm.object_format = "bc"
diff --git a/python/tvm/contrib/hexagon/tools.py b/python/tvm/contrib/hexagon/tools.py
index ce5f97c696..c72caede2a 100644
--- a/python/tvm/contrib/hexagon/tools.py
+++ b/python/tvm/contrib/hexagon/tools.py
@@ -31,6 +31,18 @@ from tvm_ffi import register_global_func
import tvm
import tvm.support.cc as cc
+from tvm.backend.hexagon.linker import (
+ HEXAGON_LINK_MAIN as HEXAGON_LINK_MAIN,
+)
+from tvm.backend.hexagon.linker import (
+ hexagon_link as hexagon_link,
+)
+from tvm.backend.hexagon.linker import (
+ link_shared as link_shared,
+)
+from tvm.backend.hexagon.linker import (
+ register_linker as register_linker,
+)
# Linking Hexagon shared libraries.
#
@@ -49,7 +61,6 @@ import tvm.support.cc as cc
HEXAGON_TOOLCHAIN = os.environ.get("HEXAGON_TOOLCHAIN", default="") # pylint: disable=invalid-name
HEXAGON_SDK_ROOT = os.environ.get("HEXAGON_SDK_ROOT", default="") # pylint: disable=invalid-name
HEXAGON_SDK_DOCKER_IMAGE = os.environ.get("HEXAGON_SDK_DOCKER_IMAGE", default="") # pylint: disable=invalid-name
-HEXAGON_LINK_MAIN = pathlib.Path(HEXAGON_TOOLCHAIN) / "bin" / "hexagon-link" # pylint: disable=invalid-name
HEXAGON_CLANG_PLUS = pathlib.Path(HEXAGON_TOOLCHAIN) / "bin" / "hexagon-clang++" # pylint: disable=invalid-name
HEXAGON_SDK_INCLUDE_DIRS = [ # pylint: disable=invalid-name
pathlib.Path(HEXAGON_SDK_ROOT) / "incs",
@@ -59,17 +70,6 @@ HEXAGON_SDK_INCLUDE_DIRS = [ # pylint: disable=invalid-name
HEXAGON_SIMULATOR_NAME = "simulator"
-def register_linker(f):
- """Register a function that will return the path to the Hexagon linker."""
- return register_global_func("tvm.contrib.hexagon.hexagon_link", f, True)
-
-
-@register_global_func("tvm.contrib.hexagon.hexagon_link")
-def hexagon_link() -> str:
- """Return path to the Hexagon linker."""
- return str(HEXAGON_LINK_MAIN)
-
-
def hexagon_clang_plus() -> str:
"""Return path to the Hexagon clang++."""
return str(HEXAGON_CLANG_PLUS)
@@ -106,79 +106,6 @@ def toolchain_version(toolchain=None) -> list[int]:
raise RuntimeError("Cannot establish toolchain version")
-@register_global_func("tvm.contrib.hexagon.link_shared")
-def link_shared(so_name, objs, extra_args=None):
- """Link shared library on Hexagon using the registered Hexagon linker.
-
- Parameters
- ----------
- so_name : str
- Name of the shared library file.
- objs : list[str, tvm.ir.StringImm]
- extra_args : dict (str->str) or Map<String,String>
- Additional arguments:
- 'hex_arch' - Hexagon architecture, e.g. v68
- 'verbose' - Print additional information if the key is present
-
- Returns
- -------
- ret_val : int
- This function returns 0 at the moment.
- """
-
- # The list of object files can be passed as built-in Python strings,
- # or as tvm.ir.StringImm's.
- def to_str(s):
- if isinstance(s, tvm.ir.StringImm):
- return s.value
- assert isinstance(s, str), 'argument "' + str(s) + '" should be a string or StrImm'
- return s
-
- objs = [to_str(s) for s in objs]
-
- if not extra_args:
- extra_args = {}
- hex_arch = extra_args.get("hex_arch") or "v68"
- linker = tvm.get_global_func("tvm.contrib.hexagon.hexagon_link")()
- if extra_args.get("verbose"):
- print("tvm.contrib.hexagon.link_shared:")
- print(" Using linker:", linker)
- print(" Library name:", so_name)
- print(" Object files:", objs)
- print(" Architecture:", hex_arch)
- if not os.access(linker, os.X_OK):
- message = 'The linker "' + linker + '" does not exist or is not executable.'
- if not os.environ.get("HEXAGON_TOOLCHAIN"):
- message += (
- " The environment variable HEXAGON_TOOLCHAIN is unset. Please export "
- + "HEXAGON_TOOLCHAIN in your environment, so that ${HEXAGON_TOOLCHAIN}/bin/"
- + "hexagon-link exists."
- )
- else:
- message += (
- " Please verify the value of the HEXAGON_LINKER environment variable "
- + '(currently set to "'
- + HEXAGON_TOOLCHAIN
- + '").'
- )
- raise Exception(message)
-
- libpath = os.path.join(HEXAGON_TOOLCHAIN, "target", "hexagon", "lib", hex_arch, "G0")
- cc.create_shared(
- so_name,
- objs,
- # pylint: disable=bad-whitespace
- options=[
- "-Bdynamic",
- "-shared",
- "-export-dynamic",
- os.path.join(libpath, "pic", "libgcc.so"),
- ],
- cc=linker,
- )
- return 0
-
-
def link_shared_macos(so_name, objs, extra_args=None):
"""Link Hexagon shared library using docker container with proper tooling.
@@ -245,7 +172,6 @@ if sys.platform == "darwin":
register_global_func("tvm.contrib.hexagon.link_shared", f=link_shared_macos, override=True)
else: # Linux and Win32
create_shared = cc.create_shared
- register_global_func("tvm.contrib.hexagon.link_shared", f=link_shared, override=True)
def create_aot_shared(so_name: str | pathlib.Path, files, hexagon_arch: str, options=None):
diff --git a/python/tvm/contrib/tvmjs.py b/python/tvm/contrib/tvmjs.py
index e06035bc4a..3c64a3917d 100644
--- a/python/tvm/contrib/tvmjs.py
+++ b/python/tvm/contrib/tvmjs.py
@@ -38,7 +38,9 @@ except ImportError:
import tvm
from tvm.runtime import DataType
-from tvm.support.emcc import create_tvmjs_wasm, find_wasm_lib
+
+from .emcc import create_wasm as create_tvmjs_wasm
+from .emcc import find_wasm_lib
def _convert_f32_to_bf16(value):
diff --git a/python/tvm/driver/build_module.py b/python/tvm/driver/build_module.py
index 8dceee2413..3e30b1a2f6 100644
--- a/python/tvm/driver/build_module.py
+++ b/python/tvm/driver/build_module.py
@@ -75,5 +75,4 @@ def compile( # pylint: disable=redefined-builtin
relax_pipeline=relax_pipeline,
tir_pipeline=tir_pipeline,
)
- lib = tvm.tirx.build(mod, target, pipeline=tir_pipeline)
- return Executable(lib)
+ return tvm.tirx.compile(mod, target, pipeline=tir_pipeline)
diff --git a/python/tvm/ir/_tensor_expr_overload.py b/python/tvm/ir/_tensor_expr_overload.py
index a08c2b9687..eaa676d619 100644
--- a/python/tvm/ir/_tensor_expr_overload.py
+++ b/python/tvm/ir/_tensor_expr_overload.py
@@ -107,3 +107,40 @@ def __call__(_value, *_args, **_kwargs):
def astype(_value, _dtype, _loc=None):
return NotImplemented
+
+
+def register_tensor_expr_operators(operators):
+ """Install a dialect's implementations of the shared tensor operator hooks.
+
+ Native scalar operations remain owned by their types. Tensor hooks return
+ NotImplemented for types outside the registering dialect.
+ """
+ supported = (
+ "__add__",
+ "__radd__",
+ "__sub__",
+ "__rsub__",
+ "__mul__",
+ "__rmul__",
+ "__div__",
+ "__rdiv__",
+ "__truediv__",
+ "__rtruediv__",
+ "__floordiv__",
+ "__rfloordiv__",
+ "__mod__",
+ "__rmod__",
+ "__pow__",
+ "__rpow__",
+ "__neg__",
+ "__lt__",
+ "__le__",
+ "__gt__",
+ "__ge__",
+ "__call__",
+ "astype",
+ )
+ for name, operator in operators.items():
+ if name not in supported or not callable(operator):
+ raise ValueError(f"Invalid tensor expression operator: {name}")
+ globals().update(operators)
diff --git a/python/tvm/ir/base.py b/python/tvm/ir/base.py
index 485c763b3a..11c8d226ed 100644
--- a/python/tvm/ir/base.py
+++ b/python/tvm/ir/base.py
@@ -30,6 +30,20 @@ class Scriptable:
"""Marker for display methods installed by the TVMScript owner."""
+def register_script_printer(*, script, show, node_repr):
+ """Attach shared printer behavior without replacing native IR wrappers."""
+ Scriptable.script = script
+ Scriptable.show = show
+ Node.__repr__ = node_repr
+
+
+def register_dialect_script_method(name, method):
+ """Register a dialect-owned private printing convenience method."""
+ if not name.startswith("_") or not name.endswith("_script") or not callable(method):
+ raise ValueError(f"Invalid dialect script method: {name}")
+ setattr(Scriptable, name, method)
+
+
class Node(Object):
"""Base class of all IR Nodes."""
diff --git a/python/tvm/relax/__init__.py b/python/tvm/relax/__init__.py
index 281cd7d994..110e347fd1 100644
--- a/python/tvm/relax/__init__.py
+++ b/python/tvm/relax/__init__.py
@@ -138,3 +138,5 @@ _register_namespace_initializer(_initialize_script_namespace, aliases=("relax",)
# Register the exported operator printer names for compiler diagnostics as well.
from .script import ir_builder as _script_ir_builder
+
+from .script import printer as _script_printer
diff --git a/python/tvm/relax/expr.py b/python/tvm/relax/expr.py
index cfafd91bc6..dbcf98b485 100644
--- a/python/tvm/relax/expr.py
+++ b/python/tvm/relax/expr.py
@@ -95,7 +95,6 @@ Type.is_base_of = _relax_type_is_base_of # type: ignore[attr-defined]
# will be registered afterwards in python/tvm/relax/op/init.py
-_op_ffi_api = None # pylint: disable=invalid-name
def _is_tensor_or_missing_type(ty: Type) -> bool:
diff --git a/python/tvm/relax/op/__init__.py b/python/tvm/relax/op/__init__.py
index 443343dd0c..72302aa00b 100644
--- a/python/tvm/relax/op/__init__.py
+++ b/python/tvm/relax/op/__init__.py
@@ -257,8 +257,6 @@ def _register_op_make():
from tvm.ir import _tensor_expr_overload
from . import _ffi_api
- expr._op_ffi_api = _ffi_api # type: ignore
-
def _add(lhs, rhs):
if isinstance(lhs.ty, expr.tvm.relax.TupleType) and isinstance(rhs, tuple):
return tuple([*lhs, *rhs])
@@ -281,47 +279,37 @@ def _register_op_make():
return NotImplemented
return expr.tvm.ir.Call(func, args, attrs=attrs)
- _tensor_expr_overload.astype = lambda lhs, dtype, _loc=None: (
- _ffi_api.astype(lhs, dtype) if expr._is_tensor_or_missing_type(lhs.ty) else NotImplemented
- )
- _tensor_expr_overload.__call__ = _call
- _tensor_expr_overload.__neg__ = lambda lhs: _unary(lhs, _ffi_api.negative)
- _tensor_expr_overload.__lt__ = lambda lhs, rhs: expr._binary_op_helper(lhs, rhs, _ffi_api.less)
- _tensor_expr_overload.__le__ = lambda lhs, rhs: expr._binary_op_helper(
- lhs, rhs, _ffi_api.less_equal
- )
- _tensor_expr_overload.__gt__ = lambda lhs, rhs: expr._binary_op_helper(
- lhs, rhs, _ffi_api.greater
- )
- _tensor_expr_overload.__ge__ = lambda lhs, rhs: expr._binary_op_helper(
- lhs, rhs, _ffi_api.greater_equal
+ operators = {}
+ operators["astype"] = (
+ lambda lhs, dtype, _loc=None: _ffi_api.astype(lhs, dtype)
+ if expr._is_tensor_or_missing_type(lhs.ty)
+ else NotImplemented
)
- _tensor_expr_overload.__add__ = _add
- _tensor_expr_overload.__radd__ = _add
- _tensor_expr_overload.__sub__ = lambda lhs, rhs: expr._binary_op_helper(
- lhs, rhs, _ffi_api.subtract
- )
- _tensor_expr_overload.__rsub__ = _rhs
- _tensor_expr_overload.__mul__ = lambda lhs, rhs: expr._binary_op_helper(
- lhs, rhs, _ffi_api.multiply
- )
- _tensor_expr_overload.__rmul__ = _tensor_expr_overload.__mul__
- _tensor_expr_overload.__div__ = lambda lhs, rhs: expr._binary_op_helper(
- lhs, rhs, _ffi_api.divide
- )
- _tensor_expr_overload.__rdiv__ = _rhs
- _tensor_expr_overload.__truediv__ = _tensor_expr_overload.__div__
- _tensor_expr_overload.__rtruediv__ = _rhs
- _tensor_expr_overload.__floordiv__ = lambda lhs, rhs: expr._binary_op_helper(
+ operators["__call__"] = _call
+ operators["__neg__"] = lambda lhs: _unary(lhs, _ffi_api.negative)
+ operators["__lt__"] = lambda lhs, rhs: expr._binary_op_helper(lhs, rhs, _ffi_api.less)
+ operators["__le__"] = lambda lhs, rhs: expr._binary_op_helper(lhs, rhs, _ffi_api.less_equal)
+ operators["__gt__"] = lambda lhs, rhs: expr._binary_op_helper(lhs, rhs, _ffi_api.greater)
+ operators["__ge__"] = lambda lhs, rhs: expr._binary_op_helper(lhs, rhs, _ffi_api.greater_equal)
+ operators["__add__"] = _add
+ operators["__radd__"] = _add
+ operators["__sub__"] = lambda lhs, rhs: expr._binary_op_helper(lhs, rhs, _ffi_api.subtract)
+ operators["__rsub__"] = _rhs
+ operators["__mul__"] = lambda lhs, rhs: expr._binary_op_helper(lhs, rhs, _ffi_api.multiply)
+ operators["__rmul__"] = operators["__mul__"]
+ operators["__div__"] = lambda lhs, rhs: expr._binary_op_helper(lhs, rhs, _ffi_api.divide)
+ operators["__rdiv__"] = _rhs
+ operators["__truediv__"] = operators["__div__"]
+ operators["__rtruediv__"] = _rhs
+ operators["__floordiv__"] = lambda lhs, rhs: expr._binary_op_helper(
lhs, rhs, _ffi_api.floor_divide
)
- _tensor_expr_overload.__rfloordiv__ = _rhs
- _tensor_expr_overload.__mod__ = lambda lhs, rhs: expr._binary_op_helper(lhs, rhs, _ffi_api.mod)
- _tensor_expr_overload.__rmod__ = _rhs
- _tensor_expr_overload.__pow__ = lambda lhs, rhs: expr._binary_op_helper(
- lhs, rhs, _ffi_api.power
- )
- _tensor_expr_overload.__rpow__ = _rhs
+ operators["__rfloordiv__"] = _rhs
+ operators["__mod__"] = lambda lhs, rhs: expr._binary_op_helper(lhs, rhs, _ffi_api.mod)
+ operators["__rmod__"] = _rhs
+ operators["__pow__"] = lambda lhs, rhs: expr._binary_op_helper(lhs, rhs, _ffi_api.power)
+ operators["__rpow__"] = _rhs
+ _tensor_expr_overload.register_tensor_expr_operators(operators)
_register_op_make()
diff --git a/python/tvm/relax/script/printer.py b/python/tvm/relax/script/printer.py
new file mode 100644
index 0000000000..15c8533d84
--- /dev/null
+++ b/python/tvm/relax/script/printer.py
@@ -0,0 +1,65 @@
+# 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.
+"""Relax-owned printing conveniences."""
+
+import tvm_ffi
+from tvm_ffi.access_path import AccessPath
+
+from tvm.ir.base import register_dialect_script_method
+from tvm.runtime import Object
+from tvm.script.printer import PrinterConfig
+
+
+def _relax_script(
+ self,
+ *,
+ name: str | None = None,
+ show_meta: bool = False,
+ int_dtype: str = "int32",
+ float_dtype: str = "void",
+ verbose_expr: bool = False,
+ indent_spaces: int = 4,
+ print_line_numbers: bool = False,
+ num_context_lines: int = -1,
+ extra_config: dict | None = None,
+ path_to_underline: list[AccessPath] | None = None,
+ path_to_annotate: dict[AccessPath, str] | None = None,
+ obj_to_underline: list[Object] | None = None,
+ obj_to_annotate: dict[Object, str] | None = None,
+) -> str:
+ merged_extra = dict(extra_config or {})
+ return tvm_ffi.get_global_func("script.printer.ReprPrintRelax")(
+ self,
+ PrinterConfig(
+ name=name,
+ show_meta=show_meta,
+ int_dtype=int_dtype,
+ float_dtype=float_dtype,
+ verbose_expr=verbose_expr,
+ indent_spaces=indent_spaces,
+ print_line_numbers=print_line_numbers,
+ num_context_lines=num_context_lines,
+ extra_config=merged_extra,
+ path_to_underline=path_to_underline,
+ path_to_annotate=path_to_annotate,
+ obj_to_underline=obj_to_underline,
+ obj_to_annotate=obj_to_annotate,
+ ),
+ )
+
+
+register_dialect_script_method("_relax_script", _relax_script)
diff --git a/python/tvm/relax/vm_build.py b/python/tvm/relax/vm_build.py
index db65ecbf68..d8e046cc10 100644
--- a/python/tvm/relax/vm_build.py
+++ b/python/tvm/relax/vm_build.py
@@ -146,7 +146,7 @@ def _vmlink(
tir_ext_libs = []
if tir_mod is not None and len(tir_mod.get_global_vars()) > 0:
tir_mod = _auto_attach_system_lib_prefix(tir_mod, target, system_lib)
- lib = tvm.tirx.build(tir_mod, target=target, pipeline=tir_pipeline)
+ lib = tvm.tirx.compile(tir_mod, target=target, pipeline=tir_pipeline).mod
for ext_mod in ext_libs:
if _is_device_module(ext_mod):
tir_ext_libs.append(ext_mod)
diff --git a/python/tvm/runtime/module.py b/python/tvm/runtime/module.py
index 82b725b72f..ccc77d6900 100644
--- a/python/tvm/runtime/module.py
+++ b/python/tvm/runtime/module.py
@@ -177,6 +177,9 @@ class Module(_Module):
module is "llvm," this is used to link all produced artifacts
into a final dynamic library.
+ For WASM/TVMJS export, explicitly pass
+ ``tvm.contrib.emcc.create_wasm`` as the compilation callback.
+
This behavior is controlled by the type of object exported.
If fcompile has attribute object_format, will compile host library
to that format. Otherwise, will use default format "o".
@@ -205,15 +208,13 @@ class Module(_Module):
If the compilation function returns an artifact it would be returned via
export_library, if any.
"""
- # NOTE: this function depends on contrib library features
- # which are only available in when TVM function is available.
+ # Exporting compiled artifacts requires the compiler runtime.
if _RUNTIME_ONLY:
raise RuntimeError("Cannot call export_library in runtime only mode")
# Extra dependencies during runtime.
from pathlib import Path
- from tvm.contrib import tvmjs as _tvmjs
from tvm.support import cc as _cc
from tvm.support import tar as _tar
from tvm.support import utils as _utils
@@ -270,8 +271,6 @@ class Module(_Module):
if not fcompile:
if file_name.endswith(".tar"):
fcompile = _tar.tar
- elif file_name.endswith(".wasm"):
- fcompile = _tvmjs.create_tvmjs_wasm
else:
fcompile = _cc.create_shared
@@ -392,32 +391,29 @@ class Module(_Module):
The function that takes same argument as func and returns a BenchmarkResult.
The ProfileResult reports `repeat` time costs in seconds.
"""
- try:
- feval = _ffi_api.ModuleTimeEvaluator(
- self,
- func_name,
- dev,
- number,
- repeat,
- min_repeat_ms,
- limit_zero_time_iterations,
- cooldown_interval_ms,
- repeats_to_cooldown,
- cache_flush_bytes,
- f_preproc,
- )
+ feval = _ffi_api.ModuleTimeEvaluator(
+ self,
+ func_name,
+ dev,
+ number,
+ repeat,
+ min_repeat_ms,
+ limit_zero_time_iterations,
+ cooldown_interval_ms,
+ repeats_to_cooldown,
+ cache_flush_bytes,
+ f_preproc,
+ )
+
+ def evaluator(*args):
+ """Internal wrapped evaluator."""
+ # Wrap feval so we can add more stats in future.
+ blob = feval(*args)
+ fmt = "@" + ("d" * repeat)
+ results = struct.unpack(fmt, blob)
+ return BenchmarkResult(results)
- def evaluator(*args):
- """Internal wrapped evaluator."""
- # Wrap feval so we can add more stats in future.
- blob = feval(*args)
- fmt = "@" + ("d" * repeat)
- results = struct.unpack(fmt, blob)
- return BenchmarkResult(results)
-
- return evaluator
- except NameError:
- raise NameError("time_evaluator is only supported when RPC is enabled")
+ return evaluator
def load_module(path):
diff --git a/python/tvm/s_tir/dlight/benchmark/bench.py b/python/tvm/s_tir/dlight/benchmark/bench.py
index 281ec129fe..8bd5cfc104 100644
--- a/python/tvm/s_tir/dlight/benchmark/bench.py
+++ b/python/tvm/s_tir/dlight/benchmark/bench.py
@@ -121,7 +121,7 @@ def benchmark(
# append scalar input tensors for rotary embedding
input_tensors.extend(scalar_input_tensors)
# build locally
- rt_mod = tvm.tirx.build(mod, target=target)
+ rt_mod = tvm.tirx.compile(mod, target=target).mod
# set up evaluator config
evaluator_config = EvaluatorConfig._normalized( # pylint: disable=protected-access
evaluator_config
diff --git a/python/tvm/s_tir/meta_schedule/builder/local_builder.py b/python/tvm/s_tir/meta_schedule/builder/local_builder.py
index 2a97a732ca..36b8825131 100644
--- a/python/tvm/s_tir/meta_schedule/builder/local_builder.py
+++ b/python/tvm/s_tir/meta_schedule/builder/local_builder.py
@@ -260,11 +260,11 @@ def default_build(mod: IRModule, target: Target, _params: dict[str, Tensor] | No
# pylint: disable=import-outside-toplevel
import tvm.s_tir.tensor_intrin # pylint: disable=unused-import
from tvm.s_tir.transform import RemoveWeightLayoutRewriteBlock
- from tvm.tirx import build as tirx_build
+ from tvm.tirx import compile as tirx_compile
# pylint: enable=import-outside-toplevel
mod = RemoveWeightLayoutRewriteBlock(skip_tensor_rewrite=True)(mod)
- return tirx_build(mod, target=target)
+ return tirx_compile(mod, target=target).mod
@register_global_func("s_tir.meta_schedule.builder.default_export")
diff --git a/python/tvm/s_tir/pipeline.py b/python/tvm/s_tir/pipeline.py
index 68fc5f0caa..0e3916a2f7 100644
--- a/python/tvm/s_tir/pipeline.py
+++ b/python/tvm/s_tir/pipeline.py
@@ -141,7 +141,7 @@ def _s_tir_pipeline(*, prepare_only=False, target=None, adreno=False):
if scheduled:
prepared.update(_lower_s_tir(adreno=adreno)(_submodule(mod, scheduled)))
if native:
- prepared.update(tir_pipeline._lower_thread_binding()(_submodule(mod, native)))
+ prepared.update(_submodule(mod, native))
# Classification sees every caller after each dialect has explicit launches.
prepared = tir_pipeline._bind_target(prepared, target)
@@ -152,9 +152,7 @@ def _s_tir_pipeline(*, prepare_only=False, target=None, adreno=False):
if scheduled:
lowered.update(_lower_s_tir_tail(adreno=adreno)(_submodule(prepared, scheduled)))
if native:
- lower_native, _, _ = tir_pipeline.tirx_pipeline(
- prepare_only=True, _bindings_lowered=True
- )
+ lower_native, _, _ = tir_pipeline.tirx_pipeline(prepare_only=True)
lowered.update(lower_native(_submodule(prepared, native)))
if prepare_only:
return lowered
@@ -175,7 +173,7 @@ def _s_tir_pipeline(*, prepare_only=False, target=None, adreno=False):
def _prepare_legacy_target_binding(mod):
"""Keep custom scheduling pipelines intact when prebinding is unambiguous."""
- scheduled, native = _partition_dialects(mod)
+ scheduled, _ = _partition_dialects(mod)
in_binding = False
in_device_region = False
@@ -218,10 +216,6 @@ def _prepare_legacy_target_binding(mod):
(tvm.ir.Call, visit_call),
],
)
- if native:
- prepared = _submodule(mod, scheduled)
- prepared.update(tir_pipeline._lower_thread_binding()(_submodule(mod, native)))
- return prepared
return mod
@@ -259,3 +253,4 @@ def _select_default_pipeline(mod, target, *, target_to_bind=None):
tir_pipeline.register_default_tir_pipeline_selector(_select_default_pipeline)
+tir_pipeline._LEGACY_TARGET_PREPARERS.append(_prepare_legacy_target_binding)
diff --git a/python/tvm/script/printer/__init__.py b/python/tvm/script/printer/__init__.py
index 77af75e0f0..ca37d066a4 100644
--- a/python/tvm/script/printer/__init__.py
+++ b/python/tvm/script/printer/__init__.py
@@ -20,4 +20,6 @@ This package provides a set of APIs to print supported TVM IR into TVMScript
in a roundtrippable way.
"""
-from .scriptable import PrinterConfig, Scriptable
+from tvm.ir.base import Scriptable as Scriptable
+
+from .scriptable import PrinterConfig
diff --git a/python/tvm/script/printer/scriptable.py b/python/tvm/script/printer/scriptable.py
index 2168d1456f..22c286a8ba 100644
--- a/python/tvm/script/printer/scriptable.py
+++ b/python/tvm/script/printer/scriptable.py
@@ -22,7 +22,7 @@ from collections.abc import Sequence
from tvm_ffi import register_object
from tvm_ffi.access_path import AccessPath
-from tvm.ir.base import Node, Scriptable
+from tvm.ir.base import register_script_printer
from tvm.runtime import Object
from . import _ffi_api
@@ -99,10 +99,6 @@ def _script(obj: Object, config: PrinterConfig) -> str:
return _ffi_api.Script(obj, config) # type: ignore # pylint: disable=no-member
-def _relax_script(obj: Object, config: PrinterConfig) -> str:
- return _ffi_api.ReprPrintRelax(obj, config)
-
-
class _ScriptableMethods:
"""A base class that enables the script() and show() method."""
@@ -195,43 +191,6 @@ class _ScriptableMethods:
),
)
- def _relax_script(
- self,
- *,
- name: str | None = None,
- show_meta: bool = False,
- int_dtype: str = "int32",
- float_dtype: str = "void",
- verbose_expr: bool = False,
- indent_spaces: int = 4,
- print_line_numbers: bool = False,
- num_context_lines: int = -1,
- extra_config: dict | None = None,
- path_to_underline: list[AccessPath] | None = None,
- path_to_annotate: dict[AccessPath, str] | None = None,
- obj_to_underline: list[Object] | None = None,
- obj_to_annotate: dict[Object, str] | None = None,
- ) -> str:
- merged_extra = dict(extra_config or {})
- return _relax_script(
- self,
- PrinterConfig(
- name=name,
- show_meta=show_meta,
- int_dtype=int_dtype,
- float_dtype=float_dtype,
- verbose_expr=verbose_expr,
- indent_spaces=indent_spaces,
- print_line_numbers=print_line_numbers,
- num_context_lines=num_context_lines,
- extra_config=merged_extra,
- path_to_underline=path_to_underline,
- path_to_annotate=path_to_annotate,
- obj_to_underline=obj_to_underline,
- obj_to_annotate=obj_to_annotate,
- ),
- )
-
def show(
self,
style: str | None = None,
@@ -351,7 +310,6 @@ def _repr(self) -> str:
# The printer supplies its explicit display capability to existing IR wrappers.
-Scriptable.script = _ScriptableMethods.script
-Scriptable.show = _ScriptableMethods.show
-Scriptable._relax_script = _ScriptableMethods._relax_script
-Node.__repr__ = _repr
+register_script_printer(
+ script=_ScriptableMethods.script, show=_ScriptableMethods.show, node_repr=_repr
+)
diff --git a/python/tvm/tirx/__init__.py b/python/tvm/tirx/__init__.py
index 1835754c7f..5c605f686d 100644
--- a/python/tvm/tirx/__init__.py
+++ b/python/tvm/tirx/__init__.py
@@ -90,7 +90,7 @@ from . import backend
from tvm.base import _RUNTIME_ONLY as _RUNTIME_ONLY_TIRX # pylint: disable=wrong-import-position
if not _RUNTIME_ONLY_TIRX:
- from .build import build
+ from ._compile import compile
from .compilation_pipeline import (
get_tir_pipeline,
get_default_tir_pipeline,
diff --git a/python/tvm/tirx/build.py b/python/tvm/tirx/_compile.py
similarity index 88%
rename from python/tvm/tirx/build.py
rename to python/tvm/tirx/_compile.py
index 4b06b881a0..b44ed86fde 100644
--- a/python/tvm/tirx/build.py
+++ b/python/tvm/tirx/_compile.py
@@ -16,18 +16,21 @@
# under the License.
# pylint: disable=invalid-name
-"""The build utils in python."""
+"""Native TIRx compilation and target code generation."""
from tvm_ffi import DLDeviceType
import tvm
from tvm import ir
from tvm.ir.module import IRModule
+from tvm.runtime import Executable
from tvm.target import Target
from tvm.tirx import Function
-def split_host_device_mods(mod: IRModule) -> tuple[IRModule, dict[tuple[Target, str], IRModule]]:
+def _split_host_device_modules(
+ mod: IRModule,
+) -> tuple[IRModule, dict[tuple[Target, str], IRModule]]:
"""Split an IRModule into host and device modules.
This function takes an IRModule containing functions with different target attributes
@@ -100,12 +103,12 @@ def split_host_device_mods(mod: IRModule) -> tuple[IRModule, dict[tuple[Target,
- Device kernel functions: use `calling_conv: 2` (kDeviceKernelLaunch)
"""
- def is_host_func(f):
+ def is_host_function(f):
target = f.attrs.get("target", tvm.target.Target("llvm"))
return target.kind.name in ["llvm", "c", "cuda_host"]
- host_mod = tvm.tirx.transform.Filter(is_host_func)(mod)
- device_mod = tvm.tirx.transform.Filter(lambda f: not is_host_func(f))(mod)
+ host_mod = tvm.tirx.transform.Filter(is_host_function)(mod)
+ device_mod = tvm.tirx.transform.Filter(lambda f: not is_host_function(f))(mod)
# Options participate in grouping even when two kernels share a Target.
groups = {}
for gv, func in device_mod.functions.items():
@@ -126,7 +129,7 @@ def split_host_device_mods(mod: IRModule) -> tuple[IRModule, dict[tuple[Target,
return host_mod, device_mod_dict
-def codegen_build(mod: IRModule, target: Target) -> tvm.runtime.Module:
+def _build_for_target(mod: IRModule, target: Target) -> tvm.runtime.Module:
"""Build a runtime module from an IRModule and a Target."""
if tvm.transform.PassContext.current().config.get("tirx.disable_assert", False):
mod = tvm.tirx.transform.SkipAssert()(mod)
@@ -137,9 +140,9 @@ def codegen_build(mod: IRModule, target: Target) -> tvm.runtime.Module:
return bf(mod, target)
-def tir_to_runtime(
+def _codegen_host_and_device(
host_mod: IRModule, device_mod_dict: dict[tuple[Target, str], IRModule], target_host: Target
-):
+) -> tvm.runtime.Module:
"""Compile each device group and attach it to the host runtime module."""
# Get the first module to get the attributes
@@ -150,27 +153,26 @@ def tir_to_runtime(
device_modules = []
for (target, _config), device_mod in device_mod_dict.items():
if len(device_mod.functions) != 0:
- device_modules.append(codegen_build(device_mod, target))
+ device_modules.append(_build_for_target(device_mod, target))
- mhost = codegen_build(mhost_all, target_host)
+ mhost = _build_for_target(mhost_all, target_host)
for dev_mod in device_modules:
if dev_mod is not None:
mhost.import_module(dev_mod)
return mhost
-def build(
+def compile( # pylint: disable=redefined-builtin
mod: Function | IRModule,
target: str | Target | None = None,
pipeline: str | tvm.transform.Pass | None = "default",
-):
- """Build a function with a signature, generating code for devices
- coupled with target information.
+) -> Executable:
+ """Compile a function or module to an executable for the target.
Parameters
----------
mod : Union[Function, IRModule]
- The input to be built.
+ The input to be compiled.
target : Optional[Union[str, Target]]
The target for compilation, including backend configuration defaults.
pipeline : Union[None, str, tvm.transform.Pass]
@@ -178,8 +180,8 @@ def build(
Returns
-------
- tvm.runtime.Module
- A module combining both host and device code.
+ tvm.runtime.Executable
+ An executable combining both host and device code.
"""
# Convert Function to IRModule
if isinstance(mod, Function):
@@ -236,7 +238,7 @@ def build(
mod = pipeline(mod)
# Step 4: Get host and device modules
- host_mod, device_mod_dict = split_host_device_mods(mod)
+ host_mod, device_mod_dict = _split_host_device_modules(mod)
# Step 5: Apply finalization passes
host_mod = finalize_host_passes()(host_mod)
@@ -246,7 +248,8 @@ def build(
}
# Convert TIR IRModules to runtime Module by calling target.build
- return tir_to_runtime(host_mod, device_mod_dict, target_host)
+ return Executable(_codegen_host_and_device(host_mod, device_mod_dict, target_host))
-tvm.register_global_func("tirx.build", build)
+# C++ compiler consumers link the generated runtime module directly.
+tvm.register_global_func("tirx.compile", lambda *args, **kwargs: compile(*args, **kwargs).mod)
diff --git a/python/tvm/tirx/compilation_pipeline.py b/python/tvm/tirx/compilation_pipeline.py
index fe0f1ab9fc..5381159c86 100644
--- a/python/tvm/tirx/compilation_pipeline.py
+++ b/python/tvm/tirx/compilation_pipeline.py
@@ -30,14 +30,12 @@ def default_tir_pipeline(*, target=None):
return tirx_pipeline(target=target)
-def tirx_pipeline(*, prepare_only=False, target=None, _bindings_lowered=False):
- """The TIRX pipeline used in tvm.tirx.build"""
+def tirx_pipeline(*, prepare_only=False, target=None):
+ """The TIRX pipeline used in tvm.tirx.compile"""
@tvm.transform.module_pass(opt_level=0)
def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.IRModule:
"""The default lowering passes for TIR backend."""
- if not _bindings_lowered:
- mod = _lower_thread_binding()(mod)
mod = _bind_target(mod, target)
pass_ctx = tvm.transform.PassContext.current()
config = pass_ctx.config
@@ -64,11 +62,6 @@ def tirx_pipeline(*, prepare_only=False, target=None, _bindings_lowered=False):
return _pipeline, finalize_host_passes, finalize_device_passes
-def _lower_thread_binding():
- """Finish S-TIR loop annotations before generic TIRx consumers."""
- return tvm.s_tir.transform.LowerThreadBinding()
-
-
def _bind_target(mod, target):
if target is None:
return mod
@@ -82,9 +75,9 @@ def _bind_target(mod, target):
def _bind_legacy_target(mod, target):
- from tvm.s_tir.pipeline import _prepare_legacy_target_binding
-
- return _bind_target(_prepare_legacy_target_binding(mod), target)
+ for prepare in _LEGACY_TARGET_PREPARERS:
+ mod = prepare(mod)
+ return _bind_target(mod, target)
def _with_target_binding(pipeline, target, *, original_mod=None, prepared_mod=None):
@@ -207,6 +200,8 @@ def get_tir_pipeline(name: str | None = None, *, target=None, **kwargs):
# Dialects register their own default selection without introducing imports here.
_DEFAULT_PIPELINE_SELECTORS = []
+# Dialects validate their own scheduling constructs before legacy target binding.
+_LEGACY_TARGET_PREPARERS = []
def register_default_tir_pipeline_selector(selector) -> None:
diff --git a/python/tvm/tirx/transform/transform.py b/python/tvm/tirx/transform/transform.py
index b50d5340d3..54af2ec41d 100644
--- a/python/tvm/tirx/transform/transform.py
+++ b/python/tvm/tirx/transform/transform.py
@@ -25,6 +25,11 @@ from . import _ffi_api
from .function_pass import function_pass
+def DecorateDeviceScope():
+ """Wrap a function body in a device scope region."""
+ return _ffi_api.DecorateDeviceScope()
+
+
def Apply(ftransform):
"""Apply ftransform to each function in the Module.
diff --git a/src/backend/cuda/tile/exec_context.h b/src/backend/cuda/tile/exec_context.h
index 35120e3fde..31585d1f4b 100644
--- a/src/backend/cuda/tile/exec_context.h
+++ b/src/backend/cuda/tile/exec_context.h
@@ -84,9 +84,10 @@ struct ExecSplit {
};
/*! \brief Initial A at Function device entry: all threads active, offsets zero. */
-TVM_DLL ActiveSet InitialActiveSet(int64_t lane_ext, int64_t warp_ext, int64_t cta_ext);
-TVM_DLL ActiveSet InitialActiveSet(int64_t lane_ext, int64_t warp_ext, int64_t cta_ext,
- const std::vector<std::pair<std::string, int64_t>>& cta_axes);
+TVM_TIRX_DLL ActiveSet InitialActiveSet(int64_t lane_ext, int64_t warp_ext, int64_t cta_ext);
+TVM_TIRX_DLL ActiveSet
+InitialActiveSet(int64_t lane_ext, int64_t warp_ext, int64_t cta_ext,
+ const std::vector<std::pair<std::string, int64_t>>& cta_axes);
/*!
* \brief Narrow A on the lane bound to ``binding``.
@@ -99,8 +100,8 @@ TVM_DLL ActiveSet InitialActiveSet(int64_t lane_ext, int64_t warp_ext, int64_t c
* not a filter target; flat thread ids are accepted only when the range can be
* represented as a rectangular lane/warp active set.
*/
-TVM_DLL bool FilterNarrow(const ActiveSet& A, ScopeBinding binding, int64_t lo, int64_t hi,
- ActiveSet* out, std::string* err);
+TVM_TIRX_DLL bool FilterNarrow(const ActiveSet& A, ScopeBinding binding, int64_t lo, int64_t hi,
+ ActiveSet* out, std::string* err);
/*!
* \brief Factor A into (inter, intra) for target scope_kind.
@@ -108,8 +109,8 @@ TVM_DLL bool FilterNarrow(const ActiveSet& A, ScopeBinding binding, int64_t lo,
* Returns false on factoring failure (warpgroup with warpid lane that
* crosses a warpgroup boundary unaligned) and writes reason to *err.
*/
-TVM_DLL bool ScopeSwitch(const ActiveSet& A, ScopeKind scope_kind, ExecSplit* out,
- std::string* err);
+TVM_TIRX_DLL bool ScopeSwitch(const ActiveSet& A, ScopeKind scope_kind, ExecSplit* out,
+ std::string* err);
/*! \brief Per-program-point ExecContext: active set + scope kind + split. */
struct ExecContext {
@@ -144,7 +145,7 @@ struct ExecContext {
* by ``DispatchContextObj::{inter, intra}``: axis name -> [extent, offset]
* for unit-stride axes, or [extent, offset, stride] for strided axes.
*/
-TVM_DLL ffi::Map<ffi::String, ffi::Array<PrimExpr>> EncodeSplitSide(
+TVM_TIRX_DLL ffi::Map<ffi::String, ffi::Array<PrimExpr>> EncodeSplitSide(
const std::unordered_map<std::string, AxisRange>& side);
} // namespace tvm::backend::cuda
diff --git a/src/backend/cuda/tile/exec_scope.h b/src/backend/cuda/tile/exec_scope.h
index 870e8934be..8bc4b35142 100644
--- a/src/backend/cuda/tile/exec_scope.h
+++ b/src/backend/cuda/tile/exec_scope.h
@@ -32,10 +32,10 @@ enum class ScopeKind : int {
};
/*! \brief Convert a ScopeKind to its string name (e.g. kThread -> "thread"). */
-TVM_DLL std::string ScopeKindToString(ScopeKind kind);
+TVM_TIRX_DLL std::string ScopeKindToString(ScopeKind kind);
/*! \brief Parse a string name to a ScopeKind. FATAL if unknown. */
-TVM_DLL ScopeKind StringToScopeKind(const ffi::String& name);
+TVM_TIRX_DLL ScopeKind StringToScopeKind(const ffi::String& name);
/*!
* \brief The binding between a parent scope and a child scope as used by a
@@ -71,10 +71,11 @@ enum class ScopeBinding : int {
};
/*! \brief Convert a ScopeBinding to its (parent, cur) string pair. */
-TVM_DLL std::pair<ffi::String, ffi::String> ScopeBindingToStringPair(ScopeBinding binding);
+TVM_TIRX_DLL std::pair<ffi::String, ffi::String> ScopeBindingToStringPair(ScopeBinding binding);
/*! \brief Parse a (parent, cur) string pair to a ScopeBinding. FATAL if unknown. */
-TVM_DLL ScopeBinding StringPairToScopeBinding(const ffi::String& parent, const ffi::String& cur);
+TVM_TIRX_DLL ScopeBinding StringPairToScopeBinding(const ffi::String& parent,
+ const ffi::String& cur);
/*!
* \brief Strict-weak "a is wider than b" on scope kinds: ``world > kernel >
diff --git a/src/backend/cuda/tile/filter_canonical.h b/src/backend/cuda/tile/filter_canonical.h
index 4fa148b8fa..e492f92aed 100644
--- a/src/backend/cuda/tile/filter_canonical.h
+++ b/src/backend/cuda/tile/filter_canonical.h
@@ -150,8 +150,8 @@ using ScopeIdPredicate = std::function<bool(const Var&)>;
* Thread safety: pure function; safe to call concurrently provided the
* callback `is_scope_id` is itself thread-safe.
*/
-TVM_DLL std::optional<CanonicalForm> TryClassifyCanonical(const PrimExpr& cond,
- const ScopeIdPredicate& is_scope_id);
+TVM_TIRX_DLL std::optional<CanonicalForm> TryClassifyCanonical(const PrimExpr& cond,
+ const ScopeIdPredicate& is_scope_id);
} // namespace tvm::backend::cuda
diff --git a/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc b/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc
index 17da581bd3..129bedb211 100644
--- a/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc
+++ b/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc
@@ -547,7 +547,7 @@ ffi::Module BuildHexagon(IRModule mod, Target target) {
const auto f = tvm::ffi::Function::GetGlobal("tvm.contrib.hexagon.link_shared");
TVM_FFI_ICHECK(f.has_value()) << "tvm.contrib.hexagon.link_shared does not to exist, "
- "do import tvm.contrib.hexagon";
+ "load the Hexagon backend linker";
ffi::Array<Expr> o_names = {StringImm(o_name)};
ffi::Map<ffi::String, ffi::String> extra_args;
diff --git a/src/backend/vulkan/codegen/codegen_spirv.cc b/src/backend/vulkan/codegen/codegen_spirv.cc
index f95d38d395..6b3823d544 100644
--- a/src/backend/vulkan/codegen/codegen_spirv.cc
+++ b/src/backend/vulkan/codegen/codegen_spirv.cc
@@ -25,7 +25,6 @@
#include <tvm/ir/prim/expr.h>
#include <tvm/ir/prim/op.h>
-#include <tvm/s_tir/stmt.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/op/abi.h>
#include <tvm/tirx/op/gpu.h>
@@ -907,7 +906,7 @@ void CodeGenSPIRV::DispatchAllocTensor(const BindNode* op, const CallNode* buffe
case runtime::StorageRank::kWMMAMatrixA:
case runtime::StorageRank::kWMMAMatrixB:
case runtime::StorageRank::kWMMAAccumulator: {
- auto shape = annotations->dict.Get(tvm::s_tir::attr::kFragmentShape);
+ auto shape = annotations->dict.Get("fragment_shape");
TVM_FFI_ICHECK(shape.has_value())
<< "Cannot find shape of the wmma fragment " << buffer.name();
fragment_info_[var_node] = {shape.value().as_or_throw<ffi::String>()};
diff --git a/src/ir/instrument.cc b/src/ir/instrument.cc
index c707909d8c..68c268ec40 100644
--- a/src/ir/instrument.cc
+++ b/src/ir/instrument.cc
@@ -109,7 +109,7 @@ class BasePassInstrument : public PassInstrument {
* \param run_before_pass_callback Callback to call before a pass run.
* \param run_after_pass_callback Callback to call after a pass run.
*/
- TVM_DLL BasePassInstrument(
+ TVM_TIRX_DLL BasePassInstrument(
ffi::String name, ffi::TypedFunction<void()> enter_pass_ctx_callback,
ffi::TypedFunction<void()> exit_pass_ctx_callback,
ffi::TypedFunction<bool(const IRModule&, const PassInfo&)> should_run_callback,
diff --git a/src/ir/prim/op_utils.h b/src/ir/prim/op_utils.h
index b33a9be44d..2913bb4b19 100644
--- a/src/ir/prim/op_utils.h
+++ b/src/ir/prim/op_utils.h
@@ -21,7 +21,7 @@
#define TVM_IR_PRIM_OP_UTILS_H_
#include <tvm/ir/prim/op.h>
namespace tvm::prim::detail {
-TVM_DLL void BinaryOpMatchTypes(PrimExpr& lhs, PrimExpr& rhs, Location loc);
+TVM_TIRX_DLL void BinaryOpMatchTypes(PrimExpr& lhs, PrimExpr& rhs, Location loc);
} // namespace tvm::prim::detail
#endif // TVM_IR_PRIM_OP_UTILS_H_
diff --git a/src/relax/backend/vm/codegen_vm.cc b/src/relax/backend/vm/codegen_vm.cc
index bf5388f534..8771a13685 100644
--- a/src/relax/backend/vm/codegen_vm.cc
+++ b/src/relax/backend/vm/codegen_vm.cc
@@ -36,7 +36,7 @@
#include <unordered_map>
#include <vector>
-#include "../../../runtime/const_loader_module.h"
+#include "../../../runtime/vm/const_loader_module.h"
#include "../../../target/source/codegen_source_base.h"
#include "../../transform/utils.h"
diff --git a/src/relax/transform/fold_constant.cc b/src/relax/transform/fold_constant.cc
index 7ca2cc8ccb..7738b067c3 100644
--- a/src/relax/transform/fold_constant.cc
+++ b/src/relax/transform/fold_constant.cc
@@ -122,7 +122,7 @@ class ConstantFolder : public ExprMutator {
// already scheduled to only work on GPU, we will need to skip this in the const folder for
// now
// TODO(Hongyi): further check and narrow the scope of foldable function
- const auto pf = tvm::ffi::Function::GetGlobalRequired("tirx.build");
+ const auto pf = tvm::ffi::Function::GetGlobalRequired("tirx.compile");
func = WithAttr(func, tvm::attr::kGlobalSymbol, ffi::String("tir_function"));
ffi::Module rt_module = pf(func, eval_cpu_target).cast<ffi::Module>();
build_func = rt_module->GetFunction("tir_function");
diff --git a/src/runtime/extra/contrib/amx/amx_config.cc b/src/runtime/extra/contrib/amx/amx_config.cc
index 9671994ce0..5d3ec351d5 100644
--- a/src/runtime/extra/contrib/amx/amx_config.cc
+++ b/src/runtime/extra/contrib/amx/amx_config.cc
@@ -24,9 +24,6 @@
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
-namespace tvm {
-namespace runtime {
-
#ifdef __linux__
#include <errno.h>
#include <fcntl.h>
@@ -37,6 +34,12 @@ namespace runtime {
#include <tvm/runtime/logging.h>
#include <unistd.h>
+#endif
+
+namespace tvm {
+namespace runtime {
+
+#ifdef __linux__
#define XFEATURE_XTILECFG 17
#define XFEATURE_XTILEDATA 18
#define XFEATURE_MASK_XTILECFG (1 << XFEATURE_XTILECFG)
diff --git a/src/runtime/rpc/rpc_module.cc b/src/runtime/rpc/rpc_module.cc
index e94aaa4422..a39b253ae9 100644
--- a/src/runtime/rpc/rpc_module.cc
+++ b/src/runtime/rpc/rpc_module.cc
@@ -33,9 +33,6 @@
#include <cstring>
#include <memory>
#include <thread>
-#if defined(_M_X64) || defined(__x86_64__)
-#include <immintrin.h>
-#endif
#include "rpc_endpoint.h"
#include "rpc_remote.h"
@@ -304,108 +301,27 @@ std::shared_ptr<RPCSession> RPCModuleGetSession(ffi::Module mod) {
return rmod->sess();
}
-/*!
- * \brief Flush the cache.
- * \param addr The address of data we want to flush
- * \param len The length of data
- */
-/*
- * When we are in the tuning of TVM, we will make TVM occupy
- * the cache fully and doesn't flush it during iteration.
- * This has problems then in e2e testing, since arrays that
- * we assume exist in cache (ie. weights) are evicted during e2e runs,
- * which leads to lower performance.
- */
-inline void CPUCacheFlushImpl(const char* addr, unsigned int len) {
-#if (defined(_M_X64) || defined(__x86_64__) || defined(__aarch64__))
-
-#if defined(__aarch64__)
- size_t ctr_el0 = 0;
- asm volatile("mrs %0, ctr_el0" : "=r"(ctr_el0));
- const size_t cache_line = 4 << ((ctr_el0 >> 16) & 15);
-#else
- const size_t cache_line = 64;
-#endif
-
- if (addr == nullptr || len <= 0) {
- return;
- }
-
- for (uintptr_t uptr = (uintptr_t)addr & ~(cache_line - 1); uptr < (uintptr_t)addr + len;
- uptr += cache_line) {
-#if defined(__aarch64__)
- asm volatile("dc civac, %0\n\t" : : "r"(reinterpret_cast<const void*>(uptr)) : "memory");
-#else
- _mm_clflush(reinterpret_cast<const void*>(uptr));
-#endif
- }
-
-#if defined(__aarch64__)
- asm volatile("dmb ishst" : : : "memory");
-#endif
-
-#endif
-}
-
-inline void CPUCacheFlush(int begin_index, const ffi::PackedArgs& args) {
- for (int i = begin_index; i < args.size(); i++) {
- CPUCacheFlushImpl(static_cast<char*>((args[i].cast<DLTensor*>()->data)),
- ffi::GetDataSize(*(args[i].cast<DLTensor*>())));
- }
-}
-
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
- refl::GlobalDef()
- .def("runtime.RPCTimeEvaluator",
- [](ffi::Optional<ffi::Module> opt_mod, std::string name, int device_type, int device_id,
- int number, int repeat, int min_repeat_ms, int limit_zero_time_iterations,
- int cooldown_interval_ms, int repeats_to_cooldown, int cache_flush_bytes,
- std::string f_preproc_name) {
- Device dev;
- dev.device_type = static_cast<DLDeviceType>(device_type);
- dev.device_id = device_id;
- if (opt_mod.has_value()) {
- ffi::Module m = opt_mod.value();
- std::string tkey = m->kind();
- if (tkey == "rpc") {
- return static_cast<RPCModuleObj*>(m.operator->())
- ->GetTimeEvaluator(name, dev, number, repeat, min_repeat_ms,
- limit_zero_time_iterations, cooldown_interval_ms,
- repeats_to_cooldown, cache_flush_bytes, f_preproc_name);
- } else {
- ffi::Function f_preproc;
- if (!f_preproc_name.empty()) {
- auto pf_preproc = tvm::ffi::Function::GetGlobal(f_preproc_name);
- TVM_FFI_ICHECK(pf_preproc.has_value())
- << "Cannot find " << f_preproc_name << " in the global function";
- f_preproc = *pf_preproc;
- }
- ffi::Optional<ffi::Function> pf = m->GetFunction(name);
- TVM_FFI_ICHECK(pf.has_value())
- << "Cannot find " << name << "` in the global registry";
- return WrapTimeEvaluator(*pf, dev, number, repeat, min_repeat_ms,
- limit_zero_time_iterations, cooldown_interval_ms,
- repeats_to_cooldown, cache_flush_bytes, f_preproc);
- }
- } else {
- auto pf = tvm::ffi::Function::GetGlobal(name);
- TVM_FFI_ICHECK(pf.has_value())
- << "Cannot find " << name << " in the global function";
- ffi::Function f_preproc;
- if (!f_preproc_name.empty()) {
- auto pf_preproc = tvm::ffi::Function::GetGlobal(f_preproc_name);
- TVM_FFI_ICHECK(pf_preproc.has_value())
- << "Cannot find " << f_preproc_name << " in the global function";
- f_preproc = *pf_preproc;
- }
- return WrapTimeEvaluator(*pf, dev, number, repeat, min_repeat_ms,
- limit_zero_time_iterations, cooldown_interval_ms,
- repeats_to_cooldown, cache_flush_bytes, f_preproc);
- }
- })
- .def_packed("cache_flush_cpu_non_first_arg",
- [](ffi::PackedArgs args, ffi::Any* rv) { CPUCacheFlush(1, args); });
+ refl::GlobalDef().def(
+ "runtime.RPCTimeEvaluator",
+ [](ffi::Optional<ffi::Module> opt_mod, std::string name, int device_type, int device_id,
+ int number, int repeat, int min_repeat_ms, int limit_zero_time_iterations,
+ int cooldown_interval_ms, int repeats_to_cooldown, int cache_flush_bytes,
+ std::string f_preproc_name) {
+ Device dev{static_cast<DLDeviceType>(device_type), device_id};
+ if (opt_mod.has_value() && std::string(opt_mod.value()->kind()) == "rpc") {
+ return static_cast<RPCModuleObj*>(opt_mod.value().operator->())
+ ->GetTimeEvaluator(name, dev, number, repeat, min_repeat_ms,
+ limit_zero_time_iterations, cooldown_interval_ms,
+ repeats_to_cooldown, cache_flush_bytes, f_preproc_name);
+ }
+ return ffi::Function::GetGlobalRequired("runtime.LocalTimeEvaluator")(
+ opt_mod, name, device_type, device_id, number, repeat, min_repeat_ms,
+ limit_zero_time_iterations, cooldown_interval_ms, repeats_to_cooldown,
+ cache_flush_bytes, f_preproc_name)
+ .cast<ffi::Function>();
+ });
}
// server function registration.
@@ -444,18 +360,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
<< "Device belongs to a different remote session";
return dev->endpoint_device;
})
- .def("runtime.ModuleTimeEvaluator",
+ .def("rpc.ModuleTimeEvaluator",
[](ffi::Module mod, std::string name, ffi::AnyView device, int number, int repeat,
int min_repeat_ms, int limit_zero_time_iterations, int cooldown_interval_ms,
int repeats_to_cooldown, int cache_flush_bytes, std::string f_preproc_name) {
- if (std::string(mod->kind()) != "rpc") {
- auto dev = device.cast<DLDevice>();
- return ffi::Function::GetGlobalRequired("runtime.RPCTimeEvaluator")(
- mod, name, static_cast<int>(dev.device_type), dev.device_id, number,
- repeat, min_repeat_ms, limit_zero_time_iterations, cooldown_interval_ms,
- repeats_to_cooldown, cache_flush_bytes, f_preproc_name)
- .cast<ffi::Function>();
- }
auto dev = device.cast<RemoteDevice>();
auto sess = RPCModuleGetSession(mod);
sess->EnsureOpen();
diff --git a/src/runtime/timer.cc b/src/runtime/timer.cc
index 99b0b253ff..fc257a1ffe 100644
--- a/src/runtime/timer.cc
+++ b/src/runtime/timer.cc
@@ -22,6 +22,7 @@
* \brief Runtime timer primitives: Timer, WrapTimeEvaluator.
*/
+#include <tvm/ffi/extra/module.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/c_backend_api.h>
@@ -34,6 +35,9 @@
#include <set>
#include <sstream>
#include <thread>
+#if defined(_M_X64) || defined(__x86_64__)
+#include <immintrin.h>
+#endif
namespace tvm {
namespace runtime {
@@ -176,5 +180,97 @@ ffi::Function WrapTimeEvaluator(ffi::Function pf, Device dev, int number, int re
return ffi::Function::FromPacked(ftimer);
}
+/*!
+ * \brief Flush the cache.
+ * \param addr The address of data we want to flush
+ * \param len The length of data
+ */
+/*
+ * When we are in the tuning of TVM, we will make TVM occupy
+ * the cache fully and doesn't flush it during iteration.
+ * This has problems then in e2e testing, since arrays that
+ * we assume exist in cache (ie. weights) are evicted during e2e runs,
+ * which leads to lower performance.
+ */
+inline void CPUCacheFlushImpl(const char* addr, unsigned int len) {
+#if (defined(_M_X64) || defined(__x86_64__) || defined(__aarch64__))
+
+#if defined(__aarch64__)
+ size_t ctr_el0 = 0;
+ asm volatile("mrs %0, ctr_el0" : "=r"(ctr_el0));
+ const size_t cache_line = 4 << ((ctr_el0 >> 16) & 15);
+#else
+ const size_t cache_line = 64;
+#endif
+
+ if (addr == nullptr || len <= 0) {
+ return;
+ }
+
+ for (uintptr_t uptr = (uintptr_t)addr & ~(cache_line - 1); uptr < (uintptr_t)addr + len;
+ uptr += cache_line) {
+#if defined(__aarch64__)
+ asm volatile("dc civac, %0\n\t" : : "r"(reinterpret_cast<const void*>(uptr)) : "memory");
+#else
+ _mm_clflush(reinterpret_cast<const void*>(uptr));
+#endif
+ }
+
+#if defined(__aarch64__)
+ asm volatile("dmb ishst" : : : "memory");
+#endif
+
+#endif
+}
+
+inline void CPUCacheFlush(int begin_index, const ffi::PackedArgs& args) {
+ for (int i = begin_index; i < args.size(); i++) {
+ CPUCacheFlushImpl(static_cast<char*>((args[i].cast<DLTensor*>()->data)),
+ ffi::GetDataSize(*(args[i].cast<DLTensor*>())));
+ }
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace refl = tvm::ffi::reflection;
+ refl::GlobalDef()
+ .def("runtime.LocalTimeEvaluator",
+ [](ffi::Optional<ffi::Module> opt_mod, std::string name, int device_type, int device_id,
+ int number, int repeat, int min_repeat_ms, int limit_zero_time_iterations,
+ int cooldown_interval_ms, int repeats_to_cooldown, int cache_flush_bytes,
+ std::string f_preproc_name) {
+ Device dev{static_cast<DLDeviceType>(device_type), device_id};
+ auto pf = opt_mod.has_value() ? opt_mod.value()->GetFunction(name)
+ : ffi::Function::GetGlobal(name);
+ TVM_FFI_ICHECK(pf.has_value()) << "Cannot find " << name;
+ ffi::Function f_preproc;
+ if (!f_preproc_name.empty()) {
+ f_preproc = ffi::Function::GetGlobalRequired(f_preproc_name);
+ }
+ return WrapTimeEvaluator(*pf, dev, number, repeat, min_repeat_ms,
+ limit_zero_time_iterations, cooldown_interval_ms,
+ repeats_to_cooldown, cache_flush_bytes, f_preproc);
+ })
+ .def("runtime.ModuleTimeEvaluator",
+ [](ffi::Module mod, std::string name, ffi::AnyView device, int number, int repeat,
+ int min_repeat_ms, int limit_zero_time_iterations, int cooldown_interval_ms,
+ int repeats_to_cooldown, int cache_flush_bytes, std::string f_preproc_name) {
+ if (std::string(mod->kind()) == "rpc") {
+ return ffi::Function::GetGlobalRequired("rpc.ModuleTimeEvaluator")(
+ mod, name, device, number, repeat, min_repeat_ms,
+ limit_zero_time_iterations, cooldown_interval_ms, repeats_to_cooldown,
+ cache_flush_bytes, f_preproc_name)
+ .cast<ffi::Function>();
+ }
+ auto dev = device.cast<DLDevice>();
+ return ffi::Function::GetGlobalRequired("runtime.LocalTimeEvaluator")(
+ mod, name, static_cast<int>(dev.device_type), dev.device_id, number, repeat,
+ min_repeat_ms, limit_zero_time_iterations, cooldown_interval_ms,
+ repeats_to_cooldown, cache_flush_bytes, f_preproc_name)
+ .cast<ffi::Function>();
+ })
+ .def_packed("cache_flush_cpu_non_first_arg",
+ [](ffi::PackedArgs args, ffi::Any* rv) { CPUCacheFlush(1, args); });
+}
+
} // namespace runtime
} // namespace tvm
diff --git a/src/runtime/const_loader_module.cc b/src/runtime/vm/const_loader_module.cc
similarity index 95%
rename from src/runtime/const_loader_module.cc
rename to src/runtime/vm/const_loader_module.cc
index 40da082fcd..20ccbdff82 100644
--- a/src/runtime/const_loader_module.cc
+++ b/src/runtime/vm/const_loader_module.cc
@@ -18,14 +18,10 @@
*/
/*!
- * \file src/runtime/const_loader_module.cc
+ * \file src/runtime/vm/const_loader_module.cc
* \brief A wrapper for initializing imported modules using constant Tensor. This
- * module is intended to be used by various runtime in the TVM stack, i.e.
- * graph executor, relax VM, AOT runtime, and various user defined runtimes. It
- * paves the way to separate the code and metedata, which makes compilation
- * and/or interpretation more convenient. In addition, the clear separation of
- * code and constants significantly reduces the efforts for handling external
- * codegen and runtimes.
+ * module belongs to the VM runtime and separates executable code from the
+ * constants used to initialize external modules.
*/
#include "const_loader_module.h"
@@ -45,7 +41,7 @@
#include <unordered_map>
#include <vector>
-#include "../support/bytes_io.h"
+#include "../../support/bytes_io.h"
namespace tvm {
namespace runtime {
diff --git a/src/runtime/const_loader_module.h b/src/runtime/vm/const_loader_module.h
similarity index 85%
rename from src/runtime/const_loader_module.h
rename to src/runtime/vm/const_loader_module.h
index 977d9bea49..5195ae58b9 100644
--- a/src/runtime/const_loader_module.h
+++ b/src/runtime/vm/const_loader_module.h
@@ -22,8 +22,8 @@
* \brief Defines an interface to use the ConstLoaderModule.
*/
-#ifndef TVM_RUNTIME_CONST_LOADER_MODULE_H_
-#define TVM_RUNTIME_CONST_LOADER_MODULE_H_
+#ifndef TVM_RUNTIME_VM_CONST_LOADER_MODULE_H_
+#define TVM_RUNTIME_VM_CONST_LOADER_MODULE_H_
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/container/map.h>
@@ -46,19 +46,19 @@ namespace runtime {
* \return The created ConstLoaderModule.
*
* Dispatches through the FFI registry ("ffi.Module.create.const_loader").
- * The creator is always available (ConstLoaderModule is a runtime-universal module).
+ * The creator is registered by the VM runtime.
*/
inline ffi::Module ConstLoaderModuleCreate(
const ffi::Map<ffi::String, ffi::Tensor>& const_var_tensor,
const ffi::Map<ffi::String, ffi::Array<ffi::String>>& const_vars_by_symbol) {
static const auto fcreate = ffi::Function::GetGlobal("ffi.Module.create.const_loader");
TVM_FFI_CHECK(fcreate.has_value(), RuntimeError)
- << "ffi.Module.create.const_loader is not registered in runtime. "
- << "Ensure libtvm_runtime is loaded.";
+ << "ffi.Module.create.const_loader is not registered by the VM runtime. "
+ << "Ensure the runtime includes VM support.";
return (*fcreate)(const_var_tensor, const_vars_by_symbol).cast<ffi::Module>();
}
} // namespace runtime
} // namespace tvm
-#endif // TVM_RUNTIME_CONST_LOADER_MODULE_H_
+#endif // TVM_RUNTIME_VM_CONST_LOADER_MODULE_H_
diff --git a/src/s_tir/meta_schedule/space_generator/space_generator.cc b/src/s_tir/meta_schedule/space_generator/space_generator.cc
index 0760fb940d..2469f9a1a5 100644
--- a/src/s_tir/meta_schedule/space_generator/space_generator.cc
+++ b/src/s_tir/meta_schedule/space_generator/space_generator.cc
@@ -19,7 +19,7 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/logging.h>
-#include "../../../target/canonicalizer/llvm/arm_aprofile.h"
+#include "../../../target/llvm/canonicalizer/arm_aprofile.h"
#include "../utils.h"
namespace tvm {
diff --git a/src/s_tir/transform/decorate_device_scope.cc b/src/s_tir/transform/decorate_device_scope.cc
index 10ee737e8f..c8333539f0 100644
--- a/src/s_tir/transform/decorate_device_scope.cc
+++ b/src/s_tir/transform/decorate_device_scope.cc
@@ -20,29 +20,22 @@
/*!
* \file decorate_device_scope.cc
*/
+#include "../../tirx/transform/decorate_device_scope.h"
+
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/prim/op.h>
#include <tvm/s_tir/transform.h>
-#include <tvm/tirx/stmt.h>
namespace tvm {
namespace s_tir {
using namespace tvm::tirx;
-Stmt DecorateDeviceScopeImpl(Stmt&& stmt) {
- static const Op device_scope = Op::Get("tirx.device_scope");
- return RegionStmt(device_scope, {}, {}, DictAttrs(), std::move(stmt));
-}
-
namespace transform {
Pass DecorateDeviceScope() {
auto pass_func = [](Function f, IRModule m, PassContext ctx) {
- if (!f->body.has_value()) return f;
- auto* n = f.CopyOnWrite();
- n->body = DecorateDeviceScopeImpl(std::move(n->body).value());
- return f;
+ return tirx::DecorateDeviceScopeImpl(std::move(f));
};
return CreateFunctionPass(pass_func, 0, "s_tir.DecorateDeviceScope");
}
diff --git a/src/sym/interval_set.h b/src/sym/interval_set.h
index 11441dbe3d..1aee24ae6e 100644
--- a/src/sym/interval_set.h
+++ b/src/sym/interval_set.h
@@ -96,7 +96,7 @@ class IntervalSet : public IntSet {
* \param max_value The maximum value in the interval.
* \return The created set.
*/
- TVM_DLL IntervalSet(PrimExpr min_value, PrimExpr max_value);
+ TVM_TIRX_DLL IntervalSet(PrimExpr min_value, PrimExpr max_value);
/*!
* \brief Create an IntervalSet that represents a single point.
@@ -127,7 +127,7 @@ class IntervalSet : public IntSet {
* \param b The second set.
* \return The result set.
*/
-TVM_DLL IntervalSet Union(AnalyzerObj* analyzer, IntervalSet a, IntervalSet b);
+TVM_TIRX_DLL IntervalSet Union(AnalyzerObj* analyzer, IntervalSet a, IntervalSet b);
/*!
* \brief Create insersection of two IntervalSets.
@@ -136,7 +136,7 @@ TVM_DLL IntervalSet Union(AnalyzerObj* analyzer, IntervalSet a, IntervalSet b);
* \param b The second set.
* \return The result set.
*/
-TVM_DLL IntervalSet Intersect(AnalyzerObj* analzyer, IntervalSet a, IntervalSet b);
+TVM_TIRX_DLL IntervalSet Intersect(AnalyzerObj* analzyer, IntervalSet a, IntervalSet b);
} // namespace sym
} // namespace tvm
diff --git a/src/sym/presburger_set.h b/src/sym/presburger_set.h
index 5fb4367de6..85074f211a 100644
--- a/src/sym/presburger_set.h
+++ b/src/sym/presburger_set.h
@@ -135,15 +135,15 @@ class PresburgerSet : public IntSet {
* \param vars The variables that the constraint describes about.
* \return The created PresburgerSet.
*/
- TVM_DLL PresburgerSet(const std::vector<IntegerRelation>& disjuncts,
- const ffi::Array<PrimVar>& vars);
+ TVM_TIRX_DLL PresburgerSet(const std::vector<IntegerRelation>& disjuncts,
+ const ffi::Array<PrimVar>& vars);
/*!
* \brief Make a new instance of PresburgerSet, collect all vars as space vars.
* \param constraint The constraint to construct the set.
* \return The created PresburgerSet.
*/
- TVM_DLL PresburgerSet(const PrimExpr& constraint);
+ TVM_TIRX_DLL PresburgerSet(const PrimExpr& constraint);
TVM_DEFINE_OBJECT_REF_COW_METHOD(PresburgerSetObj);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PresburgerSet, IntSet, PresburgerSetObj);
@@ -168,7 +168,7 @@ class PresburgerSet : public IntSet {
* \param constraint The constraint to construct the set.
* \return The created set.
*/
- TVM_DLL PresburgerSet(const PrimExpr& constraint) : IntSet(ffi::UnsafeInit{}) {
+ TVM_TIRX_DLL PresburgerSet(const PrimExpr& constraint) : IntSet(ffi::UnsafeInit{}) {
TVM_FFI_THROW(InternalError) << "MLIR is not enabled!";
}
};
diff --git a/src/sym/z3_prover.cc b/src/sym/z3_prover.cc
index 0a4a0b2b8b..137e014556 100644
--- a/src/sym/z3_prover.cc
+++ b/src/sym/z3_prover.cc
@@ -968,12 +968,12 @@ class Z3Prover::Impl : tvm::ExprFunctor<z3::expr(const Expr&)> {
}
};
-TVM_DLL bool Z3Prover::IsEnabled() const { return true; }
-TVM_DLL bool Z3Prover::CanProve(const PrimExpr& expr) { return impl_->CanProve(expr); }
-TVM_DLL void Z3Prover::Bind(const Var& var, const Range& new_range, bool allow_override) {
+TVM_TIRX_DLL bool Z3Prover::IsEnabled() const { return true; }
+TVM_TIRX_DLL bool Z3Prover::CanProve(const PrimExpr& expr) { return impl_->CanProve(expr); }
+TVM_TIRX_DLL void Z3Prover::Bind(const Var& var, const Range& new_range, bool allow_override) {
return impl_->Bind(var, new_range, allow_override);
}
-TVM_DLL void Z3Prover::Bind(const Var& var, const PrimExpr& expr, bool allow_override) {
+TVM_TIRX_DLL void Z3Prover::Bind(const Var& var, const PrimExpr& expr, bool allow_override) {
return impl_->Bind(var, expr, allow_override);
}
std::function<void()> Z3Prover::EnterConstraint(const PrimExpr& constraint, bool is_assume) {
@@ -991,12 +991,12 @@ void Z3Prover::SetRLimit(unsigned max_step) { impl_->SetRLimit(max_step); }
void Z3Prover::CopyFrom(const Z3Prover& other) { impl_->CopyFrom(*other.impl_); }
ffi::String Z3Prover::GetStats() { return impl_->GetStats(); }
ffi::String Z3Prover::GetModel(const PrimExpr& expr) { return impl_->GetModel(expr); }
-TVM_DLL int64_t Z3Prover::CountSatisfyingValues(const Var& var, int64_t max_count,
- int64_t min_consecutive) {
+TVM_TIRX_DLL int64_t Z3Prover::CountSatisfyingValues(const Var& var, int64_t max_count,
+ int64_t min_consecutive) {
return impl_->CountSatisfyingValues(var, max_count, min_consecutive);
}
Z3Prover::Z3Prover(AnalyzerObj* parent) : impl_(std::make_unique<Impl>(parent)) {}
-TVM_DLL Z3Prover::~Z3Prover() = default;
+TVM_TIRX_DLL Z3Prover::~Z3Prover() = default;
} // namespace tvm::sym
@@ -1016,10 +1016,10 @@ using namespace ffi;
// conservatively report "cannot prove" while keeping the public API available.
class Z3Prover::Impl {};
-TVM_DLL bool Z3Prover::IsEnabled() const { return false; }
-TVM_DLL bool Z3Prover::CanProve(const PrimExpr& expr) { return false; }
-TVM_DLL void Z3Prover::Bind(const Var& var, const Range& new_range, bool allow_override) {}
-TVM_DLL void Z3Prover::Bind(const Var& var, const PrimExpr& expr, bool allow_override) {}
+TVM_TIRX_DLL bool Z3Prover::IsEnabled() const { return false; }
+TVM_TIRX_DLL bool Z3Prover::CanProve(const PrimExpr& expr) { return false; }
+TVM_TIRX_DLL void Z3Prover::Bind(const Var& var, const Range& new_range, bool allow_override) {}
+TVM_TIRX_DLL void Z3Prover::Bind(const Var& var, const PrimExpr& expr, bool allow_override) {}
std::function<void()> Z3Prover::EnterConstraint(const PrimExpr& constraint, bool is_assume) {
return []() {};
}
@@ -1029,15 +1029,15 @@ ffi::String Z3Prover::GetSMTLIB2(const ffi::Optional<PrimExpr> expr) {
void Z3Prover::SetTimeoutMs(unsigned timeout_ms) {}
void Z3Prover::SetRLimit(unsigned rlimit) {}
ffi::String Z3Prover::GetModel(const PrimExpr& expr) { return "; Z3 Prover is disabled."; }
-TVM_DLL int64_t Z3Prover::CountSatisfyingValues(const Var& var, int64_t max_count,
- int64_t min_consecutive) {
+TVM_TIRX_DLL int64_t Z3Prover::CountSatisfyingValues(const Var& var, int64_t max_count,
+ int64_t min_consecutive) {
return -1; // Z3 disabled, return error
}
void Z3Prover::CopyFrom(const Z3Prover& other) {}
ffi::String Z3Prover::GetStats() { return "; Z3 Prover is disabled."; }
Z3Prover::Z3Prover(AnalyzerObj*) : impl_(nullptr) {}
-TVM_DLL Z3Prover::~Z3Prover() = default;
+TVM_TIRX_DLL Z3Prover::~Z3Prover() = default;
} // namespace tvm::sym
diff --git a/src/target/canonicalizer/llvm/arm_aprofile.cc b/src/target/llvm/canonicalizer/arm_aprofile.cc
similarity index 98%
rename from src/target/canonicalizer/llvm/arm_aprofile.cc
rename to src/target/llvm/canonicalizer/arm_aprofile.cc
index 6f6e6eb7cc..e08fccd0b4 100644
--- a/src/target/canonicalizer/llvm/arm_aprofile.cc
+++ b/src/target/llvm/canonicalizer/arm_aprofile.cc
@@ -18,7 +18,7 @@
*/
/*!
- * \file tvm/target/canonicalizer/llvm/arm_aprofile.cc
+ * \file tvm/target/llvm/canonicalizer/arm_aprofile.cc
* \brief Target canonicalizer for Arm(R) Cortex(R) A-Profile CPUs
*/
@@ -30,7 +30,7 @@
#include <string>
#include "../../../support/utils.h"
-#include "../../llvm/llvm_instance.h"
+#include "../llvm_instance.h"
namespace tvm {
namespace target {
diff --git a/src/target/canonicalizer/llvm/arm_aprofile.h b/src/target/llvm/canonicalizer/arm_aprofile.h
similarity index 85%
rename from src/target/canonicalizer/llvm/arm_aprofile.h
rename to src/target/llvm/canonicalizer/arm_aprofile.h
index 0471bedf8d..eee42b547c 100644
--- a/src/target/canonicalizer/llvm/arm_aprofile.h
+++ b/src/target/llvm/canonicalizer/arm_aprofile.h
@@ -18,12 +18,12 @@
*/
/*!
- * \file tvm/target/canonicalizer/llvm/arm_aprofile.h
+ * \file tvm/target/llvm/canonicalizer/arm_aprofile.h
* \brief Target canonicalizer for Arm(R) Cortex(R) A-Profile CPUs
*/
-#ifndef TVM_TARGET_CANONICALIZER_LLVM_ARM_APROFILE_H_
-#define TVM_TARGET_CANONICALIZER_LLVM_ARM_APROFILE_H_
+#ifndef TVM_TARGET_LLVM_CANONICALIZER_ARM_APROFILE_H_
+#define TVM_TARGET_LLVM_CANONICALIZER_ARM_APROFILE_H_
#include <tvm/target/target.h>
@@ -42,4 +42,4 @@ ffi::Map<ffi::String, ffi::Any> Canonicalize(ffi::Map<ffi::String, ffi::Any> tar
} // namespace target
} // namespace tvm
-#endif // TVM_TARGET_CANONICALIZER_LLVM_ARM_APROFILE_H_
+#endif // TVM_TARGET_LLVM_CANONICALIZER_ARM_APROFILE_H_
diff --git a/src/target/canonicalizer/llvm/arm_mprofile.cc b/src/target/llvm/canonicalizer/arm_mprofile.cc
similarity index 98%
rename from src/target/canonicalizer/llvm/arm_mprofile.cc
rename to src/target/llvm/canonicalizer/arm_mprofile.cc
index 8b8753337b..7dfe9ea6ec 100644
--- a/src/target/canonicalizer/llvm/arm_mprofile.cc
+++ b/src/target/llvm/canonicalizer/arm_mprofile.cc
@@ -18,7 +18,7 @@
*/
/*!
- * \file tvm/target/canonicalizer/llvm/arm_mprofile.cc
+ * \file tvm/target/llvm/canonicalizer/arm_mprofile.cc
* \brief Target canonicalizer for Arm(R) Cortex(R) M-Profile CPUs
*/
diff --git a/src/target/canonicalizer/llvm/arm_mprofile.h b/src/target/llvm/canonicalizer/arm_mprofile.h
similarity index 85%
rename from src/target/canonicalizer/llvm/arm_mprofile.h
rename to src/target/llvm/canonicalizer/arm_mprofile.h
index ed87c93c2f..755488bef7 100644
--- a/src/target/canonicalizer/llvm/arm_mprofile.h
+++ b/src/target/llvm/canonicalizer/arm_mprofile.h
@@ -18,12 +18,12 @@
*/
/*!
- * \file tvm/target/canonicalizer/llvm/arm_mprofile.h
+ * \file tvm/target/llvm/canonicalizer/arm_mprofile.h
* \brief Target canonicalizer for Arm(R) Cortex(R) M-Profile CPUs
*/
-#ifndef TVM_TARGET_CANONICALIZER_LLVM_ARM_MPROFILE_H_
-#define TVM_TARGET_CANONICALIZER_LLVM_ARM_MPROFILE_H_
+#ifndef TVM_TARGET_LLVM_CANONICALIZER_ARM_MPROFILE_H_
+#define TVM_TARGET_LLVM_CANONICALIZER_ARM_MPROFILE_H_
#include <tvm/target/target.h>
@@ -42,4 +42,4 @@ ffi::Map<ffi::String, ffi::Any> Canonicalize(ffi::Map<ffi::String, ffi::Any> tar
} // namespace target
} // namespace tvm
-#endif // TVM_TARGET_CANONICALIZER_LLVM_ARM_MPROFILE_H_
+#endif // TVM_TARGET_LLVM_CANONICALIZER_ARM_MPROFILE_H_
diff --git a/src/target/canonicalizer/llvm/canonicalize.cc b/src/target/llvm/canonicalizer/canonicalize.cc
similarity index 100%
rename from src/target/canonicalizer/llvm/canonicalize.cc
rename to src/target/llvm/canonicalizer/canonicalize.cc
diff --git a/src/target/canonicalizer/llvm/canonicalize.h b/src/target/llvm/canonicalizer/canonicalize.h
similarity index 84%
rename from src/target/canonicalizer/llvm/canonicalize.h
rename to src/target/llvm/canonicalizer/canonicalize.h
index 622c9767b1..9a8404b6b2 100644
--- a/src/target/canonicalizer/llvm/canonicalize.h
+++ b/src/target/llvm/canonicalizer/canonicalize.h
@@ -18,12 +18,12 @@
*/
/*!
- * \file tvm/target/canonicalizer/llvm/canonicalize.h
+ * \file tvm/target/llvm/canonicalizer/canonicalize.h
* \brief Target canonicalizer for LLVM-based CPU targets
*/
-#ifndef TVM_TARGET_CANONICALIZER_LLVM_CANONICALIZE_H_
-#define TVM_TARGET_CANONICALIZER_LLVM_CANONICALIZE_H_
+#ifndef TVM_TARGET_LLVM_CANONICALIZER_CANONICALIZE_H_
+#define TVM_TARGET_LLVM_CANONICALIZER_CANONICALIZE_H_
#include <tvm/target/target.h>
@@ -39,4 +39,4 @@ ffi::Map<ffi::String, ffi::Any> Canonicalize(ffi::Map<ffi::String, ffi::Any> tar
} // namespace target
} // namespace tvm
-#endif // TVM_TARGET_CANONICALIZER_LLVM_CANONICALIZE_H_
+#endif // TVM_TARGET_LLVM_CANONICALIZER_CANONICALIZE_H_
diff --git a/src/target/target_kind.cc b/src/target/target_kind.cc
index c9251e6787..97687beeb8 100644
--- a/src/target/target_kind.cc
+++ b/src/target/target_kind.cc
@@ -33,7 +33,7 @@
#include <algorithm>
#include "../support/utils.h"
-#include "./canonicalizer/llvm/canonicalize.h"
+#include "./llvm/canonicalizer/canonicalize.h"
namespace tvm {
diff --git a/src/tirx/ir/transform.cc b/src/tirx/ir/transform.cc
index 7e8dd35e18..bc613a2cf6 100644
--- a/src/tirx/ir/transform.cc
+++ b/src/tirx/ir/transform.cc
@@ -95,7 +95,7 @@ class FunctionPass : public Pass {
* \param pass_func The packed function which implements a pass.
* \param pass_info The pass info.
*/
- TVM_DLL FunctionPass(
+ TVM_TIRX_DLL FunctionPass(
std::function<ffi::Optional<Function>(Function, IRModule, PassContext)> pass_func,
PassInfo pass_info);
diff --git a/src/s_tir/transform/decorate_device_scope.cc b/src/tirx/transform/decorate_device_scope.cc
similarity index 71%
copy from src/s_tir/transform/decorate_device_scope.cc
copy to src/tirx/transform/decorate_device_scope.cc
index 10ee737e8f..ce83759602 100644
--- a/src/s_tir/transform/decorate_device_scope.cc
+++ b/src/tirx/transform/decorate_device_scope.cc
@@ -20,38 +20,39 @@
/*!
* \file decorate_device_scope.cc
*/
+#include "decorate_device_scope.h"
+
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/prim/op.h>
-#include <tvm/s_tir/transform.h>
#include <tvm/tirx/stmt.h>
+#include <tvm/tirx/transform.h>
namespace tvm {
-namespace s_tir {
-using namespace tvm::tirx;
+namespace tirx {
-Stmt DecorateDeviceScopeImpl(Stmt&& stmt) {
+Function DecorateDeviceScopeImpl(Function f) {
+ if (!f->body.has_value()) return f;
static const Op device_scope = Op::Get("tirx.device_scope");
- return RegionStmt(device_scope, {}, {}, DictAttrs(), std::move(stmt));
+ auto* n = f.CopyOnWrite();
+ n->body = RegionStmt(device_scope, {}, {}, DictAttrs(), std::move(n->body).value());
+ return f;
}
namespace transform {
Pass DecorateDeviceScope() {
auto pass_func = [](Function f, IRModule m, PassContext ctx) {
- if (!f->body.has_value()) return f;
- auto* n = f.CopyOnWrite();
- n->body = DecorateDeviceScopeImpl(std::move(n->body).value());
- return f;
+ return DecorateDeviceScopeImpl(std::move(f));
};
- return CreateFunctionPass(pass_func, 0, "s_tir.DecorateDeviceScope");
+ return CreateFunctionPass(pass_func, 0, "tirx.DecorateDeviceScope");
}
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("s_tir.transform.DecorateDeviceScope", DecorateDeviceScope);
+ refl::GlobalDef().def("tirx.transform.DecorateDeviceScope", DecorateDeviceScope);
}
} // namespace transform
-} // namespace s_tir
+} // namespace tirx
} // namespace tvm
diff --git a/src/ir/prim/op_utils.h b/src/tirx/transform/decorate_device_scope.h
similarity index 66%
copy from src/ir/prim/op_utils.h
copy to src/tirx/transform/decorate_device_scope.h
index b33a9be44d..f1590ee669 100644
--- a/src/ir/prim/op_utils.h
+++ b/src/tirx/transform/decorate_device_scope.h
@@ -17,11 +17,18 @@
* under the License.
*/
-#ifndef TVM_IR_PRIM_OP_UTILS_H_
-#define TVM_IR_PRIM_OP_UTILS_H_
-#include <tvm/ir/prim/op.h>
-namespace tvm::prim::detail {
-TVM_DLL void BinaryOpMatchTypes(PrimExpr& lhs, PrimExpr& rhs, Location loc);
+/*! \file decorate_device_scope.h
+ * \brief Shared implementation of native device-scope decoration.
+ */
+#ifndef TVM_TIRX_TRANSFORM_DECORATE_DEVICE_SCOPE_H_
+#define TVM_TIRX_TRANSFORM_DECORATE_DEVICE_SCOPE_H_
+
+#include <tvm/tirx/function.h>
+
+namespace tvm {
+namespace tirx {
+Function DecorateDeviceScopeImpl(Function f);
+} // namespace tirx
+} // namespace tvm
-} // namespace tvm::prim::detail
-#endif // TVM_IR_PRIM_OP_UTILS_H_
+#endif // TVM_TIRX_TRANSFORM_DECORATE_DEVICE_SCOPE_H_
diff --git a/tests/python/codegen/test_codegen_assert.py b/tests/python/codegen/test_codegen_assert.py
index d997d96550..17fb026328 100644
--- a/tests/python/codegen/test_codegen_assert.py
+++ b/tests/python/codegen/test_codegen_assert.py
@@ -32,7 +32,7 @@ def test_assert_runtime_error(codegen_target):
def func(x: T.int32):
assert x > 0, ("RuntimeError", ["Expected non-null input"])
- lib = tvm.compile(func, target=codegen_target)
+ lib = tvm.tirx.compile(func, target=codegen_target)
with pytest.raises(RuntimeError, match="Expected non-null input"):
lib(0)
@@ -44,7 +44,7 @@ def test_assert_value_error(codegen_target):
def func(x: T.int32):
assert x > 0, ("ValueError", ["Shape mismatch: expected 4 got 8"])
- lib = tvm.compile(func, target=codegen_target)
+ lib = tvm.tirx.compile(func, target=codegen_target)
with pytest.raises(ValueError, match="Shape mismatch"):
lib(0)
@@ -56,7 +56,7 @@ def test_assert_type_error(codegen_target):
def func(x: T.int32):
assert x > 0, ("TypeError", ["Expected Tensor but got int"])
- lib = tvm.compile(func, target=codegen_target)
+ lib = tvm.tirx.compile(func, target=codegen_target)
with pytest.raises(TypeError, match="Expected Tensor but got int"):
lib(0)
@@ -68,7 +68,7 @@ def test_assert_multi_part_message(codegen_target):
def func(x: T.int32):
assert x > 0, ("ValueError", ["Expected shape ", "4", " but got ", "8"])
- lib = tvm.compile(func, target=codegen_target)
+ lib = tvm.tirx.compile(func, target=codegen_target)
with pytest.raises(ValueError, match="Expected shape 4 but got 8"):
lib(0)
@@ -80,7 +80,7 @@ def test_assert_passing_condition(codegen_target):
def func(x: T.int32):
assert x > 0, ("RuntimeError", ["This should not be raised"])
- lib = tvm.compile(func, target=codegen_target)
+ lib = tvm.tirx.compile(func, target=codegen_target)
lib(1) # should pass without error
@@ -91,7 +91,7 @@ def test_assert_many_parts(codegen_target):
def func(x: T.int32):
assert x > 0, ("RuntimeError", ["p0", "p1", "p2", "p3", "p4", "p5", "p6", "p7"])
- lib = tvm.compile(func, target=codegen_target)
+ lib = tvm.tirx.compile(func, target=codegen_target)
with pytest.raises(RuntimeError, match="p0p1p2p3p4p5p6p7"):
lib(0)
@@ -103,7 +103,7 @@ def test_tvmscript_assert_preserves_kind(codegen_target):
def func(x: T.int32):
assert x > 0, ("ValueError", ["x must be positive"])
- lib = tvm.compile(func, target=codegen_target)
+ lib = tvm.tirx.compile(func, target=codegen_target)
with pytest.raises(ValueError, match="x must be positive"):
lib(0)
@@ -115,7 +115,7 @@ def test_tvmscript_assert_preserves_parts(codegen_target):
def func(x: T.int32):
assert x > 0, ("ValueError", ["x must be ", "positive"])
- lib = tvm.compile(func, target=codegen_target)
+ lib = tvm.tirx.compile(func, target=codegen_target)
with pytest.raises(ValueError, match="x must be positive"):
lib(0)
diff --git a/tests/python/codegen/test_codegen_error_handling.py b/tests/python/codegen/test_codegen_error_handling.py
index 6f5a2fad3b..c63dec7aa5 100644
--- a/tests/python/codegen/test_codegen_error_handling.py
+++ b/tests/python/codegen/test_codegen_error_handling.py
@@ -47,7 +47,7 @@ def test_wrong_argument_count_error(codegen_target):
for i in range(n0):
B[i] = A[i] + T.float32(1)
- lib = tvm.compile(func, target=codegen_target)
+ lib = tvm.tirx.compile(func, target=codegen_target)
a = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
b = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
lib(a, b) # correct input should pass
@@ -75,7 +75,7 @@ def test_type_mismatch_non_tensor(codegen_target):
for i in range(n0):
B[i] = A[i] + T.float32(1)
- lib = tvm.compile(func, target=codegen_target)
+ lib = tvm.tirx.compile(func, target=codegen_target)
a = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
b = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
lib(a, b) # correct input should pass
@@ -104,7 +104,7 @@ def test_shape_mismatch_shared_variable(codegen_target):
for i in range(n0):
B[i] = A[i] + T.float32(1)
- lib = tvm.compile(func, target=codegen_target)
+ lib = tvm.tirx.compile(func, target=codegen_target)
a = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
b = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
lib(a, b) # correct input should pass
@@ -129,7 +129,7 @@ def test_invalid_shape_fixed(codegen_target):
for i in range(128):
b[i] = a[i] + T.float32(1)
- lib = tvm.compile(func, target=codegen_target)
+ lib = tvm.tirx.compile(func, target=codegen_target)
a = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
b = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
lib(a, b) # correct input should pass
@@ -158,7 +158,7 @@ def test_ndim_mismatch_error(codegen_target):
for i, j in T.grid(4, 8):
b[i, j] = a[i, j]
- lib = tvm.compile(func, target=codegen_target)
+ lib = tvm.tirx.compile(func, target=codegen_target)
a_ok = tvm.runtime.tensor(np.zeros((4, 8), dtype="float32"))
b_ok = tvm.runtime.tensor(np.zeros((4, 8), dtype="float32"))
lib(a_ok, b_ok) # correct input should pass
@@ -187,7 +187,7 @@ def test_dtype_mismatch_error(codegen_target):
for i in range(8):
b[i] = a[i]
- lib = tvm.compile(func, target=codegen_target)
+ lib = tvm.tirx.compile(func, target=codegen_target)
a_ok = tvm.runtime.tensor(np.zeros(8, dtype="float32"))
b_ok = tvm.runtime.tensor(np.zeros(8, dtype="float32"))
lib(a_ok, b_ok) # correct input should pass
@@ -216,7 +216,7 @@ def test_data_alignment_error(codegen_target):
for i in range(128):
b[i] = a[i] + T.float32(1)
- lib = tvm.compile(func, target=codegen_target)
+ lib = tvm.tirx.compile(func, target=codegen_target)
a_ok = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
b_ok = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
lib(a_ok, b_ok) # correct input should pass
@@ -248,7 +248,7 @@ def test_strides_mismatch_transposed(codegen_target):
for i, j in T.grid(128, 128):
b[i, j] = a[i, j] + T.float32(1)
- lib = tvm.compile(func, target=codegen_target)
+ lib = tvm.tirx.compile(func, target=codegen_target)
a_ok = tvm.runtime.tensor(np.zeros((128, 128), dtype="float32"))
b_ok = tvm.runtime.tensor(np.zeros((128, 128), dtype="float32"))
lib(a_ok, b_ok) # correct input should pass
@@ -282,7 +282,7 @@ def test_device_mismatch_error():
for i in range(128):
b[i] = a[i] + T.float32(1)
- lib = tvm.compile(func, target="llvm")
+ lib = tvm.tirx.compile(func, target="llvm")
a_ok = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
b_ok = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
lib(a_ok, b_ok) # correct input should pass
@@ -317,7 +317,7 @@ def test_type_mismatch_int_parameter(codegen_target):
else:
return 20
- lib = tvm.compile(func, target=codegen_target)
+ lib = tvm.tirx.compile(func, target=codegen_target)
assert lib(5) == 10 # correct input should pass
a = tvm.runtime.tensor(np.zeros(8, dtype="float32"))
@@ -340,7 +340,7 @@ def test_type_mismatch_float_parameter(codegen_target):
else:
return 0
- lib = tvm.compile(func, target=codegen_target)
+ lib = tvm.tirx.compile(func, target=codegen_target)
assert lib(1.0) == 1 # correct input should pass
a = tvm.runtime.tensor(np.zeros(8, dtype="float32"))
@@ -363,7 +363,7 @@ def test_type_mismatch_bool_parameter(codegen_target):
else:
return 0
- lib = tvm.compile(func, target=codegen_target)
+ lib = tvm.tirx.compile(func, target=codegen_target)
assert lib(True) == 1 # correct input should pass
a = tvm.runtime.tensor(np.zeros(8, dtype="float32"))
@@ -395,7 +395,7 @@ def test_forward_reference_symbolic_shape(codegen_target):
for i in range(batch_size):
B[i] = A[i] + A[i + 1]
- lib = tvm.compile(func, target=codegen_target)
+ lib = tvm.tirx.compile(func, target=codegen_target)
# Correct inputs: A has shape (5,), B has shape (4,)
a = tvm.runtime.tensor(np.array([1, 2, 3, 4, 5], dtype="int32"))
b = tvm.runtime.tensor(np.zeros(4, dtype="int32"))
@@ -427,7 +427,7 @@ def test_invalid_arguments_mixed_params(codegen_target):
def func(a0: T.bool, a1: T.Tensor([10], "float32")) -> T.int32:
return 0
- lib = tvm.compile(func, target=codegen_target)
+ lib = tvm.tirx.compile(func, target=codegen_target)
lib(True, tvm.runtime.tensor(np.zeros(10, dtype="float32"))) # correct input should pass
with pytest.raises(
diff --git a/tests/python/codegen/test_gpu_codegen_allreduce.py b/tests/python/codegen/test_gpu_codegen_allreduce.py
index 42240ac081..8f425e3fce 100644
--- a/tests/python/codegen/test_gpu_codegen_allreduce.py
+++ b/tests/python/codegen/test_gpu_codegen_allreduce.py
@@ -88,7 +88,7 @@ def test_allreduce_sum(dims, target):
pytest.skip(f"{target} not enabled")
d1, d2, d3 = dims
mod = _reduce_sum_module(d1, d2, d3)
- f = tvm.compile(mod, target=target)
+ f = tvm.tirx.compile(mod, target=target)
# prepare input and output array
a_np = np.random.rand(1, d1, d2, d3).astype("float32")
@@ -138,7 +138,7 @@ def test_allreduce_sum_compile(optional_metal_compile_callback):
d1, d2, d3 = dims
mod = _reduce_sum_module(d1, d2, d3)
- tvm.compile(mod, target=target)
+ tvm.tirx.compile(mod, target=target)
@pytest.mark.parametrize(
@@ -153,7 +153,7 @@ def test_allreduce_max(dims, target):
pytest.skip(f"{target} not enabled")
d1, d2, d3 = dims
mod = _reduce_max_module(d1, d2, d3)
- f = tvm.compile(mod, target=target)
+ f = tvm.tirx.compile(mod, target=target)
# prepare input and output array
a_np = -np.random.rand(1, d1, d2, d3).astype("float32")
diff --git a/tests/python/codegen/test_target_codegen.py b/tests/python/codegen/test_target_codegen.py
index 23f7aa713c..090cb93162 100644
--- a/tests/python/codegen/test_target_codegen.py
+++ b/tests/python/codegen/test_target_codegen.py
@@ -42,7 +42,7 @@ def test_tensor_store_predicate_not_supported():
err_msg = "Predicated buffer store is not supported."
with pytest.raises(RuntimeError, match=err_msg):
with tvm.target.Target(target):
- tvm.compile(func)
+ tvm.tirx.compile(func)
@pytest.mark.parametrize(
@@ -77,7 +77,7 @@ def test_tensor_store_predicate_not_supported_gpu(target):
err_msg = "Predicated buffer store is not supported."
with pytest.raises(RuntimeError, match=err_msg):
with tvm.target.Target(target):
- tvm.compile(func)
+ tvm.tirx.compile(func)
def test_buffer_load_predicate_not_supported():
@@ -97,7 +97,7 @@ def test_buffer_load_predicate_not_supported():
err_msg = "Predicated buffer load is not supported."
with pytest.raises(RuntimeError, match=err_msg):
with tvm.target.Target(target):
- tvm.compile(func)
+ tvm.tirx.compile(func)
@pytest.mark.parametrize(
@@ -128,7 +128,7 @@ def test_buffer_load_predicate_not_supported_gpu(target):
err_msg = "Predicated buffer load is not supported."
with pytest.raises(RuntimeError, match=err_msg):
with tvm.target.Target(target):
- tvm.compile(func)
+ tvm.tirx.compile(func)
@pytest.mark.parametrize(
@@ -173,7 +173,7 @@ def test_codegen_loop_step(target):
for i in T.serial(3, 1024, step=96):
C[i] = A[i] + B[i]
- lib = tvm.compile(test_loop_step, target=target)
+ lib = tvm.tirx.compile(test_loop_step, target=target)
src = lib.mod.inspect_source()
if target == "c":
diff --git a/tests/python/codegen/test_target_codegen_aarch64.py b/tests/python/codegen/test_target_codegen_aarch64.py
index 9b793527fd..b85c490edd 100644
--- a/tests/python/codegen/test_target_codegen_aarch64.py
+++ b/tests/python/codegen/test_target_codegen_aarch64.py
@@ -55,7 +55,7 @@ def test_mul(dtype):
C[i] = A[i] * B[i]
with tvm.target.Target(target):
- f = tvm.tirx.build(Module)
+ f = tvm.tirx.compile(Module).mod
assembly = f.inspect_source("asm")
loads = re.findall("ld1[whdb] { z", assembly)
@@ -93,7 +93,7 @@ def test_add(dtype):
C[i] = A[i] + B[i]
with tvm.target.Target(target):
- f = tvm.tirx.build(Module)
+ f = tvm.tirx.compile(Module).mod
assembly = f.inspect_source("asm")
loads = re.findall("ld1[whdb] { z", assembly)
@@ -131,7 +131,7 @@ def test_sub(dtype):
C[i] = A[i] - B[i]
with tvm.target.Target(target):
- f = tvm.tirx.build(Module)
+ f = tvm.tirx.compile(Module).mod
assembly = f.inspect_source("asm")
loads = re.findall("ld1[whdb] { z", assembly)
@@ -170,7 +170,7 @@ def test_muladd(dtype):
D[i] = A[i] * B[i] + C[i]
with tvm.target.Target(target):
- f = tvm.tirx.build(Module)
+ f = tvm.tirx.compile(Module).mod
assembly = f.inspect_source("asm")
loads = re.findall("ld1[whdb] { z", assembly)
@@ -218,7 +218,7 @@ def test_max(dtype):
C[i] = T.max(A[i], B[i])
with tvm.target.Target(target):
- f = tvm.tirx.build(Module)
+ f = tvm.tirx.compile(Module).mod
assembly = f.inspect_source("asm")
loads = re.findall("ld1[whdb] { z", assembly)
@@ -260,7 +260,7 @@ def test_min(dtype):
C[i] = T.min(A[i], B[i])
with tvm.target.Target(target):
- f = tvm.tirx.build(Module)
+ f = tvm.tirx.compile(Module).mod
assembly = f.inspect_source("asm")
loads = re.findall("ld1[whdb] { z", assembly)
@@ -302,7 +302,7 @@ def test_div(dtype):
C[i] = tvm.tirx.div(A[i], B[i])
with tvm.target.Target(target):
- f = tvm.tirx.build(Module)
+ f = tvm.tirx.compile(Module).mod
assembly = f.inspect_source("asm")
loads = re.findall("ld1[whdb] { z", assembly)
@@ -339,7 +339,7 @@ def test_mod(dtype):
C[i] = T.floormod(A[i], B[i])
with tvm.target.Target(target):
- f = tvm.tirx.build(Module)
+ f = tvm.tirx.compile(Module).mod
assembly = f.inspect_source("asm")
loads = re.findall("ld1[whdb] { z", assembly)
@@ -377,7 +377,7 @@ def test_eq(dtype):
C[i] = A[i] == B[i]
with tvm.target.Target(target):
- f = tvm.tirx.build(Module)
+ f = tvm.tirx.compile(Module).mod
assembly = f.inspect_source("asm")
loads = re.findall("ld1[whdb] { z", assembly)
@@ -418,7 +418,7 @@ def test_neq(dtype):
C[i] = A[i] != B[i]
with tvm.target.Target(target):
- f = tvm.tirx.build(Module)
+ f = tvm.tirx.compile(Module).mod
assembly = f.inspect_source("asm")
loads = re.findall("ld1[whdb] { z", assembly)
@@ -458,7 +458,7 @@ def test_or(dtype):
C[i] = A[i] | B[i]
with tvm.target.Target(target):
- f = tvm.tirx.build(Module)
+ f = tvm.tirx.compile(Module).mod
assembly = f.inspect_source("asm")
loads = re.findall("ld1[whdb] { z", assembly)
@@ -495,7 +495,7 @@ def test_and(dtype):
C[i] = A[i] & B[i]
with tvm.target.Target(target):
- f = tvm.tirx.build(Module)
+ f = tvm.tirx.compile(Module).mod
assembly = f.inspect_source("asm")
loads = re.findall("ld1[whdb] { z", assembly)
@@ -528,7 +528,7 @@ def test_not(dtype):
C[i] = ~A[i]
with tvm.target.Target(target):
- f = tvm.tirx.build(Module)
+ f = tvm.tirx.compile(Module).mod
assembly = f.inspect_source("asm")
loads = re.findall("ld1[whdb] { z", assembly)
@@ -569,7 +569,7 @@ def test_memcpy(dtype):
C[i] = A[B[i]]
with tvm.target.Target(target):
- f = tvm.tirx.build(Module)
+ f = tvm.tirx.compile(Module).mod
assembly = f.inspect_source("asm")
loads = re.findall("ld1[whdb] { z", assembly)
@@ -608,7 +608,7 @@ def test_vscale_range_function_attribute(mattr, expect_attr):
C[i] = A[i] + T.float32(1)
with tvm.target.Target(target):
- f = tvm.tirx.build(Module)
+ f = tvm.tirx.compile(Module).mod
# Check if the vscale_range() attribute exists
ll = f.inspect_source("ll")
diff --git a/tests/python/codegen/test_target_codegen_arm.py b/tests/python/codegen/test_target_codegen_arm.py
index 1cdcf32f1b..d85a7a1b1a 100644
--- a/tests/python/codegen/test_target_codegen_arm.py
+++ b/tests/python/codegen/test_target_codegen_arm.py
@@ -38,7 +38,7 @@ def test_popcount():
for i in T.vectorized(elements):
B[i] = T.popcount(A[i])
- f = tvm.tirx.build(Module, target=target)
+ f = tvm.tirx.compile(Module, target=target).mod
# Verify we see the correct number of vpaddl and vcnt instructions in the assembly
assembly = f.inspect_source("asm")
matches = re.findall("vpaddl", assembly)
@@ -77,7 +77,7 @@ def test_vmlal_s16():
for rv in range(K):
C[n] = C[n] + T.Cast("int32", A[rv, n]) * T.Cast("int32", B[rv, n])
- f = tvm.tirx.build(Module, target=target)
+ f = tvm.tirx.compile(Module, target=target).mod
# Verify we see the correct number of vmlal.s16 instructions
assembly = f.inspect_source("asm")
@@ -105,7 +105,7 @@ def test_vmlal_s16():
for rv in range(K):
C[n] = C[n] + T.Cast("int32", A[rv, n]) * T.Cast("int32", B[rv])
- f = tvm.tirx.build(Module, target=target)
+ f = tvm.tirx.compile(Module, target=target).mod
# Verify we see the correct number of vmlal.s16 instructions
assembly = f.inspect_source("asm")
diff --git a/tests/python/codegen/test_target_codegen_blob.py b/tests/python/codegen/test_target_codegen_blob.py
index ee994df147..0b280a563b 100644
--- a/tests/python/codegen/test_target_codegen_blob.py
+++ b/tests/python/codegen/test_target_codegen_blob.py
@@ -63,8 +63,8 @@ def test_cuda_multi_lib():
temp = utils.tempdir()
target = tvm.target.Target("cuda", host="llvm")
- libA = tvm.compile(ModA, target=target)
- libB = tvm.compile(ModB, target=target)
+ libA = tvm.tirx.compile(ModA, target=target)
+ libB = tvm.tirx.compile(ModB, target=target)
pathA = temp.relpath("libA.tar")
pathB = temp.relpath("libB.tar")
diff --git a/tests/python/codegen/test_target_codegen_bool.py b/tests/python/codegen/test_target_codegen_bool.py
index d000c61998..6f9a8d5ebc 100644
--- a/tests/python/codegen/test_target_codegen_bool.py
+++ b/tests/python/codegen/test_target_codegen_bool.py
@@ -65,7 +65,7 @@ def test_cmp_load_store(target):
is_gpu = tvm.target.Target(target).kind.name != "llvm"
mod = GPUModule if is_gpu else CPUModule
- f = tvm.compile(mod, target=target)
+ f = tvm.tirx.compile(mod, target=target)
a_np = np.random.uniform(size=arr_size).astype("float32")
b_np = np.random.uniform(size=arr_size).astype("float32")
@@ -93,7 +93,7 @@ def test_bitwise_not_c(tmp_path):
for i in range(5):
output[i] = T.bitwise_not(values[i] != 0)
- built = tvm.compile(complement, target="c")
+ built = tvm.tirx.compile(complement, target="c")
library = str(tmp_path / "complement.so")
built.export_library(library)
built = tvm.runtime.load_module(library)
diff --git a/tests/python/codegen/test_target_codegen_c_host.py b/tests/python/codegen/test_target_codegen_c_host.py
index 5b3b34cb48..f90ba9c7a9 100644
--- a/tests/python/codegen/test_target_codegen_c_host.py
+++ b/tests/python/codegen/test_target_codegen_c_host.py
@@ -40,7 +40,7 @@ def test_add():
C[i0] = A[i0] + B[i0]
def check_c():
- mhost = tvm.compile(Module, target="c")
+ mhost = tvm.tirx.compile(Module, target="c")
temp = utils.tempdir()
path_dso = temp.relpath("temp.so")
mhost.export_library(path_dso)
@@ -72,7 +72,7 @@ def test_reinterpret():
B[i0] = T.reinterpret(A[i0] + 2, ty="float32")
def check_c():
- mhost = tvm.compile(Module, target="c")
+ mhost = tvm.tirx.compile(Module, target="c")
temp = utils.tempdir()
path_dso = temp.relpath("temp.so")
mhost.export_library(path_dso)
@@ -103,7 +103,7 @@ def test_ceil():
B[i0] = T.ceil(A[i0])
def check_c():
- mhost = tvm.compile(Module, target="c")
+ mhost = tvm.tirx.compile(Module, target="c")
temp = utils.tempdir()
path_dso = temp.relpath("temp.so")
mhost.export_library(path_dso)
@@ -134,7 +134,7 @@ def test_floor():
B[i0] = T.floor(A[i0])
def check_c():
- mhost = tvm.compile(Module, target="c")
+ mhost = tvm.tirx.compile(Module, target="c")
temp = utils.tempdir()
path_dso = temp.relpath("temp.so")
mhost.export_library(path_dso)
@@ -165,7 +165,7 @@ def test_round():
B[i0] = T.round(A[i0])
def check_c():
- mhost = tvm.compile(Module, target="c")
+ mhost = tvm.tirx.compile(Module, target="c")
temp = utils.tempdir()
path_dso = temp.relpath("temp.so")
mhost.export_library(path_dso)
@@ -201,7 +201,7 @@ def test_subroutine_call():
A = T.decl_tensor(A_data, ty_args=[T.Tensor(1, "float32", data=A_data)])
A[0] = 42.0
- built = tvm.tirx.build(Module, target="c")
+ built = tvm.tirx.compile(Module, target="c").mod
source = built.inspect_source()
assert source.count("__tvm_ffi_main(void*") == 2, (
@@ -226,7 +226,7 @@ def test_workspace_allocation_cast():
for i in range(256):
A[i] = workspace[i]
- built = tvm.tirx.build(Module, target="c")
+ built = tvm.tirx.compile(Module, target="c").mod
assert "((float*)TVMBackendAllocWorkspace(" in built.inspect_source()
temp = utils.tempdir()
@@ -243,7 +243,7 @@ def test_local_alloc_buffer_uses_plain_c_pointer():
B[i] = A[i] + T.float32(1)
A[i] = B[i]
- built = tvm.tirx.build(Module, target="c")
+ built = tvm.tirx.compile(Module, target="c").mod
assert "local float*" not in built.inspect_source()
temp = utils.tempdir()
@@ -262,7 +262,7 @@ def test_vector_pointer_uses_byte_offset():
body = tvm.ir.Evaluate(tvm.tirx.call_extern("void", "consume", pointer))
func = tvm.tirx.Function([buffer], body).with_attr("global_symbol", "main")
- source = tvm.tirx.build(tvm.IRModule.from_expr(func), target="c").inspect_source()
+ source = tvm.tirx.compile(tvm.IRModule.from_expr(func), target="c").mod.inspect_source()
call = next(line.strip() for line in source.splitlines() if line.strip().startswith("consume("))
assert "int32_t2" not in call
assert "float2*" in call
@@ -277,7 +277,7 @@ def test_if_then_else_avoids_extraneous_parentheses():
for i in range(8):
B[i] = T.if_then_else(i == 0, 1, A[i])
- built = tvm.tirx.build(Module, target="c")
+ built = tvm.tirx.compile(Module, target="c").mod
source = built.inspect_source()
assert "if ((" not in source, (
"Generated code contains extraneous parentheses in the if condition, "
diff --git a/tests/python/codegen/test_target_codegen_cross_llvm.py b/tests/python/codegen/test_target_codegen_cross_llvm.py
index 0b64b24832..00f33dd124 100644
--- a/tests/python/codegen/test_target_codegen_cross_llvm.py
+++ b/tests/python/codegen/test_target_codegen_cross_llvm.py
@@ -64,7 +64,7 @@ def test_llvm_add_pipeline():
print(f"Skip because {target} is not enabled..")
return
temp = utils.tempdir()
- f = tvm.tirx.build(AddModule, target=target)
+ f = tvm.tirx.compile(AddModule, target=target).mod
path = temp.relpath("myadd.o")
f.write_to_file(path)
verify_elf(path, 0x28)
diff --git a/tests/python/codegen/test_target_codegen_cuda.py b/tests/python/codegen/test_target_codegen_cuda.py
index 1b58d40f4a..55e0413826 100644
--- a/tests/python/codegen/test_target_codegen_cuda.py
+++ b/tests/python/codegen/test_target_codegen_cuda.py
@@ -55,7 +55,7 @@ def test_cuda_host_bundle(tmp_path, *, backend_config):
B[tx] = A[tx] + T.float32(1)
target = tvm.target.Target("cuda", host="cuda_host")
- built = tvm.compile(
+ built = tvm.tirx.compile(
add_one, target={**tvm.target.Target(target).export(), "backend_config": backend_config}
).mod
source = export_cuda_host(built)
@@ -116,10 +116,10 @@ def test_cuda_host_bundle_bf16_cluster(tmp_path, *, backend_config):
arch = env.cuda_arch()
target = tvm.target.Target({"kind": "cuda", "arch": arch}, host="cuda_host")
with target:
- built = tvm.compile(
+ built = tvm.tirx.compile(
tvm.IRModule({"main": main}),
target={**tvm.target.Target(target).export(), "backend_config": backend_config},
- tir_pipeline="tirx",
+ pipeline="tirx",
).mod
source = export_cuda_host(built)
assert "cudaLaunchAttributeClusterDimension" in source
@@ -170,7 +170,7 @@ def test_cuda_host_bundle_programmatic_dependent_launch(tmp_path, *, backend_con
target = tvm.target.Target("cuda", host="cuda_host")
source = export_cuda_host(
- tvm.compile(
+ tvm.tirx.compile(
add_one, target={**tvm.target.Target(target).export(), "backend_config": backend_config}
).mod
)
@@ -218,10 +218,10 @@ def test_cuda_host_bundle_tensor_map_parameter(tmp_path, monkeypatch, *, backend
target = tvm.target.Target({"kind": "cuda", "arch": "sm_90a"}, host="cuda_host")
with target:
- built = tvm.compile(
+ built = tvm.tirx.compile(
tvm.IRModule({"main": main}),
target={**tvm.target.Target(target).export(), "backend_config": backend_config},
- tir_pipeline="tirx",
+ pipeline="tirx",
).mod
source = export_cuda_host(built)
assert "(CUtensorMap*)" in source
@@ -261,7 +261,7 @@ def test_cuda_vectorize_add(*, backend_config):
if i_0 * num_thread + i_1 < n:
B[i_0 * num_thread + i_1] = A[i_0 * num_thread + i_1] + one
- fun = tvm.compile(Module, target={"kind": "cuda", "backend_config": backend_config})
+ fun = tvm.tirx.compile(Module, target={"kind": "cuda", "backend_config": backend_config})
def run_and_check():
dev = tvm.cuda(0)
@@ -325,7 +325,9 @@ def test_cuda_bf16_vectorize_add(*, backend_config):
with tvm.transform.PassContext(
disabled_pass=["tirx.BF16Promote", "tirx.BF16CastElimination", "tirx.BF16TypeLowering"]
):
- fun = tvm.compile(Module, target={"kind": "cuda", "backend_config": backend_config})
+ fun = tvm.tirx.compile(
+ Module, target={"kind": "cuda", "backend_config": backend_config}
+ )
def run_and_check():
dev = tvm.cuda(0)
@@ -377,7 +379,7 @@ def test_cuda_multiply_add(*, backend_config):
ty="int32",
)
- fun = tvm.compile(Module, target={"kind": "cuda", "backend_config": backend_config})
+ fun = tvm.tirx.compile(Module, target={"kind": "cuda", "backend_config": backend_config})
np_a = np.random.randint(low=-128, high=127, size=(n, lanes))
np_b = np.random.randint(low=-128, high=127, size=(n, lanes))
@@ -416,7 +418,7 @@ def test_cuda_vectorize_load(*, backend_config):
with T.launch_thread("threadIdx.x", num_thread) as i_1:
B[i_0 * num_thread + i_1] = A[i_0 * num_thread + i_1]
- fun = tvm.compile(Module, target={"kind": "cuda", "backend_config": backend_config})
+ fun = tvm.tirx.compile(Module, target={"kind": "cuda", "backend_config": backend_config})
def run_and_check():
dev = tvm.cuda(0)
@@ -451,7 +453,7 @@ def test_cuda_make_int8(*, backend_config):
for j in T.vectorized(lanes):
A[i, j] = const_value
- fun = tvm.compile(Module, target={"kind": "cuda", "backend_config": backend_config})
+ fun = tvm.tirx.compile(Module, target={"kind": "cuda", "backend_config": backend_config})
def run_and_check():
dev = tvm.cuda(0)
@@ -489,7 +491,7 @@ def test_cuda_inf_nan(*, backend_config):
if i_0 * 8 + i_1 < n:
C[i_0 * 8 + i_1] = inf_value
- fun = tvm.compile(Module, target={"kind": "cuda", "backend_config": backend_config})
+ fun = tvm.tirx.compile(Module, target={"kind": "cuda", "backend_config": backend_config})
def run_and_check():
dev = tvm.cuda(0)
@@ -547,7 +549,7 @@ def test_crossthread_reduction1(target, *, backend_config):
if m_0 == 0:
B[i] = reduced[0]
- fun = tvm.compile(
+ fun = tvm.tirx.compile(
Module,
target={**tvm.target.Target(target).export(), "backend_config": backend_config}
if target == "cuda"
@@ -635,7 +637,7 @@ def test_crossthread_reduction2(target, *, backend_config):
if k0_0 == 0 and k1_0 == 0:
B[i] = reduced[0]
- func = tvm.compile(
+ func = tvm.tirx.compile(
Module,
target={**tvm.target.Target(target).export(), "backend_config": backend_config}
if target == "cuda"
@@ -682,7 +684,7 @@ def test_cuda_reduction_binding(*, backend_config):
B[m_0 * 32 + m_1] = T.float32(0.0)
B[m_0 * 32 + m_1] = B[m_0 * 32 + m_1] + A[m_0 * 32 + m_1, k]
- func = tvm.compile(Module, target={"kind": "cuda", "backend_config": backend_config})
+ func = tvm.tirx.compile(Module, target={"kind": "cuda", "backend_config": backend_config})
@pytest.mark.gpu
@@ -714,7 +716,7 @@ def test_cuda_const_float_to_half(*, backend_config):
]
)
- func = tvm.compile(Module, target={"kind": "cuda", "backend_config": backend_config})
+ func = tvm.tirx.compile(Module, target={"kind": "cuda", "backend_config": backend_config})
shape = (2, 3, 4)
a_np = np.random.uniform(size=shape).astype("float16")
@@ -750,7 +752,7 @@ def test_cuda_floordiv_with_vectorization(*, backend_config):
(i_0 * 256 + i_1_0 * 4 + i_1_1) // 37
]
- func = tvm.compile(Module, target={"kind": "cuda", "backend_config": backend_config})
+ func = tvm.tirx.compile(Module, target={"kind": "cuda", "backend_config": backend_config})
def run_and_check():
dev = tvm.cuda(0)
@@ -784,7 +786,7 @@ def test_cuda_floormod_with_vectorization(*, backend_config):
(i_0 * 256 + i_1_0 * 4 + i_1_1) % 37
]
- func = tvm.compile(Module, target={"kind": "cuda", "backend_config": backend_config})
+ func = tvm.tirx.compile(Module, target={"kind": "cuda", "backend_config": backend_config})
def run_and_check():
dev = tvm.cuda(0)
@@ -867,7 +869,7 @@ def test_vectorized_casts(t0, t1, factor, *, backend_config):
t0, B[i_0 * factor + i_1]
)
- func = tvm.compile(Module, target={"kind": "cuda", "backend_config": backend_config})
+ func = tvm.tirx.compile(Module, target={"kind": "cuda", "backend_config": backend_config})
# correctness
def run_and_check():
@@ -906,7 +908,7 @@ def sched(compute_fn, dtype, n=128, *, backend_config=None):
A[i0_1_0 * 4 + i0_1_1_0 * 4 + i0_1_1_1]
)
- return tvm.compile(Module, target={"kind": "cuda", "backend_config": backend_config})
+ return tvm.tirx.compile(Module, target={"kind": "cuda", "backend_config": backend_config})
@pytest.mark.gpu
@@ -1075,7 +1077,7 @@ def test_cuda_vectorize_load_permute_pad(*, backend_config):
j < padding or l + padding <= j, zero, A[i * lanes + k, j - padding]
)
- fun = tvm.compile(Module, target={"kind": "cuda", "backend_config": backend_config})
+ fun = tvm.tirx.compile(Module, target={"kind": "cuda", "backend_config": backend_config})
np_a = np.random.randint(low=-128, high=127, size=(n, l)).astype(dtype)
np_a_reshape = np_a.reshape(n // lanes, lanes, l).transpose(0, 2, 1)
@@ -1114,7 +1116,7 @@ def test_try_unaligned_vector_load(*, backend_config):
for i_1 in T.vectorized(2):
C[i_0 * 2 + i_1] = A[i_0 * 2 + i_1 + offset]
- f = tvm.tirx.build(Module, target={"kind": "cuda", "backend_config": backend_config})
+ f = tvm.tirx.compile(Module, target={"kind": "cuda", "backend_config": backend_config}).mod
kernel_source = f.imports[0].inspect_source()
a_data = np.arange(0, N).astype("float16")
@@ -1175,7 +1177,7 @@ def test_cuda_thread_sync_inside_condition(*, backend_config):
for func in [func2, func3]:
mod = tvm.IRModule({"main": func})
- compiled = tvm.compile(mod, target={"kind": "cuda", "backend_config": backend_config})
+ compiled = tvm.tirx.compile(mod, target={"kind": "cuda", "backend_config": backend_config})
assert "__syncthreads()" in compiled.mod.imports[0].inspect_source()
@@ -1188,7 +1190,7 @@ def test_invalid_reinterpret(*, backend_config):
B[tx] = T.call_intrin("tirx.reinterpret", A[tx], ty="uint8")
with pytest.raises(RuntimeError):
- tvm.compile(func, target={"kind": "cuda", "backend_config": backend_config})
+ tvm.tirx.compile(func, target={"kind": "cuda", "backend_config": backend_config})
@pytest.mark.gpu
@@ -1209,7 +1211,7 @@ def test_cuda_tensormap(*, backend_config):
# fmt: on
mod = tvm.IRModule({"main": main})
- mod = tvm.compile(mod, target={"kind": "cuda", "backend_config": backend_config})
+ mod = tvm.tirx.compile(mod, target={"kind": "cuda", "backend_config": backend_config})
assert (
"""
extern "C" __global__ void __launch_bounds__(128) main_kernel(float* __restrict__ A_ptr, const __grid_constant__ CUtensorMap A_map) {
@@ -1240,7 +1242,7 @@ def test_cuda_device_func_call(*, backend_config):
with T.launch_thread("threadIdx.x", 1024) as tx:
C[bx, tx] = Module.add(A[bx, tx], B[bx, tx])
- lib = tvm.compile(Module, target={"kind": "cuda", "backend_config": backend_config})
+ lib = tvm.tirx.compile(Module, target={"kind": "cuda", "backend_config": backend_config})
cuda_code = lib.mod.imports[0].inspect_source()
assert 'extern "C" __device__ float add(float a, float b) {\n return (a + b);\n}' in cuda_code
@@ -1260,7 +1262,7 @@ def test_cuda_float_const_hex_format(*, backend_config):
with T.launch_thread("threadIdx.x", 1024) as tx:
A[bx, tx] = T.float32(1 / 27)
- lib = tvm.compile(Module, target={"kind": "cuda", "backend_config": backend_config})
+ lib = tvm.tirx.compile(Module, target={"kind": "cuda", "backend_config": backend_config})
cuda_code = lib.mod.imports[0].inspect_source()
assert "0x1.2f684bda12f68p-5f" in cuda_code
@@ -1291,7 +1293,7 @@ def test_device_host_call_same_func(*, backend_config):
# 2. We set a dummy mcpu value for testing purpose,
# in order to avoid checking a function is host or device based on the "cpu" substring.
target = tvm.target.Target({"kind": "cuda", "mcpu": "dummy_mcpu"}, host="c")
- lib = tvm.compile(
+ lib = tvm.tirx.compile(
Module, target={**tvm.target.Target(target).export(), "backend_config": backend_config}
)
cuda_code = lib.mod.imports[0].inspect_source()
@@ -1326,7 +1328,7 @@ def test_thread_return(monkeypatch, *, backend_config):
T.gpu_thread_return()
B[bx, tx] = A[bx, tx]
- lib = tvm.compile(
+ lib = tvm.tirx.compile(
Module,
target={
"kind": "cuda",
@@ -1354,7 +1356,7 @@ def test_cuda_loop_step(*, backend_config):
C[i] = A[i] + B[i]
target = tvm.target.Target({"kind": "cuda"})
- lib = tvm.compile(
+ lib = tvm.tirx.compile(
cuda_loop_step,
target={**tvm.target.Target(target).export(), "backend_config": backend_config},
)
@@ -1392,7 +1394,7 @@ def test_export_load_with_fallback(monkeypatch, tmp_path, *, backend_config):
B[i_0 * 32 + i_1] = A[i_0 * 32 + i_1] + 1.0
monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1")
- host_lib = tvm.compile(Module, target={"kind": "cuda", "backend_config": backend_config})
+ host_lib = tvm.tirx.compile(Module, target={"kind": "cuda", "backend_config": backend_config})
monkeypatch.delenv("TVM_COMPILE_FORCE_FALLBACK")
lib_path = str(tmp_path / "lib.so")
diff --git a/tests/python/codegen/test_target_codegen_cuda_fastmath.py b/tests/python/codegen/test_target_codegen_cuda_fastmath.py
index 1ce3c9f9c2..8d8f025948 100644
--- a/tests/python/codegen/test_target_codegen_cuda_fastmath.py
+++ b/tests/python/codegen/test_target_codegen_cuda_fastmath.py
@@ -245,7 +245,7 @@ def check_cuda_source(
) -> Executable:
"""Check the CUDA source for the given dtype and case."""
with tvm.transform.PassContext(config={"tirx.enable_fast_math": enable_fast_math}):
- executable = tvm.compile(mod, target=target)
+ executable = tvm.tirx.compile(mod, target=target)
source = executable.mod.imports[0].inspect_source()
expected = expected_intrinsic(dtype, case, enable_fast_math)
assert re.search(rf"(?<!_)\b{re.escape(expected)}\s*\(", source)
diff --git a/tests/python/codegen/test_target_codegen_cuda_fp4.py b/tests/python/codegen/test_target_codegen_cuda_fp4.py
index 75bad2b7c4..24442ef78a 100644
--- a/tests/python/codegen/test_target_codegen_cuda_fp4.py
+++ b/tests/python/codegen/test_target_codegen_cuda_fp4.py
@@ -59,7 +59,7 @@ def test_e2m1_vector_conversions(promoted_dtype):
)
target = "cuda"
- fadd = tvm.compile(Module, target=target)
+ fadd = tvm.tirx.compile(Module, target=target)
if "x" in native_dtype:
lanes = int(native_dtype.split("x")[-1])
@@ -212,7 +212,7 @@ def test_e2m1_dequantize():
else:
mod = _scalar_reinterpret_module(n, num_blocks, vector_length, num_elem_per_storage)
- tvm.compile(mod, target=target)
+ tvm.tirx.compile(mod, target=target)
@pytest.mark.gpu
@@ -240,7 +240,7 @@ def test_e2m1_scalar_buffer_offset():
B[bx * 32 + tx] = T.Cast("float16", A[bx * 32 + tx])
target = "cuda"
- fadd = tvm.compile(func, target=target)
+ fadd = tvm.tirx.compile(func, target=target)
# float4_e2m1fn: 4-bit values 0..15, two packed per byte.
# Encoding (sign | exp1 | man1 man0):
diff --git a/tests/python/codegen/test_target_codegen_cuda_fp8.py b/tests/python/codegen/test_target_codegen_cuda_fp8.py
index 23b4b24d15..f306e4c355 100644
--- a/tests/python/codegen/test_target_codegen_cuda_fp8.py
+++ b/tests/python/codegen/test_target_codegen_cuda_fp8.py
@@ -67,7 +67,7 @@ def test_fp8_conversions(input):
mod = _create_mod(dtype)
target = "cuda"
- fadd = tvm.tirx.build(mod, target=target)
+ fadd = tvm.tirx.compile(mod, target=target).mod
cuda_src = fadd.imports[0].inspect_source()
assert nv_dtype in cuda_src, f"{nv_dtype} datatype not found in generated CUDA"
@@ -116,7 +116,7 @@ def test_fp8_packing(dtype):
mod = _create_mod(native_dtype, packed_dtype, length)
target = "cuda"
- f = tvm.compile(mod, target=target)
+ f = tvm.tirx.compile(mod, target=target)
dev = tvm.cuda(0)
np_shape = (length, vector_length)
@@ -174,7 +174,7 @@ def test_fp8_vector_conversions(native_dtype, promoted_dtype, numpytype):
mod = _create_mod(native_dtype, promoted_dtype)
target = "cuda"
- fadd = tvm.tirx.build(mod, target=target)
+ fadd = tvm.tirx.compile(mod, target=target).mod
cuda_src = fadd.imports[0].inspect_source()
dev = tvm.cuda(0)
@@ -224,7 +224,7 @@ def test_half_broadcast(bcast_length):
mod = _create_mod(bcast_length, dtype)
target = "cuda"
- func = tvm.compile(mod, target=target)
+ func = tvm.tirx.compile(mod, target=target)
dev = tvm.cuda(0)
a_np = np.random.uniform(low=0, high=4, size=()).astype(dtype)
@@ -258,7 +258,7 @@ def test_half_misaligned_vector_load(vector_length):
B[i] = A[vec_index]
target = "cuda"
- f = tvm.compile(vector_load, target=target)
+ f = tvm.tirx.compile(vector_load, target=target)
dev = tvm.cuda(0)
a_np = np.random.uniform(low=0, high=1, size=(length,)).astype(dtype)
@@ -299,7 +299,7 @@ def test_half4_vector_add():
C[i_0 * 32 + i_1] = A[i_0 * 32 + i_1] + B[i_0 * 32 + i_1]
target = "cuda"
- fadd = tvm.compile(Module, target=target)
+ fadd = tvm.tirx.compile(Module, target=target)
dev = tvm.cuda(0)
a_np = np.random.uniform(-1, 1, (length, vector_length)).astype(dtype)
@@ -378,8 +378,8 @@ class BaseFP8E4M3QuantScaleOnly:
vec_model_dtype, T.reinterpret(packed[row, k], ty=vec_quantized_dtype)
) * T.Broadcast(scale[row, k * 4 // group_size], 4)
- quant_func = tvm.compile(quantize, target=target_str)
- dequant_func = tvm.compile(dequantize, target=target_str)
+ quant_func = tvm.tirx.compile(quantize, target=target_str)
+ dequant_func = tvm.tirx.compile(dequantize, target=target_str)
def quant(weight):
packed = tvm.runtime.empty(quant_weight_shape, storage_dtype, dev)
@@ -505,7 +505,7 @@ def test_const(dtype):
A[tx] = A_local[tx]
mod = tvm.IRModule({"main": func})
- tvm.compile(mod, target="cuda")
+ tvm.tirx.compile(mod, target="cuda")
@pytest.mark.gpu
@@ -535,7 +535,7 @@ def test_copy(dtype, vec_len):
B[tx, i] = A[tx, i]
mod = tvm.IRModule({"main": func})
- rtmod = tvm.compile(mod, target="cuda")
+ rtmod = tvm.tirx.compile(mod, target="cuda")
num_experts = 8
@@ -585,7 +585,7 @@ def test_moe_gemv_shfl_down_illegal_instr():
if reduction == 0:
output[expert, block * 4 + spatial] = reduced[0]
- rt_mod = tvm.compile(moe_dequantize_gemv, target="cuda")
+ rt_mod = tvm.tirx.compile(moe_dequantize_gemv, target="cuda")
x_data = np.zeros((1, reduce_size), dtype=np.float16)
indptr_data = np.zeros((1, 2), dtype=np.int32)
@@ -633,7 +633,7 @@ def test_fp8_fp16_bf16_vectorize_arith(vec_length, dtype):
mod = _create_mod(vec_length, dtype)
target = tvm.target.Target.from_device(tvm.cuda())
- f = tvm.tirx.build(mod, target=target)
+ f = tvm.tirx.compile(mod, target=target).mod
a_np = np.random.rand(128).astype("float8_e4m3fn")
b_np = np.random.rand(128).astype(dtype)
diff --git a/tests/python/codegen/test_target_codegen_device.py b/tests/python/codegen/test_target_codegen_device.py
index 10d87b1486..97c0eb84ce 100644
--- a/tests/python/codegen/test_target_codegen_device.py
+++ b/tests/python/codegen/test_target_codegen_device.py
@@ -43,7 +43,7 @@ def test_large_uint_imm():
target_kind = target["kind"] if isinstance(target, dict) else target
if not tvm.testing.device_enabled(target_kind):
return
- f = tvm.compile(Module, target=target)
+ f = tvm.tirx.compile(Module, target=target)
def run_and_check():
dev = tvm.device(target_kind, 0)
@@ -82,7 +82,7 @@ def test_add_pipeline():
if not tvm.testing.device_enabled(device) or not tvm.testing.device_enabled(host):
return
target = tvm.target.Target(device, host)
- mhost = tvm.tirx.build(Module, target=target)
+ mhost = tvm.tirx.compile(Module, target=target).mod
f = mhost.main
n = 1027
diff --git a/tests/python/codegen/test_target_codegen_extern.py b/tests/python/codegen/test_target_codegen_extern.py
index 216112bdde..88afc9b097 100644
--- a/tests/python/codegen/test_target_codegen_extern.py
+++ b/tests/python/codegen/test_target_codegen_extern.py
@@ -54,7 +54,7 @@ def test_add_pipeline():
return
mod = ModuleGPU if target in ["opencl", "cuda"] else ModuleCPU
# build and invoke the kernel.
- f = tvm.compile(mod, target=target)
+ f = tvm.tirx.compile(mod, target=target)
n = nn
def run_and_check():
@@ -92,7 +92,7 @@ def test_pack_buffer_simple():
if not tvm.testing.device_enabled(target):
return
# build and invoke the kernel.
- f = tvm.compile(Module, target=target)
+ f = tvm.tirx.compile(Module, target=target)
dev = tvm.cpu(0)
# launch the kernel.
n = nn
diff --git a/tests/python/codegen/test_target_codegen_gpu_common.py b/tests/python/codegen/test_target_codegen_gpu_common.py
index 26866f6753..26a4f9a189 100644
--- a/tests/python/codegen/test_target_codegen_gpu_common.py
+++ b/tests/python/codegen/test_target_codegen_gpu_common.py
@@ -59,7 +59,7 @@ def test_int_intrin(target, dtype):
with T.launch_thread("threadIdx.x", n) as i0:
B[i0] = tvm_intrin(A[i0])
- f = tvm.compile(Module, target=target)
+ f = tvm.tirx.compile(Module, target=target)
def run_and_check():
dev = tvm.device_from_target(target)
diff --git a/tests/python/codegen/test_target_codegen_hexagon.py b/tests/python/codegen/test_target_codegen_hexagon.py
index b17e9583a8..98c0990357 100644
--- a/tests/python/codegen/test_target_codegen_hexagon.py
+++ b/tests/python/codegen/test_target_codegen_hexagon.py
@@ -20,8 +20,8 @@ import re
import pytest
import tvm
-import tvm.contrib.hexagon as hexagon
import tvm.testing
+from tvm.backend.hexagon import linker as hexagon
from tvm.script import ir as I
from tvm.script import tirx as T
from tvm.testing import env
@@ -29,7 +29,7 @@ from tvm.testing import env
@pytest.fixture(autouse=True)
def register_linker():
- original_linker = hexagon.hexagon_link()
+ original_linker = tvm.get_global_func("tvm.contrib.hexagon.hexagon_link")
# Register a phony linker, so that we can test codegen without a Hexagon toolchain.
hexagon.register_linker(lambda: "/bin/true")
yield None
@@ -53,7 +53,7 @@ def test_basic():
for i in range(128):
C[i] = A[i] + A_1[i]
- hexm = tvm.compile(Module, target=tvm.target.Target(target, target))
+ hexm = tvm.tirx.compile(Module, target=tvm.target.Target(target, target))
asm = hexm.inspect_source("s")
vadds = re.findall(r"v[0-9]+.b = vadd\(v[0-9]+.b,v[0-9]+.b\)", asm)
assert vadds # Check that it's non-empty
@@ -71,7 +71,7 @@ def test_llvm_target_features():
for i in range(128):
C[i] = T.Cast("int32", A[i]) + 1
- m = tvm.compile(Module, target=tvm.target.Target(target, target))
+ m = tvm.tirx.compile(Module, target=tvm.target.Target(target, target))
llvm_ir = m.inspect_source("ll")
# Make sure we find +hvx-length128b in "attributes".
fs = re.findall(r"attributes.*\+hvx-length128b", llvm_ir)
@@ -102,7 +102,7 @@ def test_llvm_options():
# Check that BuildHexagon hasn't crashed because of target attribute
# type mismatch.
- tvm.compile(Module, target=tvm.target.Target(target, target))
+ tvm.tirx.compile(Module, target=tvm.target.Target(target, target))
assert re.search("-hexagon-noopt", str(target))
diff --git a/tests/python/codegen/test_target_codegen_llvm.py b/tests/python/codegen/test_target_codegen_llvm.py
index e5df6f7b7c..48e3c7b5a5 100644
--- a/tests/python/codegen/test_target_codegen_llvm.py
+++ b/tests/python/codegen/test_target_codegen_llvm.py
@@ -49,7 +49,7 @@ def test_duplicate_function_global_symbol_diagnostic():
with pytest.raises(
tvm.error.InternalError, match="Duplicate Function global_symbol 'dup_symbol'"
) as err:
- tvm.compile(Module, target="llvm")
+ tvm.tirx.compile(Module, target="llvm")
assert "first_unique_key" in str(err.value)
assert "second_unique_key" in str(err.value)
@@ -68,7 +68,7 @@ def test_unique_function_global_symbols_compile():
T.func_attr({"global_symbol": "dup_symbol_b", "tirx.noalias": True})
A[0] = T.float32(2)
- tvm.compile(Module, target="llvm")
+ tvm.tirx.compile(Module, target="llvm")
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
@@ -80,7 +80,7 @@ def test_llvm_intrin():
A_buf = T.decl_tensor(A, ty_args=[T.Tensor((4,), "float32", data=A)])
T.evaluate(T.Call("tirx.prefetch", [T.address_of(A_buf[0]), 0, 3, 1], ty="void"))
- fcode = tvm.compile(Module)
+ fcode = tvm.tirx.compile(Module)
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
@@ -92,7 +92,7 @@ def test_llvm_void_intrin():
# Create an intrinsic that returns void.
T.call_llvm_intrin("llvm.assume", T.bool(True), ty="void")
- fcode = tvm.compile(Module)
+ fcode = tvm.tirx.compile(Module)
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
@@ -119,7 +119,7 @@ def test_llvm_overloaded_intrin():
def main(A: T.Tensor((1, 1), "int32"), C: T.Tensor((1, 1), "int32")):
C[0, 0] = T.call_llvm_pure_intrin("llvm.ctlz", A[0, 0], int1_zero, ty="int32")
- f = tvm.compile(Module, target="llvm")
+ f = tvm.tirx.compile(Module, target="llvm")
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
@@ -131,7 +131,7 @@ def test_llvm_lookup_intrin():
A_buf = T.decl_tensor(A, ty_args=[T.Tensor((1,), "uint8x8", data=A)])
T.evaluate(T.call_llvm_pure_intrin("llvm.ctpop.v8i8", A_buf[0], ty="uint8x8"))
- fcode = tvm.compile(Module, None)
+ fcode = tvm.tirx.compile(Module, None)
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
@@ -146,7 +146,7 @@ def test_llvm_large_uintimm():
T.func_attr({"tirx.noalias": True})
A[()] = large_val + T.uint64(3)
- f = tvm.compile(Module, target="llvm")
+ f = tvm.tirx.compile(Module, target="llvm")
dev = tvm.cpu(0)
a = tvm.runtime.empty((), dtype="uint64", device=dev)
f(a)
@@ -171,7 +171,7 @@ def test_llvm_multi_parallel():
) * T.float32(2.0) + T.float32(2.0)
n = 128
- f = tvm.compile(Module, target="llvm")
+ f = tvm.tirx.compile(Module, target="llvm")
dev = tvm.cpu(0)
a = tvm.runtime.tensor(np.random.uniform(size=n).astype("float32"), dev)
c = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev)
@@ -192,7 +192,7 @@ def test_llvm_flip_pipeline():
if i_0 * 4 + i_1 < nn:
C[i_0 * 4 + i_1] = A[nn + base - 1 - (i_0 * 4 + i_1)]
- f = tvm.compile(Module, target="llvm")
+ f = tvm.tirx.compile(Module, target="llvm")
dev = tvm.cpu(0)
a = tvm.runtime.tensor(np.random.uniform(size=(nn + base)).astype("float32"), dev)
c = tvm.runtime.tensor(np.zeros(nn, dtype="float32"), dev)
@@ -220,7 +220,7 @@ def test_llvm_vadd_pipeline():
if i_0 * 4 + i_1 < n:
C[i_0 * 4 + i_1] = A[i_0 * 4 + i_1] + B[i_0 * 4 + i_1]
- f = tvm.compile(Module, target="llvm")
+ f = tvm.tirx.compile(Module, target="llvm")
dev = tvm.cpu(0)
n = 128
a = tvm.runtime.tensor(np.random.uniform(size=n).astype("float32"), dev)
@@ -247,7 +247,7 @@ def test_llvm_madd_pipeline():
if i_0 * 4 + i_1 < nn:
C[i_0 * 4 + i_1, j] = A[i_0 * 4 + i_1 + base, j] + T.float32(1.0)
- f = tvm.compile(Module, target="llvm")
+ f = tvm.tirx.compile(Module, target="llvm")
dev = tvm.cpu(0)
a = tvm.runtime.tensor(np.random.uniform(size=(nn + base, stride)).astype("float32"), dev)
c = tvm.runtime.tensor(np.zeros((nn, stride), dtype="float32"), dev)
@@ -275,7 +275,7 @@ def test_llvm_temp_space():
C[i] = B[i] + T.float32(1.0)
nn = 1024
- f = tvm.compile(Module, target="llvm")
+ f = tvm.tirx.compile(Module, target="llvm")
dev = tvm.cpu(0)
a = tvm.runtime.tensor(np.random.uniform(size=nn).astype("float32"), dev)
c = tvm.runtime.tensor(np.zeros(nn, dtype="float32"), dev)
@@ -304,7 +304,7 @@ def test_multiple_func():
for i in range(fadd2_n):
C[i] = A[i] + B[i]
- f = tvm.compile(Module, target="llvm")
+ f = tvm.tirx.compile(Module, target="llvm")
dev = tvm.cpu(0)
n = 10
a = tvm.runtime.tensor(np.random.uniform(size=n).astype("float32"), dev)
@@ -329,7 +329,7 @@ def test_llvm_condition():
n = 64
offset = 8
- f = tvm.compile(Module, target="llvm")
+ f = tvm.tirx.compile(Module, target="llvm")
dev = tvm.cpu(0)
a = tvm.runtime.tensor(np.random.uniform(size=(n,)).astype("float32"), dev)
c = tvm.runtime.empty((n,), "float32", dev)
@@ -350,7 +350,7 @@ def test_llvm_bool():
C[i] = T.Cast("float32", A[i] == 1)
n = 64
- f = tvm.compile(Module, target="llvm")
+ f = tvm.tirx.compile(Module, target="llvm")
dev = tvm.cpu(0)
a = tvm.runtime.tensor(np.random.randint(0, 2, size=(n,)).astype("int32"), dev)
c = tvm.runtime.empty((n,), "float32", dev)
@@ -370,7 +370,7 @@ def test_llvm_cast_float_to_bool():
C[i] = T.Cast("bool", A[i])
n = 4
- f = tvm.compile(Module, target="llvm")
+ f = tvm.tirx.compile(Module, target="llvm")
dev = tvm.cpu(0)
a = tvm.runtime.tensor(np.array([0.0, 1.0, np.nan, np.inf], dtype="float32"), dev)
c = tvm.runtime.empty((n,), dtype="bool", device=dev)
@@ -397,7 +397,7 @@ def test_rank_zero():
compute[()] = C[()] + T.float32(1.0)
n = 64
- f = tvm.compile(Module, target="llvm")
+ f = tvm.tirx.compile(Module, target="llvm")
dev = tvm.cpu(0)
a = tvm.runtime.tensor(np.random.randint(0, 2, size=(n,)).astype("float32"), dev)
sc = tvm.runtime.tensor(np.random.randint(0, 2, size=()).astype("float32"), dev)
@@ -418,7 +418,7 @@ def test_alignment():
for i_1 in T.vectorized(8):
B[i_0 * 8 + i_1] = A[i_0 * 8 + i_1] * T.float32(3.0)
- f = tvm.tirx.build(Module, target="llvm")
+ f = tvm.tirx.compile(Module, target="llvm").mod
lines = f.inspect_source().split("\n")
@@ -558,7 +558,7 @@ def test_llvm_div(start, end, dstart, dend, dtype, floor_div):
D[i, j] = div_fn(clipa(A[i]), clipb(B[j]))
M[i, j] = mod_fn(clipa(A[i]), clipb(B[j]))
- f = tvm.compile(Module, target="llvm")
+ f = tvm.tirx.compile(Module, target="llvm")
# Fill input arrays with values
A_arr = tvm.runtime.empty((a_size,), dtype)
@@ -625,7 +625,7 @@ def test_llvm_fp_math():
for i in range(n):
B[i] = T.float32(1.0) / (T.float32(1e37) * A[i])
- f_recip = tvm.compile(RecipModule, target="llvm")
+ f_recip = tvm.tirx.compile(RecipModule, target="llvm")
for n in [4, 8, 16]:
a = tvm.runtime.tensor(np.full((n,), 100, "float32"))
@@ -644,7 +644,7 @@ def test_llvm_fp_math():
for i in range(n):
B[i] = T.sigmoid(A[i])
- f_sigmoid = tvm.compile(SigmoidModule, target="llvm")
+ f_sigmoid = tvm.tirx.compile(SigmoidModule, target="llvm")
for n in [4, 8, 16]:
a = tvm.runtime.tensor(np.full((n,), -1000, "float32"))
@@ -680,7 +680,7 @@ def test_dwarf_debug_information():
"fadd2": Module["main"].with_attr("global_symbol", "fadd2"),
}
)
- m = tvm.compile(mod, target="llvm")
+ m = tvm.tirx.compile(mod, target="llvm")
temp = utils.tempdir()
o_path = temp.relpath("temp.o")
m.write_to_file(o_path)
@@ -718,7 +718,7 @@ def test_dwarf_debug_information():
"fadd2": Module["main"].with_attr("global_symbol", "fadd2"),
}
)
- m = tvm.tirx.build(mod, target={"kind": "llvm", "mtriple": "aarch64-linux-gnu"})
+ m = tvm.tirx.compile(mod, target={"kind": "llvm", "mtriple": "aarch64-linux-gnu"}).mod
ll = m.inspect_source("ll")
# On non-Darwin OS, don't explicitly specify DWARF version.
@@ -728,7 +728,9 @@ def test_dwarf_debug_information():
assert re.search(r"""llvm.dbg.value""", ll)
# Try Darwin, require DWARF-2
- m = tvm.tirx.build(mod, target={"kind": "llvm", "mtriple": "x86_64-apple-darwin-macho"})
+ m = tvm.tirx.compile(
+ mod, target={"kind": "llvm", "mtriple": "x86_64-apple-darwin-macho"}
+ ).mod
ll = m.inspect_source("ll")
assert re.search(r"""i32 4, !"Dwarf Version", i32 2""", ll)
assert re.search(r"""llvm.dbg.value""", ll)
@@ -755,7 +757,7 @@ def test_llvm_bf16():
D[x] = A[x] + B[x]
np.random.seed(122)
- module = tvm.compile(Module, target="llvm")
+ module = tvm.tirx.compile(Module, target="llvm")
npa = np.random.rand(32).astype("bfloat16")
npb = np.random.rand(32).astype("bfloat16")
res = npa + npb
@@ -785,10 +787,10 @@ def test_llvm_crt_static_lib():
for x in range(32):
C[x] = A[x] + B[x]
- module = tvm.tirx.build(
+ module = tvm.tirx.compile(
Module.with_attr("system_lib_prefix", ""),
target=tvm.target.Target("llvm"),
- )
+ ).mod
module.inspect_source()
with utils.tempdir() as temp:
module.write_to_file(temp.relpath("test.o"))
@@ -814,7 +816,7 @@ def test_llvm_order_functions():
def Kirby(v: T.float32) -> T.float32:
return T.call_extern("Fred", v, ty="float32")
- ir_text = tvm.tirx.build(Module, target="llvm").inspect_source("ll")
+ ir_text = tvm.tirx.compile(Module, target="llvm").mod.inspect_source("ll")
# Skip functions whose names start with _.
matches = re.findall(r"^define[^@]*@([a-zA-Z][a-zA-Z0-9_]*)", ir_text, re.MULTILINE)
assert matches == sorted(matches)
@@ -831,11 +833,11 @@ def test_llvm_large_stack_allocation_uses_64bit_extent(extent):
)
A[0] = B[extent - 1]
- module = tvm.tirx.build(
+ module = tvm.tirx.compile(
main,
target={"kind": "llvm", "opt-level": 0},
pipeline="tirx",
- )
+ ).mod
llvm_ir = module.inspect_source("ll")
assert re.search(rf"alloca float, i64 {extent}(?:,|$)", llvm_ir)
@@ -852,7 +854,7 @@ def test_llvm_scalar_concat():
# This will crash in LLVM codegen if CodeGenLLVM::CreateVecConcat doesn't convert
# scalars to single-lane LLVM vectors.
with tvm.transform.PassContext(config={"tirx.disable_assert": True}):
- m = tvm.compile(Module, target="llvm")
+ m = tvm.tirx.compile(Module, target="llvm")
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
@@ -867,7 +869,7 @@ def test_raise_exception_during_codegen():
B[i, j] = A[i, j] * 2.0
with pytest.raises(RuntimeError) as e:
- tvm.compile(Module, target="llvm")
+ tvm.tirx.compile(Module, target="llvm")
msg = str(e)
assert msg.find("Nested parallel loop is not supported") != -1
@@ -905,7 +907,7 @@ def test_llvm_target_attributes():
"mattr": ["+avx512f"],
}
target = tvm.target.Target(target_llvm, host=target_llvm)
- module = tvm.tirx.build(Module, target=target)
+ module = tvm.tirx.compile(Module, target=target).mod
llvm_ir = module.inspect_source()
llvm_ir_lines = llvm_ir.split("\n")
@@ -955,7 +957,7 @@ def test_llvm_assume():
B_1 = T.decl_tensor(B.data, ty_args=[T.Tensor((14,), "int32", data=B.data)])
B_1[i] = A_1[i] * 2
- m = tvm.compile(Module, target="llvm")
+ m = tvm.tirx.compile(Module, target="llvm")
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
@@ -977,7 +979,7 @@ def test_debug_symbol_for_float64():
for i in range(n):
B[i] = A[i]
- tvm.compile(Module, target="llvm")
+ tvm.tirx.compile(Module, target="llvm")
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
@@ -993,7 +995,7 @@ def test_debug_symbol_for_buffer_var():
C[i] = A[i]
B[i] = C[i]
- tvm.compile(Module, target="llvm")
+ tvm.tirx.compile(Module, target="llvm")
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
@@ -1015,7 +1017,7 @@ def test_subroutine_call():
target = "llvm"
dev = tvm.cpu()
- built = tvm.compile(Module)
+ built = tvm.tirx.compile(Module)
arr = tvm.runtime.tensor(np.zeros([1], "float32"), device=dev)
built["main"](arr)
@@ -1050,7 +1052,7 @@ def test_call_packed_returning_void():
# Error occurred during build, as part of
# CodeGenCPU::MakeCallPackedLowered.
- built = tvm.compile(Module, target="llvm")
+ built = tvm.tirx.compile(Module, target="llvm")
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
@@ -1070,7 +1072,7 @@ def test_call_packed_without_string_arg():
T.Call(tvm.ir.Op.get("tirx.call_packed"), [A.data], ty="int32")
with pytest.raises(RuntimeError):
- built = tvm.compile(Module, target="llvm")
+ built = tvm.tirx.compile(Module, target="llvm")
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
@@ -1083,7 +1085,7 @@ def test_call_extern_returning_void():
def main():
T.Call(tvm.ir.Op.get("tirx.call_extern"), ["dummy_function_name"], ty="void")
- built = tvm.compile(Module, target="llvm")
+ built = tvm.tirx.compile(Module, target="llvm")
def test_invalid_volatile_masked_buffer_load():
@@ -1099,7 +1101,7 @@ def test_invalid_volatile_masked_buffer_load():
err_msg = "The masked load intrinsic does not support declaring load as volatile."
with pytest.raises(RuntimeError, match=err_msg):
with tvm.target.Target("llvm"):
- tvm.compile(Module)
+ tvm.tirx.compile(Module)
def test_invalid_volatile_masked_decl_buffer_load():
@@ -1120,7 +1122,7 @@ def test_invalid_volatile_masked_decl_buffer_load():
err_msg = "The masked load intrinsic does not support declaring load as volatile."
with pytest.raises(RuntimeError, match=err_msg):
with tvm.target.Target("llvm"):
- tvm.compile(Module)
+ tvm.tirx.compile(Module)
def test_invalid_volatile_masked_tensor_store():
@@ -1143,7 +1145,7 @@ def test_invalid_volatile_masked_tensor_store():
err_msg = "The masked store intrinsic does not support declaring store as volatile."
with pytest.raises(RuntimeError, match=err_msg):
with tvm.target.Target("llvm"):
- tvm.compile(Module)
+ tvm.tirx.compile(Module)
def test_int_parameter():
@@ -1159,7 +1161,7 @@ def test_int_parameter():
else:
return 20
- built = tvm.compile(Module)
+ built = tvm.tirx.compile(Module)
output = built(True)
assert output == 10
@@ -1180,7 +1182,7 @@ def test_bool_parameter():
else:
return 20
- built = tvm.compile(Module)
+ built = tvm.tirx.compile(Module)
output = built(1)
assert output == 10
@@ -1201,7 +1203,7 @@ def test_bool_return_value():
T.func_attr({"target": T.target("llvm")})
return value < 10
- built = tvm.compile(Module)
+ built = tvm.tirx.compile(Module)
assert isinstance(built(0), bool)
assert built(0)
diff --git a/tests/python/codegen/test_target_codegen_llvm_vla.py b/tests/python/codegen/test_target_codegen_llvm_vla.py
index 6441d29fa5..5d28a7a913 100644
--- a/tests/python/codegen/test_target_codegen_llvm_vla.py
+++ b/tests/python/codegen/test_target_codegen_llvm_vla.py
@@ -56,7 +56,7 @@ def test_codegen_vscale(target):
A[i] = 2 * vscale
with tvm.target.Target(target):
- build_mod = tvm.tirx.build(main)
+ build_mod = tvm.tirx.compile(main).mod
llvm = build_mod.inspect_source()
assert re.findall(r"llvm.vscale.i32", llvm), "No vscale in generated LLVM."
@@ -88,7 +88,7 @@ def test_scalable_buffer_load_store(target):
B[T.ramp(0, 1, 4 * T.vscale())] = A[T.ramp(0, 1, 4 * T.vscale())]
with tvm.target.Target(target):
- mod = tvm.tirx.build(my_func)
+ mod = tvm.tirx.compile(my_func).mod
llvm = mod.inspect_source("ll")
assert re.findall(r"load <vscale x 4 x float>", llvm), "No scalable load in generated LLVM."
@@ -121,7 +121,7 @@ def test_scalable_broadcast(target):
A[T.ramp(0, 1, 4 * T.vscale())] = T.broadcast(1, 4 * T.vscale())
with tvm.target.Target(target):
- mod = tvm.tirx.build(my_func)
+ mod = tvm.tirx.compile(my_func).mod
llvm = mod.inspect_source("ll")
# Older LLVM versions print the broadcast as a shufflevector of an insertelement,
@@ -159,7 +159,7 @@ def test_get_active_lane_mask(target):
A[i : i + T.vscale() * 4] = T.get_active_lane_mask(i, 30, ty="uint1xvscalex4")
with tvm.target.Target(target):
- out = tvm.tirx.build(before)
+ out = tvm.tirx.compile(before).mod
ll = out.inspect_source("ll")
assert "get.active.lane.mask" in ll
@@ -194,7 +194,7 @@ def test_predicated_scalable_buffer(target):
B[i_0 * 4 * T.vscale() + i_1] = A[i_0 * 4 * T.vscale() + i_1] + 1.0
with tvm.target.Target(target):
- out = tvm.tirx.build(before)
+ out = tvm.tirx.compile(before).mod
ll = out.inspect_source("ll")
assert "get.active.lane.mask" in ll
diff --git a/tests/python/codegen/test_target_codegen_metal.py b/tests/python/codegen/test_target_codegen_metal.py
index ecb9069317..e444c5c787 100644
--- a/tests/python/codegen/test_target_codegen_metal.py
+++ b/tests/python/codegen/test_target_codegen_metal.py
@@ -42,7 +42,7 @@ def test_metal_inf_nan():
with T.launch_thread("threadIdx.x", 1) as i:
C[i] = T.Cast(dtype, value)
- fun = tvm.compile(Module, target=target)
+ fun = tvm.tirx.compile(Module, target=target)
def run_and_check():
dev = tvm.metal(0)
@@ -74,7 +74,7 @@ def test_unaligned_vectorize():
target = "metal"
a = (np.arange(6).reshape(2, 3)).astype("float32")
- f = tvm.compile(IRModule, target=target)
+ f = tvm.tirx.compile(IRModule, target=target)
def run_and_check():
dev = tvm.metal()
@@ -103,7 +103,7 @@ def test_metal_erf():
with T.launch_thread("threadIdx.x", 1) as i0:
C[i0] = T.erf(A[i0])
- fun = tvm.compile(Module, target=target)
+ fun = tvm.tirx.compile(Module, target=target)
def run_and_check():
dev = tvm.metal(0)
@@ -131,7 +131,7 @@ def test_ramp():
r: T.let = T.ramp(i, 3, 2)
A[0, T.ramp(0, 1, 2)] = r
- f = tvm.compile(IRModule, target=target)
+ f = tvm.tirx.compile(IRModule, target=target)
def run_and_check():
dev = tvm.metal()
@@ -158,7 +158,7 @@ def test_select_vectorize():
target = "metal"
a = np.arange(6).astype("float32")
- f = tvm.compile(IRModule, target=target)
+ f = tvm.tirx.compile(IRModule, target=target)
a.reshape(3, 2)[:, 1] = 0
def run_and_check():
@@ -181,7 +181,7 @@ def test_vectorized_uint8():
B[i * 4 + j] = T.Cast("float32", A[i * 4 + j])
a = np.arange(16).astype("uint8")
- f = tvm.compile(func, target="metal")
+ f = tvm.tirx.compile(func, target="metal")
def run_and_check():
dev = tvm.metal()
@@ -209,7 +209,7 @@ def test_func_with_trailing_pod_params():
mod = tvm.IRModule({"main": func})
- f = tvm.tirx.build(mod, target="metal")
+ f = tvm.tirx.compile(mod, target="metal").mod
src: str = f.imports[0].inspect_source()
occurrences = src.count("struct func_kernel_args_t")
assert occurrences == 1, occurrences
@@ -239,7 +239,7 @@ def test_metal_compile_callback_source_passthrough():
tvm.register_global_func("tvm_callback_metal_compile", inspect_callback, override=True)
try:
- f = tvm.compile(Module, target="metal")
+ f = tvm.tirx.compile(Module, target="metal")
dev = tvm.metal()
a = np.random.rand(n).astype("float32")
a_nd = tvm.runtime.tensor(a, dev)
@@ -289,7 +289,7 @@ def test_metal_compile_callback_mixed_formats_rejected():
tvm.register_global_func("tvm_callback_metal_compile", mixed_callback, override=True)
try:
with pytest.raises(Exception, match="inconsistent formats"):
- tvm.compile(Module, target="metal")
+ tvm.tirx.compile(Module, target="metal")
finally:
tvm_ffi.registry.remove_global_func("tvm_callback_metal_compile")
@@ -310,7 +310,7 @@ def test_export_load_with_fallback(monkeypatch, tmp_path):
B[i_0 * 32 + i_1] = A[i_0 * 32 + i_1] + 1.0
monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1")
- host_lib = tvm.compile(Module, target="metal")
+ host_lib = tvm.tirx.compile(Module, target="metal")
monkeypatch.delenv("TVM_COMPILE_FORCE_FALLBACK")
lib_path = str(tmp_path / "lib.so")
@@ -652,7 +652,7 @@ def test_pointer_byte_offsets_execute_in_threadgroup_memory():
B[1] = typed_buffer[0]
B[2] = void_buffer[0]
- executable = tvm.compile(Module, target="metal")
+ executable = tvm.tirx.compile(Module, target="metal")
source = executable.mod.imports[0].inspect_source()
assert "threadgroup float shared[16]" in source
diff --git a/tests/python/codegen/test_target_codegen_opencl.py b/tests/python/codegen/test_target_codegen_opencl.py
index 31ef34cc56..e05fbb05ec 100644
--- a/tests/python/codegen/test_target_codegen_opencl.py
+++ b/tests/python/codegen/test_target_codegen_opencl.py
@@ -45,7 +45,7 @@ def test_opencl_ternary_expression():
),
)
- fun = tvm.tirx.build(Module, target=target)
+ fun = tvm.tirx.compile(Module, target=target).mod
def run_and_check():
dev = tvm.opencl(0)
@@ -67,7 +67,7 @@ def test_opencl_ternary_expression():
T.Select(0 < T.Cast("int32", A[0]), T.Cast(dtype, 1), T.Cast(dtype, 3)),
)
- fun = tvm.tirx.build(Module, target=target)
+ fun = tvm.tirx.compile(Module, target=target).mod
def run_and_check():
dev = tvm.opencl(0)
@@ -99,7 +99,7 @@ def test_opencl_inf_nan():
with T.launch_thread("threadIdx.x", 1) as i:
C[i] = T.Cast(dtype, value)
- fun = tvm.tirx.build(Module, target=target)
+ fun = tvm.tirx.compile(Module, target=target).mod
def run_and_check():
dev = tvm.opencl(0)
@@ -129,7 +129,7 @@ def test_opencl_max():
with T.launch_thread("threadIdx.x", 1) as i:
C[i] = T.max(A[0] + T.Cast(dtype, 1), T.Cast(dtype, 0))
- fun = tvm.tirx.build(Module, target=target)
+ fun = tvm.tirx.compile(Module, target=target).mod
def run_and_check():
dev = tvm.opencl(0)
@@ -157,7 +157,7 @@ def test_opencl_erf():
with T.launch_thread("threadIdx.x", 1) as i0:
C[i0] = T.erf(A[i0])
- fun = tvm.tirx.build(Module, target=target)
+ fun = tvm.tirx.compile(Module, target=target).mod
source_str = fun.imports[0].inspect_source()
matches = re.findall("erf", source_str)
@@ -185,7 +185,7 @@ def test_opencl_type_casting():
)
def check_type_casting(n, dtype):
- fun = tvm.tirx.build(Module, target=target)
+ fun = tvm.tirx.compile(Module, target=target).mod
assembly = fun.imports[0].inspect_source()
lcond = "convert_int4(((convert_uint4(((uint4)(((convert_int(get_local_id(0))) == 3), ((convert_int(get_local_id(0))) == 3), ((convert_int(get_local_id(0))) == 3), ((convert_int(get_local_id(0))) == 3)))))"
rcond = "(convert_uint4(((((int4)(((convert_int(get_local_id(0))))+(1*0), ((convert_int(get_local_id(0))))+(1*1), ((convert_int(get_local_id(0))))+(1*2), ((convert_int(get_local_id(0))))+(1*3))) % ((int4)(3, 3, 3, 3))) == ((int4)(1, 1, 1, 1))))))))"
@@ -230,7 +230,7 @@ def test_opencl_ceil_log2(target):
with T.launch_thread("threadIdx.x", n) as i:
C[i] = T.Cast("int32", T.ceil(T.log2(T.Cast(inter_dtype, i))))
- fun = tvm.tirx.build(Module, target=target)
+ fun = tvm.tirx.compile(Module, target=target).mod
assembly = fun.imports[0].inspect_source()
if is_adreno:
pattern = "convert_float"
@@ -274,7 +274,7 @@ def test_export_load_with_fallback(monkeypatch, tmp_path):
B[i_0 * 32 + i_1] = A[i_0 * 32 + i_1] + 1.0
monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1")
- host_lib = tvm.compile(Module, target=target)
+ host_lib = tvm.tirx.compile(Module, target=target)
monkeypatch.delenv("TVM_COMPILE_FORCE_FALLBACK")
lib_path = str(tmp_path / "lib.so")
diff --git a/tests/python/codegen/test_target_codegen_riscv.py b/tests/python/codegen/test_target_codegen_riscv.py
index 5ec05b8c84..0a54680824 100644
--- a/tests/python/codegen/test_target_codegen_riscv.py
+++ b/tests/python/codegen/test_target_codegen_riscv.py
@@ -71,7 +71,7 @@ def test_rvv(target):
for j in T.vectorized(0, extent):
A[j] = 1
- f = tvm.tirx.build(load_vec, target)
+ f = tvm.tirx.compile(load_vec, target).mod
# Check RVV `vsetvli` prensence
assembly = f.inspect_source("asm")
if target_has_features("v"):
@@ -120,7 +120,7 @@ def test_rvv_vscale_llvm_dbginfo(target):
# tvm.error.InternalError: Can't fetch the lanes of a scalable vector at a compile time.
with tvm.target.Target(target):
- f = tvm.tirx.build(rvv_with_vscale, target)
+ f = tvm.tirx.compile(rvv_with_vscale, target).mod
@pytest.mark.skipif(not env.has_llvm_min_version(14), reason="need llvm >= 14")
@@ -152,7 +152,7 @@ def test_rvv_fixed_width_vectorized_loop_uses_scalable_chunks():
def check_codegen(func):
with target:
- f = tvm.tirx.build(func, target)
+ f = tvm.tirx.compile(func, target).mod
assembly = f.inspect_source("asm")
assert "vle32.v" in assembly
@@ -183,7 +183,7 @@ def test_rvv_scalable_ramp_expression():
)
with target:
- f = tvm.tirx.build(ramp_compare, target)
+ f = tvm.tirx.compile(ramp_compare, target).mod
assembly = f.inspect_source("asm")
assert "vid.v" in assembly
diff --git a/tests/python/codegen/test_target_codegen_rocm.py b/tests/python/codegen/test_target_codegen_rocm.py
index 85915a0ecb..4beb61392e 100644
--- a/tests/python/codegen/test_target_codegen_rocm.py
+++ b/tests/python/codegen/test_target_codegen_rocm.py
@@ -39,7 +39,7 @@ def test_rocm_inf_nan():
if i_0 * 128 + i_1 < 1:
C[i_0 * 128 + i_1] = T.Cast(dtype, value)
- fun = tvm.compile(Module, "rocm")
+ fun = tvm.tirx.compile(Module, "rocm")
def run_and_check():
dev = tvm.rocm(0)
@@ -94,7 +94,7 @@ def test_rocm_vectorize_add():
with T.launch_thread("threadIdx.x", 4) as i_1:
B[i_0 * 4 + i_1] = A[i_0 * 4 + i_1] + T.Broadcast(T.Cast(dtype, 1), lanes)
- fun = tvm.compile(Module, target="rocm")
+ fun = tvm.tirx.compile(Module, target="rocm")
def run_and_check():
dev = tvm.rocm(0)
@@ -125,7 +125,7 @@ def test_rocm_warp_shuffle():
A_local[0] = T.gpu_warp_shuffle(mask[0], A_local[0], 0, 32, 32)
A[tx] = A_local[0]
- mod = tvm.compile(func, target="rocm")
+ mod = tvm.tirx.compile(func, target="rocm")
def run_and_check():
dev = tvm.rocm(0)
@@ -149,7 +149,7 @@ def test_rocm_vectorized_exp():
for i in T.vectorized(0, 4):
B[i] = T.exp2(A[i])
- mod = tvm.compile(func, target="rocm")
+ mod = tvm.tirx.compile(func, target="rocm")
def run_and_check():
dev = tvm.rocm(0)
@@ -177,7 +177,7 @@ def test_export_load_with_fallback(monkeypatch, tmp_path):
B[i_0 * 32 + i_1] = A[i_0 * 32 + i_1] + 1.0
monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1")
- host_lib = tvm.compile(Module, target="rocm")
+ host_lib = tvm.tirx.compile(Module, target="rocm")
monkeypatch.delenv("TVM_COMPILE_FORCE_FALLBACK")
lib_path = str(tmp_path / "lib.so")
diff --git a/tests/python/codegen/test_target_codegen_static_init.py b/tests/python/codegen/test_target_codegen_static_init.py
index 3f30bd7e5e..eeab7ab596 100644
--- a/tests/python/codegen/test_target_codegen_static_init.py
+++ b/tests/python/codegen/test_target_codegen_static_init.py
@@ -44,7 +44,7 @@ def test_static_init():
)
mod = Module
- f = tvm.tirx.build(mod, target="llvm")
+ f = tvm.tirx.compile(mod, target="llvm").mod
a = tvm.runtime.tensor(np.zeros(10, dtype="int64"))
f(a)
diff --git a/tests/python/codegen/test_target_codegen_vulkan.py b/tests/python/codegen/test_target_codegen_vulkan.py
index f1a7c1db01..d239caa017 100644
--- a/tests/python/codegen/test_target_codegen_vulkan.py
+++ b/tests/python/codegen/test_target_codegen_vulkan.py
@@ -77,7 +77,7 @@ def test_vector_comparison(dtype):
)
# Build
- f = tvm.tirx.build(Module, target=target)
+ f = tvm.tirx.compile(Module, target=target).mod
# Verify we generate the boolx4 type declaration and the OpSelect
# v4{float,half,int} instruction
@@ -142,7 +142,7 @@ def test_array_vectorize_add(dtype):
with T.launch_thread("threadIdx.x", 4) as i_1:
B[i_0 * 4 + i_1] = A[i_0 * 4 + i_1] + one
- f = tvm.compile(Module, target=target)
+ f = tvm.tirx.compile(Module, target=target)
def run_and_check():
dev = tvm.vulkan()
@@ -174,7 +174,7 @@ def test_vulkan_bool_load():
with T.launch_thread("threadIdx.x", 128) as i_1:
B[i_0 * 128 + i_1] = T.Cast("int32", A[i_0 * 128 + i_1])
- f = tvm.compile(Module, target=target)
+ f = tvm.tirx.compile(Module, target=target)
a_np = np.random.uniform(size=arr_size) > 0.5
b_np = np.zeros((arr_size,), dtype="int32")
@@ -247,7 +247,7 @@ def test_vulkan_constant_passing(vulkan_parameter_impl, vulkan_parameter_dtype):
with T_builder.then_():
T_builder.tensor_store(B, [v_i], scalar_sum + A[v_i])
mod = ib.get()
- f_add = tvm.compile(mod, target=target)
+ f_add = tvm.tirx.compile(mod, target=target)
n = 1024
scalars = np.array([1 for _ in range(num_int_params)]).astype(dtype)
@@ -284,7 +284,7 @@ def test_vulkan_while_if():
B[0] = B[0] + iterations[0]
mod = tvm.IRModule.from_expr(while_if_gpu.with_attr("target", target))
- compiled_func = tvm.compile(mod, target=target)
+ compiled_func = tvm.tirx.compile(mod, target=target)
def run_and_check():
dev = tvm.vulkan()
@@ -319,7 +319,7 @@ def test_vulkan_local_threadidx():
B[tx + 16] = A[tx + 16]
mod = tvm.IRModule.from_expr(local_threadidx_func)
- func = tvm.compile(mod, target=target)
+ func = tvm.tirx.compile(mod, target=target)
a_np = np.arange(n).astype(dtype="int32")
b_np = np.zeros((n,), dtype="int32")
@@ -356,7 +356,7 @@ def test_vectorized_index_ramp():
bx = T.launch_thread("blockIdx.x", 1)
B[ramp_index] = A[ramp_index]
- f = tvm.compile(Module, target=target)
+ f = tvm.tirx.compile(Module, target=target)
a_np = np.random.randint(np.iinfo("int32").max, size=n).astype("int32")
b_np = np.zeros(n, dtype="int32")
@@ -394,7 +394,7 @@ def test_vectorized_index_broadcast():
bx = T.launch_thread("blockIdx.x", 1)
B[ramp_index] = A[broadcast_index]
- f = tvm.compile(Module, target=target)
+ f = tvm.tirx.compile(Module, target=target)
a_np = np.random.randint(np.iinfo("int32").max, size=n).astype("int32")
b_np = np.zeros(n, dtype="int32")
@@ -440,7 +440,7 @@ def test_negative_operand_divmod():
A[i, 0] = T.floordiv(i - offset, divisor)
A[i, 1] = T.floormod(i - offset, divisor)
- built = tvm.compile(func, target=target)
+ built = tvm.tirx.compile(func, target=target)
def run_and_check():
dev = tvm.vulkan()
@@ -493,7 +493,7 @@ def test_cooperative_matrix(out_dtype):
tgt_attrs = tvm.target.Target(target).attrs
if tgt_attrs.get("supports_cooperative_matrix"):
- f = tvm.compile(Module, target=target)
+ f = tvm.tirx.compile(Module, target=target)
def run_and_check():
dev = tvm.vulkan(0)
@@ -537,7 +537,7 @@ def test_codegen_decl_buffer():
with T.launch_thread("threadIdx.x", 1) as tx:
B_buf[tx] = A_buf[tx]
- tvm.compile(ParameterBacked, target=target)
+ tvm.tirx.compile(ParameterBacked, target=target)
@pytest.mark.gpu
@@ -556,7 +556,7 @@ def test_codegen_static_shared_memory():
A_shared[tx] = A[tx]
B[tx] = A_shared[tx]
- tvm.compile(Module, target="vulkan")
+ tvm.tirx.compile(Module, target="vulkan")
@pytest.mark.gpu
@@ -592,7 +592,7 @@ def test_unary():
B[i_0 * 64 + i_1] = tvm_intrin(A[i_0 * 64 + i_1])
target = tvm.target.Target("vulkan")
- func = tvm.compile(Module, target=target)
+ func = tvm.tirx.compile(Module, target=target)
if tvm_intrin in [tvm.tirx.asin, tvm.tirx.acos]:
data = np.random.uniform(-1.0, 1.0, size=n)
@@ -632,7 +632,7 @@ def test_export_load_with_fallback(monkeypatch, tmp_path):
B[i_0 * 32 + i_1] = A[i_0 * 32 + i_1] + 1.0
monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1")
- host_lib = tvm.compile(Module, target="vulkan")
+ host_lib = tvm.tirx.compile(Module, target="vulkan")
monkeypatch.delenv("TVM_COMPILE_FORCE_FALLBACK")
lib_path = str(tmp_path / "lib.so")
diff --git a/tests/python/codegen/test_target_codegen_webgpu.py b/tests/python/codegen/test_target_codegen_webgpu.py
index b4aedf8ed1..52f1c55b6d 100644
--- a/tests/python/codegen/test_target_codegen_webgpu.py
+++ b/tests/python/codegen/test_target_codegen_webgpu.py
@@ -35,7 +35,7 @@ def test_codegen_buffer_access_modes():
with T.launch_thread("threadIdx.x", 8) as tx:
B[tx] = A[tx]
- executable = tvm.compile(Module, target="webgpu")
+ executable = tvm.tirx.compile(Module, target="webgpu")
source = executable.mod.imports[0].inspect_source("wgsl")
assert "var<storage, read> A_ptr" in source
@@ -387,7 +387,7 @@ def test_grid_pack_guard_rejects_id_equal_to_workgroup_count():
with T.launch_thread("threadIdx.x", 1) as j: # noqa: F841
B[i] = i
- executable = tvm.compile(Module, target="webgpu")
+ executable = tvm.tirx.compile(Module, target="webgpu")
source = executable.mod.imports[0].inspect_source("wgsl")
assert re.search(r"blockIdx\.x >= \w+\.packGridDimX\) \{ return; \}", source)
diff --git a/tests/python/codegen/test_target_codegen_x86.py b/tests/python/codegen/test_target_codegen_x86.py
index e4900ebd39..eadac7fcc5 100644
--- a/tests/python/codegen/test_target_codegen_x86.py
+++ b/tests/python/codegen/test_target_codegen_x86.py
@@ -50,7 +50,7 @@ def test_fp16_to_fp32():
for i1 in T.vectorized(width):
B[i0, i1] = T.Cast("float32", A[i0, i1])
- f = tvm.tirx.build(Module, target=target)
+ f = tvm.tirx.compile(Module, target=target).mod
assembly = f.inspect_source("asm").splitlines()
if match:
diff --git a/tests/python/relax/nvshmem/test_runtime_builtin_kv_cache_transfer.py b/tests/python/relax/nvshmem/test_runtime_builtin_kv_cache_transfer.py
index 53fe84214a..bed7847fd9 100644
--- a/tests/python/relax/nvshmem/test_runtime_builtin_kv_cache_transfer.py
+++ b/tests/python/relax/nvshmem/test_runtime_builtin_kv_cache_transfer.py
@@ -177,7 +177,7 @@ def set_global_func(head_dim, dtype, target):
mod = tvm.IRModule({"main": tir_func})
with target:
mod = dl.ApplyDefaultSchedule(dl.gpu.Fallback())(mod)
- f = tvm.tirx.build(mod["main"], target=target)
+ f = tvm.tirx.compile(mod["main"], target=target).mod
builts.append(f.main)
(
diff --git a/tests/python/relax/test_frontend_nn_llm_kernel_config.py b/tests/python/relax/test_frontend_nn_llm_kernel_config.py
index 0183481693..85fb823f90 100644
--- a/tests/python/relax/test_frontend_nn_llm_kernel_config.py
+++ b/tests/python/relax/test_frontend_nn_llm_kernel_config.py
@@ -73,7 +73,7 @@ def test_non_power_of_two_head_prefill_reduces_query_tile_for_metal():
assert func.attrs["tirx.is_scheduled"]
assert _get_allocated_shared_memory(func) == 18_784
assert 18_784 <= int(target.attrs["max_shared_memory_per_block"])
- assert tvm.tirx.build(func, target=target).imports[0].inspect_source()
+ assert tvm.tirx.compile(func, target=target).mod.imports[0].inspect_source()
def test_reduced_query_tile_keeps_the_real_head_dimension():
@@ -91,7 +91,7 @@ def test_reduced_query_tile_keeps_the_real_head_dimension():
assert func.attrs["tirx.is_scheduled"]
assert _get_allocated_shared_memory(func) == 20_320
assert 20_320 <= int(target.attrs["max_shared_memory_per_block"])
- assert tvm.tirx.build(func, target=target).imports[0].inspect_source()
+ assert tvm.tirx.compile(func, target=target).mod.imports[0].inspect_source()
def test_fallback_can_select_larger_nearby_query_tile():
@@ -109,7 +109,7 @@ def test_fallback_can_select_larger_nearby_query_tile():
assert func.attrs["tirx.is_scheduled"]
assert _get_allocated_shared_memory(func) == 31_072
assert 31_072 <= int(target.attrs["max_shared_memory_per_block"])
- assert tvm.tirx.build(func, target=target).imports[0].inspect_source()
+ assert tvm.tirx.compile(func, target=target).mod.imports[0].inspect_source()
@pytest.mark.parametrize(
@@ -137,7 +137,7 @@ def test_fallback_can_expand_key_tile_for_factorization(
assert func.attrs["tirx.is_scheduled"]
assert _get_allocated_shared_memory(func) == expected_shared_memory
assert expected_shared_memory <= int(target.attrs["max_shared_memory_per_block"])
- assert tvm.tirx.build(func, target=target).imports[0].inspect_source()
+ assert tvm.tirx.compile(func, target=target).mod.imports[0].inspect_source()
def test_normal_head_prefill_keeps_existing_metal_config():
@@ -255,7 +255,10 @@ def test_merge_state_waits_before_overwriting_the_shared_lse(head_dim):
func = _merge_state_inplace(8, head_dim, "float16", target, global_symbol="merge")
lines = [
line.strip()
- for line in tvm.tirx.build(func, target=target).imports[0].inspect_source().splitlines()
+ for line in tvm.tirx.compile(func, target=target)
+ .mod.imports[0]
+ .inspect_source()
+ .splitlines()
]
loads = [i for i, line in enumerate(lines) if "= S_ptr[" in line]
diff --git a/tests/python/relax/test_frontend_nn_llm_sequence_prefill_masked.py b/tests/python/relax/test_frontend_nn_llm_sequence_prefill_masked.py
index b04d0a966c..044cc66101 100644
--- a/tests/python/relax/test_frontend_nn_llm_sequence_prefill_masked.py
+++ b/tests/python/relax/test_frontend_nn_llm_sequence_prefill_masked.py
@@ -124,7 +124,7 @@ def _build_masked_prefill(h_kv, h_q, d, dtype, target, mask_mode="padded"):
mask_mode=mask_mode,
)
mod = tvm.IRModule({"main": tir_func})
- return tvm.tirx.build(mod["main"], target=target), sm_scale
+ return tvm.tirx.compile(mod["main"], target=target).mod, sm_scale
def _run_case(
diff --git a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_cpu.py b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_cpu.py
index 6fc0b516bc..8a3ca8d6f1 100644
--- a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_cpu.py
+++ b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_cpu.py
@@ -194,7 +194,7 @@ def set_global_func(head_dim, dtype, layer_sliding_window_size=1024):
mod = tvm.IRModule({"main": tir_func})
with target:
mod = dl.ApplyDefaultSchedule(dl.gpu.Fallback())(mod)
- f = tvm.tirx.build(mod["main"], target=target)
+ f = tvm.tirx.compile(mod["main"], target=target).mod
builts.append(f.main)
builts = tuple(builts)
_COMPILED_KERNEL_CACHE[cache_key] = builts
diff --git a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_flashinfer.py b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_flashinfer.py
index 13fa9856d8..1cff41909d 100644
--- a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_flashinfer.py
+++ b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_flashinfer.py
@@ -148,7 +148,7 @@ def set_global_func(rope_mode: RopeMode, target):
mod = tvm.IRModule({"main": tir_func})
with target:
mod = dl.ApplyDefaultSchedule(dl.gpu.Fallback())(mod)
- f = tvm.tirx.build(mod["main"], target=target)
+ f = tvm.tirx.compile(mod["main"], target=target).mod
builts.append(f.main)
(
diff --git a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_flashinfer.py b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_flashinfer.py
index 2f86dedc68..183420d7a8 100644
--- a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_flashinfer.py
+++ b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_flashinfer.py
@@ -157,7 +157,7 @@ def set_global_func(dtype, target):
mod = tvm.IRModule({"main": tir_func})
with target:
mod = dl.ApplyDefaultSchedule(dl.gpu.Fallback())(mod)
- f = tvm.tirx.build(mod["main"], target=target)
+ f = tvm.tirx.compile(mod["main"], target=target).mod
builts.append(f.main)
(
diff --git a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_tir.py b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_tir.py
index a7eea62702..85eea3bd31 100644
--- a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_tir.py
+++ b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_tir.py
@@ -132,7 +132,7 @@ def set_global_func(dtype, target):
mod = tvm.IRModule({"main": tir_func})
with target:
mod = dl.ApplyDefaultSchedule(dl.gpu.Fallback())(mod)
- f = tvm.tirx.build(mod["main"], target=target)
+ f = tvm.tirx.compile(mod["main"], target=target).mod
builts.append(f.main)
(
diff --git a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py
index 7f878f118f..f4c58952ce 100644
--- a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py
+++ b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py
@@ -117,7 +117,7 @@ def test_paged_prefill_layer_sliding_window_mask(device_type):
page_size=page_size,
sliding_window_size=3,
)
- built = tvm.tirx.build(tir_func, target=target)
+ built = tvm.tirx.compile(tir_func, target=target).mod
q = np.zeros((2, num_qo_heads, head_dim), dtype="float16")
q_indptr = np.array([0, 2], dtype="int32")
@@ -269,7 +269,7 @@ def set_global_func(head_dim, dtype, target):
mod = tvm.IRModule({"main": tir_func})
with target:
mod = dl.ApplyDefaultSchedule(dl.gpu.Fallback())(mod)
- f = tvm.tirx.build(mod["main"], target=target)
+ f = tvm.tirx.compile(mod["main"], target=target).mod
builts.append(f.main)
builts = tuple(builts)
_COMPILED_KERNEL_CACHE[cache_key] = builts
diff --git a/tests/python/relax/test_runtime_builtin_rnn_state.py b/tests/python/relax/test_runtime_builtin_rnn_state.py
index 4df1f6c712..8541c20cdd 100644
--- a/tests/python/relax/test_runtime_builtin_rnn_state.py
+++ b/tests/python/relax/test_runtime_builtin_rnn_state.py
@@ -81,7 +81,7 @@ def set_global_func():
mod = tvm.IRModule({"main": tir_func})
with target:
mod = dl.ApplyDefaultSchedule(dl.gpu.Fallback())(mod) # pylint: disable=not-callable
- f = tvm.tirx.build(mod["main"], target=target)
+ f = tvm.tirx.compile(mod["main"], target=target).mod
return f.main
_f_tir_gets, _f_tir_sets = [], []
diff --git a/tests/python/runtime/test_executable.py b/tests/python/runtime/test_executable.py
index 7af211084b..512b830fb9 100644
--- a/tests/python/runtime/test_executable.py
+++ b/tests/python/runtime/test_executable.py
@@ -42,7 +42,7 @@ class MyModule:
def test_executable_init():
"""Test initialization of Executable class."""
- lib = tvm.tirx.build(MyModule, target="llvm")
+ lib = tvm.tirx.compile(MyModule, target="llvm").mod
executable = Executable(lib)
assert executable.mod is lib
@@ -51,7 +51,7 @@ def test_executable_init():
def test_executable_getitem():
"""Test __getitem__ method of Executable class."""
- lib = tvm.tirx.build(MyModule, target="llvm")
+ lib = tvm.tirx.compile(MyModule, target="llvm").mod
executable = Executable(lib)
# Jit the module first
@@ -73,7 +73,7 @@ def test_executable_getitem():
def test_executable_jit_already_jitted():
"""Test jit method when module is already jitted."""
- lib = tvm.tirx.build(MyModule, target="llvm")
+ lib = tvm.tirx.compile(MyModule, target="llvm").mod
executable = Executable(lib)
# First jit call
@@ -102,7 +102,7 @@ def test_executable_jit_already_jitted():
def test_executable_export_library():
"""Test export_library method."""
- lib = tvm.tirx.build(MyModule, target="llvm")
+ lib = tvm.tirx.compile(MyModule, target="llvm").mod
executable = Executable(lib)
# Create a temporary directory for the library
@@ -137,7 +137,7 @@ def test_executable_export_library():
def test_executable_export_library_with_workspace():
"""Test export_library method with workspace_dir."""
- lib = tvm.tirx.build(MyModule, target="llvm")
+ lib = tvm.tirx.compile(MyModule, target="llvm").mod
executable = Executable(lib)
# Create temporary directories
@@ -177,7 +177,7 @@ def test_executable_integration():
"""Integration test for Executable with a simple TVM module."""
# Create target and build
target = tvm.target.Target("llvm")
- lib = tvm.tirx.build(MyModule, target=target)
+ lib = tvm.tirx.compile(MyModule, target=target).mod
# Create an executable
executable = Executable(lib)
@@ -233,7 +233,7 @@ def test_executable_jit_force_recompile():
"""Test jit method with force_recompile=True."""
# Create target and build
target = tvm.target.Target("c")
- lib = tvm.tirx.build(MyModule, target=target)
+ lib = tvm.tirx.compile(MyModule, target=target).mod
# Create an executable
executable = Executable(lib)
diff --git a/tests/python/runtime/test_runtime_measure.py b/tests/python/runtime/test_runtime_measure.py
index 34f4b69990..29d2832e99 100644
--- a/tests/python/runtime/test_runtime_measure.py
+++ b/tests/python/runtime/test_runtime_measure.py
@@ -36,7 +36,7 @@ def test_min_repeat_ms():
fout.write("c")
X = te.compute((), lambda: tvm.tirx.call_packed("my_debug", filename))
- func = tvm.tirx.build(te.create_function([X]))
+ func = tvm.tirx.compile(te.create_function([X])).mod
x = tvm.runtime.empty((), dtype="int32")
ftimer = func.time_evaluator(func.entry_name, tvm.cpu(), number=1, repeat=1)
diff --git a/tests/python/runtime/test_runtime_module_export.py b/tests/python/runtime/test_runtime_module_export.py
index 654ae261e8..65117e958c 100644
--- a/tests/python/runtime/test_runtime_module_export.py
+++ b/tests/python/runtime/test_runtime_module_export.py
@@ -33,8 +33,8 @@ def test_import_static_library():
irmod0 = tvm.IRModule.from_expr(te.create_function([A, B]).with_attr("global_symbol", "myadd0"))
irmod1 = tvm.IRModule.from_expr(te.create_function([A, B]).with_attr("global_symbol", "myadd1"))
- mod0 = tvm.tirx.build(irmod0, target="llvm")
- mod1 = tvm.tirx.build(irmod1, target="llvm")
+ mod0 = tvm.tirx.compile(irmod0, target="llvm").mod
+ mod1 = tvm.tirx.compile(irmod1, target="llvm").mod
assert mod0.implements_function("myadd0")
assert mod1.implements_function("myadd1")
diff --git a/tests/python/runtime/test_runtime_module_load.py b/tests/python/runtime/test_runtime_module_load.py
index 5993bc6c61..dd09aa1fad 100644
--- a/tests/python/runtime/test_runtime_module_load.py
+++ b/tests/python/runtime/test_runtime_module_load.py
@@ -64,7 +64,7 @@ def test_dso_module_load():
mod = tvm.IRModule.from_expr(
tvm.tirx.Function([Ab], stmt).with_attr("global_symbol", "main")
)
- m = tvm.tirx.build(mod, target="llvm")
+ m = tvm.tirx.compile(mod, target="llvm").mod
for name in names:
m.write_to_file(name)
@@ -131,8 +131,8 @@ def test_device_module_dump():
import tvm
f1 = tvm.runtime.load_module(path_dso)
- a = tvm.runtime.tensor(np.random.uniform(size=1024).astype(A.dtype), dev)
- b = tvm.runtime.tensor(np.zeros(1024, dtype=A.dtype), dev)
+ a = tvm.runtime.tensor(np.random.uniform(size=1024).astype(str(A.dtype)), dev)
+ b = tvm.runtime.tensor(np.zeros(1024, dtype=str(A.dtype)), dev)
f1(a, b)
np.testing.assert_equal(b.numpy(), a.numpy() + 1)
@@ -153,8 +153,8 @@ def test_device_module_dump():
def run_and_check():
dev = tvm.device(device, 0)
- a = tvm.runtime.tensor(np.random.uniform(size=1024).astype(A.dtype), dev)
- b = tvm.runtime.tensor(np.zeros(1024, dtype=A.dtype), dev)
+ a = tvm.runtime.tensor(np.random.uniform(size=1024).astype(str(A.dtype)), dev)
+ b = tvm.runtime.tensor(np.zeros(1024, dtype=str(A.dtype)), dev)
f["main"](a, b)
np.testing.assert_equal(b.numpy(), a.numpy() + 1)
@@ -181,8 +181,8 @@ def test_combine_module_llvm():
def check_llvm():
dev = tvm.cpu(0)
temp = utils.tempdir()
- fadd1 = tvm.tirx.build(mod1, "llvm")
- fadd2 = tvm.tirx.build(mod2, "llvm")
+ fadd1 = tvm.tirx.compile(mod1, "llvm").mod
+ fadd2 = tvm.tirx.compile(mod2, "llvm").mod
path1 = temp.relpath("myadd1.o")
path2 = temp.relpath("myadd2.o")
path_dso = temp.relpath("mylib.so")
@@ -193,8 +193,8 @@ def test_combine_module_llvm():
m = tvm.runtime.load_module(path_dso)
fadd1 = m["myadd1"]
fadd2 = m["myadd2"]
- a = tvm.runtime.tensor(np.random.uniform(size=nn).astype(A.dtype), dev)
- b = tvm.runtime.tensor(np.zeros(nn, dtype=A.dtype), dev)
+ a = tvm.runtime.tensor(np.random.uniform(size=nn).astype(str(A.dtype)), dev)
+ b = tvm.runtime.tensor(np.zeros(nn, dtype=str(A.dtype)), dev)
fadd1(a, b)
np.testing.assert_equal(b.numpy(), a.numpy() + 1)
fadd2(a, b)
@@ -207,8 +207,8 @@ def test_combine_module_llvm():
return
temp = utils.tempdir()
print("Running popen check")
- fadd1 = tvm.tirx.build(mod1.with_attr("system_lib_prefix", ""), "llvm")
- fadd2 = tvm.tirx.build(mod2.with_attr("system_lib_prefix", ""), "llvm")
+ fadd1 = tvm.tirx.compile(mod1.with_attr("system_lib_prefix", ""), "llvm").mod
+ fadd2 = tvm.tirx.compile(mod2.with_attr("system_lib_prefix", ""), "llvm").mod
path1 = temp.relpath("myadd1.o")
path2 = temp.relpath("myadd2.o")
path_dso = temp.relpath("mylib.so")
@@ -225,8 +225,8 @@ def test_combine_module_llvm():
ctypes.CDLL(path_dso)
# Load the system wide library
mm = tvm.runtime.system_lib()
- a = tvm.runtime.tensor(np.random.uniform(size=nn).astype(A.dtype), dev)
- b = tvm.runtime.tensor(np.zeros(nn, dtype=A.dtype), dev)
+ a = tvm.runtime.tensor(np.random.uniform(size=nn).astype(str(A.dtype)), dev)
+ b = tvm.runtime.tensor(np.zeros(nn, dtype=str(A.dtype)), dev)
mm["myadd1"](a, b)
np.testing.assert_equal(b.numpy(), a.numpy() + 1)
mm["myadd2"](a, b)
diff --git a/tests/python/runtime/test_runtime_module_property.py b/tests/python/runtime/test_runtime_module_property.py
index 9ff4a77e28..1f1ea4be97 100644
--- a/tests/python/runtime/test_runtime_module_property.py
+++ b/tests/python/runtime/test_runtime_module_property.py
@@ -33,7 +33,7 @@ def create_csource_module():
def create_llvm_module():
A = te.placeholder((1024,), name="A")
B = te.compute(A.shape, lambda *i: A(*i) + 1.0, name="B")
- return tvm.tirx.build(te.create_function([A, B]), target="llvm")
+ return tvm.tirx.compile(te.create_function([A, B]), target="llvm").mod
def test_property():
diff --git a/tests/python/s_tir/base/test_compilation_pipeline.py b/tests/python/s_tir/base/test_compilation_pipeline.py
index a125ddc143..4ef79e4013 100644
--- a/tests/python/s_tir/base/test_compilation_pipeline.py
+++ b/tests/python/s_tir/base/test_compilation_pipeline.py
@@ -53,7 +53,7 @@ def test_default_pipeline_selects_dialect(mixed):
finalized_modules.append({gv.name_hint for gv in mod.functions})
with tvm.transform.PassContext(instruments=[RecordFinalization()]):
- module = tvm.tirx.build(tvm.IRModule(functions), target="llvm")
+ module = tvm.tirx.compile(tvm.IRModule(functions), target="llvm").mod
assert finalized_modules == [set(functions)]
source = np.arange(16, dtype="float32")
a = tvm.runtime.tensor(source)
diff --git a/tests/python/s_tir/base/test_s_tir_texture_scope.py b/tests/python/s_tir/base/test_s_tir_texture_scope.py
index 646106dcdb..5fb61825d9 100644
--- a/tests/python/s_tir/base/test_s_tir_texture_scope.py
+++ b/tests/python/s_tir/base/test_s_tir_texture_scope.py
@@ -25,7 +25,7 @@ from tvm.ir.module import IRModule
from tvm.s_tir.backend.adreno import pipeline as adreno_pipeline
from tvm.script import s_tir as Ts
from tvm.script import tirx as T
-from tvm.tirx.build import split_host_device_mods
+from tvm.tirx._compile import _split_host_device_modules
from tvm.tirx.compilation_pipeline import finalize_device_passes
@@ -67,7 +67,7 @@ def test_texture_scope():
lowered = tirx.transform.BindTarget(target.with_host("c"))(sch.mod)
pipeline, _, _ = adreno_pipeline.default_tir_pipeline()
lowered = pipeline(lowered)
- _, device_mods = split_host_device_mods(lowered)
+ _, device_mods = _split_host_device_modules(lowered)
assert len(device_mods) == 1
(device_target, _backend_config), device_mod = next(iter(device_mods.items()))
device_mod = finalize_device_passes()(device_mod)
diff --git a/tests/python/s_tir/dlight/test_cpu_reduction.py b/tests/python/s_tir/dlight/test_cpu_reduction.py
index 88793f48a0..7be5f27b42 100644
--- a/tests/python/s_tir/dlight/test_cpu_reduction.py
+++ b/tests/python/s_tir/dlight/test_cpu_reduction.py
@@ -141,9 +141,9 @@ def _codegen_llvm_ir(mod, target):
bound = tirx.transform.BindTarget(target.with_host(target))(mod)
pipeline, finalize_host, _ = tirx.get_tir_pipeline("s_tir")
lowered = pipeline(bound)
- from tvm.tirx.build import split_host_device_mods
+ from tvm.tirx._compile import _split_host_device_modules
- host_mod, _ = split_host_device_mods(lowered)
+ host_mod, _ = _split_host_device_modules(lowered)
host_mod = finalize_host()(host_mod)
built = tvm.target.codegen.build_module(host_mod, target)
return built.inspect_source("ll")
@@ -154,9 +154,9 @@ def _codegen_asm(mod, target):
bound = tirx.transform.BindTarget(target.with_host(target))(mod)
pipeline, finalize_host, _ = tirx.get_tir_pipeline("s_tir")
lowered = pipeline(bound)
- from tvm.tirx.build import split_host_device_mods
+ from tvm.tirx._compile import _split_host_device_modules
- host_mod, _ = split_host_device_mods(lowered)
+ host_mod, _ = _split_host_device_modules(lowered)
host_mod = finalize_host()(host_mod)
built = tvm.target.codegen.build_module(host_mod, target)
return built.inspect_source("s")
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_force_narrow_index_to_i32.py b/tests/python/s_tir/transform/test_s_tir_transform_force_narrow_index_to_i32.py
index 9e9625eaff..9cbfba3a08 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_force_narrow_index_to_i32.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_force_narrow_index_to_i32.py
@@ -170,7 +170,7 @@ def test_metal_simdgroup_matmul_builds():
mod = dl.ApplyDefaultSchedule(dl.gpu.Matmul())(tvm.IRModule({"main": main}))
assert "metal.simdgroup" in mod.script()
mod = tvm.s_tir.transform.ForceNarrowIndexToInt32()(mod)
- tvm.tirx.build(mod, target=target)
+ tvm.tirx.compile(mod, target=target).mod
if __name__ == "__main__":
diff --git a/tests/python/te/test_te_create_function.py b/tests/python/te/test_te_create_function.py
index 52fb7c4f62..35162e2d4f 100644
--- a/tests/python/te/test_te_create_function.py
+++ b/tests/python/te/test_te_create_function.py
@@ -414,7 +414,7 @@ def test_topi_float_unary_accepts_float_input(op_name, dtype):
func = te.create_function([x, out]).with_attr("target", tvm.target.Target("llvm"))
mod = tvm.IRModule({"main": func})
- compiled = tvm.tirx.build(mod, target="llvm")
+ compiled = tvm.tirx.compile(mod, target="llvm").mod
assert compiled is not None
diff --git a/tests/python/tirx-base/test_tir_base.py b/tests/python/tirx-base/test_tir_base.py
index 07a3b981fd..8efa9539fc 100644
--- a/tests/python/tirx-base/test_tir_base.py
+++ b/tests/python/tirx-base/test_tir_base.py
@@ -33,7 +33,7 @@ def build_tir_func(func):
if pass_ctx.config.get("tirx.noalias", True):
func = func.with_attr("tirx.noalias", True)
mod = tvm.IRModule({"main": func})
- func = tvm.compile(mod)
+ func = tvm.tirx.compile(mod)
return func
diff --git a/tests/python/tirx-base/test_tir_imm_values.py b/tests/python/tirx-base/test_tir_imm_values.py
index 68f2a718ba..4fbd80deb5 100644
--- a/tests/python/tirx-base/test_tir_imm_values.py
+++ b/tests/python/tirx-base/test_tir_imm_values.py
@@ -150,7 +150,7 @@ def test_tir_too_large_literal_f64():
def imm_overflow_fp64() -> T.float64:
return T.float64(1.7976e309)
- f = tvm.compile(imm_overflow_fp64, target="llvm")
+ f = tvm.tirx.compile(imm_overflow_fp64, target="llvm")
assert math.isinf(f())
@@ -273,7 +273,7 @@ def test_tir_floatimm_const_fold():
z[()] = x / y
def __wrap_build(f):
- lib = tvm.compile(f, target="llvm")
+ lib = tvm.tirx.compile(f, target="llvm")
z = tvm.runtime.tensor(np.zeros([]).astype("float32"))
def _func(x, y):
@@ -334,11 +334,11 @@ def test_tir_int8_const_fold():
def imm_floordiv(x: T.int8, y: T.int8) -> T.int8:
return T.floordiv(x, y)
- fmul = tvm.compile(imm_multiply, target="llvm")
- fadd = tvm.compile(imm_add, target="llvm")
- fsub = tvm.compile(imm_sub, target="llvm")
- ffloordiv = tvm.compile(imm_floordiv, target="llvm")
- ftruncdiv = tvm.compile(imm_truncdiv, target="llvm")
+ fmul = tvm.tirx.compile(imm_multiply, target="llvm")
+ fadd = tvm.tirx.compile(imm_add, target="llvm")
+ fsub = tvm.tirx.compile(imm_sub, target="llvm")
+ ffloordiv = tvm.tirx.compile(imm_floordiv, target="llvm")
+ ftruncdiv = tvm.tirx.compile(imm_truncdiv, target="llvm")
# overflow
check_tir_const_fold("int8", lambda x, y: x + y, fadd, 127, 1, -128)
@@ -390,11 +390,11 @@ def test_tir_uint8_const_fold():
def imm_floordiv(x: T.uint8, y: T.uint8) -> T.uint8:
return T.floordiv(x, y)
- fmul = tvm.compile(imm_multiply, target="llvm")
- fadd = tvm.compile(imm_add, target="llvm")
- fsub = tvm.compile(imm_sub, target="llvm")
- ffloordiv = tvm.compile(imm_floordiv, target="llvm")
- ftruncdiv = tvm.compile(imm_truncdiv, target="llvm")
+ fmul = tvm.tirx.compile(imm_multiply, target="llvm")
+ fadd = tvm.tirx.compile(imm_add, target="llvm")
+ fsub = tvm.tirx.compile(imm_sub, target="llvm")
+ ffloordiv = tvm.tirx.compile(imm_floordiv, target="llvm")
+ ftruncdiv = tvm.tirx.compile(imm_truncdiv, target="llvm")
# overflow
check_tir_const_fold("uint8", lambda x, y: x + y, fadd, 255, 1, 0)
@@ -463,13 +463,13 @@ def test_tir_int32_const_fold():
def imm_floormod(x: T.int32, y: T.int32) -> T.int32:
return T.floormod(x, y)
- fmul = tvm.compile(imm_multiply, target="llvm")
- fadd = tvm.compile(imm_add, target="llvm")
- fsub = tvm.compile(imm_sub, target="llvm")
- ffloordiv = tvm.compile(imm_floordiv, target="llvm")
- ffloormod = tvm.compile(imm_floormod, target="llvm")
- ftruncdiv = tvm.compile(imm_truncdiv, target="llvm")
- ftruncmod = tvm.compile(imm_truncmod, target="llvm")
+ fmul = tvm.tirx.compile(imm_multiply, target="llvm")
+ fadd = tvm.tirx.compile(imm_add, target="llvm")
+ fsub = tvm.tirx.compile(imm_sub, target="llvm")
+ ffloordiv = tvm.tirx.compile(imm_floordiv, target="llvm")
+ ffloormod = tvm.tirx.compile(imm_floormod, target="llvm")
+ ftruncdiv = tvm.tirx.compile(imm_truncdiv, target="llvm")
+ ftruncmod = tvm.tirx.compile(imm_truncmod, target="llvm")
# i32 overflow is not specified, only check for range
assert -(2**31) <= int(tirx.const(2**31 - 1, "int32") + tirx.const(1, "int32")) < 2**31
@@ -543,11 +543,11 @@ def test_tir_uint32_const_fold():
def imm_floordiv(x: T.uint32, y: T.uint32) -> T.uint32:
return T.floordiv(x, y)
- fmul = tvm.compile(imm_multiply, target="llvm")
- fadd = tvm.compile(imm_add, target="llvm")
- fsub = tvm.compile(imm_sub, target="llvm")
- ffloordiv = tvm.compile(imm_floordiv, target="llvm")
- ftruncdiv = tvm.compile(imm_truncdiv, target="llvm")
+ fmul = tvm.tirx.compile(imm_multiply, target="llvm")
+ fadd = tvm.tirx.compile(imm_add, target="llvm")
+ fsub = tvm.tirx.compile(imm_sub, target="llvm")
+ ffloordiv = tvm.tirx.compile(imm_floordiv, target="llvm")
+ ftruncdiv = tvm.tirx.compile(imm_truncdiv, target="llvm")
# u32 overflow is not specified, only check for range
assert 0 <= int(tirx.const(2**32 - 1, "uint32") + tirx.const(1, "uint32")) < 2**32
diff --git a/tests/python/tirx-base/test_tir_intrin.py b/tests/python/tirx-base/test_tir_intrin.py
index 503c727af5..5e02ee3348 100644
--- a/tests/python/tirx-base/test_tir_intrin.py
+++ b/tests/python/tirx-base/test_tir_intrin.py
@@ -66,7 +66,7 @@ def _binary_kernel(op, rhs_dtype="float32"):
def test_nearbyint():
- func = tvm.compile(_unary_kernel(tvm.tirx.nearbyint), target="llvm")
+ func = tvm.tirx.compile(_unary_kernel(tvm.tirx.nearbyint), target="llvm")
dev = tvm.cpu(0)
n = 10
@@ -92,7 +92,7 @@ def test_round_ties_to_even(target):
if target != "c" and not tvm.testing.device_enabled(target):
pytest.skip(f"{target} not enabled")
- func = tvm.compile(_unary_kernel(tvm.tirx.round), target=target)
+ func = tvm.tirx.compile(_unary_kernel(tvm.tirx.round), target=target)
dev = tvm.cpu(0)
# Midpoint values where ties-to-even and ties-away differ
@@ -134,7 +134,7 @@ def test_unary_intrin():
]
def run_test(tvm_intrin, np_func, atol=1e-5, rtol=1e-5):
- func = tvm.compile(_unary_kernel(tvm_intrin), target="llvm")
+ func = tvm.tirx.compile(_unary_kernel(tvm_intrin), target="llvm")
dev = tvm.cpu(0)
n = 10
@@ -183,7 +183,7 @@ def test_asin_acos_boundary_values():
]
def run_test(tvm_intrin, np_func):
- func = tvm.compile(_unary_kernel(tvm_intrin), target="llvm")
+ func = tvm.tirx.compile(_unary_kernel(tvm_intrin), target="llvm")
dev = tvm.cpu(0)
@@ -230,7 +230,7 @@ def test_binary_intrin():
]
def run_test(tvm_intrin, np_func):
- func = tvm.compile(_binary_kernel(tvm_intrin), target="llvm")
+ func = tvm.tirx.compile(_binary_kernel(tvm_intrin), target="llvm")
dev = tvm.cpu(0)
n = 10
@@ -245,7 +245,7 @@ def test_binary_intrin():
def test_ldexp():
- func = tvm.compile(_binary_kernel(tvm.tirx.ldexp, rhs_dtype="int32"), target="llvm")
+ func = tvm.tirx.compile(_binary_kernel(tvm.tirx.ldexp, rhs_dtype="int32"), target="llvm")
dev = tvm.cpu(0)
n = 10
@@ -281,7 +281,7 @@ def test_clz(target, dtype):
clz[np.bitwise_and(x, x - 1) == 0] -= 1
return clz
- func = tvm.compile(
+ func = tvm.tirx.compile(
_unary_kernel(tvm.tirx.clz, dtype, "int32", gpu=target.kind.name == "vulkan"), target=target
)
diff --git a/tests/python/tirx-base/test_tir_ptx_cp_async.py b/tests/python/tirx-base/test_tir_ptx_cp_async.py
index 8af3437ad7..8345385ea3 100644
--- a/tests/python/tirx-base/test_tir_ptx_cp_async.py
+++ b/tests/python/tirx-base/test_tir_ptx_cp_async.py
@@ -53,7 +53,7 @@ def ptx_cp_async(A: T.Tensor((32, 128), "float16"), B: T.Tensor((32, 128), "floa
def test_ptx_cp_async():
f = ptx_cp_async
- mod = tvm.compile(f, target="cuda")
+ mod = tvm.tirx.compile(f, target="cuda")
A_np = np.random.rand(32, 128).astype("float16")
B_np = np.zeros((32, 128)).astype("float16")
diff --git a/tests/python/tirx-base/test_tir_ptx_griddepcontrol.py b/tests/python/tirx-base/test_tir_ptx_griddepcontrol.py
index 38f51bc6de..1d1bed0ef4 100644
--- a/tests/python/tirx-base/test_tir_ptx_griddepcontrol.py
+++ b/tests/python/tirx-base/test_tir_ptx_griddepcontrol.py
@@ -38,7 +38,7 @@ def ptx_griddepcontrol(A: T.Tensor((32,), "float32"), B: T.Tensor((32,), "float3
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
def test_ptx_griddepcontrol():
f = ptx_griddepcontrol
- mod = tvm.compile(f, target="cuda")
+ mod = tvm.tirx.compile(f, target="cuda")
A_np = np.random.default_rng(0).standard_normal(32).astype("float32")
B_np = np.zeros((32,), dtype="float32")
diff --git a/tests/python/tirx-base/test_tir_ptx_ldmatrix.py b/tests/python/tirx-base/test_tir_ptx_ldmatrix.py
index c94a962689..338aeca995 100644
--- a/tests/python/tirx-base/test_tir_ptx_ldmatrix.py
+++ b/tests/python/tirx-base/test_tir_ptx_ldmatrix.py
@@ -61,7 +61,9 @@ def test_ptx_ldmatrix():
for num in [1, 2, 4]:
for trans in [False, True]:
- mod = tvm.compile(f.specialize({param_num: num, param_trans: trans}), target="cuda")
+ mod = tvm.tirx.compile(
+ f.specialize({param_num: num, param_trans: trans}), target="cuda"
+ )
A_np = np.random.rand(16, 16).astype("float16")
A_mask_np = np.zeros_like(A_np)
if num == 1:
diff --git a/tests/python/tirx-base/test_tir_ptx_mma.py b/tests/python/tirx-base/test_tir_ptx_mma.py
index 5f2fa83f63..103ce7cc7f 100644
--- a/tests/python/tirx-base/test_tir_ptx_mma.py
+++ b/tests/python/tirx-base/test_tir_ptx_mma.py
@@ -69,7 +69,7 @@ def gemm_mma_m8n8k4_row_col_fp64pf64fp64(
@pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0")
def test_gemm_mma_m8n8k4_row_col_fp64pf64fp64():
mod = tvm.IRModule.from_expr(gemm_mma_m8n8k4_row_col_fp64pf64fp64)
- cuda_mod = tvm.compile(mod, target="cuda")
+ cuda_mod = tvm.tirx.compile(mod, target="cuda")
A_np = np.random.uniform(-1, 1, [8, 4]).astype("float64")
B_np = np.random.uniform(-1, 1, [8, 4]).astype("float64")
@@ -145,7 +145,7 @@ def gemm_mma_m8n8k4_row_row_fp16fp16fp16(
@pytest.mark.skipif(not env.has_cuda_compute(7), reason="need cuda compute >= 7.0")
def test_gemm_mma_m8n8k4_row_row_fp16fp16fp16():
mod = tvm.IRModule.from_expr(gemm_mma_m8n8k4_row_row_fp16fp16fp16)
- cuda_mod = tvm.compile(mod, target="cuda")
+ cuda_mod = tvm.tirx.compile(mod, target="cuda")
A_np = np.random.uniform(-1, 1, [16, 4]).astype("float16")
B_np = np.random.uniform(-1, 1, [4, 16]).astype("float16")
@@ -228,7 +228,7 @@ def gemm_mma_m8n8k4_row_row_fp16fp16fp32(
@pytest.mark.skipif(not env.has_cuda_compute(7), reason="need cuda compute >= 7.0")
def test_gemm_mma_m8n8k4_row_row_fp16fp16fp32():
mod = tvm.IRModule.from_expr(gemm_mma_m8n8k4_row_row_fp16fp16fp32)
- cuda_mod = tvm.compile(mod, target="cuda")
+ cuda_mod = tvm.tirx.compile(mod, target="cuda")
A_np = np.random.uniform(-1, 1, [16, 4]).astype("float16")
B_np = np.random.uniform(-1, 1, [4, 16]).astype("float16")
@@ -299,7 +299,7 @@ def gemm_mma_m8n8k16_row_col_s8s8s32(
@pytest.mark.skipif(not env.has_cuda_compute(7, 5), reason="need cuda compute >= 7.5")
def test_gemm_mma_m8n8k16_row_col_s8s8s32():
mod = tvm.IRModule.from_expr(gemm_mma_m8n8k16_row_col_s8s8s32)
- cuda_mod = tvm.compile(mod, target="cuda")
+ cuda_mod = tvm.tirx.compile(mod, target="cuda")
A_np = np.random.uniform(-10, 10, [8, 16]).astype("int8")
B_np = np.random.uniform(-10, 10, [8, 16]).astype("int8")
@@ -370,7 +370,7 @@ def gemm_mma_m8n8k16_row_col_s8u8s32(
@pytest.mark.skipif(not env.has_cuda_compute(7, 5), reason="need cuda compute >= 7.5")
def test_gemm_mma_m8n8k16_row_col_s8u8s32():
mod = tvm.IRModule.from_expr(gemm_mma_m8n8k16_row_col_s8u8s32)
- cuda_mod = tvm.compile(mod, target="cuda")
+ cuda_mod = tvm.tirx.compile(mod, target="cuda")
A_np = np.random.uniform(-10, 10, [8, 16]).astype("int8")
B_np = np.random.uniform(-10, 10, [8, 16]).astype("uint8")
@@ -441,7 +441,7 @@ def gemm_mma_m8n8k32_row_col_s4s4s32(
@pytest.mark.skipif(not env.has_cuda_compute(7, 5), reason="need cuda compute >= 7.5")
def test_gemm_mma_m8n8k32_row_col_s4s4s32():
mod = tvm.IRModule.from_expr(gemm_mma_m8n8k32_row_col_s4s4s32)
- cuda_mod = tvm.compile(mod, target="cuda")
+ cuda_mod = tvm.tirx.compile(mod, target="cuda")
def run_and_check():
ctx = tvm.cuda()
@@ -506,7 +506,7 @@ def gemm_mma_m8n8k32_row_col_s4u4s32(
@pytest.mark.skipif(not env.has_cuda_compute(7, 5), reason="need cuda compute >= 7.5")
def test_gemm_mma_m8n8k32_row_col_s4u4s32():
mod = tvm.IRModule.from_expr(gemm_mma_m8n8k32_row_col_s4u4s32)
- cuda_mod = tvm.compile(mod, target="cuda")
+ cuda_mod = tvm.tirx.compile(mod, target="cuda")
def run_and_check():
ctx = tvm.cuda()
@@ -573,7 +573,7 @@ def gemm_mma_m16n8k8_row_col_fp16fp16fp32(
@pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0")
def test_gemm_mma_m16n8k8_row_col_fp16fp16fp32():
mod = tvm.IRModule.from_expr(gemm_mma_m16n8k8_row_col_fp16fp16fp32)
- cuda_mod = tvm.compile(mod, target="cuda")
+ cuda_mod = tvm.tirx.compile(mod, target="cuda")
A_np = np.random.uniform(-1, 1, [16, 8]).astype("float16")
B_np = np.random.uniform(-1, 1, [8, 8]).astype("float16")
@@ -649,7 +649,7 @@ def gemm_mma_m16n8k16_row_col_fp16fp16fp16(
@pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0")
def test_gemm_mma_m16n8k16_row_col_fp16fp16fp16():
mod = tvm.IRModule.from_expr(gemm_mma_m16n8k16_row_col_fp16fp16fp16)
- cuda_mod = tvm.compile(mod, target="cuda")
+ cuda_mod = tvm.tirx.compile(mod, target="cuda")
A_np = np.random.uniform(-1, 1, [16, 16]).astype("float16")
B_np = np.random.uniform(-1, 1, [8, 16]).astype("float16")
@@ -725,7 +725,7 @@ def gemm_mma_m16n8k16_row_col_fp16fp16fp32(
@pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0")
def test_gemm_mma_m16n8k16_row_col_fp16fp16fp32():
mod = tvm.IRModule.from_expr(gemm_mma_m16n8k16_row_col_fp16fp16fp32)
- cuda_mod = tvm.compile(mod, target="cuda")
+ cuda_mod = tvm.tirx.compile(mod, target="cuda")
A_np = np.random.uniform(-1, 1, [16, 16]).astype("float16")
B_np = np.random.uniform(-1, 1, [8, 16]).astype("float16")
@@ -801,7 +801,7 @@ def gemm_mma_m16n8k16_row_col_s8s8s32(
@pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0")
def test_gemm_mma_m16n8k16_row_col_s8s8s32():
mod = tvm.IRModule.from_expr(gemm_mma_m16n8k16_row_col_s8s8s32)
- cuda_mod = tvm.compile(mod, target="cuda")
+ cuda_mod = tvm.tirx.compile(mod, target="cuda")
A_np = np.random.uniform(-10, 10, [16, 16]).astype("int8")
B_np = np.random.uniform(-10, 10, [8, 16]).astype("int8")
@@ -877,7 +877,7 @@ def gemm_mma_m16n8k16_row_col_s8u8s32(
@pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0")
def test_gemm_mma_m16n8k16_row_col_s8u8s32():
mod = tvm.IRModule.from_expr(gemm_mma_m16n8k16_row_col_s8u8s32)
- cuda_mod = tvm.compile(mod, target="cuda")
+ cuda_mod = tvm.tirx.compile(mod, target="cuda")
A_np = np.random.uniform(-10, 10, [16, 16]).astype("int8")
B_np = np.random.uniform(-10, 10, [8, 16]).astype("uint8")
@@ -953,7 +953,7 @@ def gemm_mma_m16n8k32_row_col_s8s8s32(
@pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0")
def test_gemm_mma_m16n8k32_row_col_s8s8s32():
mod = tvm.IRModule.from_expr(gemm_mma_m16n8k32_row_col_s8s8s32)
- cuda_mod = tvm.compile(mod, target="cuda")
+ cuda_mod = tvm.tirx.compile(mod, target="cuda")
A_np = np.random.uniform(-10, 10, [16, 32]).astype("int8")
B_np = np.random.uniform(-10, 10, [8, 32]).astype("int8")
@@ -1029,7 +1029,7 @@ def gemm_mma_m16n8k32_row_col_s8u8s32(
@pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0")
def test_gemm_mma_m16n8k32_row_col_s8u8s32():
mod = tvm.IRModule.from_expr(gemm_mma_m16n8k32_row_col_s8u8s32)
- cuda_mod = tvm.compile(mod, target="cuda")
+ cuda_mod = tvm.tirx.compile(mod, target="cuda")
A_np = np.random.uniform(-10, 10, [16, 32]).astype("int8")
B_np = np.random.uniform(-10, 10, [8, 32]).astype("uint8")
@@ -1105,7 +1105,7 @@ def gemm_mma_m16n8k64_row_col_s4s4s32(
@pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0")
def test_gemm_mma_m16n8k64_row_col_s4s4s32():
mod = tvm.IRModule.from_expr(gemm_mma_m16n8k64_row_col_s4s4s32)
- cuda_mod = tvm.compile(mod, target="cuda")
+ cuda_mod = tvm.tirx.compile(mod, target="cuda")
def run_and_check():
ctx = tvm.cuda()
@@ -1175,7 +1175,7 @@ def gemm_mma_m16n8k64_row_col_s4u4s32(
@pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0")
def test_gemm_mma_m16n8k64_row_col_s4u4s32():
mod = tvm.IRModule.from_expr(gemm_mma_m16n8k64_row_col_s4u4s32)
- cuda_mod = tvm.compile(mod, target="cuda")
+ cuda_mod = tvm.tirx.compile(mod, target="cuda")
def run_and_check():
ctx = tvm.cuda()
@@ -1246,7 +1246,7 @@ def gemm_mma_m16n8k256_row_col_b1b1s32(
@pytest.mark.skipif(not env.has_cuda_compute(8), reason="need cuda compute >= 8.0")
def test_gemm_mma_m16n8k256_row_col_b1b1s32():
mod = tvm.IRModule.from_expr(gemm_mma_m16n8k256_row_col_b1b1s32)
- cuda_mod = tvm.compile(mod, target="cuda")
+ cuda_mod = tvm.tirx.compile(mod, target="cuda")
def run_and_check():
ctx = tvm.cuda()
diff --git a/tests/python/tirx-base/test_tir_ptx_mma_sp.py b/tests/python/tirx-base/test_tir_ptx_mma_sp.py
index 5635220847..191387eda5 100644
--- a/tests/python/tirx-base/test_tir_ptx_mma_sp.py
+++ b/tests/python/tirx-base/test_tir_ptx_mma_sp.py
@@ -263,7 +263,7 @@ def test_mma_sp_m16n8k16_f16():
for out_dtype in ["float16", "float32"]:
func = mma_sp_m16n8k16_f16f16f16 if out_dtype == "float16" else mma_sp_m16n8k16_f16f16f32
mod = tvm.IRModule.from_expr(func)
- cuda_mod = tvm.compile(mod, target="cuda")
+ cuda_mod = tvm.tirx.compile(mod, target="cuda")
A_np = np.random.uniform(-1, 1, [16, 8]).astype("float16")
B_np = np.random.uniform(-1, 1, [16, 8]).astype("float16")
@@ -305,7 +305,7 @@ def test_mma_sp_m16n8k32_f16():
for out_dtype in ["float16", "float32"]:
func = mma_sp_m16n8k32_f16f16f16 if out_dtype == "float16" else mma_sp_m16n8k32_f16f16f32
mod = tvm.IRModule.from_expr(func)
- cuda_mod = tvm.compile(mod, target="cuda")
+ cuda_mod = tvm.tirx.compile(mod, target="cuda")
A_np = np.random.uniform(-1, 1, [16, 16]).astype("float16")
B_np = np.random.uniform(-1, 1, [32, 8]).astype("float16")
diff --git a/tests/python/tirx-base/test_tir_ptx_scalar_f32_math.py b/tests/python/tirx-base/test_tir_ptx_scalar_f32_math.py
index 00b079df85..1d88309bf9 100644
--- a/tests/python/tirx-base/test_tir_ptx_scalar_f32_math.py
+++ b/tests/python/tirx-base/test_tir_ptx_scalar_f32_math.py
@@ -44,7 +44,7 @@ def ptx_scalar_f32_math(
@pytest.mark.skipif(not env.has_cuda_compute(7), reason="need cuda compute >= 7.0")
def test_ptx_scalar_f32_math():
f = ptx_scalar_f32_math
- mod = tvm.compile(f, target="cuda")
+ mod = tvm.tirx.compile(f, target="cuda")
rng = np.random.default_rng(0)
A_np = rng.standard_normal(32).astype("float32")
B_np = rng.standard_normal(32).astype("float32")
diff --git a/tests/python/tirx-transform/test_tir_transform_flatten_tensor.py b/tests/python/tirx-transform/test_tir_transform_flatten_tensor.py
index 8d75dce878..4d02a34534 100644
--- a/tests/python/tirx-transform/test_tir_transform_flatten_tensor.py
+++ b/tests/python/tirx-transform/test_tir_transform_flatten_tensor.py
@@ -429,7 +429,7 @@ def test_build_with_optional_unroll_explicit():
for ax1, ax2 in T.grid(5, 6):
B[ax0, ax1, ax2] = A[ax0, ax1, ax2]
... 2038 lines suppressed ...