tuning_api.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_RELAX_TUNING_API_H_
#define TVM_RELAX_TUNING_API_H_
#include <tvm/ir/module.h>
#include <tvm/ir/transform.h>
#include <tvm/meta_schedule/database.h>
 
#include <vector>
namespace tvm {
namespace relax {
 
TVM_ALWAYS_INLINE TVMRetValue CallPackedWithArgsInArray(const runtime::PackedFunc f,
                                                        const Array<ObjectRef>& args) {
  size_t num_args = args.size();
  std::vector<TVMValue> values(num_args);
  std::vector<int> codes(num_args);
  runtime::TVMArgsSetter setter(values.data(), codes.data());
  const ObjectRef* ptr = args.template as<ArrayNode>()->begin();
  for (size_t i = 0; i < num_args; ++i) {
    setter(i, *(ptr + i));
  }
 
  TVMRetValue rv;
  f.CallPacked(TVMArgs(values.data(), codes.data(), num_args), &rv);
  return rv;
}
 
class ChoiceNode : public runtime::Object {
 public:
  String transform_func_key;
  String constr_func_key;
  Array<ObjectRef> transform_func_args;
  Array<ObjectRef> constr_func_args;
 
  virtual ~ChoiceNode() = default;
 
  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("transform_func_key", &transform_func_key);
    v->Visit("transform_func_args", &transform_func_args);
    v->Visit("constr_func_key", &constr_func_key);
    v->Visit("constr_func_args", &constr_func_args);
  }
 
  const runtime::PackedFunc GetConstrFunc() {
    const auto* constr_func = tvm::runtime::Registry::Get(constr_func_key);
    ICHECK(constr_func != nullptr) << "constr_func_key is not registered: " << constr_func_key;
    return *constr_func;
  }
 
  const runtime::PackedFunc GetTransformFunc() {
    auto* transform_func = tvm::runtime::Registry::Get(transform_func_key);
    ICHECK(transform_func != nullptr)
        << "transform_func_key is not registered: " << transform_func_key;
    return *transform_func;
  }
 
  bool CheckConstr(const IRModule& mod) {
    Array<ObjectRef> args(constr_func_args);
    args.insert(args.begin(), mod);
    return CallPackedWithArgsInArray(GetConstrFunc(), args);
  }
 
  IRModule ApplyTransformFunc(IRModule mod) {
    // Apply transformation when constraint is satisfied.
    if (CheckConstr(mod)) {
      Array<ObjectRef> args(transform_func_args);
      args.insert(args.begin(), GetRef<IRModule>(mod.CopyOnWrite()));
      return CallPackedWithArgsInArray(GetTransformFunc(), args);
    }
    return mod;
  }
 
  ObjectRef AsJSON() const;
 
  static constexpr const char* _type_key = "relax.tuning_api.Choice";
  TVM_DECLARE_BASE_OBJECT_INFO(ChoiceNode, Object);
};
 
class Choice : public runtime::ObjectRef {
 public:
  TVM_DLL explicit Choice(String transform_func_key, Array<ObjectRef> transform_func_args,
                          String constr_func_key, Array<ObjectRef> constr_func_args);
  TVM_DLL static Choice FromJSON(const ObjectRef& json_obj);
  TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(Choice, ObjectRef, ChoiceNode);
};
 
class KnobNode : public runtime::Object {
 public:
  String name;
  Map<String, Choice> choices;
 
  virtual ~KnobNode() = default;
 
  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("name", &name);
    v->Visit("choices", &choices);
  }
 
  bool IsValidDecision(String decision) { return choices.count(decision) > 0; }
 
  IRModule Apply(IRModule mod, String decision) {
    ICHECK(IsValidDecision(decision)) << "Invalid choice for this knob: " << decision;
    return choices[decision]->ApplyTransformFunc(mod);
  }
 
  ObjectRef AsJSON() const;
 
  static constexpr const char* _type_key = "relax.tuning_api.Knob";
  TVM_DECLARE_BASE_OBJECT_INFO(KnobNode, Object);
};
 
class Knob : public runtime::ObjectRef {
 public:
  TVM_DLL explicit Knob(String name, Map<String, Choice> choices);
  TVM_DLL static Knob FromJSON(const ObjectRef& json_obj);
  TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(Knob, ObjectRef, KnobNode);
};
 
class TraceNode : public runtime::Object {
 public:
  IRModule in_mod;
  mutable IRModule out_mod;
  // TODO(sunggg): can we move knobs and decisions into private?
  Array<Knob> knobs;
  Array<String> decisions;
  mutable double perf = -1;
  mutable int size = 0;
  virtual ~TraceNode() = default;
 
  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("in_mod", &in_mod);
    v->Visit("out_mod", &out_mod);
    v->Visit("knobs", &knobs);
    v->Visit("decisions", &decisions);
    v->Visit("perf", &perf);
    v->Visit("size", &size);
  }
 
  bool Verify() const {
    if (knobs.size() != decisions.size()) return false;
    int n = knobs.size();
    for (int i = 0; i < n; i++) {
      if (!knobs[i]->IsValidDecision(decisions[i])) return false;
    }
    return true;
  }
 
  IRModule Add(Knob knob, String decision) {
    out_mod = knob->Apply(out_mod, decision);
    knobs.push_back(knob);
    decisions.push_back(decision);
    // perf number should be initialized after new decision is applied.
    perf = -1;
    // increment history size.
    size++;
    return out_mod;
  }
 
  ObjectRef AsJSON(bool include_in_mod = true) const;
 
  void SetPerf(double _perf) { perf = _perf; }
  void SetOutMod(IRModule mod_) { out_mod = mod_; }
 
  static constexpr const char* _type_key = "relax.tuning_api.Trace";
  TVM_DECLARE_BASE_OBJECT_INFO(TraceNode, Object);
};
 
class Trace : public runtime::ObjectRef {
 public:
  Trace();
  TVM_DLL explicit Trace(IRModule in_mod, Array<Knob> knobs, Array<String> decisions);
  TVM_DLL static Trace FromJSON(const ObjectRef& json_obj);
  TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(Trace, ObjectRef, TraceNode);
};
 
class TuningRecordNode : public runtime::Object {
 public:
  Trace trace;
  Optional<Array<FloatImm>> run_secs;
 
  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("trace", &trace);
    v->Visit("run_secs", &run_secs);
  }
 
  static constexpr const char* _type_key = "relax.tuning_api.TuningRecord";
  TVM_DECLARE_FINAL_OBJECT_INFO(TuningRecordNode, runtime::Object);
 
  ObjectRef AsJSON(bool include_irmod = false) const;
};
 
class TuningRecord : public runtime::ObjectRef {
 public:
  TVM_DLL explicit TuningRecord(Trace trace, Optional<Array<FloatImm>> run_secs);
  TVM_DLL static TuningRecord FromJSON(const ObjectRef& json_obj);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(TuningRecord, runtime::ObjectRef, TuningRecordNode);
};
 
struct WorkloadEqual {
  bool operator()(const meta_schedule::Workload& a, const meta_schedule::Workload& b) const {
    return a->shash == b->shash && tvm::StructuralEqual()(a->mod, b->mod);
  }
};
 
/* \brief The abstract interface of database. */
class DatabaseNode : public runtime::Object {
 public:
  virtual ~DatabaseNode() = default;
  virtual bool HasWorkload(const IRModule& mod) = 0;
  virtual bool HasMeasurementRecord(const meta_schedule::Workload& workload,
                                    const Target& target) = 0;
  virtual bool HasTuningRecord(const meta_schedule::Workload& workload, const Target& target) = 0;
  virtual meta_schedule::Workload CommitWorkload(const IRModule& mod) = 0;
  virtual void CommitMeasurementRecord(const meta_schedule::Workload& workload,
                                       const Target& target, const Array<FloatImm>& record) = 0;
  virtual void CommitTuningRecord(const meta_schedule::Workload& workload, const Target& target,
                                  const TuningRecord& record) = 0;
  virtual Array<TuningRecord> GetTopK(const meta_schedule::Workload& workload, const Target& target,
                                      int top_k) = 0;
  virtual Array<FloatImm> GetMeasurementRecord(const meta_schedule::Workload& workload,
                                               const Target target) = 0;
 
  static constexpr const char* _type_key = "relax.tuning_api.Database";
  TVM_DECLARE_BASE_OBJECT_INFO(DatabaseNode, runtime::Object);
};
 
class Database : public runtime::ObjectRef {
 public:
  TVM_DLL static Database JSONDatabase(String path_workload, String path_tuning_record,
                                       String path_measurement_record, bool allow_missing);
  TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(Database, runtime::ObjectRef, DatabaseNode);
};
 
}  // namespace relax
}  // namespace tvm
#endif  // TVM_RELAX_TUNING_API_H_