util.py

# 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.
"""Utility for benchmark"""

import os
import numpy as np
import sys, getopt
import tvm
from tvm import relay
from tvm.relay import testing
from tvm.contrib.download import download_testdata
from tvm.relay.op.contrib import arm_compute_lib
from PIL import Image
from tvm.contrib.download import download_testdata
import cpuinfo
import multiprocessing

def parse_cmd_options(argv):
    device='llvm'
    logfile=None

    try:
        opts, args = getopt.getopt(argv, "hd:l:")
    except getopt.GetoptError:
        print('python3 blahblah.py -d <llvm|arm_cpu>')
        sys.exit()
    for opt,arg in opts:
        if opt == '-h':
            print('python3 blahblah.py -d <llvm|arm_cpu>')
            sys.exit()
        elif opt == '-d':
            device=arg
        elif opt == '-l':
            logfile=arg
    return device, logfile

def parse_options(argv):
    device='llvm'
    logfile=None

    try:
        opts, args = getopt.getopt(argv, "hd:l:")
    except getopt.GetoptError:
        print('python3 blahblah.py -d <llvm|arm_cpu>')
        sys.exit()
    for opt,arg in opts:
        if opt == '-h':
            print('python3 blahblah.py -d <llvm|arm_cpu>')
            sys.exit()
        elif opt == '-d':
            device=arg
        elif opt == '-l':
            logfile=arg
    return device


def get_device_arch():
    arch = cpuinfo.get_cpu_info()['arch_string_raw']
    foo = cpuinfo.get_cpu_info()
    return arch

def get_cpu_count():
    cpu_count = multiprocessing.cpu_count()
    return cpu_count

def get_device_attributes():
    if get_device_arch() == "aarch64":
        if get_device_type() == 'thunderxt88':
            return '+neon,+crc,+lse'
        elif get_device_type() == 'cortex-a78':
            return '+neon,+crc'
        return '+neon'
    if get_device_arch() == "armv7l":
        return  '+neon,+vfp4'
    else:
        return  ''

def get_device_type():
    cpudictionary=cpuinfo.get_cpu_info()
    brand=cpudictionary.get('brand_raw')
    cpuversion=cpudictionary.get('cpuinfo_version_string')
    if brand == 'ThunderX 88XX':
        return 'thunderxt88'
    elif brand == 'ARMv7 Processor rev 5 (v7l)':
        return 'cortex-a7'
    elif brand is None:
        brand=cpudictionary.get('vendor_id_raw')
        if brand =='Qualcomm' :
            return 'cortex-a75'
        elif brand == 'ARM':
            return 'native'
        #elif cpuversion =='7.0.0':
        #    return 'cortex-a78'
        else:
            return 'native'
    else:
        return ' '

def get_tvm_target(device, dev_type, arch_token, attributes):
    if arch_token == "aarch64":
        arch_token = 'aarch64-unknown-linux-gnu'
    if arch_token == "armv7a" or arch_token == "armv7l":
        arch_token = "armv7a-linux-gnueabihf"
    if device in ("llvm"):
        target_string = "llvm -mcpu=" + dev_type + " -mtriple=" + arch_token + " -mattr=" + attributes
    else:
        target_string = "llvm -device=arm_cpu" + " -mcpu=" + dev_type + " -mtriple=" + arch_token + " -mattr=" + attributes
    return target_string

def download_model_zoo(model_dir, model_name, url='http://people.linaro.org/~tom.gall/model_zoo/'):
    model_url = url + model_dir + model_name
    model_path = download_testdata(model_url, model_name, module=["tf", "official"])
    model_dir = os.path.dirname(model_path)
    return model_dir
   

def extract(path):
    import tarfile
    if path.endswith("tgz") or path.endswith("gz"):
        dir_path = os.path.dirname(path)
        tar = tarfile.open(path)
        tar.extractall(path=dir_path)
        tar.close()
    else:
        raise RuntimeError('Could not decompress the file: ' + path)

def load_test_image(dtype='float32', width=224, height=224):
    image_url = 'https://github.com/dmlc/mxnet.js/blob/master/data/cat.png?raw=true'
    image_path = download_testdata(image_url, 'cat.png', module='data')
    resized_image = Image.open(image_path).resize((width, height))

    #image_data = np.asarray(resized_image).astype("float32")
    image_data = np.asarray(resized_image).astype(dtype)

    # Add a dimension to the image so that we have NHWC format layout
    image_data = np.expand_dims(image_data, axis=0)

    if (dtype=='float32'):
    	# Preprocess image as described here:
    	# https://github.com/tensorflow/models/blob/edb6ed22a801665946c63d650ab9a0b23d98e1b1/research/slim/preprocessing/inception_preprocessing.py#L243
    	image_data[:, :, :, 0] = 2.0 / 255.0 * image_data[:, :, :, 0] - 1
    	image_data[:, :, :, 1] = 2.0 / 255.0 * image_data[:, :, :, 1] - 1
    	image_data[:, :, :, 2] = 2.0 / 255.0 * image_data[:, :, :, 2] - 1    	
    print('input', image_data.shape)
    return image_data


def get_network(name, batch_size, dtype='float32'):
    """Get the symbol definition and random weight of a network

    Parameters
    ----------
    name: str
        The name of the network, can be 'resnet-18', 'resnet-50', 'vgg-16', 'inception_v3', 'mobilenet', ...
    batch_size: int
        batch size
    dtype: str
        Data type

    Returns
    -------
    net: tvm.IRModule
        The relay function of network definition
    params: dict
        The random parameters for benchmark
    input_shape: tuple
        The shape of input tensor
    output_shape: tuple
        The shape of output tensor
    """
    input_shape = (batch_size, 3, 224, 224)
    output_shape = (batch_size, 1000)

    if name == 'mobilenet':
        net, params = testing.mobilenet.get_workload(batch_size=batch_size, dtype=dtype)
    elif name == 'inception_v3':
        input_shape = (batch_size, 3, 299, 299)
        net, params = testing.inception_v3.get_workload(batch_size=batch_size, dtype=dtype)
    elif "resnet" in name:
        n_layer = int(name.split('-')[1])
        net, params = testing.resnet.get_workload(num_layers=n_layer, batch_size=batch_size, dtype=dtype)
    elif "vgg" in name:
        n_layer = int(name.split('-')[1])
        net, params = testing.vgg.get_workload(num_layers=n_layer, batch_size=batch_size, dtype=dtype)
    elif "densenet" in name:
        n_layer = int(name.split('-')[1])
        net, params = testing.densenet.get_workload(densenet_size=n_layer, batch_size=batch_size, dtype=dtype)
    elif "squeezenet" in name:
        version = name.split("_v")[1]
        net, params = testing.squeezenet.get_workload(batch_size=batch_size, version=version, dtype=dtype)
    elif name == 'mxnet':
        # an example for mxnet model
        from mxnet.gluon.model_zoo.vision import get_model
        block = get_model('resnet18_v1', pretrained=True)
        net, params = relay.frontend.from_mxnet(block, shape={'data': input_shape}, dtype=dtype)
        net = net["main"]
        net = relay.Function(net.params, relay.nn.softmax(net.body), None, net.type_params, net.attrs)
        net = tvm.IRModule.from_expr(net)
    else:
        raise ValueError("Unsupported network: " + name)

    return net, params, input_shape, output_shape

def print_progress(msg):
    """print progress message

    Parameters
    ----------
    msg: str
        The message to print
    """
    sys.stdout.write(msg + "\r")
    sys.stdout.flush()

# Experimental scaffolding code for using ACL
def build_module(mod, target, params=None, enable_acl=True, tvm_ops=0, acl_partitions=1):
    """Build module with option to build for ACL."""
    if isinstance(mod, tvm.relay.expr.Call):
        mod = tvm.IRModule.from_expr(mod)
    with tvm.transform.PassContext(opt_level=3, disabled_pass=["AlterOpLayout"]):
        if enable_acl:
            mod = arm_compute_lib.partition_for_arm_compute_lib(mod, params)
            tvm_op_count = get_cpu_op_count(mod)
            assert tvm_op_count == tvm_ops, "Got {} TVM operators, expected {}".format(
                tvm_op_count, tvm_ops
            )
            partition_count = 0
            for global_var in mod.get_global_vars():
                if "arm_compute_lib" in global_var.name_hint:
                    partition_count += 1

            assert (
                acl_partitions == partition_count
            ), "Got {} Arm Compute Library partitions, expected {}".format(
                partition_count, acl_partitions
            )
        relay.backend.compile_engine.get().clear()
        return relay.build(mod, target=target, params=params)


def update_lib(lib, device, cross_compile):
    """Export the library to the remote/local device."""
    lib_name = "mod.so"
    temp = util.tempdir()
    lib_path = temp.relpath(lib_name)
    if cross_compile:
        lib.export_library(lib_path, cc=cross_compile)
    else:
        lib.export_library(lib_path)
    device.upload(lib_path)
    lib = device.load_module(lib_name)
    return lib

def get_cpu_op_count(mod):
    """Traverse graph counting ops offloaded to TVM."""

    class Counter(tvm.relay.ExprVisitor):
        def __init__(self):
            super().__init__()
            self.count = 0

        def visit_call(self, call):
            if isinstance(call.op, tvm.ir.Op):
                self.count += 1

            super().visit_call(call)

    c = Counter()
    c.visit(mod["main"])
    return c.count