This is an automated email from the ASF dual-hosted git repository.
wuwei pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/unity by this push:
new 5808cea9af [Unity][BYOC] CoreML Scaffolding (#15556)
5808cea9af is described below
commit 5808cea9af4a6243f20fef8cd1e342e839a7f307
Author: Sunghyun Park <[email protected]>
AuthorDate: Wed Oct 25 10:27:22 2023 -0700
[Unity][BYOC] CoreML Scaffolding (#15556)
* scaffolding
* add comments
* wip
* refactor merge_composite to append the codegen name. This is helpful to
analyze the profiling results
* create tmp folder when it does not exist
* remove debug codes
* lint
* lint
* fix
* fix ci
* lint
* hide coreml dependency for ci
* lint
* lint
* ci bugfix
* fix
---
python/tvm/contrib/coreml_runtime.py | 1 +
python/tvm/relax/backend/contrib/coreml.py | 490 +++++++++++++++++++++++++++
src/runtime/contrib/coreml/coreml_runtime.mm | 2 +-
tests/python/relax/test_codegen_coreml.py | 294 ++++++++++++++++
4 files changed, 786 insertions(+), 1 deletion(-)
diff --git a/python/tvm/contrib/coreml_runtime.py
b/python/tvm/contrib/coreml_runtime.py
index b2555572ed..aa4f212799 100644
--- a/python/tvm/contrib/coreml_runtime.py
+++ b/python/tvm/contrib/coreml_runtime.py
@@ -42,6 +42,7 @@ def create(symbol, compiled_model_path, device):
fcreate = device._rpc_sess.get_function(runtime_func)
else:
fcreate = tvm._ffi.get_global_func(runtime_func)
+ assert fcreate, "Cannot find `tvm.coreml_runtime.create` function."
return CoreMLModule(fcreate(symbol, compiled_model_path))
diff --git a/python/tvm/relax/backend/contrib/coreml.py
b/python/tvm/relax/backend/contrib/coreml.py
new file mode 100644
index 0000000000..b5caa688f2
--- /dev/null
+++ b/python/tvm/relax/backend/contrib/coreml.py
@@ -0,0 +1,490 @@
+# 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-argument, import-outside-toplevel
+"""Pattern table and codegen for CoreML"""
+
+import os
+import shutil
+import tvm._ffi
+from tvm.contrib import coreml_runtime
+from tvm.contrib.xcode import compile_coreml
+
+import tvm
+from tvm.relax import transform
+from tvm.relax.struct_info import TensorStructInfo, PrimStructInfo
+from tvm.relax.expr import (
+ BindingBlock,
+ Call,
+ Function,
+ PrimValue,
+ SeqExpr,
+ Var,
+ VarBinding,
+ Constant,
+)
+from tvm.relax.dpl.pattern import is_op, wildcard
+from tvm.relax.transform import PatternCheckContext
+from ..pattern_registry import get_patterns_with_prefix, register_patterns
+from ..patterns import make_matmul_pattern
+from ...expr_functor import PyExprVisitor, visitor
+
+
+def _check_default(context: PatternCheckContext) -> bool:
+ return True
+
+
+def default_binary_patterns(op_name: str):
+ """
+ Returns a list of binary op patterns in coreML BYOC backend.
+ """
+
+ def _make_binary_pattern():
+ lhs = wildcard()
+ rhs = wildcard()
+ out = is_op("relax." + op_name)(lhs, rhs)
+ annotations = {"lhs": lhs, "rhs": rhs, "root": out}
+ return out, annotations
+
+ def _binary_pattern(pattern_name):
+ return (pattern_name, *_make_binary_pattern(), _check_default)
+
+ return [_binary_pattern("coreml." + op_name)]
+
+
+def default_unary_patterns(op_name: str):
+ """
+ Returns a list of unary op patterns in coreML BYOC backend.
+ """
+
+ def _make_unary_pattern():
+ lhs = wildcard()
+ out = is_op("relax." + op_name)(lhs)
+ annotations = {"lhs": lhs, "root": out}
+ return out, annotations
+
+ def _unary_pattern(pattern_name):
+ return (pattern_name, *_make_unary_pattern(), _check_default)
+
+ return [_unary_pattern("coreml." + op_name)]
+
+
+def conv2d_patterns():
+ """
+ Returns a list of conv2d patterns in coreML BYOC backend.
+ """
+
+ def _make_conv2d_pattern():
+ lhs = wildcard()
+ rhs = wildcard()
+ out = is_op("relax.nn.conv2d")(lhs, rhs)
+ annotations = {"lhs": lhs, "rhs": rhs, "root": out}
+ return out, annotations
+
+ def _conv2d_pattern(pattern_name):
+ return (pattern_name, *_make_conv2d_pattern(), _check_default)
+
+ return [_conv2d_pattern("coreml.nn.conv2d")]
+
+
+def matmul_patterns():
+ """
+ Returns a list of all matmul patterns in coreML BYOC backend.
+ """
+
+ def _matmul_pattern(pattern_name):
+ return (
+ pattern_name,
+ *make_matmul_pattern(),
+ _check_default,
+ )
+
+ return [_matmul_pattern("coreml.matmul")]
+
+
+def clip_patterns():
+ """
+ Returns a list of clip patterns in coreML BYOC backend.
+ """
+
+ def _make_clip_pattern():
+ arg0 = wildcard()
+ arg1 = wildcard()
+ arg2 = wildcard()
+ out = is_op("relax.clip")(arg0, arg1, arg2)
+ annotations = {"arg0": arg0, "arg1": arg1, "arg2": arg2, "root": out}
+ return out, annotations
+
+ def _conv2d_pattern(pattern_name):
+ return (pattern_name, *_make_clip_pattern(), _check_default)
+
+ return [_conv2d_pattern("coreml.clip")]
+
+
+register_patterns(
+ [
+ *default_binary_patterns(op_name="add"),
+ *default_binary_patterns(op_name="multiply"),
+ *default_unary_patterns(op_name="nn.softmax"),
+ *default_unary_patterns(op_name="nn.relu"),
+ *default_unary_patterns(op_name="expand_dims"),
+ *default_unary_patterns(op_name="nn.avg_pool2d"),
+ *conv2d_patterns(),
+ *clip_patterns(),
+ *matmul_patterns()
+ # TODO(@tvm-team): enable when relax op is implemented
+ # ("coreml.nn.batch_flatten",
is_op("relax.nn.batch_flatten")(wildcard())),
+ ]
+)
+
+
+def partition_for_coreml(mod):
+ """
+ Partition the input module into coreml-supported subgraphs.
+
+ Parameters
+ ----------
+ mod: tvm.IRModule
+ The IRModule to be partitioned.
+
+ Returns
+ -------
+ mod: tvm.IRModule
+ The resulting IRModule, containing partitioned subgraphs to be
+ offloaded to the coreml backend.
+ """
+
+ patterns = get_patterns_with_prefix("coreml")
+ mod = transform.FoldDataflowBlockOutput()(mod)
+ mod = transform.FuseOpsByPattern(patterns, bind_constants=True,
annotate_codegen=False)(mod)
+ mod = transform.MergeCompositeFunctions()(mod)
+ return mod
+
+
+# Codegen for coreml API reference:
+#
https://apple.github.io/coremltools/source/coremltools.models.neural_network.html
+def _convert_add(builder, name, inputs, outputs, args, attrs):
+ builder.add_elementwise(name=name, input_names=inputs,
output_name=outputs[0], mode="ADD")
+
+
+def _convert_multiply(builder, name, inputs, outputs, args, attrs):
+ builder.add_elementwise(name=name, input_names=inputs,
output_name=outputs[0], mode="MULTIPLY")
+
+
+def _convert_matmul(builder, name, inputs, outputs, args, attrs):
+ builder.add_batched_mat_mul(
+ name=name,
+ input_names=inputs,
+ output_name=outputs[0],
+ )
+
+
+def _convert_clip(builder, name, inputs, outputs, args, attrs):
+ builder.add_clip(
+ name=name,
+ input_name=inputs[0],
+ output_name=outputs[0],
+ min_value=inputs[1],
+ max_value=inputs[2],
+ )
+
+
+def _convert_batch_flatten(builder, name, inputs, outputs, args, attrs):
+ builder.add_flatten_to_2d(name=name, input_name=inputs[0],
output_name=outputs[0])
+
+
+def _convert_expand_dims(builder, name, inputs, outputs, args, attrs):
+ axes = [int(v) for v in attrs["axis"]]
+ builder.add_expand_dims(name=name, input_name=inputs[0],
output_name=outputs[0], axes=axes)
+
+
+def _convert_relu(builder, name, inputs, outputs, args, attrs):
+ builder.add_activation(
+ name=name, non_linearity="RELU", input_name=inputs[0],
output_name=outputs[0]
+ )
+
+
+def _convert_softmax(builder, name, inputs, outputs, args, attrs):
+ builder.add_softmax_nd(
+ name=name, input_name=inputs[0], output_name=outputs[0],
axis=int(attrs["axis"])
+ )
+
+
+def _convert_conv2d(builder, name, inputs, outputs, args, attrs):
+ weight = args[1].data.numpy()
+ oc, kc, kh, kw = weight.shape
+
+ builder.add_convolution(
+ name=name,
+ kernel_channels=kc,
+ output_channels=oc,
+ height=kh,
+ width=kw,
+ stride_height=int(attrs["strides"][0]),
+ stride_width=int(attrs["strides"][0]),
+ border_mode="valid",
+ groups=int(attrs["groups"]),
+ W=weight,
+ b=None,
+ has_bias=False,
+ input_name=inputs[0],
+ output_name=outputs[0],
+ dilation_factors=[int(v) for v in attrs["dilation"]],
+ padding_top=int(attrs["padding"][0]),
+ padding_bottom=int(attrs["padding"][2]),
+ padding_left=int(attrs["padding"][1]),
+ padding_right=int(attrs["padding"][3]),
+ )
+
+
+def _convert_avg_pool2d(builder, name, inputs, outputs, args, attrs):
+ builder.add_pooling(
+ name=name,
+ height=1,
+ width=1,
+ stride_height=1,
+ stride_width=1,
+ layer_type="AVERAGE",
+ padding_type="VALID",
+ input_name=inputs[0],
+ output_name=outputs[0],
+ )
+
+
+_convert_map = {
+ "add": _convert_add,
+ "multiply": _convert_multiply,
+ "matmul": _convert_matmul,
+ "clip": _convert_clip,
+ "expand_dims": _convert_expand_dims,
+ "nn.relu": _convert_relu,
+ # "nn.batch_flatten": _convert_batch_flatten,
+ "nn.softmax": _convert_softmax,
+ "nn.conv2d": _convert_conv2d,
+ "nn.avg_pool2d": _convert_avg_pool2d,
+}
+
+
+@visitor
+class CallNodeInfoCollector(PyExprVisitor):
+ """
+ Collect PrimValue, Constant and attributes in the inner function
+ """
+
+ def __init__(self, op_name):
+ self.primvals = []
+ self.attrs = []
+ self.consts = []
+ self.op_name = op_name
+
+ def visit_call_(self, call: Call) -> None:
+ self.attrs.append(call.attrs)
+ for arg in call.args:
+ if isinstance(arg, PrimValue):
+ self.primvals.append(arg)
+ if isinstance(arg, Constant):
+ self.consts.append(arg)
+
+ def collect(self, expr):
+ self.visit_expr(expr)
+ return self.primvals, self.attrs, self.consts
+
+
+@visitor
+class CodegenCoreML(PyExprVisitor):
+ """
+ A visitor to traverse subgraphs and build Core ML models.
+ """
+
+ def __init__(self, model_name, function):
+ import coremltools
+ from coremltools.models.neural_network import NeuralNetworkBuilder
+
+ self.model_name = model_name
+ self.function = function
+ self.out_map = {}
+ self.const_map = {} # (buffer name, object)
+ self.model_inputs_ = []
+ self.buf_idx_ = 0
+
+ getter = tvm.get_global_func("relax.analysis.get_var2val")
+ assert getter, "Cannot find `relax.analysis.get_var2val` function."
+
+ self.var2val = getter(function)
+ self.cur_binding_var = None
+
+ inputs = [
+ (
+ "",
+ coremltools.models.datatypes.Array(
+ 1,
+ ),
+ )
+ for _ in self.function.params
+ ]
+ outputs = [
+ (
+ "",
+ coremltools.models.datatypes.Array(
+ 1,
+ ),
+ )
+ ]
+ self.builder = NeuralNetworkBuilder(inputs, outputs,
disable_rank5_shape_mapping=True)
+
+ def visit_function_(self, op) -> None:
+ for var in op.params:
+ name = var.name_hint
+ sinfo = var.struct_info
+ if isinstance(sinfo, TensorStructInfo):
+ shape = [int(v) for v in list(sinfo.shape)]
+ elif isinstance(sinfo, PrimStructInfo):
+ shape = []
+ else:
+ raise Exception("Currently not supported: ", type(sinfo))
+ dtype = sinfo.dtype
+ self.model_inputs_.append((name, shape, dtype))
+
+ self.visit_expr(op.body)
+
+ def visit_var_(self, var):
+ self.out_map[var] = [var.name_hint]
+ prev_binding_var = self.cur_binding_var
+ self.cur_binding_var = var
+ if var in self.var2val:
+ self.visit_expr(self.var2val[var])
+ self.cur_binding_var = prev_binding_var
+
+ def visit_call_(self, call: Call) -> None:
+ assert isinstance(call.op, Var)
+ assert call.op in self.var2val
+ func = self.var2val[call.op]
+
+ assert "Composite" in func.attrs, "Only composite functions are
supported."
+ composite_name = func.attrs["Composite"]
+
+ # Get the op name and remove "relax." prefix.
+ op_name = composite_name[7:]
+
+ inputs = []
+ args = []
+ for arg in call.args:
+ args.append(arg)
+ super().visit_expr(arg)
+ for out in self.out_map[arg]:
+ inputs.append(out)
+
+ primvals, attrs, consts =
CallNodeInfoCollector(op_name).collect(func.body)
+ for arg in primvals:
+ args.append(arg)
+ inputs.append(arg.value.value)
+
+ for arg in consts:
+ output = "buf_" + str(self.buf_idx_)
+ self.builder.add_load_constant_nd(
+ name=output,
+ output_name=output,
+ constant_value=arg.data.numpy(),
+ shape=arg.data.shape,
+ )
+ self.buf_idx_ = self.buf_idx_ + 1
+ self.out_map[arg] = [output]
+ inputs.append(output)
+ args.append(arg)
+
+ layer_name = op_name + "_" + str(self.buf_idx_)
+
+ assert op_name in _convert_map, "{} is not supported".format(op_name)
+ outputs = ["buf_" + str(self.buf_idx_)]
+ _convert_map[op_name](self.builder, layer_name, inputs, outputs, args,
attrs[0])
+ self.buf_idx_ = self.buf_idx_ + 1
+ self.out_map[self.cur_binding_var] = outputs
+
+ def visit_var_binding_(self, binding: VarBinding) -> None:
+ # Visit var of the last binding
+ self.visit_expr(binding.var)
+
+ def visit_binding_block_(self, block: BindingBlock) -> None:
+ # We only visit the last VarBinding to retrieve
+ # target composite function
+ self.visit_binding(block.bindings[-1])
+
+ def visit_seq_expr_(self, op: SeqExpr) -> None:
+ for bb in op.blocks:
+ self.visit_binding_block_(bb)
+
+ def serialize(self, func: Function):
+ self.visit_expr(func)
+
+ def compile(self, out_dir):
+ """
+ Build a Core ML model and compile it with Xcode toolchain.
+ """
+ import coremltools
+ from coremltools.proto.Model_pb2 import ArrayFeatureType
+
+ FEATURE_TYPE_MAP = {
+ "float32": ArrayFeatureType.FLOAT32,
+ "float64": ArrayFeatureType.DOUBLE,
+ "int32": ArrayFeatureType.INT32,
+ }
+ input_names, input_dims, input_dtypes = zip(*self.model_inputs_)
+ self.builder.set_input(input_names, input_dims)
+
+ for i, dtype in enumerate(input_dtypes):
+ assert dtype in FEATURE_TYPE_MAP
+ input_desc = self.builder.spec.description.input
+ input_desc[i].type.multiArrayType.dataType =
FEATURE_TYPE_MAP[dtype]
+
+ output_dim = [int(n) for n in self.function.struct_info.ret.shape]
+
+ last_binding_var = self.function.body.blocks[0].bindings[-1].var
+ self.builder.set_output(self.out_map[last_binding_var], [output_dim])
+
+ for i, dtype in enumerate([self.function.struct_info.ret.dtype]):
+ assert dtype in FEATURE_TYPE_MAP
+ output_desc = self.builder.spec.description.output
+ output_desc[i].type.multiArrayType.dataType =
FEATURE_TYPE_MAP[dtype]
+
+ model = coremltools.models.MLModel(self.builder.spec)
+ compile_coreml(model, self.model_name, out_dir)
+
+
+@tvm._ffi.register_func("relax.ext.coreml")
+def coreml_compiler(funcs, options, constant_names):
+ """
+ Create a CoreML runtime from a Relax module.
+ """
+ compiled_funcs = []
+ for func in funcs:
+ assert isinstance(func, tvm.relax.Function)
+ model_dir = os.getcwd() + "/tmp/"
+ if not os.path.exists(model_dir):
+ os.mkdir(model_dir)
+
+ name = str(func.attrs.global_symbol)
+ builder = CodegenCoreML(name, func)
+ builder.serialize(func)
+
+ mlmodelc_path = "{}/{}.mlmodelc".format(model_dir, name)
+
+ if os.path.exists(mlmodelc_path):
+ shutil.rmtree(mlmodelc_path)
+
+ builder.compile(model_dir)
+ dev = tvm.cpu(0)
+ compiled_funcs.append(coreml_runtime.create(name, mlmodelc_path,
dev).module)
+ return compiled_funcs
diff --git a/src/runtime/contrib/coreml/coreml_runtime.mm
b/src/runtime/contrib/coreml/coreml_runtime.mm
index 0fac49a822..ab1372e228 100644
--- a/src/runtime/contrib/coreml/coreml_runtime.mm
+++ b/src/runtime/contrib/coreml/coreml_runtime.mm
@@ -157,7 +157,7 @@ PackedFunc CoreMLRuntime::GetFunction(const String& name,
const ObjectPtr<Object
ICHECK(args[i].type_code() == kTVMDLTensorHandle ||
args[i].type_code() == kTVMNDArrayHandle)
<< "Expect NDArray or DLTensor as inputs\n";
- if (args[i].type_code() == kTVMDLTensorHandle) {
+ if (args[i].type_code() == kTVMDLTensorHandle || args[i].type_code()
== kTVMNDArrayHandle) {
model_->SetInput([input_names[i] UTF8String], args[i]);
} else {
LOG(FATAL) << "Not implemented";
diff --git a/tests/python/relax/test_codegen_coreml.py
b/tests/python/relax/test_codegen_coreml.py
new file mode 100644
index 0000000000..ba8304bca7
--- /dev/null
+++ b/tests/python/relax/test_codegen_coreml.py
@@ -0,0 +1,294 @@
+# 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
+
+import tvm
+import tvm.testing
+from tvm import relax
+
+requires_coremltools = tvm.testing.requires_package("coremltools")
+target, dev = "llvm", tvm.cpu()
+
+
+def _has_xcode():
+ try:
+ tvm.contrib.xcode.xcrun([])
+ return True
+ except FileNotFoundError:
+ pass
+ return False
+
+
+pytestmark = pytest.mark.skipif(
+ not (requires_coremltools and _has_xcode()),
+ reason="coreml is not enabled.",
+)
+
+
+def verify(mod, inputs):
+ from tvm.relax.backend.contrib.coreml import partition_for_coreml
+
+ mod1 = partition_for_coreml(mod)
+ mod1 = relax.transform.RunCodegen()(mod1)
+ assert relax.analysis.well_formed(mod1)
+ assert mod1.attrs, "Should exist if offloaded successfully."
+ assert "external_mods" in mod1.attrs, "Should exist if offloaded
successfully."
+ mod1 = relax.transform.LegalizeOps()(mod1)
+ assert relax.analysis.well_formed(mod1)
+
+ ex1 = relax.build(mod1, target=target)
+ vm1 = relax.VirtualMachine(ex1, dev, profile=True)
+ out1 = vm1["main"](*inputs)
+
+ mod2 = relax.transform.LegalizeOps()(mod)
+ ex2 = relax.build(mod2, target=target)
+ vm2 = relax.VirtualMachine(ex2, dev, profile=True)
+ out2 = vm2["main"](*inputs)
+
+ tvm.testing.assert_allclose(out1.numpy(), out2.numpy(), rtol=1e-3,
atol=1e-3)
+
+
+def test_add():
+ x = relax.Var("x", relax.TensorStructInfo([10, 10], "float32"))
+ y = relax.Var("y", relax.TensorStructInfo([10, 10], "float32"))
+ bb = relax.BlockBuilder()
+ with bb.function("main", [x, y]):
+ with bb.dataflow():
+ lv0 = bb.emit(relax.op.add(x, y))
+ gv = bb.emit_output(lv0)
+ bb.emit_func_output(gv)
+ mod = bb.get()
+ x_data = tvm.nd.array(np.random.rand(10, 10).astype("float32"), dev)
+ y_data = tvm.nd.array(np.random.rand(10, 10).astype("float32"), dev)
+ verify(mod, [x_data, y_data])
+
+
+def test_add_const():
+ x = relax.Var("x", relax.TensorStructInfo([10, 10], "float32"))
+ y = relax.const(np.ones([10, 10]), "float32")
+ bb = relax.BlockBuilder()
+ with bb.function("main", [x]):
+ with bb.dataflow():
+ lv0 = bb.emit(relax.op.add(x, y))
+ gv = bb.emit_output(lv0)
+ bb.emit_func_output(gv)
+ mod = bb.get()
+ x_data = tvm.nd.array(np.random.rand(10, 10).astype("float32"), dev)
+ verify(mod, [x_data])
+
+
+def test_multiply():
+ x = relax.Var("x", relax.TensorStructInfo([10, 10], "float32"))
+ y = relax.Var("y", relax.TensorStructInfo([10, 10], "float32"))
+ bb = relax.BlockBuilder()
+ with bb.function("main", [x, y]):
+ with bb.dataflow():
+ lv0 = bb.emit(relax.op.multiply(x, y))
+ gv = bb.emit_output(lv0)
+ bb.emit_func_output(gv)
+ mod = bb.get()
+
+ x_data = tvm.nd.array(np.random.rand(10, 10).astype("float32"), dev)
+ y_data = tvm.nd.array(np.random.rand(10, 10).astype("float32"), dev)
+ verify(mod, [x_data, y_data])
+
+
+def test_matmul():
+ x = relax.Var("x", relax.TensorStructInfo([8, 10], "float32"))
+ y = relax.Constant(tvm.nd.array(np.random.rand(10, 8).astype("float32"),
dev))
+ bb = relax.BlockBuilder()
+ with bb.function("main", [x]):
+ with bb.dataflow():
+ lv0 = bb.emit(relax.op.matmul(x, y))
+ gv = bb.emit_output(lv0)
+ bb.emit_func_output(gv)
+ mod = bb.get()
+
+ x_data = tvm.nd.array(np.random.rand(8, 10).astype("float32"), dev)
+ verify(mod, [x_data])
+
+ x = relax.Var("x", relax.TensorStructInfo([8, 10], "float32"))
+ y = relax.Var("y", relax.TensorStructInfo([10, 8], "float32"))
+ bb = relax.BlockBuilder()
+ with bb.function("main", [x, y]):
+ with bb.dataflow():
+ lv0 = bb.emit(relax.op.matmul(x, y))
+ gv = bb.emit_output(lv0)
+ bb.emit_func_output(gv)
+ mod = bb.get()
+
+ x_data = tvm.nd.array(np.random.rand(8, 10).astype("float32"), dev)
+ y_data = tvm.nd.array(np.random.rand(10, 8).astype("float32"), dev)
+ verify(mod, [x_data, y_data])
+
+
+def test_clip():
+ x = relax.Var("x", relax.TensorStructInfo([10, 10], "float32"))
+ bb = relax.BlockBuilder()
+
+ with bb.function("main", [x]):
+ with bb.dataflow():
+ lv0 = bb.emit(relax.op.clip(x, 0, 4))
+ gv0 = bb.emit_output(lv0)
+ bb.emit_func_output(gv0)
+ mod = bb.get()
+
+ x_data = tvm.nd.array(np.random.rand(10, 10).astype("float32"), dev)
+ verify(mod, [x_data])
+
+ x = relax.Var("x", relax.TensorStructInfo([10, 10], "float32"))
+ bb = relax.BlockBuilder()
+
+ with bb.function("main", [x]):
+ with bb.dataflow():
+ lv0 = bb.emit(relax.op.clip(x, 0, 4))
+ lv1 = bb.emit(relax.op.clip(x, 1, 3))
+ gv0 = bb.emit_output(lv0)
+ gv1 = bb.emit_output(lv1)
+ bb.emit_func_output([gv0, gv1])
+
+ x_data = tvm.nd.array(np.random.rand(10, 10).astype("float32"), dev)
+ verify(mod, [x_data])
+
+
+def test_expand_dims():
+ def get_mod(axis):
+ x = relax.Var("x", relax.TensorStructInfo([10, 10], "float32"))
+ bb = relax.BlockBuilder()
+ with bb.function("main", [x]):
+ with bb.dataflow():
+ lv0 = bb.emit(relax.op.expand_dims(x, axis=axis))
+ gv = bb.emit_output(lv0)
+ bb.emit_func_output(gv)
+ return bb.get()
+
+ x_data = tvm.nd.array(np.random.rand(10, 10).astype("float32"), dev)
+ verify(get_mod(axis=0), [x_data])
+ verify(get_mod(axis=1), [x_data])
+
+
+def test_relu():
+ x = relax.Var("x", relax.TensorStructInfo([10, 10], "float32"))
+ bb = relax.BlockBuilder()
+ with bb.function("main", [x]):
+ with bb.dataflow():
+ lv0 = bb.emit(relax.op.nn.relu(x))
+ gv = bb.emit_output(lv0)
+ bb.emit_func_output(gv)
+ mod = bb.get()
+
+ x_data = tvm.nd.array(np.random.rand(10, 10).astype("float32"), dev)
+ verify(mod, [x_data])
+
+
[email protected]("`batch_flatten` is not implemented yet.")
+def test_batch_flatten():
+ x = relax.Var("x", relax.TensorStructInfo([10, 10, 10], "float32"))
+ bb = relax.BlockBuilder()
+ with bb.function("main", [x]):
+ with bb.dataflow():
+ lv0 = bb.emit(relax.op.nn.batch_flatten(x))
+ gv = bb.emit_output(lv0)
+ bb.emit_func_output(gv)
+ mod = bb.get()
+
+ x_data = tvm.nd.array(np.random.rand(10, 10, 10).astype("float32"), dev)
+ verify(mod, [x_data])
+
+
+@requires_coremltools
+def test_softmax():
+ x = relax.Var("x", relax.TensorStructInfo([10, 10], "float32"))
+ bb = relax.BlockBuilder()
+ with bb.function("main", [x]):
+ with bb.dataflow():
+ lv0 = bb.emit(relax.op.nn.softmax(x))
+ gv = bb.emit_output(lv0)
+ bb.emit_func_output(gv)
+ mod = bb.get()
+
+ x_data = tvm.nd.array(np.random.rand(10, 10).astype("float32"), dev)
+ verify(mod, [x_data])
+
+
+def test_conv2d():
+ x = relax.Var("x", relax.TensorStructInfo([1, 3, 224, 224], "float32"))
+ w = relax.const(np.zeros((16, 3, 3, 3), dtype="float32"))
+ bb = relax.BlockBuilder()
+ with bb.function("main", [x]):
+ with bb.dataflow():
+ lv0 = bb.emit(relax.op.nn.conv2d(x, w, strides=[2, 2], padding=[1,
1, 1, 1]))
+ gv = bb.emit_output(lv0)
+ bb.emit_func_output(gv)
+ mod = bb.get()
+ x_data = tvm.nd.array(np.random.rand(1, 3, 224, 224).astype("float32"),
dev)
+ verify(mod, [x_data])
+
+
+def test_global_avg_pool2d():
+ x = relax.Var("x", relax.TensorStructInfo([1, 1, 10, 10], "float32"))
+ bb = relax.BlockBuilder()
+ with bb.function("main", [x]):
+ with bb.dataflow():
+ lv0 = bb.emit(relax.op.nn.avg_pool2d(x))
+ gv = bb.emit_output(lv0)
+ bb.emit_func_output(gv)
+ mod = bb.get()
+ x_data = tvm.nd.array(np.random.rand(1, 1, 10, 10).astype("float32"), dev)
+ verify(mod, [x_data])
+
+
+def test_subgraph1():
+ x = relax.Var("x", relax.TensorStructInfo([10, 10], "float32"))
+ y = relax.Var("y", relax.TensorStructInfo([10, 10], "float32"))
+ bb = relax.BlockBuilder()
+ with bb.function("main", [x, y]):
+ with bb.dataflow():
+ lv0 = bb.emit(relax.op.multiply(x, y))
+ lv1 = bb.emit(relax.op.nn.softmax(lv0))
+ gv = bb.emit_output(lv1)
+ bb.emit_func_output(gv)
+ mod = bb.get()
+ x_data = tvm.nd.array(np.random.rand(10, 10).astype("float32"), dev)
+ y_data = tvm.nd.array(np.random.rand(10, 10).astype("float32"), dev)
+ verify(mod, [x_data, y_data])
+
+
+def test_subgraph2():
+ x = relax.Var("x", relax.TensorStructInfo([10, 10], "float32"))
+ y = relax.Var("y", relax.TensorStructInfo([10, 10], "float32"))
+ bb = relax.BlockBuilder()
+ with bb.function("main", [x, y]):
+ with bb.dataflow():
+ # multiply+relu will be offloaded to coreml
+ lv0 = bb.emit(relax.op.multiply(x, y))
+ lv1 = bb.emit(relax.op.nn.relu(lv0))
+ # gelu wouldn't be offloaded to coreml
+ lv2 = bb.emit(relax.op.nn.gelu(lv1))
+ # relu would be offloaded to coreml
+ lv3 = bb.emit(relax.op.nn.relu(lv2))
+ gv = bb.emit_output(lv3)
+ bb.emit_func_output(gv)
+ mod = bb.get()
+ x_data = tvm.nd.array(np.random.rand(10, 10).astype("float32"), dev)
+ y_data = tvm.nd.array(np.random.rand(10, 10).astype("float32"), dev)
+ verify(mod, [x_data, y_data])
+
+
+if __name__ == "__main__":
+ pytest.main([__file__])
