wehub-resource-sync
691 lines
25 KiB
Python
691 lines
25 KiB
Python
# Copyright (c) 2023 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 dist_api_gen
|
|
import yaml
|
|
from backward_api_gen import BackwardAPI
|
|
from dist_api_gen import DistForwardAPI
|
|
|
|
######################
|
|
# Code Gen Templates #
|
|
######################
|
|
|
|
MAIN_DIST_BRANCH_TEMPLATE = """
|
|
// Auto Parallel condition
|
|
if (run_auto_parallel) {{
|
|
// 1. InferSpmd (Infer DistAttr of Inputs&Outputs){}
|
|
// 2. Create Temporary Output & Prepare Dist and Dense Output{}
|
|
// 3. Infer DistTensor's Global Shape{}\n
|
|
// 4. Set Output Dist Attr For Default Impl{}\n
|
|
if (rank_is_in_current_mesh) {{
|
|
// 5. Select Kernel{}
|
|
// 6. Reshard Input{}\n
|
|
// 7. PrepareData (DataTransform & Prepare Dense Input){}
|
|
// 8. RecordOpInfoSupplement{}
|
|
// 9. Infer Local DenseTensor Meta{}
|
|
// 10. DenseTensor Kernel Call{}
|
|
// 11. Fallback{}
|
|
}}
|
|
// 12. Reshard Kernel Output to API output{}\n
|
|
// 13. Return
|
|
{}
|
|
}}
|
|
"""
|
|
|
|
# 1. Create API Outputs
|
|
SINGLE_OUT_CREATION_TEMPLATE_NO_SPMD = """
|
|
auto dist_out = SetKernelDistOutput({});
|
|
auto dense_out = dist_out->unsafe_mutable_value();
|
|
"""
|
|
SINGLE_OUT_CREATION_TEMPLATE_WITH_SPMD = """
|
|
std::shared_ptr<phi::distributed::DistTensor> shared_dist_out =
|
|
CreateKernelDistOutput({}, !rank_is_in_current_mesh, spmd_info.second[0]);
|
|
phi::distributed::DistTensor* dist_out = shared_dist_out.get();
|
|
phi::DenseTensor* dense_out = nullptr;
|
|
if (dist_out) {{
|
|
dense_out = dist_out->unsafe_mutable_value();
|
|
if (dense_out && !rank_is_in_current_mesh && !dist_out->defined()) {{
|
|
*dense_out = phi::DenseTensor(
|
|
std::make_shared<phi::Allocation>(nullptr, 0, phi::distributed::GetDefaultPlace()),
|
|
phi::DenseTensorMeta());
|
|
}}
|
|
}}
|
|
"""
|
|
SINGLE_OUT_CREATION_TEMPLATE = """
|
|
std::shared_ptr<phi::distributed::DistTensor> shared_dist_out =
|
|
CreateKernelDistOutput({}, !rank_is_in_current_mesh);
|
|
phi::distributed::DistTensor* dist_out = shared_dist_out.get();
|
|
phi::DenseTensor* dense_out = nullptr;
|
|
if (dist_out) {{
|
|
dense_out = dist_out->unsafe_mutable_value();
|
|
if (dense_out && !rank_is_in_current_mesh && !dist_out->defined()) {{
|
|
*dense_out = phi::DenseTensor(
|
|
std::make_shared<phi::Allocation>(nullptr, 0, phi::distributed::GetDefaultPlace()),
|
|
phi::DenseTensorMeta());
|
|
}}
|
|
}}
|
|
"""
|
|
VECTOR_OUT_CREATION_TEMPLATE_WITH_NO_SPMD = """
|
|
auto dist_out = SetKernelDistOutput({name});
|
|
std::vector<phi::DenseTensor*> dense_out(dist_out.size(), nullptr);
|
|
for (size_t i=0; i<dist_out.size(); i++) {{
|
|
if (dist_out[i]) {{
|
|
dense_out[i] = dist_out[i]->unsafe_mutable_value();
|
|
if (dense_out[i] && !rank_is_in_current_mesh && !dist_out[i]->defined()) {{
|
|
*dense_out[i] = phi::DenseTensor(
|
|
std::make_shared<phi::Allocation>(nullptr, 0, phi::distributed::GetDefaultPlace()),
|
|
phi::DenseTensorMeta());
|
|
}}
|
|
}}
|
|
}}
|
|
"""
|
|
|
|
VECTOR_OUT_CREATION_TEMPLATE_WITH_SPMD = """
|
|
auto shared_dist_out = CreateKernelDistOutput({name}, !rank_is_in_current_mesh, spmd_info.second[0]);
|
|
std::vector<phi::distributed::DistTensor*> dist_out;
|
|
for(auto& e: shared_dist_out){{
|
|
dist_out.push_back(e.get());
|
|
}}
|
|
std::vector<phi::DenseTensor*> dense_out(dist_out.size(), nullptr);
|
|
for (size_t i=0; i<dist_out.size(); i++) {{
|
|
if (dist_out[i]) {{
|
|
dense_out[i] = dist_out[i]->unsafe_mutable_value();
|
|
if (dense_out[i] && !rank_is_in_current_mesh && !dist_out[i]->defined()) {{
|
|
*dense_out[i] = phi::DenseTensor(
|
|
std::make_shared<phi::Allocation>(nullptr, 0, phi::distributed::GetDefaultPlace()),
|
|
phi::DenseTensorMeta());
|
|
}}
|
|
}}
|
|
}}
|
|
"""
|
|
|
|
|
|
VECTOR_OUT_CREATION_TEMPLATE = """
|
|
auto shared_dist_out = CreateKernelDistOutput({name}, !rank_is_in_current_mesh);
|
|
std::vector<phi::distributed::DistTensor*> dist_out;
|
|
for(auto& e: shared_dist_out){{
|
|
dist_out.push_back(e.get());
|
|
}}
|
|
std::vector<phi::DenseTensor*> dense_out(dist_out.size(), nullptr);
|
|
for (size_t i=0; i<dist_out.size(); i++) {{
|
|
if (dist_out[i]) {{
|
|
dense_out[i] = dist_out[i]->unsafe_mutable_value();
|
|
if (dense_out[i] && !rank_is_in_current_mesh && !dist_out[i]->defined()) {{
|
|
*dense_out[i] = phi::DenseTensor(
|
|
std::make_shared<phi::Allocation>(nullptr, 0, phi::distributed::GetDefaultPlace()),
|
|
phi::DenseTensorMeta());
|
|
}}
|
|
}}
|
|
}}
|
|
"""
|
|
|
|
|
|
INPLACE_OUT_CREATION_TEMPLATE = """
|
|
*{} = {};
|
|
"""
|
|
MULTI_SINGLE_OUT_CREATION_TEMPLATE_NO_SPMD = """
|
|
auto dist_out_{idx} = SetKernelDistOutput({name});
|
|
auto dense_out_{idx} = dist_out_{idx} ? dist_out_{idx}->unsafe_mutable_value() : nullptr;
|
|
if (dense_out_{idx} && !rank_is_in_current_mesh && !dist_out_{idx}->defined()) {{
|
|
*dense_out_{idx} = phi::DenseTensor(
|
|
std::make_shared<phi::Allocation>(nullptr, 0, phi::distributed::GetDefaultPlace()),
|
|
phi::DenseTensorMeta());
|
|
}}
|
|
"""
|
|
MULTI_SINGLE_OUT_CREATION_TEMPLATE_WITH_SPMD = """
|
|
std::shared_ptr<phi::distributed::DistTensor> shared_dist_out_{idx} =
|
|
CreateKernelDistOutput({name}, !rank_is_in_current_mesh, spmd_info.second[{idx}]);
|
|
phi::distributed::DistTensor* dist_out_{idx} = shared_dist_out_{idx}.get();
|
|
phi::DenseTensor* dense_out_{idx} = dist_out_{idx} ? dist_out_{idx}->unsafe_mutable_value() : nullptr;
|
|
if (dense_out_{idx} && !rank_is_in_current_mesh && !dist_out_{idx}->defined()) {{
|
|
*dense_out_{idx} = phi::DenseTensor(
|
|
std::make_shared<phi::Allocation>(nullptr, 0, phi::distributed::GetDefaultPlace()),
|
|
phi::DenseTensorMeta());
|
|
}}
|
|
"""
|
|
MULTI_SINGLE_OUT_CREATION_TEMPLATE = """
|
|
std::shared_ptr<phi::distributed::DistTensor> shared_dist_out_{idx} =
|
|
CreateKernelDistOutput({name}, !rank_is_in_current_mesh);
|
|
phi::distributed::DistTensor* dist_out_{idx} = shared_dist_out_{idx}.get();
|
|
phi::DenseTensor* dense_out_{idx} = dist_out_{idx} ? dist_out_{idx}->unsafe_mutable_value() : nullptr;
|
|
if (dense_out_{idx} && !rank_is_in_current_mesh && !dist_out_{idx}->defined()) {{
|
|
*dense_out_{idx} = phi::DenseTensor(
|
|
std::make_shared<phi::Allocation>(nullptr, 0, phi::distributed::GetDefaultPlace()),
|
|
phi::DenseTensorMeta());
|
|
}}
|
|
"""
|
|
MULTI_VECTOR_OUT_CREATION_TEMPLATE = """
|
|
auto dist_out_{i} = SetKernelDistOutput({name});
|
|
std::vector<phi::DenseTensor*> dense_out_{i}(dist_out_{i}.size(), nullptr);
|
|
for (size_t i = 0; i < dist_out_{i}.size(); i++) {{
|
|
if (dist_out_{i}[i]) {{
|
|
dense_out_{i}[i] = const_cast<phi::DenseTensor*>(&dist_out_{i}[i]->value());
|
|
if (dense_out_{i}[i] && !rank_is_in_current_mesh && !dist_out_{i}[i]->defined()) {{
|
|
*dense_out_{i}[i]= phi::DenseTensor(
|
|
std::make_shared<phi::Allocation>(nullptr, 0, phi::distributed::GetDefaultPlace()),
|
|
phi::DenseTensorMeta());
|
|
}}
|
|
}}
|
|
}}
|
|
"""
|
|
|
|
# 9. Reshard Output
|
|
RESHARD_SINGLE_OUTPUT_TEMPLATE = """
|
|
ReshardKernelOutputToApiOutput(dev_ctx, shared_dist_out, {}, "{}");"""
|
|
|
|
RESHARD_MULTI_SINGLE_OUTPUT_TEMPLATE = """
|
|
ReshardKernelOutputToApiOutput(dev_ctx, shared_dist_out_{}, {}, "{}");"""
|
|
|
|
RESHARD_VECTOR_OUTPUT_TEMPLATE = """
|
|
ReshardKernelOutputToApiOutput(dev_ctx, shared_dist_out, {}, "{}");"""
|
|
|
|
NONEED_TO_RESHARD_OUTPUT_TEMPLATE = """
|
|
// API `{}` does not need to reshard output."""
|
|
|
|
SET_LOCAL_SHAPE_TEMPLATE = """
|
|
{meta_tensor}.set_dims(phi::make_ddim(local_shape));"""
|
|
|
|
|
|
class DistBackwardAPI(DistForwardAPI, BackwardAPI):
|
|
def __init__(self, backward_item_yaml):
|
|
BackwardAPI.__init__(self, backward_item_yaml)
|
|
self.forward_config = backward_item_yaml['forward']
|
|
self.init_dist_api_members()
|
|
|
|
# override DistForwardAPI's method
|
|
def generate_output_creation_code(self) -> str:
|
|
# backward api only need to generate kernel outputs
|
|
output_num = len(self.outputs['types'])
|
|
output_creation_code = ""
|
|
output_creation_code += "\n phi::DeviceContext* dev_ctx = nullptr;"
|
|
if output_num == 1:
|
|
self.dist_output_args.append('dist_out')
|
|
self.dense_output_args.append('dense_out')
|
|
if self.outputs['types'][0] == 'Tensor':
|
|
if self.infer_meta['spmd_rule'] is not None:
|
|
output_creation_code += (
|
|
SINGLE_OUT_CREATION_TEMPLATE_WITH_SPMD.format(
|
|
self.outputs['names'][0]
|
|
)
|
|
)
|
|
elif self.generate_general_infer_spmd is True:
|
|
output_creation_code += SINGLE_OUT_CREATION_TEMPLATE.format(
|
|
self.outputs['names'][0]
|
|
)
|
|
else:
|
|
output_creation_code += (
|
|
SINGLE_OUT_CREATION_TEMPLATE_NO_SPMD.format(
|
|
self.outputs['names'][0]
|
|
)
|
|
)
|
|
elif self.outputs['types'][0] == 'std::vector<Tensor>':
|
|
if self.infer_meta['spmd_rule'] is not None:
|
|
output_creation_code += (
|
|
VECTOR_OUT_CREATION_TEMPLATE_WITH_SPMD.format(
|
|
name=self.outputs['names'][0]
|
|
)
|
|
)
|
|
elif self.generate_general_infer_spmd is True:
|
|
output_creation_code += VECTOR_OUT_CREATION_TEMPLATE.format(
|
|
name=self.outputs['names'][0]
|
|
)
|
|
else:
|
|
output_creation_code += (
|
|
VECTOR_OUT_CREATION_TEMPLATE_WITH_NO_SPMD.format(
|
|
name=self.outputs['names'][0]
|
|
)
|
|
)
|
|
else:
|
|
self.vector_output_size_assertion_check()
|
|
elif output_num > 1:
|
|
for i, out_type in enumerate(self.outputs['types']):
|
|
self.dist_output_args.append(f'dist_out_{i}')
|
|
self.dense_output_args.append(f'dense_out_{i}')
|
|
if out_type == 'Tensor':
|
|
if self.infer_meta['spmd_rule'] is not None:
|
|
output_creation_code += (
|
|
MULTI_SINGLE_OUT_CREATION_TEMPLATE_WITH_SPMD.format(
|
|
name=self.outputs['names'][i], idx=i
|
|
)
|
|
)
|
|
elif self.generate_general_infer_spmd is True:
|
|
output_creation_code += (
|
|
MULTI_SINGLE_OUT_CREATION_TEMPLATE.format(
|
|
name=self.outputs['names'][i], idx=i
|
|
)
|
|
)
|
|
else:
|
|
output_creation_code += (
|
|
MULTI_SINGLE_OUT_CREATION_TEMPLATE_NO_SPMD.format(
|
|
name=self.outputs['names'][i], idx=i
|
|
)
|
|
)
|
|
elif out_type == 'std::vector<Tensor>':
|
|
output_creation_code += (
|
|
MULTI_VECTOR_OUT_CREATION_TEMPLATE.format(
|
|
i=i, name=self.outputs['names'][i]
|
|
)
|
|
)
|
|
else:
|
|
self.vector_output_size_assertion_check()
|
|
else:
|
|
raise ValueError(
|
|
f"{self.api} : Output error: the output should not be empty."
|
|
)
|
|
|
|
return output_creation_code
|
|
|
|
def generate_bw_infer_local_shape_code(self, need_kernel=False):
|
|
arg_name = self.infer_meta['local_shape']
|
|
assert arg_name in self.outputs['names'], (
|
|
f"Auto Parallel will calculate local_shape for {arg_name} "
|
|
f"in {self.api}, but {arg_name} is not found in its outputs."
|
|
)
|
|
|
|
_, fw_inputs, fw_attrs, fw_outputs = self.parse_forward_config(
|
|
self.forward_config
|
|
)
|
|
# shape_type = self.attrs['attr_info'][shape_name][0]
|
|
# out_name = self.dist_output_args[0]
|
|
dist_out_name = self.dist_output_args[
|
|
self.outputs['names'].index(arg_name)
|
|
]
|
|
shape_type = self.get_shape_type(fw_attrs['attr_info'])
|
|
return_code = dist_api_gen.CALCULATE_LOCAL_SHAPE_TEMPLATE.format(
|
|
out_name=dist_out_name,
|
|
out_dist_attr=(
|
|
"PADDLE_GET_CONST(phi::distributed::TensorDistAttr, spmd_info.second[0]);"
|
|
if self.infer_meta['spmd_rule']
|
|
else f"phi::distributed::TensorDistAttr(common::vectorize({dist_out_name}->dims()))"
|
|
),
|
|
dtype=shape_type,
|
|
op_name=self.kernel['func'][0],
|
|
)
|
|
if need_kernel:
|
|
return (
|
|
dist_api_gen.CALCULATE_LOCAL_SHAPE_KERNEL_TEMPLATE.format(
|
|
out_grad_dist_attr=(
|
|
"PADDLE_GET_CONST(phi::distributed::TensorDistAttr, spmd_info.first[1]);"
|
|
if self.infer_meta['spmd_rule']
|
|
else "phi::distributed::TensorDistAttr(common::vectorize(out_grad.dims()))"
|
|
),
|
|
dtype=shape_type,
|
|
op_name=self.kernel['func'][0],
|
|
)
|
|
+ return_code
|
|
)
|
|
return return_code
|
|
|
|
def generate_infer_meta_code(self) -> str:
|
|
(
|
|
infer_meta_func_code,
|
|
input_args_code,
|
|
output_decl_code,
|
|
output_args_code,
|
|
) = self.generate_infer_meta_func_and_args_code()
|
|
|
|
infer_meta_code = ""
|
|
|
|
if self.infer_meta['global_shape'] is not None:
|
|
for i, out_name in enumerate(self.outputs['names']):
|
|
if out_name == self.infer_meta[
|
|
'global_shape'
|
|
] and self.need_to_generate_code_for_inplace_impl(i):
|
|
infer_meta_code += dist_api_gen.SET_DIMS_TEMPLATE.format(
|
|
dst=self.dist_output_args[i],
|
|
src=(
|
|
self.dist_output_args[i] + '_tmp'
|
|
if i > 0
|
|
else self.dist_output_args[i]
|
|
),
|
|
)
|
|
|
|
infer_meta_code = (
|
|
infer_meta_code
|
|
+ dist_api_gen.INFER_META_TEMPLATE.format(
|
|
infer_meta_func_code, input_args_code, output_args_code
|
|
)
|
|
)
|
|
# TODO(GhostScreaming): kernel like reshape need calculate local_shape
|
|
if self.infer_meta['local_shape'] is not None:
|
|
if (
|
|
self.kernel['param'] is not None
|
|
and self.infer_meta['local_shape'] not in self.kernel['param']
|
|
):
|
|
infer_meta_code += self.generate_bw_infer_local_shape_code()
|
|
else:
|
|
infer_meta_code += self.generate_bw_infer_local_shape_code(
|
|
need_kernel=True
|
|
)
|
|
infer_meta_code += SET_LOCAL_SHAPE_TEMPLATE.format(
|
|
meta_tensor="meta_" + self.dense_output_args[0]
|
|
)
|
|
|
|
return output_decl_code + infer_meta_code
|
|
|
|
# override DistForwardAPI's method
|
|
def generate_return_code(self) -> str:
|
|
return "return;"
|
|
|
|
# override BaseAPI's method
|
|
def get_api_func_name(self):
|
|
return self.api
|
|
|
|
# override BaseAPI's method
|
|
# The method lookup order are: (DistBackwardAPI.__mro__)
|
|
# <class '__main__.DistBackwardAPI'>,
|
|
# <class 'dist_api_gen.DistForwardAPI'>,
|
|
# <class 'api_gen.ForwardAPI'>,
|
|
# <class 'backward_api_gen.BackwardAPI'>,
|
|
# <class 'api_base.BaseAPI'>,
|
|
# <class 'object'>
|
|
# if don't override it, the ForwardAPI's gene_output will be called
|
|
def gene_output(
|
|
self,
|
|
out_dtype_list,
|
|
out_tensor_type_list=None,
|
|
code_indent='',
|
|
inplace_flag=False,
|
|
):
|
|
return BackwardAPI.gene_output(
|
|
self,
|
|
out_dtype_list,
|
|
out_tensor_type_list,
|
|
code_indent,
|
|
inplace_flag,
|
|
)
|
|
|
|
# override BaseAPI's method
|
|
def get_return_type(self, inplace_flag=False):
|
|
return BackwardAPI.get_return_type(self)
|
|
|
|
# override BaseAPI's method
|
|
def gene_return_code(self):
|
|
return ""
|
|
|
|
# override BaseAPI's method
|
|
def gene_api_declaration(
|
|
self, grad_flag=False, append_predefined_out=False
|
|
) -> str:
|
|
return BackwardAPI.gene_api_declaration(
|
|
self, grad_flag=grad_flag, append_predefined_out=not grad_flag
|
|
)
|
|
|
|
def generate_reshard_output_code(self):
|
|
reshard_output_code = ""
|
|
if self.generate_infer_spmd is True:
|
|
output_num = len(self.outputs['types'])
|
|
if output_num == 1:
|
|
if self.outputs['types'][0] == 'Tensor':
|
|
reshard_output_code += (
|
|
RESHARD_SINGLE_OUTPUT_TEMPLATE.format(
|
|
self.outputs['names'][0], self.outputs['names'][0]
|
|
)
|
|
)
|
|
elif self.outputs['types'][0] == 'std::vector<Tensor>':
|
|
reshard_output_code += (
|
|
RESHARD_VECTOR_OUTPUT_TEMPLATE.format(
|
|
self.outputs['names'][0], self.outputs['names'][0]
|
|
)
|
|
)
|
|
else:
|
|
self.vector_output_size_assertion_check()
|
|
elif output_num > 1:
|
|
for i, out_type in enumerate(self.outputs['types']):
|
|
if out_type == 'Tensor':
|
|
reshard_output_code += (
|
|
RESHARD_MULTI_SINGLE_OUTPUT_TEMPLATE.format(
|
|
i,
|
|
self.outputs['names'][i],
|
|
self.outputs['names'][i],
|
|
)
|
|
)
|
|
else:
|
|
self.vector_output_size_assertion_check()
|
|
else:
|
|
raise ValueError(
|
|
f"{self.api} : Output error: the output should not be empty."
|
|
)
|
|
else:
|
|
reshard_output_code += NONEED_TO_RESHARD_OUTPUT_TEMPLATE.format(
|
|
self.kernel['func'][0]
|
|
)
|
|
# do nothing
|
|
pass
|
|
|
|
return reshard_output_code
|
|
|
|
def generate_auto_parallel_branch(self) -> str:
|
|
# if no tensor input, do not generate auto parallel branch
|
|
if len(self.inputs['names']) == 0:
|
|
return ""
|
|
infer_spmd_code = self.generate_infer_spmd_code()
|
|
output_creation_code = self.generate_output_creation_code()
|
|
infer_global_shape_code = self.generate_infer_global_shape_code()
|
|
output_dist_attr_setting = self.generate_output_dist_attr_setting()
|
|
kernel_selection_code = self.generate_kernel_selection_code()
|
|
reshard_input_code = self.generate_reshard_input_code()
|
|
(
|
|
prepare_data_code,
|
|
input_name_tensor_map,
|
|
) = self.generate_prepare_data_code()
|
|
record_op_info_supplement_code = (
|
|
self.generate_record_op_info_supplement(
|
|
input_name_tensor_map, ' ', True
|
|
)
|
|
)
|
|
infer_meta_code = self.generate_infer_meta_code()
|
|
kernel_call_code = self.generate_kernel_call_code(is_forward=False)
|
|
fallback_code = self.generate_fallback_code()
|
|
reshard_output_code = self.generate_reshard_output_code()
|
|
return_code = self.generate_return_code()
|
|
|
|
return MAIN_DIST_BRANCH_TEMPLATE.format(
|
|
infer_spmd_code,
|
|
output_creation_code,
|
|
infer_global_shape_code,
|
|
output_dist_attr_setting,
|
|
kernel_selection_code,
|
|
reshard_input_code,
|
|
prepare_data_code,
|
|
record_op_info_supplement_code,
|
|
infer_meta_code,
|
|
kernel_call_code,
|
|
fallback_code,
|
|
reshard_output_code,
|
|
return_code,
|
|
)
|
|
|
|
|
|
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, fw_header_file_path):
|
|
return f"""
|
|
#include "{header_file_path}"
|
|
#include <memory>
|
|
|
|
#include "glog/logging.h"
|
|
#include "paddle/common/flags.h"
|
|
|
|
#include "paddle/phi/api/lib/api_custom_impl.h"
|
|
#include "paddle/phi/api/lib/api_gen_utils.h"
|
|
#include "paddle/phi/api/lib/data_transform.h"
|
|
#include "paddle/phi/api/lib/kernel_dispatch.h"
|
|
#include "paddle/phi/common/type_traits.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "{fw_header_file_path}"
|
|
#include "paddle/phi/infermeta/backward.h"
|
|
#include "paddle/phi/infermeta/unary.h"
|
|
#include "paddle/phi/infermeta/fusion.h"
|
|
|
|
#include "paddle/phi/api/profiler/event_tracing.h"
|
|
#include "paddle/phi/api/profiler/supplement_tracing.h"
|
|
|
|
#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
|
|
#include "paddle/phi/core/distributed/comm_context_manager.h"
|
|
#include "paddle/phi/core/distributed/nccl_comm_context.h"
|
|
#elif defined(PADDLE_WITH_XPU_BKCL)
|
|
#include "paddle/phi/core/distributed/comm_context_manager.h"
|
|
#include "paddle/phi/core/distributed/bkcl_comm_context.h"
|
|
#elif defined(PADDLE_WITH_CUSTOM_DEVICE)
|
|
#include "paddle/phi/core/distributed/comm_context_manager.h"
|
|
#include "paddle/phi/core/distributed/xccl_comm_context.h"
|
|
#endif
|
|
|
|
#ifdef PADDLE_WITH_DISTRIBUTE
|
|
#include "paddle/phi/core/distributed/store/store_utils.h"
|
|
#include "paddle/phi/infermeta/spmd_rules/rules.h"
|
|
#include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_utils.h"
|
|
#endif
|
|
|
|
PD_DECLARE_bool(conv2d_disable_cudnn);
|
|
COMMON_DECLARE_int32(low_precision_op_list);
|
|
COMMON_DECLARE_bool(benchmark);
|
|
"""
|
|
|
|
|
|
def backward_api_namespace():
|
|
return (
|
|
"""
|
|
namespace paddle {
|
|
namespace experimental {
|
|
|
|
""",
|
|
"""
|
|
|
|
} // namespace experimental
|
|
} // namespace paddle
|
|
""",
|
|
)
|
|
|
|
|
|
def generate_backward_api(
|
|
backward_yaml_path,
|
|
is_fused_backward_yaml,
|
|
header_file_path,
|
|
source_file_path,
|
|
):
|
|
bw_apis = []
|
|
for each_api_yaml in backward_yaml_path:
|
|
with open(each_api_yaml, 'r') as f:
|
|
api_list = yaml.load(f, Loader=yaml.FullLoader)
|
|
if api_list:
|
|
bw_apis.extend(api_list)
|
|
|
|
header_file = open(header_file_path, 'w')
|
|
source_file = open(source_file_path, 'w')
|
|
|
|
namespace = backward_api_namespace()
|
|
|
|
header_file.write("#pragma once\n")
|
|
header_file.write(header_include())
|
|
header_file.write(namespace[0])
|
|
|
|
include_header_file = (
|
|
"paddle/phi/api/backward/fused_backward_api_base.h"
|
|
if is_fused_backward_yaml
|
|
else "paddle/phi/api/backward/backward_api_base.h"
|
|
)
|
|
include_fw_header_file = (
|
|
"paddle/phi/api/include/fused_api.h"
|
|
if is_fused_backward_yaml
|
|
else "paddle/phi/api/include/api.h"
|
|
)
|
|
source_file.write(
|
|
source_include(include_header_file, include_fw_header_file)
|
|
)
|
|
source_file.write(namespace[0])
|
|
# not all fused ops support dygraph
|
|
if is_fused_backward_yaml is True:
|
|
new_bw_apis = [
|
|
bw_api
|
|
for bw_api in bw_apis
|
|
if "support_dygraph_mode" in bw_api
|
|
and bw_api["support_dygraph_mode"] is True
|
|
]
|
|
bw_apis = new_bw_apis
|
|
|
|
for bw_api in bw_apis:
|
|
dist_bw_api = DistBackwardAPI(bw_api)
|
|
header_file.write(dist_bw_api.gene_api_declaration())
|
|
if is_fused_backward_yaml is True:
|
|
source_file.write(dist_bw_api.gene_api_code())
|
|
else:
|
|
source_file.write(dist_bw_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++ backward API files'
|
|
)
|
|
parser.add_argument(
|
|
'--backward_yaml_path',
|
|
help='path to backward yaml file',
|
|
nargs='+',
|
|
default=['paddle/phi/ops/yaml/backward.yaml'],
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--is_fused_backward_yaml',
|
|
help='flag of fused backward yaml',
|
|
action='store_true',
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--backward_header_path',
|
|
help='output of generated backward header code file',
|
|
default='paddle/phi/api/backward/backward_api_base.h',
|
|
)
|
|
|
|
parser.add_argument(
|
|
'--backward_source_path',
|
|
help='output of generated backward source code file',
|
|
default='paddle/phi/api/lib/backward_api_base.cc',
|
|
)
|
|
|
|
options = parser.parse_args()
|
|
|
|
backward_yaml_path = options.backward_yaml_path
|
|
is_fused_backward_yaml = options.is_fused_backward_yaml
|
|
header_file_path = options.backward_header_path
|
|
source_file_path = options.backward_source_path
|
|
|
|
generate_backward_api(
|
|
backward_yaml_path,
|
|
is_fused_backward_yaml,
|
|
header_file_path,
|
|
source_file_path,
|
|
)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|