From 11e1fc2dfbc16d18380dc67177ee84c829c86d6e Mon Sep 17 00:00:00 2001
From: trahman-quic <trahman@quicinc.com>
Date: Wed, 1 Jun 2022 17:24:15 -0500
Subject: [PATCH 01/11] [UPSTREAM][HEXAGON] Slice ops added - add, subtract,
 multiply

---
 python/tvm/topi/hexagon/slice_ops/__init__.py |  22 ++
 .../slice_ops/add_subtract_multiply.py        |  60 +++++
 python/tvm/topi/hexagon/utils.py              |  75 ++++++
 .../contrib/test_hexagon/infrastructure.py    |  17 +-
 .../test_add_subtract_multiply.py             | 235 ++++++++++++++++++
 5 files changed, 408 insertions(+), 1 deletion(-)
 create mode 100755 python/tvm/topi/hexagon/slice_ops/__init__.py
 create mode 100755 python/tvm/topi/hexagon/slice_ops/add_subtract_multiply.py
 create mode 100755 python/tvm/topi/hexagon/utils.py
 create mode 100755 tests/python/contrib/test_hexagon/test_add_subtract_multiply.py

diff --git a/python/tvm/topi/hexagon/slice_ops/__init__.py b/python/tvm/topi/hexagon/slice_ops/__init__.py
new file mode 100755
index 000000000000..2c5ba9321803
--- /dev/null
+++ b/python/tvm/topi/hexagon/slice_ops/__init__.py
@@ -0,0 +1,22 @@
+# 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.
+
+""" Computes and Schedules for Hexagon slice ops. """
+
+# pylint: disable=wildcard-import
+
+from .add_subtract_multiply import *
diff --git a/python/tvm/topi/hexagon/slice_ops/add_subtract_multiply.py b/python/tvm/topi/hexagon/slice_ops/add_subtract_multiply.py
new file mode 100755
index 000000000000..43711fd5ba7e
--- /dev/null
+++ b/python/tvm/topi/hexagon/slice_ops/add_subtract_multiply.py
@@ -0,0 +1,60 @@
+from tvm.ir.module import IRModule
+from tvm import te
+from tvm import tir
+from tvm.script import tir as T
+from tvm import topi
+from ..utils import apply_transform, get_layout_transform_fn
+
+
+def add_broadcast_compute(A, B):
+    return topi.add(A, B)
+
+
+def subtract_broadcast_compute(A, B):
+    return topi.subtract(A, B)
+
+
+def multiply_broadcast_compute(A, B):
+    return topi.multiply(A, B)
+
+
+def get_layout(layout):
+    layout += "-2d"
+    return get_layout_transform_fn(layout)
+
+
+def STIR_broadcast_schedule(
+    M, A, B, output_layout: str, input_A_layout: str, input_B_layout: str, op_name: str
+):
+    func = te.create_prim_func([A, B, M])
+
+    s = tir.Schedule(func)
+
+    block_dict = {"add": "T_add", "subtract": "T_subtract", "multiply": "T_multiply"}
+
+    block = s.get_block(block_dict[op_name])
+
+    if input_A_layout == "nhwc-8h2w32c2w":
+        input_A_transformed_layout = get_layout(input_A_layout)
+        s.transform_layout(block, buffer=("read", 0), index_map=input_A_transformed_layout)
+
+    if input_B_layout == "nhwc-8h2w32c2w":
+        input_B_transformed_layout = get_layout(input_B_layout)
+        s.transform_layout(block, buffer=("read", 1), index_map=input_B_transformed_layout)
+
+    output_transformed_layout = get_layout(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/python/tvm/topi/hexagon/utils.py b/python/tvm/topi/hexagon/utils.py
new file mode 100755
index 000000000000..99e00f19028b
--- /dev/null
+++ b/python/tvm/topi/hexagon/utils.py
@@ -0,0 +1,75 @@
+# 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.
+
+from tvm import te
+
+
+def n11c_1024c_2d(n, h, w, c):
+    return [n, h, w, c // 1024, te.AXIS_SEPARATOR, c % 1024]
+
+
+def n11c_1024c_1d(n, h, w, c):
+    return [n, h, w, c // 1024, c % 1024]
+
+
+def nhwc_8h2w32c2w_2d(n, h, w, c):
+    return [n, h // 8, w // 4, c // 32, te.AXIS_SEPARATOR, h % 8, (w % 4) // 2, c % 32, w % 2]
+
+
+def nhwc_8h2w32c2w_1d(n, h, w, c):
+    return [n, h // 8, w // 4, c // 32, h % 8, (w % 4) // 2, c % 32, w % 2]
+
+
+def get_layout_transform_fn(layout):
+    if layout == "nhwc-8h2w32c2w-2d":
+        return nhwc_8h2w32c2w_2d
+    if layout == "nhwc-8h2w32c2w-1d":
+        return nhwc_8h2w32c2w_1d
+    elif layout == "n11c-1024c-2d":
+        return n11c_1024c_2d
+    elif layout == "n11c-1024c-1d":
+        return n11c_1024c_1d
+    else:
+        raise RuntimeError(f"Unexpected layout '{layout}'")
+
+
+def apply_transform(s, block, block_index: int, buffer_type: str, layout: str):
+    """Apply transform layout on a buffer
+
+    Parameters
+    ----------
+    s: Schedule
+    block : BlockRV
+        The block that accesses the target buffer
+    buffer_index: int
+        The index of the buffer in block's read or write region
+    buffer_type : str
+        Type of the buffer index, "read" or "write"
+    layout : str
+        Layout of the buffer
+    """
+    transform_fn = get_layout_transform_fn(layout)
+    if layout == "nhwc-8h2w32c2w-1d":
+        axis_separators = [4]
+    elif layout == "n11c-1024c-1d":
+        axis_separators = [2]
+    else:
+        raise RuntimeError(f"Unexpected layout '{layout}'")
+
+    s.transform_layout(block, block_index, buffer_type, transform_fn)
+    if axis_separators:
+        s.set_axis_separator(block, block_index, buffer_type, axis_separators)
\ No newline at end of file
diff --git a/tests/python/contrib/test_hexagon/infrastructure.py b/tests/python/contrib/test_hexagon/infrastructure.py
index 0c9a9478c870..c90379333361 100644
--- a/tests/python/contrib/test_hexagon/infrastructure.py
+++ b/tests/python/contrib/test_hexagon/infrastructure.py
@@ -48,7 +48,7 @@ def allocate_hexagon_array(
         for dim_i, dim_f in zip(boundaries[:-1], boundaries[1:])
     ]
 
-    arr = tvm.nd.empty(physical_shape, dtype=dtype, device=dev)
+    arr = tvm.nd.empty(physical_shape, dtype=dtype, device=dev, mem_scope=mem_scope)
 
     if data is not None:
         arr.copyfrom(data.reshape(physical_shape))
@@ -228,3 +228,18 @@ def compute(n, ho, wo, ko, hi, wi, ki):
         )
 
     return output_shape, compute
+
+
+# Transpose and reshape numpy array according to the specified layout
+def transform_numpy(arr_np, layout):
+    if layout == "nhwc":
+        return arr_np
+    elif layout == "nhwc-8h2w32c2w":
+        N, H, W, C = arr_np.shape
+        return arr_np.reshape([N, H // 8, 8, W // 4, 2, 2, C // 32, 32]).transpose(0, 1, 3, 6, 2, 4, 7, 5)
+    elif layout == "n11c-1024c":
+        N, H, W, C = arr_np.shape
+        assert (H == 1 and W == 1), "The size of H and W must be 1!"
+        return arr_np.reshape([N, C//1024, 1024]).transpose(0, 1, 2)
+    else:
+        raise RuntimeError(f"Unexpected layout '{layout}'")
diff --git a/tests/python/contrib/test_hexagon/test_add_subtract_multiply.py b/tests/python/contrib/test_hexagon/test_add_subtract_multiply.py
new file mode 100755
index 000000000000..f29d2b9c3f7b
--- /dev/null
+++ b/tests/python/contrib/test_hexagon/test_add_subtract_multiply.py
@@ -0,0 +1,235 @@
+# 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.
+
+
+import pytest
+import numpy as np
+
+np.set_printoptions(threshold=np.inf)
+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
+
+
+@tvm.testing.fixture
+def expected_output_np(input_np_A, input_np_B, op_name):
+    if op_name == "add":
+        out_ref = np.add(input_np_A, input_np_B)
+    elif op_name == "subtract":
+        out_ref = np.subtract(input_np_A, input_np_B)
+    elif op_name == "multiply":
+        out_ref = np.multiply(input_np_A, input_np_B)
+    return out_ref
+
+
+@tvm.testing.fixture
+def input_np_A(input_shape_A, dtype):
+    return np.random.random(input_shape_A).astype(dtype)
+
+
+@tvm.testing.fixture
+def input_np_B(input_shape_B, dtype):
+    return np.random.random(input_shape_B).astype(dtype)
+
+
+@tvm.testing.fixture
+def transformed_input_np_A(input_np_A, input_A_layout):
+    if input_A_layout == "nhwc-8h2w32c2w":
+        return transform_numpy(input_np_A, input_A_layout)
+    elif input_A_layout == "nhwc":
+        return input_np_A
+
+
+@tvm.testing.fixture
+def transformed_input_np_B(input_np_B, input_B_layout):
+    if input_B_layout == "nhwc-8h2w32c2w":
+        return transform_numpy(input_np_B, input_B_layout)
+    elif input_B_layout == "nhwc":
+        return input_np_B
+
+
+@tvm.testing.fixture
+def transformed_expected_output_np(expected_output_np, output_layout):
+    return transform_numpy(expected_output_np, output_layout)
+
+
+def hexagon_wrapper_allocation(
+    device, layout, axis_separators, tensor_shape=None, data=None, transformed_data=None, dtype=None
+):
+    if layout == "nhwc-8h2w32c2w":
+        data_nd = allocate_hexagon_array(
+            device,
+            tensor_shape=tensor_shape,
+            data=transformed_data,
+            dtype=dtype,
+            axis_separators=axis_separators,
+            mem_scope="global.vtcm",
+        )
+    elif layout == "nhwc":
+        data_nd = allocate_hexagon_array(
+            device,
+            data=data,
+        )
+    return data_nd
+
+
+class TestAddSubtractMultiplyBroadcast2d:
+    (
+        input_shape_A,
+        input_shape_B,
+        input_A_layout,
+        input_B_layout,
+        output_layout,
+        dtype,
+    ) = tvm.testing.parameters(
+        # no broadcast needed - short input
+        (
+            [1, 8, 4, 32],
+            [1, 8, 4, 32],
+            "nhwc-8h2w32c2w",
+            "nhwc-8h2w32c2w",
+            "nhwc-8h2w32c2w",
+            "float16",
+        ),
+        # no broadcast needed - large input
+        (
+            [1, 56, 64, 128],
+            [1, 56, 64, 128],
+            "nhwc-8h2w32c2w",
+            "nhwc-8h2w32c2w",
+            "nhwc-8h2w32c2w",
+            "float16",
+        ),
+        # one input needs broadcast
+        (
+            [1, 56, 64, 128],
+            [1, 1, 64, 1],
+            "nhwc-8h2w32c2w",
+            "nhwc",
+            "nhwc-8h2w32c2w",
+            "float16",
+        ),
+        # Both input needs broadcast
+        (
+            [1, 56, 1, 128],
+            [1, 1, 64, 1],
+            "nhwc",
+            "nhwc",
+            "nhwc-8h2w32c2w",
+            "float16",
+        ),
+        # One axis in one input needs broadcast
+        (
+            [1, 56, 20, 128],
+            [1, 56, 20, 1],
+            "nhwc-8h2w32c2w",
+            "nhwc",
+            "nhwc-8h2w32c2w",
+            "float16",
+        ),
+    )
+
+    op_name = tvm.testing.parameter("add", "subtract", "multiply")
+
+    @tvm.testing.requires_hexagon
+    def test_transform(
+        self,
+        dtype,
+        input_shape_A,
+        input_shape_B,
+        input_np_A,
+        input_np_B,
+        transformed_input_np_A,
+        transformed_input_np_B,
+        expected_output_np,
+        transformed_expected_output_np,
+        hexagon_session,
+        output_layout,
+        input_A_layout,
+        input_B_layout,
+        op_name,
+    ):
+        target_hexagon = tvm.target.hexagon("v69")
+        A = te.placeholder(input_shape_A, name="A", dtype=dtype)
+        B = te.placeholder(input_shape_B, name="B", dtype=dtype)
+        if op_name == "add":
+            M = sl.add_broadcast_compute(A, B)
+        elif op_name == "subtract":
+            M = sl.subtract_broadcast_compute(A, B)
+        elif op_name == "multiply":
+            M = sl.multiply_broadcast_compute(A, B)
+
+        tir_schedule = sl.STIR_broadcast_schedule(
+            M, A, B, output_layout, input_A_layout, input_B_layout, op_name
+        )
+        sch = tir_schedule.mod
+
+        input_axis_separator = [4]
+        if output_layout == "nhwc-8h2w32c2w":
+            output_axis_separator = [4]
+        else:
+            raise RuntimeError(f"Unexpected layout '{output_layout}'")
+
+        with tvm.transform.PassContext(opt_level=3, config={"tir.disable_assert": True}):
+            func = tvm.build(
+                sch,
+                [A, B, M],
+                tvm.target.Target(target_hexagon, host=target_hexagon),
+                name="slice_op_with_transform",
+            )
+
+        output_shape = expected_output_np.shape
+
+        A_data_nd = hexagon_wrapper_allocation(
+            hexagon_session.device,
+            layout=input_A_layout,
+            data=input_np_A,
+            transformed_data=transformed_input_np_A,
+            axis_separators=input_axis_separator,
+        )
+        B_data_nd = hexagon_wrapper_allocation(
+            hexagon_session.device,
+            layout=input_B_layout,
+            data=input_np_B,
+            transformed_data=transformed_input_np_B,
+            axis_separators=input_axis_separator,
+        )
+        M_data_nd = hexagon_wrapper_allocation(
+            hexagon_session.device,
+            layout=output_layout,
+            tensor_shape=transformed_expected_output_np.shape,
+            axis_separators=output_axis_separator,
+            dtype=dtype,
+        )
+
+        mod = hexagon_session.load_module(func)
+        mod(A_data_nd, B_data_nd, M_data_nd)
+
+        b, h, w, c = output_shape
+        # convert nd to np and reshape to fixed chunk size layout
+        if output_layout == "nhwc-8h2w32c2w":
+            M_data_np = M_data_nd.numpy().reshape([b, h // 8, w // 4, c // 32, 8, 2, 32, 2])
+
+        np.testing.assert_allclose(transformed_expected_output_np, M_data_np, rtol=1e-3, atol=1e-3)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(sys.argv))

From b8a82e17b42b294996cdac9db8ed8c648f1b54f6 Mon Sep 17 00:00:00 2001
From: trahman-quic <trahman@quicinc.com>
Date: Wed, 1 Jun 2022 17:38:49 -0500
Subject: [PATCH 02/11] reformatted

---
 python/tvm/topi/hexagon/utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/tvm/topi/hexagon/utils.py b/python/tvm/topi/hexagon/utils.py
index 99e00f19028b..a11b85d3a8b1 100755
--- a/python/tvm/topi/hexagon/utils.py
+++ b/python/tvm/topi/hexagon/utils.py
@@ -72,4 +72,4 @@ def apply_transform(s, block, block_index: int, buffer_type: str, layout: str):
 
     s.transform_layout(block, block_index, buffer_type, transform_fn)
     if axis_separators:
-        s.set_axis_separator(block, block_index, buffer_type, axis_separators)
\ No newline at end of file
+        s.set_axis_separator(block, block_index, buffer_type, axis_separators)

From 037b7e2335a234f03c34d521cd84aae0273eba0b Mon Sep 17 00:00:00 2001
From: trahman-quic <trahman@quicinc.com>
Date: Thu, 2 Jun 2022 11:34:16 -0500
Subject: [PATCH 03/11] More reformatting

---
 .../tvm/topi/hexagon/slice_ops/add_subtract_multiply.py  | 6 +++---
 python/tvm/topi/hexagon/utils.py                         | 1 +
 tests/python/contrib/test_hexagon/infrastructure.py      | 9 ++++++---
 .../contrib/test_hexagon/test_add_subtract_multiply.py   | 1 +
 4 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/python/tvm/topi/hexagon/slice_ops/add_subtract_multiply.py b/python/tvm/topi/hexagon/slice_ops/add_subtract_multiply.py
index 43711fd5ba7e..cb41cd078b28 100755
--- a/python/tvm/topi/hexagon/slice_ops/add_subtract_multiply.py
+++ b/python/tvm/topi/hexagon/slice_ops/add_subtract_multiply.py
@@ -1,9 +1,9 @@
-from tvm.ir.module import IRModule
+# pylint: disable=invalid-name
+
 from tvm import te
 from tvm import tir
-from tvm.script import tir as T
 from tvm import topi
-from ..utils import apply_transform, get_layout_transform_fn
+from ..utils import get_layout_transform_fn
 
 
 def add_broadcast_compute(A, B):
diff --git a/python/tvm/topi/hexagon/utils.py b/python/tvm/topi/hexagon/utils.py
index a11b85d3a8b1..450708f4b2a9 100755
--- a/python/tvm/topi/hexagon/utils.py
+++ b/python/tvm/topi/hexagon/utils.py
@@ -14,6 +14,7 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+# pylint: disable=invalid-name
 
 from tvm import te
 
diff --git a/tests/python/contrib/test_hexagon/infrastructure.py b/tests/python/contrib/test_hexagon/infrastructure.py
index c90379333361..e151d6986b6b 100644
--- a/tests/python/contrib/test_hexagon/infrastructure.py
+++ b/tests/python/contrib/test_hexagon/infrastructure.py
@@ -14,6 +14,7 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+# pylint: disable=invalid-name
 
 """ Hexagon testing infrastructure """
 
@@ -236,10 +237,12 @@ def transform_numpy(arr_np, layout):
         return arr_np
     elif layout == "nhwc-8h2w32c2w":
         N, H, W, C = arr_np.shape
-        return arr_np.reshape([N, H // 8, 8, W // 4, 2, 2, C // 32, 32]).transpose(0, 1, 3, 6, 2, 4, 7, 5)
+        return arr_np.reshape([N, H // 8, 8, W // 4, 2, 2, C // 32, 32]).transpose(
+            0, 1, 3, 6, 2, 4, 7, 5
+        )
     elif layout == "n11c-1024c":
         N, H, W, C = arr_np.shape
-        assert (H == 1 and W == 1), "The size of H and W must be 1!"
-        return arr_np.reshape([N, C//1024, 1024]).transpose(0, 1, 2)
+        assert H == 1 and W == 1, "The size of H and W must be 1!"
+        return arr_np.reshape([N, C // 1024, 1024]).transpose(0, 1, 2)
     else:
         raise RuntimeError(f"Unexpected layout '{layout}'")
diff --git a/tests/python/contrib/test_hexagon/test_add_subtract_multiply.py b/tests/python/contrib/test_hexagon/test_add_subtract_multiply.py
index f29d2b9c3f7b..7639c9bee497 100755
--- a/tests/python/contrib/test_hexagon/test_add_subtract_multiply.py
+++ b/tests/python/contrib/test_hexagon/test_add_subtract_multiply.py
@@ -14,6 +14,7 @@
 # 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

From 009e536c193760e51c8f9d04def400e3f3eda1b4 Mon Sep 17 00:00:00 2001
From: trahman-quic <trahman@quicinc.com>
Date: Thu, 2 Jun 2022 14:13:50 -0500
Subject: [PATCH 04/11] Address comments

---
 .../slice_ops/add_subtract_multiply.py        | 37 +++++++++++++---
 python/tvm/topi/hexagon/utils.py              | 42 ++++---------------
 .../test_add_subtract_multiply.py             |  4 +-
 3 files changed, 43 insertions(+), 40 deletions(-)

diff --git a/python/tvm/topi/hexagon/slice_ops/add_subtract_multiply.py b/python/tvm/topi/hexagon/slice_ops/add_subtract_multiply.py
index cb41cd078b28..f6e6d9671e5e 100755
--- a/python/tvm/topi/hexagon/slice_ops/add_subtract_multiply.py
+++ b/python/tvm/topi/hexagon/slice_ops/add_subtract_multiply.py
@@ -1,4 +1,26 @@
-# pylint: disable=invalid-name
+# 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, unused-variable, unused-argument, too-many-locals
+
+"""Compute and schedule for add, multiply, subtract slice op
+
+Please note the following assumptions made by the implementation:
+
+1) The inputs will be multiple of crouton layout except for the axis that needs broadcasting."""
 
 from tvm import te
 from tvm import tir
@@ -7,18 +29,22 @@
 
 
 def add_broadcast_compute(A, B):
+    """Call the add op from topi"""
     return topi.add(A, B)
 
 
 def subtract_broadcast_compute(A, B):
+    """Call the subtract op from topi"""
     return topi.subtract(A, B)
 
 
 def multiply_broadcast_compute(A, B):
+    """Call the multiply op from topi"""
     return topi.multiply(A, B)
 
 
-def get_layout(layout):
+def get_2d_layout(layout):
+    """Get the 2d layout for transformation"""
     layout += "-2d"
     return get_layout_transform_fn(layout)
 
@@ -26,6 +52,7 @@ def get_layout(layout):
 def STIR_broadcast_schedule(
     M, A, B, output_layout: str, input_A_layout: str, input_B_layout: str, op_name: str
 ):
+    """Schedule for input and output layout nhwc-8h2w32c2w considering broadcast"""
     func = te.create_prim_func([A, B, M])
 
     s = tir.Schedule(func)
@@ -35,14 +62,14 @@ def STIR_broadcast_schedule(
     block = s.get_block(block_dict[op_name])
 
     if input_A_layout == "nhwc-8h2w32c2w":
-        input_A_transformed_layout = get_layout(input_A_layout)
+        input_A_transformed_layout = get_2d_layout(input_A_layout)
         s.transform_layout(block, buffer=("read", 0), index_map=input_A_transformed_layout)
 
     if input_B_layout == "nhwc-8h2w32c2w":
-        input_B_transformed_layout = get_layout(input_B_layout)
+        input_B_transformed_layout = get_2d_layout(input_B_layout)
         s.transform_layout(block, buffer=("read", 1), index_map=input_B_transformed_layout)
 
-    output_transformed_layout = get_layout(output_layout)
+    output_transformed_layout = get_2d_layout(output_layout)
     s.transform_layout(block, buffer=("write", 0), index_map=output_transformed_layout)
 
     n, h, w, c = s.get_loops(block)
diff --git a/python/tvm/topi/hexagon/utils.py b/python/tvm/topi/hexagon/utils.py
index 450708f4b2a9..def86486dbe9 100755
--- a/python/tvm/topi/hexagon/utils.py
+++ b/python/tvm/topi/hexagon/utils.py
@@ -15,62 +15,38 @@
 # specific language governing permissions and limitations
 # under the License.
 # pylint: disable=invalid-name
-
+"""Common hexagon specific utilities"""
 from tvm import te
 
 
 def n11c_1024c_2d(n, h, w, c):
+    """Return index map for n11c_1024 2d layout"""
     return [n, h, w, c // 1024, te.AXIS_SEPARATOR, c % 1024]
 
 
 def n11c_1024c_1d(n, h, w, c):
+    """Return index map for n11c_1024 1d layout"""
     return [n, h, w, c // 1024, c % 1024]
 
 
 def nhwc_8h2w32c2w_2d(n, h, w, c):
+    """Return index map for nhwc_8h2w32c2w 2d layout"""
     return [n, h // 8, w // 4, c // 32, te.AXIS_SEPARATOR, h % 8, (w % 4) // 2, c % 32, w % 2]
 
 
 def nhwc_8h2w32c2w_1d(n, h, w, c):
+    """Return index map for nhwc_8h2w32c2w 1d layout"""
     return [n, h // 8, w // 4, c // 32, h % 8, (w % 4) // 2, c % 32, w % 2]
 
 
 def get_layout_transform_fn(layout):
+    """Return index map function as per the layout string"""
     if layout == "nhwc-8h2w32c2w-2d":
         return nhwc_8h2w32c2w_2d
     if layout == "nhwc-8h2w32c2w-1d":
         return nhwc_8h2w32c2w_1d
-    elif layout == "n11c-1024c-2d":
+    if layout == "n11c-1024c-2d":
         return n11c_1024c_2d
-    elif layout == "n11c-1024c-1d":
+    if layout == "n11c-1024c-1d":
         return n11c_1024c_1d
-    else:
-        raise RuntimeError(f"Unexpected layout '{layout}'")
-
-
-def apply_transform(s, block, block_index: int, buffer_type: str, layout: str):
-    """Apply transform layout on a buffer
-
-    Parameters
-    ----------
-    s: Schedule
-    block : BlockRV
-        The block that accesses the target buffer
-    buffer_index: int
-        The index of the buffer in block's read or write region
-    buffer_type : str
-        Type of the buffer index, "read" or "write"
-    layout : str
-        Layout of the buffer
-    """
-    transform_fn = get_layout_transform_fn(layout)
-    if layout == "nhwc-8h2w32c2w-1d":
-        axis_separators = [4]
-    elif layout == "n11c-1024c-1d":
-        axis_separators = [2]
-    else:
-        raise RuntimeError(f"Unexpected layout '{layout}'")
-
-    s.transform_layout(block, block_index, buffer_type, transform_fn)
-    if axis_separators:
-        s.set_axis_separator(block, block_index, buffer_type, axis_separators)
+    raise RuntimeError(f"Unexpected layout '{layout}'")
diff --git a/tests/python/contrib/test_hexagon/test_add_subtract_multiply.py b/tests/python/contrib/test_hexagon/test_add_subtract_multiply.py
index 7639c9bee497..1310e569a11e 100755
--- a/tests/python/contrib/test_hexagon/test_add_subtract_multiply.py
+++ b/tests/python/contrib/test_hexagon/test_add_subtract_multiply.py
@@ -20,7 +20,6 @@
 import pytest
 import numpy as np
 
-np.set_printoptions(threshold=np.inf)
 from tvm import te, topi
 
 import tvm.testing
@@ -75,6 +74,7 @@ def transformed_expected_output_np(expected_output_np, output_layout):
 def hexagon_wrapper_allocation(
     device, layout, axis_separators, tensor_shape=None, data=None, transformed_data=None, dtype=None
 ):
+    """Input layout can either be nhwc-8h2w32c2w or nhwc"""
     if layout == "nhwc-8h2w32c2w":
         data_nd = allocate_hexagon_array(
             device,
@@ -233,4 +233,4 @@ def test_transform(
 
 
 if __name__ == "__main__":
-    sys.exit(pytest.main(sys.argv))
+    tvm.testing.main()

From 8daeb730761dfad7f237dc2005775c2a456ad9d6 Mon Sep 17 00:00:00 2001
From: trahman-quic <trahman@quicinc.com>
Date: Fri, 3 Jun 2022 10:23:00 -0500
Subject: [PATCH 05/11] Change to v68

---
 tests/python/contrib/test_hexagon/test_add_subtract_multiply.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/python/contrib/test_hexagon/test_add_subtract_multiply.py b/tests/python/contrib/test_hexagon/test_add_subtract_multiply.py
index 1310e569a11e..3168eb23973a 100755
--- a/tests/python/contrib/test_hexagon/test_add_subtract_multiply.py
+++ b/tests/python/contrib/test_hexagon/test_add_subtract_multiply.py
@@ -168,7 +168,7 @@ def test_transform(
         input_B_layout,
         op_name,
     ):
-        target_hexagon = tvm.target.hexagon("v69")
+        target_hexagon = tvm.target.hexagon("v68")
         A = te.placeholder(input_shape_A, name="A", dtype=dtype)
         B = te.placeholder(input_shape_B, name="B", dtype=dtype)
         if op_name == "add":

From 7c08bd27df6bd6bb55919313f431ff29afde7f94 Mon Sep 17 00:00:00 2001
From: trahman-quic <trahman@quicinc.com>
Date: Tue, 7 Jun 2022 10:04:35 -0500
Subject: [PATCH 06/11] Change transform_numpy function call

---
 .../slice_ops/add_subtract_multiply.py        | 18 +++-----
 .../contrib/test_hexagon/infrastructure.py    | 32 +++++++-------
 .../test_add_subtract_multiply.py             | 44 ++++++++-----------
 3 files changed, 42 insertions(+), 52 deletions(-)

diff --git a/python/tvm/topi/hexagon/slice_ops/add_subtract_multiply.py b/python/tvm/topi/hexagon/slice_ops/add_subtract_multiply.py
index f6e6d9671e5e..1995dac12eb2 100755
--- a/python/tvm/topi/hexagon/slice_ops/add_subtract_multiply.py
+++ b/python/tvm/topi/hexagon/slice_ops/add_subtract_multiply.py
@@ -43,16 +43,10 @@ def multiply_broadcast_compute(A, B):
     return topi.multiply(A, B)
 
 
-def get_2d_layout(layout):
-    """Get the 2d layout for transformation"""
-    layout += "-2d"
-    return get_layout_transform_fn(layout)
-
-
 def STIR_broadcast_schedule(
     M, A, B, output_layout: str, input_A_layout: str, input_B_layout: str, op_name: str
 ):
-    """Schedule for input and output layout nhwc-8h2w32c2w considering broadcast"""
+    """Schedule for input and output layout nhwc-8h2w32c2w-2d considering broadcast"""
     func = te.create_prim_func([A, B, M])
 
     s = tir.Schedule(func)
@@ -61,15 +55,15 @@ def STIR_broadcast_schedule(
 
     block = s.get_block(block_dict[op_name])
 
-    if input_A_layout == "nhwc-8h2w32c2w":
-        input_A_transformed_layout = get_2d_layout(input_A_layout)
+    if input_A_layout == "nhwc-8h2w32c2w-2d":
+        input_A_transformed_layout = get_layout_transform_fn(input_A_layout)
         s.transform_layout(block, buffer=("read", 0), index_map=input_A_transformed_layout)
 
-    if input_B_layout == "nhwc-8h2w32c2w":
-        input_B_transformed_layout = get_2d_layout(input_B_layout)
+    if input_B_layout == "nhwc-8h2w32c2w-2d":
+        input_B_transformed_layout = get_layout_transform_fn(input_B_layout)
         s.transform_layout(block, buffer=("read", 1), index_map=input_B_transformed_layout)
 
-    output_transformed_layout = get_2d_layout(output_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)
diff --git a/tests/python/contrib/test_hexagon/infrastructure.py b/tests/python/contrib/test_hexagon/infrastructure.py
index e151d6986b6b..57a9dff8b424 100644
--- a/tests/python/contrib/test_hexagon/infrastructure.py
+++ b/tests/python/contrib/test_hexagon/infrastructure.py
@@ -231,18 +231,20 @@ def compute(n, ho, wo, ko, hi, wi, ki):
     return output_shape, compute
 
 
-# Transpose and reshape numpy array according to the specified layout
-def transform_numpy(arr_np, layout):
-    if layout == "nhwc":
-        return arr_np
-    elif layout == "nhwc-8h2w32c2w":
-        N, H, W, C = arr_np.shape
-        return arr_np.reshape([N, H // 8, 8, W // 4, 2, 2, C // 32, 32]).transpose(
-            0, 1, 3, 6, 2, 4, 7, 5
-        )
-    elif layout == "n11c-1024c":
-        N, H, W, C = arr_np.shape
-        assert H == 1 and W == 1, "The size of H and W must be 1!"
-        return arr_np.reshape([N, C // 1024, 1024]).transpose(0, 1, 2)
-    else:
-        raise RuntimeError(f"Unexpected layout '{layout}'")
+def transform_numpy(arr_np, current_layout: str, new_layout: str):
+    """Reshape and transpose numpy array according to the specified layout"""
+    if current_layout == "nhwc":
+        if new_layout == "nhwc":
+            return arr_np
+        if new_layout in ["nhwc-8h2w32c2w-2d", "nhwc-8h2w32c2w-1d"]:
+            n, h, w, c = arr_np.shape
+            return arr_np.reshape([n, h // 8, 8, w // 4, 2, 2, c // 32, 32]).transpose(
+                0, 1, 3, 6, 2, 4, 7, 5
+            )
+        if new_layout in ["n11c-1024c-2d", "n11c-1024c-1d"]:
+            n, h, w, c = arr_np.shape
+            assert h == 1 and w == 1, "The size of h and w must be 1"
+            return arr_np.reshape([n, 1, 1, c // 1024, 1024])
+
+        raise RuntimeError(f"Unexpected new_layout '{new_layout}'")
+    raise RuntimeError(f"Unexpected current_layout '{current_layout}'")
diff --git a/tests/python/contrib/test_hexagon/test_add_subtract_multiply.py b/tests/python/contrib/test_hexagon/test_add_subtract_multiply.py
index 3168eb23973a..ff0b9e921ca9 100755
--- a/tests/python/contrib/test_hexagon/test_add_subtract_multiply.py
+++ b/tests/python/contrib/test_hexagon/test_add_subtract_multiply.py
@@ -52,30 +52,24 @@ def input_np_B(input_shape_B, dtype):
 
 @tvm.testing.fixture
 def transformed_input_np_A(input_np_A, input_A_layout):
-    if input_A_layout == "nhwc-8h2w32c2w":
-        return transform_numpy(input_np_A, input_A_layout)
-    elif input_A_layout == "nhwc":
-        return input_np_A
+    return transform_numpy(input_np_A, "nhwc", input_A_layout)
 
 
 @tvm.testing.fixture
 def transformed_input_np_B(input_np_B, input_B_layout):
-    if input_B_layout == "nhwc-8h2w32c2w":
-        return transform_numpy(input_np_B, input_B_layout)
-    elif input_B_layout == "nhwc":
-        return input_np_B
+    return transform_numpy(input_np_B, "nhwc", input_B_layout)
 
 
 @tvm.testing.fixture
 def transformed_expected_output_np(expected_output_np, output_layout):
-    return transform_numpy(expected_output_np, output_layout)
+    return transform_numpy(expected_output_np, "nhwc", output_layout)
 
 
 def hexagon_wrapper_allocation(
     device, layout, axis_separators, tensor_shape=None, data=None, transformed_data=None, dtype=None
 ):
-    """Input layout can either be nhwc-8h2w32c2w or nhwc"""
-    if layout == "nhwc-8h2w32c2w":
+    """Input layout can either be nhwc-8h2w32c2w-2d or nhwc"""
+    if layout == "nhwc-8h2w32c2w-2d":
         data_nd = allocate_hexagon_array(
             device,
             tensor_shape=tensor_shape,
@@ -105,27 +99,27 @@ class TestAddSubtractMultiplyBroadcast2d:
         (
             [1, 8, 4, 32],
             [1, 8, 4, 32],
-            "nhwc-8h2w32c2w",
-            "nhwc-8h2w32c2w",
-            "nhwc-8h2w32c2w",
+            "nhwc-8h2w32c2w-2d",
+            "nhwc-8h2w32c2w-2d",
+            "nhwc-8h2w32c2w-2d",
             "float16",
         ),
         # no broadcast needed - large input
         (
             [1, 56, 64, 128],
             [1, 56, 64, 128],
-            "nhwc-8h2w32c2w",
-            "nhwc-8h2w32c2w",
-            "nhwc-8h2w32c2w",
+            "nhwc-8h2w32c2w-2d",
+            "nhwc-8h2w32c2w-2d",
+            "nhwc-8h2w32c2w-2d",
             "float16",
         ),
         # one input needs broadcast
         (
             [1, 56, 64, 128],
             [1, 1, 64, 1],
-            "nhwc-8h2w32c2w",
+            "nhwc-8h2w32c2w-2d",
             "nhwc",
-            "nhwc-8h2w32c2w",
+            "nhwc-8h2w32c2w-2d",
             "float16",
         ),
         # Both input needs broadcast
@@ -134,16 +128,16 @@ class TestAddSubtractMultiplyBroadcast2d:
             [1, 1, 64, 1],
             "nhwc",
             "nhwc",
-            "nhwc-8h2w32c2w",
+            "nhwc-8h2w32c2w-2d",
             "float16",
         ),
         # One axis in one input needs broadcast
         (
             [1, 56, 20, 128],
             [1, 56, 20, 1],
-            "nhwc-8h2w32c2w",
+            "nhwc-8h2w32c2w-2d",
             "nhwc",
-            "nhwc-8h2w32c2w",
+            "nhwc-8h2w32c2w-2d",
             "float16",
         ),
     )
@@ -168,7 +162,7 @@ def test_transform(
         input_B_layout,
         op_name,
     ):
-        target_hexagon = tvm.target.hexagon("v68")
+        target_hexagon = tvm.target.hexagon("v69")
         A = te.placeholder(input_shape_A, name="A", dtype=dtype)
         B = te.placeholder(input_shape_B, name="B", dtype=dtype)
         if op_name == "add":
@@ -184,7 +178,7 @@ def test_transform(
         sch = tir_schedule.mod
 
         input_axis_separator = [4]
-        if output_layout == "nhwc-8h2w32c2w":
+        if output_layout == "nhwc-8h2w32c2w-2d":
             output_axis_separator = [4]
         else:
             raise RuntimeError(f"Unexpected layout '{output_layout}'")
@@ -226,7 +220,7 @@ def test_transform(
 
         b, h, w, c = output_shape
         # convert nd to np and reshape to fixed chunk size layout
-        if output_layout == "nhwc-8h2w32c2w":
+        if output_layout == "nhwc-8h2w32c2w-2d":
             M_data_np = M_data_nd.numpy().reshape([b, h // 8, w // 4, c // 32, 8, 2, 32, 2])
 
         np.testing.assert_allclose(transformed_expected_output_np, M_data_np, rtol=1e-3, atol=1e-3)

From 1462a3d383fb0dfef7490681ee05e850d998d19b Mon Sep 17 00:00:00 2001
From: trahman-quic <trahman@quicinc.com>
Date: Wed, 8 Jun 2022 16:54:10 -0500
Subject: [PATCH 07/11] Do not disbale pylint errors and fix them

---
 python/tvm/topi/hexagon/slice_ops/__init__.py |  2 -
 .../slice_ops/add_subtract_multiply.py        | 49 ++++++++++---------
 python/tvm/topi/hexagon/utils.py              |  2 +-
 .../contrib/test_hexagon/infrastructure.py    |  1 -
 .../test_add_subtract_multiply.py             |  3 +-
 5 files changed, 29 insertions(+), 28 deletions(-)

diff --git a/python/tvm/topi/hexagon/slice_ops/__init__.py b/python/tvm/topi/hexagon/slice_ops/__init__.py
index 2c5ba9321803..6671b8757137 100755
--- a/python/tvm/topi/hexagon/slice_ops/__init__.py
+++ b/python/tvm/topi/hexagon/slice_ops/__init__.py
@@ -17,6 +17,4 @@
 
 """ Computes and Schedules for Hexagon slice ops. """
 
-# pylint: disable=wildcard-import
-
 from .add_subtract_multiply import *
diff --git a/python/tvm/topi/hexagon/slice_ops/add_subtract_multiply.py b/python/tvm/topi/hexagon/slice_ops/add_subtract_multiply.py
index 1995dac12eb2..b18757548387 100755
--- a/python/tvm/topi/hexagon/slice_ops/add_subtract_multiply.py
+++ b/python/tvm/topi/hexagon/slice_ops/add_subtract_multiply.py
@@ -14,7 +14,6 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-# pylint: disable=invalid-name, unused-variable, unused-argument, too-many-locals
 
 """Compute and schedule for add, multiply, subtract slice op
 
@@ -28,26 +27,32 @@
 from ..utils import get_layout_transform_fn
 
 
-def add_broadcast_compute(A, B):
+def add_broadcast_compute(input_a, input_b):
     """Call the add op from topi"""
-    return topi.add(A, B)
+    return topi.add(input_a, input_b)
 
 
-def subtract_broadcast_compute(A, B):
+def subtract_broadcast_compute(input_a, input_b):
     """Call the subtract op from topi"""
-    return topi.subtract(A, B)
+    return topi.subtract(input_a, input_b)
 
 
-def multiply_broadcast_compute(A, B):
+def multiply_broadcast_compute(input_a, input_b):
     """Call the multiply op from topi"""
-    return topi.multiply(A, B)
+    return topi.multiply(input_a, input_b)
 
 
-def STIR_broadcast_schedule(
-    M, A, B, output_layout: str, input_A_layout: str, input_B_layout: str, op_name: str
+def tir_broadcast_schedule(
+    out_m,
+    input_a,
+    input_b,
+    output_layout: str,
+    input_a_layout: str,
+    input_b_layout: str,
+    op_name: str,
 ):
     """Schedule for input and output layout nhwc-8h2w32c2w-2d considering broadcast"""
-    func = te.create_prim_func([A, B, M])
+    func = te.create_prim_func([input_a, input_b, out_m])
 
     s = tir.Schedule(func)
 
@@ -55,27 +60,27 @@ def STIR_broadcast_schedule(
 
     block = s.get_block(block_dict[op_name])
 
-    if input_A_layout == "nhwc-8h2w32c2w-2d":
-        input_A_transformed_layout = get_layout_transform_fn(input_A_layout)
-        s.transform_layout(block, buffer=("read", 0), index_map=input_A_transformed_layout)
+    if input_a_layout == "nhwc-8h2w32c2w-2d":
+        input_a_transformed_layout = get_layout_transform_fn(input_a_layout)
+        s.transform_layout(block, buffer=("read", 0), index_map=input_a_transformed_layout)
 
-    if input_B_layout == "nhwc-8h2w32c2w-2d":
-        input_B_transformed_layout = get_layout_transform_fn(input_B_layout)
-        s.transform_layout(block, buffer=("read", 1), index_map=input_B_transformed_layout)
+    if input_b_layout == "nhwc-8h2w32c2w-2d":
+        input_b_transformed_layout = get_layout_transform_fn(input_b_layout)
+        s.transform_layout(block, buffer=("read", 1), index_map=input_b_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])
+    h_o, h_i = s.split(h, [None, 8])
+    w_o, w_i = s.split(w, [None, 4])
+    c_o, c_i = s.split(c, [None, 32])
+    wio, wii = s.split(w_i, [None, 2])
 
-    s.reorder(n, ho, wo, co, hi, wio, ci, wii)
+    s.reorder(n, h_o, w_o, c_o, h_i, wio, c_i, wii)
 
-    fused = s.fuse(ci, wii)
+    fused = s.fuse(c_i, wii)
     s.vectorize(fused)
 
     return s
diff --git a/python/tvm/topi/hexagon/utils.py b/python/tvm/topi/hexagon/utils.py
index def86486dbe9..1a10f796485a 100755
--- a/python/tvm/topi/hexagon/utils.py
+++ b/python/tvm/topi/hexagon/utils.py
@@ -14,7 +14,7 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-# pylint: disable=invalid-name
+
 """Common hexagon specific utilities"""
 from tvm import te
 
diff --git a/tests/python/contrib/test_hexagon/infrastructure.py b/tests/python/contrib/test_hexagon/infrastructure.py
index 57a9dff8b424..d868679e19fd 100644
--- a/tests/python/contrib/test_hexagon/infrastructure.py
+++ b/tests/python/contrib/test_hexagon/infrastructure.py
@@ -14,7 +14,6 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-# pylint: disable=invalid-name
 
 """ Hexagon testing infrastructure """
 
diff --git a/tests/python/contrib/test_hexagon/test_add_subtract_multiply.py b/tests/python/contrib/test_hexagon/test_add_subtract_multiply.py
index ff0b9e921ca9..f9860368c396 100755
--- a/tests/python/contrib/test_hexagon/test_add_subtract_multiply.py
+++ b/tests/python/contrib/test_hexagon/test_add_subtract_multiply.py
@@ -14,7 +14,6 @@
 # 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
@@ -172,7 +171,7 @@ def test_transform(
         elif op_name == "multiply":
             M = sl.multiply_broadcast_compute(A, B)
 
-        tir_schedule = sl.STIR_broadcast_schedule(
+        tir_schedule = sl.tir_broadcast_schedule(
             M, A, B, output_layout, input_A_layout, input_B_layout, op_name
         )
         sch = tir_schedule.mod

From 3da8f073dfc5f32c134799be8d110e3965dd2767 Mon Sep 17 00:00:00 2001
From: trahman-quic <trahman@quicinc.com>
Date: Fri, 10 Jun 2022 11:21:36 -0500
Subject: [PATCH 08/11] Fix variable names

---
 .../slice_ops/add_subtract_multiply.py        |  1 +
 python/tvm/topi/hexagon/utils.py              | 35 ++++++++++++++-----
 2 files changed, 28 insertions(+), 8 deletions(-)

diff --git a/python/tvm/topi/hexagon/slice_ops/add_subtract_multiply.py b/python/tvm/topi/hexagon/slice_ops/add_subtract_multiply.py
index b18757548387..86b6adb997cb 100755
--- a/python/tvm/topi/hexagon/slice_ops/add_subtract_multiply.py
+++ b/python/tvm/topi/hexagon/slice_ops/add_subtract_multiply.py
@@ -14,6 +14,7 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+# pylint: disable=invalid-name
 
 """Compute and schedule for add, multiply, subtract slice op
 
diff --git a/python/tvm/topi/hexagon/utils.py b/python/tvm/topi/hexagon/utils.py
index 1a10f796485a..9cfbc57dafca 100755
--- a/python/tvm/topi/hexagon/utils.py
+++ b/python/tvm/topi/hexagon/utils.py
@@ -19,24 +19,43 @@
 from tvm import te
 
 
-def n11c_1024c_2d(n, h, w, c):
+def n11c_1024c_2d(batch, height, width, channel):
     """Return index map for n11c_1024 2d layout"""
-    return [n, h, w, c // 1024, te.AXIS_SEPARATOR, c % 1024]
+    return [batch, height, width, channel // 1024, te.AXIS_SEPARATOR, channel % 1024]
 
 
-def n11c_1024c_1d(n, h, w, c):
+def n11c_1024c_1d(batch, height, width, channel):
     """Return index map for n11c_1024 1d layout"""
-    return [n, h, w, c // 1024, c % 1024]
+    return [batch, height, width, channel // 1024, channel % 1024]
 
 
-def nhwc_8h2w32c2w_2d(n, h, w, c):
+def nhwc_8h2w32c2w_2d(batch, height, width, channel):
     """Return index map for nhwc_8h2w32c2w 2d layout"""
-    return [n, h // 8, w // 4, c // 32, te.AXIS_SEPARATOR, h % 8, (w % 4) // 2, c % 32, w % 2]
+    return [
+        batch,
+        height // 8,
+        width // 4,
+        channel // 32,
+        te.AXIS_SEPARATOR,
+        height % 8,
+        (width % 4) // 2,
+        channel % 32,
+        width % 2,
+    ]
 
 
-def nhwc_8h2w32c2w_1d(n, h, w, c):
+def nhwc_8h2w32c2w_1d(batch, height, width, channel):
     """Return index map for nhwc_8h2w32c2w 1d layout"""
-    return [n, h // 8, w // 4, c // 32, h % 8, (w % 4) // 2, c % 32, w % 2]
+    return [
+        batch,
+        height // 8,
+        width // 4,
+        channel // 32,
+        height % 8,
+        (width % 4) // 2,
+        channel % 32,
+        width % 2,
+    ]
 
 
 def get_layout_transform_fn(layout):

From 188fe441146732b2bf50bfaa28e56e56a522bc7f Mon Sep 17 00:00:00 2001
From: trahman-quic <trahman@quicinc.com>
Date: Fri, 10 Jun 2022 11:22:57 -0500
Subject: [PATCH 09/11] Move the test file to topi

---
 .../test_hexagon/{ => topi}/test_add_subtract_multiply.py     | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)
 rename tests/python/contrib/test_hexagon/{ => topi}/test_add_subtract_multiply.py (95%)

diff --git a/tests/python/contrib/test_hexagon/test_add_subtract_multiply.py b/tests/python/contrib/test_hexagon/topi/test_add_subtract_multiply.py
similarity index 95%
rename from tests/python/contrib/test_hexagon/test_add_subtract_multiply.py
rename to tests/python/contrib/test_hexagon/topi/test_add_subtract_multiply.py
index f9860368c396..fa2d9797a882 100755
--- a/tests/python/contrib/test_hexagon/test_add_subtract_multiply.py
+++ b/tests/python/contrib/test_hexagon/topi/test_add_subtract_multiply.py
@@ -25,7 +25,7 @@
 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
+from ..infrastructure import allocate_hexagon_array, transform_numpy
 
 
 @tvm.testing.fixture
@@ -161,7 +161,7 @@ def test_transform(
         input_B_layout,
         op_name,
     ):
-        target_hexagon = tvm.target.hexagon("v69")
+        target_hexagon = tvm.target.hexagon("v68")
         A = te.placeholder(input_shape_A, name="A", dtype=dtype)
         B = te.placeholder(input_shape_B, name="B", dtype=dtype)
         if op_name == "add":

From f29844a834236d2aae98736c0aea07117f320aaa Mon Sep 17 00:00:00 2001
From: trahman-quic <trahman@quicinc.com>
Date: Wed, 15 Jun 2022 14:25:13 -0500
Subject: [PATCH 10/11] Resolve conflict

---
 python/tvm/topi/hexagon/slice_ops/__init__.py | 20 ------
 python/tvm/topi/hexagon/utils.py              | 71 -------------------
 2 files changed, 91 deletions(-)
 delete mode 100755 python/tvm/topi/hexagon/slice_ops/__init__.py
 delete mode 100755 python/tvm/topi/hexagon/utils.py

diff --git a/python/tvm/topi/hexagon/slice_ops/__init__.py b/python/tvm/topi/hexagon/slice_ops/__init__.py
deleted file mode 100755
index 6671b8757137..000000000000
--- a/python/tvm/topi/hexagon/slice_ops/__init__.py
+++ /dev/null
@@ -1,20 +0,0 @@
-# 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.
-
-""" Computes and Schedules for Hexagon slice ops. """
-
-from .add_subtract_multiply import *
diff --git a/python/tvm/topi/hexagon/utils.py b/python/tvm/topi/hexagon/utils.py
deleted file mode 100755
index 9cfbc57dafca..000000000000
--- a/python/tvm/topi/hexagon/utils.py
+++ /dev/null
@@ -1,71 +0,0 @@
-# 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.
-
-"""Common hexagon specific utilities"""
-from tvm import te
-
-
-def n11c_1024c_2d(batch, height, width, channel):
-    """Return index map for n11c_1024 2d layout"""
-    return [batch, height, width, channel // 1024, te.AXIS_SEPARATOR, channel % 1024]
-
-
-def n11c_1024c_1d(batch, height, width, channel):
-    """Return index map for n11c_1024 1d layout"""
-    return [batch, height, width, channel // 1024, channel % 1024]
-
-
-def nhwc_8h2w32c2w_2d(batch, height, width, channel):
-    """Return index map for nhwc_8h2w32c2w 2d layout"""
-    return [
-        batch,
-        height // 8,
-        width // 4,
-        channel // 32,
-        te.AXIS_SEPARATOR,
-        height % 8,
-        (width % 4) // 2,
-        channel % 32,
-        width % 2,
-    ]
-
-
-def nhwc_8h2w32c2w_1d(batch, height, width, channel):
-    """Return index map for nhwc_8h2w32c2w 1d layout"""
-    return [
-        batch,
-        height // 8,
-        width // 4,
-        channel // 32,
-        height % 8,
-        (width % 4) // 2,
-        channel % 32,
-        width % 2,
-    ]
-
-
-def get_layout_transform_fn(layout):
-    """Return index map function as per the layout string"""
-    if layout == "nhwc-8h2w32c2w-2d":
-        return nhwc_8h2w32c2w_2d
-    if layout == "nhwc-8h2w32c2w-1d":
-        return nhwc_8h2w32c2w_1d
-    if layout == "n11c-1024c-2d":
-        return n11c_1024c_2d
-    if layout == "n11c-1024c-1d":
-        return n11c_1024c_1d
-    raise RuntimeError(f"Unexpected layout '{layout}'")

From 6bc4d4abdc6ae840a8da14e29ec9efdb8037994e Mon Sep 17 00:00:00 2001
From: trahman-quic <trahman@quicinc.com>
Date: Wed, 15 Jun 2022 14:28:23 -0500
Subject: [PATCH 11/11] Modify init

---
 python/tvm/topi/hexagon/slice_ops/__init__.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/python/tvm/topi/hexagon/slice_ops/__init__.py b/python/tvm/topi/hexagon/slice_ops/__init__.py
index b52d410676af..70531c629e4c 100644
--- a/python/tvm/topi/hexagon/slice_ops/__init__.py
+++ b/python/tvm/topi/hexagon/slice_ops/__init__.py
@@ -17,6 +17,5 @@
 
 """ Computes and Schedules for Hexagon slice ops. """
 
-# pylint: disable=wildcard-import
-
 from .avg_pool2d import avg_pool2d_compute, avg_pool2d_STIR_schedule
+from .add_subtract_multiply import *