diff options
| author |  Tong Shen <endlessroad@google.com> | 2018-09-09 01:37:02 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-09-09 01:41:14 -0700 |
| commit | a3776a234f555213aafcf41f49a42a8a9448c4ac (patch) | |
| tree | 5497f0e149f2d463b6bf5b1dcc15746044e64bc7 /tensorflow/compiler/jit | |
| parent | 1bf545492596f1d3dbaf1485de500116a2d2a25b (diff) | |
Move control flow functionalization as a graph optimization pass, instead of a step in XlaCompiler.
PiperOrigin-RevId: 212164482
Diffstat (limited to 'tensorflow/compiler/jit')
| -rw-r--r-- | tensorflow/compiler/jit/BUILD | 1 | ||||
| -rw-r--r-- | tensorflow/compiler/jit/jit_compilation_pass_registration.cc | 12 |
2 files changed, 13 insertions, 0 deletions
diff --git a/tensorflow/compiler/jit/BUILD b/tensorflow/compiler/jit/BUILD index a989f15a1c..7d5db713f6 100644 --- a/tensorflow/compiler/jit/BUILD +++ b/tensorflow/compiler/jit/BUILD @@ -265,6 +265,7 @@ cc_library( srcs = ["jit_compilation_pass_registration.cc"], deps = [ ":compilation_passes", + "//tensorflow/compiler/tf2xla:functionalize_control_flow_pass_registration", "//tensorflow/core:core_cpu_internal", ], alwayslink = 1, diff --git a/tensorflow/compiler/jit/jit_compilation_pass_registration.cc b/tensorflow/compiler/jit/jit_compilation_pass_registration.cc index c37b6112cc..5dcf754969 100644 --- a/tensorflow/compiler/jit/jit_compilation_pass_registration.cc +++ b/tensorflow/compiler/jit/jit_compilation_pass_registration.cc @@ -21,6 +21,18 @@ limitations under the License. namespace tensorflow { +// PRE_PLACEMENT passes: + +// from +// third_party/tensorflow/compiler/tf2xla/functionalize_control_flow_pass_registration.cc +// FunctionalizeControlFlowPass: 27 +// +// This pass looks at the graph and all associated FunctionDefs, and turns +// traditional control flow structure (Switch/Merge/etc.) into functional +// control flow structure (XlaIf/XlaWhile). Following passes must +// handle those FunctionDef correctly. + +// POST_REWRITE_FOR_EXEC passes: REGISTER_OPTIMIZATION(OptimizationPassRegistry::POST_REWRITE_FOR_EXEC, 10, MarkForCompilationPass); |