/* StarPU --- Runtime system for heterogeneous multicore architectures.
 *
 * Copyright (C) 2009-2020  Université de Bordeaux, CNRS (LaBRI UMR 5800), Inria
 *
 * StarPU is free software; you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published by
 * the Free Software Foundation; either version 2.1 of the License, or (at
 * your option) any later version.
 *
 * StarPU is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 *
 * See the GNU Lesser General Public License in COPYING.LGPL for more details.
 */

#include <stdio.h>
#include <unistd.h>

#include <starpu.h>
#include "../helper.h"

/*
 * Measure the memory bandwidth available to kernels depending on the number of
 * kernels and number of idle workers.
 */

#ifdef STARPU_QUICK_CHECK
static size_t size = 1024;
static unsigned cpustep = 4;
#else
/* Must be bigger than available cache size per core, 64MiB should be enough */
static size_t size = 64UL << 20;
static unsigned cpustep = 1;
#endif

static unsigned noalone = 0;
static unsigned iter = 30;
static unsigned total_ncpus;
static starpu_pthread_barrier_t barrier;
static float *result;
static void **buffers;

void bw_func(void *descr[], void *arg)
{
	void *src = buffers[starpu_worker_get_id()];
	void *dst = (void*) ((uintptr_t)src + size);
	unsigned i;
	double start, stop;
	int ret;

	memset(src, 0, size);
	memset(dst, 0, size);

	STARPU_PTHREAD_BARRIER_WAIT(&barrier);
	start = starpu_timing_now();
	for (i = 0; i < iter; i++)
	{
		memcpy(dst, src, size);
		STARPU_SYNCHRONIZE();
	}
	stop = starpu_timing_now();
	STARPU_PTHREAD_BARRIER_WAIT(&barrier);

	result[starpu_worker_get_id()] = (size*iter) / (stop - start);
}

static struct starpu_codelet bw_codelet =
{
	.cpu_funcs = {bw_func},
	.model = NULL,
	.nbuffers = 0,
};

static void usage(char **argv)
{
	fprintf(stderr, "Usage: %s [-n iter] [-s size (MB)] [-i increment] [-a]\n", argv[0]);
	fprintf(stderr, "\t-n iter\tNumber of iterations\n");
	fprintf(stderr, "\t-s size\tBuffer size in MB\n");
	fprintf(stderr, "\t-i increment\tCpu number increment\n");
	fprintf(stderr, "\t-a\tDo not run the alone test\n");
	exit(EXIT_FAILURE);
}

static void parse_args(int argc, char **argv)
{
	int c;
	while ((c = getopt(argc, argv, "n:s:c:ah")) != -1)
	switch(c)
	{
		case 'n':
			iter = atoi(optarg);
			break;
		case 's':
			size = (long)atoi(optarg) << 20;
			break;
		case 'c':
			cpustep = atoi(optarg);
			break;
		case 'a':
			noalone = 1;
			break;
		case 'h':
			usage(argv);
			break;
	}
}

static unsigned interleave(unsigned i)
{
	/* TODO: rather distribute over hierarchy */
	if (total_ncpus > 1)
		return (i % (total_ncpus/2))*2 + i / (total_ncpus/2);
	else
		return 0;
}

static float bench(int *argc, char ***argv, unsigned nbusy, unsigned ncpus, int intl)
{
	int ret;
	unsigned i;
	struct starpu_conf conf;
	float bw;

	starpu_conf_init(&conf);
	conf.precedence_over_environment_variables = 1;
	conf.ncuda = 0;
	conf.nopencl = 0;
	conf.nmic = 0;
	conf.nmpi_ms = 0;
	conf.ncpus = ncpus;

	if (intl && nbusy == ncpus)
	{
		conf.use_explicit_workers_bindid = 1;
		for (i = 0; i < ncpus; i++)
			conf.workers_bindid[i] = interleave(i);
	}

	ret = starpu_initialize(&conf, argc, argv);
	if (ret == -ENODEV) return STARPU_TEST_SKIPPED;
	STARPU_CHECK_RETURN_VALUE(ret, "starpu_init");

	STARPU_PTHREAD_BARRIER_INIT(&barrier, NULL, nbusy);

	for (i = 0; i < nbusy; i++)
	{
		struct starpu_task *task = starpu_task_create();
		task->cl = &bw_codelet;
		task->execute_on_a_specific_worker = 1;
		if (intl && nbusy != ncpus)
			task->workerid = interleave(i);
		else
			task->workerid = i;
		ret = starpu_task_submit(task);
		STARPU_CHECK_RETURN_VALUE(ret, "starpu_task_submit");
	}

	starpu_task_wait_for_all();
	starpu_shutdown();

	for (bw = 0., i = 0; i < nbusy; i++)
	{
		bw += result[i];
	}
	return bw;
}

int main(int argc, char **argv)
{
	int ret;
	unsigned n;
	struct starpu_conf conf;
	float alone, alone_int, idle, idle_int;

	parse_args(argc, argv);

	starpu_conf_init(&conf);
	conf.precedence_over_environment_variables = 1;
	conf.ncuda = 0;
	conf.nopencl = 0;
	conf.nmic = 0;
	conf.nmpi_ms = 0;

	ret = starpu_initialize(&conf, &argc, &argv);
	if (ret == -ENODEV) return STARPU_TEST_SKIPPED;
	STARPU_CHECK_RETURN_VALUE(ret, "starpu_init");
	total_ncpus = starpu_cpu_worker_get_count();
	starpu_shutdown();

	if (total_ncpus == 0)
		return STARPU_TEST_SKIPPED;

	result = malloc(total_ncpus * sizeof(result[0]));
	buffers = malloc(total_ncpus * sizeof(*buffers));
	for (n = 0; n < total_ncpus; n++)
	{
#ifdef STARPU_HAVE_POSIX_MEMALIGN
		ret = posix_memalign(&buffers[n], getpagesize(), 2*size);
		STARPU_ASSERT(ret == 0);
#else
		buffers[n] = malloc(2*size);
#endif
	}

	printf("# nw\talone\t\t+idle\t\tefficiency\talone int.l\t+idle int.l\tefficiency\n");
	for (n = cpustep; n <= total_ncpus; n += cpustep)
	{
		if (noalone)
		{
			alone = 0.;
			alone_int = 0.;
		}
		else
		{
			alone = bench(&argc, &argv, n, n, 0);
			alone_int = bench(&argc, &argv, n, n, 1);
		}
		idle = bench(&argc, &argv, n, total_ncpus, 0);
		idle_int = bench(&argc, &argv, n, total_ncpus, 1);
		printf("%d\t%f\t%f\t%f\t%f\t%f\t%f\n", n, alone/1000, idle/1000, idle*100/alone, alone_int/1000, idle_int/1000, idle_int*100/alone_int);
		fflush(stdout);
	}

	free(result);

	for (n = 0; n < total_ncpus; n++)
		free(buffers[n]);

	return EXIT_SUCCESS;
}