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_