wehub-resource-sync
397 lines
14 KiB
C++
397 lines
14 KiB
C++
// Copyright (c) 2023 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 <cstring> // for memcpy
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
#include "paddle/common/errors.h"
|
|
#include "paddle/phi/core/enforce.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/core/tensor_utils.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::fusion {
|
|
|
|
#define INIT_BASE_DEFINES \
|
|
auto x_lod = x.lod(); \
|
|
auto x_dims = x.dims(); /* T x M*/ \
|
|
auto x_mat_dims = (x_dims.size() == 3 && x_dims[1] == 1) \
|
|
? common::flatten_to_2d(x_dims, 1) \
|
|
: x_dims; \
|
|
auto wh_dims = weight_h.dims(); /* D x 3D*/ \
|
|
const int total_T = x_mat_dims[0]; \
|
|
const int D3 = wh_dims[1]
|
|
|
|
#define INIT_OTHER_DEFINES \
|
|
const int M = x_mat_dims[1]; \
|
|
const int D = wh_dims[0]; \
|
|
const int D2 = D * 2; \
|
|
const phi::jit::gru_attr_t attr(D, \
|
|
phi::jit::to_kerneltype(gate_activation), \
|
|
phi::jit::to_kerneltype(activation)); \
|
|
phi::jit::gru_t one_step; \
|
|
auto ComputeH1 = \
|
|
phi::jit::KernelFuncs<phi::jit::GRUH1Tuple<T>, CPUPlace>::Cache().At( \
|
|
attr); \
|
|
auto ComputeHtPart1 = \
|
|
phi::jit::KernelFuncs<phi::jit::GRUHtPart1Tuple<T>, CPUPlace>::Cache() \
|
|
.At(attr); \
|
|
auto ComputeHtPart2 = \
|
|
phi::jit::KernelFuncs<phi::jit::GRUHtPart2Tuple<T>, CPUPlace>::Cache() \
|
|
.At(attr); \
|
|
const T* x_data = x.data<T>(); \
|
|
const T* wx_data = weight_x.data<T>(); \
|
|
const T* wh_data = weight_h.data<T>(); \
|
|
T* xx_data = dev_ctx.template Alloc<T>(xx)
|
|
|
|
template <typename T, typename Context>
|
|
void SeqCompute(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const optional<DenseTensor>& h0,
|
|
const DenseTensor& weight_x,
|
|
const DenseTensor& weight_h,
|
|
const optional<DenseTensor>& bias,
|
|
const std::string& activation,
|
|
const std::string& gate_activation,
|
|
const bool is_reverse,
|
|
const bool use_seq,
|
|
DenseTensor* reordered_h0,
|
|
DenseTensor* xx,
|
|
DenseTensor* batched_input,
|
|
DenseTensor* batched_out,
|
|
DenseTensor* hidden) {
|
|
INIT_BASE_DEFINES;
|
|
INIT_OTHER_DEFINES;
|
|
const int N = static_cast<int>(x_lod[0].size() - 1);
|
|
const T* h0_data = h0 ? h0->data<T>() : nullptr;
|
|
const T* wh_state_data = wh_data + D * D2;
|
|
T* hidden_out_data = dev_ctx.template Alloc<T>(hidden);
|
|
|
|
auto blas = funcs::GetBlas<Context, T>(dev_ctx);
|
|
|
|
funcs::FCFunctor<Context, T> fc;
|
|
fc(dev_ctx,
|
|
total_T,
|
|
D3,
|
|
M,
|
|
x_data,
|
|
wx_data,
|
|
xx_data,
|
|
bias ? bias->data<T>() : nullptr);
|
|
|
|
int xx_offset = D3;
|
|
int gate_offset = D;
|
|
if (is_reverse) {
|
|
const int offset = (total_T - 1) * D;
|
|
xx_data = xx_data + offset * 3;
|
|
hidden_out_data = hidden_out_data + offset;
|
|
xx_offset = -D3;
|
|
gate_offset = -D;
|
|
}
|
|
auto move_step = [&]() {
|
|
xx_data = xx_data + xx_offset;
|
|
hidden_out_data = hidden_out_data + gate_offset;
|
|
};
|
|
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_hidden_data = nullptr;
|
|
int tstart = 0;
|
|
if (h0_data) {
|
|
prev_hidden_data = h0_data + bid * D;
|
|
} else {
|
|
one_step.gates = xx_data;
|
|
one_step.ht = hidden_out_data;
|
|
ComputeH1(&one_step, &attr);
|
|
prev_hidden_data = hidden_out_data;
|
|
tstart = 1;
|
|
move_step();
|
|
}
|
|
for (int step = tstart; step < seq_len; ++step) {
|
|
// gemm prev * (Wu + Wr)
|
|
blas.GEMM(CblasNoTrans,
|
|
CblasNoTrans,
|
|
1,
|
|
D2,
|
|
D,
|
|
static_cast<T>(1),
|
|
prev_hidden_data,
|
|
D,
|
|
wh_data,
|
|
D2,
|
|
static_cast<T>(1),
|
|
xx_data,
|
|
D3);
|
|
one_step.gates = xx_data;
|
|
one_step.ht_1 = prev_hidden_data;
|
|
one_step.ht = hidden_out_data;
|
|
ComputeHtPart1(&one_step, &attr);
|
|
// gemm rt * Ws
|
|
blas.GEMM(CblasNoTrans,
|
|
CblasNoTrans,
|
|
1,
|
|
D,
|
|
D,
|
|
static_cast<T>(1),
|
|
hidden_out_data,
|
|
D,
|
|
wh_state_data,
|
|
D,
|
|
static_cast<T>(1),
|
|
xx_data + D2,
|
|
D3);
|
|
ComputeHtPart2(&one_step, &attr);
|
|
// save prev
|
|
prev_hidden_data = hidden_out_data;
|
|
move_step();
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void BatchCompute(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const optional<DenseTensor>& h0,
|
|
const DenseTensor& weight_x,
|
|
const DenseTensor& weight_h,
|
|
const optional<DenseTensor>& bias,
|
|
const std::string& activation,
|
|
const std::string& gate_activation,
|
|
const bool is_reverse,
|
|
const bool use_seq,
|
|
DenseTensor* reordered_h0,
|
|
DenseTensor* xx,
|
|
DenseTensor* batched_input,
|
|
DenseTensor* batched_out,
|
|
DenseTensor* hidden) {
|
|
INIT_BASE_DEFINES;
|
|
if (x_lod[0].size() == 2) {
|
|
xx->Resize({total_T, D3});
|
|
SeqCompute<T, Context>(dev_ctx,
|
|
x,
|
|
h0,
|
|
weight_x,
|
|
weight_h,
|
|
bias,
|
|
activation,
|
|
gate_activation,
|
|
is_reverse,
|
|
use_seq,
|
|
reordered_h0,
|
|
xx,
|
|
batched_input,
|
|
batched_out,
|
|
hidden);
|
|
return;
|
|
}
|
|
INIT_OTHER_DEFINES;
|
|
T* batched_input_data = dev_ctx.template Alloc<T>(batched_input);
|
|
T* batched_out_data = dev_ctx.template Alloc<T>(batched_out);
|
|
dev_ctx.template Alloc<T>(hidden);
|
|
auto blas = funcs::GetBlas<Context, T>(dev_ctx);
|
|
funcs::DenseTensor2BatchFunctor<Context, T> to_batch;
|
|
|
|
funcs::FCFunctor<Context, T> fc;
|
|
if (M > D3) {
|
|
fc(dev_ctx,
|
|
total_T,
|
|
D3,
|
|
M,
|
|
x_data,
|
|
wx_data,
|
|
xx_data,
|
|
bias ? bias->data<T>() : nullptr);
|
|
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,
|
|
total_T,
|
|
D3,
|
|
M,
|
|
xx_data,
|
|
wx_data,
|
|
batched_input_data,
|
|
bias ? bias->data<T>() : nullptr);
|
|
}
|
|
|
|
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});
|
|
|
|
int tstart = 0;
|
|
T* prev_hidden_data = nullptr;
|
|
if (h0) {
|
|
// reorder h0
|
|
T* reordered_h0_data = dev_ctx.template Alloc<T>(reordered_h0);
|
|
const T* h0_data = h0->data<T>();
|
|
prev_hidden_data = reordered_h0_data;
|
|
size_t sz = sizeof(T) * D;
|
|
for (int i = 0; i < max_bs; ++i) {
|
|
std::memcpy(reordered_h0_data, h0_data + seq_order[i] * D, sz);
|
|
reordered_h0_data += D;
|
|
}
|
|
} else {
|
|
// compute without h0
|
|
T* cur_in_data = batched_input_data;
|
|
T* cur_out_data = batched_out_data;
|
|
// W: {W_update, W_reset; W_state}
|
|
for (int i = 0; i < max_bs; ++i) {
|
|
one_step.gates = cur_in_data;
|
|
one_step.ht = cur_out_data;
|
|
ComputeH1(&one_step, &attr);
|
|
// add offset
|
|
cur_in_data += D3;
|
|
cur_out_data += D;
|
|
}
|
|
tstart = 1;
|
|
prev_hidden_data = batched_out_data;
|
|
}
|
|
// Then start from next
|
|
const T* wh_state_data = wh_data + D * D2;
|
|
const auto& batch_starts = batched_lod[0];
|
|
const int max_seq_len = static_cast<int>(batch_starts.size() - 1);
|
|
batched_input_data = batched_input_data + tstart * max_bs * D3;
|
|
batched_out_data = batched_out_data + tstart * max_bs * D;
|
|
for (int step = tstart; step < max_seq_len; ++step) {
|
|
const int cur_bs =
|
|
static_cast<int>(batch_starts[step + 1] - batch_starts[step]);
|
|
// gemm prev * (Wu + Wr)
|
|
blas.GEMM(CblasNoTrans,
|
|
CblasNoTrans,
|
|
cur_bs,
|
|
D2,
|
|
D,
|
|
static_cast<T>(1),
|
|
prev_hidden_data,
|
|
D,
|
|
wh_data,
|
|
D2,
|
|
static_cast<T>(1),
|
|
batched_input_data,
|
|
D3);
|
|
|
|
T* cur_batched_data = batched_input_data;
|
|
T* cur_out_data = batched_out_data;
|
|
T* cur_prev_hidden_data = prev_hidden_data;
|
|
for (int i = 0; i < cur_bs; ++i) {
|
|
one_step.gates = cur_batched_data;
|
|
one_step.ht_1 = cur_prev_hidden_data;
|
|
one_step.ht = cur_out_data;
|
|
ComputeHtPart1(&one_step, &attr);
|
|
|
|
cur_batched_data += D3;
|
|
cur_prev_hidden_data += D;
|
|
cur_out_data += D;
|
|
}
|
|
|
|
cur_batched_data = batched_input_data;
|
|
cur_out_data = batched_out_data;
|
|
blas.GEMM(CblasNoTrans,
|
|
CblasNoTrans,
|
|
cur_bs,
|
|
D,
|
|
D,
|
|
static_cast<T>(1),
|
|
cur_out_data,
|
|
D,
|
|
wh_state_data,
|
|
D,
|
|
static_cast<T>(1),
|
|
cur_batched_data + D2,
|
|
D3);
|
|
|
|
cur_prev_hidden_data = prev_hidden_data;
|
|
for (int i = 0; i < cur_bs; ++i) {
|
|
one_step.gates = cur_batched_data;
|
|
one_step.ht_1 = cur_prev_hidden_data;
|
|
one_step.ht = cur_out_data;
|
|
ComputeHtPart2(&one_step, &attr);
|
|
cur_batched_data += D3;
|
|
cur_prev_hidden_data += D;
|
|
cur_out_data += D;
|
|
}
|
|
prev_hidden_data = batched_out_data;
|
|
batched_out_data = cur_out_data;
|
|
batched_input_data = cur_batched_data;
|
|
}
|
|
|
|
funcs::Batch2DenseTensorFunctor<Context, T> to_seq;
|
|
batched_out->set_lod(batched_lod);
|
|
to_seq(dev_ctx, *batched_out, hidden);
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void FusionGRUKernel(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const optional<DenseTensor>& h0,
|
|
const DenseTensor& weight_x,
|
|
const DenseTensor& weight_h,
|
|
const optional<DenseTensor>& bias,
|
|
const std::string& activation,
|
|
const std::string& gate_activation,
|
|
const bool is_reverse,
|
|
const bool use_seq,
|
|
const bool origin_mode,
|
|
const bool force_fp32_output,
|
|
DenseTensor* reordered_h0,
|
|
DenseTensor* xx,
|
|
DenseTensor* batched_input,
|
|
DenseTensor* batched_out,
|
|
DenseTensor* hidden) {
|
|
if (use_seq) {
|
|
SeqCompute<T, Context>(dev_ctx,
|
|
x,
|
|
h0,
|
|
weight_x,
|
|
weight_h,
|
|
bias,
|
|
activation,
|
|
gate_activation,
|
|
is_reverse,
|
|
use_seq,
|
|
reordered_h0,
|
|
xx,
|
|
batched_input,
|
|
batched_out,
|
|
hidden);
|
|
} else {
|
|
BatchCompute<T, Context>(dev_ctx,
|
|
x,
|
|
h0,
|
|
weight_x,
|
|
weight_h,
|
|
bias,
|
|
activation,
|
|
gate_activation,
|
|
is_reverse,
|
|
use_seq,
|
|
reordered_h0,
|
|
xx,
|
|
batched_input,
|
|
batched_out,
|
|
hidden);
|
|
}
|
|
}
|
|
|
|
} // namespace phi::fusion
|
|
|
|
PD_REGISTER_KERNEL(
|
|
fusion_gru, CPU, ALL_LAYOUT, phi::fusion::FusionGRUKernel, float, double) {}
|