xref: /aosp_15_r20/external/pytorch/c10/util/numa.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #pragma once
2 
3 #include <c10/macros/Export.h>
4 #include <c10/util/Flags.h>
5 #include <cstddef>
6 
7 C10_DECLARE_bool(caffe2_cpu_numa_enabled);
8 
9 namespace c10 {
10 
11 /**
12  * Check whether NUMA is enabled
13  */
14 C10_API bool IsNUMAEnabled();
15 
16 /**
17  * Bind to a given NUMA node
18  */
19 C10_API void NUMABind(int numa_node_id);
20 
21 /**
22  * Get the NUMA id for a given pointer `ptr`
23  */
24 C10_API int GetNUMANode(const void* ptr);
25 
26 /**
27  * Get number of NUMA nodes
28  */
29 C10_API int GetNumNUMANodes();
30 
31 /**
32  * Move the memory pointed to by `ptr` of a given size to another NUMA node
33  */
34 C10_API void NUMAMove(void* ptr, size_t size, int numa_node_id);
35 
36 /**
37  * Get the current NUMA node id
38  */
39 C10_API int GetCurrentNUMANode();
40 
41 } // namespace c10
42