wehub-resource-sync
497 lines
15 KiB
Python
497 lines
15 KiB
Python
# 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
|
|
|
|
import tvm
|
|
from tvm.contrib import coreml_runtime
|
|
from tvm.ir import Call, PrimType
|
|
from tvm.relax import transform
|
|
from tvm.relax.dpl.pattern import is_op, wildcard
|
|
from tvm.relax.expr import (
|
|
BindingBlock,
|
|
Constant,
|
|
Function,
|
|
SeqExpr,
|
|
Var,
|
|
VarBinding,
|
|
)
|
|
from tvm.relax.transform import PatternCheckContext
|
|
from tvm.relax.type import TensorType
|
|
from tvm.support.xcode import compile_coreml
|
|
|
|
from ...expr_functor import PyExprVisitor, visitor
|
|
from ..pattern_registry import get_patterns_with_prefix, register_patterns
|
|
from ..patterns import make_matmul_pattern
|
|
|
|
|
|
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"),
|
|
*default_unary_patterns(op_name="nn.batch_flatten"),
|
|
*conv2d_patterns(),
|
|
*clip_patterns(),
|
|
*matmul_patterns(),
|
|
]
|
|
)
|
|
|
|
|
|
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.CanonicalizeBindings()(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 Expr, 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 tvm.ir.is_prim_expr(arg):
|
|
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):
|
|
try:
|
|
import coremltools
|
|
from coremltools.models.neural_network import NeuralNetworkBuilder
|
|
except ImportError as err:
|
|
raise ImportError(
|
|
"coremltools is required by the CoreML backend. "
|
|
"Install it with: pip install coremltools"
|
|
) from err
|
|
|
|
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
|
|
ty = var.ty
|
|
if isinstance(ty, TensorType):
|
|
shape = [int(v) for v in list(ty.shape)]
|
|
elif isinstance(ty, PrimType):
|
|
shape = []
|
|
else:
|
|
raise Exception("Currently not supported: ", type(ty))
|
|
dtype = ty.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, f"{op_name} is not supported"
|
|
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.ty.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.ty.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_global_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 = f"{model_dir}/{name}.mlmodelc"
|
|
|
|
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
|