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

Support Session's ConfigProto in TF2 MLIR bridge

Browse Source 
Some models are using "TF2" but also using Session and passing
a ConfigProto. This is TF1 code running on the TF2 MLIR bridge.
The TF2 MLIR bridge assumed that this was not possible. This cl
updates the TF2 version of the MLIR bridge to support a ConfigProto
passed in via Session.

Disable MLIR bridge presubmit testing of saved_model_test.py because this fix reveals that the test is actually broken. It is using TF2 control flow but loading a model with Session.

PiperOrigin-RevId: 335915669
Change-Id: Ib50bef389449ce0011878dd50b73856e9c520289
This commit is contained in:
Marissa Ikonomidis 2020-10-07 11:39:06 -07:00 committed by TensorFlower Gardener
parent 6f8ee54073
commit 3708f73481
3 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -38,7 +38,7 @@ auto* mlir_bridge_gauge_v2 = monitoring::Gauge<bool, 0>::New(
// encapsulated graph to a particular device.
Status MlirBridgePass::Run(const ConfigProto& config_proto,
mlir::ModuleOp module) {
if (!config_proto.experimental().enable_mlir_bridge()) {
if (!IsEnabled(config_proto)) {
VLOG(0) << "Skipping MLIR TPU Bridge, session flag not enabled";
mlir_bridge_gauge_v2->GetCell()->Set(false);
return Status::OK();

3
tensorflow/compiler/tf2xla/mlir_bridge_pass.h
View File

@ -30,7 +30,8 @@ class MlirBridgePass : public MlirOptimizationPass {
llvm::StringRef name() const override { return "bridge"; }
bool IsEnabled(const ConfigProto& config_proto) const override {
return config_proto.experimental().enable_mlir_bridge();
return config_proto.experimental().enable_mlir_bridge() ||
tensorflow::GetMlirCommonFlags()->tf_mlir_enable_mlir_bridge;
}
// This should be used as a thin mapper around mlir::ModulePass::runOnModule

1
tensorflow/python/distribute/integration_test/BUILD
View File

@ -9,7 +9,6 @@ package(
distribute_py_test(
name = "saved_model_test",
srcs = ["saved_model_test.py"],
disable_mlir_bridge = False,
deps = [
"//tensorflow:tensorflow_py",
"//tensorflow/python:lookup_ops",