wehub-resource-sync
444 lines
17 KiB
C++
444 lines
17 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 <string>
|
|
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/kernels/funcs/blas/blas.h"
|
|
#include "paddle/phi/kernels/funcs/fc_functor.h"
|
|
#include "paddle/phi/kernels/funcs/jit/kernels.h"
|
|
#include "paddle/phi/kernels/funcs/sequence2batch.h"
|
|
|
|
namespace phi {
|
|
|
|
#define INIT_BASE_DEFINES \
|
|
auto *x = &x_in; \
|
|
auto *h0 = h0_in.get_ptr(); \
|
|
auto *c0 = c0_in.get_ptr(); \
|
|
auto *wx = &weight_x_in; \
|
|
auto *wh = &weight_h_in; \
|
|
auto *bias = &bias_in; \
|
|
auto *hidden_out = hidden; \
|
|
auto *cell_out = cell; \
|
|
auto x_dims = x->dims(); /* T x M*/ \
|
|
auto wh_dims = wh->dims(); /* D x 4D*/ \
|
|
const int M = x_dims[1]; \
|
|
const int D = wh_dims[0]; \
|
|
const int D4 = wh_dims[1]
|
|
|
|
#define INIT_OTHER_DEFINES \
|
|
const T *x_data = x->data<T>(); \
|
|
const T *wx_data = wx->data<T>(); \
|
|
const T *wh_data = wh->data<T>(); \
|
|
/* diagonal weight*/ \
|
|
const T *wp_data = bias->data<T>() + D4; \
|
|
/* for peephole only*/ \
|
|
T *checked_cell_data = nullptr; \
|
|
if (use_peepholes) { \
|
|
/* w_ic * Ct-1, w_fc * Ct-1 ; w_oc * Ct => ih*/ \
|
|
checked_cell_data = dev_ctx.template Alloc<T>(checked_cell); \
|
|
} \
|
|
const phi::jit::lstm_attr_t attr( \
|
|
D, \
|
|
phi::jit::to_kerneltype(gate_activation), \
|
|
phi::jit::to_kerneltype(candidate_activation), \
|
|
phi::jit::to_kerneltype(cell_activation), \
|
|
use_peepholes); \
|
|
phi::jit::lstm_t one_step; \
|
|
one_step.wp = wp_data; \
|
|
one_step.checked = checked_cell_data; \
|
|
auto ComputeC1H1 = \
|
|
phi::jit::KernelFuncs<phi::jit::LSTMC1H1Tuple<T>, CPUPlace>::Cache().At( \
|
|
attr); \
|
|
auto ComputeCtHt = \
|
|
phi::jit::KernelFuncs<phi::jit::LSTMCtHtTuple<T>, CPUPlace>::Cache().At( \
|
|
attr)
|
|
|
|
// Wh GEMM
|
|
#define GEMM_WH_ADDON(bs, prev, out) \
|
|
blas.GEMM(CblasNoTrans, \
|
|
CblasNoTrans, \
|
|
bs, \
|
|
D4, \
|
|
D, \
|
|
static_cast<T>(1), \
|
|
prev, \
|
|
D, \
|
|
wh_data, \
|
|
D4, \
|
|
static_cast<T>(1), \
|
|
out, \
|
|
D4)
|
|
|
|
template <typename T, typename Context>
|
|
void SeqCompute(const Context &dev_ctx,
|
|
const DenseTensor &x_in,
|
|
const DenseTensor &weight_x_in,
|
|
const DenseTensor &weight_h_in,
|
|
const DenseTensor &bias_in,
|
|
const optional<DenseTensor> &h0_in,
|
|
const optional<DenseTensor> &c0_in,
|
|
bool use_peepholes,
|
|
bool is_reverse,
|
|
bool use_seq,
|
|
const std::string &gate_activation,
|
|
const std::string &cell_activation,
|
|
const std::string &candidate_activation,
|
|
float scale_data,
|
|
float shift_data,
|
|
const std::vector<float> &scale_weights,
|
|
bool force_fp32_output,
|
|
DenseTensor *hidden,
|
|
DenseTensor *cell,
|
|
DenseTensor *xx,
|
|
DenseTensor *batched_input,
|
|
DenseTensor *batched_hidden,
|
|
DenseTensor *batched_cell,
|
|
DenseTensor *reordered_h0,
|
|
DenseTensor *reordered_c0,
|
|
DenseTensor *checked_cell) {
|
|
INIT_BASE_DEFINES;
|
|
INIT_OTHER_DEFINES;
|
|
auto x_lod = x->lod();
|
|
const int total_T = static_cast<int>(x_dims[0]);
|
|
const int N = static_cast<int>(x_lod[0].size() - 1);
|
|
const T *h0_data = h0 ? h0->data<T>() : nullptr;
|
|
const T *c0_data = c0 ? c0->data<T>() : nullptr;
|
|
T *xx_data = dev_ctx.template Alloc<T>(xx);
|
|
T *h_out_data = dev_ctx.template Alloc<T>(hidden_out);
|
|
T *c_out_data = dev_ctx.template Alloc<T>(cell_out);
|
|
auto blas = funcs::GetBlas<Context, T>(dev_ctx);
|
|
|
|
funcs::FCFunctor<Context, T> fc;
|
|
fc(dev_ctx, total_T, D4, M, x_data, wx_data, xx_data, bias->data<T>());
|
|
|
|
int xx_offset = D4;
|
|
int gate_offset = D;
|
|
if (is_reverse) {
|
|
const int offset = (total_T - 1) * D;
|
|
xx_data = xx_data + offset * 4;
|
|
h_out_data = h_out_data + offset;
|
|
c_out_data = c_out_data + offset;
|
|
xx_offset = -D4;
|
|
gate_offset = -D;
|
|
}
|
|
|
|
for (int i = 0; i < N; ++i) {
|
|
int bid = is_reverse ? N - 1 - i : i;
|
|
int seq_len = static_cast<int>(x_lod[0][bid + 1] - x_lod[0][bid]);
|
|
const T *prev_c_data = nullptr;
|
|
const T *prev_h_data = nullptr;
|
|
int tstart = 0;
|
|
if (h0_data) {
|
|
prev_h_data = h0_data + bid * D;
|
|
prev_c_data = c0_data + bid * D;
|
|
} else {
|
|
one_step.gates = xx_data;
|
|
one_step.ct = c_out_data;
|
|
one_step.ht = h_out_data;
|
|
ComputeC1H1(&one_step, &attr);
|
|
tstart = 1;
|
|
// move one step
|
|
prev_h_data = h_out_data;
|
|
prev_c_data = c_out_data;
|
|
xx_data = xx_data + xx_offset;
|
|
h_out_data = h_out_data + gate_offset;
|
|
c_out_data = c_out_data + gate_offset;
|
|
}
|
|
for (int step = tstart; step < seq_len; ++step) {
|
|
GEMM_WH_ADDON(1, prev_h_data, xx_data);
|
|
|
|
one_step.gates = xx_data;
|
|
one_step.ct_1 = prev_c_data;
|
|
one_step.ct = c_out_data;
|
|
one_step.ht = h_out_data;
|
|
ComputeCtHt(&one_step, &attr);
|
|
// move one step
|
|
prev_h_data = h_out_data;
|
|
prev_c_data = c_out_data;
|
|
xx_data = xx_data + xx_offset;
|
|
h_out_data = h_out_data + gate_offset;
|
|
c_out_data = c_out_data + gate_offset;
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void BatchCompute(const Context &dev_ctx,
|
|
const DenseTensor &x_in,
|
|
const DenseTensor &weight_x_in,
|
|
const DenseTensor &weight_h_in,
|
|
const DenseTensor &bias_in,
|
|
const optional<DenseTensor> &h0_in,
|
|
const optional<DenseTensor> &c0_in,
|
|
bool use_peepholes,
|
|
bool is_reverse,
|
|
bool use_seq,
|
|
const std::string &gate_activation,
|
|
const std::string &cell_activation,
|
|
const std::string &candidate_activation,
|
|
float scale_data,
|
|
float shift_data,
|
|
const std::vector<float> &scale_weights,
|
|
bool force_fp32_output,
|
|
DenseTensor *hidden,
|
|
DenseTensor *cell,
|
|
DenseTensor *xx,
|
|
DenseTensor *batched_input,
|
|
DenseTensor *batched_hidden,
|
|
DenseTensor *batched_cell,
|
|
DenseTensor *reordered_h0,
|
|
DenseTensor *reordered_c0,
|
|
DenseTensor *checked_cell) {
|
|
INIT_BASE_DEFINES;
|
|
if (x->lod()[0].size() == 2) {
|
|
xx->Resize({x_dims[0], D4});
|
|
SeqCompute<T, Context>(dev_ctx,
|
|
x_in,
|
|
weight_x_in,
|
|
weight_h_in,
|
|
bias_in,
|
|
h0_in,
|
|
c0_in,
|
|
use_peepholes,
|
|
is_reverse,
|
|
use_seq,
|
|
gate_activation,
|
|
cell_activation,
|
|
candidate_activation,
|
|
scale_data,
|
|
shift_data,
|
|
scale_weights,
|
|
force_fp32_output,
|
|
hidden,
|
|
cell,
|
|
xx,
|
|
batched_input,
|
|
batched_hidden,
|
|
batched_cell,
|
|
reordered_h0,
|
|
reordered_c0,
|
|
checked_cell);
|
|
return;
|
|
}
|
|
INIT_OTHER_DEFINES;
|
|
|
|
auto *batched_c_out = batched_cell;
|
|
auto *batched_h_out = batched_hidden;
|
|
T *xx_data = dev_ctx.template Alloc<T>(xx);
|
|
T *batched_input_data = dev_ctx.template Alloc<T>(batched_input);
|
|
T *batched_c_out_data = dev_ctx.template Alloc<T>(batched_c_out);
|
|
T *batched_h_out_data = dev_ctx.template Alloc<T>(batched_h_out);
|
|
dev_ctx.template Alloc<T>(hidden_out);
|
|
dev_ctx.template Alloc<T>(cell_out);
|
|
|
|
funcs::DenseTensor2BatchFunctor<Context, T> to_batch;
|
|
auto blas = funcs::GetBlas<Context, T>(dev_ctx);
|
|
funcs::FCFunctor<Context, T> fc;
|
|
if (M > D4) {
|
|
fc(dev_ctx, x_dims[0], D4, M, x_data, wx_data, xx_data, bias->data<T>());
|
|
to_batch(dev_ctx, *xx, batched_input, true, is_reverse);
|
|
} else {
|
|
to_batch(dev_ctx, *x, xx, true, is_reverse);
|
|
batched_input->set_lod(xx->lod());
|
|
fc(dev_ctx,
|
|
x_dims[0],
|
|
D4,
|
|
M,
|
|
xx_data,
|
|
wx_data,
|
|
batched_input_data,
|
|
bias->data<T>());
|
|
}
|
|
|
|
auto batched_lod = batched_input->lod();
|
|
const auto &seq_order = batched_lod[2];
|
|
const int max_bs = static_cast<int>(seq_order.size());
|
|
reordered_h0->Resize({max_bs, D});
|
|
reordered_c0->Resize({max_bs, D});
|
|
|
|
int tstart = 0;
|
|
T *prev_h_data = nullptr;
|
|
T *prev_c_data = nullptr;
|
|
if (h0) {
|
|
// reorder h0, c0
|
|
T *reordered_h0_data = dev_ctx.template Alloc<T>(reordered_h0);
|
|
T *reordered_c0_data = dev_ctx.template Alloc<T>(reordered_c0);
|
|
const T *h0_data = h0->data<T>();
|
|
const T *c0_data = c0->data<T>();
|
|
prev_h_data = reordered_h0_data;
|
|
prev_c_data = reordered_c0_data;
|
|
size_t sz = D;
|
|
for (int i = 0; i < max_bs; ++i) {
|
|
blas.VCOPY(sz, h0_data + seq_order[i] * D, reordered_h0_data);
|
|
blas.VCOPY(sz, c0_data + seq_order[i] * D, reordered_c0_data);
|
|
reordered_h0_data += D;
|
|
reordered_c0_data += D;
|
|
}
|
|
} else {
|
|
// compute without h0, c0
|
|
T *cur_in_data = batched_input_data;
|
|
T *cur_h_out_data = batched_h_out_data;
|
|
T *cur_c_out_data = batched_c_out_data;
|
|
for (int i = 0; i < max_bs; ++i) {
|
|
one_step.gates = cur_in_data;
|
|
one_step.ct = cur_c_out_data;
|
|
one_step.ht = cur_h_out_data;
|
|
ComputeC1H1(&one_step, &attr);
|
|
|
|
cur_in_data += D4;
|
|
cur_c_out_data += D;
|
|
cur_h_out_data += D;
|
|
}
|
|
tstart = 1;
|
|
prev_h_data = batched_h_out_data;
|
|
prev_c_data = batched_c_out_data;
|
|
}
|
|
|
|
// compute kernel part
|
|
const auto &batch_starts = batched_lod[0];
|
|
const int max_seq_len = static_cast<int>(batch_starts.size() - 1);
|
|
const int offset = tstart * max_bs * D;
|
|
batched_input_data = batched_input_data + offset * 4;
|
|
batched_h_out_data = batched_h_out_data + offset;
|
|
batched_c_out_data = batched_c_out_data + offset;
|
|
for (int step = tstart; step < max_seq_len; ++step) {
|
|
const int cur_bs =
|
|
static_cast<int>(batch_starts[step + 1] - batch_starts[step]);
|
|
GEMM_WH_ADDON(cur_bs, prev_h_data, batched_input_data);
|
|
T *cur_in_data = batched_input_data;
|
|
T *cur_prev_c_data = prev_c_data;
|
|
T *cur_c_out_data = batched_c_out_data;
|
|
T *cur_h_out_data = batched_h_out_data;
|
|
for (int i = 0; i < cur_bs; ++i) {
|
|
one_step.gates = cur_in_data;
|
|
one_step.ct_1 = cur_prev_c_data;
|
|
one_step.ct = cur_c_out_data;
|
|
one_step.ht = cur_h_out_data;
|
|
ComputeCtHt(&one_step, &attr);
|
|
|
|
// move one batch
|
|
cur_in_data += D4;
|
|
cur_prev_c_data += D;
|
|
cur_c_out_data += D;
|
|
cur_h_out_data += D;
|
|
}
|
|
// move one step
|
|
prev_c_data = batched_c_out_data;
|
|
prev_h_data = batched_h_out_data;
|
|
batched_c_out_data = cur_c_out_data;
|
|
batched_h_out_data = cur_h_out_data;
|
|
batched_input_data = cur_in_data;
|
|
}
|
|
|
|
funcs::Batch2DenseTensorFunctor<Context, T> to_seq;
|
|
batched_h_out->set_lod(batched_lod);
|
|
to_seq(dev_ctx, *batched_h_out, hidden_out);
|
|
batched_c_out->set_lod(batched_lod);
|
|
to_seq(dev_ctx, *batched_c_out, cell_out);
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void FusionLSTMKernel(const Context &dev_ctx,
|
|
const DenseTensor &x_in,
|
|
const DenseTensor &weight_x_in,
|
|
const DenseTensor &weight_h_in,
|
|
const DenseTensor &bias_in,
|
|
const optional<DenseTensor> &h0_in,
|
|
const optional<DenseTensor> &c0_in,
|
|
bool use_peepholes,
|
|
bool is_reverse,
|
|
bool use_seq,
|
|
const std::string &gate_activation,
|
|
const std::string &cell_activation,
|
|
const std::string &candidate_activation,
|
|
float scale_data,
|
|
float shift_data,
|
|
const std::vector<float> &scale_weights,
|
|
bool force_fp32_output,
|
|
DenseTensor *hidden,
|
|
DenseTensor *cell,
|
|
DenseTensor *xx,
|
|
DenseTensor *batched_input,
|
|
DenseTensor *batched_hidden,
|
|
DenseTensor *batched_cell,
|
|
DenseTensor *reordered_h0,
|
|
DenseTensor *reordered_c0,
|
|
DenseTensor *checked_cell) {
|
|
if (use_seq) {
|
|
SeqCompute<T, Context>(dev_ctx,
|
|
x_in,
|
|
weight_x_in,
|
|
weight_h_in,
|
|
bias_in,
|
|
h0_in,
|
|
c0_in,
|
|
use_peepholes,
|
|
is_reverse,
|
|
use_seq,
|
|
gate_activation,
|
|
cell_activation,
|
|
candidate_activation,
|
|
scale_data,
|
|
shift_data,
|
|
scale_weights,
|
|
force_fp32_output,
|
|
hidden,
|
|
cell,
|
|
xx,
|
|
batched_input,
|
|
batched_hidden,
|
|
batched_cell,
|
|
reordered_h0,
|
|
reordered_c0,
|
|
checked_cell);
|
|
} else {
|
|
BatchCompute<T, Context>(dev_ctx,
|
|
x_in,
|
|
weight_x_in,
|
|
weight_h_in,
|
|
bias_in,
|
|
h0_in,
|
|
c0_in,
|
|
use_peepholes,
|
|
is_reverse,
|
|
use_seq,
|
|
gate_activation,
|
|
cell_activation,
|
|
candidate_activation,
|
|
scale_data,
|
|
shift_data,
|
|
scale_weights,
|
|
force_fp32_output,
|
|
hidden,
|
|
cell,
|
|
xx,
|
|
batched_input,
|
|
batched_hidden,
|
|
batched_cell,
|
|
reordered_h0,
|
|
reordered_c0,
|
|
checked_cell);
|
|
}
|
|
}
|
|
|
|
#undef GEMM_WH_ADDON
|
|
#undef INIT_OTHER_DEFINES
|
|
#undef INIT_BASE_DEFINES
|
|
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(
|
|
fusion_lstm, CPU, ALL_LAYOUT, phi::FusionLSTMKernel, float, double) {}
|