wehub-resource-sync
834 lines
34 KiB
Python
834 lines
34 KiB
Python
# 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()
|