task: Auto context save

Automatic context save

This patch changes the task creation API so that users no longer need to
specify a context save unit list for context switching. The areas of LS that
need to be context saved and restored will be determined automatically by the
task module. Users can specify either 0 for a run-complete task or
MARS_TASK_CONTEXT_SAVE_SIZE_MAX that allocates the maximum context save area
and uses only the area needed during context switching. Advanced users can
specify an explicit context save size if they can know the maximum context save
area size that will be required during the context switch.

Signed-off-by: Kazunori Asayama <asayama@sm.sony.co.jp>
Signed-off-by: Yuji Mano <yuji.mano@am.sony.com>

12 files changed, 105 insertions, 176 deletions

diff --git a/task/include/common/mars/task_types.h b/task/include/common/mars/task_types.h
index a50756e..8c67ef3 100644
--- a/task/include/common/mars/task_types.h
+++ b/task/include/common/mars/task_types.h
@@ -56,19 +56,7 @@
  * \ingroup group_mars_task
  * \brief Max size of context save area
  */
-#define MARS_TASK_CONTEXT_SAVE_SIZE_MAX		0x3C000
-
-/**
- * \ingroup group_mars_task
- * \brief Max number of context save units
- */
-#define MARS_TASK_CONTEXT_SAVE_UNIT_MAX		16
-
-/**
- * \ingroup group_mars_task
- * \brief Parameter for full context switch
- */
-#define MARS_TASK_CONTEXT_SAVE_ALL		mars_task_context_save_all
+#define MARS_TASK_CONTEXT_SAVE_SIZE_MAX		0x3c000
 
 /**
  * \ingroup group_mars_task
@@ -115,23 +103,4 @@ struct mars_task_args {
 	} type;
 };
 
-/**
- * \ingroup group_mars_task
- * \brief MARS task context save unit structure
- *
- * This structure is initialized by the user and passed into
- * \ref mars_task_create for MARS task creation.
- *
- * A list of this structure determines what areas of MPU storage
- * to save and restore during a task context switch.
- */
-struct mars_task_context_save_unit {
-	/** address relative to the base address of task */
-	uint32_t addr;
-	/** size of context save unit */
-	uint32_t size;
-};
-
-extern const struct mars_task_context_save_unit mars_task_context_save_all[];
-
 #endif
diff --git a/task/include/host/mars/task.h b/task/include/host/mars/task.h
index 45d4958..a00b044 100644
--- a/task/include/host/mars/task.h
+++ b/task/include/host/mars/task.h
@@ -83,44 +83,27 @@ extern "C" {
  * by this task.
  * - The elf image must remain allocated until the task is destroyed.
 
- * \e save_units
- * - A list of structures defining the areas of MPU storage to save and restore
- * during a task context switch.
- * - If NULL is specified, then no context save area will be allocated for
- * the task and therefore the task must be a run-complete task.
+ * \e context_save_size
+ * - If 0 is specified, then no context save area will be allocated for the task
+ * and therefore the task must be a run-complete task.
  * - A run-complete task will run and occupy an MPU until it has completed
  * running and exits.
  * - A run-complete task cannot context switch, and therefore cannot call
  * functions that will enter that task into a wait state.
- * - If \ref MARS_TASK_CONTEXT_SAVE_ALL is specified, the maximum size context
- * save area will be allocated and all of the MPU storage not occuppied
- * by the MARS kernel will be saved and restored.
- * This will guarantee all necessary areas of the task context will be
- * saved and restored during a context switch.
- * However, this also may not be the most efficient, since areas of MPU
- * storage not needed by the task may be saved and restored.
- * - For maximum efficiency, the user should initialize a static list of
- * struct \ref mars_task_context_save_unit that define each memory areas
- * to save and restore during a task context switch.
- * - A maximum of \ref MARS_TASK_CONTEXT_SAVE_UNIT_MAX separate context save
- * units can be specified.
- * - If you want to specify a context save unit list of size less than
- * \ref MARS_TASK_CONTEXT_SAVE_UNIT_MAX, you must specify a terminating unit
- * with an address and size of 0.
- * - Each context save unit defines an address offset from the task program's
- * base address (\ref MARS_TASK_BASE_ADDR) and the size of the unit.
- * - The address must be aligned to a 16-byte address and the size must be a
- * multiple of 16 bytes. Otherwise \ref MARS_ERROR_PARAMS will be returned.
- * - The list does need to be kept allocated after this function returns.
- *
- * \note The efficiency and validity of the specified context save unit
- * list is left up to the responsibility of the user.
+ * - If \ref MARS_TASK_CONTEXT_SAVE_SIZE_MAX is specified, the maximum size
+ * context save area required will be allocated and all of the MPU storage area
+ * used by the task will be saved and restored.
+ * - Advanced users can specify a size between 0 and
+ * \ref MARS_TASK_CONTEXT_SAVE_SIZE_MAX in order to minimize memory usage for
+ * the context save area.
+ * - The size specified must be large enough to hold all of the task program's
+ * text, data, heap, stack and non-volatile registers 80-127.
  *
  * \param[in] mars		- pointer to MARS context
  * \param[out] id		- address of pointer to task id instance
  * \param[in] name		- name of task
  * \param[in] elf_image		- address of MPU program elf image
- * \param[in] save_units	- pointer to list of context save units
+ * \param[in] context_save_size	- [ 0 ~ \ref MARS_TASK_CONTEXT_SAVE_SIZE_MAX ]
  * \return
  *	MARS_SUCCESS		- successfully created MARS task
  * \n	MARS_ERROR_NULL		- null pointer specified
@@ -132,7 +115,7 @@ int mars_task_create(struct mars_context *mars,
 		     struct mars_task_id *id,
 		     const char *name,
 		     const void *elf_image,
-		     const struct mars_task_context_save_unit *save_units);
+		     uint32_t context_save_size);
 
 /**
  * \ingroup group_mars_task
diff --git a/task/src/common/task_internal_types.h b/task/src/common/task_internal_types.h
index fc8be54..9740950 100644
--- a/task/src/common/task_internal_types.h
+++ b/task/src/common/task_internal_types.h
@@ -45,10 +45,6 @@
 #define MARS_TASK_CONTEXT_SIZE				128
 #define MARS_TASK_CONTEXT_ALIGN				16
 #define MARS_TASK_CONTEXT_SAVE_ALIGN			128
-#define MARS_TASK_CONTEXT_SAVE_UNIT_SIZE		8
-#define MARS_TASK_CONTEXT_SAVE_UNIT_ALIGN		128
-#define MARS_TASK_CONTEXT_SAVE_UNIT_ADDR_ALIGN_MASK	0xf
-#define MARS_TASK_CONTEXT_SAVE_UNIT_SIZE_ALIGN_MASK	0xf
 
 struct mars_task_context {
 	struct mars_workload_module module;	/* workload module */
@@ -59,7 +55,8 @@ struct mars_task_context {
 	uint32_t entry;				/* entry address of exec */
 	uint32_t stack;				/* stack pointer of exec */
 	uint64_t context_save_area_ea;		/* context save area */
-	uint64_t context_save_unit_ea;		/* context save unit list */
+	uint32_t context_save_area_low_size;	/* size of low save area */
+	uint32_t context_save_area_high_size;	/* size of high save area */
 	struct mars_task_id id;			/* task id */
 	struct mars_task_args args;		/* task args */
 } __attribute__((aligned(MARS_TASK_CONTEXT_ALIGN)));
diff --git a/task/src/host/lib/task.c b/task/src/host/lib/task.c
index 48a04d4..8e30eea 100644
--- a/task/src/host/lib/task.c
+++ b/task/src/host/lib/task.c
@@ -52,12 +52,6 @@
 
 extern const unsigned char mars_task_module_entry[];
 
-const struct mars_task_context_save_unit mars_task_context_save_all[] =
-{
-	{ .addr = 0, .size = MARS_TASK_CONTEXT_SAVE_SIZE_MAX },
-	{ .addr = 0, .size = 0 },
-};
-
 static uint64_t task_exit_code_ea(uint64_t task_ea)
 {
 	return task_ea + offsetof(struct mars_task_context, exit_code);
@@ -66,7 +60,7 @@ static uint64_t task_exit_code_ea(uint64_t task_ea)
 int mars_task_create(struct mars_context *mars,
 		     struct mars_task_id *id_ret,
 		     const char *name, const void *elf_image,
-		     const struct mars_task_context_save_unit *context_save_unit)
+		     uint32_t context_save_size)
 {
 	int ret;
 
@@ -77,10 +71,6 @@ int mars_task_create(struct mars_context *mars,
 	uint64_t workload_ea;
 	struct mars_task_context *task;
 
-	const struct mars_task_context_save_unit *p = context_save_unit;
-	int context_save_area_size = 0;
-	int context_save_unit_count = 1;
-
 	/* check function params */
 	if (!mars)
 		return MARS_ERROR_NULL;
@@ -90,6 +80,8 @@ int mars_task_create(struct mars_context *mars,
 		return MARS_ERROR_NULL;
 	if (name && strlen(name) > MARS_TASK_NAME_LEN_MAX)
 		return MARS_ERROR_PARAMS;
+	if (context_save_size > MARS_TASK_CONTEXT_SAVE_SIZE_MAX)
+		return MARS_ERROR_PARAMS;
 
 	/* process elf header information */
 	ehdr_module = (Elf32_Ehdr *)mars_task_module_entry;
@@ -148,62 +140,20 @@ int mars_task_create(struct mars_context *mars,
 		task->id.name[0] = 0;
 
 	/* no context save - run complete */
-	if (!context_save_unit) {
-		task->context_save_unit_ea = 0;
+	if (!context_save_size) {
 		task->context_save_area_ea = 0;
 		goto done;
 	}
 
-	/* calculate save unit count and save area size */
-	while (p->size) {
-		/* check for valid addr/size alignment and high addr boundary */
-		if (p->addr & MARS_TASK_CONTEXT_SAVE_UNIT_ADDR_ALIGN_MASK ||
-		    p->size & MARS_TASK_CONTEXT_SAVE_UNIT_SIZE_ALIGN_MASK ||
-		    p->addr + p->size > MARS_TASK_CONTEXT_SAVE_SIZE_MAX) {
-			ret = MARS_ERROR_PARAMS;
-			goto error_context_save_unit_addr_align;
-		}
-
-		/* increment save area size and check size limit */
-		context_save_area_size += p->size;
-		if (context_save_area_size > MARS_TASK_CONTEXT_SAVE_SIZE_MAX) {
-			ret = MARS_ERROR_PARAMS;
-			goto error_context_save_area_size_limit;
-		}
-
-		/* increment save unit count and check count limit */
-		context_save_unit_count++;
-		if (context_save_unit_count == MARS_TASK_CONTEXT_SAVE_UNIT_MAX)
-			break;
-
-		/* increment save unit pointer */
-		p++;
-	}
-
-	/* allocate context save unit storage */
-	task->context_save_unit_ea =
-		mars_ea_memalign(MARS_TASK_CONTEXT_SAVE_UNIT_ALIGN,
-				 MARS_TASK_CONTEXT_SAVE_UNIT_SIZE *
-				 context_save_unit_count);
-	if (!task->context_save_unit_ea) {
-		ret = MARS_ERROR_MEMORY;
-		goto error_malloc_context_save_unit;
-	}
-
 	/* allocate context save area */
 	task->context_save_area_ea =
 		mars_ea_memalign(MARS_TASK_CONTEXT_SAVE_ALIGN,
-				 context_save_area_size);
+				 context_save_size);
 	if (!task->context_save_area_ea) {
 		ret = MARS_ERROR_MEMORY;
 		goto error_malloc_context_save_area;
 	}
 
-	/* copy the context save struct into allocated storage */
-	mars_ea_put(task->context_save_unit_ea,
-		    context_save_unit,
-		    context_save_unit_count * MARS_TASK_CONTEXT_SAVE_UNIT_SIZE);
-
 done:
 	/* update task context on EA */
 	mars_ea_put(workload_ea, task, MARS_TASK_CONTEXT_SIZE);
@@ -222,10 +172,6 @@ done:
 error_workload_queue_add_end:
 	mars_ea_free(task->context_save_area_ea);
 error_malloc_context_save_area:
-	mars_ea_free(task->context_save_unit_ea);
-error_malloc_context_save_unit:
-error_context_save_area_size_limit:
-error_context_save_unit_addr_align:
 	mars_ea_unmap(task->exec_ea, task->exec_size);
 error_map_exec:
 	mars_ea_unmap(task->module.exec_ea, task->module.exec_size);
@@ -269,10 +215,6 @@ int mars_task_destroy(struct mars_task_id *id)
 	if (task->context_save_area_ea)
 		mars_ea_free(task->context_save_area_ea);
 
-	/* free the allocated context save unit storage if it has one */
-	if (task->context_save_unit_ea)
-		mars_ea_free(task->context_save_unit_ea);
-
 	/* unmap ELFs */
 	mars_ea_unmap(task->module.exec_ea, task->module.exec_size);
 	mars_ea_unmap(task->exec_ea, task->exec_size);
diff --git a/task/src/mpu/lib/task.c b/task/src/mpu/lib/task.c
index cfdfaf2..9c42d8c 100644
--- a/task/src/mpu/lib/task.c
+++ b/task/src/mpu/lib/task.c
@@ -92,7 +92,7 @@ int mars_task_yield(void)
 	if (!task->context_save_area_ea)
 		return MARS_ERROR_FORMAT;
 
-	(*mars_task_module_syscalls->yield)();
+	(*mars_task_module_syscalls->yield)(mars_task_get_heap());
 
 	return MARS_SUCCESS;
 }
@@ -123,7 +123,8 @@ int mars_task_wait(struct mars_task_id *id, int32_t *exit_code)
 	if (!task->context_save_area_ea)
 		return MARS_ERROR_FORMAT;
 
-	(*mars_task_module_syscalls->wait)(id->workload_id);
+	(*mars_task_module_syscalls->wait)(id->workload_id,
+					   mars_task_get_heap());
 
 	/* exit code requested so get it from the task context and return it */
 	if (exit_code) {
diff --git a/task/src/mpu/lib/task_barrier.c b/task/src/mpu/lib/task_barrier.c
index e0de5a9..c272d46 100644
--- a/task/src/mpu/lib/task_barrier.c
+++ b/task/src/mpu/lib/task_barrier.c
@@ -94,7 +94,7 @@ static int notify(uint64_t barrier_ea, int try)
 			(struct mars_mutex *)&barrier);
 
 		/* wait for signal */
-		(*mars_task_module_syscalls->signal_wait)();
+		(*mars_task_module_syscalls->signal_wait)(mars_task_get_heap());
 
 		mars_mutex_lock_get(barrier_ea, (struct mars_mutex *)&barrier);
 	}
@@ -170,7 +170,7 @@ static int wait(uint64_t barrier_ea, int try)
 				      (struct mars_mutex *)&barrier);
 
 		/* wait for signal */
-		(*mars_task_module_syscalls->signal_wait)();
+		(*mars_task_module_syscalls->signal_wait)(mars_task_get_heap());
 
 		mars_mutex_lock_get(barrier_ea, (struct mars_mutex *)&barrier);
 	}
diff --git a/task/src/mpu/lib/task_event_flag.c b/task/src/mpu/lib/task_event_flag.c
index 1083d5e..3cd5f91 100644
--- a/task/src/mpu/lib/task_event_flag.c
+++ b/task/src/mpu/lib/task_event_flag.c
@@ -211,7 +211,7 @@ static int wait(uint64_t event_flag_ea,uint32_t mask, uint8_t mask_mode,
 				      (struct mars_mutex *)&event_flag);
 
 		/* wait for signal */
-		(*mars_task_module_syscalls->signal_wait)();
+		(*mars_task_module_syscalls->signal_wait)(mars_task_get_heap());
 
 		mars_mutex_lock_get(event_flag_ea,
 				    (struct mars_mutex *)&event_flag);
diff --git a/task/src/mpu/lib/task_queue.c b/task/src/mpu/lib/task_queue.c
index d31dd43..1783f09 100644
--- a/task/src/mpu/lib/task_queue.c
+++ b/task/src/mpu/lib/task_queue.c
@@ -198,7 +198,7 @@ static int push(uint64_t queue_ea, const void *data,
 		mars_mutex_unlock_put(queue_ea, (struct mars_mutex *)&queue);
 
 		/* wait for signal */
-		(*mars_task_module_syscalls->signal_wait)();
+		(*mars_task_module_syscalls->signal_wait)(mars_task_get_heap());
 
 		mars_mutex_lock_get(queue_ea, (struct mars_mutex *)&queue);
 	}
@@ -351,7 +351,7 @@ static int pop(uint64_t queue_ea, void *data,
 		mars_mutex_unlock_put(queue_ea, (struct mars_mutex *)&queue);
 
 		/* wait for signal */
-		(*mars_task_module_syscalls->signal_wait)();
+		(*mars_task_module_syscalls->signal_wait)(mars_task_get_heap());
 
 		mars_mutex_lock_get(queue_ea, (struct mars_mutex *)&queue);
 	}
diff --git a/task/src/mpu/lib/task_semaphore.c b/task/src/mpu/lib/task_semaphore.c
index 749dc9e..74386fd 100644
--- a/task/src/mpu/lib/task_semaphore.c
+++ b/task/src/mpu/lib/task_semaphore.c
@@ -87,7 +87,7 @@ int mars_task_semaphore_acquire(uint64_t semaphore_ea)
 				      (struct mars_mutex *)&semaphore);
 
 		/* wait for signal */
-		(*mars_task_module_syscalls->signal_wait)();
+		(*mars_task_module_syscalls->signal_wait)(mars_task_get_heap());
 
 		return MARS_SUCCESS;
 	}
diff --git a/task/src/mpu/lib/task_signal.c b/task/src/mpu/lib/task_signal.c
index fdadf0f..b7e05e2 100644
--- a/task/src/mpu/lib/task_signal.c
+++ b/task/src/mpu/lib/task_signal.c
@@ -64,7 +64,7 @@ int mars_task_signal_wait(void)
 	if (!task->context_save_area_ea)
 		return MARS_ERROR_FORMAT;
 
-	return (*mars_task_module_syscalls->signal_wait)();
+	return (*mars_task_module_syscalls->signal_wait)(mars_task_get_heap());
 }
 
 int mars_task_signal_try_wait(void)
diff --git a/task/src/mpu/module/task_module.c b/task/src/mpu/module/task_module.c
index bd7cb88..a4266cc 100644
--- a/task/src/mpu/module/task_module.c
+++ b/task/src/mpu/module/task_module.c
@@ -48,8 +48,11 @@
 
 #include "task_module.h"
 
-#define MARS_TASK_MODULE_DMA_TAG	31
-#define MARS_TASK_MODULE_DMA_SIZE_MAX	16384
+#define MARS_TASK_MODULE_DMA_TAG		31
+#define MARS_TASK_MODULE_DMA_SIZE_MAX		16384
+
+#define MARS_TASK_MODULE_DMA_SIZE_MASK		0x7f
+#define MARS_TASK_REGISTER_SAVE_AREA_SIZE 	(16 * (127 - 80 + 1))
 
 /* stack pointer storage */
 void *__module_stack;
@@ -57,7 +60,6 @@ void *__task_stack;
 
 /* global task structs */
 static struct mars_task_context *task;
-static struct mars_task_context_save_unit list[MARS_TASK_CONTEXT_SAVE_UNIT_MAX];
 
 /* task entry */
 typedef void (*mars_task_entry)(struct mars_task_args *args,
@@ -69,7 +71,7 @@ static void dma_large(void *ls, uint64_t ea, uint32_t size, int put)
 		unsigned int block_size;
 
 		block_size = (size < MARS_TASK_MODULE_DMA_SIZE_MAX) ?
-			size : MARS_TASK_MODULE_DMA_SIZE_MAX;
+			      size : MARS_TASK_MODULE_DMA_SIZE_MAX;
 
 		if (put)
 			mfc_put((volatile void *)ls, ea, block_size,
@@ -222,35 +224,61 @@ static void registers_restore(void)
 		: : [ptr] "r" (__task_stack));
 }
 
-static void process_context_save_list(int save)
+static void dma_context(int save, void *task_heap)
 {
-	int i, offset = 0;
-
-	mfc_get(list, task->context_save_unit_ea,
-		MARS_TASK_CONTEXT_SAVE_UNIT_SIZE *
-		MARS_TASK_CONTEXT_SAVE_UNIT_MAX,
-		MARS_TASK_MODULE_DMA_TAG, 0, 0);
-	dma_wait();
-
-	/* loop through save unit list */
-	for (i = 0; i < MARS_TASK_CONTEXT_SAVE_UNIT_MAX; i++) {
-		/* list terminator so exit loop */
-		if (!list[i].size)
-			break;
+	int offset = 0;
+	uint32_t low_size, high_size;
+
+	if (save) {
+		/* text, data and heap (low address) */
+		low_size =
+			(((uintptr_t)task_heap -
+			  MARS_TASK_BASE_ADDR +
+			  MARS_TASK_MODULE_DMA_SIZE_MASK) &
+			 ~MARS_TASK_MODULE_DMA_SIZE_MASK);
+
+		/* stack (high address) */
+		high_size =
+			((MARS_TASK_BASE_ADDR +
+			  MARS_TASK_CONTEXT_SAVE_SIZE_MAX -
+			  (uintptr_t)__task_stack +
+			  MARS_TASK_REGISTER_SAVE_AREA_SIZE +
+			  MARS_TASK_MODULE_DMA_SIZE_MASK) &
+			 ~MARS_TASK_MODULE_DMA_SIZE_MASK);
+
+		/* full context save/restore */
+		if (low_size + high_size >= MARS_TASK_CONTEXT_SAVE_SIZE_MAX) {
+			low_size = MARS_TASK_CONTEXT_SAVE_SIZE_MAX;
+			high_size = 0;
+		}
+
+		/* save save sizes to task context */
+		task->context_save_area_low_size = low_size;
+		task->context_save_area_high_size = high_size;
+	}
+	else {
+		/* restore save sizes from task context */
+		low_size = task->context_save_area_low_size;
+		high_size = task->context_save_area_high_size;
+	}
 
-		/* save or restore context save unit specified */
-		dma_large((void *)(MARS_TASK_BASE_ADDR + list[i].addr),
-			  task->context_save_area_ea + offset,
-			  list[i].size, save);
+	/* save or restore text and heap (low address) */
+	dma_large((void *)MARS_TASK_BASE_ADDR,
+		  task->context_save_area_ea + offset,
+		  low_size, save);
+	offset += low_size;
 
-		offset += list[i].size;
-	}
+	/* save or restore stack (high addres) */
+	dma_large((void *)(MARS_TASK_BASE_ADDR +
+			   MARS_TASK_CONTEXT_SAVE_SIZE_MAX - high_size),
+		  task->context_save_area_ea + offset,
+		  high_size, save);
+	offset += high_size;
 
-	/* wait for all dmas to complete */
 	dma_wait();
 }
 
-static void __attribute__((noinline)) context_save(void)
+static void __attribute__((noinline)) context_save(void *task_heap)
 {
 	/* save registers state */
 	registers_save();
@@ -258,12 +286,12 @@ static void __attribute__((noinline)) context_save(void)
 	/* save workload stack pointer */
 	task->stack = (uint32_t)__task_stack;
 
-	process_context_save_list(1);
+	dma_context(1, task_heap);
 }
 
 static void __attribute__((noinline)) context_restore(void)
 {
-	process_context_save_list(0);
+	dma_context(0, NULL);
 
 	/* restore workload stack pointer */
 	__task_stack = (void *)task->stack;
@@ -272,7 +300,7 @@ static void __attribute__((noinline)) context_restore(void)
 	registers_restore();
 }
 
-static void __attribute__((noinline)) context_wait(void)
+static void __attribute__((noinline)) context_wait(void *task_heap)
 {
 	register void *sp asm("$sp");
 
@@ -283,7 +311,7 @@ static void __attribute__((noinline)) context_wait(void)
 	sp = __module_stack;
 
 	/* save task context */
-	context_save();
+	context_save(task_heap);
 
 	/* put workload into waiting state */
 	mars_module_workload_wait();
@@ -358,7 +386,7 @@ static int task_schedule(uint16_t workload_id, struct mars_task_args *args,
 	return MARS_SUCCESS;
 }
 
-static int task_wait(uint16_t workload_id)
+static int task_wait(uint16_t workload_id, void *task_heap)
 {
 	int ret;
 
@@ -370,7 +398,7 @@ static int task_wait(uint16_t workload_id)
 	if (ret != MARS_SUCCESS)
 		return ret;
 
-	context_wait();
+	context_wait(task_heap);
 
 	return mars_module_workload_wait_reset();
 }
@@ -398,9 +426,9 @@ static int task_signal_send(uint16_t workload_id)
 	return mars_module_workload_signal_set(workload_id);
 }
 
-static int task_signal_wait(void)
+static int task_signal_wait(void *task_heap)
 {
-	context_wait();
+	context_wait(task_heap);
 
 	return mars_module_workload_signal_reset();
 }
diff --git a/task/src/mpu/module/task_module.h b/task/src/mpu/module/task_module.h
index d42fd3e..1b73dae 100644
--- a/task/src/mpu/module/task_module.h
+++ b/task/src/mpu/module/task_module.h
@@ -39,6 +39,7 @@
 #define MARS_TASK_MODULE_H
 
 #include <stdint.h>
+#include <unistd.h>
 
 #include "mars/task_types.h"
 
@@ -53,17 +54,25 @@ struct mars_task_module_syscalls {
 						(struct mars_task_id *task_id);
 
 	void	(*exit)(void);
-	void	(*yield)(void);
+	void	(*yield)(void *heap);
 	int	(*schedule)(uint16_t workload_id, struct mars_task_args *args,
 			    uint8_t priority);
-	int	(*wait)(uint16_t workload_id);
+	int	(*wait)(uint16_t workload_id, void *heap);
 	int	(*try_wait)(uint16_t workload_id);
 	int	(*signal_host)(uint64_t watch_point_ea);
 	int	(*signal_send)(uint16_t workload_id);
-	int	(*signal_wait)(void);
+	int	(*signal_wait)(void *heap);
 	int	(*signal_try_wait)(void);
 };
 
+/* This function must be linked not with 'task module' but with 'task'
+ * to return task's heap.
+ */
+static inline void *mars_task_get_heap(void)
+{
+	return sbrk(0);
+}
+
 #if defined(__cplusplus)
 extern "C" {
 #endif