enable ps

CMakeLists.txt

@@ -40,6 +40,20 @@ function(find_blas)
   endif()
 endfunction()
 
+function(find_ps)
+  set(PS_ROOT "" CACHE PATH "minervaps root path")
+  find_library(PS_LIBRARIES NAMES libminervaps.a
+    PATHS ${PS_ROOT}/build
+    DOC "PS library path")
+  if(PS_LIBRARIES)
+    set(PS_FOUND TRUE PARENT_SCOPE)
+    message(STATUS "Found PS (library: ${PS_LIBRARIES})")
+    mark_as_advanced(PS_LIBRARIES)
+  else()
+    MESSAGE(FATAL_ERROR "Failed to find PS in path: ${PS_ROOT}. Make sure you have built ps with 'make minerva'.")
+  endif()
+endfunction()
+
 project(Minerva)
 
 option(BUILD_CXX_APPS    "build C++ applications"                OFF)
@@ -103,6 +117,9 @@ set(CORE_DEPS ${CORE_DEPS} ${THIRD_LIBS})
 
 if(BUILD_WITH_PS)
   add_definitions(-DHAS_PS)
+  find_ps()
+  message(STATUS "PS enabled, forcing BUILD_TESTS to 0")
+  set(BUILD_TESTS 0)
 endif()
 
 if(BUILD_WITH_BLAS)

apps/CMakeLists.txt

@@ -6,16 +6,16 @@ include_directories(
 
 set(app_list "")
 
-if (ENABLE_PS)
+if (BUILD_WITH_PS)
   file(GLOB_RECURSE app_file_list "ps/*.cpp")
-else (ENABLE_PS)
+else (BUILD_WITH_PS)
   if (CUDA_FOUND)
     file(GLOB app_file_list "*.cpp")
   else (CUDA_FOUND)
     #For CPU only.
     file(GLOB app_file_list "mnist_mlp.cpp")
   endif (CUDA_FOUND)
-endif (ENABLE_PS)
+endif (BUILD_WITH_PS)
 
 foreach(app_src ${app_file_list})
   get_filename_component(app_name ${app_src} NAME_WE)
@@ -24,9 +24,13 @@ foreach(app_src ${app_file_list})
   if (CUDA_FOUND)
     CUDA_ADD_CUBLAS_TO_TARGET(${app_name})
   endif (CUDA_FOUND)
+  if (BUILD_WITH_PS)
+    target_link_libraries(${app_name} ${PS_LIBRARIES})
+  endif ()
   list(APPEND app_list ${app_name})
 endforeach()
 
+
 message(STATUS "Build CXX Applications:")
 foreach(app ${app_list})
   message(STATUS "  ${app}")

apps/ps/mnist_mlp_ps.cpp

@@ -1,12 +1,11 @@
-#include <glog/logging.h>
 #include <minerva.h>
 #include <fstream>
 #include "ps.h"
 
 using namespace std;
 using namespace minerva;
 
-#define LL LOG(ERROR)
+#define LL cout
 
 template <typename V>
 inline std::string arrstr(const V* data, int n) {
@@ -149,7 +148,7 @@ int MinervaWorkerMain(int rank, int size, int argc, char *argv[]) {
 #ifdef HAS_CUDA
   uint64_t gpuDevice = ms.device_manager().CreateGpuDevice(0);
 #endif
-  ms.current_device_id_ = cpuDevice;
+  ms.SetDevice(cpuDevice);
 
   weights.resize(num_layers - 1);
   bias.resize(num_layers - 1);
@@ -171,7 +170,7 @@ int MinervaWorkerMain(int rank, int size, int argc, char *argv[]) {
     DumpParams(0);
     for (int mb = rank; mb < num_mb_per_epoch; mb+=size) {
 
-      ms.current_device_id_ = cpuDevice;
+      ms.SetDevice(cpuDevice);
 
       Scale data_size{ lsize[0], mb_size };
       Scale label_size{ lsize[num_layers - 1], mb_size };
@@ -186,7 +185,7 @@ int MinervaWorkerMain(int rank, int size, int argc, char *argv[]) {
       NArray label = NArray::MakeNArray(label_size, label_ptr);
 
 #ifdef HAS_CUDA
-      ms.current_device_id_ = gpuDevice;
+      ms.SetDevice(gpuDevice);
 #endif
 
       // ff
@@ -204,24 +203,23 @@ int MinervaWorkerMain(int rank, int size, int argc, char *argv[]) {
         sens[k] = weights[k].Trans() * sens[k + 1];
         sens[k] = Elewise::Mult(sens[k], d_act);
       }
-
-      ms.current_device_id_ = cpuDevice;
+      ms.SetDevice(cpuDevice);
       for (int k = 0; k < num_layers - 1; ++k) {
         bias[k] = NArray::PushGradAndPullWeight(sens[k + 1].Sum(1) / mb_size, GetBiasName(k));
         weights[k] = NArray::PushGradAndPullWeight(sens[k + 1] * acts[k].Trans() / mb_size, GetWeightName(k));
       }
-      ms.current_device_id_ = gpuDevice;
+      ms.SetDevice(gpuDevice);
 
       if ((mb - rank) % 20 == 0) {
-        ms.current_device_id_ = cpuDevice;
+        ms.SetDevice(cpuDevice);
         PrintTrainingAccuracy(acts[num_layers - 1], label);
       }
       DumpParams(mb + 1);
     }
     data_file_in.close();
     label_file_in.close();
   }
-  ms.current_device_id_ = cpuDevice;
+  ms.SetDevice(cpuDevice);
 
   // output weights
   cout << "Write weight to files" << endl;

build.sh

@@ -29,6 +29,7 @@ CXXFLAGS="$CXXFLAGS \
   -DBUILD_CPU_ONLY=$BUILD_CPU_ONLY \
   -DBUILD_WITH_BLAS=$BUILD_WITH_BLAS \
   -DBLAS_ROOT=$BLAS_ROOT \
+  -DPS_ROOT=$PS_ROOT \
   "
 
 while [[ $# -gt 0 ]]; do

minerva/CMakeLists.txt

@@ -22,9 +22,5 @@ if (BUILD_WITH_BLAS)
   target_link_libraries(minerva ${CBLAS_LIBRARIES})
 endif ()
 
-if (BUILD_WITH_PS)
-  target_link_libraries(minerva minervaps)
-endif ()
-
 set(LIBRARY_OUTPUT_PATH ${Minerva_BINARY_DIR}/lib)