wehub/apache--tvm
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Lint / lint (push) Has been cancelled
Details
CI / MacOS (push) Has been cancelled
Details
CI / Windows (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/tvm/relax/backend/metal/__init__.py
+17
View File
@@ -0,0 +1,17 @@
# 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.
"""The Relax Metal backend compilation pipeline and other passes."""
python/tvm/relax/backend/metal/coreml.py
+496
View File
@@ -0,0 +1,496 @@
# 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