1 /* Copyright 2021 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_ASYNC_COLLECTIVE_CREATOR_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_ASYNC_COLLECTIVE_CREATOR_H_ 18 19 #include <functional> 20 21 #include "tensorflow/compiler/xla/service/hlo_pass_interface.h" 22 23 namespace xla { 24 25 // Transforms each all-reduce instruction to a pair of all-reduce-start and 26 // all-reduce-done. 27 class AsyncCollectiveCreator : public HloModulePass { 28 public: 29 struct CollectiveCreatorConfig { 30 HloPredicate convert_all_reduce = [](const HloInstruction*) { 31 return false; 32 }; 33 HloPredicate convert_all_gather = [](const HloInstruction*) { 34 return false; 35 }; 36 HloPredicate convert_collective_permute = [](const HloInstruction*) { 37 return false; 38 }; 39 HloPredicate convert_all_to_all = [](const HloInstruction*) { 40 return false; 41 }; 42 }; AsyncCollectiveCreator(CollectiveCreatorConfig creator_config)43 explicit AsyncCollectiveCreator(CollectiveCreatorConfig creator_config) 44 : convert_all_reduce_(creator_config.convert_all_reduce), 45 convert_all_gather_(creator_config.convert_all_gather), 46 convert_collective_permute_(creator_config.convert_collective_permute), 47 convert_all_to_all_(creator_config.convert_all_to_all) {} name()48 absl::string_view name() const override { return "async-collective-creator"; } 49 50 using HloPassInterface::Run; 51 StatusOr<bool> Run( 52 HloModule* module, 53 const absl::flat_hash_set<absl::string_view>& execution_threads) override; 54 55 private: 56 HloPredicate convert_all_reduce_; 57 HloPredicate convert_all_gather_; 58 HloPredicate convert_collective_permute_; 59 HloPredicate convert_all_to_all_; 60 }; 61 62 } // namespace xla 63 64 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_ASYNC_ALL_REDUCE_CREATOR_H_ 65