tvm_tirx.tirx.IndexMap

tvm_tirx.tirx.IndexMap#

class tvm_tirx.tirx.IndexMap(initial_indices, final_indices, inverse_index_map)#

Bases: Object

Map input index variables to expressions for output indices.

Use from_func() to derive a map from a Python callable, or supply the input variables and output expressions directly.

Parameters:

  • initial_indices (list[Var]) – Bound variables representing the input coordinates, in dimension order.

  • final_indices (list[Expr]) – Output coordinate expressions in terms of the input variables.

  • inverse_index_map (IndexMap or Callable or None) – A supplied inverse, or None to compute the inverse when requested. The caller is responsible for correctness of a supplied inverse.

property initial_indices#
property final_indices#
static from_func(mapping_function, ndim=None, inverse_index_map=None, *, index_dtype='int64')#

Create an index map from a function

Parameters:

  • mapping_function (Callable) – The function to map from source indices to target indices. The function should accept tirx.Var parameters and return a either a tirx.Expr, or a list of tirx.Expr. Returning a tirx.Expr is equivalent to returning a list of length 1 containing that tirx.Expr.

  • ndim (Optional[int]) – The dimensionality of the tensor to which this transformation should be applied. If mapping_function uses variadic argument *args, ndim must be specified. If mapping_function does not use variadic arguments, ndim is optional.

  • inverse_index_map (Union[Callable, Optional[IndexMap]]) – The optional pre-defined 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.

  • index_dtype (str) – The default index dtype to use for input iters in the mapping function.

Returns:

index_map (IndexMap) – Returns an IndexMap representing the mapping_function.

is_equivalent_to(other_map, analyzer=None)#

Return if the index maps are equivalent.

Parameters:

  • other_map (IndexMap) – The IndexMap to which the comparison should be made.

  • analyzer (Optional[tvm_tirx.sym.Analyzer]) – The analyzer to use while comparing the mapped indices. When provided, its accumulated bindings and constraints are reused so that maps that are only equivalent under those bindings can be proven equal.

Returns:

is_equivalent (bool) – True if the two mappings represent the same transformation, otherwise False

Return type:

bool

map_indices(indices, analyzer=None)#

Apply the index map to a set of indices

Parameters:
Returns:

result (List[Expr]) – The mapped indices

Return type:

list[Expr]

map_shape(shape, analyzer=None)#

Apply the index map to a tensor shape

Parameters:

  • shape (List[Expr]) – The tensor shape to be mapped

  • analyzer (Optional[tvm_tirx.sym.Analyzer]) – The analyzer to use while simplifying mapped shape expressions.

Returns:

result (List[Expr]) – The mapped shape

Return type:

list[Expr]

map_tensor(arr_src)#

Apply thie index map to transform the layout of the input Tensor

Parameters:

arr_src (runtime.Tensor) – The Tensor to be transformed

Returns:

arr_dst (runtime.Tensor) – The transformed Tensor

Return type:

Tensor

inverse(shape, analyzer=None)#

Return the inverse of the map

Throws an error if the function is not bijective.

Parameters:

  • shape (List[Union[Range,Expr]]) – The region over which the inverse should be determined. Used for validating that the mapping is bijective over this range.

  • analyzer (Optional[tvm_tirx.sym.Analyzer]) – The analyzer to use while deriving and validating the inverse.

Returns:

inverse (IndexMap) – The inverse

Return type:

IndexMap

non_surjective_inverse(shape, analyzer=None)#

Return the inverse of the map

Can be applied to transformations that introduce padding.

Parameters:

  • shape (List[Union[Range,Expr]]) – The region over which the inverse should be determined. Used for determining the predicate.

  • analyzer (Optional[tvm_tirx.sym.Analyzer]) – The analyzer to use while deriving the inverse and padding predicate.

Returns:

result (Tuple[IndexMap, Expr]) – The inverse, and a predicate for which the inverse maps to a valid index in the input range.

Return type:

tuple[IndexMap, Expr]

Examples

index_map = IndexMap.from_func(lambda i: [i//4, i%4])
inverse_map, predicate = index_map.non_surjective_inverse([14])
assert inverse_map.is_equivalent_to(IndexMap.from_func(lambda j,k: [4*j + k])
print(predicate) # Prints "(axis0==3) && (axis2 >= 2)"
property inverse_index_map#