Move numa.{h, cc} to c10/util (#14393)
authorJerry Zhang <jerryzh@fb.com>
Thu, 6 Dec 2018 18:56:14 +0000 (10:56 -0800)
committerFacebook Github Bot <facebook-github-bot@users.noreply.github.com>
Thu, 6 Dec 2018 19:30:13 +0000 (11:30 -0800)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/14393

att

Reviewed By: ezyang

Differential Revision: D13205604

fbshipit-source-id: 54166492d31827b0343ed070cc36a825dd86e2ed

c10/CMakeLists.txt
c10/macros/cmake_macros.h.in
c10/util/numa.cpp [new file with mode: 0644]
c10/util/numa.h [new file with mode: 0644]
caffe2/core/macros.h.in
caffe2/core/numa.cc
caffe2/core/numa.h
caffe2/python/numa_benchmark.py
caffe2/python/numa_test.py

index 8e5ddc4..e80dd79 100644 (file)
@@ -16,6 +16,8 @@ set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 set(C10_USE_GFLAGS ${USE_GFLAGS}) # used in cmake_macros.h.in
 set(C10_USE_GLOG ${USE_GLOG}) # used in cmake_macros.h.in
 set(C10_BUILD_SHARED_LIBS ${BUILD_SHARED_LIBS}) # used in cmake_macros.h.in
+set(C10_USE_NUMA ${USE_NUMA})
+set(C10_DISABLE_NUMA ${CAFFE2_DISABLE_NUMA}) # used in cmake_macros.h.in
 configure_file(
     ${CMAKE_CURRENT_LIST_DIR}/macros/cmake_macros.h.in
     ${CMAKE_BINARY_DIR}/c10/macros/cmake_macros.h)
@@ -50,6 +52,20 @@ if (${USE_GLOG})
     target_link_libraries(c10 PUBLIC glog::glog)
 endif()
 
+if (USE_NUMA)
+  if (NOT CAFFE2_DISABLE_NUMA)
+    message(STATUS "NUMA paths:")
+    message(STATUS ${Numa_INCLUDE_DIR})
+    message(STATUS ${Numa_LIBRARIES})
+    include_directories(SYSTEM ${Numa_INCLUDE_DIR})
+    target_link_libraries(c10 PRIVATE ${Numa_LIBRARIES})
+  else()
+    message(STATUS "NUMA is disabled")
+  endif()
+else()
+  message(STATUS "don't use NUMA")
+endif()
+
 if (ANDROID)
     target_link_libraries(c10 PRIVATE log)
 endif()
index 29f8705..5b20360 100644 (file)
@@ -7,5 +7,6 @@
 #cmakedefine C10_BUILD_SHARED_LIBS
 #cmakedefine C10_USE_GLOG
 #cmakedefine C10_USE_GFLAGS
+#cmakedefine C10_DISABLE_NUMA
 
 #endif // C10_MACROS_CMAKE_MACROS_H_
diff --git a/c10/util/numa.cpp b/c10/util/numa.cpp
new file mode 100644 (file)
index 0000000..9fa9a93
--- /dev/null
@@ -0,0 +1,131 @@
+#include "c10/util/numa.h"
+
+#if defined(__linux__) && !defined(C10_DISABLE_NUMA) && C10_MOBILE == 0
+#include <numa.h>
+#include <numaif.h>
+#include <unistd.h>
+#define C10_ENABLE_NUMA
+#endif
+
+namespace c10 {
+
+#ifdef C10_ENABLE_NUMA
+bool IsNUMAEnabled() {
+  return numa_available() >= 0;
+}
+
+void NUMABind(int numa_node_id) {
+  if (numa_node_id < 0) {
+    return;
+  }
+  if (!IsNUMAEnabled()) {
+    VLOG(1) << "NUMA is not enabled";
+    return;
+  }
+
+  AT_CHECK(
+      numa_node_id <= numa_max_node(),
+      "NUMA node id ", numa_node_id, " is unavailable");
+
+  auto bm = numa_allocate_nodemask();
+  numa_bitmask_setbit(bm, numa_node_id);
+  numa_bind(bm);
+  numa_bitmask_free(bm);
+}
+
+int GetNUMANode(const void* ptr) {
+  if (!IsNUMAEnabled()) {
+    VLOG(1) << "NUMA is not enabled";
+    return -1;
+  }
+  AT_ASSERT(ptr);
+
+  int numa_node = -1;
+  AT_CHECK(
+      get_mempolicy(
+          &numa_node, NULL, 0, const_cast<void*>(ptr), MPOL_F_NODE | MPOL_F_ADDR) == 0,
+      "Unable to get memory policy, errno:", errno);
+  return numa_node;
+}
+
+int GetNumNUMANodes() {
+  if (!IsNUMAEnabled()) {
+    VLOG(1) << "NUMA is not enabled";
+    return -1;
+  }
+
+  return numa_num_configured_nodes();
+}
+
+void NUMAMove(void* ptr, size_t size, int numa_node_id) {
+  if (numa_node_id < 0) {
+    return;
+  }
+  if (!IsNUMAEnabled()) {
+    VLOG(1) << "NUMA is not enabled";
+    return;
+  }
+  AT_ASSERT(ptr);
+
+  uintptr_t page_start_ptr = ((reinterpret_cast<uintptr_t>(ptr)) & ~(getpagesize() - 1));
+  ptrdiff_t offset = reinterpret_cast<uintptr_t>(ptr) - page_start_ptr;
+  // Avoid extra dynamic allocation and NUMA api calls
+  AT_ASSERT(numa_node_id >= 0 && static_cast<unsigned>(numa_node_id) < sizeof(unsigned long) * 8);
+  unsigned long mask = 1UL << numa_node_id;
+  AT_CHECK(
+      mbind(
+          reinterpret_cast<void*>(page_start_ptr),
+          size + offset,
+          MPOL_BIND,
+          &mask,
+          sizeof(mask) * 8,
+          MPOL_MF_MOVE | MPOL_MF_STRICT) == 0,
+      "Could not move memory to a NUMA node");
+}
+
+int GetCurrentNUMANode() {
+  if (!IsNUMAEnabled()) {
+    VLOG(1) << "NUMA is not enabled";
+    return -1;
+  }
+
+  auto n = numa_node_of_cpu(sched_getcpu());
+  return n;
+}
+
+#else // C10_ENABLE_NUMA
+
+bool IsNUMAEnabled() {
+  return false;
+}
+
+void NUMABind(int numa_node_id) {
+  if (numa_node_id >= 0) {
+    VLOG(1) << "NUMA is not enabled";
+  }
+}
+
+int GetNUMANode(const void* ptr) {
+  VLOG(1) << "NUMA is not enabled";
+  return -1;
+}
+
+int GetNumNUMANodes() {
+  VLOG(1) << "NUMA is not enabled";
+  return -1;
+}
+
+void NUMAMove(void* ptr, size_t size, int numa_node_id) {
+  if (numa_node_id >= 0) {
+    VLOG(1) << "NUMA is not enabled";
+  }
+}
+
+int GetCurrentNUMANode() {
+  VLOG(1) << "NUMA is not enabled";
+  return -1;
+}
+
+#endif // C10_NUMA_ENABLED
+
+} // namespace c10
diff --git a/c10/util/numa.h b/c10/util/numa.h
new file mode 100644 (file)
index 0000000..37f8546
--- /dev/null
@@ -0,0 +1,38 @@
+#pragma once
+
+#include <c10/util/Logging.h>
+#include <c10/util/Optional.h>
+
+namespace c10 {
+
+/**
+ * Check whether NUMA is enabled
+ */
+C10_API bool IsNUMAEnabled();
+
+/**
+ * Bind to a given NUMA node
+ */
+C10_API void NUMABind(int numa_node_id);
+
+/**
+ * Get the NUMA id for a given pointer `ptr`
+ */
+C10_API int GetNUMANode(const void* ptr);
+
+/**
+ * Get number of NUMA nodes
+ */
+C10_API int GetNumNUMANodes();
+
+/**
+ * Move the memory pointed to by `ptr` of a given size to another NUMA node
+ */
+C10_API void NUMAMove(void* ptr, size_t size, int numa_node_id);
+
+/**
+ * Get the current NUMA node id
+ */
+C10_API int GetCurrentNUMANode();
+
+} // namespace c10
index f8a5b4e..8c2d18f 100644 (file)
@@ -44,7 +44,6 @@ static_assert(
 #cmakedefine CAFFE2_USE_MKLDNN
 #cmakedefine CAFFE2_USE_NVTX
 #cmakedefine CAFFE2_USE_TRT
-#cmakedefine CAFFE2_DISABLE_NUMA
 
 #ifndef USE_NUMPY
 #cmakedefine USE_NUMPY
@@ -81,5 +80,5 @@ static_assert(
   {"USE_MKLDNN", "${CAFFE2_USE_MKLDNN}"}, \
   {"USE_NVTX", "${CAFFE2_USE_NVTX}"}, \
   {"USE_TRT", "${CAFFE2_USE_TRT}"}, \
-  {"DISABLE_NUMA", "${CAFFE2_DISABLE_NUMA}"}, \
+  {"DISABLE_NUMA", "${CAFFE2_DISABLE_NUMA}"},   \
 }
index 953f162..c9b548b 100644 (file)
@@ -1,132 +1 @@
 #include "caffe2/core/numa.h"
-
-C10_DEFINE_bool(caffe2_cpu_numa_enabled, false, "Use NUMA whenever possible.");
-
-#if defined(__linux__) && !defined(CAFFE2_DISABLE_NUMA) && CAFFE2_MOBILE == 0
-#include <numa.h>
-#include <numaif.h>
-#define CAFFE2_NUMA_ENABLED
-#endif
-
-namespace caffe2 {
-
-#ifdef CAFFE2_NUMA_ENABLED
-bool IsNUMAEnabled() {
-  return FLAGS_caffe2_cpu_numa_enabled && numa_available() >= 0;
-}
-
-void NUMABind(int numa_node_id) {
-  if (numa_node_id < 0) {
-    return;
-  }
-  if (!IsNUMAEnabled()) {
-    VLOG(1) << "NUMA is not enabled";
-    return;
-  }
-
-  CAFFE_ENFORCE(
-      numa_node_id <= numa_max_node(),
-      "NUMA node id " + c10::to_string(numa_node_id) + " is unavailable");
-
-  auto bm = numa_allocate_nodemask();
-  numa_bitmask_clearall(bm);
-  numa_bitmask_setbit(bm, numa_node_id);
-  numa_bind(bm);
-  numa_bitmask_free(bm);
-}
-
-int GetNUMANode(const void* ptr) {
-  if (!IsNUMAEnabled()) {
-    VLOG(1) << "NUMA is not enabled";
-    return -1;
-  }
-  CAFFE_ENFORCE(ptr);
-
-  int numa_node = -1;
-  CAFFE_ENFORCE(
-      get_mempolicy(
-          &numa_node, NULL, 0, (void*)ptr, MPOL_F_NODE | MPOL_F_ADDR) == 0,
-      "Unable to get memory policy");
-  return numa_node;
-}
-
-int GetNumNUMANodes() {
-  if (!IsNUMAEnabled()) {
-    VLOG(1) << "NUMA is not enabled";
-    return -1;
-  }
-
-  return numa_num_configured_nodes();
-}
-
-void NUMAMove(void* ptr, size_t size, int numa_node_id) {
-  if (numa_node_id < 0) {
-    return;
-  }
-  if (!IsNUMAEnabled()) {
-    VLOG(1) << "NUMA is not enabled";
-    return;
-  }
-  CAFFE_ENFORCE(ptr);
-
-  size_t page_start_ptr = (((size_t)ptr) & ~(getpagesize() - 1));
-  size_t offset = ((size_t)ptr) - page_start_ptr;
-  // Avoid extra dynamic allocation and NUMA api calls
-  CAFFE_ENFORCE(numa_node_id >= 0 && (unsigned)numa_node_id < sizeof(unsigned long) * 8);
-  unsigned long mask = 1UL << numa_node_id;
-  CAFFE_ENFORCE(
-      mbind(
-          (void*)page_start_ptr,
-          size + offset,
-          MPOL_BIND,
-          &mask,
-          sizeof(mask) * 8,
-          MPOL_MF_MOVE | MPOL_MF_STRICT) == 0,
-      "Could not move memory to a NUMA node");
-}
-
-int GetCurrentNUMANode() {
-  if (!IsNUMAEnabled()) {
-    VLOG(1) << "NUMA is not enabled";
-    return -1;
-  }
-
-  return numa_node_of_cpu(sched_getcpu());
-}
-
-#else // CAFFE2_NUMA_ENABLED
-
-bool IsNUMAEnabled() {
-  return false;
-}
-
-void NUMABind(int numa_node_id) {
-  if (numa_node_id >= 0) {
-    VLOG(1) << "NUMA is not enabled";
-  }
-}
-
-int GetNUMANode(const void* ptr) {
-  VLOG(1) << "NUMA is not enabled";
-  return -1;
-}
-
-int GetNumNUMANodes() {
-  VLOG(1) << "NUMA is not enabled";
-  return -1;
-}
-
-void NUMAMove(void* ptr, size_t size, int numa_node_id) {
-  if (numa_node_id >= 0) {
-    VLOG(1) << "NUMA is not enabled";
-  }
-}
-
-int GetCurrentNUMANode() {
-  VLOG(1) << "NUMA is not enabled";
-  return -1;
-}
-
-#endif // CAFFE2_NUMA_ENABLED
-
-} // namespace caffe2
index 1e211b0..8424d54 100644 (file)
@@ -1,24 +1,3 @@
-#ifndef CAFFE2_CORE_NUMA_H_
-#define CAFFE2_CORE_NUMA_H_
-
-#include "caffe2/core/logging.h"
-
-C10_DECLARE_bool(caffe2_cpu_numa_enabled);
-
-namespace caffe2 {
-
-CAFFE2_API bool IsNUMAEnabled();
-
-CAFFE2_API void NUMABind(int numa_node_id);
-
-CAFFE2_API int GetNUMANode(const void* ptr);
-
-CAFFE2_API int GetNumNUMANodes();
-
-CAFFE2_API void NUMAMove(void* ptr, size_t size, int numa_node_id);
-
-CAFFE2_API int GetCurrentNUMANode();
-
-} // namespace caffe2
-
-#endif // CAFFE2_CORE_NUMA_H_
+#pragma once
+#include "c10/util/numa.h"
+#include "caffe2/core/common.h"
index 21c1cb1..3ef52b3 100644 (file)
@@ -65,5 +65,4 @@ def main():
 
 
 if __name__ == '__main__':
-    core.GlobalInit(["caffe2", "--caffe2_cpu_numa_enabled=1"])
     main()
index 3178345..36093b5 100644 (file)
@@ -7,9 +7,6 @@ from caffe2.proto import caffe2_pb2
 from caffe2.python.test_util import TestCase
 import unittest
 
-core.GlobalInit(["caffe2", "--caffe2_cpu_numa_enabled=1"])
-
-
 def build_test_net(net_name):
     net = core.Net(net_name)
     net.Proto().type = "async_scheduling"