From 83de50a83326b2fc7999c60b30fa930f9fb5c0f6 Mon Sep 17 00:00:00 2001 From: srinivamd <52507740+srinivamd@users.noreply.github.com> Date: Thu, 30 Apr 2026 07:19:03 -0700 Subject: [PATCH 1/3] Fix miopenStatusUnknownError in RNN kernels: remove premature input_tensor_desc destruction PR#726 ("Fix HIP memory leaks in RNN kernels") introduced gpudnnDestroyTensorDescriptor(input_tensor_desc) in DnnRNNForward_ and DnnRNNBackward_ cleanup blocks. MIOpen requires input_tensor_desc to remain valid through gpudnnRNNBackwardWeights because it uses the descriptor as part of its execution buffer; destroying it early triggers miopenStatusUnknownError on gfx1201 (RDNA4). Remove the two premature destroy calls. The descriptor is stack-allocated so it is freed when the function returns. Fixes: ROCM-21859 (test_lstm1, test_lstm9 on gfx1201) Regressed-by: ROCm/jax#726 Equivalent to: ROCm/jax#729 (targeting v0.9.0) --- jaxlib/gpu/rnn_kernels.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/jaxlib/gpu/rnn_kernels.cc b/jaxlib/gpu/rnn_kernels.cc index 5020c9a1d36f..3b594ed2b8a4 100644 --- a/jaxlib/gpu/rnn_kernels.cc +++ b/jaxlib/gpu/rnn_kernels.cc @@ -356,11 +356,11 @@ static absl::Status DnnRNNForward_(gpuStream_t stream, void** buffers, JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnDestroyRNNDescriptor(rnn_desc))); #ifdef JAX_GPU_HIP JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpuFree(dropout_states_dev))); - JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnDestroyTensorDescriptor(input_tensor_desc))); #endif return absl::OkStatus(); } +} static absl::Status DnnRNNBackward_(gpuStream_t stream, void** buffers, const char* opaque, size_t opaque_len) { From 226a8479f504437895ea7fdd6a3f810c12b0c6ec Mon Sep 17 00:00:00 2001 From: srinivamd <52507740+srinivamd@users.noreply.github.com> Date: Thu, 30 Apr 2026 07:19:23 -0700 Subject: [PATCH 2/3] Fix: remove spurious closing brace and remaining premature input_tensor_desc destroy The previous commit accidentally introduced an extra } after DnnRNNForward_. Remove it and also remove the premature gpudnnDestroyTensorDescriptor(input_tensor_desc) from DnnRNNBackward_. --- jaxlib/gpu/rnn_kernels.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/jaxlib/gpu/rnn_kernels.cc b/jaxlib/gpu/rnn_kernels.cc index 3b594ed2b8a4..18831d4b1d9e 100644 --- a/jaxlib/gpu/rnn_kernels.cc +++ b/jaxlib/gpu/rnn_kernels.cc @@ -360,8 +360,8 @@ static absl::Status DnnRNNForward_(gpuStream_t stream, void** buffers, return absl::OkStatus(); } -} +static absl::Status DnnRNNBackward_(gpuStream_t stream, void** buffers, static absl::Status DnnRNNBackward_(gpuStream_t stream, void** buffers, const char* opaque, size_t opaque_len) { auto s = UnpackDescriptor(opaque, opaque_len); From 11ceaea056d46abf5cd679c1dacc04e069148685 Mon Sep 17 00:00:00 2001 From: srinivamd <52507740+srinivamd@users.noreply.github.com> Date: Thu, 30 Apr 2026 07:21:56 -0700 Subject: [PATCH 3/3] jaxlib/gpu/rnn_kernels.cc: remove premature input_tensor_desc destruction (ROCM-21859) PR#726 ("Fix HIP memory leaks in RNN kernels", 2026-03-05) introduced gpudnnDestroyTensorDescriptor(input_tensor_desc) in the cleanup blocks of both DnnRNNForward_ and DnnRNNBackward_. MIOpen requires input_tensor_desc to remain valid throughout gpudnnRNNBackwardWeights because it uses the descriptor as part of its execution buffer; destroying it early triggers miopenStatusUnknownError on gfx1201 (RDNA4, Navi48). Remove the two premature destroy calls. input_tensor_desc is stack-allocated so MIOpen owns no persistent reference beyond the call; it is automatically released when the function returns. Fixes test_lstm1 and test_lstm9 failures on gfx1201 (image 1317). Regressed-by: ROCm/jax#726 Equivalent-to: ROCm/jax#729 (targeting v0.9.0) --- jaxlib/gpu/rnn_kernels.cc | 2 -- 1 file changed, 2 deletions(-) diff --git a/jaxlib/gpu/rnn_kernels.cc b/jaxlib/gpu/rnn_kernels.cc index 18831d4b1d9e..52736202f015 100644 --- a/jaxlib/gpu/rnn_kernels.cc +++ b/jaxlib/gpu/rnn_kernels.cc @@ -361,7 +361,6 @@ static absl::Status DnnRNNForward_(gpuStream_t stream, void** buffers, return absl::OkStatus(); } -static absl::Status DnnRNNBackward_(gpuStream_t stream, void** buffers, static absl::Status DnnRNNBackward_(gpuStream_t stream, void** buffers, const char* opaque, size_t opaque_len) { auto s = UnpackDescriptor(opaque, opaque_len); @@ -545,7 +544,6 @@ static absl::Status DnnRNNBackward_(gpuStream_t stream, void** buffers, JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnDestroyRNNDescriptor(rnn_desc))); #ifdef JAX_GPU_HIP JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpuFree(dropout_states_dev))); - JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnDestroyTensorDescriptor(input_tensor_desc))); #endif return absl::OkStatus();