memory: remove function starpu_malloc_count and implement it by combining starpu_malloc_set/get_flags and starpu_malloc

ChangeLog

@@ -94,11 +94,6 @@ New features:
     STARPU_LIMIT_OPENCL_devid_MEM to limit memory per specific device
   * Introduce new variable STARPU_LIMIT_CPU_MEM to limit memory for
     the CPU devices
-  * Define new functions starpu_malloc_count and starpu_free_count to
-    be used for allocating memory up to the limits defined by the
-    environment variables STARPU_LIMIT_xxx (see above). When no memory
-    is left, starpu_malloc_count tries to reclaim memory from StarPU
-    and returns -ENOMEM on failure.
   * New functions starpu_malloc_set_flags and starpu_malloc_get_flags
     to define flags to be used by the memory allocator.
     - STARPU_MALLOC_PINNED specifies memory should be pinned
@@ -106,6 +101,9 @@ New features:
       the limits defined by the environment variables STARPU_LIMIT_xxx
       (see above). When no memory is left, starpu_malloc tries
       to reclaim memory from StarPU and returns -ENOMEM on failure.
+  * Define new function starpu_free_count to
+    be used for freeing memory with a specified size. This is needed
+    to indicate the memory allocator the amount of memory.
 
 Small features:
   * Add starpu_worker_get_by_type and starpu_worker_get_by_devid

doc/chapters/basic-api.texi

@@ -282,20 +282,11 @@ This function frees memory which has previously allocated with
 @code{starpu_malloc}.
 @end deftypefun
 
-@deftypefun int starpu_malloc_count (void **@var{A}, size_t @var{dim})
-This function is similar to @code{starpu_malloc}. It only allocates
-memory up to the limit defined by the environment variables
-@code{STARPU_LIMIT_CUDA_devid_MEM}, @code{STARPU_LIMIT_CUDA_MEM},
-@code{STARPU_LIMIT_OPENCL_devid_MEM}, @code{STARPU_LIMIT_OPENCL_MEM}
-and @code{STARPU_LIMIT_CPU_MEM} (@pxref{Limit memory}). If no memory
-is available, it tries to reclaim memory from StarPU.
-Memory allocated through this function needs to be freed thanks to the
-@code{starpu_free_count} function.
-@end deftypefun
-
 @deftypefun int starpu_free_count (void *@var{A}, size_t @var{dim})
-This function frees memory which has previously allocated with
-@code{starpu_malloc_count}.
+This function frees memory by specifying its size. It should be used
+to free memory which was allocated following a call to
+@code{starpu_malloc_set_flags} with the value
+@code{STARPU_MALLOC_COUNT}.
 @end deftypefun
 
 @node Workers' Properties

include/starpu_stdlib.h

@@ -35,7 +35,6 @@ void starpu_malloc_set_align(size_t align);
 int starpu_malloc(void **A, size_t dim);
 int starpu_free(void *A);
 
-int starpu_malloc_count(void **A, size_t dim);
 int starpu_free_count(void *A, size_t dim);
 
 #ifdef __cplusplus

src/datawizard/malloc.c

@@ -29,6 +29,7 @@ static int _malloc_flags = 0;
 int starpu_malloc_set_flags(int flags)
 {
 	_malloc_flags = flags;
+	return _malloc_flags;
 }
 
 int starpu_malloc_get_flags()
@@ -98,6 +99,25 @@ int starpu_malloc(void **A, size_t dim)
 
 	STARPU_ASSERT(A);
 
+	if (_malloc_flags && STARPU_MALLOC_COUNT)
+	{
+		if (_starpu_memory_manager_can_allocate_size(dim, 0) == 0)
+		{
+			size_t freed;
+			size_t reclaim = 2 * dim;
+			_STARPU_DEBUG("There is not enough memory left, we are going to reclaim %ld\n", reclaim);
+			_STARPU_TRACE_START_MEMRECLAIM(0);
+			freed = _starpu_memory_reclaim_generic(0, 0, reclaim);
+			_STARPU_TRACE_END_MEMRECLAIM(0);
+			if (freed < dim)
+			{
+				// We could not reclaim enough memory
+				*A = NULL;
+				return -ENOMEM;
+			}
+		}
+	}
+
 #ifndef STARPU_SIMGRID
 	if (_starpu_can_submit_cuda_task())
 	{
@@ -301,31 +321,10 @@ int starpu_free(void *A)
 }
 
 
-int starpu_malloc_count(void **A, size_t dim)
-{
-	if (_starpu_memory_manager_can_allocate_size(dim, 0) == 0)
-	{
-		size_t freed;
-		size_t reclaim = 2 * dim;
-		_STARPU_DEBUG("There is not enough memory left, we are going to reclaim %ld\n", reclaim);
-		_STARPU_TRACE_START_MEMRECLAIM(0);
-		freed = _starpu_memory_reclaim_generic(0, 0, reclaim);
-		_STARPU_TRACE_END_MEMRECLAIM(0);
-		if (freed < dim)
-		{
-			// We could not reclaim enough memory
-			*A = NULL;
-			return -ENOMEM;
-		}
-	}
-
-	return starpu_malloc(A, dim);
-}
-
 int starpu_free_count(void *A, size_t dim)
 {
 	_starpu_memory_manager_deallocate_size(dim, 0);
-	starpu_free(A);
+	return starpu_free(A);
 }
 
 #ifdef STARPU_SIMGRID

tests/datawizard/allocate.c

@@ -53,28 +53,35 @@ int main(int argc, char **argv)
 	STARPU_CHECK_RETURN_VALUE_IS((int)global_size, 1*1024*1024, "_starpu_memory_manager_get_global_memory_size");
 	FPRINTF(stderr, "Available memory size on node 0: %ld\n", global_size);
 
-	ret = starpu_malloc_count((void **)&buffer, 1);
-	STARPU_CHECK_RETURN_VALUE(ret, "starpu_malloc_count");
+	starpu_malloc_set_flags(STARPU_MALLOC_COUNT);
+
+	ret = starpu_malloc((void **)&buffer, 1);
+	STARPU_CHECK_RETURN_VALUE(ret, "starpu_malloc");
 	FPRINTF(stderr, "Allocation succesfull for 1 b\n");
 
-	ret = starpu_malloc_count((void **)&buffer2, 1*1024*512);
-	STARPU_CHECK_RETURN_VALUE(ret, "starpu_malloc_count");
+	ret = starpu_malloc((void **)&buffer2, 1*1024*512);
+	STARPU_CHECK_RETURN_VALUE(ret, "starpu_malloc");
 	FPRINTF(stderr, "Allocation succesfull for %d b\n", 1*1024*512);
 
-	ret = starpu_malloc_count((void **)&buffer3, 1*1024*512);
-	STARPU_CHECK_RETURN_VALUE_IS(ret, -ENOMEM, "starpu_malloc_count");
+	ret = starpu_malloc((void **)&buffer3, 1*1024*512);
+	STARPU_CHECK_RETURN_VALUE_IS(ret, -ENOMEM, "starpu_malloc");
 	FPRINTF(stderr, "Allocation failed for %d b\n", 1*1024*512);
 
+	starpu_malloc_set_flags(0);
+
 	ret = starpu_malloc((void **)&buffer3, 1*1024*512);
 	STARPU_CHECK_RETURN_VALUE(ret, "starpu_malloc");
 	FPRINTF(stderr, "Allocation successful for %d b\n", 1*1024*512);
 	starpu_free(buffer3);
 
+	starpu_malloc_set_flags(STARPU_MALLOC_COUNT);
+
 	starpu_free_count(buffer2, 1*1024*512);
 	FPRINTF(stderr, "Freeing %d b\n", 1*1024*512);
 
-	ret = starpu_malloc_count((void **)&buffer3, 1*1024*512);
-	STARPU_CHECK_RETURN_VALUE(ret, "starpu_malloc_count");
+
+	ret = starpu_malloc((void **)&buffer3, 1*1024*512);
+	STARPU_CHECK_RETURN_VALUE(ret, "starpu_malloc");
 	FPRINTF(stderr, "Allocation succesfull for %d b\n", 1*1024*512);
 
 	starpu_free_count(buffer3, 1*1024*512);