common: disable numa control when not summing locally (#157)

bytedance · Nov 20, 2019 · da1328d · da1328d
1 parent c1487f3
commit da1328d
Showing 1 changed file with 11 additions and 7 deletions.
diff --git a/byteps/common/shared_memory.cc b/byteps/common/shared_memory.cc
@@ -54,22 +54,26 @@ std::vector<void*> BytePSSharedMemory::openPcieSharedMemory(uint64_t key,
   for (int i = 0; i < BytePSGlobal::GetPcieSwitchNum(); i++) {
     auto prefix = std::string("BytePS_Pcie") + std::to_string(i) + "_Shm_";
     if (BytePSGlobal::IsDistributed()) {
-      if (i <= numa_max_node()) {
-        numa_set_preferred(i);
+      if (BytePSGlobal::IsCrossPcieSwitch()) {
+        if (i <= numa_max_node()) {
+          numa_set_preferred(i);
+          r.push_back(openSharedMemory(prefix, key, size));
+          numa_set_preferred(-1);
+        } else {
+          numa_set_preferred(numa_max_node());
+          r.push_back(openSharedMemory(prefix, key, size));
+          numa_set_preferred(-1);
+        }
       } else {
-        numa_set_preferred(numa_max_node());
+        r.push_back(openSharedMemory(prefix, key, size));
       }
-      r.push_back(openSharedMemory(prefix, key, size));
-      numa_set_preferred(-1);
     } else {
       if (BytePSGlobal::IsCrossPcieSwitch()) {
         numa_set_interleave_mask(numa_all_nodes_ptr);
         r.push_back(openSharedMemory(prefix, key, size));
         numa_set_interleave_mask(numa_no_nodes_ptr);
       } else {
-        numa_set_preferred(0);
         r.push_back(openSharedMemory(prefix, key, size));
-        numa_set_preferred(-1);
       }
     }
   }