# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # ruff: noqa: E741 import pytest import tvm.testing from tvm import relax, tirx from tvm.script import tirx as T def apply_transformations(func, suggested_transfoms, print_transformation=False): sch = tvm.s_tir.Schedule(func) for block, per_block_transformations in suggested_transfoms.items(): blockrv = sch.get_sblock(block.name_hint) for obj, index_map in per_block_transformations.items(): if isinstance(obj, tirx.SBlock): block_name = obj.name_hint if print_transformation: print("Block transformation: ", block_name, " :: ", index_map) sch.transform_block_layout(block_name, index_map) else: assert isinstance(obj, tirx.Buffer) buffer = obj if print_transformation: print("Buffer transformation: ", buffer, " :: ", index_map) sch.transform_layout(blockrv, buffer, index_map) return sch.mod["main"] def test_nested_blocks(): @T.prim_func(private=True, s_tir=True) def nested_block( arg: T.Buffer((32, 64, 224, 224), "float32"), relu: T.Buffer((32, 64, 224, 224), "float32"), ): for i, j in T.grid(32, 64): with T.sblock("outer"): v_i, v_j = T.axis.remap("SS", [i, j]) T.reads(arg[v_i, v_j, 0:224, 0:224]) T.writes(relu[v_i, v_j, 0:224, 0:224]) for k, l in T.grid(224, 224): with T.sblock("inner"): v_k, v_l = T.axis.remap("SS", [k, l]) T.reads(arg[v_i, v_j, v_k, v_l]) T.writes(relu[v_i, v_j, v_k, v_l]) relu[v_i, v_j, v_k, v_l] = T.max(arg[v_i, v_j, v_k, v_l], T.float32(0)) suggested_transforms = relax.analysis.suggest_layout_transforms( func=nested_block, write_buffer_transforms=[lambda n, c, h, w: (n, h, w, c)] ) # no suggestions for nested block. assert len(suggested_transforms.items()) == 0 def test_mismatch_transformations_and_num_params(): @T.prim_func(private=True, s_tir=True) def elemwise( arg: T.Buffer((32, 64, 224, 224), "float32"), relu: T.Buffer((32, 64, 224, 224), "float32"), ): for i0, i1, i2, i3 in T.grid(32, 64, 224, 224): with T.sblock("compute"): v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(arg[v_i0, v_i1, v_i2, v_i3]) T.writes(relu[v_i0, v_i1, v_i2, v_i3]) relu[v_i0, v_i1, v_i2, v_i3] = T.max(arg[v_i0, v_i1, v_i2, v_i3], T.float32(0)) with pytest.raises(RuntimeError, match="Incompatible PrimFunc and write_transformations"): _ = relax.analysis.suggest_layout_transforms( func=elemwise, write_buffer_transforms=[ lambda n, c, h, w: (n, h, w, c), lambda n, c, h, w: (n, h, w, c), lambda n, c, h, w: (n, h, w, c), ], ) def test_empty_write_transformations(): @T.prim_func(private=True, s_tir=True) def elemwise( arg: T.Buffer((32, 64, 224, 224), "float32"), relu: T.Buffer((32, 64, 224, 224), "float32"), ): for i0, i1, i2, i3 in T.grid(32, 64, 224, 224): with T.sblock("compute"): v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(arg[v_i0, v_i1, v_i2, v_i3]) T.writes(relu[v_i0, v_i1, v_i2, v_i3]) relu[v_i0, v_i1, v_i2, v_i3] = T.max(arg[v_i0, v_i1, v_i2, v_i3], T.float32(0)) suggested_transforms = relax.analysis.suggest_layout_transforms( func=elemwise, write_buffer_transforms=[] ) assert len(suggested_transforms.items()) == 0 def test_non_bijective_block_transform(): @T.prim_func(private=True, s_tir=True) def before( arg: T.Buffer((32, 64), "float32"), output: T.Buffer((32, 64), "float32"), ): for ax0, ax1 in T.grid(32, 64): with T.sblock("compute"): v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1]) T.reads(arg[v_ax0, v_ax1]) T.writes(output[v_ax0, v_ax1]) output[v_ax0, v_ax1] = arg[v_ax0, v_ax1] suggested_transforms = relax.analysis.suggest_layout_transforms( func=before, write_buffer_transforms=[lambda n, c: (n, c // 5, c % 5)] ) assert len(suggested_transforms.items()) == 0 def test_non_affine_access(): @T.prim_func(private=True, s_tir=True) def before( arg: T.Buffer((32, 64), "float32"), output: T.Buffer((32 * 64, 10), "float32"), ): for ax0, ax1, ax2 in T.grid(32, 64, 10): with T.sblock("compute"): v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2]) T.reads(arg[v_ax0, v_ax1]) T.writes(output[v_ax0 * v_ax1, v_ax2]) output[v_ax0 * v_ax1, v_ax2] = arg[v_ax0, v_ax1] suggested_transforms = relax.analysis.suggest_layout_transforms( func=before, write_buffer_transforms=[lambda a, b: (b, a)] ) assert len(suggested_transforms.items()) == 0 def test_unsupported_write_spatial_layout(): @T.prim_func(private=True, s_tir=True) def before( arg: T.Buffer((4, 4), "float32"), output: T.Buffer((16), "float32"), ): for ax0, ax1 in T.grid(4, 4): with T.sblock("flatten"): v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1]) T.reads(arg[v_ax0, v_ax1]) T.writes(output[v_ax0 * 4 + v_ax1]) output[v_ax0 * 4 + v_ax1] = arg[v_ax0, v_ax1] suggested_transforms = relax.analysis.suggest_layout_transforms( func=before, write_buffer_transforms=[lambda a: (a // 4, a % 4)] ) assert len(suggested_transforms.items()) == 0 def test_unpacked_iter_used_in_read_access(): @T.prim_func(private=True, s_tir=True) def before( arg: T.Buffer((8, 4), "float32"), output: T.Buffer((4, 8), "float32"), ): for ax0, ax1, ax2 in T.grid(4, 8, 4): with T.sblock("compute"): v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2]) T.reads(arg[v_ax1, v_ax2]) T.writes(output[v_ax0, v_ax1]) output[v_ax0, v_ax1] = arg[v_ax1, v_ax2] @T.prim_func(private=True, s_tir=True) def expected( arg: T.Buffer((8, 4), "float32"), output: T.Buffer((32), "float32"), ): for ax0, ax2 in T.grid(32, 4): with T.sblock("compute"): v_ax0, v_ax2 = T.axis.remap("SS", [ax0, ax2]) T.reads(arg[v_ax0 % 8, v_ax2]) T.writes(output[v_ax0]) output[v_ax0] = arg[v_ax0 % 8, v_ax2] suggested_transforms = relax.analysis.suggest_layout_transforms( func=before, write_buffer_transforms=[lambda a, b: a * 8 + b] ) after = apply_transformations(before, suggested_transforms) tvm.ir.assert_structural_equal(after, expected) def test_invalid_index_map(): @T.prim_func(private=True, s_tir=True) def elemwise( arg: T.Buffer((32, 64, 224, 224), "float32"), relu: T.Buffer((32, 64, 224, 224), "float32"), ): for i0, i1, i2, i3 in T.grid(32, 64, 224, 224): with T.sblock("compute"): v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(arg[v_i0, v_i1, v_i2, v_i3]) T.writes(relu[v_i0, v_i1, v_i2, v_i3]) relu[v_i0, v_i1, v_i2, v_i3] = T.max(arg[v_i0, v_i1, v_i2, v_i3], T.float32(0)) with pytest.raises(RuntimeError, match="Mismatch between output buffer shape and index map"): _ = relax.analysis.suggest_layout_transforms( func=elemwise, write_buffer_transforms=[lambda n, h, w: (n, w, h)] ) with pytest.raises(AssertionError): _ = relax.analysis.suggest_layout_transforms(func=elemwise, write_buffer_transforms=[2]) def test_SRSR_block(): @T.prim_func(private=True, s_tir=True) def before( arg: T.Buffer((32, 224, 64, 224), "float32"), sum: T.Buffer((32, 64), "float32"), ): for ax0, k2, ax1, k3 in T.grid(32, 224, 64, 224): with T.sblock("rxplaceholder_red"): v_ax0, v_k2, v_ax1, v_k3 = T.axis.remap("SRSR", [ax0, k2, ax1, k3]) T.reads(arg[v_ax0, v_ax1, v_k2, v_k3]) T.writes(sum[v_ax0, v_ax1]) with T.init(): sum[v_ax0, v_ax1] = T.float32(0) sum[v_ax0, v_ax1] = sum[v_ax0, v_ax1] + arg[v_ax0, v_k2, v_ax1, v_k3] @T.prim_func(private=True, s_tir=True) def expected( arg: T.Buffer((32, 224, 16, 224, 4), "float32"), sum: T.Buffer((32, 16, 4), "float32"), ): for ax0, ax1, ax2, ax3, ax4 in T.grid(32, 224, 16, 224, 4): with T.sblock("rxplaceholder_red"): v0, v1, v2, v3, v4 = T.axis.remap("SRSRS", [ax0, ax1, ax2, ax3, ax4]) T.reads(arg[v0, v1, v2, v3, v4]) T.writes(sum[v0, v2, v4]) with T.init(): sum[v0, v2, v4] = T.float32(0) sum[v0, v2, v4] = sum[v0, v2, v4] + arg[v0, v1, v2, v3, v4] suggested_transforms = relax.analysis.suggest_layout_transforms( func=before, write_buffer_transforms=[lambda n, c: (n, c // 4, c % 4)] ) after = apply_transformations(before, suggested_transforms) tvm.ir.assert_structural_equal(after, expected) def test_op_elemwise_symbolic(): @T.prim_func(private=True, s_tir=True) def before(arg: T.handle, relu: T.handle): N = T.int64() C = T.int64() H = T.int64() W = T.int64() Arg = T.match_buffer(arg, (N, C, H, W)) Relu = T.match_buffer(relu, (N, C, H, W)) for i0, i1, i2, i3 in T.grid(N, C, H, W): with T.sblock("compute"): v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(Arg[v_i0, v_i1, v_i2, v_i3]) T.writes(Relu[v_i0, v_i1, v_i2, v_i3]) Relu[v_i0, v_i1, v_i2, v_i3] = T.max(Arg[v_i0, v_i1, v_i2, v_i3], T.float32(0)) @T.prim_func(private=True, s_tir=True) def expected(arg: T.handle, relu: T.handle): N = T.int64() C = T.int64() H = T.int64() W = T.int64() Arg = T.match_buffer(arg, (N, H, W, C)) Relu = T.match_buffer(relu, (N, H, W, C)) # with T.sblock("root"): for ax0, ax1, ax2, ax3 in T.grid(N, H, W, C): with T.sblock("compute"): v0, v1, v2, v3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3]) T.reads(Arg[v0, v1, v2, v3]) T.writes(Relu[v0, v1, v2, v3]) Relu[v0, v1, v2, v3] = T.max(Arg[v0, v1, v2, v3], T.float32(0)) suggested_transforms = relax.analysis.suggest_layout_transforms( func=before, write_buffer_transforms=[lambda n, c, h, w: (n, h, w, c)] ) after = apply_transformations(before, suggested_transforms) tvm.ir.assert_structural_equal(after, expected) def test_op_elemwise(): @T.prim_func(private=True, s_tir=True) def before( arg: T.Buffer((32, 64, 224, 224), "float32"), relu: T.Buffer((32, 64, 224, 224), "float32"), ): for i0, i1, i2, i3 in T.grid(32, 64, 224, 224): with T.sblock("compute"): v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(arg[v_i0, v_i1, v_i2, v_i3]) T.writes(relu[v_i0, v_i1, v_i2, v_i3]) relu[v_i0, v_i1, v_i2, v_i3] = T.max(arg[v_i0, v_i1, v_i2, v_i3], T.float32(0)) @T.prim_func(private=True, s_tir=True) def expected( arg: T.Buffer((32, 224, 224, 64), "float32"), relu: T.Buffer((32, 224, 224, 64), "float32"), ): for ax0, ax1, ax2, ax3 in T.grid(32, 224, 224, 64): with T.sblock("compute"): v0, v1, v2, v3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3]) T.reads(arg[v0, v1, v2, v3]) T.writes(relu[v0, v1, v2, v3]) relu[v0, v1, v2, v3] = T.max(arg[v0, v1, v2, v3], T.float32(0)) suggested_transforms = relax.analysis.suggest_layout_transforms( func=before, write_buffer_transforms=[lambda n, c, h, w: (n, h, w, c)] ) after = apply_transformations(before, suggested_transforms) tvm.ir.assert_structural_equal(after, expected) def test_op_pool_nchw_nhwc(): @T.prim_func(private=True, s_tir=True) def before( arg: T.Buffer((32, 64, 224, 224), "float32"), pool_max: T.Buffer((32, 64, 111, 223), "float32"), ): for ax0, ax1, ax2, ax3, rv0, rv1 in T.grid(32, 64, 111, 223, 2, 2): with T.sblock("pool_max"): v_ax0, v_ax1, v_ax2, v_ax3, v_rv0, v_rv1 = T.axis.remap( "SSSSRR", [ax0, ax1, ax2, ax3, rv0, rv1] ) T.reads( arg[ v_ax0, v_ax1, v_ax2 * 2 + v_rv0 * 2, v_ax3 + v_rv1, ] ) T.writes(pool_max[v_ax0, v_ax1, v_ax2, v_ax3]) T.sblock_attr({"schedule_rule": "meta_schedule.pool_max"}) with T.init(): pool_max[v_ax0, v_ax1, v_ax2, v_ax3] = T.float32(-3.4028234663852886e38) pool_max[v_ax0, v_ax1, v_ax2, v_ax3] = T.max( pool_max[v_ax0, v_ax1, v_ax2, v_ax3], arg[ v_ax0, v_ax1, v_ax2 * 2 + v_rv0 * 2, v_ax3 + v_rv1, ], ) @T.prim_func(private=True, s_tir=True) def expected( arg: T.Buffer((32, 224, 224, 64), "float32"), pool_max: T.Buffer((32, 111, 223, 64), "float32"), ): # with T.sblock("root"): for ax0, ax1, ax2, ax3, ax4, ax5 in T.grid(32, 111, 223, 64, 2, 2): with T.sblock("pool_max"): v0, v1, v2, v3, v4, v5 = T.axis.remap("SSSSRR", [ax0, ax1, ax2, ax3, ax4, ax5]) T.reads(arg[v0, v1 * 2 + v4 * 2, v2 + v5, v3]) T.writes(pool_max[v0, v1, v2, v3]) T.sblock_attr({"schedule_rule": "meta_schedule.pool_max"}) with T.init(): pool_max[v0, v1, v2, v3] = T.float32(-3.4028234663852886e38) pool_max[v0, v1, v2, v3] = T.max( pool_max[v0, v1, v2, v3], arg[v0, v1 * 2 + v4 * 2, v2 + v5, v3], ) suggested_transforms = relax.analysis.suggest_layout_transforms( func=before, write_buffer_transforms=[lambda n, c, h, w: (n, h, w, c)], ) after = apply_transformations(before, suggested_transforms) tvm.ir.assert_structural_equal(after, expected) def test_op_pool_nchw16c_nhwc(): @T.prim_func(private=True, s_tir=True) def before( arg: T.Buffer( (32, 4, 224, 224, 16), "float32", ), pool_max: T.Buffer( (32, 4, 110, 220, 16), "float32", ), ): for ax0, ax1, ax2, ax3, ax4, rv0, rv1 in T.grid(32, 4, 110, 220, 16, 5, 5): with T.sblock("pool_max"): v_ax0, v_ax1, v_ax2, v_ax3, v_ax4, v_rv0, v_rv1 = T.axis.remap( "SSSSSRR", [ax0, ax1, ax2, ax3, ax4, rv0, rv1] ) T.reads(arg[v_ax0, v_ax1, v_ax2 * 2 + v_rv0, v_ax3 + v_rv1, v_ax4]) T.writes(pool_max[v_ax0, v_ax1, v_ax2, v_ax3, v_ax4]) T.sblock_attr({"schedule_rule": "meta_schedule.pool_max"}) with T.init(): pool_max[v_ax0, v_ax1, v_ax2, v_ax3, v_ax4] = T.float32(-3.4028234663852886e38) pool_max[v_ax0, v_ax1, v_ax2, v_ax3, v_ax4] = T.max( pool_max[v_ax0, v_ax1, v_ax2, v_ax3, v_ax4], arg[v_ax0, v_ax1, v_ax2 * 2 + v_rv0, v_ax3 + v_rv1, v_ax4], ) @T.prim_func(private=True, s_tir=True) def expected( arg: T.Buffer((32, 224, 224, 64), "float32"), pool_max: T.Buffer((32, 110, 220, 64), "float32"), ): for ax0, ax1, ax2, ax3, ax4, ax5 in T.grid(32, 110, 220, 64, 5, 5): with T.sblock("pool_max"): v0, v1, v2, v3, v4, v5 = T.axis.remap("SSSSRR", [ax0, ax1, ax2, ax3, ax4, ax5]) T.reads(arg[v0, v1 * 2 + v4, v2 + v5, v3]) T.writes(pool_max[v0, v1, v2, v3]) T.sblock_attr({"schedule_rule": "meta_schedule.pool_max"}) with T.init(): pool_max[v0, v1, v2, v3] = T.float32(-3.4028234663852886e38) pool_max[v0, v1, v2, v3] = T.max( pool_max[v0, v1, v2, v3], arg[v0, v1 * 2 + v4, v2 + v5, v3], ) suggested_transforms = relax.analysis.suggest_layout_transforms( func=before, write_buffer_transforms=[lambda n, C, h, w, c: (n, h, w, C * 16 + c)], ) after = apply_transformations(before, suggested_transforms) tvm.ir.assert_structural_equal(after, expected) def test_op_reduce(): @T.prim_func(private=True, s_tir=True) def before( arg: T.Buffer((32, 64, 224, 224), "float32"), sum: T.Buffer((32, 64), "float32"), ): for ax0, ax1, k2, k3 in T.grid(32, 64, 224, 224): with T.sblock("rxplaceholder_red"): v_ax0, v_ax1, v_k2, v_k3 = T.axis.remap("SSRR", [ax0, ax1, k2, k3]) T.reads(arg[v_ax0, v_ax1, v_k2, v_k3]) T.writes(sum[v_ax0, v_ax1]) with T.init(): sum[v_ax0, v_ax1] = T.float32(0) sum[v_ax0, v_ax1] = sum[v_ax0, v_ax1] + arg[v_ax0, v_ax1, v_k2, v_k3] @T.prim_func(private=True, s_tir=True) def expected( arg: T.Buffer((32, 4, 224, 224, 16), "float32"), sum: T.Buffer((32, 4, 16), "float32"), ): for ax0, ax1, ax2, ax3, ax4 in T.grid(32, 4, 224, 224, 16): with T.sblock("rxplaceholder_red"): v0, v1, v2, v3, v4 = T.axis.remap("SSRRS", [ax0, ax1, ax2, ax3, ax4]) T.reads(arg[v0, v1, v2, v3, v4]) T.writes(sum[v0, v1, v4]) with T.init(): sum[v0, v1, v4] = T.float32(0) sum[v0, v1, v4] = sum[v0, v1, v4] + arg[v0, v1, v2, v3, v4] suggested_transforms = relax.analysis.suggest_layout_transforms( func=before, write_buffer_transforms=[lambda n, c: (n, c // 16, c % 16)] ) after = apply_transformations(before, suggested_transforms) tvm.ir.assert_structural_equal(after, expected) def test_op_upsampling(): # relax materializes the layout if H, W or D dimensions are moved or tiled. @T.prim_func(private=True, s_tir=True) def before( arg: T.Buffer((32, 64, 224, 224), "float32"), resize: T.Buffer((32, 64, 202, 246), "float32"), ): for i0, i1, i2, i3 in T.grid(32, 64, 202, 246): with T.sblock("resize"): v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(arg[v_i0, v_i1, 0:224, 0:224]) T.writes(resize[v_i0, v_i1, v_i2, v_i3]) resize[v_i0, v_i1, v_i2, v_i3] = arg[ v_i0, v_i1, T.max( T.min( T.Cast( "int64", T.floor( T.float32(1.1089109182357788) * T.Cast("float32", v_i2) + T.float32(1.0000000000000001e-05) ), ), 223, ), 0, ), T.max( T.min( T.Cast( "int64", T.floor( T.float32(0.91056913137435913) * T.Cast("float32", v_i3) + T.float32(1.0000000000000001e-05) ), ), 223, ), 0, ), ] @T.prim_func(private=True, s_tir=True) def expected( arg: T.Buffer((32, 64, 224, 224), "float32"), resize: T.Buffer((32, 202, 246, 64), "float32"), ): # with T.sblock("root"): for ax0, ax1, ax2, ax3 in T.grid(32, 202, 246, 64): with T.sblock("resize"): v0, v1, v2, v3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3]) T.reads(arg[v0, v3, 0:224, 0:224]) T.writes(resize[v0, v1, v2, v3]) resize[v0, v1, v2, v3] = arg[ v0, v3, T.max( T.min( T.Cast( "int64", T.floor( T.float32(1.1089109182357788) * T.Cast("float32", v1) + T.float32(1.0000000000000001e-05) ), ), T.int64(223), ), T.int64(0), ), T.max( T.min( T.Cast( "int64", T.floor( T.float32(0.91056913137435913) * T.Cast("float32", v2) + T.float32(1.0000000000000001e-05) ), ), T.int64(223), ), T.int64(0), ), ] suggested_transforms = relax.analysis.suggest_layout_transforms( func=before, write_buffer_transforms=[lambda n, c, h, w: (n, h, w, c)] ) after = apply_transformations(before, suggested_transforms) tvm.ir.assert_structural_equal(after, expected) def test_op_strided_slice(): @T.prim_func(private=True, s_tir=True) def before( arg: T.Buffer((32, 64, 224, 224), "float32"), T_strided_slice_with_axes: T.Buffer((32, 64, 10, 8), "float32"), ): for ax0, ax1, ax2, ax3 in T.grid(32, 64, 10, 8): with T.sblock("T_strided_slice_with_axes"): v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3]) T.reads( arg[ v_ax0, v_ax1, v_ax2 * 5 + 2, v_ax3 * 7 + 4, ] ) T.writes(T_strided_slice_with_axes[v_ax0, v_ax1, v_ax2, v_ax3]) T_strided_slice_with_axes[v_ax0, v_ax1, v_ax2, v_ax3] = arg[ v_ax0, v_ax1, v_ax2 * 5 + 2, v_ax3 * 7 + 4, ] @T.prim_func(private=True, s_tir=True) def expected( arg: T.Buffer((32, 224, 224, 16, 4), "float32"), T_strided_slice_with_axes: T.Buffer((32, 10, 8, 16, 4), "float32"), ): # with T.sblock("root"): for ax0, ax1, ax2, ax3, ax4 in T.grid(32, 10, 8, 16, 4): with T.sblock("T_strided_slice_with_axes"): v0, v1, v2, v3, v4 = T.axis.remap("SSSSS", [ax0, ax1, ax2, ax3, ax4]) T.reads(arg[v0, v1 * 5 + 2, v2 * 7 + 4, v3, v4]) T.writes(T_strided_slice_with_axes[v0, v1, v2, v3, v4]) T_strided_slice_with_axes[v0, v1, v2, v3, v4] = arg[ v0, v1 * 5 + 2, v2 * 7 + 4, v3, v4 ] suggested_transforms = relax.analysis.suggest_layout_transforms( func=before, write_buffer_transforms=[lambda n, c, h, w: (n, h, w, c // 4, c % 4)] ) after = apply_transformations(before, suggested_transforms) tvm.ir.assert_structural_equal(after, expected) def test_op_binary_broadcast(): @T.prim_func(private=True, s_tir=True) def before( arg0: T.Buffer((32, 64, 224, 224), "float32"), arg1: T.Buffer((64, 224, 224), "float32"), T_add: T.Buffer((32, 64, 224, 224), "float32"), ): T.func_attr({"tirx.noalias": True}) # with T.sblock("root"): for ax0, ax1, ax2, ax3 in T.grid(32, 64, 224, 224): with T.sblock("T_add"): v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3]) T.reads( arg0[v_ax0, v_ax1, v_ax2, v_ax3], arg1[v_ax1, v_ax2, v_ax3], ) T.writes(T_add[v_ax0, v_ax1, v_ax2, v_ax3]) T_add[v_ax0, v_ax1, v_ax2, v_ax3] = ( arg0[v_ax0, v_ax1, v_ax2, v_ax3] + arg1[v_ax1, v_ax2, v_ax3] ) @T.prim_func(private=True, s_tir=True) def expected( arg0: T.Buffer((32, 224, 224, 16, 4), "float32"), arg1: T.Buffer((224, 224, 16, 4), "float32"), T_add: T.Buffer((32, 224, 224, 16, 4), "float32"), ): T.func_attr({"tirx.noalias": True}) # with T.sblock("root"): for ax0, ax1, ax2, ax3, ax4 in T.grid(32, 224, 224, 16, 4): with T.sblock("T_add"): v0, v1, v2, v3, v4 = T.axis.remap("SSSSS", [ax0, ax1, ax2, ax3, ax4]) T.reads(arg0[v0, v1, v2, v3, v4], arg1[v1, v2, v3, v4]) T.writes(T_add[v0, v1, v2, v3, v4]) T_add[v0, v1, v2, v3, v4] = arg0[v0, v1, v2, v3, v4] + arg1[v1, v2, v3, v4] suggested_transforms = relax.analysis.suggest_layout_transforms( func=before, write_buffer_transforms=[lambda n, c, h, w: (n, h, w, c // 4, c % 4)] ) after = apply_transformations(before, suggested_transforms) tvm.ir.assert_structural_equal(after, expected) def test_op_transpose(): @T.prim_func(private=True, s_tir=True) def before( arg: T.Buffer((32, 64, 224, 224), "float32"), T_transpose: T.Buffer((32, 224, 224, 64), "float32"), ): for ax0, ax1, ax2, ax3 in T.grid(32, 224, 224, 64): with T.sblock("T_transpose"): v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3]) T.reads(arg[v_ax0, v_ax3, v_ax1, v_ax2]) T.writes(T_transpose[v_ax0, v_ax1, v_ax2, v_ax3]) T_transpose[v_ax0, v_ax1, v_ax2, v_ax3] = arg[v_ax0, v_ax3, v_ax1, v_ax2] @T.prim_func(private=True, s_tir=True) def expected( arg: T.Buffer((32, 64, 224, 224), "float32"), T_transpose: T.Buffer((32, 224, 64, 224), "float32"), ): for ax0, ax1, ax2, ax3 in T.grid(32, 224, 64, 224): with T.sblock("T_transpose"): v0, v1, v2, v3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3]) T.reads(arg[v0, v2, v3, v1]) T.writes(T_transpose[v0, v1, v2, v3]) T_transpose[v0, v1, v2, v3] = arg[v0, v2, v3, v1] suggested_transforms = relax.analysis.suggest_layout_transforms( func=before, write_buffer_transforms=[lambda n, c, h, w: (n, h, w, c)] ) after = apply_transformations(before, suggested_transforms) tvm.ir.assert_structural_equal(after, expected) def test_op_pad(): @T.prim_func(private=True, s_tir=True) def before( arg: T.Buffer((32, 64, 224, 224), "float32"), PadInput: T.Buffer((32, 64, 230, 230), "float32"), ): for i0, i1, i2, i3 in T.grid(32, 64, 230, 230): with T.sblock("PadInput"): v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSS", [i0, i1, i2, i3]) T.reads(arg[v_i0, v_i1, v_i2 - 2, v_i3 - 2]) T.writes(PadInput[v_i0, v_i1, v_i2, v_i3]) PadInput[v_i0, v_i1, v_i2, v_i3] = T.if_then_else( 2 <= v_i2 and v_i2 < 226 and 2 <= v_i3 and v_i3 < 226, arg[v_i0, v_i1, v_i2 - 2, v_i3 - 2], T.float32(2), ) @T.prim_func(private=True, s_tir=True) def expected( arg: T.Buffer((32, 224, 224, 16, 4), "float32"), PadInput: T.Buffer((32, 230, 230, 16, 4), "float32"), ): for ax0, ax1, ax2, ax3, ax4 in T.grid(32, 230, 230, 16, 4): with T.sblock("PadInput"): v0, v1, v2, v3, v4 = T.axis.remap("SSSSS", [ax0, ax1, ax2, ax3, ax4]) T.reads(arg[v0, v1 - 2, v2 - 2, v3, v4]) T.writes(PadInput[v0, v1, v2, v3, v4]) PadInput[v0, v1, v2, v3, v4] = T.if_then_else( 2 <= v1 and v1 < 226 and 2 <= v2 and v2 < 226, arg[v0, v1 - 2, v2 - 2, v3, v4], T.float32(2), ) suggested_transforms = relax.analysis.suggest_layout_transforms( func=before, write_buffer_transforms=[lambda n, c, h, w: (n, h, w, c // 4, c % 4)] ) after = apply_transformations(before, suggested_transforms) tvm.ir.assert_structural_equal(after, expected) def test_op_split(): @T.prim_func(private=True, s_tir=True) def before( arg: T.Buffer((32, 64, 224, 224), "float32"), split0: T.Buffer((32, 32, 224, 224), "float32"), split1: T.Buffer((32, 32, 224, 224), "float32"), ): for ax0, ax1, ax2, ax3 in T.grid(32, 32, 224, 224): with T.sblock("T_split_sections"): v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3]) T.reads(arg[v_ax0, v_ax1, v_ax2, v_ax3]) T.writes(split0[v_ax0, v_ax1, v_ax2, v_ax3]) split0[v_ax0, v_ax1, v_ax2, v_ax3] = arg[v_ax0, v_ax1, v_ax2, v_ax3] for ax0, ax1, ax2, ax3 in T.grid(32, 32, 224, 224): with T.sblock("T_split_sections_1"): v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3]) T.reads(arg[v_ax0, v_ax1 + 32, v_ax2, v_ax3]) T.writes(split1[v_ax0, v_ax1, v_ax2, v_ax3]) split1[v_ax0, v_ax1, v_ax2, v_ax3] = arg[v_ax0, v_ax1 + 32, v_ax2, v_ax3] @T.prim_func(private=True, s_tir=True) def expected( arg: T.Buffer((32, 224, 224, 64), "float32"), split0: T.Buffer((32, 224, 224, 32), "float32"), split1: T.Buffer((32, 224, 224, 32), "float32"), ): for ax0, ax1, ax2, ax3 in T.grid(32, 224, 224, 32): with T.sblock("T_split_sections"): v0, v1, v2, v3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3]) T.reads(arg[v0, v1, v2, v3]) T.writes(split0[v0, v1, v2, v3]) split0[v0, v1, v2, v3] = arg[v0, v1, v2, v3] for ax0, ax1, ax2, ax3 in T.grid(32, 224, 224, 32): with T.sblock("T_split_sections_1"): v0, v1, v2, v3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3]) T.reads(arg[v0, v1, v2, v3 + 32]) T.writes(split1[v0, v1, v2, v3]) split1[v0, v1, v2, v3] = arg[v0, v1, v2, v3 + 32] suggested_transforms = relax.analysis.suggest_layout_transforms( func=before, write_buffer_transforms=[lambda n, c, h, w: (n, h, w, c), lambda n, c, h, w: (n, h, w, c)], ) after = apply_transformations(before, suggested_transforms) tvm.ir.assert_structural_equal(after, expected) @pytest.mark.skip("temp disable, due to minor arith regression") def test_op_split_tiling_split_dim(): @T.prim_func(private=True, s_tir=True) def before( arg: T.Buffer((32, 64, 224, 224), "float32"), split0: T.Buffer((32, 32, 224, 224), "float32"), split1: T.Buffer((32, 32, 224, 224), "float32"), ): for ax0, ax1, ax2, ax3 in T.grid(32, 32, 224, 224): with T.sblock("T_split_sections"): v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3]) T.reads(arg[v_ax0, v_ax1, v_ax2, v_ax3]) T.writes(split0[v_ax0, v_ax1, v_ax2, v_ax3]) split0[v_ax0, v_ax1, v_ax2, v_ax3] = arg[v_ax0, v_ax1, v_ax2, v_ax3] for ax0, ax1, ax2, ax3 in T.grid(32, 32, 224, 224): with T.sblock("T_split_sections_1"): v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3]) T.reads(arg[v_ax0, v_ax1 + 32, v_ax2, v_ax3]) T.writes(split1[v_ax0, v_ax1, v_ax2, v_ax3]) split1[v_ax0, v_ax1, v_ax2, v_ax3] = arg[v_ax0, v_ax1 + 32, v_ax2, v_ax3] @T.prim_func(private=True, s_tir=True) def expected( arg: T.Buffer((32, 224, 224, 16, 4), "float32"), split0: T.Buffer((32, 224, 224, 8, 4), "float32"), split1: T.Buffer((32, 224, 224, 8, 4), "float32"), ): # with T.sblock("root"): for ax0, ax1, ax2, ax3, ax4 in T.grid(32, 224, 224, 8, 4): with T.sblock("T_split_sections"): v0, v1, v2, v3, v4 = T.axis.remap("SSSSS", [ax0, ax1, ax2, ax3, ax4]) T.reads(arg[v0, v1, v2, v3, v4]) T.writes(split0[v0, v1, v2, v3, v4]) split0[v0, v1, v2, v3, v4] = arg[v0, v1, v2, v3, v4] for ax0, ax1, ax2, ax3, ax4 in T.grid(32, 224, 224, 8, 4): with T.sblock("T_split_sections_1"): v0, v1, v2, v3, v4 = T.axis.remap("SSSSS", [ax0, ax1, ax2, ax3, ax4]) T.reads(arg[v0, v1, v2, v3 + 8, v4]) T.writes(split1[v0, v1, v2, v3, v4]) split1[v0, v1, v2, v3, v4] = arg[v0, v1, v2, v3 + 8, v4] suggested_transforms = relax.analysis.suggest_layout_transforms( func=before, write_buffer_transforms=[ lambda n, c, h, w: (n, h, w, c // 4, c % 4), lambda n, c, h, w: (n, h, w, c // 4, c % 4), ], ) after = apply_transformations(before, suggested_transforms) tvm.ir.assert_structural_equal(after, expected) if __name__ == "__main__": tvm.testing.main()