|
|
|
|
@ -16,6 +16,7 @@ limitations under the License. */
|
|
|
|
|
#include "paddle/fluid/operators/cudnn_rnn_cache.h"
|
|
|
|
|
#include "paddle/fluid/operators/math/math_function.h"
|
|
|
|
|
#include "paddle/fluid/platform/cudnn_desc.h"
|
|
|
|
|
#include "paddle/fluid/platform/cudnn_helper.h"
|
|
|
|
|
|
|
|
|
|
namespace paddle {
|
|
|
|
|
namespace operators {
|
|
|
|
|
@ -55,50 +56,96 @@ class CudnnLSTMGPUKernel : public framework::OpKernel<T> {
|
|
|
|
|
int num_layers = ctx.Attr<int>("num_layers");
|
|
|
|
|
bool is_test = ctx.Attr<bool>("is_test");
|
|
|
|
|
int seed = ctx.Attr<int>("seed");
|
|
|
|
|
auto sequence_length = ctx.Attr<std::vector<int>>("sequence_length");
|
|
|
|
|
|
|
|
|
|
auto &dev_ctx = ctx.template device_context<platform::CUDADeviceContext>();
|
|
|
|
|
auto handle = dev_ctx.cudnn_handle();
|
|
|
|
|
|
|
|
|
|
CudnnRNNCache *cudnn_rnn_cache = new CudnnRNNCache();
|
|
|
|
|
int seq_length = x->dims()[0];
|
|
|
|
|
int batch_size = x->dims()[1];
|
|
|
|
|
int input_size = x->dims()[2];
|
|
|
|
|
int weight_numel = w->numel();
|
|
|
|
|
bool state_initialized = state_out->IsInitialized() ? true : false;
|
|
|
|
|
|
|
|
|
|
auto input_w_numel = w->numel();
|
|
|
|
|
auto seq_len = x->dims()[0];
|
|
|
|
|
auto batch_size = x->dims()[1];
|
|
|
|
|
auto input_dim = x->dims()[2];
|
|
|
|
|
size_t workspace_size;
|
|
|
|
|
size_t reserve_size;
|
|
|
|
|
bool state_initialized = state_out->IsInitialized() ? true : false;
|
|
|
|
|
cudnnDataType_t cudnn_type = platform::ToCudnnDataType(
|
|
|
|
|
framework::ToDataType(std::type_index(typeid(T))));
|
|
|
|
|
cudnn_rnn_cache->init(handle, ctx.GetPlace(), seq_len, batch_size,
|
|
|
|
|
input_dim, hidden_size, num_layers, dropout_prob,
|
|
|
|
|
is_bidirec, seed, input_w_numel, &reserve_size,
|
|
|
|
|
state_out, state_initialized, cudnn_type);
|
|
|
|
|
|
|
|
|
|
platform::ScopedRNNBase rnn(seq_length, batch_size, input_size, hidden_size,
|
|
|
|
|
num_layers, dropout_prob, seed, weight_numel,
|
|
|
|
|
state_initialized, is_bidirec);
|
|
|
|
|
rnn.Create<T>(handle, ctx.GetPlace(), sequence_length, &workspace_size,
|
|
|
|
|
&reserve_size, state_out);
|
|
|
|
|
|
|
|
|
|
framework::Tensor workspace_data_;
|
|
|
|
|
workspace_data_.Resize({static_cast<int64_t>(workspace_size)});
|
|
|
|
|
workspace_data_.mutable_data<uint8_t>(ctx.GetPlace());
|
|
|
|
|
|
|
|
|
|
auto *reserve_data = reserve->mutable_data<uint8_t>(
|
|
|
|
|
{static_cast<int64_t>(reserve_size)}, ctx.GetPlace());
|
|
|
|
|
|
|
|
|
|
if (is_test) {
|
|
|
|
|
// for inference
|
|
|
|
|
PADDLE_ENFORCE_CUDA_SUCCESS(platform::dynload::cudnnRNNForwardInference(
|
|
|
|
|
handle, cudnn_rnn_cache->rnn_desc_, seq_len, cudnn_rnn_cache->x_desc_,
|
|
|
|
|
x_data, cudnn_rnn_cache->hx_desc_, init_h_data,
|
|
|
|
|
cudnn_rnn_cache->cx_desc_, init_c_data, cudnn_rnn_cache->w_desc_,
|
|
|
|
|
w_data, cudnn_rnn_cache->y_desc_, out_data, cudnn_rnn_cache->hy_desc_,
|
|
|
|
|
last_h_data, cudnn_rnn_cache->cy_desc_, last_c_data,
|
|
|
|
|
cudnn_rnn_cache->workspace_data_.data<uint8_t>(),
|
|
|
|
|
cudnn_rnn_cache->workspace_size_));
|
|
|
|
|
if (sequence_length.empty()) {
|
|
|
|
|
// for inference
|
|
|
|
|
// This interface is used when the input/output is unpadded.
|
|
|
|
|
PADDLE_ENFORCE_CUDA_SUCCESS(platform::dynload::cudnnRNNForwardInference(
|
|
|
|
|
handle, rnn.rnn_desc(), seq_length, rnn.x_desc(), x_data,
|
|
|
|
|
rnn.hx_desc(), init_h_data, rnn.cx_desc(), init_c_data,
|
|
|
|
|
rnn.w_desc(), w_data, rnn.y_desc(), out_data, rnn.hy_desc(),
|
|
|
|
|
last_h_data, rnn.cy_desc(), last_c_data,
|
|
|
|
|
workspace_data_.data<uint8_t>(), workspace_size));
|
|
|
|
|
} else {
|
|
|
|
|
#if CUDNN_VERSION >= 7201
|
|
|
|
|
// for inference
|
|
|
|
|
// This interface is used when the input/output is padded.
|
|
|
|
|
PADDLE_ENFORCE_CUDA_SUCCESS(
|
|
|
|
|
platform::dynload::cudnnRNNForwardInferenceEx(
|
|
|
|
|
handle, rnn.rnn_desc(), rnn.x_seq_desc(), x_data, rnn.hx_desc(),
|
|
|
|
|
init_h_data, rnn.cx_desc(), init_c_data, rnn.w_desc(), w_data,
|
|
|
|
|
rnn.y_seq_desc(), out_data, rnn.hy_desc(), last_h_data,
|
|
|
|
|
rnn.cy_desc(), last_c_data, nullptr, nullptr, nullptr, nullptr,
|
|
|
|
|
nullptr, nullptr, nullptr, nullptr,
|
|
|
|
|
workspace_data_.data<uint8_t>(), workspace_size));
|
|
|
|
|
#else
|
|
|
|
|
PADDLE_ENFORCE_NOT_NULL(
|
|
|
|
|
nullptr, platform::errors::Unavailable(
|
|
|
|
|
"The padded input is supported by "
|
|
|
|
|
"cudnnRNNForwardInferenceEx, but it only works when "
|
|
|
|
|
"the version of cudnn is larger than 7.2.1"));
|
|
|
|
|
#endif
|
|
|
|
|
}
|
|
|
|
|
} else {
|
|
|
|
|
// for train
|
|
|
|
|
PADDLE_ENFORCE_CUDA_SUCCESS(platform::dynload::cudnnRNNForwardTraining(
|
|
|
|
|
handle, cudnn_rnn_cache->rnn_desc_, seq_len, cudnn_rnn_cache->x_desc_,
|
|
|
|
|
x_data, cudnn_rnn_cache->hx_desc_, init_h_data,
|
|
|
|
|
cudnn_rnn_cache->cx_desc_, init_c_data, cudnn_rnn_cache->w_desc_,
|
|
|
|
|
w_data, cudnn_rnn_cache->y_desc_, out_data, cudnn_rnn_cache->hy_desc_,
|
|
|
|
|
last_h_data, cudnn_rnn_cache->cy_desc_, last_c_data,
|
|
|
|
|
cudnn_rnn_cache->workspace_data_.data<uint8_t>(),
|
|
|
|
|
cudnn_rnn_cache->workspace_size_, reserve_data, reserve_size));
|
|
|
|
|
if (sequence_length.empty()) {
|
|
|
|
|
// for train
|
|
|
|
|
// This interface is used when the input/output is unpadded.
|
|
|
|
|
PADDLE_ENFORCE_CUDA_SUCCESS(platform::dynload::cudnnRNNForwardTraining(
|
|
|
|
|
handle, rnn.rnn_desc(), seq_length, rnn.x_desc(), x_data,
|
|
|
|
|
rnn.hx_desc(), init_h_data, rnn.cx_desc(), init_c_data,
|
|
|
|
|
rnn.w_desc(), w_data, rnn.y_desc(), out_data, rnn.hy_desc(),
|
|
|
|
|
last_h_data, rnn.cy_desc(), last_c_data,
|
|
|
|
|
workspace_data_.data<uint8_t>(), workspace_size, reserve_data,
|
|
|
|
|
reserve_size));
|
|
|
|
|
} else {
|
|
|
|
|
#if CUDNN_VERSION >= 7201
|
|
|
|
|
// for train
|
|
|
|
|
// This interface is used when the input/output is padded.
|
|
|
|
|
PADDLE_ENFORCE_CUDA_SUCCESS(
|
|
|
|
|
platform::dynload::cudnnRNNForwardTrainingEx(
|
|
|
|
|
handle, rnn.rnn_desc(), rnn.x_seq_desc(), x_data, rnn.hx_desc(),
|
|
|
|
|
init_h_data, rnn.cx_desc(), init_c_data, rnn.w_desc(), w_data,
|
|
|
|
|
rnn.y_seq_desc(), out_data, rnn.hy_desc(), last_h_data,
|
|
|
|
|
rnn.cy_desc(), last_c_data, nullptr, nullptr, nullptr, nullptr,
|
|
|
|
|
nullptr, nullptr, nullptr, nullptr,
|
|
|
|
|
workspace_data_.data<uint8_t>(), workspace_size, reserve_data,
|
|
|
|
|
reserve_size));
|
|
|
|
|
#else
|
|
|
|
|
PADDLE_ENFORCE_NOT_NULL(
|
|
|
|
|
nullptr, platform::errors::Unavailable(
|
|
|
|
|
"The padded input is supported by "
|
|
|
|
|
"cudnnRNNForwardTrainingEx, but it only works when "
|
|
|
|
|
"the version of cudnn is larger than 7.2.1"));
|
|
|
|
|
#endif
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
delete cudnn_rnn_cache;
|
|
|
|
|
}
|
|
|
|
|
};
|
|
|
|
|
|
|
|
|
|
@ -156,44 +203,74 @@ class CudnnLSTMGPUGradKernel : public framework::OpKernel<T> {
|
|
|
|
|
int hidden_size = ctx.Attr<int>("hidden_size");
|
|
|
|
|
int num_layers = ctx.Attr<int>("num_layers");
|
|
|
|
|
int seed = ctx.Attr<int>("seed");
|
|
|
|
|
auto sequence_length = ctx.Attr<std::vector<int>>("sequence_length");
|
|
|
|
|
|
|
|
|
|
CudnnRNNCache *cudnn_rnn_cache = new CudnnRNNCache();
|
|
|
|
|
int seq_length = input_dims[0];
|
|
|
|
|
int batch_size = input->dims()[1];
|
|
|
|
|
int input_size = input->dims()[2];
|
|
|
|
|
int weight_numel = weight->numel();
|
|
|
|
|
|
|
|
|
|
auto input_w_numel = weight->numel();
|
|
|
|
|
auto seq_len = input_dims[0];
|
|
|
|
|
auto batch_size = input->dims()[1];
|
|
|
|
|
auto input_dim = input->dims()[2];
|
|
|
|
|
size_t workspace_size;
|
|
|
|
|
size_t reserve_size;
|
|
|
|
|
cudnnDataType_t cudnn_type = platform::ToCudnnDataType(
|
|
|
|
|
framework::ToDataType(std::type_index(typeid(T))));
|
|
|
|
|
cudnn_rnn_cache->init(handle, ctx.GetPlace(), seq_len, batch_size,
|
|
|
|
|
input_dim, hidden_size, num_layers, dropout_prob,
|
|
|
|
|
is_bidirec, seed, input_w_numel, &reserve_size,
|
|
|
|
|
const_cast<Tensor *>(state_out), true, cudnn_type);
|
|
|
|
|
|
|
|
|
|
auto work_data = cudnn_rnn_cache->workspace_data_.data<uint8_t>();
|
|
|
|
|
|
|
|
|
|
platform::ScopedRNNBase rnn(seq_length, batch_size, input_size, hidden_size,
|
|
|
|
|
num_layers, dropout_prob, seed, weight_numel,
|
|
|
|
|
true, is_bidirec);
|
|
|
|
|
|
|
|
|
|
rnn.Create<T>(handle, ctx.GetPlace(), sequence_length, &workspace_size,
|
|
|
|
|
&reserve_size, const_cast<Tensor *>(state_out));
|
|
|
|
|
|
|
|
|
|
framework::Tensor workspace_data_;
|
|
|
|
|
workspace_data_.Resize({static_cast<int64_t>(workspace_size)});
|
|
|
|
|
workspace_data_.mutable_data<uint8_t>(ctx.GetPlace());
|
|
|
|
|
const uint8_t *reserve_data = reserve->data<uint8_t>();
|
|
|
|
|
|
|
|
|
|
PADDLE_ENFORCE_CUDA_SUCCESS(platform::dynload::cudnnRNNBackwardData(
|
|
|
|
|
handle, cudnn_rnn_cache->rnn_desc_, seq_len, cudnn_rnn_cache->y_desc_,
|
|
|
|
|
out_data, cudnn_rnn_cache->y_desc_, out_grad_data,
|
|
|
|
|
cudnn_rnn_cache->hy_desc_, last_h_grad_data, cudnn_rnn_cache->cy_desc_,
|
|
|
|
|
last_c_grad_data, cudnn_rnn_cache->w_desc_, weight_data,
|
|
|
|
|
cudnn_rnn_cache->hx_desc_, init_h_data, cudnn_rnn_cache->cx_desc_,
|
|
|
|
|
init_c_data, cudnn_rnn_cache->x_desc_, in_grad_data,
|
|
|
|
|
cudnn_rnn_cache->hx_desc_, init_h_grad_data, cudnn_rnn_cache->cx_desc_,
|
|
|
|
|
init_c_grad_data, work_data, cudnn_rnn_cache->workspace_size_,
|
|
|
|
|
const_cast<uint8_t *>(reserve_data), reserve_size));
|
|
|
|
|
|
|
|
|
|
PADDLE_ENFORCE_CUDA_SUCCESS(platform::dynload::cudnnRNNBackwardWeights(
|
|
|
|
|
handle, cudnn_rnn_cache->rnn_desc_, seq_len, cudnn_rnn_cache->x_desc_,
|
|
|
|
|
input->data<T>(), cudnn_rnn_cache->hx_desc_, init_h->data<T>(),
|
|
|
|
|
cudnn_rnn_cache->y_desc_, out->data<T>(),
|
|
|
|
|
cudnn_rnn_cache->workspace_data_.data<uint8_t>(),
|
|
|
|
|
cudnn_rnn_cache->workspace_size_, cudnn_rnn_cache->w_desc_,
|
|
|
|
|
weight_grad->data<T>(), const_cast<uint8_t *>(reserve_data),
|
|
|
|
|
reserve_size));
|
|
|
|
|
delete cudnn_rnn_cache;
|
|
|
|
|
if (sequence_length.empty()) {
|
|
|
|
|
// This interface is used when the input/output is unpadded.
|
|
|
|
|
PADDLE_ENFORCE_CUDA_SUCCESS(platform::dynload::cudnnRNNBackwardData(
|
|
|
|
|
handle, rnn.rnn_desc(), seq_length, rnn.y_desc(), out_data,
|
|
|
|
|
rnn.y_desc(), out_grad_data, rnn.hy_desc(), last_h_grad_data,
|
|
|
|
|
rnn.cy_desc(), last_c_grad_data, rnn.w_desc(), weight_data,
|
|
|
|
|
rnn.hx_desc(), init_h_data, rnn.cx_desc(), init_c_data, rnn.x_desc(),
|
|
|
|
|
in_grad_data, rnn.hx_desc(), init_h_grad_data, rnn.cx_desc(),
|
|
|
|
|
init_c_grad_data, workspace_data_.data<uint8_t>(), workspace_size,
|
|
|
|
|
const_cast<uint8_t *>(reserve_data), reserve_size));
|
|
|
|
|
|
|
|
|
|
PADDLE_ENFORCE_CUDA_SUCCESS(platform::dynload::cudnnRNNBackwardWeights(
|
|
|
|
|
handle, rnn.rnn_desc(), seq_length, rnn.x_desc(), input->data<T>(),
|
|
|
|
|
rnn.hx_desc(), init_h->data<T>(), rnn.y_desc(), out->data<T>(),
|
|
|
|
|
workspace_data_.data<uint8_t>(), workspace_size, rnn.w_desc(),
|
|
|
|
|
weight_grad->data<T>(), const_cast<uint8_t *>(reserve_data),
|
|
|
|
|
reserve_size));
|
|
|
|
|
} else {
|
|
|
|
|
#if CUDNN_VERSION >= 7201
|
|
|
|
|
// for train
|
|
|
|
|
// This interface is used when the input/output is padded.
|
|
|
|
|
PADDLE_ENFORCE_CUDA_SUCCESS(platform::dynload::cudnnRNNBackwardDataEx(
|
|
|
|
|
handle, rnn.rnn_desc(), rnn.y_seq_desc(), out_data, rnn.y_seq_desc(),
|
|
|
|
|
out_grad_data, nullptr, nullptr, rnn.hy_desc(), last_h_grad_data,
|
|
|
|
|
rnn.cy_desc(), last_c_grad_data, rnn.w_desc(), weight_data,
|
|
|
|
|
rnn.hx_desc(), init_h_data, rnn.cx_desc(), init_c_data,
|
|
|
|
|
rnn.x_seq_desc(), in_grad_data, rnn.hx_desc(), init_h_grad_data,
|
|
|
|
|
rnn.cx_desc(), init_c_grad_data, nullptr, nullptr,
|
|
|
|
|
workspace_data_.data<uint8_t>(), workspace_size,
|
|
|
|
|
const_cast<uint8_t *>(reserve_data), reserve_size));
|
|
|
|
|
|
|
|
|
|
PADDLE_ENFORCE_CUDA_SUCCESS(platform::dynload::cudnnRNNBackwardWeightsEx(
|
|
|
|
|
handle, rnn.rnn_desc(), rnn.x_seq_desc(), input->data<T>(),
|
|
|
|
|
rnn.hx_desc(), init_h->data<T>(), rnn.y_seq_desc(), out->data<T>(),
|
|
|
|
|
workspace_data_.data<uint8_t>(), workspace_size, rnn.w_desc(),
|
|
|
|
|
weight_grad->data<T>(), const_cast<uint8_t *>(reserve_data),
|
|
|
|
|
reserve_size));
|
|
|
|
|
#else
|
|
|
|
|
PADDLE_ENFORCE_NOT_NULL(
|
|
|
|
|
nullptr,
|
|
|
|
|
platform::errors::Unavailable(
|
|
|
|
|
"The padded input of rnn is supported by cudnnRNNBackwardDataEx, "
|
|
|
|
|
"cudnnRNNBackwardWeightsEx, but it only works when the version "
|
|
|
|
|
"of cudnn is larger than 7.2.1"));
|
|
|
|
|
#endif
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
};
|
|
|
|
|
|
|
|
|
|
|