From 0d1695247018a7db3446a6339053ff13c55e7c08 Mon Sep 17 00:00:00 2001 From: Megvii Engine Team Date: Thu, 25 Nov 2021 14:05:00 +0800 Subject: [PATCH] fix(mgb/cuda): fix conv error when the input tensor is too large GitOrigin-RevId: 1b1d693795e665630a65a28997078694a78cb214 --- dnn/src/cuda/conv_bias/cudnn_conv.cpp | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/dnn/src/cuda/conv_bias/cudnn_conv.cpp b/dnn/src/cuda/conv_bias/cudnn_conv.cpp index 2d11b7f9..5254974a 100644 --- a/dnn/src/cuda/conv_bias/cudnn_conv.cpp +++ b/dnn/src/cuda/conv_bias/cudnn_conv.cpp @@ -40,6 +40,12 @@ bool ConvBiasForwardImpl::AlgoCUDNNConv::is_available(const SizeArgs& args) cons return false; } + // In conv_args.init_conv_desc will call cudnnSetTensor4dDescriptorEx(),which can't + // been supported when total_nr_elems() > 2 ^ 31 + if (args.src_layout->total_nr_elems() > INT_MAX || + args.dst_layout->total_nr_elems() > INT_MAX) { + return false; + } auto dst_layout = *args.dst_layout; if (dst_layout.dtype.enumv() != args.bias_layout->dtype.enumv()) { dst_layout.dtype = DType();