Removed TaskExtraction node

Author: masahi

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<ExtractedTask> tasks{nullptr};
-
-  void VisitAttrs(AttrVisitor* v) { v->Visit("tasks", &tasks); }
-
-  // Inherited from base class
-  Optional<ObjectRef> Query(runtime::String task_name, IRModule mod, Target target,
-                            Optional<Array<IRModule>> 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 ****************/
 
 /*!

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"""

src/meta_schedule/integration.cc

@@ -94,25 +94,6 @@ Optional<ObjectRef> MetaScheduleContext::QueryInsideWithScope(
   return NullOpt;
 }
 
-/**************** TaskExtraction ****************/
-
-TaskExtraction::TaskExtraction() {
-  ObjectPtr<TaskExtractionNode> n = make_object<TaskExtractionNode>();
-  n->tasks = Array<ExtractedTask>();
-  data_ = n;
-}
-
-Optional<ObjectRef> TaskExtractionNode::Query(runtime::String task_name, IRModule mod,
-                                              Target target, Optional<Array<IRModule>> dispatched) {
-  ICHECK(dispatched.defined());
-  ICHECK_EQ(dispatched.value().size(), 1);
-  IRModule prim_mod = dispatched.value()[0];
-  ICHECK(HasOnlyOneFunction<tir::PrimFunc>(prim_mod)) << prim_mod;
-  ICHECK(HasOnlyOneFunction<relay::Function>(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<MetaScheduleContext>(&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);