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

[MLIR] Move functional->region to almost the end of the first part of bridge

Browse Source 
- This also means that the MLIR attached to _TPUCompileMlirOp has region based
  control flow. So introduce a region -> functional control flow transformation at the
  start of the second part of the bridge in ConvertMLIRToXlaComputation.

PiperOrigin-RevId: 331797895
Change-Id: I208bfad71dc7506ebe7cc7c8b647ec0d92f54c9b
This commit is contained in:
Rahul Joshi 2020-09-15 10:23:05 -07:00 committed by TensorFlower Gardener
parent 8d4b9e6a6c
commit 579a12c875
2 changed files with 2 additions and 1 deletions
tensorflow/compiler/mlir/tensorflow
transforms
bridge.cc
utils
compile_mlir_util.cc

2
tensorflow/compiler/mlir/tensorflow/transforms/bridge.cc
View File

@ -105,7 +105,6 @@ void CreateTPUBridgePipeline(OpPassManager &pm) {
pm.addPass(TFDevice::CreateMarkOpsForOutsideCompilationPass());
pm.addPass(CreateTPUExtractHeadTailOutsideCompilationPass());
pm.addPass(CreateTPUExtractOutsideCompilationPass());
pm.addPass(TF::CreateTFRegionControlFlowToFunctional());
pm.addNestedPass<FuncOp>(tf_executor::CreateTFExecutorConstantSinkingPass());
pm.addPass(TF::CreateResourceDeviceInferencePass());
@ -121,6 +120,7 @@ void CreateTPUBridgePipeline(OpPassManager &pm) {
pm.addNestedPass<FuncOp>(CreateTPUParallelExecuteSinkResourceWritePass());
pm.addNestedPass<FuncOp>(CreateTPUMergeVariablesWithExecutePass());
pm.addNestedPass<FuncOp>(CreateTPUColocateCompositeResourceOps());
pm.addPass(TF::CreateTFRegionControlFlowToFunctional());
pm.addPass(CreateTPUVariableReformattingPass());
}

1
tensorflow/compiler/mlir/tensorflow/utils/compile_mlir_util.cc
View File

@ -273,6 +273,7 @@ void CreateConvertMlirToXlaHloPipeline(
mlir::OpPassManager& pm, llvm::StringRef device_type,
llvm::MutableArrayRef<std::unique_ptr<mlir::Pass>>
custom_legalization_passes) {
pm.addPass(mlir::TF::CreateTFRegionControlFlowToFunctional());
pm.addNestedPass<mlir::FuncOp>(mlir::createCanonicalizerPass());
pm.addPass(mlir::TF::CreateTensorListOpsDecompositionPass());
pm.addPass(mlir::TF::CreateStackOpsDecompositionPass());