diff --git a/python/tvm/topi/hexagon/slice_ops/__init__.py b/python/tvm/topi/hexagon/slice_ops/__init__.py index b96156dc46d2..5f86e706af50 100644 --- a/python/tvm/topi/hexagon/slice_ops/__init__.py +++ b/python/tvm/topi/hexagon/slice_ops/__init__.py @@ -35,3 +35,4 @@ from .relu import relu_compute, relu_stir_schedule from .tanh import tanh_te_compute, tanhf16_schedule from .dwconv2d import * +from .depth_to_space import d2s_compute, d2s_schedule diff --git a/python/tvm/topi/hexagon/slice_ops/depth_to_space.py b/python/tvm/topi/hexagon/slice_ops/depth_to_space.py new file mode 100644 index 000000000000..aa14a97f5ee9 --- /dev/null +++ b/python/tvm/topi/hexagon/slice_ops/depth_to_space.py @@ -0,0 +1,43 @@ +# 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. + +""" Compute and schedule for depth to space slice op +""" + +from tvm import te, tir, topi +from ..utils import get_layout_transform_fn + + +def d2s_compute(inp, block_size, layout, mode): + """depth_to_space compute""" + return topi.nn.depth_to_space(inp, block_size=block_size, layout=layout, mode=mode) + + +def d2s_schedule(inp, out, input_layout, output_layout): + """Schedule for depth to space: top level function""" + if (input_layout != output_layout) or ( + output_layout not in ("nhwc-8h2w32c2w-2d", "nhwc-8h8w32c-2d") + ): + raise RuntimeError( + f"Unexpected input_layout, output_layout '{input_layout, output_layout}'" + ) + d2s_func = te.create_prim_func([inp, out]) + sch = tir.Schedule(d2s_func, debug_mask="all") + compute = sch.get_block("depth_to_space") + sch.transform_layout(compute, inp.name, get_layout_transform_fn(input_layout)) + sch.transform_layout(compute, out.name, get_layout_transform_fn(output_layout)) + return sch diff --git a/tests/python/contrib/test_hexagon/topi/test_depth_to_space.py b/tests/python/contrib/test_hexagon/topi/test_depth_to_space.py new file mode 100644 index 000000000000..f74d13f641d5 --- /dev/null +++ b/tests/python/contrib/test_hexagon/topi/test_depth_to_space.py @@ -0,0 +1,136 @@ +# 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. +# pylint: disable=line-too-long, redefined-outer-name + +"""Test depth_to_space slice op for hexagon""" + +import numpy as np +import pytest + +import tvm +from tvm import te +import tvm.testing +from tvm.topi.hexagon.slice_ops.depth_to_space import d2s_compute, d2s_schedule +from tvm.topi.testing import depth_to_space_python + +from ..infrastructure import allocate_hexagon_array, transform_numpy + + +d2s_fp16_tests = ( + ((1, 8, 8, 256), 2, "CDR", "float16", "nhwc-8h2w32c2w-2d", "nhwc-8h2w32c2w-2d"), + ((1, 8, 8, 1024), 4, "CDR", "float16", "nhwc-8h2w32c2w-2d", "nhwc-8h2w32c2w-2d"), + ((1, 16, 16, 256), 2, "CDR", "float16", "nhwc-8h2w32c2w-2d", "nhwc-8h2w32c2w-2d"), + ((1, 16, 16, 1024), 4, "CDR", "float16", "nhwc-8h2w32c2w-2d", "nhwc-8h2w32c2w-2d"), + ((1, 8, 8, 256), 2, "DCR", "float16", "nhwc-8h2w32c2w-2d", "nhwc-8h2w32c2w-2d"), + ((1, 8, 8, 1024), 4, "DCR", "float16", "nhwc-8h2w32c2w-2d", "nhwc-8h2w32c2w-2d"), + ((1, 16, 16, 256), 2, "DCR", "float16", "nhwc-8h2w32c2w-2d", "nhwc-8h2w32c2w-2d"), + ((1, 16, 16, 1024), 4, "DCR", "float16", "nhwc-8h2w32c2w-2d", "nhwc-8h2w32c2w-2d"), +) + +d2s_uint8_tests = ( + ((1, 8, 8, 256), 2, "CDR", "uint8", "nhwc-8h8w32c-2d", "nhwc-8h8w32c-2d"), + ((1, 8, 8, 1024), 4, "CDR", "uint8", "nhwc-8h8w32c-2d", "nhwc-8h8w32c-2d"), + ((1, 8, 8, 256), 2, "DCR", "uint8", "nhwc-8h8w32c-2d", "nhwc-8h8w32c-2d"), + ((1, 8, 8, 1024), 4, "DCR", "uint8", "nhwc-8h8w32c-2d", "nhwc-8h8w32c-2d"), +) + + +class TestD2SSlice: + """Test class that defines the Depth to Space slice test""" + + (input_shape, block_size, mode, dtype, input_layout, output_layout,) = tvm.testing.parameters( + *d2s_fp16_tests, + *d2s_uint8_tests, + ) + + working_scope = tvm.testing.parameter("global.vtcm") + + @tvm.testing.fixture + def input_np(self, input_shape, dtype): + return np.random.uniform(size=input_shape).astype(dtype) + + @tvm.testing.fixture + def transformed_input_np(self, input_np, input_layout): + return transform_numpy(input_np, "nhwc", input_layout) + + @tvm.testing.fixture + def ref_output_np(self, input_np, block_size, mode): + a_np = np.transpose(input_np, axes=[0, 3, 1, 2]) + ref_np = depth_to_space_python(a_np, block_size, mode=mode) + ref_np = np.transpose(ref_np, axes=[0, 2, 3, 1]) + return ref_np + + @tvm.testing.fixture + def transformed_ref_output_np(self, ref_output_np, output_layout): + return transform_numpy(ref_output_np, "nhwc", output_layout) + + @tvm.testing.requires_hexagon + def test_d2s_slice( + self, + input_shape, + block_size, + mode, + dtype, + input_layout, + output_layout, + hexagon_session, + working_scope, + transformed_input_np, + transformed_ref_output_np, + ): + """Top level testing function for depth to space""" + Input = te.placeholder(input_shape, name="Input", dtype=dtype) + + Output = d2s_compute(Input, block_size, "NHWC", mode) + + target_hexagon = tvm.target.hexagon("v69") + target = tvm.target.Target(target_hexagon, host=target_hexagon) + + tir_s = d2s_schedule(Input, Output, input_layout, output_layout) + + input_data = allocate_hexagon_array( + hexagon_session.device, + data=transformed_input_np, + axis_separators=[4], + mem_scope=working_scope, + ) + output_data = allocate_hexagon_array( + hexagon_session.device, + tensor_shape=transformed_ref_output_np.shape, + dtype=transformed_ref_output_np.dtype, + axis_separators=[4], + mem_scope=working_scope, + ) + with tvm.transform.PassContext(opt_level=3): + runtime_module = tvm.build( + tir_s.mod, [Input, Output], target=target, name="depth_to_space" + ) + mod = hexagon_session.load_module(runtime_module) + + mod(input_data, output_data) + output_np = output_data.numpy() + + tvm.testing.assert_allclose( + output_np, + transformed_ref_output_np, + 1e-3, + 1e-3, + ) + + +if __name__ == "__main__": + tvm.testing.main()