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/database.h>
#include <tvm/meta_schedule/runner.h>
#include <tvm/meta_schedule/tune_context.h>

namespace tvm {
namespace meta_schedule {

class TaskSchedulerNode : public runtime::Object {
 public:
  Array<TuneContext> tasks;
  Builder builder{nullptr};
  Runner runner{nullptr};
  Database database{nullptr};

  virtual ~TaskSchedulerNode() = default;

  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("tasks", &tasks);
    v->Visit("builder", &builder);
    v->Visit("runner", &runner);
    v->Visit("database", &database);
  }

  virtual void Tune();

  virtual void InitializeTask(int task_id);

  virtual void SetTaskStopped(int task_id);

  virtual bool IsTaskRunning(int task_id);

  virtual void 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 FSetTaskStopped = runtime::TypedPackedFunc<void(int)>;

  using FIsTaskRunning = runtime::TypedPackedFunc<bool(int)>;

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

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

  FTune f_tune;
  FInitializeTask f_initialize_task;
  FSetTaskStopped f_set_task_stopped;
  FIsTaskRunning f_is_task_running;
  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_set_task_stopped` is not visited
    // `f_is_task_running` is not visited
    // `f_join_running_task` is not visited
    // `f_next_task_id` is not visited
  }

  void Tune() final {
    if (f_tune == nullptr) {
      TaskSchedulerNode::Tune();
    } else {
      f_tune();
    }
  }

  void InitializeTask(int task_id) final {
    if (f_initialize_task == nullptr) {
      TaskSchedulerNode::InitializeTask(task_id);
    } else {
      f_initialize_task(task_id);
    }
  }

  void SetTaskStopped(int task_id) final {
    if (f_set_task_stopped == nullptr) {
      TaskSchedulerNode::SetTaskStopped(task_id);
    } else {
      f_set_task_stopped(task_id);
    }
  }

  bool IsTaskRunning(int task_id) final {
    if (f_is_task_running == nullptr) {
      return TaskSchedulerNode::IsTaskRunning(task_id);
    } else {
      return f_is_task_running(task_id);
    }
  }

  void JoinRunningTask(int task_id) final {
    if (f_join_running_task == nullptr) {
      return TaskSchedulerNode::JoinRunningTask(task_id);
    } else {
      return f_join_running_task(task_id);
    }
  }

  int NextTaskId() final {
    ICHECK(f_next_task_id != nullptr) << "PyTaskScheduler's NextTaskId method not implemented!";
    return f_next_task_id();
  }

  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,             //
                                          Database database);        //
  TVM_DLL static TaskScheduler PyTaskScheduler(
      Array<TuneContext> tasks,                                   //
      Builder builder,                                            //
      Runner runner,                                              //
      Database database,                                          //
      PyTaskSchedulerNode::FTune f_tune,                          //
      PyTaskSchedulerNode::FInitializeTask f_initialize_task,     //
      PyTaskSchedulerNode::FSetTaskStopped f_set_task_stopped,    //
      PyTaskSchedulerNode::FIsTaskRunning f_is_task_running,      //
      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_