From c1545a4b3dadcc6c1bd221fb21fa7103d7510428 Mon Sep 17 00:00:00 2001 From: liangjianzhong Date: Mon, 29 May 2023 19:59:59 +0800 Subject: [PATCH] compile 1 --- .../auto_parallel/spmd_rules/CMakeLists.txt | 10 +++ .../spmd_rules/matmul_spmd_rule.cc | 78 +++++++++++-------- .../spmd_rules/matmul_spmd_rule.h | 6 +- 3 files changed, 60 insertions(+), 34 deletions(-) diff --git a/paddle/fluid/distributed/auto_parallel/spmd_rules/CMakeLists.txt b/paddle/fluid/distributed/auto_parallel/spmd_rules/CMakeLists.txt index 975a9172ffc..e56f753c54f 100644 --- a/paddle/fluid/distributed/auto_parallel/spmd_rules/CMakeLists.txt +++ b/paddle/fluid/distributed/auto_parallel/spmd_rules/CMakeLists.txt @@ -2,3 +2,13 @@ cc_library( dist_tensor_spec SRCS dist_tensor_spec.cc DEPS dist_attr) + +cc_library( + spmd_rule_base + SRCS common.cc + DEPS dist_tensor_spec) + +cc_library( + matmul_spmd_rule + SRCS matmul_spmd_rule.cc + DEPS spmd_rule_base) diff --git a/paddle/fluid/distributed/auto_parallel/spmd_rules/matmul_spmd_rule.cc b/paddle/fluid/distributed/auto_parallel/spmd_rules/matmul_spmd_rule.cc index 899548659df..d819a8b7ae3 100644 --- a/paddle/fluid/distributed/auto_parallel/spmd_rules/matmul_spmd_rule.cc +++ b/paddle/fluid/distributed/auto_parallel/spmd_rules/matmul_spmd_rule.cc @@ -13,11 +13,12 @@ See the License for the specific language governing permissions and limitations under the License. */ #include "paddle/fluid/distributed/auto_parallel/spmd_rules/matmul_spmd_rule.h" +#include "paddle/phi/core/distributed/auto_parallel/utils.h" namespace paddle { namespace distributed { namespace auto_parallel { - +using phi::distributed::auto_parallel::str_join; std::vector MatmulSPMDRule::InferForward( const std::vector& input_specs, const paddle::framework::AttributeMap& attrs) { @@ -29,26 +30,31 @@ std::vector MatmulSPMDRule::InferForward( phi::errors::InvalidArgument( "The size of InputSpec of matmul should be 2, but got [%d].", input_specs_size)); - - int x_ndim = input_specs[0].shape.size(); - int y_ndim = input_specs[1].shape.size(); - std::vector x_dims_mapping = input_specs[0].DistAttr.dims_mapping; - std::vector y_dims_mapping = input_specs[0].DistAttr.dims_mapping; + auto x_shape = input_specs[0].get_shape(); + auto y_shape = input_specs[1].get_shape(); + int x_ndim = x_shape.size(); + int y_ndim = y_shape.size(); + auto x_dist_attr_src = input_specs[0].get_dist_attr(); + auto y_dist_attr_src = input_specs[1].get_dist_attr(); + std::vector x_dims_mapping = x_dist_attr_src.dims_mapping(); + std::vector y_dims_mapping = y_dist_attr_src.dims_mapping(); PADDLE_ENFORCE_EQ( x_ndim, - x_dims_mapping.size() phi::errors::InvalidArgument( + x_dims_mapping.size(), + phi::errors::InvalidArgument( "Mismatch of X's tensor size: [%d] and X's dims_mapping size [%d].", x_ndim, x_dims_mapping.size())); PADDLE_ENFORCE_EQ( y_ndim, - y_dims_mapping.size() phi::errors::InvalidArgument( + y_dims_mapping.size(), + phi::errors::InvalidArgument( "Mismatch of Y's tensor size: [%d] and Y's dims_mapping size [%d].", x_ndim, x_dims_mapping.size())); - bool trans_x = ExtractAttr("trans_x"); - bool trans_y = ExtractAttr("trans_y"); + bool trans_x = ExtractAttr("trans_x", attrs); + bool trans_y = ExtractAttr("trans_y", attrs); // step1: build Einsum Notation @@ -129,19 +135,19 @@ std::vector MatmulSPMDRule::InferForward( // step2.2: Infer Output's Dims Mapping. TensorDistAttr output_dist_attr_dst = - CopyTensorDistAttrForOutput(input_specs[0].DistAttr); + CopyTensorDistAttrForOutput(x_dist_attr_src); std::vector out_dims_mapping; out_dims_mapping.reserve(out_axes.size()); - for (int i = 0; i < out_axes.size(); ++i) { + for (size_t i = 0; i < out_axes.size(); ++i) { out_dims_mapping.push_back(axis_to_dim_map[out_axes.substr(i, 1)]); } output_dist_attr_dst.set_dims_mapping(out_dims_mapping); // step2.3: Merge and get Inputs' New Dims Mapping. - TensorDistAttr x_dist_attr_dst = GetInferedDistAttr( - input_specs[0].DistAttr, input_specs[0].shape, x_axes, axis_to_dim_map); - TensorDistAttr y_dist_attr_dst = GetInferedDistAttr( - input_specs[1].DistAttr, input_specs[1].shape, y_axes, axis_to_dim_map); + TensorDistAttr x_dist_attr_dst = + GetInferedDistAttr(x_dist_attr_src, x_shape, x_axes, axis_to_dim_map); + TensorDistAttr y_dist_attr_dst = + GetInferedDistAttr(y_dist_attr_src, y_shape, y_axes, axis_to_dim_map); // step2.3: Handle Partial // Step2.3.1 Output Partial @@ -149,18 +155,17 @@ std::vector MatmulSPMDRule::InferForward( ResoluteOutputPartialDimension(axis_to_dim_map, out_axes); // Step2.3.2 handle input tensor partial (TODO) - VLOG(4) << "MatmulSPMDRule InferForward: " - << "X shape: " << input_specs[0].shape - << ", src_dims_mapping: " << x_dims_mapping - << ", dst_dims_mapping: " << x_dist_attr_dst.dims_mapping - << "; Y shape: " << input_specs[1].shape - << ", src_dims_mapping: " << x_dims_mapping - << ", dst_dims_mapping: " << y_dist_attr_dst.dims_mapping - << "; Output dims_mapping: " << out_dims_mapping - << ", partial_on_dims: " << partial_on_dims; - - return { x_dist_attr_dst, y_dist_attr_dst, output_dist_attr_dst } + << "X shape: [" << str_join(x_shape) << "], src_dims_mapping: [" + << str_join(x_dims_mapping) << "], dst_dims_mapping: [" + << str_join(x_dist_attr_dst.dims_mapping()) << "]; Y shape: [" + << str_join(y_shape) << "], src_dims_mapping: [" + << str_join(x_dims_mapping) << "], dst_dims_mapping: [" + << str_join(y_dist_attr_dst.dims_mapping()) + << "]; Output dims_mapping: [" << str_join(out_dims_mapping) + << "], partial_on_dims: [" << str_join(partial_on_dims) << "]"; + + return {x_dist_attr_dst, y_dist_attr_dst, output_dist_attr_dst}; } TensorDistAttr GetInferedDistAttr( @@ -170,13 +175,19 @@ TensorDistAttr GetInferedDistAttr( const std::unordered_map& axis_to_dim_map) { TensorDistAttr dist_attr_ = CopyTensorDistAttrForOutput(origin_dist_attr); std::vector infered_dims_mapping; - infered_dims_mapping.reserve(tensor_string.size()); + infered_dims_mapping.reserve(tensor_axis.size()); - for (int i = 0; i < tensor_axis.size(); ++i) { + for (size_t i = 0; i < tensor_axis.size(); ++i) { if (shape.size() > i && shape[i] == 1) { infered_dims_mapping.push_back(-1); } else { - infered_dims_mapping.push_back(axis_to_dim_map[tensor_axis.substr(i, 1)]); + auto itr = axis_to_dim_map.find(tensor_axis.substr(i, 1)); + if (itr == axis_to_dim_map.end()) { + phi::errors::InvalidArgument( + "Tensor axis [%s] of not in axis_to_dim_map.", + tensor_axis.substr(i, 1)); + } + infered_dims_mapping.push_back(itr->second); } } @@ -186,7 +197,12 @@ TensorDistAttr GetInferedDistAttr( std::vector MatmulSPMDRule::InferBackward( const std::vector& output_specs, - const paddle::framework::AttributeMap& attrs) {} + const paddle::framework::AttributeMap& attrs) { + PADDLE_THROW(phi::errors::Unimplemented( + "InferBackward of MatmulSPMDRule is NOT implemented yet.")); + + return {}; +} } // namespace auto_parallel } // namespace distributed diff --git a/paddle/fluid/distributed/auto_parallel/spmd_rules/matmul_spmd_rule.h b/paddle/fluid/distributed/auto_parallel/spmd_rules/matmul_spmd_rule.h index 2afee21c322..8a452624b2a 100644 --- a/paddle/fluid/distributed/auto_parallel/spmd_rules/matmul_spmd_rule.h +++ b/paddle/fluid/distributed/auto_parallel/spmd_rules/matmul_spmd_rule.h @@ -33,14 +33,14 @@ TensorDistAttr GetInferedDistAttr( class MatmulSPMDRule : public SPMDRuleBase { public: - std::vector InferForward( + std::vector InferForward( const std::vector& input_specs, const paddle::framework::AttributeMap& attrs) override; - std::vector InferBackward( + std::vector InferBackward( const std::vector& output_specs, const paddle::framework::AttributeMap& attrs) override; -} +}; } // namespace auto_parallel } // namespace distributed } // namespace paddle -- GitLab