From 958168bf988763f820eb37730ade0f8b3d607d9d Mon Sep 17 00:00:00 2001 From: MORITA Kazutaka Date: Tue, 2 Jun 2020 04:17:40 +0900 Subject: [PATCH] move _register_coreml_op to python/tvm/relay/op/contrib/coreml.py --- python/tvm/contrib/target/coreml.py | 26 -------------- python/tvm/relay/op/contrib/__init__.py | 1 + python/tvm/relay/op/contrib/coreml.py | 45 +++++++++++++++++++++++++ 3 files changed, 46 insertions(+), 26 deletions(-) create mode 100644 python/tvm/relay/op/contrib/coreml.py diff --git a/python/tvm/contrib/target/coreml.py b/python/tvm/contrib/target/coreml.py index 5f9446d841667..e74457ee5378a 100644 --- a/python/tvm/contrib/target/coreml.py +++ b/python/tvm/contrib/target/coreml.py @@ -22,8 +22,6 @@ import tvm._ffi from ...relay.expr_functor import ExprVisitor -from ...relay.expr import Constant -from ...relay import op as _op from .. import xcode, coreml_runtime def _convert_add(builder, name, inputs, outputs, args, attrs): @@ -226,27 +224,3 @@ def coreml_compiler(ref): ctx = tvm.cpu(0) return coreml_runtime.create(model_dir, ctx).module - - -def _register_coreml_op(op_name): - """Register a function to check the given operator is supported by Core ML. - - Paramters - --------- - op_name : Str - The name of operator that will be registered. - - """ - def _check_supported(attrs, args): - if op_name == 'nn.conv2d': - if not isinstance(args[1], Constant): - return False - if attrs['kernel_layout'] not in ['HWIO', 'OIHW']: - return False - return True - - _op.register(op_name, "target.coremlcompiler", _check_supported) - - -for op in _convert_map: - _register_coreml_op(op) diff --git a/python/tvm/relay/op/contrib/__init__.py b/python/tvm/relay/op/contrib/__init__.py index 3a3f6d5aa3045..0e1b4b024a5aa 100644 --- a/python/tvm/relay/op/contrib/__init__.py +++ b/python/tvm/relay/op/contrib/__init__.py @@ -19,3 +19,4 @@ from .register import get_pattern_table, register_pattern_table from .dnnl import * +from .coreml import * diff --git a/python/tvm/relay/op/contrib/coreml.py b/python/tvm/relay/op/contrib/coreml.py new file mode 100644 index 0000000000000..c23e351eaf0d2 --- /dev/null +++ b/python/tvm/relay/op/contrib/coreml.py @@ -0,0 +1,45 @@ +# 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 +"""CoreML codegen supported operators.""" +from ... import op as _op +from ...expr import Constant +from ....contrib.target.coreml import _convert_map + + +def _register_coreml_op(op_name): + """Register a function to check the given operator is supported by Core ML. + + Paramters + --------- + op_name : Str + The name of operator that will be registered. + + """ + def _check_supported(attrs, args): + if op_name == 'nn.conv2d': + if not isinstance(args[1], Constant): + return False + if attrs['kernel_layout'] not in ['HWIO', 'OIHW']: + return False + return True + + _op.register(op_name, "target.coremlcompiler", _check_supported) + + +for op in _convert_map: + _register_coreml_op(op)