// 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 "paddle/phi/kernels/cudnn_lstm_grad_kernel.h" #include "paddle/phi/core/kernel_registry.h" #include "paddle/phi/kernels/gpu/cudnn_lstm_utils.h" namespace phi { template void CudnnLSTMGradKernel( const Context &dev_ctx, const DenseTensor &x, const DenseTensor &init_h, const DenseTensor &init_c, const optional> &weight_list, const optional &sequence_length, const DenseTensor &out, const DenseTensor &reserve, const DenseTensor &state_out, const DenseTensor &out_grad, const DenseTensor &last_h_grad, const DenseTensor &last_c_grad, float dropout_prob, bool is_bidirec, int hidden_size, int num_layers, bool is_test, int seed, DenseTensor *x_grad, DenseTensor *init_h_grad, DenseTensor *init_c_grad, std::vector weight_grad_list) { auto input_dims = x.dims(); auto init_h_dims = init_h.dims(); auto init_c_dims = init_c.dims(); auto *init_h_data = init_h.data(); auto *init_c_data = init_c.data(); auto *out_data = out.data(); auto *out_grad_data = out_grad.data(); auto *last_h_grad_data = last_h_grad.data(); auto *last_c_grad_data = last_c_grad.data(); auto running_weight_list = *weight_list.get_ptr(); int weight_numel = size_sum(running_weight_list); bool continuous = is_continuous>(running_weight_list); auto handle = dev_ctx.cudnn_handle(); auto place = dev_ctx.GetPlace(); auto stream = dev_ctx.stream(); DenseTensor weight_whole; T *weight_data = nullptr; if (!continuous) { weight_whole.Resize({weight_numel}); dev_ctx.template Alloc(&weight_whole); weight_to_tensor(place, stream, running_weight_list, &weight_whole); weight_data = weight_whole.data(); } else { weight_data = const_cast(running_weight_list[0]->data()); } DenseTensor weight_grad; funcs::SetConstant zero; weight_grad.Resize({weight_numel}); dev_ctx.template Alloc(&weight_grad); zero(dev_ctx, &weight_grad, static_cast(0.0)); T *weight_grad_data = weight_grad.data(); int offset = 0; for (size_t i = 0; i < weight_grad_list.size(); ++i) { size_t len = weight_grad_list[i]->numel(); auto dim = weight_grad_list[i]->dims(); weight_grad_list[i] ->ShareDataWith(weight_grad.Slice(static_cast(offset), static_cast(offset + len))) .Resize(dim); offset += len; } x_grad->Resize(input_dims); dev_ctx.template Alloc(x_grad); auto *in_grad_data = x_grad->data(); if (init_h_grad) { init_h_grad->Resize(init_h_dims); dev_ctx.template Alloc(init_h_grad); } auto *init_h_grad_data = init_h_grad ? init_h_grad->data() : nullptr; if (init_c_grad) { init_c_grad->Resize(init_c_dims); dev_ctx.template Alloc(init_c_grad); } auto *init_c_grad_data = init_c_grad ? init_c_grad->data() : nullptr; auto running_seq_length = sequence_length.get_ptr(); bool has_seq_length = running_seq_length != nullptr; std::vector SequenceLength; if (has_seq_length) { SequenceLength = GetVectorFromTensor(running_seq_length); } int seq_length = input_dims[0]; int64_t batch_size = x.dims()[1]; int64_t input_size = x.dims()[2]; // TODO(large-tensor): cudnn rnn dims not support int64 PADDLE_ENFORCE_LE_INT_MAX(batch_size, "batch_size"); PADDLE_ENFORCE_LE_INT_MAX(input_size, "input_size"); int batch_size_int = static_cast(batch_size); int input_size_int = static_cast(input_size); size_t workspace_size; size_t reserve_size; ScopedRNNBase rnn(seq_length, batch_size_int, input_size_int, hidden_size, num_layers, dropout_prob, seed, weight_numel, true, is_bidirec); rnn.Create(handle, dev_ctx.GetPlace(), SequenceLength, &workspace_size, &reserve_size, const_cast(&state_out)); DenseTensor workspace_data_; workspace_data_.Resize({static_cast(workspace_size)}); dev_ctx.template Alloc(&workspace_data_); const uint8_t *reserve_data = reserve.data(); #if CUDNN_VERSION >= 90000 PADDLE_ENFORCE_GPU_SUCCESS( dynload::cudnnRNNBackwardData_v8(handle, rnn.rnn_desc(), nullptr, rnn.y_seq_desc(), out_data, out_grad_data, rnn.x_seq_desc(), in_grad_data, rnn.init_h_desc(), init_h_data, last_h_grad_data, init_h_grad_data, rnn.init_c_desc(), init_c_data, last_c_grad_data, init_c_grad_data, rnn.weights_size(), weight_data, workspace_size, workspace_data_.data(), reserve_size, const_cast(reserve_data))); PADDLE_ENFORCE_GPU_SUCCESS( dynload::cudnnRNNBackwardWeights_v8(handle, rnn.rnn_desc(), CUDNN_WGRAD_MODE_ADD, nullptr, rnn.x_seq_desc(), x.data(), rnn.init_h_desc(), init_h.data(), rnn.y_seq_desc(), out.data(), rnn.weights_size(), weight_grad_data, workspace_size, workspace_data_.data(), reserve_size, const_cast(reserve_data))); #else if (!has_seq_length) { // This interface is used when the input/output is unpadded. #ifdef PADDLE_WITH_HIP PADDLE_ENFORCE_GPU_SUCCESS( dynload::miopenRNNBackwardData(handle, rnn.rnn_desc(), seq_length, rnn.y_descs(), out_data, rnn.y_descs(), out_grad_data, rnn.last_h_desc(), last_h_grad_data, rnn.last_c_desc(), last_c_grad_data, rnn.weight_desc(), weight_data, rnn.init_h_desc(), init_h_data, rnn.init_c_desc(), init_c_data, rnn.x_descs(), in_grad_data, rnn.init_h_desc(), init_h_grad_data, rnn.init_c_desc(), init_c_grad_data, workspace_data_.data(), workspace_size, const_cast(reserve_data), reserve_size)); PADDLE_ENFORCE_GPU_SUCCESS( dynload::miopenRNNBackwardWeights(handle, rnn.rnn_desc(), seq_length, rnn.x_descs(), x.data(), rnn.init_h_desc(), init_h.data(), rnn.y_descs(), out.data(), rnn.weight_desc(), weight_grad_data, workspace_data_.data(), workspace_size, const_cast(reserve_data), reserve_size)); #else PADDLE_ENFORCE_GPU_SUCCESS( dynload::cudnnRNNBackwardData(handle, rnn.rnn_desc(), seq_length, rnn.y_descs(), out_data, rnn.y_descs(), out_grad_data, rnn.last_h_desc(), last_h_grad_data, rnn.last_c_desc(), last_c_grad_data, rnn.weight_desc(), weight_data, rnn.init_h_desc(), init_h_data, rnn.init_c_desc(), init_c_data, rnn.x_descs(), in_grad_data, rnn.init_h_desc(), init_h_grad_data, rnn.init_c_desc(), init_c_grad_data, workspace_data_.data(), workspace_size, const_cast(reserve_data), reserve_size)); PADDLE_ENFORCE_GPU_SUCCESS( dynload::cudnnRNNBackwardWeights(handle, rnn.rnn_desc(), seq_length, rnn.x_descs(), x.data(), rnn.init_h_desc(), init_h.data(), rnn.y_descs(), out.data(), workspace_data_.data(), workspace_size, rnn.weight_desc(), weight_grad_data, const_cast(reserve_data), reserve_size)); #endif } else { #if !defined(PADDLE_WITH_HIP) && CUDNN_VERSION >= 7201 // for train // This interface is used when the input/output is padded. PADDLE_ENFORCE_GPU_SUCCESS( dynload::cudnnRNNBackwardDataEx(handle, rnn.rnn_desc(), rnn.y_seq_desc(), out_data, rnn.y_seq_desc(), out_grad_data, nullptr, nullptr, rnn.last_h_desc(), last_h_grad_data, rnn.last_c_desc(), last_c_grad_data, rnn.weight_desc(), weight_data, rnn.init_h_desc(), init_h_data, rnn.init_c_desc(), init_c_data, rnn.x_seq_desc(), in_grad_data, rnn.init_h_desc(), init_h_grad_data, rnn.init_c_desc(), init_c_grad_data, nullptr, nullptr, workspace_data_.data(), workspace_size, const_cast(reserve_data), reserve_size)); PADDLE_ENFORCE_GPU_SUCCESS( dynload::cudnnRNNBackwardWeightsEx(handle, rnn.rnn_desc(), rnn.x_seq_desc(), x.data(), rnn.init_h_desc(), init_h.data(), rnn.y_seq_desc(), out.data(), workspace_data_.data(), workspace_size, rnn.weight_desc(), weight_grad_data, const_cast(reserve_data), reserve_size)); #else PADDLE_THROW(common::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 } #endif // end CUDNN_VERSION >= 90000 } } // namespace phi #ifdef PADDLE_WITH_HIP PD_REGISTER_KERNEL( cudnn_lstm_grad, GPU, ALL_LAYOUT, phi::CudnnLSTMGradKernel, float) {} #else PD_REGISTER_KERNEL( cudnn_lstm_grad, GPU, ALL_LAYOUT, phi::CudnnLSTMGradKernel, float, double) { } #endif