This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch refactor-s1 in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 3b0c654c89c4429a844d4104866dc99c3a25e7f8 Author: tqchen <[email protected]> AuthorDate: Mon Apr 14 10:20:49 2025 -0400 Fixup MS to ensure concrete value --- src/tir/schedule/concrete_schedule.cc | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/src/tir/schedule/concrete_schedule.cc b/src/tir/schedule/concrete_schedule.cc index 9b0ce75bed..80796a243e 100644 --- a/src/tir/schedule/concrete_schedule.cc +++ b/src/tir/schedule/concrete_schedule.cc @@ -920,18 +920,26 @@ Any ConcreteScheduleNode::CheckAndGetAnnotationValue(const ffi::Any& ann_val) { if (ann_val.type_index() < ffi::TypeIndex::kTVMFFIStaticObjectBegin) { return ann_val; } - + // prefer to return int/float literals for annotations if (auto opt_intimm = ann_val.as<IntImm>()) { - return *std::move(opt_intimm); + return (*std::move(opt_intimm))->value; } if (auto opt_floatimm = ann_val.as<FloatImm>()) { - return *std::move(opt_floatimm); + return (*std::move(opt_floatimm))->value; } if (const auto* expr = ann_val.as<PrimExprNode>()) { ICHECK(!expr->IsInstance<StringImmNode>()) << "TypeError: runtime::String is expected, but gets StringImm"; - return this->Get(GetRef<PrimExpr>(expr)); + auto res_expr = this->Get(GetRef<PrimExpr>(expr)); + // prefer to return int/float literals for annotations + if (auto opt_intimm = res_expr.as<IntImm>()) { + return (*std::move(opt_intimm))->value; + } + if (auto opt_floatimm = res_expr.as<FloatImm>()) { + return (*std::move(opt_floatimm))->value; + } + return res_expr; } if (const auto* arr = ann_val.as<ArrayNode>()) { Array<Any> result;
