CMakeLists.txt 6.2 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34
cc_library(
  allocator
  SRCS allocator.cc
  DEPS place stats)
cc_library(
  cpu_allocator
  SRCS cpu_allocator.cc
  DEPS allocator)
cc_library(
  locked_allocator
  SRCS locked_allocator.cc
  DEPS allocator)
cc_library(
  buffered_allocator
  SRCS buffered_allocator.cc
  DEPS allocator)
cc_library(
  best_fit_allocator
  SRCS best_fit_allocator.cc
  DEPS allocator)
cc_library(
  naive_best_fit_allocator
  SRCS naive_best_fit_allocator.cc
  DEPS allocator buddy_allocator profiler)
cc_test(
  naive_best_fit_allocator_test
  SRCS naive_best_fit_allocator_test.cc
  DEPS naive_best_fit_allocator)
cc_test(
  buffered_allocator_test
  SRCS buffered_allocator_test.cc
  DEPS locked_allocator buffered_allocator cpu_allocator best_fit_allocator)

if(WITH_MKLDNN)
35
  set(MKLDNN_CTX_DEPS mkldnn)
36
else()
37 38 39
  set(MKLDNN_CTX_DEPS)
endif()

40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60
if(WITH_GPU)
  nv_library(
    cuda_allocator
    SRCS cuda_allocator.cc
    DEPS allocator cuda_device_guard stats)
  nv_library(
    cuda_managed_allocator
    SRCS cuda_managed_allocator.cc
    DEPS allocator cuda_device_guard gpu_info)
  nv_library(
    pinned_allocator
    SRCS pinned_allocator.cc
    DEPS allocator)
  nv_library(
    stream_safe_cuda_allocator
    SRCS stream_safe_cuda_allocator.cc
    DEPS allocator cuda_graph)
  nv_library(
    thread_local_allocator
    SRCS thread_local_allocator.cc
    DEPS allocator)
61

62 63 64 65
  cc_test(
    thread_local_allocator_test
    SRCS thread_local_allocator_test.cc
    DEPS thread_local_allocator)
66
  if(CUDA_VERSION GREATER_EQUAL 10.2)
67 68 69 70
    nv_library(
      cuda_virtual_mem_allocator
      SRCS cuda_virtual_mem_allocator.cc
      DEPS dynload_cuda)
71
  endif()
72 73
endif()

74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99
if(WITH_ROCM)
  hip_library(
    cuda_allocator
    SRCS cuda_allocator.cc
    DEPS allocator cuda_device_guard stats)
  hip_library(
    cuda_managed_allocator
    SRCS cuda_managed_allocator.cc
    DEPS allocator cuda_device_guard gpu_info)
  hip_library(
    pinned_allocator
    SRCS pinned_allocator.cc
    DEPS allocator)
  hip_library(
    stream_safe_cuda_allocator
    SRCS stream_safe_cuda_allocator.cc
    DEPS allocator)
  hip_library(
    thread_local_allocator
    SRCS thread_local_allocator.cc
    DEPS allocator)

  cc_test(
    thread_local_allocator_test
    SRCS thread_local_allocator_test.cc
    DEPS thread_local_allocator)
S
sneaxiy 已提交
100
endif()
101

102 103 104 105 106 107 108 109 110
if(WITH_ASCEND_CL)
  cc_library(
    npu_allocator
    SRCS npu_allocator.cc
    DEPS allocator npu_info)
  cc_library(
    npu_pinned_allocator
    SRCS npu_pinned_allocator.cc
    DEPS allocator npu_info)
111 112
endif()

113 114 115 116
cc_library(
  retry_allocator
  SRCS retry_allocator.cc
  DEPS allocator)
S
sneaxiy 已提交
117

118 119 120 121 122 123 124 125 126 127 128 129 130
if(WITH_GPU OR WITH_ROCM)
  set(AllocatorFacadeDeps
      gpu_info
      cuda_allocator
      cuda_managed_allocator
      pinned_allocator
      cuda_device_guard
      thread_local_allocator
      stream_safe_cuda_allocator
      device_context)
  if(CUDA_VERSION GREATER_EQUAL 10.2)
    list(APPEND AllocatorFacadeDeps cuda_virtual_mem_allocator)
  endif()
131
elseif(WITH_XPU)
132
  set(AllocatorFacadeDeps xpu_info)
J
jianghaicheng 已提交
133
elseif(WITH_IPU)
134
  set(AllocatorFacadeDeps ipu_info)
135
elseif(WITH_ASCEND)
136 137 138
  set(AllocatorFacadeDeps ascend_npu_info)
else()
  set(AllocatorFacadeDeps)
139 140
endif()

141 142 143 144 145
if(WITH_CUSTOM_DEVICE)
  cc_library(
    custom_allocator
    SRCS custom_allocator.cc
    DEPS allocator device_manager)
146 147 148
  set(AllocatorFacadeDeps ${AllocatorFacadeDeps} custom_allocator)
endif()

149 150 151 152 153 154 155 156 157 158 159 160
if(WITH_GPU)
  nv_test(
    best_fit_allocator_test
    SRCS best_fit_allocator_test.cc best_fit_allocator_test.cu
    DEPS best_fit_allocator locked_allocator cpu_allocator cuda_allocator
         device_context memcpy)
elseif(WITH_ROCM)
  hip_test(
    best_fit_allocator_test
    SRCS best_fit_allocator_test.cc best_fit_allocator_test.cu
    DEPS best_fit_allocator locked_allocator cpu_allocator cuda_allocator
         device_context memcpy)
161
else()
162 163 164 165
  cc_test(
    best_fit_allocator_test
    SRCS best_fit_allocator_test.cc
    DEPS best_fit_allocator locked_allocator cpu_allocator)
166
endif()
167

168 169 170 171 172 173 174 175 176 177 178 179
list(
  APPEND
  AllocatorFacadeDeps
  cpu_allocator
  locked_allocator
  aligned_allocator
  retry_allocator
  buffered_allocator
  naive_best_fit_allocator
  auto_growth_best_fit_allocator
  virtual_memory_auto_growth_best_fit_allocator
  best_fit_allocator)
Z
Zeng Jinle 已提交
180

181 182
if(WITH_ASCEND_CL)
  list(APPEND AllocatorFacadeDeps npu_pinned_allocator)
183 184
endif()

185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200
cc_library(
  aligned_allocator
  SRCS aligned_allocator.cc
  DEPS allocator)
cc_test(
  test_aligned_allocator
  SRCS test_aligned_allocator.cc
  DEPS aligned_allocator)
cc_library(
  allocator_strategy
  SRCS allocator_strategy.cc
  DEPS gflags ${AllocatorFacadeDeps})
cc_library(
  allocator_facade
  SRCS allocator_facade.cc
  DEPS allocator_strategy stats)
201

202
if(WITH_GPU)
203 204
  target_link_libraries(allocator_facade cuda_graph)
endif()
Y
Yu Yang 已提交
205

206 207 208 209 210 211
cc_test(
  retry_allocator_test
  SRCS retry_allocator_test.cc
  DEPS retry_allocator locked_allocator cpu_allocator)
if(WITH_TESTING)
  if((WITH_GPU OR WITH_ROCM) AND TARGET retry_allocator_test)
212 213 214
    target_link_libraries(retry_allocator_test cuda_allocator)
  endif()

215 216 217
  if(TEST retry_allocator_test)
    set_tests_properties(retry_allocator_test PROPERTIES LABELS
                                                         "RUN_TYPE=EXCLUSIVE")
218
  endif()
Z
Zeng Jinle 已提交
219 220
endif()

221 222 223 224
cc_test(
  allocator_facade_abs_flags_test
  SRCS allocator_facade_abs_flags_test.cc
  DEPS allocator_facade)
S
sneaxiy 已提交
225

226 227 228 229
cc_test(
  allocator_facade_frac_flags_test
  SRCS allocator_facade_frac_flags_test.cc
  DEPS allocator_facade)
230

231 232 233 234 235 236 237 238 239 240 241 242
cc_library(
  auto_growth_best_fit_allocator
  SRCS auto_growth_best_fit_allocator.cc
  DEPS allocator aligned_allocator flags)
cc_test(
  auto_growth_best_fit_allocator_facade_test
  SRCS auto_growth_best_fit_allocator_facade_test.cc
  DEPS cpu_allocator auto_growth_best_fit_allocator)
cc_test(
  auto_growth_best_fit_allocator_test
  SRCS auto_growth_best_fit_allocator_test.cc
  DEPS auto_growth_best_fit_allocator)
243

244 245 246 247
cc_library(
  virtual_memory_auto_growth_best_fit_allocator
  SRCS virtual_memory_auto_growth_best_fit_allocator.cc
  DEPS allocator aligned_allocator)
248

249
if(NOT WIN32)
250 251 252 253 254 255 256 257 258 259 260 261 262
  cc_library(
    mmap_allocator
    SRCS mmap_allocator.cc
    DEPS allocator)
  cc_test(
    mmap_allocator_test
    SRCS mmap_allocator_test.cc
    DEPS mmap_allocator allocator)
  if(WITH_GPU)
    cc_library(
      cuda_ipc_allocator
      SRCS cuda_ipc_allocator.cc
      DEPS allocator)
263
  endif()
264
endif(NOT WIN32)