loop_state.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_LOOP_STATE_H_
#define TVM_AUTO_SCHEDULER_LOOP_STATE_H_

#include <dmlc/common.h>
#include <tvm/auto_scheduler/transform_step.h>

#include <functional>
#include <unordered_map>
#include <utility>
#include <vector>

namespace tvm {
namespace auto_scheduler {

using namespace tvm::tir;

class ComputeDAG;

enum class StageKind : int {
  kPlaceholder = 0,
  kCompute = 1
};

enum class ComputeAtKind : int {
  kRoot = 0,
  kInlined = 1,
  kIter = 2,
};

struct StageAttributes {
  int auto_unroll_max_step;
  int storage_offset;
};

class StageNode : public Object {
 public:
  te::Operation op;
  Array<Iterator> iters;
  StageKind op_type;
  ComputeAtKind compute_at;
  StageAttributes attrs;

  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("op", &op);
    v->Visit("iters", &iters);
    v->Visit("op_type", &op_type);
    v->Visit("compute_at", &compute_at);
  }

  static constexpr const char* _type_key = "auto_scheduler.Stage";
  TVM_DECLARE_FINAL_OBJECT_INFO(StageNode, Object);
};

class Stage : public ObjectRef {
 public:
  explicit Stage(te::Operation op);
  Stage(te::Operation op, StageKind op_type, const Array<Iterator>& iters, ComputeAtKind compute_at,
        StageAttributes attrs);

  TVM_DEFINE_OBJECT_REF_METHODS(Stage, ObjectRef, StageNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(StageNode);
};

using StageKey = int;
using IterKey = std::pair<int, int>;

class AttachMapNode : public Object {
 public:
  struct IterKeyHash {
    std::size_t operator()(const IterKey& k) const {
      return ::dmlc::HashCombine(std::hash<int>()(k.first), std::hash<int>()(k.second));
    }
  };

  std::unordered_map<StageKey, IterKey> stage_to_attach_iter;
  std::unordered_map<IterKey, std::vector<StageKey>, IterKeyHash> iter_to_attached_stages;

  static constexpr const char* _type_key = "auto_scheduler.AttachMap";
  TVM_DECLARE_FINAL_OBJECT_INFO(AttachMapNode, Object);
};

class AttachMap : public ObjectRef {
 public:
  void SetComputeAtIter(int stage_id, int target_stage_id, int target_iter_id);

  void DeleteStage(int stage_id);

  void UpdateIters(const std::vector<IterKey>& original_iters,
                   const std::vector<IterKey>& new_iters);

  AttachMap ApplyStageIdOffset(int start_id, int offset = 1) const;

  TVM_DEFINE_OBJECT_REF_METHODS(AttachMap, ObjectRef, AttachMapNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(AttachMapNode);

 private:
  static void DeleteStageEntry(AttachMapNode* pnode, int stage_id);
};

class StateNode : public Object {
 public:
  Array<Stage> stages;
  Array<Step> transform_steps;
  AttachMap attach_map;
  Optional<ObjectRef> current_compute_dag;
  bool concrete;

  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("stages", &stages);
    v->Visit("transform_steps", &transform_steps);
    v->Visit("concrete", &concrete);
  }

  static constexpr const char* _type_key = "auto_scheduler.State";
  TVM_DECLARE_FINAL_OBJECT_INFO(StateNode, Object);
};

class State : public ObjectRef {
 public:
  explicit State(const Array<te::Operation>& ops);

  String ToStr(bool delete_trivial_loop = true) const;

  /********** Step APIs working on a single stage **********/
  TVM_DLL Iterator bind(int stage_id, const Iterator& it, IteratorAnnotation thread_type);
  TVM_DLL Iterator parallel(int stage_id, const Iterator& it);
  TVM_DLL Iterator unroll(int stage_id, const Iterator& it, int max_unroll = -1);
  TVM_DLL Iterator vectorize(int stage_id, const Iterator& it);
  TVM_DLL Iterator fuse(int stage_id, const Array<Iterator>& iters);
  TVM_DLL void pragma(int stage_id, const Iterator& it, const String& pragma_type);
  TVM_DLL void reorder(int stage_id, const Array<Iterator>& order);
  TVM_DLL Array<Iterator> split(int stage_id, const Iterator& it,
                                const Array<Optional<Integer>>& lengths,
                                bool inner_to_outer = true);
  TVM_DLL Array<Iterator> follow_split(int stage_id, const Iterator& it, int src_step_id,
                                       int n_split);
  TVM_DLL Array<Iterator> follow_fused_split(int stage_id, const Iterator& it,
                                             const Array<Integer>& src_step_ids, int level,
                                             bool factor_or_nparts);
  TVM_DLL void storage_align(int stage_id, const Iterator& it, int factor, int offset);

  /********** Step APIs working on multiple stages **********/
  TVM_DLL void compute_at(int stage_id, int target_stage_id, const Iterator& target_iter);
  TVM_DLL void compute_inline(int stage_id);
  TVM_DLL void compute_root(int stage_id);

  /********** Step APIs adding new stages **********/
  TVM_DLL int cache_read(int stage_id, const String& scope_name,
                         const Array<Integer>& reader_stage_ids, const ComputeDAG& dag);
  TVM_DLL int cache_write(int stage_id, const String& scope_name, const ComputeDAG& dag);
  TVM_DLL int rfactor(int stage_id, const Iterator& it, int factor_iter_id, const ComputeDAG& dag);

  TVM_DEFINE_OBJECT_REF_METHODS(State, ObjectRef, StateNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(StateNode);
};

}  // namespace auto_scheduler
}  // namespace tvm

// Hash and equal function for State
namespace std {

template <>
struct equal_to<::tvm::auto_scheduler::State> {
  bool operator()(const ::tvm::auto_scheduler::State& lhs,
                  const ::tvm::auto_scheduler::State& rhs) const {
    return lhs.ToStr() == rhs.ToStr();
  }
};

template <>
struct hash<::tvm::auto_scheduler::State> {
  std::size_t operator()(const ::tvm::auto_scheduler::State& state) const {
    return tvm::runtime::ObjectHash()(state.ToStr());
  }
};

}  // namespace std

#endif  // TVM_AUTO_SCHEDULER_LOOP_STATE_H_