1 /* Copyright 2018 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_HLO_MODULE_GROUP_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_MODULE_GROUP_H_ 18 19 #include <iosfwd> 20 #include <string> 21 #include <vector> 22 23 #include "absl/strings/string_view.h" 24 #include "absl/types/span.h" 25 #include "tensorflow/compiler/xla/service/hlo.pb.h" 26 #include "tensorflow/compiler/xla/service/hlo_module.h" 27 28 namespace xla { 29 30 // An abstraction representing a ordered set of HLO module built to run 31 // concurrently across different devices. 32 class HloModuleGroup { 33 public: 34 // Construct an empty module group. HloModuleGroup(absl::string_view name)35 explicit HloModuleGroup(absl::string_view name) : name_(name) {} 36 37 // Construct a module group containing a single module. 38 explicit HloModuleGroup(std::unique_ptr<HloModule> module); 39 40 // Construct a module group containing any number of modules. 41 HloModuleGroup(absl::string_view name, 42 absl::Span<std::unique_ptr<HloModule>> modules); 43 HloModuleGroup(absl::string_view name, 44 std::vector<std::unique_ptr<HloModule>>&& modules); 45 46 // Returns the modules contained in the group. modules()47 const std::vector<HloModule*>& modules() const { return module_ptrs_; } 48 49 // Returns a module at a particular index. module(int index)50 HloModule& module(int index) const { return *module_ptrs_.at(index); } 51 52 // Add a module to the back of vector of modules in the group. 53 void push_back(std::unique_ptr<HloModule> module); 54 55 // Replaces the existing module at the given index with the given module. The 56 // existing module is discarded. 57 void ReplaceModule(int index, std::unique_ptr<HloModule> module); 58 59 // Moves all modules from the group into the returned vector. After this 60 // method runs, the module group will be empty. 61 std::vector<std::unique_ptr<HloModule>> ConsumeModules(); 62 name()63 string name() const { return name_; } 64 65 string ToString() const; 66 67 // Serialize the module group to/from a proto. 68 HloModuleGroupProto ToProto() const; 69 static StatusOr<HloModuleGroup> CreateFromProto( 70 const HloModuleGroupProto& proto, 71 absl::Span<const HloModuleConfig> module_configs); 72 73 // Returns the number of modules in the module group. size()74 int size() const { return modules_.size(); } 75 76 // Returns true if there are no modules in the module group. empty()77 bool empty() const { return modules_.empty(); } 78 79 private: 80 string name_; 81 82 // Vector of modules as std::unique_ptrs. 83 std::vector<std::unique_ptr<HloModule>> modules_; 84 85 // Vector of modules as normal pointers. This vector is kept in sync with 86 // modules_ as modules are added to the group with push_back. 87 std::vector<HloModule*> module_ptrs_; 88 }; 89 90 std::ostream& operator<<(std::ostream& out, const HloModuleGroup& group); 91 92 } // namespace xla 93 94 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_MODULE_GROUP_H_ 95