diff --git a/python/tvm/contrib/target/android_nnapi/__init__.py b/python/tvm/contrib/target/android_nnapi/__init__.py new file mode 100644 index 000000000000..358851f55b10 --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/__init__.py @@ -0,0 +1,112 @@ +# 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. +"""BYOC External Compiler Implementation for Android NNAPI target.""" +import tvm +from .compiler import Compiler + + +def _get_c_type(tipe): + """Get matching C type for Relay types.""" + dtype = str(tipe.dtype) + if dtype == "float32": + return "float" + if dtype == "float16": + return "uint16_t" + if dtype == "int32": + return "int32_t" + assert dtype == "int64", f"{dtype} is unsupported" + return "int64_t" + + +@tvm.register_func("relay.ext.android_nnapi") +def _codegen(func): + """Codegen Relay IR to Android NNAPI. + + Parameters + ---------- + func: tvm.relay.Function + The Relay IR function to be codegened. + + Returns + ------- + mod: runtime.CSourceModule + The resulting Android NNAPI in C++ source code. + + Notes + ----- + Certain function attributes should be configured: + + * func.attrs.NnapiTargetVersion: (int) The targeting API level of Android. + """ + assert isinstance(func, tvm.relay.Function), "Only Function can be codegened to Android NNAPI" + code = """#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { +""" + + sid = str(func.attrs.global_symbol) + class_name = sid + "_class" + options = { + "class": { + "self": { + "name": class_name, + }, + }, + "target": { + "api_level": int(func.attrs.NnapiTargetVersion), + }, + } + code += Compiler(options).codegen(func) + code += "\n" + + instance_name = sid + "_model" + code += f" {class_name} {instance_name};\n" + + sid_impl_name = sid + "_" + code += f" void {sid_impl_name}" + code += "(::tvm::runtime::TVMArgs args, ::tvm::runtime::TVMRetValue *rv) {\n" + code += f" CHECK_EQ(args.num_args, {len(func.params) + 1})" + code += f'<< "num_args is expected to be {len(func.params) + 1}";\n' + code += f" {instance_name}.execute(\n" + for i, p in enumerate(func.params): + assert isinstance( + p.checked_type, tvm.relay.TensorType + ), "Function parameter is expected to be a tensor" + code += f" reinterpret_cast< {_get_c_type(p.checked_type)}* >" + code += f"(args[{i}].operator DLTensor*()->data), \n" + assert isinstance( + func.body.checked_type, tvm.relay.TensorType + ), "Function output is expected to be a tensor" + code += f" reinterpret_cast< {_get_c_type(func.body.checked_type)}* >" + code += f"(args[{len(func.params)}].operator DLTensor*()->data)\n" + code += f" );\n" + code += " *rv = 0;\n" + code += f" }} // {sid_impl_name}\n" + code += "} // anonymous namespace\n" + code += f"TVM_DLL_EXPORT_PACKED_FUNC({sid}, {sid_impl_name});\n" + + return tvm.get_global_func("runtime.CSourceModuleCreate")(code, "c", [sid], []) diff --git a/python/tvm/contrib/target/android_nnapi/_export_object/__init__.py b/python/tvm/contrib/target/android_nnapi/_export_object/__init__.py new file mode 100644 index 000000000000..a2ef2ca816eb --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/_export_object/__init__.py @@ -0,0 +1,18 @@ +# 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. +"""Internal namespaces of ExportObject.""" +from .json_analyzer import JSONAnalyzer diff --git a/python/tvm/contrib/target/android_nnapi/_export_object/json_analyzer.py b/python/tvm/contrib/target/android_nnapi/_export_object/json_analyzer.py new file mode 100644 index 000000000000..a61dd4574ce0 --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/_export_object/json_analyzer.py @@ -0,0 +1,147 @@ +# 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. +"""Namespace for methods that analyzes the exported JSON.""" + + +class JSONAnalyzer: + """Analyzing methods of the JSON format of Android NNAPI model.""" + + class _Operand: + """Android NNAPI Operand-related analyzing methods on the exported JSON.""" + + def __init__(self, export_json): + self._export_json = export_json + + def get_dtype(self, idx): + """Get operand dtype. + + Parameters + ---------- + idx: int + operand to be queried. + + Returns + ------- + dtype: str + dtype of the queried operand. + """ + return self._export_json["types"][self._export_json["operands"][idx]["type"]]["type"] + + def get_shape(self, idx): + """Get operand shape. + + Parameters + ---------- + idx: int + operand to be queried. + + Returns + ------- + shape: tuple of int or None + shape of the queried operand. None if operand has no shape. + """ + return self._export_json["types"][self._export_json["operands"][idx]["type"]].get( + "shape", None + ) + + def get_rank(self, idx): + """Get operand rank. + + Parameters + ---------- + idx: int + operand to be queried. + + Returns + ------- + rank: int + rank of the queried operand. + """ + shape = self.get_shape(idx) + if shape is None: + return 0 + return len(shape) + + def get_value(self, idx): + """Get operand value. + + Parameters + ---------- + idx: int + operand to be queried. + + Returns + ------- + value: + value of the queried operand. None if there's no value. + """ + value_dict = self._export_json["operands"][idx].get("value", None) + if value_dict is None: + return None + + if value_dict["type"] == "constant_idx": + return self._export_json["constants"][value_dict["value"]]["value"] + assert value_dict["type"] == "memory_ptr" + return value_dict["value"] + + def get_constant(self, idx): + """Get operand constant. + + Parameters + ---------- + idx: int + operand to be queried. + + Returns + ------- + obj: dict + constant dict of the queried operand. None if there's no value. + """ + value_dict = self._export_json["operands"][idx].get("value", None) + if value_dict is None or value_dict["type"] != "constant_idx": + return None + return self._export_json["constants"][value_dict["value"]] + + def is_fuse_code(self, idx): + """Check whether the operand pointed by idx is a FuseCode + + Parameters + ---------- + idx: int + the index of the queried operand. + + Returns + ------- + b: bool + the queried operand is a FuseCode or not. + """ + dtype = self.get_dtype(idx) + if dtype != "INT32": + return False + shape = self.get_shape(idx) + if shape is not None: + return False + value = self.get_value(idx) + return value in { + "ANEURALNETWORKS_FUSED_NONE", + "ANEURALNETWORKS_FUSED_RELU", + "ANEURALNETWORKS_FUSED_RELU1", + "ANEURALNETWORKS_FUSED_RELU6", + } + + def __init__(self, export_json): + self.operand = JSONAnalyzer._Operand(export_json) diff --git a/python/tvm/contrib/target/android_nnapi/compiler.py b/python/tvm/contrib/target/android_nnapi/compiler.py new file mode 100644 index 000000000000..a791482a2fb5 --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/compiler.py @@ -0,0 +1,101 @@ +# 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. +"""Compile a Relay IR Function into Android NNAPI C++ class.""" +import copy +import tvm +from . import transform +from . import json_to_nnapi +from .function_to_json_compiler import FunctionToJsonCompiler + + +class Compiler: + """Compile a Relay IR Function into Android NNAPI C++ class. + + Parameters + ---------- + options: dict + The compiler option dict. See below for available options. + + options["class"]["self"]["name"]: str + The name of the C++ class wrapping the Android NNAPI model. Defaults to "AnnGraph". + + options["target"]["api_level"]: int + The targeting Android API level. Defaults to 29. + """ + + DEFAULT_OPTIONS = { + "class": { + "self": { + "name": "AnnGraph", + }, + }, + "target": { + "api_level": 29, + }, + } + + def __init__(self, options): + self._options = self._expand_options(options) + + def codegen(self, func): + """Compile a Relay IR Function into Android NNAPI C++ class source code + + Parameters + ---------- + func: tvm.relay.Function + The Relay IR Function to be compiled + + Returns + ------- + code: str + The C++ class source code describing func in Android NNAPI + """ + assert isinstance(func, tvm.relay.Function) + func = transform.FixIllegalPatternForNnapi()(func) + + mod = tvm.IRModule({"main": func}) + export_obj = FunctionToJsonCompiler(self._options)(mod["main"]) + + ret = json_to_nnapi.codegen( + export_json=export_obj.asjson(), + options={ + "class": { + "name": self._options["class"]["self"]["name"], + }, + }, + ) + return ret + + @classmethod + def _expand_options(cls, options): + ret = copy.deepcopy(options) + + def _recursive_merge(cur_opts, def_opts): + for k, v in def_opts.items(): + if k in cur_opts: + if isinstance(v, dict): + assert isinstance(cur_opts[k], dict) + _recursive_merge(cur_opts[k], v) + else: + # type(cur_opts[k]) should be a basic type + assert isinstance(cur_opts[k], (float, int, str)) + else: # option k does not exist in current options, so copy from default options + cur_opts[k] = copy.deepcopy(v) + + _recursive_merge(ret, cls.DEFAULT_OPTIONS) + + return ret diff --git a/python/tvm/contrib/target/android_nnapi/error.py b/python/tvm/contrib/target/android_nnapi/error.py new file mode 100644 index 000000000000..6f75bda92028 --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/error.py @@ -0,0 +1,52 @@ +# 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. +"""Implements the errors and assertions function for Android NNAPI Compiler.""" + + +class AndroidNNAPICompilerError(RuntimeError): + """Android NNAPI compiler error base class. + + Parameters + ---------- + msg: str + The error message. + """ + + +class AndroidNNAPICompilerIncompatibleError(AndroidNNAPICompilerError): + """Error caused by parsing unsupported Relay AST. + + Parameters + ---------- + msg: str + The error message. + """ + + +def assert_anc_compatibility(boolean, *msg): + """Check for True or raise an AndroidNNAPICompilerIncompatibleError. + + Parameters + ---------- + boolean: bool + The checking condition. + + msg: str + Optional string message to be raised. + """ + if not boolean: + raise AndroidNNAPICompilerIncompatibleError(*msg) diff --git a/python/tvm/contrib/target/android_nnapi/export_object.py b/python/tvm/contrib/target/android_nnapi/export_object.py new file mode 100644 index 000000000000..7a2fb32b623a --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/export_object.py @@ -0,0 +1,335 @@ +# 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. +"""ExportObject, a dict-like structure providing infrastructure for +Android NNAPI codegen.""" +import struct +import copy +from .error import assert_anc_compatibility +from ._export_object import JSONAnalyzer as _JSONAnalyzer + + +class ExportObject: + """A dict-like structure providing infrastructure for Android NNAPI codegen. + + Parameters + ---------- + options: dict + The compiler option dict. + """ + + _SCALAR_RELAY_NNAPI_TYPE_MAP = { + "bool": "BOOL", + "float16": "FLOAT16", + "float32": "FLOAT32", + "int32": "INT32", + "uint32": "UINT32", + } + + _TENSOR_RELAY_NNAPI_TYPE_MAP = { + "bool": "TENSOR_BOOL", + "float16": "TENSOR_FLOAT16", + "float32": "TENSOR_FLOAT32", + "int32": "TENSOR_INT32", + "uint32": "TENSOR_UINT32", + } + + def __init__(self, options): + self._node_to_operand_idxs_map = {} + self._type_to_idx_map = {} + self._json = { + "constants": [], + "inputs": [], + "memories": [], + "operands": [], + "operations": [], + "outputs": [], + "types": [], + } + self.json_analyzer = _JSONAnalyzer(self._json) + self._options = options + + def __getitem__(self, key): + return self._json[key] + + def __setitem__(self, key, value): + self._json[key] = value + + def asjson(self): + """Return the content of ExportObject as a primitive Python dict. + + Returns + ------- + json: dict + The content of ExportObject as a primitive Python dict. + """ + return copy.deepcopy(self._json) + + def get_type_idx(self, tipe): + """Register and lookup type index in export_obj["types"]. + + Parameters + ---------- + tipe: ((int, ...), str) + type (shape, dtype) to look up. + + Returns + ------- + index: int + type index in export object. + """ + tipe = (tuple(map(int, tipe[0])), str(tipe[1])) # canonicalize + shape, dtype = tipe + assert_anc_compatibility( + dtype in ["bool", "float16", "float32", "int32", "uint32"], + f"Unsupported data type {dtype}", + ) + + if tipe not in self._type_to_idx_map: # create new type + if dtype == "bool": + assert_anc_compatibility( + self._options["target"]["api_level"] >= 29, + f"Boolean is not supported for Android API{self._options['target']['api_level']}", # pylint: disable=line-too-long + ) + + new_type = {} + if len(shape) == 0: + new_type["type"] = self._SCALAR_RELAY_NNAPI_TYPE_MAP[dtype] + else: + new_type["shape"] = list(shape) + new_type["type"] = self._TENSOR_RELAY_NNAPI_TYPE_MAP[dtype] + + self["types"].append(new_type) + self._type_to_idx_map[tipe] = len(self["types"]) - 1 + return self._type_to_idx_map[tipe] + + def register_node_operand_idxs(self, node, idxs): + """Register in the internal symbol table about the Android NNAPI + Operand indices of a given node. + + Parameters + ---------- + node: tvm.relay.Node + The node to be registered. + + idxs: list of int + The corresponding Android NNAPI Operand indices of the node. + """ + assert node not in self._node_to_operand_idxs_map + self._node_to_operand_idxs_map[node] = copy.deepcopy(idxs) + + def get_node_operand_idxs(self, node): + """Query the internal symbol table to find Android NNAPI Operand indices for a given node. + + Parameters + ---------- + node: tvm.relay.Node + The node to be queried. + + Returns + ------- + idxs: list of int + The indices which is mapped to the queried node. + """ + assert node in self._node_to_operand_idxs_map, f"Node {node} not found in the symbol table" + return self._node_to_operand_idxs_map[node] + + @staticmethod + def _canonicalize_scalar_constant(dtype, val): + # skip canonicalizing strings as they may carry specific meanings, + # e.g. macro-defined values + if not isinstance(val, str): + if dtype == "float16": + assert isinstance(val, float) + val = hex( + struct.unpack("H", struct.pack("e", val))[0] + ) # for float16 we use uint16_t in C, hence the conversion + elif dtype == "float32": + val = float(val) + elif dtype == "int32": + val = int(val) + elif dtype == "uint32": + val = int(val) + else: + assert dtype == "bool" + val = bool(val) + return val + + def add_scalar_constant(self, val, dtype): + """Add scalar constant to export object. + + Parameters + ---------- + val: numerical or str + value of the constant. Can be defined constant in the NNAPI framework. + + dtype: str + data type of the constant. + + Returns + ------- + index: int + index of the constant in export object constants array. + """ + # canonicalize + dtype = str(dtype) + assert_anc_compatibility( + dtype in ["float16", "float32", "int32", "uint32", "bool"], + f"Unsupported data type {dtype}", + ) + val = self._canonicalize_scalar_constant(dtype, val) + + new_const = { + "type": "scalar", + "dtype": dtype, + "value": val, + } + if new_const in self["constants"]: + return self["constants"].index(new_const) + + self["constants"].append(new_const) + return len(self["constants"]) - 1 + + def add_array_constant(self, vals, dtype): + """Add array constant to export object. + + Parameters + ---------- + vals: list of values in dtype + values of array. + + dtype: string + data type of array. + + Returns + ------- + index: int + index of added constant in export_obj["constants"]. + """ + # canonicalize + dtype = str(dtype) + assert_anc_compatibility( + dtype in ["float16", "float32", "int32", "uint32", "bool"], + f"Unsupported data type { dtype }", + ) + assert vals, "Array constant should not be empty" + vals = [self._canonicalize_scalar_constant(dtype, v) for v in vals] + + new_const = { + "type": "array", + "dtype": dtype, + "value": vals, + } + if new_const in self["constants"]: + return self["constants"].index(new_const) + + self["constants"].append(new_const) + return len(self["constants"]) - 1 + + def add_operand(self, type_idx, **kwargs): + """Add node to export_obj["operands"] and return its index. + + Parameters + ---------- + type_idx: int + index of node type in export_obj["types"]. + + kwargs["value"]: dict + dict representing node value. See below for more info. + + kwargs["value"]["type"]: str + type of value. Can be "constant_idx", "memory_ptr". + + kwargs["value"]["value"]: dict + value of initialized value. Should correspond to `kwargs["value"]["type"]`. + + kwargs["node"]: relay.Node + node to add. Use `None` to prevent operand being added to `node_to_operand_idxs_map`. + + Returns + ------- + indices: array of int + indices of node in export_obj["operands"]. + """ + node = kwargs.get("node", None) + value = kwargs.get("value", None) + + new_op = { + "type": type_idx, + } + + if value is not None: + new_op["value"] = copy.deepcopy(value) + + if node is not None and node in self._node_to_operand_idxs_map: + old_node_idxs = self.get_node_operand_idxs(node) + assert ( + len(old_node_idxs) == 1 + ) # Nodes registered with add_operand should be single indexed + assert self["operands"][old_node_idxs[0]] == new_op + return old_node_idxs + + self["operands"].append(new_op) + ret = [len(self["operands"]) - 1] + if node is not None: + self.register_node_operand_idxs(node, ret) + return ret + + def add_operation(self, nnapi_op_name, inputs, outputs): + """Add operation to export_obj["operations"]. + + Parameters + ---------- + nnapi_op_name: str + name of operator to be added in NNAPI. + + inputs: list of int + indices of input operands. + + outputs: list of int + indices of output operands. + """ + new_op = { + "input": copy.deepcopy(inputs), + "op": nnapi_op_name, + "output": copy.deepcopy(outputs), + } + self["operations"].append(new_op) + + def add_ann_memory(self, file_name, size): + """Add memory to export_obj["memories"]. + + Parameters + ---------- + file_name: str + file name or relative path to the underlying file of memory. + + size: int + size in bytes of the underlying file. + + Returns + ------- + idx: int + the index of the new memory. + """ + new_mem = { + "file_name": file_name, + "size": size, + } + if new_mem not in self["memories"]: + self["memories"].append(new_mem) + + return self["memories"].index(new_mem) diff --git a/python/tvm/contrib/target/android_nnapi/function_to_json_compiler.py b/python/tvm/contrib/target/android_nnapi/function_to_json_compiler.py new file mode 100644 index 000000000000..1e436cc65646 --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/function_to_json_compiler.py @@ -0,0 +1,176 @@ +# 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=wildcard-import,unused-wildcard-import +"""Compile a Relay IR Function to its Android NNAPI equivalence.""" +import tvm +import tvm.relay +from .error import * +from .operation_utils import relay_op +from .export_object import ExportObject + + +class FunctionToJsonCompiler(tvm.relay.ExprVisitor): + """Compile a Relay IR Function to an imtermediate JSON format for json2nnapi. + + Parameters + ---------- + options: dict + The compiler option dict. + """ + + def __init__(self, options): + super().__init__() + self._options = options + self._export_obj = ExportObject(self._options) + + def __call__(self, func): + """Compile a Relay IR Function to an imtermediate JSON format for json2nnapi. + + Parameters + ---------- + func: tvm.relay.Function + The Relay IR Function to be compiled. + + Returns + ------- + json: dict + A Python dict acting as the resulting JSON of the conversion. + """ + assert isinstance(func, tvm.relay.Function) + self.visit(func.body) + + # identify Android NNAPI model inputs + for p in func.params: + for i in self._export_obj.get_node_operand_idxs( + p + ): # param may be a tuple, which results in multiple indices + if i not in self._export_obj["inputs"]: + self._export_obj["inputs"].append(i) + + # identify Android NNAPI model outputs + for i in self._export_obj.get_node_operand_idxs( + func.body + ): # again, the output may be a tuple, which results in multiple indices + if i not in self._export_obj["outputs"]: + self._export_obj["outputs"].append(i) + # for now, let's force the function to return a single value, + # i.e. denying tuple as return type + assert len(self._export_obj["outputs"]) == 1 + + # set resulting memory for outputs + for i, op_i in enumerate(self._export_obj["outputs"]): + op = self._export_obj["operands"][op_i] + assert "value" not in op + op["value"] = { + "type": "memory_ptr", + "value": "out", # no formatting since len(outs) == 1 + } + + return self._export_obj + + @property + def export_obj(self): + """The associated ExportObject of this compiler instance.""" + return self._export_obj + + @property + def options(self): + """The associated compiler option dict.""" + return self._options + + def visit(self, expr): + assert_anc_compatibility( + isinstance( + expr, + ( + tvm.relay.Call, + tvm.relay.Var, + tvm.relay.Tuple, + tvm.relay.TupleGetItem, + tvm.relay.Constant, + ), + ), + f"{type(expr)} is not supported", + ) + return super().visit(expr) + + def visit_call(self, call): + if isinstance(call.op, tvm.ir.Op): + op_handler_module = relay_op + for namespace in call.op.name.split("."): # lookup the handler dynamically + op_handler_module = getattr(op_handler_module, namespace, None) + assert_anc_compatibility( + op_handler_module is not None, f"Relay IR Op { call.op } not implemented" + ) + op_handler_module.handler(self, call) + else: + raise AndroidNNAPICompilerIncompatibleError( + f"Conversion of { call.op.type_key } not supported" + ) + + def visit_var(self, var): + self._export_obj.add_operand( + type_idx=self._export_obj.get_type_idx( + (var.checked_type.shape, var.checked_type.dtype) + ), + node=var, + value={ + "type": "memory_ptr", + "value": var.name_hint, + }, + ) + + def visit_tuple(self, tup): + field_idxs = [] + for f in tup.fields: + self.visit(f) + field_idxs += self._export_obj.get_node_operand_idxs(f) + self._export_obj.register_node_operand_idxs(tup, field_idxs) + + def visit_tuple_getitem(self, t): + assert_anc_compatibility( + isinstance(t.tuple_value, tvm.relay.Tuple), + f"Getting tuple item from {type(t.tuple_value)} is not supported", + ) + self.visit(t.tuple_value) + self._export_obj.register_node_operand_idxs( + t, [self._export_obj.get_node_operand_idxs(t.tuple_value)[t.index]] + ) + + def visit_constant(self, const): + assert_anc_compatibility( + isinstance(const.checked_type, tvm.relay.TensorType), + f"Unsupported type {type(const.checked_type)}", + ) + shape, dtype = const.data.shape, const.data.dtype + type_idx = self._export_obj.get_type_idx((shape, dtype)) + + if shape == (): + const_idx = self._export_obj.add_scalar_constant(const.data.asnumpy().item(), dtype) + else: + assert_anc_compatibility(len(shape) == 1, "Only flat array constants are supported") + constants = [i.item() for i in const.data.asnumpy()] + const_idx = self._export_obj.add_array_constant(constants, dtype) + + self._export_obj.add_operand( + type_idx=type_idx, + value={ + "type": "constant_idx", + "value": const_idx, + }, + node=const, + ) diff --git a/python/tvm/contrib/target/android_nnapi/json_to_nnapi/__init__.py b/python/tvm/contrib/target/android_nnapi/json_to_nnapi/__init__.py new file mode 100644 index 000000000000..d18fe9557f6c --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/json_to_nnapi/__init__.py @@ -0,0 +1,116 @@ +# 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. +"""Codegen a JSON object to Android NNAPI source code.""" +import copy +from .stages import STAGES as _STAGES + + +_DEFAULT_OPTIONS = { + "class": { + "base_path": "/sdcard/nnapi_result", + "name": "AnnGraph", + }, + "model": { + "name": "model", + }, + "compilation": { + "name": "compilation", + }, + "execution": { + "name": "run", + "end_event_name": "run_end", + }, +} + + +def codegen(export_json, options={}): # pylint: disable=dangerous-default-value + """Codegen export_json to NNAPI codes. + + Parameters + ---------- + export_json: dict + The json representation of a NNAPI model. + + options["class"]["base_path"]: str + The base path of file accesses. Defaults to "/sdcard/nnapi_result". + + options["class"]["name"]: str + The name of the generated C++ class wrapping around NNAPI codes. Defaults to "AnnGraph". + + options["model"]["name"]: str + The name of the `ANeuralNetworksModel*` created. Defaults to "model". + + options["compilation"]["name"]: str + The name of the `ANeuralNetworksCompilation*` created. Defaults to "compilation". + + options["execution"]["name"]: str + The name of the `ANeuralNetworksExecution*` created. Defaults to "run". + + options["execution"]["end_event_name"]: str + The name of the `ANeuralNetworksEvent*` used to wait for execution completion. + Defaults to "run_end". + + Returns + ------- + code: str + The generated code. + """ + lines = { + "tmp": { + "model_creation": [], + "set_execution_io": [], + "wrapper_class": [], + }, + "result": "", + } + options = _set_options(options) + _export_json = copy.deepcopy(export_json) + + for s in _STAGES: + lines, _export_json = s(lines, _export_json, options) + + return lines["result"] + + +def _set_options(options): + """Set options. + + Parameters + ---------- + options: dict + The options to be set. + + Returns + ------- + options: dict + The updated options. + """ + + def _recursive_merge(cur_opts, def_opts): + for k, v in def_opts.items(): + if k in cur_opts: + if isinstance(v, dict): + assert isinstance(cur_opts[k], dict) + _recursive_merge(cur_opts[k], v) + else: + assert isinstance(cur_opts[k], (float, int, str)) + else: + cur_opts[k] = copy.deepcopy(v) + + _recursive_merge(options, _DEFAULT_OPTIONS) + + return options diff --git a/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/__init__.py b/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/__init__.py new file mode 100644 index 000000000000..f1ade4729faa --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/__init__.py @@ -0,0 +1,46 @@ +# 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 steps and pipeline of Android NNAPI codegen.""" +from .symbolize import symbolize +from .declare_types import declare_types +from .declare_operands import declare_operands +from .declare_constants import declare_constants +from .declare_memories import declare_memories +from .initialize_operands import initialize_operands +from .declare_operations import declare_operations +from .declare_inputs_outputs import declare_inputs_outputs +from .declare_wrapper_class import declare_wrapper_class +from .set_execution_inputs_outputs import set_execution_inputs_outputs +from .finalize import finalize + + +STAGES = [ + # model creation + symbolize, + declare_types, + declare_operands, + declare_constants, + declare_memories, + initialize_operands, + declare_operations, + declare_inputs_outputs, + # set execution io + set_execution_inputs_outputs, + # finalize + declare_wrapper_class, + finalize, +] diff --git a/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/declare_constants.py b/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/declare_constants.py new file mode 100644 index 000000000000..a6a6f33f1f8d --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/declare_constants.py @@ -0,0 +1,50 @@ +# 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. +"""Declare and define C constants used to set operand values.""" +from .. import templates + +C_TYPES_MAP = { + "int32": "int32_t", + "uint32": "uint32_t", + "float16": "uint16_t", + "float32": "float", + "bool": "bool", +} + + +def declare_constants(lines, export_obj, options): # pylint: disable=unused-argument + """Declare and define C constants used to set operand values.""" + for c in export_obj["constants"]: + tipe = c["type"] + c_dtype = C_TYPES_MAP[c["dtype"]] + if tipe == "scalar": + data = { + "dtype": c_dtype, + "name": c["name"], + "value": c["value"], + } + elif tipe == "array": + data = { + "dtype": c_dtype, + "name": c["name"], + "length": len(c["value"]), + "value": "{" + ", ".join([str(v) for v in c["value"]]) + "}", + } + else: + raise RuntimeError("Unknown constant type {}".format(tipe)) + lines["tmp"]["model_creation"].append(templates.declare_constant[tipe].substitute(**data)) + return lines, export_obj diff --git a/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/declare_inputs_outputs.py b/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/declare_inputs_outputs.py new file mode 100644 index 000000000000..413b585f81f5 --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/declare_inputs_outputs.py @@ -0,0 +1,37 @@ +# 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. +"""Specify Android NNAPI model input and output operands.""" +from .. import templates + + +def declare_inputs_outputs(lines, export_obj, options): + """Specify Android NNAPI model input and output operands.""" + inputs = export_obj["inputs"] + outputs = export_obj["outputs"] + data = { + "inputs": { + "length": len(inputs), + "str": "{" + ", ".join([str(i) for i in inputs]) + "}", + }, + "outputs": { + "length": len(outputs), + "str": "{" + ", ".join([str(i) for i in outputs]) + "}", + }, + "model": options["model"]["name"], + } + lines["tmp"]["model_creation"].append(templates.declare_inputs_outputs.substitute(**data)) + return lines, export_obj diff --git a/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/declare_memories.py b/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/declare_memories.py new file mode 100644 index 000000000000..e288e80a5dfe --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/declare_memories.py @@ -0,0 +1,29 @@ +# 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. +"""Create Android NNAPI memories.""" +from .. import templates + + +def declare_memories(lines, export_obj, options): + """Create Android NNAPI memories.""" + for m in export_obj["memories"]: + data = { + "file_path": "{}/{}".format(options["class"]["base_path"], m["file_name"]), + "mem_size": m["size"], + } + lines["tmp"]["model_creation"].append(templates.declare_memory.substitute(**data)) + return lines, export_obj diff --git a/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/declare_operands.py b/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/declare_operands.py new file mode 100644 index 000000000000..46b56444d8cc --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/declare_operands.py @@ -0,0 +1,31 @@ +# 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. +"""Declare Android NNAPI operands.""" +from .. import templates + + +def declare_operands(lines, export_obj, options): + """Declare Android NNAPI operands.""" + for i, op in enumerate(export_obj["operands"]): + op_type = export_obj["types"][op["type"]] + data = { + "model": options["model"]["name"], + "type": op_type["name"], + "index": i, + } + lines["tmp"]["model_creation"].append(templates.declare_operand.substitute(**data)) + return lines, export_obj diff --git a/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/declare_operations.py b/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/declare_operations.py new file mode 100644 index 000000000000..b764589e2b03 --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/declare_operations.py @@ -0,0 +1,37 @@ +# 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. +"""Declare Android NNAPI Operations.""" +from .. import templates + + +def declare_operations(lines, export_obj, options): + """Declare Android NNAPI Operations.""" + for op in export_obj["operations"]: + data = { + "inputs": { + "length": len(op["input"]), + "str": "{" + ", ".join([str(i) for i in op["input"]]) + "}", + }, + "outputs": { + "length": len(op["output"]), + "str": "{" + ", ".join([str(i) for i in op["output"]]) + "}", + }, + "model": options["model"]["name"], + "op_code": templates.ANN_PREFIX + op["op"], + } + lines["tmp"]["model_creation"].append(templates.declare_operation.substitute(**data)) + return lines, export_obj diff --git a/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/declare_types.py b/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/declare_types.py new file mode 100644 index 000000000000..20ff9c579076 --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/declare_types.py @@ -0,0 +1,35 @@ +# 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. +"""Declare and define Android NNAPI ANeuralNetworksOperandType.""" +from .. import templates + + +def declare_types(lines, export_obj, options): # pylint: disable=unused-argument + """Declare and define Android NNAPI ANeuralNetworksOperandType.""" + for t in export_obj["types"]: + tipe = { + "name": t["name"], + "type": templates.ANN_PREFIX + t["type"], + } + if "shape" in t: + tipe["dim_name"] = tipe["name"] + "_dims" + tipe["shape"] = { + "rank": len(t["shape"]), + "str": "{" + ", ".join([str(i) for i in t["shape"]]) + "}", + } + lines["tmp"]["model_creation"].append(templates.declare_type.substitute(tipe=tipe)) + return lines, export_obj diff --git a/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/declare_wrapper_class.py b/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/declare_wrapper_class.py new file mode 100644 index 000000000000..48085134ad8d --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/declare_wrapper_class.py @@ -0,0 +1,77 @@ +# 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. +"""Wraps the Android NNAPI model in a class.""" +from .. import templates + + +# NOTICE: make sure TVM maps type A to type B before modifying this table!! +C_TYPES_MAP = { + "BOOL": "bool", + "FLOAT32": "float", + "INT32": "int", + "TENSOR_BOOL8": "bool", + "TENSOR_FLOAT16": "uint16_t", + "TENSOR_FLOAT32": "float", + "TENSOR_INT32": "int", +} + + +def declare_wrapper_class(lines, export_obj, options): + """Wraps the Android NNAPI model in a class.""" + data = { + "class": { + "self": { + "name": options["class"]["name"], + }, + "model": { + "name": options["model"]["name"], + }, + "compilation": { + "name": options["compilation"]["name"], + }, + "execution": { + "name": options["execution"]["name"], + "end_event_name": options["execution"]["end_event_name"], + }, + }, + "codes": { + "model_creation": "\n".join( + [" " + s for s in "\n".join(lines["tmp"]["model_creation"]).split("\n")] + ), + "set_execution_io": "\n".join( + [" " + s for s in "\n".join(lines["tmp"]["set_execution_io"]).split("\n")] + ), + }, + } + + def _scope(): + var_decls = [] + for inp in export_obj["inputs"]: + op = export_obj["operands"][inp] + assert op["value"]["type"] == "memory_ptr" + tipe = export_obj["types"][op["type"]] + var_decls.append("{}* {}".format(C_TYPES_MAP[tipe["type"]], op["value"]["value"])) + for outp in export_obj["outputs"]: + op = export_obj["operands"][outp] + assert op["value"]["type"] == "memory_ptr" + tipe = export_obj["types"][op["type"]] + var_decls.append("{}* {}".format(C_TYPES_MAP[tipe["type"]], op["value"]["value"])) + data["class"]["execution"]["func_params_decl_str"] = ", ".join(var_decls) + + _scope() + lines["tmp"]["wrapper_class"].append(templates.declare_wrapper_class.substitute(**data)) + return lines, export_obj diff --git a/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/finalize.py b/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/finalize.py new file mode 100644 index 000000000000..ade176b2d08b --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/finalize.py @@ -0,0 +1,24 @@ +# 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. +"""Produce codegen result from intermediate results.""" + + +def finalize(lines, export_obj, options): # pylint: disable=unused-argument + """Produce codegen result from intermediate results.""" + lines["result"] = "\n".join(lines["tmp"]["wrapper_class"]) + lines["result"] = "\n".join([s for s in lines["result"].split("\n") if s.strip()]) + return lines, export_obj diff --git a/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/initialize_operands.py b/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/initialize_operands.py new file mode 100644 index 000000000000..cead1cd1dfe3 --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/initialize_operands.py @@ -0,0 +1,57 @@ +# 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. +"""Set initialized value to Android NNAPI operands.""" +from .. import templates + + +def initialize_operands(lines, export_obj, options): + """Set initialized value to Android NNAPI operands.""" + for i, op in enumerate(export_obj["operands"]): + value = op.get("value", None) + if value is None: + continue + + data = { + "model": options["model"]["name"], + "op_idx": i, + } + if value["type"] == "constant_idx": + const = export_obj["constants"][value["value"]] + data["memory_size"] = "sizeof({})".format(const["name"]) + if const["type"] == "scalar": + data["memory_ptr"] = "&" + const["name"] + elif const["type"] == "array": + data["memory_ptr"] = const["name"] + else: + raise RuntimeError( + "Unknown const type ({}) for operand {}".format(const["type"], i) + ) + lines["tmp"]["model_creation"].append( + templates.initialize_operand["memory_ptr"].substitute(**data) + ) + elif value["type"] == "memory_ptr": + pass + elif value["type"] == "ann_memory": + memory = export_obj["memories"][value["value"]] + data["memory_idx"] = value["value"] + data["length"] = memory["size"] + lines["tmp"]["model_creation"].append( + templates.initialize_operand["ann_memory"].substitute(**data) + ) + else: + raise RuntimeError("Unknown value type ({}) for operand {}".format(value["type"], i)) + return lines, export_obj diff --git a/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/set_execution_inputs_outputs.py b/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/set_execution_inputs_outputs.py new file mode 100644 index 000000000000..f374ed192dd7 --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/set_execution_inputs_outputs.py @@ -0,0 +1,69 @@ +# 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. +"""Sets the inputs and outputs for the generated Android NNAPI +model.""" +import re +from functools import reduce +from .. import templates + + +def set_execution_inputs_outputs(lines, export_obj, options): + """Sets the inputs and outputs for the generated Android NNAPI + model.""" + for i, op_i in enumerate(export_obj["inputs"]): + op = export_obj["operands"][op_i] + value = op["value"] + assert value["type"] == "memory_ptr" + + data = { + "execution": options["execution"]["name"], + "input_idx": i, + } + tipe = export_obj["types"][op["type"]] + nnapi_dtype = tipe["type"] + nbits = int((lambda s: s if s != "" else "8")(re.sub(r"^[^0-9]+", "", nnapi_dtype))) + assert (nbits != 0) and (nbits % 8 == 0) + data["memory_ptr"] = value["value"] + if nnapi_dtype.startswith("TENSOR"): + data["memory_size"] = reduce(lambda a, b: a * b, tipe["shape"], 1) * nbits // 8 + else: + data["memory_size"] = nbits // 8 + lines["tmp"]["set_execution_io"].append(templates.set_execution_input.substitute(**data)) + + def _outputs(): + assert len(export_obj["outputs"]) == 1 + op = export_obj["operands"][export_obj["outputs"][0]] + value = op["value"] + assert value["type"] == "memory_ptr" + + data = { + "execution": options["execution"]["name"], + "output_idx": 0, + } + tipe = export_obj["types"][op["type"]] + nnapi_dtype = tipe["type"] + nbits = int((lambda s: s if s != "" else "8")(re.sub(r"^[^0-9]+", "", nnapi_dtype))) + assert (nbits != 0) and (nbits % 8 == 0) + data["memory_ptr"] = value["value"] + if nnapi_dtype.startswith("TENSOR"): + data["memory_size"] = reduce(lambda a, b: a * b, tipe["shape"], 1) * nbits // 8 + else: + data["memory_size"] = nbits // 8 + lines["tmp"]["set_execution_io"].append(templates.set_execution_output.substitute(**data)) + + _outputs() + return lines, export_obj diff --git a/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/symbolize.py b/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/symbolize.py new file mode 100644 index 000000000000..4e30b36a3749 --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/json_to_nnapi/stages/symbolize.py @@ -0,0 +1,47 @@ +# 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. +"""Prepare JSON object for Android NNAPI codegen.""" + + +def symbolize(lines, export_obj, options): # pylint: disable=unused-argument + """Assign C symbols to JSON objects.""" + + def _symbolize_types(types): + cnts = { + "tensor": 0, + "scalar": 0, + } + for t in types: + if t["type"].startswith("TENSOR_"): + t["name"] = "tensor" + str(cnts["tensor"]) + cnts["tensor"] += 1 + else: + t["name"] = "scalar" + str(cnts["scalar"]) + cnts["scalar"] += 1 + + _symbolize_types(export_obj["types"]) + + def _symbolize_consts(consts): + cnt = 0 + for c in consts: + c["name"] = "const_val" + str(cnt) + cnt += 1 + + if "constants" in export_obj: + _symbolize_consts(export_obj["constants"]) + + return lines, export_obj diff --git a/python/tvm/contrib/target/android_nnapi/json_to_nnapi/templates.py b/python/tvm/contrib/target/android_nnapi/json_to_nnapi/templates.py new file mode 100644 index 000000000000..b2d47efd6f21 --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/json_to_nnapi/templates.py @@ -0,0 +1,276 @@ +# 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,missing-class-docstring,missing-function-docstring +"""The string templates for Android NNAPI codegen.""" +import string + +ANN_PREFIX = "ANEURALNETWORKS_" + + +class declare_type: + @staticmethod + def substitute(**kwargs): + tipe = kwargs["tipe"] + ret = "" + ret += f"""ANeuralNetworksOperandType {tipe["name"]}; +{tipe["name"]}.type = {tipe["type"]}; +{tipe["name"]}.scale = 0.f; +{tipe["name"]}.zeroPoint = 0; +""" + if "shape" in tipe: + ret += f"""{tipe["name"]}.dimensionCount = {tipe["shape"]["rank"]}; +static uint32_t {tipe["dim_name"]}[{tipe["shape"]["rank"]}] = {tipe["shape"]["str"]}; +{tipe["name"]}.dimensions = {tipe["dim_name"]}; +""" + else: + ret += f"""{tipe["name"]}.dimensionCount = 0; +{tipe["name"]}.dimensions = NULL; +""" + + return ret + + +declare_operand = string.Template( + """JSON2NNAPI_CHECK_EQ( + ANeuralNetworksModel_addOperand( + ${model}, + &${type} + ), + ANEURALNETWORKS_NO_ERROR +); // Operand ${index} +""" +) + +declare_constant = { + "scalar": string.Template( + """static ${dtype} ${name} = ${value}; +""" + ), + "array": string.Template( + """static ${dtype} ${name}[${length}] = ${value}; +""" + ), +} + + +class declare_memory: + @staticmethod + def substitute(**kwargs): + file_path = kwargs["file_path"] + mem_size = kwargs["mem_size"] + ret = f"""{{ + ANeuralNetworksMemory* mem = nullptr; + int fd = open("{file_path}", O_RDONLY); + JSON2NNAPI_CHECK_NE(fd, -1); + JSON2NNAPI_CHECK_EQ( + ANeuralNetworksMemory_createFromFd( + {mem_size}, + PROT_READ, + fd, + 0, + &mem + ), + ANEURALNETWORKS_NO_ERROR + ); + this->memories_.push_back({{ fd, mem }}); +}} +""" + return ret + + +initialize_operand = { + "memory_ptr": string.Template( + """JSON2NNAPI_CHECK_EQ( + ANeuralNetworksModel_setOperandValue( + ${model}, + ${op_idx}, + ${memory_ptr}, + ${memory_size} + ), + ANEURALNETWORKS_NO_ERROR +); +""" + ), + "ann_memory": string.Template( + """JSON2NNAPI_CHECK_EQ( + ANeuralNetworksModel_setOperandValueFromMemory( + ${model}, + ${op_idx}, + std::get< 1 >(this->memories_[${memory_idx}]), + 0, + ${length} + ), + ANEURALNETWORKS_NO_ERROR +); +""" + ), +} + + +class declare_operation: + @staticmethod + def substitute(**kwargs): + inputs = kwargs["inputs"] + outputs = kwargs["outputs"] + model = kwargs["model"] + op_code = kwargs["op_code"] + ret = f"""{{ + static uint32_t inputIndexes[{inputs["length"]}] = {inputs["str"]}; + static uint32_t outputIndexes[{outputs["length"]}] = {outputs["str"]}; + JSON2NNAPI_CHECK_EQ( + ANeuralNetworksModel_addOperation( + {model}, + {op_code}, + {inputs["length"]}, + inputIndexes, + {outputs["length"]}, + outputIndexes + ), + ANEURALNETWORKS_NO_ERROR + ); +}} +""" + return ret + + +class declare_inputs_outputs: + @staticmethod + def substitute(**kwargs): + model = kwargs["model"] + inputs = kwargs["inputs"] + outputs = kwargs["outputs"] + ret = f"""static uint32_t {model}InputIndexes[{inputs["length"]}] = {inputs["str"]}; +static uint32_t {model}OutputIndexes[{outputs["length"]}] = {outputs["str"]}; +JSON2NNAPI_CHECK_EQ( + ANeuralNetworksModel_identifyInputsAndOutputs( + {model}, + {inputs["length"]}, + {model}InputIndexes, + {outputs["length"]}, + {model}OutputIndexes + ), + ANEURALNETWORKS_NO_ERROR +); +""" + return ret + + +class declare_wrapper_class: + @staticmethod + def substitute(**kwargs): + clas = kwargs["class"] + codes = kwargs["codes"] + ret = f"""#define JSON2NNAPI_CHECK_EQ(a, b) {{ assert((a) == (b)); }} +#define JSON2NNAPI_CHECK_NE(a, b) {{ assert((a) != (b)); }} +class {clas["self"]["name"]} +{{ +public: + {clas["self"]["name"]}() + {{ + JSON2NNAPI_CHECK_EQ(ANeuralNetworksModel_create(&this->{clas["model"]["name"]}), ANEURALNETWORKS_NO_ERROR); + this->createAnnModel(); + JSON2NNAPI_CHECK_EQ(ANeuralNetworksModel_finish(this->{clas["model"]["name"]}), ANEURALNETWORKS_NO_ERROR); +#if __ANDROID_API__ >= 29 && defined(JSON2NNAPI_FORCE_CPU_FALLBACK) + uint32_t num_nnapi_devices; + JSON2NNAPI_CHECK_EQ(ANeuralNetworks_getDeviceCount(&num_nnapi_devices), ANEURALNETWORKS_NO_ERROR); + ANeuralNetworksDevice * nnapi_fallback_dev; + for (int i = 0; i < num_nnapi_devices; i++) + {{ + JSON2NNAPI_CHECK_EQ(ANeuralNetworks_getDevice(i, &nnapi_fallback_dev), ANEURALNETWORKS_NO_ERROR); + int32_t dev_type; + JSON2NNAPI_CHECK_EQ(ANeuralNetworksDevice_getType(nnapi_fallback_dev, &dev_type), ANEURALNETWORKS_NO_ERROR); + if (dev_type == ANEURALNETWORKS_DEVICE_CPU) + {{ + break; + }} + }} + {{ + const ANeuralNetworksDevice * const dev_list[] = {{ nnapi_fallback_dev }}; + JSON2NNAPI_CHECK_EQ(ANeuralNetworksCompilation_createForDevices(this->{clas["model"]["name"]}, dev_list, 1, &this->{clas["compilation"]["name"]}), ANEURALNETWORKS_NO_ERROR); + }} +#else // #if __ANDROID_API__ >= 29 && defined(JSON2NNAPI_FORCE_CPU_FALLBACK) + JSON2NNAPI_CHECK_EQ(ANeuralNetworksCompilation_create(this->{clas["model"]["name"]}, &this->{clas["compilation"]["name"]}), ANEURALNETWORKS_NO_ERROR); +#endif // #if __ANDROID_API__ >= 29 && defined(JSON2NNAPI_FORCE_CPU_FALLBACK) + JSON2NNAPI_CHECK_EQ(ANeuralNetworksCompilation_finish(this->{clas["compilation"]["name"]}), ANEURALNETWORKS_NO_ERROR); + }} + + ~{clas["self"]["name"]}() + {{ + ANeuralNetworksCompilation_free(this->{clas["compilation"]["name"]}); + ANeuralNetworksModel_free(this->{clas["model"]["name"]}); + for (const auto &t: this->memories_) + {{ + ANeuralNetworksMemory_free(std::get< 1 >(t)); + close(std::get< 0 >(t)); + }} + }} + + void createAnnModel() + {{ +{codes["model_creation"]} + }} + + void execute({clas["execution"]["func_params_decl_str"]}) + {{ + ANeuralNetworksExecution* {clas["execution"]["name"]} = nullptr; + JSON2NNAPI_CHECK_EQ(ANeuralNetworksExecution_create(this->{clas["compilation"]["name"]}, &{clas["execution"]["name"]}), ANEURALNETWORKS_NO_ERROR); + +{codes["set_execution_io"]} + + ANeuralNetworksEvent* {clas["execution"]["end_event_name"]} = nullptr; + JSON2NNAPI_CHECK_EQ(ANeuralNetworksExecution_startCompute({clas["execution"]["name"]}, &{clas["execution"]["end_event_name"]}), ANEURALNETWORKS_NO_ERROR); + JSON2NNAPI_CHECK_EQ(ANeuralNetworksEvent_wait({clas["execution"]["end_event_name"]}), ANEURALNETWORKS_NO_ERROR); + ANeuralNetworksEvent_free({clas["execution"]["end_event_name"]}); + ANeuralNetworksExecution_free({clas["execution"]["name"]}); + }} + +private: + ANeuralNetworksModel* {clas["model"]["name"]} = nullptr; + ANeuralNetworksCompilation* {clas["compilation"]["name"]} = nullptr; + std::vector< std::tuple< int, ANeuralNetworksMemory* > > memories_; +}}; +""" + return ret + + +set_execution_input = string.Template( + """JSON2NNAPI_CHECK_EQ( + ANeuralNetworksExecution_setInput( + ${execution}, + ${input_idx}, + nullptr, + ${memory_ptr}, + ${memory_size} + ), + ANEURALNETWORKS_NO_ERROR +); +""" +) + +set_execution_output = string.Template( + """JSON2NNAPI_CHECK_EQ( + ANeuralNetworksExecution_setOutput( + ${execution}, + ${output_idx}, + nullptr, + ${memory_ptr}, + ${memory_size} + ), + ANEURALNETWORKS_NO_ERROR +); +""" +) diff --git a/python/tvm/contrib/target/android_nnapi/operation_utils/__init__.py b/python/tvm/contrib/target/android_nnapi/operation_utils/__init__.py new file mode 100644 index 000000000000..1567019d976d --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/operation_utils/__init__.py @@ -0,0 +1,19 @@ +# 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. +"""Namespace for compiling tvm.relay.Call to Android NNAPI Operation.""" +from . import relay_op +from . import nnapi_op diff --git a/python/tvm/contrib/target/android_nnapi/operation_utils/_utils.py b/python/tvm/contrib/target/android_nnapi/operation_utils/_utils.py new file mode 100644 index 000000000000..1d0be3c90159 --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/operation_utils/_utils.py @@ -0,0 +1,37 @@ +# 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. +"""Utilities for compiling tvm.relay.Call to Android NNAPI Operations.""" + + +def name_args(args, arg_names): + """Put arguments into dict for convenient lookup. + + Parameters + ---------- + args: array of relay.Expr + args of relay.Call. + + arg_names: array of string + names of args. + + Returns + ------- + args_map: dict of string to relay.Expr + named args dict. + """ + assert len(args) == len(arg_names) + return dict(zip(arg_names, args)) diff --git a/python/tvm/contrib/target/android_nnapi/operation_utils/nnapi_op/__init__.py b/python/tvm/contrib/target/android_nnapi/operation_utils/nnapi_op/__init__.py new file mode 100644 index 000000000000..b1c82b838efc --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/operation_utils/nnapi_op/__init__.py @@ -0,0 +1,22 @@ +# 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. +"""Namespace for Android NNAPI operation checkers.""" +from . import cast +from . import conv_2d +from . import depthwise_conv_2d +from . import grouped_conv_2d +from . import transpose diff --git a/python/tvm/contrib/target/android_nnapi/operation_utils/nnapi_op/cast.py b/python/tvm/contrib/target/android_nnapi/operation_utils/nnapi_op/cast.py new file mode 100644 index 000000000000..d79a96b5630e --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/operation_utils/nnapi_op/cast.py @@ -0,0 +1,71 @@ +# 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=wildcard-import,unused-wildcard-import +"""Add an ANEURALNETWORKS_CAST operation with checking.""" +from .error import * + + +def add_operation(compiler, inputs, outputs): + """Add an ANEURALNETWORKS_CAST operation with checking. + + Parameters + ---------- + compiler: FunctionToJsonCompiler + the compiler object holding export_obj. + + inputs: list of int + inputs to the operation. + + outputs: list of int + outputs of the operation. + """ + api_level = compiler.options["target"]["api_level"] + assert_anc_compatibility( + api_level >= 29, + f"Target Android API level { api_level } is too low to support the operation", + ) + + # check inputs + assert_nnapi_op_check(len(inputs) == 1) + ins = [{}] + + # check inputs[0] + ins[0] = {} + ins[0]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[0]) + assert_nnapi_op_check( + ins[0]["dtype"] == "TENSOR_FLOAT16" + or ins[0]["dtype"] == "TENSOR_FLOAT32" + or ins[0]["dtype"] == "TENSOR_INT32" + ) + ins[0]["shape"] = compiler.export_obj.json_analyzer.operand.get_shape(inputs[0]) + + # check outputs + assert_nnapi_op_check(len(outputs) == 1) + outs = [{}] + + # check outputs[0] + outs[0] = {} + outs[0]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(outputs[0]) + assert_nnapi_op_check( + outs[0]["dtype"] == "TENSOR_FLOAT16" + or outs[0]["dtype"] == "TENSOR_FLOAT32" + or outs[0]["dtype"] == "TENSOR_INT32" + ) + outs[0]["shape"] = compiler.export_obj.json_analyzer.operand.get_shape(outputs[0]) + assert_nnapi_op_check(outs[0]["shape"] == ins[0]["shape"]) + + compiler.export_obj.add_operation("CAST", inputs, outputs) diff --git a/python/tvm/contrib/target/android_nnapi/operation_utils/nnapi_op/conv_2d.py b/python/tvm/contrib/target/android_nnapi/operation_utils/nnapi_op/conv_2d.py new file mode 100644 index 000000000000..854371a39f0e --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/operation_utils/nnapi_op/conv_2d.py @@ -0,0 +1,210 @@ +# 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=wildcard-import,unused-wildcard-import +"""Add an ANEURALNETWORKS_CONV_2D operation with checking.""" +from .error import * + + +def add_operation(compiler, inputs, outputs): + """Add an ANEURALNETWORKS_CONV_2D operation with checking. + + Parameters + ---------- + compiler: FunctionToJsonCompiler + the compiler object holding export_obj. + + inputs: list of int + inputs to the operation. + + outputs: list of int + outputs of the operation. + """ + api_level = compiler.options["target"]["api_level"] + assert_anc_compatibility( + api_level >= 27, + f"Target Android API level { api_level } is too low to support the operation", + ) + + # check inputs + if api_level >= 29: + assert_nnapi_op_check(len(inputs) == 13) + else: + assert_nnapi_op_check(len(inputs) == 10) + ins = [{} for i in range(len(inputs))] + + # check inputs[0] + ins[0] = {} + ins[0]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[0]) + if ins[0]["dtype"] == "TENSOR_FLOAT16": + assert_nnapi_op_check(api_level >= 29) + else: + assert_nnapi_op_check(ins[0]["dtype"] == "TENSOR_FLOAT32") + ins[0]["rank"] = compiler.export_obj.json_analyzer.operand.get_rank(inputs[0]) + assert_nnapi_op_check(ins[0]["rank"] == 4) + ins[0]["shape"] = compiler.export_obj.json_analyzer.operand.get_shape(inputs[0]) + if ins[0]["shape"][0] == 0: + assert_nnapi_op_check(api_level >= 29) + + # check inputs[1] + ins[1] = {} + ins[1]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[1]) + if ins[1]["dtype"] == "TENSOR_FLOAT16": + assert_nnapi_op_check(api_level >= 29) + else: + assert_nnapi_op_check(ins[1]["dtype"] == "TENSOR_FLOAT32") + assert_nnapi_op_check(ins[1]["dtype"] == ins[0]["dtype"]) + ins[1]["rank"] = compiler.export_obj.json_analyzer.operand.get_rank(inputs[1]) + assert_nnapi_op_check(ins[1]["rank"] == 4) + ins[1]["shape"] = compiler.export_obj.json_analyzer.operand.get_shape(inputs[1]) + felter = dict(zip(["do", "fh", "fw", "di"], ins[1]["shape"])) + + # check inputs[2] + ins[2] = {} + ins[2]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[2]) + assert_nnapi_op_check(ins[2]["dtype"] == ins[1]["dtype"]) + ins[2]["rank"] = compiler.export_obj.json_analyzer.operand.get_rank(inputs[2]) + assert_nnapi_op_check(ins[2]["rank"] == 1) + ins[2]["constant"] = compiler.export_obj.json_analyzer.operand.get_constant(inputs[2]) + assert_nnapi_op_check( + ins[2]["constant"]["type"] == "array" and len(ins[2]["constant"]["value"]) == felter["do"] + ) + + # check inputs[3] + ins[3] = {} + ins[3]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[3]) + assert_nnapi_op_check(ins[3]["dtype"] == "INT32") + ins[3]["value"] = compiler.export_obj.json_analyzer.operand.get_value(inputs[3]) + assert_nnapi_op_check(ins[3]["value"] >= 0) + padding = {} + padding["l"] = ins[3]["value"] + + # check inputs[4] + ins[4] = {} + ins[4]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[4]) + assert_nnapi_op_check(ins[4]["dtype"] == "INT32") + ins[4]["value"] = compiler.export_obj.json_analyzer.operand.get_value(inputs[4]) + assert_nnapi_op_check(ins[4]["value"] >= 0) + padding["r"] = ins[4]["value"] + + # check inputs[5] + ins[5] = {} + ins[5]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[5]) + assert_nnapi_op_check(ins[5]["dtype"] == "INT32") + ins[5]["value"] = compiler.export_obj.json_analyzer.operand.get_value(inputs[5]) + assert_nnapi_op_check(ins[5]["value"] >= 0) + padding["t"] = ins[5]["value"] + + # check inputs[6] + ins[6] = {} + ins[6]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[6]) + assert_nnapi_op_check(ins[6]["dtype"] == "INT32") + ins[6]["value"] = compiler.export_obj.json_analyzer.operand.get_value(inputs[6]) + assert_nnapi_op_check(ins[6]["value"] >= 0) + padding["b"] = ins[6]["value"] + + # check inputs[7] + ins[7] = {} + ins[7]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[7]) + assert_nnapi_op_check(ins[7]["dtype"] == "INT32") + ins[7]["value"] = compiler.export_obj.json_analyzer.operand.get_value(inputs[7]) + assert_nnapi_op_check(ins[7]["value"] >= 0) + stride = {} + stride["w"] = ins[7]["value"] + + # check inputs[8] + ins[8] = {} + ins[8]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[8]) + assert_nnapi_op_check(ins[8]["dtype"] == "INT32") + ins[8]["value"] = compiler.export_obj.json_analyzer.operand.get_value(inputs[8]) + assert_nnapi_op_check(ins[8]["value"] >= 0) + stride["h"] = ins[8]["value"] + + # check inputs[9] + assert_nnapi_op_check(compiler.export_obj.json_analyzer.operand.is_fuse_code(inputs[9])) + + if api_level >= 29: + # check inputs[10] + ins[10] = {} + ins[10]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[10]) + assert_nnapi_op_check(ins[10]["dtype"] == "BOOL") + ins[10]["value"] = compiler.export_obj.json_analyzer.operand.get_value(inputs[10]) + assert_nnapi_op_check(ins[10]["value"] == "false" or ins[10]["value"] == "true") + + # check inputs[11] + ins[11] = {} + ins[11]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[11]) + assert_nnapi_op_check(ins[11]["dtype"] == "INT32") + ins[11]["value"] = compiler.export_obj.json_analyzer.operand.get_value(inputs[11]) + assert_nnapi_op_check(ins[11]["value"] >= 1) + + # check inputs[12] + ins[12] = {} + ins[12]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[12]) + assert_nnapi_op_check(ins[12]["dtype"] == "INT32") + ins[12]["value"] = compiler.export_obj.json_analyzer.operand.get_value(inputs[12]) + assert_nnapi_op_check(ins[12]["value"] >= 1) + + # check shapes + if api_level >= 29 and ins[10]["value"] == "true": + data_shape = { + "n": ins[0]["shape"][0], + "c": ins[0]["shape"][1], + "h": ins[0]["shape"][2], + "w": ins[0]["shape"][3], + } + else: + data_shape = { + "n": ins[0]["shape"][0], + "h": ins[0]["shape"][1], + "w": ins[0]["shape"][2], + "c": ins[0]["shape"][3], + } + + assert_nnapi_op_check(felter["di"] == data_shape["c"]) + + # check outputs + assert_nnapi_op_check(len(outputs) == 1) + outs = [{}] + + # check outputs[0] + outs[0] = {} + outs[0]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(outputs[0]) + assert_nnapi_op_check(outs[0]["dtype"] == ins[0]["dtype"]) + outs[0]["shape"] = compiler.export_obj.json_analyzer.operand.get_shape(outputs[0]) + + if api_level >= 29 and ins[10]["value"] == "true": + out_data_shape = { + "n": outs[0]["shape"][0], + "c": outs[0]["shape"][1], + "h": outs[0]["shape"][2], + "w": outs[0]["shape"][3], + } + else: + out_data_shape = { + "n": outs[0]["shape"][0], + "h": outs[0]["shape"][1], + "w": outs[0]["shape"][2], + "c": outs[0]["shape"][3], + } + total_h = data_shape["h"] + padding["t"] + padding["b"] + total_w = data_shape["w"] + padding["l"] + padding["r"] + assert_nnapi_op_check(out_data_shape["n"] == data_shape["n"]) + assert_nnapi_op_check(out_data_shape["h"] == ((total_h - felter["fh"]) // stride["h"] + 1)) + assert_nnapi_op_check(out_data_shape["w"] == ((total_w - felter["fw"]) // stride["w"] + 1)) + assert_nnapi_op_check(out_data_shape["c"] == felter["do"]) + + compiler.export_obj.add_operation("CONV_2D", inputs, outputs) diff --git a/python/tvm/contrib/target/android_nnapi/operation_utils/nnapi_op/depthwise_conv_2d.py b/python/tvm/contrib/target/android_nnapi/operation_utils/nnapi_op/depthwise_conv_2d.py new file mode 100644 index 000000000000..3ce487d5039b --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/operation_utils/nnapi_op/depthwise_conv_2d.py @@ -0,0 +1,219 @@ +# 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=wildcard-import,unused-wildcard-import +"""Add an ANEURALNETWORKS_DEPTHWISE_CONV_2D operation with checking.""" +from .error import * + + +def add_operation(compiler, inputs, outputs): + """Add an ANEURALNETWORKS_DEPTHWISE_CONV_2D operation with checking. + + Parameters + ---------- + compiler: FunctionToJsonCompiler + the compiler object holding export_obj. + + inputs: list of int + inputs to the operation. + + outputs: list of int + outputs of the operation. + """ + api_level = compiler.options["target"]["api_level"] + assert_anc_compatibility( + api_level >= 27, + f"Target Android API level { api_level } is too low to support the operation", + ) + + # check inputs + if api_level >= 29: + assert_nnapi_op_check(len(inputs) == 14) + else: + assert_nnapi_op_check(len(inputs) == 11) + ins = [{} for i in range(len(inputs))] + + # check inputs[0] + ins[0] = {} + ins[0]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[0]) + if ins[0]["dtype"] == "TENSOR_FLOAT16": + assert_nnapi_op_check(api_level >= 29) + else: + assert_nnapi_op_check(ins[0]["dtype"] == "TENSOR_FLOAT32") + ins[0]["rank"] = compiler.export_obj.json_analyzer.operand.get_rank(inputs[0]) + assert_nnapi_op_check(ins[0]["rank"] == 4) + ins[0]["shape"] = compiler.export_obj.json_analyzer.operand.get_shape(inputs[0]) + + # check inputs[1] + ins[1] = {} + ins[1]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[1]) + if ins[1]["dtype"] == "TENSOR_FLOAT16": + assert_nnapi_op_check(api_level >= 29) + else: + assert_nnapi_op_check(ins[1]["dtype"] == "TENSOR_FLOAT32") + assert_nnapi_op_check(ins[1]["dtype"] == ins[0]["dtype"]) + ins[1]["rank"] = compiler.export_obj.json_analyzer.operand.get_rank(inputs[1]) + assert_nnapi_op_check(ins[1]["rank"] == 4) + ins[1]["shape"] = compiler.export_obj.json_analyzer.operand.get_shape(inputs[1]) + felter = dict(zip(["di", "fh", "fw", "do"], ins[1]["shape"])) + assert_nnapi_op_check(felter["di"] == 1) + + # check inputs[2] + ins[2] = {} + ins[2]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[2]) + assert_nnapi_op_check(ins[2]["dtype"] == ins[1]["dtype"] and ins[2]["dtype"] == ins[0]["dtype"]) + ins[2]["rank"] = compiler.export_obj.json_analyzer.operand.get_rank(inputs[2]) + assert_nnapi_op_check(ins[2]["rank"] == 1) + ins[2]["constant"] = compiler.export_obj.json_analyzer.operand.get_constant(inputs[2]) + assert_nnapi_op_check( + ins[2]["constant"]["type"] == "array" and len(ins[2]["constant"]["value"]) == felter["do"] + ) + + # check inputs[3] + ins[3] = {} + ins[3]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[3]) + assert_nnapi_op_check(ins[3]["dtype"] == "INT32") + ins[3]["value"] = compiler.export_obj.json_analyzer.operand.get_value(inputs[3]) + assert_nnapi_op_check(ins[3]["value"] >= 0) + padding = {} + padding["l"] = ins[3]["value"] + + # check inputs[4] + ins[4] = {} + ins[4]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[4]) + assert_nnapi_op_check(ins[4]["dtype"] == "INT32") + ins[4]["value"] = compiler.export_obj.json_analyzer.operand.get_value(inputs[4]) + assert_nnapi_op_check(ins[4]["value"] >= 0) + padding["r"] = ins[4]["value"] + + # check inputs[5] + ins[5] = {} + ins[5]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[5]) + assert_nnapi_op_check(ins[5]["dtype"] == "INT32") + ins[5]["value"] = compiler.export_obj.json_analyzer.operand.get_value(inputs[5]) + assert_nnapi_op_check(ins[5]["value"] >= 0) + padding["t"] = ins[5]["value"] + + # check inputs[6] + ins[6] = {} + ins[6]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[6]) + assert_nnapi_op_check(ins[6]["dtype"] == "INT32") + ins[6]["value"] = compiler.export_obj.json_analyzer.operand.get_value(inputs[6]) + assert_nnapi_op_check(ins[6]["value"] >= 0) + padding["b"] = ins[6]["value"] + + # check inputs[7] + ins[7] = {} + ins[7]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[7]) + assert_nnapi_op_check(ins[7]["dtype"] == "INT32") + ins[7]["value"] = compiler.export_obj.json_analyzer.operand.get_value(inputs[7]) + assert_nnapi_op_check(ins[7]["value"] >= 0) + stride = {} + stride["w"] = ins[7]["value"] + + # check inputs[8] + ins[8] = {} + ins[8]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[8]) + assert_nnapi_op_check(ins[8]["dtype"] == "INT32") + ins[8]["value"] = compiler.export_obj.json_analyzer.operand.get_value(inputs[8]) + assert_nnapi_op_check(ins[8]["value"] >= 0) + stride["h"] = ins[8]["value"] + + # check inputs[9] + ins[9] = {} + ins[9]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[9]) + assert_nnapi_op_check(ins[9]["dtype"] == "INT32") + ins[9]["value"] = compiler.export_obj.json_analyzer.operand.get_value(inputs[9]) + depth_multiplier = ins[9]["value"] + assert_nnapi_op_check(depth_multiplier >= 0) + + # check inputs[10] + assert_nnapi_op_check(compiler.export_obj.json_analyzer.operand.is_fuse_code(inputs[10])) + + if api_level >= 29: + # check inputs[11] + ins[11] = {} + ins[11]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[11]) + assert_nnapi_op_check(ins[11]["dtype"] == "BOOL") + ins[11]["value"] = compiler.export_obj.json_analyzer.operand.get_value(inputs[11]) + assert_nnapi_op_check(ins[11]["value"] == "false" or ins[11]["value"] == "true") + + # check inputs[12] + ins[12] = {} + ins[12]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[12]) + assert_nnapi_op_check(ins[12]["dtype"] == "INT32") + ins[12]["value"] = compiler.export_obj.json_analyzer.operand.get_value(inputs[12]) + assert_nnapi_op_check(ins[12]["value"] >= 1) + + # check inputs[13] + ins[13] = {} + ins[13]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[13]) + assert_nnapi_op_check(ins[13]["dtype"] == "INT32") + ins[13]["value"] = compiler.export_obj.json_analyzer.operand.get_value(inputs[13]) + assert_nnapi_op_check(ins[13]["value"] >= 1) + + # check shapes + if api_level >= 29 and ins[11]["value"] == "true": + data_shape = { + "n": ins[0]["shape"][0], + "c": ins[0]["shape"][1], + "h": ins[0]["shape"][2], + "w": ins[0]["shape"][3], + } + else: + data_shape = { + "n": ins[0]["shape"][0], + "h": ins[0]["shape"][1], + "w": ins[0]["shape"][2], + "c": ins[0]["shape"][3], + } + + assert_nnapi_op_check(felter["do"] == data_shape["c"] * depth_multiplier) + + # check outputs + assert_nnapi_op_check(len(outputs) == 1) + outs = [{}] + + # check outputs[0] + outs[0] = {} + outs[0]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(outputs[0]) + assert_nnapi_op_check( + outs[0]["dtype"] == ins[0]["dtype"] and outs[0]["dtype"] == ins[1]["dtype"] + ) + outs[0]["shape"] = compiler.export_obj.json_analyzer.operand.get_shape(outputs[0]) + + if api_level >= 29 and ins[11]["value"] == "true": + out_data_shape = { + "n": outs[0]["shape"][0], + "c": outs[0]["shape"][1], + "h": outs[0]["shape"][2], + "w": outs[0]["shape"][3], + } + else: + out_data_shape = { + "n": outs[0]["shape"][0], + "h": outs[0]["shape"][1], + "w": outs[0]["shape"][2], + "c": outs[0]["shape"][3], + } + total_h = data_shape["h"] + padding["t"] + padding["b"] + total_w = data_shape["w"] + padding["l"] + padding["r"] + assert_nnapi_op_check(out_data_shape["n"] == data_shape["n"]) + assert_nnapi_op_check(out_data_shape["h"] == ((total_h - felter["fh"]) // stride["h"] + 1)) + assert_nnapi_op_check(out_data_shape["w"] == ((total_w - felter["fw"]) // stride["w"] + 1)) + assert_nnapi_op_check(out_data_shape["c"] == felter["do"]) + + compiler.export_obj.add_operation("DEPTHWISE_CONV_2D", inputs, outputs) diff --git a/python/tvm/contrib/target/android_nnapi/operation_utils/nnapi_op/error.py b/python/tvm/contrib/target/android_nnapi/operation_utils/nnapi_op/error.py new file mode 100644 index 000000000000..a3cf0c378670 --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/operation_utils/nnapi_op/error.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,wildcard-import,unused-wildcard-import +"""Namespace for errors encountered during checks of outputting +Android NNAPI operations.""" +from ...error import * + + +class AndroidNNAPICompilerBadNNAPIOperationError(AndroidNNAPICompilerError): + """Error caused by unexpected parse result of the Relay AST. + + Parameters + ---------- + msg: str + The error message. + """ + + +def assert_nnapi_op_check(boolean, *msg): + """Check for True or raise an AndroidNNAPICompilerBadNNAPIOperationError. + + Parameters + ---------- + boolean: bool + The condition to be checked + + msg: str + Optional error message to be raised + """ + if not boolean: + raise AndroidNNAPICompilerBadNNAPIOperationError(*msg) diff --git a/python/tvm/contrib/target/android_nnapi/operation_utils/nnapi_op/grouped_conv_2d.py b/python/tvm/contrib/target/android_nnapi/operation_utils/nnapi_op/grouped_conv_2d.py new file mode 100644 index 000000000000..0636fa3d37a5 --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/operation_utils/nnapi_op/grouped_conv_2d.py @@ -0,0 +1,199 @@ +# 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=wildcard-import,unused-wildcard-import +"""Add an ANEURALNETWORKS_GROUPED_CONV_2D operation with checking.""" +from .error import * + + +def add_operation(compiler, inputs, outputs): + """Add an ANEURALNETWORKS_GROUPED_CONV_2D operation with checking. + + Parameters + ---------- + compiler: FunctionToJsonCompiler + the compiler object holding export_obj. + + inputs: list of int + inputs to the operation. + + outputs: list of int + outputs of the operation. + """ + api_level = compiler.options["target"]["api_level"] + assert_anc_compatibility( + api_level >= 29, + f"Target Android API level { api_level } is too low to support the operation", + ) + + # check inputs + assert_nnapi_op_check(len(inputs) == 12) + ins = [{} for i in range(len(inputs))] + + # check inputs[0] + ins[0] = {} + ins[0]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[0]) + assert_nnapi_op_check( + ins[0]["dtype"] == "TENSOR_FLOAT32" or ins[0]["dtype"] == "TENSOR_FLOAT16" + ) + ins[0]["rank"] = compiler.export_obj.json_analyzer.operand.get_rank(inputs[0]) + assert_nnapi_op_check(ins[0]["rank"] == 4) + ins[0]["shape"] = compiler.export_obj.json_analyzer.operand.get_shape(inputs[0]) + + # check inputs[1] + ins[1] = {} + ins[1]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[1]) + assert_nnapi_op_check( + ins[1]["dtype"] == "TENSOR_FLOAT32" or ins[1]["dtype"] == "TENSOR_FLOAT16" + ) + assert_nnapi_op_check(ins[1]["dtype"] == ins[0]["dtype"]) + ins[1]["rank"] = compiler.export_obj.json_analyzer.operand.get_rank(inputs[1]) + assert_nnapi_op_check(ins[1]["rank"] == 4) + ins[1]["shape"] = compiler.export_obj.json_analyzer.operand.get_shape(inputs[1]) + felter = dict(zip(["do", "fh", "fw", "dg"], ins[1]["shape"])) + + # check inputs[2] + ins[2] = {} + ins[2]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[2]) + assert_nnapi_op_check(ins[2]["dtype"] == ins[1]["dtype"] and ins[2]["dtype"] == ins[0]["dtype"]) + ins[2]["rank"] = compiler.export_obj.json_analyzer.operand.get_rank(inputs[2]) + assert_nnapi_op_check(ins[2]["rank"] == 1) + ins[2]["constant"] = compiler.export_obj.json_analyzer.operand.get_constant(inputs[2]) + assert_nnapi_op_check( + ins[2]["constant"]["type"] == "array" and len(ins[2]["constant"]["value"]) == felter["do"] + ) + + # check inputs[3] + ins[3] = {} + ins[3]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[3]) + assert_nnapi_op_check(ins[3]["dtype"] == "INT32") + ins[3]["value"] = compiler.export_obj.json_analyzer.operand.get_value(inputs[3]) + assert_nnapi_op_check(ins[3]["value"] >= 0) + padding = {} + padding["l"] = ins[3]["value"] + + # check inputs[4] + ins[4] = {} + ins[4]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[4]) + assert_nnapi_op_check(ins[4]["dtype"] == "INT32") + ins[4]["value"] = compiler.export_obj.json_analyzer.operand.get_value(inputs[4]) + assert_nnapi_op_check(ins[4]["value"] >= 0) + padding["r"] = ins[4]["value"] + + # check inputs[5] + ins[5] = {} + ins[5]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[5]) + assert_nnapi_op_check(ins[5]["dtype"] == "INT32") + ins[5]["value"] = compiler.export_obj.json_analyzer.operand.get_value(inputs[5]) + assert_nnapi_op_check(ins[5]["value"] >= 0) + padding["t"] = ins[5]["value"] + + # check inputs[6] + ins[6] = {} + ins[6]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[6]) + assert_nnapi_op_check(ins[6]["dtype"] == "INT32") + ins[6]["value"] = compiler.export_obj.json_analyzer.operand.get_value(inputs[6]) + assert_nnapi_op_check(ins[6]["value"] >= 0) + padding["b"] = ins[6]["value"] + + # check inputs[7] + ins[7] = {} + ins[7]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[7]) + assert_nnapi_op_check(ins[7]["dtype"] == "INT32") + ins[7]["value"] = compiler.export_obj.json_analyzer.operand.get_value(inputs[7]) + assert_nnapi_op_check(ins[7]["value"] >= 0) + stride = {} + stride["w"] = ins[7]["value"] + + # check inputs[8] + ins[8] = {} + ins[8]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[8]) + assert_nnapi_op_check(ins[8]["dtype"] == "INT32") + ins[8]["value"] = compiler.export_obj.json_analyzer.operand.get_value(inputs[8]) + assert_nnapi_op_check(ins[8]["value"] >= 0) + stride["h"] = ins[8]["value"] + + # check inputs[9] + ins[9] = {} + ins[9]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[9]) + assert_nnapi_op_check(ins[9]["dtype"] == "INT32") + ins[9]["value"] = compiler.export_obj.json_analyzer.operand.get_value(inputs[9]) + num_groups = ins[9]["value"] + assert_nnapi_op_check(num_groups >= 0) + assert_nnapi_op_check(felter["do"] % num_groups == 0) + + # check inputs[10] + assert_nnapi_op_check(compiler.export_obj.json_analyzer.operand.is_fuse_code(inputs[10])) + + # check inputs[11] + ins[11] = {} + ins[11]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[11]) + assert_nnapi_op_check(ins[11]["dtype"] == "BOOL") + ins[11]["value"] = compiler.export_obj.json_analyzer.operand.get_value(inputs[11]) + assert_nnapi_op_check(ins[11]["value"] == "false" or ins[11]["value"] == "true") + + # check shapes + if api_level >= 29 and ins[11]["value"] == "true": + data_shape = { + "n": ins[0]["shape"][0], + "c": ins[0]["shape"][1], + "h": ins[0]["shape"][2], + "w": ins[0]["shape"][3], + } + else: + data_shape = { + "n": ins[0]["shape"][0], + "h": ins[0]["shape"][1], + "w": ins[0]["shape"][2], + "c": ins[0]["shape"][3], + } + + assert_nnapi_op_check(data_shape["c"] == num_groups * felter["dg"]) + + # check outputs + assert_nnapi_op_check(len(outputs) == 1) + outs = [{}] + + # check outputs[0] + outs[0] = {} + outs[0]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(outputs[0]) + assert_nnapi_op_check( + outs[0]["dtype"] == ins[0]["dtype"] and outs[0]["dtype"] == ins[1]["dtype"] + ) + outs[0]["shape"] = compiler.export_obj.json_analyzer.operand.get_shape(outputs[0]) + + if api_level >= 29 and ins[11]["value"] == "true": + out_data_shape = { + "n": outs[0]["shape"][0], + "c": outs[0]["shape"][1], + "h": outs[0]["shape"][2], + "w": outs[0]["shape"][3], + } + else: + out_data_shape = { + "n": outs[0]["shape"][0], + "h": outs[0]["shape"][1], + "w": outs[0]["shape"][2], + "c": outs[0]["shape"][3], + } + total_h = data_shape["h"] + padding["t"] + padding["b"] + total_w = data_shape["w"] + padding["l"] + padding["r"] + assert_nnapi_op_check(out_data_shape["n"] == data_shape["n"]) + assert_nnapi_op_check(out_data_shape["h"] == ((total_h - felter["fh"]) // stride["h"] + 1)) + assert_nnapi_op_check(out_data_shape["w"] == ((total_w - felter["fw"]) // stride["w"] + 1)) + assert_nnapi_op_check(out_data_shape["c"] == felter["do"]) + + compiler.export_obj.add_operation("GROUPED_CONV_2D", inputs, outputs) diff --git a/python/tvm/contrib/target/android_nnapi/operation_utils/nnapi_op/transpose.py b/python/tvm/contrib/target/android_nnapi/operation_utils/nnapi_op/transpose.py new file mode 100644 index 000000000000..047f5483a533 --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/operation_utils/nnapi_op/transpose.py @@ -0,0 +1,80 @@ +# 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=wildcard-import,unused-wildcard-import +"""Add an ANEURALNETWORKS_TRANSPOSE operation with checking.""" +from .error import * + + +def add_operation(compiler, inputs, outputs): + """Add an ANEURALNETWORKS_TRANSPOSE operation with checking. + + Parameters + ---------- + compiler: FunctionToJsonCompiler + the compiler object holding export_obj. + + inputs: list of int + inputs to the operation. + + outputs: list of int + outputs of the operation. + """ + api_level = compiler.options["target"]["api_level"] + assert_anc_compatibility( + api_level >= 28, + f"Target Android API level { api_level } is too low to support the operation", + ) + + # check inputs + assert_nnapi_op_check(len(inputs) == 2) + ins = [{}, {}] + + # check inputs[0] + ins[0] = {} + ins[0]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[0]) + if ins[0]["dtype"] == "TENSOR_FLOAT16": + assert_nnapi_op_check(api_level >= 29) + else: + assert_nnapi_op_check(ins[0]["dtype"] == "TENSOR_FLOAT32") + ins[0]["shape"] = compiler.export_obj.json_analyzer.operand.get_shape(inputs[0]) + ins[0]["rank"] = compiler.export_obj.json_analyzer.operand.get_rank(inputs[0]) + assert_nnapi_op_check(ins[0]["rank"] <= 4) + + # check inputs[1] + ins[1] = {} + ins[1]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(inputs[1]) + assert_nnapi_op_check(ins[1]["dtype"] == "TENSOR_INT32") + ins[1]["rank"] = compiler.export_obj.json_analyzer.operand.get_rank(inputs[1]) + assert_nnapi_op_check(ins[1]["rank"] == 1) + ins[1]["constant"] = compiler.export_obj.json_analyzer.operand.get_constant(inputs[1]) + assert_nnapi_op_check( + ins[1]["constant"]["type"] == "array" and len(ins[1]["constant"]["value"]) == ins[0]["rank"] + ) + ins[1]["value"] = compiler.export_obj.json_analyzer.operand.get_value(inputs[1]) + + # check outputs + assert_nnapi_op_check(len(outputs) == 1) + outs = [{}] + + # check outputs[0] + outs[0] = {} + outs[0]["dtype"] = compiler.export_obj.json_analyzer.operand.get_dtype(outputs[0]) + assert_nnapi_op_check(outs[0]["dtype"] == ins[0]["dtype"]) + outs[0]["shape"] = compiler.export_obj.json_analyzer.operand.get_shape(outputs[0]) + assert_nnapi_op_check(outs[0]["shape"] == [ins[0]["shape"][i] for i in ins[1]["value"]]) + + compiler.export_obj.add_operation("TRANSPOSE", inputs, outputs) diff --git a/python/tvm/contrib/target/android_nnapi/operation_utils/relay_op/__init__.py b/python/tvm/contrib/target/android_nnapi/operation_utils/relay_op/__init__.py new file mode 100644 index 000000000000..e6181c2e63b2 --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/operation_utils/relay_op/__init__.py @@ -0,0 +1,18 @@ +# 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. +"""Namespace for tvm.relay.Call handler(parser).""" +from . import nn diff --git a/python/tvm/contrib/target/android_nnapi/operation_utils/relay_op/nn/__init__.py b/python/tvm/contrib/target/android_nnapi/operation_utils/relay_op/nn/__init__.py new file mode 100644 index 000000000000..430f6d6188a6 --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/operation_utils/relay_op/nn/__init__.py @@ -0,0 +1,18 @@ +# 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. +"""Namespace for tvm.relay.Call handler(parser) for nn.* operators.""" +from . import conv2d diff --git a/python/tvm/contrib/target/android_nnapi/operation_utils/relay_op/nn/conv2d.py b/python/tvm/contrib/target/android_nnapi/operation_utils/relay_op/nn/conv2d.py new file mode 100644 index 000000000000..f9cac70f1eaa --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/operation_utils/relay_op/nn/conv2d.py @@ -0,0 +1,929 @@ +# 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=wildcard-import,unused-wildcard-import +"""Relay IR handler(parser) for tvm.relay.nn.conv2d.""" +from ....error import * +from ... import _utils +from ... import nnapi_op + + +def handler(compiler, node): + """Handler for tvm.relay.nn.conv2d. + + Parameters + ---------- + compiler: FunctionToJsonCompiler + the compiler object holding export_obj. + + node: relay.Call + operation call node. + """ + args = _utils.name_args(node.args, ["data", "weight"]) + attrs = node.attrs + ngroups = int(attrs.groups) + channel_dims = int(args["data"].checked_type.shape[attrs.data_layout.index("C")]) + output_dims = int(args["weight"].checked_type.shape[attrs.kernel_layout.index("O")]) + input_dims = int(args["weight"].checked_type.shape[attrs.kernel_layout.index("I")]) + if ngroups == 1: # classic convolution + _1_group_handler(compiler, node) + elif ngroups == channel_dims and channel_dims == output_dims and input_dims == 1: + _depthwise_handler(compiler, node) + else: + _grouped_handler(compiler, node) + + +def _1_group_handler(compiler, node): + api_level = compiler.options["target"]["api_level"] + args = _utils.name_args(node.args, ["data", "weight"]) + attrs = node.attrs + nnapi = {} + + # START: handle inputs + # use explicit padding of ANEURALNETWORKS_CONV_2D + nnapi["inputs"] = [] + + # START: handle input[0] + # check compatibility + assert_anc_compatibility( + args["data"].checked_type.dtype == "float32" or args["data"].checked_type.dtype == "float16" + ) + + # generate nnapi node of "data" + compiler.visit(args["data"]) + + # change layout of "data" to NNAPI's NHWC + assert_anc_compatibility( + len(attrs.data_layout) == 4, f"Unrecognized layout {attrs.data_layout}" + ) + if attrs.data_layout == "NHWC" or (api_level >= 29 and attrs.data_layout == "NCHW"): + nnapi["inputs"] += compiler.export_obj.get_node_operand_idxs(args["data"]) + else: + # START: add TRANSPOSE + transpose_idxs = list(map(attrs.data_layout.index, ["N", "H", "W", "C"])) + inputs = [] + inputs += compiler.export_obj.get_node_operand_idxs(args["data"]) + inputs += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx(((4,), "int32")), + value={ + "type": "constant_idx", + "value": compiler.export_obj.add_array_constant( + vals=transpose_idxs, + dtype="int32", + ), + }, + ) + outputs = [] + outputs += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx( + ( + tuple(map(lambda ele: args["data"].checked_type.shape[ele], transpose_idxs)), + args["data"].checked_type.dtype, + ) + ) + ) + nnapi_op.transpose.add_operation(compiler, inputs, outputs) + nnapi["inputs"] += outputs + # END: add TRANSPOSE + # END: handle input[0] + + # START: handle input[1] + # check compatibility + assert_anc_compatibility(args["weight"].checked_type.dtype == args["data"].checked_type.dtype) + + # generate nnapi node for weight + compiler.visit(args["weight"]) + + # change layout of "weight" to NNAPI's OHWI + assert_anc_compatibility( + len(attrs.kernel_layout) == 4, f"Unrecognized layout {attrs.kernel_layout}" + ) + if attrs.kernel_layout == "OHWI": + nnapi["inputs"] += compiler.export_obj.get_node_operand_idxs(args["weight"]) + else: + # START: add TRANSPOSE + transpose_idxs = list(map(attrs.kernel_layout.index, ["O", "H", "W", "I"])) + inputs = [] + inputs += compiler.export_obj.get_node_operand_idxs(args["weight"]) + inputs += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx(((4,), "int32")), + value={ + "type": "constant_idx", + "value": compiler.export_obj.add_array_constant( + vals=transpose_idxs, + dtype="int32", + ), + }, + ) + outputs = [] + outputs += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx( + ( + tuple(map(lambda ele: args["weight"].checked_type.shape[ele], transpose_idxs)), + args["weight"].checked_type.dtype, + ) + ) + ) + nnapi_op.transpose.add_operation(compiler, inputs, outputs) + nnapi["inputs"] += outputs + # END: add TRANSPOSE + # END: handle input[1] + + # START: handle input[2] + # add empty bias since CONV_2D needs it + bias_shape = (compiler.export_obj.json_analyzer.operand.get_shape(nnapi["inputs"][1])[0],) + if args["data"].checked_type.dtype == "float32" or args["data"].checked_type.dtype == "float16": + bias_dtype = args["data"].checked_type.dtype + else: + raise AndroidNNAPICompilerIncompatibleError( + f"Unable to determine bias data type for CONV_2D. \ + args['data'].dtype was {args['data'].checked_type.dtype}" + ) + bias_type = (bias_shape, bias_dtype) + nnapi["inputs"] += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx(bias_type), + value={ + "type": "constant_idx", + "value": compiler.export_obj.add_array_constant( + vals=[0.0] * bias_shape[0], + dtype=bias_dtype, + ), + }, + ) + # END: handle input[2] + + # START: handle input[3:7] + def _add_int32_scalar_constant(ele): + return compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx(((), "int32")), + value={ + "type": "constant_idx", + "value": compiler.export_obj.add_scalar_constant( + val=int(ele), + dtype="int32", + ), + }, + )[0] + + relay_paddings = list(map(_add_int32_scalar_constant, attrs.padding)) + if len(relay_paddings) == 2: + nnapi["inputs"] += [ + relay_paddings[1], + relay_paddings[1], + relay_paddings[0], + relay_paddings[0], + ] + elif len(relay_paddings) == 4: + nnapi["inputs"] += [ + relay_paddings[1], + relay_paddings[3], + relay_paddings[0], + relay_paddings[2], + ] + else: + raise AndroidNNAPICompilerIncompatibleError(f"Unexpected padding format {attrs.padding}") + # END: handle input[3:7] + + # START: handle input[7:9] + relay_strides = list(map(_add_int32_scalar_constant, attrs.strides)) + nnapi["inputs"] += [relay_strides[1], relay_strides[0]] + # END: handle input[7:9] + + # START: handle input[9] + # add ANEURALNETWORKS_FUSED_NONE activation since CONV_2D needs it + nnapi["inputs"] += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx(((), "int32")), + value={ + "type": "constant_idx", + "value": compiler.export_obj.add_scalar_constant( + val="ANEURALNETWORKS_FUSED_NONE", + dtype="int32", + ), + }, + ) + # END: handle input[9] + + nnapi_output_layout = "NHWC" + if api_level >= 29: + # START: handle input[10] + if attrs.data_layout == "NCHW": + nnapi["inputs"] += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx(((), "bool")), + value={ + "type": "constant_idx", + "value": compiler.export_obj.add_scalar_constant( + val="true", + dtype="bool", + ), + }, + ) + nnapi_output_layout = "NCHW" + else: + nnapi["inputs"] += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx(((), "bool")), + value={ + "type": "constant_idx", + "value": compiler.export_obj.add_scalar_constant( + val="false", + dtype="bool", + ), + }, + ) + # END: handle input[10] + + # START: handle input[11:] + # unpack dilation + relay_dilations = list(map(_add_int32_scalar_constant, attrs.dilation)) + nnapi["inputs"] += [relay_dilations[1], relay_dilations[0]] + # END: handle input[11:] + # END: handle inputs + + # START: handle outputs + nnapi["outputs"] = [] + + # START: handle output[0] + attrs_out_layout = attrs.data_layout if attrs.out_layout == "" else attrs.out_layout + attrs_out_dtype = args["data"].checked_type.dtype if attrs.out_dtype == "" else attrs.out_dtype + if ( + attrs_out_dtype == args["data"].checked_type.dtype + and attrs_out_layout == nnapi_output_layout + ): + nnapi["outputs"] += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx( + (node.checked_type.shape, node.checked_type.dtype) + ) + ) + node_operands = nnapi["outputs"] + else: + if attrs_out_layout == nnapi_output_layout: + nnapi["outputs"] += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx( + (node.checked_type.shape, args["data"].checked_type.dtype) + ) + ) + last_outputs = nnapi["outputs"] + else: + transpose_idxs = list(map(attrs_out_layout.index, ["N", "H", "W", "C"])) + nhwc_shape = tuple(map(lambda ele: node.checked_type.shape[ele], transpose_idxs)) + nnapi["outputs"] += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx( + (nhwc_shape, args["data"].checked_type.dtype) + ) + ) + + # START: add TRANSPOSE + rev_transpose_idxs = list(map("NHWC".index, attrs_out_layout)) + inputs = [] + inputs += nnapi["outputs"] + inputs += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx(((4,), "int32")), + value={ + "type": "constant_idx", + "value": compiler.export_obj.add_array_constant( + vals=rev_transpose_idxs, + dtype="int32", + ), + }, + ) + outputs = [] + outputs += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx( + (node.checked_type.shape, args["data"].checked_type.dtype) + ) + ) + nnapi_op.transpose.add_operation(compiler, inputs, outputs) + # END: add TRANSPOSE + + last_outputs = outputs + + if attrs_out_dtype == args["data"].checked_type.dtype: + node_operands = last_outputs + else: + # START: add CAST + inputs = [] + inputs += last_outputs + outputs = [] + outputs += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx( + (node.checked_type.shape, node.checked_type.dtype) + ) + ) + nnapi_op.cast.add_operation(compiler, inputs, outputs) + # END: add CAST + + node_operands = outputs + + # register operands to node + compiler.export_obj.register_node_operand_idxs(node, node_operands) + # END: handle output[0] + # END: handle outputs + + nnapi_op.conv_2d.add_operation(compiler, nnapi["inputs"], nnapi["outputs"]) + + +def _depthwise_handler(compiler, node): + api_level = compiler.options["target"]["api_level"] + args = _utils.name_args(node.args, ["data", "weight"]) + attrs = node.attrs + nnapi = {} + + # START: handle inputs + # use explicit padding + nnapi["inputs"] = [] + + # START: handle input[0] + # generate nnapi node of "data" + compiler.visit(args["data"]) + + # change layout of "data" to NNAPI's NHWC + assert_anc_compatibility( + len(attrs.data_layout) == 4, f"Unrecognized layout {attrs.data_layout}" + ) + if attrs.data_layout == "NHWC" or (api_level >= 29 and attrs.data_layout == "NCHW"): + nnapi["inputs"] += compiler.export_obj.get_node_operand_idxs(args["data"]) + else: + # START: add TRANSPOSE + transpose_idxs = list(map(attrs.data_layout.index, ["N", "H", "W", "C"])) + inputs = [] + inputs += compiler.export_obj.get_node_operand_idxs(args["data"]) + inputs += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx(((4,), "int32")), + value={ + "type": "constant_idx", + "value": compiler.export_obj.add_array_constant( + vals=transpose_idxs, + dtype="int32", + ), + }, + ) + outputs = [] + outputs += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx( + ( + tuple(map(lambda ele: args["data"].checked_type.shape[ele], transpose_idxs)), + args["data"].checked_type.dtype, + ) + ) + ) + nnapi_op.transpose.add_operation(compiler, inputs, outputs) + nnapi["inputs"] += outputs + # END: add TRANSPOSE + # END: handle input[0] + + # START: handle input[1] + # check compatibility + assert_anc_compatibility(args["weight"].checked_type.dtype == args["data"].checked_type.dtype) + + # generate nnapi node for weight + compiler.visit(args["weight"]) + + # change layout of "weight" to NNAPI's IHWO + assert_anc_compatibility( + len(attrs.kernel_layout) == 4, f"Unrecognized layout {attrs.kernel_layout}" + ) + if attrs.kernel_layout == "IHWO": + nnapi["inputs"] += compiler.export_obj.get_node_operand_idxs(args["weight"]) + else: + # START: add TRANSPOSE + transpose_idxs = list(map(attrs.kernel_layout.index, ["I", "H", "W", "O"])) + inputs = [] + inputs += compiler.export_obj.get_node_operand_idxs(args["weight"]) + inputs += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx(((4,), "int32")), + value={ + "type": "constant_idx", + "value": compiler.export_obj.add_array_constant( + vals=transpose_idxs, + dtype="int32", + ), + }, + ) + outputs = [] + outputs += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx( + ( + tuple(map(lambda ele: args["weight"].checked_type.shape[ele], transpose_idxs)), + args["weight"].checked_type.dtype, + ) + ) + ) + nnapi_op.transpose.add_operation(compiler, inputs, outputs) + nnapi["inputs"] += outputs + # END: add TRANSPOSE + # END: handle input[1] + + # START: handle input[2] + # add empty bias + bias_shape = (compiler.export_obj.json_analyzer.operand.get_shape(nnapi["inputs"][1])[3],) + if args["data"].checked_type.dtype == "float32" or args["data"].checked_type.dtype == "float16": + bias_dtype = args["data"].checked_type.dtype + else: + raise AndroidNNAPICompilerIncompatibleError( + f"Unable to determine bias data type for \ + DEPTHWISE_CONV_2D. args['data'].dtype was {args['data'].checked_type.dtype}" + ) + bias_type = (bias_shape, bias_dtype) + nnapi["inputs"] += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx(bias_type), + value={ + "type": "constant_idx", + "value": compiler.export_obj.add_array_constant( + vals=[0.0] * bias_shape[0], + dtype=bias_dtype, + ), + }, + ) + # END: handle input[2] + + # START: handle input[3:7] + def _add_int32_scalar_constant(ele): + return compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx(((), "int32")), + value={ + "type": "constant_idx", + "value": compiler.export_obj.add_scalar_constant( + val=int(ele), + dtype="int32", + ), + }, + )[0] + + relay_paddings = list(map(_add_int32_scalar_constant, attrs.padding)) + if len(relay_paddings) == 2: + nnapi["inputs"] += [ + relay_paddings[1], + relay_paddings[1], + relay_paddings[0], + relay_paddings[0], + ] + elif len(relay_paddings) == 4: + nnapi["inputs"] += [ + relay_paddings[1], + relay_paddings[3], + relay_paddings[0], + relay_paddings[2], + ] + else: + raise AndroidNNAPICompilerIncompatibleError(f"Unexpected padding format {attrs.padding}") + # END: handle input[3:7] + + # START: handle input[7:9] + relay_strides = list(map(_add_int32_scalar_constant, attrs.strides)) + nnapi["inputs"] += [relay_strides[1], relay_strides[0]] + # END: handle input[7:9] + + # START: handle input[9] + def _scope(): + if api_level >= 29 and attrs.data_layout == "NCHW": + depth_in = compiler.export_obj.json_analyzer.operand.get_shape(nnapi["inputs"][0])[1] + else: + depth_in = compiler.export_obj.json_analyzer.operand.get_shape(nnapi["inputs"][0])[3] + depth_out = compiler.export_obj.json_analyzer.operand.get_shape(nnapi["inputs"][1])[3] + assert depth_out % depth_in == 0 + depth_multiplier = int(depth_out // depth_in) + nnapi["inputs"] += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx(((), "int32")), + value={ + "type": "constant_idx", + "value": compiler.export_obj.add_scalar_constant( + val=depth_multiplier, + dtype="int32", + ), + }, + ) + + _scope() + # END: handle input[9] + + # START: handle input[10] + # add ANEURALNETWORKS_FUSED_NONE activation since CONV_2D needs it + nnapi["inputs"] += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx(((), "int32")), + value={ + "type": "constant_idx", + "value": compiler.export_obj.add_scalar_constant( + val="ANEURALNETWORKS_FUSED_NONE", + dtype="int32", + ), + }, + ) + # END: handle input[10] + + nnapi_output_layout = "NHWC" + if api_level >= 29: + # START: handle input[11] + if attrs.data_layout == "NCHW": + nnapi["inputs"] += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx(((), "bool")), + value={ + "type": "constant_idx", + "value": compiler.export_obj.add_scalar_constant( + val="true", + dtype="bool", + ), + }, + ) + nnapi_output_layout = "NCHW" + else: + nnapi["inputs"] += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx(((), "bool")), + value={ + "type": "constant_idx", + "value": compiler.export_obj.add_scalar_constant( + val="false", + dtype="bool", + ), + }, + ) + # END: handle input[11] + + # START: handle input[12:] + # unpack dilation + relay_dilations = list(map(_add_int32_scalar_constant, attrs.dilation)) + nnapi["inputs"] += [relay_dilations[1], relay_dilations[0]] + # END: handle input[12:] + # END: handle inputs + + # START: handle outputs + nnapi["outputs"] = [] + + # START: handle output[0] + attrs_out_layout = attrs.data_layout if attrs.out_layout == "" else attrs.out_layout + attrs_out_dtype = args["data"].checked_type.dtype if attrs.out_dtype == "" else attrs.out_dtype + if ( + attrs_out_dtype == args["data"].checked_type.dtype + and attrs_out_layout == nnapi_output_layout + ): + nnapi["outputs"] += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx( + (node.checked_type.shape, node.checked_type.dtype) + ) + ) + node_operands = nnapi["outputs"] + else: + if attrs_out_layout == nnapi_output_layout: + nnapi["outputs"] += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx( + (node.checked_type.shape, args["data"].checked_type.dtype) + ) + ) + last_outputs = nnapi["outputs"] + else: + transpose_idxs = list(map(attrs_out_layout.index, ["N", "H", "W", "C"])) + nhwc_shape = tuple(map(lambda ele: node.checked_type.shape[ele], transpose_idxs)) + nnapi["outputs"] += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx( + (nhwc_shape, args["data"].checked_type.dtype) + ) + ) + + # START: add TRANSPOSE + rev_transpose_idxs = list(map("NHWC".index, attrs_out_layout)) + inputs = [] + inputs += nnapi["outputs"] + inputs += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx(((4,), "int32")), + value={ + "type": "constant_idx", + "value": compiler.export_obj.add_array_constant( + vals=rev_transpose_idxs, + dtype="int32", + ), + }, + ) + outputs = [] + outputs += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx( + (node.checked_type.shape, args["data"].checked_type.dtype) + ) + ) + nnapi_op.transpose.add_operation(compiler, inputs, outputs) + # END: add TRANSPOSE + + last_outputs = outputs + + if attrs_out_dtype == args["data"].checked_type.dtype: + node_operands = last_outputs + else: + # START: add CAST + inputs = [] + inputs += last_outputs + outputs = [] + outputs += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx( + (node.checked_type.shape, node.checked_type.dtype) + ) + ) + nnapi_op.cast.add_operation(compiler, inputs, outputs) + # END: add CAST + + node_operands = outputs + + # register operands to node + compiler.export_obj.register_node_operand_idxs(node, node_operands) + # END: handle output[0] + # END: handle outputs + + nnapi_op.depthwise_conv_2d.add_operation(compiler, nnapi["inputs"], nnapi["outputs"]) + + +def _grouped_handler(compiler, node): + api_level = compiler.options["target"]["api_level"] + args = _utils.name_args(node.args, ["data", "weight"]) + attrs = node.attrs + nnapi = {} + + # START: handle inputs + # use explicit padding + nnapi["inputs"] = [] + + # START: handle input[0] + # generate nnapi node of "data" + compiler.visit(args["data"]) + + # change layout of "data" to NNAPI's NHWC + assert_anc_compatibility( + len(attrs.data_layout) == 4, f"Unrecognized layout {attrs.data_layout}" + ) + if attrs.data_layout == "NHWC" or (api_level >= 29 and attrs.data_layout == "NCHW"): + nnapi["inputs"] += compiler.export_obj.get_node_operand_idxs(args["data"]) + else: + # START: add TRANSPOSE + transpose_idxs = list(map(attrs.data_layout.index, ["N", "H", "W", "C"])) + inputs = [] + inputs += compiler.export_obj.get_node_operand_idxs(args["data"]) + inputs += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx(((4,), "int32")), + value={ + "type": "constant_idx", + "value": compiler.export_obj.add_array_constant( + vals=transpose_idxs, + dtype="int32", + ), + }, + ) + outputs = [] + outputs += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx( + ( + tuple(map(lambda ele: args["data"].checked_type.shape[ele], transpose_idxs)), + args["data"].checked_type.dtype, + ) + ) + ) + nnapi_op.transpose.add_operation(compiler, inputs, outputs) + nnapi["inputs"] += outputs + # END: add TRANSPOSE + # END: handle input[0] + + # START: handle input[1] + # check compatibility + assert_anc_compatibility(args["weight"].checked_type.dtype == args["data"].checked_type.dtype) + + # generate nnapi node for weight + compiler.visit(args["weight"]) + + # change layout of "weight" to NNAPI's OHWI + assert_anc_compatibility( + len(attrs.kernel_layout) == 4, f"Unrecognized layout {attrs.kernel_layout}" + ) + if attrs.kernel_layout == "OHWI": + nnapi["inputs"] += compiler.export_obj.get_node_operand_idxs(args["weight"]) + else: + # START: add TRANSPOSE + transpose_idxs = list(map(attrs.kernel_layout.index, ["O", "H", "W", "I"])) + inputs = [] + inputs += compiler.export_obj.get_node_operand_idxs(args["weight"]) + inputs += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx(((4,), "int32")), + value={ + "type": "constant_idx", + "value": compiler.export_obj.add_array_constant( + vals=transpose_idxs, + dtype="int32", + ), + }, + ) + outputs = [] + outputs += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx( + ( + tuple(map(lambda ele: args["weight"].checked_type.shape[ele], transpose_idxs)), + args["weight"].checked_type.dtype, + ) + ) + ) + nnapi_op.transpose.add_operation(compiler, inputs, outputs) + nnapi["inputs"] += outputs + # END: add TRANSPOSE + # END: handle input[1] + + # START: handle input[2] + # add empty bias + bias_shape = (compiler.export_obj.json_analyzer.operand.get_shape(nnapi["inputs"][1])[0],) + if args["data"].checked_type.dtype == "float32" or args["data"].checked_type.dtype == "float16": + bias_dtype = args["data"].checked_type.dtype + else: + raise AndroidNNAPICompilerIncompatibleError( + f"Unable to determine bias type for GROUPED_CONV_2D. \ + args['data'].dtype was {args['data'].checked_type.dtype}" + ) + bias_type = (bias_shape, bias_dtype) + nnapi["inputs"] += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx(bias_type), + value={ + "type": "constant_idx", + "value": compiler.export_obj.add_array_constant( + vals=[0.0] * bias_shape[0], + dtype=bias_dtype, + ), + }, + ) + # END: handle input[2] + + # START: handle input[3:7] + def _add_int32_scalar_constant(ele): + return compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx(((), "int32")), + value={ + "type": "constant_idx", + "value": compiler.export_obj.add_scalar_constant( + val=int(ele), + dtype="int32", + ), + }, + )[0] + + relay_paddings = list(map(_add_int32_scalar_constant, attrs.padding)) + if len(relay_paddings) == 2: + nnapi["inputs"] += [ + relay_paddings[1], + relay_paddings[1], + relay_paddings[0], + relay_paddings[0], + ] + elif len(relay_paddings) == 4: + nnapi["inputs"] += [ + relay_paddings[1], + relay_paddings[3], + relay_paddings[0], + relay_paddings[2], + ] + else: + raise AndroidNNAPICompilerIncompatibleError(f"Unexpected padding format {attrs.padding}") + # END: handle input[3:7] + + # START: handle input[7:9] + relay_strides = list(map(_add_int32_scalar_constant, attrs.strides)) + nnapi["inputs"] += [relay_strides[1], relay_strides[0]] + # END: handle input[7:9] + + # START: handle input[9] + nnapi["inputs"] += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx(((), "int32")), + value={ + "type": "constant_idx", + "value": compiler.export_obj.add_scalar_constant( + val=int(attrs.groups), + dtype="int32", + ), + }, + ) + # END: handle input[9] + + # START: handle input[10] + # add ANEURALNETWORKS_FUSED_NONE activation since CONV_2D needs it + nnapi["inputs"] += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx(((), "int32")), + value={ + "type": "constant_idx", + "value": compiler.export_obj.add_scalar_constant( + val="ANEURALNETWORKS_FUSED_NONE", + dtype="int32", + ), + }, + ) + # END: handle input[10] + + # START: handle input[11] + nnapi_output_layout = "NHWC" + if attrs.data_layout == "NCHW": + nnapi["inputs"] += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx(((), "bool")), + value={ + "type": "constant_idx", + "value": compiler.export_obj.add_scalar_constant( + val="true", + dtype="bool", + ), + }, + ) + nnapi_output_layout = "NCHW" + else: + nnapi["inputs"] += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx(((), "bool")), + value={ + "type": "constant_idx", + "value": compiler.export_obj.add_scalar_constant( + val="false", + dtype="bool", + ), + }, + ) + # END: handle input[11] + # END: handle inputs + + # START: handle outputs + nnapi["outputs"] = [] + + # START: handle output[0] + attrs_out_layout = attrs.data_layout if attrs.out_layout == "" else attrs.out_layout + attrs_out_dtype = args["data"].checked_type.dtype if attrs.out_dtype == "" else attrs.out_dtype + if ( + attrs_out_dtype == args["data"].checked_type.dtype + and attrs_out_layout == nnapi_output_layout + ): + nnapi["outputs"] += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx( + (node.checked_type.shape, node.checked_type.dtype) + ) + ) + node_operands = nnapi["outputs"] + else: + if attrs_out_layout == nnapi_output_layout: + nnapi["outputs"] += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx( + (node.checked_type.shape, args["data"].checked_type.dtype) + ) + ) + last_outputs = nnapi["outputs"] + else: + transpose_idxs = list(map(attrs_out_layout.index, ["N", "H", "W", "C"])) + nhwc_shape = tuple(map(lambda ele: node.checked_type.shape[ele], transpose_idxs)) + nnapi["outputs"] += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx( + (nhwc_shape, args["data"].checked_type.dtype) + ) + ) + + # START: add TRANSPOSE + rev_transpose_idxs = list(map("NHWC".index, attrs_out_layout)) + inputs = [] + inputs += nnapi["outputs"] + inputs += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx(((4,), "int32")), + value={ + "type": "constant_idx", + "value": compiler.export_obj.add_array_constant( + vals=rev_transpose_idxs, + dtype="int32", + ), + }, + ) + outputs = [] + outputs += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx( + (node.checked_type.shape, args["data"].checked_type.dtype) + ) + ) + nnapi_op.transpose.add_operation(compiler, inputs, outputs) + # END: add TRANSPOSE + + last_outputs = outputs + + if attrs_out_dtype == args["data"].checked_type.dtype: + node_operands = last_outputs + else: + # START: add CAST + inputs = [] + inputs += last_outputs + outputs = [] + outputs += compiler.export_obj.add_operand( + type_idx=compiler.export_obj.get_type_idx( + (node.checked_type.shape, node.checked_type.dtype) + ) + ) + nnapi_op.cast.add_operation(compiler, inputs, outputs) + # END: add CAST + + node_operands = outputs + + # register operands to node + compiler.export_obj.register_node_operand_idxs(node, node_operands) + # END: handle output[0] + # END: handle outputs + + nnapi_op.grouped_conv_2d.add_operation(compiler, nnapi["inputs"], nnapi["outputs"]) diff --git a/python/tvm/contrib/target/android_nnapi/transform/__init__.py b/python/tvm/contrib/target/android_nnapi/transform/__init__.py new file mode 100644 index 000000000000..f3ffc8108c7d --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/transform/__init__.py @@ -0,0 +1,19 @@ +# 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. +"""Transform Relay IR constructs that's not suitable to lower to +Android NNAPI.""" +from .fix_illegal_pattern_for_nnapi import FixIllegalPatternForNnapi diff --git a/python/tvm/contrib/target/android_nnapi/transform/fix_illegal_pattern_for_nnapi/__init__.py b/python/tvm/contrib/target/android_nnapi/transform/fix_illegal_pattern_for_nnapi/__init__.py new file mode 100644 index 000000000000..0e5c1e6b981d --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/transform/fix_illegal_pattern_for_nnapi/__init__.py @@ -0,0 +1,33 @@ +# 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. +"""Transform Relay IR patterns that's not suitable to lower to Android +NNAPI.""" +import tvm.relay +from .convert_scalar_to_tensor_for_broadcast_operators import ( + ConvertScalarToTensorForBroadcastOperators, +) + + +class FixIllegalPatternForNnapi: + def __call__(self, func): + assert isinstance(func, tvm.relay.Function) + passes = [ConvertScalarToTensorForBroadcastOperators()] + func = tvm.relay.transform.InferType()(tvm.IRModule({"main": func}))["main"] + for p in passes: + func = p(func) + func = tvm.relay.transform.InferType()(tvm.IRModule({"main": func}))["main"] + return func diff --git a/python/tvm/contrib/target/android_nnapi/transform/fix_illegal_pattern_for_nnapi/convert_scalar_to_tensor_for_broadcast_operators.py b/python/tvm/contrib/target/android_nnapi/transform/fix_illegal_pattern_for_nnapi/convert_scalar_to_tensor_for_broadcast_operators.py new file mode 100644 index 000000000000..ac7f049ab297 --- /dev/null +++ b/python/tvm/contrib/target/android_nnapi/transform/fix_illegal_pattern_for_nnapi/convert_scalar_to_tensor_for_broadcast_operators.py @@ -0,0 +1,91 @@ +# 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. +"""Convert scalar arguments to a broadcasting operator to its tensor equivalent +for Android NNAPI conversion.""" +import tvm.relay + + +ZEROS_OP = tvm.relay.op.get("zeros") +ADD_OP = tvm.relay.op.get("add") +SUBTRACT_OP = tvm.relay.op.get("subtract") +MULTIPLY_OP = tvm.relay.op.get("multiply") +DIVIDE_OP = tvm.relay.op.get("divide") + + +class ConvertScalarToTensorForBroadcastOperators(tvm.relay.ExprMutator): + """Convert scalar arguments to a broadcasting operator to its tensor equivalent + for Android NNAPI conversion.""" + + def __init__(self): + super().__init__() + self._call_op_stack = [] + + def __call__(self, expr): + return self.visit(expr) + + def visit_call(self, call): + self._call_op_stack.append(call) + if self._parent_is_transform_target() and self._is_scalar(call): + assert ( + isinstance(call.op, tvm.ir.Op) and call.op == zeros + ), "Only tvm.relay.zeros are supported for \ + tvm.relay.Call scalar to tensor transformation" + self._call_op_stack.pop() + return tvm.relay.zeros(shape=(1,), dtype=call.checked_type.dtype) + + ret = super().visit_call(call) + self._call_op_stack.pop() + return ret + + def visit_constant(self, const): + if self._parent_is_transform_target() and self._is_scalar(const): + return tvm.relay.Constant( + tvm.nd.array( + const.data.asnumpy().reshape( + [ + 1, + ] + ) + ) + ) + return super().visit_constant(const) + + def visit_var(self, var): + # due to the need to also transform the parameter dict, + # we only transform scalar variables + assert not self._parent_is_transform_target() or not self._is_scalar( + var + ), "Transforming variable scalar is not supported" + return super().visit_var(var) + + def _parent_is_transform_target(self): + if len(self._call_op_stack) == 0: + return False + + last_call = self._call_op_stack[-1] + if not isinstance(last_call, tvm.ir.Op): + return False + + return last_call.op in { + ADD_OP, + SUBTRACT_OP, + MULTIPLY_OP, + DIVIDE_OP, + } # only these ops are supported for the fix for now + + def _is_scalar(self, node): + return len(node.checked_type.shape) == 0 diff --git a/python/tvm/relay/op/contrib/android_nnapi/__init__.py b/python/tvm/relay/op/contrib/android_nnapi/__init__.py new file mode 100644 index 000000000000..e0200761d591 --- /dev/null +++ b/python/tvm/relay/op/contrib/android_nnapi/__init__.py @@ -0,0 +1,19 @@ +# 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. +"""BYOC partition namespace for Android NNAPI.""" +from ._partitioner.byoc import byoc_partition as byoc_partition_for_android_nnapi +from ._partitioner.rpc import rpc_partition as rpc_partition_for_android_nnapi diff --git a/python/tvm/relay/op/contrib/android_nnapi/_partitioner/__init__.py b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/__init__.py new file mode 100644 index 000000000000..aba4a9fc44ca --- /dev/null +++ b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/__init__.py @@ -0,0 +1,19 @@ +# 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. +"""Internal namespace for BYOC partitioning for Android NNAPI.""" +from . import byoc +from . import rpc diff --git a/python/tvm/relay/op/contrib/android_nnapi/_partitioner/_base/__init__.py b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/_base/__init__.py new file mode 100644 index 000000000000..9f8c9c60ecd9 --- /dev/null +++ b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/_base/__init__.py @@ -0,0 +1,81 @@ +# 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. +"""Common utilities for all Android NNAPI partitioning.""" +import tvm +from . import transform as _transform + + +def pre_partition_transform(mod): + """Perform pre-partition transforms on modules. + + Parameters + ---------- + mod: tvm.IRModule + The module to be transformed. + + Returns + ------- + mod: tvm.IRModule + The transformed module. + """ + mod = tvm.relay.transform.ToGraphNormalForm()(mod) + mod = tvm.relay.transform.RemoveUnusedFunctions()(mod) + mod = tvm.relay.transform.SimplifyInference()(mod) + mod = tvm.relay.transform.DeadCodeElimination(inline_once=True)(mod) + mod = tvm.relay.transform.FoldConstant()(mod) + mod = tvm.relay.transform.EliminateCommonSubexpr()(mod) + mod = tvm.relay.transform.InferType()(mod) + mod = _transform.PruneInferenceAgnosticOperators()(mod) + mod = _transform.TransformRelayOpForNnapi()(mod) + return mod + + +def post_partition_transform( + mod, params, android_nnapi_level=29, external_compiler="android_nnapi" +): + """Perform post-partition transforms on modules. + + Parameters + ---------- + mod: tvm.IRModule + The module to be transformed. + + params: dict of str to tvm.ndarray + The params dict associated to the module. + + android_nnapi_level: int + The targeted Android API level. + + external_compiler: str + The name of the external Relay compiler. + + Returns + ------- + mod: tvm.IRModule + The transformed module. + + params: dict of str to NDArray + The transformed params. + """ + mod = _transform.AnnotateNnapiFunctionAttributes( + external_compiler=external_compiler, android_nnapi_level=android_nnapi_level + )(mod) + mod, params = _transform.TransformConv2dWeightLayout( + external_compiler=external_compiler, target_layout="OHWI" + )(mod, params) + mod = tvm.relay.transform.LambdaLift()(mod) + return mod, params diff --git a/python/tvm/relay/op/contrib/android_nnapi/_partitioner/_base/transform/__init__.py b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/_base/transform/__init__.py new file mode 100644 index 000000000000..f95ddb39c31a --- /dev/null +++ b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/_base/transform/__init__.py @@ -0,0 +1,22 @@ +# 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. +"""Namespace for transformation for Android NNAPI that is suitable +to do during the partitioning step.""" +from .annotate_nnapi_function_attributes import AnnotateNnapiFunctionAttributes +from .transform_conv2d_weight_layout import TransformConv2dWeightLayout +from .transform_relay_op_for_nnapi import TransformRelayOpForNnapi +from .prune_inference_agnostic_operators import PruneInferenceAgnosticOperators diff --git a/python/tvm/relay/op/contrib/android_nnapi/_partitioner/_base/transform/annotate_nnapi_function_attributes.py b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/_base/transform/annotate_nnapi_function_attributes.py new file mode 100644 index 000000000000..fb01bbc71a9c --- /dev/null +++ b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/_base/transform/annotate_nnapi_function_attributes.py @@ -0,0 +1,80 @@ +# 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. +"""Annotate Android NNAPI functions (in Relay IR) for additional +attributes required for lowering.""" +import tvm +import tvm.relay + + +class AnnotateNnapiFunctionAttributes: + """Tag Android NNAPI compiler-specific attributes to exported Relay IR Functions. + + Parameters + ---------- + external_compiler: str + The name of the BYOC external compiler. + + android_nnapi_level: int + The targeted Android API level. + """ + + def __init__(self, external_compiler, android_nnapi_level): + super().__init__() + self._external_compiler = external_compiler + self._android_nnapi_level = android_nnapi_level + + def __call__(self, mod): + """Tag Android NNAPI compiler-specific attributes to exported Relay IR Functions. + + Parameters + ---------- + mod: tvm.IRModule + The module containing exported functions to be tagged. + + Returns + ------- + mod: tvm.IRModule + The tagged module. + """ + assert isinstance(mod, tvm.IRModule) + ret = tvm.IRModule() + gvs = mod.get_global_vars() + for gvar in gvs: + func = mod[gvar] + func = self._Annotator(self._external_compiler, self._android_nnapi_level).annotate( + func + ) + ret[gvar] = func + return ret + + class _Annotator(tvm.relay.ExprMutator): + def __init__(self, external_compiler, android_nnapi_level): + super().__init__() + self._external_compiler = external_compiler + self._android_nnapi_level = android_nnapi_level + + def annotate(self, func): + assert isinstance(func, tvm.relay.Function) + return self.visit(func) + + def visit_function(self, fn): + new_func = super().visit_function(fn) + if getattr(new_func.attrs, "Compiler", None) == self._external_compiler: + new_func = new_func.with_attr( + "NnapiTargetVersion", tvm.tir.IntImm("int32", self._android_nnapi_level) + ) + return new_func diff --git a/python/tvm/relay/op/contrib/android_nnapi/_partitioner/_base/transform/prune_inference_agnostic_operators.py b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/_base/transform/prune_inference_agnostic_operators.py new file mode 100644 index 000000000000..2725bdb4fba0 --- /dev/null +++ b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/_base/transform/prune_inference_agnostic_operators.py @@ -0,0 +1,56 @@ +# 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. +"""Remove operators that does not change inference results.""" +import tvm + +NN_DROPOUT_OP = tvm.relay.op.get("nn.dropout") + + +class PruneInferenceAgnosticOperators: + """Remove operators that does not change inference results.""" + + class _OperatorPruner(tvm.relay.ExprMutator): + def visit_tuple_getitem(self, op): + if ( + isinstance(op.tuple_value, tvm.relay.Call) + and op.tuple_value.op == NN_DROPOUT_OP + and op.index == 0 + ): + return super().visit(op.tuple_value.args[0]) + return super().visit_tuple_getitem(op) + + def __call__(self, mod): + """Remove operators that does not change inference results. + + Parameters + ---------- + mod: tvm.IRModule + The module to be pruned. + + Returns + ------- + mod: tvm.IRModule + The pruned module. + """ + assert isinstance(mod, tvm.IRModule) + ret = tvm.IRModule() + gvs = mod.get_global_vars() + for gvar in gvs: + func = mod[gvar] + func = PruneInferenceAgnosticOperators._OperatorPruner().visit(func) + ret[gvar] = func + return ret diff --git a/python/tvm/relay/op/contrib/android_nnapi/_partitioner/_base/transform/transform_conv2d_weight_layout.py b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/_base/transform/transform_conv2d_weight_layout.py new file mode 100644 index 000000000000..52e9b6108e14 --- /dev/null +++ b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/_base/transform/transform_conv2d_weight_layout.py @@ -0,0 +1,186 @@ +# 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. +"""Transform the layout of nn.conv2d weights to preferred layout for exported subgraphs.""" +import numpy as np +import tvm + +NN_CONV2D_OP = tvm.relay.op.get("nn.conv2d") + + +class TransformConv2dWeightLayout(tvm.relay.ExprMutator): + """Transform the layout of nn.conv2d weights to preferred layout for exported subgraphs. + + Parameters + ---------------------- + external_compiler: str + The name of BYOC external compiler. + + target_layout: str + The target layout for nn.conv2d weights. + """ + + def __init__(self, external_compiler, target_layout): + super().__init__() + self._external_compiler = external_compiler + self._target_layout = target_layout + + def __call__(self, mod, params): + """Transform the layout of nn.conv2d weights to preferred layout for exported subgraphs. + + Parameters + ---------- + mod: tvm.IRModule + The transform target module. + + params: dict of str to tvm.runtime.NDArray + The corresponding parameter inputs to mod. + + Returns + ------- + mod: tvm.IRModule + The transformed mod. + + params: dict of str to tvm.runtime.NDArray + The transformed params. + """ + assert isinstance(mod, tvm.IRModule) + assert isinstance(params, dict) + + self._mod = mod + self._params = params + self._call_stack = [] + self._transformed_vars = [] + self._in_export_func = False + self._mod["main"] = self.visit(mod["main"]) + self._mod = tvm.relay.transform.InferType()(self._mod) + + return self._mod, self._params + + def visit_call(self, call): + if isinstance(call.op, tvm.ir.Op): + if ( + call.op == NN_CONV2D_OP + and call.attrs["kernel_layout"] != self._target_layout + and self._in_export_func + ): + transpose_idx = [call.attrs["kernel_layout"].index(d) for d in self._target_layout] + assert len(self._call_stack) % 2 == 1 + + self._call_stack.append(call) + + # Check if kernel is straight var + weight_arg_idx = 1 + for i in range(len(self._call_stack) - 1, -1, -1): + func = self._call_stack[i] + if i % 2 == 0: + assert isinstance(func, tvm.relay.Function) + weight_arg_idx = list(func.params).index(weight_arg) + else: + assert isinstance(func, tvm.relay.Call) + weight_arg = func.args[weight_arg_idx] + if not isinstance(weight_arg, tvm.relay.Var): + self._call_stack.pop() + return super().visit_call(call) + call = self._call_stack.pop() + + # Mutate + new_attrs = dict(call.attrs) + new_attrs["kernel_layout"] = self._target_layout + call = tvm.relay.nn.conv2d(call.args[0], call.args[1], **new_attrs) + self._call_stack.append(call) + weight_arg_idx = 1 + for i in range(len(self._call_stack) - 1, -1, -1): + func = self._call_stack[i] + if i % 2 == 0: # Function + weight_param = weight_arg + weight_param_idx = list(func.params).index(weight_param) + + new_params = [ + p for p_i, p in enumerate(func.params) if p_i != weight_param_idx + ] + new_weight_param = self.memo_map[weight_param] + new_params.insert(weight_param_idx, new_weight_param) + func = tvm.relay.Function( + params=list(new_params), + body=func.body, + ret_type=func.ret_type, + type_params=func.type_params, + attrs=func.attrs, + ) + + weight_arg_idx = weight_param_idx + else: # Call + weight_arg = func.args[weight_arg_idx] + if ( + weight_arg in self._transformed_vars + ): # visited, perhaps it's a shared weight? + call = self._call_stack.pop() + return super().visit_call(call) + new_args = [a for a_i, a in enumerate(func.args) if a_i != weight_arg_idx] + new_shape = tuple( + [weight_arg.checked_type.shape[idx] for idx in transpose_idx] + ) + new_weight_arg = tvm.relay.var( + name_hint=weight_arg.name_hint, + shape=new_shape, + dtype=weight_arg.checked_type.dtype, + ) + self.memo_map[weight_arg] = new_weight_arg + self._transformed_vars.append(weight_arg) + new_args.insert(weight_arg_idx, new_weight_arg) + func = tvm.relay.Call( + op=func.op, + args=new_args, + attrs=func.attrs, + type_args=func.type_args, + ) + self._call_stack[i] = func + call = self._call_stack.pop() + + main_weight_param = str(self._mod["main"].params[weight_arg_idx].name_hint) + if main_weight_param in self._params: + weight_ndarray = self._params[main_weight_param] + self._params[main_weight_param] = tvm.runtime.ndarray.array( + np.transpose(weight_ndarray.asnumpy(), transpose_idx) + ) + elif isinstance(call.op, (tvm.relay.Function, tvm.relay.GlobalVar)): + self._call_stack.append(call) + self.visit(call.op) + call = self._call_stack.pop() + return super().visit_call(call) + + def visit_function(self, fn): + is_export_fn = getattr(fn.attrs, "Compiler", "") == self._external_compiler + self._call_stack.append(fn) + if is_export_fn: + assert not self._in_export_func + self._in_export_func = True + self.visit(fn.body) + if is_export_fn: + assert self._in_export_func + self._in_export_func = False + fn = self._call_stack.pop() + return super().visit_function(fn) + + def visit_global_var(self, gvar): + if isinstance(self._mod[gvar], tvm.relay.Function): + self._mod[gvar] = self.visit_function(self._mod[gvar]) + return super().visit_global_var(gvar) + + def visit_var(self, var): + assert var not in self._transformed_vars + return super().visit_var(var) diff --git a/python/tvm/relay/op/contrib/android_nnapi/_partitioner/_base/transform/transform_relay_op_for_nnapi/__init__.py b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/_base/transform/transform_relay_op_for_nnapi/__init__.py new file mode 100644 index 000000000000..ff2e12fbfa07 --- /dev/null +++ b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/_base/transform/transform_relay_op_for_nnapi/__init__.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. +"""Convert Relay operators into mathematically equivalent forms +so that Android NNAPI supports it.""" +import tvm +from .expand_batch_norm import ExpandBatchNorm as _ExpandBatchNorm +from .expand_split import ExpandSplit as _ExpandSplit + + +class TransformRelayOpForNnapi: + """Convert Relay operators into mathematically equivalent forms so + that Android NNAPI supports it.""" + + def __call__(self, mod): + """Convert Relay operators into mathematically equivalent forms + so that Android NNAPI supports it. + + Parameters + ---------- + mod: tvm.IRModule + The module to be transformed. + + Returns + ------- + mod: tvm.IRModule + The transformed module. + """ + assert isinstance(mod, tvm.IRModule) + for pazz in [_ExpandBatchNorm(), _ExpandSplit()]: + mod = pazz(mod) + return mod diff --git a/python/tvm/relay/op/contrib/android_nnapi/_partitioner/_base/transform/transform_relay_op_for_nnapi/expand_batch_norm.py b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/_base/transform/transform_relay_op_for_nnapi/expand_batch_norm.py new file mode 100644 index 000000000000..49009346e9e0 --- /dev/null +++ b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/_base/transform/transform_relay_op_for_nnapi/expand_batch_norm.py @@ -0,0 +1,91 @@ +# 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. +"""Expand Relay IR batch_norm for mapping to Android NNAPI.""" +import tvm + + +NN_BATCH_NORM_OP = tvm.relay.op.get("nn.batch_norm") + + +class ExpandBatchNorm(tvm.relay.ExprMutator): + """Expand Relay IR batch_norm for mapping to Android NNAPI.""" + + def __call__(self, mod): + assert isinstance(mod, tvm.IRModule) + ret = tvm.IRModule() + gvs = mod.get_global_vars() + for gvar in gvs: + func = mod[gvar] + func = self.visit(func) + ret[gvar] = func + ret = tvm.relay.transform.InferType()(ret) + return ret + + def _expand_batch_norm(self, tgi): + batch_norm = {} + + def _scope(): + new_args = [self.visit(a) for a in tgi.tuple_value.args] + batch_norm["args"] = dict( + zip(["data", "gamma", "beta", "moving_mean", "moving_var"], new_args) + ) + + _scope() + batch_norm["attrs"] = tgi.tuple_value.attrs + + assert all( + [ + len(batch_norm["args"][f].checked_type.shape) == 1 + for f in ["gamma", "beta", "moving_mean", "moving_var"] + ] + ) + + # reshape args + data_type = tgi.tuple_value.args[0].checked_type + data_rank = len(data_type.shape) + for arg_name in ["gamma", "beta", "moving_mean", "moving_var"]: + target_shape = ( + [1 for i in range(0, batch_norm["attrs"]["axis"])] + + list([int(i) for i in batch_norm["args"][arg_name].checked_type.shape]) + + [1 for i in range(batch_norm["attrs"]["axis"] + 1, data_rank)] + ) + batch_norm["args"][arg_name] = tvm.relay.reshape( + batch_norm["args"][arg_name], target_shape + ) + + # start expanding + step = [] + step.append( + batch_norm["args"]["moving_var"] + + tvm.relay.const(batch_norm["attrs"]["epsilon"], data_type.dtype) + ) + step.append(tvm.relay.sqrt(step[0])) + step.append(batch_norm["args"]["data"] - batch_norm["args"]["moving_mean"]) + step.append(step[2] / step[1]) + step.append(step[3] * batch_norm["args"]["gamma"]) + step.append(step[4] + batch_norm["args"]["beta"]) + + return step[-1] + + def visit_tuple_getitem(self, op): + if ( + isinstance(op.tuple_value, tvm.relay.Call) + and op.tuple_value.op == NN_BATCH_NORM_OP + and op.index == 0 + ): + return self._expand_batch_norm(op) + return super().visit_tuple_getitem(op) diff --git a/python/tvm/relay/op/contrib/android_nnapi/_partitioner/_base/transform/transform_relay_op_for_nnapi/expand_split.py b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/_base/transform/transform_relay_op_for_nnapi/expand_split.py new file mode 100644 index 000000000000..b08b4a6bda95 --- /dev/null +++ b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/_base/transform/transform_relay_op_for_nnapi/expand_split.py @@ -0,0 +1,75 @@ +# 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. +"""Expand Relay IR split for mapping to Android NNAPI.""" +import tvm + + +SPLIT_OP = tvm.relay.op.get("split") + + +class ExpandSplit(tvm.relay.ExprMutator): + """Expand Relay IR split for mapping to Android NNAPI.""" + + def __call__(self, mod): + assert isinstance(mod, tvm.IRModule) + ret = tvm.IRModule() + gvs = mod.get_global_vars() + for gvar in gvs: + func = mod[gvar] + func = self.visit(func) + ret[gvar] = func + ret = tvm.relay.transform.InferType()(ret) + return ret + + def visit_tuple_getitem(self, op): + if isinstance(op.tuple_value, tvm.relay.Call) and op.tuple_value.op == SPLIT_OP: + split = op.tuple_value + data = split.args[0] + new_strided_slice = {} + new_strided_slice["data"] = self.visit(data) + if isinstance(split.attrs["indices_or_sections"], (int, tvm.tir.IntImm)): + section_size = int(data.checked_type.shape[split.attrs["axis"]]) // int( + split.attrs["indices_or_sections"] + ) + indices = [section_size * i for i in range(int(split.attrs["indices_or_sections"]))] + else: + indices = [0] + indices += list(map(int, split.attrs["indices_or_sections"])) + + split_attrs_axis = ( + int(split.attrs["axis"]) + if split.attrs["axis"] >= 0 + else int(len(data.checked_type.shape) + split.attrs["axis"]) + ) + new_strided_slice["begin"] = [ + (0 if i != split_attrs_axis else indices[op.index]) + for i in range(len(data.checked_type.shape)) + ] + new_strided_slice["end"] = [ + ( + int(data.checked_type.shape[i]) + if i != split_attrs_axis + else ( + indices[op.index + 1] + if op.index < len(indices) - 1 + else int(data.checked_type.shape[split.attrs["axis"]]) + ) + ) + for i in range(len(data.checked_type.shape)) + ] + return tvm.relay.strided_slice(**new_strided_slice) + return super().visit_tuple_getitem(op) diff --git a/python/tvm/relay/op/contrib/android_nnapi/_partitioner/byoc.py b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/byoc.py new file mode 100644 index 000000000000..3fdbd9b9751d --- /dev/null +++ b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/byoc.py @@ -0,0 +1,216 @@ +# 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. +"""Implement the official BYOC partitioning flow for Android NNAPI.""" +import tvm +import tvm.relay +import tvm.relay.op.contrib.register +from . import _base + +# avoid re-registering byoc annotation rules +_BYOC_ANNOTATION_RULES_REGISTERED = False + + +def _register_byoc_annotation_rules(external_compiler, android_nnapi_level): + global _BYOC_ANNOTATION_RULES_REGISTERED + # avoid re-registering byoc annotation rules + if _BYOC_ANNOTATION_RULES_REGISTERED: + return + _BYOC_ANNOTATION_RULES_REGISTERED = True + + from tvm.contrib.target.android_nnapi import ( # pylint: disable=import-outside-toplevel + Compiler as RelayFunctionToAndroidNNAPICompiler, + ) + from tvm.contrib.target.android_nnapi.error import ( # pylint: disable=line-too-long,import-outside-toplevel + AndroidNNAPICompilerIncompatibleError, + ) + import tvm.contrib.target.android_nnapi.operation_utils.relay_op as relay_op_handler_root # pylint: disable=line-too-long,import-outside-toplevel + + def _isolate_op_call_node(call, compiler): + func_params = [] + new_call_args = [] + for i, arg in enumerate(call.args): + if isinstance(arg.checked_type, tvm.relay.TupleType): + tuple_param_fields = [ + tvm.relay.var(f"arg{ i }.{ j }", type_annotation=f) + for j, f in enumerate(arg.checked_type.fields) + ] + func_params += tuple_param_fields + tuple_arg = tvm.relay.Tuple( + [tvm.relay.annotation.compiler_begin(f, compiler) for f in tuple_param_fields] + ) + new_call_args.append(tuple_arg) + elif isinstance(arg.checked_type, tvm.relay.TensorType): + func_params.append(tvm.relay.var(f"arg{ i }", type_annotation=arg.checked_type)) + new_call_args.append(tvm.relay.annotation.compiler_begin(func_params[-1], compiler)) + else: + raise NotImplementedError(arg.checked_type) + new_call = tvm.relay.annotation.compiler_end( + tvm.relay.Call(call.op, new_call_args, call.attrs, call.type_args), compiler + ) + return tvm.relay.Function(func_params, new_call) + + def _check_call_support(call): + assert isinstance(call, tvm.relay.Call) + mod = tvm.IRModule.from_expr(_isolate_op_call_node(call, external_compiler)) + mod = tvm.relay.transform.PartitionGraph()(mod) + mod, _ = _base.post_partition_transform( + mod, {}, android_nnapi_level=android_nnapi_level, external_compiler=external_compiler + ) + external_func = (lambda op: op if isinstance(op, tvm.relay.Function) else mod[op])( + mod["main"].body.op + ) # op may be a GlobalVar, hence the if + options = { + "target": {"api_level": android_nnapi_level}, + } + assert isinstance(external_func, tvm.relay.Function) + try: + RelayFunctionToAndroidNNAPICompiler(options).codegen(external_func) + except AndroidNNAPICompilerIncompatibleError: + return False + return True + + tvm.ir.register_op_attr( + "annotation.compiler_begin", "target.android_nnapi", lambda expr: False + ) # create "target.android_nnapi" in OpAttrMap + + def _recursive_register(cur_namespace, handle): + if ( + hasattr(handle, "handler") + and tvm.relay.op.get(cur_namespace).get_attr("target.android_nnapi") is None + ): # avoid overriding user-registered rules + tvm.ir.register_op_attr(cur_namespace, "target.android_nnapi", _check_call_support) + else: # non-leaf + for attr_name in dir(handle): + if not attr_name.startswith("_"): + _recursive_register( + f"{ cur_namespace }.{ attr_name }" if cur_namespace != "" else attr_name, + getattr(handle, attr_name), + ) + + _recursive_register("", relay_op_handler_root) + + +def _prune_android_nnapi_subgraphs(mod, external_compiler): + """Prune a IRModule for subgraphs that are not suitable to be offloaded + to Android NNAPI. + + Parameters + ---------- + mod: tvm.IRModule + The TVM Module to be pruned. + + external_compiler: str + The name of the Android NNAPI external compiler. + + Returns + ------- + mod: tvm.IRModule + The pruned TVM Module. + """ + + def _func_should_be_pruned(func): + if getattr(func.attrs, "Compiler", None) != external_compiler: + return False + + def _scope(): + visitor = tvm.relay.ExprVisitor() + visitor.visit(func) + return {expr for expr in visitor.memo_map if isinstance(expr, tvm.ir.Op)} + + all_ops = _scope() + if any([wanted_op in all_ops for wanted_op in [tvm.relay.op.get("nn.conv2d")]]): + return False + return True + + subgraphs_to_prune = { + gv.name_hint for gv in mod.get_global_vars() if _func_should_be_pruned(mod[gv]) + } + if len(subgraphs_to_prune) == 0: + return mod + + def _remove_subgraphs(mod, subgraphs_to_prune): + class InlineSubgraphs(tvm.relay.ExprMutator): + """Inline subgraphs back to the invocation place.""" + + def __init__(self, subgraphs_to_prune): + super().__init__() + self._subgraphs_to_prune = subgraphs_to_prune + + def __call__(self, mod): + self._mod = mod + new_mod = tvm.IRModule() + gvs = mod.get_global_vars() + for gvar in gvs: + new_mod[gvar] = self.visit(mod[gvar]) + return new_mod + + def visit_call(self, call): + if ( + isinstance(call.op, tvm.relay.GlobalVar) + and call.op.name_hint in self._subgraphs_to_prune + ): + gfunc = self._mod[call.op] + bind_map = {} + assert len(gfunc.params) == len(call.args) + for i in range(len(call.args)): + bind_map[gfunc.params[i]] = self.visit(call.args[i]) + return tvm.relay.bind(gfunc.body, bind_map) + return super().visit_call(call) + + mod = InlineSubgraphs(subgraphs_to_prune)(mod) + return tvm.IRModule( + {gv: mod[gv] for gv in mod.get_global_vars() if gv.name_hint not in subgraphs_to_prune} + ) + + return _remove_subgraphs(mod, subgraphs_to_prune) + + +def byoc_partition(mod, params, android_nnapi_level): + """Partition a IRModule using rules registered with TVM BYOC. + + Parameters + ---------- + mod: tvm.IRModule + The TVM Module to be partitioned. + + params: dict of str to tvm.runtime.NDArray + The parameters to mod. + + android_nnapi_level: int + The targeted Android API level. + + Returns + ------- + mod: tvm.IRModule + The partitioned module. + + params: dict of str to tvm.runtime.NDArray + The transformed parameters to mod. + """ + assert isinstance(mod, tvm.IRModule) + + external_compiler = "android_nnapi" + _register_byoc_annotation_rules(external_compiler, android_nnapi_level) + pattern_table = tvm.relay.op.contrib.register.get_pattern_table(external_compiler) + if pattern_table is not None: + mod = tvm.relay.transform.MergeComposite(pattern_table)(mod) + mod = tvm.relay.transform.AnnotateTarget([external_compiler])(mod) + mod = tvm.relay.transform.MergeCompilerRegions()(mod) + mod = tvm.relay.transform.PartitionGraph()(mod) + mod = _prune_android_nnapi_subgraphs(mod, external_compiler) + mod, params = _base.post_partition_transform(mod, params, android_nnapi_level) + return mod, params diff --git a/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/__init__.py b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/__init__.py new file mode 100644 index 000000000000..22c660a06b52 --- /dev/null +++ b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/__init__.py @@ -0,0 +1,50 @@ +# 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. +"""Partition Relay IR graph for Android NNAPI based on RPC profiling.""" +from .partitioner import Partitioner as _Partitioner + + +def rpc_partition(mod, params, tracker, options={}): # pylint: disable=dangerous-default-value + """Partition Relay IR graph into NNAPI compilable graph. + + Parameters + ---------- + mod: tvm.IRModule + The graph to be partitioned. + + trackers: tvm.rpc.TrackerSession + The tracker client managing RPC device sessions. + + options["target"]["api_level"]: int + The targeting API level of Android. Defaults to 29. + + options["target"]["llvm_triple"]: str + The LLVM triple describing the target. Defaults to "aarch64-linux-android29". + + options["tvm"]["rpc"]["remote_key"]: str + The key under which the profiling device is registered in the tracker. + Defaults to "android". + + options["tvm"]["rpc"]["profile_run"]: int + The remote profile cycle count for an operation. Defaults to 10. + + Returns + ------- + mod: tvm.IRModule + The partitioned graph. + """ + return _Partitioner(tracker, options).partition(mod, params) diff --git a/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/__init__.py b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/__init__.py new file mode 100644 index 000000000000..35274f8ada77 --- /dev/null +++ b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/__init__.py @@ -0,0 +1,18 @@ +# 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. +"""Partition a TVM Module.""" +from .partition_module import PartitionModule diff --git a/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/annotate_for_relay_compiler.py b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/annotate_for_relay_compiler.py new file mode 100644 index 000000000000..edf4b33894fd --- /dev/null +++ b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/annotate_for_relay_compiler.py @@ -0,0 +1,76 @@ +# 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. +"""Insert annotation.compiler_begin/compiler_end according to the +coloring of the Relay IR nodes.""" +import tvm +from tvm.relay.op.annotation import compiler_begin, compiler_end + + +class AnnotateForRelayCompiler(tvm.relay.ExprMutator): + """Annotate the graph with `annotation.compiler_begin` and `annotation.compiler_end`. + + Parameters + ---------- + options: dict + The partitioner option dict. + + edm: ExportDecisionMaker + A object telling whether a Relay node should be exported. + """ + + def __init__(self, options, edm): + super().__init__() + self._options = options + self._compiler = self._options["tvm"]["external_compiler"] + self._edm = edm + self._in_graph = False + + def annotate(self, func): + """Annotate the graph with `annotation.compiler_begin` and `annotation.compiler_end`. + + Parameters + ---------- + func: tvm.relay.Function + The function to be annotated. + + Returns + ------- + func: tvm.relay.Function + The annotated function. + """ + assert isinstance(func, tvm.relay.Function) + return self.visit(func) + + def visit(self, expr): + export_result = self._edm.node_is_exported(expr, self._compiler) + if export_result == self._edm.EXPORT_RESULT["YES"]: + if not self._in_graph: + self._in_graph = True + new_expr = super().visit(expr) + assert self._in_graph + self._in_graph = False # subgraph should exit here when returning from children + return compiler_end(new_expr, self._compiler) + elif export_result == self._edm.EXPORT_RESULT["NO"]: + if self._in_graph: + self._in_graph = False + new_expr = super().visit(expr) + assert not self._in_graph + self._in_graph = True # restore `self._in_graph` in case other siblings needs it + return compiler_begin(new_expr, self._compiler) + + ret = super().visit(expr) + return ret diff --git a/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/collect_branching_nodes.py b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/collect_branching_nodes.py new file mode 100644 index 000000000000..3f3938b911fc --- /dev/null +++ b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/collect_branching_nodes.py @@ -0,0 +1,72 @@ +# 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. +"""Collect nodes that have more than a single child (branching) from a Relay graph.""" +import tvm + + +class CollectBranchingNodes: + """Collect nodes that have more than a single child (branching) from a Relay graph.""" + + class _BranchingNodeCollector(tvm.relay.ExprVisitor): + def __init__(self): + super().__init__() + self._branching_nodes = set() + + def collect(self, expr): + self.visit(expr) + return self._branching_nodes + + def visit(self, expr): + if (not isinstance(expr, tvm.ir.Op)) and (expr in self.memo_map): + self._branching_nodes.add(expr) + return super().visit(expr) + + def visit_function(self, fn): + self.visit(fn.body) + + class _RelayTopologicalSorter(tvm.relay.ExprVisitor): + def __init__(self, expr_root): + super().__init__() + self._expr_root = expr_root + + def sort(self, branching_nodes_set): + self._branching_nodes_set = branching_nodes_set + self._ret = [] + self.visit(self._expr_root) + return self._ret + + def visit(self, expr): + super().visit(expr) + if expr in self._branching_nodes_set: + self._ret.append(expr) + self._branching_nodes_set.remove(expr) + + def collect(self, expr): + """Collect nodes that have more than a single child (branching) from a Relay graph. + + Parameters + ---------- + expr: tvm.relay.Expr + The expression whose branching children are to be collected. + + Returns + ------- + branching_nodes: list of tvm.relay.Expr + The expressions where branching happens. + """ + branching_nodes_set = self._BranchingNodeCollector().collect(expr) + return self._RelayTopologicalSorter(expr).sort(branching_nodes_set) diff --git a/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/export_decision_marker.py b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/export_decision_marker.py new file mode 100644 index 000000000000..9a3291117a60 --- /dev/null +++ b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/export_decision_marker.py @@ -0,0 +1,146 @@ +# 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. +"""Color Relay IR nodes to indicate the designated device of +execution.""" +import tvm +from .platform_simulator import compute_device + + +class ExportDecisionMarker(tvm.relay.ExprVisitor): + """A blackbox object telling whether a Relay node should be exported to the queried compiler. + + Parameters + ---------- + options: dict + The partitioner option dict. + + node_transfers: (Internal Format) + The artifact of the partitioning algorithm. + """ + + EXPORT_RESULT = { + "NO": 0, + "YES": 1, + "UNSURE": 2, + } + + _DEVICE_COMPILERS = { + compute_device.NnapiDevice.DEV_NAME: "android_nnapi", + compute_device.TvmDevice.DEV_NAME: "tvm", + } + + def __init__(self, options, node_transfers): + super().__init__() + self._options = options + self._node_transfers = node_transfers + assert ( + self._options["tvm"]["external_compiler"] + == self._DEVICE_COMPILERS[compute_device.NnapiDevice.DEV_NAME] + ) + + def mark(self, func): + assert isinstance(func, tvm.relay.Function) + self._node_compiler_map = {func: compute_device.TvmDevice.DEV_NAME} + self.memo_map[func] = None + self._saved_devs = [] + self._parent_dev = compute_device.TvmDevice.DEV_NAME + self.visit(func.body) + + def _set_parent(self, dev): + self._saved_devs.append(self._parent_dev) + self._parent_dev = dev + + def _restore_parent(self): + self._parent_dev = self._saved_devs.pop() + + def node_is_exported(self, node, compiler): + """Report whether a node is marked as exported. + + Parameters + ---------- + node: tvm.relay.Node + The queried node. + + compiler: str + The compiler used to export. + + Returns + ------- + exported: self.EXPORT_RESULT + Whether the node is marked as exported with the compiler. + """ + if isinstance(node, tvm.ir.Op): + return self.EXPORT_RESULT["UNSURE"] + + verdict = self._node_compiler_map[node] + if len(verdict) == 1 and verdict[0] == compiler: + return self.EXPORT_RESULT["YES"] + if compiler in verdict: + return self.EXPORT_RESULT["UNSURE"] + return self.EXPORT_RESULT["NO"] + + def visit(self, expr): + if isinstance(expr, tvm.ir.Op): + return super().visit(expr) + + next_dev = self._node_transfers[self._parent_dev][expr] + next_compiler = self._DEVICE_COMPILERS[next_dev] + self._node_compiler_map[expr] = [next_compiler] + + self._set_parent(next_dev) + ret = super().visit(expr) + self._restore_parent() + return ret + + def visit_var(self, var): + assert self._node_compiler_map[var] == ["tvm"] + super().visit_var(var) + + def visit_let(self, let): + raise NotImplementedError(let.type_key) + + def visit_function(self, fn): + assert self._node_compiler_map[fn] == ["tvm"] + super().visit_function(f) + + def visit_if(self, i): + assert self._node_compiler_map[i] == ["tvm"] + super().visit_if(i) + + def visit_global_var(self, gv): + assert self._node_compiler_map[gv] == ["tvm"] + super().visit_global_var(gv) + + def visit_ref_create(self, r): + raise NotImplementedError(r.type_key) + + def visit_ref_read(self, r): + raise NotImplementedError(r.type_key) + + def visit_ref_write(self, r): + raise NotImplementedError(r.type_key) + + def visit_tuple_getitem(self, t): + if isinstance(t.tuple_value, tvm.relay.Call): + assert self._node_compiler_map[t] == ["tvm"] + super().visit_tuple_getitem(t) + + def visit_constructor(self, c): + raise NotImplementedError(c.type_key) + + def visit_match(self, m): + raise NotImplementedError(m.type_key) diff --git a/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/partition_module.py b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/partition_module.py new file mode 100644 index 000000000000..e16b25982525 --- /dev/null +++ b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/partition_module.py @@ -0,0 +1,72 @@ +# 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. +"""Partition graphs in a given Relay module into those for tvm/android_nnapi compilers.""" +import tvm +from .collect_branching_nodes import CollectBranchingNodes +from .platform_simulator import PlatformSimulator +from .export_decision_marker import ExportDecisionMarker +from .annotate_for_relay_compiler import AnnotateForRelayCompiler + + +class PartitionModule: + """Partition graphs in a given Relay module into those for tvm/android_nnapi compilers. + + Parameters + ---------- + tracker: tvm.rpc.TrackerSession + The tracker client managing RPC device sessions. + + options: dict + The partitioner option dict. + """ + + def __init__(self, tracker, options): + self._tracker = tracker + self._options = options + + def __call__(self, mod): + """Partition graphs in a given Relay module into those for tvm/android_nnapi compilers. + + Parameters + ---------- + mod: tvm.IRModule + The partition target module. + + Returns + ------- + mod: tvm.IRModule + The partitioned module. + """ + assert isinstance(mod, tvm.IRModule) + gvs = mod.get_global_vars() + for gvar in gvs: + func = mod[gvar] + branching_nodes = CollectBranchingNodes().collect(func) + psim = PlatformSimulator(self._tracker, self._options, branching_nodes) + psim.calculate_cost(func) + edm = ExportDecisionMarker(self._options, psim.node_transfers) + edm.mark(func) + assert all( + [ + edm.node_is_exported(n, "tvm") == edm.EXPORT_RESULT["YES"] + for n in branching_nodes + ] + ) + func = AnnotateForRelayCompiler(self._options, edm).annotate(func) + mod[gvar] = func + mod = tvm.relay.transform.PartitionGraph()(mod) + return mod diff --git a/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/platform_simulator/__init__.py b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/platform_simulator/__init__.py new file mode 100644 index 000000000000..6b969ba525c0 --- /dev/null +++ b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/platform_simulator/__init__.py @@ -0,0 +1,18 @@ +# 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. +"""Platform simulator for cost calculation.""" +from .platform_simulator import PlatformSimulator diff --git a/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/platform_simulator/_utils.py b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/platform_simulator/_utils.py new file mode 100644 index 000000000000..4f35a21201b1 --- /dev/null +++ b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/platform_simulator/_utils.py @@ -0,0 +1,53 @@ +# 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. +"""Utilities for PlatformSimulator.""" +import functools +import re +import tvm + + +def get_type_size(tipe): + """Get node size in bytes. + + Parameters + ---------- + tipe: tvm.relay.Type + The Relay type whose size is to be calculated. + """ + if isinstance(tipe, tvm.ir.type.TupleType): + return sum([_get_type_size(f) for f in tipe.fields]) + + dtype = str(tipe.dtype) + shape = list([int(i) for i in tipe.shape]) + nbits = (lambda s: int(s) if s != "" else 8)(re.sub("[a-z]", "", dtype)) + assert nbits % 8 == 0 + return functools.reduce( + lambda x, y: x * y, + shape, + nbits / 8, # use byte as basic unit + ) + + +def get_node_size(node): + """Get node size in bytes. + + Parameters + ---------- + node: tvm.relay.Expr + The Relay expression whose size is to be calculated. + """ + return get_type_size(node.checked_type) diff --git a/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/platform_simulator/compute_device/__init__.py b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/platform_simulator/compute_device/__init__.py new file mode 100644 index 000000000000..ba976dbb8075 --- /dev/null +++ b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/platform_simulator/compute_device/__init__.py @@ -0,0 +1,19 @@ +# 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. +"""Available ComputeDevices.""" +from .nnapi_device import NnapiDevice +from .tvm_device import TvmDevice diff --git a/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/platform_simulator/compute_device/_compute_device.py b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/platform_simulator/compute_device/_compute_device.py new file mode 100644 index 000000000000..7ce1d826b14c --- /dev/null +++ b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/platform_simulator/compute_device/_compute_device.py @@ -0,0 +1,68 @@ +# 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. +"""Base class for computation device.""" +import abc + + +class ComputeDevice(abc.ABC): + """Base class for computation device.""" + + @abc.abstractmethod + def estimate_call_op_cost(self, call): + """Estimate the runtime cost of executing a given call. + + Parameters + ---------- + call: tvm.relay.Call + The Relay call expression whose runtime cost is to be estimated. + """ + + @abc.abstractmethod + def estimate_memory_read_cost(self, dtype, size): + """Estimate cost of memory read from this device. + + Parameters + ---------- + dtype: str + The dtype to be read. + + size: int + The amount of memory to be read. + + Returns + ------- + cost: int + The cost of performing the read. + """ + + @abc.abstractmethod + def estimate_memory_write_cost(self, dtype, size): + """Estimate cost of memory write from this device. + + Parameters + ---------- + dtype: str + The dtype to be write. + + size: int + The amount of memory to be write. + + Returns + ------- + cost: int + The cost of performing the write. + """ diff --git a/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/platform_simulator/compute_device/_error.py b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/platform_simulator/compute_device/_error.py new file mode 100644 index 000000000000..7a296e06874b --- /dev/null +++ b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/platform_simulator/compute_device/_error.py @@ -0,0 +1,35 @@ +# 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. +"""Error encountered during RPC profiling.""" +from tvm.contrib.target.android_nnapi.error import ( + AndroidNNAPICompilerError, +) + + +class AndroidNNAPICompilerProfilingError(AndroidNNAPICompilerError): + """Error caused by profiling failure + + Parameters + ---------- + msg: str + An optional error message + + Notes + ----- + This error is used internally in the partitioner and does not intend to be + handled by other modules. + """ diff --git a/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/platform_simulator/compute_device/_rpc_device.py b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/platform_simulator/compute_device/_rpc_device.py new file mode 100644 index 000000000000..c7a6bd9034c2 --- /dev/null +++ b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/platform_simulator/compute_device/_rpc_device.py @@ -0,0 +1,43 @@ +# 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. +"""Base class for RPC-based ComputeDevice +""" +from ._compute_device import ComputeDevice + + +class RPCDevice(ComputeDevice): # pylint: disable=abstract-method + """Base class for RPC-based ComputeDevice. + + Parameters + ---------- + options: dict + The partitioner options dict. + + tracker: tvm.rpc.TrackerSession + The tracker managing RPC devices used for profiling. + """ + + def __init__(self, options, tracker): + super().__init__() + self._options = options + self._tracker = tracker + + self._remote_key = options["tvm"]["rpc"]["remote_key"] + self._remote_profile_run = options["tvm"]["rpc"]["profile_run"] + + self._target_triple = options["target"]["llvm_triple"] + self._tvm_target = f"llvm -mtriple={ self._target_triple }" diff --git a/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/platform_simulator/compute_device/_utils.py b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/platform_simulator/compute_device/_utils.py new file mode 100644 index 000000000000..856b7c8d46e5 --- /dev/null +++ b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/platform_simulator/compute_device/_utils.py @@ -0,0 +1,52 @@ +# 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. +"""Utilities for ComputeDevices.""" +import tvm +import tvm.relay +from .._utils import get_node_size # pylint: disable=unused-import + + +def get_function_output_buffer(func, device): + """Get a NDArray for buffering the function output. + + Parameters + ---------- + func: tvm.relay.Function + The function for which the buffer is generated. + + device: tvm.runtime.Device + The device on which the generated buffer is allocated. + + Returns + ------- + buf: tvm.runtime.NDArray + The generated NDArray buffer. + """ + assert isinstance(func, tvm.relay.Function) + + def _get_ndarray(ttype): + assert isinstance(ttype, tvm.relay.TensorType) + return tvm.nd.empty( + shape=tuple([int(i) for i in ttype.shape]), dtype=ttype.dtype, device=device + ) + + ret_type = func.ret_type + if isinstance(ret_type, tvm.relay.TensorType): + return _get_ndarray(ret_type) + if isinstance(ret_type, tvm.relay.TupleType): + return tvm.runtime.container.tuple_object([_get_ndarray(t) for t in ret_type.fields]) + raise NotImplementedError(ret_type) diff --git a/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/platform_simulator/compute_device/nnapi_device.py b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/platform_simulator/compute_device/nnapi_device.py new file mode 100644 index 000000000000..78fc250aa064 --- /dev/null +++ b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/platform_simulator/compute_device/nnapi_device.py @@ -0,0 +1,225 @@ +# 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. +"""NNAPI ComputeDevice specialization.""" +import numpy as np +import tvm +from tvm.contrib.target.android_nnapi import Compiler as RelayFunctionToAndroidNNAPICompiler +from tvm.contrib.target.android_nnapi.error import ( + AndroidNNAPICompilerIncompatibleError, +) +from ....._base import post_partition_transform +from ._rpc_device import RPCDevice +from ._error import AndroidNNAPICompilerProfilingError +from . import _utils + + +def _isolate_op_call_node(call, compiler): + func_params = [] + new_call_args = [] + for i, arg in enumerate(call.args): + if isinstance(arg.checked_type, tvm.relay.TupleType): + tuple_param_fields = [ + tvm.relay.var(f"arg{ i }.{ j }", type_annotation=f) + for j, f in enumerate(arg.checked_type.fields) + ] + func_params += tuple_param_fields + tuple_arg = tvm.relay.Tuple( + [tvm.relay.annotation.compiler_begin(f, compiler) for f in tuple_param_fields] + ) + new_call_args.append(tuple_arg) + elif isinstance(arg.checked_type, tvm.relay.TensorType): + func_params.append(tvm.relay.var(f"arg{ i }", type_annotation=arg.checked_type)) + new_call_args.append(tvm.relay.annotation.compiler_begin(func_params[-1], compiler)) + else: + raise NotImplementedError(arg.checked_type) + new_call = tvm.relay.annotation.compiler_end( + tvm.relay.Call(call.op, new_call_args, call.attrs, call.type_args), compiler + ) + return tvm.relay.Function(func_params, new_call) + + +class NnapiDevice(RPCDevice): + """NNAPI ComputeDevice specialization.""" + + DEV_NAME = "nnapi" + + def __init__(self, options, tracker): + super().__init__(options, tracker) + self._api_level = options["target"]["api_level"] + self._compiler_name = options["tvm"]["external_compiler"] + self._cached_memory_op_coefficient = {} + + def estimate_call_op_cost(self, call): + assert isinstance(call.op, tvm.ir.Op) + + # prepare the module to run + mod = tvm.IRModule({"main": _isolate_op_call_node(call, self._compiler_name)}) + mod = tvm.relay.transform.PartitionGraph()(mod) + + # get runtime on device (or failure) + try: + return self._get_runtime_on_device(mod) + except AndroidNNAPICompilerProfilingError: + return None + + def _get_runtime_on_device(self, mod): + assert isinstance(mod, tvm.IRModule) + + mod = tvm.relay.transform.InferType()(mod) + if isinstance(mod["main"].ret_type, tvm.relay.TensorType): + # prepare params + params = { + p.name_hint: tvm.nd.array( + np.random.uniform(size=tuple([int(i) for i in p.checked_type.shape])).astype( + str(p.checked_type.dtype) + ), + tvm.cpu(0), + ) + for p in mod["main"].params + } + + # run some post partition transformation and fixes + # here we try to mimic the result of an partition + mod, params = post_partition_transform( + mod, + params, + android_nnapi_level=self._options["target"]["api_level"], + external_compiler=self._options["tvm"]["external_compiler"], + ) + + external_func = (lambda op: op if isinstance(op, tvm.relay.Function) else mod[op])( + mod["main"].body.op + ) # op may be a GlobalVar, hence the if + assert isinstance(external_func, tvm.relay.Function) + + # try compiling first to see if there's any problem + # if there's any incompatible case, an error would be thrown + options = { + "target": { + "api_level": self._options["target"]["api_level"], + }, + } + try: + RelayFunctionToAndroidNNAPICompiler(options).codegen(external_func) + except AndroidNNAPICompilerIncompatibleError as err: + raise AndroidNNAPICompilerProfilingError( + f"Relay operator unsupported by Android NNAPI compiler: { str(err) }" + ) + + # build binary + mod = tvm.relay.transform.InferType()(mod) + with tvm.transform.PassContext(opt_level=3): + exe = tvm.relay.vm.compile(mod, target=self._tvm_target) + _, lib = exe.save() + assert lib + + temp_dir = tvm.contrib.utils.tempdir() + temp_lib_path = temp_dir.relpath("lib.so") + + def _scope(): + kwargs = {} + kwargs["options"] = [ + "--target={}".format(self._target_triple), + "-O3", + "-lneuralnetworks", + "-shared", + "-fPIC", + ] + lib.export_library(temp_lib_path, fcompile=tvm.contrib.ndk.create_shared, **kwargs) + + _scope() + + # push binary + remote = self._tracker.request(self._remote_key) + remote.upload(temp_lib_path) + remote_mod = remote.load_module("lib.so") + + # run + device = remote.cpu() + args = [params[p.name_hint] for p in mod["main"].params] + args.append( + _utils.get_function_output_buffer(external_func, device) + ) # arg contains an additional output buffer at the end + remote_func = remote_mod.time_evaluator( + str(external_func.attrs.global_symbol), device, number=self._remote_profile_run + ) + ret = remote_func(*args).mean + elif isinstance(mod["main"].ret_type, tvm.relay.TupleType): + # Tuple(ADT) is not supported by RPC (and NNAPI!) + raise AndroidNNAPICompilerProfilingError(f"Relay tuple-typed operator is unsupported") + else: + raise NotImplementedError(str(mod["main"].ret_type)) + + return ret + + def estimate_memory_read_cost(self, dtype, size): + scale, init = self._memory_op_coefficient(str(dtype)) + return max(scale * size + init, 0) + + def estimate_memory_write_cost(self, dtype, size): + scale, init = self._memory_op_coefficient(str(dtype)) + return max(scale * size + init, 0) + + def _memory_op_coefficient(self, benchmark_dtype): + if benchmark_dtype in self._cached_memory_op_coefficient: + return self._cached_memory_op_coefficient[benchmark_dtype] + + def _scope(): + img = tvm.relay.var("img", shape=[32, 512, 512, 1], dtype=benchmark_dtype) + ann_img = tvm.relay.annotation.compiler_begin(img, self._compiler_name) + weight_0 = tvm.relay.var("weight_0", shape=[1, 1, 1, 1], dtype=benchmark_dtype) + ann_weight_0 = tvm.relay.annotation.compiler_begin(weight_0, self._compiler_name) + conv_0 = tvm.relay.nn.conv2d( + ann_img, ann_weight_0, data_layout="NHWC", kernel_layout="OHWI" + ) + ann_conv_0 = tvm.relay.annotation.compiler_end(conv_0, self._compiler_name) + single_conv_f = tvm.relay.Function([img, weight_0], ann_conv_0) + mod = tvm.IRModule({"main": single_conv_f}) + mod = tvm.relay.transform.PartitionGraph()(mod) + + mod = tvm.relay.transform.InferType()(mod) + size = _utils.get_node_size(mod["main"].body) + time = self._get_runtime_on_device(mod) / 2 + return size, time + + size1, time1 = _scope() + + def _scope(): + img = tvm.relay.var("img", shape=[32, 256, 256, 1], dtype=benchmark_dtype) + ann_img = tvm.relay.annotation.compiler_begin(img, self._compiler_name) + weight_0 = tvm.relay.var("weight_0", shape=[1, 1, 1, 1], dtype=benchmark_dtype) + ann_weight_0 = tvm.relay.annotation.compiler_begin(weight_0, self._compiler_name) + conv_0 = tvm.relay.nn.conv2d( + ann_img, ann_weight_0, data_layout="NHWC", kernel_layout="OHWI" + ) + ann_conv_0 = tvm.relay.annotation.compiler_end(conv_0, self._compiler_name) + single_conv_f = tvm.relay.Function([img, weight_0], ann_conv_0) + mod = tvm.IRModule({"main": single_conv_f}) + mod = tvm.relay.transform.PartitionGraph()(mod) + + mod = tvm.relay.transform.InferType()(mod) + size = _utils.get_node_size(mod["main"].body) + time = self._get_runtime_on_device(mod) / 2 + return size, time + + size2, time2 = _scope() + + # solve time = scale * size + init for scale, init + scale = (time1 - time2) / (size1 - size2) + init = time1 - scale * size1 + self._cached_memory_op_coefficient[benchmark_dtype] = (scale, init) + return scale, init diff --git a/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/platform_simulator/compute_device/tvm_device.py b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/platform_simulator/compute_device/tvm_device.py new file mode 100644 index 000000000000..1bae0df8a9ce --- /dev/null +++ b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/platform_simulator/compute_device/tvm_device.py @@ -0,0 +1,122 @@ +# 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. +"""TVM ComputeDevice specialization.""" +import numpy as np +import tvm +from ._rpc_device import RPCDevice +from . import _utils + + +def _isolate_op_call_node(call): + func_params = [] + new_call_args = [] + for i, arg in enumerate(call.args): + if isinstance(arg.checked_type, tvm.relay.TupleType): + tuple_param_fields = [ + tvm.relay.var(f"arg{ i }.{ j }", type_annotation=f) + for j, f in enumerate(arg.checked_type.fields) + ] + func_params += tuple_param_fields + tuple_arg = tvm.relay.Tuple(tuple_param_fields) + new_call_args.append(tuple_arg) + elif isinstance(arg.checked_type, tvm.relay.TensorType): + func_params.append(tvm.relay.var(f"arg{ i }", type_annotation=arg.checked_type)) + new_call_args.append(func_params[-1]) + else: + raise NotImplementedError(arg.checked_type) + new_call = tvm.relay.Call(call.op, new_call_args, call.attrs, call.type_args) + return tvm.relay.Function(func_params, new_call) + + +class TvmDevice(RPCDevice): + """TVM ComputeDevice specialization.""" + + DEV_NAME = "tvm" + + def estimate_call_op_cost(self, call): + assert isinstance(call.op, tvm.ir.Op) + + mod = tvm.IRModule({"main": _isolate_op_call_node(call)}) + mod = tvm.relay.transform.InferType()(mod) + + return self._get_runtime_on_device(mod) + + def _get_runtime_on_device(self, mod): + assert isinstance(mod, tvm.IRModule) + + mod = tvm.relay.transform.InferType()(mod) + if isinstance(mod["main"].ret_type, tvm.relay.TensorType): + with tvm.transform.PassContext(opt_level=3, disabled_pass=["AlterOpLayout"]): + exe = tvm.relay.vm.compile(mod, target=self._tvm_target) + _, lib = exe.save() + + if not lib: + return 0 + + temp_dir = tvm.contrib.utils.tempdir() + temp_lib_path = temp_dir.relpath("lib.so") + + def _scope(): + kwargs = {} + kwargs["options"] = [ + "--target={}".format(self._target_triple), + "-O3", + "-shared", + "-fPIC", + ] + lib.export_library(temp_lib_path, fcompile=tvm.contrib.ndk.create_shared, **kwargs) + + _scope() + + remote = self._tracker.request(self._remote_key) + remote.upload(temp_lib_path) + remote_mod = remote.load_module("lib.so") + + device = remote.cpu() + args = [ + tvm.nd.array( + np.random.uniform(size=tuple([int(i) for i in p.checked_type.shape])).astype( + str(p.checked_type.dtype) + ), + device, + ) + for p in mod["main"].params + ] + args.append(_utils.get_function_output_buffer(mod["main"], device)) # output buffer + + def _scope(): + primitives = exe.primitive_ops + assert len(primitives) == 1 + return primitives[0] + + main_sym = _scope() + remote_func = remote_mod.time_evaluator( + main_sym, device, number=self._remote_profile_run + ) + ret = remote_func(*args).mean + elif isinstance(mod["main"].ret_type, tvm.relay.TupleType): + # Tuple(ADT) is not supported by RPC + ret = 0 + else: + raise NotImplementedError(mod["main"].ret_type) + return ret + + def estimate_memory_read_cost(self, dtype, size): + return 0 + + def estimate_memory_write_cost(self, dtype, size): + return 0 diff --git a/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/platform_simulator/platform_simulator.py b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/platform_simulator/platform_simulator.py new file mode 100644 index 000000000000..f654b8343ebf --- /dev/null +++ b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partition_module/platform_simulator/platform_simulator.py @@ -0,0 +1,245 @@ +# 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. +"""Simulate computation platform and compute runtime costs for a given Relay IR Function.""" +import tvm +from . import compute_device +from . import _utils + + +class PlatformSimulator(tvm.relay.ExprVisitor): + """Simulate computation platform and compute runtime costs for a given Relay IR Function. + + Parameters + ---------- + tracker: tvm.rpc.TrackerSession + The tracker client managing RPC device sessions. + + options: dict + The partitioner option dict. + """ + + ENABLED_DEVICES = [compute_device.TvmDevice.DEV_NAME, compute_device.NnapiDevice.DEV_NAME] + + def __init__(self, tracker, options, branching_nodes): + super().__init__() + self._tracker = tracker + self._options = options + + # DP artifacts + self._node_costs = {dev: {} for dev in self.ENABLED_DEVICES} + self._node_transfers = {dev: {} for dev in self.ENABLED_DEVICES} + + # node assignment exceptions + self._pinned_nodes = {n: compute_device.TvmDevice.DEV_NAME for n in branching_nodes} + + # init platform components + def _scope(): + self._compute_devices = { + compute_device.TvmDevice.DEV_NAME: compute_device.TvmDevice(options, self._tracker), + compute_device.NnapiDevice.DEV_NAME: compute_device.NnapiDevice( + options, self._tracker + ), + } + + _scope() + assert all([dev in self._compute_devices for dev in self.ENABLED_DEVICES]) + + @property + def node_costs(self): + return self._node_costs + + @property + def node_transfers(self): + return self._node_transfers + + def calculate_cost(self, func): + """Compute runtime costs for a given Relay IR Function. + + Parameters + ---------- + func: tvm.relay.Function + The function whose cost is to be evaluated. + """ + self.visit(func) + + def visit_tuple(self, tup): + super().visit_tuple(tup) + + for tdev in self.ENABLED_DEVICES: + if self._skip_node_on_dev(tup, tdev): + continue + t_cost = 0 + for f in tup.fields: + if f in self._node_costs[tdev]: + t_cost += self._node_costs[tdev][f] + else: + t_cost = None + break + if t_cost is None: + continue + self._node_costs[tdev][tup] = t_cost + self._node_transfers[tdev][tup] = tdev + + def visit_call(self, call): + super().visit_call(call) + + for tdev in self.ENABLED_DEVICES: + c_cost = None + for cdev in self.ENABLED_DEVICES: # compute device + if self._skip_node_on_dev(call, cdev): + continue + cost = 0 + for a in call.args: + if a in self._node_costs[cdev]: + cost += self._node_costs[cdev][a] + else: + cost = None + break + if cost is None: + continue + + if isinstance(call.op, tvm.ir.Op): + op_cost = self._compute_devices[cdev].estimate_call_op_cost(call) + if op_cost is None: + continue + cost += op_cost + elif isinstance(call.op, (tvm.relay.Function, tvm.relay.GlobalVar)): + if call.op not in self._node_costs[cdev]: + continue + cost += self._node_costs[cdev][call.op] + else: + raise NotImplementedError(call.op.type_key) + cost += self.get_transfer_cost(call, cdev, tdev) + if c_cost is None or c_cost > cost: + c_cost = cost + if isinstance(call.op, (tvm.relay.Function, tvm.relay.GlobalVar)): + assert cdev == compute_device.TvmDevice.DEV_NAME + self._node_transfers[tdev][call] = cdev + if isinstance(call.op, tvm.ir.Op): + self._node_transfers[tdev][call.op] = cdev + assert c_cost is not None + self._node_costs[tdev][call] = c_cost + + def visit_var(self, var): + super().visit_var(var) + if isinstance(var.checked_type, tvm.relay.TupleType): + self._node_costs[compute_device.TvmDevice.DEV_NAME][var] = 0 + self._node_transfers[compute_device.TvmDevice.DEV_NAME][ + var + ] = compute_device.TvmDevice.DEV_NAME + else: + for tdev in self.ENABLED_DEVICES: + if self._skip_node_on_dev(var, tdev): + continue + self._node_costs[tdev][var] = self.get_transfer_cost( + var, compute_device.TvmDevice.DEV_NAME, tdev + ) + self._node_transfers[tdev][var] = compute_device.TvmDevice.DEV_NAME + + def visit_let(self, let): + raise NotImplementedError(let.type_key) + + def visit_function(self, fn): + super().visit_function(fn) + assert fn not in self._pinned_nodes + f_cost = None + for sdev in self.ENABLED_DEVICES: + if fn.body in self._node_costs[sdev]: + cost = self._node_costs[sdev][fn.body] + self.get_transfer_cost( + fn.body, sdev, compute_device.TvmDevice.DEV_NAME + ) + if f_cost is None or f_cost > cost: + f_cost = cost + fb_dev = sdev + assert f_cost is not None + self._node_costs[compute_device.TvmDevice.DEV_NAME][fn] = f_cost + self._node_transfers[compute_device.TvmDevice.DEV_NAME][fn] = fb_dev + + def visit_if(self, i): + raise NotImplementedError(i.type_key) + + def visit_global_var(self, gv): + super().visit_global_var(gv) + assert gv not in self._pinned_nodes + self._node_costs[compute_device.TvmDevice.DEV_NAME][gv] = 0 + self._node_transfers[compute_device.TvmDevice.DEV_NAME][ + gv + ] = compute_device.TvmDevice.DEV_NAME + + def visit_constructor(self, c): + raise NotImplementedError(c.type_key) + + def visit_constant(self, const): + for tdev in self.ENABLED_DEVICES: + if self._skip_node_on_dev(const, tdev): + continue + self._node_costs[tdev][const] = 0 + # force constants to be cut out by compiler annotations + # the PartitionGraph pass is smart enough to export constants if it should be exported + self._node_transfers[tdev][const] = compute_device.TvmDevice.DEV_NAME + + def visit_ref_create(self, r): + raise NotImplementedError(r.type_key) + + def visit_ref_read(self, r): + raise NotImplementedError(r.type_key) + + def visit_ref_write(self, r): + raise NotImplementedError(r.type_key) + + def visit_tuple_getitem(self, t): + super().visit_tuple_getitem(t) + if isinstance(t.tuple_value, tvm.relay.Tuple): + for tdev in self.ENABLED_DEVICES: + if self._skip_node_on_dev(t, tdev): + continue + self._node_costs[tdev][t] = self._node_costs[tdev][t.tuple_value] + self._node_transfers[tdev][t] = tdev + elif isinstance(t.tuple_value, (tvm.relay.Call, tvm.relay.Var)): + for tdev in self.ENABLED_DEVICES: + self._node_costs[tdev][t] = self._node_costs[compute_device.TvmDevice.DEV_NAME][ + t.tuple_value + ] + self.get_transfer_cost(t, compute_device.TvmDevice.DEV_NAME, tdev) + self._node_transfers[tdev][t] = compute_device.TvmDevice.DEV_NAME + else: + raise NotImplementedError(t.tuple_value.type_key) + + def visit_match(self, m): + raise NotImplementedError(m.type_key) + + def get_transfer_cost(self, node, sdev, tdev): + if sdev == tdev: + return 0 + return self.get_transfer_cost_typed(node.checked_type, sdev, tdev) + + def get_transfer_cost_typed(self, tipe, sdev, tdev): + if sdev == tdev: + return 0 + if isinstance(tipe, tvm.relay.TensorType): + size = _utils.get_type_size(tipe) + return self._compute_devices[sdev].estimate_memory_read_cost( + tipe.dtype, size + ) + self._compute_devices[tdev].estimate_memory_write_cost(tipe.dtype, size) + assert isinstance(tipe, tvm.relay.TupleType) + return sum([self.get_transfer_cost_typed(f, sdev, tdev) for f in tipe.fields]) + + def _skip_node_on_dev(self, node, dev): + if node in self._pinned_nodes: + if self._pinned_nodes[node] == dev: + return False + return True + return False diff --git a/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partitioner.py b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partitioner.py new file mode 100644 index 000000000000..4098523c3854 --- /dev/null +++ b/python/tvm/relay/op/contrib/android_nnapi/_partitioner/rpc/partitioner.py @@ -0,0 +1,104 @@ +# 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. +"""Partition a Relay IR graph into subgraphs compiled by +TVM/Android NNAPI compilers using RPC profiling.""" +import copy +import tvm.relay.transform +from .. import _base +from .partition_module import PartitionModule + + +class Partitioner: + """Partition a Relay IR graph into subgraphs compiled by + TVM/Android NNAPI compilers using RPC profiling. + + Parameters + ---------- + tracker: tvm.rpc.TrackerSession + The tracker client managing RPC device sessions. + + options: dict + The partitioner option dict. + """ + + DEFAULT_OPTIONS = { + "target": { + "api_level": 29, + "llvm_triple": "aarch64-linux-android29", + }, + "tvm": { + "external_compiler": "android_nnapi", + "rpc": { + "profile_run": 10, + "remote_key": "android", + }, + }, + } + + def __init__(self, tracker, options): + self._tracker = tracker + self._options = self._expand_options(options) + + def partition(self, mod, params): + """Partition a Relay IR graph + + Parameters + ---------- + mod: tvm.IRModule + The graph to be partitioned + + params: dict of str to tvm.runtime.NDArray + The input parameters to the graph + + Returns + ------- + mod: tvm.IRModule + The partitioned graph + + params: dict of str to tvm.runtime.NDArray + The transformed input parameters to the graph + + """ + assert isinstance(mod, tvm.IRModule) + mod = _base.pre_partition_transform(mod) + mod = PartitionModule(self._tracker, self._options)(mod) + mod, params = _base.post_partition_transform( + mod, + params, + android_nnapi_level=self._options["target"]["api_level"], + external_compiler=self._options["tvm"]["external_compiler"], + ) + return mod, params + + @staticmethod + def _expand_options(options): + ret = options + + def _recursive_merge(cur_opts, def_opts): + for k, v in def_opts.items(): + if k in cur_opts: + if isinstance(v, dict): + assert isinstance(cur_opts[k], dict) + _recursive_merge(cur_opts[k], v) + else: + assert isinstance(cur_opts[k], (float, int, str)) + else: + cur_opts[k] = copy.deepcopy(v) + + _recursive_merge(ret, Partitioner.DEFAULT_OPTIONS) + + return ret diff --git a/tests/python/contrib/test_android_nnapi/__init__.py b/tests/python/contrib/test_android_nnapi/__init__.py new file mode 100644 index 000000000000..13a83393a912 --- /dev/null +++ b/tests/python/contrib/test_android_nnapi/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/tests/python/contrib/test_android_nnapi/infrastructure.py b/tests/python/contrib/test_android_nnapi/infrastructure.py new file mode 100644 index 000000000000..ab2e1df56e04 --- /dev/null +++ b/tests/python/contrib/test_android_nnapi/infrastructure.py @@ -0,0 +1,76 @@ +# 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 re +import tvm + + +def annotate_for_android_nnapi(mod, android_api_level): + """Annotate Relay IR Function with attrs required by the Android NNAPI compiler. + + Parameters + ---------- + mod: tvm.IRModule + The module to be annotated. + + android_api_level: int + The target Android API level. + + Returns + ------- + mod: tvm.IRModule + The annotated module. + """ + ret = tvm.IRModule() + gvs = mod.get_global_vars() + for gv in gvs: + func = mod[gv] + func = func.with_attr("NnapiTargetVersion", tvm.tir.IntImm("int32", android_api_level)) + ret[gv] = func + return ret + + +def is_compilable(mod, android_api_level): + """Check if a module is compilable. + + Parameters + ---------- + mod: runtime.Module + The module to be checked for compilability. + + android_api_level: int + The targeting Android API level for testing of compilability. + + Returns + ------- + result: bool + Whether the module is compilable. + """ + tempdir = tvm.contrib.utils.tempdir() + temp_lib_path = tempdir.relpath("lib.so") + kwargs = {} + kwargs["options"] = [ + "--target={}".format( + f"aarch64-linux-android{android_api_level}" + ), # use aarch64 for testing + "-O0", # disable opt for testing + "-lneuralnetworks", + "-shared", + "-fPIC", + ] + mod.export_library(temp_lib_path, fcompile=tvm.contrib.ndk.create_shared, **kwargs) + return True diff --git a/tests/python/contrib/test_android_nnapi/test_byoc_partition.py b/tests/python/contrib/test_android_nnapi/test_byoc_partition.py new file mode 100644 index 000000000000..d602de40f11f --- /dev/null +++ b/tests/python/contrib/test_android_nnapi/test_byoc_partition.py @@ -0,0 +1,46 @@ +# 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 unittest.mock +import tvm +import tvm.relay +import tvm.relay.op.contrib.android_nnapi +import tvm.contrib.target.android_nnapi.operation_utils.relay_op as relay_op_handler_root + + +def test_byoc_partition(): + data = tvm.relay.var("data", shape=(1, 3, 4, 4), dtype="float32") + kernel = tvm.relay.var("kernel", shape=(2, 3, 4, 4), dtype="float32") + bias = tvm.relay.var("bias", shape=(2,), dtype="float32") + mod = tvm.IRModule.from_expr(tvm.relay.nn.bias_add(tvm.relay.nn.conv2d(data, kernel), bias)) + mock_root_handler = lambda: None + mock_root_handler.nn = lambda: None + mock_root_handler.nn.conv2d = lambda: None + mock_root_handler.nn.conv2d.handler = relay_op_handler_root.nn.conv2d.handler + with unittest.mock.patch( + "tvm.contrib.target.android_nnapi.operation_utils.relay_op", + new=mock_root_handler, + ): + mod, _ = tvm.relay.op.contrib.android_nnapi.byoc_partition_for_android_nnapi(mod, {}, 29) + assert len(mod.get_global_vars()) == 2 + ext_func_gv = next(filter(lambda gv: gv.name_hint != "main", mod.get_global_vars())) + ext_func = mod[ext_func_gv] + assert ext_func.body.op == tvm.relay.op.get("nn.conv2d") + + +if __name__ == "__main__": + test_byoc_partition() diff --git a/tests/python/contrib/test_android_nnapi/test_nn_conv2d.py b/tests/python/contrib/test_android_nnapi/test_nn_conv2d.py new file mode 100644 index 000000000000..a72aa496a0ab --- /dev/null +++ b/tests/python/contrib/test_android_nnapi/test_nn_conv2d.py @@ -0,0 +1,70 @@ +# 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 tvm +import tvm.relay +import tvm.contrib.target.android_nnapi +from . import infrastructure + + +def test_codegen_nchw_conv2d(): + data_t = tvm.relay.TensorType((1, 1, 4, 4), "float32") + data_v = tvm.relay.var("data", data_t) + data_a = tvm.relay.annotation.compiler_begin(data_v, "android_nnapi") + weight_t = tvm.relay.TensorType((1, 1, 2, 2), "float32") + weight_v = tvm.relay.var("weight", weight_t) + weight_a = tvm.relay.annotation.compiler_begin(weight_v, "android_nnapi") + conv_c = tvm.relay.nn.conv2d(data=data_a, weight=weight_a) + conv_a = tvm.relay.annotation.compiler_end(conv_c, "android_nnapi") + func = tvm.relay.Function([data_v, weight_v], conv_a) + mod = tvm.IRModule.from_expr(func) + mod = tvm.relay.transform.PartitionGraph()(mod) + mod = infrastructure.annotate_for_android_nnapi(mod, 28) + + exe = tvm.relay.backend.vm.compile( + mod, target="llvm -mtriple=aarch64-linux-android28", params={} + ) + _, lib = exe.save() + c_mod = lib.imported_modules[1] + assert infrastructure.is_compilable(c_mod, 28) + + +def test_codegen_nchw_conv2d_on_api29(): + data_t = tvm.relay.TensorType((1, 1, 4, 4), "float32") + data_v = tvm.relay.var("data", data_t) + data_a = tvm.relay.annotation.compiler_begin(data_v, "android_nnapi") + weight_t = tvm.relay.TensorType((1, 1, 2, 2), "float32") + weight_v = tvm.relay.var("weight", weight_t) + weight_a = tvm.relay.annotation.compiler_begin(weight_v, "android_nnapi") + conv_c = tvm.relay.nn.conv2d(data=data_a, weight=weight_a) + conv_a = tvm.relay.annotation.compiler_end(conv_c, "android_nnapi") + func = tvm.relay.Function([data_v, weight_v], conv_a) + mod = tvm.IRModule.from_expr(func) + mod = tvm.relay.transform.PartitionGraph()(mod) + mod = infrastructure.annotate_for_android_nnapi(mod, 29) + + exe = tvm.relay.backend.vm.compile( + mod, target="llvm -mtriple=aarch64-linux-android29", params={} + ) + _, lib = exe.save() + c_mod = lib.imported_modules[1] + assert infrastructure.is_compilable(c_mod, 29) + + +if __name__ == "__main__": + test_codegen_nchw_conv2d() + test_codegen_nchw_conv2d_on_api29() diff --git a/tests/python/contrib/test_android_nnapi/test_rpc_partition.py b/tests/python/contrib/test_android_nnapi/test_rpc_partition.py new file mode 100644 index 000000000000..d0e7ce3016fc --- /dev/null +++ b/tests/python/contrib/test_android_nnapi/test_rpc_partition.py @@ -0,0 +1,123 @@ +# 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 mmap +import os +import tvm +import tvm.relay +import tvm.relay.op.contrib.android_nnapi + + +class RPCTestingTracker: + def request(self, key): + return RPCTestingSession() + + +class RPCTestingSession: + def __init__(self): + self._remote_fs = {} + + def cpu(self, *args, **kwargs): + return tvm.cpu(*args, **kwargs) + + def load_module(self, remote_fpath): + return RPCTestingModule(self._remote_fs[remote_fpath]) + + def upload(self, local_fpath): + self._remote_fs[os.path.basename(local_fpath)] = local_fpath + + +class RPCTestingModule: + def __init__(self, module_fpath): + self._module_fpath = module_fpath + + def time_evaluator(self, fname, *args, **kwargs): + return RPCTestingFunction(self._module_fpath, fname) + + +class RPCTestingFunction: + def __init__(self, module_fpath, fname): + fname = fname.lower() + fd = os.open(module_fpath, os.O_RDONLY) + with mmap.mmap(fd, 0, access=mmap.ACCESS_READ) as mcontent: + assert mcontent.find(fname.encode()) != -1 + if mcontent.find(b"ANEURALNETWORKS") != -1: # mod is built with android nnapi + # this cost structure should put nn.conv2d on android nnapi and add on tvm + if mcontent.find(b"CONV_2D") != -1: + self.mean = 100 + else: + self.mean = 10 + else: + if mcontent.find(b"nn_conv2d") != -1: + self.mean = 1000 + else: + self.mean = 1 + os.close(fd) + + def __call__(self, *args, **kwargs): + return self + + +def test_rpc_partition(): + def _scope(): + data_t = tvm.relay.TensorType((1, 1, 4, 4), "float32") + data = tvm.relay.var("data", data_t) + weight_t = tvm.relay.TensorType((1, 1, 2, 2), "float32") + weight = tvm.relay.var("weight", weight_t) + conv = tvm.relay.nn.conv2d(data=data, weight=weight) + bias_t = tvm.relay.TensorType((1,), "float32") + bias = tvm.relay.var("bias", bias_t) + func_body = conv + bias + func = tvm.relay.Function([data, weight, bias], func_body) + mod = tvm.IRModule({"main": func}) + mod, _ = tvm.relay.op.contrib.android_nnapi.rpc_partition_for_android_nnapi( + mod=mod, params={}, tracker=RPCTestingTracker(), options={} + ) + return mod + + res = _scope() + + def _scope(): + data_t = tvm.relay.TensorType((1, 1, 4, 4), "float32") + data = tvm.relay.var("data", data_t) + data_a = tvm.relay.annotation.compiler_begin(data, "android_nnapi") + weight_t = tvm.relay.TensorType((1, 2, 2, 1), "float32") + weight = tvm.relay.var("weight", weight_t) + weight_a = tvm.relay.annotation.compiler_begin(weight, "android_nnapi") + conv = tvm.relay.nn.conv2d(data=data_a, weight=weight_a, kernel_layout="OHWI") + conv_a = tvm.relay.annotation.compiler_end(conv, "android_nnapi") + bias_t = tvm.relay.TensorType((1,), "float32") + bias = tvm.relay.var("bias", bias_t) + func_body = conv_a + bias + func = tvm.relay.Function([data, weight, bias], func_body) + mod = tvm.IRModule({"main": func}) + mod = tvm.relay.transform.PartitionGraph()(mod) + gvs = mod.get_global_vars() + for gv in gvs: + fn = mod[gv] + if getattr(fn.attrs, "Compiler", None) == "android_nnapi": + fn = fn.with_attr("NnapiTargetVersion", 29) + mod[gv] = fn + return mod + + ans = _scope() + + tvm.ir.assert_structural_equal(ans, res, map_free_vars=True) + + +if __name__ == "__main__": + test_rpc_partition() diff --git a/tests/scripts/task_python_integration.sh b/tests/scripts/task_python_integration.sh index 00b63af48646..86617a7fc092 100755 --- a/tests/scripts/task_python_integration.sh +++ b/tests/scripts/task_python_integration.sh @@ -63,6 +63,8 @@ run_pytest ctypes ${TVM_INTEGRATION_TESTSUITE_NAME} tests/python/integration if python -c "import tvm; from tvm.relay.op.contrib.ethosn import ethosn_available; print(ethosn_available().name)" -eq "SW_ONLY"; then ETHOSN_VARIANT_CONFIG=ETHOSN78_1TOPS_4PLE_448KSRAM run_pytest ctypes ${TVM_INTEGRATION_TESTSUITE_NAME}-contrib-test_ethosn tests/python/contrib/test_ethosn fi +# Set Android NDK Clang for Android NNAPI testing +export TVM_NDK_CC="${ANDROID_NDK_HOME}"/toolchains/llvm/prebuilt/linux-x86_64/bin/clang++ run_pytest ctypes ${TVM_INTEGRATION_TESTSUITE_NAME}-contrib tests/python/contrib # forked is needed because the global registry gets contaminated