[PHI]add check for cudnn big tensor (#73194)

Hongqing-work · web-flow · commit b898ab7948b6 · 2025-06-10T11:00:21.000+08:00
* [PHI]add check for cudnn big tensor

* fix hip compile
diff --git a/paddle/phi/backends/gpu/cuda/cudnn_workspace_helper.h b/paddle/phi/backends/gpu/cuda/cudnn_workspace_helper.h
@@ -18,6 +18,18 @@ namespace phi {
 namespace backends {
 namespace gpu {
 
+#define CUDNN_ENFORCE_TENSOR_SIZE_SUPPORTED(tensor)                            \
+  do {                                                                         \
+    int64_t largest = (1LL << 31) - 1;                                         \
+    PADDLE_ENFORCE_LE(                                                         \
+        tensor.numel(),                                                        \
+        largest,                                                               \
+        ::common::errors::PreconditionNotMet(                                  \
+            "The element size of " #tensor " should be <= INT_MAX(2147483647)" \
+            ", but got %lld",                                                  \
+            tensor.numel()));                                                  \
+  } while (0)
+
 static constexpr int kDefaultConvWorkspaceSizeLimitMB = 512;
 
 int GetDefaultConvWorkspaceSizeLimitMB();
diff --git a/paddle/phi/kernels/gpudnn/conv_grad_kernel.cu b/paddle/phi/kernels/gpudnn/conv_grad_kernel.cu
@@ -641,6 +641,9 @@ void ConvCudnnGradKernel(const Context& dev_ctx,
                  ? phi::backends::gpu::DataLayout::kNDHWC
                  : phi::backends::gpu::DataLayout::kNCDHW;
   }
+  CUDNN_ENFORCE_TENSOR_SIZE_SUPPORTED(transformed_input);
+  CUDNN_ENFORCE_TENSOR_SIZE_SUPPORTED(transformed_filter_channel);
+  CUDNN_ENFORCE_TENSOR_SIZE_SUPPORTED(transformed_output_grad_channel);
 
 #ifdef PADDLE_WITH_CUDNN_FRONTEND
   if (dynload::IsCudnnFrontendEnabled() && (groups == 1))
diff --git a/paddle/phi/kernels/gpudnn/conv_kernel.cu b/paddle/phi/kernels/gpudnn/conv_kernel.cu
@@ -495,6 +495,9 @@ void ConvCudnnKernel(const Context& dev_ctx,
                  : phi::backends::gpu::DataLayout::kNCDHW;
   }
 
+  CUDNN_ENFORCE_TENSOR_SIZE_SUPPORTED(transformed_input);
+  CUDNN_ENFORCE_TENSOR_SIZE_SUPPORTED(transformed_filter_channel);
+  CUDNN_ENFORCE_TENSOR_SIZE_SUPPORTED(transformed_output);
 #ifdef PADDLE_WITH_CUDNN_FRONTEND
   if (dynload::IsCudnnFrontendEnabled() && (groups == 1))
     ConvCudnnKernelImplV8<T>(&transformed_input,
diff --git a/paddle/phi/kernels/gpudnn/conv_transpose_grad_kernel.cu b/paddle/phi/kernels/gpudnn/conv_transpose_grad_kernel.cu
@@ -34,6 +34,7 @@ limitations under the License. */
 #include "paddle/phi/kernels/gpudnn/conv_miopen_helper.h"
 #else
 #include "paddle/phi/backends/gpu/cuda/cudnn_helper.h"
+#include "paddle/phi/backends/gpu/cuda/cudnn_workspace_helper.h"
 #include "paddle/phi/kernels/gpudnn/conv_cudnn_v7.h"
 #endif
 
@@ -201,6 +202,9 @@ void ConvTransposeGradRawGPUDNNKernel(const Context& dev_ctx,
   SearchResult<miopenConvFwdAlgorithm_t> fwd_result;
   SearchResult<miopenConvBwdWeightsAlgorithm_t> filter_result;
 #else
+  CUDNN_ENFORCE_TENSOR_SIZE_SUPPORTED(transformed_dout);
+  CUDNN_ENFORCE_TENSOR_SIZE_SUPPORTED(filter);
+  CUDNN_ENFORCE_TENSOR_SIZE_SUPPORTED(x_transpose);
   SearchResult<cudnnConvolutionFwdAlgo_t> fwd_result;
   SearchResult<cudnnConvolutionBwdFilterAlgo_t> filter_result;
 #endif
diff --git a/paddle/phi/kernels/gpudnn/conv_transpose_kernel.cu b/paddle/phi/kernels/gpudnn/conv_transpose_kernel.cu
@@ -32,6 +32,7 @@ limitations under the License. */
 #include "paddle/phi/kernels/gpudnn/conv_miopen_helper.h"
 #else
 #include "paddle/phi/backends/gpu/cuda/cudnn_helper.h"
+#include "paddle/phi/backends/gpu/cuda/cudnn_workspace_helper.h"
 #include "paddle/phi/kernels/gpudnn/conv_cudnn_v7.h"
 #endif
 
@@ -228,6 +229,9 @@ void ConvTransposeRawGPUDNNKernel(const Context& dev_ctx,
   bwd_result.algo =
       search::Find<T>(args, false, deterministic, workspace_size, dev_ctx);
 #else
+  CUDNN_ENFORCE_TENSOR_SIZE_SUPPORTED(transformed_x);
+  CUDNN_ENFORCE_TENSOR_SIZE_SUPPORTED(filter);
+  CUDNN_ENFORCE_TENSOR_SIZE_SUPPORTED(transformed_out);
   SearchResult<cudnnConvolutionBwdDataAlgo_t> bwd_result;
   using search = SearchAlgorithm<ConvKind::kBackwardData>;
   bwd_result = search::Find<T>(dev_ctx, args, false, deterministic, false);

Original file line number	Diff line number	Diff line change
`@@ -641,6 +641,9 @@ void ConvCudnnGradKernel(const Context& dev_ctx,`
`641`	`641`	`? phi::backends::gpu::DataLayout::kNDHWC`
`642`	`642`	`: phi::backends::gpu::DataLayout::kNCDHW;`
`643`	`643`	`}`
	`644`	`+ CUDNN_ENFORCE_TENSOR_SIZE_SUPPORTED(transformed_input);`
	`645`	`+ CUDNN_ENFORCE_TENSOR_SIZE_SUPPORTED(transformed_filter_channel);`
	`646`	`+ CUDNN_ENFORCE_TENSOR_SIZE_SUPPORTED(transformed_output_grad_channel);`
`644`	`647`
`645`	`648`	`#ifdef PADDLE_WITH_CUDNN_FRONTEND`
`646`	`649`	`if (dynload::IsCudnnFrontendEnabled() && (groups == 1))`
Original file line number	Diff line number	Diff line change
`@@ -495,6 +495,9 @@ void ConvCudnnKernel(const Context& dev_ctx,`
`495`	`495`	`: phi::backends::gpu::DataLayout::kNCDHW;`
`496`	`496`	`}`
`497`	`497`
	`498`	`+ CUDNN_ENFORCE_TENSOR_SIZE_SUPPORTED(transformed_input);`
	`499`	`+ CUDNN_ENFORCE_TENSOR_SIZE_SUPPORTED(transformed_filter_channel);`
	`500`	`+ CUDNN_ENFORCE_TENSOR_SIZE_SUPPORTED(transformed_output);`
`498`	`501`	`#ifdef PADDLE_WITH_CUDNN_FRONTEND`
`499`	`502`	`if (dynload::IsCudnnFrontendEnabled() && (groups == 1))`
`500`	`503`	`ConvCudnnKernelImplV8<T>(&transformed_input,`