| 2224 | } |
| 2225 | |
| 2226 | port::Status CudnnSupport::DoCtcLossImpl( |
| 2227 | Stream* stream, const CudnnRnnStateTensorDescriptor& probs_desc, |
| 2228 | const DeviceMemory<float>& probs_data, |
| 2229 | const absl::Span<const int32>& labels_data, |
| 2230 | const absl::Span<const int32>& labels_lengths_data, |
| 2231 | const absl::Span<const int32>& input_lengths_data, |
| 2232 | DeviceMemory<float>* costs_data, |
| 2233 | const CudnnRnnStateTensorDescriptor& grads_desc, |
| 2234 | DeviceMemory<float>* grads_data, |
| 2235 | const CudnnCtcLossDescriptor& ctc_loss_desc, |
| 2236 | ScratchAllocator* workspace_allocator) { |
| 2237 | auto cudnn = cudnn_->GetHandle(parent_, stream); |
| 2238 | |
| 2239 | SE_ASSIGN_OR_RETURN(DeviceMemory<uint8> workspace, |
| 2240 | CreateCtcLossWorkspace(stream, cudnn, ctc_loss_desc, |
| 2241 | probs_desc, grads_desc, |
| 2242 | labels_data, labels_lengths_data, |
| 2243 | input_lengths_data, |
| 2244 | workspace_allocator)); |
| 2245 | int kNumTimestamps = probs_desc.num_layers(); |
| 2246 | int kBatchSize = probs_desc.batch_size(); |
| 2247 | int kNumLabels = probs_desc.data_size(); |
| 2248 | int total_size = kNumLabels * kNumTimestamps * kBatchSize; |
| 2249 | |
| 2250 | #if CUDNN_VERSION >= 7604 |
| 2251 | RETURN_IF_CUDNN_ERROR(cudnnCTCLoss( |
| 2252 | /*handle=*/cudnn.handle(), /*probsDesc=*/probs_desc.handle(), |
| 2253 | /*probs=*/probs_data.opaque(), /*labels=*/labels_data.data(), |
| 2254 | /*labelsLengths=*/labels_lengths_data.data(), |
| 2255 | /*inputLengths=*/input_lengths_data.data(), |
| 2256 | /*costs=*/costs_data->opaque(), /*gradientsDesc=*/grads_desc.handle(), |
| 2257 | /*gradients=*/grads_data->opaque(), |
| 2258 | /*algo=*/CUDNN_CTC_LOSS_ALGO_NON_DETERMINISTIC, |
| 2259 | /*ctcLossDesc=*/ctc_loss_desc.handle(), |
| 2260 | /*workspace=*/workspace.opaque(), |
| 2261 | /*workSpaceSizeInBytes=*/workspace.size())); |
| 2262 | #else |
| 2263 | return port::Status(port::error::INVALID_ARGUMENT, |
| 2264 | "No supported cudnnCTCLoss when " |
| 2265 | "CUDNN_VERSION < 7.6.3"); |
| 2266 | #endif |
| 2267 | |
| 2268 | return port::Status::OK(); |
| 2269 | } |
| 2270 | |
| 2271 | port::StatusOr<std::unique_ptr<dnn::RnnDescriptor>> |
| 2272 | CudnnSupport::createRnnDescriptor( |
nothing calls this directly
no test coverage detected