10
10
#include < cerrno>
11
11
#include < algorithm>
12
12
13
+ #ifdef USE_LIBNUMA
14
+ #include < numa.h>
15
+ #include < numaif.h>
16
+ #include < sched.h>
17
+ #endif
18
+
13
19
#ifdef __has_include
14
20
#if __has_include(<unistd.h>)
15
21
#include < unistd.h>
@@ -273,6 +279,27 @@ struct llama_mmap::impl {
273
279
#ifdef _POSIX_MAPPED_FILES
274
280
std::vector<std::pair<size_t , size_t >> mapped_fragments;
275
281
282
+ #ifdef USE_LIBNUMA
283
+ static void move_pages (void *addr, size_t size) {
284
+ int cpu, ret;
285
+ struct bitmask *nodemask = numa_allocate_nodemask ();
286
+
287
+ /* Get memory policy of the calling thread. */
288
+ ret = get_mempolicy (nullptr , nodemask->maskp , nodemask->size , nullptr , 0 );
289
+ if (ret || numa_bitmask_weight (nodemask) == 0 ) {
290
+ cpu = sched_getcpu ();
291
+ if (cpu >= 0 ) {
292
+ numa_bitmask_clearall (nodemask);
293
+ numa_bitmask_setbit (nodemask, numa_node_of_cpu (cpu));
294
+ }
295
+ }
296
+ if (numa_bitmask_weight (nodemask) == 1 ) {
297
+ mbind (addr, size, MPOL_BIND, nodemask->maskp , nodemask->size , MPOL_MF_MOVE);
298
+ }
299
+ numa_free_nodemask (nodemask);
300
+ }
301
+ #endif
302
+
276
303
impl (struct llama_file * file, size_t prefetch, bool numa) {
277
304
size = file->size ();
278
305
int fd = file->file_id ();
@@ -291,6 +318,17 @@ struct llama_mmap::impl {
291
318
}
292
319
293
320
if (prefetch > 0 ) {
321
+ #ifdef USE_LIBNUMA
322
+ /*
323
+ * Given that we already pre-fault all memory when prefetch > 0, it is
324
+ * necessary to move any page cache pages that might have been
325
+ * instantiated during previous runs on different NUMA nodes. This call
326
+ * to move_pages() ensures that all memory-mapped pages are relocated
327
+ * according to the calling thread's memory policy or the CPU on which
328
+ * it is running.
329
+ */
330
+ move_pages (addr, file->size ());
331
+ #endif
294
332
if (posix_madvise (addr, std::min (file->size (), prefetch), POSIX_MADV_WILLNEED)) {
295
333
LLAMA_LOG_WARN (" warning: posix_madvise(.., POSIX_MADV_WILLNEED) failed: %s\n " ,
296
334
strerror (errno));
0 commit comments