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_