experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Move per pass logging to vlog 2

Browse Source 
Have before and after logging at vlog 1 and per pass at 2. Also aimed to have the file name closer to the function name.

PiperOrigin-RevId: 357215310
Change-Id: I7e5a3d2ecb95b8c8f8fc400ccee58c36ea88d739
This commit is contained in:
Jacques Pienaar 2021-02-12 10:11:21 -08:00 committed by TensorFlower Gardener
parent ed1e414aac
commit 164f5b2a5b
2 changed files with 26 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
tensorflow/compiler/mlir/tensorflow/transforms/bridge.cc
View File

@ -59,7 +59,10 @@ tensorflow::Status RunTPUBridge(
llvm::function_ref<void(OpPassManager &pm)> pipeline_builder) {
PassManager bridge(module.getContext());
::tensorflow::applyTensorflowAndCLOptions(bridge);
if (enable_logging) EnableLogging(&bridge);
if (enable_logging || VLOG_IS_ON(1)) {
tensorflow::DumpMlirOpToFile("tpu_bridge_before", module);
if (VLOG_IS_ON(2)) EnableLogging(&bridge);
}
// Populate a passmanager with the list of passes that implement the bridge.
pipeline_builder(bridge);
@ -72,6 +75,8 @@ tensorflow::Status RunTPUBridge(
mlir::StatusScopedDiagnosticHandler diag_handler(module.getContext());
LogicalResult result = bridge.run(module);
(void)result;
if (enable_logging || VLOG_IS_ON(1))
tensorflow::DumpMlirOpToFile("tpu_bridge_after", module);
return diag_handler.ConsumeStatus();
}
} // namespace
@ -163,7 +168,10 @@ tensorflow::Status RunBridgeWithStandardPipeline(ModuleOp module,
bool enable_logging,
bool enable_inliner) {
PassManager bridge(module.getContext());
if (enable_logging) EnableLogging(&bridge);
if (enable_logging || VLOG_IS_ON(1)) {
tensorflow::DumpMlirOpToFile("standard_pipeline_before", module);
if (VLOG_IS_ON(2)) EnableLogging(&bridge);
}
StandardPipelineOptions pipeline_options;
pipeline_options.enable_inliner.setValue(enable_inliner);
@ -171,6 +179,8 @@ tensorflow::Status RunBridgeWithStandardPipeline(ModuleOp module,
mlir::StatusScopedDiagnosticHandler diag_handler(module.getContext());
LogicalResult result = bridge.run(module);
(void)result;
if (enable_logging || VLOG_IS_ON(1))
tensorflow::DumpMlirOpToFile("standard_pipeline_after", module);
return diag_handler.ConsumeStatus();
}

22
tensorflow/compiler/mlir/tensorflow/utils/compile_mlir_util.cc
View File

@ -345,7 +345,9 @@ Status LegalizeToHlo(mlir::ModuleOp module_op, llvm::StringRef device_type,
CreateConvertMlirToXlaHloPipeline(tf2xla, device_type,
custom_legalization_passes);
if (VLOG_IS_ON(1)) {
if (VLOG_IS_ON(1))
tensorflow::DumpMlirOpToFile("legalize_hlo_before", module_op);
if (VLOG_IS_ON(2)) {
// Print the whole module after each pass which requires disabling
// multi-threading as well.
module_op.getContext()->disableMultithreading();
@ -364,7 +366,7 @@ Status LegalizeToHlo(mlir::ModuleOp module_op, llvm::StringRef device_type,
}
if (VLOG_IS_ON(1))
tensorflow::DumpMlirOpToFile("mlir_compile_legalize_hlo", module_op);
tensorflow::DumpMlirOpToFile("legalize_hlo_after", module_op);
return Status::OK();
}
@ -406,8 +408,8 @@ Status CompileMlirSetup(
// Use arg_shapes to improve the mlir type information of `main` in module_op.
TF_RETURN_IF_ERROR(RefineShapes(arg_shapes, module_op));
if (VLOG_IS_ON(1))
tensorflow::DumpMlirOpToFile("mlir_compile_shape_refiner", module_op);
if (VLOG_IS_ON(2))
tensorflow::DumpMlirOpToFile("compile_mlir_shape_refiner", module_op);
if (!*shape_representation_fn)
*shape_representation_fn = IdentityShapeRepresentationFn();
@ -422,8 +424,8 @@ Status BuildHloFromTf(mlir::ModuleOp module_op, xla::XlaBuilder& builder,
llvm::StringRef device_type,
llvm::MutableArrayRef<std::unique_ptr<mlir::Pass>>
custom_legalization_passes) {
if (VLOG_IS_ON(1))
tensorflow::DumpMlirOpToFile("mlir_compile_before_build_hlo_tf", module_op);
if (VLOG_IS_ON(2))
tensorflow::DumpMlirOpToFile("build_hlo_tf_before", module_op);
XlaHelpers::ShapeRepresentationFn shape_representation_fn;
TF_RETURN_IF_ERROR(
@ -434,8 +436,8 @@ Status BuildHloFromTf(mlir::ModuleOp module_op, xla::XlaBuilder& builder,
returns, device_type,
custom_legalization_passes));
if (VLOG_IS_ON(1))
tensorflow::DumpMlirOpToFile("mlir_compile_after_build_hlo_tf", module_op);
if (VLOG_IS_ON(2))
tensorflow::DumpMlirOpToFile("build_hlo_tf_after", module_op);
return Status::OK();
}
@ -600,8 +602,12 @@ Status CompileGraphSetup(
mlir::TF::StandardPipelineOptions tf_options;
mlir::TF::CreateTFStandardPipeline(pm, tf_options);
if (VLOG_IS_ON(1))
tensorflow::DumpMlirOpToFile("compile_graph_setup_before", module_op);
mlir::StatusScopedDiagnosticHandler diag_handler(module_op.getContext());
if (failed(pm.run(module_op))) return diag_handler.ConsumeStatus();
if (VLOG_IS_ON(1))
tensorflow::DumpMlirOpToFile("compile_graph_setup_after", module_op);
return Status::OK();
}