From c1cf16ed5d078b69db865c633e446a6038d22cea Mon Sep 17 00:00:00 2001 From: Ryan Hill <38674843+RyanUnderhill@users.noreply.github.com> Date: Tue, 1 Mar 2022 01:31:57 -0800 Subject: [PATCH] Conv node bug, cached state was incoherent (#10041) * Moved the init earlier to keep the cache coherent * Move setting of w_desc later, and zero shape check later to catch all cacheable changes. * Add comment --- onnxruntime/core/providers/cuda/nn/conv.cc | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/onnxruntime/core/providers/cuda/nn/conv.cc b/onnxruntime/core/providers/cuda/nn/conv.cc index c4c4272c46..a9c67f2297 100644 --- a/onnxruntime/core/providers/cuda/nn/conv.cc +++ b/onnxruntime/core/providers/cuda/nn/conv.cc @@ -176,9 +176,6 @@ Status Conv::UpdateState(OpKernelContext* context, bool bias_expected) const s_.slice_axes = slice_axes; s_.Y = context->Output(0, TensorShape(s_.y_dims)); - if (s_.Y->Shape().Size() == 0) { - return Status::OK(); - } if (post_slicing_required) { // Post slicing needed. Create and fill in the Conv results in an intermediate buffer. s_.memory_for_cudnn_conv_results = GetScratchBuffer(TensorShape(y_dims_with_adjusted_pads).Size() * s_.element_size); @@ -206,9 +203,14 @@ Status Conv::UpdateState(OpKernelContext* context, bool bias_expected) const dilations.push_back(1); } - if (w_dims_changed) { + if (w_dims_changed) ORT_RETURN_IF_ERROR(s_.w_desc.Set(w_dims, CudnnTensor::GetDataType())); + + // We must delay returning early until here so that the weight dims have been cached properly + if (s_.Y->Shape().Size() == 0) { + return Status::OK(); } + ORT_RETURN_IF_ERROR(s_.x_tensor.Set(x_dims_cudnn, CudnnTensor::GetDataType())); ORT_RETURN_IF_ERROR(s_.y_tensor.Set(y_dims_cudnn, CudnnTensor::GetDataType())); ORT_RETURN_IF_ERROR(s_.conv_desc.Set(kernel_shape.size(), pads, strides, dilations,