mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-18 21:06:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Use stream-synchronized copy in rnn_kernels.cc.

Browse Source 
May fix flaky wrong outputs sometimes seen in CI.

Also check for errors in another use of gpuStreamSynchronize().

PiperOrigin-RevId: 530391917
This commit is contained in:
Peter Hawkins 2023-05-08 13:25:58 -07:00 committed by jax authors
parent 821b38da12
commit 6b9a109939
2 changed files with 5 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
jaxlib/gpu/rnn_kernels.cc
View File

@ -201,9 +201,11 @@ static absl::Status DnnRNNForward_(gpuStream_t stream, void** buffers,
auto seq_lengths_buf = buffers[4];
std::vector<int32_t> seq_length_vector(d.batch_size, 0);
int32_t* seq_length_array = &seq_length_vector[0];
JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpuMemcpy(
JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpuMemcpyAsync(
seq_length_array, seq_lengths_buf,
seq_length_vector.size() * sizeof(int32_t), gpuMemcpyDeviceToHost)));
seq_length_vector.size() * sizeof(int32_t), gpuMemcpyDeviceToHost,
stream)));
JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpuStreamSynchronize(stream)));
cudnnRNNDataDescriptor_t input_data_desc;
JAX_RETURN_IF_ERROR(

2
jaxlib/gpu/solver_kernels.cc
View File

@ -465,7 +465,7 @@ static absl::Status Syevd_(gpuStream_t stream, void** buffers,
reinterpret_cast<const std::int64_t*>(buffers[1]),
sizeof(batch), gpuMemcpyDeviceToHost,
stream);
gpuStreamSynchronize(stream);
JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpuStreamSynchronize(stream)));
output_idx = 2;
}
JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpuMemcpyAsync(