361 lines
12 KiB
Python
361 lines
12 KiB
Python
#!/usr/bin/env python3
|
|
# -----------------------------------------------------------------------------------------------------------
|
|
# Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
|
# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
|
# CANN Open Software License Agreement Version 2.0 (the "License").
|
|
# Please refer to the License for details. You may not use this file except in compliance with the License.
|
|
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
|
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
|
# See LICENSE in the root of the software repository for the full text of the License.
|
|
# -----------------------------------------------------------------------------------------------------------
|
|
|
|
import argparse
|
|
import glob
|
|
import json
|
|
import os
|
|
import sys
|
|
from typing import NamedTuple
|
|
|
|
import const_var
|
|
|
|
|
|
class OpConfig(NamedTuple):
|
|
op_type: str
|
|
support_info: dict
|
|
core_type: str
|
|
task_ration: str
|
|
obj_file: str
|
|
|
|
|
|
def load_json(json_file: str):
|
|
with open(json_file, encoding="utf-8") as file:
|
|
json_content = json.load(file)
|
|
return json_content
|
|
|
|
|
|
def get_specified_suffix_file(root_dir, suffix):
|
|
specified_suffix = os.path.join(root_dir, "**/*{}".format(suffix))
|
|
all_suffix_files = glob.glob(specified_suffix, recursive=True)
|
|
return sorted(all_suffix_files)
|
|
|
|
|
|
def add_dict_key(dict_to_add, key, value):
|
|
if value is None:
|
|
return
|
|
dict_to_add[key] = value
|
|
|
|
|
|
def correct_format_mode(format_mode):
|
|
if format_mode == "FormatDefault":
|
|
return "nd_agnostic"
|
|
if format_mode == "FormatAgnostic":
|
|
return "static_nd_agnostic"
|
|
if format_mode == "FormatFixed":
|
|
return "normal"
|
|
return format_mode
|
|
|
|
|
|
def get_input_or_output_config(in_or_out):
|
|
param_dict = {}
|
|
name = in_or_out.get("name")
|
|
index = in_or_out.get("index")
|
|
param_type = in_or_out.get("paramType")
|
|
|
|
format_match_mode = in_or_out.get("format_match_mode")
|
|
format_mode = correct_format_mode(format_match_mode)
|
|
|
|
dtype_mode = in_or_out.get("dtype_match_mode")
|
|
if dtype_mode == "DtypeByte":
|
|
dtype_mode = "bit"
|
|
|
|
add_dict_key(param_dict, "name", name)
|
|
add_dict_key(param_dict, "index", index)
|
|
add_dict_key(param_dict, "paramType", param_type)
|
|
add_dict_key(param_dict, "dtypeMode", dtype_mode)
|
|
add_dict_key(param_dict, "formatMode", format_mode)
|
|
return param_dict
|
|
|
|
|
|
def get_inputs_or_outputs_config(inputs_or_outputs):
|
|
if inputs_or_outputs is None:
|
|
return None
|
|
inputs_or_outputs_list = []
|
|
|
|
for in_or_out in inputs_or_outputs:
|
|
if isinstance(in_or_out, dict):
|
|
dict_param_config = get_input_or_output_config(in_or_out)
|
|
inputs_or_outputs_list.append(dict_param_config)
|
|
elif isinstance(in_or_out, list):
|
|
param_info = in_or_out[0]
|
|
list_param_config = get_input_or_output_config(param_info)
|
|
tmp_list = [list_param_config]
|
|
inputs_or_outputs_list.append(tmp_list)
|
|
return inputs_or_outputs_list
|
|
|
|
|
|
def gen_attrs_config(attrs):
|
|
attrs_list = []
|
|
for attr in attrs:
|
|
attrs_dict = {}
|
|
name = attr.get("name")
|
|
mode = attr.get("mode")
|
|
add_dict_key(attrs_dict, "name", name)
|
|
add_dict_key(attrs_dict, "mode", mode)
|
|
attrs_list.append(attrs_dict)
|
|
return attrs_list
|
|
|
|
|
|
def get_params_config(support_info):
|
|
params_dict = {}
|
|
|
|
inputs = support_info.get("inputs")
|
|
inputs_list = get_inputs_or_outputs_config(inputs)
|
|
params_dict["inputs"] = inputs_list
|
|
|
|
outputs = support_info.get("outputs")
|
|
outputs_list = get_inputs_or_outputs_config(outputs)
|
|
params_dict["outputs"] = outputs_list
|
|
|
|
attrs = support_info.get("attrs")
|
|
if attrs is not None:
|
|
attrs_list = gen_attrs_config(attrs)
|
|
params_dict["attrs"] = attrs_list
|
|
|
|
return params_dict
|
|
|
|
|
|
def add_simplified_config(op_info, binary_info_config, config):
|
|
simplified_key = op_info.support_info.get("simplifiedKey")
|
|
|
|
json_path = op_info.obj_file.split(".")[0] + ".json"
|
|
|
|
simple_cfg = config.get(binary_info_config)
|
|
op_cfg = simple_cfg.get(op_info.op_type)
|
|
if not op_cfg:
|
|
op_cfg = {"dynamicRankSupport": True}
|
|
|
|
simplified_key_mode = op_info.support_info.get("simplifiedKeyMode")
|
|
add_dict_key(op_cfg, "simplifiedKeyMode", simplified_key_mode)
|
|
|
|
optional_input_mode = op_info.support_info.get("optionalInputMode")
|
|
optional_output_mode = op_info.support_info.get("optionalOutputMode")
|
|
add_dict_key(op_cfg, "optionalInputMode", optional_input_mode)
|
|
if optional_output_mode is not None:
|
|
add_dict_key(op_cfg, "optionalOutputMode", optional_output_mode)
|
|
|
|
params_info = get_params_config(op_info.support_info)
|
|
op_cfg["params"] = params_info
|
|
op_cfg["binaryList"] = []
|
|
simple_cfg[op_info.op_type] = op_cfg
|
|
|
|
bin_list = op_cfg.get("binaryList")
|
|
if op_info.core_type == 0 and op_info.task_ration == "tilingKey":
|
|
bin_list.append(
|
|
{
|
|
"coreType": op_info.core_type,
|
|
"simplifiedKey": simplified_key,
|
|
"multiKernelType": 1,
|
|
"binPath": op_info.obj_file,
|
|
"jsonPath": json_path,
|
|
}
|
|
)
|
|
else:
|
|
bin_list.append(
|
|
{
|
|
"coreType": op_info.core_type,
|
|
"simplifiedKey": simplified_key,
|
|
"binPath": op_info.obj_file,
|
|
"jsonPath": json_path,
|
|
}
|
|
)
|
|
|
|
|
|
def add_op_config(op_file, bin_info, config):
|
|
op_cfg = config.get(op_file)
|
|
if not op_cfg:
|
|
op_cfg = {"binList": []}
|
|
config[op_file] = op_cfg
|
|
op_cfg.get("binList").append(bin_info)
|
|
|
|
|
|
def gen_ops_config(json_file, soc, binary_info_config, config):
|
|
core_type_map = {"MIX": 0, "AiCore": 1, "VectorCore": 2, "MIX_AICORE": 3, "MIX_VECTOR_CORE": 4, "MIX_AIV": 4}
|
|
contents = load_json(json_file)
|
|
if ("binFileName" not in contents) or ("supportInfo" not in contents):
|
|
return
|
|
json_base_name = os.path.basename(json_file)
|
|
op_dir = os.path.basename(os.path.dirname(json_file))
|
|
|
|
support_info = contents.get("supportInfo")
|
|
bin_name = contents.get("binFileName")
|
|
bin_suffix = contents.get("binFileSuffix")
|
|
core_type = contents.get("coreType")
|
|
task_ration = contents.get("taskRation")
|
|
core_type = core_type_map.get(core_type, -1)
|
|
if core_type == -1 and soc != "ascend310b":
|
|
raise Exception("[ERROR]: must set coreType in json when soc version is {soc}.")
|
|
|
|
bin_file_name = bin_name + bin_suffix
|
|
op_type = bin_name.split("_")[0]
|
|
op_file = op_dir + ".json"
|
|
bin_info = {}
|
|
|
|
add_dict_key(bin_info, "implMode", support_info.get("implMode"))
|
|
add_dict_key(bin_info, "int64Mode", support_info.get("int64Mode"))
|
|
add_dict_key(bin_info, "simplifiedKeyMode", support_info.get("simplifiedKeyMode"))
|
|
|
|
simplified_key = support_info.get("simplifiedKey")
|
|
if simplified_key is not None:
|
|
bin_info["simplifiedKey"] = simplified_key
|
|
obj_file = os.path.join(soc, op_dir, bin_file_name)
|
|
op_info = OpConfig(
|
|
op_type=op_type,
|
|
support_info=support_info,
|
|
core_type=core_type,
|
|
task_ration=task_ration,
|
|
obj_file=obj_file,
|
|
)
|
|
add_simplified_config(op_info, binary_info_config, config)
|
|
|
|
add_dict_key(bin_info, "dynamicParamMode", support_info.get("dynamicParamMode"))
|
|
bin_info["staticKey"] = support_info.get("staticKey")
|
|
bin_info["inputs"] = support_info.get("inputs")
|
|
bin_info["outputs"] = support_info.get("outputs")
|
|
if support_info.get("attrs"):
|
|
bin_info["attrs"] = support_info.get("attrs")
|
|
|
|
add_dict_key(bin_info, "opMode", support_info.get("opMode"))
|
|
add_dict_key(bin_info, "optionalInputMode", support_info.get("optionalInputMode"))
|
|
add_dict_key(bin_info, "deterministic", support_info.get("deterministic"))
|
|
if support_info.get("optionalOutputMode") is not None:
|
|
add_dict_key(bin_info, "optionalOutputMode", support_info.get("optionalOutputMode"))
|
|
|
|
bin_info["binInfo"] = {"jsonFilePath": os.path.join(soc, op_dir, json_base_name)}
|
|
add_op_config(op_file, bin_info, config)
|
|
|
|
|
|
def check_single_op_is_void(root_dir):
|
|
for root, dirs, _ in os.walk(root_dir):
|
|
for sub_dir in dirs:
|
|
dir_path = os.path.join(root, sub_dir)
|
|
if len(os.listdir(dir_path)) == 0:
|
|
print(f"[ERROR] op {sub_dir}: not any obj compile success")
|
|
sys.exit(1)
|
|
|
|
|
|
def write_jsons(out_dir, file_list, config):
|
|
for json_name in file_list:
|
|
json_file = os.path.join(out_dir, json_name)
|
|
with os.fdopen(os.open(json_file, const_var.WFLAGS, const_var.WMODES), "w") as fd:
|
|
json.dump(config.get(json_name), fd, indent=" ")
|
|
|
|
|
|
def generate_operator_cfg_file(json_files, binary_info_config, soc, out_dir, gen_json_status):
|
|
if not json_files:
|
|
return
|
|
|
|
if gen_json_status == "not_generated":
|
|
return
|
|
|
|
json_files.sort()
|
|
config = {binary_info_config: {}}
|
|
for _json in json_files:
|
|
gen_ops_config(_json, soc, binary_info_config, config)
|
|
|
|
if gen_json_status == "single_json":
|
|
file_list = [json_file for json_file in config if json_file != binary_info_config]
|
|
elif gen_json_status == "summary_json":
|
|
file_list = [binary_info_config]
|
|
else:
|
|
file_list = config.keys()
|
|
|
|
write_jsons(out_dir, file_list, config)
|
|
|
|
|
|
def gen_all_config(root_dir, soc, out_dir, skip_binary_info_config, op_range="all"):
|
|
if op_range != "relocatable":
|
|
check_single_op_is_void(root_dir)
|
|
all_json_files = get_specified_suffix_file(root_dir, ".json")
|
|
relocatable_json_files = get_specified_suffix_file(root_dir, "_relocatable.json")
|
|
normal_json_files = list(set(all_json_files) - set(relocatable_json_files))
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
|
|
if op_range != "relocatable":
|
|
for _json in all_json_files:
|
|
file_path = soc + _json.split(soc, maxsplit=1)[1]
|
|
with open(_json, "r+") as f:
|
|
data = json.load(f)
|
|
data["filePath"] = file_path
|
|
f.seek(0)
|
|
json.dump(data, f, indent=" ")
|
|
f.truncate()
|
|
|
|
if skip_binary_info_config:
|
|
gen_normale_json = "single_json"
|
|
gen_relocatable_json = "not_generated"
|
|
else:
|
|
gen_normale_json = "all_json"
|
|
gen_relocatable_json = "summary_json"
|
|
|
|
# normal kernel
|
|
if op_range == "all" or op_range == "normal":
|
|
binary_info_config = "binary_info_config.json"
|
|
generate_operator_cfg_file(normal_json_files, binary_info_config, soc, out_dir, gen_normale_json)
|
|
|
|
# relocatable kernel
|
|
if op_range == "all" or op_range == "relocatable":
|
|
binary_info_config = "relocatable_kernel_info_config.json"
|
|
generate_operator_cfg_file(relocatable_json_files, binary_info_config, soc, out_dir, gen_relocatable_json)
|
|
|
|
|
|
# Parse multiple soc_versions ops in single path.
|
|
def gen_all_soc_config(all_path):
|
|
soc_roots = glob.glob(os.path.join(all_path, "ascend*"))
|
|
|
|
for soc_root in soc_roots:
|
|
soc = os.path.basename(soc_root)
|
|
gen_all_config(soc_root, soc, soc_root, True)
|
|
cfg_files = glob.glob(os.path.join(soc_root, "*.json"))
|
|
cfg_path = os.path.join(all_path, "config", soc)
|
|
os.makedirs(cfg_path, exist_ok=True)
|
|
for cfg_file in cfg_files:
|
|
new_file = os.path.join(cfg_path, os.path.basename(cfg_file))
|
|
os.rename(cfg_file, new_file)
|
|
|
|
|
|
def args_prase():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("-p", "--path", nargs="?", required=True, help="Parse the path of the json file.")
|
|
|
|
parser.add_argument("-s", "--soc", nargs="?", required=True, help="Parse the soc_version of ops.")
|
|
|
|
parser.add_argument("-o", "--out", nargs="?", help="Output directory.")
|
|
|
|
parser.add_argument(
|
|
"--skip-binary-info-config", action="store_true", help="binary_info_config.json file is not parsed."
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--op-range",
|
|
type=str,
|
|
choices=["all", "normal", "relocatable"],
|
|
default="all",
|
|
help="all operators/normal operators/relocatable operators.",
|
|
)
|
|
|
|
return parser.parse_args()
|
|
|
|
|
|
def main():
|
|
args = args_prase()
|
|
if args.out is None:
|
|
out_dir = args.path
|
|
else:
|
|
out_dir = args.out
|
|
|
|
gen_all_config(args.path, args.soc, out_dir, args.skip_binary_info_config, args.op_range)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|