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)