| # 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 numpy as np |
| import pytest |
| from unittest import mock |
| |
| import tvm |
| import tvm.testing |
| from tvm import relay |
| from tvm.relay import transform |
| from tvm.contrib.target import coreml as _coreml |
| |
| requires_coremltools = tvm.testing.requires_package("coremltools") |
| |
| |
| def _has_xcode(): |
| try: |
| tvm.contrib.xcode.xcrun([]) |
| return True |
| except FileNotFoundError: |
| pass |
| |
| return False |
| |
| |
| def _create_graph(): |
| shape = (10, 10) |
| mod = tvm.IRModule() |
| |
| x = relay.var("x", shape=shape) |
| y = relay.var("y", shape=shape) |
| z = x + x |
| p = y * y |
| func = relay.Function([x, y], p - z) |
| mod["main"] = func |
| |
| return mod |
| |
| |
| def _create_graph_annotated(): |
| shape = (10, 10) |
| target = "coremlcompiler" |
| mod = tvm.IRModule() |
| |
| # function 0 |
| f0_i0 = relay.var(target + "_0_i0", shape=shape) |
| func0 = relay.Function([f0_i0], f0_i0 * f0_i0) |
| |
| func0 = func0.with_attr("Primitive", tvm.tir.IntImm("int32", 1)) |
| func0 = func0.with_attr("Inline", tvm.tir.IntImm("int32", 1)) |
| func0 = func0.with_attr("Compiler", target) |
| func0 = func0.with_attr("global_symbol", target + "_0") |
| gv0 = relay.GlobalVar(target + "_0") |
| mod[gv0] = func0 |
| |
| # function 2 |
| f2_i0 = relay.var(target + "_2_i0", shape=shape) |
| func2 = relay.Function([f2_i0], f2_i0 + f2_i0) |
| |
| func2 = func2.with_attr("Primitive", tvm.tir.IntImm("int32", 1)) |
| func2 = func2.with_attr("Inline", tvm.tir.IntImm("int32", 1)) |
| func2 = func2.with_attr("Compiler", target) |
| func2 = func2.with_attr("global_symbol", target + "_2") |
| gv2 = relay.GlobalVar(target + "_2") |
| mod[gv2] = func2 |
| mod = relay.transform.InferType()(mod) |
| |
| # body |
| x = relay.var("x", shape=shape) |
| y = relay.var("y", shape=shape) |
| func = relay.Function([x, y], gv0(y) - gv2(x)) |
| mod["main"] = func |
| mod = relay.transform.InferType()(mod) |
| |
| return mod |
| |
| |
| @pytest.mark.xfail( |
| reason="Currently failing test. See tracking issue https://github.com/apache/tvm/issues/8901" |
| ) |
| @tvm.testing.uses_gpu |
| @requires_coremltools |
| def test_annotate(): |
| mod = _create_graph() |
| mod = transform.AnnotateTarget("coremlcompiler")(mod) |
| mod = transform.PartitionGraph()(mod) |
| |
| expected = _create_graph_annotated() |
| assert tvm.ir.structural_equal(mod, expected, map_free_vars=True) |
| |
| |
| @pytest.mark.skipif(not _has_xcode(), reason="Xcode is not available") |
| @tvm.testing.uses_gpu |
| @requires_coremltools |
| def test_compile_and_run(): |
| dev = tvm.cpu() |
| target = "llvm" |
| tol = 1e-3 |
| |
| with relay.build_config(opt_level=3): |
| lib = relay.build(_create_graph_annotated(), target=target) |
| m = tvm.contrib.graph_executor.GraphModule(lib["default"](dev)) |
| |
| shape = (10, 10) |
| x_data = np.random.rand(*shape).astype("float32") |
| y_data = np.random.rand(*shape).astype("float32") |
| |
| m.set_input("x", x_data) |
| m.set_input("y", y_data) |
| m.run() |
| out = tvm.nd.empty(shape, device=dev) |
| out = m.get_output(0, out) |
| |
| expected = (y_data * y_data) - (x_data + x_data) |
| tvm.testing.assert_allclose(out.numpy(), expected, rtol=tol, atol=tol) |
| |
| |
| @mock.patch("tvm.contrib.coreml_runtime.create") |
| @mock.patch("tvm.contrib.xcode.compile_coreml") |
| def _construct_model(func, m1, m2): |
| mod = tvm.IRModule() |
| mod["main"] = func |
| mod = transform.AnnotateTarget("coremlcompiler")(mod) |
| mod = transform.PartitionGraph()(mod) |
| |
| fcompile = tvm._ffi.get_global_func("relay.ext.coremlcompiler") |
| |
| for var, func in mod.functions.items(): |
| if func.attrs and "Compiler" in func.attrs and func.attrs["Compiler"] == "coremlcompiler": |
| fcompile(func) |
| |
| |
| @tvm.testing.uses_gpu |
| @requires_coremltools |
| def test_add(): |
| shape = (10, 10) |
| x = relay.var("x", shape=shape) |
| y = x + x |
| func = relay.Function([x], y) |
| _construct_model(func) |
| |
| |
| @tvm.testing.uses_gpu |
| @requires_coremltools |
| def test_multiply(): |
| shape = (10, 10) |
| x = relay.var("x", shape=shape) |
| y = x * x |
| func = relay.Function([x], y) |
| _construct_model(func) |
| |
| |
| @tvm.testing.uses_gpu |
| @requires_coremltools |
| def test_clip(): |
| shape = (10, 10) |
| x = relay.var("x", shape=shape) |
| y = relay.clip(x, a_min=0.0, a_max=1.0) |
| func = relay.Function([x], y) |
| _construct_model(func) |
| |
| |
| @tvm.testing.uses_gpu |
| @requires_coremltools |
| def test_batch_flatten(): |
| shape = (10, 10, 10) |
| x = relay.var("x", shape=shape) |
| y = relay.nn.batch_flatten(x) |
| func = relay.Function([x], y) |
| _construct_model(func) |
| |
| |
| @tvm.testing.uses_gpu |
| @requires_coremltools |
| def test_expand_dims(): |
| shape = (10, 10) |
| x = relay.var("x", shape=shape) |
| y = relay.expand_dims(x, axis=0) |
| func = relay.Function([x], y) |
| _construct_model(func) |
| |
| y = relay.expand_dims(x, axis=-1) |
| func = relay.Function([x], y) |
| _construct_model(func) |
| |
| |
| @tvm.testing.uses_gpu |
| @requires_coremltools |
| def test_relu(): |
| shape = (10, 10) |
| x = relay.var("x", shape=shape) |
| y = relay.nn.relu(x) |
| func = relay.Function([x], y) |
| _construct_model(func) |
| |
| |
| @tvm.testing.uses_gpu |
| @requires_coremltools |
| def test_softmax(): |
| shape = (10, 10) |
| x = relay.var("x", shape=shape) |
| y = relay.nn.softmax(x, axis=1) |
| func = relay.Function([x], y) |
| _construct_model(func) |
| |
| |
| @tvm.testing.uses_gpu |
| @requires_coremltools |
| def test_conv2d(): |
| x = relay.var("x", shape=(1, 3, 224, 224)) |
| w = relay.const(np.zeros((16, 3, 3, 3), dtype="float32")) |
| y = relay.nn.conv2d(x, w, strides=[2, 2], padding=[1, 1, 1, 1], kernel_size=[3, 3]) |
| func = relay.Function([x], y) |
| _construct_model(func) |
| |
| |
| @tvm.testing.uses_gpu |
| @requires_coremltools |
| def test_global_avg_pool2d(): |
| shape = (10, 10, 10, 10) |
| x = relay.var("x", shape=shape) |
| y = relay.nn.global_avg_pool2d(x) |
| func = relay.Function([x], y) |
| _construct_model(func) |
| |
| |
| if __name__ == "__main__": |
| test_annotate() |
| test_compile_and_run() |
| test_add() |
| test_multiply() |
| test_clip() |
| test_expand_dims() |
| test_relu() |
| test_batch_flatten() |
| test_softmax() |
| test_conv2d() |
| test_global_avg_pool2d() |