fixed bug in setup.py

Tixxx · Aug 23, 2019 · 4491b32 · 4491b32
1 parent eda4e4e
commit 4491b32
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 1 deletion.
diff --git a/MANIFEST.in b/MANIFEST.in
@@ -1,4 +1,4 @@
-recursive-include * *.h *.hpp *.cc *.md
+recursive-include * *.h *.hpp *.cc *.md *.ipp
 
 include LICENSE horovod.lds horovod.exp
 prune .eggs
@@ -19,3 +19,6 @@ exclude third_party/eigen/Eigen/src/SparseCholesky/*
 graft third_party/gloo/cmake
 recursive-include third_party/gloo CMakeLists.txt
 recursive-include third_party/gloo *.in
+
+# include cmake related files for msallreduce kernel
+graft horovod/common/ops/cuda/
diff --git a/horovod/common/ops/msallreduce_cuda_operations.cc b/horovod/common/ops/msallreduce_cuda_operations.cc
@@ -78,6 +78,9 @@ cudaDataType_t CublasContext::GetCublasDataType(T* variable) {
 }
 
 //thread_local cublasHandle_t MsCudaAllreduceOp::cublas_Handle;
+thread_local double* MsCudaAllreduceOp::device_normsq_memory_a;
+thread_local double* MsCudaAllreduceOp::device_normsq_memory_b;
+thread_local double* MsCudaAllreduceOp::device_dot_product_memory;
 
 MsCudaAllreduceOp::MsCudaAllreduceOp(MPIContext* mpi_context, CUDAContext* cuda_context, HorovodGlobalState* global_state)
     : MsAllreduceOp(mpi_context, global_state), cuda_context_(cuda_context) {

diff --git a/setup.py b/setup.py
@@ -661,6 +661,7 @@ def get_common_options(build_ext):
     if have_cuda:
         MACROS += [('HAVE_CUDA', '1')]
         INCLUDES += cuda_include_dirs
+        INCLUDES += ['horovod/common/ops/cuda']
         SOURCES += ['horovod/common/ops/cuda_operations.cc',
                     'horovod/common/ops/mpi_cuda_operations.cc',
                     'horovod/common/ops/msallreduce_cuda_operations.cc']