[ROCm] fix missing dcu kernel in operator.cmake, test=develop (Paddle…

…Paddle#39480)
liqitong-a · Feb 14, 2022 · 55da934 · 55da934
1 parent ddb1e23
commit 55da934
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/cmake/operators.cmake b/cmake/operators.cmake
@@ -335,6 +335,17 @@ function(op_library TARGET)
         endif()
     endforeach()
 
+    # pybind USE_OP_DEVICE_KERNEL for ROCm
+    list (APPEND hip_srcs ${hip_cc_srcs})
+    # message("hip_srcs ${hip_srcs}")
+    foreach(hip_src ${hip_srcs})
+        set(op_name "")
+        find_register(${hip_src} "REGISTER_OP_CUDA_KERNEL" op_name)
+        if(NOT ${op_name} EQUAL "")
+            file(APPEND ${pybind_file} "USE_OP_DEVICE_KERNEL(${op_name}, CUDA);\n")
+            set(pybind_flag 1)
+        endif()
+    endforeach()
 
     # pybind USE_OP_DEVICE_KERNEL for CUDNN/MIOPEN
     list(APPEND cudnn_cu_srcs ${cudnn_cu_cc_srcs})