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/13 12:08:05 UTC
(tvm) branch main updated: [REFACTOR][Python] Place pass and testing helpers with their owning modules (#20692)
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 c89ebf33e2 [REFACTOR][Python] Place pass and testing helpers with their owning modules (#20692)
c89ebf33e2 is described below
commit c89ebf33e27103ed2765bf5667c8614a3918a6e1
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Tue Oct 13 08:07:59 2026 -0400
[REFACTOR][Python] Place pass and testing helpers with their owning modules (#20692)
Move Python pass infrastructure under tvm.ir and domain testing helpers
into Relax and schedulable TensorIR. Reduce core testing dependencies by
replacing thin wrappers with direct calls and removing unused fixture,
target-enumeration, and request-hook machinery.
---
docs/arch/index.rst | 2 +-
docs/arch/pass_infra.rst | 32 +-
docs/conf.py | 8 +-
docs/contribute/code_guide.rst | 6 +-
docs/contribute/testing.rst | 21 -
.../relax/tutorials/relax_transformation.py | 6 +-
.../deep_dive/tensor_ir/tutorials/meta_schedule.py | 4 +-
docs/how_to/tutorials/bring_your_own_codegen.py | 8 +-
docs/how_to/tutorials/customize_opt.py | 8 +-
docs/how_to/tutorials/e2e_opt_model.py | 2 +-
docs/how_to/tutorials/optimize_llm.py | 6 +-
docs/reference/api/python/instrument.rst | 6 +-
docs/reference/api/python/transform.rst | 6 +-
include/tvm/ir/type.h | 2 +-
python/tvm/__init__.py | 3 +-
python/tvm/backend/cuda/transforms/__init__.py | 2 +-
python/tvm/backend/trn/pipeline.py | 8 +-
python/tvm/backend/trn/transform/__init__.py | 2 +-
python/tvm/contrib/hexagon/_ci_env_check.py | 55 --
python/tvm/contrib/hexagon/generate_take_op.py | 2 +-
python/tvm/driver/build_module.py | 8 +-
python/tvm/exec/gpu_memory_bandwidth.py | 2 +-
python/tvm/ir/__init__.py | 4 +-
python/tvm/{ => ir}/transform/__init__.py | 0
.../tvm/{ => ir}/transform/_ffi_instrument_api.py | 2 +-
.../tvm/{ => ir}/transform/_ffi_transform_api.py | 2 +-
python/tvm/{ => ir}/transform/core.py | 14 +-
python/tvm/{ => ir}/transform/instrument.py | 2 +-
python/tvm/ir/type_relation.py | 76 ---
python/tvm/relax/backend/adreno/clml.py | 12 +-
python/tvm/relax/backend/adreno/pipeline.py | 6 +-
.../relax/backend/adreno/transform/transform.py | 8 +-
python/tvm/relax/backend/contrib/tensorrt.py | 2 +-
python/tvm/relax/backend/cpu_generic/pipeline.py | 6 +-
python/tvm/relax/backend/cuda/cudnn.py | 4 +-
python/tvm/relax/backend/cuda/cutlass.py | 4 +-
python/tvm/relax/backend/cuda/pipeline.py | 6 +-
python/tvm/relax/backend/dispatch_sampling.py | 2 +-
python/tvm/relax/backend/dispatch_sort_scan.py | 2 +-
python/tvm/relax/backend/gpu_generic/pipeline.py | 6 +-
python/tvm/relax/backend/rocm/pipeline.py | 6 +-
.../tvm/relax/distributed/transform/transform.py | 16 +-
python/tvm/relax/expr_functor.py | 2 +-
python/tvm/relax/frontend/nn/core.py | 2 +-
python/tvm/relax/frontend/torch/dynamo.py | 4 +-
python/tvm/relax/ir/instrument.py | 2 +-
python/tvm/relax/pipeline.py | 26 +-
python/tvm/relax/testing/gradient.py | 168 ++++++
python/tvm/relax/testing/transform.py | 2 +-
python/tvm/relax/training/setup_trainer.py | 4 +-
python/tvm/relax/training/utils.py | 2 +-
python/tvm/relax/transform/_ffi_api.py | 2 +-
.../tvm/relax/transform/attach_external_modules.py | 2 +-
python/tvm/relax/transform/fast_math.py | 4 +-
.../fold_batch_norm_to_conv2d_for_inference.py | 2 +-
.../tvm/relax/transform/fuse_transpose_matmul.py | 4 +-
.../tvm/relax/transform/ipc_allreduce_rewrite.py | 4 +-
.../tvm/relax/transform/lazy_transform_params.py | 4 +-
.../relax/transform/lower_gpu_ipc_alloc_storage.py | 4 +-
.../relax/transform/optimize_layout_transform.py | 2 +-
.../relax/transform/remove_redundant_reshape.py | 2 +-
python/tvm/relax/transform/transform.py | 234 ++++----
python/tvm/relax/utils.py | 24 +-
python/tvm/relax/vm_build.py | 6 +-
python/tvm/s_tir/analysis/__init__.py | 6 +-
.../s_tir/backend/adreno/transform/transform.py | 4 +-
python/tvm/s_tir/dlight/base/transform.py | 2 +-
python/tvm/s_tir/dlight/benchmark/bench.py | 2 +-
.../s_tir/meta_schedule/builder/local_builder.py | 2 +-
.../s_tir/meta_schedule/cost_model/mlp_model.py | 2 +-
.../s_tir/meta_schedule/cost_model/random_model.py | 2 +-
.../s_tir/meta_schedule/cost_model/xgb_model.py | 2 +-
.../feature_extractor/random_feature_extractor.py | 2 +-
.../tvm/s_tir/meta_schedule/relax_integration.py | 10 +-
.../tvm/s_tir/meta_schedule/runner/local_runner.py | 2 +-
.../tvm/s_tir/meta_schedule/runner/rpc_runner.py | 2 +-
.../s_tir/meta_schedule/testing/dummy_object.py | 2 +-
python/tvm/s_tir/meta_schedule/testing/tune_te.py | 2 +-
.../meta_schedule/testing/validate_database.py | 2 +-
python/tvm/s_tir/pipeline.py | 16 +-
.../testing/__init__.py} | 6 +-
python/tvm/{ => s_tir}/testing/runner.py | 0
python/tvm/{ => s_tir}/testing/tir.py | 0
python/tvm/s_tir/testing/utils.py | 43 ++
python/tvm/s_tir/transform/transform.py | 58 +-
python/tvm/{ir => s_tir}/utils.py | 2 +-
python/tvm/support/popen_pool.py | 2 +-
python/tvm/{exec => support}/popen_worker.py | 0
python/tvm/testing/__init__.py | 1 -
python/tvm/testing/env.py | 66 ++-
python/tvm/testing/utils.py | 632 +--------------------
python/tvm/tirx/_compile.py | 6 +-
python/tvm/tirx/compilation_pipeline.py | 20 +-
python/tvm/tirx/transform/function_pass.py | 2 +-
python/tvm/tirx/transform/transform.py | 64 +--
tests/nightly/python/test_nnapi/test_network.py | 2 +-
tests/nightly/python/test_nnapi/test_ops.py | 2 +-
tests/python/ci/test_ci.py | 6 +-
tests/python/ci/test_tvmbot.py | 8 +-
tests/python/codegen/test_target_codegen_cuda.py | 2 +-
.../codegen/test_target_codegen_cuda_fastmath.py | 4 +-
.../python/codegen/test_target_codegen_cuda_fp8.py | 76 +--
tests/python/codegen/test_target_codegen_llvm.py | 4 +-
tests/python/conftest.py | 29 -
tests/python/contrib/test_example_npu.py | 2 +-
tests/python/contrib/test_random.py | 37 +-
tests/python/disco/test_callback.py | 2 +-
tests/python/driver/test_compile.py | 5 +-
tests/python/ir/test_node_reflection.py | 8 +-
tests/python/ir/test_pass_instrument.py | 6 +-
tests/python/relax/conftest.py | 4 +-
tests/python/relax/test_codegen_cublas.py | 2 +-
tests/python/relax/test_codegen_cudnn.py | 2 +-
tests/python/relax/test_codegen_cutlass.py | 2 +-
tests/python/relax/test_codegen_dnnl.py | 4 +-
tests/python/relax/test_codegen_hipblas.py | 2 +-
tests/python/relax/test_codegen_tensorrt.py | 8 +-
tests/python/relax/test_contrib_vllm.py | 2 +-
tests/python/relax/test_dataflow_rewriter.py | 4 +-
tests/python/relax/test_frontend_nn_jit.py | 2 +-
tests/python/relax/test_frontend_onnx.py | 10 +-
tests/python/relax/test_frontend_onnx_backend.py | 2 +-
tests/python/relax/test_op_gradient_numeric.py | 2 +-
.../python/relax/test_optimize_layout_transform.py | 2 +-
.../relax/test_tensorrt_engine_serialization.py | 6 +-
tests/python/relax/test_transform_codegen_pass.py | 8 +-
tests/python/relax/test_transform_cse.py | 2 +-
.../relax/test_transform_error_enrichment.py | 4 +-
.../relax/test_transform_fuse_ops_by_pattern.py | 2 +-
.../relax/test_transform_fuse_transpose_matmul.py | 4 +-
.../relax/test_transform_gradient_numeric.py | 2 +-
.../test_transform_legalize_ops_manipulate.py | 2 +-
.../relax/test_transform_lift_transform_params.py | 2 +-
.../relax/test_transform_meta_schedule_tuning.py | 5 +-
.../relax/test_transform_rewrite_cuda_graph.py | 8 +-
.../test_transform_rewrite_dataflow_reshape.py | 6 +-
.../relax/test_vm_alloc_storage_with_scope.py | 2 +-
tests/python/relax/test_vm_build.py | 2 +-
tests/python/relax/test_vm_cuda_graph.py | 4 +-
tests/python/relax/texture/test_texture_nd.py | 4 +-
tests/python/request_hook.py | 56 --
tests/python/runtime/test_runtime_rpc.py | 3 +-
.../python/s_tir/base/test_compilation_pipeline.py | 4 +-
.../meta_schedule/test_meta_schedule_cost_model.py | 2 +-
.../meta_schedule/test_meta_schedule_database.py | 2 +-
.../test_meta_schedule_feature_extractor.py | 2 +-
.../test_meta_schedule_measure_callback.py | 2 +-
.../test_meta_schedule_mma_tensorize.py | 2 +-
.../test_meta_schedule_post_order_apply.py | 2 +-
.../meta_schedule/test_meta_schedule_runner.py | 5 +-
.../test_meta_schedule_search_strategy.py | 2 +-
.../test_meta_schedule_space_generator.py | 2 +-
.../test_meta_schedule_task_scheduler.py | 2 +-
.../meta_schedule/test_meta_schedule_tune_tir.py | 2 +-
..._tir_schedule_tensorize_ldmatrix_mma_numeric.py | 2 +-
.../test_tir_schedule_tensorize_mfma_numeric.py | 2 +-
.../test_s_tir_transform_compact_tensor_region.py | 4 +-
.../test_s_tir_transform_default_gpu_schedule.py | 20 +-
.../test_s_tir_transform_inject_ptx_async_copy.py | 13 +-
...est_s_tir_transform_inject_software_pipeline.py | 11 +-
.../test_s_tir_transform_loop_partition.py | 22 +-
.../transform/test_s_tir_transform_thread_sync.py | 2 +-
tests/python/sym/test_sym_deduce_bound.py | 56 +-
tests/python/sym/test_sym_detect_clip_bound.py | 26 +-
.../python/sym/test_sym_detect_linear_equation.py | 12 +-
tests/python/sym/test_sym_iter_affine_map.py | 16 +-
tests/python/testing/test_env.py | 145 +++++
tests/python/tirx-base/test_tir_base.py | 2 +-
tests/python/tirx-base/test_tir_intrin.py | 2 +-
.../test_tir_transform_flatten_tensor.py | 2 +-
.../test_tir_transform_function_pass.py | 4 +-
.../test_tir_transform_lower_intrin.py | 2 +-
.../test_tir_transform_remove_no_op.py | 2 +-
.../tirx-transform/test_tir_transform_simplify.py | 2 +-
.../test_tir_transform_unroll_loop.py | 12 +-
.../tirx-transform/test_tir_transform_vectorize.py | 10 +-
tests/python/tirx/iket/test_iket_profiler.py | 4 +-
tests/python/tirx/test_cuda_backend_config.py | 8 +-
tests/python/tirx/tile/cuda/copy/test_reg.py | 4 +-
179 files changed, 1093 insertions(+), 1531 deletions(-)
diff --git a/docs/arch/index.rst b/docs/arch/index.rst
index c1a05b5a34..77d9538465 100644
--- a/docs/arch/index.rst
+++ b/docs/arch/index.rst
@@ -320,7 +320,7 @@ The following code snippet gives an example of PassContext configuration.
.. code-block:: python
# configure the behavior of the tirx.UnrollLoop pass
- with tvm.transform.PassContext(config={"tirx.UnrollLoop": { "auto_max_step": 10 }}):
+ with tvm.ir.transform.PassContext(config={"tirx.UnrollLoop": { "auto_max_step": 10 }}):
# code affected by the pass context
diff --git a/docs/arch/pass_infra.rst b/docs/arch/pass_infra.rst
index 3171f96609..a0b4b28388 100644
--- a/docs/arch/pass_infra.rst
+++ b/docs/arch/pass_infra.rst
@@ -228,7 +228,7 @@ creating a sequence. For example:
.. code:: python
- pipeline = tvm.transform.Sequential([
+ pipeline = tvm.ir.transform.Sequential([
relax.transform.Normalize(),
relax.transform.FoldConstant(),
])
@@ -329,7 +329,7 @@ of the ``PassContext`` instance.
``InstrumentExitPassContext`` is called when leaving the scope of ``PassContext``,
or exceptions occur during the execution of passes.
-This method is also called when instruments is being overridden by ``override_instruments`` in :py:class:`tvm.transform.PassContext`.
+This method is also called when instruments is being overridden by ``override_instruments`` in :py:class:`tvm.ir.transform.PassContext`.
See :ref:`pass_instrument_overriden`.
``InstrumentBeforePass`` is called before execution.
@@ -438,19 +438,19 @@ There are several built-in instruments.
* Profile the execution time of passes.
-- PrintBeforeAll (see `python/tvm/transform/instrument.py`_)
+- PrintBeforeAll (see `python/tvm/ir/transform/instrument.py`_)
* Print the IR module and pass info before each pass executes.
-- PrintAfterAll (see `python/tvm/transform/instrument.py`_)
+- PrintAfterAll (see `python/tvm/ir/transform/instrument.py`_)
* Print the IR module and pass info after each pass executes.
-- PassPrintingInstrument (see `python/tvm/transform/instrument.py`_)
+- PassPrintingInstrument (see `python/tvm/ir/transform/instrument.py`_)
* Selectively print the IR module before or after specific named passes.
-- DumpIR (see `python/tvm/transform/instrument.py`_)
+- DumpIR (see `python/tvm/ir/transform/instrument.py`_)
* Dump the IR module to files after each pass executes.
@@ -460,7 +460,7 @@ Python Frontend
Only some simple APIs are needed for the frontend side. For example, we can
provide users the following APIs to create and execute a pass (full
implementation is provided in `python/tvm/relax/transform/transform.py`_ and
-`python/tvm/transform/core.py`_). The backend
+`python/tvm/ir/transform/core.py`_). The backend
receives the information and decides which function it should use to create
a Pass object.
@@ -519,8 +519,8 @@ Pass Instrument
^^^^^^^^^^^^^^^
One can implement a ``PassInstrument`` by using the ``pass_instrument``
-decorator (`python/tvm/transform/core.py`_) or by subclassing
-``tvm.transform.PassInstrument``. Implement any of the following callbacks:
+decorator (`python/tvm/ir/transform/core.py`_) or by subclassing
+``tvm.ir.transform.PassInstrument``. Implement any of the following callbacks:
- ``enter_pass_ctx``
@@ -544,11 +544,11 @@ decorator (`python/tvm/transform/core.py`_) or by subclassing
* This method is run right after a pass has been executed.
``PassInstrument`` instances can be registered through ``instruments`` argument in
-:py:class:`tvm.transform.PassContext`.
+:py:class:`tvm.ir.transform.PassContext`.
-Core interfaces are exported from ``tvm.transform``. Concrete tools, including
-``PassTimingInstrument`` and ``DumpIR``, live in ``tvm.transform.instrument``;
-see `python/tvm/transform/instrument.py`_ for examples.
+Core interfaces are exported from ``tvm.ir.transform``. Concrete tools, including
+``PassTimingInstrument`` and ``DumpIR``, live in ``tvm.ir.transform.instrument``;
+see `python/tvm/ir/transform/instrument.py`_ for examples.
.. _pass_instrument_overriden:
@@ -561,7 +561,7 @@ one can still register ``PassInstrument`` into the global ``PassContext`` by:
.. code:: python
- cur_pass_ctx = tvm.transform.PassContext.current()
+ cur_pass_ctx = tvm.ir.transform.PassContext.current()
# override PassInstrument instances
cur_pass_ctx.override_instruments([pass_inst])
mod = pass_seq(mod)
@@ -587,9 +587,9 @@ new ``PassInstrument`` are called.
.. _include/tvm/relax/transform.h: https://github.com/apache/tvm/blob/main/include/tvm/relax/transform.h
-.. _python/tvm/transform/core.py: https://github.com/apache/tvm/blob/main/python/tvm/transform/core.py
+.. _python/tvm/ir/transform/core.py: https://github.com/apache/tvm/blob/main/python/tvm/ir/transform/core.py
-.. _python/tvm/transform/instrument.py: https://github.com/apache/tvm/blob/main/python/tvm/transform/instrument.py
+.. _python/tvm/ir/transform/instrument.py: https://github.com/apache/tvm/blob/main/python/tvm/ir/transform/instrument.py
.. _src/tirx/transform/unroll_loop.cc: https://github.com/apache/tvm/blob/main/src/tirx/transform/unroll_loop.cc
diff --git a/docs/conf.py b/docs/conf.py
index c226de1c40..599e03ca38 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -66,7 +66,7 @@ os.environ["TVM_BUILD_DOC"] = "1"
# Version information.
import tvm
-from tvm import te, testing, topi
+from tvm import te, topi
# The version is derived from the Git tag by setuptools_scm at build time and exposed
# as tvm.__version__ (see [tool.setuptools_scm] in pyproject.toml).
@@ -214,10 +214,6 @@ def rst2md(text, gallery_conf, target_dir, heading_levels, real_func):
return text
-def install_request_hook(gallery_conf, fname):
- testing.utils.install_request_hook(tvm_path.resolve() / "tests" / "python" / "request_hook.py")
-
-
INSTALL_TVM_DEV = """\
%%shell
# Installs the latest dev build of TVM from PyPI. If you wish to build
@@ -540,7 +536,7 @@ sphinx_gallery_conf = {
"download_all_examples": False,
"min_reported_time": 60,
"expected_failing_examples": [],
- "reset_modules": ("matplotlib", "seaborn", force_gc, install_request_hook),
+ "reset_modules": ("matplotlib", "seaborn", force_gc),
"promote_jupyter_magic": True,
# Drop the "Gallery generated by Sphinx-Gallery" signature line on generated pages.
"show_signature": False,
diff --git a/docs/contribute/code_guide.rst b/docs/contribute/code_guide.rst
index eb4ba5420d..070db8c4f2 100644
--- a/docs/contribute/code_guide.rst
+++ b/docs/contribute/code_guide.rst
@@ -162,10 +162,8 @@ In CI, downloading files from the Internet is a big source of flaky test failure
server can go down or be slow), so try to avoid using the network at all during tests. In some cases
this isn't a reasonable proposition (e.g. the docs tutorials which need to download models).
-New network downloads are rejected by the CI `request hook
-<https://github.com/apache/tvm/blob/main/tests/python/request_hook.py>`_. Prefer
-checked-in fixtures or generated data. If a download is unavoidable, arrange a stable
-project-managed mirror with the maintainers and add an explicit ``URL_MAP`` entry.
+Prefer checked-in fixtures or generated data. If a download is unavoidable, arrange a
+stable project-managed mirror with the maintainers.
Handle Integer Constant Expression
diff --git a/docs/contribute/testing.rst b/docs/contribute/testing.rst
index 8dd71ca4ca..44a996da01 100644
--- a/docs/contribute/testing.rst
+++ b/docs/contribute/testing.rst
@@ -185,27 +185,6 @@ persistent lock file that every contender can write, or enforce exclusivity
through the runner. A per-user lock path cannot protect a GPU shared across
users because each user would lock a different file.
-There also exists a ``tvm.testing.enabled_targets()`` that returns
-all targets that are enabled and runnable on the current machine,
-based on the environment variable ``TVM_TEST_TARGETS``, the build
-configuration, and the physical hardware present. Some legacy tests
-explicitly loop over the targets returned from ``enabled_targets()``,
-but this style should not be used for new tests. The pytest output
-for this style silently skips runtimes that are disabled in
-``config.cmake``, or do not have a device on which they can run. In
-addition, the test halts on the first target to fail, which is
-ambiguous as to whether the error occurs on a particular target, or on
-every target.
-
-.. code-block:: python
-
- # Old style, do not use.
- def test_function():
- for target, dev in tvm.testing.enabled_targets():
- # Test code goes here
-
-
-
Running Locally
~~~~~~~~~~~~~~~
diff --git a/docs/deep_dive/relax/tutorials/relax_transformation.py b/docs/deep_dive/relax/tutorials/relax_transformation.py
index b119652100..8abaa2d75e 100644
--- a/docs/deep_dive/relax/tutorials/relax_transformation.py
+++ b/docs/deep_dive/relax/tutorials/relax_transformation.py
@@ -73,7 +73,7 @@ mod.show()
# in ML compilers. Note that in relax, fusion optimizations are done with the collaboration of
# a set of passes. We can apply them in a sequence.
-mod = tvm.transform.Sequential(
+mod = tvm.ir.transform.Sequential(
[
tvm.relax.transform.AnnotateTIROpPattern(),
tvm.relax.transform.FuseOps(),
@@ -115,9 +115,9 @@ class ReluRewriter(PyExprMutator):
# Then we can write a pass to apply the mutator to the whole module.
-@tvm.transform.module_pass(opt_level=0, name="ReluToGelu")
+@tvm.ir.transform.module_pass(opt_level=0, name="ReluToGelu")
class ReluToGelu: # pylint: disable=too-few-public-methods
- def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+ def transform_module(self, mod: IRModule, _ctx: tvm.ir.transform.PassContext) -> IRModule:
"""IRModule-level transformation"""
rewriter = ReluRewriter(mod)
for g_var, func in mod.functions_items():
diff --git a/docs/deep_dive/tensor_ir/tutorials/meta_schedule.py b/docs/deep_dive/tensor_ir/tutorials/meta_schedule.py
index b27b0d6b51..584a05bcf2 100644
--- a/docs/deep_dive/tensor_ir/tutorials/meta_schedule.py
+++ b/docs/deep_dive/tensor_ir/tutorials/meta_schedule.py
@@ -129,7 +129,7 @@ target = tvm.target.Target.from_device(device)
if os.getenv("CI", "") != "true":
with target, tempfile.TemporaryDirectory() as tmp_dir:
- tuned_mod = tvm.transform.Sequential(
+ tuned_mod = tvm.ir.transform.Sequential(
[
relax.get_pipeline("zero"),
relax.transform.MetaScheduleTuneTIR(
@@ -173,7 +173,7 @@ for i, task in enumerate(tasks):
# .. code-block:: python
#
# with target:
-# mod = tvm.transform.Sequential([
+# mod = tvm.ir.transform.Sequential([
# relax.transform.MetaScheduleTuneIRMod(
# params={},
# work_dir="./tuning_logs",
diff --git a/docs/how_to/tutorials/bring_your_own_codegen.py b/docs/how_to/tutorials/bring_your_own_codegen.py
index 0587a040b1..96d94b3806 100644
--- a/docs/how_to/tutorials/bring_your_own_codegen.py
+++ b/docs/how_to/tutorials/bring_your_own_codegen.py
@@ -166,7 +166,7 @@ weight_np = np.random.randn(16, 3, 3, 3).astype("float32")
if has_example_npu:
npu_mod = RunCodegen()(npu_mod)
- with tvm.transform.PassContext(opt_level=3):
+ with tvm.ir.transform.PassContext(opt_level=3):
npu_exec = relax.build(npu_mod, tvm.target.Target("llvm"))
npu_vm = relax.VirtualMachine(npu_exec, tvm.cpu())
@@ -207,7 +207,7 @@ print(trt_mod)
if has_tensorrt and has_cuda:
dev = tvm.cuda(0)
- with tvm.transform.PassContext(opt_level=3):
+ with tvm.ir.transform.PassContext(opt_level=3):
trt_exec = relax.build(RunCodegen()(trt_mod), "cuda")
trt_out = relax.VirtualMachine(trt_exec, dev)["main"](tvm.runtime.tensor(data_np, dev)).numpy()
@@ -235,7 +235,7 @@ if has_tensorrt and has_cuda:
fp16_mod = partition_for_tensorrt(
relax.transform.BindParams("main", {"weight": weight_np})(ConvReLU)
)
- with tvm.transform.PassContext(
+ with tvm.ir.transform.PassContext(
opt_level=3, config={"relax.ext.tensorrt.options": {"use_fp16": True}}
):
fp16_exec = relax.build(RunCodegen()(fp16_mod), "cuda")
@@ -313,7 +313,7 @@ if has_torch and has_tensorrt and has_cuda:
print(torch_mod)
torch_dev = tvm.cuda(0)
- with tvm.transform.PassContext(opt_level=3):
+ with tvm.ir.transform.PassContext(opt_level=3):
torch_exec = relax.build(RunCodegen()(torch_mod), "cuda")
deployed = relax.VirtualMachine(torch_exec, torch_dev)["main"](
tvm.runtime.tensor(example_input.numpy(), torch_dev)
diff --git a/docs/how_to/tutorials/customize_opt.py b/docs/how_to/tutorials/customize_opt.py
index 01c1f4fe8e..e6e83c33c4 100644
--- a/docs/how_to/tutorials/customize_opt.py
+++ b/docs/how_to/tutorials/customize_opt.py
@@ -119,9 +119,9 @@ except ImportError as e:
# Define a new pass for CUBLAS dispatch
-@tvm.transform.module_pass(opt_level=0, name="CublasDispatch")
+@tvm.ir.transform.module_pass(opt_level=0, name="CublasDispatch")
class CublasDispatch:
- def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+ def transform_module(self, mod: IRModule, _ctx: tvm.ir.transform.PassContext) -> IRModule:
# Check if CUBLAS is enabled
if not tvm.get_global_func("relax.ext.cublas", True):
raise Exception("CUBLAS is not enabled.")
@@ -168,7 +168,7 @@ target = tvm.target.Target.from_device(device)
if os.getenv("CI", "") != "true":
trials = 2000
with target, tempfile.TemporaryDirectory() as tmp_dir:
- mod = tvm.transform.Sequential(
+ mod = tvm.ir.transform.Sequential(
[
relax.get_pipeline("zero"),
relax.transform.MetaScheduleTuneTIR(work_dir=tmp_dir, max_trials_global=trials),
@@ -190,7 +190,7 @@ from tvm.s_tir import dlight as dl
# Apply DLight rules
with target:
- mod = tvm.transform.Sequential(
+ mod = tvm.ir.transform.Sequential(
[
relax.get_pipeline("zero"),
dl.ApplyDefaultSchedule( # pylint: disable=not-callable
diff --git a/docs/how_to/tutorials/e2e_opt_model.py b/docs/how_to/tutorials/e2e_opt_model.py
index 5c12ad8e0e..5de3d50485 100644
--- a/docs/how_to/tutorials/e2e_opt_model.py
+++ b/docs/how_to/tutorials/e2e_opt_model.py
@@ -91,7 +91,7 @@ if not IS_IN_CI:
# ---------------------
# Apache TVM provides a flexible way to optimize the IRModule. Everything centered
# around IRModule optimization can be composed with existing pipelines. Note that each
-# transformation can be combined as an optimization pipeline via ``tvm.transform.Sequential``.
+# transformation can be combined as an optimization pipeline via ``tvm.ir.transform.Sequential``.
#
# In this tutorial, we focus on the end-to-end optimization of the model via auto-tuning. We
# leverage MetaSchedule to tune the model and store the tuning logs to the database. We also
diff --git a/docs/how_to/tutorials/optimize_llm.py b/docs/how_to/tutorials/optimize_llm.py
index 9e83a0c12a..7faa4f43bc 100644
--- a/docs/how_to/tutorials/optimize_llm.py
+++ b/docs/how_to/tutorials/optimize_llm.py
@@ -382,9 +382,9 @@ def _pipeline( # pylint: disable=too-many-arguments
):
ext_mods = ext_mods or []
- @tvm.transform.module_pass(opt_level=0)
- def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.IRModule:
- seq = tvm.transform.Sequential(
+ @tvm.ir.transform.module_pass(opt_level=0)
+ def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.ir.transform.PassContext) -> tvm.ir.IRModule:
+ seq = tvm.ir.transform.Sequential(
[
# Phase 1. Passes on high-level operator graph
# We can enable cublas for further optimization
diff --git a/docs/reference/api/python/instrument.rst b/docs/reference/api/python/instrument.rst
index 30a83fbb6e..46d193ce15 100644
--- a/docs/reference/api/python/instrument.rst
+++ b/docs/reference/api/python/instrument.rst
@@ -15,7 +15,7 @@
specific language governing permissions and limitations
under the License.
-tvm.transform.instrument
-========================
-.. automodule:: tvm.transform.instrument
+tvm.ir.transform.instrument
+===========================
+.. automodule:: tvm.ir.transform.instrument
:members:
diff --git a/docs/reference/api/python/transform.rst b/docs/reference/api/python/transform.rst
index e634cef4b4..aebdbdbad5 100644
--- a/docs/reference/api/python/transform.rst
+++ b/docs/reference/api/python/transform.rst
@@ -15,8 +15,8 @@
specific language governing permissions and limitations
under the License.
-tvm.transform
-=============
-.. automodule:: tvm.transform
+tvm.ir.transform
+================
+.. automodule:: tvm.ir.transform
:members:
:imported-members:
diff --git a/include/tvm/ir/type.h b/include/tvm/ir/type.h
index 19c9ecdab5..f74b718ea3 100644
--- a/include/tvm/ir/type.h
+++ b/include/tvm/ir/type.h
@@ -199,7 +199,7 @@ inline bool IsVoidType(const Type& type) {
* We support polymorphic function type.
* This can be roughly viewed as template function in C++.
*
- * \sa FuncType, TypeConstraint
+ * \sa FuncType
*/
class FuncTypeNode : public TypeNode {
public:
diff --git a/python/tvm/__init__.py b/python/tvm/__init__.py
index aec2212adf..13a6744480 100644
--- a/python/tvm/__init__.py
+++ b/python/tvm/__init__.py
@@ -39,8 +39,7 @@ from .runtime import DataType, DataTypeCode
from . import error
# tvm.ir
-from .ir import IRModule
-from .ir import transform
+from .ir import IRModule, transform
from . import ir
# tvm.script — must be imported before any dialect package so that
diff --git a/python/tvm/backend/cuda/transforms/__init__.py b/python/tvm/backend/cuda/transforms/__init__.py
index fd37f294cb..3e36291c4e 100644
--- a/python/tvm/backend/cuda/transforms/__init__.py
+++ b/python/tvm/backend/cuda/transforms/__init__.py
@@ -32,7 +32,7 @@ def LowerIket():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass.
"""
return get_global_func("tirx.backend.cuda.transforms.LowerIket")()
diff --git a/python/tvm/backend/trn/pipeline.py b/python/tvm/backend/trn/pipeline.py
index 0b8daf57ca..46d7a1fe22 100644
--- a/python/tvm/backend/trn/pipeline.py
+++ b/python/tvm/backend/trn/pipeline.py
@@ -26,8 +26,8 @@ from . import transform as trn_transform
def trn_pipeline(*, target=None):
"""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:
+ @tvm.ir.transform.module_pass(opt_level=0)
+ def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.ir.transform.PassContext) -> tvm.ir.IRModule:
"""Lower TIRx for the Trainium backend."""
mod = _bind_target(mod, target)
passes = [
@@ -44,14 +44,14 @@ def trn_pipeline(*, target=None):
tirx.transform.SplitHostDevice(),
tirx.transform.MakePackedAPI(),
]
- return tvm.transform.Sequential(passes)(mod)
+ return tvm.ir.transform.Sequential(passes)(mod)
return _pipeline, finalize_host_passes, finalize_device_passes_trn
def finalize_device_passes_trn(): # pylint: disable=unused-argument
"""The finalization passes for the Trainium backend."""
- return tvm.transform.Sequential([tirx.transform.StmtSimplify()])
+ return tvm.ir.transform.Sequential([tirx.transform.StmtSimplify()])
__all__ = ["finalize_device_passes_trn", "trn_pipeline"]
diff --git a/python/tvm/backend/trn/transform/__init__.py b/python/tvm/backend/trn/transform/__init__.py
index d612a0858d..5be506bb50 100644
--- a/python/tvm/backend/trn/transform/__init__.py
+++ b/python/tvm/backend/trn/transform/__init__.py
@@ -35,7 +35,7 @@ def LowerTrainiumLayout():
def LowerTIRx():
"""Lower TIRx tile primitive calls for the Trainium backend."""
- return tvm.transform.Sequential(
+ return tvm.ir.transform.Sequential(
[tirx.transform.TileDispatch(), LowerTrainiumLayout()],
name="tirx.backend.trn.LowerTIRx",
)
diff --git a/python/tvm/contrib/hexagon/_ci_env_check.py b/python/tvm/contrib/hexagon/_ci_env_check.py
deleted file mode 100644
index ed7beb246a..0000000000
--- a/python/tvm/contrib/hexagon/_ci_env_check.py
+++ /dev/null
@@ -1,55 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-"""Hexagon environment checks for CI usage
-
-These are required by tvm.testing and live here to avoid a circular dependency.
-"""
-
-import os
-
-import tvm
-
-ANDROID_SERIAL_NUMBER = "ANDROID_SERIAL_NUMBER"
-HEXAGON_TOOLCHAIN = "HEXAGON_TOOLCHAIN"
-
-
-def _compile_time_check():
- """Return True if compile-time support for Hexagon is present, otherwise
- error string.
-
- Backs :func:`tvm.testing.env.has_hexagon_toolchain`.
- """
- if tvm.runtime.enabled("llvm") and tvm.target.codegen.llvm_version_major() < 7:
- return "Hexagon requires LLVM 7 or later"
-
- if "HEXAGON_TOOLCHAIN" not in os.environ:
- return f"Missing environment variable {HEXAGON_TOOLCHAIN}."
-
- return True
-
-
-def _run_time_check():
- """Return True if run-time support for Hexagon is present, otherwise
- error string.
-
- Backs :func:`tvm.testing.env.has_hexagon`.
- """
- if ANDROID_SERIAL_NUMBER not in os.environ:
- return f"Missing environment variable {ANDROID_SERIAL_NUMBER}."
-
- return True
diff --git a/python/tvm/contrib/hexagon/generate_take_op.py b/python/tvm/contrib/hexagon/generate_take_op.py
index 20e1ff1aba..c77f08c81f 100644
--- a/python/tvm/contrib/hexagon/generate_take_op.py
+++ b/python/tvm/contrib/hexagon/generate_take_op.py
@@ -90,7 +90,7 @@ class Tanh2TakeReplace(tvm.relax.PyExprMutator):
return call_node
-@tvm.transform.module_pass(opt_level=2, name="replace_tanh_take")
+@tvm.ir.transform.module_pass(opt_level=2, name="replace_tanh_take")
class PassReplaceWithTakeOpFunctions:
def transform_module(self, mod, ctx):
return Tanh2TakeReplace(mod).transform()
diff --git a/python/tvm/driver/build_module.py b/python/tvm/driver/build_module.py
index 3e30b1a2f6..15cc37c4e9 100644
--- a/python/tvm/driver/build_module.py
+++ b/python/tvm/driver/build_module.py
@@ -40,8 +40,8 @@ def compile( # pylint: disable=redefined-builtin
mod: Function | IRModule,
target: Target | None = None,
*,
- relax_pipeline: tvm.transform.Pass | Callable | str | None = "default",
- tir_pipeline: tvm.transform.Pass | Callable | str | None = "default",
+ relax_pipeline: tvm.ir.transform.Pass | Callable | str | None = "default",
+ tir_pipeline: tvm.ir.transform.Pass | Callable | str | None = "default",
) -> Executable:
"""
Compile an IRModule to a runtime executable.
@@ -56,10 +56,10 @@ def compile( # pylint: disable=redefined-builtin
TIR or Relax functions.
target : Optional[Target]
The target platform and its backend compilation configuration.
- relax_pipeline : Optional[Union[tvm.transform.Pass, Callable, str]]
+ relax_pipeline : Optional[Union[tvm.ir.transform.Pass, Callable, str]]
The compilation pipeline to use for Relax functions.
Only used if the module contains Relax functions.
- tir_pipeline : Optional[Union[tvm.transform.Pass, Callable, str]]
+ tir_pipeline : Optional[Union[tvm.ir.transform.Pass, Callable, str]]
The compilation pipeline to use for TIR functions.
Returns
diff --git a/python/tvm/exec/gpu_memory_bandwidth.py b/python/tvm/exec/gpu_memory_bandwidth.py
index 5330e5fff5..2b7b87e3e9 100644
--- a/python/tvm/exec/gpu_memory_bandwidth.py
+++ b/python/tvm/exec/gpu_memory_bandwidth.py
@@ -25,7 +25,7 @@ import numpy as np
import tvm
from tvm import te
from tvm.s_tir.meta_schedule.runner import EvaluatorConfig, RPCConfig
-from tvm.testing import local_run, rpc_run
+from tvm.s_tir.testing.runner import local_run, rpc_run
def _parse_args() -> argparse.Namespace:
diff --git a/python/tvm/ir/__init__.py b/python/tvm/ir/__init__.py
index ddf7c73068..e7a69a26a7 100644
--- a/python/tvm/ir/__init__.py
+++ b/python/tvm/ir/__init__.py
@@ -86,5 +86,5 @@ from .stmt import (
from tvm_ffi import Array, Map
-# Shared transformation infrastructure is available through the IR namespace.
-from .. import transform as transform
+# Load passes after the common IR classes used by transformations are registered.
+from . import transform
diff --git a/python/tvm/transform/__init__.py b/python/tvm/ir/transform/__init__.py
similarity index 100%
rename from python/tvm/transform/__init__.py
rename to python/tvm/ir/transform/__init__.py
diff --git a/python/tvm/transform/_ffi_instrument_api.py b/python/tvm/ir/transform/_ffi_instrument_api.py
similarity index 94%
copy from python/tvm/transform/_ffi_instrument_api.py
copy to python/tvm/ir/transform/_ffi_instrument_api.py
index 6499a5f0cf..d4aeb6dff0 100644
--- a/python/tvm/transform/_ffi_instrument_api.py
+++ b/python/tvm/ir/transform/_ffi_instrument_api.py
@@ -14,7 +14,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""FFI APIs for tvm.transform.instrument"""
+"""FFI APIs for tvm.ir.transform.instrument"""
import tvm_ffi
diff --git a/python/tvm/transform/_ffi_transform_api.py b/python/tvm/ir/transform/_ffi_transform_api.py
similarity index 95%
rename from python/tvm/transform/_ffi_transform_api.py
rename to python/tvm/ir/transform/_ffi_transform_api.py
index 64de643e54..2378b606c9 100644
--- a/python/tvm/transform/_ffi_transform_api.py
+++ b/python/tvm/ir/transform/_ffi_transform_api.py
@@ -14,7 +14,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""FFI APIs for tvm.transform"""
+"""FFI APIs for tvm.ir.transform"""
import tvm_ffi
diff --git a/python/tvm/transform/core.py b/python/tvm/ir/transform/core.py
similarity index 98%
rename from python/tvm/transform/core.py
rename to python/tvm/ir/transform/core.py
index c29abfb461..213969d2c0 100644
--- a/python/tvm/transform/core.py
+++ b/python/tvm/ir/transform/core.py
@@ -114,7 +114,7 @@ class PassInstrument(tvm.runtime.Object):
The module on which an optimization pass is being run.
- info : tvm.transform.PassInfo
+ info : tvm.ir.transform.PassInfo
The pass information.
@@ -137,7 +137,7 @@ class PassInstrument(tvm.runtime.Object):
The module on which an optimization pass is being run.
- info : tvm.transform.PassInfo
+ info : tvm.ir.transform.PassInfo
The pass information.
@@ -158,7 +158,7 @@ class PassInstrument(tvm.runtime.Object):
The module on which an optimization pass is being run.
- info : tvm.transform.PassInfo
+ info : tvm.ir.transform.PassInfo
The pass information.
@@ -204,7 +204,7 @@ def pass_instrument(pi_cls=None):
.. code-block:: python
- @tvm.transform.pass_instrument
+ @tvm.ir.transform.pass_instrument
class SkipPass:
def __init__(self, skip_pass_name):
self.skip_pass_name = skip_pass_name
@@ -232,7 +232,7 @@ def pass_instrument(pi_cls=None):
# pass
skip_annotate = SkipPass("AnnotateSpans")
- with tvm.transform.PassContext(instruments=[skip_annotate]):
+ with tvm.ir.transform.PassContext(instruments=[skip_annotate]):
tvm.compile(mod, "llvm")
"""
@@ -474,7 +474,7 @@ def module_pass(pass_func=None, opt_level=None, name=None):
.. code-block:: python
- @tvm.transform.module_pass
+ @tvm.ir.transform.module_pass
class CustomPipeline:
def __init__(self, enable_fold):
self.enable_fold = enable_fold
@@ -496,7 +496,7 @@ def module_pass(pass_func=None, opt_level=None, name=None):
.. code-block:: python
- @tvm.transform.module_pass(opt_level=2)
+ @tvm.ir.transform.module_pass(opt_level=2)
def transform(mod, ctx):
return relax.transform.FoldConstant(mod)
diff --git a/python/tvm/transform/instrument.py b/python/tvm/ir/transform/instrument.py
similarity index 98%
rename from python/tvm/transform/instrument.py
rename to python/tvm/ir/transform/instrument.py
index 6282baa4e8..37a2e96026 100644
--- a/python/tvm/transform/instrument.py
+++ b/python/tvm/ir/transform/instrument.py
@@ -44,7 +44,7 @@ class PassTimingInstrument(PassInstrument):
.. code-block:: python
timing_inst = PassTimingInstrument()
- with tvm.transform.PassContext(instruments=[timing_inst]):
+ with tvm.ir.transform.PassContext(instruments=[timing_inst]):
relax_mod = relax.transform.FuseOps()(relax_mod)
# before exiting the context, get profile results.
profiles = timing_inst.render()
diff --git a/python/tvm/ir/type_relation.py b/python/tvm/ir/type_relation.py
deleted file mode 100644
index 1160192c03..0000000000
--- a/python/tvm/ir/type_relation.py
+++ /dev/null
@@ -1,76 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-"""Type relation and function for type checking."""
-
-import tvm_ffi
-
-from . import _ffi_api
-from .type import Type, TypeConstraint
-
-
-@tvm_ffi.register_object("TypeCall")
-class TypeCall(Type):
- """Type function application.
-
- Parameters
- ----------
- func: tvm.ir.Type
- The function.
-
- args: List[tvm.ir.Type]
- The arguments.
-
- Returns
- -------
- type_call: TypeCall
- The type function application.
- """
-
- def __init__(self, func, args):
- self.__init_handle_by_constructor__(_ffi_api.TypeCall, func, args)
-
-
-@tvm_ffi.register_object("TypeRelation")
-class TypeRelation(TypeConstraint):
- """User defined type relation, it is an input-output relation on types.
-
- TypeRelation is more generalized than TypeCall as it allows inference
- of both inputs and outputs.
-
- Parameters
- ----------
- func : EnvFunc
- User defined relation function.
-
- args : [tvm.ir.Type]
- List of types to the func.
-
- num_inputs : int
- Number of input arguments in args,
- this act as a hint for type inference.
-
- attrs : Attrs
- The attribute attached to the relation information
-
- Returns
- -------
- type_relation : tvm.ir.TypeRelation
- The type relation.
- """
-
- def __init__(self, func, args, num_inputs, attrs):
- self.__init_handle_by_constructor__(_ffi_api.TypeRelation, func, args, num_inputs, attrs)
diff --git a/python/tvm/relax/backend/adreno/clml.py b/python/tvm/relax/backend/adreno/clml.py
index 0755770c5c..f7f0884bdf 100644
--- a/python/tvm/relax/backend/adreno/clml.py
+++ b/python/tvm/relax/backend/adreno/clml.py
@@ -19,6 +19,7 @@
import tvm
from tvm import IRModule, relax, tirx
+from tvm.ir.transform import PassContext, module_pass
from tvm.relax import transform
from tvm.relax.dpl.pattern import (
GlobalVarPattern,
@@ -31,7 +32,6 @@ from tvm.relax.dpl.pattern import (
from tvm.relax.expr import TupleGetItem, VarBinding
from tvm.relax.expr_functor import PyExprMutator, mutator
from tvm.relax.transform import PatternCheckContext
-from tvm.transform import PassContext, module_pass
from ..pattern_registry import register_patterns
@@ -82,7 +82,7 @@ class AppendReshapeToBNRewriter(PyExprMutator):
@transform.function_pass(opt_level=0, name="AppendReshapeToBN")
class AppendReshapeToBNRewriterPass:
def transform_function(
- self, func: relax.Function, mod: IRModule, _ctx: tvm.transform.PassContext
+ self, func: relax.Function, mod: IRModule, _ctx: tvm.ir.transform.PassContext
) -> relax.Function:
updated_func = AppendReshapeToBNRewriter(mod).visit_expr(func)
updated_func = relax.analysis.remove_all_unused(updated_func)
@@ -612,7 +612,7 @@ class OpenCLMLOffLoad:
"relax.nn.conv2d": ["NCHW", "OIHW"],
"relax.nn.conv2d_transpose": ["NCHW", "OIHW"],
}
- seq = tvm.transform.Sequential(
+ seq = tvm.ir.transform.Sequential(
[
transform.ConvertLayout(clml_layouts),
transform.Normalize(),
@@ -695,7 +695,7 @@ clml_llm_patterns = [
register_patterns(clml_llm_patterns)
-@tvm.transform.module_pass(opt_level=0, name="OpenCLMLOffLoadForLLM")
+@tvm.ir.transform.module_pass(opt_level=0, name="OpenCLMLOffLoadForLLM")
class OpenCLMLOffLoadForLLM:
"""A compiler pass that partition the graph with dequant Matmul to CLML backend offload."""
@@ -711,12 +711,12 @@ class OpenCLMLOffLoadForLLM:
def transform_module(
self,
mod: IRModule,
- _ctx: tvm.transform.PassContext,
+ _ctx: tvm.ir.transform.PassContext,
) -> IRModule:
"""Apply required passed to transform"""
if "adreno" in self.target.keys and (clml_sdk_version() >= 5):
- mod = tvm.transform.Sequential(
+ mod = tvm.ir.transform.Sequential(
[
transform.Normalize(),
transform.FuseOpsByPattern(clml_llm_patterns, annotate_codegen=True),
diff --git a/python/tvm/relax/backend/adreno/pipeline.py b/python/tvm/relax/backend/adreno/pipeline.py
index a58274f75c..b99bbdd473 100644
--- a/python/tvm/relax/backend/adreno/pipeline.py
+++ b/python/tvm/relax/backend/adreno/pipeline.py
@@ -123,10 +123,10 @@ def finalize_passes(target: tvm.target.Target): # pylint: disable=unused-argume
def get_default_pipeline(target: tvm.target.Target):
"""Return the default compilation pipeline for Adreno GPU."""
- @tvm.transform.module_pass(opt_level=0)
- def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext):
+ @tvm.ir.transform.module_pass(opt_level=0)
+ def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.ir.transform.PassContext):
with target:
- seq = tvm.transform.Sequential(
+ seq = tvm.ir.transform.Sequential(
library_dispatch_passes(target)
+ legalize_passes(target)
+ dataflow_lower_passes(target)
diff --git a/python/tvm/relax/backend/adreno/transform/transform.py b/python/tvm/relax/backend/adreno/transform/transform.py
index 99de83b1b4..5b591bbf46 100644
--- a/python/tvm/relax/backend/adreno/transform/transform.py
+++ b/python/tvm/relax/backend/adreno/transform/transform.py
@@ -23,27 +23,27 @@ from tvm.target import Target
from . import _ffi_api
-def AnnotateCustomMemoryScope(target: Target | None = None) -> tvm.transform.Pass:
+def AnnotateCustomMemoryScope(target: Target | None = None) -> tvm.ir.transform.Pass:
"""Allocate the memory scope information. This is Adreno specific pass to annotate
The memory scope information and realize the same with RealizeVDevice pass followed by
updating the Functiontion var_buffer mapping using SpecializeFunctionBasedOnCallSite.
Returns
-------
- ret: tvm.transform.Pass
+ ret: tvm.ir.transform.Pass
The registered pass for allocating workspace.
"""
return _ffi_api.AnnotateCustomMemoryScope(target) # type: ignore
-def FoldVDeviceScopeChange() -> tvm.transform.Pass:
+def FoldVDeviceScopeChange() -> tvm.ir.transform.Pass:
"""This pass is a texture specific pass that can optimize unnecessary to_device copies.
Like texture_scope -> ToVDevice -> global scope. In this case the producer can directly
store into global scope avoiding unnecessary device copy.
Returns
-------
- ret: tvm.transform.Pass
+ ret: tvm.ir.transform.Pass
The registered pass for allocating workspace.
"""
return _ffi_api.FoldVDeviceScopeChange() # type: ignore
diff --git a/python/tvm/relax/backend/contrib/tensorrt.py b/python/tvm/relax/backend/contrib/tensorrt.py
index 3711f29de1..b3a3c2b9cd 100644
--- a/python/tvm/relax/backend/contrib/tensorrt.py
+++ b/python/tvm/relax/backend/contrib/tensorrt.py
@@ -234,7 +234,7 @@ def partition_for_tensorrt(mod: IRModule) -> IRModule:
from tvm.relax.backend.contrib.tensorrt import partition_for_tensorrt
partitioned = partition_for_tensorrt(mod)
- with tvm.transform.PassContext(
+ with tvm.ir.transform.PassContext(
config={
"relax.ext.tensorrt.options": {
"build_at_compile_time": True,
diff --git a/python/tvm/relax/backend/cpu_generic/pipeline.py b/python/tvm/relax/backend/cpu_generic/pipeline.py
index d0b819cea7..36a1798f67 100644
--- a/python/tvm/relax/backend/cpu_generic/pipeline.py
+++ b/python/tvm/relax/backend/cpu_generic/pipeline.py
@@ -65,10 +65,10 @@ def finalize_passes(target: tvm.target.Target): # pylint: disable=unused-argume
def get_default_pipeline(target: tvm.target.Target):
"""Return the default compilation pipeline for CPU."""
- @tvm.transform.module_pass(opt_level=0)
- def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext):
+ @tvm.ir.transform.module_pass(opt_level=0)
+ def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.ir.transform.PassContext):
with target:
- seq = tvm.transform.Sequential(
+ seq = tvm.ir.transform.Sequential(
library_dispatch_passes(target)
+ legalize_passes(target)
+ dataflow_lower_passes(target)
diff --git a/python/tvm/relax/backend/cuda/cudnn.py b/python/tvm/relax/backend/cuda/cudnn.py
index 5c3c37f08a..523a651592 100644
--- a/python/tvm/relax/backend/cuda/cudnn.py
+++ b/python/tvm/relax/backend/cuda/cudnn.py
@@ -153,7 +153,7 @@ def partition_for_cudnn(mod):
"""
patterns = get_patterns_with_prefix("cudnn")
- return tvm.transform.Sequential(
+ return tvm.ir.transform.Sequential(
[
transform.FuseOpsByPattern(patterns, bind_constants=False, annotate_codegen=True),
annotate_workspace,
@@ -200,7 +200,7 @@ class WorkspaceAnnotator(PyExprMutator):
return f
-@tvm.transform.module_pass(opt_level=0)
+@tvm.ir.transform.module_pass(opt_level=0)
def annotate_workspace(mod, _):
"""Pass to annotate a workspace requirement for each cuDNN-offloaded function."""
annotator = WorkspaceAnnotator(mod)
diff --git a/python/tvm/relax/backend/cuda/cutlass.py b/python/tvm/relax/backend/cuda/cutlass.py
index 30178a8292..7ef24ad5b5 100644
--- a/python/tvm/relax/backend/cuda/cutlass.py
+++ b/python/tvm/relax/backend/cuda/cutlass.py
@@ -560,7 +560,7 @@ class WorkspaceAnnotator(PyExprMutator):
return f
-@tvm.transform.module_pass(opt_level=0)
+@tvm.ir.transform.module_pass(opt_level=0)
def annotate_workspace(mod, _):
"""Pass to annotate a workspace requirement for each CUTLASS-offloaded function."""
annotator = WorkspaceAnnotator(mod)
@@ -609,7 +609,7 @@ def partition_for_cutlass(mod, annotate_codegen=True, use_flash_mqa=True):
mod[func_name] = func
patterns = get_patterns_with_prefix("cutlass")
- return tvm.transform.Sequential(
+ return tvm.ir.transform.Sequential(
[
transform.FuseOpsByPattern(
patterns, bind_constants=False, annotate_codegen=annotate_codegen
diff --git a/python/tvm/relax/backend/cuda/pipeline.py b/python/tvm/relax/backend/cuda/pipeline.py
index e3d8b66c12..8073cf80f7 100644
--- a/python/tvm/relax/backend/cuda/pipeline.py
+++ b/python/tvm/relax/backend/cuda/pipeline.py
@@ -75,10 +75,10 @@ def finalize_passes(target: tvm.target.Target): # pylint: disable=unused-argume
def get_default_pipeline(target: tvm.target.Target):
"""Return the default compilation pipeline for CUDA."""
- @tvm.transform.module_pass(opt_level=0)
- def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext):
+ @tvm.ir.transform.module_pass(opt_level=0)
+ def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.ir.transform.PassContext):
with target:
- seq = tvm.transform.Sequential(
+ seq = tvm.ir.transform.Sequential(
library_dispatch_passes(target)
+ legalize_passes(target)
+ dataflow_lower_passes(target)
diff --git a/python/tvm/relax/backend/dispatch_sampling.py b/python/tvm/relax/backend/dispatch_sampling.py
index 3e7c75936c..bf9b079dc5 100644
--- a/python/tvm/relax/backend/dispatch_sampling.py
+++ b/python/tvm/relax/backend/dispatch_sampling.py
@@ -20,8 +20,8 @@
from tvm import relax
from tvm.ir import Op
from tvm.ir.module import IRModule
+from tvm.ir.transform import PassContext, module_pass
from tvm.relax import expr_functor
-from tvm.transform import PassContext, module_pass
from .utils import BackendDispatcher
diff --git a/python/tvm/relax/backend/dispatch_sort_scan.py b/python/tvm/relax/backend/dispatch_sort_scan.py
index 10902f8fe7..f91728daa0 100644
--- a/python/tvm/relax/backend/dispatch_sort_scan.py
+++ b/python/tvm/relax/backend/dispatch_sort_scan.py
@@ -24,9 +24,9 @@ from tvm import DataType, relax, topi
from tvm.contrib.thrust import can_use_thrust
from tvm.ir import GlobalVar, Op
from tvm.ir.module import IRModule
+from tvm.ir.transform import PassContext, module_pass
from tvm.relax import expr_functor
from tvm.target import Target
-from tvm.transform import PassContext, module_pass
from .utils import BackendDispatcher
diff --git a/python/tvm/relax/backend/gpu_generic/pipeline.py b/python/tvm/relax/backend/gpu_generic/pipeline.py
index a86ec6f311..8ded70313d 100644
--- a/python/tvm/relax/backend/gpu_generic/pipeline.py
+++ b/python/tvm/relax/backend/gpu_generic/pipeline.py
@@ -74,10 +74,10 @@ def finalize_passes(target: tvm.target.Target): # pylint: disable=unused-argume
def get_default_pipeline(target: tvm.target.Target):
"""Return the default compilation pipeline for generic GPU."""
- @tvm.transform.module_pass(opt_level=0)
- def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext):
+ @tvm.ir.transform.module_pass(opt_level=0)
+ def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.ir.transform.PassContext):
with target:
- seq = tvm.transform.Sequential(
+ seq = tvm.ir.transform.Sequential(
library_dispatch_passes(target)
+ legalize_passes(target)
+ dataflow_lower_passes(target)
diff --git a/python/tvm/relax/backend/rocm/pipeline.py b/python/tvm/relax/backend/rocm/pipeline.py
index a57172829a..f031f617da 100644
--- a/python/tvm/relax/backend/rocm/pipeline.py
+++ b/python/tvm/relax/backend/rocm/pipeline.py
@@ -74,10 +74,10 @@ def finalize_passes(target: tvm.target.Target): # pylint: disable=unused-argume
def get_default_pipeline(target: tvm.target.Target):
"""Return the default compilation pipeline for ROCm."""
- @tvm.transform.module_pass(opt_level=0)
- def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext):
+ @tvm.ir.transform.module_pass(opt_level=0)
+ def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.ir.transform.PassContext):
with target:
- seq = tvm.transform.Sequential(
+ seq = tvm.ir.transform.Sequential(
library_dispatch_passes(target)
+ legalize_passes(target)
+ dataflow_lower_passes(target)
diff --git a/python/tvm/relax/distributed/transform/transform.py b/python/tvm/relax/distributed/transform/transform.py
index e19af3b60f..87721e0b11 100644
--- a/python/tvm/relax/distributed/transform/transform.py
+++ b/python/tvm/relax/distributed/transform/transform.py
@@ -22,47 +22,47 @@ import tvm.ir
from . import _ffi_api
-def PropagateSharding() -> tvm.transform.Pass:
+def PropagateSharding() -> tvm.ir.transform.Pass:
"""Propagate sharding information.
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
The registered pass
"""
return _ffi_api.PropagateSharding() # type: ignore
-def LowerGlobalViewToLocalView() -> tvm.transform.Pass:
+def LowerGlobalViewToLocalView() -> tvm.ir.transform.Pass:
"""Lower global view TIR to local view
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
The registered pass
"""
return _ffi_api.LowerGlobalViewToLocalView() # type: ignore
-def LegalizeRedistribute() -> tvm.transform.Pass:
+def LegalizeRedistribute() -> tvm.ir.transform.Pass:
"""Legalize redistribute op to ccl op.
S->R: R.ccl.allgather
R->S: R.dist.redistribute_replica_to_shard
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
The registered pass
"""
return _ffi_api.LegalizeRedistribute() # type: ignore
-def LowerDistIR() -> tvm.transform.Pass:
+def LowerDistIR() -> tvm.ir.transform.Pass:
"""Lower DistIR to Relax
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
The registered pass
"""
return _ffi_api.LowerDistIR() # type: ignore
diff --git a/python/tvm/relax/expr_functor.py b/python/tvm/relax/expr_functor.py
index c6c600d55f..bfd77b176e 100644
--- a/python/tvm/relax/expr_functor.py
+++ b/python/tvm/relax/expr_functor.py
@@ -25,8 +25,8 @@ import tvm_ffi
from tvm import te as _te
from tvm import tirx as _tirx
from tvm.ir import Call, DataTypeImm, GenericConst, Op, StringImm, TensorLoad, is_prim_expr
-from tvm.ir.utils import derived_object
from tvm.runtime import Object
+from tvm.s_tir.utils import derived_object
from ..ir.module import IRModule
from . import _ffi_api
diff --git a/python/tvm/relax/frontend/nn/core.py b/python/tvm/relax/frontend/nn/core.py
index 73388732e9..e93152403f 100644
--- a/python/tvm/relax/frontend/nn/core.py
+++ b/python/tvm/relax/frontend/nn/core.py
@@ -37,11 +37,11 @@ import numpy as np # type: ignore
import tvm.runtime
from tvm import tirx
from tvm.ir import IRModule
+from tvm.ir.transform import Pass
from tvm.runtime import Device
from tvm.runtime import device as as_device
from tvm.runtime.vm import VirtualMachine
from tvm.target import Target
-from tvm.transform import Pass
from .... import relax as rx
from ...block_builder import BlockBuilder
diff --git a/python/tvm/relax/frontend/torch/dynamo.py b/python/tvm/relax/frontend/torch/dynamo.py
index c8fde62be2..e39dabd34e 100644
--- a/python/tvm/relax/frontend/torch/dynamo.py
+++ b/python/tvm/relax/frontend/torch/dynamo.py
@@ -37,12 +37,12 @@ def device_from_inputs(example_inputs):
return None
-def relax_dynamo(pipeline: tvm.transform.Pass | None = None):
+def relax_dynamo(pipeline: tvm.ir.transform.Pass | None = None):
"""A helper function to create a relax backend.
Parameters
----------
- pipeline : Optional[tvm.transform.Pass]
+ pipeline : Optional[tvm.ir.transform.Pass]
The pipeline to be applied to the relax module before sent to build.
Returns
diff --git a/python/tvm/relax/ir/instrument.py b/python/tvm/relax/ir/instrument.py
index 8690893817..9f61bc6f08 100644
--- a/python/tvm/relax/ir/instrument.py
+++ b/python/tvm/relax/ir/instrument.py
@@ -20,7 +20,7 @@ import tvm
from tvm import relax
-@tvm.transform.pass_instrument
+@tvm.ir.transform.pass_instrument
class WellFormedInstrument:
"""An instrument that checks the input/output IRModule of the Pass
is well formed. It will skip specific passes, like Normalize.
diff --git a/python/tvm/relax/pipeline.py b/python/tvm/relax/pipeline.py
index 668c2d9050..87226a3ce5 100644
--- a/python/tvm/relax/pipeline.py
+++ b/python/tvm/relax/pipeline.py
@@ -43,8 +43,8 @@ def zero_pipeline(*, enable_warning: bool = False):
the database. By default we don't print warning.
"""
- @tvm.transform.module_pass(opt_level=0)
- def f_zero_pipeline(mod: tvm.ir.IRModule, ctx: tvm.transform.PassContext) -> tvm.ir.IRModule:
+ @tvm.ir.transform.module_pass(opt_level=0)
+ def f_zero_pipeline(mod: tvm.ir.IRModule, ctx: tvm.ir.transform.PassContext) -> tvm.ir.IRModule:
"""Pipeline that applies pre-tuned logs.
Parameters
@@ -52,7 +52,7 @@ def zero_pipeline(*, enable_warning: bool = False):
mod : tvm.ir.IRModule
Input IRModule.
- ctx : tvm.transform.PassContext
+ ctx : tvm.ir.transform.PassContext
The pass context
Returns
@@ -60,7 +60,7 @@ def zero_pipeline(*, enable_warning: bool = False):
mod: tvm.ir.IRModule
The result transformed module.
"""
- seq = tvm.transform.Sequential(
+ seq = tvm.ir.transform.Sequential(
[
transform.LegalizeOps(enable_warning=enable_warning),
transform.AnnotateTIROpPattern(),
@@ -80,9 +80,9 @@ def zero_pipeline(*, enable_warning: bool = False):
def default_build_pipeline():
"""The default compilation pipeline used in tvm.compile"""
- @tvm.transform.module_pass(opt_level=0)
- def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.IRModule:
- seq = tvm.transform.Sequential(
+ @tvm.ir.transform.module_pass(opt_level=0)
+ def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.ir.transform.PassContext) -> tvm.ir.IRModule:
+ seq = tvm.ir.transform.Sequential(
[
backend.DispatchSampling(),
backend.DispatchSortScan(),
@@ -168,8 +168,8 @@ def static_shape_tuning_pipeline(
out = vm["main"](input_data, *params).numpy()
"""
- @tvm.transform.module_pass(opt_level=0)
- def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.IRModule:
+ @tvm.ir.transform.module_pass(opt_level=0)
+ def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.ir.transform.PassContext) -> tvm.ir.IRModule:
if cpu_weight_prepack:
pre_tuning_layout_rewrite = [transform.AttachAttrLayoutFreeTensors()]
post_tuning_layout_rewrite = [
@@ -182,7 +182,7 @@ def static_shape_tuning_pipeline(
post_tuning_layout_rewrite = []
with tvm.target.Target(target):
- mod = tvm.transform.Sequential(
+ mod = tvm.ir.transform.Sequential(
[
transform.DecomposeOpsForInference(),
transform.CanonicalizeBindings(),
@@ -197,7 +197,7 @@ def static_shape_tuning_pipeline(
max_trials_per_task=max_trials_per_task,
)
if total_trials > 0
- else tvm.transform.Sequential([])
+ else tvm.ir.transform.Sequential([])
),
transform.MetaScheduleApplyDatabase(work_dir),
*post_tuning_layout_rewrite,
@@ -218,7 +218,7 @@ PIPELINE_MAP = {
}
-def get_pipeline(name: str = "zero", **kwargs) -> tvm.transform.Pass:
+def get_pipeline(name: str = "zero", **kwargs) -> tvm.ir.transform.Pass:
"""Get pre-build pipeline by name
Parameters
@@ -230,7 +230,7 @@ def get_pipeline(name: str = "zero", **kwargs) -> tvm.transform.Pass:
Returns
-------
- pipeline: tvm.transform.Pass
+ pipeline: tvm.ir.transform.Pass
The transformation pipeline.
"""
diff --git a/python/tvm/relax/testing/gradient.py b/python/tvm/relax/testing/gradient.py
new file mode 100644
index 0000000000..bb89d71b54
--- /dev/null
+++ b/python/tvm/relax/testing/gradient.py
@@ -0,0 +1,168 @@
+# 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.
+# ruff: noqa: E501
+
+# pylint: disable=invalid-name,unnecessary-comprehension,redefined-outer-name
+"""Numerical checks for Relax gradients."""
+
+import logging
+
+import numpy as np
+
+
+def check_numerical_grads(
+ function, input_values, grad_values, function_value=None, delta=1e-3, atol=1e-2, rtol=0.1
+):
+ """A helper function that checks that numerical gradients of a function are
+ equal to gradients computed in some different way (analytical gradients).
+
+ Numerical gradients are computed using finite difference approximation. To
+ reduce the number of function evaluations, the number of points used is
+ gradually increased if the error value is too high (up to 5 points).
+
+ Parameters
+ ----------
+ function
+ A function that takes inputs either as positional or as keyword
+ arguments (either `function(*input_values)` or `function(**input_values)`
+ should be correct) and returns a scalar result. Should accept numpy
+ ndarrays.
+
+ input_values : Dict[str, numpy.ndarray] or List[numpy.ndarray]
+ A list of values or a dict assigning values to variables. Represents the
+ point at which gradients should be computed.
+
+ grad_values : Dict[str, numpy.ndarray] or List[numpy.ndarray]
+ Gradients computed using a different method.
+
+ function_value : float, optional
+ Should be equal to `function(**input_values)`.
+
+ delta : float, optional
+ A small number used for numerical computation of partial derivatives.
+ The default 1e-3 is a good choice for float32.
+
+ atol : float, optional
+ Absolute tolerance. Gets multiplied by `sqrt(n)` where n is the size of a
+ gradient.
+
+ rtol : float, optional
+ Relative tolerance.
+ """
+ # If input_values is a list then function accepts positional arguments
+ # In this case transform it to a function taking kwargs of the form {"0": ..., "1": ...}
+ if not isinstance(input_values, dict):
+ input_len = len(input_values)
+ input_values = {str(idx): val for idx, val in enumerate(input_values)}
+
+ def _function(_input_len=input_len, _orig_function=function, **kwargs):
+ return _orig_function(*(kwargs[str(i)] for i in range(input_len)))
+
+ function = _function
+
+ grad_values = {str(idx): val for idx, val in enumerate(grad_values)}
+
+ if function_value is None:
+ function_value = function(**input_values)
+
+ # a helper to modify j-th element of val by a_delta
+ def modify(val, j, a_delta):
+ val = val.copy()
+ val.reshape(-1)[j] = val.reshape(-1)[j] + a_delta
+ return val
+
+ # numerically compute a partial derivative with respect to j-th element of the var `name`
+ def derivative(x_name, j, a_delta):
+ modified_values = {
+ n: modify(val, j, a_delta) if n == x_name else val for n, val in input_values.items()
+ }
+ return (function(**modified_values) - function_value) / a_delta
+
+ def compare_derivative(j, n_der, grad):
+ der = grad.reshape(-1)[j]
+ return np.abs(n_der - der) < atol + rtol * np.abs(n_der)
+
+ for x_name, grad in grad_values.items():
+ if grad.shape != input_values[x_name].shape:
+ raise AssertionError(
+ f"Gradient wrt '{x_name}' has unexpected shape {grad.shape}, expected {input_values[x_name].shape} "
+ )
+
+ ngrad = np.zeros_like(grad)
+
+ wrong_positions = []
+
+ # compute partial derivatives for each position in this variable
+ for j in range(np.prod(grad.shape)):
+ # forward difference approximation
+ nder = derivative(x_name, j, delta)
+
+ # if the derivative is not equal to the analytical one, try to use more
+ # precise and expensive methods
+ if not compare_derivative(j, nder, grad):
+ # central difference approximation
+ nder = (derivative(x_name, j, -delta) + nder) / 2
+
+ if not compare_derivative(j, nder, grad):
+ # central difference approximation using h = delta/2
+ cnder2 = (
+ derivative(x_name, j, delta / 2) + derivative(x_name, j, -delta / 2)
+ ) / 2
+ # five-point derivative
+ nder = (4 * cnder2 - nder) / 3
+
+ # if the derivatives still don't match, add this position to the
+ # list of wrong positions
+ if not compare_derivative(j, nder, grad):
+ wrong_positions.append(np.unravel_index(j, grad.shape))
+
+ ngrad.reshape(-1)[j] = nder
+
+ wrong_percentage = int(100 * len(wrong_positions) / np.prod(grad.shape))
+
+ dist = np.sqrt(np.sum((ngrad - grad) ** 2))
+ grad_norm = np.sqrt(np.sum(ngrad**2))
+
+ if not (np.isfinite(dist) and np.isfinite(grad_norm)):
+ raise ValueError(
+ f"NaN or infinity detected during numerical gradient checking wrt '{x_name}'\n"
+ f"analytical grad = {grad}\n numerical grad = {ngrad}\n"
+ )
+
+ # we multiply atol by this number to make it more universal for different sizes
+ sqrt_n = np.sqrt(float(np.prod(grad.shape)))
+
+ if dist > atol * sqrt_n + rtol * grad_norm:
+ raise AssertionError(
+ f"Analytical and numerical grads wrt '{x_name}' differ too much\n"
+ f"analytical grad = {grad}\n numerical grad = {ngrad}\n"
+ f"{wrong_percentage}% of elements differ, first 10 of wrong positions: {wrong_positions[:10]}\n"
+ "distance > atol*sqrt(n) + rtol*grad_norm\n"
+ f"distance {dist} > {atol}*{sqrt_n} + {rtol}*{grad_norm}"
+ )
+
+ max_diff = np.max(np.abs(ngrad - grad))
+ avg_diff = np.mean(np.abs(ngrad - grad))
+ logging.info(
+ "Numerical grad test wrt '%s' of shape %s passes, "
+ "dist = %f, max_diff = %f, avg_diff = %f",
+ x_name,
+ grad.shape,
+ dist,
+ max_diff,
+ avg_diff,
+ )
diff --git a/python/tvm/relax/testing/transform.py b/python/tvm/relax/testing/transform.py
index f5512632b3..0094acc49f 100644
--- a/python/tvm/relax/testing/transform.py
+++ b/python/tvm/relax/testing/transform.py
@@ -29,7 +29,7 @@ from tvm.relax.expr import DataflowBlock, Var
from tvm.runtime import Object
-def ApplyEmptyCppMutator() -> tvm.transform.Pass:
+def ApplyEmptyCppMutator() -> tvm.ir.transform.Pass:
"""Create empty cpp mutator"""
packed_func = tvm.get_global_func("relax.testing.transform.ApplyEmptyCppMutator")
return packed_func()
diff --git a/python/tvm/relax/training/setup_trainer.py b/python/tvm/relax/training/setup_trainer.py
index bcbea96403..82db692cc1 100644
--- a/python/tvm/relax/training/setup_trainer.py
+++ b/python/tvm/relax/training/setup_trainer.py
@@ -30,7 +30,7 @@ from .loss import Loss
from .optimizer import Optimizer
-@tvm.transform.module_pass(opt_level=0, name="SetupTrainer")
+@tvm.ir.transform.module_pass(opt_level=0, name="SetupTrainer")
class SetupTrainer:
"""Transform a backbone module to a complete, legalized trainer module.
@@ -170,7 +170,7 @@ class SetupTrainer:
"tuple of length more than the number of states"
)
- def transform_module(self, mod: IRModule, ctx: tvm.transform.PassContext) -> IRModule:
+ def transform_module(self, mod: IRModule, ctx: tvm.ir.transform.PassContext) -> IRModule:
"""Transform the backbone module into a trainer module."""
self._check_well_formed(mod)
diff --git a/python/tvm/relax/training/utils.py b/python/tvm/relax/training/utils.py
index acb9392249..59391285a1 100644
--- a/python/tvm/relax/training/utils.py
+++ b/python/tvm/relax/training/utils.py
@@ -35,7 +35,7 @@ def AppendLoss(
loss_function: Function,
num_backbone_outputs: int = 1,
new_func_name: str | None = None,
-) -> tvm.transform.Pass:
+) -> tvm.ir.transform.Pass:
"""Append the loss function to the backbone function specified by `func_name`. Generally, the
loss function is generated by instances of `relax.training.Loss`.
diff --git a/python/tvm/relax/transform/_ffi_api.py b/python/tvm/relax/transform/_ffi_api.py
index 637ca3e6ea..a930142bf8 100644
--- a/python/tvm/relax/transform/_ffi_api.py
+++ b/python/tvm/relax/transform/_ffi_api.py
@@ -13,7 +13,7 @@
# "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
-"""FFI APIs for tvm.transform"""
+"""FFI APIs for tvm.relax.transform"""
import tvm_ffi
diff --git a/python/tvm/relax/transform/attach_external_modules.py b/python/tvm/relax/transform/attach_external_modules.py
index 2bf6ccd022..7becf3fd68 100644
--- a/python/tvm/relax/transform/attach_external_modules.py
+++ b/python/tvm/relax/transform/attach_external_modules.py
@@ -22,7 +22,7 @@ Note: "external modules" here refers to `relax.frontend.nn.ExternModule`.
from typing import TYPE_CHECKING
from tvm.ir import IRModule
-from tvm.transform import PassContext, module_pass
+from tvm.ir.transform import PassContext, module_pass
if TYPE_CHECKING:
from tvm.relax.frontend.nn import ExternModule
diff --git a/python/tvm/relax/transform/fast_math.py b/python/tvm/relax/transform/fast_math.py
index 088e778062..e28afbfd0a 100644
--- a/python/tvm/relax/transform/fast_math.py
+++ b/python/tvm/relax/transform/fast_math.py
@@ -50,13 +50,13 @@ class FastMathCodeGenerator(PyExprMutator):
return super().visit_call_(call)
-@tvm.transform.module_pass(opt_level=0, name="FastMathTransform")
+@tvm.ir.transform.module_pass(opt_level=0, name="FastMathTransform")
class FastMathTransform:
"""
Pass to convert the expensive non linear functions to their fast but approximate counterparts.
"""
- def transform_module(self, mod: IRModule, ctx: tvm.transform.PassContext) -> IRModule:
+ def transform_module(self, mod: IRModule, ctx: tvm.ir.transform.PassContext) -> IRModule:
fast_math_codegen = FastMathCodeGenerator(mod)
for gv, func in mod.functions_items():
if isinstance(func, tvm.relax.Function):
diff --git a/python/tvm/relax/transform/fold_batch_norm_to_conv2d_for_inference.py b/python/tvm/relax/transform/fold_batch_norm_to_conv2d_for_inference.py
index 7e25f95ae1..47056d3837 100644
--- a/python/tvm/relax/transform/fold_batch_norm_to_conv2d_for_inference.py
+++ b/python/tvm/relax/transform/fold_batch_norm_to_conv2d_for_inference.py
@@ -19,9 +19,9 @@
from tvm import relax, tirx
from tvm.ir.module import IRModule
+from tvm.ir.transform import PassContext
from tvm.relax import Expr
from tvm.relax.dpl import TupleGetItemPattern, is_const, is_op, rewrite_call, wildcard
-from tvm.transform import PassContext
from . import function_pass
diff --git a/python/tvm/relax/transform/fuse_transpose_matmul.py b/python/tvm/relax/transform/fuse_transpose_matmul.py
index 049a282524..342c390d13 100644
--- a/python/tvm/relax/transform/fuse_transpose_matmul.py
+++ b/python/tvm/relax/transform/fuse_transpose_matmul.py
@@ -29,11 +29,11 @@ from tvm.relax.dpl.pattern import is_op, wildcard
from tvm.relax.expr_functor import PyExprMutator, mutator
-@tvm.transform.module_pass(opt_level=0, name="FuseTransposeMatmul")
+@tvm.ir.transform.module_pass(opt_level=0, name="FuseTransposeMatmul")
class FuseTransposeMatmul: # pylint: disable=too-few-public-methods
"""A compiler pass that fuses transpose + matmul."""
- def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+ def transform_module(self, mod: IRModule, _ctx: tvm.ir.transform.PassContext) -> IRModule:
"""IRModule-level transformation"""
mod = relax.transform.FuseOpsByPattern(
[
diff --git a/python/tvm/relax/transform/ipc_allreduce_rewrite.py b/python/tvm/relax/transform/ipc_allreduce_rewrite.py
index d86b2c4351..191d1b77a6 100644
--- a/python/tvm/relax/transform/ipc_allreduce_rewrite.py
+++ b/python/tvm/relax/transform/ipc_allreduce_rewrite.py
@@ -25,7 +25,7 @@ from tvm.relax.expr import Expr, Var
from tvm.relax.expr_functor import PyExprMutator, PyExprVisitor, mutator, visitor
-@tvm.transform.module_pass(opt_level=0, name="IPCAllReduceRewrite")
+@tvm.ir.transform.module_pass(opt_level=0, name="IPCAllReduceRewrite")
class IPCAllReduceRewrite:
"""Rewrite all-reduce operation to customized all-reduce impl with IPC memory."""
@@ -40,7 +40,7 @@ class IPCAllReduceRewrite:
"""
self.allreduce_strategy = allreduce_strategy
- def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+ def transform_module(self, mod: IRModule, _ctx: tvm.ir.transform.PassContext) -> IRModule:
"""IRModule-level transformation"""
fcustom_allreduce = tvm.get_global_func(
"runtime.disco.cuda_ipc.custom_allreduce", allow_missing=True
diff --git a/python/tvm/relax/transform/lazy_transform_params.py b/python/tvm/relax/transform/lazy_transform_params.py
index 3ce65c494b..8d0c68b6c8 100644
--- a/python/tvm/relax/transform/lazy_transform_params.py
+++ b/python/tvm/relax/transform/lazy_transform_params.py
@@ -340,7 +340,7 @@ class LazyOutputMutator(PyExprMutator):
)
-@tvm.transform.module_pass(opt_level=0, name="LazyTransformParams")
+@tvm.ir.transform.module_pass(opt_level=0, name="LazyTransformParams")
class LazyTransformParams:
"""
Convert transform_params functions into a lazy version.
@@ -379,7 +379,7 @@ class LazyTransformParams:
self.fset_item = fset_item
self.extra_set_item_params = [] if extra_set_item_params is None else extra_set_item_params
- def transform_module(self, mod: IRModule, ctx: tvm.transform.PassContext) -> IRModule:
+ def transform_module(self, mod: IRModule, ctx: tvm.ir.transform.PassContext) -> IRModule:
lazy_mutator = LazyTransformParamsFuncCreator(
self.fget_item,
self.fset_item,
diff --git a/python/tvm/relax/transform/lower_gpu_ipc_alloc_storage.py b/python/tvm/relax/transform/lower_gpu_ipc_alloc_storage.py
index 7472c75a9f..3873411d33 100644
--- a/python/tvm/relax/transform/lower_gpu_ipc_alloc_storage.py
+++ b/python/tvm/relax/transform/lower_gpu_ipc_alloc_storage.py
@@ -25,11 +25,11 @@ from tvm.relax.expr import Expr
from tvm.relax.expr_functor import PyExprMutator, mutator
-@tvm.transform.module_pass(opt_level=0, name="LowerGPUIPCAllocStorage")
+@tvm.ir.transform.module_pass(opt_level=0, name="LowerGPUIPCAllocStorage")
class LowerGPUIPCAllocStorage:
"""Lower the storage/tensor allocation on IPC memory."""
- def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+ def transform_module(self, mod: IRModule, _ctx: tvm.ir.transform.PassContext) -> IRModule:
"""IRModule-level transformation"""
return _Rewriter(mod).transform()
diff --git a/python/tvm/relax/transform/optimize_layout_transform.py b/python/tvm/relax/transform/optimize_layout_transform.py
index 071a1e5c60..0dc00b2f4f 100644
--- a/python/tvm/relax/transform/optimize_layout_transform.py
+++ b/python/tvm/relax/transform/optimize_layout_transform.py
@@ -20,9 +20,9 @@
import tvm_ffi
from tvm.ir.module import IRModule
+from tvm.ir.transform import PassContext
from tvm.relax import Expr
from tvm.relax.dpl import TuplePattern, is_op, rewrite_call, wildcard
-from tvm.transform import PassContext
from . import function_pass
diff --git a/python/tvm/relax/transform/remove_redundant_reshape.py b/python/tvm/relax/transform/remove_redundant_reshape.py
index c27aa68888..0cc1e11900 100644
--- a/python/tvm/relax/transform/remove_redundant_reshape.py
+++ b/python/tvm/relax/transform/remove_redundant_reshape.py
@@ -20,9 +20,9 @@
import tvm_ffi
from tvm import IRModule, relax
+from tvm.ir.transform import PassContext
from tvm.relax import Expr
from tvm.relax.dpl import is_op, rewrite_call, wildcard
-from tvm.transform import PassContext
from . import function_pass
diff --git a/python/tvm/relax/transform/transform.py b/python/tvm/relax/transform/transform.py
index 03b671a6c6..abc6af1876 100644
--- a/python/tvm/relax/transform/transform.py
+++ b/python/tvm/relax/transform/transform.py
@@ -41,20 +41,20 @@ from .legalize_ops.common import LegalizeFunc
@tvm_ffi.register_object("relax.FunctionPass")
-class FunctionPass(tvm.transform.Pass):
+class FunctionPass(tvm.ir.transform.Pass):
"""A pass that works on each tvm.relax.Function in a module. A function
pass class should be created through `function_pass`.
"""
@tvm_ffi.register_object("relax.DataflowBlockPass")
-class DataflowBlockPass(tvm.transform.Pass):
+class DataflowBlockPass(tvm.ir.transform.Pass):
"""A pass that works on each tvm.relax.DataflowBlock in a module."""
def Gradient(
func_name: str, require_grads: Var | list[Var] | None = None, target_index: int = 0
-) -> tvm.transform.Pass:
+) -> tvm.ir.transform.Pass:
"""Reverse-mode automatic differentiation.
This pass will differentiate one function in the IRModule. Now the input function must have only
@@ -102,7 +102,7 @@ def Gradient(
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
The Pass.
Examples
@@ -229,17 +229,17 @@ def Gradient(
return _ffi_api.Gradient(func_name, require_grads, target_index) # type: ignore
-def ToNonDataflow() -> tvm.transform.Pass:
+def ToNonDataflow() -> tvm.ir.transform.Pass:
"""Transform all dataflow structure to non-dataflow version.
Returns
-------
- ret: tvm.transform.Pass
+ ret: tvm.ir.transform.Pass
"""
return _ffi_api.ToNonDataflow() # type: ignore
-def TopologicalSort(order="depth-first", direction="from-inputs") -> tvm.transform.Pass:
+def TopologicalSort(order="depth-first", direction="from-inputs") -> tvm.ir.transform.Pass:
"""Sort bindings in relax.Dataflow blocks in the order specified
Parameters
@@ -256,13 +256,13 @@ def TopologicalSort(order="depth-first", direction="from-inputs") -> tvm.transfo
Returns
-------
- ret: tvm.transform.Pass
+ ret: tvm.ir.transform.Pass
"""
return _ffi_api.TopologicalSort(order, direction) # type: ignore
-def RemovePurityChecking() -> tvm.transform.Pass:
+def RemovePurityChecking() -> tvm.ir.transform.Pass:
"""Activate relax.force_pure on all pure functions in the module
and unwrap all pure override ops into the normal versions.
@@ -271,7 +271,7 @@ def RemovePurityChecking() -> tvm.transform.Pass:
Returns
-------
- ret: tvm.transform.Pass
+ ret: tvm.ir.transform.Pass
The Pass.
Note
@@ -281,7 +281,7 @@ def RemovePurityChecking() -> tvm.transform.Pass:
return _ffi_api.RemovePurityChecking() # type: ignore
-def DataflowUseInplaceCalls() -> tvm.transform.Pass:
+def DataflowUseInplaceCalls() -> tvm.ir.transform.Pass:
"""
Pass that changes calls to operators that can be done in-place
(generally, these are elementwise operations) into in-place implementations.
@@ -293,23 +293,23 @@ def DataflowUseInplaceCalls() -> tvm.transform.Pass:
Returns
-------
- ret: tvm.transform.Pass
+ ret: tvm.ir.transform.Pass
The pass
"""
return _ffi_api.DataflowUseInplaceCalls()
-def LambdaLift() -> tvm.transform.Pass:
+def LambdaLift() -> tvm.ir.transform.Pass:
"""A pass that lifts local functions into global.
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
"""
return _ffi_api.LambdaLift()
-def LazyGetInput() -> tvm.transform.Pass:
+def LazyGetInput() -> tvm.ir.transform.Pass:
"""A pass that requests inputs lazily.
In many cases, the size of the model weights exceeds the available
@@ -341,13 +341,13 @@ def LazyGetInput() -> tvm.transform.Pass:
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
"""
return _ffi_api.LazyGetInput()
-def LazySetOutput() -> tvm.transform.Pass:
+def LazySetOutput() -> tvm.ir.transform.Pass:
"""A pass that sets function outputs when available
In many cases, the size of the model weights exceeds the available
@@ -383,13 +383,13 @@ def LazySetOutput() -> tvm.transform.Pass:
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
"""
return _ffi_api.LazySetOutput()
-def ConvertToDataflow(min_size: int = 2) -> tvm.transform.Pass:
+def ConvertToDataflow(min_size: int = 2) -> tvm.ir.transform.Pass:
"""A pass that converts consecutive dataflow operations
inside binding blocks into dataflow blocks.
@@ -403,13 +403,13 @@ def ConvertToDataflow(min_size: int = 2) -> tvm.transform.Pass:
Returns
-------
- ret: tvm.transform.Pass
+ ret: tvm.ir.transform.Pass
The pass.
"""
return _ffi_api.ConvertToDataflow(min_size)
-def CallTIRRewrite() -> tvm.transform.Pass:
+def CallTIRRewrite() -> tvm.ir.transform.Pass:
"""Perform explicit tensor allocation for call_tir and call_dps_packed.
Run ToNonDataflow and RemovePurityChecking first when lowering checked
@@ -418,23 +418,23 @@ def CallTIRRewrite() -> tvm.transform.Pass:
Returns
-------
- ret: tvm.transform.Pass
+ ret: tvm.ir.transform.Pass
"""
return _ffi_api.CallTIRRewrite() # type: ignore
-def Normalize() -> tvm.transform.Pass:
+def Normalize() -> tvm.ir.transform.Pass:
"""Transforming Relax IR to normal form, i.e., the expressions are normalized(no nesting
and hence the AST is in ANF), and all `ty` fields of expressions are available.
Returns
-------
- ret: tvm.transform.Pass
+ ret: tvm.ir.transform.Pass
"""
return _ffi_api.Normalize() # type: ignore
-def NormalizeGlobalVar() -> tvm.transform.Pass:
+def NormalizeGlobalVar() -> tvm.ir.transform.Pass:
"""Possibly rename the GlobalVar in an IRModule to ensure these properties:
1. (Invariant) First ensure every public function has the same name as its "global_symbol"
@@ -444,12 +444,12 @@ def NormalizeGlobalVar() -> tvm.transform.Pass:
Returns
-------
- ret: tvm.transform.Pass
+ ret: tvm.ir.transform.Pass
"""
return _ffi_api.NormalizeGlobalVar() # type: ignore
-def CanonicalizeBindings() -> tvm.transform.Pass:
+def CanonicalizeBindings() -> tvm.ir.transform.Pass:
"""
Canonicalizes variable definitions
(e.g., if there is y = x and z = y, it replaces uses of y and z with x).
@@ -464,7 +464,7 @@ def CanonicalizeBindings() -> tvm.transform.Pass:
Returns
-------
- ret: tvm.transform.Pass
+ ret: tvm.ir.transform.Pass
"""
return _ffi_api.CanonicalizeBindings() # type: ignore
@@ -481,13 +481,13 @@ def EliminateCommonSubexpr(call_only=False) -> FunctionPass:
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
The registered pass that eliminates common subexpressions.
"""
return _ffi_api.EliminateCommonSubexpr(call_only) # type: ignore
-def UpdateVDevice(new_vdevice: VDevice, index: int) -> tvm.transform.Pass:
+def UpdateVDevice(new_vdevice: VDevice, index: int) -> tvm.ir.transform.Pass:
"""Update virtual device.
Parameters
@@ -499,13 +499,13 @@ def UpdateVDevice(new_vdevice: VDevice, index: int) -> tvm.transform.Pass:
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
The registered pass that modifies the virtual device.
"""
return _ffi_api.UpdateVDevice(new_vdevice, index) # type: ignore
-def RewriteDataflowReshape() -> tvm.transform.Pass:
+def RewriteDataflowReshape() -> tvm.ir.transform.Pass:
"""Convert all reshape-like call_tir to VM reshape operator call.
The VM reshape operator calls will be further lowered to a CreateView
operation at runtime, instead of doing real data copy.
@@ -515,12 +515,12 @@ def RewriteDataflowReshape() -> tvm.transform.Pass:
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
"""
return _ffi_api.RewriteDataflowReshape() # type: ignore
-def StaticPlanBlockMemory() -> tvm.transform.Pass:
+def StaticPlanBlockMemory() -> tvm.ir.transform.Pass:
"""The static memory planning pass on BindingBlock level.
The pass will reuse allocated memory to its best effort, in order to
reduce the total amount of allocated memory size.
@@ -540,12 +540,12 @@ def StaticPlanBlockMemory() -> tvm.transform.Pass:
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
"""
return _ffi_api.StaticPlanBlockMemory() # type: ignore
-def LowerAllocTensor() -> tvm.transform.Pass:
+def LowerAllocTensor() -> tvm.ir.transform.Pass:
"""Lower remaining instances of R.builtin.alloc_tensor
The static memory planner removes static instances of
@@ -560,22 +560,22 @@ def LowerAllocTensor() -> tvm.transform.Pass:
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
"""
return _ffi_api.LowerAllocTensor() # type: ignore
-def KillAfterLastUse() -> tvm.transform.Pass:
+def KillAfterLastUse() -> tvm.ir.transform.Pass:
"""Drop all tensor/storage objects after last use
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
"""
return _ffi_api.KillAfterLastUse() # type: ignore
-def ComputePrimValue() -> tvm.transform.Pass:
+def ComputePrimValue() -> tvm.ir.transform.Pass:
"""Compute symbolic primitive expressions
While high-level relax can include expressions in terms of its
@@ -588,23 +588,23 @@ def ComputePrimValue() -> tvm.transform.Pass:
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
"""
return _ffi_api.ComputePrimValue() # type: ignore
-def LowerRuntimeBuiltin() -> tvm.transform.Pass:
+def LowerRuntimeBuiltin() -> tvm.ir.transform.Pass:
"""Lowering generic intrinsic to VM intrinsics.
Returns
-------
- ret: tvm.transform.Pass
+ ret: tvm.ir.transform.Pass
"""
return _ffi_api.LowerRuntimeBuiltin() # type: ignore
-def VMShapeLower(*, emit_err_ctx: bool = True) -> tvm.transform.Pass:
+def VMShapeLower(*, emit_err_ctx: bool = True) -> tvm.ir.transform.Pass:
"""Lower the symbolic shape and argument and match-cast structinfo matching.
Parameters
@@ -614,17 +614,17 @@ def VMShapeLower(*, emit_err_ctx: bool = True) -> tvm.transform.Pass:
Returns
-------
- ret: tvm.transform.Pass
+ ret: tvm.ir.transform.Pass
"""
return _ffi_api.VMShapeLower(emit_err_ctx) # type: ignore
-def AttachGlobalSymbol() -> tvm.transform.Pass:
+def AttachGlobalSymbol() -> tvm.ir.transform.Pass:
"""Attach global_symbol to Relax functions and TIR Functions for codegen.
Returns
-------
- ret: tvm.transform.Pass
+ ret: tvm.ir.transform.Pass
"""
return _ffi_api.AttachGlobalSymbol() # type: ignore
@@ -632,7 +632,7 @@ def AttachGlobalSymbol() -> tvm.transform.Pass:
def BindParams(
func_name: str,
params: dict[str | Var, tvm.runtime.Tensor | np.ndarray],
-) -> tvm.transform.Pass:
+) -> tvm.ir.transform.Pass:
"""Bind params of function of the module to constant tensors.
Parameters
@@ -645,7 +645,7 @@ def BindParams(
Returns
-------
- ret: tvm.transform.Pass
+ ret: tvm.ir.transform.Pass
"""
tvm_params = {}
for k, v in params.items():
@@ -663,7 +663,7 @@ def BindParams(
def BindSymbolicVars(
binding_map: Mapping[str | tvm.tirx.Var, tvm.tirx.Expr],
func_name: str | None = None,
-) -> tvm.transform.Pass:
+) -> tvm.ir.transform.Pass:
"""Bind params of function of the module to constant tensors.
Parameters
@@ -677,7 +677,7 @@ def BindSymbolicVars(
Returns
-------
- ret: tvm.transform.Pass
+ ret: tvm.ir.transform.Pass
"""
# Relax uses int64 for symbolic variables, but the FFI
# converts python integers into int32.
@@ -691,7 +691,7 @@ def BindSymbolicVars(
def RunCodegen(
target_options: dict | None = None,
entry_functions: list[str] | None = None,
-) -> tvm.transform.Pass:
+) -> tvm.ir.transform.Pass:
"""Produce the runtime::Module with an annotated codegen and global symbol.
Parameters
@@ -703,7 +703,7 @@ def RunCodegen(
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
The registered pass to remove unused functions.
"""
if entry_functions is None:
@@ -716,69 +716,69 @@ def RunCodegen(
return _ffi_api.RunCodegen(target_options, entry_functions) # type: ignore
-def FoldConstant() -> tvm.transform.Pass:
+def FoldConstant() -> tvm.ir.transform.Pass:
"""Fold constant expressions within dataflow blocks.
Note: ConvertToDataflow may need to be called first to provide dataflow blocks.
Returns
-------
- ret: tvm.transform.Pass
+ ret: tvm.ir.transform.Pass
"""
return _ffi_api.FoldConstant() # type: ignore
-def ExpandTupleArguments() -> tvm.transform.Pass:
+def ExpandTupleArguments() -> tvm.ir.transform.Pass:
"""Expand tuple arguments to internal functions
Returns
-------
- ret: tvm.transform.Pass
+ ret: tvm.ir.transform.Pass
"""
return _ffi_api.ExpandTupleArguments() # type: ignore
-def RemoveUnusedParameters() -> tvm.transform.Pass:
+def RemoveUnusedParameters() -> tvm.ir.transform.Pass:
"""Remove unused arguments to internal functions
Returns
-------
- ret: tvm.transform.Pass
+ ret: tvm.ir.transform.Pass
"""
return _ffi_api.RemoveUnusedParameters() # type: ignore
-def RemoveUnusedOutputs() -> tvm.transform.Pass:
+def RemoveUnusedOutputs() -> tvm.ir.transform.Pass:
"""Remove unused outputs from internal functions
Returns
-------
- ret: tvm.transform.Pass
+ ret: tvm.ir.transform.Pass
"""
return _ffi_api.RemoveUnusedOutputs() # type: ignore
-def InlinePrivateFunctions() -> tvm.transform.Pass:
+def InlinePrivateFunctions() -> tvm.ir.transform.Pass:
"""Inline all private relax functions
Returns
-------
- ret: tvm.transform.Pass
+ ret: tvm.ir.transform.Pass
"""
return _ffi_api.InlinePrivateFunctions() # type: ignore
-def AnnotateTIROpPattern() -> tvm.transform.Pass:
+def AnnotateTIROpPattern() -> tvm.ir.transform.Pass:
"""Annotate Op Pattern Kind for TIR functions
Returns
-------
- ret: tvm.transform.Pass
+ ret: tvm.ir.transform.Pass
"""
return _ffi_api.AnnotateTIROpPattern() # type: ignore
-def FuseOps(fuse_opt_level=-1) -> tvm.transform.Pass:
+def FuseOps(fuse_opt_level=-1) -> tvm.ir.transform.Pass:
"""This pass groups bindings in a dataflow block of Relax functions and generate a new grouped
Relax function for each group, according to the fusion algorithm described in the pass
implementation. By grouping bindings into new Relax functions, we substitute the bindings in
@@ -796,18 +796,18 @@ def FuseOps(fuse_opt_level=-1) -> tvm.transform.Pass:
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
The registered pass for operator fusion.
"""
return _ffi_api.FuseOps(fuse_opt_level) # type: ignore
-def FuseTIR() -> tvm.transform.Pass:
+def FuseTIR() -> tvm.ir.transform.Pass:
"""Fuse primitive relax function into a larger TIR function if possible
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
The registered pass for tirx fusion.
"""
return _ffi_api.FuseTIR() # type: ignore
@@ -896,7 +896,7 @@ def FuseOpsByPattern(
bind_constants: bool = True,
annotate_codegen: bool = False,
entry_functions: list[str] | None = None,
-) -> tvm.transform.Pass:
+) -> tvm.ir.transform.Pass:
"""Apply pattern matching to each function in the given module, and group matched expressions
into a new function.
@@ -930,7 +930,7 @@ def FuseOpsByPattern(
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
The registered pass for pattern-based fusion.
"""
@@ -951,20 +951,20 @@ def FuseOpsByPattern(
) # type: ignore
-def MergeCompositeFunctions() -> tvm.transform.Pass:
+def MergeCompositeFunctions() -> tvm.ir.transform.Pass:
"""Group one or multiple composite functions created by FuseOpsByPattern into a new function.
The new function will be annotated with "Codegen" and "global_symbol" attributes, and it
is intented to be offloaded to an external backend.
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
The registered pass for merging composite functions.
"""
return _ffi_api.MergeCompositeFunctions() # type: ignore
-def AttachAttrLayoutFreeTensors() -> tvm.transform.Pass:
+def AttachAttrLayoutFreeTensors() -> tvm.ir.transform.Pass:
"""Attach layout free tensors to the tirx::Function.
This pass is used to attach layout free tensors to the tirx::Function according to
@@ -975,25 +975,25 @@ def AttachAttrLayoutFreeTensors() -> tvm.transform.Pass:
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
The registered pass for attaching layout free tensors.
"""
return _ffi_api.AttachAttrLayoutFreeTensors() # type: ignore
-def SplitLayoutRewritePreproc() -> tvm.transform.Pass:
+def SplitLayoutRewritePreproc() -> tvm.ir.transform.Pass:
"""Split the TIR layout rewrite into multiple TIR functions.
This pass is used in the prepack weight after meta_schedule tuning.
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
The registered pass for splitting TIR layout rewrite.
"""
return _ffi_api.SplitLayoutRewritePreproc() # type: ignore
-def LiftTransformParams(shared_transform: bool | list[str] = False) -> tvm.transform.Pass:
+def LiftTransformParams(shared_transform: bool | list[str] = False) -> tvm.ir.transform.Pass:
"""Lift transformation of the parameters of a function.
When some inputs of the function is marked as 'parameters' (the model weights), this pass
@@ -1025,13 +1025,13 @@ def LiftTransformParams(shared_transform: bool | list[str] = False) -> tvm.trans
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
The registered pass for lifting transformation of parameters.
"""
return _ffi_api.LiftTransformParams(shared_transform) # type: ignore
-def BundleModelParams(param_tuple_name: str | None = None) -> tvm.transform.Pass:
+def BundleModelParams(param_tuple_name: str | None = None) -> tvm.ir.transform.Pass:
"""Bundle several model parameters into a single parameter tuple
For each function, if the function has the attribute "num_input",
@@ -1048,7 +1048,7 @@ def BundleModelParams(param_tuple_name: str | None = None) -> tvm.transform.Pass
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
The registered pass for bundling model parameters.
"""
return _ffi_api.BundleModelParams(param_tuple_name) # type: ignore
@@ -1093,7 +1093,7 @@ def LegalizeOps(
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
The registered pass
Examples
@@ -1168,12 +1168,12 @@ def LegalizeOps(
return _ffi_api.LegalizeOps(customize_legalize_map, skip_ops, enable_warning) # type: ignore
-def RealizeVDevice() -> tvm.transform.Pass:
+def RealizeVDevice() -> tvm.ir.transform.Pass:
"""Propagate virtual device information.
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
The registered pass
"""
@@ -1182,7 +1182,7 @@ def RealizeVDevice() -> tvm.transform.Pass:
def MetaScheduleApplyDatabase(
work_dir: str | None = None, enable_warning: bool = False
-) -> tvm.transform.Pass:
+) -> tvm.ir.transform.Pass:
"""Apply the best schedule from tuning database.
Parameters
@@ -1196,7 +1196,7 @@ def MetaScheduleApplyDatabase(
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
The registered pass
"""
return _ffi_api.MetaScheduleApplyDatabase(work_dir, enable_warning) # type: ignore
@@ -1205,7 +1205,7 @@ def MetaScheduleApplyDatabase(
def MetaScheduleTuneTIR(
work_dir: str,
max_trials_global: int,
-) -> tvm.transform.Pass:
+) -> tvm.ir.transform.Pass:
"""Tune TIR with MetaSchedule.
Parameters
@@ -1216,7 +1216,7 @@ def MetaScheduleTuneTIR(
maximum number of total trials allowed for tuning
Returns
-------
- ret: tvm.transform.Pass
+ ret: tvm.ir.transform.Pass
"""
return _ffi_api.MetaScheduleTuneTIR(work_dir, max_trials_global) # type: ignore
@@ -1227,7 +1227,7 @@ def MetaScheduleTuneIRMod(
max_trials_global: int,
max_trials_per_task: int | None = None,
op_names: list[str] | None = None,
-) -> tvm.transform.Pass:
+) -> tvm.ir.transform.Pass:
"""Tune Relax IRModule with MetaSchedule.
Parameters
@@ -1246,14 +1246,14 @@ def MetaScheduleTuneIRMod(
Returns
-------
- ret: tvm.transform.Pass
+ ret: tvm.ir.transform.Pass
"""
return _ffi_api.MetaScheduleTuneIRMod(
params, work_dir, max_trials_global, max_trials_per_task, op_names
) # type: ignore
-def DecomposeOpsForInference(func_name: str | None = None) -> tvm.transform.Pass:
+def DecomposeOpsForInference(func_name: str | None = None) -> tvm.ir.transform.Pass:
"""Decompose composite operators that are composed by other operators during inference.
For example, the result of batch norm (a triple) will be simplified. Attention, tensor_to_shape,
etc. can be also decomposed into a number of simplified operators as well.
@@ -1266,13 +1266,13 @@ def DecomposeOpsForInference(func_name: str | None = None) -> tvm.transform.Pass
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
The registered pass
"""
return _ffi_api.DecomposeOpsForInference(func_name) # type: ignore
-def DecomposeOpsForTraining(func_name: str | None = None) -> tvm.transform.Pass:
+def DecomposeOpsForTraining(func_name: str | None = None) -> tvm.ir.transform.Pass:
"""Decompose composite operators that are composed by other operators during training.
For example, the result of batch norm (a triple) will be simplified. Attention, tensor_to_shape,
etc. can be also decomposed into a number of simplified operators as well.
@@ -1285,7 +1285,7 @@ def DecomposeOpsForTraining(func_name: str | None = None) -> tvm.transform.Pass:
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
The registered pass
"""
return _ffi_api.DecomposeOpsForTraining(func_name) # type: ignore
@@ -1309,7 +1309,7 @@ def AlterOpImpl(
op_kind to layout transformation map for each of the tensors
Returns
-------
- ret: tvm.transform.Pass
+ ret: tvm.ir.transform.Pass
"""
for operator_name, transform_list in op_tensor_transforms.items():
l = []
@@ -1328,7 +1328,7 @@ def AlterOpImpl(
def ConvertLayout(
desired_layouts: dict[str, list[str]],
layout_cb: Callable | None = None,
-) -> tvm.transform.Pass:
+) -> tvm.ir.transform.Pass:
"""Automatic layout conversion pass.
Parameters
@@ -1344,13 +1344,13 @@ def ConvertLayout(
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
The registered pass for layout conversion.
"""
return _ffi_api.ConvertLayout(desired_layouts, layout_cb) # type: ignore
-def DeadCodeElimination(entry_functions: list[str] | None = None) -> tvm.transform.Pass:
+def DeadCodeElimination(entry_functions: list[str] | None = None) -> tvm.ir.transform.Pass:
"""Remove dead code in the IRModule.
Currently it removes:
@@ -1372,7 +1372,7 @@ def DeadCodeElimination(entry_functions: list[str] | None = None) -> tvm.transfo
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
The registered pass.
"""
if entry_functions is None:
@@ -1382,7 +1382,7 @@ def DeadCodeElimination(entry_functions: list[str] | None = None) -> tvm.transfo
def ToMixedPrecision(
out_dtype="float32", fp16_input_names: list[str] | None = None
-) -> tvm.transform.Pass:
+) -> tvm.ir.transform.Pass:
"""Automatic mixed precision pass. Currently the pass assumes the input module to be fp32
only, and will automatically cast fp32 to fp16 for certain ops.
@@ -1398,13 +1398,13 @@ def ToMixedPrecision(
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
The registered pass for mixed precision.
"""
return _ffi_api.ToMixedPrecision(out_dtype, fp16_input_names) # type: ignore
-def SplitCallTIRByPattern(patterns: list[Function], fcodegen: Callable) -> tvm.transform.Pass:
+def SplitCallTIRByPattern(patterns: list[Function], fcodegen: Callable) -> tvm.ir.transform.Pass:
"""Split a Function into 2 parts: the first part is a TIR Function which is
matched with some pattern, and the second part is the rest of the original
Function. It will call fcodegen to generate the code for the matched pattern
@@ -1420,7 +1420,7 @@ def SplitCallTIRByPattern(patterns: list[Function], fcodegen: Callable) -> tvm.t
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
The registered pass for splitting call_tir.
"""
return _ffi_api.SplitCallTIRByPattern(patterns, fcodegen) # type: ignore
@@ -1443,7 +1443,7 @@ def UpdateParamType(ty_func: Callable[[Var], Type | None]):
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
The corresponding pass.
"""
@@ -1460,7 +1460,7 @@ def AdjustMatmulOrder():
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
The corresponding pass.
"""
@@ -1480,7 +1480,7 @@ def ExpandMatmulOfSum():
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
The corresponding pass.
"""
@@ -1500,7 +1500,7 @@ def ReorderPermuteDimsAfterConcat():
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
The corresponding pass.
"""
@@ -1515,7 +1515,7 @@ def ReorderTakeAfterMatmul():
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
The corresponding pass.
"""
@@ -1538,7 +1538,7 @@ def CombineParallelMatmul(check=None):
Returns
-------
- ret : tvm.transform.Pass
+ ret : tvm.ir.transform.Pass
The corresponding pass.
"""
if check is None:
@@ -1546,19 +1546,19 @@ def CombineParallelMatmul(check=None):
return _ffi_api.CombineParallelMatmul(check) # type: ignore
-def RewriteCUDAGraph() -> tvm.transform.Pass:
+def RewriteCUDAGraph() -> tvm.ir.transform.Pass:
"""Rewrite a Relax module for executing with CUDA graph. This pass identifies the regions that
can be executed with CUDA graph and lifts them into new functions for runtime graph capturing.
Returns
-------
- ret: tvm.transform.Pass
+ ret: tvm.ir.transform.Pass
The registered pass for rewriting cuda graph
"""
return _ffi_api.RewriteCUDAGraph() # type: ignore
-def AllocateWorkspace() -> tvm.transform.Pass:
+def AllocateWorkspace() -> tvm.ir.transform.Pass:
"""Allocate a workspace, represented by a tensor of size big enough for all external
functions that require a temporary storage, and append it to the arguments of external
functions.
@@ -1567,20 +1567,20 @@ def AllocateWorkspace() -> tvm.transform.Pass:
Returns
-------
- ret: tvm.transform.Pass
+ ret: tvm.ir.transform.Pass
The registered pass for allocating workspace.
"""
return _ffi_api.AllocateWorkspace() # type: ignore
-def SpecializeFunctionBasedOnCallSite() -> tvm.transform.Pass:
+def SpecializeFunctionBasedOnCallSite() -> tvm.ir.transform.Pass:
"""This pass updates the var_tensor mapping of Functiontions from the call_tir info.
Primarily used to update the VDevice information if any changes occurred from the caller.
This pass recreates the tensors and updates the map.
Returns
-------
- ret: tvm.transform.Pass
+ ret: tvm.ir.transform.Pass
The registered pass for specializing Functions based on call site.
"""
return _ffi_api.SpecializeFunctionBasedOnCallSite() # type: ignore
@@ -1710,7 +1710,7 @@ def function_pass(
def create_function_pass(pass_arg):
"""Internal function that creates a function pass"""
fname = name if name else pass_arg.__name__
- info = tvm.transform.PassInfo(opt_level, fname)
+ info = tvm.ir.transform.PassInfo(opt_level, fname)
if inspect.isclass(pass_arg):
return _wrap_class_function_pass(pass_arg, info)
if not isinstance(pass_arg, types.FunctionType | types.LambdaType):
@@ -1850,7 +1850,7 @@ def dataflowblock_pass(pass_func=None, opt_level=None, name=None) -> Callable |
def create_dataflowblock_pass(pass_arg):
"""Internal function that creates a dataflowblock pass"""
fname = name if name else pass_arg.__name__
- info = tvm.transform.PassInfo(opt_level, fname)
+ info = tvm.ir.transform.PassInfo(opt_level, fname)
if inspect.isclass(pass_arg):
return _wrap_class_dataflowblock_pass(pass_arg, info)
if not isinstance(pass_arg, types.FunctionType | types.LambdaType):
diff --git a/python/tvm/relax/utils.py b/python/tvm/relax/utils.py
index 3f8fa37f7e..0d7c135d99 100644
--- a/python/tvm/relax/utils.py
+++ b/python/tvm/relax/utils.py
@@ -42,14 +42,6 @@ from .expr import Tuple as rx_Tuple
from .type import ShapeType, TensorType
-def _substitute_tir_vars(value, var_map):
- return tvm_ffi.structural_map(
- value,
- (tvm.ir.Var, lambda var: var_map.get(var, var)),
- order="post",
- )
-
-
def metadata_partitioner(rx_txt: str) -> list[str]:
"""Extract Relax program and metadata section.
@@ -230,7 +222,9 @@ def gen_call_tir_inputs(
if tvm.ir.is_prim_expr(arg):
_copy_undefined_var(arg)
- new_arg = _substitute_tir_vars(arg, tir_var_map)
+ new_arg = tvm_ffi.structural_map(
+ arg, (tvm.ir.Var, lambda var: tir_var_map.get(var, var)), order="post"
+ )
extra_tir_args_list.append(new_arg)
return new_arg
@@ -332,7 +326,12 @@ def gen_call_tir_inputs(
shape_values: list[tirx.Expr], tir_var_inverse_map: dict[tirx.Var, tirx.Expr]
):
return ShapeExpr(
- [_substitute_tir_vars(value, tir_var_inverse_map) for value in shape_values]
+ [
+ tvm_ffi.structural_map(
+ value, (tvm.ir.Var, lambda var: tir_var_inverse_map.get(var, var)), order="post"
+ )
+ for value in shape_values
+ ]
)
function_attrs = kwargs.pop("function_attrs", None)
@@ -377,7 +376,10 @@ def gen_call_tir_inputs(
output_ty = output_ty[0] if len(output_ty) == 1 else tvm.ir.TupleType(output_ty)
call_tir_args.extend(
- _substitute_tir_vars(value, tir_var_inverse_map) for value in unbound_tir_vars
+ tvm_ffi.structural_map(
+ value, (tvm.ir.Var, lambda var: tir_var_inverse_map.get(var, var)), order="post"
+ )
+ for value in unbound_tir_vars
)
return (tir_func, call_tir_args, output_ty)
diff --git a/python/tvm/relax/vm_build.py b/python/tvm/relax/vm_build.py
index d8e046cc10..fcd5c1c3ce 100644
--- a/python/tvm/relax/vm_build.py
+++ b/python/tvm/relax/vm_build.py
@@ -98,7 +98,7 @@ def _vmlink(
builder: "relax.ExecBuilder",
target: str | tvm.target.Target | None,
tir_mod: tvm.IRModule | None = None,
- tir_pipeline: str | tvm.transform.Pass | None = "default",
+ tir_pipeline: str | tvm.ir.transform.Pass | None = "default",
ext_libs: list[tvm.runtime.Module] | None = None,
params: dict[str, list] | None = None,
*,
@@ -165,8 +165,8 @@ def build(
mod: tvm.IRModule,
target: str | tvm.target.Target | None = None,
params: dict[str, list] | None = None,
- relax_pipeline: None | str | tvm.transform.Pass = "default",
- tir_pipeline: None | str | tvm.transform.Pass = "default",
+ relax_pipeline: None | str | tvm.ir.transform.Pass = "default",
+ tir_pipeline: None | str | tvm.ir.transform.Pass = "default",
*,
system_lib: bool | None = None,
) -> Executable:
diff --git a/python/tvm/s_tir/analysis/__init__.py b/python/tvm/s_tir/analysis/__init__.py
index ffd9c42213..97b2e578ed 100644
--- a/python/tvm/s_tir/analysis/__init__.py
+++ b/python/tvm/s_tir/analysis/__init__.py
@@ -188,19 +188,19 @@ def OOBChecker():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.OOBChecker() # type: ignore
-def get_vtcm_compaction_passes() -> list[tvm.transform.Pass]:
+def get_vtcm_compaction_passes() -> list[tvm.ir.transform.Pass]:
"""Utility function to get the list of lowering passes to be applied to calculate the compacted
VTCM allocation size
Returns
-------
- result : List[tvm.transform.Pass]
+ result : List[tvm.ir.transform.Pass]
returns list of passes
"""
return _ffi_api.get_vtcm_compaction_passes() # type: ignore # pylint: disable=no-member
diff --git a/python/tvm/s_tir/backend/adreno/transform/transform.py b/python/tvm/s_tir/backend/adreno/transform/transform.py
index 03cc162a4b..0fa478818e 100644
--- a/python/tvm/s_tir/backend/adreno/transform/transform.py
+++ b/python/tvm/s_tir/backend/adreno/transform/transform.py
@@ -25,7 +25,7 @@ def InjectTextureAlloc():
via alloc_nd/alloc_free calls
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.InjectTextureAlloc() # type: ignore
@@ -40,7 +40,7 @@ def TextureFlatten():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.TextureFlatten() # type: ignore
diff --git a/python/tvm/s_tir/dlight/base/transform.py b/python/tvm/s_tir/dlight/base/transform.py
index 1983815fa2..ab2743a427 100644
--- a/python/tvm/s_tir/dlight/base/transform.py
+++ b/python/tvm/s_tir/dlight/base/transform.py
@@ -21,8 +21,8 @@ or a space for MetaSchedule tuning
from tvm import s_tir, tirx
from tvm.ir import IRModule
+from tvm.ir.transform import PassContext, module_pass
from tvm.target import Target
-from tvm.transform import PassContext, module_pass
from .schedule_rule import ScheduleRule
diff --git a/python/tvm/s_tir/dlight/benchmark/bench.py b/python/tvm/s_tir/dlight/benchmark/bench.py
index 8bd5cfc104..8f1fc146c8 100644
--- a/python/tvm/s_tir/dlight/benchmark/bench.py
+++ b/python/tvm/s_tir/dlight/benchmark/bench.py
@@ -139,7 +139,7 @@ def benchmark(
),
)(*input_tensors)
else:
- from tvm.testing import rpc_run # pylint: disable=import-outside-toplevel
+ from tvm.s_tir.testing.runner import rpc_run # pylint: disable=import-outside-toplevel
_, profile_result = rpc_run(
rt_mod,
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 48edcba1c3..c4fc0577bd 100644
--- a/python/tvm/s_tir/meta_schedule/builder/local_builder.py
+++ b/python/tvm/s_tir/meta_schedule/builder/local_builder.py
@@ -25,9 +25,9 @@ from typing import Optional, Union
from tvm_ffi import register_global_func
from tvm.ir import IRModule
-from tvm.ir.utils import derived_object
from tvm.runtime import Module, Tensor
from tvm.runtime.vm import load_param_dict, save_param_dict
+from tvm.s_tir.utils import derived_object
from tvm.support.popen_pool import MapResult, PopenPoolExecutor, StatusKind
from tvm.target import Target
diff --git a/python/tvm/s_tir/meta_schedule/cost_model/mlp_model.py b/python/tvm/s_tir/meta_schedule/cost_model/mlp_model.py
index 52bc9ce99d..e2a4826a4a 100644
--- a/python/tvm/s_tir/meta_schedule/cost_model/mlp_model.py
+++ b/python/tvm/s_tir/meta_schedule/cost_model/mlp_model.py
@@ -31,7 +31,7 @@ from typing import NamedTuple
import numpy as np # type: ignore
import torch # type: ignore
-from tvm.ir.utils import derived_object
+from tvm.s_tir.utils import derived_object
from tvm.support.tar import tar, untar
from ....runtime import Tensor
diff --git a/python/tvm/s_tir/meta_schedule/cost_model/random_model.py b/python/tvm/s_tir/meta_schedule/cost_model/random_model.py
index 86df91d58d..a3891de31e 100644
--- a/python/tvm/s_tir/meta_schedule/cost_model/random_model.py
+++ b/python/tvm/s_tir/meta_schedule/cost_model/random_model.py
@@ -18,7 +18,7 @@
Random cost model
"""
-from tvm.ir.utils import derived_object
+from tvm.s_tir.utils import derived_object
from ..cost_model import PyCostModel
from ..runner import RunnerResult
diff --git a/python/tvm/s_tir/meta_schedule/cost_model/xgb_model.py b/python/tvm/s_tir/meta_schedule/cost_model/xgb_model.py
index 8d6aa49b10..a071f6ce85 100644
--- a/python/tvm/s_tir/meta_schedule/cost_model/xgb_model.py
+++ b/python/tvm/s_tir/meta_schedule/cost_model/xgb_model.py
@@ -25,7 +25,7 @@ from typing import TYPE_CHECKING, Any, Literal, NamedTuple, Optional
import numpy as np # type: ignore
-from tvm.ir.utils import derived_object
+from tvm.s_tir.utils import derived_object
from tvm.support.tar import tar, untar
from ....runtime import Tensor
diff --git a/python/tvm/s_tir/meta_schedule/feature_extractor/random_feature_extractor.py b/python/tvm/s_tir/meta_schedule/feature_extractor/random_feature_extractor.py
index fc42b36604..07a478f523 100644
--- a/python/tvm/s_tir/meta_schedule/feature_extractor/random_feature_extractor.py
+++ b/python/tvm/s_tir/meta_schedule/feature_extractor/random_feature_extractor.py
@@ -19,7 +19,7 @@
import numpy as np # type: ignore
import tvm.runtime
-from tvm.ir.utils import derived_object
+from tvm.s_tir.utils import derived_object
from ..feature_extractor import PyFeatureExtractor
from ..search_strategy import MeasureCandidate
diff --git a/python/tvm/s_tir/meta_schedule/relax_integration.py b/python/tvm/s_tir/meta_schedule/relax_integration.py
index dddab539ce..a8d6d0d76b 100644
--- a/python/tvm/s_tir/meta_schedule/relax_integration.py
+++ b/python/tvm/s_tir/meta_schedule/relax_integration.py
@@ -29,10 +29,10 @@ from typing import Literal
from tvm_ffi import get_global_func, register_global_func
from tvm.ir import IRModule
+from tvm.ir.transform import PassContext
from tvm.runtime import Tensor
from tvm.target import Target
from tvm.tirx.expr import IntImm
-from tvm.transform import PassContext
from .builder import Builder
from .cost_model import CostModel
@@ -465,8 +465,8 @@ def compile_relax(
is_gpu_target = relax_pipeline_mod.BackendDispatcher.is_gpu_target(target)
- @tvm.transform.module_pass(opt_level=3)
- def _ms_pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.IRModule:
+ @tvm.ir.transform.module_pass(opt_level=3)
+ def _ms_pipeline(mod: tvm.ir.IRModule, _ctx: tvm.ir.transform.PassContext) -> tvm.ir.IRModule:
fuse_seq = [
*dispatch_passes,
relax.transform.LegalizeOps(enable_warning=enable_warning),
@@ -475,7 +475,7 @@ def compile_relax(
relax.transform.FuseOps(),
relax.transform.FuseTIR(),
]
- mod = tvm.transform.Sequential(fuse_seq)(mod)
+ mod = tvm.ir.transform.Sequential(fuse_seq)(mod)
mod = MetaScheduleApplyDatabase(enable_warning=enable_warning)(mod)
# DLight handles functions not covered by the database.
# GPU rules apply only for GPU targets.
@@ -487,7 +487,7 @@ def compile_relax(
dl.gpu.GeneralReduction(),
dl.gpu.Fallback(),
)(mod)
- mod = tvm.transform.Sequential(lower_passes + finalize_passes)(mod)
+ mod = tvm.ir.transform.Sequential(lower_passes + finalize_passes)(mod)
return mod
with target, database, PassContext(opt_level=3):
diff --git a/python/tvm/s_tir/meta_schedule/runner/local_runner.py b/python/tvm/s_tir/meta_schedule/runner/local_runner.py
index b56cd6613c..11e15b000a 100644
--- a/python/tvm/s_tir/meta_schedule/runner/local_runner.py
+++ b/python/tvm/s_tir/meta_schedule/runner/local_runner.py
@@ -22,7 +22,7 @@ from collections.abc import Callable
from contextlib import contextmanager
import tvm
-from tvm.ir.utils import derived_object
+from tvm.s_tir.utils import derived_object
from tvm.support.popen_pool import PopenPoolExecutor
from ....runtime import Device, Module
diff --git a/python/tvm/s_tir/meta_schedule/runner/rpc_runner.py b/python/tvm/s_tir/meta_schedule/runner/rpc_runner.py
index 898e2ffe7c..ea87c75018 100644
--- a/python/tvm/s_tir/meta_schedule/runner/rpc_runner.py
+++ b/python/tvm/s_tir/meta_schedule/runner/rpc_runner.py
@@ -21,9 +21,9 @@ import os.path as osp
from collections.abc import Callable
from contextlib import contextmanager
-from tvm.ir.utils import derived_object
from tvm.rpc import RPCSession
from tvm.runtime import Device, Module
+from tvm.s_tir.utils import derived_object
from tvm.support.popen_pool import PopenPoolExecutor
from ..logging import get_logger
diff --git a/python/tvm/s_tir/meta_schedule/testing/dummy_object.py b/python/tvm/s_tir/meta_schedule/testing/dummy_object.py
index d3e0d55a93..a5d0643832 100644
--- a/python/tvm/s_tir/meta_schedule/testing/dummy_object.py
+++ b/python/tvm/s_tir/meta_schedule/testing/dummy_object.py
@@ -18,8 +18,8 @@
import random
-from tvm.ir.utils import derived_object
from tvm.s_tir.schedule import Trace
+from tvm.s_tir.utils import derived_object
from ..builder import BuilderInput, BuilderResult, PyBuilder
from ..mutator import PyMutator
diff --git a/python/tvm/s_tir/meta_schedule/testing/tune_te.py b/python/tvm/s_tir/meta_schedule/testing/tune_te.py
index 959155c3a5..75a7bf77db 100644
--- a/python/tvm/s_tir/meta_schedule/testing/tune_te.py
+++ b/python/tvm/s_tir/meta_schedule/testing/tune_te.py
@@ -22,8 +22,8 @@ import logging
import tvm
from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.meta_schedule.testing.te_workload import create_te_workload
+from tvm.s_tir.testing import strtobool
from tvm.support import describe
-from tvm.testing.utils import strtobool
def _parse_args():
diff --git a/python/tvm/s_tir/meta_schedule/testing/validate_database.py b/python/tvm/s_tir/meta_schedule/testing/validate_database.py
index 410779178a..1b2095c3b4 100644
--- a/python/tvm/s_tir/meta_schedule/testing/validate_database.py
+++ b/python/tvm/s_tir/meta_schedule/testing/validate_database.py
@@ -37,9 +37,9 @@ from tvm.s_tir.meta_schedule.testing.tune_utils import generate_input_data
from tvm.s_tir.meta_schedule.utils import remove_build_dir
from tvm.s_tir.schedule import Trace
from tvm.s_tir.tensor_intrin import * # type: ignore # pylint: disable=wildcard-import,unused-wildcard-import
+from tvm.s_tir.testing import strtobool
from tvm.support import describe
from tvm.target import Target
-from tvm.testing.utils import strtobool
DELIMITOR = "\n" + "-" * 30 + "\n"
diff --git a/python/tvm/s_tir/pipeline.py b/python/tvm/s_tir/pipeline.py
index 0e3916a2f7..e64a2823b2 100644
--- a/python/tvm/s_tir/pipeline.py
+++ b/python/tvm/s_tir/pipeline.py
@@ -58,13 +58,13 @@ def _lower_s_tir(*, adreno=False):
s_tir.transform.LowerThreadBinding(),
]
)
- return tvm.transform.Sequential(passes)
+ return tvm.ir.transform.Sequential(passes)
def _lower_s_tir_tail(*, adreno=False):
- @tvm.transform.module_pass(opt_level=0)
+ @tvm.ir.transform.module_pass(opt_level=0)
def _pipeline(mod, _ctx):
- config = tvm.transform.PassContext.current().config
+ config = tvm.ir.transform.PassContext.current().config
passes = []
if adreno:
passes.append(s_tir.backend.adreno.transform.InjectTextureAlloc())
@@ -115,7 +115,7 @@ def _lower_s_tir_tail(*, adreno=False):
passes.append(s_tir.transform.InjectPTXAsyncCopy())
passes.append(s_tir.transform.LowerSynchronization())
passes.append(s_tir.transform.MergeSharedMemoryAllocations())
- return tvm.transform.Sequential(passes)(mod)
+ return tvm.ir.transform.Sequential(passes)(mod)
return _pipeline
@@ -134,7 +134,7 @@ def _submodule(mod, functions):
def _s_tir_pipeline(*, prepare_only=False, target=None, adreno=False):
- @tvm.transform.module_pass(opt_level=0)
+ @tvm.ir.transform.module_pass(opt_level=0)
def _pipeline(mod, _ctx):
scheduled, native = _partition_dialects(mod)
prepared = _submodule(mod, {})
@@ -158,7 +158,7 @@ def _s_tir_pipeline(*, prepare_only=False, target=None, adreno=False):
return lowered
if native:
return tir_pipeline.finalize_tir_pipeline()(lowered)
- return tvm.transform.Sequential(
+ return tvm.ir.transform.Sequential(
[
tirx.transform.AnnotateEntryFunc(),
tirx.transform.SplitHostDevice(),
@@ -225,7 +225,7 @@ def finalize_host_passes(): # pylint: disable=unused-argument
tirx.transform.LowerTVMBuiltin(),
tirx.transform.LowerIntrin(),
]
- return tvm.transform.Sequential(host_pass_list)
+ return tvm.ir.transform.Sequential(host_pass_list)
def finalize_device_passes(): # pylint: disable=unused-argument
@@ -235,7 +235,7 @@ def finalize_device_passes(): # pylint: disable=unused-argument
tirx.transform.StmtSimplify(),
tirx.transform.LowerIntrin(),
]
- return tvm.transform.Sequential(device_pass_list)
+ return tvm.ir.transform.Sequential(device_pass_list)
tir_pipeline.PIPELINE_MAP["s_tir"] = default_s_tir_pipeline
diff --git a/python/tvm/transform/_ffi_instrument_api.py b/python/tvm/s_tir/testing/__init__.py
similarity index 81%
rename from python/tvm/transform/_ffi_instrument_api.py
rename to python/tvm/s_tir/testing/__init__.py
index 6499a5f0cf..b4896ffe30 100644
--- a/python/tvm/transform/_ffi_instrument_api.py
+++ b/python/tvm/s_tir/testing/__init__.py
@@ -14,8 +14,8 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""FFI APIs for tvm.transform.instrument"""
-import tvm_ffi
+# pylint: disable=invalid-name,unnecessary-comprehension,redefined-outer-name
+"""Testing helpers for schedulable TensorIR."""
-tvm_ffi.init_ffi_api("transform.instrument", __name__)
+from .utils import is_ampere_or_newer, strtobool
diff --git a/python/tvm/testing/runner.py b/python/tvm/s_tir/testing/runner.py
similarity index 100%
rename from python/tvm/testing/runner.py
rename to python/tvm/s_tir/testing/runner.py
diff --git a/python/tvm/testing/tir.py b/python/tvm/s_tir/testing/tir.py
similarity index 100%
rename from python/tvm/testing/tir.py
rename to python/tvm/s_tir/testing/tir.py
diff --git a/python/tvm/s_tir/testing/utils.py b/python/tvm/s_tir/testing/utils.py
new file mode 100644
index 0000000000..426acc1c8f
--- /dev/null
+++ b/python/tvm/s_tir/testing/utils.py
@@ -0,0 +1,43 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+# pylint: disable=invalid-name,unnecessary-comprehension,redefined-outer-name
+"""Small helpers for S-TIR tests and tuning scripts."""
+
+from tvm.support import nvcc
+
+
+def is_ampere_or_newer():
+ """Check if the target environment has an NVIDIA Ampere GPU or newer."""
+ arch = nvcc.get_target_compute_version()
+ major, minor = nvcc.parse_compute_version(arch)
+ return major >= 8 and minor != 9
+
+
+def strtobool(val):
+ """Convert a string representation of truth to true (1) or false (0).
+ True values are 'y', 'yes', 't', 'true', 'on', and '1'; false values
+ are 'n', 'no', 'f', 'false', 'off', and '0'. Raises ValueError if
+ 'val' is anything else.
+ """
+ val = val.lower()
+ if val in ("y", "yes", "t", "true", "on", "1"):
+ return 1
+ elif val in ("n", "no", "f", "false", "off", "0"):
+ return 0
+ else:
+ raise ValueError(f"invalid truth value {val!r}")
diff --git a/python/tvm/s_tir/transform/transform.py b/python/tvm/s_tir/transform/transform.py
index 991ccecde9..6fa217d011 100644
--- a/python/tvm/s_tir/transform/transform.py
+++ b/python/tvm/s_tir/transform/transform.py
@@ -37,7 +37,7 @@ def CanonicalizeLoop():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.CanonicalizeLoop() # type: ignore
@@ -49,7 +49,7 @@ def LowerCrossThreadReduction():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.LowerCrossThreadReduction() # type: ignore
@@ -60,7 +60,7 @@ def LowerInitBlock():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.LowerInitBlock() # type: ignore
@@ -73,7 +73,7 @@ def PlanAndUpdateTensorAllocationLocation():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.PlanAndUpdateTensorAllocationLocation() # type: ignore
@@ -86,7 +86,7 @@ def ConvertBlocksToOpaque():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.ConvertBlocksToOpaque() # type: ignore
@@ -97,7 +97,7 @@ def LiftThreadBinding():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.LiftThreadBinding() # type: ignore
@@ -116,7 +116,7 @@ def CompactTensorAllocation(is_strict: bool = True):
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.CompactTensorAllocation(is_strict) # type: ignore
@@ -127,7 +127,7 @@ def LowerMatchTensor():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.LowerMatchTensor() # type: ignore
@@ -140,7 +140,7 @@ def LowerOpaqueBlock():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.LowerOpaqueBlock() # type: ignore
@@ -151,7 +151,7 @@ def TransformMmaTensorLayout():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.TransformMmaTensorLayout() # type: ignore
@@ -162,7 +162,7 @@ def InjectPermutedLayout():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.InjectPermutedLayout() # type: ignore
@@ -176,7 +176,7 @@ def LowerThreadBinding():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.LowerThreadBinding() # type: ignore
@@ -187,7 +187,7 @@ def InjectSoftwarePipeline():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.InjectSoftwarePipeline() # type: ignore
@@ -198,7 +198,7 @@ def LowerAutoCopy():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.LowerAutoCopy() # type: ignore
@@ -209,7 +209,7 @@ def ManifestSharedMemoryLocalStage():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.ManifestSharedMemoryLocalStage() # type: ignore
@@ -220,7 +220,7 @@ def AnnotateIrregularLoop():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.AnnotateIrregularLoop() # type: ignore
@@ -240,7 +240,7 @@ def LoopPartition():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.LoopPartition() # type: ignore
@@ -251,7 +251,7 @@ def InjectVirtualThread():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.InjectVirtualThread() # type: ignore
@@ -262,7 +262,7 @@ def RewriteUnsafeSelect():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.RewriteUnsafeSelect() # type: ignore
@@ -280,7 +280,7 @@ def VerifyVTCMLimit(default_target=None):
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.VerifyVTCMLimit(default_target) # type: ignore
@@ -291,7 +291,7 @@ def LowerVtcmAlloc():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.LowerVtcmAlloc() # type: ignore
@@ -307,7 +307,7 @@ def ThreadSync(storage_scope):
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.ThreadSync(storage_scope) # type: ignore
@@ -323,7 +323,7 @@ def InferFragment():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.InferFragment() # type: ignore
@@ -334,7 +334,7 @@ def LowerThreadAllreduce():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.LowerThreadAllreduce() # type: ignore
@@ -345,7 +345,7 @@ def InjectPTXAsyncCopy():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.InjectPTXAsyncCopy() # type: ignore
@@ -357,7 +357,7 @@ def MergeSharedMemoryAllocations():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.MergeSharedMemoryAllocations() # type: ignore
@@ -368,7 +368,7 @@ def DefaultGPUSchedule():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.DefaultGPUSchedule() # type: ignore
@@ -384,7 +384,7 @@ def RemoveWeightLayoutRewriteBlock(skip_tensor_rewrite=False):
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.RemoveWeightLayoutRewriteBlock(skip_tensor_rewrite) # type: ignore
@@ -395,7 +395,7 @@ def RemoveStoreUndef():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.RemoveStoreUndef() # type: ignore
diff --git a/python/tvm/ir/utils.py b/python/tvm/s_tir/utils.py
similarity index 99%
rename from python/tvm/ir/utils.py
rename to python/tvm/s_tir/utils.py
index e5d012220f..de7d66b234 100644
--- a/python/tvm/ir/utils.py
+++ b/python/tvm/s_tir/utils.py
@@ -14,7 +14,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Utilities shared across TVM IR packages."""
+"""Utilities for schedulable TensorIR objects."""
from typing import TypeVar
diff --git a/python/tvm/support/popen_pool.py b/python/tvm/support/popen_pool.py
index 8e7371f13c..60573a9874 100644
--- a/python/tvm/support/popen_pool.py
+++ b/python/tvm/support/popen_pool.py
@@ -166,7 +166,7 @@ class PopenWorker:
main_read, worker_write = os.pipe()
worker_read, main_write = os.pipe()
- cmd = [sys.executable, "-m", "tvm.exec.popen_worker"]
+ cmd = [sys.executable, "-m", "tvm.support.popen_worker"]
if sys.platform == "win32":
# pylint: disable=import-outside-toplevel
import msvcrt
diff --git a/python/tvm/exec/popen_worker.py b/python/tvm/support/popen_worker.py
similarity index 100%
rename from python/tvm/exec/popen_worker.py
rename to python/tvm/support/popen_worker.py
diff --git a/python/tvm/testing/__init__.py b/python/tvm/testing/__init__.py
index 91fa954643..11936b2e6d 100644
--- a/python/tvm/testing/__init__.py
+++ b/python/tvm/testing/__init__.py
@@ -24,7 +24,6 @@ from ._ffi_api import (
object_use_count,
test_raise_error,
)
-from .runner import local_run, rpc_run
from .locking import run_with_gpu_lock
from .utils import *
from . import env
diff --git a/python/tvm/testing/env.py b/python/tvm/testing/env.py
index b3375f18d2..19cc7af219 100644
--- a/python/tvm/testing/env.py
+++ b/python/tvm/testing/env.py
@@ -113,6 +113,56 @@ def build_flag_enabled(flag: str) -> bool:
return False
+DEFAULT_TEST_TARGETS = [
+ "llvm",
+ "cuda",
+ "nvptx",
+ {"kind": "vulkan", "from_device": 0},
+ "opencl",
+ {"kind": "opencl", "device": "mali"},
+ {"kind": "opencl", "device": "intel_graphics"},
+ "metal",
+ "rocm",
+ "hexagon",
+]
+
+
+def _parse_target_entry(entry):
+ """Parse a target entry from TVM_TEST_TARGETS env var.
+
+ Entries can be plain kind names (e.g. "llvm") or JSON dicts
+ (e.g. '{"kind": "opencl", "device": "mali"}').
+ """
+ entry = entry.strip()
+ if entry.startswith("{"):
+ import json # pylint: disable=import-outside-toplevel
+
+ return json.loads(entry)
+ return entry
+
+
+def _tvm_test_targets():
+ target_str = os.environ.get("TVM_TEST_TARGETS", "").strip()
+ if target_str:
+ # De-duplicate while preserving order. dict items can't be hashed
+ # directly, so use their str() form as the dedup key.
+ targets = []
+ seen = set()
+ for t in target_str.split(";"):
+ t = t.strip()
+ if not t:
+ continue
+ parsed = _parse_target_entry(t)
+ key = str(parsed)
+ if key in seen:
+ continue
+ seen.add(key)
+ targets.append(parsed)
+ return targets
+
+ return DEFAULT_TEST_TARGETS
+
+
@functools.cache
def _target_enabled(kind: str) -> bool:
"""True if ``kind`` is selected by ``TVM_TEST_TARGETS`` (or the default set).
@@ -122,8 +172,6 @@ def _target_enabled(kind: str) -> bool:
even when a device is physically present.
"""
try:
- from tvm.testing.utils import _tvm_test_targets # pylint: disable=import-outside-toplevel
-
for target in _tvm_test_targets():
k = target["kind"] if isinstance(target, dict) else str(target).split()[0]
if k == kind:
@@ -396,11 +444,11 @@ def has_cudagraph() -> bool:
def has_hexagon_toolchain() -> bool:
"""True if the Hexagon toolchain is available for compilation."""
try:
- from tvm.contrib.hexagon import ( # pylint: disable=import-outside-toplevel
- _ci_env_check,
+ return (
+ build_flag_enabled("USE_HEXAGON")
+ and (not _runtime_enabled("llvm") or _llvm_version_major() >= 7)
+ and "HEXAGON_TOOLCHAIN" in os.environ
)
-
- return build_flag_enabled("USE_HEXAGON") and _ci_env_check._compile_time_check() is True
except Exception: # pylint: disable=broad-except
return False
@@ -409,11 +457,7 @@ def has_hexagon_toolchain() -> bool:
def has_hexagon() -> bool:
"""True if Hexagon can both compile and run (toolchain + attached device)."""
try:
- from tvm.contrib.hexagon import ( # pylint: disable=import-outside-toplevel
- _ci_env_check,
- )
-
- return has_hexagon_toolchain() and _ci_env_check._run_time_check() is True
+ return has_hexagon_toolchain() and "ANDROID_SERIAL_NUMBER" in os.environ
except Exception: # pylint: disable=broad-except
return False
diff --git a/python/tvm/testing/utils.py b/python/tvm/testing/utils.py
index 64186cb539..ea2c3f3ddd 100644
--- a/python/tvm/testing/utils.py
+++ b/python/tvm/testing/utils.py
@@ -14,7 +14,6 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# ruff: noqa: E501
# pylint: disable=invalid-name,unnecessary-comprehension,redefined-outer-name
"""TVM testing utilities
@@ -64,50 +63,20 @@ variable, so the CI can run different targets on different testing nodes.
"""
-import copy
-import copyreg
-import ctypes
-import functools
import inspect
-import logging
import os
-import pickle
-import platform
-import runpy
import sys
import time
-from pathlib import Path
import ml_dtypes
import numpy as np
import pytest
-import tvm_ffi
import tvm
-import tvm.support.utils
-import tvm.sym
-import tvm.te
-import tvm.tirx
-from tvm.contrib import cudnn
-from tvm.support import nvcc
-
-SKIP_SLOW_TESTS = os.getenv("SKIP_SLOW_TESTS", "").lower() in {"true", "1", "yes"}
-IS_IN_CI = os.getenv("CI", "") == "true"
-_REQUEST_HOOK_INITIALIZERS = {}
-
-
-def _substitute_tir_vars(value, var_map):
- return tvm_ffi.structural_map(
- value,
- (tvm.ir.Var, lambda var: var_map.get(var, var)),
- order="post",
- )
+from . import env
-skip_if_wheel_test = pytest.mark.skipif(
- os.getenv("WHEEL_TEST", "").lower() in {"true", "1", "yes"},
- reason="Test not supported in wheel.",
-)
+SKIP_SLOW_TESTS = os.getenv("SKIP_SLOW_TESTS", "").lower() in {"true", "1", "yes"}
def assert_allclose(actual, desired, rtol=1e-7, atol=1e-7, verbose=True):
@@ -124,293 +93,6 @@ def assert_allclose(actual, desired, rtol=1e-7, atol=1e-7, verbose=True):
np.testing.assert_allclose(actual, desired, rtol=rtol, atol=atol, verbose=verbose)
-def check_numerical_grads(
- function, input_values, grad_values, function_value=None, delta=1e-3, atol=1e-2, rtol=0.1
-):
- """A helper function that checks that numerical gradients of a function are
- equal to gradients computed in some different way (analytical gradients).
-
- Numerical gradients are computed using finite difference approximation. To
- reduce the number of function evaluations, the number of points used is
- gradually increased if the error value is too high (up to 5 points).
-
- Parameters
- ----------
- function
- A function that takes inputs either as positional or as keyword
- arguments (either `function(*input_values)` or `function(**input_values)`
- should be correct) and returns a scalar result. Should accept numpy
- ndarrays.
-
- input_values : Dict[str, numpy.ndarray] or List[numpy.ndarray]
- A list of values or a dict assigning values to variables. Represents the
- point at which gradients should be computed.
-
- grad_values : Dict[str, numpy.ndarray] or List[numpy.ndarray]
- Gradients computed using a different method.
-
- function_value : float, optional
- Should be equal to `function(**input_values)`.
-
- delta : float, optional
- A small number used for numerical computation of partial derivatives.
- The default 1e-3 is a good choice for float32.
-
- atol : float, optional
- Absolute tolerance. Gets multiplied by `sqrt(n)` where n is the size of a
- gradient.
-
- rtol : float, optional
- Relative tolerance.
- """
- # If input_values is a list then function accepts positional arguments
- # In this case transform it to a function taking kwargs of the form {"0": ..., "1": ...}
- if not isinstance(input_values, dict):
- input_len = len(input_values)
- input_values = {str(idx): val for idx, val in enumerate(input_values)}
-
- def _function(_input_len=input_len, _orig_function=function, **kwargs):
- return _orig_function(*(kwargs[str(i)] for i in range(input_len)))
-
- function = _function
-
- grad_values = {str(idx): val for idx, val in enumerate(grad_values)}
-
- if function_value is None:
- function_value = function(**input_values)
-
- # a helper to modify j-th element of val by a_delta
- def modify(val, j, a_delta):
- val = val.copy()
- val.reshape(-1)[j] = val.reshape(-1)[j] + a_delta
- return val
-
- # numerically compute a partial derivative with respect to j-th element of the var `name`
- def derivative(x_name, j, a_delta):
- modified_values = {
- n: modify(val, j, a_delta) if n == x_name else val for n, val in input_values.items()
- }
- return (function(**modified_values) - function_value) / a_delta
-
- def compare_derivative(j, n_der, grad):
- der = grad.reshape(-1)[j]
- return np.abs(n_der - der) < atol + rtol * np.abs(n_der)
-
- for x_name, grad in grad_values.items():
- if grad.shape != input_values[x_name].shape:
- raise AssertionError(
- f"Gradient wrt '{x_name}' has unexpected shape {grad.shape}, expected {input_values[x_name].shape} "
- )
-
- ngrad = np.zeros_like(grad)
-
- wrong_positions = []
-
- # compute partial derivatives for each position in this variable
- for j in range(np.prod(grad.shape)):
- # forward difference approximation
- nder = derivative(x_name, j, delta)
-
- # if the derivative is not equal to the analytical one, try to use more
- # precise and expensive methods
- if not compare_derivative(j, nder, grad):
- # central difference approximation
- nder = (derivative(x_name, j, -delta) + nder) / 2
-
- if not compare_derivative(j, nder, grad):
- # central difference approximation using h = delta/2
- cnder2 = (
- derivative(x_name, j, delta / 2) + derivative(x_name, j, -delta / 2)
- ) / 2
- # five-point derivative
- nder = (4 * cnder2 - nder) / 3
-
- # if the derivatives still don't match, add this position to the
- # list of wrong positions
- if not compare_derivative(j, nder, grad):
- wrong_positions.append(np.unravel_index(j, grad.shape))
-
- ngrad.reshape(-1)[j] = nder
-
- wrong_percentage = int(100 * len(wrong_positions) / np.prod(grad.shape))
-
- dist = np.sqrt(np.sum((ngrad - grad) ** 2))
- grad_norm = np.sqrt(np.sum(ngrad**2))
-
- if not (np.isfinite(dist) and np.isfinite(grad_norm)):
- raise ValueError(
- f"NaN or infinity detected during numerical gradient checking wrt '{x_name}'\n"
- f"analytical grad = {grad}\n numerical grad = {ngrad}\n"
- )
-
- # we multiply atol by this number to make it more universal for different sizes
- sqrt_n = np.sqrt(float(np.prod(grad.shape)))
-
- if dist > atol * sqrt_n + rtol * grad_norm:
- raise AssertionError(
- f"Analytical and numerical grads wrt '{x_name}' differ too much\n"
- f"analytical grad = {grad}\n numerical grad = {ngrad}\n"
- f"{wrong_percentage}% of elements differ, first 10 of wrong positions: {wrong_positions[:10]}\n"
- "distance > atol*sqrt(n) + rtol*grad_norm\n"
- f"distance {dist} > {atol}*{sqrt_n} + {rtol}*{grad_norm}"
- )
-
- max_diff = np.max(np.abs(ngrad - grad))
- avg_diff = np.mean(np.abs(ngrad - grad))
- logging.info(
- "Numerical grad test wrt '%s' of shape %s passes, "
- "dist = %f, max_diff = %f, avg_diff = %f",
- x_name,
- grad.shape,
- dist,
- max_diff,
- avg_diff,
- )
-
-
-def assert_prim_expr_equal(lhs, rhs):
- """Assert lhs and rhs equals to each iother.
-
- Parameters
- ----------
- lhs : tvm.tirx.Expr
- The left operand.
-
- rhs : tvm.tirx.Expr
- The left operand.
- """
- ana = tvm.sym.Analyzer()
- if not ana.can_prove_equal(lhs, rhs):
- raise ValueError(f"{lhs} and {rhs} are not equal")
-
-
-def check_bool_expr_is_true(bool_expr, vranges, cond=None):
- """Check that bool_expr holds given the condition cond
- for every value of free variables from vranges.
-
- For example, ``2x > 4y`` solves to ``x > 2y`` given ``x in (0, 10)``
- and ``y in (0, 10)``. Here bool_expr is ``x > 2y``,
- vranges is ``{x: (0, 10), y: (0, 10)}``, cond is ``2x > 4y``.
- We create iterations to check::
-
- for x in range(10):
- for y in range(10):
- assert !(2x > 4y) || (x > 2y)
-
- Parameters
- ----------
- bool_expr : tvm.ir.Expr
- Boolean expression to check
- vranges: Dict[tvm.tirx.expr.Var, tvm.ir.Range]
- Free variables and their ranges
- cond: tvm.ir.Expr
- extra conditions needs to be satisfied.
- """
- if cond is not None:
- bool_expr = tvm.te.any(tvm.tirx.Not(cond), bool_expr)
-
- def _run_expr(expr, vranges):
- """Evaluate expr for every value of free variables
- given by vranges and return the tensor of results.
- """
-
- def _compute_body(*us):
- vmap = {v: u + r.min for (v, r), u in zip(vranges.items(), us)}
- return _substitute_tir_vars(expr, vmap)
-
- A = tvm.te.compute([r.extent.value for v, r in vranges.items()], _compute_body)
- args = [tvm.runtime.empty(A.shape, A.dtype)]
- mod = tvm.compile(tvm.IRModule.from_expr(tvm.te.create_function([A])))
- mod(*args)
- return args[0].numpy()
-
- res = _run_expr(bool_expr, vranges)
- if not np.all(res):
- indices = list(np.argwhere(res == 0)[0])
- counterex = [(str(v), i + r.min) for (v, r), i in zip(vranges.items(), indices)]
- counterex = sorted(counterex, key=lambda x: x[0])
- counterex = ", ".join([v + " = " + str(i) for v, i in counterex])
- ana = tvm.sym.Analyzer()
- raise AssertionError(
- f"Expression {ana.simplify(bool_expr)}\nis not true on {vranges}\n"
- f"Counterexample: {counterex}"
- )
-
-
-def _get_targets(target_names=None):
- if target_names is None:
- target_names = _tvm_test_targets()
-
- if not target_names:
- target_names = DEFAULT_TEST_TARGETS
-
- targets = []
- for target in target_names:
- if isinstance(target, dict):
- target_kind = target["kind"]
- else:
- target_kind = target.split()[0]
-
- if target_kind == "cuda" and "cudnn" in tvm.target.Target(target).attrs.get("libs", []):
- is_enabled = tvm.support.libinfo().get("USE_CUDNN", "OFF").lower() in [
- "on",
- "true",
- "1",
- ]
- is_runnable = is_enabled and cudnn.exists()
- elif target_kind == "hexagon":
- is_enabled = tvm.support.libinfo().get("USE_HEXAGON", "OFF").lower() in [
- "on",
- "true",
- "1",
- ]
- # If Hexagon has compile-time support, we can always fall back
- is_runnable = is_enabled and "ANDROID_SERIAL_NUMBER" in os.environ
- else:
- is_enabled = tvm.runtime.enabled(target_kind)
- is_runnable = is_enabled and tvm.device_from_target(target).exist
-
- targets.append(
- {
- "target": target,
- "target_kind": target_kind,
- "is_enabled": is_enabled,
- "is_runnable": is_runnable,
- }
- )
-
- if all(not t["is_runnable"] for t in targets):
- if tvm.runtime.enabled("llvm"):
- logging.warning(
- "None of the following targets are supported by this build of TVM: %s."
- " Try setting TVM_TEST_TARGETS to a supported target. Defaulting to llvm.",
- target_names,
- )
- return _get_targets(["llvm"])
-
- raise RuntimeError(
- "None of the following targets are supported by this build of TVM: %s."
- " Try setting TVM_TEST_TARGETS to a supported target."
- " Cannot default to llvm, as it is not enabled." % target_names
- )
-
- return targets
-
-
-DEFAULT_TEST_TARGETS = [
- "llvm",
- "cuda",
- "nvptx",
- {"kind": "vulkan", "from_device": 0},
- "opencl",
- {"kind": "opencl", "device": "mali"},
- {"kind": "opencl", "device": "intel_graphics"},
- "metal",
- "rocm",
- "hexagon",
-]
-
-
def device_enabled(target):
"""Check if a target should be used when testing.
@@ -450,81 +132,23 @@ def device_enabled(target):
assert isinstance(target, str), "device_enabled requires a target as a string"
# Target strings may include extra flags; only compare the kind.
target_kind = target.split(" ")[0]
- return any(target_kind == t["target_kind"] for t in _get_targets() if t["is_runnable"])
-
-
-def enabled_targets():
- """Get all enabled targets with associated devices.
-
- In most cases, parametrize over the specific targets you need with
- ``@pytest.mark.parametrize`` instead of iterating this function.
-
- In this context, enabled means that TVM was built with support for
- this target, the target name appears in the TVM_TEST_TARGETS
- environment variable, and a suitable device for running this
- target exists. If TVM_TEST_TARGETS is not set, it defaults to
- variable DEFAULT_TEST_TARGETS in this module.
-
- If you use this function in a test, you **must** mark the test with
- ``@pytest.mark.gpu`` (otherwise it will never be run on the gpu).
-
- Returns
- -------
- targets: list
- A list of pairs of all enabled devices and the associated context
-
- """
- return [
- (t["target"], tvm.device_from_target(t["target"]))
- for t in _get_targets()
- if t["is_runnable"]
- ]
-
-
-def _parse_target_entry(entry):
- """Parse a target entry from TVM_TEST_TARGETS env var.
-
- Entries can be plain kind names (e.g. "llvm") or JSON dicts
- (e.g. '{"kind": "opencl", "device": "mali"}').
- """
- entry = entry.strip()
- if entry.startswith("{"):
- import json # pylint: disable=import-outside-toplevel
-
- return json.loads(entry)
- return entry
-
-
-def _tvm_test_targets():
- target_str = os.environ.get("TVM_TEST_TARGETS", "").strip()
- if target_str:
- # De-duplicate while preserving order. dict items can't be hashed
- # directly, so use their str() form as the dedup key.
- targets = []
- seen = set()
- for t in target_str.split(";"):
- t = t.strip()
- if not t:
- continue
- parsed = _parse_target_entry(t)
- key = str(parsed)
- if key in seen:
- continue
- seen.add(key)
- targets.append(parsed)
- return targets
-
- return DEFAULT_TEST_TARGETS
-
-
-def _compose(args, decs):
- """Helper to apply multiple markers"""
- if len(args) > 0:
- f = args[0]
- for d in reversed(decs):
- f = d(f)
- return f
- return decs
+ # Check only explicitly selected candidates of the requested kind. In
+ # particular, an unavailable selection must not silently enable LLVM.
+ for candidate in env._tvm_test_targets():
+ kind = candidate["kind"] if isinstance(candidate, dict) else candidate.split()[0]
+ if kind != target_kind:
+ continue
+ if kind == "hexagon":
+ return env.build_flag_enabled("USE_HEXAGON") and "ANDROID_SERIAL_NUMBER" in os.environ
+ if not env._runtime_enabled(kind):
+ continue
+ if kind == "cuda" and "cudnn" in tvm.target.Target(candidate).attrs.get("libs", []):
+ exists = tvm.get_global_func("tvm.contrib.cudnn.exists", allow_missing=True)
+ if env.has_cudnn() and exists is not None and exists():
+ return True
+ elif tvm.device_from_target(candidate).exist:
+ return True
+ return False
slow = pytest.mark.skipif(
@@ -533,16 +157,6 @@ slow = pytest.mark.skipif(
)
-def skip_if_32bit(reason):
- def decorator(*args):
- if "32bit" in platform.architecture()[0]:
- return _compose(args, [pytest.mark.skip(reason=reason)])
-
- return _compose(args, [])
-
- return decorator
-
-
def parameter(*values, ids=None, by_dict=None):
"""Convenience function to define pytest parametrized fixtures.
@@ -550,7 +164,7 @@ def parameter(*values, ids=None, by_dict=None):
parametrized pytest fixture that can be used by test
functions. This is intended for cases that have no setup cost,
such as strings, integers, tuples, etc. For cases that have a
- significant setup cost, please use :py:func:`tvm.testing.fixture`
+ significant setup cost, use an ordinary pytest fixture
instead.
If a test function accepts multiple parameters defined using
@@ -622,172 +236,6 @@ def parameter(*values, ids=None, by_dict=None):
return as_fixture
-def fixture(func=None, *, cache_return_value=False):
- """Convenience function to define pytest fixtures.
-
- This should be used as a decorator to mark functions that set up
- state before a function. The return value of that fixture
- function is then accessible by test functions as that accept it as
- a parameter.
-
- Fixture functions can accept parameters defined with
- :py:func:`tvm.testing.parameter`.
-
- By default, the setup will be performed once for each unit test
- that uses a fixture, to ensure that unit tests are independent.
- If the setup is expensive to perform, then the
- cache_return_value=True argument can be passed to cache the setup.
- The fixture function will be run only once (or once per parameter,
- if used with tvm.testing.parameter). The cached setup value is
- retained for the lifetime of the test process, and each test receives
- an independent copy. If the environment variable TVM_TEST_DISABLE_CACHE
- is set to a non-zero value, it will disable this feature and no caching
- will be performed.
-
- Example
- -------
- >>> @tvm.testing.fixture
- >>> def cheap_setup():
- >>> return 5 # Setup code here.
- >>>
- >>> def test_feature_x(target, dev, cheap_setup)
- >>> assert(cheap_setup == 5) # Run test here
-
- Or
-
- >>> size = tvm.testing.parameter(1, 10, 100)
- >>>
- >>> @tvm.testing.fixture
- >>> def cheap_setup(size):
- >>> return 5*size # Setup code here, based on size.
- >>>
- >>> def test_feature_x(cheap_setup):
- >>> assert(cheap_setup in [5, 50, 500])
-
- Or
-
- >>> @tvm.testing.fixture(cache_return_value=True)
- >>> def expensive_setup():
- >>> time.sleep(10) # Setup code here
- >>> return 5
- >>>
- >>> def test_feature_x(target, dev, expensive_setup):
- >>> assert(expensive_setup == 5)
-
- """
-
- force_disable_cache = bool(int(os.environ.get("TVM_TEST_DISABLE_CACHE", "0")))
- cache_return_value = cache_return_value and not force_disable_cache
-
- def wraps(func):
- if cache_return_value:
- func = _fixture_cache(func)
- func = pytest.fixture(func, scope="function")
- return func
-
- if func is None:
- return wraps
-
- return wraps(func)
-
-
-class _DeepCopyAllowedClasses(dict):
- def __init__(self, allowed_class_list):
- self.allowed_class_list = allowed_class_list
- super().__init__()
-
- def get(self, key, *args, **kwargs):
- """Overrides behavior of copy.deepcopy to avoid implicit copy.
-
- By default, copy.deepcopy uses a dict of id->object to track
- all objects that it has seen, which is passed as the second
- argument to all recursive calls. This class is intended to be
- passed in instead, and inspects the type of all objects being
- copied.
-
- Where copy.deepcopy does a best-effort attempt at copying an
- object, for unit tests we would rather have all objects either
- be copied correctly, or to throw an error. Classes that
- define an explicit method to perform a copy are allowed, as
- are any explicitly listed classes. Classes that would fall
- back to using object.__reduce__, and are not explicitly listed
- as safe, will throw an exception.
-
- """
- obj = ctypes.cast(key, ctypes.py_object).value
- cls = type(obj)
- if (
- cls in copy._deepcopy_dispatch
- or issubclass(cls, type)
- or getattr(obj, "__deepcopy__", None)
- or copyreg.dispatch_table.get(cls)
- or cls.__reduce__ is not object.__reduce__
- or cls.__reduce_ex__ is not object.__reduce_ex__
- or cls in self.allowed_class_list
- ):
- return super().get(key, *args, **kwargs)
-
- rfc_url = (
- "https://github.com/apache/tvm-rfcs/blob/main/rfcs/0007-parametrized-unit-tests.md"
- )
- raise TypeError(
- f"Cannot copy fixture of type {cls.__name__}. TVM fixture caching "
- "is limited to objects that explicitly provide the ability "
- "to be copied (e.g. through __deepcopy__, __getstate__, or __setstate__),"
- "and forbids the use of the default `object.__reduce__` and "
- "`object.__reduce_ex__`. For third-party classes that are "
- "safe to use with copy.deepcopy, please add the class to "
- "the arguments of _DeepCopyAllowedClasses in tvm.testing._fixture_cache.\n"
- "\n"
- f"For discussion on this restriction, please see {rfc_url}."
- )
-
-
-def _fixture_cache(func):
- cache = {}
-
- # Using functools.lru_cache would require the function arguments
- # to be hashable, which wouldn't allow caching fixtures that
- # depend on numpy arrays. For example, a fixture that takes a
- # numpy array as input, then calculates uses a slow method to
- # compute a known correct output for that input. Therefore,
- # including a fallback for serializable types.
- def get_cache_key(*args, **kwargs):
- try:
- hash((args, kwargs))
- return (args, kwargs)
- except TypeError:
- pass
-
- try:
- return pickle.dumps((args, kwargs))
- except TypeError as e:
- raise TypeError(
- "TVM caching of fixtures requires arguments to the fixture "
- "to be either hashable or serializable"
- ) from e
-
- @functools.wraps(func)
- def wrapper(*args, **kwargs):
- cache_key = get_cache_key(*args, **kwargs)
-
- try:
- cached_value = cache[cache_key]
- except KeyError:
- cached_value = cache[cache_key] = func(*args, **kwargs)
-
- return copy.deepcopy(
- cached_value,
- # allowed_class_list should be a list of classes that
- # are safe to copy using copy.deepcopy, but do not
- # implement __deepcopy__, __reduce__, or
- # __reduce_ex__.
- _DeepCopyAllowedClasses(allowed_class_list=[]),
- )
-
- return wrapper
-
-
def identity_after(x, sleep):
"""Testing function to return identity after sleep
@@ -814,46 +262,6 @@ def terminate_self():
sys.exit(-1)
-def is_ampere_or_newer():
- """Check if the target environment has an NVIDIA Ampere GPU or newer."""
- arch = nvcc.get_target_compute_version()
- major, minor = nvcc.parse_compute_version(arch)
- return major >= 8 and minor != 9
-
-
-def install_request_hook(hook_script: Path) -> None:
- """Add a wrapper around urllib.request for CI tests."""
- if not IS_IN_CI:
- return
-
- hook_script = Path(hook_script).resolve()
- if not hook_script.is_file():
- raise RuntimeError(f"Request hook {hook_script} does not exist")
-
- # Load the exact hook file without exposing the test root as an import path.
- # Cache its initializer because Sphinx invokes this once per gallery example.
- try:
- init = _REQUEST_HOOK_INITIALIZERS[hook_script]
- except KeyError:
- init = _REQUEST_HOOK_INITIALIZERS[hook_script] = runpy.run_path(str(hook_script))["init"]
- init()
-
-
-def strtobool(val):
- """Convert a string representation of truth to true (1) or false (0).
- True values are 'y', 'yes', 't', 'true', 'on', and '1'; false values
- are 'n', 'no', 'f', 'false', 'off', and '0'. Raises ValueError if
- 'val' is anything else.
- """
- val = val.lower()
- if val in ("y", "yes", "t", "true", "on", "1"):
- return 1
- elif val in ("n", "no", "f", "false", "off", "0"):
- return 0
- else:
- raise ValueError(f"invalid truth value {val!r}")
-
-
def main():
test_file = inspect.getsourcefile(sys._getframe(1))
sys.exit(pytest.main([test_file, *sys.argv[1:]]))
diff --git a/python/tvm/tirx/_compile.py b/python/tvm/tirx/_compile.py
index b44ed86fde..1007ec3b7a 100644
--- a/python/tvm/tirx/_compile.py
+++ b/python/tvm/tirx/_compile.py
@@ -131,7 +131,7 @@ def _split_host_device_modules(
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):
+ if tvm.ir.transform.PassContext.current().config.get("tirx.disable_assert", False):
mod = tvm.tirx.transform.SkipAssert()(mod)
build_f_name = "target.build." + target.kind.name
bf = tvm.get_global_func(build_f_name)
@@ -165,7 +165,7 @@ def _codegen_host_and_device(
def compile( # pylint: disable=redefined-builtin
mod: Function | IRModule,
target: str | Target | None = None,
- pipeline: str | tvm.transform.Pass | None = "default",
+ pipeline: str | tvm.ir.transform.Pass | None = "default",
) -> Executable:
"""Compile a function or module to an executable for the target.
@@ -175,7 +175,7 @@ def compile( # pylint: disable=redefined-builtin
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]
+ pipeline : Union[None, str, tvm.ir.transform.Pass]
The pipeline to use for compilation.
Returns
diff --git a/python/tvm/tirx/compilation_pipeline.py b/python/tvm/tirx/compilation_pipeline.py
index 5381159c86..8f63fae2b9 100644
--- a/python/tvm/tirx/compilation_pipeline.py
+++ b/python/tvm/tirx/compilation_pipeline.py
@@ -33,11 +33,11 @@ def default_tir_pipeline(*, target=None):
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:
+ @tvm.ir.transform.module_pass(opt_level=0)
+ def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.ir.transform.PassContext) -> tvm.ir.IRModule:
"""The default lowering passes for TIR backend."""
mod = _bind_target(mod, target)
- pass_ctx = tvm.transform.PassContext.current()
+ pass_ctx = tvm.ir.transform.PassContext.current()
config = pass_ctx.config
passes = [
tirx.transform.LowerTIRx(),
@@ -56,7 +56,7 @@ def tirx_pipeline(*, prepare_only=False, target=None):
tirx.transform.LowerThreadAllreduce(),
]
)
- mod = tvm.transform.Sequential(passes)(mod)
+ mod = tvm.ir.transform.Sequential(passes)(mod)
return mod if prepare_only else finalize_tir_pipeline()(mod)
return _pipeline, finalize_host_passes, finalize_device_passes
@@ -65,7 +65,7 @@ def tirx_pipeline(*, prepare_only=False, target=None):
def _bind_target(mod, target):
if target is None:
return mod
- return tvm.transform.Sequential(
+ return tvm.ir.transform.Sequential(
[
tirx.transform.BindTarget(target),
cuda_transforms.BindBackendConfig(),
@@ -92,7 +92,7 @@ def _with_target_binding(pipeline, target, *, original_mod=None, prepared_mod=No
else ()
)
- @tvm.transform.module_pass(opt_level=0)
+ @tvm.ir.transform.module_pass(opt_level=0)
def _pipeline(mod, _ctx):
if (
prepared_mod is not None
@@ -119,7 +119,7 @@ def _with_target_binding(pipeline, target, *, original_mod=None, prepared_mod=No
def finalize_tir_pipeline():
"""Lower module-wide calling conventions after all dialects are prepared."""
- return tvm.transform.Sequential(
+ return tvm.ir.transform.Sequential(
[
tirx.transform.AnnotateEntryFunc(),
tirx.transform.SplitHostDevice(),
@@ -140,7 +140,7 @@ def finalize_host_passes(): # pylint: disable=unused-argument
tirx.transform.LowerTVMBuiltin(),
tirx.transform.LowerIntrin(),
]
- return tvm.transform.Sequential(host_pass_list)
+ return tvm.ir.transform.Sequential(host_pass_list)
def finalize_device_passes(): # pylint: disable=unused-argument
@@ -150,13 +150,13 @@ def finalize_device_passes(): # pylint: disable=unused-argument
tirx.transform.StmtSimplify(),
tirx.transform.LowerIntrin(),
]
- return tvm.transform.Sequential(device_pass_list)
+ return tvm.ir.transform.Sequential(device_pass_list)
def finalize_device_passes_tirx(): # pylint: disable=unused-argument
"""The TIRx finalization passes for TIR backend."""
device_pass_list = [tirx.transform.LowerIntrin()]
- return tvm.transform.Sequential(device_pass_list)
+ return tvm.ir.transform.Sequential(device_pass_list)
# global map of pre-built pipelines
diff --git a/python/tvm/tirx/transform/function_pass.py b/python/tvm/tirx/transform/function_pass.py
index 582f54757c..ce20f91ebd 100644
--- a/python/tvm/tirx/transform/function_pass.py
+++ b/python/tvm/tirx/transform/function_pass.py
@@ -22,7 +22,7 @@ from collections.abc import Callable
import tvm_ffi
-from tvm.transform import Pass, PassInfo
+from tvm.ir.transform import Pass, PassInfo
from . import _ffi_api
diff --git a/python/tvm/tirx/transform/transform.py b/python/tvm/tirx/transform/transform.py
index 892d945bee..826ceb8bf8 100644
--- a/python/tvm/tirx/transform/transform.py
+++ b/python/tvm/tirx/transform/transform.py
@@ -42,7 +42,7 @@ def Apply(ftransform):
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
@@ -69,7 +69,7 @@ def VectorizeLoop(enable_vectorize: bool = True):
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.VectorizeLoop(enable_vectorize) # type: ignore
@@ -84,7 +84,7 @@ def StorageRewrite():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.StorageRewrite() # type: ignore
@@ -98,7 +98,7 @@ def PointerValueTypeRewrite():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.PointerValueTypeRewrite() # type: ignore
@@ -116,7 +116,7 @@ def UnrollLoop():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.UnrollLoop() # type: ignore
@@ -132,7 +132,7 @@ def RemoveNoOp():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.RemoveNoOp() # type: ignore
@@ -143,7 +143,7 @@ def HoistIf():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass.
"""
return _ffi_api.HoistIf()
@@ -154,7 +154,7 @@ def RemoveAssume():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.RemoveAssume() # type: ignore
@@ -165,7 +165,7 @@ def BF16ComputeLegalize():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.BF16ComputeLegalize() # type: ignore
@@ -181,7 +181,7 @@ def FP8ComputeLegalize(promote_dtype: str = "float32"):
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.FP8ComputeLegalize(promote_dtype) # type: ignore
@@ -192,7 +192,7 @@ def BF16StorageLegalize():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.BF16StorageLegalize() # type: ignore
@@ -203,7 +203,7 @@ def FP8StorageLegalize():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.FP8StorageLegalize() # type: ignore
@@ -214,7 +214,7 @@ def CommonSubexprElim():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.CommonSubexprElim() # type: ignore
@@ -225,7 +225,7 @@ def RenormalizeSplitPattern():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.RenormalizeSplitPattern() # type: ignore
@@ -241,7 +241,7 @@ def StmtSimplify():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.StmtSimplify() # type: ignore
@@ -252,7 +252,7 @@ def LowerThreadAllreduce():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass.
"""
return _ffi_api.LowerThreadAllreduce()
@@ -269,7 +269,7 @@ def ConvertSSA():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
@@ -300,7 +300,7 @@ def MakePackedAPI():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.MakePackedAPI() # type: ignore
@@ -315,7 +315,7 @@ def SplitHostDevice():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.SplitHostDevice() # type: ignore
@@ -326,7 +326,7 @@ def SkipAssert():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.SkipAssert() # type: ignore
@@ -337,7 +337,7 @@ def LowerWarpMemory():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.LowerWarpMemory() # type: ignore
@@ -348,7 +348,7 @@ def LowerTVMBuiltin():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.LowerTVMBuiltin() # type: ignore
@@ -359,7 +359,7 @@ def LowerIntrin():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.LowerIntrin() # type: ignore
@@ -375,7 +375,7 @@ def NarrowDataType(target_bits: int):
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
Note
@@ -392,7 +392,7 @@ def ForceNarrowIndexToInt32():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
Note
@@ -407,7 +407,7 @@ def VerifyMemory():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.VerifyMemory() # type: ignore
@@ -419,7 +419,7 @@ def FlattenTensor():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.FlattenTensor() # type: ignore
@@ -438,7 +438,7 @@ def BindTarget(target):
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.BindTarget(target) # type: ignore
@@ -449,7 +449,7 @@ def AnnotateEntryFunc():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.AnnotateEntryFunc() # type: ignore
@@ -461,7 +461,7 @@ def Filter(fcond: Callable):
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.Filter(fcond) # type: ignore
@@ -472,7 +472,7 @@ def TileDispatch():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.TileDispatch() # type: ignore
@@ -483,7 +483,7 @@ def LowerTIRx():
Returns
-------
- fpass : tvm.transform.Pass
+ fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.LowerTIRx() # type: ignore
diff --git a/tests/nightly/python/test_nnapi/test_network.py b/tests/nightly/python/test_nnapi/test_network.py
index 88d8854801..30320b1d82 100644
--- a/tests/nightly/python/test_nnapi/test_network.py
+++ b/tests/nightly/python/test_nnapi/test_network.py
@@ -35,7 +35,7 @@ def _build_and_run_network(remote_obj, tracker, mod, input_data):
"""Helper function to build and run a network."""
def execute_on_host(mod, inputs):
- with tvm.transform.PassContext(opt_level=3):
+ with tvm.ir.transform.PassContext(opt_level=3):
ex = tvm.compile(mod, target="llvm")
dev = tvm.cpu(0)
vm = tvm.relax.VirtualMachine(ex, device=dev)
diff --git a/tests/nightly/python/test_nnapi/test_ops.py b/tests/nightly/python/test_nnapi/test_ops.py
index 6181add25b..a8cfec2b78 100644
--- a/tests/nightly/python/test_nnapi/test_ops.py
+++ b/tests/nightly/python/test_nnapi/test_ops.py
@@ -30,7 +30,7 @@ def _build_and_run_network(remote_obj, tracker, mod, input_data):
"""Helper function to build and run a network."""
def execute_on_host(mod, inputs):
- with tvm.transform.PassContext(opt_level=3):
+ with tvm.ir.transform.PassContext(opt_level=3):
ex = tvm.compile(mod, target="llvm")
dev = tvm.cpu(0)
vm = tvm.relax.VirtualMachine(ex, device=dev)
diff --git a/tests/python/ci/test_ci.py b/tests/python/ci/test_ci.py
index cc410712f3..c716bf1490 100644
--- a/tests/python/ci/test_ci.py
+++ b/tests/python/ci/test_ci.py
@@ -18,6 +18,7 @@
"""Test various CI scripts and GitHub Actions workflows"""
import json
+import os
import subprocess
from pathlib import Path
@@ -169,7 +170,10 @@ def assert_in(needle: str, haystack: str):
raise AssertionError(f"item not found:\n{needle}\nin:\n{haystack}")
-@tvm.testing.skip_if_wheel_test
+@pytest.mark.skipif(
+ os.getenv("WHEEL_TEST", "").lower() in {"true", "1", "yes"},
+ reason="Test not supported in wheel.",
+)
@parameterize_named(
same_tags=dict(
tlcpackstaging_body={
diff --git a/tests/python/ci/test_tvmbot.py b/tests/python/ci/test_tvmbot.py
index 362a87d4fa..ddd98b2a29 100644
--- a/tests/python/ci/test_tvmbot.py
+++ b/tests/python/ci/test_tvmbot.py
@@ -20,9 +20,12 @@ Test the @tvm-bot merge code
"""
import json
+import os
from pathlib import Path
from typing import Any
+import pytest
+
import tvm
from .test_utils import GITHUB_SCRIPT_ROOT, TempGit, run_script
@@ -47,7 +50,10 @@ class _TvmBotTest:
"""
return data
- @tvm.testing.skip_if_wheel_test
+ @pytest.mark.skipif(
+ os.getenv("WHEEL_TEST", "").lower() in {"true", "1", "yes"},
+ reason="Test not supported in wheel.",
+ )
def test(self, tmpdir_factory):
"""
Run the tvm-bot script using the data from preprocess_data
diff --git a/tests/python/codegen/test_target_codegen_cuda.py b/tests/python/codegen/test_target_codegen_cuda.py
index 8c95e63696..c674527c4c 100644
--- a/tests/python/codegen/test_target_codegen_cuda.py
+++ b/tests/python/codegen/test_target_codegen_cuda.py
@@ -322,7 +322,7 @@ def test_cuda_bf16_vectorize_add(*, 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] + one
- with tvm.transform.PassContext(
+ with tvm.ir.transform.PassContext(
disabled_pass=["tirx.BF16Promote", "tirx.BF16CastElimination", "tirx.BF16TypeLowering"]
):
fun = tvm.tirx.compile(
diff --git a/tests/python/codegen/test_target_codegen_cuda_fastmath.py b/tests/python/codegen/test_target_codegen_cuda_fastmath.py
index 8d8f025948..4416de138f 100644
--- a/tests/python/codegen/test_target_codegen_cuda_fastmath.py
+++ b/tests/python/codegen/test_target_codegen_cuda_fastmath.py
@@ -228,7 +228,7 @@ def check_lowered_ir(
) -> tuple[tvm.target.Target, IRModule]:
"""Check the lowered IR for the given dtype and case."""
target, mod = make_mod(dtype, case, enable_fast_math)
- with tvm.transform.PassContext(config={"tirx.enable_fast_math": enable_fast_math}):
+ with tvm.ir.transform.PassContext(config={"tirx.enable_fast_math": enable_fast_math}):
lowered_mod = tvm.tirx.transform.LowerIntrin()(mod)
script = lowered_mod.script(show_meta=False)
expected = expected_intrinsic(dtype, case, enable_fast_math)
@@ -244,7 +244,7 @@ def check_cuda_source(
enable_fast_math: bool,
) -> Executable:
"""Check the CUDA source for the given dtype and case."""
- with tvm.transform.PassContext(config={"tirx.enable_fast_math": enable_fast_math}):
+ with tvm.ir.transform.PassContext(config={"tirx.enable_fast_math": enable_fast_math}):
executable = tvm.tirx.compile(mod, target=target)
source = executable.mod.imports[0].inspect_source()
expected = expected_intrinsic(dtype, case, enable_fast_math)
diff --git a/tests/python/codegen/test_target_codegen_cuda_fp8.py b/tests/python/codegen/test_target_codegen_cuda_fp8.py
index f306e4c355..46d7169f64 100644
--- a/tests/python/codegen/test_target_codegen_cuda_fp8.py
+++ b/tests/python/codegen/test_target_codegen_cuda_fp8.py
@@ -399,69 +399,27 @@ class TestFP8e4x4QuantDequantScale(BaseFP8E4M3QuantScaleOnly):
# weight_shape = tvm.testing.parameter((32000, 4096), (4096, 14336))
weight_shape = tvm.testing.parameter((128, 256), (128, 64))
- @tvm.testing.fixture
- def group_size(self):
- return 64
-
- @tvm.testing.fixture
- def axis(self):
- return 1
-
- @tvm.testing.fixture
- def model_dtype(self):
- return "float16"
-
- @tvm.testing.fixture
- def storage_dtype(self):
- return "uint32"
-
- @tvm.testing.fixture
- def quantize_dtype(self):
- return "float8_e4m3fn"
-
- @tvm.testing.fixture
- def num_el_per_storage(self):
- return 4
-
- @tvm.testing.fixture
- def max_int_value(self):
- return 448
-
- @tvm.testing.fixture
- def target_str(self):
- return "cuda"
-
- @tvm.testing.fixture
- def scale_shape(self, weight_shape, group_size, axis):
- return [
+ @pytest.mark.gpu
+ @pytest.mark.skipif(not env.has_cuda_compute(8, 9), reason="need cuda compute >= 8.9")
+ def test_main(self, weight_shape):
+ group_size = 64
+ axis = 1
+ model_dtype = "float16"
+ storage_dtype = "uint32"
+ quantize_dtype = "float8_e4m3fn"
+ num_el_per_storage = 4
+ max_int_value = 448
+ target_str = "cuda"
+ scale_shape = [
(d + group_size - 1) // group_size if axis == i else d
for i, d in enumerate(weight_shape)
]
-
- @tvm.testing.fixture
- def quant_weight_shape(self, weight_shape, num_el_per_storage, axis):
- return [
+ quant_weight_shape = [
(d + num_el_per_storage - 1) // num_el_per_storage if axis == i else d
for i, d in enumerate(weight_shape)
]
-
- @tvm.testing.fixture
- def compiled_functions(
- self,
- weight_shape,
- scale_shape,
- quant_weight_shape,
- model_dtype,
- quantize_dtype,
- storage_dtype,
- group_size,
- num_el_per_storage,
- max_int_value,
- axis,
- target_str,
- ):
dev = tvm.cuda(0)
- return self.compile_quant_and_dequant_by_scale(
+ quant, dequant = self.compile_quant_and_dequant_by_scale(
weight_shape,
scale_shape,
quant_weight_shape,
@@ -476,12 +434,6 @@ class TestFP8e4x4QuantDequantScale(BaseFP8E4M3QuantScaleOnly):
dev,
)
- @pytest.mark.gpu
- @pytest.mark.skipif(not env.has_cuda_compute(8, 9), reason="need cuda compute >= 8.9")
- def test_main(self, weight_shape, model_dtype, target_str, compiled_functions):
- quant, dequant = compiled_functions
- dev = tvm.cuda(0)
-
weight_np = np.random.uniform(-100, 100, weight_shape).astype(model_dtype)
weight = tvm.runtime.tensor(weight_np, device=dev)
quant_weight, scales = quant(weight)
diff --git a/tests/python/codegen/test_target_codegen_llvm.py b/tests/python/codegen/test_target_codegen_llvm.py
index fa9577b4df..632bf7bd74 100644
--- a/tests/python/codegen/test_target_codegen_llvm.py
+++ b/tests/python/codegen/test_target_codegen_llvm.py
@@ -259,7 +259,7 @@ def test_llvm_madd_pipeline():
check_llvm(64, 0, 2)
check_llvm(4, 0, 1)
- with tvm.transform.PassContext(config={"tirx.noalias": False}):
+ with tvm.ir.transform.PassContext(config={"tirx.noalias": False}):
check_llvm(4, 0, 3)
@@ -863,7 +863,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}):
+ with tvm.ir.transform.PassContext(config={"tirx.disable_assert": True}):
m = tvm.tirx.compile(Module, target="llvm")
diff --git a/tests/python/conftest.py b/tests/python/conftest.py
deleted file mode 100644
index b3a7dc8378..0000000000
--- a/tests/python/conftest.py
+++ /dev/null
@@ -1,29 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-"""Configure pytest for TVM's Python test suite."""
-
-import os
-from pathlib import Path
-
-
-def pytest_sessionstart():
- if os.getenv("CI", "") == "true":
- from tvm.testing.utils import (
- install_request_hook, # pylint: disable=import-outside-toplevel
- )
-
- install_request_hook(Path(__file__).with_name("request_hook.py"))
diff --git a/tests/python/contrib/test_example_npu.py b/tests/python/contrib/test_example_npu.py
index 217d50d11f..90a5ce52a1 100644
--- a/tests/python/contrib/test_example_npu.py
+++ b/tests/python/contrib/test_example_npu.py
@@ -245,7 +245,7 @@ def test_example_npu_runtime_execution():
# Build the module
target = tvm.target.Target("llvm")
- with tvm.transform.PassContext(opt_level=3):
+ with tvm.ir.transform.PassContext(opt_level=3):
built = relax.build(mod, target)
# Create VM and run
diff --git a/tests/python/contrib/test_random.py b/tests/python/contrib/test_random.py
index 94e6913eae..2055216d59 100644
--- a/tests/python/contrib/test_random.py
+++ b/tests/python/contrib/test_random.py
@@ -104,8 +104,28 @@ def test_normal():
@pytest.mark.gpu
-def test_random_fill():
+@pytest.mark.parametrize(
+ "target",
+ [
+ "llvm",
+ "cuda",
+ "nvptx",
+ {"kind": "vulkan", "from_device": 0},
+ "opencl",
+ {"kind": "opencl", "device": "mali"},
+ {"kind": "opencl", "device": "intel_graphics"},
+ "metal",
+ "rocm",
+ "hexagon",
+ ],
+)
+def test_random_fill(target):
"""Tests random_fill function"""
+ if not tvm.testing.device_enabled(target):
+ pytest.skip("target is unavailable or excluded by TVM_TEST_TARGETS")
+ if not tvm.get_global_func("tvm.contrib.random.random_fill", True):
+ pytest.skip("random_fill is not built")
+ dev = tvm.device_from_target(target)
def test_local(dev, dtype):
if not tvm.get_global_func("tvm.contrib.random.random_fill", True):
@@ -158,11 +178,10 @@ def test_random_fill():
"float32",
"float64",
]:
- for target, dev in tvm.testing.enabled_targets():
- if tvm.target.Target(target).kind.name == "llvm":
- test_local(dev, dtype)
- else:
- tvm.testing.run_with_gpu_lock(test_local, dev, dtype)
+ if tvm.target.Target(target).kind.name == "llvm":
+ test_local(dev, dtype)
+ else:
+ tvm.testing.run_with_gpu_lock(test_local, dev, dtype)
test_rpc(dtype)
@@ -193,8 +212,4 @@ def test_random_fill_mt():
if __name__ == "__main__":
- test_randint()
- test_uniform()
- test_normal()
- test_random_fill()
- test_random_fill_mt()
+ tvm.testing.main()
diff --git a/tests/python/disco/test_callback.py b/tests/python/disco/test_callback.py
index 2997ced1c9..8c79c1b64d 100644
--- a/tests/python/disco/test_callback.py
+++ b/tests/python/disco/test_callback.py
@@ -92,7 +92,7 @@ def test_callback():
return (A, B)
- pipeline = tvm.transform.Sequential(
+ pipeline = tvm.ir.transform.Sequential(
[
tvm.relax.transform.LegalizeOps(),
tvm.s_tir.dlight.ApplyDefaultSchedule(tvm.s_tir.dlight.gpu.Fallback()),
diff --git a/tests/python/driver/test_compile.py b/tests/python/driver/test_compile.py
index da0e63208c..ac8dbafd99 100644
--- a/tests/python/driver/test_compile.py
+++ b/tests/python/driver/test_compile.py
@@ -15,7 +15,10 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
+import platform
+
import numpy as np
+import pytest
import tvm
import tvm.testing
@@ -86,7 +89,7 @@ def test_compile_relax():
tvm.testing.assert_allclose(z.numpy(), x_np + y_np)
-@tvm.testing.skip_if_32bit(reason="skipping test for i386.")
+@pytest.mark.skipif("32bit" in platform.architecture()[0], reason="skipping test for i386.")
def test_compile_mixed_module():
@tvm.script.ir_module
class MyModule:
diff --git a/tests/python/ir/test_node_reflection.py b/tests/python/ir/test_node_reflection.py
index cb49632cab..9d63b228d3 100644
--- a/tests/python/ir/test_node_reflection.py
+++ b/tests/python/ir/test_node_reflection.py
@@ -287,7 +287,7 @@ def test_string():
def test_pass_config():
- cfg = tvm.transform.PassContext(
+ cfg = tvm.ir.transform.PassContext(
opt_level=1,
config={
"tirx.UnrollLoop": {
@@ -303,15 +303,15 @@ def test_pass_config():
# schema checking for specific config key
with pytest.raises(TypeError):
- cfg = tvm.transform.PassContext(config={"tirx.UnrollLoop": {"invalid": 1}})
+ cfg = tvm.ir.transform.PassContext(config={"tirx.UnrollLoop": {"invalid": 1}})
# schema check for un-registered config
with pytest.raises(AttributeError):
- cfg = tvm.transform.PassContext(config={"inavlid-opt": True})
+ cfg = tvm.ir.transform.PassContext(config={"inavlid-opt": True})
# schema check for wrong type
with pytest.raises(AttributeError):
- cfg = tvm.transform.PassContext(config={"tirx.UnrollLoop": 1})
+ cfg = tvm.ir.transform.PassContext(config={"tirx.UnrollLoop": 1})
def test_dict():
diff --git a/tests/python/ir/test_pass_instrument.py b/tests/python/ir/test_pass_instrument.py
index 007d4e1d39..e469a75b55 100644
--- a/tests/python/ir/test_pass_instrument.py
+++ b/tests/python/ir/test_pass_instrument.py
@@ -19,11 +19,11 @@
import tvm
from tvm import relax
+from tvm.ir.transform.instrument import PrintAfterAll, PrintBeforeAll
from tvm.script import ir as I
from tvm.script import relax as R
from tvm.script import s_tir as Ts
from tvm.script import tirx as T
-from tvm.transform.instrument import PrintAfterAll, PrintBeforeAll
# pylint: disable=invalid-name,missing-function-docstring,no-value-for-parameter
@@ -38,7 +38,7 @@ def test_tir_print_all_passes(capsys):
vi, vj, vk, vl = Ts.axis.remap("SSSS", [i, j, k, l])
B[vi, vj, vk, vl] = A[vi, vj, vk, vl] * 2.0
- with tvm.transform.PassContext(opt_level=3, instruments=[PrintBeforeAll(), PrintAfterAll()]):
+ with tvm.ir.transform.PassContext(opt_level=3, instruments=[PrintBeforeAll(), PrintAfterAll()]):
tvm.compile(func)
all_passes_output = capsys.readouterr().out
assert "Before Running Pass:" in all_passes_output
@@ -56,7 +56,7 @@ def test_relax_print_all_passes(capsys):
return y
pipeline = relax.get_pipeline("default_build")
- with tvm.transform.PassContext(opt_level=3, instruments=[PrintBeforeAll(), PrintAfterAll()]):
+ with tvm.ir.transform.PassContext(opt_level=3, instruments=[PrintBeforeAll(), PrintAfterAll()]):
pipeline(Module)
all_passes_output = capsys.readouterr().out
assert "Before Running Pass:" in all_passes_output
diff --git a/tests/python/relax/conftest.py b/tests/python/relax/conftest.py
index a8db5d7788..0b1f155190 100644
--- a/tests/python/relax/conftest.py
+++ b/tests/python/relax/conftest.py
@@ -64,7 +64,7 @@ def apply_instrument_well_formed(unit_test_marks):
validate_before_transform = "skip_well_formed_check_before_transform" not in unit_test_marks
validate_after_transform = "skip_well_formed_check_after_transform" not in unit_test_marks
- current = tvm.transform.PassContext.current()
+ current = tvm.ir.transform.PassContext.current()
instruments = list(current.instruments)
if validate_before_transform or validate_after_transform:
@@ -72,7 +72,7 @@ def apply_instrument_well_formed(unit_test_marks):
WellFormedInstrument(validate_before_transform=validate_before_transform)
)
- override = tvm.transform.PassContext(
+ override = tvm.ir.transform.PassContext(
# With the new WellFormedInstrument appended
instruments=instruments,
# Forward all other parameters
diff --git a/tests/python/relax/test_codegen_cublas.py b/tests/python/relax/test_codegen_cublas.py
index 0f50fd5d82..f56f33c9f8 100644
--- a/tests/python/relax/test_codegen_cublas.py
+++ b/tests/python/relax/test_codegen_cublas.py
@@ -50,7 +50,7 @@ pytestmark = [
def build_and_run(mod, inputs_np, target, legalize=False, cuda_graph=False):
- with tvm.transform.PassContext(
+ with tvm.ir.transform.PassContext(
config={
"relax.backend.use_cuda_graph": cuda_graph,
"relax.transform.apply_legalize_ops": legalize,
diff --git a/tests/python/relax/test_codegen_cudnn.py b/tests/python/relax/test_codegen_cudnn.py
index 003c72011a..ed623ad873 100644
--- a/tests/python/relax/test_codegen_cudnn.py
+++ b/tests/python/relax/test_codegen_cudnn.py
@@ -111,7 +111,7 @@ def get_result_with_relax_cudnn_offload(mod, np_inputs, cuda_graph=False):
def build_and_run(mod, inputs_np, target, legalize=False, cuda_graph=False):
- with tvm.transform.PassContext(
+ with tvm.ir.transform.PassContext(
config={
"relax.backend.use_cuda_graph": cuda_graph,
"relax.transform.apply_legalize_ops": legalize,
diff --git a/tests/python/relax/test_codegen_cutlass.py b/tests/python/relax/test_codegen_cutlass.py
index 9ad2eea92d..29e0701724 100644
--- a/tests/python/relax/test_codegen_cutlass.py
+++ b/tests/python/relax/test_codegen_cutlass.py
@@ -93,7 +93,7 @@ pytestmark = [
def build_and_run(mod, inputs_np, target, legalize=True, cuda_graph=False):
- with tvm.transform.PassContext(
+ with tvm.ir.transform.PassContext(
config={
"relax.backend.use_cuda_graph": cuda_graph,
"relax.transform.apply_legalize_ops": legalize,
diff --git a/tests/python/relax/test_codegen_dnnl.py b/tests/python/relax/test_codegen_dnnl.py
index e172f84351..83183b3f0d 100644
--- a/tests/python/relax/test_codegen_dnnl.py
+++ b/tests/python/relax/test_codegen_dnnl.py
@@ -56,7 +56,7 @@ def build_and_run(mod, inputs, legalize=False):
dev = tvm.cpu()
inputs = [tvm.runtime.tensor(inp, dev) for inp in inputs]
- with tvm.transform.PassContext(config={"relax.transform.apply_legalize_ops": legalize}):
+ with tvm.ir.transform.PassContext(config={"relax.transform.apply_legalize_ops": legalize}):
ex = tvm.compile(mod, target)
vm = relax.VirtualMachine(ex, dev)
f = vm["main"]
@@ -68,7 +68,7 @@ def test_dnnl_offload():
"relax.nn.conv2d", with_bias=False, activation="relax.nn.relu"
)
- seq = tvm.transform.Sequential(
+ seq = tvm.ir.transform.Sequential(
[
relax.transform.FuseOpsByPattern([("dnnl.conv2d_relu", pat)]),
relax.transform.MergeCompositeFunctions(),
diff --git a/tests/python/relax/test_codegen_hipblas.py b/tests/python/relax/test_codegen_hipblas.py
index 77ed14fbac..41402e2afa 100644
--- a/tests/python/relax/test_codegen_hipblas.py
+++ b/tests/python/relax/test_codegen_hipblas.py
@@ -47,7 +47,7 @@ pytestmark = [
def build_and_run(mod, inputs_np, target, legalize=False):
- with tvm.transform.PassContext(config={"relax.transform.apply_legalize_ops": legalize}):
+ with tvm.ir.transform.PassContext(config={"relax.transform.apply_legalize_ops": legalize}):
ex = tvm.compile(mod, target)
def run_and_check():
diff --git a/tests/python/relax/test_codegen_tensorrt.py b/tests/python/relax/test_codegen_tensorrt.py
index 4a3c35ebd7..364015d648 100644
--- a/tests/python/relax/test_codegen_tensorrt.py
+++ b/tests/python/relax/test_codegen_tensorrt.py
@@ -66,7 +66,7 @@ pytestmark = [
def build_and_run(mod, inputs_np, target, legalize=False):
- with tvm.transform.PassContext(config={"relax.transform.apply_legalize_ops": legalize}):
+ with tvm.ir.transform.PassContext(config={"relax.transform.apply_legalize_ops": legalize}):
ex = tvm.compile(mod, target)
def run_and_check():
@@ -107,7 +107,7 @@ def test_tensorrt_offload():
params_np = {"weight1": inputs[1], "weight2": inputs[2]}
- mod = tvm.transform.Sequential(
+ mod = tvm.ir.transform.Sequential(
[
relax.transform.BindParams("main", params_np),
relax.transform.FuseOpsByPattern(patterns),
@@ -133,7 +133,7 @@ def _offload_and_compare(mod, params_np, patterns, data_np, rtol=1e-2, atol=1e-2
if patterns is None:
partitioned = partition_for_tensorrt(mod)
else:
- partitioned = tvm.transform.Sequential(
+ partitioned = tvm.ir.transform.Sequential(
[relax.transform.FuseOpsByPattern(patterns), relax.transform.MergeCompositeFunctions()]
)(mod)
# Guard against a silent false pass: if no pattern matched, nothing is offloaded and the
@@ -323,7 +323,7 @@ def test_tensorrt_int8_calibration(monkeypatch):
ref = build_and_run(Conv2dInt8, [data, weight], "llvm", legalize=True)
patterns = [("tensorrt.nn.conv2d", is_op("relax.nn.conv2d")(wildcard(), wildcard()))]
- offloaded = tvm.transform.Sequential(
+ offloaded = tvm.ir.transform.Sequential(
[
relax.transform.BindParams("main", {"weight": weight}),
relax.transform.FuseOpsByPattern(patterns),
diff --git a/tests/python/relax/test_contrib_vllm.py b/tests/python/relax/test_contrib_vllm.py
index d287e4899b..51de3ec895 100644
--- a/tests/python/relax/test_contrib_vllm.py
+++ b/tests/python/relax/test_contrib_vllm.py
@@ -48,7 +48,7 @@ def build_and_run(mod, inputs_np, target, legalize=True):
with tvm.target.Target("cuda"):
mod = tvm.s_tir.transform.DefaultGPUSchedule()(mod)
- with tvm.transform.PassContext():
+ with tvm.ir.transform.PassContext():
ex = tvm.compile(mod, target)
dev = tvm.cuda(0)
diff --git a/tests/python/relax/test_dataflow_rewriter.py b/tests/python/relax/test_dataflow_rewriter.py
index c92753ccad..57d2abb548 100644
--- a/tests/python/relax/test_dataflow_rewriter.py
+++ b/tests/python/relax/test_dataflow_rewriter.py
@@ -195,7 +195,7 @@ def test_rewriter_may_be_applied_to_ir_module():
def test_rewriter_may_be_used_as_ir_transform():
- """A rewriter may be used as a tvm.transform.Pass"""
+ """A rewriter may be used as a tvm.ir.transform.Pass"""
@R.rewriter
class Rewriter:
@@ -225,7 +225,7 @@ def test_rewriter_may_be_used_as_ir_transform():
)
return out
- After = tvm.transform.Sequential([Rewriter])(Before)
+ After = tvm.ir.transform.Sequential([Rewriter])(Before)
tvm.ir.assert_structural_equal(Expected, After)
diff --git a/tests/python/relax/test_frontend_nn_jit.py b/tests/python/relax/test_frontend_nn_jit.py
index a70808b4ca..270678b60b 100644
--- a/tests/python/relax/test_frontend_nn_jit.py
+++ b/tests/python/relax/test_frontend_nn_jit.py
@@ -84,7 +84,7 @@ def test_jit_with_effect(debug):
}
mod = Layer()
- with tvm.transform.PassContext(opt_level=3):
+ with tvm.ir.transform.PassContext(opt_level=3):
model = mod.jit(spec=forward_spec, debug=debug)
x0 = torch.rand((1, 10, 5), dtype=torch.float32)
diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py
index 81db2cd25e..4bbc24268b 100644
--- a/tests/python/relax/test_frontend_onnx.py
+++ b/tests/python/relax/test_frontend_onnx.py
@@ -155,7 +155,7 @@ def check_correctness(
# Separate model from parameters.
tvm_model, params = relax.frontend.detach_params(tvm_model)
# Compile the relax graph into a VM then run.
- with tvm.transform.PassContext(opt_level=3):
+ with tvm.ir.transform.PassContext(opt_level=3):
ex = tvm.compile(tvm_model, target="llvm")
vm = relax.VirtualMachine(ex, tvm.cpu())
# Prepare inputs.
@@ -235,7 +235,7 @@ def run_in_tvm(
tvm_model = relax.transform.LegalizeOps()(tvm_model)
tvm_model, params = relax.frontend.detach_params(tvm_model)
- with tvm.transform.PassContext(opt_level=3):
+ with tvm.ir.transform.PassContext(opt_level=3):
ex = tvm.compile(tvm_model, target="llvm")
vm = relax.VirtualMachine(ex, tvm.cpu())
@@ -2452,7 +2452,7 @@ def test_scatter_dynamic_shape():
tvm_model = from_onnx(model, shape_dict=shape_dict, opset=9)
tvm_model = relax.transform.DecomposeOpsForInference()(tvm_model)
tvm_model = relax.transform.LegalizeOps()(tvm_model)
- with tvm.transform.PassContext(opt_level=3):
+ with tvm.ir.transform.PassContext(opt_level=3):
ex = tvm.compile(tvm_model, target="llvm")
vm = relax.VirtualMachine(ex, tvm.cpu())
vm.set_input("main", data, indices, updates)
@@ -4013,7 +4013,7 @@ def test_where_shape_expr_broadcast(cond_vals, y_vals, expected):
model = _make_where_shape_expr_model(t1_shape, cond_vals, y_vals)
feeds = {"__t": np.ones(t1_shape, dtype="float32")}
mod = from_onnx(model, shape_dict={"__t": list(t1_shape)})
- with tvm.transform.PassContext(opt_level=3):
+ with tvm.ir.transform.PassContext(opt_level=3):
ex = tvm.compile(mod, target="llvm")
vm = relax.VirtualMachine(ex, tvm.cpu())
out = vm["main"](feeds["__t"])
@@ -4769,7 +4769,7 @@ def test_softplus_float64_accuracy():
tvm_model = relax.transform.LegalizeOps()(tvm_model)
tvm_model, params = relax.frontend.detach_params(tvm_model)
- with tvm.transform.PassContext(opt_level=3):
+ with tvm.ir.transform.PassContext(opt_level=3):
ex = tvm.compile(tvm_model, target="llvm")
vm = relax.VirtualMachine(ex, tvm.cpu())
diff --git a/tests/python/relax/test_frontend_onnx_backend.py b/tests/python/relax/test_frontend_onnx_backend.py
index 45386fe57c..94aef5c648 100644
--- a/tests/python/relax/test_frontend_onnx_backend.py
+++ b/tests/python/relax/test_frontend_onnx_backend.py
@@ -58,7 +58,7 @@ class TVMRelaxBackendRep(BackendRep):
self._func_param_names = func_param_names
self._graph_input_names = graph_input_names
- with tvm.transform.PassContext(opt_level=3):
+ with tvm.ir.transform.PassContext(opt_level=3):
ex = tvm.compile(mod, target="llvm")
self._vm = relax.VirtualMachine(ex, tvm.cpu())
diff --git a/tests/python/relax/test_op_gradient_numeric.py b/tests/python/relax/test_op_gradient_numeric.py
index 66e60641f9..6357eb1ef9 100644
--- a/tests/python/relax/test_op_gradient_numeric.py
+++ b/tests/python/relax/test_op_gradient_numeric.py
@@ -27,9 +27,9 @@ import tvm.testing
from tvm import relax
from tvm.ir import Call
from tvm.ir.op import Op
+from tvm.relax.testing.gradient import check_numerical_grads
from tvm.relax.transform import LegalizeOps
from tvm.relax.type import TensorType, TupleType
-from tvm.testing.utils import check_numerical_grads
def relax_check_gradients(
diff --git a/tests/python/relax/test_optimize_layout_transform.py b/tests/python/relax/test_optimize_layout_transform.py
index c82ffb18c5..c5bb7f0b71 100644
--- a/tests/python/relax/test_optimize_layout_transform.py
+++ b/tests/python/relax/test_optimize_layout_transform.py
@@ -31,7 +31,7 @@ from tvm.script import tirx as T
def _run_pass_compare_output(Before, Expected):
- After = tvm.transform.Sequential(
+ After = tvm.ir.transform.Sequential(
[
OptimizeLayoutTransform(),
DeadCodeElimination(),
diff --git a/tests/python/relax/test_tensorrt_engine_serialization.py b/tests/python/relax/test_tensorrt_engine_serialization.py
index a527c63880..401fc05e61 100644
--- a/tests/python/relax/test_tensorrt_engine_serialization.py
+++ b/tests/python/relax/test_tensorrt_engine_serialization.py
@@ -107,7 +107,7 @@ def _relu_partition(shape=(2, 4)):
with builder.dataflow():
output = builder.emit_output(relax.op.nn.relu(data))
builder.emit_func_output(output)
- return tvm.transform.Sequential(
+ return tvm.ir.transform.Sequential(
[
relax.transform.FuseOpsByPattern(
[("tensorrt.nn.relu", is_op("relax.nn.relu")(wildcard()))]
@@ -121,7 +121,7 @@ def _codegen(partitioned, eager=None):
config = {}
if eager is not None:
config["relax.ext.tensorrt.options"] = {"build_at_compile_time": eager}
- with tvm.transform.PassContext(config=config):
+ with tvm.ir.transform.PassContext(config=config):
return relax.transform.RunCodegen()(partitioned)
@@ -212,7 +212,7 @@ def _two_partitions(dtype):
("tensorrt.add", is_op("relax.add")(wildcard(), wildcard())),
("tensorrt.multiply", is_op("relax.multiply")(wildcard(), wildcard())),
]
- partitioned = tvm.transform.Sequential(
+ partitioned = tvm.ir.transform.Sequential(
[
relax.transform.BindParams("main", parameters),
# Leave sin on the TVM side to force two TensorRT partitions.
diff --git a/tests/python/relax/test_transform_codegen_pass.py b/tests/python/relax/test_transform_codegen_pass.py
index 409584f59e..26f604bdd8 100644
--- a/tests/python/relax/test_transform_codegen_pass.py
+++ b/tests/python/relax/test_transform_codegen_pass.py
@@ -87,7 +87,7 @@ def gen_ground_truth(mod, target):
# Lower and run tuning
# Since there is no default schedule for GPU in MS yet, this is necessary
with target:
- seq = tvm.transform.Sequential(
+ seq = tvm.ir.transform.Sequential(
[relax.transform.LegalizeOps(), s_tir.transform.DefaultGPUSchedule()]
)
new_mod = seq(mod)
@@ -148,7 +148,7 @@ def test_tensorrt_only(entry_func_name):
("tensorrt.add", is_op("relax.add")(wildcard(), wildcard())),
]
- new_mod = tvm.transform.Sequential(
+ new_mod = tvm.ir.transform.Sequential(
[
relax.transform.FuseOpsByPattern(patterns),
relax.transform.MergeCompositeFunctions(),
@@ -176,8 +176,8 @@ def test_mix_use_tensorrt_and_tvm():
# Run Codegen pass
with tempfile.TemporaryDirectory() as work_dir:
- with target, tvm.transform.PassContext(opt_level=0):
- new_mod = tvm.transform.Sequential(
+ with target, tvm.ir.transform.PassContext(opt_level=0):
+ new_mod = tvm.ir.transform.Sequential(
[
relax.transform.FuseOpsByPattern(patterns),
relax.transform.MergeCompositeFunctions(),
diff --git a/tests/python/relax/test_transform_cse.py b/tests/python/relax/test_transform_cse.py
index d45edcc42c..3c6816b4bc 100644
--- a/tests/python/relax/test_transform_cse.py
+++ b/tests/python/relax/test_transform_cse.py
@@ -428,7 +428,7 @@ def test_call_tir_tuple_arg():
# normalization of those expressions may produce additional
# variables bindings. This test case should be agnostic to those
# additional bindings, so DCE is applied after CSE.
- After = tvm.transform.Sequential(
+ After = tvm.ir.transform.Sequential(
[
EliminateCommonSubexpr(),
tvm.relax.transform.DeadCodeElimination(),
diff --git a/tests/python/relax/test_transform_error_enrichment.py b/tests/python/relax/test_transform_error_enrichment.py
index 97a954cfcd..d7cd7812d8 100644
--- a/tests/python/relax/test_transform_error_enrichment.py
+++ b/tests/python/relax/test_transform_error_enrichment.py
@@ -75,7 +75,7 @@ def test_sequential_does_not_double_append():
"""Running the failing pass inside a Sequential must not enrich twice — the
Sequential wrapper does not guard, only the leaf pass does."""
mod = _bad_matmul_module()
- seq = tvm.transform.Sequential([relax.transform.Normalize()])
+ seq = tvm.ir.transform.Sequential([relax.transform.Normalize()])
with pytest.raises(ValueError) as excinfo:
seq(mod)
msg = str(excinfo.value)
@@ -89,7 +89,7 @@ def test_error_without_resolvable_node_is_not_masked():
"""A pass that throws an error whose node is not findable in the module must
surface the original message without raising a printer/render error."""
- @tvm.transform.module_pass(opt_level=0, name="ThrowUnresolvable")
+ @tvm.ir.transform.module_pass(opt_level=0, name="ThrowUnresolvable")
class ThrowUnresolvable:
def transform_module(self, mod, ctx):
# A bare error with no VisitErrorContext payload -> nothing to resolve.
diff --git a/tests/python/relax/test_transform_fuse_ops_by_pattern.py b/tests/python/relax/test_transform_fuse_ops_by_pattern.py
index 30448c0cbe..e6a14ddef8 100644
--- a/tests/python/relax/test_transform_fuse_ops_by_pattern.py
+++ b/tests/python/relax/test_transform_fuse_ops_by_pattern.py
@@ -511,7 +511,7 @@ def test_cyclic_dependency():
def test_bind_params():
weight_np = np.random.randn(64, 64, 3, 3).astype("float32")
- mod = tvm.transform.Sequential(
+ mod = tvm.ir.transform.Sequential(
[
relax.transform.BindParams("main", {"weight1": weight_np}),
relax.transform.FuseOpsByPattern(
diff --git a/tests/python/relax/test_transform_fuse_transpose_matmul.py b/tests/python/relax/test_transform_fuse_transpose_matmul.py
index d88997ee5b..a6b3e6b839 100644
--- a/tests/python/relax/test_transform_fuse_transpose_matmul.py
+++ b/tests/python/relax/test_transform_fuse_transpose_matmul.py
@@ -72,7 +72,7 @@ def test_transform_fuse_transpose_matmul():
R.output(gv)
return gv
- after = tvm.transform.Sequential(
+ after = tvm.ir.transform.Sequential(
[
relax.transform.FuseTransposeMatmul(),
relax.transform.FuseTIR(), # Only used for remove unused primitive function
@@ -125,7 +125,7 @@ def test_transform_fuse_transpose_matmul_const():
R.output(gv)
return gv
- after = tvm.transform.Sequential(
+ after = tvm.ir.transform.Sequential(
[
relax.transform.FuseTransposeMatmul(),
relax.transform.FuseTIR(), # Only used for remove unused primitive function
diff --git a/tests/python/relax/test_transform_gradient_numeric.py b/tests/python/relax/test_transform_gradient_numeric.py
index 1d5466a165..6a7a76346b 100644
--- a/tests/python/relax/test_transform_gradient_numeric.py
+++ b/tests/python/relax/test_transform_gradient_numeric.py
@@ -21,10 +21,10 @@ import pytest
import tvm
import tvm.testing
from tvm import relax
+from tvm.relax.testing.gradient import check_numerical_grads
from tvm.script.parser import ir as I
from tvm.script.parser import relax as R
from tvm.testing import assert_allclose
-from tvm.testing.utils import check_numerical_grads
def rand(dtype, *shape):
diff --git a/tests/python/relax/test_transform_legalize_ops_manipulate.py b/tests/python/relax/test_transform_legalize_ops_manipulate.py
index cc7eb189e2..8d9bd874b6 100644
--- a/tests/python/relax/test_transform_legalize_ops_manipulate.py
+++ b/tests/python/relax/test_transform_legalize_ops_manipulate.py
@@ -1811,7 +1811,7 @@ def test_func_ty_of_legalized_layout_transform():
R.output(gv)
return gv
- After = tvm.transform.Sequential(
+ After = tvm.ir.transform.Sequential(
[
relax.transform.LegalizeOps(),
relax.transform.ToNonDataflow(),
diff --git a/tests/python/relax/test_transform_lift_transform_params.py b/tests/python/relax/test_transform_lift_transform_params.py
index 5b941bef46..ae6d5f4a0d 100644
--- a/tests/python/relax/test_transform_lift_transform_params.py
+++ b/tests/python/relax/test_transform_lift_transform_params.py
@@ -226,7 +226,7 @@ def test_basic(consume_params):
mod = Before
expected = Expected if not consume_params else ExpectedConsumeParams
- with tvm.transform.PassContext(
+ with tvm.ir.transform.PassContext(
config={"relax.lift_transform_params.consume_params": consume_params}
):
after = relax.transform.LiftTransformParams()(mod)
diff --git a/tests/python/relax/test_transform_meta_schedule_tuning.py b/tests/python/relax/test_transform_meta_schedule_tuning.py
index 8f4d2728c1..7865efa65d 100644
--- a/tests/python/relax/test_transform_meta_schedule_tuning.py
+++ b/tests/python/relax/test_transform_meta_schedule_tuning.py
@@ -39,12 +39,13 @@ import tvm_ffi
import tvm
import tvm.s_tir.meta_schedule as ms
import tvm.testing
-from tvm import relax, transform
+from tvm import relax
+from tvm.ir import transform
from tvm.ir.module import IRModule
+from tvm.ir.transform import PassContext
from tvm.script import relax as R
from tvm.script import s_tir as Ts
from tvm.script import tirx as T
-from tvm.transform import PassContext
target = tvm.target.Target({"kind": "llvm", "num-cores": 16})
diff --git a/tests/python/relax/test_transform_rewrite_cuda_graph.py b/tests/python/relax/test_transform_rewrite_cuda_graph.py
index 133d6d4e7a..2ff3dd794f 100644
--- a/tests/python/relax/test_transform_rewrite_cuda_graph.py
+++ b/tests/python/relax/test_transform_rewrite_cuda_graph.py
@@ -31,7 +31,7 @@ from tvm.script import tirx as T
@pytest.fixture(autouse=True)
def enable_cuda_graph():
"""Enable cuda graph transform for all tests in this file"""
- with tvm.transform.PassContext(config={"relax.backend.use_cuda_graph": True}):
+ with tvm.ir.transform.PassContext(config={"relax.backend.use_cuda_graph": True}):
yield
@@ -657,7 +657,7 @@ def test_capture_fixed_inputs():
_4: R.Tuple = R.memory.kill_storage(storage1)
return gv_1
- mod = tvm.transform.Sequential(
+ mod = tvm.ir.transform.Sequential(
[
relax.pipeline.get_pipeline(),
relax.transform.LiftTransformParams(),
@@ -698,9 +698,9 @@ def test_transform_is_no_op_when_disabled():
alloc3 = R.memory.alloc_tensor(storage, 0, R.shape([8]), "float32")
return ()
- with tvm.transform.PassContext(config={"relax.backend.use_cuda_graph": True}):
+ with tvm.ir.transform.PassContext(config={"relax.backend.use_cuda_graph": True}):
AfterWhenEnabled = relax.transform.RewriteCUDAGraph()(Before)
- with tvm.transform.PassContext(config={"relax.backend.use_cuda_graph": False}):
+ with tvm.ir.transform.PassContext(config={"relax.backend.use_cuda_graph": False}):
AfterWhenDisabled = relax.transform.RewriteCUDAGraph()(Before)
assert not tvm_ffi.structural_equal(Before, AfterWhenEnabled)
diff --git a/tests/python/relax/test_transform_rewrite_dataflow_reshape.py b/tests/python/relax/test_transform_rewrite_dataflow_reshape.py
index 832b0fd529..42fa84af5e 100644
--- a/tests/python/relax/test_transform_rewrite_dataflow_reshape.py
+++ b/tests/python/relax/test_transform_rewrite_dataflow_reshape.py
@@ -639,7 +639,7 @@ def test_rewrite_static_reshape():
i, j = Ts.axis.remap("SS", iters)
z[i, j] = y1[i, j] + y2[i, j]
- After = tvm.transform.Sequential(
+ After = tvm.ir.transform.Sequential(
[
# Lower both R.reshape and R.add from Relax to TIR
relax.transform.LegalizeOps(),
@@ -700,7 +700,7 @@ def test_rewrite_static_reshape():
# i, j = Ts.axis.remap("SS", iters)
# z[i, j] = y1[i, j] + y2[i, j]
-# After = tvm.transform.Sequential(
+# After = tvm.ir.transform.Sequential(
# [
# # Lower both R.reshape and R.add from Relax to TIR
# relax.transform.LegalizeOps(),
@@ -763,7 +763,7 @@ def test_rewrite_dynamic_reshape():
i, j = Ts.axis.remap("SS", iters)
z[i, j] = y1[i, j] + y2[i, j]
- After = tvm.transform.Sequential(
+ After = tvm.ir.transform.Sequential(
[
# Lower both R.reshape and R.add from Relax to TIR
relax.transform.LegalizeOps(),
diff --git a/tests/python/relax/test_vm_alloc_storage_with_scope.py b/tests/python/relax/test_vm_alloc_storage_with_scope.py
index a7ceb5a4bf..3d3fd97028 100644
--- a/tests/python/relax/test_vm_alloc_storage_with_scope.py
+++ b/tests/python/relax/test_vm_alloc_storage_with_scope.py
@@ -62,7 +62,7 @@ def test_alloc_storage_with_scope_global():
output_ref = arg0 + arg0
mod = Module
target = "llvm"
- with tvm.transform.PassContext(opt_level=3):
+ with tvm.ir.transform.PassContext(opt_level=3):
lib = tvm.relax.build(mod, target=target)
dev = tvm.cpu()
diff --git a/tests/python/relax/test_vm_build.py b/tests/python/relax/test_vm_build.py
index ddc358dad9..0fd605de5d 100644
--- a/tests/python/relax/test_vm_build.py
+++ b/tests/python/relax/test_vm_build.py
@@ -1269,7 +1269,7 @@ def test_relax_module_with_multiple_targets():
C = R.add(A, B)
return C
- seq = tvm.transform.Sequential(
+ seq = tvm.ir.transform.Sequential(
[
tvm.relax.transform.LegalizeOps(),
tvm.s_tir.dlight.ApplyDefaultSchedule(tvm.s_tir.dlight.gpu.Fallback()),
diff --git a/tests/python/relax/test_vm_cuda_graph.py b/tests/python/relax/test_vm_cuda_graph.py
index e63683385f..b07fe69ddd 100644
--- a/tests/python/relax/test_vm_cuda_graph.py
+++ b/tests/python/relax/test_vm_cuda_graph.py
@@ -153,8 +153,8 @@ def test_capture_error_is_recoverable():
D = R.add(C, C)
return D
- with target, tvm.transform.PassContext(config={"relax.backend.use_cuda_graph": True}):
- Module = tvm.transform.Sequential(
+ with target, tvm.ir.transform.PassContext(config={"relax.backend.use_cuda_graph": True}):
+ Module = tvm.ir.transform.Sequential(
[
tvm.relax.transform.LegalizeOps(),
tvm.s_tir.transform.DefaultGPUSchedule(),
diff --git a/tests/python/relax/texture/test_texture_nd.py b/tests/python/relax/texture/test_texture_nd.py
index a627c0cb1f..aaedcccf12 100644
--- a/tests/python/relax/texture/test_texture_nd.py
+++ b/tests/python/relax/texture/test_texture_nd.py
@@ -62,7 +62,7 @@ def get_rpc():
def preprocess_pipeline(mod: IRModule) -> IRModule:
desired_layouts = {"relax.nn.conv2d": ["NCHW16c", "OIHW16o", "NCHW16c"]}
- seq = tvm.transform.Sequential(
+ seq = tvm.ir.transform.Sequential(
[
tvm.tirx.transform.BindTarget(Target.current(allow_none=False)),
tvm.relax.transform.FoldConstant(),
@@ -89,7 +89,7 @@ def preprocess_pipeline(mod: IRModule) -> IRModule:
def postprocess_pipeline(mod: IRModule) -> IRModule:
- seq = tvm.transform.Sequential(
+ seq = tvm.ir.transform.Sequential(
[
tvm.relax.transform.ToNonDataflow(),
tvm.relax.transform.RemovePurityChecking(),
diff --git a/tests/python/request_hook.py b/tests/python/request_hook.py
deleted file mode 100644
index f3faf77a19..0000000000
--- a/tests/python/request_hook.py
+++ /dev/null
@@ -1,56 +0,0 @@
-#!/usr/bin/env bash
-
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-import logging
-import urllib.request
-
-LOGGER = None
-
-
-URL_MAP = {}
-
-
-class TvmRequestHook(urllib.request.Request):
- def __init__(self, url, *args, **kwargs):
- LOGGER.info(f"Caught access to {url}")
- url = url.strip()
- if url not in URL_MAP:
- # Disallow network accesses without an explicitly maintained mirror.
- msg = (
- f"Uncaught URL found in CI: {url}. "
- "Avoid network access or arrange a stable project-managed mirror, "
- "then add it to URL_MAP in tests/python/request_hook.py."
- )
- raise RuntimeError(msg)
-
- new_url = URL_MAP[url]
- LOGGER.info(f"Mapped URL {url} to {new_url}")
- super().__init__(new_url, *args, **kwargs)
-
-
-def init():
- global LOGGER
- if urllib.request.Request is TvmRequestHook:
- return
- urllib.request.Request = TvmRequestHook
- LOGGER = logging.getLogger("tvm_request_hook")
- LOGGER.setLevel(logging.DEBUG)
- fh = logging.FileHandler("redirected_urls.log")
- fh.setLevel(logging.DEBUG)
- LOGGER.addHandler(fh)
diff --git a/tests/python/runtime/test_runtime_rpc.py b/tests/python/runtime/test_runtime_rpc.py
index 965af262ac..0bfbcfe36e 100644
--- a/tests/python/runtime/test_runtime_rpc.py
+++ b/tests/python/runtime/test_runtime_rpc.py
@@ -19,6 +19,7 @@
import gc
import multiprocessing
import os
+import platform
import stat
import sys
import tempfile
@@ -175,7 +176,7 @@ def test_rpc_large_array():
check_remote()
-@tvm.testing.skip_if_32bit(reason="skipping test for i386.")
+@pytest.mark.skipif("32bit" in platform.architecture()[0], reason="skipping test for i386.")
@pytest.mark.skipif(
tvm.get_global_func("rpc.LocalSession", allow_missing=True) is None, reason="need rpc"
)
diff --git a/tests/python/s_tir/base/test_compilation_pipeline.py b/tests/python/s_tir/base/test_compilation_pipeline.py
index 4ef79e4013..2b6bbee42b 100644
--- a/tests/python/s_tir/base/test_compilation_pipeline.py
+++ b/tests/python/s_tir/base/test_compilation_pipeline.py
@@ -46,13 +46,13 @@ def test_default_pipeline_selects_dialect(mixed):
functions["native"] = native
finalized_modules = []
- @tvm.transform.pass_instrument
+ @tvm.ir.transform.pass_instrument
class RecordFinalization:
def run_before_pass(self, mod, info):
if info.name == "tirx.MakePackedAPI":
finalized_modules.append({gv.name_hint for gv in mod.functions})
- with tvm.transform.PassContext(instruments=[RecordFinalization()]):
+ with tvm.ir.transform.PassContext(instruments=[RecordFinalization()]):
module = tvm.tirx.compile(tvm.IRModule(functions), target="llvm").mod
assert finalized_modules == [set(functions)]
source = np.arange(16, dtype="float32")
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_cost_model.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_cost_model.py
index bec87fe531..537e518ef4 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_cost_model.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_cost_model.py
@@ -28,7 +28,6 @@ import pytest
import tvm
import tvm.testing
-from tvm.ir.utils import derived_object
from tvm.s_tir.meta_schedule.cost_model import PyCostModel, RandomModel, XGBModel
from tvm.s_tir.meta_schedule.cost_model.xgb_model import PackSum, _get_custom_call_back
from tvm.s_tir.meta_schedule.feature_extractor import RandomFeatureExtractor
@@ -36,6 +35,7 @@ from tvm.s_tir.meta_schedule.runner import RunnerResult
from tvm.s_tir.meta_schedule.search_strategy import MeasureCandidate
from tvm.s_tir.meta_schedule.tune_context import TuneContext
from tvm.s_tir.schedule.schedule import Schedule
+from tvm.s_tir.utils import derived_object
from tvm.script import s_tir as Ts
from tvm.script import tirx as T
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_database.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_database.py
index 0f2f61080a..27cc1462fa 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_database.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_database.py
@@ -30,10 +30,10 @@ import tvm
import tvm.testing
from tvm import tirx
from tvm.ir.module import IRModule
-from tvm.ir.utils import derived_object
from tvm.s_tir import Schedule
from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.meta_schedule.database import TuningRecord, Workload
+from tvm.s_tir.utils import derived_object
from tvm.script import s_tir as Ts
from tvm.script import tirx as T
from tvm.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_feature_extractor.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_feature_extractor.py
index 602a6ddc30..8be765b672 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_feature_extractor.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_feature_extractor.py
@@ -19,10 +19,10 @@
import numpy as np
import tvm.runtime
-from tvm.ir.utils import derived_object
from tvm.s_tir.meta_schedule import TuneContext
from tvm.s_tir.meta_schedule.feature_extractor import PyFeatureExtractor
from tvm.s_tir.meta_schedule.search_strategy import MeasureCandidate
+from tvm.s_tir.utils import derived_object
def test_meta_schedule_feature_extractor():
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_measure_callback.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_measure_callback.py
index ae239eca4a..deff80c248 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_measure_callback.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_measure_callback.py
@@ -20,9 +20,9 @@ import tempfile
import pytest
import tvm
-from tvm.ir.utils import derived_object
from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.schedule import Schedule
+from tvm.s_tir.utils import derived_object
from tvm.script import s_tir as Ts
from tvm.script import tirx as T
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_mma_tensorize.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_mma_tensorize.py
index 35f62e1e8d..49c3006e89 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_mma_tensorize.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_mma_tensorize.py
@@ -74,7 +74,7 @@ def test_run_target(mod=None, tgt_str=None, in_dtype="float16", out_dtype="float
return
tgt_str = tgt_str or "cuda"
target = tvm.target.Target(target=tgt_str)
- with tvm.transform.PassContext(opt_level=3):
+ with tvm.ir.transform.PassContext(opt_level=3):
lib: tvm.runtime.Module = tvm.compile(mod, target=target)
a_np = np.random.rand(M, K).astype(in_dtype)
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_post_order_apply.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_post_order_apply.py
index cd0529daa8..ddc7ff5f36 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_post_order_apply.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_post_order_apply.py
@@ -28,11 +28,11 @@ import tvm
import tvm.testing
from tvm import te
from tvm.ir.module import IRModule
-from tvm.ir.utils import derived_object
from tvm.s_tir.meta_schedule import TuneContext
from tvm.s_tir.meta_schedule.schedule_rule import PyScheduleRule
from tvm.s_tir.meta_schedule.space_generator import PostOrderApply
from tvm.s_tir.schedule import SBlockRV, Schedule
+from tvm.s_tir.utils import derived_object
from tvm.script import s_tir as Ts
from tvm.script import tirx as T
from tvm.target import Target
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_runner.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_runner.py
index 4c4d2a95d6..e1a710e81d 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_runner.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_runner.py
@@ -18,6 +18,7 @@
"""Test Meta Schedule Runner"""
import itertools
+import platform
import sys
import time
from typing import Any
@@ -30,7 +31,6 @@ pytest.importorskip("tornado") # tvm.rpc.tracker (LocalRPC) requires tornado
import tvm
import tvm.testing
-from tvm.ir.utils import derived_object
from tvm.rpc import RPCSession
from tvm.runtime import Device, Module
from tvm.s_tir.meta_schedule.arg_info import TensorInfo
@@ -58,6 +58,7 @@ from tvm.s_tir.meta_schedule.testing.local_rpc import LocalRPC
from tvm.s_tir.meta_schedule.utils import (
get_global_func_with_default_on_worker,
)
+from tvm.s_tir.utils import derived_object
from tvm.script import s_tir as Ts
from tvm.script import tirx as T
from tvm.target import Target
@@ -409,7 +410,7 @@ def test_meta_schedule_py_runner():
@pytest.mark.skip("Tuning test - launches runner")
-@tvm.testing.skip_if_32bit(reason="skipping test for i386.")
+@pytest.mark.skipif("32bit" in platform.architecture()[0], reason="skipping test for i386.")
def test_meta_schedule_rpc_runner_time_out():
"""Test meta schedule RPC Runner time out by using a super large workload"""
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_search_strategy.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_search_strategy.py
index e26a02967a..d3df48067f 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_search_strategy.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_search_strategy.py
@@ -23,10 +23,10 @@ import pytest
import tvm
import tvm.testing
-from tvm.ir.utils import derived_object
from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.meta_schedule.testing.dummy_object import DummyMutator
from tvm.s_tir.schedule import Schedule, Trace
+from tvm.s_tir.utils import derived_object
from tvm.script import s_tir as Ts
from tvm.script import tirx as T
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_space_generator.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_space_generator.py
index 6f631bce0d..3952ef4440 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_space_generator.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_space_generator.py
@@ -23,7 +23,6 @@ import pytest
import tvm
import tvm.testing
-from tvm.ir.utils import derived_object
from tvm.s_tir.meta_schedule.space_generator import (
PySpaceGenerator,
ScheduleFn,
@@ -31,6 +30,7 @@ from tvm.s_tir.meta_schedule.space_generator import (
)
from tvm.s_tir.meta_schedule.tune_context import TuneContext
from tvm.s_tir.schedule import Schedule
+from tvm.s_tir.utils import derived_object
from tvm.script import s_tir as Ts
from tvm.script import tirx as T
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_task_scheduler.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_task_scheduler.py
index 13ff890b40..f17c3d54a6 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_task_scheduler.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_task_scheduler.py
@@ -24,10 +24,10 @@ import pytest
import tvm
import tvm.testing
-from tvm.ir.utils import derived_object
from tvm.s_tir import Schedule
from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.meta_schedule.testing.dummy_object import DummyBuilder, DummyRunner
+from tvm.s_tir.utils import derived_object
from tvm.script import s_tir as Ts
from tvm.script import tirx as T
diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_tune_tir.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_tune_tir.py
index 83c87877ed..7debafbbca 100644
--- a/tests/python/s_tir/meta_schedule/test_meta_schedule_tune_tir.py
+++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_tune_tir.py
@@ -25,11 +25,11 @@ pytest.importorskip("tornado") # tvm.rpc.tracker (LocalRPC) requires tornado
import tvm
import tvm.testing
-from tvm.ir.utils import derived_object
from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.meta_schedule.testing.custom_builder_runner import run_module_via_rpc
from tvm.s_tir.meta_schedule.testing.local_rpc import LocalRPC
from tvm.s_tir.schedule import SBlockRV, Schedule
+from tvm.s_tir.utils import derived_object
from tvm.script import s_tir as Ts
from tvm.script import tirx as T
from tvm.target import Target
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_tensorize_ldmatrix_mma_numeric.py b/tests/python/s_tir/schedule/test_tir_schedule_tensorize_ldmatrix_mma_numeric.py
index f9b1212f59..685c800519 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_tensorize_ldmatrix_mma_numeric.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_tensorize_ldmatrix_mma_numeric.py
@@ -50,8 +50,8 @@ from tvm.s_tir.tensor_intrin.cuda import (
shared_16x32_to_ldmatrix_32x16_layout,
shared_32x16_to_ldmatrix_32x16_layout,
)
+from tvm.s_tir.testing.tir import mma_schedule
from tvm.testing import env
-from tvm.testing.tir import mma_schedule
M = 4096
N = 4096
diff --git a/tests/python/s_tir/schedule/test_tir_schedule_tensorize_mfma_numeric.py b/tests/python/s_tir/schedule/test_tir_schedule_tensorize_mfma_numeric.py
index 325ab5428e..0121b0e4e1 100644
--- a/tests/python/s_tir/schedule/test_tir_schedule_tensorize_mfma_numeric.py
+++ b/tests/python/s_tir/schedule/test_tir_schedule_tensorize_mfma_numeric.py
@@ -41,8 +41,8 @@ from tvm.s_tir.tensor_intrin.rocm import (
shared_16x16_to_local_64x4_layout_B,
shared_16x16_to_local_64x4_layout_C,
)
+from tvm.s_tir.testing.tir import mfma_schedule
from tvm.testing import env
-from tvm.testing.tir import mfma_schedule
M = 1024
N = 1024
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_compact_tensor_region.py b/tests/python/s_tir/transform/test_s_tir_transform_compact_tensor_region.py
index fb121876c7..024fc5bf28 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_compact_tensor_region.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_compact_tensor_region.py
@@ -39,7 +39,7 @@ class BaseCompactTest:
before = tvm.IRModule.from_expr(self.before.with_attr("global_symbol", "main"))
expected = tvm.IRModule.from_expr(self.expected.with_attr("global_symbol", "main"))
- simplify = tvm.transform.Sequential(
+ simplify = tvm.ir.transform.Sequential(
[s_tir.transform.StmtSimplify(), tirx.transform.RemoveNoOp()]
)
after = simplify(s_tir.transform.CompactTensorAllocation(is_strict=is_strict)(before))
@@ -1031,7 +1031,7 @@ class TestTileAwareCompaction(BaseCompactTest):
# Get partitioned workload to compact
mod = tvm.IRModule.from_expr(main)
- with tvm.transform.PassContext(
+ with tvm.ir.transform.PassContext(
config={"s_tir.LoopPartition": {"partition_const_loop": True}}
):
mod = tvm.s_tir.transform.LowerOpaqueBlock()(mod)
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_default_gpu_schedule.py b/tests/python/s_tir/transform/test_s_tir_transform_default_gpu_schedule.py
index 1c3a4beb6a..f1d8349c47 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_default_gpu_schedule.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_default_gpu_schedule.py
@@ -66,7 +66,7 @@ def test_broadcast_to_symbolic():
# fmt: on
# pylint: enable=no-self-argument,missing-class-docstring,line-too-long
target = tvm.target.Target("nvidia/geforce-rtx-3070")
- with target, tvm.transform.PassContext(opt_level=3):
+ with target, tvm.ir.transform.PassContext(opt_level=3):
After = DefaultGPUSchedule()(Before)
tvm.ir.assert_structural_equal(After, Expected)
@@ -195,7 +195,7 @@ def test_matmul():
# fmt: on
# pylint: enable=no-self-argument,missing-class-docstring,line-too-long
target = tvm.target.Target("nvidia/geforce-rtx-3070")
- with target, tvm.transform.PassContext(opt_level=3):
+ with target, tvm.ir.transform.PassContext(opt_level=3):
After = DefaultGPUSchedule()(Before)
tvm.ir.assert_structural_equal(After, Expected)
@@ -265,7 +265,7 @@ def test_add():
# fmt: on
# pylint: enable=no-self-argument,missing-class-docstring,line-too-long
target = tvm.target.Target("nvidia/geforce-rtx-3070")
- with target, tvm.transform.PassContext(opt_level=3):
+ with target, tvm.ir.transform.PassContext(opt_level=3):
After = DefaultGPUSchedule()(Before)
tvm.ir.assert_structural_equal(After, Expected)
@@ -312,7 +312,7 @@ def test_full():
# fmt: on
# pylint: enable=no-self-argument,missing-class-docstring,line-too-long
target = tvm.target.Target("nvidia/geforce-rtx-3070")
- with target, tvm.transform.PassContext(opt_level=3):
+ with target, tvm.ir.transform.PassContext(opt_level=3):
After = DefaultGPUSchedule()(Before)
tvm.ir.assert_structural_equal(After, Expected)
@@ -348,7 +348,7 @@ def test_scheduled():
# fmt: on
# pylint: enable=no-self-argument,missing-class-docstring,line-too-long
target = tvm.target.Target("nvidia/geforce-rtx-3070")
- with target, tvm.transform.PassContext(opt_level=3):
+ with target, tvm.ir.transform.PassContext(opt_level=3):
# should do nothing
After = DefaultGPUSchedule()(Scheduled)
tvm.ir.assert_structural_equal(After, Scheduled)
@@ -450,7 +450,7 @@ def test_multiple():
# fmt: on
# pylint: enable=no-self-argument,missing-class-docstring,line-too-long
target = tvm.target.Target("nvidia/geforce-rtx-3070")
- with target, tvm.transform.PassContext(opt_level=3):
+ with target, tvm.ir.transform.PassContext(opt_level=3):
After = DefaultGPUSchedule()(Before)
tvm.ir.assert_structural_equal(After, Expected)
@@ -488,7 +488,7 @@ def test_add_on_metal():
# fmt: on
# pylint: enable=no-self-argument,missing-class-docstring,line-too-long
target = tvm.target.Target("apple/m1-gpu")
- with target, tvm.transform.PassContext(opt_level=0):
+ with target, tvm.ir.transform.PassContext(opt_level=0):
mod = DefaultGPUSchedule()(Before)
tvm.ir.assert_structural_equal(mod, Expected)
@@ -523,7 +523,7 @@ def test_scalar_add():
# fmt: on
# pylint: enable=no-self-argument,missing-class-docstring,line-too-long
target = tvm.target.Target("nvidia/geforce-rtx-3070")
- with target, tvm.transform.PassContext(opt_level=0):
+ with target, tvm.ir.transform.PassContext(opt_level=0):
mod = DefaultGPUSchedule()(Before)
tvm.ir.assert_structural_equal(mod, Expected)
@@ -562,7 +562,7 @@ def test_sum():
# fmt: on
# pylint: enable=no-self-argument,missing-class-docstring,line-too-long
target = tvm.target.Target("nvidia/geforce-rtx-3070")
- with target, tvm.transform.PassContext(opt_level=0):
+ with target, tvm.ir.transform.PassContext(opt_level=0):
mod = DefaultGPUSchedule()(Before)
tvm.ir.assert_structural_equal(mod, Expected)
@@ -596,7 +596,7 @@ def test_scalar_block_no_loops():
# fmt: on
# pylint: enable=no-self-argument,missing-class-docstring,line-too-long
target = tvm.target.Target("nvidia/geforce-rtx-3070")
- with target, tvm.transform.PassContext(opt_level=0):
+ with target, tvm.ir.transform.PassContext(opt_level=0):
mod = DefaultGPUSchedule()(Before)
tvm.ir.assert_structural_equal(mod, Expected)
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py b/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py
index 0b4d50643e..d58a5c1f39 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py
@@ -21,6 +21,7 @@ import pytest
import tvm_ffi
import tvm
+import tvm.s_tir.testing
import tvm.testing
from tvm import s_tir
from tvm.script import ir as I
@@ -157,10 +158,10 @@ def test_inject_async_copy():
assert count_cp_async(mod["main"].body) == 1
- if not tvm.testing.is_ampere_or_newer():
+ if not tvm.s_tir.testing.is_ampere_or_newer():
continue
- with tvm.transform.PassContext(config={"tirx.use_async_copy": 1}):
+ with tvm.ir.transform.PassContext(config={"tirx.use_async_copy": 1}):
mod = tvm.compile(tvm.IRModule.from_expr(f), target="cuda")
A_np = np.random.rand(32, 128).astype(dtype)
@@ -191,10 +192,10 @@ def test_inject_async_copy_shared_dyn():
assert count_cp_async(mod["main"].body) == 2
- if not tvm.testing.is_ampere_or_newer():
+ if not tvm.s_tir.testing.is_ampere_or_newer():
return
- with tvm.transform.PassContext(config={"tirx.use_async_copy": 1}):
+ with tvm.ir.transform.PassContext(config={"tirx.use_async_copy": 1}):
mod = tvm.compile(tvm.IRModule.from_expr(f), target="cuda")
A_np = np.random.rand(32, 128).astype("float16")
@@ -407,7 +408,7 @@ def test_cp_async_in_if_then_else(postproc_if_missing_async_support):
C[tx, i] = A_shared[tx, 0] + B_shared[tx, 0]
mod = tvm.IRModule.from_expr(simple_compute)
- with tvm.transform.PassContext(config={"tirx.use_async_copy": 1}):
+ with tvm.ir.transform.PassContext(config={"tirx.use_async_copy": 1}):
tvm.compile(mod, target="cuda")
generated_code = postproc_if_missing_async_support()
print(generated_code)
@@ -870,7 +871,7 @@ def test_vectorize_cp_async_in_if_then_else(postproc_if_missing_async_support):
)
mod = tvm.IRModule.from_expr(complex_compute)
- with tvm.transform.PassContext(config={"tirx.use_async_copy": 1}):
+ with tvm.ir.transform.PassContext(config={"tirx.use_async_copy": 1}):
tvm.compile(mod, target="cuda")
generated_code = postproc_if_missing_async_support()
# generated_code must contain " setp.ne.b32 p, %0, 0;"
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_inject_software_pipeline.py b/tests/python/s_tir/transform/test_s_tir_transform_inject_software_pipeline.py
index 10a0fee3c4..ddbfd09f4b 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_inject_software_pipeline.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_inject_software_pipeline.py
@@ -22,6 +22,7 @@ import pytest
import tvm
import tvm.s_tir.tensor_intrin.cuda
+import tvm.s_tir.testing
import tvm.testing
from tvm import te, tirx
from tvm.s_tir.meta_schedule.testing import te_workload
@@ -33,10 +34,10 @@ from tvm.s_tir.tensor_intrin.cuda import (
MMA_store_16x16_f32_global_INTRIN,
shared_16x16_to_ldmatrix_32x8_layout,
)
+from tvm.s_tir.testing.tir import mma_schedule
from tvm.script import s_tir as Ts
from tvm.script import tirx as T
from tvm.testing import env
-from tvm.testing.tir import mma_schedule
def _check(original, transformed):
@@ -1514,8 +1515,8 @@ def get_mma_schedule():
def build_and_run(sch):
- if tvm.testing.is_ampere_or_newer():
- with tvm.transform.PassContext(config={"tirx.use_async_copy": 1}):
+ if tvm.s_tir.testing.is_ampere_or_newer():
+ with tvm.ir.transform.PassContext(config={"tirx.use_async_copy": 1}):
f = tvm.compile(sch.mod["main"], target="cuda")
a_np = np.random.uniform(size=(N, K)).astype("float16")
@@ -1544,7 +1545,7 @@ def test_async_pipelined_mma_gemm_simple():
sch.annotate(k0, ann_key="software_pipeline_order", ann_val=[0, 1, 2])
sch.annotate(k0, ann_key="software_pipeline_async_stages", ann_val=[0])
- seq = tvm.transform.Sequential(
+ seq = tvm.ir.transform.Sequential(
[
tvm.s_tir.transform.PlanAndUpdateTensorAllocationLocation(),
tvm.s_tir.transform.ConvertBlocksToOpaque(),
@@ -1596,7 +1597,7 @@ def test_async_nested_pipeline_mma_gemm_ideal_annotation():
sch.annotate(k1, ann_key="software_pipeline_stage", ann_val=[0, 0, 1])
sch.annotate(k1, ann_key="software_pipeline_order", ann_val=[0, 1, 2])
- seq = tvm.transform.Sequential(
+ seq = tvm.ir.transform.Sequential(
[
tvm.s_tir.transform.PlanAndUpdateTensorAllocationLocation(),
tvm.s_tir.transform.ConvertBlocksToOpaque(),
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_loop_partition.py b/tests/python/s_tir/transform/test_s_tir_transform_loop_partition.py
index 3e3895c844..bb5901f724 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_loop_partition.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_loop_partition.py
@@ -93,7 +93,9 @@ def test_condition_EQ():
T.evaluate(T.Select(T.likely(i == 5), m, n))
mod = tvm.IRModule.from_expr(func.with_attr("global_symbol", "main"))
- with tvm.transform.PassContext(config={"s_tir.LoopPartition": {"partition_const_loop": True}}):
+ with tvm.ir.transform.PassContext(
+ config={"s_tir.LoopPartition": {"partition_const_loop": True}}
+ ):
mod = tvm.s_tir.transform.LoopPartition()(mod)
stmt = tvm.tirx.transform.StmtSimplify()(mod)["main"].body
@@ -139,7 +141,9 @@ def test_oneD_pool():
mod = tvm.IRModule.from_expr(func.with_attr("global_symbol", "main"))
- with tvm.transform.PassContext(config={"s_tir.LoopPartition": {"partition_const_loop": True}}):
+ with tvm.ir.transform.PassContext(
+ config={"s_tir.LoopPartition": {"partition_const_loop": True}}
+ ):
mod = tvm.s_tir.transform.LoopPartition()(mod)
stmt = tvm.tirx.transform.StmtSimplify()(mod)["main"].body
@@ -160,7 +164,9 @@ def test_cce_loop_1():
)
mod = tvm.IRModule.from_expr(func.with_attr("global_symbol", "main"))
- with tvm.transform.PassContext(config={"s_tir.LoopPartition": {"partition_const_loop": True}}):
+ with tvm.ir.transform.PassContext(
+ config={"s_tir.LoopPartition": {"partition_const_loop": True}}
+ ):
mod = tvm.s_tir.transform.LoopPartition()(mod)
stmt = tvm.tirx.transform.StmtSimplify()(mod)["main"].body
@@ -181,7 +187,9 @@ def test_cce_loop_2():
T.evaluate(T.call_extern("cce_intrisic", i * tile, i * tile + tile, ty="float32"))
mod = tvm.IRModule.from_expr(func.with_attr("global_symbol", "main"))
- with tvm.transform.PassContext(config={"s_tir.LoopPartition": {"partition_const_loop": True}}):
+ with tvm.ir.transform.PassContext(
+ config={"s_tir.LoopPartition": {"partition_const_loop": True}}
+ ):
mod = tvm.s_tir.transform.LoopPartition()(mod)
stmt = tvm.tirx.transform.StmtSimplify()(mod)["main"].body
@@ -202,7 +210,9 @@ def test_cce_loop_3():
mod = tvm.IRModule.from_expr(func.with_attr("global_symbol", "main"))
- with tvm.transform.PassContext(config={"s_tir.LoopPartition": {"partition_const_loop": True}}):
+ with tvm.ir.transform.PassContext(
+ config={"s_tir.LoopPartition": {"partition_const_loop": True}}
+ ):
mod = tvm.s_tir.transform.LoopPartition()(mod)
stmt = tvm.tirx.transform.StmtSimplify()(mod)["main"].body
@@ -221,7 +231,7 @@ def partitioned_concat(
def partition_from_scheduled_tir(function, pass_cfg, do_flatten=True):
- with tvm.transform.PassContext(config=pass_cfg):
+ with tvm.ir.transform.PassContext(config=pass_cfg):
mod = IRModule.from_expr(function.with_attr("global_symbol", "main"))
mod = tvm.s_tir.transform.LoopPartition()(mod)
mod = tvm.s_tir.transform.LowerOpaqueBlock()(mod)
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_thread_sync.py b/tests/python/s_tir/transform/test_s_tir_transform_thread_sync.py
index 93db8d27d5..4c8d9168ee 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_thread_sync.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_thread_sync.py
@@ -268,7 +268,7 @@ def test_hoist_if_preserves_sync_placement(loop):
# Each iteration must keep its own guard so synchronization can be inserted
# between shared-memory accesses, outside conditional control flow.
- finish = tvm.transform.Sequential(
+ finish = tvm.ir.transform.Sequential(
[tvm.tirx.transform.UnrollLoop(), tvm.s_tir.transform.ThreadSync("shared")]
)
Expected = finish(Before)
diff --git a/tests/python/sym/test_sym_deduce_bound.py b/tests/python/sym/test_sym_deduce_bound.py
index 3b840543d4..e6f34343bf 100644
--- a/tests/python/sym/test_sym_deduce_bound.py
+++ b/tests/python/sym/test_sym_deduce_bound.py
@@ -38,30 +38,30 @@ def test_deduce():
e0 = (-b) * a + c - d
res0 = tvm.sym.deduce_bound(a, e0 >= 0, {b: b_s, c: c_s, d: d_s}, {})
ans0 = fdiv(d - c, b * -1)
- tvm.testing.assert_prim_expr_equal(res0.max_value, ans0)
+ assert tvm.sym.Analyzer().can_prove_equal(res0.max_value, ans0)
# expression containing variable a is on rhs
res0 = tvm.sym.deduce_bound(a, zero <= e0, {b: b_s, c: c_s, d: d_s}, {})
- tvm.testing.assert_prim_expr_equal(res0.max_value, ans0)
+ assert tvm.sym.Analyzer().can_prove_equal(res0.max_value, ans0)
e0 = d * a + c - d
res0 = tvm.sym.deduce_bound(a, e0 >= 0, {b: b_s, c: c_s, d: d_s}, {})
ans0 = fdiv(d - c, d)
- tvm.testing.assert_prim_expr_equal(res0.max_value, ans0)
+ assert tvm.sym.Analyzer().can_prove_equal(res0.max_value, ans0)
# expression containing variable a is on rhs
res0 = tvm.sym.deduce_bound(a, zero <= e0, {b: b_s, c: c_s, d: d_s}, {})
- tvm.testing.assert_prim_expr_equal(res0.max_value, ans0)
+ assert tvm.sym.Analyzer().can_prove_equal(res0.max_value, ans0)
e1 = a * 4 + b < c
res1 = tvm.sym.deduce_bound(a, e1, {b: b_s, c: c_s, d: d_s}, {})
ans1 = fdiv(c - 1 - b, 4)
- tvm.testing.assert_prim_expr_equal(res1.max_value, ans1)
+ assert tvm.sym.Analyzer().can_prove_equal(res1.max_value, ans1)
# expression containing variable a is on rhs
e1 = c > a * 4 + b
res1 = tvm.sym.deduce_bound(a, e1, {b: b_s, c: c_s, d: d_s}, {})
- tvm.testing.assert_prim_expr_equal(res1.max_value, ans1)
+ assert tvm.sym.Analyzer().can_prove_equal(res1.max_value, ans1)
e2 = tvm.tirx.max(5, a * 4) < 0
res2 = tvm.sym.deduce_bound(a, e2, {b: b_s, c: c_s, d: d_s}, {})
@@ -77,15 +77,15 @@ def test_deduce():
e3 = (-b) + a * c - d
res3 = tvm.sym.deduce_bound(a, e3 >= 0, {b: b_s, c: c_s, d: d_s}, {b: b_s, d: d_s})
ans3 = fdiv(2, c) + 1
- tvm.testing.assert_prim_expr_equal(res3.min_value, ans3)
+ assert tvm.sym.Analyzer().can_prove_equal(res3.min_value, ans3)
res3 = tvm.sym.deduce_bound(a, zero <= e3, {b: b_s, c: c_s, d: d_s}, {b: b_s, d: d_s})
- tvm.testing.assert_prim_expr_equal(res3.min_value, ans3)
+ assert tvm.sym.Analyzer().can_prove_equal(res3.min_value, ans3)
# tests for `EQ` op
res4 = tvm.sym.deduce_bound(a, a == b, {}, {})
- tvm.testing.assert_prim_expr_equal(res4.max_value, b)
- tvm.testing.assert_prim_expr_equal(res4.min_value, b)
+ assert tvm.sym.Analyzer().can_prove_equal(res4.max_value, b)
+ assert tvm.sym.Analyzer().can_prove_equal(res4.min_value, b)
# Unsatisfiable `EQ`, variable as one of the Operand
res5 = tvm.sym.deduce_bound(a, (a == b), {b: b_s}, {b: b_s})
@@ -94,20 +94,20 @@ def test_deduce():
# variable `a` on the RHS side
res6 = tvm.sym.deduce_bound(a, 10 == a, {}, {})
- tvm.testing.assert_prim_expr_equal(res6.max_value, 10)
- tvm.testing.assert_prim_expr_equal(res6.min_value, 10)
+ assert tvm.sym.Analyzer().can_prove_equal(res6.max_value, 10)
+ assert tvm.sym.Analyzer().can_prove_equal(res6.min_value, 10)
# Add, Sub in `EQ`
e4 = (a - c) == (b + d)
ans4 = b + d + c
res7 = tvm.sym.deduce_bound(a, e4, {b: b_s, c: c_s, d: d_s}, {})
- tvm.testing.assert_prim_expr_equal(res7.max_value, ans4)
- tvm.testing.assert_prim_expr_equal(res7.min_value, ans4)
+ assert tvm.sym.Analyzer().can_prove_equal(res7.max_value, ans4)
+ assert tvm.sym.Analyzer().can_prove_equal(res7.min_value, ans4)
# Satisfiable Mul in `EQ` with negative sign
res8 = tvm.sym.deduce_bound(a, (5 * a == -10), {}, {})
- tvm.testing.assert_prim_expr_equal(res8.max_value, -2)
- tvm.testing.assert_prim_expr_equal(res8.min_value, -2)
+ assert tvm.sym.Analyzer().can_prove_equal(res8.max_value, -2)
+ assert tvm.sym.Analyzer().can_prove_equal(res8.min_value, -2)
# Unsatisfiable Mul in `EQ`
e5 = 4 * a == b
@@ -117,8 +117,8 @@ def test_deduce():
res10 = tvm.sym.deduce_bound(a, (b * a == b), {b: b_s}, {})
# simplifier is now able to prove symbolic relation (b * a % b == 0)
- tvm.testing.assert_prim_expr_equal(res10.max_value, 1)
- tvm.testing.assert_prim_expr_equal(res10.min_value, 1)
+ assert tvm.sym.Analyzer().can_prove_equal(res10.max_value, 1)
+ assert tvm.sym.Analyzer().can_prove_equal(res10.min_value, 1)
def test_check():
@@ -153,22 +153,22 @@ def test_deduce_basic():
res1 = tvm.sym.deduce_bound(a, e0 < 17, {b: b_s}, {b: b_s})
[x, y] = [res1.max_value, b_s.max_value] if coff > 0 else [res1.min_value, b_s.min_value]
- tvm.testing.assert_prim_expr_equal((x * coff + 3 + y) < 17, True)
+ assert tvm.sym.Analyzer().can_prove_equal((x * coff + 3 + y) < 17, True)
# expression containing variable a is on rhs
res1 = tvm.sym.deduce_bound(a, tvm.tirx.const(17, "int32") < e0, {b: b_s}, {b: b_s})
[x, y] = [res1.max_value, b_s.max_value] if coff < 0 else [res1.min_value, b_s.min_value]
- tvm.testing.assert_prim_expr_equal((x * coff + 3 + y) > 17, True)
+ assert tvm.sym.Analyzer().can_prove_equal((x * coff + 3 + y) > 17, True)
# expression containing variable a is on rhs
res1 = tvm.sym.deduce_bound(a, tvm.tirx.const(17, "int32") >= e0, {b: b_s}, {b: b_s})
[x, y] = [res1.max_value, b_s.max_value] if coff > 0 else [res1.min_value, b_s.min_value]
- tvm.testing.assert_prim_expr_equal((x * coff + 3 + y) <= 17, True)
+ assert tvm.sym.Analyzer().can_prove_equal((x * coff + 3 + y) <= 17, True)
res1 = tvm.sym.deduce_bound(a, e0 >= 17, {b: b_s}, {b: b_s})
[x, y] = [res1.max_value, b_s.max_value] if coff < 0 else [res1.min_value, b_s.min_value]
- tvm.testing.assert_prim_expr_equal((x * coff + 3 + y) >= 17, True)
+ assert tvm.sym.Analyzer().can_prove_equal((x * coff + 3 + y) >= 17, True)
test_basic(0, 4, 4)
test_basic(1, 5, 4)
@@ -187,21 +187,21 @@ def test_deduce_complex():
res1 = tvm.sym.deduce_bound(a, e0 < 63, {b: b_s}, {b: b_s})
[t, x] = [res1.max_value, b_s.max_value] if coff > 0 else [res1.min_value, b_s.min_value]
- tvm.testing.assert_prim_expr_equal(((x * 3 + t * coff) * 4) < 63, True)
+ assert tvm.sym.Analyzer().can_prove_equal(((x * 3 + t * coff) * 4) < 63, True)
# expression containing variable a is on rhs
res1 = tvm.sym.deduce_bound(a, tvm.tirx.const(63, "int32") >= e0, {b: b_s}, {b: b_s})
[t, x] = [res1.max_value, b_s.max_value] if coff > 0 else [res1.min_value, b_s.min_value]
- tvm.testing.assert_prim_expr_equal(((x * 3 + t * coff) * 4) <= 63, True)
+ assert tvm.sym.Analyzer().can_prove_equal(((x * 3 + t * coff) * 4) <= 63, True)
res1 = tvm.sym.deduce_bound(a, e0 > 63, {b: b_s}, {b: b_s})
[t, x] = [res1.max_value, b_s.max_value] if coff < 0 else [res1.min_value, b_s.min_value]
- tvm.testing.assert_prim_expr_equal(((x * 3 + t * coff) * 4) > 63, True)
+ assert tvm.sym.Analyzer().can_prove_equal(((x * 3 + t * coff) * 4) > 63, True)
# expression containing variable a is on rhs
res1 = tvm.sym.deduce_bound(a, tvm.tirx.const(63, "int32") <= e0, {b: b_s}, {b: b_s})
[t, x] = [res1.max_value, b_s.max_value] if coff < 0 else [res1.min_value, b_s.min_value]
- tvm.testing.assert_prim_expr_equal(((x * 3 + t * coff) * 4) >= 63, True)
+ assert tvm.sym.Analyzer().can_prove_equal(((x * 3 + t * coff) * 4) >= 63, True)
test_complex(0, 4, 4)
test_complex(0, 4, -4)
@@ -239,11 +239,11 @@ def test_deduce_floordiv():
if isinstance(expect_min, str):
assert res.min_value.name == expect_min
else:
- tvm.testing.assert_prim_expr_equal(res.min_value, expect_min)
+ assert tvm.sym.Analyzer().can_prove_equal(res.min_value, expect_min)
if isinstance(expect_max, str):
assert res.max_value.name == expect_max
else:
- tvm.testing.assert_prim_expr_equal(res.max_value, expect_max)
+ assert tvm.sym.Analyzer().can_prove_equal(res.max_value, expect_max)
# test basic cases
do_test(lambda a: a // 8 > 3, {}, 32, "pos_inf")
diff --git a/tests/python/sym/test_sym_detect_clip_bound.py b/tests/python/sym/test_sym_detect_clip_bound.py
index eca33c54f2..31bf175bf3 100644
--- a/tests/python/sym/test_sym_detect_clip_bound.py
+++ b/tests/python/sym/test_sym_detect_clip_bound.py
@@ -23,32 +23,32 @@ def test_basic():
b = tvm.tirx.Var("b", "int32")
c = tvm.tirx.Var("c", "int32")
m = tvm.sym.detect_clip_bound(tvm.tirx.all(a * 1 < b * 6, a - 1 > 0), [a])
- tvm.testing.assert_prim_expr_equal(m[1], b * 6 - 1)
+ assert tvm.sym.Analyzer().can_prove_equal(m[1], b * 6 - 1)
assert m[0].value == 2
m = tvm.sym.detect_clip_bound(tvm.tirx.all(a * 1 < b * 6, a - 1 > 0), [a, b])
assert len(m) == 0
m = tvm.sym.detect_clip_bound(tvm.tirx.all(a + 10 * c <= 20, b - 1 > 0), [a, b])
- tvm.testing.assert_prim_expr_equal(m[1], 20 - 10 * c)
- tvm.testing.assert_prim_expr_equal(m[2], 2)
+ assert tvm.sym.Analyzer().can_prove_equal(m[1], 20 - 10 * c)
+ assert tvm.sym.Analyzer().can_prove_equal(m[2], 2)
m = tvm.sym.detect_clip_bound(tvm.tirx.all(tvm.tirx.Not(a * 1 > b * 6), a - 1 > 0), [a])
- tvm.testing.assert_prim_expr_equal(m[1], b * 6)
+ assert tvm.sym.Analyzer().can_prove_equal(m[1], b * 6)
m = tvm.sym.detect_clip_bound(tvm.tirx.all(tvm.tirx.Min(a, b) > 3, a - 10 < 0), [a, b])
- tvm.testing.assert_prim_expr_equal(m[0], 4)
- tvm.testing.assert_prim_expr_equal(m[1], 9)
- tvm.testing.assert_prim_expr_equal(m[2], 4)
+ assert tvm.sym.Analyzer().can_prove_equal(m[0], 4)
+ assert tvm.sym.Analyzer().can_prove_equal(m[1], 9)
+ assert tvm.sym.Analyzer().can_prove_equal(m[2], 4)
def test_trivial_eq():
a = tvm.tirx.Var("a", "int32")
b = tvm.tirx.Var("b", "int32")
m = tvm.sym.detect_clip_bound(b == 3, [a, b])
- tvm.testing.assert_prim_expr_equal(m[2], 3)
- tvm.testing.assert_prim_expr_equal(m[3], 3)
+ assert tvm.sym.Analyzer().can_prove_equal(m[2], 3)
+ assert tvm.sym.Analyzer().can_prove_equal(m[3], 3)
m = tvm.sym.detect_clip_bound(tvm.tirx.all(a == 4, b == 3), [a, b])
- tvm.testing.assert_prim_expr_equal(m[0], 4)
- tvm.testing.assert_prim_expr_equal(m[1], 4)
- tvm.testing.assert_prim_expr_equal(m[2], 3)
- tvm.testing.assert_prim_expr_equal(m[3], 3)
+ assert tvm.sym.Analyzer().can_prove_equal(m[0], 4)
+ assert tvm.sym.Analyzer().can_prove_equal(m[1], 4)
+ assert tvm.sym.Analyzer().can_prove_equal(m[2], 3)
+ assert tvm.sym.Analyzer().can_prove_equal(m[3], 3)
if __name__ == "__main__":
diff --git a/tests/python/sym/test_sym_detect_linear_equation.py b/tests/python/sym/test_sym_detect_linear_equation.py
index 9e3154b6e9..2f6ab64174 100644
--- a/tests/python/sym/test_sym_detect_linear_equation.py
+++ b/tests/python/sym/test_sym_detect_linear_equation.py
@@ -23,14 +23,14 @@ def test_basic():
b = tvm.tirx.Var("b", "int32")
m = tvm.sym.detect_linear_equation(a * 4 + b * 6 + 7, [a])
assert m[0].value == 4
- tvm.testing.assert_prim_expr_equal(m[1], b * 6 + 7)
+ assert tvm.sym.Analyzer().can_prove_equal(m[1], b * 6 + 7)
m = tvm.sym.detect_linear_equation(a * 4 * (a + 1) + b * 6 + 7, [a])
assert len(m) == 0
m = tvm.sym.detect_linear_equation(a * 4 + (a + 1) + b * 6 + 7, [a])
assert m[0].value == 5
- tvm.testing.assert_prim_expr_equal(m[1], b * 6 + 7 + 1)
+ assert tvm.sym.Analyzer().can_prove_equal(m[1], b * 6 + 7 + 1)
m = tvm.sym.detect_linear_equation(a * b + 7, [a])
assert m[0] == b
@@ -40,7 +40,7 @@ def test_basic():
m = tvm.sym.detect_linear_equation(b * 7, [])
assert len(m) == 1
- tvm.testing.assert_prim_expr_equal(m[0], b * 7)
+ assert tvm.sym.Analyzer().can_prove_equal(m[0], b * 7)
c = tvm.tirx.Var("c", "uint32")
m = tvm.sym.detect_linear_equation(128 - c, [c])
@@ -52,7 +52,7 @@ def test_multivariate():
b = tvm.tirx.Var("b", "int32")
m = tvm.sym.detect_linear_equation(v[0] * (b + 4) + v[0] + v[1] * 8, v)
- tvm.testing.assert_prim_expr_equal(m[0], b + 5)
+ assert tvm.sym.Analyzer().can_prove_equal(m[0], b + 5)
assert m[1].value == 8
@@ -70,11 +70,11 @@ def test_multivariate():
m = tvm.sym.detect_linear_equation((v[0] - v[1]), [v[2]])
assert m[0].value == 0
- tvm.testing.assert_prim_expr_equal(m[1], v[0] - v[1])
+ assert tvm.sym.Analyzer().can_prove_equal(m[1], v[0] - v[1])
m = tvm.sym.detect_linear_equation((v[0] - v[1]), [])
assert len(m) == 1
- tvm.testing.assert_prim_expr_equal(m[0], v[0] - v[1])
+ assert tvm.sym.Analyzer().can_prove_equal(m[0], v[0] - v[1])
if __name__ == "__main__":
diff --git a/tests/python/sym/test_sym_iter_affine_map.py b/tests/python/sym/test_sym_iter_affine_map.py
index ff11396ca5..8e7417050c 100644
--- a/tests/python/sym/test_sym_iter_affine_map.py
+++ b/tests/python/sym/test_sym_iter_affine_map.py
@@ -76,9 +76,9 @@ def assert_iter_sum_pattern(
assert len(sum_expr.args) == 0
else:
assert len(sum_expr.args) == 1
- tvm.testing.assert_prim_expr_equal(sum_expr.args[0].extent, extent)
- tvm.testing.assert_prim_expr_equal(sum_expr.args[0].scale, scale)
- tvm.testing.assert_prim_expr_equal(sum_expr.base, base)
+ assert tvm.sym.Analyzer().can_prove_equal(sum_expr.args[0].extent, extent)
+ assert tvm.sym.Analyzer().can_prove_equal(sum_expr.args[0].scale, scale)
+ assert tvm.sym.Analyzer().can_prove_equal(sum_expr.base, base)
if expect_iter is not None:
if not isinstance(expect_iter, tvm.sym.IterMapExpr):
sum_expr = convert_iter_expr(sum_expr)
@@ -1344,11 +1344,11 @@ def assert_normalize_to_iter_sum(index, input_iters, args, base):
if isinstance(item, tvm.sym.IterSplitExpr):
tvm.ir.assert_structural_equal(split, item)
continue
- tvm.testing.assert_prim_expr_equal(split.scale, item[1])
- tvm.testing.assert_prim_expr_equal(
+ assert tvm.sym.Analyzer().can_prove_equal(split.scale, item[1])
+ assert tvm.sym.Analyzer().can_prove_equal(
tvm.sym.normalize_iter_map_to_expr(split), item[0] * item[1]
)
- tvm.testing.assert_prim_expr_equal(res.base, base)
+ assert tvm.sym.Analyzer().can_prove_equal(res.base, base)
def test_normalize_to_iter_sum():
@@ -1432,8 +1432,8 @@ def test_normalize_to_iter_sum_accepts_external_analyzer():
res = tvm.sym.normalize_to_iter_sum(i // tile, input_iters, analyzer=analyzer)
assert len(res.args) == 1
- tvm.testing.assert_prim_expr_equal(res.args[0].lower_factor, tile)
- tvm.testing.assert_prim_expr_equal(res.args[0].extent, T.int32(4))
+ assert tvm.sym.Analyzer().can_prove_equal(res.args[0].lower_factor, tile)
+ assert tvm.sym.Analyzer().can_prove_equal(res.args[0].extent, T.int32(4))
def test_detect_iter_map_with_bufferload_recursion():
diff --git a/tests/python/testing/test_env.py b/tests/python/testing/test_env.py
index 89b398285a..089144d19d 100644
--- a/tests/python/testing/test_env.py
+++ b/tests/python/testing/test_env.py
@@ -16,6 +16,11 @@
# under the License.
"""Tests for the thin ``tvm.testing.env`` capability probes."""
+import subprocess
+import sys
+import textwrap
+from types import SimpleNamespace
+
import pytest
import tvm
@@ -166,6 +171,146 @@ def test_probes_are_memoized():
assert info.hits + info.misses >= 1
+@pytest.fixture
+def clear_env_cache():
+ probes = (
+ env._target_enabled,
+ env._runtime_enabled,
+ env._device_exists,
+ env.build_flag_enabled,
+ env._llvm_version_major,
+ env.has_hexagon_toolchain,
+ env.has_hexagon,
+ )
+ for probe in probes:
+ probe.cache_clear()
+ yield
+ for probe in probes:
+ probe.cache_clear()
+
+
+def test_unavailable_selected_target_does_not_enable_llvm(monkeypatch, clear_env_cache):
+ monkeypatch.setenv("TVM_TEST_TARGETS", "cuda")
+ monkeypatch.setattr(tvm.runtime, "enabled", lambda kind: kind == "llvm")
+
+ def unexpected_device_lookup(target):
+ pytest.fail(f"Excluded or unavailable target was probed: {target}")
+
+ monkeypatch.setattr(tvm, "device_from_target", unexpected_device_lookup)
+ assert not tvm.testing.device_enabled("cuda")
+ assert not tvm.testing.device_enabled("llvm")
+
+
+@pytest.mark.parametrize("target_form", ["string", "dict", "target"])
+def test_device_enabled_accepts_target_forms(monkeypatch, clear_env_cache, target_form):
+ monkeypatch.setenv("TVM_TEST_TARGETS", '{"kind": "llvm"};cuda')
+ monkeypatch.setattr(tvm.runtime, "enabled", lambda kind: kind == "llvm")
+ monkeypatch.setattr(tvm, "device_from_target", lambda target: SimpleNamespace(exist=True))
+ target = {
+ "string": "llvm -num-cores=2",
+ "dict": {"kind": "llvm", "num-cores": 2},
+ "target": tvm.target.Target({"kind": "llvm", "num-cores": 2}),
+ }[target_form]
+ assert tvm.testing.device_enabled(target)
+
+
+@pytest.mark.parametrize("kind", ["llvm", "cuda"])
+@pytest.mark.parametrize(
+ ("runtime_enabled", "device_exists", "expected"),
+ [(False, True, False), (True, False, False), (True, True, True)],
+)
+def test_device_enabled_requires_runtime_and_device(
+ monkeypatch, clear_env_cache, kind, runtime_enabled, device_exists, expected
+):
+ # In particular, the CPU can exist even when LLVM was not built.
+ monkeypatch.setenv("TVM_TEST_TARGETS", kind)
+ monkeypatch.setattr(tvm.runtime, "enabled", lambda target: runtime_enabled)
+ monkeypatch.setattr(
+ tvm, "device_from_target", lambda target: SimpleNamespace(exist=device_exists)
+ )
+ assert tvm.testing.device_enabled(kind) is expected
+
+
+@pytest.mark.parametrize("cudnn_available", [None, False, True])
+def test_device_enabled_cudnn_runtime_probe(monkeypatch, clear_env_cache, cudnn_available):
+ monkeypatch.setenv("TVM_TEST_TARGETS", '{"kind": "cuda", "libs": ["cudnn"]}')
+ monkeypatch.setattr(tvm.runtime, "enabled", lambda kind: True)
+ monkeypatch.setattr(tvm, "device", lambda *args: SimpleNamespace(exist=True))
+ monkeypatch.setattr(tvm.support, "libinfo", lambda: {"USE_CUDNN": "ON"})
+ exists = None if cudnn_available is None else lambda: cudnn_available
+ monkeypatch.setattr(tvm, "get_global_func", lambda name, allow_missing: exists)
+ assert tvm.testing.device_enabled("cuda") is (cudnn_available is True)
+
+
+@pytest.mark.parametrize(
+ ("built", "toolchain", "llvm_version", "serial", "can_compile", "can_run"),
+ [
+ (False, True, 7, True, False, False),
+ (True, False, 7, True, False, False),
+ (True, True, 6, True, False, False),
+ (True, True, 7, False, True, False),
+ (True, True, 7, True, True, True),
+ (True, True, None, True, True, True),
+ ],
+)
+def test_hexagon_environment_requirements(
+ monkeypatch, clear_env_cache, built, toolchain, llvm_version, serial, can_compile, can_run
+):
+ # LLVM selection for tests is independent of whether its codegen can support Hexagon.
+ monkeypatch.setenv("TVM_TEST_TARGETS", "hexagon")
+ for variable, present in (
+ ("HEXAGON_TOOLCHAIN", toolchain),
+ ("ANDROID_SERIAL_NUMBER", serial),
+ ):
+ if present:
+ monkeypatch.setenv(variable, "test-value")
+ else:
+ monkeypatch.delenv(variable, raising=False)
+ monkeypatch.setattr(tvm.support, "libinfo", lambda: {"USE_HEXAGON": "ON" if built else "OFF"})
+ monkeypatch.setattr(tvm.runtime, "enabled", lambda kind: llvm_version is not None)
+ monkeypatch.setattr(tvm.target.codegen, "llvm_version_major", lambda: llvm_version)
+ assert env.has_hexagon_toolchain() is can_compile
+ assert env.has_hexagon() is can_run
+
+
+def test_core_testing_imports_without_domain_modules():
+ # Let TVM finish its own startup, then make the domain modules unavailable to testing.
+ code = textwrap.dedent(
+ """
+ import os
+ import sys
+ import tvm
+
+ for module in list(sys.modules):
+ if module == "tvm.testing" or module.startswith("tvm.testing."):
+ del sys.modules[module]
+ if hasattr(tvm, "testing"):
+ del tvm.testing
+
+ for name in ("te", "contrib", "s_tir"):
+ prefix = "tvm." + name
+ for module in list(sys.modules):
+ if module == prefix or module.startswith(prefix + "."):
+ sys.modules[module] = None
+ sys.modules[prefix] = None
+ if hasattr(tvm, name):
+ delattr(tvm, name)
+
+ import tvm.testing
+ from tvm.testing import env
+
+ tvm.testing.assert_allclose([1], [1])
+ tvm.support.libinfo = lambda: {"USE_HEXAGON": "ON"}
+ tvm.runtime.enabled = lambda kind: False
+ os.environ["HEXAGON_TOOLCHAIN"] = "test-toolchain"
+ os.environ["ANDROID_SERIAL_NUMBER"] = "test-device"
+ assert env.has_hexagon_toolchain()
+ assert env.has_hexagon()
+ """
+ )
+ subprocess.run([sys.executable, "-c", code], check=True, capture_output=True, text=True)
+
+
# --- demonstration of the target idiom -------------------------------------
#
# The standard gating idiom: a plain registered pytest marker (for ``-m``
diff --git a/tests/python/tirx-base/test_tir_base.py b/tests/python/tirx-base/test_tir_base.py
index 8efa9539fc..2ca2484113 100644
--- a/tests/python/tirx-base/test_tir_base.py
+++ b/tests/python/tirx-base/test_tir_base.py
@@ -23,8 +23,8 @@ import tvm_ffi
import tvm
from tvm import tirx
+from tvm.ir.transform import PassContext
from tvm.script import tirx as T
-from tvm.transform import PassContext
def build_tir_func(func):
diff --git a/tests/python/tirx-base/test_tir_intrin.py b/tests/python/tirx-base/test_tir_intrin.py
index 4d16a35dc6..5121132c57 100644
--- a/tests/python/tirx-base/test_tir_intrin.py
+++ b/tests/python/tirx-base/test_tir_intrin.py
@@ -332,7 +332,7 @@ class Module:
def test_fma():
- opt = tvm.transform.Sequential(
+ opt = tvm.ir.transform.Sequential(
[
tvm.tirx.transform.Apply(lambda f: f.with_attr("target", tvm.target.Target("llvm"))),
tvm.tirx.transform.LowerIntrin(),
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 cd2e104089..14fc6f8865 100644
--- a/tests/python/tirx-transform/test_tir_transform_flatten_tensor.py
+++ b/tests/python/tirx-transform/test_tir_transform_flatten_tensor.py
@@ -23,7 +23,7 @@ from tvm.script import tirx as T
def _transform():
- return tvm.transform.Sequential(
+ return tvm.ir.transform.Sequential(
[
tvm.tirx.transform.FlattenTensor(),
tvm.tirx.transform.StmtSimplify(),
diff --git a/tests/python/tirx-transform/test_tir_transform_function_pass.py b/tests/python/tirx-transform/test_tir_transform_function_pass.py
index 49ae2d655f..2c52fc2bc1 100644
--- a/tests/python/tirx-transform/test_tir_transform_function_pass.py
+++ b/tests/python/tirx-transform/test_tir_transform_function_pass.py
@@ -57,7 +57,9 @@ def test_cow_pass():
del func
# copy on write
mod_hash = mod.__hash__()
- mod = tvm.transform.Sequential([pidentity, tvm.tirx.transform.NarrowDataType(32)])(mod._move())
+ mod = tvm.ir.transform.Sequential([pidentity, tvm.tirx.transform.NarrowDataType(32)])(
+ mod._move()
+ )
assert mod_hash == mod.__hash__()
assert func_hash == mod["main"].__hash__()
diff --git a/tests/python/tirx-transform/test_tir_transform_lower_intrin.py b/tests/python/tirx-transform/test_tir_transform_lower_intrin.py
index cc53d0feec..a65bf63382 100644
--- a/tests/python/tirx-transform/test_tir_transform_lower_intrin.py
+++ b/tests/python/tirx-transform/test_tir_transform_lower_intrin.py
@@ -30,7 +30,7 @@ def lower_intrin(params, stmt):
mod = tvm.IRModule.from_expr(
tvm.tirx.Function(params, stmt).with_attr("target", tvm.target.Target("llvm"))
)
- mod = tvm.transform.Sequential(
+ mod = tvm.ir.transform.Sequential(
[tvm.tirx.transform.StmtSimplify(), tvm.tirx.transform.LowerIntrin()]
)(mod)
func = mod["main"]
diff --git a/tests/python/tirx-transform/test_tir_transform_remove_no_op.py b/tests/python/tirx-transform/test_tir_transform_remove_no_op.py
index 4a356d9674..a7d50bafb0 100644
--- a/tests/python/tirx-transform/test_tir_transform_remove_no_op.py
+++ b/tests/python/tirx-transform/test_tir_transform_remove_no_op.py
@@ -91,7 +91,7 @@ def _apply_remove_no_op(mod, max_simplification_steps=0):
"max_simplification_steps": max_simplification_steps,
}
}
- with tvm.transform.PassContext(config=config):
+ with tvm.ir.transform.PassContext(config=config):
mod = tvm.tirx.transform.RemoveNoOp()(mod)
return mod
diff --git a/tests/python/tirx-transform/test_tir_transform_simplify.py b/tests/python/tirx-transform/test_tir_transform_simplify.py
index f388c0ffbd..35dc9f1998 100644
--- a/tests/python/tirx-transform/test_tir_transform_simplify.py
+++ b/tests/python/tirx-transform/test_tir_transform_simplify.py
@@ -122,7 +122,7 @@ def _apply_simplify(
}
}
mod = tvm.IRModule.from_expr(func)
- with tvm.transform.PassContext(config=config):
+ with tvm.ir.transform.PassContext(config=config):
mod = tvm.tirx.transform.StmtSimplify()(mod)
return mod["main"]
diff --git a/tests/python/tirx-transform/test_tir_transform_unroll_loop.py b/tests/python/tirx-transform/test_tir_transform_unroll_loop.py
index 6833c4538a..9c3b0cfdf6 100644
--- a/tests/python/tirx-transform/test_tir_transform_unroll_loop.py
+++ b/tests/python/tirx-transform/test_tir_transform_unroll_loop.py
@@ -35,15 +35,15 @@ def test_unroll_loop():
assert isinstance(stmt, tvm.ir.For)
- with tvm.transform.PassContext(config={"tirx.UnrollLoop": {"auto_max_step": 16}}):
+ with tvm.ir.transform.PassContext(config={"tirx.UnrollLoop": {"auto_max_step": 16}}):
ret = tvm.tirx.transform.UnrollLoop()(mod)["main"].body
assert not any(isinstance(stmt, tvm.ir.For) for stmt in ret)
- with tvm.transform.PassContext(config={"tirx.UnrollLoop": {"auto_max_step": 15}}):
+ with tvm.ir.transform.PassContext(config={"tirx.UnrollLoop": {"auto_max_step": 15}}):
ret = tvm.tirx.transform.UnrollLoop()(mod)["main"].body
assert isinstance(ret[0], tvm.ir.For)
- with tvm.transform.PassContext(
+ with tvm.ir.transform.PassContext(
config={"tirx.UnrollLoop": {"auto_max_step": 16, "explicit_unroll": False}}
):
ret = tvm.tirx.transform.UnrollLoop()(mod)["main"].body
@@ -61,7 +61,7 @@ def test_unroll_loop():
for j in T.unroll(8):
Ab[j + 1] = Ab[i] + T.int64(1)
- with tvm.transform.PassContext(
+ with tvm.ir.transform.PassContext(
config={"tirx.UnrollLoop": {"auto_max_depth": 8, "explicit_unroll": False}}
):
ret = tvm.tirx.transform.UnrollLoop()(ModuleWithPolicy)["main"].body
@@ -81,7 +81,7 @@ def test_unroll_fake_loop():
for j in T.serial(10):
Ab[j + 1] = Ab[i] + 1
- with tvm.transform.PassContext(
+ with tvm.ir.transform.PassContext(
config={
"tirx.UnrollLoop": {"auto_max_depth": 8, "auto_max_extent": 1, "explicit_unroll": False}
}
@@ -136,7 +136,7 @@ def test_unroll_local_access():
A_local[2] = T.float32(2)
A_local[3] = T.float32(3)
- with tvm.transform.PassContext(
+ with tvm.ir.transform.PassContext(
config={
"tirx.UnrollLoop": {
"auto_max_depth": 0,
diff --git a/tests/python/tirx-transform/test_tir_transform_vectorize.py b/tests/python/tirx-transform/test_tir_transform_vectorize.py
index c8d48b45b9..f79bac64d5 100644
--- a/tests/python/tirx-transform/test_tir_transform_vectorize.py
+++ b/tests/python/tirx-transform/test_tir_transform_vectorize.py
@@ -568,7 +568,7 @@ def test_vectorize_and_predicate_all_tensor_loads_stores():
)
mod = tvm.IRModule.from_expr(before)
- with tvm.transform.PassContext(config={"tirx.enable_tensor_level_predication": True}):
+ with tvm.ir.transform.PassContext(config={"tirx.enable_tensor_level_predication": True}):
after = tvm.tirx.transform.VectorizeLoop()(mod)["main"]
tvm.ir.assert_structural_equal(after, expected)
@@ -592,7 +592,7 @@ def test_vectorize_and_predicate_some_tensor_loads_stores():
B[i_0 * 4 + i_1_s] = A[i_0] + T.float32(1)
mod = tvm.IRModule.from_expr(before)
- with tvm.transform.PassContext(config={"tirx.enable_tensor_level_predication": True}):
+ with tvm.ir.transform.PassContext(config={"tirx.enable_tensor_level_predication": True}):
after = tvm.tirx.transform.VectorizeLoop()(mod)["main"]
tvm.ir.assert_structural_equal(after, expected)
@@ -633,7 +633,7 @@ def test_vectorize_and_predicate_multiple_access_statements():
)
before_mod = tvm.IRModule.from_expr(before)
- with tvm.transform.PassContext(config={"tirx.enable_tensor_level_predication": True}):
+ with tvm.ir.transform.PassContext(config={"tirx.enable_tensor_level_predication": True}):
after = tvm.tirx.transform.VectorizeLoop()(before_mod)["main"]
tvm.ir.assert_structural_equal(after, expected)
@@ -694,7 +694,7 @@ def test_vectorize_and_predicate_invalid_conditions():
A[i_0 * 4 + i_1_s] = T.float32(2)
before_mod = tvm.IRModule.from_expr(before)
- with tvm.transform.PassContext(config={"tirx.enable_tensor_level_predication": True}):
+ with tvm.ir.transform.PassContext(config={"tirx.enable_tensor_level_predication": True}):
after = tvm.tirx.transform.VectorizeLoop()(before_mod)["main"]
tvm.ir.assert_structural_equal(after, expected)
@@ -719,7 +719,7 @@ def test_vectorize_with_explicitly_disabled_tensor_level_predication():
B[i_0 * 4 + i_1_s] = A[i_0 * 4 + i_1_s] + T.float32(1)
mod = tvm.IRModule.from_expr(before)
- with tvm.transform.PassContext(config={"tirx.enable_tensor_level_predication": False}):
+ with tvm.ir.transform.PassContext(config={"tirx.enable_tensor_level_predication": False}):
after = tvm.tirx.transform.VectorizeLoop()(mod)["main"]
tvm.ir.assert_structural_equal(after, expected)
diff --git a/tests/python/tirx/iket/test_iket_profiler.py b/tests/python/tirx/iket/test_iket_profiler.py
index 1ee85257b9..5f48cf8015 100644
--- a/tests/python/tirx/iket/test_iket_profiler.py
+++ b/tests/python/tirx/iket/test_iket_profiler.py
@@ -487,7 +487,7 @@ def test_tirx_pipelines_immediately_lower_iket(module_name, factory_name):
factory = getattr(importlib.import_module(module_name), factory_name)
visited = []
- @tvm.transform.pass_instrument
+ @tvm.ir.transform.pass_instrument
class RecordPasses:
def run_before_pass(self, mod, info):
visited.append(("before", info.name))
@@ -496,7 +496,7 @@ def test_tirx_pipelines_immediately_lower_iket(module_name, factory_name):
visited.append(("after", info.name))
pipeline, _, _ = factory()
- with tvm.transform.PassContext(instruments=[RecordPasses()]):
+ with tvm.ir.transform.PassContext(instruments=[RecordPasses()]):
pipeline(tvm.IRModule())
split = visited.index(("after", tvm.tirx.transform.SplitHostDevice().info.name))
diff --git a/tests/python/tirx/test_cuda_backend_config.py b/tests/python/tirx/test_cuda_backend_config.py
index 871cae752e..c2677284d5 100644
--- a/tests/python/tirx/test_cuda_backend_config.py
+++ b/tests/python/tirx/test_cuda_backend_config.py
@@ -464,7 +464,7 @@ def test_device_helpers_are_cloned_per_configuration():
assert set(groups[("sm_90a", False)]).isdisjoint(groups[("sm_100a", True)])
dispatched = set()
- @tvm.transform.pass_instrument
+ @tvm.ir.transform.pass_instrument
class CheckHelperTargets:
def run_before_pass(self, lowered, info):
if info.name == "tirx.TileDispatch":
@@ -473,7 +473,7 @@ def test_device_helpers_are_cloned_per_configuration():
dispatched.add(fn.attrs["target"].arch)
assert gv.name_hint not in ("helper", "leaf")
- with tvm.transform.PassContext(instruments=[CheckHelperTargets()]):
+ with tvm.ir.transform.PassContext(instruments=[CheckHelperTargets()]):
built = tvm.tirx.compile(Module)
assert dispatched == {"sm_90a", "sm_100a"}
assert len(built.mod.imports) == 2
@@ -727,12 +727,12 @@ def test_build_snapshots_config_before_passes(monkeypatch):
T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=32))
A[T.cuda.thread_idx("x")] = T.float32(1)
- @tvm.transform.pass_instrument
+ @tvm.ir.transform.pass_instrument
class ChangeCallerConfig:
def run_before_pass(self, mod, info):
config["cuda"]["nvrtc"] = ["--ftz=true"]
- with tvm.transform.PassContext(instruments=[ChangeCallerConfig()]):
+ with tvm.ir.transform.PassContext(instruments=[ChangeCallerConfig()]):
module = tvm.tirx.compile(main, target={"kind": "cuda", "backend_config": config})
resolved = parse_backend_config(module.mod.imports[0].inspect_source("backend_config"))
assert config["cuda"]["nvrtc"] == ["--ftz=true"]
diff --git a/tests/python/tirx/tile/cuda/copy/test_reg.py b/tests/python/tirx/tile/cuda/copy/test_reg.py
index 859353f173..981b941dc2 100644
--- a/tests/python/tirx/tile/cuda/copy/test_reg.py
+++ b/tests/python/tirx/tile/cuda/copy/test_reg.py
@@ -462,14 +462,14 @@ def test_reg_copy_linear_shared_hoists_thread_base():
dispatched = []
- @tvm.transform.pass_instrument
+ @tvm.ir.transform.pass_instrument
class CaptureDispatch:
def run_after_pass(self, mod, info):
if info.name == "tirx.TileDispatch":
dispatched.append(mod["main"])
target = tvm.target.Target("cuda")
- with target, tvm.transform.PassContext(instruments=[CaptureDispatch()]):
+ with target, tvm.ir.transform.PassContext(instruments=[CaptureDispatch()]):
ex = tvm.tirx.compile(tvm.IRModule({"main": kernel}), target=target, pipeline="tirx")
src = ex.mod.imports[0].inspect_source()