Use white list to control pslite compilation

cmake/mind_expression.cmake

@@ -30,7 +30,7 @@ include(${CMAKE_SOURCE_DIR}/cmake/external_libs/flatbuffers.cmake)
 if(USE_GLOG)
     include(${CMAKE_SOURCE_DIR}/cmake/external_libs/glog.cmake)
 endif()
-if (NOT ${CMAKE_SYSTEM_NAME} MATCHES "Windows" AND NOT ENABLE_GE)
+if (ENABLE_CPU AND (ENABLE_D OR ENABLE_GPU))
     include(${CMAKE_SOURCE_DIR}/cmake/external_libs/zeromq.cmake)
     include(${CMAKE_SOURCE_DIR}/cmake/external_libs/pslite.cmake)
 endif()

cmake/options.cmake

@@ -70,6 +70,10 @@ if (ENABLE_GPU)
     add_compile_definitions(ENABLE_GPU_COLLECTIVE)
 endif()
 
+if (ENABLE_CPU)
+    add_compile_definitions(ENABLE_CPU)
+endif()
+
 if (ENABLE_GE)
     add_compile_definitions(ENABLE_GE)
     add_compile_definitions(CUSTOM_OP)
@@ -119,7 +123,3 @@ endif()
 if(ENABLE_DEBUGGER)
     add_compile_definitions(ENABLE_DEBUGGER)
 endif()
-
-if(ENABLE_TESTCASES)
-    add_compile_definitions(ENABLE_TESTCASES)
-endif()

mindspore/ccsrc/CMakeLists.txt

@@ -252,7 +252,7 @@ if (CMAKE_SYSTEM_NAME MATCHES "Windows")
     target_link_libraries(mindspore mindspore_gvar)
     target_link_libraries(_c_expression PRIVATE -Wl,--whole-archive mindspore -Wl,--no-whole-archive)
 else ()
-    if (NOT ENABLE_GE)
+    if (ENABLE_CPU AND (ENABLE_D OR ENABLE_GPU))
         target_link_libraries(mindspore mindspore::pslite mindspore::protobuf ${zeromq_DIRPATH}/zmq_install/lib/libzmq.a)
         if (${ENABLE_IBVERBS} STREQUAL "ON")
             target_link_libraries(mindspore ibverbs rdmacm)

mindspore/ccsrc/backend/kernel_compiler/CMakeLists.txt

@@ -33,7 +33,7 @@ if (ENABLE_CPU)
     endif ()
 endif ()
 
-if (${CMAKE_SYSTEM_NAME} MATCHES "Windows" OR ENABLE_GE)
+if (NOT (ENABLE_CPU AND (ENABLE_D OR ENABLE_GPU)))
     list(REMOVE_ITEM CPU_SRC_LIST "cpu/ps/apply_momentum_ps_kernel.cc")
     list(REMOVE_ITEM CPU_SRC_LIST "cpu/ps/embedding_look_up_proxy_kernel.cc")
     list(REMOVE_ITEM CPU_SRC_LIST "cpu/ps/embedding_look_up_ps_kernel.cc")

mindspore/ccsrc/backend/session/ascend_session.cc

@@ -353,7 +353,7 @@ GraphId AscendSession::CompileGraph(NotNull<FuncGraphPtr> func_graph) {
   RootGraphExecutorValidate(NOT_NULL(root_graph));
   // adjust kernel
   AdjustKernel(root_graph);
-#if (!_WIN32 && !ENABLE_GE && !ENABLE_TESTCASES)
+#if (ENABLE_CPU && (ENABLE_D || ENABLE_GPU))
   // Assign parameter keys.
   AssignParamKey(root_graph);
 #endif
@@ -515,7 +515,7 @@ void AscendSession::RunGraph(const GraphId &graph_id, const std::vector<tensor::
   }
   // load input data from user input
   LoadInputData(kernel_graph, inputs);
-#if (!_WIN32 && !ENABLE_GE && !ENABLE_TESTCASES)
+#if (ENABLE_CPU && (ENABLE_D || ENABLE_GPU))
   // Initialize parameter server
   if (!ps_init_) {
     InitPSParamAndOptim(kernel_graph, inputs);

mindspore/ccsrc/backend/session/cpu_session.cc

@@ -32,7 +32,7 @@
 #ifdef ENABLE_DEBUGGER
 #include "debug/debugger/debugger.h"
 #endif
-#if (!_WIN32 && !ENABLE_GE && !ENABLE_TESTCASES)
+#if (ENABLE_CPU && (ENABLE_D || ENABLE_GPU))
 #include "frontend/parallel/ps/util.h"
 #endif
 
@@ -73,7 +73,7 @@ GraphId CPUSession::CompileGraph(const AnfNodePtrList &lst, const AnfNodePtrList
   MS_EXCEPTION_IF_NULL(graph);
   MS_LOG(INFO) << "Set kernel info";
   SetKernelInfo(graph.get());
-#if (!_WIN32 && !ENABLE_GE && !ENABLE_TESTCASES)
+#if (ENABLE_CPU && (ENABLE_D || ENABLE_GPU))
   AssignParamKey(graph);
   if (parallel::ps::Util::IsRoleOfWorker()) {
     Optimize(graph);
@@ -90,7 +90,7 @@ GraphId CPUSession::CompileGraph(const AnfNodePtrList &lst, const AnfNodePtrList
 void CPUSession::RunGraph(const GraphId &graph_id, const std::vector<tensor::TensorPtr> &inputs, VectorRef *outputs) {
   auto &kernel_graph = graphs_[graph_id];
   MS_EXCEPTION_IF_NULL(kernel_graph);
-#if (!_WIN32 && !ENABLE_GE && !ENABLE_TESTCASES)
+#if (ENABLE_CPU && (ENABLE_D || ENABLE_GPU))
   // Initialize parameter server
   if (!ps_init_) {
     InitPSParamAndOptim(kernel_graph, inputs);

mindspore/ccsrc/backend/session/gpu_session.cc

@@ -182,7 +182,7 @@ GraphId GPUSession::CompileGraph(const AnfNodePtrList &lst, const AnfNodePtrList
   Optimize(graph);
   // Select kernel build info
   SelectKernel(graph);
-#if (!_WIN32 && !ENABLE_GE && !ENABLE_TESTCASES)
+#if (ENABLE_CPU && (ENABLE_D || ENABLE_GPU))
   // Assign parameter keys.
   AssignParamKey(graph);
 #endif
@@ -231,10 +231,12 @@ void GPUSession::RunGraph(const GraphId &graph_id, const std::vector<tensor::Ten
   auto &kernel_graph = graphs_[graph_id];
   // Load input data from user input
   LoadInputData(kernel_graph, inputs);
+#if (ENABLE_CPU && (ENABLE_D || ENABLE_GPU))
   // Initialize parameter server
   if (!ps_init_) {
     InitPSParamAndOptim(kernel_graph, inputs);
   }
+#endif
   MS_EXCEPTION_IF_NULL(kernel_graph);
   // Convert inputs to model
   predictmodel::StepConvertWeight(inputs);

mindspore/ccsrc/backend/session/session_basic.cc

@@ -35,7 +35,7 @@
 #include "ir/dtype.h"
 #include "ir/anf.h"
 #include "ir/func_graph_cloner.h"
-#if (!_WIN32 && !ENABLE_GE && !ENABLE_TESTCASES)
+#if (ENABLE_CPU && (ENABLE_D || ENABLE_GPU))
 #include "frontend/parallel/ps/worker.h"
 #include "frontend/parallel/ps/common.h"
 #include "frontend/parallel/ps/util.h"
@@ -1141,7 +1141,7 @@ AnfNodePtr SessionBasic::FindPullNode(const AnfNodePtr &push_node, const std::ve
   return nullptr;
 }
 
-#if (!_WIN32 && !ENABLE_GE && !ENABLE_TESTCASES)
+#if (ENABLE_CPU && (ENABLE_D || ENABLE_GPU))
 void SessionBasic::AssignParamKey(const KernelGraphPtr &kernel_graph) {
   if (!parallel::ps::Util::IsRoleOfWorker()) {
     MS_LOG(INFO) << "Not parameter server mode.";

mindspore/ccsrc/frontend/parallel/CMakeLists.txt

@@ -1,6 +1,6 @@
 file(GLOB_RECURSE _PARALLEL_SRC_FILES RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} "*.cc")
 
-if (${CMAKE_SYSTEM_NAME} MATCHES "Windows" OR ENABLE_GE)
+if (NOT (ENABLE_CPU AND (ENABLE_D OR ENABLE_GPU)))
     list(REMOVE_ITEM _PARALLEL_SRC_FILES "ps/optimizer_info_builder.cc")
     list(REMOVE_ITEM _PARALLEL_SRC_FILES "ps/optimizer_info.cc")
     list(REMOVE_ITEM _PARALLEL_SRC_FILES "ps/scheduler.cc")

mindspore/ccsrc/minddata/dataset/CMakeLists.txt

@@ -146,7 +146,7 @@ if (${CMAKE_SYSTEM_NAME} MATCHES "Windows")
     target_link_libraries(_c_dataengine PRIVATE _c_mindrecord ${MINDRECORD_LINK_OBJECT} mindspore::sqlite)
 else()
     target_link_libraries(_c_dataengine PRIVATE _c_mindrecord)
-    if (NOT ENABLE_GE)
+    if (ENABLE_CPU AND (ENABLE_D OR ENABLE_GPU))
         target_link_libraries(_c_dataengine PRIVATE mindspore::pslite mindspore::protobuf ${zeromq_DIRPATH}/zmq_install/lib/libzmq.a)
         if (${ENABLE_IBVERBS} STREQUAL "ON")
             target_link_libraries(_c_dataengine PRIVATE ibverbs rdmacm)

mindspore/ccsrc/pipeline/jit/action.cc

@@ -40,7 +40,7 @@
 #include "vm/transform.h"
 #include "parse/python_adapter.h"
 #include "frontend/optimizer/py_pass_manager.h"
-#if (!_WIN32 && !ENABLE_GE && !ENABLE_TESTCASES)
+#if (ENABLE_CPU && (ENABLE_D || ENABLE_GPU))
 #include "frontend/parallel/ps/parameter_server.h"
 #include "frontend/parallel/ps/scheduler.h"
 #include "frontend/parallel/ps/worker.h"
@@ -379,7 +379,7 @@ bool ExecuteAction(const ResourcePtr &res) {
   return true;
 }
 
-#if (!_WIN32 && !ENABLE_GE && !ENABLE_TESTCASES)
+#if (ENABLE_CPU && (ENABLE_D || ENABLE_GPU))
 bool StartPSWorkerAction(const ResourcePtr &res) {
   parallel::ps::Worker<float>::GetInstance().Run();
   return true;
@@ -505,7 +505,7 @@ std::vector<ActionItem> VmPipeline() {
   actions.emplace_back(std::make_pair("py_opt", OptActionPyStub));
 
   actions.emplace_back(std::make_pair("validate", ValidateAction));
-#if (!_WIN32 && !ENABLE_GE && !ENABLE_TESTCASES)
+#if (ENABLE_CPU && (ENABLE_D || ENABLE_GPU))
   if (parallel::ps::Util::IsRoleOfWorker()) {
     actions.emplace_back(std::make_pair("worker", StartPSWorkerAction));
   }
@@ -519,7 +519,7 @@ std::vector<ActionItem> VmPipeline() {
   return actions;
 }
 
-#if (!_WIN32 && !ENABLE_GE && !ENABLE_TESTCASES)
+#if (ENABLE_CPU && (ENABLE_D || ENABLE_GPU))
 std::vector<ActionItem> PServerPipeline() {
   auto actions = CommonPipeline();
   actions.emplace_back(std::make_pair("optimize", VmOptimizeAction));

mindspore/ccsrc/pipeline/jit/pipeline.cc

@@ -42,7 +42,7 @@
 #include "frontend/optimizer/py_pass_manager.h"
 #include "pybind_api/pybind_patch.h"
 
-#if (!_WIN32 && !ENABLE_GE && !ENABLE_TESTCASES)
+#if (ENABLE_CPU && (ENABLE_D || ENABLE_GPU))
 #include "frontend/parallel/ps/common.h"
 #include "frontend/parallel/ps/util.h"
 #endif
@@ -403,7 +403,7 @@ std::vector<ActionItem> GetPipline(const ResourcePtr &resource, const std::strin
 
   std::string backend = MsContext::GetInstance()->backend_policy();
 
-#if (!_WIN32 && !ENABLE_GE && !ENABLE_TESTCASES)
+#if (ENABLE_CPU && (ENABLE_D || ENABLE_GPU))
   if (mindspore::parallel::ps::Util::IsParamServerMode()) {
     mindspore::parallel::ps::Util::SetInternalEnvVar();
   }