wehub-resource-sync
248 lines
9.5 KiB
C++
248 lines
9.5 KiB
C++
// 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/phi/infermeta/spmd_rules/squeeze.h"
|
|
#include <algorithm>
|
|
#include <numeric>
|
|
|
|
#include "glog/logging.h"
|
|
|
|
#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
|
|
#include "paddle/phi/core/distributed/auto_parallel/inferspmd_utils.h"
|
|
#include "paddle/phi/core/distributed/auto_parallel/utils.h"
|
|
#include "paddle/phi/infermeta/spmd_rules/dim_trans.h"
|
|
#include "paddle/phi/infermeta/spmd_rules/reshape.h"
|
|
#include "paddle/phi/infermeta/spmd_rules/utils.h"
|
|
|
|
namespace phi::distributed {
|
|
|
|
void MakeSqueezeDimTransWithoutAxis(
|
|
const std::vector<int64_t>& x_shape,
|
|
std::vector<int64_t>* out_shape,
|
|
std::vector<std::shared_ptr<DimTrans>>* trans) {
|
|
for (int64_t i = 0, n = static_cast<int64_t>(x_shape.size()); i < n; i++) {
|
|
if (x_shape[i] != 1) {
|
|
trans->emplace_back(std::make_shared<InputDim>(i));
|
|
out_shape->emplace_back(x_shape[i]);
|
|
}
|
|
}
|
|
}
|
|
|
|
void MakeSqueezeDimTransWithAxis(
|
|
const std::vector<int64_t>& x_shape,
|
|
std::vector<int64_t>* out_shape,
|
|
const std::vector<int64_t>& axis,
|
|
std::vector<std::shared_ptr<DimTrans>>* trans) {
|
|
for (int64_t i = 0, n = static_cast<int64_t>(x_shape.size()); i < n; i++) {
|
|
if (x_shape[i] == 1) {
|
|
auto it = find(axis.begin(), axis.end(), i);
|
|
if (it == axis.end()) {
|
|
trans->emplace_back(std::make_shared<Singleton>());
|
|
out_shape->emplace_back(1);
|
|
}
|
|
} else {
|
|
trans->emplace_back(std::make_shared<InputDim>(i));
|
|
out_shape->emplace_back(x_shape[i]);
|
|
}
|
|
}
|
|
}
|
|
|
|
void MakeSqueezeDimTransReverseWithoutAxis(
|
|
const std::vector<int64_t>& x_shape,
|
|
std::vector<std::shared_ptr<DimTrans>>* trans) {
|
|
for (int64_t i = 0, j = 0, n = static_cast<int64_t>(x_shape.size()); i < n;
|
|
i++) {
|
|
if (x_shape[i] != 1) {
|
|
trans->emplace_back(std::make_shared<InputDim>(j++));
|
|
} else {
|
|
trans->emplace_back(std::make_shared<Singleton>());
|
|
}
|
|
}
|
|
}
|
|
|
|
void MakeSqueezeDimTransReverseWithAxis(
|
|
const std::vector<int64_t>& x_shape,
|
|
const std::vector<int64_t>& out_shape,
|
|
const std::vector<int64_t>& axis,
|
|
std::vector<std::shared_ptr<DimTrans>>* trans) {
|
|
for (int64_t i = 0, j = 0, n = static_cast<int64_t>(x_shape.size()); i < n;
|
|
i++) {
|
|
if (x_shape[i] == 1) {
|
|
trans->emplace_back(std::make_shared<Singleton>());
|
|
|
|
auto it = find(axis.begin(), axis.end(), i);
|
|
if (it == axis.end()) {
|
|
j++;
|
|
}
|
|
} else {
|
|
trans->emplace_back(std::make_shared<InputDim>(j++));
|
|
}
|
|
}
|
|
}
|
|
|
|
SpmdInfo SqueezeInferSpmd(const DistMetaTensor& x,
|
|
const std::vector<int64_t>& axis) {
|
|
// Step0: Verify input args based on squeeze logic
|
|
auto x_shape = vectorize(x.dims());
|
|
int x_ndim = x_shape.size();
|
|
auto x_dist_attr_src = x.dist_attr();
|
|
std::vector<int64_t> x_dims_mapping = x_dist_attr_src.dims_mapping();
|
|
|
|
PADDLE_ENFORCE_EQ(
|
|
x_ndim,
|
|
x_dims_mapping.size(),
|
|
common::errors::InvalidArgument("The Tensor X's rank [%d] and X's "
|
|
"dims_mapping size [%d] are not matched.",
|
|
x_ndim,
|
|
x_dims_mapping.size()));
|
|
|
|
// Step1: Build the transformation from
|
|
// the original shape to the target shape
|
|
|
|
std::vector<std::shared_ptr<DimTrans>> trans;
|
|
std::vector<int64_t> out_shape;
|
|
|
|
if (static_cast<int64_t>(axis.size()) == 0) {
|
|
MakeSqueezeDimTransWithoutAxis(x_shape, &out_shape, &trans);
|
|
} else {
|
|
std::vector<int64_t> axis_copy(axis);
|
|
for (auto& v : axis_copy) {
|
|
if (v < 0) {
|
|
v += x_ndim;
|
|
}
|
|
}
|
|
MakeSqueezeDimTransWithAxis(x_shape, &out_shape, axis_copy, &trans);
|
|
}
|
|
|
|
// Step2: Infer the dims mapping of input (if reshard is
|
|
// needed) and output from the dimension transformation.
|
|
const auto& dims_mapping_vec = InferFromDimTrans(x, trans);
|
|
const auto& input_dims_mapping = std::get<0>(dims_mapping_vec);
|
|
const auto& output_dims_mapping = std::get<1>(dims_mapping_vec);
|
|
|
|
// Step3: Update the dist attributes of input
|
|
// and output with the inferred dims mapping.
|
|
TensorDistAttr x_dist_attr_dst = CopyTensorDistAttrForOutput(x_dist_attr_src);
|
|
x_dist_attr_dst.set_dims_mapping(input_dims_mapping);
|
|
if (x_dist_attr_dst.dynamic_dims().size() !=
|
|
x_dist_attr_dst.dims_mapping().size()) {
|
|
VLOG(4) << "SqueezeInferSPMD change x dist attr dynamic dims";
|
|
x_dist_attr_dst.set_default_dynamic_dims(x_dist_attr_dst.dims_mapping());
|
|
}
|
|
TensorDistAttr out_dist_attr = CopyTensorDistAttrForOutput(x_dist_attr_src);
|
|
out_dist_attr.set_dims_mapping(output_dims_mapping);
|
|
if (out_dist_attr.dynamic_dims().size() !=
|
|
out_dist_attr.dims_mapping().size()) {
|
|
VLOG(4) << "SqueezeInferSPMD change output dist attr dynamic dims";
|
|
out_dist_attr.set_default_dynamic_dims(out_dist_attr.dims_mapping());
|
|
}
|
|
VLOG(4) << "SqueezeInferSpmd: X shape: [" << str_join(x_shape)
|
|
<< "] Out shape: [" << str_join(out_shape) << "]";
|
|
VLOG(4) << "Transformation from input to output:";
|
|
for (int64_t i = 0, n = static_cast<int64_t>(trans.size()); i < n; i++) {
|
|
VLOG(4) << "\tOut axis[" << i << "]: " << trans[i]->to_string();
|
|
}
|
|
VLOG(4) << "X dims_mapping_src: [" << str_join(x_dims_mapping)
|
|
<< "] dims_mapping_dst: [" << str_join(input_dims_mapping)
|
|
<< "]\n Out dims_mapping: [" << str_join(output_dims_mapping)
|
|
<< "]\n\n";
|
|
|
|
return {{x_dist_attr_dst}, {out_dist_attr}};
|
|
}
|
|
|
|
SpmdInfo SqueezeInferSpmdReverse(const DistMetaTensor& x,
|
|
const DistMetaTensor& out,
|
|
const std::vector<int64_t>& axis) {
|
|
// Step0: Verify input args based on squeeze logic
|
|
auto x_shape = vectorize(x.dims());
|
|
int x_ndim = x_shape.size();
|
|
auto out_shape = vectorize(out.dims());
|
|
int out_ndim = out_shape.size();
|
|
auto out_dist_attr_src = out.dist_attr();
|
|
std::vector<int64_t> out_dims_mapping = out_dist_attr_src.dims_mapping();
|
|
PADDLE_ENFORCE_EQ(
|
|
out_ndim,
|
|
out_dims_mapping.size(),
|
|
common::errors::InvalidArgument("The Tensor Out's rank [%d] and Out's "
|
|
"dims_mapping size [%d] are not matched.",
|
|
out_ndim,
|
|
out_dims_mapping.size()));
|
|
|
|
// Step1: Build the transformation from the output shape
|
|
// to original shape. This function infers the dims mapping
|
|
// from output to input, we first get the transformation
|
|
// from output to input so that we can infer the dims mapping
|
|
// with the map from output axes to input axes.
|
|
|
|
std::vector<std::shared_ptr<DimTrans>> trans;
|
|
|
|
if (static_cast<int64_t>(axis.size()) == 0) {
|
|
MakeSqueezeDimTransReverseWithoutAxis(x_shape, &trans);
|
|
} else {
|
|
std::vector<int64_t> axis_copy(axis);
|
|
for (auto& v : axis_copy) {
|
|
if (v < 0) {
|
|
v += x_ndim;
|
|
}
|
|
}
|
|
MakeSqueezeDimTransReverseWithAxis(x_shape, out_shape, axis_copy, &trans);
|
|
}
|
|
|
|
// Step2: Infer the dims mapping of input with
|
|
// output's dims_mapping and the transformation.
|
|
const auto& dims_mapping_vec = InferFromDimTrans(out, trans);
|
|
const auto& input_dims_mapping = std::get<0>(dims_mapping_vec);
|
|
const auto& output_dims_mapping = std::get<1>(dims_mapping_vec);
|
|
|
|
// Step3: Update the dist attributes of input
|
|
// and output with the inferred dims mapping
|
|
TensorDistAttr out_dist_attr_dst =
|
|
CopyTensorDistAttrForOutput(out_dist_attr_src);
|
|
out_dist_attr_dst.set_dims_mapping(input_dims_mapping);
|
|
if (out_dist_attr_dst.dynamic_dims().size() !=
|
|
out_dist_attr_dst.dims_mapping().size()) {
|
|
VLOG(4) << "SqueezeInferSPMD change output dist attr dynamic dims";
|
|
out_dist_attr_dst.set_default_dynamic_dims(
|
|
out_dist_attr_dst.dims_mapping());
|
|
}
|
|
TensorDistAttr x_dist_attr = CopyTensorDistAttrForOutput(x.dist_attr());
|
|
x_dist_attr.set_dims_mapping(output_dims_mapping);
|
|
if (x_dist_attr.dynamic_dims().size() != x_dist_attr.dims_mapping().size()) {
|
|
VLOG(4) << "SqueezeInferSPMD change x dist attr dynamic dims";
|
|
x_dist_attr.set_default_dynamic_dims(x_dist_attr.dims_mapping());
|
|
}
|
|
|
|
VLOG(4) << "SqueezeInferSpmdReverse: Out shape: [" << str_join(out_shape)
|
|
<< "] X shape: [" << str_join(x_shape) << "]";
|
|
VLOG(4) << "Transformation from output to input:";
|
|
for (int64_t i = 0, n = trans.size(); i < n; i++) {
|
|
VLOG(4) << "\tX axis[" << i << "]: " << trans[i]->to_string();
|
|
}
|
|
VLOG(4) << "Out dims_mapping_src: [" << str_join(out_dims_mapping) << "] "
|
|
<< "dims_mapping_dst: [" << str_join(input_dims_mapping) << "]";
|
|
VLOG(4) << "X dims_mapping: [" << str_join(output_dims_mapping) << "]\n\n";
|
|
|
|
return {{x_dist_attr}, {out_dist_attr_dst}};
|
|
}
|
|
|
|
SpmdInfo SqueezeGradInferSpmd(const DistMetaTensor& x,
|
|
const DistMetaTensor& out_grad,
|
|
const IntArray& axis) {
|
|
auto shape = vectorize(x.dims());
|
|
const auto& spmd = ReshapeInferSpmd(out_grad, shape);
|
|
return {{x.dist_attr(), spmd.first[0]}, {spmd.second[0]}};
|
|
}
|
|
|
|
} // namespace phi::distributed
|