task_scheduler.h

Go to the documentation of this file.

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
#ifndef TVM_META_SCHEDULE_TASK_SCHEDULER_H_
#define TVM_META_SCHEDULE_TASK_SCHEDULER_H_

#include <tvm/meta_schedule/builder.h>
#include <tvm/meta_schedule/cost_model.h>
#include <tvm/meta_schedule/database.h>
#include <tvm/meta_schedule/measure_callback.h>
#include <tvm/meta_schedule/runner.h>
#include <tvm/meta_schedule/tune_context.h>
#include <tvm/node/reflection.h>
#include <tvm/runtime/container/array.h>
#include <tvm/runtime/container/optional.h>
#include <tvm/runtime/object.h>
#include <tvm/runtime/packed_func.h>
#include <tvm/support/random_engine.h>

namespace tvm {
namespace meta_schedule {

class TaskSchedulerNode : public runtime::Object {
 public:
  Array<TuneContext> tasks;
  Builder builder{nullptr};
  Runner runner{nullptr};
  Optional<Database> database;
  Optional<CostModel> cost_model;
  Array<MeasureCallback> measure_callbacks;
  int max_trials;
  int num_trials_already;
  PackedFunc logging_func;

  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
  }

  virtual void Tune();

  virtual void InitializeTask(int task_id);

  virtual void TouchTask(int task_id);

  virtual Array<RunnerResult> JoinRunningTask(int task_id);

  virtual int NextTaskId() = 0;

  static constexpr const char* _type_key = "meta_schedule.TaskScheduler";
  TVM_DECLARE_BASE_OBJECT_INFO(TaskSchedulerNode, Object);
};

class TaskScheduler;

class PyTaskSchedulerNode : public TaskSchedulerNode {
 public:
  using FTune = runtime::TypedPackedFunc<void()>;

  using FInitializeTask = runtime::TypedPackedFunc<void(int)>;

  using FTouchTask = runtime::TypedPackedFunc<void(int)>;

  using FJoinRunningTask = runtime::TypedPackedFunc<Array<RunnerResult>(int)>;

  using FNextTaskId = runtime::TypedPackedFunc<int()>;

  FTune f_tune;
  FInitializeTask f_initialize_task;
  FTouchTask f_touch_task;
  FJoinRunningTask f_join_running_task;
  FNextTaskId f_next_task_id;

  void VisitAttrs(tvm::AttrVisitor* v) {
    // `f_tune` is not visited
    // `f_initialize_task` is not visited
    // `f_touch_task` is not visited
    // `f_join_running_task` is not visited
    // `f_next_task_id` is not visited
  }

  void Tune() final;
  void InitializeTask(int task_id) final;
  void TouchTask(int task_id) final;
  Array<RunnerResult> JoinRunningTask(int task_id) final;
  int NextTaskId() final;

  static constexpr const char* _type_key = "meta_schedule.PyTaskScheduler";
  TVM_DECLARE_FINAL_OBJECT_INFO(PyTaskSchedulerNode, TaskSchedulerNode);
};

class TaskScheduler : public runtime::ObjectRef {
 public:
  TVM_DLL static TaskScheduler RoundRobin(Array<TuneContext> tasks,                            //
                                          Builder builder,                                     //
                                          Runner runner,                                       //
                                          Optional<Database> database,                         //
                                          Optional<CostModel> cost_model,                      //
                                          Optional<Array<MeasureCallback>> measure_callbacks,  //
                                          int max_trials,                                      //
                                          PackedFunc logging_func);
  TVM_DLL static TaskScheduler GradientBased(Array<TuneContext> tasks,
                                             Array<FloatImm> task_weights,                        //
                                             Builder builder,                                     //
                                             Runner runner,                                       //
                                             Optional<Database> database,                         //
                                             Optional<CostModel> cost_model,                      //
                                             Optional<Array<MeasureCallback>> measure_callbacks,  //
                                             int max_trials,                                      //
                                             PackedFunc logging_func,                             //
                                             double alpha,                                        //
                                             int window_size,                                     //
                                             support::LinearCongruentialEngine::TRandState seed);
  TVM_DLL static TaskScheduler PyTaskScheduler(
      Array<TuneContext> tasks,                                   //
      Builder builder,                                            //
      Runner runner,                                              //
      Optional<Database> database,                                //
      Optional<CostModel> cost_model,                             //
      Optional<Array<MeasureCallback>> measure_callbacks,         //
      int max_trials,                                             //
      PackedFunc logging_func,                                    //
      PyTaskSchedulerNode::FTune f_tune,                          //
      PyTaskSchedulerNode::FInitializeTask f_initialize_task,     //
      PyTaskSchedulerNode::FTouchTask f_touch_task,               //
      PyTaskSchedulerNode::FJoinRunningTask f_join_running_task,  //
      PyTaskSchedulerNode::FNextTaskId f_next_task_id);
  TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(TaskScheduler, ObjectRef, TaskSchedulerNode);
};

}  // namespace meta_schedule
}  // namespace tvm

#endif  // TVM_META_SCHEDULE_TASK_SCHEDULER_H_