Move TensorArray decomposition pass before region to functional control flow conversion.

Some unit tests fail if we remove tf-region-control-flow-to-functional pass completely. So we still need this conversion before control flow legalization. Created a new bug to track it.

PiperOrigin-RevId: 338355165
Change-Id: I31bed4d9bd98bfa483cbce3f595aed6cf19c6ea9
This commit is contained in:
Prakalp Srivastava 2020-10-21 15:41:43 -07:00 committed by TensorFlower Gardener
parent be10b788f7
commit eb68be4a42

View File

@ -278,10 +278,6 @@ void CreateConvertMlirToXlaHloPipeline(
pm.addNestedPass<mlir::FuncOp>(mlir::createCanonicalizerPass());
pm.addPass(mlir::TF::CreateTensorListOpsDecompositionPass());
pm.addPass(mlir::TF::CreateStackOpsDecompositionPass());
// TODO(b/159127949): TensorArray decomposition passes does not handle region
// based control flow yet. So convert back to functional control flow.
pm.addPass(mlir::TF::CreateTFRegionControlFlowToFunctional());
pm.addPass(mlir::TF::CreateTensorArrayOpsDecompositionPass());
pm.addPass(mlir::TFDevice::CreateDecomposeResourceOpsPass());
pm.addPass(mlir::TF::CreatePromoteResourcesToArgsPass());
@ -289,6 +285,10 @@ void CreateConvertMlirToXlaHloPipeline(
// Guarantee all functions have one use, which enables shape inference.
pm.addPass(mlir::TF::CreateGuaranteeAllFuncsOneUsePass());
pm.addPass(mlir::TF::CreateTFShapeInferencePass());
// TODO(b/171426148): We cannot completely remove region to functional control
// flow conversion from this pipeline yet as it causes some unit tests to
// fail.
pm.addPass(mlir::TF::CreateTFRegionControlFlowToFunctional());
// LegalizeTFControlFlow encapsulates arguments for control flow operations
// with a tuple argument which break the assumption of resource lifting
// inside PromoteResourcesToArgs.