229 lines
9.0 KiB
C++
229 lines
9.0 KiB
C++
// Copyright (c) 2024 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/core/device_context.h"
|
|
#include "paddle/phi/kernels/impl/sequence_conv_kernel_impl.h"
|
|
|
|
namespace phi {
|
|
template <typename T, typename Context>
|
|
void SequenceConvGradXPUKernel(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const DenseTensor& padding_data,
|
|
const DenseTensor& filter,
|
|
const DenseTensor& out_grad,
|
|
int context_length,
|
|
bool padding_trainable,
|
|
int context_start,
|
|
int context_stride,
|
|
DenseTensor* x_grad,
|
|
DenseTensor* padding_data_grad,
|
|
DenseTensor* filter_grad) {
|
|
auto* in_g = x_grad;
|
|
auto* out_g = &out_grad;
|
|
auto* filter_g = filter_grad;
|
|
auto* in = &x;
|
|
auto* filter_p = &filter;
|
|
|
|
PADDLE_ENFORCE_EQ(
|
|
in->lod().empty(),
|
|
false,
|
|
common::errors::InvalidArgument("Input(X) DenseTensor of SequenceConvOp "
|
|
"does not contain LoD information."));
|
|
PADDLE_ENFORCE_EQ(
|
|
in->lod().size(),
|
|
1UL,
|
|
common::errors::InvalidArgument(
|
|
"Only support input sequence with lod level equal to 1 at "
|
|
"present. But received: lod level %u.",
|
|
in->lod().size()));
|
|
|
|
PADDLE_ENFORCE_EQ(
|
|
padding_trainable,
|
|
false,
|
|
common::errors::InvalidArgument("Only support padding_trainable "
|
|
"equal false."));
|
|
|
|
int up_pad = std::max(0, -context_start);
|
|
int down_pad = std::max(0, context_start + context_length - 1);
|
|
PADDLE_ENFORCE_EQ(
|
|
up_pad,
|
|
2,
|
|
common::errors::InvalidArgument("Only support up_pad equal 2."));
|
|
PADDLE_ENFORCE_EQ(
|
|
down_pad,
|
|
2,
|
|
common::errors::InvalidArgument("Only support down_pad equal 2."));
|
|
|
|
auto lod_level_0 = in->lod()[0];
|
|
int lod_size = lod_level_0.size();
|
|
PADDLE_ENFORCE_LE(
|
|
lod_size,
|
|
257,
|
|
common::errors::InvalidArgument("Only support batch size <= 256."));
|
|
|
|
std::vector<int> cpu_lodx(lod_size);
|
|
for (int i = 0; i < lod_size; i++) {
|
|
cpu_lodx[i] = lod_level_0[i];
|
|
}
|
|
xpu::VectorParam<int> lodx = {
|
|
cpu_lodx.data(), static_cast<int64_t>(cpu_lodx.size()), nullptr};
|
|
|
|
auto* xpu_context = dev_ctx.x_context();
|
|
int64_t sequence_width = in->dims()[1];
|
|
DDim col_shape = {in->dims()[0], context_length * sequence_width};
|
|
xpu::ctx_guard RAII_GUARD(xpu_context);
|
|
int64_t col_numel = col_shape[0] * col_shape[1];
|
|
T* col_data = RAII_GUARD.alloc_l3_or_gm<T>(col_numel);
|
|
PADDLE_ENFORCE_NOT_NULL(col_data,
|
|
common::errors::Fatal("XPU memory is not enough"));
|
|
|
|
if (in_g || filter_g) {
|
|
bool trans_a = false;
|
|
bool trans_b = true;
|
|
int64_t m = out_g->dims()[0];
|
|
int64_t k = out_g->dims()[1];
|
|
int64_t n = filter_p->dims()[0];
|
|
int64_t k1 = filter_p->dims()[1];
|
|
PADDLE_ENFORCE_EQ(k,
|
|
k1,
|
|
common::errors::InvalidArgument(
|
|
"The shape of FC in SequenceConvGradOp is invalid."
|
|
"The k of matrix A is %d, k1 of matrix B is %d."
|
|
"But expect k == k1",
|
|
k,
|
|
k1));
|
|
int64_t lda = (!trans_a) ? k : m;
|
|
int64_t ldb = (!trans_b) ? n : k;
|
|
int64_t ldc = n;
|
|
T alpha = static_cast<T>(1.0);
|
|
T beta = static_cast<T>(0.0);
|
|
const T* data_a = out_g->data<T>();
|
|
const T* data_b = filter_p->data<T>();
|
|
T* data_c = col_data;
|
|
|
|
int r = xpu::fc_fusion<T, T, T, int32_t>(xpu_context,
|
|
data_a,
|
|
data_b,
|
|
data_c,
|
|
m,
|
|
n,
|
|
k,
|
|
trans_a,
|
|
trans_b,
|
|
nullptr,
|
|
nullptr,
|
|
nullptr,
|
|
lda,
|
|
ldb,
|
|
ldc,
|
|
alpha,
|
|
beta,
|
|
nullptr,
|
|
xpu::Activation_t::LINEAR);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "fc_fusion");
|
|
}
|
|
|
|
if (in_g) {
|
|
PADDLE_ENFORCE_LT(
|
|
sequence_width,
|
|
512,
|
|
common::errors::InvalidArgument("Only support sequence_width < 512."));
|
|
|
|
dev_ctx.template Alloc<T>(in_g);
|
|
in_g->set_lod(in->lod());
|
|
|
|
int r = xpu::sequence_context_projection_grad<T, int>(xpu_context,
|
|
in_g->data<T>(),
|
|
col_data,
|
|
nullptr,
|
|
lodx,
|
|
sequence_width,
|
|
context_start,
|
|
context_length,
|
|
context_stride,
|
|
{2, 2});
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "sequence_context_projection_grad");
|
|
}
|
|
|
|
if (filter_g) {
|
|
dev_ctx.template Alloc<T>(filter_g);
|
|
|
|
int r = xpu::sequence_context_projection<T, int>(xpu_context,
|
|
in->data<T>(),
|
|
col_data,
|
|
nullptr,
|
|
lodx,
|
|
sequence_width,
|
|
context_start,
|
|
context_length,
|
|
context_stride,
|
|
{2, 2});
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "sequence_context_projection");
|
|
|
|
bool trans_a = true;
|
|
bool trans_b = false;
|
|
int64_t k = col_shape[0];
|
|
int64_t m = col_shape[1];
|
|
int64_t k1 = out_g->dims()[0];
|
|
int64_t n = out_g->dims()[1];
|
|
PADDLE_ENFORCE_EQ(k,
|
|
k1,
|
|
common::errors::InvalidArgument(
|
|
"The shape of FC in SequenceConvGradOp is invalid."
|
|
"The k of matrix A is %d, k1 of matrix B is %d."
|
|
"But expect k == k1",
|
|
k,
|
|
k1));
|
|
int64_t lda = (!trans_a) ? k : m;
|
|
int64_t ldb = (!trans_b) ? n : k;
|
|
int64_t ldc = n;
|
|
T alpha = static_cast<T>(1.0);
|
|
T beta = static_cast<T>(0.0);
|
|
const T* data_a = col_data;
|
|
const T* data_b = out_g->data<T>();
|
|
T* data_c = filter_g->data<T>();
|
|
|
|
r = xpu::fc_fusion<T, T, T, int32_t>(xpu_context,
|
|
data_a,
|
|
data_b,
|
|
data_c,
|
|
m,
|
|
n,
|
|
k,
|
|
trans_a,
|
|
trans_b,
|
|
nullptr,
|
|
nullptr,
|
|
nullptr,
|
|
lda,
|
|
ldb,
|
|
ldc,
|
|
alpha,
|
|
beta,
|
|
nullptr,
|
|
xpu::Activation_t::LINEAR);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "fc_fusion");
|
|
if (xpu_context->xpu_stream != nullptr) {
|
|
xpu_wait(xpu_context->xpu_stream);
|
|
}
|
|
}
|
|
}
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(sequence_conv_grad,
|
|
XPU,
|
|
ALL_LAYOUT,
|
|
phi::SequenceConvGradXPUKernel,
|
|
float) {}
|