| 1879 | } |
| 1880 | |
| 1881 | port::StatusOr<DeviceMemory<uint8>> CreateCtcLossWorkspace( |
| 1882 | Stream* stream, const CudnnHandle& cudnn, |
| 1883 | const CudnnCtcLossDescriptor& ctc_loss_desc, |
| 1884 | const CudnnRnnStateTensorDescriptor& probs_desc, |
| 1885 | const CudnnRnnStateTensorDescriptor& grads_desc, |
| 1886 | const absl::Span<const int32>& labels_data, |
| 1887 | const absl::Span<const int32>& labels_lengths_data, |
| 1888 | const absl::Span<const int32>& input_lengths_data, |
| 1889 | ScratchAllocator* workspace_allocator) { |
| 1890 | // Query the workspace size. |
| 1891 | size_t workspace_size_in_bytes = 0; |
| 1892 | #if CUDNN_VERSION >= 7604 |
| 1893 | RETURN_IF_CUDNN_ERROR(cudnnGetCTCLossWorkspaceSize( |
| 1894 | /*handle=*/cudnn.handle(), /*probsDesc=*/probs_desc.handle(), |
| 1895 | /*gradientsDesc=*/grads_desc.handle(), |
| 1896 | /*labels=*/labels_data.data(), |
| 1897 | /*labelLengths=*/labels_lengths_data.data(), |
| 1898 | /*inputLengths=*/input_lengths_data.data(), |
| 1899 | /*algo=*/CUDNN_CTC_LOSS_ALGO_NON_DETERMINISTIC, |
| 1900 | /*ctcLossDesc=*/ctc_loss_desc.handle(), |
| 1901 | /*sizeInBytes=*/&workspace_size_in_bytes)); |
| 1902 | #else |
| 1903 | return port::Status(port::error::INVALID_ARGUMENT, |
| 1904 | "No supported cudnnGetCTCLossWorkspaceSize when " |
| 1905 | "CUDNN_VERSION < 7.6.3"); |
| 1906 | #endif |
| 1907 | // Allocate the workspace. |
| 1908 | if (workspace_size_in_bytes == 0) { |
| 1909 | return DeviceMemory<uint8>(); |
| 1910 | } |
| 1911 | return workspace_allocator->AllocateBytes(workspace_size_in_bytes); |
| 1912 | } |
| 1913 | #endif |
| 1914 | |
| 1915 | } // namespace |
no test coverage detected