search_policy.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_SEARCH_POLICY_H_
#define TVM_AUTO_SCHEDULER_SEARCH_POLICY_H_
 
#include <tvm/auto_scheduler/measure.h>
#include <tvm/auto_scheduler/search_task.h>
#include <tvm/node/node.h>
 
#include <string>
#include <unordered_set>
#include <utility>
#include <vector>
 
namespace tvm {
namespace auto_scheduler {
 
class ProgramMeasurer;
class SearchPolicyNode;
 
class SearchCallbackNode : public Object {
 public:
  virtual void Callback(SearchPolicyNode* policy) = 0;
 
  static constexpr const char* _type_key = "auto_scheduler.SearchCallback";
  TVM_DECLARE_BASE_OBJECT_INFO(SearchCallbackNode, Object);
};
 
class SearchCallback : public ObjectRef {
 public:
  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(SearchCallback, ObjectRef, SearchCallbackNode);
};
 
class PreloadMeasuredStatesNode : public SearchCallbackNode {
 public:
  String filename;
 
  void Callback(SearchPolicyNode* policy) final;
 
  static constexpr const char* _type_key = "auto_scheduler.PreloadMeasuredStates";
  TVM_DECLARE_FINAL_OBJECT_INFO(PreloadMeasuredStatesNode, SearchCallbackNode);
};
 
class PreloadMeasuredStates : public SearchCallback {
 public:
  explicit PreloadMeasuredStates(String filename);
 
  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(PreloadMeasuredStates, SearchCallback,
                                        PreloadMeasuredStatesNode);
};
 
struct SearchPolicyKey {
  static constexpr const char* always_unroll_inner = "auto_scheduler_always_unroll_inner";
  static constexpr const char* no_split_at_inner = "auto_scheduler_no_split_at_inner";
  static constexpr const char* simplify_const_tensor_indices =
      "auto_scheduler_simplify_const_tensor_indices";
};
 
class SearchPolicyNode : public Object {
 public:
  SearchTask search_task;
  int verbose;
 
  void VisitAttrs(AttrVisitor* v) {
    v->Visit("search_task", &search_task);
    v->Visit("verbose", &verbose);
  }
 
  virtual State Search(int num_measure_trials, int early_stopping, int num_measures_per_round,
                       ProgramMeasurer measurer) = 0;
 
  virtual std::pair<Array<MeasureInput>, Array<MeasureResult>> ContinueSearchOneRound(
      int num_measure, ProgramMeasurer measurer) = 0;
 
  void PreloadMeasuredStates(const String& log_file);
 
  void RunCallbacks(const Array<SearchCallback>& callbacks);
 
  static constexpr const char* _type_key = "auto_scheduler.SearchPolicy";
  TVM_DECLARE_BASE_OBJECT_INFO(SearchPolicyNode, Object);
 
 protected:
  std::unordered_set<std::string> measured_states_set_;
  std::vector<State> measured_states_vector_;
  std::vector<float> measured_states_throughputs_;
};
 
class SearchPolicy : public ObjectRef {
 public:
  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(SearchPolicy, ObjectRef, SearchPolicyNode);
};
 
}  // namespace auto_scheduler
}  // namespace tvm
 
#endif  // TVM_AUTO_SCHEDULER_SEARCH_POLICY_H_