Searched refs:cluster_to_core_index (Results 1 – 3 of 3) sorted by relevance
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/ |
| H A D | xla_sharding_util.cc | 370 llvm::SmallVector<llvm::SmallVector<int, 4>, 4> cluster_to_core_index, in LookupClusterToCoreIndex() argument
374 cluster_to_core_index[core_id][cluster_func_output_index]; in LookupClusterToCoreIndex()
392 llvm::SmallVector<llvm::SmallVector<int, 4>, 4> cluster_to_core_index, in GetTileShardedOutputsToMerge() argument
403 location, cluster_to_core_index, logical_device_id, in GetTileShardedOutputsToMerge()
419 llvm::SmallVector<llvm::SmallVector<int, 4>, 4> cluster_to_core_index, in HandleTileShardedOutputs() argument
432 cluster_to_core_index, cluster_idx, new_parallel_execute, in HandleTileShardedOutputs()
515 llvm::SmallVectorImpl<int>* cluster_to_core_index) { in GetOutputTypesForLogicalDeviceComputation() argument
539 cluster_to_core_index->emplace_back(core_index++); in GetOutputTypesForLogicalDeviceComputation()
543 cluster_to_core_index->emplace_back(core_index++); in GetOutputTypesForLogicalDeviceComputation()
546 cluster_to_core_index->emplace_back(-1); in GetOutputTypesForLogicalDeviceComputation()
[all …]
|
| H A D | xla_sharding_util.h | 61 llvm::SmallVectorImpl<int>* cluster_to_core_index);
69 llvm::SmallVector<llvm::SmallVector<int, 4>, 4> cluster_to_core_index,
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/ |
| H A D | tpu_rewrite_pass.cc | 444 llvm::SmallVector<int, 4> cluster_to_core_index; in BuildExecuteOp() local
447 &cluster_to_core_index); in BuildExecuteOp()
509 llvm::SmallVectorImpl<llvm::SmallVector<int, 4>>* cluster_to_core_index, in AddToParallelExecuteOp() argument
525 cluster_to_core_index->emplace_back(llvm::SmallVector<int, 4>()); in AddToParallelExecuteOp()
529 &(*cluster_to_core_index)[core]); in AddToParallelExecuteOp()
829 llvm::SmallVector<llvm::SmallVector<int, 4>, 4> cluster_to_core_index; in Rewrite() local
830 cluster_to_core_index.reserve(num_cores_per_replica); in Rewrite()
833 &cluster_to_core_index, compile_op, cluster_func, builder, in Rewrite()
842 cluster_func.getLoc(), output_shardings, cluster_to_core_index, in Rewrite()
|