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 <sun...@umich.edu> 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]) + + +@pytest.mark.skip("`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__])