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 #ifndef TENSORFLOW_COMPILER_MLIR_TENSORFLOW_UTILS_DUMP_GRAPH_H_ 16 #define TENSORFLOW_COMPILER_MLIR_TENSORFLOW_UTILS_DUMP_GRAPH_H_ 17 18 #include <string> 19 20 #include "mlir/IR/OperationSupport.h" // from @llvm-project 21 #include "tensorflow/core/framework/function.h" 22 #include "tensorflow/core/graph/graph.h" 23 #include "tensorflow/core/platform/status.h" 24 25 namespace tensorflow { 26 27 struct MlirDumpConfig; 28 29 // Dumps 'graph_def' to a file, as textual IR. Returns the file name chosen. 30 // 31 // Note: This is for debugging use and is not optimized for performance. 32 Status DumpTextualIRToFile(const MlirDumpConfig& config, const Graph& graph, 33 const FunctionLibraryDefinition* flib_def, 34 WritableFile* file); 35 36 // Config of the textual dump. 37 struct MlirDumpConfig { 38 enum class Dialect { 39 // Tensorflow Graph Dialect 40 kTFG, 41 }; 42 43 // The limit of element size that gets printed. 44 MlirDumpConfig& elide_large_attributes(int large_element_limit = 16) { 45 this->op_printing_flags.elideLargeElementsAttrs(large_element_limit); 46 return *this; 47 } 48 49 // Enable printing of debug information. If 'pretty_form' is set to true, 50 // debug information is printed in a more readable 'pretty' form but this 51 // pretty form is not parsable (so only for human readability). 52 MlirDumpConfig& emit_location_information(bool pretty_form = false) { 53 this->op_printing_flags.enableDebugInfo(pretty_form); 54 return *this; 55 } 56 emit_dialectMlirDumpConfig57 MlirDumpConfig& emit_dialect(Dialect dialect) { 58 this->dialect = dialect; 59 return *this; 60 } 61 62 // Op printing flags. 63 mlir::OpPrintingFlags op_printing_flags = llvm::None; 64 65 // The target MLIR dialect. 66 Dialect dialect = Dialect::kTFG; 67 }; 68 69 // Change DumpGraphToFile to dump MLIR textual IR instead of protobuf. 70 void UseMlirForGraphDump(const MlirDumpConfig& = {}); 71 72 } // namespace tensorflow 73 74 #endif // TENSORFLOW_COMPILER_MLIR_TENSORFLOW_UTILS_DUMP_GRAPH_H_ 75