commit 4e118127252147ab755b1f92c8127f93564205f8
author Rahul Joshi <jurahul@google.com> Fri Jul 31 14:12:22 2020 -0700
committer TensorFlower Gardener <gardener@tensorflow.org> Fri Jul 31 14:51:34 2020 -0700
tree 5ec79bfa2d5cef1d411d6a9effcdc076cfc19fc3
parent 1edbacaacc88d39bd96f774f07c141a216c6a52b

Introduce functional<->region conversion passes around extract outside compilation

- Follow functional->region transformation with a inlining pass to make sure calls
  generated by the transform get inlined.

PiperOrigin-RevId: 324282111
Change-Id: Ifaacec3d8919f390fdeda8ca9af129d8e7dce086

tensorflow/compiler/mlir/tensorflow/transforms/bridge.cc

diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/bridge.cc b/tensorflow/compiler/mlir/tensorflow/transforms/bridge.cc
index cb1dd23..ed0528a 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/bridge.cc
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/bridge.cc
@@ -82,15 +82,23 @@
   // Run shape inference so that tf_executor/tf_device ops created later will
   // likely to inherit more concrete types.
   pm.addPass(TF::CreateTFShapeInferencePass());
-  OpPassManager &func_pm = pm.nest<FuncOp>();
-  func_pm.addPass(CreateTPUClusterFormationPass());
-  // Place DecomposeResourceOpsPass before TFExecutorConstantSinking pass
-  // because DecomposeResourceOpsPass uses pattern rewriter which hoists
-  // changed constants out of tf_device.Launch.
-  func_pm.addPass(TFDevice::CreateDecomposeResourceOpsPass());
-  func_pm.addPass(CreateTPUHostComputationExpansionPass());
-  pm.addNestedPass<FuncOp>(CreateTPUUpdateEmbeddingEnqueueOpInputsPass());
+  // Encode this in its own scope so that func_pm is not mistakenly used
+  // later on.
+  {
+    OpPassManager &func_pm = pm.nest<FuncOp>();
+    func_pm.addPass(CreateTPUClusterFormationPass());
+    // Place DecomposeResourceOpsPass before TFExecutorConstantSinking pass
+    // because DecomposeResourceOpsPass uses pattern rewriter which hoists
+    // changed constants out of tf_device.Launch.
+    func_pm.addPass(TFDevice::CreateDecomposeResourceOpsPass());
+    func_pm.addPass(CreateTPUHostComputationExpansionPass());
+    func_pm.addPass(CreateTPUUpdateEmbeddingEnqueueOpInputsPass());
+  }
+  pm.addPass(TF::CreateTFFunctionalControlFlowToRegions());
+  pm.addPass(mlir::createInlinerPass());
   pm.addPass(CreateTPUExtractHeadTailOutsideCompilationPass());
+  pm.addPass(TF::CreateTFRegionControlFlowToFunctional());
+
   // Run another shape inference pass because resource decomposition might have
   // created new partial types.
   pm.addPass(TF::CreateTFShapeInferencePass());