From f3d68eea37269333af7220149ab977ac9fef0066 Mon Sep 17 00:00:00 2001
From: chuntao hong <chuntao.hong@gmail.com>
Date: Wed, 3 Jun 2015 14:45:19 +0800
Subject: [PATCH] enable ps

---
 CMakeLists.txt           | 17 +++++++++++++++++
 apps/CMakeLists.txt      | 10 +++++++---
 apps/ps/mnist_mlp_ps.cpp | 18 ++++++++----------
 build.sh                 |  1 +
 minerva/CMakeLists.txt   |  4 ----
 5 files changed, 33 insertions(+), 17 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 496a638d..2ac5f298 100644
--- a/CMakeLists.txt
+++ b/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)
diff --git a/apps/CMakeLists.txt b/apps/CMakeLists.txt
index ff0088bd..33d76afc 100644
--- a/apps/CMakeLists.txt
+++ b/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}")
diff --git a/apps/ps/mnist_mlp_ps.cpp b/apps/ps/mnist_mlp_ps.cpp
index fdd30b52..4972a800 100755
--- a/apps/ps/mnist_mlp_ps.cpp
+++ b/apps/ps/mnist_mlp_ps.cpp
@@ -1,4 +1,3 @@
-#include <glog/logging.h>
 #include <minerva.h>
 #include <fstream>
 #include "ps.h"
@@ -6,7 +5,7 @@
 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,16 +203,15 @@ 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);
@@ -221,7 +219,7 @@ int MinervaWorkerMain(int rank, int size, int argc, char *argv[]) {
     data_file_in.close();
     label_file_in.close();
   }
-  ms.current_device_id_ = cpuDevice;
+  ms.SetDevice(cpuDevice);
 
   // output weights
   cout << "Write weight to files" << endl;
diff --git a/build.sh b/build.sh
index aad5fb1f..bd72e593 100755
--- a/build.sh
+++ b/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
diff --git a/minerva/CMakeLists.txt b/minerva/CMakeLists.txt
index 5976c0ff..e956b55b 100644
--- a/minerva/CMakeLists.txt
+++ b/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)