| 1155 | CudnnRnnDescriptor(CudnnRnnDescriptor&& other) = default; |
| 1156 | |
| 1157 | static port::StatusOr<CudnnRnnDescriptor> Create( |
| 1158 | const CudnnHandle& cudnn, int num_layers, int hidden_size, int input_size, |
| 1159 | int cell_size, int batch_size, cudnnRNNInputMode_t input_mode, |
| 1160 | cudnnDirectionMode_t direction_mode, cudnnRNNMode_t rnn_mode, |
| 1161 | cudnnDataType_t data_type, cudnnDataType_t compute_type, |
| 1162 | const dnn::AlgorithmConfig& algorithm_config, float dropout, uint64 seed, |
| 1163 | ScratchAllocator* state_allocator, bool use_padded_io) { |
| 1164 | SE_ASSIGN_OR_RETURN( |
| 1165 | CudnnDropoutDescriptor dropout_desc, |
| 1166 | CudnnDropoutDescriptor::Create(cudnn, dropout, seed, state_allocator)); |
| 1167 | |
| 1168 | gpu::RnnDescriptor rnn_desc = CreateRnnDescriptor(); |
| 1169 | cudnnRNNAlgo_t rnn_algo = ToCudnnRNNAlgo(algorithm_config.algorithm()); |
| 1170 | |
| 1171 | // TODO: allow the user to choose an algorithm. |
| 1172 | int unified_size = hidden_size; |
| 1173 | bool use_projection = cell_size != 0 && hidden_size < cell_size; |
| 1174 | if (use_projection) { |
| 1175 | unified_size = cell_size; |
| 1176 | } |
| 1177 | RETURN_IF_CUDNN_ERROR(cudnnSetRNNDescriptor_v6( |
| 1178 | cudnn.handle(), /*rnnDesc=*/rnn_desc.get(), |
| 1179 | /*hiddenSize=*/unified_size, /*numLayers=*/num_layers, |
| 1180 | /*dropoutDesc=*/dropout_desc.handle(), /*inputMode=*/input_mode, |
| 1181 | /*direction=*/direction_mode, /*mode=*/rnn_mode, /*algo=*/rnn_algo, |
| 1182 | /*dataType=*/compute_type)); |
| 1183 | if (use_projection) { |
| 1184 | #if CUDNN_VERSION >= 7101 |
| 1185 | RETURN_IF_CUDNN_ERROR(cudnnSetRNNProjectionLayers( |
| 1186 | cudnn.handle(), /*rnnDesc=*/rnn_desc.get(), |
| 1187 | /*recProjSize=*/hidden_size, /*outProjSize=*/0)); |
| 1188 | #else |
| 1189 | return port::Status(port::error::INVALID_ARGUMENT, |
| 1190 | "No supported cudnnSetRNNProjectionLayers when " |
| 1191 | "CUDNN_VERSION < 7.1.1"); |
| 1192 | #endif |
| 1193 | } |
| 1194 | |
| 1195 | // TODO: For now, we only use cudnnRNN**Ex API to process padded inputs. |
| 1196 | // But in the future if these APIs are used to process full length arrays, |
| 1197 | // we need to distinguish when to set it. |
| 1198 | #if CUDNN_VERSION >= 7201 |
| 1199 | if (use_padded_io) { |
| 1200 | RETURN_IF_CUDNN_ERROR( |
| 1201 | cudnnSetRNNPaddingMode(rnn_desc.get(), CUDNN_RNN_PADDED_IO_ENABLED)); |
| 1202 | } |
| 1203 | #endif |
| 1204 | |
| 1205 | port::StatusOr<PersistentRnnPlan> rnn_plan_wrapper; |
| 1206 | PersistentRnnPlan rnn_plan; |
| 1207 | if (rnn_algo == CUDNN_RNN_ALGO_PERSIST_DYNAMIC) { |
| 1208 | CHECK_GE(batch_size, 0); |
| 1209 | rnn_plan_wrapper = |
| 1210 | CreatePersistentRnnPlan(rnn_desc.get(), batch_size, data_type); |
| 1211 | if (!rnn_plan_wrapper.ok()) { |
| 1212 | return port::StatusOr<CudnnRnnDescriptor>(rnn_plan_wrapper.status()); |
| 1213 | } else { |
| 1214 | rnn_plan = rnn_plan_wrapper.ConsumeValueOrDie(); |
nothing calls this directly
no test coverage detected