Skip to content

[Fix][MetaSchedule][RISC-V] Use TuneContext target for default RVV rules #20313

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ZephyrLi-pro wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[Fix][MetaSchedule][RISC-V] Use TuneContext target for default RVV rules #20313

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 8 additions & 2 deletions include/tvm/s_tir/meta_schedule/schedule_rule.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,9 @@
#include <tvm/s_tir/schedule/schedule.h>

namespace tvm {

class Target;

namespace s_tir {
using namespace tvm::prim;
namespace meta_schedule {
Expand Down Expand Up @@ -306,8 +309,11 @@ class ScheduleRule : public ffi::ObjectRef {
TVM_DLL static ffi::Array<ScheduleRule, void> DefaultHexagon();
/*! \brief Create default schedule rules for ARM CPU (NEON and DOTPROD) */
TVM_DLL static ffi::Array<ScheduleRule, void> DefaultARM(const ffi::String& type);
/*! \brief Create default schedule rules for RISCV CPU (RVV) */
TVM_DLL static ffi::Array<ScheduleRule, void> 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<ScheduleRule, void> DefaultRISCV(const Target& target);

TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ScheduleRule, ffi::ObjectRef, ScheduleRuleNode);
};
Expand Down
10 changes: 5 additions & 5 deletions src/s_tir/meta_schedule/schedule_rule/schedule_rule.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
*/
#include <tvm/ffi/dtype.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/target/target.h>

#include "../utils.h"

Expand Down Expand Up @@ -305,7 +306,7 @@ ffi::Array<ScheduleRule> ScheduleRule::DefaultHexagon() {
};
}

ffi::Array<ScheduleRule> ScheduleRule::DefaultRISCV(const int vlen) {
ffi::Array<ScheduleRule> ScheduleRule::DefaultRISCV(const Target& target) {
ffi::Array<ScheduleRule> rules;
rules.push_back(ScheduleRule::ApplyCustomRule());
rules.push_back(ScheduleRule::InlineConstantScalars());
Expand All @@ -320,15 +321,14 @@ ffi::Array<ScheduleRule> ScheduleRule::DefaultRISCV(const int vlen) {
rules.push_back(ScheduleRule::AddRFactor(
/*max_jobs_per_core=*/16,
/*max_innermost_factor=*/static_cast<int64_t>(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<ffi::Map<ffi::String, int>>();
const auto rvv_kernels_inventory =
reg_rvv_intrinsics(target, /*inventory_only=*/true).cast<ffi::Map<ffi::String, int>>();
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,
Expand Down
5 changes: 1 addition & 4 deletions src/s_tir/meta_schedule/space_generator/space_generator.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<int>();
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") {
Expand Down
26 changes: 26 additions & 0 deletions tests/python/s_tir/meta_schedule/test_meta_schedule_post_order_apply.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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",
[
Expand Down
Loading