1 /* Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_GPU_THUNK_SCHEDULE_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_THUNK_SCHEDULE_H_ 18 19 #include <list> 20 #include <memory> 21 #include <unordered_map> 22 #include <vector> 23 24 #include "tensorflow/compiler/xla/service/gpu/stream_assignment.h" 25 #include "tensorflow/compiler/xla/service/gpu/thunk.h" 26 #include "tensorflow/compiler/xla/service/hlo_instruction.h" 27 #include "tensorflow/compiler/xla/types.h" 28 29 namespace xla { 30 namespace gpu { 31 32 // Encapsulates in which order and on which streams the thunks are executed. A 33 // schedule contains 34 // 35 // 1. A stream assignment indicating which stream each thunk is executed on. 36 // 37 // 2. A total order of all thunks. If A is ordered before B and they are 38 // assigned to the same stream, then A completes before B starts. If A is 39 // ordered before B and they are on different streams, their actual execution 40 // order is not determined. 41 // 42 // 3. A set of dependency edges. If A and B are scheduled on different streams 43 // and A has to complete before B starts (e.g. A produces an input of B), then B 44 // "depends" on A. 45 class ThunkSchedule { 46 public: 47 ThunkSchedule(std::unique_ptr<ThunkSequence> thunks, 48 std::unique_ptr<StreamAssignment> stream_assignment, 49 const std::vector<const HloInstruction*>& hlo_total_order); 50 51 // Returns the total order of executing all the thunks. 52 const std::vector<Thunk*>& TotalOrder() const { return thunk_total_order_; } 53 54 // Thunks that `thunk` depends on. 55 const std::list<const Thunk*>& DependsOn(const Thunk* thunk) const; 56 // Whether `thunk` is depended by another thunk. 57 bool Depended(const Thunk* thunk) const { return depended_by_.count(thunk); } 58 59 // Delegates to StreamAssignment. 60 int StreamCount() const { return stream_assignment_->StreamCount(); } 61 int StreamNumberForHlo(const HloInstruction& hlo) const { 62 return stream_assignment_->StreamNumberForHlo(hlo); 63 } 64 65 string ToString() const; 66 67 private: 68 void RemoveRedundantDependencyEdges(); 69 70 // Adds `operand` and its transitive operands to the dependency list of 71 // `thunk`. 72 // 73 // Precondition: `operand` is a non-trivial (i.e. excluding 74 // thunk.hlo_instruction() itself) transitive operand of 75 // thunk.hlo_instruction(). 76 void AddDependenciesOnTransitiveOperands( 77 const Thunk& thunk, const HloInstruction& operand, 78 const std::unordered_map<const HloInstruction*, Thunk*>& hlo_to_thunk); 79 80 std::unique_ptr<ThunkSequence> thunks_; 81 std::vector<Thunk*> thunk_total_order_; 82 83 std::unordered_map<const Thunk*, std::list<const Thunk*>> depends_on_; 84 std::set<const Thunk*> depended_by_; 85 std::list<const Thunk*> empty_thunk_list_; 86 87 std::unique_ptr<StreamAssignment> stream_assignment_; 88 }; 89 90 } // namespace gpu 91 } // namespace xla 92 93 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_THUNK_SCHEDULE_H_ 94
