CMakeLists.txt 6.11 KB
Newer Older
1
# Learn a lot from the MLC - LLM Project
2
# https://github.com/mlc-ai/mlc-llm/blob/main/CMakeLists.txt
3

4
cmake_minimum_required(VERSION 3.26)
5
6
project(TILE_LANG C CXX)

7
8
set(CMAKE_CXX_STANDARD 17)
set(CMAKE_POSITION_INDEPENDENT_CODE ON)
9
set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
10

11
set(CMAKE_MODULE_PATH ${CMAKE_MODULE_PATH} ${CMAKE_CURRENT_SOURCE_DIR}/cmake)
12

13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/.gitmodules" AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/.git")
  find_package(Git QUIET)
  if(Git_FOUND)
    execute_process(
      COMMAND ${GIT_EXECUTABLE} submodule update --init --recursive
      WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
      RESULT_VARIABLE TILELANG_GIT_SUBMODULE_RESULT
    )
    if(NOT TILELANG_GIT_SUBMODULE_RESULT EQUAL 0)
      message(
        FATAL_ERROR
          "Failed to initialize git submodules. Please run "
          "`git submodule update --init --recursive` and re-run CMake."
      )
    endif()
  else()
    message(
      FATAL_ERROR
        "Git is required to initialize TileLang submodules. "
        "Please install git or fetch the submodules manually."
    )
  endif()
endif()

37
38
39
40
41
find_program(CCACHE_PROGRAM ccache)
if(CCACHE_PROGRAM)
  set(CMAKE_C_COMPILER_LAUNCHER "${CCACHE_PROGRAM}" CACHE STRING "C compiler launcher")
  set(CMAKE_CXX_COMPILER_LAUNCHER "${CCACHE_PROGRAM}" CACHE STRING "CXX compiler launcher")
  set(CMAKE_CUDA_COMPILER_LAUNCHER "${CCACHE_PROGRAM}" CACHE STRING "CUDA compiler launcher")
42
43
endif()

44
45
46
47
48
49
# Configs
set(USE_CUDA OFF)
set(USE_ROCM OFF)
set(USE_METAL OFF)
set(PREBUILD_CYTHON ON)
# Configs end
50

51
include(cmake/load_tvm.cmake)
52

53
54
if(EXISTS ${TVM_SOURCE}/cmake/config.cmake)
  include(${TVM_SOURCE}/cmake/config.cmake)
55
else()
56
  message(FATAL_ERROR "Nor tvm provided or submodule checkout-ed.")
57
58
endif()

59
60
61
# Include directories for TileLang
set(TILE_LANG_INCLUDES ${TVM_INCLUDES})

62
# Collect source files
63
file(GLOB TILE_LANG_SRCS
64
65
66
67
68
  src/*.cc
  src/layout/*.cc
  src/transform/*.cc
  src/op/*.cc
  src/target/utils.cc
69
70
  src/target/codegen_cpp.cc
  src/target/rt_mod_cpp.cc
71
72
  # webgpu doesn't have system dependency
  src/target/codegen_webgpu.cc
73
74
  # intrin_rule doesn't have system dependency
  src/target/intrin_rule*.cc
75
76
)

77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
# Backend-specific checks and configs
if($ENV{USE_METAL})
  set(USE_METAL ON)
elseif(APPLE)
  message(STATUS "Enable Metal support by default.")
  set(USE_METAL ON)
elseif($ENV{USE_ROCM})
  set(USE_ROCM ON)
else()
  if($ENV{USE_CUDA})
    set(USE_CUDA ON)
  elseif(DEFINED ENV{USE_CUDA} AND NOT $ENV{USE_CUDA})
    # Build CPU-only when we explicitly disable CUDA
    set(USE_CUDA OFF)
  else()
    message(STATUS "Enable CUDA support by default.")
    set(USE_CUDA ON)
  endif()
95
96
endif()

97
if(USE_METAL)
98
  file(GLOB TILE_LANG_METAL_SRCS
99
100
101
    src/target/rt_mod_metal.cc
  )
  list(APPEND TILE_LANG_SRCS ${TILE_LANG_METAL_SRCS})
102
103
  # FIXME: CIBW failed with backtrace, why???
  set(TVM_FFI_USE_LIBBACKTRACE OFF)
104
105
106
107
108
elseif(USE_ROCM)
  set(CMAKE_HIP_STANDARD 17)
  include(${TVM_SOURCE}/cmake/utils/FindROCM.cmake)
  find_rocm($ENV{USE_ROCM})
  add_compile_definitions(__HIP_PLATFORM_AMD__ __HIP_PLATFORM_HCC__=1)
109

110
111
112
113
114
115
116
117
  file(GLOB TILE_LANG_HIP_SRCS
    src/target/codegen_hip.cc
    src/target/rt_mod_hip.cc
  )
  list(APPEND TILE_LANG_SRCS ${TILE_LANG_HIP_SRCS})
  list(APPEND TILE_LANG_INCLUDES ${ROCM_INCLUDE_DIRS})
elseif(USE_CUDA)
  set(CMAKE_CUDA_STANDARD 17)
118
  find_package(CUDAToolkit REQUIRED)
119
  add_compile_definitions("CUDA_MAJOR_VERSION=${CUDAToolkit_VERSION_MAJOR}")
120

121
122
  # Set `USE_CUDA=/usr/local/cuda-x.y`
  cmake_path(GET CUDAToolkit_BIN_DIR PARENT_PATH USE_CUDA)
123

124
125
126
127
128
129
130
  file(GLOB TILE_LANG_CUDA_SRCS
    src/runtime/*.cc
    src/target/ptx.cc
    src/target/codegen_cuda.cc
    src/target/rt_mod_cuda.cc
  )
  list(APPEND TILE_LANG_SRCS ${TILE_LANG_CUDA_SRCS})
131

132
  list(APPEND TILE_LANG_INCLUDES ${CUDAToolkit_INCLUDE_DIRS})
133
endif()
134

135
136
# Include tvm after configs have been populated
add_subdirectory(${TVM_SOURCE} tvm EXCLUDE_FROM_ALL)
137

138
139
# Resolve compile warnings in tvm
add_compile_definitions(DMLC_USE_LOGGING_LIBRARY=<tvm/runtime/logging.h>)
140

141
142
143
144
add_library(tilelang_objs OBJECT ${TILE_LANG_SRCS})
if(CMAKE_BUILD_TYPE STREQUAL "Debug")
  target_compile_definitions(tilelang_objs PRIVATE "TVM_LOG_DEBUG")
endif()
145
146
147
148

target_include_directories(tilelang_objs PRIVATE ${TILE_LANG_INCLUDES})

add_library(tilelang SHARED $<TARGET_OBJECTS:tilelang_objs>)
149
add_library(tilelang_module SHARED $<TARGET_OBJECTS:tilelang_objs>)
150
target_link_libraries(tilelang PUBLIC tvm_runtime)
151
152
153
target_link_libraries(tilelang_module PUBLIC tvm)
if(APPLE)
  # FIXME: libtilelang should only link against tvm runtime
154
155
  target_link_libraries(tilelang PUBLIC tvm)
endif()
156
157
# Build cython extension
find_package(Python REQUIRED COMPONENTS Interpreter Development.Module ${SKBUILD_SABI_COMPONENT})
158

159
160
161
162
163
164
165
166
167
add_custom_command(
  OUTPUT "${CMAKE_BINARY_DIR}/cython_wrapper.cpp"
  COMMENT
    "Cythoning tilelang/jit/adapter/cython/cython_wrapper.pyx"
  COMMAND Python::Interpreter -m cython
          "${CMAKE_CURRENT_SOURCE_DIR}/tilelang/jit/adapter/cython/cython_wrapper.pyx"
          --cplus --output-file "${CMAKE_BINARY_DIR}/cython_wrapper.cpp"
  DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/tilelang/jit/adapter/cython/cython_wrapper.pyx"
  VERBATIM)
168

169
170
if(NOT "${SKBUILD_SABI_VERSION}" STREQUAL "")
  set(USE_SABI USE_SABI ${SKBUILD_SABI_VERSION})
171
172
endif()

173
174
175
python_add_library(cython_wrapper MODULE "${CMAKE_BINARY_DIR}/cython_wrapper.cpp" ${USE_SABI} WITH_SOABI)
# Install to site dir to support direct import
install(TARGETS cython_wrapper LIBRARY DESTINATION .)
176

177
178
179
180
181
182
183
# let libtilelang to search tvm/tvm_runtime in same dir
if(APPLE)
  set_target_properties(tilelang PROPERTIES INSTALL_RPATH "@loader_path")
  set_target_properties(tilelang_module PROPERTIES INSTALL_RPATH "@loader_path")
else()
  set_target_properties(tilelang PROPERTIES INSTALL_RPATH "\$ORIGIN")
  set_target_properties(tilelang_module PROPERTIES INSTALL_RPATH "\$ORIGIN")
184
185
endif()

186
install(TARGETS tvm tvm_runtime tilelang_module tilelang LIBRARY DESTINATION tilelang/lib)
187

188
189
190
191
192
# Copy tvm cython ext for wheels
# TODO: not necessary for editable builds
if(TVM_BUILD_FROM_SOURCE)
  add_dependencies(tilelang tvm_cython)
  install(FILES "${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/tvm/python/tvm/ffi/core.abi3.so" DESTINATION tilelang/3rdparty/tvm/python/tvm/ffi/)
193
endif()