CMakeLists.txt 2.03 KB
Newer Older
czkkkkkk's avatar
czkkkkkk committed
1
# Find PyTorch cmake files and PyTorch versions with the python interpreter $TORCH_PYTHON_INTERPS
2
# ("python3" or "python" if empty)
czkkkkkk's avatar
czkkkkkk committed
3
4
if(NOT TORCH_PYTHON_INTERPS)
  find_program(TORCH_PYTHON_INTERPS NAMES python3 python)
5
endif()
czkkkkkk's avatar
czkkkkkk committed
6
message(STATUS "Using Python interpreter: ${TORCH_PYTHON_INTERPS}")
7
8
file(TO_NATIVE_PATH ${CMAKE_CURRENT_SOURCE_DIR}/find_cmake.py FIND_CMAKE_PY)
execute_process(
czkkkkkk's avatar
czkkkkkk committed
9
  COMMAND ${TORCH_PYTHON_INTERPS} ${FIND_CMAKE_PY}
10
11
12
13
14
15
  OUTPUT_VARIABLE TORCH_PREFIX_VER
  OUTPUT_STRIP_TRAILING_WHITESPACE)
message(STATUS "find_cmake.py output: ${TORCH_PREFIX_VER}")
list(GET TORCH_PREFIX_VER 0 TORCH_PREFIX)
list(GET TORCH_PREFIX_VER 1 TORCH_VER)
message(STATUS "Configuring for PyTorch ${TORCH_VER}")
16
17
18
string(REPLACE "." ";" TORCH_VERSION_LIST ${TORCH_VER})
list(GET TORCH_VERSION_LIST 0 TORCH_VERSION_MAJOR)
list(GET TORCH_VERSION_LIST 1 TORCH_VERSION_MINOR)
19
20
21
22
23
24
25
26
27
28
29

if(USE_CUDA)
  add_definitions(-DDGL_USE_CUDA)
endif()

set(Torch_DIR "${TORCH_PREFIX}/Torch")
message(STATUS "Setting directory to ${Torch_DIR}")
find_package(Torch REQUIRED)

set(LIB_DGL_SPARSE_NAME "dgl_sparse_pytorch_${TORCH_VER}")

czkkkkkk's avatar
czkkkkkk committed
30
31
set(SPARSE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/src")
set(SPARSE_INCLUDE "${CMAKE_CURRENT_SOURCE_DIR}/include")
32
33
34
35
36
37
38
39
file(GLOB SPARSE_HEADERS ${SPARSE_INCLUDE})
file(GLOB SPARSE_SRC
  ${SPARSE_DIR}/*.cc
)
add_library(${LIB_DGL_SPARSE_NAME} SHARED ${SPARSE_SRC} ${SPARSE_HEADERS})
target_include_directories(
  ${LIB_DGL_SPARSE_NAME} PRIVATE ${SPARSE_DIR} ${SPARSE_HEADERS})
target_link_libraries(${LIB_DGL_SPARSE_NAME} "${TORCH_LIBRARIES}")
40
41
target_compile_definitions(${LIB_DGL_SPARSE_NAME} PRIVATE TORCH_VERSION_MAJOR=${TORCH_VERSION_MAJOR})
target_compile_definitions(${LIB_DGL_SPARSE_NAME} PRIVATE TORCH_VERSION_MINOR=${TORCH_VERSION_MINOR})
42

czkkkkkk's avatar
czkkkkkk committed
43
target_include_directories(${LIB_DGL_SPARSE_NAME} PRIVATE "${CMAKE_SOURCE_DIR}/third_party/dmlc-core/include")
44
45
target_link_libraries(${LIB_DGL_SPARSE_NAME} dmlc)
set(GOOGLE_TEST 0) # Turn off dmlc-core test
czkkkkkk's avatar
czkkkkkk committed
46
47
48
49


target_include_directories(${LIB_DGL_SPARSE_NAME} PRIVATE ${DGL_INCLUDE})
target_link_libraries(${LIB_DGL_SPARSE_NAME} dgl)