From 8b856a888acdfc88e80c1630783eabd53f80d558 Mon Sep 17 00:00:00 2001 From: Zephyr Li Date: Fri, 11 Sep 2026 14:19:39 +0800 Subject: [PATCH] [Fix][MetaSchedule][RISC-V] Use TuneContext target for default RVV rules Signed-off-by: Zephyr Li --- .../tvm/s_tir/meta_schedule/schedule_rule.h | 10 +++++-- .../schedule_rule/schedule_rule.cc | 10 +++---- .../space_generator/space_generator.cc | 5 +--- .../test_meta_schedule_post_order_apply.py | 26 +++++++++++++++++++ 4 files changed, 40 insertions(+), 11 deletions(-) diff --git a/include/tvm/s_tir/meta_schedule/schedule_rule.h b/include/tvm/s_tir/meta_schedule/schedule_rule.h index b6b164df7be4..6da9d5f11aeb 100644 --- a/include/tvm/s_tir/meta_schedule/schedule_rule.h +++ b/include/tvm/s_tir/meta_schedule/schedule_rule.h @@ -31,6 +31,9 @@ #include namespace tvm { + +class Target; + namespace s_tir { using namespace tvm::prim; namespace meta_schedule { @@ -306,8 +309,11 @@ class ScheduleRule : public ffi::ObjectRef { TVM_DLL static ffi::Array DefaultHexagon(); /*! \brief Create default schedule rules for ARM CPU (NEON and DOTPROD) */ TVM_DLL static ffi::Array DefaultARM(const ffi::String& type); - /*! \brief Create default schedule rules for RISCV CPU (RVV) */ - TVM_DLL static ffi::Array DefaultRISCV(int vlen); + /*! + * \brief Create default schedule rules for RISC-V CPU (RVV) + * \param target The RISC-V target used to query and register RVV tensor intrinsics + */ + TVM_DLL static ffi::Array DefaultRISCV(const Target& target); TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ScheduleRule, ffi::ObjectRef, ScheduleRuleNode); }; diff --git a/src/s_tir/meta_schedule/schedule_rule/schedule_rule.cc b/src/s_tir/meta_schedule/schedule_rule/schedule_rule.cc index 8908bee1b9e7..ea76f301de2b 100644 --- a/src/s_tir/meta_schedule/schedule_rule/schedule_rule.cc +++ b/src/s_tir/meta_schedule/schedule_rule/schedule_rule.cc @@ -18,6 +18,7 @@ */ #include #include +#include #include "../utils.h" @@ -305,7 +306,7 @@ ffi::Array ScheduleRule::DefaultHexagon() { }; } -ffi::Array ScheduleRule::DefaultRISCV(const int vlen) { +ffi::Array ScheduleRule::DefaultRISCV(const Target& target) { ffi::Array rules; rules.push_back(ScheduleRule::ApplyCustomRule()); rules.push_back(ScheduleRule::InlineConstantScalars()); @@ -320,15 +321,14 @@ ffi::Array ScheduleRule::DefaultRISCV(const int vlen) { rules.push_back(ScheduleRule::AddRFactor( /*max_jobs_per_core=*/16, /*max_innermost_factor=*/static_cast(64))); - auto current_target = tvm::Target::Current(); const auto reg_rvv_intrinsics = tvm::ffi::Function::GetGlobalRequired("tirx.tensor_intrin.register_rvv_isa_intrinsics"); - const auto rvv_kernels_inventory = reg_rvv_intrinsics(current_target, /* inventory_only */ true) - .cast>(); + const auto rvv_kernels_inventory = + reg_rvv_intrinsics(target, /*inventory_only=*/true).cast>(); for (const auto& intrin : rvv_kernels_inventory) { if (!tirx::TensorIntrin::Get(intrin.first, /*allow_missing*/ true)) { // on demand intrinsic register - reg_rvv_intrinsics(current_target, /* inventory_only */ false); + reg_rvv_intrinsics(target, /*inventory_only=*/false); } rules.push_back(ScheduleRule::MultiLevelTilingWithIntrin( /*intrin_name=*/intrin.first, diff --git a/src/s_tir/meta_schedule/space_generator/space_generator.cc b/src/s_tir/meta_schedule/space_generator/space_generator.cc index 6879c9d660a4..f717cad34385 100644 --- a/src/s_tir/meta_schedule/space_generator/space_generator.cc +++ b/src/s_tir/meta_schedule/space_generator/space_generator.cc @@ -125,10 +125,7 @@ void SpaceGeneratorNode::InitializeWithTuneContext(const TuneContext& context) { default_postprocs = Postproc::DefaultCPUTensorization(); default_mutator_probs = Mutator::DefaultLLVM(); } else if (kind == "rvv") { - static auto llvm_get_vector_width = - tvm::ffi::Function::GetGlobalRequired("target.llvm_get_vector_width"); - const int vlen = llvm_get_vector_width(context->target.value()).cast(); - default_sch_rules = ScheduleRule::DefaultRISCV(vlen); + default_sch_rules = ScheduleRule::DefaultRISCV(context->target.value()); default_postprocs = Postproc::DefaultRISCV(); default_mutator_probs = Mutator::DefaultLLVM(); } else if (kind == "asimd") { 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 9ec11c077bd0..85be07a4120a 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 @@ -35,6 +35,7 @@ from tvm.s_tir.schedule import SBlockRV, Schedule from tvm.script import tirx as T from tvm.target import Target +from tvm.testing import env # pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument, # fmt: off @@ -429,6 +430,31 @@ def _get_sch(filter_fn): assert len(schs) == 8 +@pytest.mark.skipif(not env.has_llvm_min_version(14), reason="need llvm >= 14") +def test_meta_schedule_default_rvv_rules_without_target_context(): + target = Target( + { + "kind": "llvm", + "device": "riscv_cpu", + "mtriple": "riscv64-linux-gnu", + "mcpu": "generic-rv64", + "mattr": ["+64bit", "+a", "+c", "+d", "+f", "+m", "+v"], + "num-cores": 2, + } + ) + mod = IRModule({"main": get_matmul_packed(16, 16, 16, "int8", "int8", "int32")}) + + assert Target.current(allow_none=True) is None + context = TuneContext( + mod=mod, + target=target, + task_name="RISC-V RVV Default Rules", + space_generator=PostOrderApply(), + ) + + assert len(context.space_generator.sch_rules) > 0 + + @pytest.mark.parametrize( "target,mod,expected_intr", [