add hexagon schedule tests

Author: mehrdadh

python/tvm/script/tir/__init__.pyi

@@ -226,6 +226,7 @@ def alloc_buffer(
 """
 special_stmt - Reads/Writes
 """
+
 @overload
 def reads(read_regions: List[BufferSlice]) -> None: ...
 @overload
@@ -337,6 +338,7 @@ def Assert(condition: Union[PrimExpr, builtins.bool], message: str) -> PrimExpr:
 """
 Scope handler - Loops
 """
+
 @overload
 def serial(
     begin: Union[PrimExpr, int],

src/runtime/hexagon/rpc/hexagon/rpc_server.cc

@@ -41,7 +41,7 @@ extern "C" {
 #include "hexagon_rpc.h"
 
 // TODO(mehrdadh): make this configurable.
-#define TVM_HEXAGON_RPC_BUFF_SIZE_BYTES 2 * 1024 * 1024
+#define TVM_HEXAGON_RPC_BUFF_SIZE_BYTES 5 * 1024 * 1024
 
 // TODO(csulivan,adstraw,kparzysz-quic) This should be set on a TVM-wide basis.
 #if defined(__hexagon__)

tests/python/contrib/test_hexagon/test_batch_matmul.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.
+"""Test code for matmul"""
+import numpy as np
+import pytest
+import sys
+
+import tvm
+from tvm import topi
+from tvm import te
+import tvm.topi.testing
+from tvm.topi.utils import get_const_tuple
+
+from .conftest import requires_hexagon_toolchain
+
+dtype = tvm.testing.parameter(
+    "float32",
+    "float16",
+)
+
+x_batch, y_batch, M, N, K = tvm.testing.parameters(
+    (1, 1, 16, 16, 32),
+    (5, 5, 16, 16, 32),
+    (5, 5, 16, 20, 32),
+    (30, 30, 16, 20, 32),
+    # Test batch broadcasting.
+    (1, 5, 16, 16, 32),
+    (5, 1, 16, 16, 32),
+)
+
+# TODO(mehrdadh): add dynamic testing
+@requires_hexagon_toolchain
+def test_batch_matmul(hexagon_session, x_batch, y_batch, M, N, K, dtype):
+    if dtype == "float16":
+        pytest.xfail("float16 is not supported.")
+
+    x = te.placeholder((x_batch, M, K), name="x")
+    y = te.placeholder((y_batch, N, K), name="y")
+
+    def get_ref_data():
+        a_np = np.random.uniform(size=(x_batch, M, K)).astype(dtype)
+        b_np = np.random.uniform(size=(y_batch, N, K)).astype(dtype)
+        c_np = tvm.topi.testing.batch_matmul(a_np, b_np)
+        return (a_np, b_np, c_np)
+
+    # get the test data
+    a_np, b_np, c_np = get_ref_data()
+
+    target_hexagon = tvm.target.hexagon("v68")
+    with tvm.target.Target(target_hexagon):
+        fcompute = topi.nn.batch_matmul
+        fschedule = topi.hexagon.schedule_batch_matmul
+        out = fcompute(x, y)
+        s = fschedule([out])
+        out_shape = out.shape
+
+    func = tvm.build(
+        s, [x, y, out], tvm.target.Target(target_hexagon, host=target_hexagon), name="batch_matmul"
+    )
+    mod = hexagon_session.load_module(func)
+
+    dev = hexagon_session.device
+    a = tvm.nd.array(a_np, dev)
+    b = tvm.nd.array(b_np, dev)
+    c = tvm.nd.array(np.zeros(get_const_tuple(out_shape), dtype=dtype), dev)
+    mod["batch_matmul"](a, b, c)
+
+    tvm.testing.assert_allclose(c.numpy(), c_np, rtol=1e-5)
+
+
+x_batch_1, y_batch_1, M_1, N_1, K_1 = tvm.testing.parameters(
+    (1, 1, 2, 3, 1),
+    (1, 1, 16, 24, 32),
+    (5, 5, 24, 16, 32),
+    (30, 30, 16, 20, 32),
+    (1, 5, 16, 16, 32),
+    (5, 1, 16, 16, 32),
+)
+
+
+@requires_hexagon_toolchain
+def test_batch_matmul_int8(hexagon_session, x_batch_1, y_batch_1, M_1, N_1, K_1):
+    dtype = "int8"
+    out_dtype = "int8"
+    assert x_batch_1 == y_batch_1 or x_batch_1 == 1 or y_batch_1 == 1
+    x = te.placeholder((x_batch_1, M_1, K_1), name="x", dtype=dtype)
+    y = te.placeholder((y_batch_1, N_1, K_1), name="y", dtype=dtype)
+
+    def get_ref_data():
+        a_np = np.random.randint(low=-128, high=127, size=(x_batch_1, M_1, K_1)).astype(dtype)
+        b_np = np.random.randint(low=-128, high=127, size=(y_batch_1, N_1, K_1)).astype(dtype)
+        c_np = tvm.topi.testing.batch_matmul(a_np, b_np, out_dtype=out_dtype)
+        return (a_np, b_np, c_np)
+
+    # get the test data
+    a_np, b_np, c_np = get_ref_data()
+
+    target_hexagon = tvm.target.hexagon("v68")
+    with tvm.target.Target(target_hexagon):
+        fcompute = topi.nn.batch_matmul
+        fschedule = topi.hexagon.schedule_batch_matmul
+        out = fcompute(x, y)
+        s = fschedule([out])
+
+    func = tvm.build(
+        s,
+        [x, y, out],
+        tvm.target.Target(target_hexagon, host=target_hexagon),
+        name="batch_matmul_int8",
+    )
+    mod = hexagon_session.load_module(func)
+
+    dev = hexagon_session.device
+    a = tvm.nd.array(a_np, dev)
+    b = tvm.nd.array(b_np, dev)
+    c = tvm.nd.array(np.zeros(get_const_tuple(out.shape), dtype=out_dtype), dev)
+    mod["batch_matmul_int8"](a, b, c)
+    tvm.testing.assert_allclose(c.numpy(), c_np, rtol=1e-5)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(sys.argv))

tests/python/contrib/test_hexagon/test_dense.py

@@ -0,0 +1,112 @@
+# 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.
+"""Test code for dense"""
+import numpy as np
+import pytest
+import sys
+
+import tvm
+from tvm import topi
+from tvm import te
+import tvm.topi.testing
+from tvm.topi.utils import get_const_tuple
+
+from .conftest import requires_hexagon_toolchain
+
+random_seed = tvm.testing.parameter(0)
+
+use_bias = tvm.testing.parameter(True, False)
+
+# batch_size more than 8 would break
+batch_size = tvm.testing.parameter(1, 2, 8)
+
+in_dim, out_dim = tvm.testing.parameters((1024, 1000))
+
+in_dtype, out_dtype = tvm.testing.parameters(
+    ("float32", "float32"),
+    ("float16", "float32"),
+    ("int8", "int32"),
+)
+
+
+@tvm.testing.fixture(cache_return_value=True)
+def dense_ref_data(random_seed, batch_size, in_dim, out_dim, use_bias, in_dtype, out_dtype):
+    np.random.seed(random_seed)
+
+    if "float" in in_dtype:
+        a_np = np.random.uniform(size=(batch_size, in_dim)).astype(in_dtype)
+        b_np = np.random.uniform(size=(out_dim, in_dim)).astype(in_dtype)
+        c_np = np.random.uniform(size=(out_dim,)).astype(out_dtype)
+    elif in_dtype == "int8":
+        a_np = np.random.randint(low=-128, high=127, size=(batch_size, in_dim)).astype(in_dtype)
+        b_np = np.random.randint(low=-128, high=127, size=(out_dim, in_dim)).astype(in_dtype)
+        c_np = np.random.randint(low=-128, high=127, size=(out_dim,)).astype(out_dtype)
+    else:
+        raise ValueError("No method to generate test data for data type '{}'".format(in_dtype))
+
+    matmul = np.dot(a_np.astype(out_dtype), b_np.T.astype(out_dtype))
+
+    if use_bias:
+        matmul += c_np
+
+    d_np = np.maximum(matmul, 0)
+    return (a_np, b_np, c_np, d_np)
+
+
+@requires_hexagon_toolchain
+def test_dense(
+    hexagon_session, batch_size, in_dim, out_dim, use_bias, in_dtype, out_dtype, dense_ref_data
+):
+    if in_dtype == "float16":
+        pytest.xfail("float16 is not supported.")
+
+    if "int" in in_dtype:
+        tol = {"atol": 0, "rtol": 0}
+    elif in_dtype == "float32":
+        tol = {"rtol": 1e-5, "atol": 1e-5}
+
+    A = te.placeholder((batch_size, in_dim), name="A", dtype=in_dtype)
+    B = te.placeholder((out_dim, in_dim), name="B", dtype=in_dtype)
+    C = te.placeholder((out_dim,), name="C", dtype=out_dtype)
+
+    a_np, b_np, c_np, d_np = dense_ref_data
+
+    fcompute = topi.nn.dense
+    fschedule = topi.hexagon.schedule_dense
+
+    target_hexagon = tvm.target.hexagon("v68")
+    with tvm.target.Target(target_hexagon):
+        D = fcompute(A, B, C if use_bias else None, out_dtype)
+        D = topi.nn.relu(D)
+        s = fschedule([D])
+
+    func = tvm.build(
+        s, [A, B, C, D], tvm.target.Target(target_hexagon, host=target_hexagon), name="dense"
+    )
+    mod = hexagon_session.load_module(func)
+
+    dev = hexagon_session.device
+    a = tvm.nd.array(a_np, dev)
+    b = tvm.nd.array(b_np, dev)
+    c = tvm.nd.array(c_np, dev)
+    d = tvm.nd.array(np.zeros(get_const_tuple(D.shape), dtype=out_dtype), dev)
+    mod["dense"](a, b, c, d)
+    tvm.testing.assert_allclose(d.numpy(), d_np, **tol)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(sys.argv))

tests/python/contrib/test_hexagon/test_launcher.py

@@ -16,11 +16,9 @@
 # under the License.
 
 import os
-import pathlib
 import sys
 import pytest
 import numpy as np
-import logging
 
 import tvm.testing
 from tvm import te