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

Instrument the number of times the MLIR-based TF Bridge is enabled.

Browse Source 
PiperOrigin-RevId: 311633792
Change-Id: Iba286e1c82900833b5cf9f69a697a312e51f3156
This commit is contained in:
Lucy Fox 2020-05-14 16:59:59 -07:00 committed by TensorFlower Gardener
parent e6c2a5a212
commit 90077f8c7c
2 changed files with 13 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
tensorflow/compiler/tf2xla/BUILD
View File

@ -704,12 +704,8 @@ cc_library(
deps = [
"//tensorflow/compiler/mlir:mlir_graph_optimization_pass",
"//tensorflow/compiler/mlir/tensorflow",
"//tensorflow/compiler/mlir/tensorflow:convert_graphdef",
"//tensorflow/compiler/mlir/tensorflow:device_util",
"//tensorflow/compiler/mlir/tensorflow:dump_mlir_util",
"//tensorflow/compiler/mlir/tensorflow:mlir_roundtrip_flags",
"//tensorflow/core:core_cpu",
"@com_google_absl//absl/container:flat_hash_set",
"//tensorflow/core:lib",
"@llvm-project//llvm:support",
],
alwayslink = 1,

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

@ -18,10 +18,18 @@ limitations under the License.
#include <string>
#include "tensorflow/compiler/mlir/tensorflow/transforms/bridge.h"
#include "tensorflow/core/lib/monitoring/gauge.h"
#include "tensorflow/core/public/session_options.h"
namespace tensorflow {
auto* mlir_bridge_gauge_v1 = monitoring::Gauge<bool, 0>::New(
"/tensorflow/config/experimental/enable_mlir_bridge_gauge_v1",
"Tracks usage of the MLIR-based TF2XLA bridge among TF1 models");
auto* mlir_bridge_gauge_v2 = monitoring::Gauge<bool, 0>::New(
"/tensorflow/config/experimental/enable_mlir_bridge_gauge_v2",
"Tracks usage of the MLIR-based TF2XLA bridge among TF2 models");
// 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
@ -32,10 +40,12 @@ Status MlirBridgePass::Run(const ConfigProto& config_proto,
mlir::ModuleOp module) {
if (!config_proto.experimental().enable_mlir_bridge()) {
VLOG(0) << "Skipping MLIR TPU Bridge, session flag not enabled";
mlir_bridge_gauge_v2->GetCell()->Set(false);
return Status::OK();
}
VLOG(0) << "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)));
@ -48,10 +58,12 @@ Status MlirBridgeV1CompatPass::Run(const GraphOptimizationPassOptions& options,
if (!options.session_options->config.experimental().enable_mlir_bridge()) {
VLOG(0) << "Skipping MLIR TPU Bridge V1 Compat, session flag not enabled";
mlir_bridge_gauge_v1->GetCell()->Set(false);
return Status::OK();
}
VLOG(0) << "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)));