Stage, contains scheduling for a stage of computation. More...

#include <schedule.h>

Inheritance diagram for tvm::te::Stage:

Collaboration diagram for tvm::te::Stage:

Public Types
using	ContainerType = StageNode

Public Types inherited from tvm::runtime::ObjectRef
using	ContainerType = Object
	type indicate the container type. More...

Public Member Functions
	Stage ()

	Stage (ObjectPtr< Object > n)

	Stage (Operation op, const ScheduleNode *sch)
	create a new schedule for op. More...

const StageNode *	operator-> () const
	access the internal node container More...

StageNode *	operator-> ()
	access the internal node container More...

Stage &	set_scope (std::string scope)
	set the memory scope of the stage More...

Stage &	compute_at (Stage parent, IterVar scope)
	specify the schedule to be computed at the parent schedule's scope. More...

Stage &	compute_inline ()
	Compute the function inline. More...

Stage &	compute_root ()
	Compute the function at group root. More...

Stage &	bind (IterVar ivar, IterVar thread_ivar)
	Bind the IterVar to thread index. More...

Stage &	set_store_predicate (PrimExpr predicate)
	Set the predicate to determine whether a store to the array should be performed. Use this when there are multiple threads performing the same store and we only need one of them to do the store. More...

Stage &	env_threads (Array< IterVar > threads)
	Specify environment threads that launched around the group's scope. This can only be used in group stage. More...

Stage &	split (IterVar parent, PrimExpr factor, IterVar p_outer, IterVar p_inner, bool disable_predication=false)
	Split the parent by factor, generate. More...

Stage &	split_by_nparts (IterVar parent, PrimExpr nparts, IterVar p_outer, IterVar p_inner, bool disable_predication=false)
	Split the iteration with given number of parts. More...

Stage &	fuse (IterVar outer, IterVar inner, IterVar *p_target)
	Fuse the inner outer domain to the target. More...

Stage &	fuse (const Array< IterVar > &axes, IterVar *p_target)
	Fuse all the axes together into a single axis. More...

Stage &	reorder (const Array< IterVar > &order)
	Reorder the iteration. More...

Stage &	tile (IterVar x_parent, IterVar y_parent, PrimExpr x_factor, PrimExpr y_factor, IterVar p_x_outer, IterVar p_y_outer, IterVar p_x_inner, IterVar p_y_inner)
	Perform tiling on two dimensions The final loop order from outmost to inner most are [x_outer, y_outer, x_inner, y_inner]. More...

Stage &	vectorize (IterVar var)
	Vectorize iteration. More...

Stage &	tensorize (IterVar var, TensorIntrin f)
	Replace computation of the current stage by tensor intrinsic f. More...

Stage &	unroll (IterVar var)
	Unroll iteration. More...

Stage &	parallel (IterVar var)
	Parallelize iteration. More...

Stage &	pragma (IterVar var, const std::string &pragma_type, const PrimExpr &pragma_value=PrimExpr())
	Annotate the iteration with pragma. More...

Stage &	prefetch (const Tensor &domain, IterVar var, PrimExpr offset)
	Fetch data in advance. More...

Stage &	storage_align (IterVar axis, int factor, int offset)
	Set alignment requirement for specific dimension. More...

Stage &	double_buffer ()
	Compute current stage with double buffering. More...

Stage &	rolling_buffer ()
	Compute current stage with rolling buffering. More...

Stage &	transform_layout (const Array< Var > &initial_indices, const Array< PrimExpr > &final_indices, Array< IterVar > *out_iter_vars=nullptr)
	Defines a layout transformation to be applied to the buffer. More...

Stage &	set_axis_separators (const Array< IntImm > &axis_separators)
	Defines separators between groups of axes. More...

bool	is_scheduled () const
	whether the stage has been scheduled. More...

Stage	GetAttachSpec () const
	Get attachment spec of current stage. If the stage compute at Group root, this function will traverse the group function to get the final spec from the group. More...

Public Member Functions inherited from tvm::runtime::ObjectRef
	ObjectRef ()=default
	default constructor More...

	ObjectRef (ObjectPtr< Object > data)
	Constructor from existing object ptr. More...

bool	same_as (const ObjectRef &other) const
	Comparator. More...

bool	operator== (const ObjectRef &other) const
	Comparator. More...

bool	operator!= (const ObjectRef &other) const
	Comparator. More...

bool	operator< (const ObjectRef &other) const
	Comparator. More...

bool	defined () const

const Object *	get () const

const Object *	operator-> () const

bool	unique () const

int	use_count () const

template<typename ObjectType , typename = std::enable_if_t<std::is_base_of_v<Object, ObjectType>>>
const ObjectType *	as () const
	Try to downcast the internal Object to a raw pointer of a corresponding type. More...

template<typename ObjectRefType , typename = std::enable_if_t<std::is_base_of_v<ObjectRef, ObjectRefType>>>
Optional< ObjectRefType >	as () const
	Try to downcast the ObjectRef to a Optional<T> of the requested type. More...

Additional Inherited Members
Static Public Attributes inherited from tvm::runtime::ObjectRef
static constexpr bool	_type_is_nullable = true

Protected Member Functions inherited from tvm::runtime::ObjectRef
Object *	get_mutable () const

Static Protected Member Functions inherited from tvm::runtime::ObjectRef
template<typename T >
static T	DowncastNoCheck (ObjectRef ref)
	Internal helper function downcast a ref without check. More...

static void	FFIClearAfterMove (ObjectRef *ref)
	Clear the object ref data field without DecRef after we successfully moved the field. More...

template<typename ObjectType >
static ObjectPtr< ObjectType >	GetDataPtr (const ObjectRef &ref)
	Internal helper function get data_ as ObjectPtr of ObjectType. More...

Protected Attributes inherited from tvm::runtime::ObjectRef
ObjectPtr< Object >	data_
	Internal pointer that backs the reference. More...

Detailed Description

Stage, contains scheduling for a stage of computation.

Member Typedef Documentation

◆ ContainerType

using tvm::te::Stage::ContainerType = StageNode

Constructor & Destructor Documentation

◆ Stage() [1/3]

tvm::te::Stage::Stage ( )

inline

◆ Stage() [2/3]

tvm::te::Stage::Stage ( ObjectPtr< Object > n )

inlineexplicit

◆ Stage() [3/3]

tvm::te::Stage::Stage	(	Operation	op,
		const ScheduleNode *	sch
	)

explicit

create a new schedule for op.

Parameters

op	The operator in the schedule
sch	The schedule which current stage belongs to

Member Function Documentation

◆ bind()

Stage& tvm::te::Stage::bind	(	IterVar	ivar,
		IterVar	thread_ivar
	)

Bind the IterVar to thread index.

Parameters

ivar	The IterVar to be bound.
thread_ivar	The thread axis to be bound.

Returns: reference to self.

◆ compute_at()

Stage& tvm::te::Stage::compute_at	(	Stage	parent,
		IterVar	scope
	)

specify the schedule to be computed at the parent schedule's scope.

Parameters

parent	The parent schedule.
scope	The iteration point to carry the schedule.

Returns: reference to self.

◆ compute_inline()

Stage& tvm::te::Stage::compute_inline ( )

Compute the function inline.

Returns: reference to self.

◆ compute_root()

Stage& tvm::te::Stage::compute_root ( )

Compute the function at group root.

Returns: reference to self.

◆ double_buffer()

Stage& tvm::te::Stage::double_buffer ( )

Compute current stage with double buffering.

Returns: reference to self.

◆ env_threads()

Stage& tvm::te::Stage::env_threads ( Array< IterVar > threads )

Specify environment threads that launched around the group's scope. This can only be used in group stage.

Parameters

threads The threads to be launched around the scope.

Note: Each thread can only appear in one env_threads. This is a beta feature.

Returns: reference to self.

◆ fuse() [1/2]

Stage& tvm::te::Stage::fuse	(	const Array< IterVar > &	axes,
		IterVar *	p_target
	)

Fuse all the axes together into a single axis.

Parameters

axes	All the axes to be fused.
p_target	The result target domain.

Note: axes can be an empty array, in that case, a singleton IterVar is created and inserted to the outermost loop. The fuse of empty array is used to support zero-dimension tensors.

Returns: reference to self.

◆ fuse() [2/2]

Stage& tvm::te::Stage::fuse	(	IterVar	outer,
		IterVar	inner,
		IterVar *	p_target
	)

Fuse the inner outer domain to the target.

Parameters

outer	The outer domain to be fused.
inner	The inner domain to be fused
p_target	The result target domain.

Returns: reference to self.

◆ GetAttachSpec()

Stage tvm::te::Stage::GetAttachSpec ( ) const

Get attachment spec of current stage. If the stage compute at Group root, this function will traverse the group function to get the final spec from the group.

Returns: A stage representing the attach spec of the group.

◆ is_scheduled()

bool tvm::te::Stage::is_scheduled ( ) const

whether the stage has been scheduled.

Returns: whether the stage has been scheduled.

◆ operator->() [1/2]

StageNode * tvm::te::Stage::operator-> ( )

inline

access the internal node container

Returns: the pointer to the internal node container

◆ operator->() [2/2]

const StageNode * tvm::te::Stage::operator-> ( ) const

inline

access the internal node container

Returns: the pointer to the internal node container

◆ parallel()

Stage& tvm::te::Stage::parallel ( IterVar var )

Parallelize iteration.

Parameters

var	The axis to be parallelized.

Returns: reference to self.

◆ pragma()

Stage& tvm::te::Stage::pragma	(	IterVar	var,
		const std::string &	pragma_type,
		const PrimExpr &	pragma_value = `PrimExpr()`
	)

Annotate the iteration with pragma.

Parameters

var	The axis to be parallelized.
pragma_type	The pragma type.
pragma_value	The pragma value

Returns: reference to self.

◆ prefetch()

Stage& tvm::te::Stage::prefetch	(	const Tensor &	domain,
		IterVar	var,
		PrimExpr	offset
	)

Fetch data in advance.

Parameters

domain	the tensor to be prefetched
var	the iteration point at which to apply prefetching
offset	the number of iterations be to fetched in advance

Returns: reference to self

◆ reorder()

Stage& tvm::te::Stage::reorder ( const Array< IterVar > & order )

Reorder the iteration.

Parameters

order The order of iteration variable.

Returns: reference to self.

◆ rolling_buffer()

Stage& tvm::te::Stage::rolling_buffer ( )

Compute current stage with rolling buffering.

Returns: reference to self.

◆ set_axis_separators()

Stage& tvm::te::Stage::set_axis_separators ( const Array< IntImm > & axis_separators )

Defines separators between groups of axes.

Used to define BufferNode::axis_separators, which has additional details.

Parameters

axis_separators A list of axis separators.

◆ set_scope()

Stage& tvm::te::Stage::set_scope ( std::string scope )

set the memory scope of the stage

Parameters

scope The memory scope.

◆ set_store_predicate()

Stage& tvm::te::Stage::set_store_predicate ( PrimExpr predicate )

Set the predicate to determine whether a store to the array should be performed. Use this when there are multiple threads performing the same store and we only need one of them to do the store.

Note: This is a dangerous scheduling primitive that can change behavior of program. Only do when we are certain that thare are duplicated stores.

Parameters

predicate The condition to be checked.

Returns: reference to self.

◆ split()

Stage& tvm::te::Stage::split	(	IterVar	parent,
		PrimExpr	factor,
		IterVar *	p_outer,
		IterVar *	p_inner,
		bool	disable_predication = `false`
	)

Split the parent by factor, generate.

Parameters

parent	The parent iteration domain.
factor	The split factor of the loop.
p_outer	The result outer domain
p_inner	The result inner domain.
disable_predication	If enabled, don't create a predicate for guarding the loop. This can be useful when splitting with scalable factors that the schedule writer knows are divisible by the loop bound. Warning: enabling this feature may result in incorrect code generation if not used carefully.

Returns: reference to self.

◆ split_by_nparts()

Stage& tvm::te::Stage::split_by_nparts	(	IterVar	parent,
		PrimExpr	nparts,
		IterVar *	p_outer,
		IterVar *	p_inner,
		bool	disable_predication = `false`
	)

Split the iteration with given number of parts.

Parameters

parent	The parent domain.
nparts	The number of parts in the outer domain.
p_outer	The result outer domain.
p_inner	The result inner domain.
disable_predication	If enabled, don't create a predicate for guarding the loop. This can be useful when splitting with scalable factors that the schedule writer knows are divisible by the loop bound. Warning: enabling this feature may result in incorrect code generation if not used carefully.

Returns: reference to self.

◆ storage_align()

Stage& tvm::te::Stage::storage_align	(	IterVar	axis,
		int	factor,
		int	offset
	)

Set alignment requirement for specific dimension.

Such that stride[axis] == k * factor + offset for some k.

Parameters

axis	The dimension to be specified for alignment.
factor	The factor multiple of alignment
offset	The required offset factor.

Returns: reference to self

◆ tensorize()

Stage& tvm::te::Stage::tensorize	(	IterVar	var,
		TensorIntrin	f
	)

Replace computation of the current stage by tensor intrinsic f.

Parameters

var	The axis marks beginning of tensorization. Every operations inside the axis(include axis itself is tensorized).
f	The Tensor compute intrinsics.

Returns: reference to self.

◆ tile()

Stage& tvm::te::Stage::tile	(	IterVar	x_parent,
		IterVar	y_parent,
		PrimExpr	x_factor,
		PrimExpr	y_factor,
		IterVar *	p_x_outer,
		IterVar *	p_y_outer,
		IterVar *	p_x_inner,
		IterVar *	p_y_inner
	)

Perform tiling on two dimensions The final loop order from outmost to inner most are [x_outer, y_outer, x_inner, y_inner].

Parameters

x_parent	The original x dimension
y_parent	The original y dimension
x_factor	The stride factor on x axis
y_factor	The stride factor on y axis
p_x_outer	Outer axis of x dimension
p_y_outer	Outer axis of y dimension
p_x_inner	Inner axis of x dimension
p_y_inner	Inner axis of y dimension

Returns: reference to self.

◆ transform_layout()

Stage& tvm::te::Stage::transform_layout	(	const Array< Var > &	initial_indices,
		const Array< PrimExpr > &	final_indices,
		Array< IterVar > *	out_iter_vars = `nullptr`
	)

Defines a layout transformation to be applied to the buffer.

The map from initial_index to final_index must be an invertible affine transformation.

Parameters

initial_indices	An array of variables to represent a value's location in the tensor, using the pre-transformation layout. These variables are used as binding occurrences to represent the initial indices when applying the initial->final mapping, and should not occur elsewhere in the Schedule. (i.e. Pass in newly constructed variables, not the initial IterVar::var)
final_indices	An array of expressions, giving the value's location in the tensor, using the post-transformation layout. Expressions should be in terms of the variables given in initial_indices.
out_iter_vars	An optional output location for the updated loop iteration variables.

Returns: reference to self

◆ unroll()

Stage& tvm::te::Stage::unroll ( IterVar var )

Unroll iteration.

Parameters

var	The axis to be unrolled.

Returns: reference to self.

◆ vectorize()

Stage& tvm::te::Stage::vectorize ( IterVar var )

Vectorize iteration.

Parameters

var	The axis to be vectorized.

Returns: reference to self.

The documentation for this class was generated from the following file:

include/tvm/te/schedule.h

Public Types

Public Member Functions

Additional Inherited Members

Detailed Description

Member Typedef Documentation

◆ ContainerType

Constructor & Destructor Documentation

◆ Stage() [1/3]

◆ Stage() [2/3]

◆ Stage() [3/3]

Member Function Documentation

◆ bind()

◆ compute_at()

◆ compute_inline()

◆ compute_root()

◆ double_buffer()

◆ env_threads()

◆ fuse() [1/2]

◆ fuse() [2/2]

◆ GetAttachSpec()

◆ is_scheduled()

◆ operator->() [1/2]

◆ operator->() [2/2]

◆ parallel()

◆ pragma()

◆ prefetch()

◆ reorder()

◆ rolling_buffer()

◆ set_axis_separators()

◆ set_scope()

◆ set_store_predicate()

◆ split()

◆ split_by_nparts()

◆ storage_align()

◆ tensorize()

◆ tile()

◆ transform_layout()

◆ unroll()

◆ vectorize()