update aclnn kernekmod gen

This commit is contained in:
VectorSL 2023-11-29 16:05:57 +08:00
parent 397321735f
commit 58139118ea
4 changed files with 130 additions and 24 deletions

2
.gitignore vendored
View File

@ -137,3 +137,5 @@ mindspore/ccsrc/kernel/pyboost/auto_generate/
mindspore/ccsrc/plugin/device/cpu/kernel/pyboost/auto_generate/
mindspore/ccsrc/plugin/device/gpu/kernel/pyboost/auto_generate/
mindspore/ccsrc/plugin/device/ascend/kernel/pyboost/auto_generate/
mindspore/ccsrc/plugin/device/ascend/kernel/opapi/aclnn_kernel_register_auto.cc
mindspore/ccsrc/plugin/device/ascend/kernel/opapi/aclnn_auto_gen/

View File

@ -18,17 +18,20 @@ Generate aclnn kernelmod or call func by input name in ops.yaml
"""
import argparse
import os
import pathlib
import subprocess
import logging
import gen_utils
from pyboost_utils import AclnnUtils
from pyboost_utils import AclnnUtils, get_dtypes
auto_gen = ''
def gen_h(op_name, aclnn_name, op_yaml, kernelmod_h_path, need_update_shape):
"""generate h files"""
kernelmod_name = op_yaml.get('dispatch').get("Ascend")
h_head = f"""
#ifndef MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_{op_name.upper()}_ACLNN_KERNEL_MOD_H_
#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_{op_name.upper()}_ACLNN_KERNEL_MOD_H_
#ifndef MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_{op_name.upper()}_ACLNN{auto_gen.upper()}_KERNEL_MOD_H_
#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_{op_name.upper()}_ACLNN{auto_gen.upper()}_KERNEL_MOD_H_
#include <vector>
#include "ops/base_operator.h"
#include "plugin/device/ascend/kernel/opapi/aclnn_kernel_mod.h"
@ -56,7 +59,7 @@ class {kernelmod_name} : public AclnnKernelMod {{
}} // namespace kernel
}} // namespace mindspore
#endif // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_{op_name.upper()}_ACLNN_KERNEL_MOD_H_
#endif // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_{op_name.upper()}_ACLNN{auto_gen.upper()}_KERNEL_MOD_H_
"""
fd = os.open(kernelmod_h_path, os.O_WRONLY | os.O_CREAT, 0o644)
h_file = os.fdopen(fd, 'w')
@ -67,7 +70,7 @@ def gen_cc(op_name, class_name, op_yaml, kernelmod_cc_path, need_update_shape):
"""generate cc files"""
kernelmod_name = op_yaml.get('dispatch').get("Ascend")
cc_head = f"""
#include "plugin/device/ascend/kernel/opapi/aclnn/{op_name}_aclnn_kernel.h"
#include "plugin/device/ascend/kernel/opapi/aclnn{auto_gen}/{op_name}_aclnn_kernel.h"
#include <algorithm>
#include <vector>
#include <map>
@ -77,30 +80,51 @@ def gen_cc(op_name, class_name, op_yaml, kernelmod_cc_path, need_update_shape):
#include "runtime/stream.h"
#include "runtime/device/kernel_runtime.h"
#include "transform/acl_ir/acl_helper.h"
#include "transform/acl_ir/op_api_convert.h"
#include "abstract/ops/primitive_infer_map.h"
namespace mindspore {{
namespace kernel {{
"""
inputs_num = len(op_yaml.get("args"))
outputs_num = len(op_yaml.get("returns"))
inputs = ""
for i in range(inputs_num):
inputs += "inputs[kIndex" + str(i) + "], "
for i in range(outputs_num):
inputs += "outputs[kIndex" + str(i) + "], "
tuple_tensor_not_supported = f"""
It is not supported for {op_name} with tuple[tensor] inputs when using auto generate.
Please provide a KernelMod name in yaml and using python gen_aclnn_implement.py -n xx manually."""
input_templete = ''
inputs = ''
input_dtypes, output_dtypes, _ = get_dtypes(op_yaml)
for idx, n in enumerate(input_dtypes):
input_name = "inputs[kIndex" + str(idx) + "], "
dtype = input_dtypes.get(n)
if dtype != 'tensor':
input_templete += "auto {} = transform::ConvertKernelTensor<{}>(inputs[kIndex{}]);".format(n, dtype, idx)
input_name = n + ", "
if dtype == 'tuple[tensor]' and auto_gen == "_auto_gen":
raise NotImplementedError(tuple_tensor_not_supported)
inputs += input_name
for idx, n in enumerate(output_dtypes):
output_name = "outputs[kIndex" + str(idx) + "], "
dtype = output_dtypes.get(n)
if dtype != 'tensor':
input_templete += "auto {} = transform::ConvertKernelTensor<{}>(outputs[kIndex{}]);".format(n, dtype, idx)
output_name = n + ", "
if dtype == 'tuple[tensor]' and auto_gen == "_auto_gen":
raise NotImplementedError(tuple_tensor_not_supported)
inputs += output_name
inputs = inputs[:-2]
workspace_info = f"""
void {kernelmod_name}::GetWorkSpaceInfo(const std::vector<KernelTensor *> &inputs,
const std::vector<KernelTensor *> &outputs) {{
const std::vector<KernelTensor *> &outputs) {{
{input_templete}
auto return_value = GEN_EXECUTOR(op_type_, {inputs});
UpdateWorkspace(return_value);
}}
"""
launch = f"""
bool {kernelmod_name}::Launch(const std::vector<KernelTensor *> &inputs, const std::vector<KernelTensor *> &workspace,
const std::vector<KernelTensor *> &outputs, void *stream_ptr) {{
const std::vector<KernelTensor *> &outputs, void *stream_ptr) {{
MS_EXCEPTION_IF_NULL(stream_ptr);
{input_templete}
ParseGenExecutor(GEN_EXECUTOR(op_type_, {inputs}));
RunOp(stream_ptr, workspace);
return true;
@ -127,7 +151,7 @@ MS_ACLLNN_KERNEL_FACTORY_REG({class_name}, {kernelmod_name});
cc_file.write(gen_utils.cc_license_str + cc_head + workspace_info + launch + update_shape + reg)
cc_file.close()
def gen_nnacl_kernelmod(op_name, class_name, op_yaml, h_and_cc, need_update_shape):
def generate(op_name, class_name, op_yaml, h_and_cc, need_update_shape):
"""generate cc and h files"""
kernelmod_h_path = h_and_cc[0]
kernelmod_cc_path = h_and_cc[1]
@ -135,32 +159,71 @@ def gen_nnacl_kernelmod(op_name, class_name, op_yaml, h_and_cc, need_update_shap
gen_h(op_name, aclnn_name, op_yaml, kernelmod_h_path, need_update_shape)
gen_cc(op_name, class_name, op_yaml, kernelmod_cc_path, need_update_shape)
def main(op_name, need_update_shape):
"""main function"""
def gen_aclnn_kernel(op_name, need_update_shape=False, auto=False):
"""gen_aclnn_kernel function"""
if check_op_registed(op_name):
logging.warning("Kernel {%s} is already registered.", op_name)
return
current_path = os.path.dirname(os.path.abspath(__file__))
work_path = os.path.join(current_path, '../../../../')
# get ops yaml
ops_yaml_path = os.path.join(work_path, 'mindspore/python/mindspore/ops_generate/ops.yaml')
aclnn_path = 'mindspore/ccsrc/plugin/device/ascend/kernel/opapi/aclnn/'
kernelmod_cc_path = os.path.join(work_path, aclnn_path + '{}_aclnn_kernel.cc'.format(op_name))
kernelmod_h_path = os.path.join(work_path, aclnn_path + '{}_aclnn_kernel.h'.format(op_name))
yaml_str = gen_utils.safe_load_yaml(ops_yaml_path)
op_yaml = yaml_str.get(op_name)
class_name = ''.join(word.capitalize() for word in op_name.split('_'))
if op_yaml is None:
raise ValueError("Input op {} is not find in ops.yaml.".format(op_name))
dispatch = op_yaml.get("dispatch")
if not dispatch or not dispatch.get("enable"):
raise ValueError("Op {} is not enabled dispatch, please check.".format(op_name))
global auto_gen
if auto:
auto_gen = "_auto_gen"
dispatch['Ascend'] = class_name + "Ascend"
aclnn_path = 'mindspore/ccsrc/plugin/device/ascend/kernel/opapi/aclnn_auto_gen/'
pathlib.Path(os.path.join(work_path, aclnn_path)).mkdir(parents=True, exist_ok=True)
if dispatch.get("Ascend") is None:
raise ValueError("KernelMod {} is auto generated. If need achieve it, "
"please provide the KernelMod name in dispatch.".format(op_name))
class_name = ''.join(word.capitalize() for word in op_name.split('_'))
op_class = op_yaml.get("class")
if op_class is not None and op_class.get("name") is not None:
class_name = op_class.get("name")
kernelmod_cc_path = os.path.join(work_path, aclnn_path + '{}_aclnn_kernel.cc'.format(op_name))
kernelmod_h_path = os.path.join(work_path, aclnn_path + '{}_aclnn_kernel.h'.format(op_name))
h_and_cc = [kernelmod_h_path, kernelmod_cc_path]
gen_nnacl_kernelmod(op_name, class_name, op_yaml, h_and_cc, need_update_shape)
generate(op_name, class_name, op_yaml, h_and_cc, need_update_shape)
def get_registed_ops():
'''get registered ops by search files'''
# search in 'mindspore/ccsrc/plugin/device/ascend/kernel/opapi/'
current_path = os.path.dirname(os.path.abspath(__file__))
work_path = os.path.join(current_path, '../../../../')
search_path = os.path.join(work_path, 'mindspore/ccsrc/plugin/device/ascend/kernel/opapi/')
ret = []
try:
result = subprocess.run(["grep", "-rn", "KERNEL_FACTORY_REG", search_path],
timeout=5, text=True, capture_output=True, check=False)
except OSError:
logging.warning("Something wrong in check op registered.")
return ret
res = result.stdout.split('KERNEL_FACTORY_REG(')
for line in res:
ret.append(line.split(',')[0])
return ret
registed_ops = get_registed_ops()
def check_op_registed(op_name):
'''if op already registered return true'''
global registed_ops
class_name = ''.join(word.capitalize() for word in op_name.split('_'))
return class_name in registed_ops
def main(op_name, need_update_shape):
'''main func'''
gen_aclnn_kernel(op_name, need_update_shape)
parser = argparse.ArgumentParser(description="Generate aclnn KernelMod.")

View File

@ -21,9 +21,10 @@ import shutil
import pathlib
import gen_utils
from gen_utils import py_licence_str, cc_license_str, check_change_and_replace_file, merge_files, safe_load_yaml
from pyboost_utils import get_pyboost_name, is_pyboost_enable, AclnnUtils
from pyboost_utils import get_pyboost_name, is_pyboost_enable, AclnnUtils, get_dtypes
from template import CppTemplate
from gen_pyboost_func import gen_pyboost_code
from gen_aclnn_implement import gen_aclnn_kernel
def get_op_name(operator_name, class_def):
@ -713,6 +714,10 @@ namespace kernel {{
Ascend = dispatch.get("Ascend")
if Ascend is not None: # KernelMod is provided by yaml, don't auto generate it.
continue
_, _, none_tensor_exist = get_dtypes(operator_data)
if none_tensor_exist:
gen_aclnn_kernel(operator_name, auto=True)
continue
class_name = ''.join(word.capitalize() for word in operator_name.split('_'))
op_class = operator_data.get("class")
if op_class and op_class.get("name") is not None:
@ -720,7 +725,7 @@ namespace kernel {{
inputs_outputs_num = len(operator_data.get("args")) + len(operator_data.get("returns"))
aclnn_name = AclnnUtils.get_aclnn_interface(class_name)
reg_code += f"""
MS_ACLLNN_COMMON_KERNEL_FACTORY_REG({class_name}, {aclnn_name}, {inputs_outputs_num})"""
MS_ACLLNN_COMMON_KERNEL_FACTORY_REG({class_name}, {aclnn_name}, {inputs_outputs_num});"""
reg_code += f"""
}} // namespace kernel
}} // namespace mindspore
@ -733,7 +738,7 @@ def generate_aclnn_reg_file(work_path, yaml_str):
Generate nnacl kernelmod register
"""
tmp_register_file = work_path + 'mindspore/ccsrc/plugin/device/ascend/kernel/opapi/tmp_aclnn_kernel_register.cc'
register_file = work_path + 'mindspore/ccsrc/plugin/device/ascend/kernel/opapi/auto_aclnn_kernel_register.cc'
register_file = work_path + 'mindspore/ccsrc/plugin/device/ascend/kernel/opapi/aclnn_kernel_register_auto.cc'
reg_code = generate_aclnn_reg_code(yaml_str)
with open(tmp_register_file, 'w') as reg_file:
reg_file.write(cc_license_str + reg_code)

View File

@ -15,6 +15,7 @@
"""pyboost utils."""
import os
import logging
from gen_utils import safe_load_yaml
@ -268,6 +269,41 @@ def is_pyboost_enable(operator_data):
return False
def convert_types(inputs):
'''convert type to acl type'''
inputs_dtypes = {}
flag = False
for i in inputs:
inputs_dtypes[i] = inputs.get(i).get('dtype')
if inputs_dtypes[i] != 'tensor':
flag = True
if 'tuple' in inputs_dtypes[i]:
data_type = inputs_dtypes[i].split('[')[1].strip(']')
if data_type == 'tensor':
logging.info("Not support tuple[tensor] input.")
elif data_type == 'int':
inputs_dtypes[i] = 'transform::aclIntArray *'
elif data_type == 'float':
inputs_dtypes[i] = 'transform::aclFloatArray *'
elif data_type == 'bool':
inputs_dtypes[i] = 'transform::aclBoolArray *'
else:
logging.warning("Not support tuple[%s]] input.", data_type)
if inputs_dtypes[i] == 'Number':
inputs_dtypes[i] = 'transform::aclScalar *'
return inputs_dtypes, flag
def get_dtypes(op_yaml):
"""get op inputs and outputs dtypes"""
inputs = op_yaml.get('args')
outputs = op_yaml.get('returns')
inputs_dtypes, flag_in = convert_types(inputs)
outputs_dtypes, flag_out = convert_types(outputs)
none_tensor_exist = (flag_in or flag_out)
return inputs_dtypes, outputs_dtypes, none_tensor_exist
class AclnnUtils:
"""
aclnn utils