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:
parent
fbac644835
commit
c745bb6716
|
@ -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>
|
||||||
|
|
|
@ -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());
|
||||||
|
|
Loading…
Reference in New Issue