1 /* Copyright 2019 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_DUMP_H_
17 #define TENSORFLOW_COMPILER_XLA_SERVICE_DUMP_H_
18
19 #include "absl/strings/string_view.h"
20 #include "mlir/IR/Operation.h" // from @llvm-project
21 #include "tensorflow/compiler/xla/service/hlo_module.h"
22 #include "tensorflow/compiler/xla/status.h"
23 #include "tensorflow/compiler/xla/xla.pb.h"
24
25 // Consolidated utilities for logging information during compilation, usually
26 // based on the options specified in the DebugOptions proto.
27 //
28 // Most functions here take an HloModule and read the DebugOptions from the
29 // module's config.
30
31 namespace xla {
32
33 class BufferAssignment;
34 class HloExecutionProfile;
35 class HloSnapshot;
36
37 // Get a timestamp which we can use as a filename prefix specific to this
38 // module.
39 std::string TimestampFor(const HloModule& module);
40
41 // Create the filename we will use to dump in DumpToFileInDir.
42 std::string FilenameFor(const HloModule& module, absl::string_view prefix,
43 absl::string_view suffix);
44
45 // Writes the given string to a file in the xla_dump_to directory specified by
46 // module's DebugOptions.
47 //
48 // If module doesn't have an xla_dump_to directory, does nothing.
49 void DumpToFileInDir(const HloModule& module, absl::string_view file_prefix,
50 absl::string_view file_suffix, absl::string_view contents);
51 void DumpToFileInDir(const DebugOptions& debug_options,
52 absl::string_view filename, absl::string_view contents);
53
54 // Like DumpToFileInDir, except if module doesn't have an xla_dump_to directory
55 // specified, or if that directory is equal to "-", writes to stdout instead.
56 void DumpToFileInDirOrStdout(const HloModule& module,
57 absl::string_view file_prefix,
58 absl::string_view file_suffix,
59 absl::string_view contents);
60
61 // Like DumpToFileInDir, except if debug_options doesn't have an xla_dump_to
62 // directory specified, or if that directory is equal to "-", writes to stdout
63 // instead.
64 void DumpToFileInDirOrStdout(const DebugOptions& debug_options, int unique_id,
65 absl::string_view module_name,
66 absl::string_view file_prefix,
67 absl::string_view file_suffix,
68 absl::string_view contents);
69
70 // Writes the given op to a file in the xla_dump_to directory specified by
71 // module's DebugOptions. Sets the op's source locations to that file.
72 //
73 // If module doesn't have an xla_dump_to directory, does nothing.
74 void DumpToFileInDirOrStdout(const HloModule& module,
75 absl::string_view file_prefix,
76 mlir::Operation* op);
77
78 // Dumps the given protobuf to the given filename if dumping is enabled.
79 // Exactly where and in what formats it's dumped is determined by the debug
80 // options.
81 void DumpProtobufToFile(const tensorflow::protobuf::Message& proto,
82 const DebugOptions& debug_options,
83 absl::string_view filename);
84
85 // Similar to above, but the filename depends on module's information and the
86 // given name.
87 void DumpPerModuleProtobufToFile(const HloModule& module,
88 const tensorflow::protobuf::Message& proto,
89 const DebugOptions& debug_options,
90 absl::string_view name);
91
92 // Dumps the given HLO module if dumping is enabled for the module. Exactly
93 // where and in what formats it's dumped is determined by the module's config.
94 //
95 // If you pass an HloExecutionProfile, note that currently only DOT-based output
96 // formats (i.e. --xla_dump_as_{dot,html,url}) are able to incorporate it into
97 // their output. Other formats will just ignore the profile.
98 void DumpHloModuleIfEnabled(const HloModule& module, absl::string_view name);
99 void DumpHloModuleIfEnabled(const HloModule& module,
100 const BufferAssignment& buffer_assn,
101 absl::string_view name);
102 void DumpHloModuleIfEnabled(const HloModule& module,
103 const HloExecutionProfile& profile,
104 absl::string_view name);
105
106 // Dumps the given HLO module after running one HLO pass and before running
107 // another, if that's enabled. Returns the full file paths of all dumps of the
108 // module, or an empty vector if nothing was dumped.
109 std::vector<std::string> DumpHloModuleBetweenPassesIfEnabled(
110 absl::string_view pipeline_name, absl::string_view before_pass_name,
111 absl::string_view after_pass_name, const HloModule& module);
112
113 // Dumps the given HLO module during the given HLO pass, if that's enabled.
114 //
115 // "step" is a human-readable description of where we are in the middle of this
116 // pass. For example, "before-assigning-layouts".
117 void DumpHloModuleDuringPassIfEnabled(absl::string_view pass_name,
118 absl::string_view step,
119 const HloModule& module);
120
121 // Dumps the given HloSnapshot to the module's xla_dump_dir, if this is enabled.
122 //
123 // Prefer the first overload below, as this will give filenames that are
124 // consistent with the other methods here. The second overload (which doesn't
125 // take an HloModule) is useful in the cases when you're dumping an HloSnapshot
126 // and simply don't have an HloModule.
127 void DumpHloSnapshotIfEnabled(const HloModule& module,
128 const HloSnapshot& snapshot);
129 void DumpHloSnapshotIfEnabled(const HloSnapshot& snapshot,
130 const DebugOptions& opts);
131
132 void DumpHloModuleMetadataIfEnabled(const std::vector<HloModule*>& modules);
133
134 // Returns true if we should dump data for an HloModule. This is useful if you
135 // want to check if DumpToFileInDir{,OrStdout} will do anything before
136 // generating an expensive string.
137 bool DumpingEnabledForHloModule(absl::string_view hlo_module_name,
138 const DebugOptions& opts);
DumpingEnabledForHloModule(const HloModule & module)139 inline bool DumpingEnabledForHloModule(const HloModule& module) {
140 return DumpingEnabledForHloModule(module.name(),
141 module.config().debug_options());
142 }
143
144 // Returns true if DumpToFileInDirOrStdout and DumpHloModuleIfEnabled will write
145 // to stdout, rather than to a file on disk.
146 //
147 // This is useful if you want to do something different when writing to stdout.
148 // For example, maybe you have (almost-)duplicate data that you wouldn't mind
149 // writing to two files, but you don't want to print twice.
150 bool DumpingToStdout(const DebugOptions& opts);
151
152 } // namespace xla
153
154 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_DUMP_H_
155