Function test_slice

tests/python/tirx/test_layout.py:1548–1655

()

def test_slice():
    def verify_slice(layout, shape, region, sliced):
        r_shape = [r[1] - r[0] for r in region]
        r_size = functools.reduce(operator.mul, [r[1] - r[0] for r in region])

        def get_region_coord(u):
            coord = []
            for r in reversed(region):
                coord.append(u % (r[1] - r[0]))
                u //= r[1] - r[0]
            return coord[::-1]

        def get_shape_coord(r_coord, region):
            return [region[i][0] + r_coord[i] for i in range(len(region))]

        analyzer = Analyzer()

        for u in range(r_size):
            r_coord = get_region_coord(u)
            s_coord = get_shape_coord(r_coord, region)
            a = layout.apply(*s_coord, shape=shape)["m"]
            b = sliced.apply(*r_coord, shape=r_shape)["m"]
            assert analyzer.simplify(a == b)

    def case1():
        layout = TileLayout(S[(8, 8) : (8, 1)])
        shape = [64]
        region = [(5, 8)]
        sliced = layout.slice(shape, region).canonicalize()
        assert sliced is not None
        verify_slice(layout, shape, region, sliced)

        region = [tvm.ir.Range(5, 8)]
        sliced_2 = layout.slice(shape, region).canonicalize()
        assert sliced_2 is not None
        assert_structural_equal(sliced, sliced_2)

    case1()

    def case2():
        # Choose begin and extent to satisfy midpoint condition
        layout = TileLayout(S[(4, 4, 4, 4) : (64, 4, 16, 1)])
        shape = [16, 16]
        region = [(2, 3), (6, 10)]
        sliced = layout.slice(shape, region).canonicalize()
        assert sliced is not None
        verify_slice(layout, shape, region, sliced)

    case2()

    def case3():
        layout = TileLayout(S[(2, 8, 3, 8) : (192, 8, 64, 1)])
        shape = [16, 24]
        region = [(2, 6), (4, 12)]
        sliced = layout.slice(shape, region).canonicalize()
        assert sliced is not None
        verify_slice(layout, shape, region, sliced)