chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,361 @@
|
||||
# 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.
|
||||
# ruff: noqa: F841
|
||||
"""NNAPI integration operator tests."""
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.script
|
||||
import tvm.script.relax as R
|
||||
import tvm.script.tirx as T
|
||||
from test_nnapi.conftest import remote
|
||||
from test_nnapi.infrastructure import build_and_run
|
||||
|
||||
|
||||
def _build_and_run_network(remote_obj, tracker, mod, input_data):
|
||||
"""Helper function to build and run a network."""
|
||||
|
||||
def execute_on_host(mod, inputs):
|
||||
with tvm.transform.PassContext(opt_level=3):
|
||||
ex = tvm.compile(mod, target="llvm")
|
||||
dev = tvm.cpu(0)
|
||||
vm = tvm.relax.VirtualMachine(ex, device=dev)
|
||||
output = vm["main"](*inputs)
|
||||
return output.numpy()
|
||||
|
||||
outputs = []
|
||||
for nnapi in [True, False]:
|
||||
if nnapi:
|
||||
outputs.append(
|
||||
build_and_run(
|
||||
remote_obj,
|
||||
tracker,
|
||||
mod,
|
||||
input_data,
|
||||
enable_nnapi=nnapi,
|
||||
)
|
||||
)
|
||||
else:
|
||||
outputs.append(execute_on_host(mod, input_data))
|
||||
return outputs
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"op",
|
||||
[
|
||||
R.exp,
|
||||
R.log,
|
||||
R.negative,
|
||||
R.sqrt,
|
||||
R.rsqrt,
|
||||
R.floor,
|
||||
R.nn.relu,
|
||||
R.nn.softmax,
|
||||
R.sigmoid,
|
||||
R.tanh,
|
||||
R.abs,
|
||||
],
|
||||
)
|
||||
def test_unary(op, input_shape=(1, 2, 8, 5)):
|
||||
remote_obj, tracker = remote()
|
||||
|
||||
def create_model() -> tvm.IRModule:
|
||||
@tvm.script.ir_module
|
||||
class Module:
|
||||
@R.function
|
||||
def main(i0: R.Tensor((1, 2, 8, 5), "float32")) -> R.Tensor((1, 2, 8, 5), "float32"):
|
||||
with R.dataflow():
|
||||
t0 = op(i0)
|
||||
R.output(t0)
|
||||
return t0
|
||||
|
||||
return Module
|
||||
|
||||
mod = create_model()
|
||||
verify(
|
||||
remote_obj,
|
||||
tracker,
|
||||
mod,
|
||||
inputs=[np.random.uniform(size=(1, 2, 8, 5)).astype("float32")],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"op",
|
||||
[
|
||||
R.power,
|
||||
R.greater,
|
||||
R.add,
|
||||
R.multiply,
|
||||
R.subtract,
|
||||
R.equal,
|
||||
R.less,
|
||||
R.less_equal,
|
||||
R.not_equal,
|
||||
R.maximum,
|
||||
R.minimum,
|
||||
R.greater_equal,
|
||||
],
|
||||
)
|
||||
def test_elementwise_binary(op, input_shape=(1, 2, 8, 5)):
|
||||
remote_obj, tracker = remote()
|
||||
|
||||
def create_model() -> tvm.IRModule:
|
||||
@tvm.script.ir_module
|
||||
class Module:
|
||||
@R.function
|
||||
def main(
|
||||
i0: R.Tensor((1, 2, 8, 5), "float32"),
|
||||
i1: R.Tensor((1, 2, 8, 5), "float32"),
|
||||
) -> R.Tensor((1, 2, 8, 5), "float32"):
|
||||
with R.dataflow():
|
||||
t0 = op(i0, i1)
|
||||
R.output(t0)
|
||||
return t0
|
||||
|
||||
return Module
|
||||
|
||||
mod = create_model()
|
||||
verify(
|
||||
remote_obj,
|
||||
tracker,
|
||||
mod,
|
||||
inputs=[
|
||||
np.random.uniform(size=input_shape).astype("float32"),
|
||||
np.random.uniform(size=input_shape).astype("float32"),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def test_divide(input_shape=(1, 2, 8, 5)):
|
||||
remote_obj, tracker = remote()
|
||||
|
||||
def create_model(input_shape) -> tvm.IRModule:
|
||||
@tvm.script.ir_module
|
||||
class Module:
|
||||
@R.function
|
||||
def main(
|
||||
i0: R.Tensor((1, 2, 8, 5), "float32"),
|
||||
i1: R.Tensor((1, 2, 8, 5), "float32"),
|
||||
) -> R.Tensor((1, 2, 8, 5), "float32"):
|
||||
with R.dataflow():
|
||||
t0 = R.divide(i0, i1)
|
||||
R.output(t0)
|
||||
return t0
|
||||
|
||||
return Module
|
||||
|
||||
mod = create_model(input_shape)
|
||||
verify(
|
||||
remote_obj,
|
||||
tracker,
|
||||
mod,
|
||||
inputs=[
|
||||
np.random.uniform(size=input_shape).astype("float32"),
|
||||
np.random.uniform(size=input_shape).astype("float32") + np.ones(input_shape, "float32"),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def test_matmul():
|
||||
remote_obj, tracker = remote()
|
||||
|
||||
def create_model() -> tvm.IRModule:
|
||||
@tvm.script.ir_module
|
||||
class Module:
|
||||
@R.function
|
||||
def main(
|
||||
i0: R.Tensor((5, 3, 4), "float32"),
|
||||
i1: R.Tensor((5, 4, 8), "float32"),
|
||||
) -> R.Tensor((5, 3, 8), "float32"):
|
||||
with R.dataflow():
|
||||
t0 = R.matmul(i0, i1)
|
||||
R.output(t0)
|
||||
return t0
|
||||
|
||||
return Module
|
||||
|
||||
mod = create_model()
|
||||
verify(
|
||||
remote_obj,
|
||||
tracker,
|
||||
mod,
|
||||
inputs=[
|
||||
np.random.random(size=(5, 3, 4)).astype("float32"),
|
||||
np.random.random(size=(5, 4, 8)).astype("float32"),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def test_permute_dims():
|
||||
remote_obj, tracker = remote()
|
||||
|
||||
def create_model() -> tvm.IRModule:
|
||||
@tvm.script.ir_module
|
||||
class Module:
|
||||
@R.function
|
||||
def main(
|
||||
i0: R.Tensor((5, 4, 8), "float32"),
|
||||
) -> R.Tensor((8, 5, 4), "float32"):
|
||||
with R.dataflow():
|
||||
t0 = R.permute_dims(i0, axes=[2, 0, 1])
|
||||
R.output(t0)
|
||||
return t0
|
||||
|
||||
return Module
|
||||
|
||||
mod = create_model()
|
||||
verify(
|
||||
remote_obj,
|
||||
tracker,
|
||||
mod,
|
||||
inputs=[
|
||||
np.random.random(size=(5, 4, 8)).astype("float32"),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def test_astype():
|
||||
remote_obj, tracker = remote()
|
||||
|
||||
def create_model() -> tvm.IRModule:
|
||||
@tvm.script.ir_module
|
||||
class Module:
|
||||
@R.function
|
||||
def main(
|
||||
i0: R.Tensor((8, 10, 15), "float32"),
|
||||
) -> R.Tensor((8, 10, 15), "float16"):
|
||||
with R.dataflow():
|
||||
t0: R.Tensor((8, 10, 15), "float16") = R.astype(i0, dtype="float16")
|
||||
R.output(t0)
|
||||
return t0
|
||||
|
||||
return Module
|
||||
|
||||
mod = create_model()
|
||||
verify(
|
||||
remote_obj,
|
||||
tracker,
|
||||
mod,
|
||||
inputs=[
|
||||
tvm.runtime.tensor(np.random.uniform(size=(8, 10, 15)).astype("float32")),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def test_mean():
|
||||
remote_obj, tracker = remote()
|
||||
|
||||
def create_model() -> tvm.IRModule:
|
||||
@tvm.script.ir_module
|
||||
class Module:
|
||||
@R.function
|
||||
def main(
|
||||
i0: R.Tensor((1, 10, 15), "float32"),
|
||||
) -> R.Tensor((1, 10, 1), "float32"):
|
||||
n = T.int64()
|
||||
with R.dataflow():
|
||||
t0: R.Tensor((1, 10, 1), "float32") = R.mean(i0, axis=[-1], keepdims=True)
|
||||
R.output(t0)
|
||||
return t0
|
||||
|
||||
return Module
|
||||
|
||||
mod = create_model()
|
||||
verify(
|
||||
remote_obj,
|
||||
tracker,
|
||||
mod,
|
||||
inputs=[
|
||||
tvm.runtime.tensor(np.random.uniform(size=(1, 10, 15)).astype("float32")),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def test_conv2d():
|
||||
remote_obj, tracker = remote()
|
||||
|
||||
def create_model() -> tvm.IRModule:
|
||||
@tvm.script.ir_module
|
||||
class Module:
|
||||
@R.function
|
||||
def main(
|
||||
i0: R.Tensor((1, 3, 224, 224), "float32"),
|
||||
i1: R.Tensor((64, 3, 3, 3), "float32"),
|
||||
i2: R.Tensor((1, 64, 1, 1), "float32"),
|
||||
):
|
||||
with R.dataflow():
|
||||
t0 = R.nn.conv2d(i0, i1, strides=(1, 1), padding=(1, 1))
|
||||
t0 = R.add(i2, t0)
|
||||
R.output(t0)
|
||||
return t0
|
||||
|
||||
return Module
|
||||
|
||||
mod = create_model()
|
||||
verify(
|
||||
remote_obj,
|
||||
tracker,
|
||||
mod,
|
||||
inputs=[
|
||||
np.random.random(size=(1, 3, 224, 224)).astype("float32"),
|
||||
np.random.random(size=(64, 3, 3, 3)).astype("float32"),
|
||||
np.random.random(size=(1, 64, 1, 1)).astype("float32"),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def test_max_pool2d():
|
||||
remote_obj, tracker = remote()
|
||||
|
||||
def create_model() -> tvm.IRModule:
|
||||
@tvm.script.ir_module
|
||||
class Module:
|
||||
@R.function
|
||||
def main(
|
||||
i0: R.Tensor((1, 1, 28, 28), "float32"),
|
||||
):
|
||||
with R.dataflow():
|
||||
t0 = R.nn.max_pool2d(i0, pool_size=(1, 1), strides=(1, 1), padding=(0, 0))
|
||||
R.output(t0)
|
||||
return t0
|
||||
|
||||
return Module
|
||||
|
||||
mod = create_model()
|
||||
verify(
|
||||
remote_obj,
|
||||
tracker,
|
||||
mod,
|
||||
inputs=[
|
||||
np.random.random(size=(1, 1, 28, 28)).astype("float32"),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def verify(remote_obj, tracker, mod, inputs):
|
||||
inputs_tvm: list[tvm.runtime.Tensor] = [tvm.runtime.tensor(v) for v in inputs]
|
||||
outputs = _build_and_run_network(remote_obj, tracker, mod, inputs_tvm)
|
||||
nnapi_out = outputs[0]
|
||||
expected_out = outputs[1]
|
||||
tvm.testing.assert_allclose(nnapi_out, expected_out, rtol=1e-4, atol=1e-5)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
Reference in New Issue
Block a user