| 1770 | } |
| 1771 | |
| 1772 | port::StatusOr<DeviceMemory<uint8>> CreateRnnWorkspace( |
| 1773 | Stream* stream, const CudnnHandle& cudnn, |
| 1774 | const CudnnRnnDescriptor& rnn_desc, |
| 1775 | const CudnnRnnSequenceTensorDescriptor& input_desc, |
| 1776 | ScratchAllocator* workspace_allocator) { |
| 1777 | // Query the workspace size. |
| 1778 | size_t workspace_size_in_bytes = 0; |
| 1779 | RETURN_IF_CUDNN_ERROR(cudnnGetRNNWorkspaceSize( |
| 1780 | /*handle=*/cudnn.handle(), /*rnnDesc=*/rnn_desc.handle(), |
| 1781 | /*seqLength=*/input_desc.max_seq_length(), /*xDesc=*/input_desc.handles(), |
| 1782 | /*sizeInBytes=*/&workspace_size_in_bytes)); |
| 1783 | // Allocate the workspace. |
| 1784 | if (workspace_size_in_bytes == 0) { |
| 1785 | return DeviceMemory<uint8>(); |
| 1786 | } |
| 1787 | return workspace_allocator->AllocateBytes(workspace_size_in_bytes); |
| 1788 | } |
| 1789 | |
| 1790 | #if CUDNN_VERSION >= 7402 |
| 1791 | port::StatusOr<size_t> GetBatchNormalizationReserveSpaceSizeInternal( |
no test coverage detected