/* Link with "-lnuma"  */

/* Check that the allocation is done on the same node
   as the CPU executing the thread ('parition' = 'nearest').  */

#include <stdio.h>

#if !__has_include(<sched.h>) || !__has_include(<numaif.h>) || !__has_include(<dlfcn.h>)
int
main ()
{
  fprintf (stderr, "SKIPPED actual execution: (g)libc and/or libnuma header files not found\n");
  return 0;
}

#else

/* Assume that when numaif.h is available, the system also has a GLIBC/Kernel
   supporting 'sched_getcpu'.  */

#define __USE_GNU
#include <sched.h>   /* sched_getcpu - added in ec08f13dad for GLIBC_2.6. */
#include <dlfcn.h>   /* dlopen/dlcose  */
#include <numa.h>    /* numa_node_of_cpu  */
#include <numaif.h>  /* get_mempolicy  */
#include <assert.h>
#include <omp.h>

int main()
{
  const omp_alloctrait_t traits[]
    = { { omp_atk_fallback, omp_atv_abort_fb },
	{ omp_atk_partition, omp_atv_nearest } };
  int has_libnuma, has_memkind;

  void *handle = dlopen ("libnuma.so.1", RTLD_LAZY);
  has_libnuma = handle != NULL;
  dlclose (handle);

  handle = dlopen ("libmemkind.so.0", RTLD_LAZY);
  has_memkind = handle != NULL;
  dlclose (handle);

  if (!has_libnuma)
    {
      fprintf (stderr, "SKIPPED actual execution: dlopen of libnuma.so.1 failed (not installed?)\n");
      return 0;
    }
  fprintf (stderr, "Note: libmemkind.so.0 does %sexist\n", has_memkind ? "" : "NOT ");

  omp_allocator_handle_t a1, a2, a3, a4;

  a1 = omp_init_allocator (omp_default_mem_space, 2, traits);
  assert (a1 != omp_null_allocator);

  a2 = omp_init_allocator (omp_low_lat_mem_space, 2, traits);
  assert (a2 != omp_null_allocator);

  if (!has_memkind)
    {
      a3 = omp_init_allocator (omp_large_cap_mem_space, 2, traits);
      assert (a3 != omp_null_allocator);

      a4 = omp_init_allocator (omp_high_bw_mem_space, 2, traits);
      assert (a4 != omp_null_allocator);
    }

  #pragma omp parallel  /* Increase the chance that multiple nodes are used. */
  {
    unsigned node_cpu;
    int node_memory;
    int *i1, *i2, *i3, *i4;

    node_cpu = numa_node_of_cpu (sched_getcpu ());

    i1 = (int*) omp_alloc (sizeof(int)*1025*1, a1);
    assert (i1 != NULL);
    i1[1024] = 1;
    i2 = (int*) omp_alloc (sizeof(int)*1025*2, a2);
    assert (i2 != NULL);
    i2[1024] = 2;

    if (!has_memkind)
      {
        i3 = (int*) omp_alloc (sizeof(int)*1025*3, a3);
	assert (i3 != NULL);
        i3[1024] = 3;
        i4 = (int*) omp_alloc (sizeof(int)*1025*4, a4);
	assert (i3 != NULL);
        i4[1024] = 4;
      }

    node_memory = -1;
    if (get_mempolicy (&node_memory, NULL, 0, i1, MPOL_F_ADDR|MPOL_F_NODE) < 0)
      assert (0 && "get_mempolicy for i1");
    assert (node_memory == node_cpu);

    node_memory = -1;
    if (get_mempolicy (&node_memory, NULL, 0, i2, MPOL_F_ADDR|MPOL_F_NODE) < 0)
      assert (0 && "get_mempolicy for i2");
    assert (node_memory == node_cpu);

    if (!has_memkind)
      {
	node_memory = -1;
	if (get_mempolicy (&node_memory, NULL, 0, i3, MPOL_F_ADDR|MPOL_F_NODE) < 0)
	  assert (0 && "get_mempolicy for i3");
	assert (node_memory == node_cpu);

	node_memory = -1;
	if (get_mempolicy (&node_memory, NULL, 0, i4, MPOL_F_ADDR|MPOL_F_NODE) < 0)
	  assert (0 && "get_mempolicy for i4");
	assert (node_memory == node_cpu);
      }

     omp_free (i1, a1);
     omp_free (i2, a2);
    if (!has_memkind)
      {
	omp_free (i3, a3);
	omp_free (i4, a4);
      }
  }

  omp_destroy_allocator (a1);
  omp_destroy_allocator (a2);
  if (!has_memkind)
    {
      omp_destroy_allocator (a3);
      omp_destroy_allocator (a4);
    }
  return 0;
}

#endif