measure.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_AUTO_SCHEDULER_MEASURE_H_
#define TVM_AUTO_SCHEDULER_MEASURE_H_
 
#include <tvm/auto_scheduler/loop_state.h>
#include <tvm/auto_scheduler/search_task.h>
 
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
 
namespace tvm {
namespace auto_scheduler {
 
class SearchPolicy;
class MeasureInput;
class MeasureResult;
 
enum class MeasureErrorNO : int {
  kNoError = 0,
  kInstantiationError = 1,
  kCompileHostError = 2,
  kCompileDeviceError = 3,
  kRuntimeDeviceError = 4,
  kWrongAnswerError = 5,
  kBuildTimeoutError = 6,
  kRunTimeoutError = 7,
  kUnknownError = 8,
};
 
// Inputs and results of one measurement
 
class MeasureInputNode : public Object {
 public:
  SearchTask task;
  State state;
 
  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("task", &task);
    v->Visit("state", &state);
  }
 
  MeasureInput copy() const;
 
  static constexpr const char* _type_key = "auto_scheduler.MeasureInput";
  TVM_DECLARE_FINAL_OBJECT_INFO(MeasureInputNode, Object);
};
 
class MeasureInput : public ObjectRef {
 public:
  MeasureInput(SearchTask task, State state);
 
  TVM_DEFINE_OBJECT_REF_METHODS(MeasureInput, ObjectRef, MeasureInputNode);
};
 
class BuildResultNode : public Object {
 public:
  String filename;
  Array<te::Tensor> args;
  int error_no;
  String error_msg;
  double time_cost;
 
  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("filename", &filename);
    v->Visit("args", &args);
    v->Visit("error_no", &error_no);
    v->Visit("error_msg", &error_msg);
    v->Visit("time_cost", &time_cost);
  }
 
  static constexpr const char* _type_key = "auto_scheduler.BuildResult";
  TVM_DECLARE_FINAL_OBJECT_INFO(BuildResultNode, Object);
};
 
class BuildResult : public ObjectRef {
 public:
  BuildResult(String filename, Array<te::Tensor> args, int error_no, String error_msg,
              double time_cost);
  TVM_DEFINE_OBJECT_REF_METHODS(BuildResult, ObjectRef, BuildResultNode);
};
 
class MeasureResultNode : public Object {
 public:
  Array<PrimExpr> costs;
  int error_no;
  String error_msg;
  double all_cost;
  double timestamp;
 
  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("costs", &costs);
    v->Visit("error_no", &error_no);
    v->Visit("error_msg", &error_msg);
    v->Visit("all_cost", &all_cost);
    v->Visit("timestamp", &timestamp);
  }
 
  MeasureResult copy() const;
 
  static constexpr const char* _type_key = "auto_scheduler.MeasureResult";
  TVM_DECLARE_FINAL_OBJECT_INFO(MeasureResultNode, Object);
};
 
class MeasureResult : public ObjectRef {
 public:
  MeasureResult(Array<PrimExpr> costs, int error_no, String error_msg, double all_cost,
                double timestamp);
 
  TVM_DEFINE_OBJECT_REF_METHODS(MeasureResult, ObjectRef, MeasureResultNode);
};
 
class MeasureCallbackNode : public Object {
 public:
  virtual void Callback(const SearchPolicy& policy, const Array<MeasureInput>& inputs,
                        const Array<MeasureResult>& results) = 0;
  static constexpr const char* _type_key = "auto_scheduler.MeasureCallback";
  TVM_DECLARE_BASE_OBJECT_INFO(MeasureCallbackNode, Object);
};
 
class MeasureCallback : public ObjectRef {
 public:
  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(MeasureCallback, ObjectRef, MeasureCallbackNode);
};
 
class PythonBasedMeasureCallbackNode : public MeasureCallbackNode {
 public:
  PackedFunc callback_func;
 
  void Callback(const SearchPolicy& policy, const Array<MeasureInput>& inputs,
                const Array<MeasureResult>& results) final;
  static constexpr const char* _type_key = "auto_scheduler.PythonBasedMeasureCallback";
  TVM_DECLARE_FINAL_OBJECT_INFO(PythonBasedMeasureCallbackNode, MeasureCallbackNode);
};
 
class PythonBasedMeasureCallback : public MeasureCallback {
 public:
  explicit PythonBasedMeasureCallback(PackedFunc callback_func);
 
  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(PythonBasedMeasureCallback, MeasureCallback,
                                        PythonBasedMeasureCallbackNode);
};
 
// The base class of ProgramBuilders and ProgramRunners.
 
class ProgramBuilderNode : public Object {
 public:
  int n_parallel;
  int timeout;
 
  virtual Array<BuildResult> Build(const Array<MeasureInput>& inputs, int verbose) = 0;
 
  static constexpr const char* _type_key = "auto_scheduler.ProgramBuilder";
  TVM_DECLARE_BASE_OBJECT_INFO(ProgramBuilderNode, Object);
};
 
class ProgramBuilder : public ObjectRef {
 public:
  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(ProgramBuilder, ObjectRef, ProgramBuilderNode);
};
 
class ProgramRunnerNode : public Object {
 public:
  int timeout;
  int number;
  int repeat;
  int min_repeat_ms;
  double cooldown_interval;
  bool enable_cpu_cache_flush;
  int device;
 
  virtual Array<MeasureResult> Run(const Array<MeasureInput>& inputs,
                                   const Array<BuildResult>& build_results, int verbose) = 0;
 
  static constexpr const char* _type_key = "auto_scheduler.ProgramRunner";
  TVM_DECLARE_BASE_OBJECT_INFO(ProgramRunnerNode, Object);
};
 
class ProgramRunner : public ObjectRef {
 public:
  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(ProgramRunner, ObjectRef, ProgramRunnerNode);
};
 
// Implementation of various builders and runners
 
class LocalBuilderNode : public ProgramBuilderNode {
 public:
  String build_func;
 
  Array<BuildResult> Build(const Array<MeasureInput>& inputs, int verbose) final;
 
  static constexpr const char* _type_key = "auto_scheduler.LocalBuilder";
  TVM_DECLARE_FINAL_OBJECT_INFO(LocalBuilderNode, ProgramBuilderNode);
};
 
class LocalBuilder : public ProgramBuilder {
 public:
  LocalBuilder(int timeout, int n_parallel, const String& build_func);
 
  TVM_DEFINE_OBJECT_REF_METHODS(LocalBuilder, ProgramBuilder, LocalBuilderNode);
};
 
class LocalRunnerNode : public ProgramRunnerNode {
 public:
  Array<MeasureResult> Run(const Array<MeasureInput>& inputs,
                           const Array<BuildResult>& build_results, int verbose) final;
 
  static constexpr const char* _type_key = "auto_scheduler.LocalRunner";
  TVM_DECLARE_FINAL_OBJECT_INFO(LocalRunnerNode, ProgramRunnerNode);
};
 
class LocalRunner : public ProgramRunner {
 public:
  LocalRunner(int timeout, int number, int repeat, int min_repeat_ms, double cooldown_interval,
              bool enable_cpu_cache_flush, int device);
 
  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(LocalRunner, ProgramRunner, LocalRunnerNode);
};
 
class RPCRunnerNode : public ProgramRunnerNode {
 public:
  String key;
  String host;
  int port;
  int priority;
  int n_parallel;
 
  Array<MeasureResult> Run(const Array<MeasureInput>& inputs,
                           const Array<BuildResult>& build_results, int verbose) final;
 
  static constexpr const char* _type_key = "auto_scheduler.RPCRunner";
  TVM_DECLARE_FINAL_OBJECT_INFO(RPCRunnerNode, ProgramRunnerNode);
};
 
class RPCRunner : public ProgramRunner {
 public:
  RPCRunner(const String& key, const String& host, int port, int priority, int n_parallel,
            int timeout, int number, int repeat, int min_repeat_ms, double cooldown_interval,
            bool enable_cpu_cache_flush, int device);
 
  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(RPCRunner, ProgramRunner, RPCRunnerNode);
};
 
class ProgramMeasurerNode : public Object {
 public:
  int ct;
  int error_ct;
  std::unordered_map<std::string, double> best_flops;
  std::unordered_map<std::string, State> best_state;
  std::unordered_map<std::string, int> best_ct;
  std::unordered_set<std::string> has_valid;
  ProgramBuilder builder;
  ProgramRunner runner;
  Optional<Array<MeasureCallback>> callbacks;
  int verbose;
  int max_continuous_error;
 
  void Reset();
 
  Array<MeasureResult> Measure(const SearchTask& task, const SearchPolicy& policy,
                               const Array<MeasureInput>& inputs, int batch_size = -1);
  void SilentMeasure(const SearchTask& task, const Array<MeasureInput>& inputs,
                     Array<MeasureResult>* results);
 
  static const int DEFAULT_MAX_CONTINUOUS_ERROR = 150;
 
  static constexpr const char* _type_key = "auto_scheduler.ProgramMeasurer";
  TVM_DECLARE_FINAL_OBJECT_INFO(ProgramMeasurerNode, Object);
};
 
class ProgramMeasurer : public ObjectRef {
 public:
  ProgramMeasurer(ProgramBuilder builder, ProgramRunner runner,
                  Optional<Array<MeasureCallback>> callbacks, int verbose,
                  int max_continuous_error = -1);
 
  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(ProgramMeasurer, ObjectRef, ProgramMeasurerNode);
};
 
}  // namespace auto_scheduler
}  // namespace tvm
 
#endif  // TVM_AUTO_SCHEDULER_MEASURE_H_