Run symbol-dce after TPU Rewrite to remove unreachable functions.

After TPU Rewrite encapsulates device cluster function into TPUCompileMlirOp, run symbolDCE to remove leftover unused functions.

PiperOrigin-RevId: 305692030
Change-Id: I21831824f29e684fd0c3fbe755fe86159c25e637
This commit is contained in:
Prakalp Srivastava 2020-04-09 08:49:40 -07:00 committed by TensorFlower Gardener
parent fbac644835
commit c745bb6716
2 changed files with 2 additions and 0 deletions

View File

@ -5,6 +5,7 @@ module attributes {tf.devices = ["/job:localhost/replica:0/task:0/device:CPU:0",
// CHECK: std.constant // CHECK: std.constant
// CHECK: TPUCompile // CHECK: TPUCompile
// CHECK: TPUExecute // CHECK: TPUExecute
// CHECK-NOT: func @_func
tf_executor.graph { tf_executor.graph {
%outputs, %control = tf_executor.island wraps "std.constant"() {value = dense<2.000000e+00> : tensor<f32>} : () -> tensor<f32> %outputs, %control = tf_executor.island wraps "std.constant"() {value = dense<2.000000e+00> : tensor<f32>} : () -> tensor<f32>
%outputs_0, %control_1 = tf_executor.island wraps "std.constant"() {value = dense<3.000000e+00> : tensor<f32>} : () -> tensor<f32> %outputs_0, %control_1 = tf_executor.island wraps "std.constant"() {value = dense<3.000000e+00> : tensor<f32>} : () -> tensor<f32>

View File

@ -97,6 +97,7 @@ void CreateTPUBridgePipeline(OpPassManager &pm) {
pm.addPass(CreateTPUShardingIdentificationPass()); pm.addPass(CreateTPUShardingIdentificationPass());
pm.addPass(TFDevice::CreateAnnotateParameterReplicationPass()); pm.addPass(TFDevice::CreateAnnotateParameterReplicationPass());
pm.addPass(CreateTPURewritePass()); pm.addPass(CreateTPURewritePass());
pm.addPass(createSymbolDCEPass());
pm.addNestedPass<FuncOp>(TFDevice::CreateReplicateInvariantOpHoistingPass()); pm.addNestedPass<FuncOp>(TFDevice::CreateReplicateInvariantOpHoistingPass());
pm.addNestedPass<FuncOp>(CreateTPUDynamicLayoutPass()); pm.addNestedPass<FuncOp>(CreateTPUDynamicLayoutPass());
pm.addNestedPass<FuncOp>(CreateTPUMergeVariablesWithExecutePass()); pm.addNestedPass<FuncOp>(CreateTPUMergeVariablesWithExecutePass());