diff --git a/paddle/fluid/distributed/auto_parallel/CMakeLists.txt b/paddle/fluid/distributed/auto_parallel/CMakeLists.txt index 9d9cb97d855549dc098e2a79ac0ceffa1121c787..d4c74fe946c1283c653fc464ec0718f5c9499079 100644 --- a/paddle/fluid/distributed/auto_parallel/CMakeLists.txt +++ b/paddle/fluid/distributed/auto_parallel/CMakeLists.txt @@ -5,3 +5,7 @@ cc_library( phi_enforce) add_subdirectory(test) +add_subdirectory(spmd_rules) + +cc_library(auto_parallel DEPS device_mesh process_mesh dist_attr dist_mapper + dist_tensor_spec) diff --git a/paddle/fluid/distributed/auto_parallel/spmd_rules/CMakeLists.txt b/paddle/fluid/distributed/auto_parallel/spmd_rules/CMakeLists.txt new file mode 100644 index 0000000000000000000000000000000000000000..975a9172ffcde9d18712922fd7e5617e0d333686 --- /dev/null +++ b/paddle/fluid/distributed/auto_parallel/spmd_rules/CMakeLists.txt @@ -0,0 +1,4 @@ +cc_library( + dist_tensor_spec + SRCS dist_tensor_spec.cc + DEPS dist_attr) diff --git a/paddle/fluid/distributed/auto_parallel/spmd_rules/dist_tensor_spec.cc b/paddle/fluid/distributed/auto_parallel/spmd_rules/dist_tensor_spec.cc new file mode 100644 index 0000000000000000000000000000000000000000..5775e72527a7591082b30c8f6ed90d8c99331855 --- /dev/null +++ b/paddle/fluid/distributed/auto_parallel/spmd_rules/dist_tensor_spec.cc @@ -0,0 +1,72 @@ +/* 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. */ + +#include "paddle/fluid/distributed/auto_parallel/spmd_rules/dist_tensor_spec.h" +#include "paddle/fluid/distributed/auto_parallel/process_mesh.h" + +namespace paddle { +namespace distributed { +namespace auto_parallel { + +DistTensorSpec::DistTensorSpec(const std::vector& shape, + const TensorDistAttr& dist_attr) { + shape_.assign(shape.begin(), shape.end()); + // we should merge the new distributed attributes with the original one + // after inferencing, thus we get a copy of the original one + dist_attr_.copy_from(dist_attr); +} + +DistTensorSpec::~DistTensorSpec() {} + +DistTensorSpec::DistTensorSpec(const Tensor& tensor) { + shape_ = tensor.shape(); + + std::vector pm_shape, pm_ids; + pm_shape = {4}; + pm_ids = {0, 1, 2, 3}; + std::vector dim_name = {"mp"}; + + ProcessMesh pm(pm_shape, pm_ids, dim_name); + std::vector dims_mapping = {-1, 0}; + TensorDistAttr dist_attr; + dist_attr.set_process_mesh(pm); + dist_attr.set_dims_mapping(dims_mapping); + + dist_attr_.copy_from(dist_attr); + + std::cout << dist_attr_; +} + +const std::vector& DistTensorSpec::get_dims_mapping() { + return dist_attr_.dims_mapping(); +} + +void DistTensorSpec::set_dims_mapping( + const std::vector& dims_mapping) { + dist_attr_.set_dims_mapping(dims_mapping); +} + +const ProcessMesh& DistTensorSpec::get_process_mesh() { + return dist_attr_.process_mesh(); +} + +void DistTensorSpec::set_process_mesh(const ProcessMesh& process_mesh) { + dist_attr_.set_process_mesh(process_mesh); +} + +const std::vector& DistTensorSpec::get_shape() { return shape_; } + +} // namespace auto_parallel +} // namespace distributed +} // namespace paddle diff --git a/paddle/fluid/distributed/auto_parallel/spmd_rules/dist_tensor_spec.h b/paddle/fluid/distributed/auto_parallel/spmd_rules/dist_tensor_spec.h new file mode 100644 index 0000000000000000000000000000000000000000..2e79148ab0efb6a291c21731c00771dc82cd09e4 --- /dev/null +++ b/paddle/fluid/distributed/auto_parallel/spmd_rules/dist_tensor_spec.h @@ -0,0 +1,58 @@ +/* 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. */ + +#pragma once + +#include "paddle/fluid/distributed/auto_parallel/dist_attr.h" +#include "paddle/phi/api/include/tensor.h" + +namespace paddle { +namespace distributed { +namespace auto_parallel { + +/** + * A unified data class for inferring distributed attributes + * in both dygraph mode and static mode + */ +class DistTensorSpec { + public: + DistTensorSpec(const std::vector& shape, + const TensorDistAttr& dist_attr); + + explicit DistTensorSpec(const Tensor& tensor); + + ~DistTensorSpec(); + + // get dims_mapping from dist_attr_ + const std::vector& get_dims_mapping(); + + // set dims_mapping in dist_attr_ + void set_dims_mapping(const std::vector& dims_mapping); + + // get process_mesh from dist_attr_ + const ProcessMesh& get_process_mesh(); + + // set process_mesh in dist_attr_ + void set_process_mesh(const ProcessMesh& process_mesh); + + const std::vector& get_shape(); + + private: + std::vector shape_; + // distributed attributes of the corresponding tensor + TensorDistAttr dist_attr_; +}; +} // namespace auto_parallel +} // namespace distributed +} // namespace paddle diff --git a/paddle/phi/api/lib/CMakeLists.txt b/paddle/phi/api/lib/CMakeLists.txt index 03b75ee6760d7e0765e48b60bfc6d6e52435daf3..00febea52f19422ba0e5712e3bad39a7689e1ca0 100644 --- a/paddle/phi/api/lib/CMakeLists.txt +++ b/paddle/phi/api/lib/CMakeLists.txt @@ -395,7 +395,8 @@ cc_library( phi_data_transform api_custom_impl api_tensor_utils - phi_profiler) + phi_profiler + dist_tensor_spec) cc_library( phi_bw_function_api SRCS ${bw_api_source_file} ${fused_bw_api_source_file} diff --git a/paddle/phi/api/yaml/generator/api_base.py b/paddle/phi/api/yaml/generator/api_base.py index e4ac5726b9c62535ad3c6f75b09f58df630eaeca..23d3be56a11ca16cd735b7bf4b94407f800c2595 100644 --- a/paddle/phi/api/yaml/generator/api_base.py +++ b/paddle/phi/api/yaml/generator/api_base.py @@ -1278,6 +1278,17 @@ PADDLE_API {self.get_return_type(inplace_flag=True)} {api_func_name}({self.get_d }} """ + def gen_dist_tensor_code(self): + # define the DistTensorSpec vector for input and output tensors + api_code = " \nstd::vector input_specs;\n" + + # get DistTensorSpec for each input tensor + for tensor_name in self.inputs['names']: + api_code += f" input_specs.emplace_back(paddle::distributed::auto_parallel::DistTensorSpec({tensor_name}));\n" + api_code += "\n" + + return api_code + def gene_base_api_code(self, inplace_flag=False): api_func_name = self.get_api_func_name() if inplace_flag and api_func_name[-1] != '_': @@ -1286,6 +1297,8 @@ PADDLE_API {self.get_return_type(inplace_flag=True)} {api_func_name}({self.get_d PADDLE_API {self.get_return_type(inplace_flag)} {api_func_name}({self.get_define_args(inplace_flag)}) {{ {self.gene_kernel_select()} """ + if api_func_name == 'matmul': + api_code += self.gen_dist_tensor_code() if len(self.kernel['func']) > 1: kernel_dispatch_code = '' diff --git a/paddle/phi/api/yaml/generator/api_gen.py b/paddle/phi/api/yaml/generator/api_gen.py index 71285de7b24723d14ef2f96a0731c3b9f08940b1..7c7109859e06761eecc0acd30abfc54559983fef 100644 --- a/paddle/phi/api/yaml/generator/api_gen.py +++ b/paddle/phi/api/yaml/generator/api_gen.py @@ -379,6 +379,8 @@ def source_include(header_file_path): #include "paddle/phi/api/profiler/event_tracing.h" #include "paddle/phi/api/profiler/supplement_tracing.h" +#include "paddle/fluid/distributed/auto_parallel/spmd_rules/dist_tensor_spec.h" + DECLARE_bool(conv2d_disable_cudnn); DECLARE_int32(low_precision_op_list); """ diff --git a/python/paddle/distributed/auto_parallel/utils.py b/python/paddle/distributed/auto_parallel/utils.py index f4dfb8d9c20d52942082bb75a786d6d3931beb44..1a3299e20a48a9b116baed255c079c60d4726a83 100644 --- a/python/paddle/distributed/auto_parallel/utils.py +++ b/python/paddle/distributed/auto_parallel/utils.py @@ -2355,3 +2355,52 @@ def is_dep_skip_op(op): return True return False + + +# def wrap_data_for_completion( +# dist_op: DistributedOperator, +# input_names: list, +# output_names: list, +# attr_names: list +# ): +# """ +# Get data used in inferring distributed attributes, including: +# 1. DistTensorSpec for each input and output tensor of this dist_op. +# 2. Operator attributes of this dist_op, e.g. transpose_x in matmul op. +# +# Args: +# dist_op: the DistributedOperator +# input_names: list, name of the dist_op's input tensors +# output_names: list, name of the dist_op's output tensors +# attr_names: list, attribute name of the dist_op's corresponding serial op +# +# Returns: +# input_specs: list, DistTensorSpec for each input tensor of the dist_op +# output_specs: list, DistTensorSpec for each output tensor of the dist_op +# attrs: dict, attribute map of the dist op +# """ +# +# input_specs = [] +# output_specs = [] +# attrs = {} +# +# serial_op = dist_op.serial_op +# +# # Construct each input tensor's DistTensorSpec with shape and dist_attr +# for name in input_names: +# tensor_dist_attr = dist_op.dist_attr.get_input_dist_attr(name) +# var = serial_op.block._var_recursive(name) +# tensor_shape = var.shape +# dist_spec = DistTensorSpec(tensor_shape, tensor_dist_attr) +# input_specs.append(dist_spec) +# +# # Construct each output tensor's DistTensorSpec with shape and dist_attr +# for name in output_names: +# tensor_dist_attr = dist_op.dist_attr.get_output_dist_attr(name) +# var = serial_op.block._var_recursive(name) +# tensor_shape = var.shape +# dist_spec = DistTensorSpec(tensor_shape, tensor_dist_attr) +# output_specs.append(dist_spec) +# +# for attr_name in attr_names: +# attrs[attr_name] = serial_op.desc.attr(attr_name)