/* StarPU --- Runtime system for heterogeneous multicore architectures.
 *
 * Copyright (C) 2010-2011, 2013-2014, 2020-2021  Université de Bordeaux 1
 * Copyright (C) 2010  Mehdi Juhoor <mjuhoor@gmail.com>
 * Copyright (C) 2010, 2011, 2012, 2013, 2014  Centre National de la Recherche Scientifique
 *
 * 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.
 */

/*
 * This example shows a simple implementation of a blocked matrix
 * multiplication. Note that this is NOT intended to be an efficient
 * implementation of sgemm! In this example, we show:
 *  - how to declare dense matrices (starpu_matrix_data_register)
 *  - how to manipulate matrices within codelets (eg. descr[0].blas.ld)
 *  - how to use filters to partition the matrices into blocks
 *    (starpu_data_partition and starpu_data_map_filters)
 *  - how to unpartition data (starpu_data_unpartition) and how to stop
 *    monitoring data (starpu_data_unregister)
 *  - how to manipulate subsets of data (starpu_data_get_sub_data)
 *  - how to construct an autocalibrated performance model (starpu_perfmodel)
 *  - how to submit asynchronous tasks
 */

#include <string.h>
#include <math.h>
#include <sys/types.h>
#include <sys/time.h>
#include <signal.h>

#include <starpu.h>

#include "mult.h"

unsigned nslicesx = 4;
unsigned nslicesy = 4;
unsigned xdim = 1024;
unsigned ydim = 1024;
unsigned zdim = 512;

/*
   We here benchmark the mult kernel implementation, both in terms of time and
   energy (when available).

   Energy measurement is usually not very precise (e.g. every 10ms), so we need
   to submit enough tasks to make the measurement long enough.
*/

int main(int argc, char **argv)
{
	int ret;
	unsigned ncpus, n, i;
	unsigned dim;

	parse_args(argc, argv);

	/* start the runtime */
	struct starpu_conf conf;
	starpu_conf_init(&conf);
	if (!getenv("STARPU_CALIBRATE"))
		conf.calibrate = 1;
	ret = starpu_init(&conf);
	if (ret == -ENODEV)
		return 77;
	STARPU_CHECK_RETURN_VALUE(ret, "starpu_init");

	ncpus = starpu_cpu_worker_get_count();
	if (ncpus == 0)
		return 77;

	for (dim = 16; dim <= 8192; dim *= 2)
	{
		xdim = ydim = dim;
		printf("calibrating tile size %dx%d\n", dim / nslicesx, dim / nslicesy);

		n = (8 << 20) / (dim * dim * ncpus);
		if (n == 0)
			n = 1;

		float *A[n], *B[n], *C[n];
		starpu_data_handle_t A_handle[n], B_handle[n], C_handle[n];

		for (i = 0; i < n; i++)
		{
			/* initialize matrices A, B and C and register them to StarPU */
			init_problem_data(&A[i], &B[i], &C[i]);

			/* partition matrices into blocks that can be manipulated by the
			 * codelets */
			partition_mult_data(A[i], B[i], C[i], &A_handle[i], &B_handle[i], &C_handle[i]);
		}

		/* Prepare task specimen */
		struct starpu_task *task = starpu_task_create();
		task->cl = &mult_cl;
		task->handles[0] = starpu_data_get_sub_data(A_handle[0], 1, 0);
		task->handles[1] = starpu_data_get_sub_data(B_handle[0], 1, 0);
		task->handles[2] = starpu_data_get_sub_data(C_handle[0], 2, 0, 0);
		task->flops = 2 * (xdim / nslicesx) * (ydim / nslicesy) * zdim;

		starpu_energy_start(-1, STARPU_CPU_WORKER);

		for (i = 0; i < n; i++)
		{
			/* submit all tasks in an asynchronous fashion */
			ret = launch_tasks(A_handle[i], B_handle[i], C_handle[i]);
			if (ret == -ENODEV) goto enodev;
		}

		/* wait for termination */
		starpu_task_wait_for_all();

		starpu_energy_stop(mult_cl.energy_model, task, 0,
				n * nslicesx * nslicesy, // Number of tasks,
				-1, STARPU_CPU_WORKER);

		for (i = 0; i < n; i++)
		{
			unpartition_mult_data(A_handle[i], B_handle[i], C_handle[i]);

			free(A[i]);
			free(B[i]);
			free(C[i]);
		}
	}

	starpu_shutdown();

	return 0;

enodev:
	starpu_shutdown();
	return 77;
}