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 adda179705 [REFACTOR][S-TIR] Lift dlight into s_tir namespace (#18734)
adda179705 is described below
commit adda179705192c422e29fedb12f4a3901281b296
Author: Tianqi Chen <[email protected]>
AuthorDate: Sun Feb 8 20:20:11 2026 -0500
[REFACTOR][S-TIR] Lift dlight into s_tir namespace (#18734)
This PR migrates dlight into s_tir namespace, so s_tir related
components are closely grouped together.
---
docs/get_started/tutorials/ir_module.py | 2 +-
docs/how_to/tutorials/cross_compilation_and_rpc.py | 2 +-
docs/how_to/tutorials/customize_opt.py | 2 +-
docs/how_to/tutorials/optimize_llm.py | 3 ++-
docs/reference/api/python/index.rst | 5 +----
docs/reference/api/python/{ => s_tir}/dlight.rst | 6 +++---
python/tvm/relax/backend/adreno/pipeline.py | 3 ++-
python/tvm/relax/backend/cuda/pipeline.py | 3 ++-
python/tvm/relax/backend/dispatch_sort_scan.py | 4 +++-
python/tvm/relax/backend/gpu_generic/pipeline.py | 3 ++-
python/tvm/relax/backend/rocm/pipeline.py | 3 ++-
python/tvm/s_tir/__init__.py | 9 ++++++++-
python/tvm/{ => s_tir}/dlight/__init__.py | 0
python/tvm/{ => s_tir}/dlight/adreno/__init__.py | 0
python/tvm/{ => s_tir}/dlight/adreno/base.py | 0
python/tvm/{ => s_tir}/dlight/adreno/convolution.py | 0
python/tvm/{ => s_tir}/dlight/adreno/fallback.py | 0
python/tvm/{ => s_tir}/dlight/adreno/layout_transform.py | 0
python/tvm/{ => s_tir}/dlight/adreno/pool.py | 0
python/tvm/{ => s_tir}/dlight/adreno/utils.py | 0
python/tvm/{ => s_tir}/dlight/analysis/__init__.py | 0
python/tvm/{ => s_tir}/dlight/analysis/common_analysis.py | 0
python/tvm/{ => s_tir}/dlight/analysis/gemv.py | 0
python/tvm/{ => s_tir}/dlight/base/__init__.py | 0
python/tvm/{ => s_tir}/dlight/base/common_schedules.py | 0
python/tvm/{ => s_tir}/dlight/base/schedule_rule.py | 0
python/tvm/{ => s_tir}/dlight/base/transform.py | 0
python/tvm/{ => s_tir}/dlight/base/utils.py | 0
python/tvm/{ => s_tir}/dlight/benchmark/__init__.py | 0
python/tvm/{ => s_tir}/dlight/benchmark/bench.py | 0
python/tvm/{ => s_tir}/dlight/benchmark/extract.py | 2 +-
python/tvm/{ => s_tir}/dlight/benchmark/utils.py | 0
python/tvm/{ => s_tir}/dlight/cpu/__init__.py | 0
python/tvm/{ => s_tir}/dlight/cpu/base.py | 0
python/tvm/{ => s_tir}/dlight/cpu/gemv.py | 0
python/tvm/{ => s_tir}/dlight/gpu/__init__.py | 2 +-
python/tvm/{ => s_tir}/dlight/gpu/base.py | 0
python/tvm/{ => s_tir}/dlight/gpu/fallback.py | 0
python/tvm/{ => s_tir}/dlight/gpu/gemv.py | 0
python/tvm/{ => s_tir}/dlight/gpu/general_reduction.py | 0
python/tvm/{ => s_tir}/dlight/gpu/low_batch_gemv.py | 0
python/tvm/{ => s_tir}/dlight/gpu/matmul.py | 0
python/tvm/{ => s_tir}/dlight/gpu/reduction.py | 0
python/tvm/{ => s_tir}/dlight/gpu/rmsnorm.py | 1 +
python/tvm/{ => s_tir}/dlight/gpu/transpose.py | 0
python/tvm/script/parser/core/entry.py | 4 ++--
python/tvm/script/parser/ir/entry.py | 10 ++++++++--
python/tvm/script/parser/ir/parser.py | 5 ++++-
tests/python/codegen/test_target_codegen_cuda_fp8.py | 14 +++++++-------
tests/python/disco/test_callback.py | 2 +-
tests/python/disco/test_ccl.py | 2 +-
tests/python/disco/test_loader.py | 2 +-
tests/python/relax/backend/clml/utils.py | 2 +-
.../nvshmem/test_runtime_builtin_kv_cache_transfer.py | 2 +-
tests/python/relax/test_backend_dispatch_sort_scan.py | 3 ++-
.../test_runtime_builtin_paged_attention_kv_cache_cpu.py | 2 +-
..._runtime_builtin_paged_attention_kv_cache_flashinfer.py | 2 +-
...time_builtin_paged_attention_kv_cache_mla_flashinfer.py | 2 +-
...est_runtime_builtin_paged_attention_kv_cache_mla_tir.py | 2 +-
.../test_runtime_builtin_paged_attention_kv_cache_tir.py | 2 +-
tests/python/relax/test_runtime_builtin_rnn_state.py | 2 +-
tests/python/relax/test_vm_build.py | 2 +-
tests/python/{ => s_tir}/dlight/test_benchmark.py | 2 +-
tests/python/{ => s_tir}/dlight/test_cpu_gemv.py | 2 +-
tests/python/{ => s_tir}/dlight/test_gpu_conv.py | 2 +-
tests/python/{ => s_tir}/dlight/test_gpu_fallback.py | 2 +-
tests/python/{ => s_tir}/dlight/test_gpu_gemv.py | 2 +-
.../{ => s_tir}/dlight/test_gpu_general_reduction.py | 2 +-
tests/python/{ => s_tir}/dlight/test_gpu_low_batch_gemv.py | 2 +-
tests/python/{ => s_tir}/dlight/test_gpu_matmul.py | 2 +-
.../python/{ => s_tir}/dlight/test_gpu_matmul_tensorize.py | 2 +-
tests/python/{ => s_tir}/dlight/test_gpu_reduction.py | 2 +-
tests/python/{ => s_tir}/dlight/test_gpu_rmsnorm.py | 2 +-
tests/python/{ => s_tir}/dlight/test_gpu_transpose.py | 2 +-
tests/python/{ => s_tir}/dlight/test_primitives.py | 0
tests/scripts/release/make_notes.py | 2 +-
tests/scripts/task_python_unittest.sh | 1 +
tests/scripts/unity/task_python_relax.sh | 1 -
78 files changed, 80 insertions(+), 58 deletions(-)
diff --git a/docs/get_started/tutorials/ir_module.py
b/docs/get_started/tutorials/ir_module.py
index 8bb8fb77a4..91aedaa682 100644
--- a/docs/get_started/tutorials/ir_module.py
+++ b/docs/get_started/tutorials/ir_module.py
@@ -252,7 +252,7 @@ print(cpu_out)
# the details of ``DLight``.
#
-from tvm import dlight as dl
+from tvm.s_tir import dlight as dl
with tvm.target.Target("cuda"):
gpu_mod = dl.ApplyDefaultSchedule(
diff --git a/docs/how_to/tutorials/cross_compilation_and_rpc.py
b/docs/how_to/tutorials/cross_compilation_and_rpc.py
index b6b06e2676..229f40f938 100644
--- a/docs/how_to/tutorials/cross_compilation_and_rpc.py
+++ b/docs/how_to/tutorials/cross_compilation_and_rpc.py
@@ -513,7 +513,7 @@ def run_pytorch_model_via_rpc():
#
# .. code-block:: python
#
- # from tvm import dlight as dl
+ # from tvm.s_tir import dlight as dl
# with target:
# mod = dl.ApplyDefaultSchedule()(mod)
#
diff --git a/docs/how_to/tutorials/customize_opt.py
b/docs/how_to/tutorials/customize_opt.py
index 2e2747d61f..43ce71343f 100644
--- a/docs/how_to/tutorials/customize_opt.py
+++ b/docs/how_to/tutorials/customize_opt.py
@@ -174,7 +174,7 @@ if os.getenv("CI", "") != "true":
# e.g. language model, DLight provides excellent performance, while for
generic models,
# it achieves a balance between performance and compilation time.
-from tvm import dlight as dl
+from tvm.s_tir import dlight as dl
# Apply DLight rules
with target:
diff --git a/docs/how_to/tutorials/optimize_llm.py
b/docs/how_to/tutorials/optimize_llm.py
index 0e82b05559..cf8ec5fe8c 100644
--- a/docs/how_to/tutorials/optimize_llm.py
+++ b/docs/how_to/tutorials/optimize_llm.py
@@ -63,7 +63,8 @@ from pprint import pprint
from typing import List, Optional
import tvm
-from tvm import dlight, relax, te, tir
+from tvm import relax, te, tir
+from tvm.s_tir import dlight
from tvm.relax import register_pipeline
from tvm.relax.frontend import nn
from tvm.relax.frontend.nn import Tensor, op
diff --git a/docs/reference/api/python/index.rst
b/docs/reference/api/python/index.rst
index 38ae6a81b1..1752b9d99c 100644
--- a/docs/reference/api/python/index.rst
+++ b/docs/reference/api/python/index.rst
@@ -63,6 +63,7 @@ Python API
:caption: tvm.s_tir
s_tir/schedule
+ s_tir/dlight
.. toctree::
:maxdepth: 1
@@ -77,11 +78,7 @@ Python API
meta_schedule
-.. toctree::
- :maxdepth: 1
- :caption: tvm.dlight
- dlight
.. toctree::
:maxdepth: 1
diff --git a/docs/reference/api/python/dlight.rst
b/docs/reference/api/python/s_tir/dlight.rst
similarity index 92%
rename from docs/reference/api/python/dlight.rst
rename to docs/reference/api/python/s_tir/dlight.rst
index 37859ed790..777840055f 100644
--- a/docs/reference/api/python/dlight.rst
+++ b/docs/reference/api/python/s_tir/dlight.rst
@@ -15,8 +15,8 @@
specific language governing permissions and limitations
under the License.
-tvm.dlight
-----------
-.. automodule:: tvm.dlight
+tvm.s_tir.dlight
+----------------
+.. automodule:: tvm.s_tir.dlight
:members:
:imported-members:
diff --git a/python/tvm/relax/backend/adreno/pipeline.py
b/python/tvm/relax/backend/adreno/pipeline.py
index cbbb17a618..3567753a88 100644
--- a/python/tvm/relax/backend/adreno/pipeline.py
+++ b/python/tvm/relax/backend/adreno/pipeline.py
@@ -16,7 +16,6 @@
# under the License.
"""The Relax Adreno GPU backend compilation pipeline and other passes."""
import tvm
-from tvm import dlight as dl
from tvm import relax
from tvm.relax.transform.legalize_ops import adreno as legalize_adreno
@@ -83,6 +82,8 @@ def legalize_passes(target: tvm.target.Target): # pylint:
disable=unused-argume
relax.transform.SpecializePrimFuncBasedOnCallSite(),
]
)
+ from tvm.s_tir import dlight as dl # pylint:
disable=import-outside-toplevel
+
pass_list.extend([relax.transform.Normalize()])
pass_list.extend(
[
diff --git a/python/tvm/relax/backend/cuda/pipeline.py
b/python/tvm/relax/backend/cuda/pipeline.py
index d5c4c08561..4319d76d03 100644
--- a/python/tvm/relax/backend/cuda/pipeline.py
+++ b/python/tvm/relax/backend/cuda/pipeline.py
@@ -16,7 +16,6 @@
# under the License.
"""The Relax CUDA backend compilation pipeline and other passes."""
import tvm
-from tvm import dlight as dl
from tvm import relax
@@ -30,6 +29,8 @@ def library_dispatch_passes(target: tvm.target.Target): #
pylint: disable=unuse
def legalize_passes(target: tvm.target.Target): # pylint:
disable=unused-argument
"""The default legalization passes for CUDA backend."""
+ from tvm.s_tir import dlight as dl # pylint:
disable=import-outside-toplevel
+
return [
tvm.relax.transform.LegalizeOps(),
tvm.relax.transform.AnnotateTIROpPattern(),
diff --git a/python/tvm/relax/backend/dispatch_sort_scan.py
b/python/tvm/relax/backend/dispatch_sort_scan.py
index 1dac0bf230..62290cf333 100644
--- a/python/tvm/relax/backend/dispatch_sort_scan.py
+++ b/python/tvm/relax/backend/dispatch_sort_scan.py
@@ -21,7 +21,7 @@ from functools import reduce
from operator import mul
from typing import Dict
-from tvm import DataType, dlight, relax, topi
+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
@@ -46,6 +46,8 @@ class SortScanDispatcher(BackendDispatcher):
self,
) -> None:
"""Apply DLight rules for all the calls that need to be updated."""
+ from tvm.s_tir import dlight # pylint: disable=import-outside-toplevel
+
for gvar, target in self.calls_to_update.items():
func = self.builder_.get()[gvar]
sch = dlight.base.transform._apply_rules(
diff --git a/python/tvm/relax/backend/gpu_generic/pipeline.py
b/python/tvm/relax/backend/gpu_generic/pipeline.py
index 86c60114c6..99a723cb2c 100644
--- a/python/tvm/relax/backend/gpu_generic/pipeline.py
+++ b/python/tvm/relax/backend/gpu_generic/pipeline.py
@@ -16,7 +16,6 @@
# under the License.
"""The Relax generic GPU backend compilation pipeline and other passes."""
import tvm
-from tvm import dlight as dl
from tvm import relax
@@ -30,6 +29,8 @@ def library_dispatch_passes(target: tvm.target.Target): #
pylint: disable=unuse
def legalize_passes(target: tvm.target.Target): # pylint:
disable=unused-argument
"""The default legalization passes for generic GPU backend."""
+ from tvm.s_tir import dlight as dl # pylint:
disable=import-outside-toplevel
+
return [
tvm.relax.transform.LegalizeOps(),
tvm.relax.transform.AnnotateTIROpPattern(),
diff --git a/python/tvm/relax/backend/rocm/pipeline.py
b/python/tvm/relax/backend/rocm/pipeline.py
index e74039ca86..a3f2462f97 100644
--- a/python/tvm/relax/backend/rocm/pipeline.py
+++ b/python/tvm/relax/backend/rocm/pipeline.py
@@ -16,7 +16,6 @@
# under the License.
"""The Relax ROCm backend compilation pipeline and other passes."""
import tvm
-from tvm import dlight as dl
from tvm import relax
@@ -30,6 +29,8 @@ def library_dispatch_passes(target: tvm.target.Target): #
pylint: disable=unuse
def legalize_passes(target: tvm.target.Target): # pylint:
disable=unused-argument
"""The default legalization passes for ROCm backend."""
+ from tvm.s_tir import dlight as dl # pylint:
disable=import-outside-toplevel
+
return [
tvm.relax.transform.LegalizeOps(),
tvm.relax.transform.AnnotateTIROpPattern(),
diff --git a/python/tvm/s_tir/__init__.py b/python/tvm/s_tir/__init__.py
index 5f233990db..fc2fdab196 100644
--- a/python/tvm/s_tir/__init__.py
+++ b/python/tvm/s_tir/__init__.py
@@ -19,10 +19,17 @@
from tvm.tir.function import TensorIntrin
+# dlight depends on compiler-only C++ functions (e.g.
s_tir.schedule.GetSBlockRealize),
+# so skip it in runtime-only builds.
+from tvm.base import _RUNTIME_ONLY
+
from . import backend
from . import pipeline
from . import transform
from . import schedule
from .schedule import StmtSRef, SBlockScope, ScheduleState, Schedule,
ScheduleError, Trace
from .block_dependence_info import SBlockDependenceInfo
-from . import meta_schedule
+
+if not _RUNTIME_ONLY:
+ from . import meta_schedule
+ from . import dlight
diff --git a/python/tvm/dlight/__init__.py b/python/tvm/s_tir/dlight/__init__.py
similarity index 100%
rename from python/tvm/dlight/__init__.py
rename to python/tvm/s_tir/dlight/__init__.py
diff --git a/python/tvm/dlight/adreno/__init__.py
b/python/tvm/s_tir/dlight/adreno/__init__.py
similarity index 100%
rename from python/tvm/dlight/adreno/__init__.py
rename to python/tvm/s_tir/dlight/adreno/__init__.py
diff --git a/python/tvm/dlight/adreno/base.py
b/python/tvm/s_tir/dlight/adreno/base.py
similarity index 100%
rename from python/tvm/dlight/adreno/base.py
rename to python/tvm/s_tir/dlight/adreno/base.py
diff --git a/python/tvm/dlight/adreno/convolution.py
b/python/tvm/s_tir/dlight/adreno/convolution.py
similarity index 100%
rename from python/tvm/dlight/adreno/convolution.py
rename to python/tvm/s_tir/dlight/adreno/convolution.py
diff --git a/python/tvm/dlight/adreno/fallback.py
b/python/tvm/s_tir/dlight/adreno/fallback.py
similarity index 100%
rename from python/tvm/dlight/adreno/fallback.py
rename to python/tvm/s_tir/dlight/adreno/fallback.py
diff --git a/python/tvm/dlight/adreno/layout_transform.py
b/python/tvm/s_tir/dlight/adreno/layout_transform.py
similarity index 100%
rename from python/tvm/dlight/adreno/layout_transform.py
rename to python/tvm/s_tir/dlight/adreno/layout_transform.py
diff --git a/python/tvm/dlight/adreno/pool.py
b/python/tvm/s_tir/dlight/adreno/pool.py
similarity index 100%
rename from python/tvm/dlight/adreno/pool.py
rename to python/tvm/s_tir/dlight/adreno/pool.py
diff --git a/python/tvm/dlight/adreno/utils.py
b/python/tvm/s_tir/dlight/adreno/utils.py
similarity index 100%
rename from python/tvm/dlight/adreno/utils.py
rename to python/tvm/s_tir/dlight/adreno/utils.py
diff --git a/python/tvm/dlight/analysis/__init__.py
b/python/tvm/s_tir/dlight/analysis/__init__.py
similarity index 100%
rename from python/tvm/dlight/analysis/__init__.py
rename to python/tvm/s_tir/dlight/analysis/__init__.py
diff --git a/python/tvm/dlight/analysis/common_analysis.py
b/python/tvm/s_tir/dlight/analysis/common_analysis.py
similarity index 100%
rename from python/tvm/dlight/analysis/common_analysis.py
rename to python/tvm/s_tir/dlight/analysis/common_analysis.py
diff --git a/python/tvm/dlight/analysis/gemv.py
b/python/tvm/s_tir/dlight/analysis/gemv.py
similarity index 100%
rename from python/tvm/dlight/analysis/gemv.py
rename to python/tvm/s_tir/dlight/analysis/gemv.py
diff --git a/python/tvm/dlight/base/__init__.py
b/python/tvm/s_tir/dlight/base/__init__.py
similarity index 100%
rename from python/tvm/dlight/base/__init__.py
rename to python/tvm/s_tir/dlight/base/__init__.py
diff --git a/python/tvm/dlight/base/common_schedules.py
b/python/tvm/s_tir/dlight/base/common_schedules.py
similarity index 100%
rename from python/tvm/dlight/base/common_schedules.py
rename to python/tvm/s_tir/dlight/base/common_schedules.py
diff --git a/python/tvm/dlight/base/schedule_rule.py
b/python/tvm/s_tir/dlight/base/schedule_rule.py
similarity index 100%
rename from python/tvm/dlight/base/schedule_rule.py
rename to python/tvm/s_tir/dlight/base/schedule_rule.py
diff --git a/python/tvm/dlight/base/transform.py
b/python/tvm/s_tir/dlight/base/transform.py
similarity index 100%
rename from python/tvm/dlight/base/transform.py
rename to python/tvm/s_tir/dlight/base/transform.py
diff --git a/python/tvm/dlight/base/utils.py
b/python/tvm/s_tir/dlight/base/utils.py
similarity index 100%
rename from python/tvm/dlight/base/utils.py
rename to python/tvm/s_tir/dlight/base/utils.py
diff --git a/python/tvm/dlight/benchmark/__init__.py
b/python/tvm/s_tir/dlight/benchmark/__init__.py
similarity index 100%
rename from python/tvm/dlight/benchmark/__init__.py
rename to python/tvm/s_tir/dlight/benchmark/__init__.py
diff --git a/python/tvm/dlight/benchmark/bench.py
b/python/tvm/s_tir/dlight/benchmark/bench.py
similarity index 100%
rename from python/tvm/dlight/benchmark/bench.py
rename to python/tvm/s_tir/dlight/benchmark/bench.py
diff --git a/python/tvm/dlight/benchmark/extract.py
b/python/tvm/s_tir/dlight/benchmark/extract.py
similarity index 99%
rename from python/tvm/dlight/benchmark/extract.py
rename to python/tvm/s_tir/dlight/benchmark/extract.py
index bc98ca1d87..1fa27e874c 100644
--- a/python/tvm/dlight/benchmark/extract.py
+++ b/python/tvm/s_tir/dlight/benchmark/extract.py
@@ -32,7 +32,7 @@ import tvm
from tvm import relax
from tvm.script import tir as T
-from tvm.dlight.benchmark import benchmark_prim_func
+from tvm.s_tir.dlight.benchmark import benchmark_prim_func
MODEL_NAME = "{model_name}"
RELAX_FUNC_NAME = "{relax_func_name}"
diff --git a/python/tvm/dlight/benchmark/utils.py
b/python/tvm/s_tir/dlight/benchmark/utils.py
similarity index 100%
rename from python/tvm/dlight/benchmark/utils.py
rename to python/tvm/s_tir/dlight/benchmark/utils.py
diff --git a/python/tvm/dlight/cpu/__init__.py
b/python/tvm/s_tir/dlight/cpu/__init__.py
similarity index 100%
rename from python/tvm/dlight/cpu/__init__.py
rename to python/tvm/s_tir/dlight/cpu/__init__.py
diff --git a/python/tvm/dlight/cpu/base.py b/python/tvm/s_tir/dlight/cpu/base.py
similarity index 100%
rename from python/tvm/dlight/cpu/base.py
rename to python/tvm/s_tir/dlight/cpu/base.py
diff --git a/python/tvm/dlight/cpu/gemv.py b/python/tvm/s_tir/dlight/cpu/gemv.py
similarity index 100%
rename from python/tvm/dlight/cpu/gemv.py
rename to python/tvm/s_tir/dlight/cpu/gemv.py
diff --git a/python/tvm/dlight/gpu/__init__.py
b/python/tvm/s_tir/dlight/gpu/__init__.py
similarity index 91%
rename from python/tvm/dlight/gpu/__init__.py
rename to python/tvm/s_tir/dlight/gpu/__init__.py
index 077fdcaeb0..cbd506b6db 100644
--- a/python/tvm/dlight/gpu/__init__.py
+++ b/python/tvm/s_tir/dlight/gpu/__init__.py
@@ -16,7 +16,7 @@
# under the License.
"""
GPU-generic schedule rules.
-For CUDA/ROCm/Vulkan/Metal-specific rules, use
`tvm.dlight.cuda/rocm/vulkan/metal` instead
+For CUDA/ROCm/Vulkan/Metal-specific rules, use
`tvm.s_tir.dlight.cuda/rocm/vulkan/metal` instead
"""
from .gemv import GEMV
from .low_batch_gemv import LowBatchGEMV
diff --git a/python/tvm/dlight/gpu/base.py b/python/tvm/s_tir/dlight/gpu/base.py
similarity index 100%
rename from python/tvm/dlight/gpu/base.py
rename to python/tvm/s_tir/dlight/gpu/base.py
diff --git a/python/tvm/dlight/gpu/fallback.py
b/python/tvm/s_tir/dlight/gpu/fallback.py
similarity index 100%
rename from python/tvm/dlight/gpu/fallback.py
rename to python/tvm/s_tir/dlight/gpu/fallback.py
diff --git a/python/tvm/dlight/gpu/gemv.py b/python/tvm/s_tir/dlight/gpu/gemv.py
similarity index 100%
rename from python/tvm/dlight/gpu/gemv.py
rename to python/tvm/s_tir/dlight/gpu/gemv.py
diff --git a/python/tvm/dlight/gpu/general_reduction.py
b/python/tvm/s_tir/dlight/gpu/general_reduction.py
similarity index 100%
rename from python/tvm/dlight/gpu/general_reduction.py
rename to python/tvm/s_tir/dlight/gpu/general_reduction.py
diff --git a/python/tvm/dlight/gpu/low_batch_gemv.py
b/python/tvm/s_tir/dlight/gpu/low_batch_gemv.py
similarity index 100%
rename from python/tvm/dlight/gpu/low_batch_gemv.py
rename to python/tvm/s_tir/dlight/gpu/low_batch_gemv.py
diff --git a/python/tvm/dlight/gpu/matmul.py
b/python/tvm/s_tir/dlight/gpu/matmul.py
similarity index 100%
rename from python/tvm/dlight/gpu/matmul.py
rename to python/tvm/s_tir/dlight/gpu/matmul.py
diff --git a/python/tvm/dlight/gpu/reduction.py
b/python/tvm/s_tir/dlight/gpu/reduction.py
similarity index 100%
rename from python/tvm/dlight/gpu/reduction.py
rename to python/tvm/s_tir/dlight/gpu/reduction.py
diff --git a/python/tvm/dlight/gpu/rmsnorm.py
b/python/tvm/s_tir/dlight/gpu/rmsnorm.py
similarity index 99%
rename from python/tvm/dlight/gpu/rmsnorm.py
rename to python/tvm/s_tir/dlight/gpu/rmsnorm.py
index 19fff039a4..c983a534d8 100644
--- a/python/tvm/dlight/gpu/rmsnorm.py
+++ b/python/tvm/s_tir/dlight/gpu/rmsnorm.py
@@ -16,6 +16,7 @@
# under the License.
# pylint: disable=missing-docstring
"""A RMS norm schedule rule for GPU operators."""
+from __future__ import annotations
import tvm
from tvm import tir
diff --git a/python/tvm/dlight/gpu/transpose.py
b/python/tvm/s_tir/dlight/gpu/transpose.py
similarity index 100%
rename from python/tvm/dlight/gpu/transpose.py
rename to python/tvm/s_tir/dlight/gpu/transpose.py
diff --git a/python/tvm/script/parser/core/entry.py
b/python/tvm/script/parser/core/entry.py
index a6be751b0d..de2d416206 100644
--- a/python/tvm/script/parser/core/entry.py
+++ b/python/tvm/script/parser/core/entry.py
@@ -19,7 +19,6 @@ import inspect
from typing import Any, Dict, Union
import tvm
-from tvm.relax import ExternFunc
from ....ir.module import IRModule
from ...ir_builder import IRBuilder
from . import doc
@@ -36,6 +35,7 @@ WELL_FORMED_ERROR_MESSAGE = (
def _default_globals() -> Dict[str, Any]:
+ # lazy import here to avoid circular deps
from tvm.script.parser import ir # pylint: disable=import-outside-toplevel
from tvm.script.parser import relax # pylint:
disable=import-outside-toplevel
from tvm.script.parser import tir # pylint:
disable=import-outside-toplevel
@@ -169,7 +169,7 @@ def _attach_pyfuncs_to_irmodule(irmodule, all_pyfuncs):
irmodule.pyfuncs = {}
for global_var, func in irmodule.functions_items():
- if not isinstance(func, ExternFunc):
+ if not isinstance(func, tvm.relax.ExternFunc):
continue
if not func.attrs.get("is_pyfunc", False):
continue
diff --git a/python/tvm/script/parser/ir/entry.py
b/python/tvm/script/parser/ir/entry.py
index 0e2adeebe3..245714d0ee 100644
--- a/python/tvm/script/parser/ir/entry.py
+++ b/python/tvm/script/parser/ir/entry.py
@@ -14,14 +14,13 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
+# pylint: disable=import-outside-toplevel
"""The entry point of TVM parser for ir module."""
import inspect
from typing import Callable, Optional, Type
from tvm.ir import IRModule, GlobalVar
-from tvm.relax.expr import ExternFunc
-from tvm.relax.base_py_module import BasePyModule
from tvm import cpu, ir
from .._core import parse, utils
@@ -57,6 +56,13 @@ def ir_module(mod: Optional[Type] = None, check_well_formed:
bool = True) -> IRM
m = parse(mod, utils.inspect_class_capture(mod),
check_well_formed=check_well_formed)
if base_py_module_inherited:
+ # Lazy import: tvm.relax cannot be imported at module level in
tvm.script.parser
+ # because tvm.script is loaded before tvm.relax during tvm
initialization.
+ from tvm.relax.expr import ExternFunc # pylint:
disable=import-outside-toplevel
+ from tvm.relax.base_py_module import (
+ BasePyModule,
+ )
+
# Collect pyfunc methods
pyfunc_methods = [
name
diff --git a/python/tvm/script/parser/ir/parser.py
b/python/tvm/script/parser/ir/parser.py
index 80d2db87ab..848472404d 100644
--- a/python/tvm/script/parser/ir/parser.py
+++ b/python/tvm/script/parser/ir/parser.py
@@ -18,7 +18,6 @@
"""The base parser for ir module"""
from tvm.ir import GlobalVar
-from tvm.relax import ExternFunc
from ...ir_builder import ir as I
from .._core import Parser, dispatch, doc
@@ -154,6 +153,10 @@ def visit_tvm_declare_function(self: Parser, node:
doc.FunctionDef) -> GlobalVar
f"Class '{current_class}' does not inherit from BasePyModule.",
)
+ # Lazy import: tvm.relax cannot be imported at module level in
tvm.script.parser
+ # because tvm.script is loaded before tvm.relax during tvm initialization.
+ from tvm.relax import ExternFunc # pylint: disable=import-outside-toplevel
+
# Create ExternFunc with proper attributes for Python functions
func = ExternFunc(node.name)
func = func.with_attr("is_pyfunc", True)
diff --git a/tests/python/codegen/test_target_codegen_cuda_fp8.py
b/tests/python/codegen/test_target_codegen_cuda_fp8.py
index 229839add9..6c3a03f280 100644
--- a/tests/python/codegen/test_target_codegen_cuda_fp8.py
+++ b/tests/python/codegen/test_target_codegen_cuda_fp8.py
@@ -24,7 +24,7 @@ import pytest
import tvm
import tvm.testing
from tvm import DataType, DataTypeCode, IRModule
-from tvm import dlight as dl
+from tvm.s_tir import dlight as dl
from tvm import relax, te, tir, topi
from tvm.script import ir as I
from tvm.script import relax as R
@@ -926,12 +926,12 @@ def test_moe_gemv_shfl_down_illegal_instr():
seq = tvm.transform.Sequential(
[
tvm.relax.transform.LegalizeOps(),
- tvm.dlight.ApplyDefaultSchedule(
- tvm.dlight.gpu.Matmul(),
- tvm.dlight.gpu.GEMV(),
- tvm.dlight.gpu.Reduction(),
- tvm.dlight.gpu.GeneralReduction(),
- tvm.dlight.gpu.Fallback(),
+ dl.ApplyDefaultSchedule(
+ dl.gpu.Matmul(),
+ dl.gpu.GEMV(),
+ dl.gpu.Reduction(),
+ dl.gpu.GeneralReduction(),
+ dl.gpu.Fallback(),
),
]
)
diff --git a/tests/python/disco/test_callback.py
b/tests/python/disco/test_callback.py
index 8e78058331..c8f4c88e75 100644
--- a/tests/python/disco/test_callback.py
+++ b/tests/python/disco/test_callback.py
@@ -56,7 +56,7 @@ def test_callback():
pipeline = tvm.ir.transform.Sequential(
[
tvm.relax.transform.LegalizeOps(),
- tvm.dlight.ApplyDefaultSchedule(tvm.dlight.gpu.Fallback()),
+
tvm.s_tir.dlight.ApplyDefaultSchedule(tvm.s_tir.dlight.gpu.Fallback()),
],
name="pipeline",
)
diff --git a/tests/python/disco/test_ccl.py b/tests/python/disco/test_ccl.py
index 8a1518765f..bca28936d0 100644
--- a/tests/python/disco/test_ccl.py
+++ b/tests/python/disco/test_ccl.py
@@ -24,7 +24,7 @@ import pytest
import tvm
import tvm.testing
-from tvm import dlight as dl
+from tvm.s_tir import dlight as dl
from tvm import get_global_func
from tvm import relax as rx
from tvm.runtime import disco as di
diff --git a/tests/python/disco/test_loader.py
b/tests/python/disco/test_loader.py
index a68f539176..ddcd001b3e 100644
--- a/tests/python/disco/test_loader.py
+++ b/tests/python/disco/test_loader.py
@@ -23,7 +23,7 @@ import numpy as np
import tvm
import tvm.testing
-from tvm import dlight as dl
+from tvm.s_tir import dlight as dl
from tvm import relax as rx
from tvm_ffi import register_global_func
from tvm.contrib import tvmjs
diff --git a/tests/python/relax/backend/clml/utils.py
b/tests/python/relax/backend/clml/utils.py
index 0f3516f961..a5b7ebe9cf 100644
--- a/tests/python/relax/backend/clml/utils.py
+++ b/tests/python/relax/backend/clml/utils.py
@@ -24,7 +24,7 @@ import copy
from tvm import relax, rpc
from tvm.relax import transform
-from tvm import dlight as dl
+from tvm.s_tir import dlight as dl
from tvm.contrib import utils, ndk
from tvm.relax.backend.adreno.clml import OpenCLMLOffLoad
diff --git
a/tests/python/relax/nvshmem/test_runtime_builtin_kv_cache_transfer.py
b/tests/python/relax/nvshmem/test_runtime_builtin_kv_cache_transfer.py
index 5c994028ac..7f20381b48 100644
--- a/tests/python/relax/nvshmem/test_runtime_builtin_kv_cache_transfer.py
+++ b/tests/python/relax/nvshmem/test_runtime_builtin_kv_cache_transfer.py
@@ -25,7 +25,7 @@ import torch
import tvm
import tvm.testing
-from tvm import dlight as dl
+from tvm.s_tir import dlight as dl
from tvm.relax.frontend.nn.llm.kv_cache import (
AttnKind,
RopeMode,
diff --git a/tests/python/relax/test_backend_dispatch_sort_scan.py
b/tests/python/relax/test_backend_dispatch_sort_scan.py
index d48227fc62..564ea05274 100644
--- a/tests/python/relax/test_backend_dispatch_sort_scan.py
+++ b/tests/python/relax/test_backend_dispatch_sort_scan.py
@@ -21,7 +21,8 @@ import pytest
import tvm
import tvm.script
import tvm.testing
-from tvm import dlight, relax, tir, topi
+from tvm import relax, tir, topi
+from tvm.s_tir import dlight
from tvm.contrib.thrust import can_use_thrust
from tvm.ir.base import assert_structural_equal
from tvm.relax.backend import DispatchSortScan
diff --git
a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_cpu.py
b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_cpu.py
index 970cf38260..0c2350a5b2 100644
--- a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_cpu.py
+++ b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_cpu.py
@@ -24,7 +24,7 @@ import scipy.special
import tvm
import tvm.testing
-from tvm import dlight as dl
+from tvm.s_tir import dlight as dl
from tvm.relax.frontend.nn.llm.kv_cache import (
AttnKind,
_attention_decode_cpu,
diff --git
a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_flashinfer.py
b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_flashinfer.py
index 4aae9dec59..158e0fa800 100644
---
a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_flashinfer.py
+++
b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_flashinfer.py
@@ -21,7 +21,7 @@ import torch
import tvm
import tvm.testing
-from tvm import dlight as dl
+from tvm.s_tir import dlight as dl
from tvm import relax
from tvm.relax.frontend.nn.llm.kv_cache import (
AttnKind,
diff --git
a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_flashinfer.py
b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_flashinfer.py
index cd76f9ce20..c448a13a89 100644
---
a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_flashinfer.py
+++
b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_flashinfer.py
@@ -23,7 +23,7 @@ import torch
import tvm
import tvm.testing
-from tvm import dlight as dl
+from tvm.s_tir import dlight as dl
from tvm import relax
from tvm.relax.frontend.nn.llm.kv_cache import (
AttnKind,
diff --git
a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_tir.py
b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_tir.py
index efc0a5694c..48c9740032 100644
---
a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_tir.py
+++
b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_mla_tir.py
@@ -23,7 +23,7 @@ import torch
import tvm
import tvm.testing
-from tvm import dlight as dl
+from tvm.s_tir import dlight as dl
from tvm.relax.frontend.nn.llm.kv_cache import (
AttnKind,
RopeMode,
diff --git
a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py
b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py
index b80bd1acb7..d7e05a8943 100644
--- a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py
+++ b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py
@@ -22,7 +22,7 @@ import torch
import tvm
import tvm.testing
-from tvm import dlight as dl
+from tvm.s_tir import dlight as dl
from tvm.relax.frontend.nn.llm.kv_cache import (
AttnKind,
RopeMode,
diff --git a/tests/python/relax/test_runtime_builtin_rnn_state.py
b/tests/python/relax/test_runtime_builtin_rnn_state.py
index c34c957ce7..37666c7a7a 100644
--- a/tests/python/relax/test_runtime_builtin_rnn_state.py
+++ b/tests/python/relax/test_runtime_builtin_rnn_state.py
@@ -22,7 +22,7 @@ import pytest
import tvm
import tvm.testing
-from tvm import dlight as dl
+from tvm.s_tir import dlight as dl
from tvm import tir
from tvm.runtime import ShapeTuple
from tvm.script import tir as T
diff --git a/tests/python/relax/test_vm_build.py
b/tests/python/relax/test_vm_build.py
index 748fd8fe10..5078de1d39 100644
--- a/tests/python/relax/test_vm_build.py
+++ b/tests/python/relax/test_vm_build.py
@@ -1279,7 +1279,7 @@ def test_relax_module_with_multiple_targets(exec_mode):
seq = tvm.ir.transform.Sequential(
[
tvm.relax.transform.LegalizeOps(),
- tvm.dlight.ApplyDefaultSchedule(tvm.dlight.gpu.Fallback()),
+
tvm.s_tir.dlight.ApplyDefaultSchedule(tvm.s_tir.dlight.gpu.Fallback()),
],
name="LegalizeAndSchedule",
)
diff --git a/tests/python/dlight/test_benchmark.py
b/tests/python/s_tir/dlight/test_benchmark.py
similarity index 99%
rename from tests/python/dlight/test_benchmark.py
rename to tests/python/s_tir/dlight/test_benchmark.py
index 9679b32e9e..410d06ed75 100644
--- a/tests/python/dlight/test_benchmark.py
+++ b/tests/python/s_tir/dlight/test_benchmark.py
@@ -26,7 +26,7 @@ from tvm.script import tir as T
from tvm.script import relax as R
-from tvm.dlight.benchmark import (
+from tvm.s_tir.dlight.benchmark import (
benchmark,
benchmark_prim_func,
benchmark_relax_func,
diff --git a/tests/python/dlight/test_cpu_gemv.py
b/tests/python/s_tir/dlight/test_cpu_gemv.py
similarity index 99%
rename from tests/python/dlight/test_cpu_gemv.py
rename to tests/python/s_tir/dlight/test_cpu_gemv.py
index b5821b8089..e6e0ccf668 100644
--- a/tests/python/dlight/test_cpu_gemv.py
+++ b/tests/python/s_tir/dlight/test_cpu_gemv.py
@@ -18,7 +18,7 @@
import pytest
import tvm.testing
-from tvm import dlight as dl
+from tvm.s_tir import dlight as dl
from tvm.script import tir as T
from tvm.target import Target
diff --git a/tests/python/dlight/test_gpu_conv.py
b/tests/python/s_tir/dlight/test_gpu_conv.py
similarity index 99%
rename from tests/python/dlight/test_gpu_conv.py
rename to tests/python/s_tir/dlight/test_gpu_conv.py
index c69a3541ff..3f0d312b67 100644
--- a/tests/python/dlight/test_gpu_conv.py
+++ b/tests/python/s_tir/dlight/test_gpu_conv.py
@@ -17,7 +17,7 @@
# pylint: disable=missing-docstring
import tvm
import tvm.testing
-from tvm import dlight as dl
+from tvm.s_tir import dlight as dl
from tvm.script import tir as T
from tvm.target import Target
diff --git a/tests/python/dlight/test_gpu_fallback.py
b/tests/python/s_tir/dlight/test_gpu_fallback.py
similarity index 99%
rename from tests/python/dlight/test_gpu_fallback.py
rename to tests/python/s_tir/dlight/test_gpu_fallback.py
index d0fdfee0c5..d5a35c0e66 100644
--- a/tests/python/dlight/test_gpu_fallback.py
+++ b/tests/python/s_tir/dlight/test_gpu_fallback.py
@@ -16,7 +16,7 @@
# under the License.
# pylint: disable=missing-docstring
import tvm.testing
-from tvm import dlight as dl
+from tvm.s_tir import dlight as dl
from tvm.ir import assert_structural_equal
from tvm.script import ir as I
from tvm.script import tir as T
diff --git a/tests/python/dlight/test_gpu_gemv.py
b/tests/python/s_tir/dlight/test_gpu_gemv.py
similarity index 99%
rename from tests/python/dlight/test_gpu_gemv.py
rename to tests/python/s_tir/dlight/test_gpu_gemv.py
index af86ed1dbb..a23521c7e6 100644
--- a/tests/python/dlight/test_gpu_gemv.py
+++ b/tests/python/s_tir/dlight/test_gpu_gemv.py
@@ -17,7 +17,7 @@
# pylint: disable=missing-docstring
import tvm
import tvm.testing
-from tvm import dlight as dl
+from tvm.s_tir import dlight as dl
from tvm.script import tir as T
from tvm.target import Target
diff --git a/tests/python/dlight/test_gpu_general_reduction.py
b/tests/python/s_tir/dlight/test_gpu_general_reduction.py
similarity index 99%
rename from tests/python/dlight/test_gpu_general_reduction.py
rename to tests/python/s_tir/dlight/test_gpu_general_reduction.py
index 5586cc7409..2ccc6a4007 100644
--- a/tests/python/dlight/test_gpu_general_reduction.py
+++ b/tests/python/s_tir/dlight/test_gpu_general_reduction.py
@@ -17,7 +17,7 @@
# pylint: disable=missing-docstring
import tvm
import tvm.testing
-from tvm import dlight as dl
+from tvm.s_tir import dlight as dl
from tvm.ir import IRModule, assert_structural_equal
from tvm.script import ir as I
from tvm.script import tir as T
diff --git a/tests/python/dlight/test_gpu_low_batch_gemv.py
b/tests/python/s_tir/dlight/test_gpu_low_batch_gemv.py
similarity index 99%
rename from tests/python/dlight/test_gpu_low_batch_gemv.py
rename to tests/python/s_tir/dlight/test_gpu_low_batch_gemv.py
index ecfc6b5241..1dabd406bc 100644
--- a/tests/python/dlight/test_gpu_low_batch_gemv.py
+++ b/tests/python/s_tir/dlight/test_gpu_low_batch_gemv.py
@@ -17,7 +17,7 @@
# pylint: disable=missing-docstring
import tvm.testing
-from tvm import dlight as dl
+from tvm.s_tir import dlight as dl
from tvm.script import tir as T
from tvm.target import Target
diff --git a/tests/python/dlight/test_gpu_matmul.py
b/tests/python/s_tir/dlight/test_gpu_matmul.py
similarity index 99%
rename from tests/python/dlight/test_gpu_matmul.py
rename to tests/python/s_tir/dlight/test_gpu_matmul.py
index b4de74c18a..30f999ac72 100644
--- a/tests/python/dlight/test_gpu_matmul.py
+++ b/tests/python/s_tir/dlight/test_gpu_matmul.py
@@ -17,7 +17,7 @@
# pylint: disable=missing-docstring
import tvm
import tvm.testing
-from tvm import dlight as dl
+from tvm.s_tir import dlight as dl
from tvm.script import tir as T
from tvm.target import Target
diff --git a/tests/python/dlight/test_gpu_matmul_tensorize.py
b/tests/python/s_tir/dlight/test_gpu_matmul_tensorize.py
similarity index 99%
rename from tests/python/dlight/test_gpu_matmul_tensorize.py
rename to tests/python/s_tir/dlight/test_gpu_matmul_tensorize.py
index 4ae10f5d91..dd8fc9e631 100644
--- a/tests/python/dlight/test_gpu_matmul_tensorize.py
+++ b/tests/python/s_tir/dlight/test_gpu_matmul_tensorize.py
@@ -18,7 +18,7 @@
# flake8: noqa: E501
import tvm
import tvm.testing
-from tvm import dlight as dl
+from tvm.s_tir import dlight as dl
from tvm.script import tir as T
from tvm.target import Target
diff --git a/tests/python/dlight/test_gpu_reduction.py
b/tests/python/s_tir/dlight/test_gpu_reduction.py
similarity index 99%
rename from tests/python/dlight/test_gpu_reduction.py
rename to tests/python/s_tir/dlight/test_gpu_reduction.py
index ca51e1e40f..00d1dbb3d9 100644
--- a/tests/python/dlight/test_gpu_reduction.py
+++ b/tests/python/s_tir/dlight/test_gpu_reduction.py
@@ -17,7 +17,7 @@
# pylint:
disable=missing-docstring,line-too-long,invalid-name,too-few-public-methods,too-many-locals
import tvm.testing
-from tvm import dlight as dl
+from tvm.s_tir import dlight as dl
from tvm.ir import assert_structural_equal
from tvm.script import ir as I
from tvm.script import tir as T
diff --git a/tests/python/dlight/test_gpu_rmsnorm.py
b/tests/python/s_tir/dlight/test_gpu_rmsnorm.py
similarity index 99%
rename from tests/python/dlight/test_gpu_rmsnorm.py
rename to tests/python/s_tir/dlight/test_gpu_rmsnorm.py
index 6003eed3e9..8cbb792767 100644
--- a/tests/python/dlight/test_gpu_rmsnorm.py
+++ b/tests/python/s_tir/dlight/test_gpu_rmsnorm.py
@@ -18,7 +18,7 @@
import tvm.testing
from tvm.ir import IRModule, assert_structural_equal
-from tvm import dlight as dl
+from tvm.s_tir import dlight as dl
from tvm.script import ir as I
from tvm.target import Target
from tvm.script import tir as T
diff --git a/tests/python/dlight/test_gpu_transpose.py
b/tests/python/s_tir/dlight/test_gpu_transpose.py
similarity index 99%
rename from tests/python/dlight/test_gpu_transpose.py
rename to tests/python/s_tir/dlight/test_gpu_transpose.py
index bdd79e8f5e..576aaa2384 100644
--- a/tests/python/dlight/test_gpu_transpose.py
+++ b/tests/python/s_tir/dlight/test_gpu_transpose.py
@@ -16,7 +16,7 @@
# under the License.
# pylint: disable=missing-docstring
import tvm
-from tvm import dlight as dl
+from tvm.s_tir import dlight as dl
from tvm.ir import IRModule, assert_structural_equal
from tvm.script import ir as I
from tvm.script import tir as T
diff --git a/tests/python/dlight/test_primitives.py
b/tests/python/s_tir/dlight/test_primitives.py
similarity index 100%
rename from tests/python/dlight/test_primitives.py
rename to tests/python/s_tir/dlight/test_primitives.py
diff --git a/tests/scripts/release/make_notes.py
b/tests/scripts/release/make_notes.py
index f09a90649a..11626973d1 100644
--- a/tests/scripts/release/make_notes.py
+++ b/tests/scripts/release/make_notes.py
@@ -87,7 +87,7 @@ TAG_DICT = {
"unity": "Relax",
"transform": "Relax",
"kvcache": "Relax",
- "dlight": "Dlight",
+ "s_tir": "S-TIR",
"disco": "Disco",
"tvmscript": "TVMScript",
"tvmscripts": "TVMScript",
diff --git a/tests/scripts/task_python_unittest.sh
b/tests/scripts/task_python_unittest.sh
index 2bfe8823fe..3619b51e79 100755
--- a/tests/scripts/task_python_unittest.sh
+++ b/tests/scripts/task_python_unittest.sh
@@ -49,6 +49,7 @@ TEST_FILES=(
"testing"
"s_tir/base"
"s_tir/schedule"
+ "s_tir/dlight"
"tir-analysis"
"tir-base"
"tir-transform"
diff --git a/tests/scripts/unity/task_python_relax.sh
b/tests/scripts/unity/task_python_relax.sh
index c25cc6ec65..4af875830e 100755
--- a/tests/scripts/unity/task_python_relax.sh
+++ b/tests/scripts/unity/task_python_relax.sh
@@ -30,7 +30,6 @@ python3 -m pip install -v --target=python ./3rdparty/tvm-ffi/
# Run Relax tests
TVM_TEST_TARGETS="${TVM_RELAY_TEST_TARGETS:-llvm}" pytest tests/python/relax
-TVM_TEST_TARGETS="${TVM_RELAY_TEST_TARGETS:-llvm}" pytest tests/python/dlight
# Run Relax examples
# python3 ./apps/relax_examples/mlp.py
