tvm::tir::IndexMapNode Class Reference

Defines a mapping between two representations of indices into a buffer. More...

#include <index_map.h>

Inheritance diagram for tvm::tir::IndexMapNode:

Collaboration diagram for tvm::tir::IndexMapNode:

Public Member Functions
	IndexMapNode ()
	Default constructor. More...
Array< PrimExpr >	MapIndices (const Array< PrimExpr > &indices, arith::Analyzer *analyzer) const
	Map indices to the output space. More...
Array< Range >	MapRanges (const Array< Range > &ranges, arith::Analyzer *analyzer) const
	Map a memory range to the output space. More...
Array< PrimExpr >	MapShape (const Array< PrimExpr > &shape, arith::Analyzer *analyzer) const
	Map a buffer shape to the output space. More...
runtime::NDArray	MapNDArray (runtime::NDArray arr_src) const
String	ToPythonString (const std::function< Optional< String >(const Var &var)> &f_name_map=nullptr) const
	Convert to string representation in Python. More...
	TVM_DECLARE_FINAL_OBJECT_INFO (IndexMapNode, Object)

Static Public Member Functions
static void	RegisterReflection ()

Public Attributes
Array< Var >	initial_indices
	Variables representing the indices prior to remapping. More...
Array< PrimExpr >	final_indices
	Expressions defining the indices after remapping. More...
Optional< ObjectRef >	inverse_index_map
	The inverse index map. More...

Static Public Attributes
static constexpr TVMFFISEqHashKind	_type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode
static constexpr const char *	_type_key = "tir.IndexMap"

Detailed Description

Defines a mapping between two representations of indices into a buffer.

This is primarily used for layout transformations of Buffer objects.

Constructor & Destructor Documentation

IndexMapNode()

tvm::tir::IndexMapNode::IndexMapNode ( )

inline

Default constructor.

Defines the mapping as an identity function, with initial_indices equal to the final indices.

Member Function Documentation

MapIndices()

Array<PrimExpr> tvm::tir::IndexMapNode::MapIndices	(	const Array< PrimExpr > &	indices,
		arith::Analyzer *	analyzer
	)		const

Map indices to the output space.

Parameters

indices	The indices in the input space. Should contain one value for each variable in initial_indices.
analyzer	An optional analyzer to be used to simplify the resulting expressions. If null, will use a fresh analyzer.

Returns

The indices in the output space. Contains one value for each expression in final_indices.

MapNDArray()

runtime::NDArray tvm::tir::IndexMapNode::MapNDArray

(

runtime::NDArray

arr_src

)

const

MapRanges()

Array<Range> tvm::tir::IndexMapNode::MapRanges	(	const Array< Range > &	ranges,
		arith::Analyzer *	analyzer
	)		const

Map a memory range to the output space.

If contiguous memory locations in the input space are not necessarily contiguous in the output space (e.g. lambda i: [8*(i%8) + (i//8)]), then this will return the smallest range such that all valid indices are contained within the given range.

Parameters

ranges	The ranges in the input space. Should contain one value for each variable in initial_indices.
analyzer	An optional analyzer to be used to simplify the resulting expressions. If null, will use a fresh analyzer.

Returns

The ranges in the output space. Contains one value for each expression in final_indices.

MapShape()

Array<PrimExpr> tvm::tir::IndexMapNode::MapShape	(	const Array< PrimExpr > &	shape,
		arith::Analyzer *	analyzer
	)		const

Map a buffer shape to the output space.

Parameters

shape	The buffer shape in the input space. Should contain one value for each variable in initial_indices.
analyzer	An optional analyzer to be used to simplify the resulting expressions. If null, will use a fresh analyzer.

Returns

The buffer shape in the output space. Contains one value for each expression in final_indices.

RegisterReflection()

static void tvm::tir::IndexMapNode::RegisterReflection ( )

inlinestatic

ToPythonString()

String tvm::tir::IndexMapNode::ToPythonString

(

const std::function< Optional< String >(const Var &var)> &

f_name_map = nullptr

)

const

Convert to string representation in Python.

Parameters

f_name_map

Optional function to specify the stringified name of the variables.

Returns

The stringified lambda expression in Python.

TVM_DECLARE_FINAL_OBJECT_INFO()

tvm::tir::IndexMapNode::TVM_DECLARE_FINAL_OBJECT_INFO	(	IndexMapNode	,
		Object
	)

Member Data Documentation

_type_key

constexpr const char* tvm::tir::IndexMapNode::_type_key = "tir.IndexMap"

staticconstexpr

_type_s_eq_hash_kind

constexpr TVMFFISEqHashKind tvm::tir::IndexMapNode::_type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode

staticconstexpr

final_indices

Array<PrimExpr> tvm::tir::IndexMapNode::final_indices

Expressions defining the indices after remapping.

These expressions should only be in terms of the initial_indices, and must be expressible as an IterSumExpr. The mapping from initial_indices to final_indices must be injective.

If final_indices is empty, then initial_indices should also be empty, and the map is an identity function.

initial_indices

Array<Var> tvm::tir::IndexMapNode::initial_indices

Variables representing the indices prior to remapping.

If initial_indices is empty, then final_indices should also be empty, and no mapping is applied.

inverse_index_map

Optional<ObjectRef> tvm::tir::IndexMapNode::inverse_index_map

The inverse index map.

When this is defined, IndexMap::Inverse will return the pre-defined inverse index map. Otherwise, the inverse index map will be computed on the fly. It is the user's responsibility to ensure the correctness of the pre-defined inverse index map.

Note

ObjectRef is used here instead of IndexMap to avoid circular reference.

---

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

• include/tvm/tir/index_map.h