Searched refs:new_parallel_execute (Results 1 – 5 of 5) sorted by relevance
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/ |
| H A D | xla_sharding_util.cc | 393 int cluster_idx, mlir::tf_device::ParallelExecuteOp new_parallel_execute, in GetTileShardedOutputsToMerge() argument
407 new_parallel_execute.GetRegionOutputs( in GetTileShardedOutputsToMerge()
421 int cluster_idx, mlir::tf_device::ParallelExecuteOp new_parallel_execute, in HandleTileShardedOutputs() argument
425 builder->setInsertionPointAfter(new_parallel_execute); in HandleTileShardedOutputs()
432 cluster_to_core_index, cluster_idx, new_parallel_execute, in HandleTileShardedOutputs()
558 mlir::tf_device::ParallelExecuteOp new_parallel_execute, in RemapOutputsFromLogicalDevices() argument
596 new_parallel_execute.GetRegionOutputs( in RemapOutputsFromLogicalDevices()
606 cluster_to_core_index, cluster_idx, new_parallel_execute, in RemapOutputsFromLogicalDevices()
620 new_parallel_execute, builder))) in RemapOutputsFromLogicalDevices()
638 new_parallel_execute.GetRegionOutputs( in RemapOutputsFromLogicalDevices()
|
| H A D | xla_sharding_util.h | 71 mlir::tf_device::ParallelExecuteOp new_parallel_execute,
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/ |
| H A D | tpu_rewrite_pass.cc | 468 tf_device::ParallelExecuteOp* new_parallel_execute) { in MovePreservedParallelExecuteChildren() argument
471 *new_parallel_execute = builder->create<tf_device::ParallelExecuteOp>( in MovePreservedParallelExecuteChildren()
493 new_parallel_execute->getRegions()[new_idx]->takeBody( in MovePreservedParallelExecuteChildren()
512 tf_device::ParallelExecuteOp* new_parallel_execute, int* cluster_idx) { in AddToParallelExecuteOp() argument
537 old_parallel_execute, new_parallel_execute); in AddToParallelExecuteOp()
543 builder->setInsertionPoint(*new_parallel_execute); in AddToParallelExecuteOp()
553 new_parallel_execute->GetRegionBlockWithIndex((*cluster_idx) + core); in AddToParallelExecuteOp()
827 tf_device::ParallelExecuteOp new_parallel_execute; in Rewrite() local
834 old_parallel_execute, &new_parallel_execute, &cluster_idx); in Rewrite()
843 old_parallel_execute, cluster_idx, new_parallel_execute, builder); in Rewrite()
[all …]
|
| H A D | tpu_parallel_execute_sink_resource_write.cc | 132 auto new_parallel_execute = builder.create<tf_device::ParallelExecuteOp>( in SinkResourceWritesIntoParallelExecute() local
135 for (auto region : llvm::zip(new_parallel_execute.getRegions(), in SinkResourceWritesIntoParallelExecute()
140 llvm::zip(results_to_remap, new_parallel_execute.getResults())) in SinkResourceWritesIntoParallelExecute()
|
| H A D | tpu_merge_variables_with_execute.cc | 384 auto new_parallel_execute = builder->create<tf_device::ParallelExecuteOp>( in ReplaceParallelExecute() local
389 Operation* new_parallel_execute_op = new_parallel_execute.getOperation(); in ReplaceParallelExecute()
|