TaskGraphRunner Group support

cc/raster/task_graph_work_queue.cc

Index: cc/raster/task_graph_work_queue.cc
diff --git a/cc/raster/task_graph_work_queue.cc b/cc/raster/task_graph_work_queue.cc
index 655a005b024804117fc9052a16ca20dd523868ca..b906c244658c82eed9a6e641c4b79ac0288826c1 100644
--- a/cc/raster/task_graph_work_queue.cc
+++ b/cc/raster/task_graph_work_queue.cc
@@ -5,11 +5,41 @@
 #include "cc/raster/task_graph_work_queue.h"
 
 #include <algorithm>
+#include <map>
 #include <utility>
 
 #include "base/trace_event/trace_event.h"
 
 namespace cc {
+namespace {
+
+bool CompareTaskPriority(const TaskGraphWorkQueue::PrioritizedTask& a,
+                         const TaskGraphWorkQueue::PrioritizedTask& b) {
+  // In this system, numerically lower priority is run first.
+  return a.priority > b.priority;
+}
+
+class CompareTaskNamespacePriority {
+ public:
+  explicit CompareTaskNamespacePriority(uint16_t category)
+      : category_(category) {}
+
+  bool operator()(const TaskGraphWorkQueue::TaskNamespace* a,
+                  const TaskGraphWorkQueue::TaskNamespace* b) {
+    DCHECK_GT(a->ready_to_run_tasks.count(category_), 0u);
+    DCHECK_GT(b->ready_to_run_tasks.count(category_), 0u);
+
+    // Compare based on task priority of the ready_to_run_tasks heap .front()
+    // will hold the max element of the heap, except after pop_heap, when max
+    // element is moved to .back().
+    return CompareTaskPriority(a->ready_to_run_tasks.at(category_).front(),
+                               b->ready_to_run_tasks.at(category_).front());

[reveman, 2015/12/04 20:55:51]

ready_to_run_tasks[category_] instead?

[ericrk, 2015/12/09 22:56:43]

Can't call operator[] on const maps (as it may need to create an entry).

+  }
+
+ private:
+  uint16_t category_;
+};
+}  // namespace
 
 TaskGraphWorkQueue::TaskNamespace::TaskNamespace() {}
 
@@ -66,14 +96,17 @@ void TaskGraphWorkQueue::ScheduleTasks(NamespaceToken token, TaskGraph* graph) {
                   node.task) != task_namespace.running_tasks.end())
       continue;
 
-    task_namespace.ready_to_run_tasks.push_back(
-        PrioritizedTask(node.task, &task_namespace, node.priority));
+    task_namespace.ready_to_run_tasks[node.category].push_back(PrioritizedTask(
+        node.task, &task_namespace, node.category, node.priority));
   }
 
-  // Rearrange the elements in |ready_to_run_tasks| in such a way that they
-  // form a heap.
-  std::make_heap(task_namespace.ready_to_run_tasks.begin(),
-                 task_namespace.ready_to_run_tasks.end(), CompareTaskPriority);
+  // Rearrange the elements in each vector within |ready_to_run_tasks| in such a
+  // way that they form a heap.
+  for (auto& it : task_namespace.ready_to_run_tasks) {
+    auto& ready_to_run_tasks = it.second;
+    std::make_heap(ready_to_run_tasks.begin(), ready_to_run_tasks.end(),
+                   CompareTaskPriority);
+  }
 
   // Swap task graph.
   task_namespace.graph.Swap(graph);
@@ -101,40 +134,61 @@ void TaskGraphWorkQueue::ScheduleTasks(NamespaceToken token, TaskGraph* graph) {
 
   // Build new "ready to run" task namespaces queue.
   ready_to_run_namespaces_.clear();
-  for (auto& it : namespaces_) {
-    if (!it.second.ready_to_run_tasks.empty())
-      ready_to_run_namespaces_.push_back(&it.second);
+  for (auto& namespace_it : namespaces_) {
+    auto& task_namespace = namespace_it.second;
+    for (auto& tasks_it : task_namespace.ready_to_run_tasks) {

[reveman, 2015/12/04 20:55:51]

nit: ready_to_run_tasks_it to be perfectly consistent

[ericrk, 2015/12/09 22:56:43]

Done.

+      uint16_t category = tasks_it.first;
+      ready_to_run_namespaces_[category].push_back(&task_namespace);
+    }
   }
 
   // Rearrange the task namespaces in |ready_to_run_namespaces| in such a
   // way that they form a heap.
-  std::make_heap(ready_to_run_namespaces_.begin(),
-                 ready_to_run_namespaces_.end(), CompareTaskNamespacePriority);
+  for (auto& it : ready_to_run_namespaces_) {
+    uint16_t category = it.first;
+    auto& task_namespace = it.second;
+    std::make_heap(task_namespace.begin(), task_namespace.end(),
+                   CompareTaskNamespacePriority(category));
+  }
 }
 
-TaskGraphWorkQueue::PrioritizedTask TaskGraphWorkQueue::GetNextTaskToRun() {
-  DCHECK(!ready_to_run_namespaces_.empty());
+TaskGraphWorkQueue::PrioritizedTask TaskGraphWorkQueue::GetNextTaskToRun(
+    uint16_t category) {
+  DCHECK_GT(ready_to_run_namespaces_.count(category), 0u);

[reveman, 2015/12/04 20:55:51]

nit: this failing isn't really a problem but
ready_to_run_namespaces_[category].empty() being true is. should we check that
instead?

[ericrk, 2015/12/09 22:56:43]

yup = this made more sense when I was erasing things, updated.

+  TaskNamespace::Vector& ready_to_run_namespaces =
+      ready_to_run_namespaces_[category];
 
-  // Take top priority TaskNamespace from |ready_to_run_namespaces_|.
-  std::pop_heap(ready_to_run_namespaces_.begin(),
-                ready_to_run_namespaces_.end(), CompareTaskNamespacePriority);
-  TaskNamespace* task_namespace = ready_to_run_namespaces_.back();
-  ready_to_run_namespaces_.pop_back();
-  DCHECK(!task_namespace->ready_to_run_tasks.empty());
+  // Take top priority TaskNamespace from |ready_to_run_namespaces|.
+  std::pop_heap(ready_to_run_namespaces.begin(), ready_to_run_namespaces.end(),
+                CompareTaskNamespacePriority(category));
+  TaskNamespace* task_namespace = ready_to_run_namespaces.back();
+  ready_to_run_namespaces.pop_back();
+
+  DCHECK_GT(task_namespace->ready_to_run_tasks.count(category), 0u);

[reveman, 2015/12/04 20:55:51]

same here, DCHECK(!empty()) instead?

[ericrk, 2015/12/09 22:56:43]

done.

+  PrioritizedTask::Vector& ready_to_run_tasks =
+      task_namespace->ready_to_run_tasks[category];
 
   // Take top priority task from |ready_to_run_tasks|.
-  std::pop_heap(task_namespace->ready_to_run_tasks.begin(),
-                task_namespace->ready_to_run_tasks.end(), CompareTaskPriority);
-  PrioritizedTask task = task_namespace->ready_to_run_tasks.back();
-  task_namespace->ready_to_run_tasks.pop_back();
+  std::pop_heap(ready_to_run_tasks.begin(), ready_to_run_tasks.end(),
+                CompareTaskPriority);
+  PrioritizedTask task = ready_to_run_tasks.back();
+  ready_to_run_tasks.pop_back();
 
-  // Add task namespace back to |ready_to_run_namespaces_| if not empty after
+  // Add task namespace back to |ready_to_run_namespaces| if not empty after
   // taking top priority task.
-  if (!task_namespace->ready_to_run_tasks.empty()) {
-    ready_to_run_namespaces_.push_back(task_namespace);
-    std::push_heap(ready_to_run_namespaces_.begin(),
-                   ready_to_run_namespaces_.end(),
-                   CompareTaskNamespacePriority);
+  if (!ready_to_run_tasks.empty()) {
+    ready_to_run_namespaces.push_back(task_namespace);
+    std::push_heap(ready_to_run_namespaces.begin(),
+                   ready_to_run_namespaces.end(),
+                   CompareTaskNamespacePriority(category));
+  } else {
+    // Task namespace is empty for this category, remove the entry from the map.
+    task_namespace->ready_to_run_tasks.erase(category);
+  }
+
+  // If |ready_to_run_namespaces| is now empty, remove it from the map.
+  if (ready_to_run_namespaces.empty()) {
+    ready_to_run_namespaces_.erase(category);

[reveman, 2015/12/04 20:55:51]

Will this and the above erase not cause a large amount of heap allocations in
the common use case? existing code has been tuned to not require any heap
allocation for the common case where we schedule a task graph that is the same
as the previous one and it might be best to keep it that way. what if we let the
map grow as needed and instead use the !vector.empty() to determine if category
has ready to run tasks? seems like that would be fine as we expect a limited
number of categories to be used..

Either way, can you post the results from running task graph perf tests in a
release build before and after this change to this code review? that would
useful to know independent of how we solve this.

[ericrk, 2015/12/09 22:56:43]

Got rid of the erasing - took some benchmarks and it was a 50% regression in
many cases. - without erasing we're at a 30% regression max - will add some
numbers in the reply.

   }
 
   // Add task to |running_tasks|.
@@ -164,19 +218,26 @@ void TaskGraphWorkQueue::CompleteTask(const PrioritizedTask& completed_task) {
     dependent_node.dependencies--;
     // Task is ready if it has no dependencies. Add it to |ready_to_run_tasks_|.
     if (!dependent_node.dependencies) {
-      bool was_empty = task_namespace->ready_to_run_tasks.empty();
-      task_namespace->ready_to_run_tasks.push_back(PrioritizedTask(
-          dependent_node.task, task_namespace, dependent_node.priority));
-      std::push_heap(task_namespace->ready_to_run_tasks.begin(),
-                     task_namespace->ready_to_run_tasks.end(),
+      PrioritizedTask::Vector& ready_to_run_tasks =
+          task_namespace->ready_to_run_tasks[dependent_node.category];
+
+      bool was_empty = ready_to_run_tasks.empty();
+      ready_to_run_tasks.push_back(
+          PrioritizedTask(dependent_node.task, task_namespace,
+                          dependent_node.category, dependent_node.priority));
+      std::push_heap(ready_to_run_tasks.begin(), ready_to_run_tasks.end(),
                      CompareTaskPriority);
+
       // Task namespace is ready if it has at least one ready to run task. Add
       // it to |ready_to_run_namespaces_| if it just become ready.
       if (was_empty) {
-        DCHECK(std::find(ready_to_run_namespaces_.begin(),
-                         ready_to_run_namespaces_.end(),
-                         task_namespace) == ready_to_run_namespaces_.end());
-        ready_to_run_namespaces_.push_back(task_namespace);
+        TaskNamespace::Vector& ready_to_run_namespaces =
+            ready_to_run_namespaces_[dependent_node.category];
+
+        DCHECK(std::find(ready_to_run_namespaces.begin(),
+                         ready_to_run_namespaces.end(),
+                         task_namespace) == ready_to_run_namespaces.end());
+        ready_to_run_namespaces.push_back(task_namespace);
       }
       ready_to_run_namespaces_has_heap_properties = false;
     }
@@ -185,9 +246,13 @@ void TaskGraphWorkQueue::CompleteTask(const PrioritizedTask& completed_task) {
   // Rearrange the task namespaces in |ready_to_run_namespaces_| in such a way
   // that they yet again form a heap.
   if (!ready_to_run_namespaces_has_heap_properties) {
-    std::make_heap(ready_to_run_namespaces_.begin(),
-                   ready_to_run_namespaces_.end(),
-                   CompareTaskNamespacePriority);
+    for (auto& it : ready_to_run_namespaces_) {
+      uint16_t category = it.first;
+      auto& ready_to_run_namespaces = it.second;
+      std::make_heap(ready_to_run_namespaces.begin(),
+                     ready_to_run_namespaces.end(),
+                     CompareTaskNamespacePriority(category));
+    }
   }
 
   // Finally add task to |completed_tasks_|.
@@ -209,7 +274,7 @@ void TaskGraphWorkQueue::CollectCompletedTasks(NamespaceToken token,
 
   // Remove namespace if finished running tasks.
   DCHECK_EQ(0u, task_namespace.completed_tasks.size());
-  DCHECK_EQ(0u, task_namespace.ready_to_run_tasks.size());
+  DCHECK(task_namespace.ready_to_run_tasks.empty());
   DCHECK_EQ(0u, task_namespace.running_tasks.size());
   namespaces_.erase(it);
 }
@@ -228,4 +293,10 @@ bool TaskGraphWorkQueue::DependencyMismatch(const TaskGraph* graph) {
   return false;
 }
 
+void TaskGraphWorkQueue::UncategorizeTaskGraph(TaskGraph* graph) {
+  for (auto& node : graph->nodes) {
+    node.category = 0u;
+  }
+}
+
 }  // namespace cc