converter.py

# Copyright 2018 Xiaomi, Inc.  All rights reserved.
#
# Licensed 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 argparse
import filelock
import hashlib
import os
import re
import sh
import subprocess
import sys
import urllib
import yaml

from enum import Enum

import sh_commands
from sh_commands import BuildType

from common import CaffeEnvType
from common import DeviceType
from common import mace_check
from common import MaceLogger
from common import StringFormatter

################################
# common definitions
################################
BUILD_OUTPUT_DIR = 'build'
PHONE_DATA_DIR = "/data/local/tmp/mace_run"
MODEL_OUTPUT_DIR_NAME = 'model'
BUILD_TMP_DIR_NAME = '_tmp'
BUILD_TMP_GENERAL_OUTPUT_DIR_NAME = 'general'
OUTPUT_LIBRARY_DIR_NAME = 'library'
OUTPUT_OPENCL_BINARY_DIR_NAME = 'opencl'
OUTPUT_OPENCL_BINARY_FILE_NAME = 'compiled_opencl_kernel.bin'
CL_COMPILED_BINARY_FILE_NAME = "mace_cl_compiled_program.bin"
CODEGEN_BASE_DIR = 'mace/codegen'
MODEL_CODEGEN_DIR = CODEGEN_BASE_DIR + '/models'
MACE_RUN_TARGET = "//mace/tools/validation:mace_run"
ALL_SOC_TAG = 'all'

ABITypeStrs = [
    'armeabi-v7a',
    'arm64-v8a',
    'host',
]


class ABIType(object):
    armeabi_v7a = 'armeabi-v7a'
    arm64_v8a = 'arm64-v8a'
    host = 'host'


PlatformTypeStrs = [
    "tensorflow",
    "caffe",
]
PlatformType = Enum('PlatformType', [(ele, ele) for ele in PlatformTypeStrs],
                    type=str)

RuntimeTypeStrs = [
    "cpu",
    "gpu",
    "dsp",
    "cpu+gpu"
]


class RuntimeType(object):
    cpu = 'cpu'
    gpu = 'gpu'
    dsp = 'dsp'
    cpu_gpu = 'cpu+gpu'


CPUDataTypeStrs = [
    "fp32",
]

CPUDataType = Enum('CPUDataType', [(ele, ele) for ele in CPUDataTypeStrs],
                   type=str)

GPUDataTypeStrs = [
    "fp16_fp32",
    "fp32_fp32",
]

GPUDataType = Enum('GPUDataType', [(ele, ele) for ele in GPUDataTypeStrs],
                   type=str)

DSPDataTypeStrs = [
    "uint8",
]

DSPDataType = Enum('DSPDataType', [(ele, ele) for ele in DSPDataTypeStrs],
                   type=str)


class DefaultValues(object):
    omp_num_threads = -1,
    cpu_affinity_policy = 1,
    gpu_perf_hint = 3,
    gpu_priority_hint = 3,


class YAMLKeyword(object):
    library_name = 'library_name'
    target_abis = 'target_abis'
    target_socs = 'target_socs'
    build_type = 'build_type'
    embed_model_data = 'embed_model_data'
    models = 'models'
    platform = 'platform'
    model_file_path = 'model_file_path'
    model_sha256_checksum = 'model_sha256_checksum'
    weight_file_path = 'weight_file_path'
    weight_sha256_checksum = 'weight_sha256_checksum'
    subgraphs = 'subgraphs'
    input_tensors = 'input_tensors'
    input_shapes = 'input_shapes'
    output_tensors = 'output_tensors'
    output_shapes = 'output_shapes'
    runtime = 'runtime'
    data_type = 'data_type'
    limit_opencl_kernel_time = 'limit_opencl_kernel_time'
    nnlib_graph_mode = 'nnlib_graph_mode'
    obfuscate = 'obfuscate'
    winograd = 'winograd'
    validation_inputs_data = 'validation_inputs_data'


class ModuleName(object):
    YAML_CONFIG = 'YAML CONFIG'
    MODEL_CONVERTER = 'Model Converter'
    RUN = 'RUN'
    BENCHMARK = 'Benchmark'


CPP_KEYWORDS = [
    'alignas', 'alignof', 'and', 'and_eq', 'asm', 'atomic_cancel',
    'atomic_commit', 'atomic_noexcept', 'auto', 'bitand', 'bitor',
    'bool', 'break', 'case', 'catch', 'char', 'char16_t', 'char32_t',
    'class', 'compl', 'concept', 'const', 'constexpr', 'const_cast',
    'continue', 'co_await', 'co_return', 'co_yield', 'decltype', 'default',
    'delete', 'do', 'double', 'dynamic_cast', 'else', 'enum', 'explicit',
    'export', 'extern', 'false', 'float', 'for', 'friend', 'goto', 'if',
    'import', 'inline', 'int', 'long', 'module', 'mutable', 'namespace',
    'new', 'noexcept', 'not', 'not_eq', 'nullptr', 'operator', 'or', 'or_eq',
    'private', 'protected', 'public', 'register', 'reinterpret_cast',
    'requires', 'return', 'short', 'signed', 'sizeof', 'static',
    'static_assert', 'static_cast', 'struct', 'switch', 'synchronized',
    'template', 'this', 'thread_local', 'throw', 'true', 'try', 'typedef',
    'typeid', 'typename', 'union', 'unsigned', 'using', 'virtual', 'void',
    'volatile', 'wchar_t', 'while', 'xor', 'xor_eq', 'override', 'final',
    'transaction_safe', 'transaction_safe_dynamic', 'if', 'elif', 'else',
    'endif', 'defined', 'ifdef', 'ifndef', 'define', 'undef', 'include',
    'line', 'error', 'pragma',
]


################################
# common functions
################################
def parse_device_type(runtime):
    device_type = ""

    if runtime == RuntimeType.dsp:
        device_type = DeviceType.HEXAGON
    elif runtime == RuntimeType.gpu:
        device_type = DeviceType.GPU
    elif runtime == RuntimeType.cpu:
        device_type = DeviceType.CPU

    return device_type


def get_hexagon_mode(configs):
    runtime_list = []
    for model_name in configs[YAMLKeyword.models]:
        model_runtime =\
            configs[YAMLKeyword.models][model_name].get(
                YAMLKeyword.runtime, "")
        runtime_list.append(model_runtime.lower())

    if RuntimeType.dsp in runtime_list:
        return True
    return False


def md5sum(str):
    md5 = hashlib.md5()
    md5.update(str)
    return md5.hexdigest()


def sha256_checksum(fname):
    hash_func = hashlib.sha256()
    with open(fname, "rb") as f:
        for chunk in iter(lambda: f.read(4096), b""):
            hash_func.update(chunk)
    return hash_func.hexdigest()


def format_model_config(config_file_path):
    with open(config_file_path) as f:
        configs = yaml.load(f)

    library_name = configs.get(YAMLKeyword.library_name, "")
    mace_check(len(library_name) > 0,
               ModuleName.YAML_CONFIG, "library name should not be empty")

    target_abis = configs.get(YAMLKeyword.target_abis, [])
    mace_check((isinstance(target_abis, list) and len(target_abis) > 0),
               ModuleName.YAML_CONFIG, "target_abis list is needed")
    for abi in target_abis:
        mace_check(abi in ABITypeStrs,
                   ModuleName.YAML_CONFIG,
                   "target_abis must be in " + str(ABITypeStrs))

    target_socs = configs.get(YAMLKeyword.target_socs, "")
    if not target_socs:
        configs[YAMLKeyword.target_socs] = []
    elif not isinstance(target_socs, list):
        configs[YAMLKeyword.target_socs] = [target_socs]

    configs[YAMLKeyword.target_socs] = \
        [soc.lower() for soc in configs[YAMLKeyword.target_socs]]

    if ABIType.armeabi_v7a in target_abis \
            or ABIType.arm64_v8a in target_abis:
        available_socs = sh_commands.adb_get_all_socs()
        target_socs = configs[YAMLKeyword.target_socs]
        if ALL_SOC_TAG in target_socs:
            mace_check(available_socs,
                       ModuleName.YAML_CONFIG,
                       "Build for all SOCs plugged in computer, "
                       "you at least plug in one phone")
        else:
            for soc in target_socs:
                mace_check(soc in available_socs,
                           ModuleName.YAML_CONFIG,
                           "Build specified SOC library, "
                           "you must plug in a phone using the SOC")

    build_type = BuildType.code
    build_type_str = configs.get(YAMLKeyword.build_type, "")
    if build_type_str == BuildType.proto:
        build_type = BuildType.proto
    elif build_type_str == BuildType.code:
        build_type = BuildType.code
    else:
        MaceLogger.error(ModuleName.YAML_CONFIG,
                         "Invalid build type " + build_type_str
                         + ". only support [proto|code] format, "
                         + "proto for converting model to ProtoBuf file, "
                         + "code for converting model to c++ code.")

    embed_model_data = configs.get(YAMLKeyword.embed_model_data, "")
    if embed_model_data == "" or not isinstance(embed_model_data, int) or \
       embed_model_data < 0 or embed_model_data > 1:
        MaceLogger.error(ModuleName.YAML_CONFIG,
                         "embed_model_data must be 0 or 1. "
                         "0 for embed model data to code, 1 not.")
    if build_type == BuildType.proto:
        configs[YAMLKeyword.embed_model_data] = 0

    model_names = configs.get(YAMLKeyword.models, [])
    mace_check(len(model_names) > 0, ModuleName.YAML_CONFIG,
               "no model found in config file")

    model_name_reg = re.compile(r'^[a-zA-Z0-9_]+$')
    for model_name in model_names:
        # check model_name legality
        mace_check(model_name not in CPP_KEYWORDS,
                   ModuleName.YAML_CONFIG,
                   "model name should not be c++ keyword.")
        mace_check((model_name[0] == '_' or model_name[0].isalpha())
                   and bool(model_name_reg.match(model_name)),
                   ModuleName.YAML_CONFIG,
                   "model name should Meet the c++ naming convention"
                   " which start with '_' or alpha"
                   " and only contain alpha, number and '_'")

        model_config = configs[YAMLKeyword.models][model_name]
        platform = model_config.get(YAMLKeyword.platform, "")
        mace_check(platform in PlatformTypeStrs,
                   ModuleName.YAML_CONFIG,
                   "'platform' must be in " + str(PlatformTypeStrs))

        for key in [YAMLKeyword.model_file_path,
                    YAMLKeyword.model_sha256_checksum]:
            value = model_config.get(key, "")
            mace_check(value != "", ModuleName.YAML_CONFIG,
                       "'%s' is necessary" % key)

        weight_file_path = model_config.get(YAMLKeyword.weight_file_path, "")
        if weight_file_path:
            weight_checksum =\
                model_config.get(YAMLKeyword.weight_sha256_checksum, "")
            mace_check(weight_checksum != "", ModuleName.YAML_CONFIG,
                       "'%s' is necessary" %
                       YAMLKeyword.weight_sha256_checksum)
        else:
            model_config[YAMLKeyword.weight_sha256_checksum] = ""

        runtime = model_config.get(YAMLKeyword.runtime, "")
        mace_check(runtime in RuntimeTypeStrs,
                   ModuleName.YAML_CONFIG,
                   "'runtime' must be in " + str(RuntimeTypeStrs))
        if ABIType.host in target_abis:
            mace_check(runtime == RuntimeType.cpu,
                       ModuleName.YAML_CONFIG,
                       "host only support cpu runtime now.")

        data_type = model_config.get(YAMLKeyword.data_type, "")
        if runtime == RuntimeType.cpu_gpu and data_type not in GPUDataTypeStrs:
            model_config[YAMLKeyword.data_type] = \
                GPUDataType.fp16_fp32.value
        elif runtime == RuntimeType.cpu:
            if len(data_type) > 0:
                mace_check(data_type in CPUDataTypeStrs,
                           ModuleName.YAML_CONFIG,
                           "'data_type' must be in " + str(CPUDataTypeStrs)
                           + " for cpu runtime")
            else:
                model_config[YAMLKeyword.data_type] = \
                    CPUDataType.fp32.value
        elif runtime == RuntimeType.gpu:
            if len(data_type) > 0:
                mace_check(data_type in GPUDataTypeStrs,
                           ModuleName.YAML_CONFIG,
                           "'data_type' must be in " + str(GPUDataTypeStrs)
                           + " for gpu runtime")
            else:
                model_config[YAMLKeyword.data_type] =\
                    GPUDataType.fp16_fp32.value
        elif runtime == RuntimeType.dsp:
            if len(data_type) > 0:
                mace_check(data_type in DSPDataTypeStrs,
                           ModuleName.YAML_CONFIG,
                           "'data_type' must be in " + str(DSPDataTypeStrs)
                           + " for dsp runtime")
            else:
                model_config[YAMLKeyword.data_type] = \
                    DSPDataType.uint8.value

        subgraphs = model_config.get(YAMLKeyword.subgraphs, "")
        mace_check(len(subgraphs) > 0, ModuleName.YAML_CONFIG,
                   "at least one subgraph is needed")

        for subgraph in subgraphs:
            for key in [YAMLKeyword.input_tensors,
                        YAMLKeyword.input_shapes,
                        YAMLKeyword.output_tensors,
                        YAMLKeyword.output_shapes]:
                value = subgraph.get(key, "")
                mace_check(value != "", ModuleName.YAML_CONFIG,
                           "'%s' is necessary in subgraph" % key)
                if not isinstance(value, list):
                    subgraph[key] = [value]

        for key in [YAMLKeyword.limit_opencl_kernel_time,
                    YAMLKeyword.nnlib_graph_mode,
                    YAMLKeyword.obfuscate,
                    YAMLKeyword.winograd]:
            value = model_config.get(key, "")
            if value == "":
                model_config[key] = 0

        validation_inputs_data = model_config.get(
            YAMLKeyword.validation_inputs_data, [])
        if not isinstance(validation_inputs_data, list):
            model_config[YAMLKeyword.validation_inputs_data] = [
                validation_inputs_data]
        else:
            model_config[YAMLKeyword.validation_inputs_data] = \
                validation_inputs_data

        weight_file_path = model_config.get(YAMLKeyword.weight_file_path, "")
        model_config[YAMLKeyword.weight_file_path] = weight_file_path

    return configs


def get_build_binary_dir(library_name, target_abi, target_soc,
                         serial_num):
    if not target_soc or not serial_num:
        binary_path_digest = md5sum(target_abi)
        binary_path_digest = "%s_%s" % (target_abi, binary_path_digest)
    else:
        device_name = sh_commands.adb_get_device_name_by_serialno(serial_num)
        binary_path_digest = md5sum(target_abi + target_soc + serial_num)
        binary_path_digest = "%s_%s_%s" % \
                             (device_name, target_soc, binary_path_digest)
    return "%s/%s/%s/%s" % (
        BUILD_OUTPUT_DIR, library_name, BUILD_TMP_DIR_NAME, binary_path_digest)


def get_build_model_dirs(library_name, model_name, target_abi, target_soc,
                         serial_num, model_file_path):
    model_path_digest = md5sum(model_file_path)
    model_output_base_dir = "%s/%s/%s/%s/%s" % (
        BUILD_OUTPUT_DIR, library_name, BUILD_TMP_DIR_NAME,
        model_name, model_path_digest)

    if target_abi == ABIType.host:
        model_output_dir = "%s/%s" % (model_output_base_dir, target_abi)
    elif not target_soc or not serial_num:
        model_output_dir = "%s/%s/%s" % (
            model_output_base_dir, BUILD_TMP_GENERAL_OUTPUT_DIR_NAME,
            target_abi)
    else:
        device_name = \
            sh_commands.adb_get_device_name_by_serialno(serial_num)
        model_output_dir = "%s/%s_%s/%s" % (
            model_output_base_dir, device_name,
            target_soc, target_abi)

    mace_model_dir = \
        '%s/%s/%s' % (BUILD_OUTPUT_DIR, library_name, MODEL_OUTPUT_DIR_NAME)

    return model_output_base_dir, model_output_dir, mace_model_dir


def get_opencl_binary_output_path(library_name):
    return '%s/%s/%s/%s' % (BUILD_OUTPUT_DIR,
                            library_name,
                            OUTPUT_OPENCL_BINARY_DIR_NAME,
                            OUTPUT_OPENCL_BINARY_FILE_NAME)


################################
# build
################################
def pull_opencl_binary_and_tuning_param(target_abi,
                                        serialno,
                                        model_output_dirs):
    sh_commands.pull_binaries(target_abi, serialno, model_output_dirs,
                              CL_COMPILED_BINARY_FILE_NAME)


def print_configuration(flags, configs):
    title = "Common Configuration"
    header = ["key", "value"]
    data = list()
    data.append([YAMLKeyword.library_name,
                 configs[YAMLKeyword.library_name]])
    data.append([YAMLKeyword.target_abis,
                 configs[YAMLKeyword.target_abis]])
    data.append([YAMLKeyword.target_socs,
                 configs[YAMLKeyword.target_socs]])
    data.append([YAMLKeyword.build_type,
                 configs[YAMLKeyword.build_type]])
    data.append([YAMLKeyword.embed_model_data,
                 configs[YAMLKeyword.embed_model_data]])
    data.append(["Tuning", flags.tuning])
    MaceLogger.summary(StringFormatter.table(header, data, title))


def download_model_files(model_file_path,
                         model_output_dir,
                         weight_file_path=""):
    if model_file_path.startswith("http://") or \
            model_file_path.startswith("https://"):
        model_file = model_output_dir + "/model.pb"
        urllib.urlretrieve(model_file_path, model_file)

    if weight_file_path.startswith("http://") or \
            weight_file_path.startswith("https://"):
        weight_file = model_output_dir + "/model.caffemodel"
        urllib.urlretrieve(weight_file_path, weight_file)


def get_model_files_path(model_file_path,
                         model_output_dir,
                         weight_file_path=""):
    if model_file_path.startswith("http://") or \
            model_file_path.startswith("https://"):
        model_file = model_output_dir + "/model.pb"
    else:
        model_file = model_file_path

    if weight_file_path.startswith("http://") or \
            weight_file_path.startswith("https://"):
        weight_file = model_output_dir + "/model.caffemodel"
    else:
        weight_file = weight_file_path

    return model_file, weight_file


def convert_model(configs):
    # Remove previous output dirs
    library_name = configs[YAMLKeyword.library_name]
    if not os.path.exists(BUILD_OUTPUT_DIR):
        os.makedirs(BUILD_OUTPUT_DIR)
    elif os.path.exists(os.path.join(BUILD_OUTPUT_DIR, library_name)):
        sh.rm("-rf", os.path.join(BUILD_OUTPUT_DIR, library_name))
    os.makedirs(os.path.join(BUILD_OUTPUT_DIR, library_name))

    model_output_dir = \
        '%s/%s/%s' % (BUILD_OUTPUT_DIR, library_name, MODEL_OUTPUT_DIR_NAME)
    if os.path.exists(model_output_dir):
        sh.rm("-rf", model_output_dir)
    os.makedirs(model_output_dir)

    embed_model_data = configs[YAMLKeyword.embed_model_data]

    sh_commands.clear_model_codegen()
    for model_name in configs[YAMLKeyword.models]:
        MaceLogger.header(
            StringFormatter.block("Convert %s model" % model_name))
        model_config = configs[YAMLKeyword.models][model_name]
        runtime = model_config[YAMLKeyword.runtime]

        # Create model build directory
        model_path_digest = md5sum(
            model_config[YAMLKeyword.model_file_path])

        model_output_base_dir = "%s/%s/%s/%s/%s" % (
            BUILD_OUTPUT_DIR, library_name, BUILD_TMP_DIR_NAME,
            model_name, model_path_digest)

        if os.path.exists(model_output_base_dir):
            sh.rm("-rf", model_output_base_dir)
        os.makedirs(model_output_base_dir)

        download_model_files(
            model_config[YAMLKeyword.model_file_path],
            model_output_base_dir,
            model_config[YAMLKeyword.weight_file_path])

        model_file_path, weight_file_path = get_model_files_path(
            model_config[YAMLKeyword.model_file_path],
            model_output_base_dir,
            model_config[YAMLKeyword.weight_file_path])

        if sha256_checksum(model_file_path) != \
                model_config[YAMLKeyword.model_sha256_checksum]:
            MaceLogger.error(ModuleName.MODEL_CONVERTER,
                             "model file sha256checksum not match")

        if weight_file_path:
            if sha256_checksum(weight_file_path) != \
                    model_config[YAMLKeyword.weight_sha256_checksum]:
                MaceLogger.error(ModuleName.MODEL_CONVERTER,
                                 "weight file sha256checksum not match")

        data_type = model_config[YAMLKeyword.data_type]
        # TODO(liuqi): support multiple subgraphs
        subgraphs = model_config[YAMLKeyword.subgraphs]

        model_codegen_dir = "%s/%s" % (MODEL_CODEGEN_DIR, model_name)
        sh_commands.gen_model_code(
            model_codegen_dir,
            model_config[YAMLKeyword.platform],
            model_file_path,
            weight_file_path,
            model_config[YAMLKeyword.model_sha256_checksum],
            model_config[YAMLKeyword.weight_sha256_checksum],
            ",".join(subgraphs[0][YAMLKeyword.input_tensors]),
            ",".join(subgraphs[0][YAMLKeyword.output_tensors]),
            runtime,
            model_name,
            ":".join(subgraphs[0][YAMLKeyword.input_shapes]),
            model_config[YAMLKeyword.nnlib_graph_mode],
            embed_model_data,
            model_config[YAMLKeyword.winograd],
            model_config[YAMLKeyword.obfuscate],
            configs[YAMLKeyword.build_type],
            data_type)

        # mv pb and data file to build/model_name/model
        if not embed_model_data:
            sh_commands.mv_model_file_to_output_dir(
                model_build_type=configs[YAMLKeyword.build_type],
                model_codegen_dir=model_codegen_dir,
                model_name=model_name,
                output_dir=model_output_dir
            )

        MaceLogger.summary(
            StringFormatter.block("Model %s converted" % model_name))


def build_specific_lib(target_abi, target_soc, serial_num,
                       configs, tuning, enable_openmp,
                       address_sanitizer):
    library_name = configs[YAMLKeyword.library_name]
    build_type = configs[YAMLKeyword.build_type]
    embed_model_data = configs[YAMLKeyword.embed_model_data]
    hexagon_mode = get_hexagon_mode(configs)
    model_output_dirs = []

    build_tmp_binary_dir = get_build_binary_dir(library_name, target_abi,
                                                target_soc, serial_num)
    if os.path.exists(build_tmp_binary_dir):
        sh.rm("-rf", build_tmp_binary_dir)
    os.makedirs(build_tmp_binary_dir)

    sh_commands.gen_tuning_param_code(model_output_dirs)
    sh_commands.bazel_build(
        MACE_RUN_TARGET,
        abi=target_abi,
        hexagon_mode=hexagon_mode,
        enable_openmp=enable_openmp,
        address_sanitizer=address_sanitizer
    )
    sh_commands.update_mace_run_lib(build_tmp_binary_dir)
    binary_changed = False

    for model_name in configs[YAMLKeyword.models]:
        model_config = configs[YAMLKeyword.models][model_name]
        model_runtime = model_config[YAMLKeyword.runtime]
        # Create model build directory
        model_output_base_dir, model_output_dir, mace_model_dir = \
            get_build_model_dirs(library_name, model_name, target_abi,
                                 target_soc, serial_num,
                                 model_config[YAMLKeyword.model_file_path])

        model_output_dirs.append(model_output_dir)

        if os.path.exists(model_output_dir):
            sh.rm("-rf", model_output_dir)
        os.makedirs(model_output_dir)

        # build for specified soc
        if not address_sanitizer and target_abi != ABIType.host \
                and target_soc is not None and \
                model_runtime in [RuntimeType.gpu, RuntimeType.cpu_gpu]:
            sh_commands.clear_phone_data_dir(serial_num, PHONE_DATA_DIR)

            subgraphs = model_config[YAMLKeyword.subgraphs]
            # generate input data
            sh_commands.gen_random_input(
                model_output_dir,
                subgraphs[0][YAMLKeyword.input_tensors],
                subgraphs[0][YAMLKeyword.input_shapes],
                subgraphs[0][YAMLKeyword.validation_inputs_data])

            device_type = parse_device_type(RuntimeType.gpu)
            sh_commands.tuning_run(
                abi=target_abi,
                serialno=serial_num,
                mace_run_dir=build_tmp_binary_dir,
                vlog_level=0,
                embed_model_data=embed_model_data,
                model_output_dir=model_output_dir,
                input_nodes=subgraphs[0][YAMLKeyword.input_tensors],
                output_nodes=subgraphs[0][YAMLKeyword.output_tensors],
                input_shapes=subgraphs[0][YAMLKeyword.input_shapes],
                output_shapes=subgraphs[0][YAMLKeyword.output_shapes],
                mace_model_dir=mace_model_dir,
                model_tag=model_name,
                device_type=device_type,
                running_round=0,
                restart_round=1,
                limit_opencl_kernel_time=model_config[YAMLKeyword.limit_opencl_kernel_time],  # noqa
                tuning=tuning,
                out_of_range_check=False,
                phone_data_dir=PHONE_DATA_DIR,
                build_type=build_type,
                opencl_binary_file="",
            )

            pull_opencl_binary_and_tuning_param(target_abi, serial_num,
                                                [model_output_dir])
            sh_commands.touch_tuned_file_flag(build_tmp_binary_dir)
            binary_changed = True

    if binary_changed:
        sh_commands.merge_opencl_binaries(
            model_output_dirs, CL_COMPILED_BINARY_FILE_NAME,
            get_opencl_binary_output_path(library_name))
        sh_commands.gen_tuning_param_code(model_output_dirs)
        sh_commands.bazel_build(
            MACE_RUN_TARGET,
            abi=target_abi,
            hexagon_mode=hexagon_mode,
            enable_openmp=enable_openmp,
            address_sanitizer=address_sanitizer
        )
        sh_commands.update_mace_run_lib(build_tmp_binary_dir)

    if target_abi == ABIType.host:
        sh_commands.build_host_libraries(build_type, target_abi)

    # build benchmark_model binary
    sh_commands.build_benchmark_model(target_abi,
                                      build_tmp_binary_dir,
                                      hexagon_mode)

    # generate library
    sh_commands.merge_libs(target_soc,
                           serial_num,
                           target_abi,
                           library_name,
                           BUILD_OUTPUT_DIR,
                           OUTPUT_LIBRARY_DIR_NAME,
                           build_type,
                           hexagon_mode)


def generate_library(configs, tuning, enable_openmp, address_sanitizer):
    MaceLogger.header(StringFormatter.block("Building library"))
    # generate source
    MaceLogger.info('* generate common source files...')
    sh_commands.gen_mace_version()
    sh_commands.gen_encrypted_opencl_source()
    sh_commands.gen_mace_engine_factory_source(
        configs[YAMLKeyword.models].keys(),
        configs[YAMLKeyword.build_type])
    MaceLogger.info('generate common source files done')

    # create build dirs
    library_name = configs[YAMLKeyword.library_name]
    if not os.path.exists(BUILD_OUTPUT_DIR):
        os.makedirs(BUILD_OUTPUT_DIR)
    tmp_build_dir = os.path.join(BUILD_OUTPUT_DIR, library_name,
                                 BUILD_TMP_DIR_NAME)
    if not os.path.exists(tmp_build_dir):
        os.makedirs(tmp_build_dir)
    library_out_dir = os.path.join(BUILD_OUTPUT_DIR, library_name,
                                   OUTPUT_LIBRARY_DIR_NAME)
    if os.path.exists(library_out_dir):
        sh.rm('-rf', library_out_dir)

    target_socs = configs[YAMLKeyword.target_socs]
    for target_abi in configs[YAMLKeyword.target_abis]:
        if not target_socs or target_abi == ABIType.host:
            build_specific_lib(target_abi, None, None, configs,
                               tuning, enable_openmp, address_sanitizer)
        else:
            if ALL_SOC_TAG in target_socs:
                target_socs = sh_commands.adb_get_all_socs()
            for target_soc in target_socs:
                serial_nums = \
                    sh_commands.get_target_socs_serialnos([target_soc])
                for serial_num in serial_nums:
                    with sh_commands.device_lock(serial_num):
                        build_specific_lib(target_abi, target_soc, serial_num,
                                           configs, tuning, enable_openmp,
                                           address_sanitizer)

    # package library
    sh_commands.packaging_lib(BUILD_OUTPUT_DIR,
                              configs[YAMLKeyword.library_name])


def print_library_summary(configs):
    library_name = configs[YAMLKeyword.library_name]
    title = "Library"
    header = ["key", "value"]
    data = list()
    data.append(["library package",
                 "%s/%s/libmace_%s.tar.gz"
                 % (BUILD_OUTPUT_DIR, library_name, library_name)])
    MaceLogger.summary(StringFormatter.table(header, data, title))


def build_library(flags):
    configs = format_model_config(flags.config)

    print_configuration(flags, configs)

    convert_model(configs)

    generate_library(configs, flags.tuning,
                     flags.enable_openmp, flags.address_sanitizer)

    print_library_summary(configs)


################################
# run
################################
def report_run_statistics(stdout,
                          abi,
                          serialno,
                          model_name,
                          device_type,
                          output_dir,
                          tuned):
    metrics = [0] * 3
    for line in stdout.split('\n'):
        line = line.strip()
        parts = line.split()
        if len(parts) == 4 and parts[0].startswith("time"):
            metrics[0] = str(float(parts[1]))
            metrics[1] = str(float(parts[2]))
            metrics[2] = str(float(parts[3]))
            break

    device_name = ""
    target_soc = ""
    if abi != "host":
        props = sh_commands.adb_getprop_by_serialno(serialno)
        device_name = props.get("ro.product.model", "")
        target_soc = props.get("ro.board.platform", "")

    report_filename = output_dir + "/report.csv"
    if not os.path.exists(report_filename):
        with open(report_filename, 'w') as f:
            f.write("model_name,device_name,soc,abi,runtime,"
                    "init,warmup,run_avg,tuned\n")

    data_str = "{model_name},{device_name},{soc},{abi},{device_type}," \
               "{init},{warmup},{run_avg},{tuned}\n" \
        .format(model_name=model_name,
                device_name=device_name,
                soc=target_soc,
                abi=abi,
                device_type=device_type,
                init=metrics[0],
                warmup=metrics[1],
                run_avg=metrics[2],
                tuned=tuned,
                )
    with open(report_filename, 'a') as f:
        f.write(data_str)


def run_specific_target(flags, configs, target_abi,
                        target_soc, serial_num):
    library_name = configs[YAMLKeyword.library_name]
    build_type = configs[YAMLKeyword.build_type]
    embed_model_data = configs[YAMLKeyword.embed_model_data]
    if not configs[YAMLKeyword.target_socs]:
        build_tmp_binary_dir = get_build_binary_dir(library_name, target_abi,
                                                    None, None)
    else:
        build_tmp_binary_dir = get_build_binary_dir(library_name, target_abi,
                                                    target_soc, serial_num)
    mace_check(os.path.exists(build_tmp_binary_dir),
               ModuleName.RUN,
               'You should build before run.')

    for model_name in configs[YAMLKeyword.models]:
        if target_abi == ABIType.host:
            device_name = ABIType.host
        else:
            device_name =\
                sh_commands.adb_get_device_name_by_serialno(serial_num)
        MaceLogger.header(
            StringFormatter.block(
                "Run model %s on %s" % (model_name, device_name)))
        model_config = configs[YAMLKeyword.models][model_name]
        model_runtime = model_config[YAMLKeyword.runtime]
        subgraphs = model_config[YAMLKeyword.subgraphs]

        if not configs[YAMLKeyword.target_socs]:
            model_output_base_dir, model_output_dir, mace_model_dir = \
                get_build_model_dirs(library_name, model_name, target_abi,
                                     None, None,
                                     model_config[YAMLKeyword.model_file_path])
        else:
            model_output_base_dir, model_output_dir, mace_model_dir = \
                get_build_model_dirs(library_name, model_name, target_abi,
                                     target_soc, serial_num,
                                     model_config[YAMLKeyword.model_file_path])
        mace_check(os.path.exists(model_output_dir)
                   and os.path.exists(mace_model_dir),
                   ModuleName.RUN,
                   'You should build before run.')
        if target_abi != ABIType.host:
            sh_commands.clear_phone_data_dir(serial_num, PHONE_DATA_DIR)

        # generate input data
        sh_commands.gen_random_input(
            model_output_dir,
            subgraphs[0][YAMLKeyword.input_tensors],
            subgraphs[0][YAMLKeyword.input_shapes],
            subgraphs[0][YAMLKeyword.validation_inputs_data])
        runtime_list = []
        if target_abi == ABIType.host:
            runtime_list.extend([RuntimeType.cpu])
        elif model_runtime == RuntimeType.cpu_gpu:
            runtime_list.extend([RuntimeType.cpu, RuntimeType.gpu])
        else:
            runtime_list.extend([model_runtime])
        for runtime in runtime_list:
            device_type = parse_device_type(runtime)
            run_output = sh_commands.tuning_run(
                abi=target_abi,
                serialno=serial_num,
                mace_run_dir=build_tmp_binary_dir,
                vlog_level=flags.vlog_level,
                embed_model_data=embed_model_data,
                model_output_dir=model_output_dir,
                input_nodes=subgraphs[0][YAMLKeyword.input_tensors],
                output_nodes=subgraphs[0][YAMLKeyword.output_tensors],
                input_shapes=subgraphs[0][YAMLKeyword.input_shapes],
                output_shapes=subgraphs[0][YAMLKeyword.output_shapes],
                mace_model_dir=mace_model_dir,
                model_tag=model_name,
                device_type=device_type,
                running_round=flags.round,
                restart_round=flags.restart_round,
                limit_opencl_kernel_time=model_config[YAMLKeyword.limit_opencl_kernel_time],  # noqa
                tuning=False,
                out_of_range_check=flags.check_gpu_out_of_memory,
                phone_data_dir=PHONE_DATA_DIR,
                build_type=build_type,
                omp_num_threads=flags.omp_num_threads,
                cpu_affinity_policy=flags.cpu_affinity_policy,
                gpu_perf_hint=flags.gpu_perf_hint,
                gpu_priority_hint=flags.gpu_priority_hint,
                runtime_failure_ratio=flags.runtime_failure_ratio,
                address_sanitizer=flags.address_sanitizer,
                opencl_binary_file=get_opencl_binary_output_path(library_name),
            )
            if flags.validate:
                model_file_path, weight_file_path = get_model_files_path(
                    model_config[YAMLKeyword.model_file_path],
                    model_output_base_dir,
                    model_config[YAMLKeyword.weight_file_path])

                sh_commands.validate_model(
                    abi=target_abi,
                    serialno=serial_num,
                    model_file_path=model_file_path,
                    weight_file_path=weight_file_path,
                    platform=model_config[YAMLKeyword.platform],
                    device_type=device_type,
                    input_nodes=subgraphs[0][YAMLKeyword.input_tensors],
                    output_nodes=subgraphs[0][YAMLKeyword.output_tensors],
                    input_shapes=subgraphs[0][YAMLKeyword.input_shapes],
                    output_shapes=subgraphs[0][YAMLKeyword.output_shapes],
                    model_output_dir=model_output_dir,
                    phone_data_dir=PHONE_DATA_DIR,
                    caffe_env=flags.caffe_env)
            if flags.report and flags.round > 0:
                report_run_statistics(
                    run_output, target_abi, serial_num,
                    model_name, device_type, flags.report_dir,
                    sh_commands.is_binary_tuned(build_tmp_binary_dir))


def run_mace(flags):
    configs = format_model_config(flags.config)

    target_socs = configs[YAMLKeyword.target_socs]
    if not target_socs or ALL_SOC_TAG in target_socs:
        target_socs = sh_commands.adb_get_all_socs()

    for target_abi in configs[YAMLKeyword.target_abis]:
        if target_abi == ABIType.host:
            run_specific_target(flags, configs, target_abi, None, None)
        else:
            for target_soc in target_socs:
                serial_nums = \
                    sh_commands.get_target_socs_serialnos([target_soc])
                mace_check(serial_nums,
                           ModuleName.RUN,
                           'There is no device with soc: ' + target_soc)
                for serial_num in serial_nums:
                    with sh_commands.device_lock(serial_num):
                        run_specific_target(flags, configs, target_abi,
                                            target_soc, serial_num)


################################
#  benchmark model
################################
def bm_specific_target(flags, configs, target_abi, target_soc, serial_num):
    library_name = configs[YAMLKeyword.library_name]
    build_type = configs[YAMLKeyword.build_type]
    embed_model_data = configs[YAMLKeyword.embed_model_data]
    if not configs[YAMLKeyword.target_socs]:
        build_tmp_binary_dir = get_build_binary_dir(library_name, target_abi,
                                                    None, None)
    else:
        build_tmp_binary_dir = get_build_binary_dir(library_name, target_abi,
                                                    target_soc, serial_num)
    mace_check(os.path.exists(build_tmp_binary_dir),
               ModuleName.BENCHMARK,
               'You should build before benchmark.')

    for model_name in configs[YAMLKeyword.models]:
        if target_abi == ABIType.host:
            device_name = ABIType.host
        else:
            device_name = \
                sh_commands.adb_get_device_name_by_serialno(serial_num)
        MaceLogger.header(
            StringFormatter.block(
                "Benchmark model %s on %s" % (model_name, device_name)))
        model_config = configs[YAMLKeyword.models][model_name]
        model_runtime = model_config[YAMLKeyword.runtime]
        subgraphs = model_config[YAMLKeyword.subgraphs]

        if not configs[YAMLKeyword.target_socs]:
            model_output_base_dir, model_output_dir, mace_model_dir = \
                get_build_model_dirs(library_name, model_name, target_abi,
                                     None, None,
                                     model_config[YAMLKeyword.model_file_path])
        else:
            model_output_base_dir, model_output_dir, mace_model_dir = \
                get_build_model_dirs(library_name, model_name, target_abi,
                                     target_soc, serial_num,
                                     model_config[YAMLKeyword.model_file_path])
        mace_check(os.path.exists(model_output_dir)
                   and os.path.exists(mace_model_dir),
                   ModuleName.BENCHMARK,
                   'You should build before benchmark.')
        if target_abi != ABIType.host:
            sh_commands.clear_phone_data_dir(serial_num, PHONE_DATA_DIR)

        sh_commands.gen_random_input(
            model_output_dir,
            subgraphs[0][YAMLKeyword.input_tensors],
            subgraphs[0][YAMLKeyword.input_shapes],
            subgraphs[0][YAMLKeyword.validation_inputs_data])
        runtime_list = []
        if target_abi == ABIType.host:
            runtime_list.extend([RuntimeType.cpu])
        elif model_runtime == RuntimeType.cpu_gpu:
            runtime_list.extend([RuntimeType.cpu, RuntimeType.gpu])
        else:
            runtime_list.extend([model_runtime])
        for runtime in runtime_list:
            device_type = parse_device_type(runtime)
            sh_commands.benchmark_model(
                abi=target_abi,
                serialno=serial_num,
                benchmark_binary_dir=build_tmp_binary_dir,
                vlog_level=0,
                embed_model_data=embed_model_data,
                model_output_dir=model_output_dir,
                input_nodes=subgraphs[0][YAMLKeyword.input_tensors],
                output_nodes=subgraphs[0][YAMLKeyword.output_tensors],
                input_shapes=subgraphs[0][YAMLKeyword.input_shapes],
                output_shapes=subgraphs[0][YAMLKeyword.output_shapes],
                mace_model_dir=mace_model_dir,
                model_tag=model_name,
                device_type=device_type,
                phone_data_dir=PHONE_DATA_DIR,
                build_type=build_type,
                omp_num_threads=flags.omp_num_threads,
                cpu_affinity_policy=flags.cpu_affinity_policy,
                gpu_perf_hint=flags.gpu_perf_hint,
                gpu_priority_hint=flags.gpu_priority_hint,
                opencl_binary_file=get_opencl_binary_output_path(library_name))


def benchmark_model(flags):
    configs = format_model_config(flags.config)

    target_socs = configs[YAMLKeyword.target_socs]
    if not target_socs or ALL_SOC_TAG in target_socs:
        target_socs = sh_commands.adb_get_all_socs()

    for target_abi in configs[YAMLKeyword.target_abis]:
        if target_abi == ABIType.host:
            bm_specific_target(flags, configs, target_abi, None, None)
        else:
            for target_soc in target_socs:
                serial_nums = \
                    sh_commands.get_target_socs_serialnos([target_soc])
                mace_check(serial_nums,
                           ModuleName.BENCHMARK,
                           'There is no device with soc: ' + target_soc)
                for serial_num in serial_nums:
                    with sh_commands.device_lock(serial_num):
                        bm_specific_target(flags, configs, target_abi,
                                           target_soc, serial_num)


################################
# Parsing arguments
################################
def str2bool(v):
    if v.lower() in ('yes', 'true', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected.')


def str_to_caffe_env_type(v):
    if v.lower() == 'docker':
        return CaffeEnvType.DOCKER
    elif v.lower() == 'local':
        return CaffeEnvType.LOCAL
    else:
        raise argparse.ArgumentTypeError('[docker | local] expected.')


def parse_args():
    """Parses command line arguments."""
    all_type_parent_parser = argparse.ArgumentParser(add_help=False)
    all_type_parent_parser.add_argument(
        '--config',
        type=str,
        default="",
        required=True,
        help="model yaml configuration file path")
    build_run_parent_parser = argparse.ArgumentParser(add_help=False)
    build_run_parent_parser.add_argument(
        '--address_sanitizer',
        action="store_true",
        help="Whether to use valgrind to check memory error")
    run_bm_parent_parser = argparse.ArgumentParser(add_help=False)
    run_bm_parent_parser.add_argument(
        "--omp_num_threads",
        type=int,
        default=DefaultValues.omp_num_threads,
        help="num of openmp threads")
    run_bm_parent_parser.add_argument(
        "--cpu_affinity_policy",
        type=int,
        default=DefaultValues.cpu_affinity_policy,
        help="0:AFFINITY_NONE/1:AFFINITY_BIG_ONLY/2:AFFINITY_LITTLE_ONLY")
    run_bm_parent_parser.add_argument(
        "--gpu_perf_hint",
        type=int,
        default=DefaultValues.gpu_perf_hint,
        help="0:DEFAULT/1:LOW/2:NORMAL/3:HIGH")
    run_bm_parent_parser.add_argument(
        "--gpu_priority_hint",
        type=int,
        default=DefaultValues.gpu_priority_hint,
        help="0:DEFAULT/1:LOW/2:NORMAL/3:HIGH")

    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers()
    build = subparsers.add_parser(
        'build',
        parents=[all_type_parent_parser, build_run_parent_parser],
        help='build model library and test tools')
    build.set_defaults(func=build_library)
    build.add_argument(
        '--tuning',
        action="store_true",
        help="tuning gpu parameters for specified SOC")
    build.add_argument(
        "--enable_openmp",
        action="store_false",
        help="Enable openmp for multiple thread.")
    run = subparsers.add_parser(
        'run',
        parents=[all_type_parent_parser, run_bm_parent_parser,
                 build_run_parent_parser],
        help='run model in command line')
    run.set_defaults(func=run_mace)
    run.add_argument(
        "--round",
        type=int,
        default=1,
        help="The model running round.")
    run.add_argument(
        "--validate",
        action="store_true",
        help="validate result by comparing mace output and platform's output.")
    run.add_argument(
        "--caffe_env",
        type=str_to_caffe_env_type,
        default='docker',
        help="[docker | local] caffe environment.")
    run.add_argument(
        "--vlog_level",
        type=int,
        default=0,
        help="VLOG level: [1~5].")
    run.add_argument(
        "--check_gpu_out_of_memory",
        action="store_true",
        help="Enable out of memory check for gpu.")
    run.add_argument(
        "--restart_round",
        type=int,
        default=1,
        help="restart round for run.")
    run.add_argument(
        "--report",
        action="store_true",
        help="print run statistics report.")
    run.add_argument(
        "--report_dir",
        type=str,
        default="",
        help="print run statistics report.")
    run.add_argument(
        "--runtime_failure_ratio",
        type=float,
        default=0.0,
        help="[mock runtime failure ratio].")
    benchmark = subparsers.add_parser(
        'benchmark',
        parents=[all_type_parent_parser, run_bm_parent_parser,
                 build_run_parent_parser],
        help='benchmark model for detail information')
    benchmark.set_defaults(func=benchmark_model)
    return parser.parse_known_args()


if __name__ == "__main__":
    flags, unparsed = parse_args()
    flags.func(flags)