Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

add strided_slice_grad op for npu #35204

Merged
merged 2 commits into from
Sep 1, 2021
Merged

add strided_slice_grad op for npu #35204

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f802bc6
add strided_slice_grad op for npu
baoachun Aug 27, 2021
70fa9ab
add strided_slice_grad op for npu
baoachun Aug 28, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
202 changes: 196 additions & 6 deletions paddle/fluid/operators/strided_slice_op_npu.cc
100755 → 100644
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -226,14 +226,204 @@ class StridedSliceNPUKernel : public framework::OpKernel<T> {
}
};

template <typename DeviceContext, typename T>
class StridedSliceGradNPUKernel : public framework::OpKernel<T> {
public:
void Compute(const framework::ExecutionContext& ctx) const override {
const Variable* input_var = ctx.InputVar("Input");
bool is_tensor_array = input_var->IsType<LoDTensorArray>();
PADDLE_ENFORCE_EQ(is_tensor_array, false,
platform::errors::InvalidArgument(
"Tensor array as input is not supported."));
int rank = ctx.Input<framework::Tensor>("Input")->dims().size();

switch (rank) {
case 1:
StridedSliceGradCompute<1>(ctx);
break;
case 2:
StridedSliceGradCompute<2>(ctx);
break;
case 3:
StridedSliceGradCompute<3>(ctx);
break;
case 4:
StridedSliceGradCompute<4>(ctx);
break;
case 5:
StridedSliceGradCompute<5>(ctx);
break;
case 6:
StridedSliceGradCompute<6>(ctx);
break;
default:
PADDLE_THROW(platform::errors::InvalidArgument(
"The rank of input is supported up to 6."));
break;
}
}

private:
template <size_t D>
void StridedSliceGradCompute(const framework::ExecutionContext& ctx) const {
auto place = ctx.GetPlace();
auto& dev_ctx =
ctx.template device_context<paddle::platform::NPUDeviceContext>();

auto* input = ctx.Input<framework::Tensor>("Input");
auto input_dims = input->dims();
auto* dout = ctx.Input<framework::Tensor>(framework::GradVarName("Out"));
auto* dx = ctx.Output<Tensor>(framework::GradVarName("Input"));
dx->mutable_data<T>(input_dims, place);

auto starts_int = ctx.Attr<std::vector<int>>("starts");
auto ends_int = ctx.Attr<std::vector<int>>("ends");
auto strides_int = ctx.Attr<std::vector<int>>("strides");

std::vector<int64_t> starts(starts_int.begin(), starts_int.end());
std::vector<int64_t> ends(ends_int.begin(), ends_int.end());
std::vector<int64_t> strides(strides_int.begin(), strides_int.end());

auto axes = ctx.Attr<std::vector<int>>("axes");
auto infer_flags = ctx.Attr<std::vector<int>>("infer_flags");
auto decrease_axis = ctx.Attr<std::vector<int>>("decrease_axis");

auto list_new_ends_tensor =
ctx.MultiInput<framework::Tensor>("EndsTensorList");
auto list_new_starts_tensor =
ctx.MultiInput<framework::Tensor>("StartsTensorList");
auto list_new_strides_tensor =
ctx.MultiInput<framework::Tensor>("StridesTensorList");

if (list_new_starts_tensor.size() > 0) {
starts = GetDataFromTensorList<int64_t>(list_new_starts_tensor);
} else if (ctx.HasInput("StartsTensor")) {
auto* starts_tensor = ctx.Input<framework::Tensor>("StartsTensor");
starts = GetDataFromTensor<int64_t>(starts_tensor);
}

if (list_new_ends_tensor.size() > 0) {
ends = GetDataFromTensorList<int64_t>(list_new_ends_tensor);
} else if (ctx.HasInput("EndsTensor")) {
auto* ends_tensor = ctx.Input<framework::Tensor>("EndsTensor");
ends = GetDataFromTensor<int64_t>(ends_tensor);
}

if (list_new_strides_tensor.size() > 0) {
strides = GetDataFromTensorList<int64_t>(list_new_strides_tensor);
} else if (ctx.HasInput("StridesTensor")) {
auto* strides_tensor = ctx.Input<framework::Tensor>("StridesTensor");
strides = GetDataFromTensor<int64_t>(strides_tensor);
}

std::vector<int64_t> out_dims_vector(input_dims.size(), -1);
StridedSliceOutDims(starts, ends, strides, axes, infer_flags, input_dims,
decrease_axis, out_dims_vector.data(), axes.size(),
false);

std::vector<int> reverse_vector(starts.size(), 0);
StridedSliceFunctor(starts.data(), ends.data(), strides.data(), axes.data(),
reverse_vector.data(), input_dims, infer_flags,
decrease_axis, starts.size());

std::vector<int64_t> starts_indices_vector(D, 0);
std::vector<int64_t> ends_indices_vector(out_dims_vector.begin(),
out_dims_vector.end());
std::vector<int64_t> strides_indices_vector(D, 1);

for (size_t axis = 0; axis < axes.size(); axis++) {
int axis_index = axes[axis];
starts_indices_vector[axis_index] = starts[axis];
ends_indices_vector[axis_index] = ends[axis];
strides_indices_vector[axis_index] = strides[axis];
}

Tensor starts_indices_tensor;
Tensor ends_indices_tensor;
Tensor strides_indices_tensor;

starts_indices_tensor.mutable_data<int64_t>({D}, place);
ends_indices_tensor.mutable_data<int64_t>({D}, place);
strides_indices_tensor.mutable_data<int64_t>({D}, place);

TensorFromVector(starts_indices_vector, dev_ctx, &starts_indices_tensor);
TensorFromVector(ends_indices_vector, dev_ctx, &ends_indices_tensor);
TensorFromVector(strides_indices_vector, dev_ctx, &strides_indices_tensor);

std::vector<int64_t> input_dims_vector;
for (int i = 0; i < input_dims.size(); i++) {
input_dims_vector.push_back(input_dims[i]);
}
Tensor input_dims_tensor;
TensorFromVector(input_dims_vector, dev_ctx, &input_dims_tensor);

bool need_reverse = false;
for (size_t axis = 0; axis < axes.size(); axis++) {
if (reverse_vector[axis] == 1) {
need_reverse = true;
break;
}
}

auto stream = dev_ctx.stream();
framework::NPUAttributeMap attr_input = {{"begin_mask", 0},
{"end_mask", 0},
{"ellipsis_mask", 0},
{"new_axis_mask", 0},
{"shrink_axis_mask", 0}};

if (need_reverse) {
Tensor reverse_axis;
std::vector<int> reverse_axis_vector;
for (size_t axis = 0; axis < axes.size(); axis++) {
if (reverse_vector[axis] == 1) {
reverse_axis_vector.push_back(axes[axis]);
}
}
reverse_axis.mutable_data<int>(
{static_cast<int>(reverse_axis_vector.size())}, place);
TensorFromVector(reverse_axis_vector, dev_ctx, &reverse_axis);

Tensor dout_tmp;
dout_tmp.mutable_data<T>(dout->dims(), place);
const auto& runner_reverse =
NpuOpRunner("ReverseV2", {*dout, reverse_axis}, {dout_tmp});
runner_reverse.Run(stream);

const auto& runner =
NpuOpRunner("StridedSliceGrad",
{input_dims_tensor, starts_indices_tensor,
ends_indices_tensor, strides_indices_tensor, dout_tmp},
{*dx}, attr_input);
runner.Run(stream);
} else {
const auto& runner =
NpuOpRunner("StridedSliceGrad",
{input_dims_tensor, starts_indices_tensor,
ends_indices_tensor, strides_indices_tensor, *dout},
{*dx}, attr_input);
runner.Run(stream);
}
}
};

} // namespace operators
} // namespace paddle

namespace ops = paddle::operators;
namespace plat = paddle::platform;

REGISTER_OP_NPU_KERNEL(
strided_slice, ops::StridedSliceNPUKernel<plat::NPUDeviceContext, bool>,
ops::StridedSliceNPUKernel<plat::NPUDeviceContext, int>,
ops::StridedSliceNPUKernel<plat::NPUDeviceContext, int64_t>,
ops::StridedSliceNPUKernel<plat::NPUDeviceContext, float>,
ops::StridedSliceNPUKernel<plat::NPUDeviceContext, double>);

REGISTER_OP_NPU_KERNEL(
strided_slice,
ops::StridedSliceNPUKernel<paddle::platform::NPUDeviceContext, bool>,
ops::StridedSliceNPUKernel<paddle::platform::NPUDeviceContext, int>,
ops::StridedSliceNPUKernel<paddle::platform::NPUDeviceContext, int64_t>,
ops::StridedSliceNPUKernel<paddle::platform::NPUDeviceContext, float>,
ops::StridedSliceNPUKernel<paddle::platform::NPUDeviceContext, double>);
strided_slice_grad,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

前向反向的type没有对齐，如果是因为Ascend库函数前反向类型没有对齐的原因，就可以忽略

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

刚看到资料里前向上没有bool类型的，前向注册类型可能需要改下。
图片

ops::StridedSliceGradNPUKernel<plat::NPUDeviceContext, plat::float16>,
ops::StridedSliceGradNPUKernel<plat::NPUDeviceContext, float>,
ops::StridedSliceGradNPUKernel<plat::NPUDeviceContext, double>,
ops::StridedSliceGradNPUKernel<plat::NPUDeviceContext, int>,
ops::StridedSliceGradNPUKernel<plat::NPUDeviceContext, int64_t>);
Loading