This is an automated email from the ASF dual-hosted git repository. tlopex pushed a commit to branch tests/ci-test-pressure-audit in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 885347abbaa70229be614ff664f3780f739b5064 Author: tlopex <[email protected]> AuthorDate: Fri Jul 17 15:26:43 2026 -0400 [Tests][MetaSchedule] Reduce redundant tuning work --- .../relax/test_meta_schedule_relax_integration.py | 55 ++++++++-------------- .../relax/test_transform_meta_schedule_tuning.py | 25 +--------- 2 files changed, 20 insertions(+), 60 deletions(-) diff --git a/tests/python/relax/test_meta_schedule_relax_integration.py b/tests/python/relax/test_meta_schedule_relax_integration.py index e2d80f7f49..59f15205cd 100644 --- a/tests/python/relax/test_meta_schedule_relax_integration.py +++ b/tests/python/relax/test_meta_schedule_relax_integration.py @@ -17,10 +17,7 @@ # ruff: noqa: E501, F401, F841 """Integration test for MetaSchedule""" -import tempfile - import numpy as np -import pytest import tvm import tvm.testing @@ -81,42 +78,28 @@ def test_extracting_tasks(): assert len(extracted_tasks) == count -def test_compile_relax_with_database(): - """End-to-end test: tune with MetaSchedule then compile_relax with the database. +def test_compile_relax_queries_fused_tir_workloads(): + """The database must be queried after Relax operators are fused into TIR workloads. - Verifies that the pipeline ordering in compile_relax is correct: tasks are - extracted and tuned against fused-TIR keys, and compile_relax produces those - same keys (by running LegalizeOps + FuseOps + FuseTIR before applying the - database), so the scheduled kernels are actually picked up. + Applying the database to the raw Relax module yields no queries, so recording + the fused workload names directly verifies the compile_relax pipeline order. """ - pytest.importorskip("cloudpickle") # needed by meta_schedule popen workers - target = tvm.target.Target({"kind": "llvm", "num-cores": 1}) - - # Prepare the fused module whose TIR keys will populate the database. - fused_mod = Module0 - fused_mod = relax.transform.LegalizeOps()(fused_mod) - fused_mod = relax.transform.AnnotateTIROpPattern()(fused_mod) - fused_mod = relax.transform.FuseOps()(fused_mod) - fused_mod = relax.transform.FoldConstant()(fused_mod) - fused_mod = relax.transform.FuseTIR()(fused_mod) - - with tempfile.TemporaryDirectory() as work_dir: - database = ms.relax_integration.tune_relax( - fused_mod, - params={}, - target=target, - work_dir=work_dir, - max_trials_global=4, - ) - # compile_relax takes the raw module and builds the fused-TIR pipeline - # internally; the database keys must therefore match the ones above. - exe = ms.relax_integration.compile_relax( - database=database, - mod=Module0, - target=target, - params=None, - ) + queried_workloads = [] + + def record_workload(sch): + queried_workloads.append(str(sch.mod.attrs["task_name"])) + # Report a hit so compile_relax applies the returned no-op schedule. + return True + + database = ms.database.ScheduleFnDatabase(record_workload) + exe = ms.relax_integration.compile_relax( + database=database, + mod=Module0, + target=target, + params=None, + ) + assert sorted(queried_workloads) == ["conv2d", "fused_conv2d_add"] dev = tvm_cpu() vm = VirtualMachine(exe.jit(), dev) diff --git a/tests/python/relax/test_transform_meta_schedule_tuning.py b/tests/python/relax/test_transform_meta_schedule_tuning.py index 9863f0e702..7148d87dfe 100644 --- a/tests/python/relax/test_transform_meta_schedule_tuning.py +++ b/tests/python/relax/test_transform_meta_schedule_tuning.py @@ -119,32 +119,9 @@ def test_ms_tuning_irmodule(): assert not tvm_ffi.structural_equal(mod, out_mod) -def test_ms_tuning_primfunc(): +def test_ms_tuning_irmodule_op_names(): mod = InputModule assert isinstance(mod, IRModule) - with tempfile.TemporaryDirectory() as work_dir: - """ - # TODO(@sunggg): revisit when ready - with target, PassContext(trace=Trace(mod), opt_level=0): - tuning_pass = relax.transform.MetaScheduleTuneTIR( - work_dir=work_dir, max_trials_global=4 - ) - out_mod = tuning_pass(mod) - assert PassContext.current().get_trace_stack_size() == 1 - # TODO (@sunggg): Need to determine how to track subgraph-level tuning traces. - # Currently, we don't track this so the trace size. Revisit this later. - tvm.ir.assert_structural_equal(mod, out_mod) - """ - with target, PassContext(opt_level=0): - tuning_pass = relax.transform.MetaScheduleTuneIRMod( - params={}, work_dir=work_dir, max_trials_global=4 - ) - out_mod = tuning_pass(mod) - - application_pass = relax.transform.MetaScheduleApplyDatabase(work_dir) - out_mod = application_pass(mod) - assert not tvm_ffi.structural_equal(mod, out_mod) - with tempfile.TemporaryDirectory() as work_dir: with target, PassContext(opt_level=0): tuning_pass = relax.transform.MetaScheduleTuneIRMod(
