CMakeLists.txt 790 字节
Newer Older
W
Wu Yi 已提交
1
include(operators)
P
peizhilin 已提交
2 3 4 5 6
if(WITH_GPU)
    register_operators(DEPS cub)
else()
    register_operators()
endif()
W
Wu Yi 已提交
7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24

if(WITH_GPU)
    file(GLOB OPS RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}" "*.part.cu")
    string(REPLACE ".part.cu" "" OPS "${OPS}")

    foreach(src ${OPS})
        if (EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/${src}.part.cu)
            set(CUDA_KERNEL_FILE ${CMAKE_CURRENT_SOURCE_DIR}/${src}.part.cu)
            file(READ ${CUDA_KERNEL_FILE} TARGET_CONTENT)
            string(REGEX MATCH "REGISTER_OP_CUDA_KERNEL\\(\\n?([^,]+),.*" MATCHED ${TARGET_CONTENT})
            if (MATCHED)
                string(STRIP ${CMAKE_MATCH_1} MATCHED)
                file(APPEND ${pybind_file} "USE_OP_DEVICE_KERNEL(${MATCHED}, CUDA);\n")
            endif()

        endif()
    endforeach()
endif()