Fix bn performance degradation (#50287)

* fix bn performance degradation

Fix bn performance degradation (#50287)
* fix bn performance degradation
6f1ec935 · zhangkaihuo · GitHub · 8c14b02b · 6f1ec935
显示空白变更内容
内联并排

Showing with 2 addition and 1 deletion

paddle/phi/kernels/gpu/batch_norm_grad_kernel.cu paddle/phi/kernels/gpu/batch_norm_grad_kernel.cu +2 -1

未找到文件。
--- a/paddle/phi/kernels/gpu/batch_norm_grad_kernel.cu
+++ b/paddle/phi/kernels/gpu/batch_norm_grad_kernel.cu
@@ -783,7 +783,8 @@ void BatchNormGradRawKernel(const Context &ctx,
    }
    // CUDNN only support small batch size
    bool use_native_nhwc =
-        d_x ? (x_dims.size() == 4 && compute_format == DataLayout::kNHWC)
+        d_x ? (x_dims.size() == 4 && compute_format == DataLayout::kNHWC &&
+               H * W >= CUDNN_SPATIAL_THRESHOLD_EVAL)
            : false;
    const bool use_native_kernel =
        ((x_dims.size() == 2 && N >= CUDNN_PER_ACTIVATION_THRESHOLD) ||