| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/ |
| D | tpu_rewrite_device_util_test.cc | 334 const auto& tpu_devices = tpu_device_assignment.tpu_devices; in TEST() local
335 ASSERT_EQ(tpu_devices.size(), 8); in TEST()
336 for (const auto& replica_tpu_devices : tpu_devices) in TEST()
339 EXPECT_EQ(tpu_devices[0][0].device, in TEST()
341 EXPECT_EQ(tpu_devices[0][0].host, in TEST()
343 EXPECT_EQ(tpu_devices[1][0].device, in TEST()
345 EXPECT_EQ(tpu_devices[1][0].host, in TEST()
347 EXPECT_EQ(tpu_devices[2][0].device, in TEST()
349 EXPECT_EQ(tpu_devices[2][0].host, in TEST()
351 EXPECT_EQ(tpu_devices[3][0].device, in TEST()
[all …]
|
| D | tpu_rewrite_device_util.h | 65 TPUDevicesAndHosts&& tpu_devices) in TPUDeviceAssignment()
67 tpu_devices(std::move(tpu_devices)) {} in TPUDeviceAssignment()
70 TPUDevicesAndHosts&& tpu_devices, in TPUDeviceAssignment()
73 tpu_devices(std::move(tpu_devices)), in TPUDeviceAssignment()
77 TPUDevicesAndHosts tpu_devices; member
|
| D | tpu_rewrite_device_util.cc | 119 llvm::SmallVectorImpl<llvm::SmallVector<Device, 8>>* tpu_devices) { in GetTPUDevices() argument
120 tpu_devices->reserve(system_devices.size()); in GetTPUDevices()
142 tpu_devices->push_back(std::move(host_tpu_devices)); in GetTPUDevices()
155 tpu_devices->push_back(std::move(host_tpu_devices)); in GetTPUDevices()
185 llvm::ArrayRef<llvm::SmallVector<Device, 8>> tpu_devices) { in GetFullMeshTPUExecutionDeviceAssignment() argument
186 const int num_tasks = tpu_devices.size(); in GetFullMeshTPUExecutionDeviceAssignment()
187 const int num_tpus_per_task = tpu_devices[0].size(); in GetFullMeshTPUExecutionDeviceAssignment()
204 const auto& tpu_device = tpu_devices[task][device]; in GetFullMeshTPUExecutionDeviceAssignment()
344 llvm::ArrayRef<llvm::SmallVector<Device, 8>> tpu_devices, in GetGeneralTPUExecutionDeviceAssignment() argument
347 const int num_tasks = tpu_devices.size(); in GetGeneralTPUExecutionDeviceAssignment()
[all …]
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/ |
| D | tpu_rewrite_pass.cc | 404 tpu_devices, in AssignDevicesToReplicate() argument
408 const int num_replicas = tpu_devices.size(); in AssignDevicesToReplicate()
409 const int num_cores_per_replica = tpu_devices.front().size(); in AssignDevicesToReplicate()
416 devices_by_core.push_back(tpu_devices[replica][core].device); in AssignDevicesToReplicate()
428 hosts.push_back(tpu_devices[replica][0].host); in AssignDevicesToReplicate()
507 tpu_devices, in AddToParallelExecuteOp() argument
513 const int num_cores_per_replica = tpu_devices.front().size(); in AddToParallelExecuteOp()
548 const bool replicated = tpu_devices.size() != 1; in AddToParallelExecuteOp()
574 : tpu_devices.front()[core].device; in AddToParallelExecuteOp()
588 tpu_devices, in AssignDevicesToReplicatedExecute() argument
[all …]
|
| /external/tensorflow/tensorflow/core/tpu/graph_rewrite/ |
| D | distributed_tpu_rewrite_helpers.cc | 152 int* num_tpus_per_host, std::vector<std::vector<Device*>>* tpu_devices) { in GetTPUDevices() argument
167 tpu_devices->clear(); in GetTPUDevices()
168 tpu_devices->reserve(host_system_devices.size()); in GetTPUDevices()
183 if (tpu_devices->empty()) { in GetTPUDevices()
194 tpu_devices->push_back(std::move(host_tpu_devices)); in GetTPUDevices()
|
| D | update_tpu_embedding_ops_passes.cc | 55 std::vector<Device*> tpu_devices; in Run() local
59 options.device_set->FindMatchingDevices(tpu_device_spec, &tpu_devices); in Run()
60 if (tpu_devices.empty()) { in Run()
81 tpu_devices[0]->parsed_name(); in Run()
88 for (const auto* tpu_device : tpu_devices) { in Run()
|
| D | distributed_tpu_rewrite_helpers.h | 64 std::vector<std::vector<Device*>>* tpu_devices);
|
| D | distributed_tpu_rewrite_pass.cc | 1569 std::vector<std::vector<Device*>>* tpu_devices) { in GetTPUDeviceNames() argument
1586 replication_spec, device_set, num_tpus_per_task, tpu_devices)); in GetTPUDeviceNames()
1713 int num_replicas, const std::vector<std::vector<Device*>>& tpu_devices, in BuildFullMeshDeviceAssignment() argument
1726 (*tf_device_assignment)[i] = {tpu_devices[task][device]->name()}; in BuildFullMeshDeviceAssignment()
1737 const std::vector<std::vector<Device*>>& tpu_devices, in BuildGeneralDeviceAssignment() argument
1765 TF_RET_CHECK(task >= 0 && task < tpu_devices.size()); in BuildGeneralDeviceAssignment()
1766 TF_RET_CHECK(device >= 0 && device < tpu_devices[task].size()); in BuildGeneralDeviceAssignment()
1768 tpu_devices[task][device]->name()); in BuildGeneralDeviceAssignment()
1769 devices_to_lock->push_back((task * tpu_devices[task].size()) + device); in BuildGeneralDeviceAssignment()
1777 const std::vector<std::vector<Device*>>& tpu_devices, int num_replicas, in BuildDeviceAssignment() argument
[all …]
|
| D | distributed_tpu_rewrite_pass.h | 276 const std::vector<std::vector<Device*>>& tpu_devices, int num_replicas,
|
| /external/tensorflow/tensorflow/core/tpu/kernels/ |
| D | tpu_functional_ops.cc | 1219 std::vector<Device*> tpu_devices; in ComputeAsync() local
1457 const std::vector<string>& tpu_devices) { in InitializeShardedVarOnTPU() argument
1467 string device = tpu_devices[i]; in InitializeShardedVarOnTPU()
1861 std::vector<string> tpu_devices; in ReplaceAndPartitionXLAShardingVariable() local
1876 tpu_devices.push_back(tpu_device); in ReplaceAndPartitionXLAShardingVariable()
1970 library_runtime_->device_mgr()->LookupDevice(tpu_devices[i], &d)); in ReplaceAndPartitionXLAShardingVariable()
1981 InitializeShardedVarOnTPU(ctx, var, ndefs, split_dim, tpu_devices)); in ReplaceAndPartitionXLAShardingVariable()
1987 << " device: " << tpu_devices[i]; in ReplaceAndPartitionXLAShardingVariable()
|
| D | tpu_functional_ops.h | 177 const std::vector<string>& tpu_devices)
|
| /external/tensorflow/tensorflow/python/distribute/ |
| D | tpu_strategy.py | 878 tpu_devices = []
890 tpu_devices.append(replica_devices)
891 self._tpu_devices = np.array(tpu_devices, dtype=object)
|