expose memory manager functions to applications

ChangeLog

@@ -102,6 +102,8 @@ Small features:
     working implementations
   * Add STARPU_MALLOC_NORECLAIM flag to allocate without running a reclaim if
     the node is out of memory.
+  * Add starpu_memory_allocate and _deallocate to let the application declare
+    its own allocation to the reclaiming engine.
   * New flag STARPU_DATA_MODE_ARRAY for the function family
     starpu_task_insert to allow to define a array of data handles
     along with their access modes.

doc/doxygen/chapters/api/standard_memory_library.doxy

@@ -87,4 +87,38 @@ If a memory limit is defined on the given node (see Section \ref
 HowToLimitMemoryPerNode), return the amount of available memory
 on the node. Otherwise return -1.
 
+\fn int starpu_memory_allocate(unsigned node, size_t size, int flags)
+\ingroup API_Standard_Memory_Library
+If a memory limit is defined on the given node (see Section \ref
+HowToLimitMemoryPerNode), try to allocate some of it. This does not actually
+allocate memory, but only accounts for it. This can be useful when the
+application allocates data another way, but want StarPU to be aware of the
+allocation size e.g. for memory reclaiming.
+By default, the function returns -ENOMEM if there is not enough room on
+the given node. \p flags can be either STARPU_MEMORY_MANAGER_WAIT or
+STARPU_MEMORY_MANAGER_OVERFLOW to change this.
+
+\fn int starpu_memory_deallocate(unsigned node, size_t size)
+\ingroup API_Standard_Memory_Library
+If a memory limit is defined on the given node (see Section \ref
+HowToLimitMemoryPerNode), free some of it. This does not actually free memory,
+but only accounts for it, like starpu_memory_allocate(). The amount does not
+have to be exactly the same as what was passed to starpu_memory_allocate(),
+only the eventual amount needs to be the same, i.e. one call to
+starpu_memory_allocate() can be followed by several calls to
+starpu_memory_deallocate() to declare the deallocation piece by piece.
+
+\def STARPU_MEMORY_MANAGER_WAIT
+\ingroup API_Standard_Memory_Library
+Value passed to starpu_memory_allocate() to specify that the function should
+wait for the requested amount of memory to become available, and atomically
+allocate it.
+
+\def STARPU_MEMORY_MANAGER_OVERFLOW
+\ingroup API_Standard_Memory_Library
+Value passed to starpu_memory_allocate() to specify that the function should
+allocate the amount of memory, even if that means overflowing the total size of
+the memory node.
+
+
 */

include/starpu_stdlib.h

@@ -40,6 +40,27 @@ int starpu_free_flags(void *A, size_t dim, int flags);
 starpu_ssize_t starpu_memory_get_total(unsigned node);
 starpu_ssize_t starpu_memory_get_available(unsigned node);
 
+#define STARPU_MEMORY_WAIT (1)
+#define STARPU_MEMORY_OVERFLOW (2)
+
+/**
+ * Try to allocate memory on the given node
+ *
+ * @param size amount of memory to allocate
+ * @param node node where the memory is to be allocated
+ * @return 1 if the given amount of memory was allocated on the given node
+ */
+int starpu_memory_allocate(unsigned node, size_t size, int flags);
+
+/**
+ * Indicates the given amount of memory is going to be deallocated from the given node
+ *
+ * @param size amount of memory to be deallocated
+ * @param node node where the memory is going to be deallocated
+ */
+void starpu_memory_deallocate(unsigned node, size_t size);
+
+
 #ifdef __cplusplus
 }
 #endif

src/datawizard/malloc.c

@@ -93,7 +93,7 @@ int starpu_malloc_flags(void **A, size_t dim, int flags)
 	if (flags & STARPU_MALLOC_COUNT)
 	{
 		if (!(flags & STARPU_MALLOC_NORECLAIM))
-			while (_starpu_memory_manager_can_allocate_size(dim, STARPU_MAIN_RAM) == 0)
+			while (starpu_memory_allocate(STARPU_MAIN_RAM, dim, 0) != 0)
 			{
 				size_t freed;
 				size_t reclaim = 2 * dim;
@@ -109,7 +109,7 @@ int starpu_malloc_flags(void **A, size_t dim, int flags)
 				}
 			}
 		else
-			_starpu_memory_manager_allocate_size(dim, STARPU_MAIN_RAM);
+			starpu_memory_allocate(STARPU_MAIN_RAM, dim, STARPU_MEMORY_OVERFLOW);
 	}
 
 	if (flags & STARPU_MALLOC_PINNED && starpu_get_env_number("STARPU_DISABLE_PINNING") <= 0 && RUNNING_ON_VALGRIND == 0)
@@ -231,7 +231,7 @@ end:
 	}
 	else if (flags & STARPU_MALLOC_COUNT)
 	{
-		_starpu_memory_manager_deallocate_size(dim, 0);
+		starpu_memory_deallocate(STARPU_MAIN_RAM, dim);
 	}
 
 	return ret;
@@ -359,7 +359,7 @@ int starpu_free_flags(void *A, size_t dim, int flags)
 out:
 	if (flags & STARPU_MALLOC_COUNT)
 	{
-		_starpu_memory_manager_deallocate_size(dim, STARPU_MAIN_RAM);
+		starpu_memory_deallocate(STARPU_MAIN_RAM, dim);
 	}
 
 	return 0;
@@ -384,7 +384,7 @@ _starpu_malloc_on_node(unsigned dst_node, size_t size)
 	cudaError_t status;
 #endif
 
-	if (_starpu_memory_manager_can_allocate_size(size, dst_node) == 0)
+	if (starpu_memory_allocate(dst_node, size, 0) != 0)
 		return 0;
 
 	switch(starpu_node_get_kind(dst_node))
@@ -497,7 +497,7 @@ _starpu_malloc_on_node(unsigned dst_node, size_t size)
 		file = strrchr(__FILE__,'/');							
 		file += sizeof(char);										
 		_STARPU_TRACE_MEMORY_FULL(size);
-		_starpu_memory_manager_deallocate_size(size, dst_node);
+		starpu_memory_deallocate(dst_node, size);
 	}
 	return addr;
 }
@@ -588,7 +588,7 @@ _starpu_free_on_node(unsigned dst_node, uintptr_t addr, size_t size)
 		default:
 			STARPU_ABORT();
 	}
-	_starpu_memory_manager_deallocate_size(size, dst_node);
+	starpu_memory_deallocate(dst_node, size);
 
 }
 

src/datawizard/memalloc.c

@@ -1119,7 +1119,7 @@ get_better_disk_can_accept_size(starpu_data_handle_t handle, unsigned node)
 	for (i = 0; i < nnodes; i++)
 	{
 		if (starpu_node_get_kind(i) == STARPU_DISK_RAM && i != node &&
-		    (_starpu_memory_manager_test_allocate_size_(_starpu_data_get_size(handle), i) == 1 ||
+		    (_starpu_memory_manager_test_allocate_size(i, _starpu_data_get_size(handle)) == 1 ||
 		     handle->per_node[i].allocated))
 		{
 			/* if we can write on the disk */
@@ -1150,7 +1150,7 @@ choose_target(starpu_data_handle_t handle, unsigned node)
 		if(starpu_node_get_kind(handle->home_node) == STARPU_DISK_RAM && node != STARPU_MAIN_RAM)
 		{
 			if (handle->per_node[STARPU_MAIN_RAM].allocated || 
-			    _starpu_memory_manager_test_allocate_size_(size_handle, STARPU_MAIN_RAM) == 1)
+			    _starpu_memory_manager_test_allocate_size(STARPU_MAIN_RAM, size_handle) == 1)
 			{
 				target = STARPU_MAIN_RAM;
 			}
@@ -1176,7 +1176,7 @@ choose_target(starpu_data_handle_t handle, unsigned node)
 		/* node != 0 */
 		/* try to push data to RAM if we can before to push on disk*/
 		else if (handle->per_node[STARPU_MAIN_RAM].allocated || 
-			 _starpu_memory_manager_test_allocate_size_(size_handle, STARPU_MAIN_RAM) == 1)
+			 _starpu_memory_manager_test_allocate_size(STARPU_MAIN_RAM, size_handle) == 1)
 		{
 			target = STARPU_MAIN_RAM;
 		}

src/datawizard/memory_manager.c

@@ -18,10 +18,21 @@
 #include <common/utils.h>
 #include <common/thread.h>
 #include <datawizard/memory_manager.h>
+#include <starpu_stdlib.h>
 
 static size_t global_size[STARPU_MAXNODES];
 static size_t used_size[STARPU_MAXNODES];
+
+/* This is used as an optimization to avoid to wake up allocating threads for
+ * each and every deallocation, only to find that there is still not enough
+ * room.  */
+/* Minimum amount being waited for */
+static size_t min_waiting_size[STARPU_MAXNODES];
+/* Number of waiters */
+static int waiters[STARPU_MAXNODES];
+
 static starpu_pthread_mutex_t lock_nodes[STARPU_MAXNODES];
+static starpu_pthread_cond_t cond_nodes[STARPU_MAXNODES];
 
 int _starpu_memory_manager_init()
 {
@@ -31,7 +42,9 @@ int _starpu_memory_manager_init()
 	{
 		global_size[i] = 0;
 		used_size[i] = 0;
+		min_waiting_size[i] = 0;
 		STARPU_PTHREAD_MUTEX_INIT(&lock_nodes[i], NULL);
+		STARPU_PTHREAD_COND_INIT(&cond_nodes[i], NULL);
 	}
 	return 0;
 }
@@ -48,41 +61,58 @@ size_t _starpu_memory_manager_get_global_memory_size(unsigned node)
 }
 
 
-int _starpu_memory_manager_can_allocate_size(size_t size, unsigned node)
+int starpu_memory_allocate(unsigned node, size_t size, int flags)
 {
 	int ret;
 
 	STARPU_PTHREAD_MUTEX_LOCK(&lock_nodes[node]);
-	if (global_size[node] == 0)
+	if (flags == STARPU_MEMORY_WAIT)
 	{
-		// We do not have information on the available size, let's suppose it is going to fit
+		waiters[node]++;
+
+		/* Tell deallocators we need this amount */
+		if (!min_waiting_size[node] || size < min_waiting_size[node])
+			min_waiting_size[node] = size;
+
+		/* Wait for it */
+		while (used_size[node] + size > global_size[node])
+			STARPU_PTHREAD_COND_WAIT(&cond_nodes[node], &lock_nodes[node]);
+
+		/* And take it */
 		used_size[node] += size;
-		ret = 1;
+		ret = 0;
+
+		if (!--waiters[node])
+			/* Nobody is waiting any more, we can reset the minimum
+			 */
+			min_waiting_size[node] = 0;
 	}
-	else if (used_size[node] + size <= global_size[node])
+	else if (flags == STARPU_MEMORY_OVERFLOW
+			|| global_size[node] == 0
+			|| used_size[node] + size <= global_size[node])
 	{
 		used_size[node] += size;
-		ret = 1;
+		ret = 0;
 	}
 	else
 	{
-		ret = 0;
+		ret = -ENOMEM;
 	}
 	STARPU_PTHREAD_MUTEX_UNLOCK(&lock_nodes[node]);
 	return ret;
 }
 
-void _starpu_memory_manager_allocate_size(size_t size, unsigned node)
+void starpu_memory_deallocate(unsigned node, size_t size)
 {
 	STARPU_PTHREAD_MUTEX_LOCK(&lock_nodes[node]);
-	used_size[node] += size;
-	STARPU_PTHREAD_MUTEX_UNLOCK(&lock_nodes[node]);
-}
 
-void _starpu_memory_manager_deallocate_size(size_t size, unsigned node)
-{
-	STARPU_PTHREAD_MUTEX_LOCK(&lock_nodes[node]);
 	used_size[node] -= size;
+
+	/* If there's now room for waiters, wake them */
+	if (min_waiting_size[node] &&
+		global_size[node] - used_size[node] >= min_waiting_size[node])
+		STARPU_PTHREAD_COND_BROADCAST(&cond_nodes[node]);
+
 	STARPU_PTHREAD_MUTEX_UNLOCK(&lock_nodes[node]);
 }
 
@@ -102,7 +132,7 @@ starpu_ssize_t starpu_memory_get_available(unsigned node)
 		return global_size[node] - used_size[node];
 }
 
-int _starpu_memory_manager_test_allocate_size_(size_t size, unsigned node)
+int _starpu_memory_manager_test_allocate_size(unsigned node, size_t size)
 {
 	int ret;
 

src/datawizard/memory_manager.h

@@ -41,33 +41,7 @@ void _starpu_memory_manager_set_global_memory_size(unsigned node, size_t size);
  */
 size_t _starpu_memory_manager_get_global_memory_size(unsigned node);
 
-/**
- * Try to allocate memory on the given node
- *
- * @param size amount of memory to allocate
- * @param node node where the memory is to be allocated
- * @return 1 if the given amount of memory was allocated on the given node
- */
-int _starpu_memory_manager_can_allocate_size(size_t size, unsigned node) STARPU_WARN_UNUSED_RESULT;
-
-/**
- * Allocate memory on the given node, without caring about overflowing
- *
- * @param size amount of memory to allocate
- * @param node node where the memory is to be allocated
- */
-void _starpu_memory_manager_allocate_size(size_t size, unsigned node);
-
-/**
- * Indicates the given amount of memory is going to be deallocated from the given node
- *
- * @param size amount of memory to be deallocated
- * @param node node where the memory is going to be deallocated
- */
-void _starpu_memory_manager_deallocate_size(size_t size, unsigned node);
-
-
-int _starpu_memory_manager_test_allocate_size_(size_t size, unsigned node);
+int _starpu_memory_manager_test_allocate_size(unsigned node, size_t size);
 
 #ifdef __cplusplus
 }