1 /* Copyright 2016 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_FUSION_MERGER_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_FUSION_MERGER_H_ 18 19 #include "tensorflow/compiler/xla/service/hlo_module.h" 20 #include "tensorflow/compiler/xla/service/hlo_pass_interface.h" 21 22 namespace xla { 23 namespace gpu { 24 25 // An HLO pass that attempts to merge fusion instructions to reduce memory 26 // bandwidth requirements and kernel launch overhead. 27 // 28 // Consider the example below. On the left-hand side, op A is the producer and 29 // ops B and C are its consumers. FusionMerger duplicates producer ops and fuses 30 // them into all consumers. The result is depicted on the right-hand side below. 31 // 32 // p p 33 // | / \ 34 // v / \ 35 // A +fusion+ +fusion+ 36 // / \ | A' | | A" | 37 // | | | | | | | | 38 // v v | v | | v | 39 // B C | B | | C | 40 // +------+ +------+ 41 // 42 // Op A has been cloned twice and fused with B and C. The kernel launch overhead 43 // is reduced from 3 to 2. The memory bandwidth requirements may be reduced. 44 // We trade 1 read of input(A) + 1 write and 2 reads of output(A) for 2 reads of 45 // input(A). In general the achieveable savings in memory bandwidth depend on 46 // the differences in memory read and written and the number of consumers. The 47 // FusionMeger pass takes this into account when making fusion decisions. 48 // 49 // The pass traverses the HLO module in reverse post-order (defs before uses). 50 // Fusion instructions are merged into their users if some conditions are met: 51 // * The result of merging the fusion instruction into its users would not 52 // increase bytes transferred. 53 // * Producer ops are fusible with _all_ consumers. If they are not fusible with 54 // at least one consumers, they won't be fused at all. 55 // * Producers are kLoop fusion ops. 56 // 57 // None of these restrictions are necessary for correctness. In fact, lifting 58 // the latter two could be beneficial. 59 60 class FusionMerger : public HloModulePass { 61 public: name()62 absl::string_view name() const override { return "fusion_merger"; } 63 64 StatusOr<bool> Run(HloModule* module) override; 65 }; 66 67 } // namespace gpu 68 } // namespace xla 69 70 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_FUSION_MERGER_H_ 71