1 files changed, 93 insertions, 0 deletions
diff --git a/tensorflow/compiler/xla/service/gpu/thunk_schedule.h b/tensorflow/compiler/xla/service/gpu/thunk_schedule.h
new file mode 100644
index 0000000000..d3352994f8
--- /dev/null
+++ b/tensorflow/compiler/xla/service/gpu/thunk_schedule.h
@@ -0,0 +1,93 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_GPU_THUNK_SCHEDULE_H_
+#define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_THUNK_SCHEDULE_H_
+
+#include <list>
+#include <memory>
+#include <unordered_map>
+#include <vector>
+
+#include "tensorflow/compiler/xla/service/gpu/stream_assignment.h"
+#include "tensorflow/compiler/xla/service/gpu/thunk.h"
+#include "tensorflow/compiler/xla/service/hlo_instruction.h"
+#include "tensorflow/compiler/xla/types.h"
+
+namespace xla {
+namespace gpu {
+
+// Encapsulates in which order and on which streams the thunks are executed. A
+// schedule contains
+//
+// 1. A stream assignment indicating which stream each thunk is executed on.
+//
+// 2. A total order of all thunks. If A is ordered before B and they are
+// assigned to the same stream, then A completes before B starts. If A is
+// ordered before B and they are on different streams, their actual execution
+// order is not determined.
+//
+// 3. A set of dependency edges. If A and B are scheduled on different streams
+// and A has to complete before B starts (e.g. A produces an input of B), then B
+// "depends" on A.
+class ThunkSchedule {
+ public:
+  ThunkSchedule(std::unique_ptr<ThunkSequence> thunks,
+                std::unique_ptr<StreamAssignment> stream_assignment,
+                const std::vector<const HloInstruction*>& hlo_total_order);
+
+  // Returns the total order of executing all the thunks.
+  const std::vector<Thunk*>& TotalOrder() const { return thunk_total_order_; }
+
+  // Thunks that `thunk` depends on.
+  const std::list<const Thunk*>& DependsOn(const Thunk* thunk) const;
+  // Whether `thunk` is depended by another thunk.
+  bool Depended(const Thunk* thunk) const { return depended_by_.count(thunk); }
+
+  // Delegates to StreamAssignment.
+  int StreamCount() const { return stream_assignment_->StreamCount(); }
+  int StreamNumberForHlo(const HloInstruction& hlo) const {
+    return stream_assignment_->StreamNumberForHlo(hlo);
+  }
+
+  string ToString() const;
+
+ private:
+  void RemoveRedundantDependencyEdges();
+
+  // Adds `operand` and its transitive operands to the dependency list of
+  // `thunk`.
+  //
+  // Precondition: `operand` is a non-trivial (i.e. excluding
+  // thunk.hlo_instruction() itself) transitive operand of
+  // thunk.hlo_instruction().
+  void AddDependenciesOnTransitiveOperands(
+      const Thunk& thunk, const HloInstruction& operand,
+      const std::unordered_map<const HloInstruction*, Thunk*>& hlo_to_thunk);
+
+  std::unique_ptr<ThunkSequence> thunks_;
+  std::vector<Thunk*> thunk_total_order_;
+
+  std::unordered_map<const Thunk*, std::list<const Thunk*>> depends_on_;
+  std::set<const Thunk*> depended_by_;
+  std::list<const Thunk*> empty_thunk_list_;
+
+  std::unique_ptr<StreamAssignment> stream_assignment_;
+};
+
+}  // namespace gpu
+}  // namespace xla
+
+#endif  // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_THUNK_SCHEDULE_H_