MindSpore-AKG/python/akg/ms/message.py

# Copyright 2019-2021 Huawei Technologies Co., Ltd
#
# 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.

"""message"""
import importlib.util
import json
import json.decoder as jd
import logging
import traceback
import os
from pathlib import Path

import akg.tvm
from akg.utils import kernel_exec as utils
from akg.utils import validation_check as vc_util
from akg import composite
from akg.tvm import _api_internal
from . import cce
from . import gpu
from . import op_build


def should_use_poly(kernel_info):
    if os.getenv('MS_AKG_USE_POLY') != "off":
        return True
    for desc in kernel_info['op_desc']:
        if desc['name'].startswith('Reduce'):
            return True
    return False

def _pragma_rmselfdep(kernel_info):
    for op in kernel_info["op_desc"]:
        if op['name'] == "MatMul":
            return False
    return True

@vc_util.check_input_type(str)
def compilewithjson_to_func(json_str):
    """compile with json."""
    try:
        kernel_info = json.loads(json_str)
    except jd.JSONDecodeError:
        logging.error(traceback.format_exc())
        return False

    supported_processors = ['cuda', 'aicore']
    processor = 'cuda'
    if 'process' in kernel_info:
        processor = kernel_info['process']
    if processor not in supported_processors:
        logging.error("supported processors: {}, current processor: {}".format(supported_processors, processor))
        return False

    if 'composite' in kernel_info and kernel_info['composite'] is True:
        try:
            if processor == 'cuda':
                use_poly = should_use_poly(kernel_info) 
                enable_atomic_add = composite.should_enable_atomic_add(kernel_info)
                _ = composite._build(json_str, kernel_info, attrs={
                                     "target": "cuda", "enable_akg_reduce_lib": True, "enable_atomic_add": enable_atomic_add}, poly=use_poly)
                return True
            else:
                pragma_rmselfdep = _pragma_rmselfdep(kernel_info)
                mod = composite._build_to_func(json_str, kernel_info, attr = {"pragma_rmselfdep": pragma_rmselfdep})
                return mod
        except Exception:
            logging.error(traceback.format_exc())
            return False

    op_name = kernel_info['name']
    op_func = None
    # get custom ops implementation first.
    if 'impl_path' in kernel_info and kernel_info['impl_path'] is not None:
        impl_path = os.path.realpath(kernel_info['impl_path'])
        if os.path.isfile(impl_path):
            custom_mod_name = Path(impl_path).resolve().stem
            mod_spec = importlib.util.spec_from_file_location(
                custom_mod_name, impl_path)
            custom_mod = importlib.util.module_from_spec(mod_spec)
            mod_spec.loader.exec_module(custom_mod)
            op_func = getattr(custom_mod, op_name, None)

    # get built-in ops.
    if op_func is None:
        if processor == 'cuda':
            op_func = getattr(gpu, op_name, None)
            if op_func is not None:
                input_shapes = []
                input_types = []
                for input_desc in kernel_info['input_desc']:
                    input_shapes.append(input_desc[0]['shape'])
                    input_types.append(input_desc[0]['data_type'])
                op_attrs = []
                if kernel_info['attr']:
                    for ext_arg in kernel_info['attr']:
                        op_attrs.append(ext_arg['value'])
                dump_ir = os.getenv('MS_AKG_DUMP_IR') == "on"
                dump_code = os.getenv('MS_AKG_DUMP_CODE') == "on"
                mod = utils.op_build(op_func, input_shapes, input_types, op_attrs, kernel_info['op'], dump_ir=dump_ir,
                                     dump_code=dump_code)
                return True
        else:
            op_func = getattr(cce, op_name, None)

    if op_func is None:
        logging.error(
            "this op not support by akg, please check op name %s", str(op_name))
        return False

    args = {}
    tsr = []
    for input_desc in kernel_info['input_desc']:
        if len(input_desc) == 1:
            tensor_shape = input_desc[0]['shape']
            tensor_shape = (1,) if not tensor_shape else tensor_shape
            vc_util.shape_dtype_max_size_check(
                tensor_shape, input_desc[0]['data_type'])
            args[input_desc[0]['name']] = akg.tvm.placeholder(
                shape=tensor_shape, name=input_desc[0]['tensor_name'], dtype=input_desc[0]['data_type'])
            tsr.append(args[input_desc[0]['name']])
        else:
            tmp_input = []
            for tmp_desc in input_desc:
                tensor_shape = tmp_desc['shape']
                tensor_shape = (1,) if not tensor_shape else tensor_shape
                vc_util.shape_dtype_max_size_check(
                    tensor_shape, tmp_desc['data_type'])
                tmp_input.append(akg.tvm.placeholder(
                    shape=tensor_shape, name=tmp_desc['tensor_name'], dtype=tmp_desc['data_type']))
            args[input_desc[0]['name']] = tmp_input
            tsr = tsr + tmp_input

    if kernel_info['attr']:
        for ext_arg in kernel_info['attr']:
            args[ext_arg['name']] = ext_arg['value']

    output = op_func(**args)
    schedule_func = None
    attrs = {}
    if isinstance(output, (list, tuple)):
        from inspect import isfunction
        tmp_outputs = []
        for elem in output:
            if isfunction(elem):
                schedule_func = elem
            elif isinstance(elem, dict):
                for key, value in elem.items():
                    if key not in attrs or not attrs[key]:
                        attrs[key] = value
            else:
                tmp_outputs.append(elem)

        output = tmp_outputs
    else:
        output = [output]

    tsr = tsr + [i for i in output if utils.TensorUtils.is_output_value(i)]
    return op_build([op_name], output, tsr, schedule_func, processor, kernel_info['op'], attrs)


def compilewithjson(json_str):
    tmp_rst = compilewithjson_to_func(json_str)
    if isinstance(tmp_rst, bool):
        return tmp_rst
    return _api_internal._BuildToModule(tmp_rst)