xref: /aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.h (revision b6fb3261f9314811a0f4371741dbb8839866f948)
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_MLIR_TENSORFLOW_UTILS_DUMP_MLIR_UTIL_H_
17 #define TENSORFLOW_COMPILER_MLIR_TENSORFLOW_UTILS_DUMP_MLIR_UTIL_H_
18 
19 #include <string>
20 
21 #include "absl/strings/string_view.h"
22 #include "llvm/ADT/StringRef.h"
23 #include "mlir/IR/Operation.h"  // from @llvm-project
24 #include "mlir/Pass/PassManager.h"  // from @llvm-project
25 #include "tensorflow/core/platform/status.h"
26 
27 namespace tensorflow {
28 
29 inline constexpr absl::string_view kCrashReproducerStdErr = "-";
30 inline constexpr absl::string_view kCrashReproducerCrashAnalysis =
31     "crash_analysis";
32 
33 // Creates a file to use for dumping and returns success if a file could be
34 // created. The opened file is placed in 'os' and the path of the file used is
35 // placed in 'filepath'.
36 //
37 // If the TF_DUMP_GRAPH_PREFIX environment variable is kCrashReproducerStdErr,
38 // then the LOG(INFO) macro is used instead.
39 //
40 // This will create a file name via prefixing `name` with the value of the
41 // TF_DUMP_GRAPH_PREFIX environment variable if `dirname` is empty and
42 // suffixing `name` with ".mlir".
43 Status CreateFileForDumping(llvm::StringRef name,
44                             std::unique_ptr<llvm::raw_ostream>* os,
45                             std::string* filepath,
46                             llvm::StringRef dirname = "");
47 
48 // Dumps the configuration of the pass pipeline and the MLIR module to a file
49 // and returns the file name used. The file will be in the same format of an
50 // MLIR crash reproducer.
51 //
52 // If the TF_DUMP_GRAPH_PREFIX environment variable is kCrashReproducerStdErr,
53 // then the MLIR operation will be logged (using the LOG(INFO) macro) instead.
54 //
55 // This will create a file name via prefixing `name` with the value of the
56 // TF_DUMP_GRAPH_PREFIX environment variable if `dirname` is empty and
57 // suffixing `name` with ".mlir".
58 
59 std::string DumpCrashReproducerToFile(llvm::StringRef name,
60                                       const mlir::PassManager& pm,
61                                       mlir::Operation* op,
62                                       llvm::StringRef dirname = "");
63 
64 // Dumps MLIR operation to a file and returns the file name used.
65 //
66 // If the TF_DUMP_GRAPH_PREFIX environment variable is kCrashReproducerStdErr,
67 // then the MLIR operation will be logged (using the LOG(INFO) macro) instead.
68 //
69 // This will create a file name via prefixing `name` with the value of the
70 // TF_DUMP_GRAPH_PREFIX environment variable if `dirname` is empty and
71 // suffixing `name` with ".mlir".
72 // If `pass_manager` is provided, prints a header with the pass pipeline.
73 std::string DumpMlirOpToFile(llvm::StringRef name, mlir::Operation* op,
74                              llvm::StringRef dirname = "",
75                              const mlir::PassManager* pass_manager = nullptr);
76 
77 // Reads the directory to dump the MLIR module from environment variables.
78 // Default is reading from TF_DUMP_GRAPH_PREFIX, and if the string is 'sponge'
79 // read from TEST_UNDECLARED_OUTPUTS_DIR. Returns nullptr if the directory
80 // cannot be determined and generates a warning message.
81 std::string GetDumpDirFromEnvVar();
82 
83 // Dumps a raw string to a file and returns the file name used.
84 //
85 // This will create a file name via prefixing `name` with the value of the
86 // TF_DUMP_GRAPH_PREFIX environment variable if `dirname` is empty and
87 // suffixing `name` with ".mlir".
88 std::string DumpRawStringToFile(llvm::StringRef name, llvm::StringRef content,
89                                 llvm::StringRef dirname = "");
90 
91 // Enable the crash reproducer on the provided PassManager to the provided
92 // directory path.
93 // If the provided path is empty, it is retrieved from the
94 // environment variable `MLIR_CRASH_REPRODUCER_DIRECTORY`.
95 // If the provided path is the string "sponge", the file will be included
96 // in the sponge "Output Files" by looking up the environment to infer
97 // the directory path.
98 // If the provided path is the string kCrashReproducerStdErr, the data is
99 // dumped into the stderr.
100 // If the provided path is the string kCrashReproducerCrashAnalysis, the data
101 // is dumped to the crash analysis system. Note, environment var
102 // `MLIR_CRASH_REPRODUCER_DIRECTORY` can be used to override
103 // kCrashReproducerCrashAnalysis settings.
104 void SetCrashReproducer(mlir::PassManager& pm, llvm::StringRef dir_path = "");
105 
106 // This applies both the PassManagerCLOptions provided by MLIR along with any
107 // tensorflow specific options.
108 //
109 // Note that this function should be in a more appropriate file, but it is
110 // unclear what a proper file would be as no other functions would currently be
111 // in the file also.
112 void applyTensorflowAndCLOptions(mlir::PassManager& pm,
113                                  llvm::StringRef dir_path = "");
114 
115 }  // namespace tensorflow
116 
117 #endif  // TENSORFLOW_COMPILER_MLIR_TENSORFLOW_UTILS_DUMP_MLIR_UTIL_H_
118