alloc.cpp

// Compile with g++ -O3 alloc.cpp -lnuma
#include <vector>
#include <memory>
#include <iostream>
#include <unistd.h>
#include <assert.h>

#include <numa.h>
#include <numaif.h>

int USE_MBIND = 0;

static float *
myalloc(size_t sz) {
  float* arr = (float *)malloc(sz * sizeof(float));
  assert(arr);

  if (USE_MBIND) {
    void *ptr = arr;
    int node = numa_node_of_cpu(sched_getcpu());
    uintptr_t page_start_ptr =
        ((reinterpret_cast<uintptr_t>(ptr)) & ~(getpagesize() - 1));
    ptrdiff_t offset = reinterpret_cast<uintptr_t>(ptr) - page_start_ptr;
    unsigned long mask = 1UL << node;
    int err = mbind((void*)page_start_ptr,
      sz + offset,
      MPOL_BIND,
      &mask,
      sizeof(mask) * 8,
      MPOL_MF_MOVE | MPOL_MF_STRICT);
    assert(err == 0);
  }
 
  std::fill(&arr[0], &arr[sz], 1.0f);
  return arr;
}

int main(int argc, char* argv[]) {
  USE_MBIND = argc > 1 && strcmp(argv[1], "mbind") == 0;
  printf("USE_MBIND: %d\n", USE_MBIND);

  std::cout << "allocating memory\n";
  std::vector<float*> pile;
  // big allocation gets returned to OS
  pile.emplace_back(myalloc(1500*4*256*130));

  // but not small allocs if using mbind
  for (int i = 0; i < 1500; i++) {
    for (int n = 0; n < 4; n++) {
      pile.emplace_back(myalloc(256 * 64));
      pile.emplace_back(myalloc(256 * 64));
      pile.emplace_back(myalloc(256));
      pile.emplace_back(myalloc(256));
    }
  }

  std::cout << "waiting two seconds\n";
  sleep(2);
 
  std::cout << "deallocating memory\n";
  for (float* ptr : pile) {
    assert(ptr);
    free(ptr);
  }
  pile.clear();

  std::cout << "waiting forever\n";
  for (;;) {
    sleep(1);
  }
}