commit b47ae9810c1a645f4942737ab4a58b2b1407e7bd
author Michael Dagitses <mikeyd@fb.com> Tue Nov 30 04:19:53 2021 -0800
committer Facebook GitHub Bot <facebook-github-bot@users.noreply.github.com> Tue Nov 30 04:24:55 2021 -0800
tree 8f8cd74c75afb50854f23e1204fac1a41e849c22
parent 174eea8a05fc0f6ed11f1948c48ced2173f6ae72

Split cuda: list cpp files that go in _cu library explicitly (#67216)

Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/67216

Test Plan: Imported from OSS

Reviewed By: zou3519

Differential Revision: D32099294

Pulled By: dagitses

fbshipit-source-id: 8a3582944b6b48af1ac31c5df09a7e6e838892c4

caffe2/CMakeLists.txt

diff --git a/caffe2/CMakeLists.txt b/caffe2/CMakeLists.txt
index ae387ac..6d70085 100644
--- a/caffe2/CMakeLists.txt
+++ b/caffe2/CMakeLists.txt
@@ -176,12 +176,47 @@
 
 if(BUILD_SPLIT_CUDA)
   # Splitting the source files that'll be in torch_cuda between torch_cuda_cu and torch_cuda_cpp
+  set(Caffe2_GPU_SRCS_CU
+      "${PROJECT_SOURCE_DIR}/aten/src/ATen/cuda/CUDABlas.cpp"
+      "${PROJECT_SOURCE_DIR}/aten/src/ATen/cuda/CUDASolver.cpp"
+      "${PROJECT_SOURCE_DIR}/aten/src/ATen/cuda/CUDASparseBlas.cpp"
+      "${PROJECT_SOURCE_DIR}/aten/src/ATen/cuda/CublasHandlePool.cpp"
+      "${PROJECT_SOURCE_DIR}/aten/src/ATen/cuda/CusolverDnHandlePool.cpp"
+      "${PROJECT_SOURCE_DIR}/aten/src/ATen/native/cuda/Activation.cpp"
+      "${PROJECT_SOURCE_DIR}/aten/src/ATen/native/cuda/BatchLinearAlgebra.cpp"
+      "${PROJECT_SOURCE_DIR}/aten/src/ATen/native/cuda/BatchLinearAlgebraLib.cpp"
+      "${PROJECT_SOURCE_DIR}/aten/src/ATen/native/cuda/Blas.cpp"
+      "${PROJECT_SOURCE_DIR}/aten/src/ATen/native/cuda/Equal.cpp"
+      "${PROJECT_SOURCE_DIR}/aten/src/ATen/native/cuda/IndexKernel.cpp"
+      "${PROJECT_SOURCE_DIR}/aten/src/ATen/native/cuda/ReduceOps.cpp"
+      "${PROJECT_SOURCE_DIR}/aten/src/ATen/native/cuda/ScanKernels.cpp"
+      "${PROJECT_SOURCE_DIR}/aten/src/ATen/native/cuda/Sort.cpp"
+      "${PROJECT_SOURCE_DIR}/aten/src/ATen/native/cuda/Sorting.cpp"
+      "${PROJECT_SOURCE_DIR}/aten/src/ATen/native/cuda/TensorModeKernel.cpp"
+      "${PROJECT_SOURCE_DIR}/aten/src/ATen/native/cuda/TensorShapeCUDA.cpp"
+      "${PROJECT_SOURCE_DIR}/aten/src/ATen/native/cuda/TensorTopK.cpp"
+      "${PROJECT_SOURCE_DIR}/aten/src/ATen/native/sparse/cuda/SparseBlas.cpp"
+      "${PROJECT_SOURCE_DIR}/aten/src/ATen/native/sparse/cuda/SparseBlasImpl.cpp"
+      "${PROJECT_SOURCE_DIR}/aten/src/ATen/native/sparse/cuda/SparseBlasLegacy.cpp"
+      "${PROJECT_SOURCE_DIR}/aten/src/ATen/native/sparse/cuda/SparseCUDABlas.cpp"
+      "${PROJECT_SOURCE_DIR}/aten/src/THC/THCStorage.cpp"
+      "${PROJECT_SOURCE_DIR}/aten/src/THC/THCTensor.cpp"
+
+      "${PROJECT_BINARY_DIR}/aten/src/ATen/RegisterCUDA.cpp"
+      "${PROJECT_BINARY_DIR}/aten/src/ATen/RegisterQuantizedCUDA.cpp"
+      "${PROJECT_BINARY_DIR}/aten/src/ATen/RegisterSparseCUDA.cpp"
+      "${PROJECT_BINARY_DIR}/aten/src/ATen/RegisterSparseCsrCUDA.cpp"
+   )
+
   foreach(tmp ${Caffe2_GPU_SRCS})
-    if("${tmp}" MATCHES "(.*aten.*\\.cu|.*(b|B)las.*|.*((s|S)olver|Register.*CUDA|Legacy|THC|TensorShapeCUDA|BatchLinearAlgebra|ReduceOps|Equal|Activation|ScanKernels|Sort|TensorTopK|TensorModeKernel|IndexKernel).*\\.cpp)" AND NOT "${tmp}" MATCHES ".*(THC((CachingHost)?Allocator|General)).*")
+    if("${tmp}" MATCHES ".*aten.*\\.cu")
       # Currently, torch_cuda_cu will have all the .cu files in aten, as well as some others that depend on those files
-      list(APPEND Caffe2_GPU_SRCS_CU ${tmp})
+      list(APPEND Caffe2_GPU_SRCS_CU "${tmp}")
     else()
-      list(APPEND Caffe2_GPU_SRCS_CPP ${tmp})
+      # Exclude special cases listed above that need to link to torch_cuda_cu
+      if(NOT "${tmp}" IN_LIST Caffe2_GPU_SRCS_CU)
+        list(APPEND Caffe2_GPU_SRCS_CPP "${tmp}")
+      endif()
     endif()
   endforeach()