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

Jump to bottom

[NPU] add broadcast supporting for elementwise_add_op_npu #34057

Merged
merged 7 commits into from
Aug 9, 2021
+655 −96
Merged

[NPU] add broadcast supporting for elementwise_add_op_npu #34057

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
5704ba8
add broadcast supporting for elementwise_add
ronny1996 Jul 12, 2021
3bbf0dc
add broadcast supporting for elementwise_add
ronny1996 Jul 12, 2021
1541110
add more tests
ronny1996 Jul 13, 2021
7c48829
remove the redundant code
ronny1996 Jul 13, 2021
5d4f4ec
update
ronny1996 Jul 14, 2021
8d8d87c
fix place error in unittest
ronny1996 Jul 15, 2021
79a4d2a
remove skip.If
ronny1996 Jul 21, 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
178 changes: 85 additions & 93 deletions paddle/fluid/operators/elementwise/elementwise_add_op_npu.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ limitations under the License. */

#include "paddle/fluid/framework/tensor_util.h"
#include "paddle/fluid/operators/elementwise/elementwise_add_op.h"
#include "paddle/fluid/operators/elementwise/elementwise_npu.h"
#include "paddle/fluid/operators/npu_op_runner.h"

namespace paddle {
Expand All @@ -27,12 +28,37 @@ template <typename T>
class ElementwiseAddNPUKernel : public framework::OpKernel<T> {
public:
void Compute(const framework::ExecutionContext& ctx) const override {
auto& dev_ctx =
ctx.template device_context<paddle::platform::NPUDeviceContext>();
auto* x = ctx.Input<framework::LoDTensor>("X");
auto* y = ctx.Input<framework::LoDTensor>("Y");
auto* out = ctx.Output<framework::LoDTensor>("Out");
out->mutable_data<T>(ctx.GetPlace());

const auto& runner = NpuOpRunner("Add", {*x, *y}, {*out}, {});
int axis = ctx.Attr<int>("axis");

bool direct_compute = false;
auto x_dims = x->dims();
auto y_dims = y->dims();
axis = (axis == -1 ? std::abs(x_dims.size() - y_dims.size()) : axis);
if (x_dims.size() >= y_dims.size()) {
direct_compute =
y_dims == framework::slice_ddim(x_dims, axis, x_dims.size());
} else {
direct_compute =
x_dims == framework::slice_ddim(y_dims, axis, y_dims.size());
}

Tensor transformed_x, transformed_y;
if (direct_compute) {
transformed_x.ShareDataWith(*x);
transformed_y.ShareDataWith(*y);
} else {
NpuElementWiseOpBroadcast<T>(dev_ctx, x, y, axis, &transformed_x,
&transformed_y);
}
const auto& runner =
NpuOpRunner("Add", {transformed_x, transformed_y}, {*out}, {});
auto stream =
ctx.template device_context<paddle::platform::NPUDeviceContext>()
.stream();
Expand All @@ -44,109 +70,75 @@ template <typename T>
class ElementwiseAddGradNPUKernel : public framework::OpKernel<T> {
public:
void Compute(const framework::ExecutionContext& ctx) const override {
auto* dout = ctx.Input<Tensor>(framework::GradVarName("Out"));
auto* dx = ctx.Output<Tensor>(framework::GradVarName("X"));
auto* dy = ctx.Output<Tensor>(framework::GradVarName("Y"));

auto stream =
ctx.template device_context<paddle::platform::NPUDeviceContext>()
.stream();

// NOTE(zhiqiu): It seems Ascend Sub follow the broadcast sematics with
// default axis=-1?
// So, the sub_grad should do reduce if needed.
// For example, the shape of each variable in elementwise_sub:
// x, dx: [2, 3, 5]
// y, dy: [1, 5]
// out, dout: [2, 3, 5]
// Then, out = x - y => dx = dout, dy = -dout
// And, the shape of dy can be computed by two stages reduce,
// 1. [2, 3, 5] => [3, 5], ReduceSumD on axis = 0, keep_dims = false.
// 2. [3, 5] => [1, 5], ReduceSumD on axis = 0, keep_dims = true.

auto& dev_ctx =
ctx.template device_context<paddle::platform::NPUDeviceContext>();
auto* x = ctx.Input<framework::Tensor>("X");
auto* y = ctx.Input<framework::Tensor>("Y");
auto* dout = ctx.Input<framework::Tensor>(framework::GradVarName("Out"));
auto* dx = ctx.Output<framework::Tensor>(framework::GradVarName("X"));
auto* dy = ctx.Output<framework::Tensor>(framework::GradVarName("Y"));
int axis = ctx.Attr<int>("axis");

axis = (axis == -1 ? std::abs(x->dims().size() - y->dims().size()) : axis);
auto stream = dev_ctx.stream();
if (dx) {
dx->mutable_data<T>(ctx.GetPlace());
// For dx
// stage 1
auto reduce_ndim = dout->dims().size() - dx->dims().size();
std::vector<int> axes;
for (auto i = 0; i < reduce_ndim; ++i) {
axes.push_back(i);
}
Tensor* tmp_dout = const_cast<Tensor*>(dout);
Tensor reduced_dout(dx->type());
if (axes.size() != 0) {
std::vector<int64_t> reduced_dout_dims;
for (auto i = reduce_ndim; i < dout->dims().size(); ++i) {
reduced_dout_dims.push_back(dout->dims()[i]);
if (dx->dims() != dout->dims()) {
std::vector<int> dst_dims_vec;
std::vector<int> reduce_axes;
auto src_dims = dx->dims();
auto dout_dims = dout->dims();

int src_axis = (src_dims.size() < dout_dims.size() ? axis : 0);
for (int ax = 0; ax < dout_dims.size(); ++ax) {
if ((ax < src_axis || ax >= src_axis + src_dims.size()) ||
(dout_dims[ax] > 1 && src_dims[ax - src_axis] == 1)) {
reduce_axes.push_back(ax);
} else {
dst_dims_vec.push_back(dout_dims[ax]);
}
}
reduced_dout.Resize(framework::make_ddim(reduced_dout_dims));
reduced_dout.mutable_data<T>(ctx.GetPlace());
const auto& runner =
NpuOpRunner("ReduceSumD", {*dout}, {reduced_dout},
{{"axes", axes}, {"keep_dims", false}});
runner.Run(stream);
tmp_dout = &reduced_dout;
}

// stage 2
axes.clear();
for (auto i = 0; i < dx->dims().size(); ++i) {
if (dx->dims()[i] == 1) {
axes.push_back(i);
if (!reduce_axes.empty()) {
Tensor tmp;
tmp.ShareDataWith(*dx);
tmp.Resize(framework::make_ddim(dst_dims_vec));
const auto& runner =
NpuOpRunner("ReduceSumD", {*dout}, {tmp},
{{"axes", reduce_axes}, {"keep_dims", false}});
runner.Run(stream);
}
}
if (axes.size() != 0) {
const auto& runner = NpuOpRunner("ReduceSumD", {*tmp_dout}, {*dx},
{{"axes", axes}, {"keep_dims", true}});
runner.Run(stream);
} else {
framework::TensorCopy(
*tmp_dout, ctx.GetPlace(),
ctx.template device_context<platform::DeviceContext>(), dx);
framework::TensorCopy(*dout, ctx.GetPlace(), dev_ctx, dx);
}
}

if (dy) {
// For dy
// stage 1
auto reduce_ndim = dout->dims().size() - dy->dims().size();
std::vector<int> axes;
for (auto i = 0; i < reduce_ndim; ++i) {
axes.push_back(i);
}
Tensor* tmp_dout = const_cast<Tensor*>(dout);
Tensor reduced_dout(dout->type());
if (axes.size() != 0) {
std::vector<int64_t> reduced_dout_dims;
for (auto i = reduce_ndim; i < dout->dims().size(); ++i) {
reduced_dout_dims.push_back(dout->dims()[i]);
dy->mutable_data<T>(ctx.GetPlace());
if (dy->dims() != dout->dims()) {
std::vector<int> dst_dims_vec;
std::vector<int> reduce_axes;
auto src_dims = dy->dims();
auto dout_dims = dout->dims();

int src_axis = (src_dims.size() < dout_dims.size() ? axis : 0);
for (int ax = 0; ax < dout_dims.size(); ++ax) {
if ((ax < src_axis || ax >= src_axis + src_dims.size()) ||
(dout_dims[ax] > 1 && src_dims[ax - src_axis] == 1)) {
reduce_axes.push_back(ax);
} else {
dst_dims_vec.push_back(dout_dims[ax]);
}
}
reduced_dout.Resize(framework::make_ddim(reduced_dout_dims));
reduced_dout.mutable_data<T>(ctx.GetPlace());
const auto& runner =
NpuOpRunner("ReduceSumD", {*dout}, {reduced_dout},
{{"axes", axes}, {"keep_dims", false}});
runner.Run(stream);
tmp_dout = &reduced_dout;
}

// stage 2
axes.clear();
for (auto i = 0; i < dy->dims().size(); ++i) {
if (dy->dims()[i] == 1) {
axes.push_back(i);
if (!reduce_axes.empty()) {
Tensor tmp;
tmp.ShareDataWith(*dy);
tmp.Resize(framework::make_ddim(dst_dims_vec));
const auto& runner =
NpuOpRunner("ReduceSumD", {*dout}, {tmp},
{{"axes", reduce_axes}, {"keep_dims", false}});
runner.Run(stream);
}
}
if (axes.size() != 0) {
dy->mutable_data<T>(ctx.GetPlace());
const auto& runner = NpuOpRunner("ReduceSumD", {*tmp_dout}, {*dy},
{{"axes", axes}, {"keep_dims", true}});
runner.Run(stream);
} else {
framework::TensorCopy(
*tmp_dout, ctx.GetPlace(),
ctx.template device_context<platform::DeviceContext>(), dy);
framework::TensorCopy(*dout, ctx.GetPlace(), dev_ctx, dy);
}
}
}
Expand Down
135 changes: 135 additions & 0 deletions paddle/fluid/operators/elementwise/elementwise_npu.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,135 @@
/* Copyright (c) 2021 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/framework/tensor_util.h"
#include "paddle/fluid/operators/elementwise/elementwise_op.h"
#include "paddle/fluid/operators/elementwise/elementwise_op_function.h"
#include "paddle/fluid/operators/npu_op_runner.h"

namespace paddle {
namespace operators {
using Tensor = framework::Tensor;

template <typename T>
void NpuBroadcast(const platform::NPUDeviceContext& dev_ctx, const Tensor* src,
int axis, const framework::DDim& dst_dims,
Tensor* transformed_src) {
auto stream = dev_ctx.stream();

// 1. expand the axis with dim 1
auto src_dims = src->dims();
Tensor tmp_src;
tmp_src.ShareDataWith(*src);
tmp_src.Resize(src_dims);
for (int i = 0; i < src_dims.size(); ++i) {
if (src_dims[i] == 1 && dst_dims[i + axis] > 1) {
Tensor tmp_tensor;
auto tmp_tensor_dims = tmp_src.dims();
tmp_tensor_dims[i] = dst_dims[i + axis];
tmp_tensor.mutable_data<T>(tmp_tensor_dims, dev_ctx.GetPlace());
const auto& runner =
NpuOpRunner("TileWithAxis", {tmp_src}, {tmp_tensor},
{{"axis", static_cast<int64_t>(i)},
{"tiles", static_cast<int64_t>(dst_dims[i + axis])}});
runner.Run(stream);
tmp_src.ShareDataWith(tmp_tensor);
tmp_src.Resize(tmp_tensor_dims);
}
}

// 2.expand the ahead axis
auto prev = framework::product(framework::slice_ddim(dst_dims, 0, axis));
if (prev > 1) {
Tensor tmp_tensor;
auto tmp_tensor_dims =
framework::slice_ddim(dst_dims, 0, axis + src_dims.size());
tmp_tensor.mutable_data<T>(tmp_tensor_dims, dev_ctx.GetPlace());
const auto& runner = NpuOpRunner(
"ExpandD", {tmp_src}, {tmp_tensor},
{{"shape", framework::vectorize<int64_t>(tmp_tensor_dims)}});
runner.Run(stream);
tmp_src.ShareDataWith(tmp_tensor);
tmp_src.Resize(tmp_tensor_dims);
} else {
tmp_src.Resize(framework::slice_ddim(dst_dims, 0, axis + src_dims.size()));
}

// 3.expand the tail axis
auto post = framework::product(
framework::slice_ddim(dst_dims, axis + src_dims.size(), dst_dims.size()));
if (post > 1) {
auto src_dims_vec = framework::vectorize<int>(tmp_src.dims());
src_dims_vec.push_back(1);
tmp_src.Resize(framework::make_ddim(src_dims_vec));

Tensor tmp_tensor;
tmp_tensor.mutable_data<T>(dst_dims, dev_ctx.GetPlace());
const auto& runner =
NpuOpRunner("TileWithAxis", {tmp_src}, {tmp_tensor},
{{"axis", static_cast<int64_t>(axis + src_dims.size())},
{"tiles", static_cast<int64_t>(post)}});
runner.Run(stream);
tmp_src.ShareDataWith(tmp_tensor);
}
tmp_src.Resize(dst_dims);
framework::TensorCopy(tmp_src, dev_ctx.GetPlace(), transformed_src);
}

template <typename T>
void NpuElementWiseOpBroadcast(const platform::NPUDeviceContext& dev_ctx,
const Tensor* x, const Tensor* y, int axis,
Tensor* transformed_x, Tensor* transformed_y) {
auto x_dims = x->dims();
auto y_dims = y->dims();
bool is_xsize_larger = true;
int max_dim = x_dims.size();
std::vector<int> dst_dims_vec = framework::vectorize<int>(x_dims);

if (x_dims.size() < y_dims.size()) {
is_xsize_larger = false;
max_dim = y_dims.size();
dst_dims_vec = framework::vectorize<int>(y_dims);
}

axis = (axis == -1 ? std::abs(x_dims.size() - y_dims.size()) : axis);
int x_axis = is_xsize_larger ? 0 : axis;
int y_axis = is_xsize_larger ? axis : 0;

PADDLE_ENFORCE_GE(
axis, 0,
platform::errors::InvalidArgument(
"Axis should be great than or equal to 0, but received axis is %d.",
axis));
PADDLE_ENFORCE_LT(axis, max_dim,
platform::errors::InvalidArgument(
"Axis should be less than %d, but received axis is %d.",
max_dim, axis));

for (int i = 0; i < x_dims.size(); ++i) {
dst_dims_vec[i + x_axis] =
std::max(dst_dims_vec[i + x_axis], static_cast<int>(x_dims[i]));
}
for (int i = 0; i < y_dims.size(); ++i) {
dst_dims_vec[i + y_axis] =
std::max(dst_dims_vec[i + y_axis], static_cast<int>(y_dims[i]));
}

auto dst_dims = framework::make_ddim(dst_dims_vec);
NpuBroadcast<T>(dev_ctx, x, x_axis, dst_dims, transformed_x);
NpuBroadcast<T>(dev_ctx, y, y_axis, dst_dims, transformed_y);
}

} // namespace operators
} // namespace paddle
Loading