wehub-resource-sync
440 lines
19 KiB
C++
440 lines
19 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/slice.h"
|
|
|
|
#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/utils.h"
|
|
|
|
namespace phi::distributed {
|
|
|
|
std::vector<int64_t> BuildOutputAxisToInputAxisMap(
|
|
const std::vector<int64_t>& decrease_axis, int input_ndim) {
|
|
std::vector<int64_t> output_axis_to_input_axis(input_ndim -
|
|
decrease_axis.size());
|
|
int index = 0;
|
|
for (int i = 0; i < input_ndim; ++i) {
|
|
if (std::find(decrease_axis.begin(), decrease_axis.end(), i) ==
|
|
decrease_axis.end()) {
|
|
output_axis_to_input_axis[index] = i;
|
|
++index;
|
|
}
|
|
}
|
|
return output_axis_to_input_axis;
|
|
}
|
|
|
|
SpmdInfo SliceInferSpmdBase(const DistMetaTensor& input,
|
|
const std::vector<int64_t>& axes,
|
|
const std::vector<int64_t>& decrease_axis,
|
|
const std::vector<int>& starts,
|
|
const std::vector<int>& ends,
|
|
const std::vector<int>& strides) {
|
|
// Step0: Verify input args based on slice logic
|
|
auto input_shape = vectorize(input.dims());
|
|
int input_ndim = static_cast<int>(input_shape.size());
|
|
int output_ndim = input_ndim - static_cast<int>(decrease_axis.size());
|
|
auto input_dist_attr_src = input.dist_attr();
|
|
std::vector<int64_t> input_dims_mapping = input_dist_attr_src.dims_mapping();
|
|
PADDLE_ENFORCE_EQ(input_ndim,
|
|
input_dims_mapping.size(),
|
|
common::errors::InvalidArgument(
|
|
"The Tensor Input's rank [%d] and Input's "
|
|
"dims_mapping size [%d] are not matched.",
|
|
input_ndim,
|
|
input_dims_mapping.size()));
|
|
|
|
// Step1: Build Einsum Notation
|
|
std::string alphabet = "abcdefghijklmnopqrstuvwxyz";
|
|
|
|
// get einsum notation for input
|
|
std::string input_axes = alphabet.substr(0, input_ndim);
|
|
|
|
auto output_input_axis_mapping =
|
|
BuildOutputAxisToInputAxisMap(decrease_axis, input_ndim);
|
|
// get einsum notation for output
|
|
std::string out_axes = alphabet.substr(0, output_ndim);
|
|
for (int i = 0; i < output_ndim; i++) {
|
|
auto input_axis = output_input_axis_mapping[i];
|
|
out_axes[i] = input_axes[input_axis];
|
|
}
|
|
|
|
// Step2.3 get new dist attribute for input. the sliced
|
|
// cannot be sharded, if it is sharded, set it to replicated.
|
|
std::vector<int64_t> input_process_mesh =
|
|
input_dist_attr_src.process_mesh().shape();
|
|
TensorDistAttr input_dist_attr_dst =
|
|
CopyTensorDistAttrForOutput(input_dist_attr_src);
|
|
for (auto axe : axes) {
|
|
int axis = axe < 0 ? axe + input_ndim : axe;
|
|
if (!(axis == (input_ndim - 1) && strides.size() == 1 &&
|
|
(input_shape[axis] / input_process_mesh[input_dims_mapping[axis]]) %
|
|
strides[0] ==
|
|
0)) {
|
|
input_dims_mapping[axis] = -1;
|
|
}
|
|
}
|
|
input_dist_attr_dst.set_dims_mapping(input_dims_mapping);
|
|
|
|
std::vector<int64_t> out_dims_mapping(output_ndim);
|
|
for (int i = 0; i < output_ndim; i++) {
|
|
auto input_axis = output_input_axis_mapping[i];
|
|
out_dims_mapping[i] = input_dims_mapping[input_axis];
|
|
}
|
|
TensorDistAttr out_dist_attr =
|
|
CopyTensorDistAttrForOutput(input_dist_attr_src);
|
|
out_dist_attr.set_dims_mapping(out_dims_mapping);
|
|
|
|
VLOG(4) << "SliceInferSpmd:";
|
|
VLOG(4) << "Einsum Notation: " << input_axes << "-->" << out_axes;
|
|
VLOG(4) << "Input shape: [" << str_join(input_shape) << "] "
|
|
<< "axes: [" << str_join(axes) << "] "
|
|
<< "src_dims_mapping: ["
|
|
<< str_join(input_dist_attr_src.dims_mapping()) << "] "
|
|
<< "dst_dims_mapping: [" << str_join(input_dims_mapping) << "]";
|
|
VLOG(4) << "Output"
|
|
<< " dims_mapping: [" << str_join(out_dims_mapping) << "]";
|
|
VLOG(4) << std::endl;
|
|
|
|
return {{input_dist_attr_dst}, {out_dist_attr}};
|
|
}
|
|
|
|
SpmdInfo SliceInferSpmd(const DistMetaTensor& input,
|
|
const std::vector<int64_t>& axes,
|
|
const std::vector<int>& starts,
|
|
const std::vector<int>& ends,
|
|
const std::vector<int64_t>& infer_flags,
|
|
const std::vector<int64_t>& decrease_axis) {
|
|
// starts, ends, infer_flags and decrease_axis have no impact on the
|
|
// derivation, only to align with the definition in phi api
|
|
return SliceInferSpmdBase(input, axes, {}, starts, ends, {});
|
|
}
|
|
|
|
SpmdInfo SliceInferSpmdReverseBase(const DistMetaTensor& input,
|
|
const DistMetaTensor& output,
|
|
const std::vector<int64_t>& axes,
|
|
const std::vector<int64_t>& decrease_axis,
|
|
const std::vector<int>& starts,
|
|
const std::vector<int>& ends,
|
|
const std::vector<int>& strides) {
|
|
auto output_shape = vectorize(output.dims());
|
|
int out_ndim = output_shape.size();
|
|
auto out_dist_attr = output.dist_attr();
|
|
int out_dims_mapping_size =
|
|
static_cast<int>(out_dist_attr.dims_mapping().size());
|
|
auto input_shape = vectorize(input.dims());
|
|
int input_ndim = static_cast<int>(input_shape.size());
|
|
auto input_dist_attr = input.dist_attr();
|
|
std::vector<int64_t> input_dims_mapping = input_dist_attr.dims_mapping();
|
|
|
|
int decrease_axis_num = decrease_axis.size();
|
|
|
|
PADDLE_ENFORCE_EQ(input_ndim,
|
|
out_ndim + decrease_axis_num,
|
|
common::errors::InvalidArgument(
|
|
"The Tensor Input's rank [%d] is not equal "
|
|
"to the Tensor Output's rank [%d]",
|
|
input_ndim,
|
|
out_ndim));
|
|
|
|
PADDLE_ENFORCE_EQ(
|
|
out_ndim,
|
|
out_dims_mapping_size,
|
|
common::errors::InvalidArgument("The Tensor Output's rank [%d] and Its "
|
|
"dims_mapping size [%d] are not matched.",
|
|
out_ndim,
|
|
out_dims_mapping_size));
|
|
|
|
// Step1: Build Einsum Notation
|
|
std::string alphabet = "abcdefghijklmnopqrstuvwxyz";
|
|
|
|
// get einsum notation for input
|
|
std::string input_axes = alphabet.substr(0, input_ndim);
|
|
|
|
auto output_input_axis_mapping =
|
|
BuildOutputAxisToInputAxisMap(decrease_axis, input_ndim);
|
|
// get einsum notation for output
|
|
std::string out_axes = alphabet.substr(0, out_ndim);
|
|
for (int i = 0; i < out_ndim; i++) {
|
|
auto input_axis = output_input_axis_mapping[i];
|
|
out_axes[i] = input_axes[input_axis];
|
|
}
|
|
|
|
std::vector<int64_t> input_process_mesh =
|
|
input_dist_attr.process_mesh().shape();
|
|
for (auto axe : axes) {
|
|
int axis = axe < 0 ? axe + input_ndim : axe;
|
|
// the sliced axis cannot be sharded, set its notation
|
|
// with the special '1' to set its dim mapping to -1.
|
|
if (!(axis == (input_ndim - 1) && strides.size() == 1 &&
|
|
(input_shape[axis] / input_process_mesh[input_dims_mapping[axis]]) %
|
|
strides[0] ==
|
|
0)) {
|
|
input_axes[axis] = '1';
|
|
}
|
|
}
|
|
|
|
// Step2: Sharding Propagation
|
|
// Step2.1: merge output shardings
|
|
std::vector<std::pair<std::string, std::vector<int64_t>>> axes_sharding_info;
|
|
std::vector<int64_t> out_dims_mapping = output.dist_attr().dims_mapping();
|
|
axes_sharding_info.emplace_back(out_axes, out_dims_mapping);
|
|
|
|
std::unordered_map<std::string, int64_t> axis_to_dim_map =
|
|
ShardingMergeForTensors(axes_sharding_info);
|
|
|
|
// Step2.2: infer input dims mapping from output dims mapping. the sliced
|
|
// cannot be sharded, if it is sharded, set it to replicated.
|
|
input_dims_mapping = GetDimsMappingForAxes(input_axes, axis_to_dim_map, true);
|
|
|
|
auto input_dist_attr_dst = CopyTensorDistAttrForOutput(input_dist_attr);
|
|
input_dist_attr_dst.set_dims_mapping(input_dims_mapping);
|
|
|
|
// step2.3 get new dist attribute for output. the sliced
|
|
// cannot be sharded, if it is sharded, set it to replicated.
|
|
out_dims_mapping = GetDimsMappingForAxes(out_axes, axis_to_dim_map, true);
|
|
std::vector<int64_t> output_process_mesh =
|
|
out_dist_attr.process_mesh().shape();
|
|
for (auto axe : axes) {
|
|
int axis = axe < 0 ? axe + input_ndim : axe;
|
|
if (!(axis == (out_ndim - 1) && strides.size() == 1 &&
|
|
(output_shape[axis] / output_process_mesh[out_dims_mapping[axis]]) %
|
|
strides[0] ==
|
|
0)) {
|
|
out_dims_mapping[axis] = -1;
|
|
}
|
|
}
|
|
auto out_dist_attr_dst = CopyTensorDistAttrForOutput(out_dist_attr);
|
|
out_dist_attr_dst.set_dims_mapping(out_dims_mapping);
|
|
|
|
VLOG(4) << "SliceInferSpmdReverse:";
|
|
VLOG(4) << "Einsum Notation: " << input_axes << "-->" << out_axes;
|
|
VLOG(4) << "Output"
|
|
<< " shape: [" << str_join(vectorize(output.dims())) << "] "
|
|
<< "axes: [" << str_join(axes) << "] "
|
|
<< "src_dims_mapping: ["
|
|
<< str_join(output.dist_attr().dims_mapping()) << "] "
|
|
<< "dst_dims_mapping: [" << str_join(out_dist_attr_dst.dims_mapping())
|
|
<< "]";
|
|
VLOG(4) << "Input shape: [" << str_join(input_shape) << "] "
|
|
<< "dims_mapping: [" << str_join(input_dims_mapping) << "]\n\n";
|
|
|
|
return {{input_dist_attr_dst}, {out_dist_attr_dst}};
|
|
}
|
|
|
|
SpmdInfo SliceInferSpmdReverse(const DistMetaTensor& input,
|
|
const DistMetaTensor& output,
|
|
const std::vector<int64_t>& axes,
|
|
const std::vector<int>& starts,
|
|
const std::vector<int>& ends,
|
|
const std::vector<int64_t>& infer_flags,
|
|
const std::vector<int64_t>& decrease_axis) {
|
|
// starts, ends, infer_flags and decrease_axis have no impact on the
|
|
// derivation, only to align with the definition in phi api
|
|
return SliceInferSpmdReverseBase(input, output, axes, {}, starts, ends, {});
|
|
}
|
|
|
|
SpmdInfo SliceInferSpmdDynamic(const DistMetaTensor& input,
|
|
const std::vector<int64_t>& axes,
|
|
const IntArray& starts,
|
|
const IntArray& ends,
|
|
const std::vector<int64_t>& infer_flags,
|
|
const std::vector<int64_t>& decrease_axis) {
|
|
// starts, ends, infer_flags and decrease_axis have no impact on the
|
|
// derivation, only to align with the definition in phi api
|
|
std::vector<int> start_indexes(starts.GetData().begin(),
|
|
starts.GetData().end());
|
|
std::vector<int> end_indexes(ends.GetData().begin(), ends.GetData().end());
|
|
return SliceInferSpmdBase(input, axes, decrease_axis, {}, {}, {});
|
|
}
|
|
|
|
SpmdInfo ViewSliceInferSpmd(const DistMetaTensor& input,
|
|
int64_t begin_idx,
|
|
int64_t end_idx) {
|
|
auto input_dist_attr_src = input.dist_attr();
|
|
std::vector<int64_t> input_dims_mapping = input_dist_attr_src.dims_mapping();
|
|
input_dims_mapping[0] = -1;
|
|
input_dist_attr_src.set_dims_mapping(input_dims_mapping);
|
|
return {{input_dist_attr_src}, {input_dist_attr_src}};
|
|
}
|
|
|
|
SpmdInfo SliceGradInferBase(const DistMetaTensor& input,
|
|
const DistMetaTensor& out_grad,
|
|
const std::vector<int64_t>& axes,
|
|
const std::vector<int64_t>& decrease_axis) {
|
|
// Step0: Verify input args based on slice logic
|
|
auto input_shape = vectorize(input.dims());
|
|
int input_ndim = static_cast<int>(input_shape.size());
|
|
auto input_dist_attr = input.dist_attr();
|
|
|
|
input_dist_attr = UnShardTensorDims(input_dist_attr, axes);
|
|
std::vector<int64_t> input_dims_mapping = input_dist_attr.dims_mapping();
|
|
|
|
auto output_axis_to_input_axis_mapping =
|
|
BuildOutputAxisToInputAxisMap(decrease_axis, input_ndim);
|
|
std::unordered_map<int, int> reverse_output_axis_to_input_axis_mapping;
|
|
for (size_t i = 0; i < output_axis_to_input_axis_mapping.size(); ++i) {
|
|
reverse_output_axis_to_input_axis_mapping
|
|
[output_axis_to_input_axis_mapping[i]] = i;
|
|
}
|
|
std::vector<int64_t> mapped_axes;
|
|
for (const auto& axe : axes) {
|
|
int axis = axe < 0 ? axe + input_ndim : axe;
|
|
if (reverse_output_axis_to_input_axis_mapping.count(axis) > 0) {
|
|
mapped_axes.push_back(reverse_output_axis_to_input_axis_mapping[axis]);
|
|
}
|
|
}
|
|
auto out_dist_attr = out_grad.dist_attr();
|
|
out_dist_attr = UnShardTensorDims(out_dist_attr, mapped_axes);
|
|
auto output_shape = vectorize(out_grad.dims());
|
|
int out_ndim = output_shape.size();
|
|
int out_dims_mapping_size =
|
|
static_cast<int>(out_dist_attr.dims_mapping().size());
|
|
int decrease_axis_num = decrease_axis.size();
|
|
|
|
PADDLE_ENFORCE_EQ(input_ndim,
|
|
out_ndim + decrease_axis_num,
|
|
common::errors::InvalidArgument(
|
|
"The Tensor Input's rank [%d] is not equal "
|
|
"to the Tensor Output's rank [%d]",
|
|
input_ndim,
|
|
out_ndim));
|
|
|
|
PADDLE_ENFORCE_EQ(
|
|
out_ndim,
|
|
out_dims_mapping_size,
|
|
common::errors::InvalidArgument("The Tensor Output's rank [%d] and Its "
|
|
"dims_mapping size [%d] are not matched.",
|
|
out_ndim,
|
|
out_dims_mapping_size));
|
|
|
|
// Step1: Build Einsum Notation
|
|
std::string alphabet = "abcdefghijklmnopqrstuvwxyz";
|
|
|
|
// get einsum notation for input
|
|
std::string input_axes = alphabet.substr(0, input_ndim);
|
|
|
|
// get einsum notation for output
|
|
std::string out_axes(out_ndim, ' ');
|
|
for (int i = 0; i < out_ndim; ++i) {
|
|
out_axes[i] = input_axes[output_axis_to_input_axis_mapping[i]];
|
|
}
|
|
|
|
// Step2: Sharding Propagation
|
|
// Step2.1: merge input shardings
|
|
std::vector<std::pair<std::string, std::vector<int64_t>>> axes_sharding_info;
|
|
axes_sharding_info.emplace_back(out_axes, out_dist_attr.dims_mapping());
|
|
axes_sharding_info.emplace_back(input_axes, input_dist_attr.dims_mapping());
|
|
std::unordered_map<std::string, int64_t> axis_to_dim_map =
|
|
ShardingMergeForTensors(axes_sharding_info);
|
|
|
|
// Step2.2: infer output dims mapping from merged input dims mapping
|
|
auto input_dim_mapping_dst =
|
|
GetDimsMappingForAxes(input_axes, axis_to_dim_map, true);
|
|
|
|
// get the dist attributes for output
|
|
TensorDistAttr input_grad = CopyTensorDistAttrForOutput(input_dist_attr);
|
|
input_dist_attr.set_dims_mapping(input_dim_mapping_dst);
|
|
input_grad.set_dims_mapping(input_dim_mapping_dst);
|
|
|
|
auto out_grad_dims_mapping =
|
|
GetDimsMappingForAxes(out_axes, axis_to_dim_map, true);
|
|
out_dist_attr.set_dims_mapping(out_grad_dims_mapping);
|
|
|
|
VLOG(4) << "SliceGradInfer:";
|
|
|
|
VLOG(4) << "input"
|
|
<< " shape: [" << str_join(input_shape) << "] "
|
|
<< "src_dims_mapping: [" << str_join(input.dist_attr().dims_mapping())
|
|
<< "] "
|
|
<< "dst_dims_mapping: [" << str_join(input_dist_attr.dims_mapping())
|
|
<< "]";
|
|
|
|
VLOG(4) << "Output Grad"
|
|
<< " shape: [" << str_join(output_shape) << "] "
|
|
<< "src_dims_mapping: ["
|
|
<< str_join(out_grad.dist_attr().dims_mapping()) << "] "
|
|
<< "dst_dims_mapping: [" << str_join(out_dist_attr.dims_mapping())
|
|
<< "]";
|
|
|
|
VLOG(4) << "input Grad"
|
|
<< " shape: [" << str_join(output_shape) << "] "
|
|
<< "dims_mapping: [" << str_join(input_grad.dims_mapping()) << "] ";
|
|
|
|
return {{input_dist_attr, out_dist_attr}, {input_grad}};
|
|
}
|
|
|
|
SpmdInfo SliceGradInferSpmdDynamic(const DistMetaTensor& input,
|
|
const DistMetaTensor& out_grad,
|
|
const std::vector<int64_t>& axes,
|
|
const IntArray& starts,
|
|
const IntArray& ends,
|
|
const std::vector<int64_t>& infer_flags,
|
|
const std::vector<int64_t>& decrease_axis) {
|
|
// starts, ends, infer_flags and decrease_axis have no impact on the
|
|
// derivation, only to align with the definition in phi api
|
|
return SliceGradInferBase(input, out_grad, axes, decrease_axis);
|
|
}
|
|
|
|
SpmdInfo StridedSliceInferSpmd(const DistMetaTensor& input,
|
|
const std::vector<int>& axes,
|
|
const std::vector<int>& starts,
|
|
const std::vector<int>& ends,
|
|
const std::vector<int>& strides) {
|
|
// starts, ends and strides have no impact on the derivation,
|
|
// only to align with the definition in phi api
|
|
std::vector<int64_t> axes_bridge(axes.begin(), axes.end());
|
|
return SliceInferSpmdBase(input, axes_bridge, {}, starts, ends, strides);
|
|
}
|
|
|
|
SpmdInfo StridedSliceGradInferSpmd(const DistMetaTensor& input,
|
|
const DistMetaTensor& out_grad,
|
|
const std::vector<int>& axes,
|
|
const std::vector<int>& starts,
|
|
const std::vector<int>& ends,
|
|
const std::vector<int>& strides) {
|
|
// starts, ends and strides have no impact on the derivation,
|
|
// only to align with the definition in phi api
|
|
std::vector<int64_t> axes_bridge(axes.begin(), axes.end());
|
|
return SliceGradInferBase(input, out_grad, axes_bridge, {});
|
|
}
|
|
|
|
SpmdInfo StridedSliceInferSpmdDynamic(const DistMetaTensor& input,
|
|
const std::vector<int>& axes,
|
|
const IntArray& starts,
|
|
const IntArray& ends,
|
|
const IntArray& strides) {
|
|
// starts, ends and strides have no impact on the derivation,
|
|
// only to align with the definition in phi api
|
|
std::vector<int64_t> axes_bridge(axes.begin(), axes.end());
|
|
return SliceInferSpmdBase(input, axes_bridge, {}, {}, {}, {});
|
|
}
|
|
|
|
SpmdInfo StridedSliceGradInferSpmdDynamic(const DistMetaTensor& input,
|
|
const DistMetaTensor& out_grad,
|
|
const std::vector<int>& axes,
|
|
const IntArray& starts,
|
|
const IntArray& ends,
|
|
const IntArray& strides) {
|
|
// starts, ends and strides have no impact on the derivation,
|
|
// only to align with the definition in phi api
|
|
std::vector<int64_t> axes_bridge(axes.begin(), axes.end());
|
|
return SliceGradInferBase(input, out_grad, axes_bridge, {});
|
|
}
|
|
|
|
} // namespace phi::distributed
|