junrushao commented on code in PR #12895:
URL: https://github.com/apache/tvm/pull/12895#discussion_r985190584
##########
include/tvm/meta_schedule/task_scheduler.h:
##########
@@ -73,66 +127,71 @@ namespace meta_schedule {
*/
class TaskSchedulerNode : public runtime::Object {
public:
- /*! \brief The tasks to be tuned */
- Array<TuneContext> tasks;
- /*! \brief The builder of the scheduler. */
- Builder builder{nullptr};
- /*! \brief The runner of the scheduler. */
- Runner runner{nullptr};
- /*! \brief The database of the scheduler. */
- Optional<Database> database;
- /*! \brief The cost model of the scheduler. */
- Optional<CostModel> cost_model;
+ /*! \brief The tuning task's logging function. */
+ PackedFunc logger;
+ /*! \brief Records for each task */
+ Array<TaskRecord> tasks_;
/*! \brief The list of measure callbacks of the scheduler. */
- Array<MeasureCallback> measure_callbacks;
- /*! \brief The maximum number of trials allowed. */
- int max_trials;
- /*! \brief The number of trials already conducted. */
- int num_trials_already;
- /*! \brief The tuning task's logging function. t*/
- PackedFunc logging_func;
+ Array<MeasureCallback> measure_callbacks_;
+ /*! \brief The database used in tuning */
+ Optional<Database> database_;
+ /*! \brief The cost model used in tuning */
+ Optional<CostModel> cost_model_;
+ /*! \brief The number of remaining tasks to be tuned. */
+ int remaining_tasks_;
/*! \brief The default destructor. */
virtual ~TaskSchedulerNode() = default;
void VisitAttrs(tvm::AttrVisitor* v) {
- v->Visit("tasks", &tasks);
- v->Visit("builder", &builder);
- v->Visit("runner", &runner);
- v->Visit("database", &database);
- v->Visit("cost_model", &cost_model);
- v->Visit("measure_callbacks", &measure_callbacks);
- v->Visit("max_trials", &max_trials);
- v->Visit("num_trials_already", &num_trials_already);
- // `logging_func` is not visited
+ // `logger` is not visited
+ v->Visit("tasks_", &tasks_);
+ v->Visit("measure_callbacks_", &measure_callbacks_);
+ v->Visit("database_", &database_);
+ v->Visit("cost_model_", &cost_model_);
+ v->Visit("remaining_tasks_", &remaining_tasks_);
}
- /*! \brief Auto-tuning. */
- virtual void Tune();
-
- /*!
- * \brief Initialize modules of the given task.
- * \param task_id The task id to be initialized.
- */
- virtual void InitializeTask(int task_id);
-
/*!
- * \brief Touch the task and update its status
- * \param task_id The task id to be checked.
+ * \brief Fetch the next task id.
+ * \return The next task id.
*/
- virtual void TouchTask(int task_id);
-
+ virtual int NextTaskId() = 0;
/*!
* \brief Wait until the task is finished.
* \param task_id The task id to be joined.
+ * \return The results from the runner.
*/
virtual Array<RunnerResult> JoinRunningTask(int task_id);
-
/*!
- * \brief Fetch the next task id.
- * \return The next task id.
+ * \brief Jointly tune a given list of tasks.
+ * \param tasks The tasks to be tuned
+ * \param task_weights The weight of each task
+ * \param max_trials_global The maximum number of trials to be performed
globally
+ * \param max_trials_per_task The maximum number of trials to be performed
for each task
+ * \param num_trials_per_iter The number of trials to be performed in each
iteration
+ * \param builder The MetaSchedule builder
+ * \param runner The MetaSchedule runner
+ * \param measure_callbacks The callbacks to be called after each measurement
+ * \param database The database used in tuning
+ * \param cost_model The cost model used in tuning
*/
- virtual int NextTaskId() = 0;
+ virtual void Tune(Array<TuneContext> tasks, Array<FloatImm> task_weights,
int max_trials_global,
Review Comment:
Well, sure! Certainly it's personal preference but indeed your suggestion
makes sense
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
