未验证 提交 a0e3a175 编写于 作者: X xiongkun 提交者: GitHub

[XPU] transfer concat kernel (#45463)

* transfer concat kernel

* test=kunlun

* test=kunlun

* test=kunlun

* test=kunlun
上级 cfd5d40f
/* Copyright (c) 2016 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. */
#ifdef PADDLE_WITH_XPU
#include <memory>
#include <string>
#include <vector>
#include "paddle/fluid/operators/concat_op.h"
#include "paddle/fluid/platform/device/xpu/xpu_header.h"
#include "paddle/phi/core/lod_utils.h"
namespace paddle {
namespace operators {
using Tensor = framework::Tensor;
template <typename DeviceContext, typename T>
class ConcatXPUKernel : public framework::OpKernel<T> {
using XPUType = typename XPUTypeTrait<T>::Type;
public:
void Compute(const framework::ExecutionContext& ctx) const override {
auto ins = ctx.MultiInput<framework::LoDTensor>("X");
framework::LoDTensor* out = ctx.Output<framework::LoDTensor>("Out");
int axis = ctx.Attr<int>("axis");
PADDLE_ENFORCE_NE(
ins[0],
nullptr,
platform::errors::InvalidArgument("The input should not be null."));
PADDLE_ENFORCE_NE(ctx.HasInput("AxisTensor"),
true,
platform::errors::InvalidArgument(
"XPU donot surpport AxisTensor for now"));
axis = ComputeAxis(static_cast<int64_t>(axis),
static_cast<int64_t>(ins[0]->dims().size()));
PADDLE_ENFORCE_GE(axis,
0,
platform::errors::InvalidArgument(
"concat: axis should be larger than or "
"equal to 0, but received axis is %d.",
axis));
PADDLE_ENFORCE_LT(axis,
ins[0]->dims().size(),
platform::errors::InvalidArgument(
"concat: axis should be less than ins[0]->dims()!"
"But received axis is %d, while ins[0]->dims()"
"size is %d.",
axis,
ins[0]->dims().size()));
// If axis is 0, the lod of the output is not the same as inputs.
if (axis == 0 && ins[0]->lod().size() > 0) {
size_t lod_size_0 = ins[0]->lod().size();
size_t lod_size = lod_size_0;
for (size_t i = 1; i < ins.size(); ++i) {
if (ins[i]->lod().size() > 0) {
PADDLE_ENFORCE_EQ(
ins[i]->lod().size(),
lod_size_0,
platform::errors::Unimplemented(
"The lod level of all input LoDTensors should be same. "
"Maybe different lod level of input LoDTensors can concat,"
"it is not supported currently. The lod level of %dth input "
"is %d and first input is %d.",
i,
ins[i]->lod().size(),
lod_size_0));
} else {
lod_size = 0;
break;
}
}
if (lod_size) {
auto* out_lod = out->mutable_lod();
for (size_t i = 1; i < ins.size(); ++i) {
auto in_lod = phi::ConvertToLengthBasedLoD(ins[i]->lod());
phi::AppendLoD(out_lod, in_lod);
}
}
}
auto place = ctx.GetPlace();
out->mutable_data<T>(place);
std::vector<std::vector<int>> xdims_list;
std::vector<const XPUType*> ptrs;
for (unsigned int i = 0; i < ins.size(); ++i) {
if (ins[i] && ins[i]->numel() > 0) {
ptrs.push_back(reinterpret_cast<const XPUType*>(ins[i]->data<T>()));
int size = ins[i]->dims().size();
std::vector<int> tmp_dims(size);
for (int j = 0; j < size; ++j) {
tmp_dims[j] = ins[i]->dims()[j];
}
xdims_list.push_back(tmp_dims);
}
}
PADDLE_ENFORCE_GT(
xdims_list.size(),
0,
platform::errors::InvalidArgument("No tensor need concat"));
auto& dev_ctx = ctx.template device_context<DeviceContext>();
int r = xpu::concat<XPUType>(dev_ctx.x_context(),
ptrs,
reinterpret_cast<XPUType*>(out->data<T>()),
xdims_list,
axis);
PADDLE_ENFORCE_EQ(r,
XPU_SUCCESS,
platform::errors::External(
"XPU concat kernel return wrong value[%d %s]",
r,
XPUAPIErrorMsg[r]));
}
};
template <typename DeviceContext, typename T>
class ConcatGradXPUKernel : public framework::OpKernel<T> {
using XPUType = typename XPUTypeTrait<T>::Type;
public:
void Compute(const framework::ExecutionContext& ctx) const {
auto* out_grad =
ctx.Input<framework::Tensor>(framework::GradVarName("Out"));
auto ins = ctx.MultiInput<framework::LoDTensor>("X");
auto out_var_names = ctx.OutputNames(framework::GradVarName("X"));
auto outs =
ctx.MultiOutput<framework::LoDTensor>(framework::GradVarName("X"));
{
auto dx = outs;
auto x = ins;
for (size_t i = 0; i < dx.size(); ++i) {
if (dx[i] != nullptr) {
dx[i]->set_lod(x[i]->lod());
}
}
}
PADDLE_ENFORCE_NE(
ins[0],
nullptr,
platform::errors::InvalidArgument("The input should not be null."));
auto axis = ctx.Attr<int>("axis");
if (ctx.HasInput("AxisTensor")) {
auto* axis_tensor = ctx.Input<framework::Tensor>("AxisTensor");
axis = GetDataFromTensor<int>(axis_tensor)[0];
}
axis = ComputeAxis(static_cast<int64_t>(axis),
static_cast<int64_t>(ins[0]->dims().size()));
// get output tensor that the name is not kEmptyVarName
std::vector<XPUType*> ptrs(outs.size());
for (size_t j = 0; j < outs.size(); ++j) {
if (out_var_names[j] != framework::kEmptyVarName &&
outs[j]->numel() != 0UL) {
outs[j]->mutable_data<T>(ctx.GetPlace());
ptrs[j] = reinterpret_cast<XPUType*>(outs[j]->data<T>());
} else {
ptrs[j] = nullptr;
}
}
PADDLE_ENFORCE_GE(axis,
0,
platform::errors::InvalidArgument(
"concat_grad: axis should be larger than or "
"equal to 0, but received axis is %d.",
axis));
PADDLE_ENFORCE_LT(
axis,
out_grad->dims().size(),
platform::errors::InvalidArgument(
"concat_grad: axis should be less than ins[0]->dims()!"
"But received axis is %d, while ins[0]->dims()"
"size is %d.",
axis,
out_grad->dims().size()));
auto input_dims = ins[0]->dims();
std::vector<int> split_list(ins.size());
std::vector<int> xdims_list(input_dims.size());
int total_length = 0;
for (size_t i = 0; i < ins.size(); ++i) {
split_list[i] = ins[i]->dims()[axis];
total_length += ins[i]->dims()[axis];
}
for (int i = 0; i < input_dims.size(); ++i) {
if (i == axis) {
continue;
}
xdims_list[i] = input_dims[i];
}
xdims_list[axis] = total_length;
auto& dev_ctx = ctx.template device_context<DeviceContext>();
int r = xpu::split<XPUType>(
dev_ctx.x_context(),
reinterpret_cast<const XPUType*>(out_grad->data<T>()),
ptrs,
xdims_list,
split_list,
axis);
PADDLE_ENFORCE_EQ(
r,
XPU_SUCCESS,
platform::errors::External(
"XPU API return wrong value[%d], please check whether "
"Baidu Kunlun Card is properly installed.",
r));
}
};
} // namespace operators
} // namespace paddle
namespace ops = paddle::operators;
REGISTER_OP_XPU_KERNEL(
concat,
ops::ConcatXPUKernel<paddle::platform::XPUDeviceContext, float>,
ops::ConcatXPUKernel<paddle::platform::XPUDeviceContext,
paddle::platform::float16>);
REGISTER_OP_XPU_KERNEL(
concat_grad,
ops::ConcatGradXPUKernel<paddle::platform::XPUDeviceContext, float>,
ops::ConcatGradXPUKernel<paddle::platform::XPUDeviceContext,
paddle::platform::float16>);
#endif
// Copyright (c) 2022 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/kernels/concat_grad_kernel.h"
#include "paddle/phi/backends/xpu/enforce_xpu.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/funcs/axis_utils.h"
#include "paddle/phi/kernels/funcs/concat_funcs.h"
namespace phi {
template <typename T, typename Context>
void ConcatGradKernel(const Context& dev_ctx,
const std::vector<const DenseTensor*>& x,
const DenseTensor& out_grad,
const Scalar& axis_scalar,
std::vector<DenseTensor*> x_grad) {
using XPUType = typename XPUTypeTrait<T>::Type;
auto outs = x_grad;
{
auto dx = outs;
for (size_t i = 0; i < dx.size(); ++i) {
if (dx[i] != nullptr) {
dx[i]->set_lod(x[i]->lod());
}
}
}
PADDLE_ENFORCE_NE(
x[0],
nullptr,
phi::errors::InvalidArgument("The input should not be null."));
auto axis = axis_scalar.to<int>();
axis = phi::funcs::ComputeAxis(static_cast<int64_t>(axis),
static_cast<int64_t>(x[0]->dims().size()));
// get output tensor that the name is not kEmptyVarName
std::vector<XPUType*> ptrs(outs.size());
for (size_t j = 0; j < outs.size(); ++j) {
if (outs[j] && outs[j]->numel() != 0UL) {
dev_ctx.template Alloc<T>(outs[j]);
ptrs[j] = reinterpret_cast<XPUType*>(outs[j]->data<T>());
} else {
ptrs[j] = nullptr;
}
}
PADDLE_ENFORCE_GE(
axis,
0,
phi::errors::InvalidArgument("concat_grad: axis should be larger than or "
"equal to 0, but received axis is %d.",
axis));
PADDLE_ENFORCE_LT(axis,
out_grad.dims().size(),
phi::errors::InvalidArgument(
"concat_grad: axis should be less than x[0]->dims()!"
"But received axis is %d, while x[0]->dims()"
"size is %d.",
axis,
out_grad.dims().size()));
auto input_dims = x[0]->dims();
std::vector<int> split_list(x.size());
std::vector<int> xdims_list(input_dims.size());
int total_length = 0;
for (size_t i = 0; i < x.size(); ++i) {
split_list[i] = x[i]->dims()[axis];
total_length += x[i]->dims()[axis];
}
for (int i = 0; i < input_dims.size(); ++i) {
if (i == axis) {
continue;
}
xdims_list[i] = input_dims[i];
}
xdims_list[axis] = total_length;
int r =
xpu::split<XPUType>(dev_ctx.x_context(),
reinterpret_cast<const XPUType*>(out_grad.data<T>()),
ptrs,
xdims_list,
split_list,
axis);
PADDLE_ENFORCE_XDNN_SUCCESS(r, "concat_grad");
}
} // namespace phi
PD_REGISTER_KERNEL(concat_grad,
XPU,
ALL_LAYOUT,
phi::ConcatGradKernel,
float,
phi::dtype::float16) {}
// Copyright (c) 2022 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/kernels/concat_kernel.h"
#include "paddle/phi/backends/xpu/enforce_xpu.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/core/lod_utils.h"
#include "paddle/phi/kernels/funcs/axis_utils.h"
#include "paddle/phi/kernels/funcs/concat_funcs.h"
namespace phi {
template <typename T, typename Context>
void ConcatKernel(const Context& dev_ctx,
const std::vector<const DenseTensor*>& x,
const Scalar& axis_scalar,
DenseTensor* out) {
using XPUType = typename XPUTypeTrait<T>::Type;
int64_t axis = axis_scalar.to<int64_t>();
PADDLE_ENFORCE_NE(
x[0],
nullptr,
phi::errors::InvalidArgument("The input should not be null."));
axis = phi::funcs::ComputeAxis(axis, x[0]->dims().size());
PADDLE_ENFORCE_GE(
axis,
0,
phi::errors::InvalidArgument("concat: axis should be larger than or "
"equal to 0, but received axis is %d.",
axis));
PADDLE_ENFORCE_LT(axis,
x[0]->dims().size(),
phi::errors::InvalidArgument(
"concat: axis should be less than x[0]->dims()!"
"But received axis is %d, while x[0]->dims()"
"size is %d.",
axis,
x[0]->dims().size()));
// If axis is 0, the lod of the output is not the same as inputs.
if (axis == 0 && x[0]->lod().size() > 0) {
size_t lod_size_0 = x[0]->lod().size();
size_t lod_size = lod_size_0;
for (size_t i = 1; i < x.size(); ++i) {
if (x[i]->lod().size() > 0) {
PADDLE_ENFORCE_EQ(
x[i]->lod().size(),
lod_size_0,
phi::errors::Unimplemented(
"The lod level of all input LoDTensors should be same. "
"Maybe different lod level of input LoDTensors can concat,"
"it is not supported currently. The lod level of %dth input "
"is %d and first input is %d.",
i,
x[i]->lod().size(),
lod_size_0));
} else {
lod_size = 0;
break;
}
}
if (lod_size) {
auto* out_lod = out->mutable_lod();
for (size_t i = 1; i < x.size(); ++i) {
auto in_lod = phi::ConvertToLengthBasedLoD(x[i]->lod());
phi::AppendLoD(out_lod, in_lod);
}
}
}
dev_ctx.template Alloc<T>(out);
std::vector<std::vector<int>> xdims_list;
std::vector<const XPUType*> ptrs;
for (unsigned int i = 0; i < x.size(); ++i) {
if (x[i] && x[i]->numel() > 0) {
ptrs.push_back(reinterpret_cast<const XPUType*>(x[i]->data<T>()));
int size = x[i]->dims().size();
std::vector<int> tmp_dims(size);
for (int j = 0; j < size; ++j) {
tmp_dims[j] = x[i]->dims()[j];
}
xdims_list.push_back(tmp_dims);
}
}
PADDLE_ENFORCE_GT(xdims_list.size(),
0,
phi::errors::InvalidArgument("No tensor need concat"));
int r = xpu::concat<XPUType>(dev_ctx.x_context(),
ptrs,
reinterpret_cast<XPUType*>(out->data<T>()),
xdims_list,
axis);
PADDLE_ENFORCE_XDNN_SUCCESS(r, "concat");
}
} // namespace phi
PD_REGISTER_KERNEL(
concat, XPU, ALL_LAYOUT, phi::ConcatKernel, float, phi::dtype::float16) {}
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册