This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch tvm-further-cleanup-python-tests-followup in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 540957cb27e49fa4360183e7d71bc2df38b9bae0 Author: Tianqi Chen <[email protected]> AuthorDate: Sun Jul 5 18:51:40 2026 +0000 [FIX][MetaSchedule] Accept scalar parallel annotations --- src/s_tir/meta_schedule/mutator/mutate_parallel.cc | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/s_tir/meta_schedule/mutator/mutate_parallel.cc b/src/s_tir/meta_schedule/mutator/mutate_parallel.cc index a0e3639ff9..2bdfa60b50 100644 --- a/src/s_tir/meta_schedule/mutator/mutate_parallel.cc +++ b/src/s_tir/meta_schedule/mutator/mutate_parallel.cc @@ -248,7 +248,11 @@ bool FindParallelDecision(const Trace& trace, TRandState* rand_state, get_sblock_insts.at(ann_inst->inputs[0].as_or_throw<s_tir::SBlockRV>().get()); TVM_FFI_ICHECK_EQ(get_sblock_inst->attrs.size(), 2); candidate->inst = ffi::GetRef<Instruction>(ann_inst); - candidate->parallel_extent = ann_inst->inputs[1].as_or_throw<IntImm>()->value; + if (auto ann_val = ann_inst->inputs[1].as<IntImm>()) { + candidate->parallel_extent = ann_val.value()->value; + } else { + candidate->parallel_extent = ann_inst->inputs[1].cast<int64_t>(); + } candidate->block_name = get_sblock_inst->attrs[0].as_or_throw<ffi::String>(); candidate->func_name = get_sblock_inst->attrs[1].as_or_throw<ffi::String>(); return true;
