|
1 | | -if(WITH_GPU) |
2 | | - nv_library(add_op SRCS add_op.cc add_op.cu DEPS operator op_registry glog ddim) |
3 | | -else() |
4 | | - cc_library(add_op SRCS add_op.cc DEPS operator op_registry glog ddim) |
5 | | -endif() |
| 1 | +function(op_library TARGET) |
| 2 | + # op_library is a function to create op library. The interface is same as |
| 3 | + # cc_library. But it handle split GPU/CPU code and link some common library |
| 4 | + # for ops. |
| 5 | + set(cc_srcs) |
| 6 | + set(cu_srcs) |
| 7 | + set(op_common_deps operator op_registry) |
| 8 | + set(options "") |
| 9 | + set(oneValueArgs "") |
| 10 | + set(multiValueArgs SRCS DEPS) |
| 11 | + cmake_parse_arguments(op_library "${options}" "${oneValueArgs}" |
| 12 | + "${multiValueArgs}" ${ARGN}) |
| 13 | + |
| 14 | + foreach(src ${op_library_SRCS}) |
| 15 | + if (${src} MATCHES ".*\\.cu$") |
| 16 | + list(APPEND cu_srcs ${src}) |
| 17 | + elseif(${src} MATCHES ".*\\.cc$") |
| 18 | + list(APPEND cc_srcs ${src}) |
| 19 | + else() |
| 20 | + message(FATAL_ERROR "${TARGET} Source file ${src} should only be .cc or .cu") |
| 21 | + endif() |
| 22 | + endforeach() |
| 23 | + |
| 24 | + list(LENGTH cc_srcs cc_srcs_len) |
| 25 | + if (${cc_srcs_len} EQUAL 0) |
| 26 | + message(FATAL_ERROR "The op library ${TARGET} should contains at least one .cc file") |
| 27 | + endif() |
| 28 | + |
| 29 | + list(LENGTH cu_srcs cu_srcs_len) |
| 30 | + if (${cu_srcs_len} EQUAL 0) |
| 31 | + message(WARNING "The op library ${TARGET} not support GPU!") |
| 32 | + endif() |
| 33 | + |
| 34 | + if (WITH_GPU) |
| 35 | + nv_library(${TARGET} SRCS ${cc_srcs} ${cu_srcs} DEPS ${op_library_DEPS} |
| 36 | + ${op_common_deps}) |
| 37 | + else() |
| 38 | + cc_library(${TARGET} SRCS ${cc_srcs} DEPS ${op_library_DEPS} |
| 39 | + ${op_common_deps}) |
| 40 | + endif() |
| 41 | +endfunction() |
| 42 | + |
| 43 | +op_library(add_op SRCS add_op.cc add_op.cu) |
6 | 44 | cc_test(add_op_test SRCS add_op_test.cc DEPS add_op) |
0 commit comments