loop_state.h File Reference

The definition of the "state" in the search. More...

#include <dmlc/common.h>
#include <tvm/auto_scheduler/transform_step.h>
#include <functional>
#include <unordered_map>
#include <utility>
#include <vector>

Include dependency graph for loop_state.h:

This graph shows which files directly or indirectly include this file:

Go to the source code of this file.

Classes
struct	tvm::auto_scheduler::StageAttributes
	Stage-level attributes. More...
class	tvm::auto_scheduler::StageNode
	A op stage in the compute declaration. Similar to te::Stage in include/tvm/te/schedule.h. More...
class	tvm::auto_scheduler::Stage
	Managed reference to StageNode. More...
class	tvm::auto_scheduler::AttachMapNode
	stores the compute_at relation between stages This stores a bi-directional mapping from stages and iter: More...
struct	tvm::auto_scheduler::AttachMapNode::IterKeyHash
class	tvm::auto_scheduler::AttachMap
	Managed reference to AttachMapNode. More...
class	tvm::auto_scheduler::StateNode
	A state in the search process. It consists of the current loop structure and a list of transformation steps used to construct it. Each State corresponds to a specific schedule for its ComputeDAG. More...
class	tvm::auto_scheduler::State
	Managed reference to StateNode. More...

Namespaces
	tvm
	runtime implementation for LibTorch/TorchScript.
	tvm::auto_scheduler

Typedefs
using	tvm::auto_scheduler::StageKey = int
	Use stage_id to represent a stage. More...
using	tvm::auto_scheduler::IterKey = std::pair< int, int >
	Use stage_id and iter_id to represent a iterator. More...

Enumerations
enum class	tvm::auto_scheduler::StageKind : int { tvm::auto_scheduler::kPlaceholder = 0 , tvm::auto_scheduler::kCompute = 1 }
	The type of a stage. More...
enum class	tvm::auto_scheduler::ComputeAtKind : int { tvm::auto_scheduler::kRoot = 0 , tvm::auto_scheduler::kInlined = 1 , tvm::auto_scheduler::kIter = 2 }
	The type of compute location. More...

Detailed Description

The definition of the "state" in the search.

Each LoopState corresponds to a schedule for its ComputeDAG. A LoopState consists of: 1. a current loop structure; 2. a list of transformation steps used to construct the loop structure. The loop structure keeps a preview of how the schedule will finally look like after lowering the current state (e.g. number of iterators, the extent of each iterator, the compute_at locations ...). During the schedule search process, the loop structure can provide search policy with necessary information on how to manipulate the current state. The transform history is a sequence of TransformStep which will finally be mapped to TVM schedule primitives. The steps are also used for the serialization of a state.

The LoopState can be seen as a lightweight loop structure IR specifically for schedule search. We don't use the existing TVM IR but to extend a new structure on it is because:

1. We want fast incremental change to the loop structures. The search policy needs to get the immediate loop structures update rather than after TVM lowering;
2. We want serializable transform history for replay, backtracking, and mutation;
3. We may create some macro schedule primitives that represent the combination of several TVM schedule primitives.

When the search is finished, we will lower the state to TVM IR with TVM's schedule primitives. Since we share a lot of common objects during search, the transformation is implemented in copy on write style. All objects are immutable, which is similar to TVM IR.