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