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

[mlir] Add more logging on MLIR bridge execution.

Browse Source 
In the default case, we do not want to flood log with the "Running MLIR ..." messages. Log if bridge was executed once. User can enable v=1 option for more detailed logging.

PiperOrigin-RevId: 358199172
Change-Id: I03267b05a9ee4e05e7b0f926809a0002d6f60b42
This commit is contained in:
Roman Dzhabarov 2021-02-18 09:43:37 -08:00 committed by TensorFlower Gardener
parent 6e05909d98
commit 10e5d5cdb1
1 changed files with 24 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

32
tensorflow/compiler/tf2xla/mlir_bridge_pass.cc
View File

@ -106,6 +106,20 @@ MlirOptimizationPassState MlirBridgePass::GetPassState(
}
}
namespace {
// Log just once by default (on default log level), and let the user adjust
// the log level for more detailed logging.
void LogAtLeastOnce(const std::string& log_message) {
if (VLOG_IS_ON(1)) {
VLOG(1) << log_message;
} else {
LOG_FIRST_N(INFO, 1) << log_message;
}
}
} // namespace
// This runs the first phase of the "bridge", transforming the graph in a form
// that can be executed with delegation of some computations to an accelerator.
// This builds on the model of XLA where a subset of the graph is encapsulated
@ -118,19 +132,19 @@ Status MlirBridgePass::Run(const ConfigProto& config_proto,
// based on the devices in the module.
if (GetPassState(/*device_set=*/nullptr, config_proto, graph) ==
MlirOptimizationPassState::Disabled) {
VLOG(1) << "Skipping MLIR TPU Bridge, session flag not enabled";
LogAtLeastOnce("Skipping MLIR TPU Bridge, session flag not enabled");
mlir_bridge_gauge_v2->GetCell()->Set(false);
return Status::OK();
}
// Skip MLIR TPU Bridge if no TPU devices or TPU ops found.
if (!HasTPUDevicesAndOps(module)) {
VLOG(1) << "Skipping MLIR TPU Bridge, no TPU devices or TPU ops found";
LogAtLeastOnce("Skipping MLIR TPU Bridge, no TPU devices or TPU ops found");
return Status::OK();
}
// TODO(b/178633630): Revisit whether to use LOG_FIRST_N.
VLOG(1) << "Running MLIR TPU Bridge";
LogAtLeastOnce("Running MLIR TPU Bridge");
mlir_bridge_gauge_v2->GetCell()->Set(true);
TF_RETURN_IF_ERROR(
mlir::TFTPU::TPUBridge(module, /*enable_logging=*/VLOG_IS_ON(1)));
@ -162,19 +176,21 @@ Status MlirBridgeV1CompatPass::Run(const GraphOptimizationPassOptions& options,
// based on the devices in the module.
if (!IsEnabled(/*device_set=*/nullptr, options.session_options->config,
**options.graph)) {
VLOG(1) << "Skipping MLIR TPU Bridge V1 Compat, session flag not enabled";
LogAtLeastOnce(
"Skipping MLIR TPU Bridge V1 Compat, session flag not enabled");
mlir_bridge_gauge_v1->GetCell()->Set(false);
return Status::OK();
}
// Skip MLIR TPU Bridge if no TPU devices or TPU ops found.
if (!HasTPUDevicesAndOps(module)) {
VLOG(1) << "Skipping MLIR TPU Bridge V1 Compat, no TPU devices or TPU ops "
"found";
LogAtLeastOnce(
"Skipping MLIR TPU Bridge V1 Compat, no TPU devices or TPU ops found");
return Status::OK();
}
VLOG(1) << "Running MLIR TPU Bridge V1 Compat";
LogAtLeastOnce("Running MLIR TPU Bridge V1 Compat");
mlir_bridge_gauge_v1->GetCell()->Set(true);
TF_RETURN_IF_ERROR(
mlir::TFTPU::TPUBridgeV1Compat(module, /*enable_logging=*/VLOG_IS_ON(1)));