1 /* Copyright 2020 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_HORIZONTAL_INPUT_FUSION_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_HORIZONTAL_INPUT_FUSION_H_ 18 19 #include "tensorflow/compiler/xla/service/hlo_computation.h" 20 #include "tensorflow/compiler/xla/service/hlo_instruction.h" 21 #include "tensorflow/compiler/xla/service/hlo_module.h" 22 #include "tensorflow/compiler/xla/service/hlo_pass_interface.h" 23 24 namespace xla { 25 namespace gpu { 26 27 // This optimization pass horizontally fuses kInput fusions to both reduce the 28 // kernel launch overhead and increase parallelism degree. See 29 // GpuHorizontalFusion for general description and motivation about horizontal 30 // fusion. GpuHorizontalFusion deals with kLoop fusions while this pass deals 31 // with kInput fusions. 32 // 33 // Following GpuHorizontalFusion, a simple yet effective heuristic is used 34 // to search the fusion candidates while avoiding creating cycles. That is, 35 // we simply search for fusion candidates by looking for instructions whose 36 // outputs are all consumed by the same instruction. This catches the typical 37 // target cases; often, the candidate instructions are just consumed by the 38 // ROOT tuple of the entry computation. 39 class GpuHorizontalInputFusion : public HloModulePass { 40 public: GpuHorizontalInputFusion()41 GpuHorizontalInputFusion() {} 42 name()43 absl::string_view name() const override { 44 return "gpu_horizontal_input_fusion"; 45 } 46 47 using HloPassInterface::Run; 48 StatusOr<bool> Run( 49 HloModule* module, 50 const absl::flat_hash_set<absl::string_view>& execution_threads) override; 51 52 private: 53 StatusOr<bool> RunOnComputation(HloComputation*); 54 }; 55 56 } // namespace gpu 57 } // namespace xla 58 59 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_HORIZONTAL_INPUT_FUSION_H_ 60