def test_block_access_region_detector():
block = func.body.block.body.block
alloc_buffers = func.body.block.alloc_buffers
buffer_var_map = {buf.data: buf for buf in alloc_buffers}
ret = tir.analysis.get_block_access_region(block, buffer_var_map)
tvm.ir.assert_structural_equal(block.reads, ret[0])
tvm.ir.assert_structural_equal(block.writes, ret[1])
D = alloc_buffers[-1]
tvm.ir.assert_structural_equal(
[tvm.tir.BufferRegion(D, [Range(0, 128), Range(0, 128)])], ret[2])
def _get_region(tslice):
region = []
for idx in tslice.indices:
if isinstance(idx, slice):
assert idx.step is None
region.append(Range(idx.start, idx.stop))
else:
if isinstance(idx, tvm.tir.IterVar):
begin = idx.var
else:
begin = idx
region.append(Range.make_by_min_extent(begin, 1))
return region
def non_surjective_inverse(
self, shape: List[Union[Range,
PrimExpr]]) -> Tuple["IndexMap", PrimExpr]:
"""Return the inverse of the map
Can be applied to transformations that introduce padding.
Parameters
----------
shape: List[Union[Range,PrimExpr]]
The region over which the inverse should be determined.
Used for determining the predicate.
Returns
-------
result : Tuple[IndexMap, PrimExpr]
The inverse, and a predicate for which the inverse maps to
a valid index in the input range.
Examples
--------
.. code-block:: python
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)"
"""
shape = [
dim if isinstance(dim, Range) else Range(0, dim) for dim in shape
]
return _ffi_api.IndexMapNonSurjectiveInverse(self, shape)
def inverse(self, shape: List[Union[Range, PrimExpr]]) -> "IndexMap":
"""Return the inverse of the map
Throws an error if the function is not bijective.
Parameters
----------
shape: List[Union[Range,PrimExpr]]
The region over which the inverse should be determined.
Used for validating that the mapping is bijective over
this range.
Returns
-------
inverse : IndexMap
The inverse
"""
shape = [
dim if isinstance(dim, Range) else Range(0, dim) for dim in shape
]
return _ffi_api.IndexMapInverse(self, shape)