Add UfoBaseScheduler::set_gpu_nodes

This allows the user to specify which GPUs to use from a particular arch graph.

9 files changed, 73 insertions, 21 deletions

diff --git a/ufo/ufo-arch-graph.c b/ufo/ufo-arch-graph.c
index 7ae0809..3667a30 100644
--- a/ufo/ufo-arch-graph.c
+++ b/ufo/ufo-arch-graph.c
@@ -219,6 +219,7 @@ ufo_arch_graph_constructed (GObject *object)
 
     for (guint i = 0; i < priv->n_gpus; i++) {
         priv->gpu_nodes[i] = ufo_gpu_node_new (g_list_nth_data (cmd_queues, i));
+        g_debug ("Create new UfoGpuNode-%p", (gpointer) priv->gpu_nodes[i]);
     }
 
     /* Create remote nodes */
diff --git a/ufo/ufo-base-scheduler.c b/ufo/ufo-base-scheduler.c
index e666f4f..d2b5f5e 100644
--- a/ufo/ufo-base-scheduler.c
+++ b/ufo/ufo-base-scheduler.c
@@ -50,6 +50,7 @@ G_DEFINE_TYPE_WITH_CODE (UfoBaseScheduler, ufo_base_scheduler, G_TYPE_OBJECT,
 struct _UfoBaseSchedulerPrivate {
     GError          *construct_error;
     UfoArchGraph    *arch;
+    GList           *gpu_nodes;
     gboolean         expand;
     gboolean         trace;
     gboolean         rerun;
@@ -119,6 +120,50 @@ ufo_base_scheduler_get_arch (UfoBaseScheduler *scheduler)
     return scheduler->priv->arch;
 }
 
+/**
+ * ufo_base_scheduler_set_gpu_nodes:
+ * @scheduler: A #UfoBaseScheduler
+ * @arch: A #UfoArchGraph from which the nodes come from
+ * @gpu_nodes: (element-type Ufo.GpuNode): A list of #UfoGpuNode objects.
+ *
+ * Sets the GPU nodes that @scheduler can only use. Note, that the #UfoGpuNode
+ * objects must be from the same #UfoArchGraph that is returned by
+ * ufo_base_scheduler_get_arch.
+ */
+void
+ufo_base_scheduler_set_gpu_nodes (UfoBaseScheduler *scheduler,
+                                  UfoArchGraph *arch,
+                                  GList *gpu_nodes)
+{
+    g_return_if_fail (UFO_IS_BASE_SCHEDULER (scheduler));
+    g_return_if_fail (arch != NULL);
+
+    if (scheduler->priv->arch != NULL)
+        g_object_unref (scheduler->priv->arch);
+
+    scheduler->priv->arch = g_object_ref (arch);
+    scheduler->priv->gpu_nodes = g_list_copy (gpu_nodes);
+}
+
+/**
+ * ufo_base_scheduler_get_gpu_nodes:
+ * @scheduler: A #UfoBaseScheduler
+ *
+ * Get the GPU nodes that @scheduler can use for execution.
+ *
+ * Returns: (transfer none) (element-type Ufo.GpuNode): A list of #UfoGpuNode objects.
+ */
+GList *
+ufo_base_scheduler_get_gpu_nodes (UfoBaseScheduler *scheduler)
+{
+    g_return_val_if_fail (UFO_IS_BASE_SCHEDULER (scheduler), NULL);
+
+    if (scheduler->priv->gpu_nodes != NULL)
+        return scheduler->priv->gpu_nodes;
+
+    return ufo_arch_graph_get_gpu_nodes (ufo_base_scheduler_get_arch (scheduler));
+}
+
 static void
 ufo_base_scheduler_run_real (UfoBaseScheduler *scheduler,
                              UfoTaskGraph *graph,
@@ -304,4 +349,5 @@ ufo_base_scheduler_init (UfoBaseScheduler *scheduler)
     priv->rerun = FALSE;
     priv->ran = FALSE;
     priv->time = 0.0;
+    priv->gpu_nodes = NULL;
 }
diff --git a/ufo/ufo-base-scheduler.h b/ufo/ufo-base-scheduler.h
index 7d53773..063b0a3 100644
--- a/ufo/ufo-base-scheduler.h
+++ b/ufo/ufo-base-scheduler.h
@@ -76,6 +76,10 @@ void            ufo_base_scheduler_run              (UfoBaseScheduler   *schedul
                                                      UfoTaskGraph       *task_graph,
                                                      GError            **error);
 UfoArchGraph   *ufo_base_scheduler_get_arch         (UfoBaseScheduler   *scheduler);
+void            ufo_base_scheduler_set_gpu_nodes    (UfoBaseScheduler   *scheduler,
+                                                     UfoArchGraph       *arch,
+                                                     GList              *gpu_nodes);
+GList          *ufo_base_scheduler_get_gpu_nodes    (UfoBaseScheduler   *scheduler);
 GType           ufo_base_scheduler_get_type         (void);
 GQuark          ufo_base_scheduler_error_quark      (void);
 
diff --git a/ufo/ufo-fixed-scheduler.c b/ufo/ufo-fixed-scheduler.c
index f790639..13d6414 100644
--- a/ufo/ufo-fixed-scheduler.c
+++ b/ufo/ufo-fixed-scheduler.c
@@ -407,7 +407,7 @@ append_if_not_existing (GList *list, UfoTask *task)
 
 static ProcessData *
 setup_tasks (UfoGraph *graph,
-             UfoArchGraph *arch,
+             UfoBaseScheduler *scheduler,
              UfoResources *resources,
              GError **error)
 {
@@ -422,7 +422,7 @@ setup_tasks (UfoGraph *graph,
     data->tasks = NULL;
 
     nodes = ufo_graph_get_nodes (graph);
-    gpu_nodes = ufo_arch_graph_get_gpu_nodes (arch);
+    gpu_nodes = ufo_base_scheduler_get_gpu_nodes (scheduler);
 
     g_list_for (nodes, it) {
         UfoNode *source_node;
@@ -502,7 +502,7 @@ ufo_fixed_scheduler_run (UfoBaseScheduler *scheduler,
 
     arch = ufo_base_scheduler_get_arch (scheduler);
     resources = ufo_arch_graph_get_resources (arch);
-    pdata = setup_tasks (UFO_GRAPH (task_graph), arch, resources, &tmp_error);
+    pdata = setup_tasks (UFO_GRAPH (task_graph), scheduler, resources, &tmp_error);
 
     if (tmp_error != NULL) {
         g_propagate_error (error, tmp_error);
diff --git a/ufo/ufo-group-scheduler.c b/ufo/ufo-group-scheduler.c
index ba6dced..e62dc7c 100644
--- a/ufo/ufo-group-scheduler.c
+++ b/ufo/ufo-group-scheduler.c
@@ -100,7 +100,7 @@ ufo_group_scheduler_new (void)
 }
 
 static gboolean
-expand_group_graph (UfoGraph *graph, UfoArchGraph *arch, GError **error)
+expand_group_graph (UfoBaseScheduler *scheduler, UfoGraph *graph, GError **error)
 {
     GList *nodes;
     GList *it;
@@ -108,7 +108,7 @@ expand_group_graph (UfoGraph *graph, UfoArchGraph *arch, GError **error)
     guint n_gpus;
     gboolean success = TRUE;
 
-    gpu_nodes = ufo_arch_graph_get_gpu_nodes (arch);
+    gpu_nodes = ufo_base_scheduler_get_gpu_nodes (scheduler);
     n_gpus = g_list_length (gpu_nodes);
 
     nodes = ufo_graph_get_nodes (graph);
@@ -225,7 +225,7 @@ build_group_graph (UfoBaseScheduler *scheduler, UfoTaskGraph *graph, UfoArchGrap
     g_list_free (nodes);
     g_hash_table_destroy (tasks_to_groups);
 
-    if (!expand_group_graph (result, arch, error)) {
+    if (!expand_group_graph (scheduler, result, error)) {
         g_object_unref (result);
         return NULL;
     }
diff --git a/ufo/ufo-local-scheduler.c b/ufo/ufo-local-scheduler.c
index a37bb8d..f9cc549 100644
--- a/ufo/ufo-local-scheduler.c
+++ b/ufo/ufo-local-scheduler.c
@@ -362,7 +362,7 @@ ufo_local_scheduler_run (UfoBaseScheduler *scheduler,
 
     arch = ufo_base_scheduler_get_arch (scheduler);
     resources = ufo_arch_graph_get_resources (arch);
-    gpu_nodes = ufo_arch_graph_get_gpu_nodes (arch);
+    gpu_nodes = ufo_base_scheduler_get_gpu_nodes (scheduler);
     pp = ufo_pp_new (gpu_nodes);
     g_list_free (gpu_nodes);
 
diff --git a/ufo/ufo-scheduler.c b/ufo/ufo-scheduler.c
index 5b7acbe..c41bcba 100644
--- a/ufo/ufo-scheduler.c
+++ b/ufo/ufo-scheduler.c
@@ -591,6 +591,7 @@ ufo_scheduler_run (UfoBaseScheduler *scheduler,
     UfoSchedulerPrivate *priv;
     UfoArchGraph *arch;
     UfoTaskGraph *graph;
+    GList *gpu_nodes;
     GList *groups;
     guint n_nodes;
     GThread **threads;
@@ -624,6 +625,7 @@ ufo_scheduler_run (UfoBaseScheduler *scheduler,
         return;
 
     arch = ufo_base_scheduler_get_arch (scheduler);
+    gpu_nodes = ufo_base_scheduler_get_gpu_nodes (scheduler);
 
     if (priv->mode == UFO_REMOTE_MODE_REPLICATE) {
         replicate_task_graph (graph, arch);
@@ -631,11 +633,11 @@ ufo_scheduler_run (UfoBaseScheduler *scheduler,
 
     if (expand) {
         gboolean expand_remote = priv->mode == UFO_REMOTE_MODE_STREAM;
-        ufo_task_graph_expand (graph, arch, expand_remote);
+        ufo_task_graph_expand (graph, arch, g_list_length (gpu_nodes), expand_remote);
     }
 
     propagate_partition (graph);
-    ufo_task_graph_map (graph, arch);
+    ufo_task_graph_map (graph, gpu_nodes);
 
     /* Prepare task structures */
     tlds = setup_tasks (scheduler, graph, error);
diff --git a/ufo/ufo-task-graph.c b/ufo/ufo-task-graph.c
index 02a8dfb..61db0e7 100644
--- a/ufo/ufo-task-graph.c
+++ b/ufo/ufo-task-graph.c
@@ -400,6 +400,7 @@ expand_remotes (UfoTaskGraph *task_graph,
 void
 ufo_task_graph_expand (UfoTaskGraph *task_graph,
                        UfoArchGraph *arch_graph,
+                       guint n_gpus,
                        gboolean expand_remote)
 {
     GList *path;
@@ -413,7 +414,6 @@ ufo_task_graph_expand (UfoTaskGraph *task_graph,
     if (path != NULL && g_list_length (path) > 1) {
         GList *predecessors;
         GList *successors;
-        guint n_gpus;
 
         g_object_unref (UFO_NODE (g_list_first(path)->data));
         g_object_unref (UFO_NODE (g_list_last(path)->data));
@@ -446,7 +446,6 @@ ufo_task_graph_expand (UfoTaskGraph *task_graph,
             g_list_free (remotes);
         }
 
-        n_gpus = ufo_arch_graph_get_num_gpus (arch_graph);
         g_debug ("Expand for %i GPU nodes", n_gpus);
 
         for (guint i = 1; i < n_gpus; i++)
@@ -485,9 +484,9 @@ map_proc_node (UfoGraph *graph,
     if ((ufo_task_uses_gpu (UFO_TASK (node)) || UFO_IS_INPUT_TASK (node)) &&
         (!ufo_task_node_get_proc_node (UFO_TASK_NODE (node)))) {
 
-        g_debug ("Mapping GPU %i to %s-%p",
-                 proc_index, G_OBJECT_TYPE_NAME (node),
-                 (gpointer) node);
+        g_debug ("Mapping UfoGpuNode-%p to %s-%p",
+                 (gpointer) proc_node,
+                 G_OBJECT_TYPE_NAME (node), (gpointer) node);
 
         ufo_task_node_set_proc_node (UFO_TASK_NODE (node), proc_node);
     }
@@ -505,24 +504,21 @@ map_proc_node (UfoGraph *graph,
     g_list_free (successors);
 }
 
-
 /**
  * ufo_task_graph_map:
  * @task_graph: A #UfoTaskGraph
- * @arch_graph: A #UfoArchGraph to which @task_graph's nodes are mapped onto
+ * @gpu_nodes: List of #UfoGpuNode objects
  *
  * Map task nodes of @task_graph to the processing nodes of @arch_graph. Not
  * doing this could break execution of @task_graph.
  */
 void
 ufo_task_graph_map (UfoTaskGraph *task_graph,
-                    UfoArchGraph *arch_graph)
+                    GList *gpu_nodes)
 {
-    GList *gpu_nodes;
     GList *roots;
     GList *it;
 
-    gpu_nodes = ufo_arch_graph_get_gpu_nodes (arch_graph);
     roots = ufo_graph_get_roots (UFO_GRAPH (task_graph));
 
     g_list_for (roots, it) {
@@ -530,7 +526,6 @@ ufo_task_graph_map (UfoTaskGraph *task_graph,
     }
 
     g_list_free (roots);
-    g_list_free (gpu_nodes);
 }
 
 /**
@@ -877,6 +872,7 @@ ufo_task_graph_init (UfoTaskGraph *self)
 
     priv->manager = NULL;
     priv->remote_tasks = NULL;
+
     priv->json_nodes = g_hash_table_new_full (g_str_hash, g_str_equal,
                                               g_free, NULL);
 
diff --git a/ufo/ufo-task-graph.h b/ufo/ufo-task-graph.h
index 75cb1b3..eba7dce 100644
--- a/ufo/ufo-task-graph.h
+++ b/ufo/ufo-task-graph.h
@@ -86,10 +86,13 @@ void         ufo_task_graph_save_to_json        (UfoTaskGraph       *graph,
                                                  GError            **error);
 gchar       *ufo_task_graph_get_json_data       (UfoTaskGraph       *graph,
                                                  GError            **error);
+void         ufo_task_graph_set_gpu_nodes       (UfoTaskGraph       *task_graph,
+                                                 GList              *gpu_nodes);
 void         ufo_task_graph_map                 (UfoTaskGraph       *task_graph,
-                                                 UfoArchGraph       *arch_graph);
+                                                 GList              *gpu_nodes);
 void         ufo_task_graph_expand              (UfoTaskGraph       *task_graph,
                                                  UfoArchGraph       *arch_graph,
+                                                 guint               n_gpus,
                                                  gboolean            expand_remote);
 void         ufo_task_graph_connect_nodes       (UfoTaskGraph       *graph,
                                                  UfoTaskNode        *n1,