Searched refs:tpu_cluster (Results 1 – 6 of 6) sorted by relevance
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/ |
| D | tpu_cluster_util.cc | 41 tf_device::ClusterOp tpu_cluster) { in WalkReachableFromTpuCluster() argument
45 pending_call_nodes.emplace_back(node, tpu_cluster); in WalkReachableFromTpuCluster()
50 auto result = module.walk([&](tf_device::ClusterOp tpu_cluster) { in WalkReachableFromTpuCluster() argument
55 devices, tpu_cluster, &host_device_value))) in WalkReachableFromTpuCluster()
59 cluster_to_host[tpu_cluster.getOperation()] = host_device; in WalkReachableFromTpuCluster()
60 return tpu_cluster.walk([&](Operation* op) { in WalkReachableFromTpuCluster()
61 insert_pending_op(op, tpu_cluster); in WalkReachableFromTpuCluster()
62 return callback(op, tpu_cluster, in WalkReachableFromTpuCluster()
63 cluster_to_host[tpu_cluster.getOperation()]); in WalkReachableFromTpuCluster()
76 tf_device::ClusterOp tpu_cluster = pair.second; in WalkReachableFromTpuCluster() local
[all …]
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/ |
| D | tpu_extract_outside_compilation.cc | 309 tf_device::ClusterOp tpu_cluster, in GetStaticExternalOperands() argument
320 tpu_cluster->isAncestor(defining_op) && in GetStaticExternalOperands()
358 tf_device::ClusterOp tpu_cluster, in GetExternalOperands() argument
366 return GetStaticExternalOperands(tpu_cluster, cluster_ops); in GetExternalOperands()
590 LogicalResult MoveOpsToHost(tf_device::ClusterOp tpu_cluster, Block* src, in MoveOpsToHost() argument
611 GetExternalOperands(tpu_cluster, clustered_ops); in MoveOpsToHost()
629 GetExternalOperands(tpu_cluster, clustered_ops); in MoveOpsToHost()
646 LogicalResult DecomposeControlFlow(tf_device::ClusterOp tpu_cluster, in DecomposeControlFlow() argument
649 auto result = tpu_cluster.GetBody().walk([&](Operation* op) { in DecomposeControlFlow()
654 if (failed(MoveOpsToHost(tpu_cluster, &if_op.then_branch().front(), in DecomposeControlFlow()
[all …]
|
| D | outside_compiled_to_host_launch.cc | 62 auto traverse_op = [&](Operation* op, tf_device::ClusterOp tpu_cluster, in runOnOperation()
67 tpu_cluster.emitOpError( in runOnOperation()
|
| D | tpu_cluster_cleanup_attributes.cc | 40 auto traverse_op = [&](Operation* op, tf_device::ClusterOp tpu_cluster) { in runOnOperation() argument
|
| D | host_launch_to_outside_compiled.cc | 69 auto traverse_op = [&](Operation* op, tf_device::ClusterOp tpu_cluster, in runOnOperation()
|
| D | mark_ops_for_outside_compilation.cc | 274 void MarkVariantInputsOutputs(tf_device::ClusterOp tpu_cluster) { in MarkVariantInputsOutputs() argument
276 tpu_cluster.walk([&](Operation* op) { in MarkVariantInputsOutputs()
|