[Hexagon] Make allocate_hexagon_array a hexagon contrib API (apache#1…

…3336)

Make 'allocate_hexagon_array' a hexagon contrib API

python/tvm/contrib/hexagon/tools.py

@@ -20,6 +20,7 @@
 import os
 import pathlib
 from typing import Union
+import numpy
 
 import tvm
 import tvm.contrib.cc as cc
@@ -203,3 +204,41 @@ def export_module(module, out_dir, binary_name="test_binary.so"):
     binary_path = pathlib.Path(out_dir) / binary_name
     module.save(str(binary_path))
     return binary_path
+
+
+def allocate_hexagon_array(
+    dev, tensor_shape=None, dtype=None, data=None, axis_separators=None, mem_scope=None
+):
+    """
+    Allocate a hexagon array which could be a 2D array
+    on physical memory defined by axis_separators
+    """
+    if tensor_shape is None:
+        assert data is not None, "Must provide either tensor shape or numpy data array"
+        tensor_shape = data.shape
+    elif data is not None:
+        assert (
+            tensor_shape == data.shape
+        ), "Mismatch between provided tensor shape and numpy data array shape"
+
+    if dtype is None:
+        assert data is not None, "Must provide either dtype or numpy data array"
+        dtype = data.dtype.name
+    elif data is not None:
+        assert dtype == data.dtype, "Mismatch between provided dtype and numpy data array dtype"
+
+    if axis_separators is None:
+        axis_separators = []
+
+    boundaries = [0, *axis_separators, len(tensor_shape)]
+    physical_shape = [
+        numpy.prod(tensor_shape[dim_i:dim_f])
+        for dim_i, dim_f in zip(boundaries[:-1], boundaries[1:])
+    ]
+
+    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))
+
+    return arr._create_view(tensor_shape)

tests/python/contrib/test_hexagon/infrastructure.py

@@ -23,44 +23,6 @@
 from tvm import te
 
 
-def allocate_hexagon_array(
-    dev, tensor_shape=None, dtype=None, data=None, axis_separators=None, mem_scope=None
-):
-    """
-    Allocate a hexagon array which could be a 2D array
-    on physical memory defined by axis_separators
-    """
-    if tensor_shape is None:
-        assert data is not None, "Must provide either tensor shape or numpy data array"
-        tensor_shape = data.shape
-    elif data is not None:
-        assert (
-            tensor_shape == data.shape
-        ), "Mismatch between provided tensor shape and numpy data array shape"
-
-    if dtype is None:
-        assert data is not None, "Must provide either dtype or numpy data array"
-        dtype = data.dtype.name
-    elif data is not None:
-        assert dtype == data.dtype, "Mismatch between provided dtype and numpy data array dtype"
-
-    if axis_separators is None:
-        axis_separators = []
-
-    boundaries = [0, *axis_separators, len(tensor_shape)]
-    physical_shape = [
-        numpy.prod(tensor_shape[dim_i:dim_f])
-        for dim_i, dim_f in zip(boundaries[:-1], boundaries[1:])
-    ]
-
-    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))
-
-    return arr._create_view(tensor_shape)
-
-
 def ceildiv(o, d):
     assert o >= 0
     assert d >= 0

tests/python/contrib/test_hexagon/test_2d_physical_buffers.py

@@ -31,8 +31,9 @@
 from tvm import te
 from tvm.contrib.hexagon.pytest_plugin import requires_hexagon_toolchain
 from tvm.tir.stmt_functor import post_order_visit
+from tvm.contrib.hexagon import allocate_hexagon_array
 
-from .infrastructure import allocate_hexagon_array, get_hexagon_target
+from .infrastructure import get_hexagon_target
 
 # Disabling invalid name as pylint assumes global variables as constants and
 # expects them to be all upper-case. Since these are used as

tests/python/contrib/test_hexagon/test_benchmark_maxpool2d.py

@@ -53,9 +53,9 @@
 from tvm import te, topi, tir
 from tvm.topi import testing
 from tvm.contrib.hexagon.session import Session
+from tvm.contrib.hexagon import allocate_hexagon_array
 
-
-from .infrastructure import allocate_hexagon_array, get_hexagon_target
+from .infrastructure import get_hexagon_target
 from . import benchmark_util as bu
 
 # Pytest seems to require that fixture names exist in the current module.

tests/python/contrib/test_hexagon/test_memory_alloc.py

@@ -20,8 +20,9 @@
 
 import tvm
 from tvm.script import tir as T
+from tvm.contrib.hexagon import allocate_hexagon_array
 
-from .infrastructure import allocate_hexagon_array, get_hexagon_target
+from .infrastructure import get_hexagon_target
 
 
 def generated_func(shape: tuple, dtype: str, axis_separators: list):

tests/python/contrib/test_hexagon/test_sigmoid.py

@@ -23,8 +23,9 @@
 from tvm import te
 from tvm import tir
 from tvm import topi
+from tvm.contrib.hexagon import allocate_hexagon_array
 
-from .infrastructure import allocate_hexagon_array, get_hexagon_target
+from .infrastructure import get_hexagon_target
 
 
 def sigmoid_compute(sigmoid_input):

tests/python/contrib/test_hexagon/topi/slice_op/test_argmax_slice.py

@@ -22,7 +22,9 @@
 from tvm import te
 import tvm.topi.hexagon.slice_ops as sl
 import tvm.contrib.hexagon
-from ...infrastructure import allocate_hexagon_array, transform_numpy, get_hexagon_target
+from tvm.contrib.hexagon import allocate_hexagon_array
+
+from ...infrastructure import transform_numpy, get_hexagon_target
 
 
 class TestArgMaxSlice:

tests/python/contrib/test_hexagon/topi/slice_op/test_avg_pool2d_slice.py

@@ -23,8 +23,8 @@
 from tvm.contrib.hexagon.session import Session
 import tvm.topi.hexagon.slice_ops as sl
 import tvm.topi.hexagon.qnn as qn
+from tvm.contrib.hexagon import allocate_hexagon_array
 from ...infrastructure import (
-    allocate_hexagon_array,
     transform_numpy,
     quantize_np,
     get_hexagon_target,

tests/python/contrib/test_hexagon/topi/slice_op/test_cast_slice.py

@@ -22,8 +22,9 @@
 import tvm.testing
 from tvm import te
 import tvm.topi.hexagon.slice_ops as sl
+from tvm.contrib.hexagon import allocate_hexagon_array
 
-from ...infrastructure import allocate_hexagon_array, transform_numpy, get_hexagon_target
+from ...infrastructure import transform_numpy, get_hexagon_target
 
 
 class TestCastF16F32Slice2d:

tests/python/contrib/test_hexagon/topi/slice_op/test_clip_slice.py

@@ -22,8 +22,9 @@
 from tvm import te
 import tvm.testing
 import tvm.topi.hexagon.slice_ops as sl
+from tvm.contrib.hexagon import allocate_hexagon_array
 
-from ...infrastructure import allocate_hexagon_array, transform_numpy, get_hexagon_target
+from ...infrastructure import transform_numpy, get_hexagon_target
 
 input_layout = tvm.testing.parameter(
     "nhwc-8h2w32c2w-2d",

tests/python/contrib/test_hexagon/topi/slice_op/test_conv2d_slice.py

@@ -24,8 +24,9 @@
 import tvm.testing
 from tvm.topi.hexagon.slice_ops.conv2d import conv2d_compute, conv2d_schedule
 from tvm.topi.testing import conv2d_nhwc_python
+from tvm.contrib.hexagon import allocate_hexagon_array
 
-from ...infrastructure import allocate_hexagon_array, transform_numpy, get_hexagon_target
+from ...infrastructure import transform_numpy, get_hexagon_target
 
 input_layout = tvm.testing.parameter(
     "nhwc-8h2w32c2w-2d",

tests/python/contrib/test_hexagon/topi/slice_op/test_depthwise_conv2d_slice.py

@@ -25,8 +25,9 @@
 import tvm.topi.hexagon.qnn as qn
 from tvm.topi.testing import depthwise_conv2d_python_nhwc
 from tvm.topi.hexagon.slice_ops.dwconv2d import dwconv2d_compute, dwconv2d_schedule
+from tvm.contrib.hexagon import allocate_hexagon_array
 
-from ...infrastructure import allocate_hexagon_array, transform_numpy, quantize_np
+from ...infrastructure import transform_numpy, quantize_np
 
 
 @tvm.testing.fixture

tests/python/contrib/test_hexagon/topi/slice_op/test_dequantize_slice.py

@@ -23,8 +23,8 @@
 import tvm.testing
 from tvm import te
 from tvm.topi.hexagon import qnn
+from tvm.contrib.hexagon import allocate_hexagon_array
 from ...infrastructure import (
-    allocate_hexagon_array,
     transform_numpy,
     quantize_np,
     get_hexagon_target,

tests/python/contrib/test_hexagon/topi/slice_op/test_max_pool2d_slice.py

@@ -22,8 +22,9 @@
 import tvm.testing
 from tvm.contrib.hexagon.session import Session
 import tvm.topi.hexagon.slice_ops as sl
+from tvm.contrib.hexagon import allocate_hexagon_array
 
-from ...infrastructure import allocate_hexagon_array, transform_numpy, get_hexagon_target
+from ...infrastructure import transform_numpy, get_hexagon_target
 from ...pytest_util import (
     get_multitest_ids,
     create_populated_numpy_ndarray,

tests/python/contrib/test_hexagon/topi/slice_op/test_relu_slice.py

@@ -21,8 +21,9 @@
 import tvm.testing
 from tvm.topi.hexagon.slice_ops.relu import relu_compute, relu_stir_schedule
 from tvm import te
+from tvm.contrib.hexagon import allocate_hexagon_array
 
-from ...infrastructure import allocate_hexagon_array, transform_numpy, get_hexagon_target
+from ...infrastructure import transform_numpy, get_hexagon_target
 
 
 @tvm.testing.fixture

tests/python/contrib/test_hexagon/topi/slice_op/test_softmax_slice.py

@@ -20,8 +20,7 @@
 from tvm import te
 from tvm.topi.testing import softmax_python
 import tvm.topi.hexagon.slice_ops as sl
-
-from ...infrastructure import allocate_hexagon_array
+from tvm.contrib.hexagon import allocate_hexagon_array
 
 
 def transform_numpy(arr_np, layout):

tests/python/contrib/test_hexagon/topi/slice_op/test_tanh_slice.py

@@ -22,7 +22,9 @@
 from tvm import te
 import tvm.topi.hexagon.slice_ops as sl
 import tvm.contrib.hexagon
-from ...infrastructure import allocate_hexagon_array, transform_numpy, get_hexagon_target
+from tvm.contrib.hexagon import allocate_hexagon_array
+
+from ...infrastructure import transform_numpy, get_hexagon_target
 
 # pylint: disable=invalid-name
 

tests/python/contrib/test_hexagon/topi/test_add_subtract_multiply.py

@@ -21,8 +21,8 @@
 from tvm import te
 import tvm.topi.hexagon.slice_ops as sl
 import tvm.topi.hexagon.qnn as qn
+from tvm.contrib.hexagon import allocate_hexagon_array
 from ..infrastructure import (
-    allocate_hexagon_array,
     transform_numpy,
     quantize_np,
     get_hexagon_target,

tests/python/contrib/test_hexagon/topi/test_depth_to_space.py

@@ -25,8 +25,9 @@
 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 tvm.contrib.hexagon import allocate_hexagon_array
 
-from ..infrastructure import allocate_hexagon_array, transform_numpy, get_hexagon_target
+from ..infrastructure import transform_numpy, get_hexagon_target
 
 
 class TestD2SSlice:

tests/python/contrib/test_hexagon/topi/test_quantize.py

@@ -20,8 +20,8 @@
 import tvm
 from tvm import te
 import tvm.topi.hexagon.qnn as s1
+from tvm.contrib.hexagon import allocate_hexagon_array
 from ..infrastructure import (
-    allocate_hexagon_array,
     transform_numpy,
     quantize_np,
     get_hexagon_target,

tests/python/contrib/test_hexagon/topi/test_reshape.py

@@ -21,8 +21,9 @@
 import tvm.testing
 import tvm.topi.hexagon.slice_ops as sl
 from tvm import te
+from tvm.contrib.hexagon import allocate_hexagon_array
 
-from ..infrastructure import allocate_hexagon_array, transform_numpy, get_hexagon_target
+from ..infrastructure import transform_numpy, get_hexagon_target
 
 BATCH_FLATTEN_FP16_TESTS = (
     ([1, 1, 1, 2048], [1, 2048], "nhwc-1024c-2d", "nc-1024-2d", "float16"),

tests/python/contrib/test_hexagon/topi/test_resize2d.py

@@ -22,7 +22,9 @@
 from tvm import te
 from tvm.topi.testing import resize2d_python
 import tvm.topi.hexagon as s1
-from ..infrastructure import allocate_hexagon_array, transform_numpy, get_hexagon_target
+from tvm.contrib.hexagon import allocate_hexagon_array
+
+from ..infrastructure import transform_numpy, get_hexagon_target
 
 
 class TestResize2d: