stoa commented on a change in pull request #7742: URL: https://github.com/apache/tvm/pull/7742#discussion_r713992580
########## File path: python/tvm/contrib/stm32/emitter.py ########## @@ -0,0 +1,1618 @@ +# 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=line-too-long + +"""Code emission for the STM32 targets.""" + +import json +import os +import shutil +import tarfile +import re + +from datetime import datetime + +import numpy as np + +import tvm +from tvm.contrib import utils + +AI_TOOLS_VERSION_MAJOR = 1 +AI_TOOLS_VERSION_MINOR = 8 +AI_TOOLS_VERSION_MICRO = 0 + +AI_API_VERSION_MAJOR = 1 +AI_API_VERSION_MINOR = 0 +AI_API_VERSION_MICRO = 0 + +AI_TOOLS_REVISION = "v1" + +DBAR = "=" * 60 + +# ========================================================== +# _fix_name +# ========================================================== + + +def _fix_name(node_name): + """ Replace ':' with '_' in names like 'InputImg:0' """ + return node_name.replace(":", "_") + + +# ========================================================== +# get_input_tensor_name +# ========================================================== + + +def get_input_tensor_name(node_name): + return _fix_name(node_name) + + +# ========================================================== +# get_output_tensor_name +# ========================================================== + + +def get_output_tensor_name(node_name, idx): + return _fix_name(node_name) + "_" + str(idx) + + +# ========================================================== +# _get_node_args_name +# ========================================================== + + +def _get_node_args_name(node_name): + return _fix_name(node_name) + "_args" + + +# ========================================================== +# _get_node_arg_types_name +# ========================================================== + + +def _get_node_arg_types_name(node_name): + return _fix_name(node_name) + "_arg_type_ids" + + +# ========================================================== +# _get_type_size +# ========================================================== + + +def _get_type_size(dltype): + + if dltype in ("uint64", "int64"): + return 8 + if dltype in ("uint32", "int32", "float32"): + return 4 + if dltype in ("uint16", "int16"): + return 2 + if dltype in ("uint8", "int8"): + return 1 + + raise ValueError(f"Data type {dltype} is not supported") + + +# ========================================================== +# _get_type_data +# ========================================================== + + +def _get_type_data(dltype): + + if dltype == "uint64": + return "kDLUInt, 64, 1" + if dltype == "int64": + return "kDLInt, 64, 1" + if dltype == "float32": + return "kDLFloat, 32, 1" + if dltype == "uint32": + return "kDLUInt, 32, 1" + if dltype == "int32": + return "kDLInt, 32, 1" + if dltype == "uint16": + return "kDLUInt, 16, 1" + if dltype == "int16": + return "kDLInt, 16, 1" + if dltype == "uint8": + return "kDLUInt, 8, 1" + if dltype == "int8": + return "kDLInt, 8, 1" + + raise ValueError(f"Data type {dltype} is not supported") + + +# ========================================================== +# _get_aligned_offset +# ========================================================== + + +def _get_aligned_offset(offset, dltype): + + align = _get_type_size(dltype) + + if offset % align != 0: + offset = offset + (align - offset % align) + + return offset + + +# ========================================================== +# _get_tensor_elts +# ========================================================== + + +def _get_tensor_elts(dims): + size = 1 + for dim in dims: + size = size * dim + return size + + +# ========================================================== +# _get_tensor_size +# ========================================================== + + +def _get_tensor_size(dims, dltype): + size = _get_tensor_elts(dims) + return size * _get_type_size(dltype) + + +# ================================================================== +# _preprocess_code +# ================================================================== + + +def _preprocess_code(src): + """ Hack the C code implementing the model. """ + + dst = "#include <stdio.h>\n" "#include <math.h>\n" '#include "stm32lib.h"\n\n' "" + for line in src.splitlines(): + # + # This is sort of hacking - when AoT is available, we will be + # able to clean this ... + # + dst = dst + line + "\n" + + return dst + + +# ========================================================== +# CodeEmitter +# ========================================================== + + +class CodeEmitter(object): + """Code emitter class/utility.""" + + def __init__(self, include_activations=True, include_inputs=True, include_outputs=True): + """Initialize the Emitter instance. + + Parameters + ---------- + include_activations: + The Emitter allocates the storage for the activations data + and places it in a specific data section. If Falsr, the + main application is responsible for allocating the activations + storage. Default: True. + + include_inputs/include_outputs: + The Emitter allocates the storage for the input/output data. + This storage is shared with the activations and placed in the + specific activations data section. If False, the main + application is responsible for allocating the input/output + data storage. Default: True. + + Returns + ------- + CodeEmitter object. + + """ + + # + # Constants: + # + self.data_alignment = 8 + + # + # Static model: activations placed into a nn_data_act section + # Dynamic model: activations need to be malloc'ed by the + # applications. + # + if include_activations is True: + self.activations_static = 1 + else: + self.activations_static = 0 + + # + # Inputs/outputs may be allocated within the activations or + # separately. + # TODO: Separate the inputs from activations inside TVM. + # + if include_inputs is True: + assert ( + self.activations_static == 1 + ), "###Error: Static inputs are not allowed without activations." + self.inputs_static = 1 + else: + self.inputs_static = 0 + + if include_outputs is True: + assert ( + self.activations_static == 1 + ), "###Error: Static outputs are not allowed without activations." + self.outputs_static = 1 + else: + self.outputs_static = 0 + + # + # Parsed graph + # + self.nodes_ = [] + self.arg_nodes_ = [] + self.outputs_ = [] + self.attrs_ = {} + self.node_row_ptr_ = [] + # + # Parameters + # + self.params_ = {} + # + # Filled by data_placement() + # + self.weights_ = {} + self.activations_ = {} + self.input_data_ = {} + self.output_data_ = {} + self.nodes_size_ = 0 + self.weights_size_ = 0 + self.activations_size_ = 0 + + self.quantization_ = {} + + # ================================================================== + # __extract_quantization_info + # ================================================================== + + def __extract_quantization_info(self, quantization): + """ Build dictionary with quantization infos.""" + + for dl_tensor_name in self.input_data_: + if dl_tensor_name in quantization: + self.quantization_[dl_tensor_name] = quantization[dl_tensor_name] + + # + # Matching outputs is more difficult because TVM does not preserve + # output tensor names. + # We only support models with a single output now. + # + assert len(self.output_data_) == 1, "Multiple outputs models are not yet supported." + + for dl_tensor_name in self.output_data_: + for name in quantization: + if name not in self.input_data_: + self.quantization_["output"] = quantization[name] + break + + # ========================================================== + # _get_node_arg_name + # ========================================================== + + def __get_node_arg_name(self, arg): + arg_nid = arg[0] + arg_idx = arg[1] + arg_node = self.nodes_[arg_nid] + arg_name = self.nodes_[arg_nid]["name"] + if arg_node["op"] == "null": + # parameter + dl_tensor_name = get_input_tensor_name(arg_name) + elif arg_node["name"] == "reshape_nop": + # Handle __nop + src = arg_node["inputs"][0] + dl_tensor_name = self.__get_node_arg_name(src) + else: + # activation + dl_tensor_name = get_output_tensor_name(arg_name, arg_idx) + return dl_tensor_name + + # ========================================================== + # __tensor_is_output + # ========================================================== + + def __tensor_is_output(self, nid, idx): + + for out in self.outputs_: + out_nid = out[0] + out_idx = out[1] + if out_nid == nid and out_idx == idx: + return True + return False + + # ========================================================== + # __get_tensor_from_node + # ========================================================== + + def __get_tensor_from_node(self, nid, idx): + # + # 'eid' is index into the dltype', 'shape', etc. + # + eid = self.node_row_ptr_[nid] + idx + + dltype = self.attrs_["dltype"][1][eid] + dims = self.attrs_["shape"][1][eid] + storage_id = self.attrs_["storage_id"][1][eid] + ndim = len(dims) + # + # Get tensor size + # + size = _get_tensor_size(dims, dltype) + + tensor = { + "dltype": dltype, + "ndim": ndim, + "dims": dims, + "strides": None, + "storage_id": storage_id, + # + # What is this byte_offset really ? + # + "byte_offset": 0, + "offset": 0, + "size": size, + } + + return tensor + + # ================================================================== + # __compute_data_placement + # ================================================================== + + def __compute_data_placement(self): + """ Compute inputs, outputs, weight, activation sizes""" + + self.inputs_ = self.arg_nodes_.copy() + + # + # weights: + # + + offset = 0 + + for key in self.params_: + # + # First, find the node in graph + # + nid = 0 + for node in self.nodes_: + if node["name"] == key: + break + nid += 1 + + dl_tensor_name = get_input_tensor_name(key) + tensor = self.__get_tensor_from_node(nid, 0) + # + # Compute the offset + # + dltype = tensor["dltype"] + aligned_offset = _get_aligned_offset(offset, dltype) + tensor["offset"] = aligned_offset + + for idx in self.arg_nodes_: + node = self.nodes_[idx] + node_name = node["name"] + if node_name == key: + self.inputs_.remove(idx) + + self.weights_[dl_tensor_name] = tensor + + # + # Next offset + # + offset = aligned_offset + tensor["size"] + + self.weights_size_ = offset + + # + # activations: + # + + buffer_list_ = {} + + nid = 0 + + for node in self.nodes_: + + if node["op"] == "null": + nid += 1 + continue + + if node["op"] != "tvm_op": + raise ValueError(f"Only TVM ops are supported") + + node_name = node["name"] + node_attrs = node["attrs"] + func_name = node_attrs["func_name"] + num_outputs = int(node_attrs["num_outputs"]) + + if func_name == "__nop": + assert node_name == "reshape_nop", f"Unsupported __nop operator {node_name}." + assert num_outputs == 1 + assert not self.__tensor_is_output(nid, 0) + nid += 1 + continue + + for idx in range(num_outputs): + # + # Do not count the 'outputs_' + # + if self.__tensor_is_output(nid, idx): + continue + + dl_tensor_name = get_output_tensor_name(node_name, idx) + tensor = self.__get_tensor_from_node(nid, idx) + # + # Remember this tensor with the storage id + # + storage_id = tensor["storage_id"] + if storage_id not in buffer_list_: + buffer_list_[storage_id] = [] + buffer_entry = buffer_list_[storage_id] + buffer_entry.append(tensor) + + self.activations_[dl_tensor_name] = tensor + + self.nodes_size_ = self.nodes_size_ + 1 + + nid += 1 + + # + # Compute 'input_data_' + # + offset = 0 + + for nid in self.inputs_: + node = self.nodes_[nid] + node_name = node["name"] + # + # Arthur: I suppose that input nodes only have a single + # output dependency + # + dl_tensor_name = get_input_tensor_name(node_name) + # + # This tensor is at some index inside 'input_data_' dictionary + # depending on the 'inputs_' list order. We refer to this position + # when generating the XXX.h file. + # + tensor = self.__get_tensor_from_node(nid, 0) + + if self.inputs_static == 1: + # + # Remember this tensor with the storage id + # + storage_id = tensor["storage_id"] + if storage_id not in buffer_list_: + buffer_list_[storage_id] = [] + buffer_entry = buffer_list_[storage_id] + buffer_entry.append(tensor) + else: + # + # Compute the offset + # + dltype = tensor["dltype"] + aligned_offset = _get_aligned_offset(offset, dltype) + tensor["offset"] = aligned_offset + + self.input_data_[dl_tensor_name] = tensor + + # + # Next offset + # + offset = aligned_offset + tensor["size"] + + # + # Compute 'output_data_' + # + offset = 0 + + for output in self.outputs_: + + nid = output[0] + idx = output[1] + + node = self.nodes_[nid] + node_name = node["name"] + + dl_tensor_name = get_output_tensor_name(node_name, idx) + + tensor = self.__get_tensor_from_node(nid, idx) + + if self.outputs_static == 1: + # + # Remember this tensor with the storage id + # + storage_id = tensor["storage_id"] + if storage_id not in buffer_list_: + buffer_list_[storage_id] = [] + buffer_entry = buffer_list_[storage_id] + buffer_entry.append(tensor) + else: + # + # Compute the offset + # + dltype = tensor["dltype"] + aligned_offset = _get_aligned_offset(offset, dltype) + tensor["offset"] = aligned_offset + + self.output_data_[dl_tensor_name] = tensor + + # + # Next offset + # + offset = aligned_offset + tensor["size"] + + # + # Go over all storage IDs and compute offsets and activations_size_ + # + offset = 0 + + for storage_id in buffer_list_: + buffer_entry = buffer_list_[storage_id] + + new_offset = offset + for tensor in buffer_entry: + assert tensor["storage_id"] == storage_id + dltype = tensor["dltype"] + aligned_offset = _get_aligned_offset(offset, dltype) + tensor["offset"] = aligned_offset + size = tensor["size"] + if (aligned_offset + size) > new_offset: + new_offset = aligned_offset + size + offset = new_offset + + self.activations_size_ = offset + + # ========================================================== + # _parse_model + # ========================================================== + + def _parse_model(self, quantization=None): + """Parse the module. Build internal data structures. + + Parameters + ---------- + module : TVM module or ModuleLibraryFormat object + The module to parse + + quantization: Dictionary + The quantization information for model inputs/outputs. + """ + + for key in self.graph_: + if key == "nodes": + self.nodes_ = self.graph_["nodes"] + elif key == "arg_nodes": + self.arg_nodes_ = self.graph_["arg_nodes"] + elif key == "node_row_ptr": + self.node_row_ptr_ = self.graph_["node_row_ptr"] + elif key == "heads": + self.outputs_ = self.graph_["heads"] + elif key == "attrs": + self.attrs_ = self.graph_["attrs"] + elif key == "metadata": + continue + else: + print("### Error: JSON key {} not supported".format(key)) + assert False + + # + # Build all tensor lists + # + self.__compute_data_placement() + + # + # Extract quantization info for inputs/outputs + # + if quantization is not None: + self.__extract_quantization_info(quantization) + + # ========================================================== + # parse_library_format + # ========================================================== + + def parse_library_format(self, model_library_format_path, quantization=None): + """Parse the module. Build internal data structures. + + Parameters + ---------- + model_library_format_path : + The module to parse + + quantization: Dictionary + The quantization information for model inputs/outputs. + """ + + temp_dir = utils.tempdir() + extract_path = temp_dir.relpath("extract") + os.mkdir(extract_path) + with tarfile.TarFile(model_library_format_path) as f: + f.extractall(extract_path) + + # + # Extract informations from the Model Library Format + # + graph_file = os.path.join(extract_path, "runtime-config", "graph", "graph.json") + with open(graph_file, "r") as f: + # returns JSON object as a dictionary + graph_dict = json.load(f) + + params_dict = {} + param_file = os.path.join(extract_path, "parameters", "default.params") + with open(param_file, "rb") as f: + params = tvm.runtime.load_param_dict(f.read()) + # + # Map -> Python Dict + # + tmp_dict = {} + for (k, v) in params.items(): + tmp_dict[k] = v + + # + # Sort params for debugging + # + for k in sorted(tmp_dict.keys()): + params_dict[k] = tmp_dict[k] + + src_dir = os.path.join(extract_path, "codegen", "host", "src") + # List of strings from Model Library Format C files + src_files = [] + for filename in os.listdir(src_dir): + with open(os.path.join(src_dir, filename), "r") as fin: + src = fin.read() + src_files.append(src) + + self.graph_ = graph_dict + self.params_ = params_dict + self.lib_ = src_files + + self._parse_model(quantization) + + # ========================================================== + # parse_module + # ========================================================== + + def parse_module(self, module, quantization=None): + """Parse the module. Build internal data structures. + + Parameters + ---------- + module : TVM module Review comment: I kept 2 versions. Docstring fixed. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
