| 1970 | } |
| 1971 | |
| 1972 | bool CreateRnnWorkspace(Stream* stream, miopenHandle_t miopen_handle, |
| 1973 | const MIOpenRnnDescriptor& rnn_desc, |
| 1974 | const MIOpenRnnSequenceTensorDescriptor& input_desc, |
| 1975 | ScratchAllocator* workspace_allocator, |
| 1976 | DeviceMemory<uint8>* workspace) { |
| 1977 | // Query the workspace size. |
| 1978 | size_t workspace_size_in_bytes = 0; |
| 1979 | auto status = wrap::miopenGetRNNWorkspaceSize( |
| 1980 | miopen_handle /*handle*/, rnn_desc.handle() /*rnnDesc*/, |
| 1981 | input_desc.seq_length() /*seqLength*/, input_desc.handles() /*xDesc*/, |
| 1982 | &workspace_size_in_bytes /*sizeInBytes*/); |
| 1983 | if (status != miopenStatusSuccess) { |
| 1984 | LOG(ERROR) << "Unable to query workspace size: " << ToString(status); |
| 1985 | return false; |
| 1986 | } |
| 1987 | // Allocate the workspace. |
| 1988 | if (workspace_size_in_bytes > 0) { |
| 1989 | auto allocated = |
| 1990 | workspace_allocator->AllocateBytes(workspace_size_in_bytes); |
| 1991 | if (!allocated.ok() || (*workspace = allocated.ValueOrDie()) == nullptr) { |
| 1992 | LOG(ERROR) << "Failed to allocate RNN workspace"; |
| 1993 | |
| 1994 | return false; |
| 1995 | } |
| 1996 | stream->ThenMemZero(workspace, workspace_size_in_bytes); |
| 1997 | } else { |
| 1998 | *workspace = DeviceMemory<uint8>(); |
| 1999 | } |
| 2000 | return true; |
| 2001 | } |
| 2002 | |
| 2003 | } // namespace |
| 2004 |
no test coverage detected