Fix dnn related tests for custom device (PaddlePaddle#75609)

* fix cudnn related tests for custom device

* fix compile error

* add dnn available for custom device

* add header

* fix index_add

Author: YqGe585

paddle/phi/backends/custom/custom_device.cc

@@ -648,6 +648,16 @@ class CustomDevice : public DeviceInterface {
     return supported;
   }
 
+  bool IsDnnAvailable(size_t dev_id) override {
+    const auto device = &devices_pool[dev_id];
+    bool supported = false;
+    if (pimpl_->is_dnn_supported) {
+      pimpl_->is_dnn_supported(device, &supported);
+    }
+    VLOG(10) << Type() << " is dnn available: " << supported;
+    return supported;
+  }
+
   void* InitEigenDevice(const Place& place,
                         phi::stream::stream_t stream,
                         phi::Allocator* allocator) override {

paddle/phi/backends/device_base.cc

@@ -83,6 +83,11 @@ bool DeviceInterface::IsBFloat16Supported(size_t dev_id) {
   return false;
 }
 
+bool DeviceInterface::IsDnnAvailable(size_t dev_id) {
+  VLOG(10) << Type() << " is dnn available: " << false;
+  return false;
+}
+
 void* DeviceInterface::InitEigenDevice(const Place& place,
                                        phi::stream::stream_t stream,
                                        phi::Allocator* allocator) {

paddle/phi/backends/device_base.h

@@ -83,6 +83,8 @@ class DeviceInterface {  // Driver / Runtime
 
   virtual bool IsBFloat16Supported(size_t dev_id);
 
+  virtual bool IsDnnAvailable(size_t dev_id);
+
   virtual void* InitEigenDevice(const Place& place,
                                 phi::stream::stream_t stream,
                                 phi::Allocator* allocator);

paddle/phi/backends/device_ext.h

@@ -608,6 +608,13 @@ struct C_DeviceInterface {
    */
   C_Status (*is_bfloat16_supported)(const C_Device device, bool* supported);
 
+  /**
+   * @brief Is dnn supported
+   *
+   * @param[C_Device, bool*]     device, supported
+   */
+  C_Status (*is_dnn_supported)(const C_Device device, bool* supported);
+
   /**
    * @brief init eigen device
    *

paddle/phi/backends/device_manager.cc

@@ -537,6 +537,13 @@ bool DeviceManager::IsBFloat16Supported(const Place& place) {
   return dev_impl->IsBFloat16Supported(device_id);
 }
 
+bool DeviceManager::IsDnnAvailable(const Place& place) {
+  auto device_type = place.GetDeviceType();
+  auto device_id = place.GetDeviceId();
+  auto dev_impl = GetDeviceInterfaceWithType(device_type);
+  return dev_impl->IsDnnAvailable(device_id);
+}
+
 void* DeviceManager::InitEigenDevice(const Place& place,
                                      phi::stream::stream_t stream,
                                      phi::Allocator* allocator) {

paddle/phi/backends/device_manager.h

@@ -190,6 +190,8 @@ class PADDLE_API DeviceManager {
 
   static bool IsBFloat16Supported(const Place& place);
 
+  static bool IsDnnAvailable(const Place& place);
+
   static void* InitEigenDevice(const Place& place,
                                phi::stream::stream_t stream,
                                phi::Allocator* allocator);

paddle/phi/kernels/gpu/grid_sample_utils.h

@@ -16,6 +16,9 @@
 
 #include <limits.h>
 
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+#include "paddle/phi/backends/device_manager.h"
+#endif
 namespace phi {
 
 enum class Mode {
@@ -42,7 +45,22 @@ static __forceinline__ __device__ bool InBounds3D(
 }
 
 inline bool cudnnIsAvailable() {
-#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+#if defined(PADDLE_WITH_CUSTOM_DEVICE)
+  // Get all custom device types
+  auto custom_device_types = phi::DeviceManager::GetAllCustomDeviceTypes();
+
+  // Use the first custom device type
+  if (!custom_device_types.empty()) {
+    const std::string& device_type = custom_device_types[0];
+    // Get current device ID for this device type
+    int device_id = phi::DeviceManager::GetDevice(device_type);
+    // Create place for the current device
+    phi::Place place(phi::CustomPlace(device_type, device_id));
+    // Check if this device has DNN support
+    return phi::DeviceManager::IsDnnAvailable(place);
+  }
+  return false;
+#elif defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
   // cuDNN/MIOpen version > 0 means DNN lib loaded; require v7+ for sampler
   return phi::backends::gpu::DnnVersion() >= 7000;
 #else

test/legacy_test/op_test.py

@@ -467,6 +467,60 @@ def is_custom_device():
     return False
 
 
+def check_cudnn_version_and_compute_capability(
+    min_cudnn_version=None, min_device_capability=None
+):
+    """
+    Check if the current environment meets the specified cuDNN version and device capability requirements.
+
+    Args:
+        min_cudnn_version (int, optional): Minimum required cuDNN version. If None, cuDNN version check is skipped.
+        min_device_capability (int, optional): Minimum required device capability. If None, device capability check is skipped.
+
+    Returns:
+        bool: True if the environment meets the requirements or if using custom device, False otherwise.
+    """
+    if is_custom_device():
+        return True
+
+    if not core.is_compiled_with_cuda():
+        return False
+
+    # Check cuDNN version if specified
+    cudnn_check = True
+    if min_cudnn_version is not None:
+        cudnn_check = core.cudnn_version() >= min_cudnn_version
+
+    # Check device capability if specified
+    device_check = True
+    if min_device_capability is not None:
+        device_check = (
+            paddle.device.cuda.get_device_capability()[0]
+            >= min_device_capability
+        )
+
+    return cudnn_check and device_check
+
+
+def get_cuda_version():
+    if paddle.is_compiled_with_cuda():
+        import re
+
+        result = os.popen("nvcc --version").read()
+        regex = r'release (\S+),'
+        match = re.search(regex, result)
+        if match:
+            num = str(match.group(1))
+            integer, decimal = num.split('.')
+            return int(integer) * 1000 + int(float(decimal) * 10)
+        else:
+            return -1
+    elif is_custom_device():
+        return 13000
+    else:
+        return -1
+
+
 @contextmanager
 def auto_parallel_test_guard(test_info_path, generated_test_file_path):
     test_info_file, generated_test_file = None, None

test/legacy_test/test_elementwise_max_op.py

@@ -17,13 +17,12 @@
 import numpy as np
 from op_test import (
     OpTest,
+    check_cudnn_version_and_compute_capability,
     convert_float_to_uint16,
-    is_custom_device,
     skip_check_grad_ci,
 )
 
 import paddle
-from paddle.base import core
 
 
 class TestElementwiseOp(OpTest):
@@ -169,12 +168,8 @@ def init_data(self):
 
 
 @unittest.skipIf(
-    (core.is_compiled_with_cuda() or is_custom_device())
-    and (
-        core.cudnn_version() < 8100
-        or paddle.device.cuda.get_device_capability()[0] < 8
-    ),
-    "run test when gpu is available and the minimum cudnn version is 8.1.0 and gpu's compute capability is at least 8.0.",
+    not check_cudnn_version_and_compute_capability(8100, 8),
+    "only support compiled with CUDA or custom device, and for CUDA cudnn version need larger than 8.1.0 and device's compute capability is at least 8.0",
 )
 class TestElementwiseBF16Op(OpTest):
     def init_data(self):

test/legacy_test/test_elementwise_tensor_split.py

@@ -32,7 +32,7 @@ def test_float16_sub(self):
         if not (core.is_compiled_with_cuda() or is_custom_device()):
             return
 
-        gpu_info = paddle.device.cuda.get_device_properties()
+        gpu_info = paddle.device.get_device_properties()
 
         gpu_name = gpu_info.name
         try: