database.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_DATABASE_H_
#define TVM_META_SCHEDULE_DATABASE_H_
 
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/string.h>
#include <tvm/ir/expr.h>
#include <tvm/ir/module.h>
#include <tvm/meta_schedule/arg_info.h>
#include <tvm/runtime/object.h>
#include <tvm/target/target.h>
#include <tvm/tir/schedule/schedule.h>
#include <tvm/tir/schedule/trace.h>
 
#include <memory>
 
namespace tvm {
namespace meta_schedule {
 
class ModuleEquality;
 
class WorkloadNode : public runtime::Object {
 public:
  using THashCode = size_t;
  IRModule mod;
  THashCode shash;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<WorkloadNode>().def_ro("mod", &WorkloadNode::mod);
  }
 
  static constexpr const char* _type_key = "meta_schedule.Workload";
 
  TVM_DECLARE_FINAL_OBJECT_INFO(WorkloadNode, runtime::Object);
 
  ObjectRef AsJSON() const;
};
 
class Workload : public runtime::ObjectRef {
 public:
  using THashCode = WorkloadNode::THashCode;
  TVM_DLL explicit Workload(IRModule mod);
  TVM_DLL explicit Workload(IRModule mod, THashCode shash);
  TVM_DLL static Workload FromJSON(const ObjectRef& json_obj);
 
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(Workload, runtime::ObjectRef, WorkloadNode);
};
 
struct WorkloadHash {
  size_t operator()(const Workload& a) const { return a->shash; }
};
 
struct WorkloadEqual {
  explicit WorkloadEqual(const ModuleEquality& mod_eq) : mod_eq_(mod_eq) {}
 
  bool operator()(const Workload& a, const Workload& b) const;
 
 private:
  const ModuleEquality& mod_eq_;
};
 
class MeasureCandidate;
 
class TuningRecordNode : public runtime::Object {
 public:
  tir::Trace trace;
  Workload workload{nullptr};
  Optional<Array<FloatImm>> run_secs;
  Optional<Target> target;
  Optional<Array<ArgInfo>> args_info;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<TuningRecordNode>()
        .def_ro("trace", &TuningRecordNode::trace)
        .def_ro("workload", &TuningRecordNode::workload)
        .def_ro("run_secs", &TuningRecordNode::run_secs)
        .def_ro("target", &TuningRecordNode::target)
        .def_ro("args_info", &TuningRecordNode::args_info);
  }
 
  static constexpr const char* _type_key = "meta_schedule.TuningRecord";
 
  TVM_DECLARE_FINAL_OBJECT_INFO(TuningRecordNode, runtime::Object);
 
  MeasureCandidate AsMeasureCandidate() const;
  ObjectRef AsJSON() const;
  bool IsValid() const;
};
 
class TuningRecord : public runtime::ObjectRef {
 public:
  TVM_DLL explicit TuningRecord(tir::Trace trace, Workload workload,
                                Optional<Array<FloatImm>> run_secs, Optional<Target> target,
                                Optional<Array<ArgInfo>> args_info);
  TVM_DLL static TuningRecord FromJSON(const ObjectRef& json_obj, const Workload& workload);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(TuningRecord, runtime::ObjectRef, TuningRecordNode);
};
 
class Database;
 
/* \brief The abstract interface of database. */
class DatabaseNode : public runtime::Object {
 public:
  explicit DatabaseNode(String mod_eq_name = "structural");
 
  virtual ~DatabaseNode();
  virtual bool HasWorkload(const IRModule& mod) = 0;
  virtual Workload CommitWorkload(const IRModule& mod) = 0;
  virtual void CommitTuningRecord(const TuningRecord& record) = 0;
  virtual Array<TuningRecord> GetTopK(const Workload& workload, int top_k) = 0;
  virtual Array<TuningRecord> GetAllTuningRecords() = 0;
  virtual int64_t Size() = 0;
  virtual Optional<TuningRecord> QueryTuningRecord(const IRModule& mod, const Target& target,
                                                   const String& workload_name);
  virtual Optional<tir::Schedule> QuerySchedule(const IRModule& mod, const Target& target,
                                                const String& workload_name);
  virtual Optional<IRModule> QueryIRModule(const IRModule& mod, const Target& target,
                                           const String& workload_name);
  void DumpPruned(Database destination);
  const ModuleEquality& GetModuleEquality() const {
    ICHECK(mod_eq_);
    return *mod_eq_;
  }
 
  static constexpr const char* _type_key = "meta_schedule.Database";
  TVM_DECLARE_BASE_OBJECT_INFO(DatabaseNode, runtime::Object);
 
 private:
  std::unique_ptr<ModuleEquality> mod_eq_;
};
 
class PyDatabaseNode : public DatabaseNode {
 public:
  explicit PyDatabaseNode(String mod_eq_name = "structural");
 
  using FHasWorkload = ffi::TypedFunction<bool(const IRModule&)>;
  using FCommitWorkload = ffi::TypedFunction<Workload(const IRModule&)>;
  using FCommitTuningRecord = ffi::TypedFunction<void(const TuningRecord&)>;
  using FGetTopK = ffi::TypedFunction<Array<TuningRecord>(const Workload&, int)>;
  using FGetAllTuningRecords = ffi::TypedFunction<Array<TuningRecord>()>;
  using FQueryTuningRecord =
      ffi::TypedFunction<Optional<TuningRecord>(const IRModule&, const Target&, const String&)>;
  using FQuerySchedule =
      ffi::TypedFunction<Optional<tir::Schedule>(const IRModule&, const Target&, const String&)>;
  using FQueryIRModule =
      ffi::TypedFunction<Optional<IRModule>(const IRModule&, const Target&, const String&)>;
  using FSize = ffi::TypedFunction<int64_t()>;
 
  FHasWorkload f_has_workload;
  FCommitWorkload f_commit_workload;
  FCommitTuningRecord f_commit_tuning_record;
  FGetTopK f_get_top_k;
  FGetAllTuningRecords f_get_all_tuning_records;
  FQueryTuningRecord f_query_tuning_record;
  FQuerySchedule f_query_schedule;
  FQueryIRModule f_query_ir_module;
  FSize f_size;
 
  static void RegisterReflection() {
    // ffi::Functions are all not registered, because the reflection system doesn't take care of
    // them, so it cannot be accessible on the python side. If there is such need from the future,
    // we can then add corresponding accessor methods to help access on python.
    // `f_has_workload` is not registered
    // `f_commit_workload` is not registered
    // `f_commit_tuning_record` is not registered
    // `f_get_top_k` is not registered
    // `f_get_all_tuning_records` is not registered
    // `f_query_tuning_record` is not registered
    // `f_query_schedule` is not registered
    // `f_query_ir_module` is not registered
    // `f_size` is not registered
  }
 
  bool HasWorkload(const IRModule& mod) final {
    ICHECK(f_has_workload != nullptr) << "PyDatabase's HasWorkload method not implemented!";
    return f_has_workload(mod);
  }
 
  Workload CommitWorkload(const IRModule& mod) final {
    ICHECK(f_commit_workload != nullptr) << "PyDatabase's CommitWorkload method not implemented!";
    return f_commit_workload(mod);
  }
 
  void CommitTuningRecord(const TuningRecord& record) final {
    ICHECK(f_commit_tuning_record != nullptr)
        << "PyDatabase's CommitTuningRecord method not implemented!";
    f_commit_tuning_record(record);
  }
 
  Array<TuningRecord> GetTopK(const Workload& workload, int top_k) final {
    ICHECK(f_get_top_k != nullptr) << "PyDatabase's GetTopK method not implemented!";
    return f_get_top_k(workload, top_k);
  }
 
  Array<TuningRecord> GetAllTuningRecords() final {
    ICHECK(f_get_all_tuning_records != nullptr)
        << "PyDatabase's GetAllTuningRecords method not implemented!";
    return f_get_all_tuning_records();
  }
 
  Optional<TuningRecord> QueryTuningRecord(const IRModule& mod, const Target& target,
                                           const String& workload_name) final {
    if (f_query_tuning_record == nullptr) {
      return DatabaseNode::QueryTuningRecord(mod, target, workload_name);
    } else {
      return f_query_tuning_record(mod, target, workload_name);
    }
  }
 
  Optional<tir::Schedule> QuerySchedule(const IRModule& mod, const Target& target,
                                        const String& workload_name) final {
    if (f_query_schedule == nullptr) {
      return DatabaseNode::QuerySchedule(mod, target, workload_name);
    } else {
      return f_query_schedule(mod, target, workload_name);
    }
  }
 
  Optional<IRModule> QueryIRModule(const IRModule& mod, const Target& target,
                                   const String& workload_name) final {
    if (f_query_ir_module == nullptr) {
      return DatabaseNode::QueryIRModule(mod, target, workload_name);
    } else {
      return f_query_ir_module(mod, target, workload_name);
    }
  }
 
  int64_t Size() final {
    ICHECK(f_size != nullptr) << "PyDatabase's Size method not implemented!";
    return f_size();
  }
 
  static constexpr const char* _type_key = "meta_schedule.PyDatabase";
  TVM_DECLARE_FINAL_OBJECT_INFO(PyDatabaseNode, DatabaseNode);
};
 
class Database : public runtime::ObjectRef {
 public:
  TVM_DLL static Database MemoryDatabase(String mod_eq_name = "structural");
  TVM_DLL static Database ScheduleFnDatabase(ffi::TypedFunction<bool(tir::Schedule)> schedule_fn,
                                             String mod_eq_name = "structural");
  TVM_DLL static Database JSONDatabase(String path_workload, String path_tuning_record,
                                       bool allow_missing, String mod_eq_name = "structural");
  TVM_DLL static Database UnionDatabase(Array<Database, void> databases);
  TVM_DLL static Database OrderedUnionDatabase(Array<Database, void> databases);
  TVM_DLL static Database PyDatabase(PyDatabaseNode::FHasWorkload f_has_workload,
                                     PyDatabaseNode::FCommitWorkload f_commit_workload,
                                     PyDatabaseNode::FCommitTuningRecord f_commit_tuning_record,
                                     PyDatabaseNode::FGetTopK f_get_top_k,
                                     PyDatabaseNode::FGetAllTuningRecords f_get_all_tuning_records,
                                     PyDatabaseNode::FQueryTuningRecord f_query_tuning_record,
                                     PyDatabaseNode::FQuerySchedule f_query_schedule,
                                     PyDatabaseNode::FQueryIRModule f_query_ir_module,
                                     PyDatabaseNode::FSize f_size,
                                     String mod_eq_name = "structural");
  static Optional<Database> Current();
  void EnterWithScope();
  void ExitWithScope();
 
  TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(Database, runtime::ObjectRef, DatabaseNode);
};
 
}  // namespace meta_schedule
}  // namespace tvm
 
#endif  // TVM_META_SCHEDULE_DATABASE_H_