1 /* Copyright 2022 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_HLO_FUSION_STATS_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_HLO_FUSION_STATS_H_ 18 19 #include <cstdint> 20 #include <string> 21 22 #include "absl/strings/string_view.h" 23 #include "tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h" 24 #include "tensorflow/compiler/xla/service/hlo_module.h" 25 #include "tensorflow/compiler/xla/service/hlo_pass_interface.h" 26 #include "tensorflow/compiler/xla/statusor.h" 27 28 // Read-only pass logging statistics about HLO fusion ops in the module. Enabled 29 // at VLOG level 1 only. 30 namespace xla { 31 namespace gpu { 32 33 class HloOpcodeHistogram : public std::map<std::set<std::string>, int64_t> { 34 public: 35 std::string ToString(); 36 }; 37 38 class HloFusionStatsVisitor : public ConstDfsHloVisitorWithDefault { 39 public: 40 Status RunOnModule(HloModule* module); 41 42 std::string ToString(); 43 44 protected: 45 Status DefaultAction(const xla::HloInstruction* instr) final; 46 47 Status HandleFusion(const HloInstruction* fusion) override; 48 49 private: 50 int64_t num_fusions_ = 0; 51 int64_t num_loop_fusions_ = 0; 52 int64_t num_input_fusions_ = 0; 53 HloOpcodeHistogram loop_fusion_opcode_histogram_; 54 HloOpcodeHistogram input_fusion_opcode_histogram_; 55 }; 56 57 } // namespace gpu 58 } // namespace xla 59 60 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_HLO_FUSION_STATS_H_ 61