| 1124 | } // namespace |
| 1125 | |
| 1126 | class CudnnRnnDescriptor : public dnn::RnnDescriptor { |
| 1127 | CudnnRnnDescriptor(const CudnnHandle& cudnn, gpu::RnnDescriptor rnn_desc, |
| 1128 | PersistentRnnPlan rnn_plan, int num_layers, |
| 1129 | int hidden_size, int input_size, int cell_size, |
| 1130 | int batch_size, cudnnRNNInputMode_t input_mode, |
| 1131 | cudnnDirectionMode_t direction_mode, |
| 1132 | cudnnRNNMode_t rnn_mode, cudnnDataType_t data_type, |
| 1133 | cudnnDataType_t compute_type, |
| 1134 | const dnn::AlgorithmConfig& algorithm_config, |
| 1135 | CudnnDropoutDescriptor dropout_desc, |
| 1136 | CudnnRnnParamsDescriptor params_desc) |
| 1137 | : rnn_desc_(std::move(rnn_desc)), |
| 1138 | rnn_plan_(std::move(rnn_plan)), |
| 1139 | num_layers_(num_layers), |
| 1140 | hidden_size_(hidden_size), |
| 1141 | input_size_(input_size), |
| 1142 | cell_size_(cell_size), |
| 1143 | batch_size_(batch_size), |
| 1144 | rnn_algo_(ToCudnnRNNAlgo(algorithm_config.algorithm())), |
| 1145 | input_mode_(input_mode), |
| 1146 | direction_mode_(direction_mode), |
| 1147 | rnn_mode_(rnn_mode), |
| 1148 | data_type_(data_type), |
| 1149 | compute_type_(compute_type), |
| 1150 | algorithm_config_(algorithm_config), |
| 1151 | dropout_desc_(std::move(dropout_desc)), |
| 1152 | params_desc_(std::move(params_desc)) {} |
| 1153 | |
| 1154 | public: |
| 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) { |