Split kernel into libphi.so and libphi_gpu_kernel.so (#64944)

risemeup1 · silverling · web-flow · commit 62531707a01e · 2024-06-06T15:44:07.000+08:00
* [HACKATHON 6th] split GPU kernels

* [HACKATHON 6th] split phi/kernels/gpu only

* [HACKATHON 6th] fix setup.py

* [HACKATHON 6th] fix setup.py typo :(

* [HACKATHON 6th] fix inference error

* [HACKATHON 6th] fix inference demo DSO missing error

* [HACKATHON 6th] fix fluid eager test

* [HACKATHON 6th] fix tests

* [HACKATHON 6th] fix static library link dependency

* [HACKATHON 6th] fix setup.py.in

---------

Co-authored-by: Silver Ling &lt;silver.ling@outlook.com&gt;
diff --git a/cmake/inference_lib.cmake b/cmake/inference_lib.cmake
@@ -285,6 +285,14 @@ else()
       inference_lib_dist
       SRCS ${paddle_phi_lib}
       DSTS ${PADDLE_INFERENCE_INSTALL_DIR}/paddle/lib)
+    if(WITH_GPU OR WITH_ROCM)
+      set(paddle_phi_kernel_gpu_lib
+          ${PADDLE_BINARY_DIR}/paddle/phi/libphi_kernel_gpu.*)
+      copy(
+        inference_lib_dist
+        SRCS ${paddle_phi_kernel_gpu_lib}
+        DSTS ${PADDLE_INFERENCE_INSTALL_DIR}/paddle/lib)
+    endif()
   endif()
 endif()
 
diff --git a/paddle/fluid/eager/CMakeLists.txt b/paddle/fluid/eager/CMakeLists.txt
@@ -13,6 +13,10 @@ set(eager_deps
     grad_tensor_holder
     custom_operator_node)
 
+if(WITH_GPU OR WITH_ROCM)
+  set(eager_deps ${eager_deps} phi_kernel_gpu)
+endif()
+
 if(NOT (NOT WITH_PYTHON AND ON_INFER))
   set(eager_deps ${eager_deps} accumulation_node prim_utils)
 endif()
diff --git a/paddle/fluid/inference/api/demo_ci/CMakeLists.txt b/paddle/fluid/inference/api/demo_ci/CMakeLists.txt
@@ -212,6 +212,12 @@ if(NOT WIN32)
     set(DEPS
         ${DEPS} ${PADDLE_LIB}/paddle/lib/libphi${CMAKE_SHARED_LIBRARY_SUFFIX}
         ${PADDLE_LIB}/paddle/lib/libcommon${CMAKE_SHARED_LIBRARY_SUFFIX})
+    if(WITH_GPU OR WITH_ROCM)
+      set(DEPS
+          ${DEPS}
+          ${PADDLE_LIB}/paddle/lib/libphi_kernel_gpu${CMAKE_SHARED_LIBRARY_SUFFIX}
+      )
+    endif()
   endif()
 else()
   set(DEPS
diff --git a/paddle/phi/CMakeLists.txt b/paddle/phi/CMakeLists.txt
@@ -112,6 +112,8 @@ set(PHI_SRCS
     ${infermeta_srcs}
     ${capi_srcs})
 
+set(PHI_KERNEL_GPU_SRCS ${kernels_gpu_srcs})
+
 if(WITH_SHARED_PHI)
   set(PHI_BUILD_TYPE
       SHARED
@@ -205,6 +207,40 @@ set(PHI_LIB
     "${CMAKE_CURRENT_BINARY_DIR}/${PHI_NAME}"
     CACHE FILEPATH "PHI Library" FORCE)
 
+# NOTE(silverling): what we are doing here is to build a library `phi_kernel_gpu`
+# that contains all GPU kernels implementation. This can allow paddle be built
+# with more CUDA archs and reduce the binary size of `phi` library.
+if(WITH_GPU OR WITH_ROCM)
+  if(WITH_GPU)
+    nv_library(
+      phi_kernel_gpu ${PHI_BUILD_TYPE}
+      SRCS ${PHI_KERNEL_GPU_SRCS}
+      DEPS ${PHI_DEPS})
+  elseif(WITH_ROCM)
+    hip_library(
+      phi_kernel_gpu ${PHI_BUILD_TYPE}
+      SRCS ${PHI_KERNEL_GPU_SRCS}
+      DEPS ${PHI_DEPS})
+  endif()
+
+  # NOTE(silverling): making library `phi` depend on `phi_kernel_gpu` (even `phi` does not use it)
+  # will make targets that depend on `phi` also automatically depend on `phi_kernel_gpu`.
+  # This will make users' life easier.
+  target_link_libraries(phi phi_kernel_gpu)
+
+  # NOTE(silverling): `phi_kernel_gpu` needs symbols from `phi`.
+  # When it's a shared library, it will work with no problem.
+  # But when it's a static library, it must be linked to `phi` at link time explicitly.
+  if(NOT WITH_SHARED_PHI)
+    target_link_libraries(phi_kernel_gpu phi)
+  endif()
+
+  string(REPLACE "phi" "phi_kernel_gpu" PHI_KERNEL_GPU_NAME ${PHI_NAME})
+  set(PHI_KERNEL_GPU_LIB
+      "${CMAKE_CURRENT_BINARY_DIR}/${PHI_KERNEL_GPU_NAME}"
+      CACHE FILEPATH "PHI Kernel GPU Library" FORCE)
+endif()
+
 if(MKL_FOUND AND WITH_ONEMKL)
   target_include_directories(phi PRIVATE ${MKL_INCLUDE})
 endif()
diff --git a/paddle/phi/kernels/CMakeLists.txt b/paddle/phi/kernels/CMakeLists.txt
@@ -29,8 +29,6 @@ file(GLOB kernel_primitive_h "primitive/*.h")
 file(
   GLOB kernel_cu
   RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
-  "gpu/*.cu"
-  "gpu/*.cu.cc"
   "gpudnn/*.cu"
   "kps/*.cu"
   "legacy/kps/*.cu"
@@ -40,18 +38,24 @@ file(
   "strings/gpu/*.cu"
   "fusion/gpu/*.cu")
 
+file(
+  GLOB kernel_gpu
+  RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
+  "gpu/*.cu" "gpu/*.cu.cc")
+
 if(APPLE OR WIN32)
   list(REMOVE_ITEM kernel_cu "fusion/gpu/fusion_group_kernel.cu")
   list(REMOVE_ITEM kernel_cu "sparse/gpu/conv_kernel_igemm.cu")
 endif()
 
 if(NOT WITH_DGC)
-  list(REMOVE_ITEM kernel_cu "gpu/dgc_kernel.cu")
+  list(REMOVE_ITEM kernel_gpu "gpu/dgc_kernel.cu")
 endif()
 
 if(DEFINED REDUCE_INFERENCE_LIB_SIZE)
-  list(FILTER kernel_cu EXCLUDE REGEX ".*_grad_kernel\\.cc$")
   list(FILTER kernel_cu EXCLUDE REGEX ".*_grad_kernel\\.cu$")
+  list(FILTER kernel_gpu EXCLUDE REGEX ".*_grad_kernel\\.cc$")
+  list(FILTER kernel_gpu EXCLUDE REGEX ".*_grad_kernel\\.cu$")
 endif()
 
 if(WITH_CUTLASS)
@@ -216,6 +220,15 @@ if(WITH_ROCM)
   list(
     REMOVE_ITEM
     kernel_cu
+    "gpudnn/mha_cudnn_frontend.cu"
+    "fusion/gpu/blha_get_max_len.cu"
+    "fusion/gpu/block_multi_head_attention_kernel.cu"
+    "fusion/gpu/fused_bn_add_activation_grad_kernel.cu"
+    "fusion/gpu/fused_bn_add_activation_kernel.cu"
+    "fusion/gpu/fusion_transpose_flatten_concat_kernel.cu")
+  list(
+    REMOVE_ITEM
+    kernel_gpu
     "gpu/affine_grid_grad_kernel.cu"
     "gpu/apply_per_channel_scale_kernel.cu"
     "gpu/cholesky_solve_kernel.cu"
@@ -228,13 +241,7 @@ if(WITH_ROCM)
     "gpu/put_along_axis_grad_kernel.cu"
     "gpu/put_along_axis_kernel.cu"
     "gpu/qr_kernel.cu"
-    "gpu/svd_kernel.cu"
-    "gpudnn/mha_cudnn_frontend.cu"
-    "fusion/gpu/blha_get_max_len.cu"
-    "fusion/gpu/block_multi_head_attention_kernel.cu"
-    "fusion/gpu/fused_bn_add_activation_grad_kernel.cu"
-    "fusion/gpu/fused_bn_add_activation_kernel.cu"
-    "fusion/gpu/fusion_transpose_flatten_concat_kernel.cu")
+    "gpu/svd_kernel.cu")
 endif()
 
 set(cc_search_pattern
@@ -291,6 +298,8 @@ file(
 if(WITH_GPU OR WITH_ROCM)
   collect_srcs(kernels_srcs SRCS ${kernel_cu})
   kernel_declare("${kernel_cu}")
+  collect_srcs(kernels_gpu_srcs SRCS ${kernel_gpu})
+  kernel_declare("${kernel_gpu}")
 endif()
 
 if(WITH_XPU)
diff --git a/python/env_dict.py.in b/python/env_dict.py.in
@@ -14,6 +14,8 @@ env_dict={
     'FLUID_CORE_NAME':'@FLUID_CORE_NAME@',
     'PHI_LIB':'@PHI_LIB@',
     'PHI_NAME':'@PHI_NAME@',
+    'PHI_KERNEL_GPU_LIB':'@PHI_KERNEL_GPU_LIB@',
+    'PHI_KERNEL_GPU_NAME':'@PHI_KERNEL_GPU_NAME@',
     'WITH_SHARED_PHI':'@WITH_SHARED_PHI@',
     'IR_LIB':'@IR_LIB@',
     'IR_NAME':'@IR_NAME@',
diff --git a/python/setup.py.in b/python/setup.py.in
@@ -710,6 +710,12 @@ package_data['paddle.libs']= []
 if('${WITH_SHARED_PHI}' == 'ON'):
     package_data['paddle.libs'] += [('libphi' if os.name != 'nt' else 'phi') + ext_name]
     shutil.copy('${PHI_LIB}', libs_path)
+    if('${PHI_KERNEL_GPU_LIB}'):
+        package_data['paddle.libs'] += [
+            ('libphi_kernel_gpu' if os.name != 'nt' else 'phi_kernel_gpu')
+            + ext_name
+        ]
+        shutil.copy('${PHI_KERNEL_GPU_LIB}', libs_path)
 
 if('${WITH_SHARED_IR}' == 'ON'):
     package_data['paddle.libs'] += [('libpir' if os.name != 'nt' else 'pir') + ext_name]
diff --git a/setup.py b/setup.py
@@ -1061,6 +1061,12 @@ def get_package_data_and_package_dir():
             ('libphi' if os.name != 'nt' else 'phi') + ext_suffix
         ]
         shutil.copy(env_dict.get("PHI_LIB"), libs_path)
+        if env_dict.get("PHI_KERNEL_GPU_LIB"):
+            package_data['paddle.libs'] += [
+                ('libphi_kernel_gpu' if os.name != 'nt' else 'phi_kernel_gpu')
+                + ext_suffix
+            ]
+            shutil.copy(env_dict.get("PHI_KERNEL_GPU_LIB"), libs_path)
 
     if env_dict.get("WITH_SHARED_IR") == "ON":
         package_data['paddle.libs'] += [
diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt
@@ -227,6 +227,10 @@ if(${len} GREATER_EQUAL 1)
       target_link_libraries(${test_name} $<TARGET_LINKER_FILE:${paddle_lib}>)
       if(WITH_SHARED_PHI)
         target_link_libraries(${test_name} $<TARGET_LINKER_FILE:phi>)
+        if(WITH_GPU OR WITH_ROCM)
+          target_link_libraries(${test_name}
+                                $<TARGET_LINKER_FILE:phi_kernel_gpu>)
+        endif()
       endif()
       if(WITH_SHARED_IR)
         target_link_libraries(${test_name} $<TARGET_LINKER_FILE:pir>)
