diff --git a/include/tvm/meta_schedule/integration.h b/include/tvm/meta_schedule/integration.h index 9a8699b2fab94..bb384ba3ae837 100644 --- a/include/tvm/meta_schedule/integration.h +++ b/include/tvm/meta_schedule/integration.h @@ -145,38 +145,6 @@ class MetaScheduleContext : public runtime::ObjectRef { void ExitWithScope(); }; -/**************** TaskExtraction ****************/ - -/*! - * \brief An integration context for task extraction - */ -class TaskExtractionNode : public MetaScheduleContextNode { - public: - /*! \brief The extracted tasks */ - Array tasks{nullptr}; - - void VisitAttrs(AttrVisitor* v) { v->Visit("tasks", &tasks); } - - // Inherited from base class - Optional Query(runtime::String task_name, IRModule mod, Target target, - Optional> dispatched) final; - - static constexpr const char* _type_key = "meta_schedule.TaskExtraction"; - TVM_DECLARE_FINAL_OBJECT_INFO(TaskExtractionNode, MetaScheduleContextNode); -}; - -/*! - * \brief Managed reference to TaskExtractionNode - * \sa TaskExtractionNode - */ -class TaskExtraction : public MetaScheduleContext { - public: - /*! \brief The path to a cache file storing extracted tasks */ - TaskExtraction(); - TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(TaskExtraction, MetaScheduleContext, - TaskExtractionNode); -}; - /**************** ApplyHistoryBest ****************/ /*! diff --git a/python/tvm/meta_schedule/integration.py b/python/tvm/meta_schedule/integration.py index 2477aa11536cd..b39f318e0f01c 100644 --- a/python/tvm/meta_schedule/integration.py +++ b/python/tvm/meta_schedule/integration.py @@ -179,17 +179,6 @@ def __exit__(self, ptype, value, trace) -> None: _ffi_api.MetaScheduleContextExitScope(self) # type: ignore # pylint: disable=no-member -@register_object("meta_schedule.TaskExtraction") -class TaskExtraction(MetaScheduleContext): - """An integration context for task extraction""" - - tasks: List[ExtractedTask] - """The extracted tasks""" - - def __init__(self) -> None: - self.__init_handle_by_constructor__(_ffi_api.TaskExtraction) # type: ignore # pylint: disable=no-member - - @register_object("meta_schedule.ApplyHistoryBest") class ApplyHistoryBest(MetaScheduleContext): """An integration context that allows application of historically best record from database""" diff --git a/src/meta_schedule/integration.cc b/src/meta_schedule/integration.cc index 4f9055bf5bbad..1fba65493f07b 100644 --- a/src/meta_schedule/integration.cc +++ b/src/meta_schedule/integration.cc @@ -94,25 +94,6 @@ Optional MetaScheduleContext::QueryInsideWithScope( return NullOpt; } -/**************** TaskExtraction ****************/ - -TaskExtraction::TaskExtraction() { - ObjectPtr n = make_object(); - n->tasks = Array(); - data_ = n; -} - -Optional TaskExtractionNode::Query(runtime::String task_name, IRModule mod, - Target target, Optional> dispatched) { - ICHECK(dispatched.defined()); - ICHECK_EQ(dispatched.value().size(), 1); - IRModule prim_mod = dispatched.value()[0]; - ICHECK(HasOnlyOneFunction(prim_mod)) << prim_mod; - ICHECK(HasOnlyOneFunction(mod)) << mod; - tasks.push_back(ExtractedTask(task_name, mod, target, {prim_mod})); - return NullOpt; -} - /**************** ApplyHistoryBest ****************/ ApplyHistoryBest::ApplyHistoryBest(Database database) { @@ -158,7 +139,6 @@ class MetaScheduleContextInternal { TVM_REGISTER_NODE_TYPE(ExtractedTaskNode); TVM_REGISTER_OBJECT_TYPE(MetaScheduleContextNode); -TVM_REGISTER_NODE_TYPE(TaskExtractionNode); TVM_REGISTER_NODE_TYPE(ApplyHistoryBestNode); TVM_REGISTER_GLOBAL("meta_schedule.ExtractedTask") @@ -176,9 +156,6 @@ TVM_REGISTER_GLOBAL("meta_schedule.MetaScheduleContextQueryInsideWithScope") .set_body_typed(MetaScheduleContext::QueryInsideWithScope); TVM_REGISTER_GLOBAL("meta_schedule.MetaScheduleContextQuery") .set_body_method(&MetaScheduleContextNode::Query); -TVM_REGISTER_GLOBAL("meta_schedule.TaskExtraction").set_body_typed([]() -> TaskExtraction { - return TaskExtraction(); -}); TVM_REGISTER_GLOBAL("meta_schedule.ApplyHistoryBest") .set_body_typed([](Database database) -> ApplyHistoryBest { return ApplyHistoryBest(database);