CUDA Kernel on Custom Device - FIX 1 (#72668)

YqGe585 · web-flow · commit fc34fb01446b · 2025-05-21T20:33:51.000+08:00
* seperate blas

fix GetDeviceProperties

small fix

fix dcu error in gpu_types.h

modified PD_REGISTER_KERNEL_FOR_ALL_BACKEND_DTYPE 250513

* fix stack_kernel compile error and modified 3 custom kernel
diff --git a/paddle/phi/backends/gpu/cuda/cuda_info.cc b/paddle/phi/backends/gpu/cuda/cuda_info.cc
@@ -11,6 +11,7 @@
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.
+#include <mutex>
 
 #include "paddle/phi/backends/gpu/gpu_info.h"
 
@@ -24,10 +25,12 @@ static std::vector<phi::gpuDeviceProp> g_device_props;
 
 namespace phi::backends::gpu {
 
+#ifndef PADDLE_WITH_CUSTOM_DEVICE
 int DnnVersion() {
   if (!dynload::HasCUDNN()) return -1;
   return dynload::cudnnGetVersion();  // NOLINT
 }
+#endif
 
 static int GetGPUDeviceCountImpl() {
   int driverVersion = 0;
diff --git a/paddle/phi/backends/gpu/gpu_context.h b/paddle/phi/backends/gpu/gpu_context.h
@@ -18,6 +18,7 @@ limitations under the License. */
 #ifdef PADDLE_WITH_CUSTOM_DEVICE
 #include "paddle/phi/backends/custom/custom_context.h"
 #include "paddle/phi/backends/gpu/gpu_helper.h"
+#include "paddle/phi/backends/gpu/gpu_info.h"
 #else
 #if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP) || \
     defined(PADDLE_WITH_XPU_KP)
diff --git a/paddle/phi/backends/gpu/gpu_types.h b/paddle/phi/backends/gpu/gpu_types.h
@@ -23,8 +23,12 @@
 #include "paddle/phi/backends/dynload/miopen.h"
 #include "paddle/phi/backends/dynload/rocblas.h"
 #else  // PADDLE_WITH_CUDA
+#ifndef PADDLE_WITH_CUSTOM_DEVICE
 #include "paddle/phi/backends/dynload/cublas.h"
 #include "paddle/phi/backends/dynload/cudnn.h"
+#else
+#include <cuda_runtime.h>
+#endif
 #endif
 
 namespace phi {
@@ -45,6 +49,7 @@ namespace phi {
 DECLARE_TYPE_FOR_GPU(gpuError_t, cudaError_t, hipError_t);
 DECLARE_TYPE_FOR_GPU(gpuMemcpyKind, cudaMemcpyKind, hipMemcpyKind);
 DECLARE_TYPE_FOR_GPU(gpuDeviceProp, cudaDeviceProp, hipDeviceProp_t);
+#ifndef PADDLE_WITH_CUSTOM_DEVICE
 DECLARE_TYPE_FOR_GPU(dnnDataType_t, cudnnDataType_t, miopenDataType_t);
 DECLARE_TYPE_FOR_GPU(dnnPoolingMode_t, cudnnPoolingMode_t, miopenPoolingMode_t);
 DECLARE_TYPE_FOR_GPU(dnnTensorFormat_t,
@@ -53,6 +58,7 @@ DECLARE_TYPE_FOR_GPU(dnnTensorFormat_t,
 DECLARE_TYPE_FOR_GPU(dnnActivationMode_t,
                      cudnnActivationMode_t,
                      miopenActivationMode_t);
+#endif
 DECLARE_TYPE_FOR_GPU(gpuGraph_t, cudaGraph_t, hipGraph_t);
 DECLARE_TYPE_FOR_GPU(gpuFunction_t, cudaFunction_t, hipFunction_t);
 DECLARE_TYPE_FOR_GPU(gpuGraphExec_t, cudaGraphExec_t, hipGraphExec_t);
diff --git a/paddle/phi/core/kernel_registry.h b/paddle/phi/core/kernel_registry.h
@@ -1308,13 +1308,49 @@ struct KernelRegistrar {
  * Used to register a instantiated kernel function
  * for all backend with one template argument.
  */
+#if (defined(PADDLE_WITH_CUSTOM_DEVICE) && defined(PADDLE_WITH_CUDA))
+#define PD_REGISTER_KERNEL_FOR_ALL_BACKEND_DTYPE(                        \
+    kernel_name, layout, meta_kernel_fn)                                 \
+  template decltype(meta_kernel_fn<::phi::CustomContext>)                \
+      meta_kernel_fn<::phi::CustomContext>;                              \
+  static void                                                            \
+      __FAKE_PD_KERNEL_args_def_FN_##kernel_name##_##backend##_##layout( \
+          const ::phi::KernelKey kernel_key UNUSED,                      \
+          ::phi::Kernel* kernel UNUSED)
+
+#define PD_CUSTOM_KERNEL_REGISTER_FOR_ALL_DTYPE(                             \
+    kernel_name, backend, layout, meta_kernel_fn)                            \
+  PD_STATIC_ASSERT_GLOBAL_NAMESPACE(                                         \
+      PD_REGISTER_nt_kernel_ns_check_##kernel_name##_##layout,               \
+      "PD_REGISTER_KERNEL_FOR_ALL_BACKEND_DTYPE must be called in global "   \
+      "namespace.");                                                         \
+  static void __PD_KERNEL_args_def_FN_##kernel_name##_##layout(              \
+      const ::phi::KernelKey& kernel_key, ::phi::Kernel* kernel);            \
+  static const ::phi::KernelRegistrar                                        \
+      __reg_phi_kernel_##kernel_name##_##backend##_##layout(                 \
+          ::phi::RegType::OUTER,                                             \
+          #kernel_name,                                                      \
+          #backend,                                                          \
+          DATA_LAYOUT(layout),                                               \
+          ::phi::KernelArgsParseFunctor<                                     \
+              decltype(&meta_kernel_fn<::phi::CustomContext>)>::Parse,       \
+          &__PD_KERNEL_args_def_FN_##kernel_name##_##layout,                 \
+          PHI_KERNEL(meta_kernel_fn<::phi::CustomContext>),                  \
+          PHI_VARIADIC_KERNEL(meta_kernel_fn<::phi::CustomContext>));        \
+  TEST_API int TouchKernelSymbolFor_##kernel_name##_##backend##_##layout() { \
+    return 0;                                                                \
+  }                                                                          \
+  void __PD_KERNEL_args_def_FN_##kernel_name##_##layout(                     \
+      const ::phi::KernelKey& kernel_key UNUSED, ::phi::Kernel* kernel UNUSED)
+#else
 #define PD_REGISTER_KERNEL_FOR_ALL_BACKEND_DTYPE(                  \
     kernel_name, layout, meta_kernel_fn)                           \
   _PD_REGISTER_KERNEL_FOR_ALL_BACKEND_DTYPE(::phi::RegType::INNER, \
                                             kernel_name,           \
                                             layout,                \
                                             meta_kernel_fn,        \
                                             BACKEND_LIST)
+#endif
 
 #define PD_REGISTER_KERNEL_FOR_ALL_BACKEND_DTYPE_EXCEPT_CUSTOM(    \
     kernel_name, layout, meta_kernel_fn)                           \
diff --git a/paddle/phi/core/platform/device_context.h b/paddle/phi/core/platform/device_context.h
@@ -108,7 +108,8 @@ class IPUDeviceContext
 namespace xpu = baidu::xpu::api;
 #endif
 
-#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+#if (defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)) && \
+    !defined(PADDLE_WITH_CUSTOM_DEVICE)
 using CUDAPinnedDeviceContext = phi::GPUPinnedContext;
 #endif
 
@@ -137,7 +138,8 @@ struct DefaultDeviceContextType<phi::IPUPlace> {
 };
 #endif
 
-#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+#if (defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)) && \
+    !defined(PADDLE_WITH_CUSTOM_DEVICE)
 template <>
 struct DefaultDeviceContextType<phi::GPUPinnedPlace> {
   using TYPE = paddle::platform::CUDAPinnedDeviceContext;
diff --git a/paddle/phi/kernels/empty_kernel.cc b/paddle/phi/kernels/empty_kernel.cc
@@ -149,7 +149,7 @@ PD_REGISTER_KERNEL(empty_like,
 }
 #endif
 
-#ifdef PADDLE_WITH_CUSTOM_DEVICE
+#if defined(PADDLE_WITH_CUSTOM_DEVICE) && !defined(PADDLE_WITH_CUDA)
 PD_REGISTER_KERNEL(empty,
                    Custom,
                    ALL_LAYOUT,
diff --git a/paddle/phi/kernels/flatten_grad_kernel.cc b/paddle/phi/kernels/flatten_grad_kernel.cc
@@ -84,7 +84,7 @@ PD_REGISTER_KERNEL(flatten_grad,
 
 #endif
 
-#ifdef PADDLE_WITH_CUSTOM_DEVICE
+#if defined(PADDLE_WITH_CUSTOM_DEVICE) && !defined(PADDLE_WITH_CUDA)
 PD_REGISTER_KERNEL(flatten_grad,
                    Custom,
                    ALL_LAYOUT,
diff --git a/paddle/phi/kernels/flatten_kernel.cc b/paddle/phi/kernels/flatten_kernel.cc
@@ -139,7 +139,7 @@ PD_REGISTER_KERNEL(flatten_with_xshape,
                    bool) {}
 #endif
 
-#ifdef PADDLE_WITH_CUSTOM_DEVICE
+#if defined(PADDLE_WITH_CUSTOM_DEVICE) && !defined(PADDLE_WITH_CUDA)
 PD_REGISTER_KERNEL(flatten,
                    Custom,
                    ALL_LAYOUT,
diff --git a/paddle/phi/kernels/funcs/math_function.cu b/paddle/phi/kernels/funcs/math_function.cu
@@ -20,9 +20,14 @@ limitations under the License. */
 #include "paddle/phi/common/data_type.h"
 #include "paddle/phi/common/float16.h"
 #include "paddle/phi/common/memory_utils.h"
-#include "paddle/phi/kernels/funcs/blas/blas.h"
 #include "paddle/phi/kernels/funcs/math_function.h"
 #include "paddle/phi/kernels/funcs/math_function_impl.h"
+#ifndef PADDLE_WITH_CUSTOM_DEVICE
+#include "paddle/phi/kernels/funcs/blas/blas.h"
+#include "paddle/phi/kernels/funcs/math_function_blas_impl.h"
+#else
+#include "paddle/phi/backends/gpu/gpu_info.h"
+#endif
 
 namespace phi {
 namespace funcs {
@@ -157,6 +162,7 @@ template struct SetConstant<phi::GPUContext, bool>;
 template struct SetConstant<phi::GPUContext, phi::dtype::complex<float>>;
 template struct SetConstant<phi::GPUContext, phi::dtype::complex<double>>;
 
+#ifndef PADDLE_WITH_CUSTOM_DEVICE
 template struct SetConstant<phi::GPUPinnedContext, float16>;
 template struct SetConstant<phi::GPUPinnedContext, bfloat16>;
 template struct SetConstant<phi::GPUPinnedContext, float>;
@@ -169,6 +175,7 @@ template struct SetConstant<phi::GPUPinnedContext, int64_t>;
 template struct SetConstant<phi::GPUPinnedContext, bool>;
 template struct SetConstant<phi::GPUPinnedContext, phi::dtype::complex<float>>;
 template struct SetConstant<phi::GPUPinnedContext, phi::dtype::complex<double>>;
+#endif
 
 #define DEFINE_GPU_TRANS(RANK)                                     \
   template struct Transpose<phi::GPUContext, bool, RANK>;          \
@@ -422,79 +429,8 @@ template struct RowwiseAdd<phi::GPUContext, double>;
 template struct ColwiseSum<phi::GPUContext, float>;
 template struct ColwiseSum<phi::GPUContext, int>;
 template struct ColwiseSum<phi::GPUContext, int64_t>;
-// template struct ColwiseSum<phi::GPUContext, double>;
-// The ColwiseSum<phi::GPUContext, double> failed in debug
-// mode,
-// and only failed for this case. So reimplemented it.
-template <>
-void ColwiseSum<phi::GPUContext, double>::operator()(
-    const phi::GPUContext& context,
-    const phi::DenseTensor& input,
-    phi::DenseTensor* vector) {
-  auto in_dims = input.dims();
-  auto size = input.numel() / in_dims[0];
-  PADDLE_ENFORCE_EQ(vector->numel(),
-                    size,
-                    common::errors::InvalidArgument(
-                        "The size of input vector"
-                        " should be equal to the size of input tensor column"
-                        " dimension. Expected vector size=%d, but received %d",
-                        size,
-                        vector->numel()));
-  phi::DenseTensor one;
-  one.Resize({in_dims[0]});
-  context.template Alloc<double>(&one);
-
-  SetConstant<phi::GPUContext, double> set;
-  set(context, &one, static_cast<double>(1.0));
-  phi::funcs::GetBlas<phi::GPUContext, double>(context).GEMV(
-      true,
-      static_cast<int>(in_dims[0]),
-      static_cast<int>(in_dims[1]),
-      1.0,
-      input.data<double>(),
-      one.data<double>(),
-      0.0,
-      vector->data<double>());
-}
 
 template struct RowwiseSum<phi::GPUContext, float>;
-// template struct RowwiseSum<phi::GPUContext, double>;
-// TODO(zcd): Following ColwiseSum format, need to confirm.
-// The RowwiseSum<phi::GPUContext, double> failed in debug
-// mode,
-// and only failed for this case. So reimplemented it.
-template <>
-void RowwiseSum<phi::GPUContext, double>::operator()(
-    const phi::GPUContext& context,
-    const phi::DenseTensor& input,
-    phi::DenseTensor* vector) {
-  auto in_dims = input.dims();
-  auto size = input.numel() / in_dims[0];
-  PADDLE_ENFORCE_EQ(vector->numel(),
-                    in_dims[0],
-                    common::errors::InvalidArgument(
-                        "The size of input vector"
-                        " should be equal to the size of input tensor row"
-                        " dimension. Expected vector size=%d, but received %d",
-                        in_dims[0],
-                        vector->numel()));
-  phi::DenseTensor one;
-  one.Resize({size});
-  context.template Alloc<double>(&one);
-
-  SetConstant<phi::GPUContext, double> set;
-  set(context, &one, static_cast<double>(1.0));
-  phi::funcs::GetBlas<phi::GPUContext, double>(context).GEMV(
-      true,
-      static_cast<int>(in_dims[1]),
-      static_cast<int>(in_dims[0]),
-      1.0,
-      one.data<double>(),
-      input.data<double>(),
-      0.0,
-      vector->data<double>());
-}
 
 template struct RowwiseMean<phi::GPUContext, float>;
 template struct RowwiseMean<phi::GPUContext, double>;
diff --git a/paddle/phi/kernels/funcs/math_function_blas_impl.h b/paddle/phi/kernels/funcs/math_function_blas_impl.h
@@ -0,0 +1,102 @@
+/* Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+#pragma once
+
+#include <algorithm>
+#include <vector>
+
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/common/bfloat16.h"
+#include "paddle/phi/common/data_type.h"
+#include "paddle/phi/common/float16.h"
+#include "paddle/phi/common/memory_utils.h"
+#include "paddle/phi/kernels/funcs/blas/blas.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
+
+namespace phi {
+namespace funcs {
+// template struct ColwiseSum<phi::GPUContext, double>;
+// The ColwiseSum<phi::GPUContext, double> failed in debug
+// mode,
+// and only failed for this case. So reimplemented it.
+template <>
+void ColwiseSum<phi::GPUContext, double>::operator()(
+    const phi::GPUContext& context,
+    const phi::DenseTensor& input,
+    phi::DenseTensor* vector) {
+  auto in_dims = input.dims();
+  auto size = input.numel() / in_dims[0];
+  PADDLE_ENFORCE_EQ(vector->numel(),
+                    size,
+                    common::errors::InvalidArgument(
+                        "The size of input vector"
+                        " should be equal to the size of input tensor column"
+                        " dimension. Expected vector size=%d, but received %d",
+                        size,
+                        vector->numel()));
+  phi::DenseTensor one;
+  one.Resize({in_dims[0]});
+  context.template Alloc<double>(&one);
+
+  SetConstant<phi::GPUContext, double> set;
+  set(context, &one, static_cast<double>(1.0));
+  phi::funcs::GetBlas<phi::GPUContext, double>(context).GEMV(
+      true,
+      static_cast<int>(in_dims[0]),
+      static_cast<int>(in_dims[1]),
+      1.0,
+      input.data<double>(),
+      one.data<double>(),
+      0.0,
+      vector->data<double>());
+}
+
+// template struct RowwiseSum<phi::GPUContext, double>;
+// TODO(zcd): Following ColwiseSum format, need to confirm.
+// The RowwiseSum<phi::GPUContext, double> failed in debug
+// mode,
+template <>
+void RowwiseSum<phi::GPUContext, double>::operator()(
+    const phi::GPUContext& context,
+    const phi::DenseTensor& input,
+    phi::DenseTensor* vector) {
+  auto in_dims = input.dims();
+  auto size = input.numel() / in_dims[0];
+  PADDLE_ENFORCE_EQ(vector->numel(),
+                    in_dims[0],
+                    common::errors::InvalidArgument(
+                        "The size of input vector"
+                        " should be equal to the size of input tensor row"
+                        " dimension. Expected vector size=%d, but received %d",
+                        in_dims[0],
+                        vector->numel()));
+  phi::DenseTensor one;
+  one.Resize({size});
+  context.template Alloc<double>(&one);
+
+  SetConstant<phi::GPUContext, double> set;
+  set(context, &one, static_cast<double>(1.0));
+  phi::funcs::GetBlas<phi::GPUContext, double>(context).GEMV(
+      true,
+      static_cast<int>(in_dims[1]),
+      static_cast<int>(in_dims[0]),
+      1.0,
+      one.data<double>(),
+      input.data<double>(),
+      0.0,
+      vector->data<double>());
+}
+
+}  // namespace funcs
+}  // namespace phi
diff --git a/paddle/phi/kernels/funcs/selected_rows_functor.cc b/paddle/phi/kernels/funcs/selected_rows_functor.cc
@@ -11,14 +11,14 @@ distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License. */
-
-#include "paddle/phi/kernels/funcs/selected_rows_functor.h"
-
 #include <algorithm>
 #include <map>
 #include <set>
 #include <vector>
 
+#include "paddle/phi/kernels/funcs/blas/blas.h"
+#include "paddle/phi/kernels/funcs/selected_rows_functor.h"
+
 #include "paddle/common/ddim.h"
 #include "paddle/phi/core/mixed_vector.h"
 
diff --git a/paddle/phi/kernels/funcs/selected_rows_functor.h b/paddle/phi/kernels/funcs/selected_rows_functor.h
@@ -19,7 +19,6 @@ limitations under the License. */
 #include "paddle/phi/backends/all_context.h"
 #include "paddle/phi/core/mixed_vector.h"
 #include "paddle/phi/core/selected_rows.h"
-#include "paddle/phi/kernels/funcs/blas/blas.h"
 #include "paddle/phi/kernels/funcs/eigen/common.h"
 #include "paddle/phi/kernels/funcs/math_function.h"
 
diff --git a/paddle/phi/kernels/impl/activation_impl.h b/paddle/phi/kernels/impl/activation_impl.h
diff --git a/paddle/phi/kernels/shape_kernel.cc b/paddle/phi/kernels/shape_kernel.cc

Original file line number	Diff line number	Diff line change
`@@ -149,7 +149,7 @@ PD_REGISTER_KERNEL(empty_like,`
`149`	`149`	`}`
`150`	`150`	`#endif`
`151`	`151`
`152`		`-#ifdef PADDLE_WITH_CUSTOM_DEVICE`
	`152`	`+#if defined(PADDLE_WITH_CUSTOM_DEVICE) && !defined(PADDLE_WITH_CUDA)`
`153`	`153`	`PD_REGISTER_KERNEL(empty,`
`154`	`154`	`Custom,`
`155`	`155`	`ALL_LAYOUT,`