transform_step.h File Reference

Transformation steps. These steps are used to manipulate LoopState. They are similar to the schedule primitives in te::Stage. More...

#include <dmlc/common.h>
#include <dmlc/json.h>
#include <tvm/node/node.h>
#include <tvm/te/schedule.h>
#include <vector>

Include dependency graph for transform_step.h:

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

Go to the source code of this file.

Classes
class	tvm::auto_scheduler::IteratorNode
	An iterator of a for-loop Similar to tvm::IterVar in include/tvm/tir/expr.h More...
class	tvm::auto_scheduler::Iterator
	Managed reference to IteratorNode. More...
class	tvm::auto_scheduler::StepNode
	The base class of transformation steps. Each step has its corresponding tvm.te schedule primitives. More...
class	tvm::auto_scheduler::Step
	Managed reference to StepNode. More...
class	tvm::auto_scheduler::AnnotationStepNode
	Annotation step that corresponds to vectorize, parallel, unroll and thread binding. (i.e. te::Stage::vectorize, te::Stage::parallel, te::Stage::vectorize, te::Stage::bind) More...
class	tvm::auto_scheduler::AnnotationStep
	Managed reference to AnnotationStepNode. More...
class	tvm::auto_scheduler::FuseStepNode
	Fuse step that corresponds to te::Stage::fuse. More...
class	tvm::auto_scheduler::FuseStep
	Managed reference to FuseStepNode. More...
class	tvm::auto_scheduler::PragmaStepNode
	Pragma step that corresponds to te::Stage::pragma. More...
class	tvm::auto_scheduler::PragmaStep
	Managed reference to PragmaStepNode. More...
class	tvm::auto_scheduler::ReorderStepNode
	Reorder step that corresponds to te::Stage::reorder. More...
class	tvm::auto_scheduler::ReorderStep
	Managed reference to ReorderStepNode. More...
class	tvm::auto_scheduler::SplitStepNode
	Split step that corresponds to te::Stage::split with additional support of multiple-level of factors. More...
class	tvm::auto_scheduler::SplitStep
	Managed reference to SplitStepNode. More...
class	tvm::auto_scheduler::FollowSplitStepNode
	Similar to SplitStepNode, but uses split factors from another step (i.e. Follow another split step) More...
class	tvm::auto_scheduler::FollowSplitStep
	Managed reference to FollowSplitStepNode. More...
class	tvm::auto_scheduler::FollowFusedSplitStepNode
	Similar to FollowSplitStep, but uses split factors from multiple steps. More...
class	tvm::auto_scheduler::FollowFusedSplitStep
	Managed reference to FollowFusedSplitStepNode. More...
class	tvm::auto_scheduler::StorageAlignStepNode
	Storage align step that corresponds to te::Stage::storage_align. More...
class	tvm::auto_scheduler::StorageAlignStep
	Managed reference to StorageAlignStepNode. More...
class	tvm::auto_scheduler::ComputeAtStepNode
	Compute at step that corresponds to te::Stage::compute_at. More...
class	tvm::auto_scheduler::ComputeAtStep
	Managed reference to ComputeAtStepNode. More...
class	tvm::auto_scheduler::ComputeInlineStepNode
	Compute inline step that corresponds to te::Stage::compute_inline. More...
class	tvm::auto_scheduler::ComputeInlineStep
	Managed reference to ComputeInlineStepNode. More...
class	tvm::auto_scheduler::ComputeRootStepNode
	Compute root step that corresponds to te::Stage::compute_root. More...
class	tvm::auto_scheduler::ComputeRootStep
	Managed reference to ComputeRootStepNode. More...
class	tvm::auto_scheduler::CacheReadStepNode
	Cache read step that corresponds to te::Schedule::cache_read. More...
class	tvm::auto_scheduler::CacheReadStep
	Managed reference to CacheReadStepNode. More...
class	tvm::auto_scheduler::CacheWriteStepNode
	Cache write step that corresponds to te::Schedule::cache_write. More...
class	tvm::auto_scheduler::CacheWriteStep
	Managed reference to CacheWriteStepNode. More...
class	tvm::auto_scheduler::RfactorStepNode
	Reduction factor step that corresponds to te::Schedule::rfactor. More...
class	tvm::auto_scheduler::RfactorStep
	Managed reference to RfactorStepNode. More...

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

Typedefs
typedef Map< tvm::te::Stage, Array< tir::IterVar >, ObjectHash, ObjectEqual >	tvm::auto_scheduler::StageToAxesMap

Enumerations
enum	tvm::auto_scheduler::IteratorKind : int { tvm::auto_scheduler::IteratorKind::kSpatial = 0, tvm::auto_scheduler::IteratorKind::kReduction = 1, tvm::auto_scheduler::IteratorKind::kMixed = 2, tvm::auto_scheduler::IteratorKind::kSpecial = 3 }
	The type of an iterator. More...
enum	tvm::auto_scheduler::IteratorAnnotation : int {   tvm::auto_scheduler::IteratorAnnotation::kNone = 0, tvm::auto_scheduler::IteratorAnnotation::kUnroll = 1, tvm::auto_scheduler::IteratorAnnotation::kVectorize = 2, tvm::auto_scheduler::IteratorAnnotation::kParallel = 3,   tvm::auto_scheduler::IteratorAnnotation::kVThread = 4, tvm::auto_scheduler::IteratorAnnotation::kBlockX = 5, tvm::auto_scheduler::IteratorAnnotation::kThreadX = 6, tvm::auto_scheduler::IteratorAnnotation::kBlockY = 7,   tvm::auto_scheduler::IteratorAnnotation::kThreadY = 8, tvm::auto_scheduler::IteratorAnnotation::kBlockZ = 9, tvm::auto_scheduler::IteratorAnnotation::kThreadZ = 10, tvm::auto_scheduler::IteratorAnnotation::kTensorize = 11 }
	The type of an iterator's annotation. More...

Functions
void	tvm::auto_scheduler::UpdateStageToAxesMap (const te::Stage &stage, StageToAxesMap *stage_to_axes)
	Update the current stage IterVar information to StageToAxesMap. More...
Step	tvm::auto_scheduler::StepReadFromRecord (dmlc::JSONReader *reader)
	Read a step record from JSONReader and create the corresponding step. More...
void	tvm::auto_scheduler::StepApplyToState (const Step &step, State *state, const ComputeDAG &dag)
	Apply a general step to a State with runtime dynamic dispatching. More...
void	tvm::auto_scheduler::StepApplyToSchedule (const Step &step, Array< te::Stage > *stages, StageToAxesMap *stage_to_axes, te::Schedule *schedule, const Array< Step > &transform_steps)
	Apply a general step to tvm.schedule with runtime dynamic dispatching. More...
String	tvm::auto_scheduler::StepPrintAsPythonAPI (const Step &step, Array< te::Stage > *stages, StageToAxesMap *stage_to_axes, te::Schedule *schedule, const Array< Step > &transform_steps)
	Print a general step as equivalent python schedule API with runtime dynamic dispatching. More...

Variables
const char *	tvm::auto_scheduler::IteratorAnnotationString []

Detailed Description

Transformation steps. These steps are used to manipulate LoopState. They are similar to the schedule primitives in te::Stage.

Note

How to add a new transform step: Take fuse step for example:

1. Define class FuseStepNode, FuseStep in transform_steps.h, and implement its first construction function FuseStep::FuseStep() in transform_steps.cc.
2. Implement FuseStepNode::ApplyToSchedule() and FuseStepNode::PrintAsPythonAPI().
   • In these two functions you need to lower this step with tvm's te schedule API
3. Implement FuseStepNode::ApplyToState and the state API State::fuse.
   • In these two functions you need to incrementally update all data structures in State with CopyOnWrite style.
4. Add your step to StepApplyToState, StepApplyToSchedule, and StepPrintAsPythonAPI.
5. Log record serialization support:
   • Add FuseStepNode::WriteToRecord which takes a mutable JSONWriter pointer as input and output the record to it.
   • Add another construction function that takes a mutable JSONReader as input, this will get a step record from the reader and create the step.
   • Add the step implementation to StepReadFromRecord.
6. Add its corresponding Python API to loop_state.py with necessary unit tests. The test should at lease cover two parts: the functional test and the record serialization test.