wehub-resource-sync
123 lines
4.6 KiB
Plaintext
123 lines
4.6 KiB
Plaintext
// 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/kernels/gpu/gru_kernel.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/kernels/impl/gru_kernel_impl.h"
|
|
|
|
namespace phi {
|
|
|
|
template <typename T, typename Context>
|
|
void GRUKernel(const Context &dev_ctx,
|
|
const DenseTensor &input,
|
|
const optional<DenseTensor> &h0,
|
|
const DenseTensor &weight,
|
|
const optional<DenseTensor> &bias,
|
|
const std::string &activation,
|
|
const std::string &gate_activation,
|
|
bool is_reverse,
|
|
bool origin_mode,
|
|
bool is_test,
|
|
DenseTensor *param_batch_gate,
|
|
DenseTensor *param_batch_reset_hidden_prev,
|
|
DenseTensor *param_batch_hidden,
|
|
DenseTensor *hidden) {
|
|
const T *weight_data = weight.data<T>();
|
|
dev_ctx.template Alloc<T>(hidden);
|
|
|
|
auto input_dims = input.dims();
|
|
auto hidden_dims = hidden->dims();
|
|
|
|
DenseTensor *batch_gate;
|
|
DenseTensor *batch_reset_hidden_prev;
|
|
DenseTensor *batch_hidden;
|
|
DenseTensor batch_gate_tmp, batch_reset_hidden_prev_tmp, batch_hidden_tmp;
|
|
if (is_test) {
|
|
batch_gate = &batch_gate_tmp;
|
|
batch_gate->Resize(input_dims);
|
|
|
|
batch_reset_hidden_prev = &batch_reset_hidden_prev_tmp;
|
|
batch_reset_hidden_prev->Resize(hidden_dims);
|
|
|
|
batch_hidden = &batch_hidden_tmp;
|
|
batch_hidden->Resize(hidden_dims);
|
|
} else {
|
|
batch_gate = param_batch_gate;
|
|
batch_hidden = param_batch_hidden;
|
|
batch_reset_hidden_prev = param_batch_reset_hidden_prev;
|
|
}
|
|
dev_ctx.template Alloc<T>(batch_gate);
|
|
dev_ctx.template Alloc<T>(batch_reset_hidden_prev);
|
|
dev_ctx.template Alloc<T>(batch_hidden);
|
|
|
|
funcs::DenseTensor2BatchFunctor<Context, T> to_batch;
|
|
to_batch(dev_ctx, input, batch_gate, true, is_reverse);
|
|
|
|
if (bias) {
|
|
funcs::RowwiseAdd<Context, T> add_bias;
|
|
add_bias(dev_ctx, *batch_gate, *bias, batch_gate);
|
|
}
|
|
|
|
int frame_size = hidden_dims[1];
|
|
funcs::GRUMetaValue<T> gru_value;
|
|
gru_value.gate_weight = const_cast<T *>(weight_data);
|
|
gru_value.state_weight =
|
|
const_cast<T *>(weight_data + 2 * frame_size * frame_size);
|
|
DenseTensor ordered_h0;
|
|
|
|
Vector<size_t> order(batch_gate->lod()[2]);
|
|
|
|
if (h0) {
|
|
// Since the batch computing for GRU reorders the input sequences
|
|
// according to their length. The initialized cell state also needs
|
|
// to reorder.
|
|
ReorderInitState<Context, T>(dev_ctx, *h0, order, &ordered_h0, true);
|
|
gru_value.prev_out_value = ordered_h0.data<T>();
|
|
} else {
|
|
gru_value.prev_out_value = nullptr;
|
|
}
|
|
auto batch_starts = batch_gate->lod()[0];
|
|
size_t num_batch = batch_starts.size() - 1;
|
|
auto active_node = funcs::detail::GetActivationType(activation);
|
|
auto active_gate = funcs::detail::GetActivationType(gate_activation);
|
|
for (size_t n = 0; n < num_batch; n++) {
|
|
int bstart = static_cast<int>(batch_starts[n]);
|
|
int bend = static_cast<int>(batch_starts[n + 1]);
|
|
int cur_batch_size = bend - bstart;
|
|
|
|
DenseTensor gate_t = batch_gate->Slice(bstart, bend);
|
|
DenseTensor reset_hidden_prev_t =
|
|
batch_reset_hidden_prev->Slice(bstart, bend);
|
|
DenseTensor hidden_t = batch_hidden->Slice(bstart, bend);
|
|
gru_value.output_value = hidden_t.data<T>();
|
|
gru_value.gate_value = gate_t.data<T>();
|
|
gru_value.reset_output_value = reset_hidden_prev_t.data<T>();
|
|
funcs::GRUUnitFunctor<Context, T>::compute(dev_ctx, // NOLINT
|
|
gru_value,
|
|
frame_size,
|
|
cur_batch_size,
|
|
active_node,
|
|
active_gate,
|
|
origin_mode);
|
|
gru_value.prev_out_value = gru_value.output_value;
|
|
}
|
|
|
|
funcs::Batch2DenseTensorFunctor<Context, T> to_seq;
|
|
batch_hidden->set_lod(batch_gate->lod());
|
|
to_seq(dev_ctx, *batch_hidden, hidden);
|
|
}
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(gru, GPU, ALL_LAYOUT, phi::GRUKernel, float, double) {}
|