diff --git a/python/tvm/topi/hexagon/slice_ops/__init__.py b/python/tvm/topi/hexagon/slice_ops/__init__.py old mode 100644 new mode 100755 index 5b3ef530b0c0..87af3a767c38 --- a/python/tvm/topi/hexagon/slice_ops/__init__.py +++ b/python/tvm/topi/hexagon/slice_ops/__init__.py @@ -20,3 +20,4 @@ from .avg_pool2d import avg_pool2d_compute, avg_pool2d_STIR_schedule from .add_subtract_multiply import * from .softmax_slice import * +from .clip import * diff --git a/python/tvm/topi/hexagon/slice_ops/clip.py b/python/tvm/topi/hexagon/slice_ops/clip.py new file mode 100755 index 000000000000..2beb2df643bb --- /dev/null +++ b/python/tvm/topi/hexagon/slice_ops/clip.py @@ -0,0 +1,66 @@ +# 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=invalid-name + +""" +Clip the elements in `A` between `A_min` and `A_max`. +""" + +from tvm import te, tir, topi +from ..utils import get_layout_transform_fn + + +def clip_compute(A, A_min, A_max): + """ + Use topi clip implementation + """ + return topi.clip(A, A_min, A_max) + + +def clip_schedule(outs, ins, output_layout: str, input_layout: str): + """ + Hexagon clip schedule + """ + A = ins + M = outs + + func = te.create_prim_func([A, M]) + + s = tir.Schedule(func) + + block = s.get_block("compute") + + input_transformed_layout = get_layout_transform_fn(input_layout) + s.transform_layout(block, buffer=("read", 0), index_map=input_transformed_layout) + + output_transformed_layout = get_layout_transform_fn(output_layout) + s.transform_layout(block, buffer=("write", 0), index_map=output_transformed_layout) + + n, h, w, c = s.get_loops(block) + + ho, hi = s.split(h, [None, 8]) + wo, wi = s.split(w, [None, 4]) + co, ci = s.split(c, [None, 32]) + wio, wii = s.split(wi, [None, 2]) + + s.reorder(n, ho, wo, co, hi, wio, ci, wii) + + fused = s.fuse(ci, wii) + s.vectorize(fused) + + return s diff --git a/tests/python/contrib/test_hexagon/topi/test_clip.py b/tests/python/contrib/test_hexagon/topi/test_clip.py new file mode 100755 index 000000000000..37146b55dc1e --- /dev/null +++ b/tests/python/contrib/test_hexagon/topi/test_clip.py @@ -0,0 +1,128 @@ +# 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=invalid-name + +import pytest +import numpy as np + +from tvm import te, topi + +import tvm.testing +from tvm.topi import testing +from tvm.contrib.hexagon.build import HexagonLauncher +import tvm.topi.hexagon.slice_ops as sl +from ..infrastructure import allocate_hexagon_array, transform_numpy + +input_layout = tvm.testing.parameter( + "nhwc-8h2w32c2w-2d", +) + + +@tvm.testing.fixture +def input_np(input_shape, dtype): + return np.random.random(input_shape).astype(dtype) + + +@tvm.testing.fixture +def transformed_expected_output_np(expected_output_np, output_layout): + return transform_numpy(expected_output_np, "nhwc", output_layout) + + +@tvm.testing.fixture +def transformed_input_np(input_np, input_layout): + return transform_numpy(input_np, "nhwc", input_layout) + + +class TestClipSlice: + input_shape, output_shape, A_min, A_max, output_layout, dtype = tvm.testing.parameters( + ([1, 8, 4, 32], [1, 8, 4, 32], 0.1, 0.5, "nhwc-8h2w32c2w-2d", "float16") + ) + + @tvm.testing.fixture + def expected_output_np(self, input_np, A_min, A_max): + ref_np = np.clip(input_np, A_min, A_max) + return ref_np + + @tvm.testing.requires_hexagon + def test_clip_slice( + self, + input_shape, + output_shape, + input_np, + input_layout, + output_layout, + dtype, + A_min, + A_max, + transformed_input_np, + transformed_expected_output_np, + hexagon_session, + ): + # establish target and input placeholder + target_hexagon = tvm.target.hexagon("v69") + A = te.placeholder(input_shape, name="A", dtype=dtype) + + # get the compute function and schedule + M = sl.clip_compute(A, A_min, A_max) + + # Assume layout is nhwc-8h2w32c2w-2d + tir_schedule = sl.clip_schedule(M, A, output_layout, input_layout) + + # build the function + with tvm.transform.PassContext(opt_level=3): + func = tvm.build( + tir_schedule.mod, + target=tvm.target.Target(target_hexagon, host=target_hexagon), + name="clip", + ) + + # allocate input and output nd arrays + axis_separators = [4] + input_arr = allocate_hexagon_array( + hexagon_session.device, + data=transformed_input_np, + dtype=dtype, + axis_separators=axis_separators, + mem_scope="global.vtcm", + ) + + output_arr = allocate_hexagon_array( + hexagon_session.device, + transformed_expected_output_np.shape, + dtype=dtype, + axis_separators=axis_separators, + mem_scope="global.vtcm", + ) + + # execute + mod = hexagon_session.load_module(func) + mod(input_arr, output_arr) + + # convert output nd array to numpy array + output_np = output_arr.numpy() + b, h, w, c = output_shape + reshaped_output_np = np.reshape(output_np, [b, h // 8, w // 4, c // 32, 8, 2, 32, 2]) + + # test results + np.testing.assert_allclose( + reshaped_output_np, transformed_expected_output_np, rtol=1e-3, atol=1e-3 + ) + + +if __name__ == "__main__": + tvm.testing.main()