Update the pointer -> handle mapping for lazily-allocated data.

src/datawizard/interfaces/data_interface.c

@@ -53,6 +53,22 @@ void _starpu_data_interface_shutdown()
 	registered_handles = NULL;
 }
 
+/* Register the mapping from PTR to HANDLE.  */
+void _starpu_data_register_local_pointer(starpu_data_handle handle, void *ptr)
+{
+	struct handle_entry *entry;
+
+	entry = malloc(sizeof(*entry));
+	STARPU_ASSERT(entry != NULL);
+
+	entry->pointer = ptr;
+	entry->handle = handle;
+
+	_starpu_spin_lock(&registered_handles_lock);
+	HASH_ADD_PTR(registered_handles, pointer, entry);
+	_starpu_spin_unlock(&registered_handles_lock);
+}
+
 starpu_data_handle starpu_data_lookup(const void *ptr)
 {
 	starpu_data_handle result;
@@ -190,18 +206,7 @@ static void _starpu_register_new_data(starpu_data_handle handle,
 	ptr = starpu_handle_to_pointer(handle);
 	if (ptr != NULL)
 	{
-		/* Register the mapping from PTR to HANDLE.  */
-		struct handle_entry *entry;
-
-		entry = malloc(sizeof(*entry));
-		STARPU_ASSERT(entry != NULL);
-
-		entry->pointer = ptr;
-		entry->handle = handle;
-
-		_starpu_spin_lock(&registered_handles_lock);
-		HASH_ADD_PTR(registered_handles, pointer, entry);
-		_starpu_spin_unlock(&registered_handles_lock);
+		_starpu_data_register_local_pointer(handle, ptr);
 	}
 }
 

src/datawizard/interfaces/data_interface.h

@@ -29,4 +29,8 @@ void _starpu_data_free_interfaces(starpu_data_handle handle)
 extern void _starpu_data_interface_init(void) STARPU_ATTRIBUTE_INTERNAL;
 extern void _starpu_data_interface_shutdown(void) STARPU_ATTRIBUTE_INTERNAL;
 
+extern void _starpu_data_register_local_pointer(starpu_data_handle handle,
+						void *ptr)
+	STARPU_ATTRIBUTE_INTERNAL;
+
 #endif // __DATA_INTERFACE_H__

src/datawizard/memalloc.c

@@ -746,6 +746,15 @@ int _starpu_allocate_memory_on_node(starpu_data_handle handle, struct starpu_dat
 	replicate->allocated = 1;
 	replicate->automatically_allocated = 1;
 
+	if (dst_node == _starpu_get_local_memory_node())
+	{
+		void *ptr = starpu_handle_to_pointer(handle);
+		if (ptr != NULL)
+		{
+			_starpu_data_register_local_pointer(handle, ptr);
+		}
+	}
+
 	return 0;
 }
 

tests/datawizard/data_lookup.c

@@ -21,7 +21,62 @@
 #include <stdlib.h>
 #include <sys/types.h>
 
-#include <stdio.h>
+static void task(void **buffers, void *args)
+{
+	float *numbers;
+	size_t size, i;
+
+	numbers = (float *) STARPU_VECTOR_GET_PTR(buffers[0]);
+	starpu_unpack_cl_args (args, &size);
+	for(i = 0; i < size; i++)
+	{
+		numbers[i] = i;
+	}
+}
+
+static starpu_codelet cl = {
+	.where = STARPU_CPU,
+	.cpu_func = task,
+	.nbuffers = 1
+};
+
+static void test_lazy_allocation()
+{
+	static const size_t count = 123;
+
+	size_t i;
+	void *pointer;
+	starpu_data_handle handle;
+
+	/* Lazily-allocated vector.  */
+	starpu_vector_data_register(&handle, -1, 0 /* NULL */,
+				    count, sizeof(float));
+
+	starpu_insert_task(&cl,
+			   STARPU_W, handle,
+			   STARPU_VALUE, &count, sizeof(float),
+			   0);
+
+	/* Acquire the handle, forcing a local allocation.  */
+	starpu_data_acquire(handle, STARPU_R);
+
+	/* Make sure we have a local pointer to it.  */
+	pointer = starpu_handle_to_pointer(handle);
+	assert(pointer != NULL);
+	for(i = 0; i < count; i++)
+	{
+		float *numbers = (float *)pointer;
+		assert(numbers[i] == i);
+	}
+
+	/* Make sure the pointer/handle mapping is up-to-date.  */
+	assert(starpu_data_lookup(pointer) == handle);
+
+	starpu_data_release(handle);
+	starpu_data_unregister(handle);
+
+	assert(starpu_data_lookup(pointer) == NULL);
+}
 
 #define VECTOR_COUNT    12
 #define VARIABLE_COUNT  42
@@ -108,6 +163,8 @@ int main(int argc, char *argv[])
 		starpu_free(vectors[i]);
 	}
 
+	test_lazy_allocation();
+
 	starpu_shutdown();
 
 	return EXIT_SUCCESS;