| 1916 | |
| 1917 | template <class T> |
| 1918 | port::Status CudnnSupport::DoRnnForwardImpl( |
| 1919 | Stream* stream, const CudnnRnnDescriptor& rnn_desc, |
| 1920 | const CudnnRnnSequenceTensorDescriptor& input_desc, |
| 1921 | const DeviceMemory<T>& input_data, |
| 1922 | const CudnnRnnStateTensorDescriptor& input_h_desc, |
| 1923 | const DeviceMemory<T>& input_h_data, |
| 1924 | const CudnnRnnStateTensorDescriptor& input_c_desc, |
| 1925 | const DeviceMemory<T>& input_c_data, const DeviceMemory<T>& params, |
| 1926 | const CudnnRnnSequenceTensorDescriptor& output_desc, |
| 1927 | DeviceMemory<T>* output_data, |
| 1928 | const CudnnRnnStateTensorDescriptor& output_h_desc, |
| 1929 | DeviceMemory<T>* output_h_data, |
| 1930 | const CudnnRnnStateTensorDescriptor& output_c_desc, |
| 1931 | DeviceMemory<T>* output_c_data, bool is_training, |
| 1932 | ScratchAllocator* reserve_space_allocator, |
| 1933 | ScratchAllocator* workspace_allocator, |
| 1934 | dnn::ProfileResult* output_profile_result) { |
| 1935 | SE_ASSIGN_OR_RETURN( |
| 1936 | RnnModelDims model_dims, |
| 1937 | ExtractAndCheckRnnForward( |
| 1938 | rnn_desc, input_desc, input_data, input_h_desc, input_h_data, |
| 1939 | input_c_desc, input_c_data, params, output_desc, *output_data, |
| 1940 | output_h_desc, *output_h_data, output_c_desc, *output_c_data)); |
| 1941 | |
| 1942 | auto cudnn = cudnn_->GetHandle(parent_, stream); |
| 1943 | |
| 1944 | SE_RETURN_IF_ERROR(CheckRNNParameterSize(cudnn, rnn_desc, input_desc)); |
| 1945 | SE_ASSIGN_OR_RETURN(DeviceMemory<uint8> workspace, |
| 1946 | CreateRnnWorkspace(stream, cudnn, rnn_desc, input_desc, |
| 1947 | workspace_allocator)) |
| 1948 | |
| 1949 | // query the reserve space size |
| 1950 | // allocate the reserve space |
| 1951 | DeviceMemory<uint8> reserve_space; |
| 1952 | if (is_training) { |
| 1953 | size_t reserve_space_size_in_bytes = 0; |
| 1954 | RETURN_IF_CUDNN_ERROR(cudnnGetRNNTrainingReserveSize( |
| 1955 | /*handle=*/cudnn.handle(), /*rnnDesc=*/rnn_desc.handle(), |
| 1956 | /*seqLength=*/model_dims.max_seq_length, /*xDesc=*/input_desc.handles(), |
| 1957 | /*sizeInBytes=*/&reserve_space_size_in_bytes)); |
| 1958 | |
| 1959 | if (reserve_space_size_in_bytes > 0) { |
| 1960 | SE_ASSIGN_OR_RETURN(reserve_space, reserve_space_allocator->AllocateBytes( |
| 1961 | reserve_space_size_in_bytes)); |
| 1962 | } |
| 1963 | } |
| 1964 | |
| 1965 | std::unique_ptr<GpuTimer, GpuTimerDeleter> timer; |
| 1966 | const bool is_profiling = output_profile_result != nullptr; |
| 1967 | if (is_profiling) { |
| 1968 | timer.reset(new GpuTimer(parent_)); |
| 1969 | // The start and stop of the timer should be as close to the Cudnn call as |
| 1970 | // possible. It is still possible for other threads to issue workload on |
| 1971 | // to this stream. So it could take multiple profiling measurements. |
| 1972 | if (!timer->Init() || !timer->Start(AsGpuStream(stream))) { |
| 1973 | return port::Status(port::error::INTERNAL, "Failed to start timer"); |
| 1974 | } |
| 1975 | } |
nothing calls this directly
no test coverage detected