wehub-resource-sync
432 lines
15 KiB
Python
432 lines
15 KiB
Python
# Copyright (c) 2022 PaddlePaddle Authors. 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 yaml
|
|
from api_gen import ForwardAPI
|
|
|
|
PREFIX_TENSOR_NAME = 'input_'
|
|
PREFIX_META_TENSOR_NAME = 'meta_'
|
|
|
|
|
|
class StringsAPI(ForwardAPI):
|
|
def __init__(self, api_item_yaml):
|
|
super().__init__(api_item_yaml)
|
|
|
|
def get_api_func_name(self):
|
|
return self.api
|
|
|
|
def gene_api_declaration(self):
|
|
return f"""
|
|
// {", ".join(self.outputs['names'])}
|
|
{super().gene_api_declaration(append_predefined_out=False)}
|
|
"""
|
|
|
|
def get_kernel_tensor_out_type(self, output_name):
|
|
strings_type = 'TensorType::DENSE_TENSOR'
|
|
if output_name.endswith('@StringTensor'):
|
|
strings_type = 'TensorType::STRING_TENSOR'
|
|
return strings_type
|
|
|
|
def get_tensor_type(self, kernel_tensor_out_type):
|
|
tensor_type_dict = {
|
|
"TensorType::DENSE_TENSOR": "phi::DenseTensor",
|
|
"TensorType::STRING_TENSOR": "phi::StringTensor",
|
|
}
|
|
return tensor_type_dict[kernel_tensor_out_type]
|
|
|
|
def gene_output(
|
|
self,
|
|
out_dtype_list,
|
|
out_tensor_type_list=None,
|
|
code_indent='',
|
|
inplace_flag=False,
|
|
):
|
|
kernel_output = []
|
|
output_names = []
|
|
output_create = ""
|
|
return_type = self.get_return_type(inplace_flag)
|
|
|
|
if len(out_dtype_list) == 1:
|
|
kernel_output.append('kernel_out')
|
|
output_names.append('kernel_out')
|
|
kernel_tensor_out_type = self.get_kernel_tensor_out_type(
|
|
self.outputs['names'][0]
|
|
)
|
|
tensor_type = self.get_tensor_type(kernel_tensor_out_type)
|
|
inplace_assign = (
|
|
" = " + self.inplace_map[self.outputs['names'][0]]
|
|
if inplace_flag
|
|
and self.inplace_map is not None
|
|
and self.outputs['names'][0] in self.inplace_map
|
|
else ""
|
|
)
|
|
output_create = f"""
|
|
{return_type} api_output{inplace_assign};
|
|
{tensor_type}* kernel_out = dynamic_cast<{tensor_type}*>(SetStringsKernelOutput(&api_output, {kernel_tensor_out_type}));"""
|
|
|
|
elif len(out_dtype_list) > 1:
|
|
output_create = f"""
|
|
{return_type} api_output;"""
|
|
|
|
for i in range(len(out_dtype_list)):
|
|
kernel_output.append(f'kernel_out_{i}')
|
|
output_names.append(f'kernel_out_{i}')
|
|
kernel_tensor_out_type = self.get_kernel_tensor_out_type(
|
|
self.outputs['names'][i]
|
|
)
|
|
tensor_type = self.get_tensor_type(kernel_tensor_out_type)
|
|
if (
|
|
inplace_flag
|
|
and self.inplace_map is not None
|
|
and self.outputs['names'][i] in self.inplace_map
|
|
):
|
|
output_create = (
|
|
output_create
|
|
+ f"""
|
|
std::get<{i}>(api_output) = {self.inplace_map[self.outputs['names'][i]]};"""
|
|
)
|
|
|
|
output_create = (
|
|
output_create
|
|
+ f"""
|
|
{tensor_type}* kernel_out_{i} = dynamic_cast<{tensor_type}*>(SetStringsKernelOutput(&std::get<{i}>(api_output), {kernel_tensor_out_type}));"""
|
|
)
|
|
|
|
else:
|
|
raise ValueError(
|
|
f"{self.api} : Output error: the output should not be empty."
|
|
)
|
|
|
|
return kernel_output, output_names, output_create
|
|
|
|
def get_kernel_args(self, code_indent):
|
|
input_trans_map = {
|
|
'const Tensor&': 'const phi::StringTensor&',
|
|
'const std::vector<Tensor>&': 'const std::vector<const phi::StringTensor*>&',
|
|
'const paddle::optional<Tensor>&': 'paddle::optional<const phi::StringTensor&>',
|
|
'const paddle::optional<std::vector<Tensor>>&': 'paddle::optional<const std::vector<phi::StringTensor>&>',
|
|
}
|
|
out_trans_map = {
|
|
'Tensor': 'phi::StringTensor*',
|
|
'std::vector<Tensor>': 'std::vector<phi::StringTensor*>&',
|
|
}
|
|
input_names = self.inputs['names']
|
|
input_infos = self.inputs['input_info']
|
|
kernel_args_type_list = ['const phi::DeviceContext&']
|
|
|
|
attr_names = self.attrs['names']
|
|
kernel_param = self.kernel['param']
|
|
if kernel_param is None:
|
|
kernel_param = input_names + attr_names
|
|
input_tensor_code = ""
|
|
# set input_tensor_code
|
|
for i, input_name in enumerate(input_names):
|
|
input_tensor_code = (
|
|
input_tensor_code
|
|
+ f"""
|
|
{code_indent} auto {PREFIX_TENSOR_NAME}{input_name} = TensorToStringTensor({input_name});"""
|
|
)
|
|
|
|
# set kernel_args
|
|
kernel_args = "*dev_ctx, "
|
|
for param in kernel_param:
|
|
if param in input_names:
|
|
if param in self.optional_vars:
|
|
kernel_args = (
|
|
kernel_args + PREFIX_TENSOR_NAME + param + ", "
|
|
)
|
|
else:
|
|
if self.inputs['input_info'][param] == "const Tensor&":
|
|
kernel_args = (
|
|
kernel_args
|
|
+ "*"
|
|
+ PREFIX_TENSOR_NAME
|
|
+ param
|
|
+ ", "
|
|
)
|
|
elif (
|
|
self.inputs['input_info'][input_name]
|
|
== "const std::vector<Tensor>&"
|
|
):
|
|
kernel_args = (
|
|
kernel_args + PREFIX_TENSOR_NAME + param + ", "
|
|
)
|
|
else:
|
|
# do nothing
|
|
pass
|
|
kernel_in_type = input_trans_map[input_infos[param]]
|
|
kernel_args_type_list.append(kernel_in_type)
|
|
elif param in attr_names:
|
|
# set attr for kernel_context
|
|
if 'IntArray' in self.attrs['attr_info'][param][0]:
|
|
kernel_args_type_list.append('const phi::IntArray&')
|
|
param = 'phi::IntArray(' + param + ')'
|
|
elif 'Scalar' in self.attrs['attr_info'][param][0]:
|
|
kernel_args_type_list.append('const phi::Scalar&')
|
|
param = 'phi::Scalar(' + param + ')'
|
|
else:
|
|
kernel_args_type_list.append(
|
|
self.attrs['attr_info'][param][0]
|
|
)
|
|
kernel_args = kernel_args + param + ", "
|
|
elif isinstance(param, bool):
|
|
kernel_args = kernel_args + str(param).lower() + ", "
|
|
else:
|
|
kernel_args = kernel_args + str(param) + ", "
|
|
|
|
for out_type in self.outputs['types']:
|
|
kernel_args_type_list.append(out_trans_map[out_type])
|
|
|
|
# set kernel_signature
|
|
kernel_signature = "void(*)(" + ", ".join(kernel_args_type_list) + ")"
|
|
|
|
return input_tensor_code, kernel_args[:-2], kernel_signature
|
|
|
|
def gen_string_tensor_kernel_code(self, inplace_flag=False, code_indent=""):
|
|
input_tensors, kernel_args, kernel_signature = self.get_kernel_args(
|
|
code_indent
|
|
)
|
|
outputs_args, kernel_output_names, output_create = self.gene_output(
|
|
self.outputs['types'], None, '', inplace_flag
|
|
)
|
|
|
|
return f"""
|
|
// 1. Get kernel signature and kernel
|
|
VLOG(6) << "{self.api} api strings kernel key: [" << kernel_backend << ", " << kernel_layout << ", "<< kernel_data_type << "]";
|
|
auto kernel_result = phi::KernelFactory::Instance().SelectKernelOrThrowError(
|
|
"{self.kernel['func'][0]}", {{kernel_backend, kernel_layout, kernel_data_type}});
|
|
if (FLAGS_low_precision_op_list) {{
|
|
phi::KernelFactory::Instance().AddToLowPrecisionKernelList("{self.api}", kernel_data_type);
|
|
}}
|
|
const auto& kernel = kernel_result.kernel;
|
|
VLOG(6) << "{self.api} api strings kernel: " << kernel;
|
|
|
|
// 2. Get Device Context and input
|
|
auto* dev_ctx = GetDeviceContextByBackend(kernel_result.has_fallback_cpu ? Backend::CPU : kernel_backend);
|
|
{input_tensors}
|
|
|
|
// 3. Set output
|
|
{output_create}
|
|
{self.gene_infer_meta(kernel_output_names, code_indent)}
|
|
|
|
// 4. run kernel
|
|
|
|
{code_indent} using kernel_signature = {kernel_signature};
|
|
{code_indent} auto* kernel_fn = kernel.GetVariadicKernelFn<kernel_signature>();
|
|
{code_indent} (*kernel_fn)({kernel_args}, {", ".join(outputs_args)});
|
|
{code_indent} if (FLAGS_benchmark) {{
|
|
{code_indent} dev_ctx->Wait();
|
|
{code_indent} std::cout << \"{self.api} kernel run finish.\" << std::endl;
|
|
{code_indent} }}
|
|
|
|
{code_indent} {self.gene_return_code()}"""
|
|
|
|
def gene_kernel_select(self) -> str:
|
|
api = self.api
|
|
input_names = self.inputs['names']
|
|
attrs = self.attrs
|
|
kernel = self.kernel
|
|
|
|
kernel_key_item_init = """
|
|
Backend kernel_backend = Backend::UNDEFINED;
|
|
DataLayout kernel_layout = DataLayout::PSTRING_UNION;
|
|
DataType kernel_data_type = DataType::PSTRING;
|
|
"""
|
|
# Check the tensor options
|
|
attr_backend_count = 0
|
|
attr_layout_count = 0
|
|
attr_data_type_count = 0
|
|
for attr_name in attrs['names']:
|
|
if attrs['attr_info'][attr_name][0] == 'Backend':
|
|
assert kernel['backend'] is not None, (
|
|
f"{api} api: When there is a parameter with 'Backend' type in attributes, you must set backend of kernel manually."
|
|
)
|
|
attr_backend_count = attr_backend_count + 1
|
|
|
|
# preprocess kernel configures
|
|
kernel_select_code = ""
|
|
if kernel['backend'] is not None:
|
|
if '>' in kernel['backend']:
|
|
vars_list = kernel['backend'].split('>')
|
|
assert len(vars_list) == 2, (
|
|
f"{api} api: The number of params to set backend with '>' only allows 2, but received {len(vars_list)}."
|
|
)
|
|
assert (vars_list[0].strip() in attrs['names']) and (
|
|
attrs['attr_info'][vars_list[0].strip()][0]
|
|
== 'const Place&'
|
|
), (
|
|
f"{api} api: When use '>' to set kernel backend, the first param should be an attribute with Place type."
|
|
)
|
|
kernel_select_code = (
|
|
kernel_select_code
|
|
+ f"""
|
|
kernel_backend = ParseBackendWithInputOrder({vars_list[0].strip()}, {vars_list[1].strip()});
|
|
"""
|
|
)
|
|
|
|
else:
|
|
args_str = ""
|
|
for ele in kernel['backend'].split(','):
|
|
args_str = args_str + ele.strip() + ', '
|
|
kernel_select_code = (
|
|
kernel_select_code
|
|
+ f"""
|
|
kernel_backend = ParseBackend({args_str[:-2]});
|
|
"""
|
|
)
|
|
|
|
kernel_select_args = ""
|
|
for input_name in input_names:
|
|
kernel_select_args = kernel_select_args + input_name + ", "
|
|
|
|
if len(kernel_select_args) > 2:
|
|
kernel_select_args = kernel_select_args[:-2]
|
|
|
|
kernel_select_code = kernel_key_item_init + kernel_select_code
|
|
|
|
if len(input_names) > 0:
|
|
kernel_select_code = (
|
|
kernel_select_code
|
|
+ f"""
|
|
auto kernel_key_set = ParseKernelKeyByInputArgs({kernel_select_args});
|
|
auto kernel_key = kernel_key_set.GetHighestPriorityKernelKey();
|
|
kernel_backend = kernel_key.backend();"""
|
|
)
|
|
|
|
return kernel_select_code
|
|
|
|
def gene_base_api_code(
|
|
self, inplace_flag=False, grad_flag=False, append_predefined_out=False
|
|
):
|
|
api_func_name = self.get_api_func_name()
|
|
return f"""
|
|
PADDLE_API {self.get_return_type(inplace_flag)} {api_func_name}({self.get_define_args(inplace_flag, grad_flag=grad_flag, append_predefined_out=False)}) {{
|
|
{self.gene_kernel_select()}
|
|
{self.gen_string_tensor_kernel_code(inplace_flag)}
|
|
}}
|
|
"""
|
|
|
|
|
|
def header_include():
|
|
return """
|
|
#include <tuple>
|
|
|
|
#include "paddle/phi/api/include/tensor.h"
|
|
#include "paddle/phi/common/scalar.h"
|
|
#include "paddle/phi/common/int_array.h"
|
|
#include "paddle/utils/optional.h"
|
|
"""
|
|
|
|
|
|
def source_include(header_file_path):
|
|
return f"""
|
|
#include "{header_file_path}"
|
|
|
|
#include "glog/logging.h"
|
|
#include "paddle/common/flags.h"
|
|
|
|
#include "paddle/phi/api/lib/api_gen_utils.h"
|
|
#include "paddle/phi/core/kernel_context.h"
|
|
#include "paddle/phi/core/string_tensor.h"
|
|
#include "paddle/phi/infermeta/strings/nullary.h"
|
|
#include "paddle/phi/infermeta/strings/unary.h"
|
|
#include "paddle/phi/api/lib/kernel_dispatch.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
|
|
COMMON_DECLARE_int32(low_precision_op_list);
|
|
COMMON_DECLARE_bool(benchmark);
|
|
"""
|
|
|
|
|
|
def api_namespace():
|
|
return (
|
|
"""
|
|
namespace paddle {
|
|
namespace experimental {
|
|
namespace strings {
|
|
|
|
""",
|
|
"""
|
|
|
|
} // namespace strings
|
|
} // namespace experimental
|
|
} // namespace paddle
|
|
""",
|
|
)
|
|
|
|
|
|
def generate_api(api_yaml_path, header_file_path, source_file_path):
|
|
with open(api_yaml_path, 'r') as f:
|
|
apis = yaml.load(f, Loader=yaml.FullLoader)
|
|
header_file = open(header_file_path, 'w')
|
|
source_file = open(source_file_path, 'w')
|
|
|
|
namespace = api_namespace()
|
|
|
|
header_file.write("#pragma once\n")
|
|
header_file.write(header_include())
|
|
header_file.write(namespace[0])
|
|
|
|
include_header_file = "paddle/phi/api/include/strings_api.h"
|
|
source_file.write(source_include(include_header_file))
|
|
source_file.write(namespace[0])
|
|
|
|
for api in apis:
|
|
strings_api = StringsAPI(api)
|
|
header_file.write(strings_api.gene_api_declaration())
|
|
source_file.write(strings_api.gene_api_code())
|
|
|
|
header_file.write(namespace[1])
|
|
source_file.write(namespace[1])
|
|
|
|
header_file.close()
|
|
source_file.close()
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description='Generate PaddlePaddle C++ Strings API files'
|
|
)
|
|
parser.add_argument(
|
|
'--api_yaml_path',
|
|
help='path to sparse api yaml file',
|
|
default='paddle/phi/ops/yaml/strings_ops.yaml',
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--api_header_path',
|
|
help='output of generated api header code file',
|
|
default='paddle/phi/api/include/strings_api.h',
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--api_source_path',
|
|
help='output of generated api source code file',
|
|
default='paddle/phi/api/lib/strings_api.cc',
|
|
)
|
|
|
|
options = parser.parse_args()
|
|
|
|
api_yaml_path = options.api_yaml_path
|
|
header_file_path = options.api_header_path
|
|
source_file_path = options.api_source_path
|
|
|
|
generate_api(api_yaml_path, header_file_path, source_file_path)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|