implement starpu_pack_cl_args and use it in starpu_insert_task and starpu_mpi_insert_task

mpi/starpu_mpi_insert_task.c

@@ -104,6 +104,7 @@ int starpu_mpi_insert_task(MPI_Comm comm, starpu_codelet *codelet, ...)
         va_list varg_list;
         int me, do_execute;
 	size_t arg_buffer_size = 0;
+	char *arg_buffer;
         int dest, execute, inconsistent_execute;
         int mpi_tag = 100;
 
@@ -121,7 +122,10 @@ int starpu_mpi_insert_task(MPI_Comm comm, starpu_codelet *codelet, ...)
 
         /* Get the number of buffers and the size of the arguments */
 	va_start(varg_list, codelet);
-        arg_buffer_size = starpu_insert_task_get_arg_size(varg_list);
+        arg_buffer_size = _starpu_insert_task_get_arg_size(varg_list);
+
+	va_start(varg_list, codelet);
+	_starpu_pack_cl_args(arg_buffer_size, &arg_buffer, varg_list);
 
         /* Finds out if the property STARPU_EXECUTE is specified */
         execute = -1;
@@ -293,7 +297,7 @@ int starpu_mpi_insert_task(MPI_Comm comm, starpu_codelet *codelet, ...)
                 _STARPU_MPI_DEBUG("Execution of the codelet %p\n", codelet);
                 va_start(varg_list, codelet);
                 struct starpu_task *task = starpu_task_create();
-                int ret = starpu_insert_task_create_and_submit(arg_buffer_size, codelet, &task, varg_list);
+                int ret = _starpu_insert_task_create_and_submit(arg_buffer, codelet, &task, varg_list);
                 _STARPU_MPI_DEBUG("ret: %d\n", ret);
                 STARPU_ASSERT(ret==0);
         }

src/util/starpu_insert_task.c

@@ -23,13 +23,25 @@
 #include <stdarg.h>
 #include <util/starpu_insert_task_utils.h>
 
+void starpu_pack_cl_args(char **arg_buffer, size_t *arg_buffer_size, ...)
+{
+	va_list varg_list;
+
+	/* Compute the size */
+	va_start(varg_list, arg_buffer_size);
+        *arg_buffer_size = _starpu_insert_task_get_arg_size(varg_list);
+
+	va_start(varg_list, arg_buffer_size);
+	_starpu_pack_cl_args(*arg_buffer_size, arg_buffer, varg_list);
+}
+
 void starpu_unpack_cl_args(void *_cl_arg, ...)
 {
 	unsigned char *cl_arg = _cl_arg;
-
 	unsigned current_arg_offset = 0;
 	va_list varg_list;
 
+	assert(cl_arg);
 	va_start(varg_list, _cl_arg);
 
 	/* We fill the different pointers with the appropriate arguments */
@@ -54,16 +66,17 @@ void starpu_insert_task(starpu_codelet *cl, ...)
 {
 	va_list varg_list;
 
-	/* The buffer will contain : nargs, {size, content} (x nargs)*/
-
 	/* Compute the size */
 	size_t arg_buffer_size = 0;
+	va_start(varg_list, cl);
+        arg_buffer_size = _starpu_insert_task_get_arg_size(varg_list);
 
 	va_start(varg_list, cl);
-        arg_buffer_size = starpu_insert_task_get_arg_size(varg_list);
+	char *arg_buffer;
+	_starpu_pack_cl_args(arg_buffer_size, &arg_buffer, varg_list);
 
 	va_start(varg_list, cl);
         struct starpu_task *task = starpu_task_create();
-        starpu_insert_task_create_and_submit(arg_buffer_size, cl, &task, varg_list);
+        _starpu_insert_task_create_and_submit(arg_buffer, cl, &task, varg_list);
 
 }

src/util/starpu_insert_task_utils.c

@@ -39,7 +39,7 @@ void starpu_task_insert_callback_wrapper(void *_cl_arg_wrapper)
 	free(cl_arg_wrapper->arg_stack);
 }
 
-size_t starpu_insert_task_get_arg_size(va_list varg_list)
+size_t _starpu_insert_task_get_arg_size(va_list varg_list)
 {
 	int arg_type;
         size_t arg_buffer_size;
@@ -77,20 +77,65 @@ size_t starpu_insert_task_get_arg_size(va_list varg_list)
         return arg_buffer_size;
 }
 
-int starpu_insert_task_create_and_submit(size_t arg_buffer_size, starpu_codelet *cl, struct starpu_task **task, va_list varg_list) {
+int _starpu_pack_cl_args(size_t arg_buffer_size, char **arg_buffer, va_list varg_list)
+{
         int arg_type;
-	unsigned current_buffer = 0;
+	unsigned current_arg_offset = 0;
 	unsigned char nargs = 0;
 
-	/* TODO use a single malloc to allocate the memory for arg_buffer and
-	 * the callback argument wrapper */
-	char *arg_buffer = malloc(arg_buffer_size);
-	STARPU_ASSERT(arg_buffer);
-	unsigned current_arg_offset = 0;
+	/* The buffer will contain : nargs, {size, content} (x nargs)*/
+
+	*arg_buffer = malloc(arg_buffer_size);
 
 	/* We will begin the buffer with the number of args (which is stored as a char) */
 	current_arg_offset += sizeof(char);
 
+	while((arg_type = va_arg(varg_list, int)) != 0)
+	{
+		if (arg_type==STARPU_R || arg_type==STARPU_W || arg_type==STARPU_RW || arg_type == STARPU_SCRATCH)
+		{
+			va_arg(varg_list, starpu_data_handle);
+		}
+		else if (arg_type==STARPU_VALUE)
+		{
+			/* We have a constant value: this should be followed by a pointer to the cst value and the size of the constant */
+			void *ptr = va_arg(varg_list, void *);
+			size_t cst_size = va_arg(varg_list, size_t);
+
+			*(size_t *)(&(*arg_buffer)[current_arg_offset]) = cst_size;
+			current_arg_offset += sizeof(size_t);
+
+			memcpy(&(*arg_buffer)[current_arg_offset], ptr, cst_size);
+			current_arg_offset += cst_size;
+
+			nargs++;
+			STARPU_ASSERT(current_arg_offset <= arg_buffer_size);
+		}
+		else if (arg_type==STARPU_CALLBACK)
+		{
+			va_arg(varg_list, void (*)(void *));
+		}
+		else if (arg_type==STARPU_CALLBACK_ARG) {
+			va_arg(varg_list, void *);
+		}
+		else if (arg_type==STARPU_PRIORITY)
+		{
+			va_arg(varg_list, int);
+		}
+		else if (arg_type==STARPU_EXECUTE) {
+			va_arg(varg_list, int);
+		}
+	}
+
+	(*arg_buffer)[0] = nargs;
+	va_end(varg_list);
+	return 0;
+}
+
+int _starpu_insert_task_create_and_submit(char *arg_buffer, starpu_codelet *cl, struct starpu_task **task, va_list varg_list) {
+        int arg_type;
+	unsigned current_buffer = 0;
+
 	struct insert_task_cb_wrapper *cl_arg_wrapper = malloc(sizeof(struct insert_task_cb_wrapper));
 	STARPU_ASSERT(cl_arg_wrapper);
 
@@ -113,18 +158,8 @@ int starpu_insert_task_create_and_submit(size_t arg_buffer_size, starpu_codelet
 		}
 		else if (arg_type==STARPU_VALUE)
 		{
-			/* We have a constant value: this should be followed by a pointer to the cst value and the size of the constant */
-			void *ptr = va_arg(varg_list, void *);
-			size_t cst_size = va_arg(varg_list, size_t);
-
-			*(size_t *)(&arg_buffer[current_arg_offset]) = cst_size;
-			current_arg_offset += sizeof(size_t);
-
-			memcpy(&arg_buffer[current_arg_offset], ptr, cst_size);
-			current_arg_offset += cst_size;
-
-			nargs++;
-			STARPU_ASSERT(current_arg_offset <= arg_buffer_size);
+			va_arg(varg_list, void *);
+			va_arg(varg_list, size_t);
 		}
 		else if (arg_type==STARPU_CALLBACK)
 		{
@@ -149,8 +184,6 @@ int starpu_insert_task_create_and_submit(size_t arg_buffer_size, starpu_codelet
 
 	va_end(varg_list);
 
-	arg_buffer[0] = nargs;
-
 	STARPU_ASSERT(current_buffer == cl->nbuffers);
 
 	(*task)->cl = cl;

src/util/starpu_insert_task_utils.h

@@ -21,8 +21,9 @@
 #include <stdarg.h>
 #include <starpu.h>
 
-size_t starpu_insert_task_get_arg_size(va_list varg_list);
-int starpu_insert_task_create_and_submit(size_t arg_buffer_size, starpu_codelet *cl, struct starpu_task **task, va_list varg_list);
+size_t _starpu_insert_task_get_arg_size(va_list varg_list);
+int _starpu_pack_cl_args(size_t arg_buffer_size, char **arg_buffer, va_list varg_list);
+int _starpu_insert_task_create_and_submit(char *arg_buffer, starpu_codelet *cl, struct starpu_task **task, va_list varg_list);
 
 #endif // __STARPU_INSERT_TASK_UTILS_H__