Searched refs:topology_attr (Results 1 – 6 of 6) sorted by relevance
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/ |
| D | tpu_rewrite_device_util.cc | 261 llvm::StringRef topology_attr, int num_tasks, int num_tpus_per_task) { in ParseTopologyAttr() argument
263 if (!topology_proto.ParseFromString(topology_attr.str())) in ParseTopologyAttr()
345 llvm::StringRef topology_attr, in GetGeneralTPUExecutionDeviceAssignment() argument
350 TF_ASSIGN_OR_RETURN(auto topology, ParseTopologyAttr(topology_attr, num_tasks, in GetGeneralTPUExecutionDeviceAssignment()
447 llvm::StringRef topology_attr, in GetTPUCompilationAndExecutionDevices() argument
459 if (topology_attr.empty()) { in GetTPUCompilationAndExecutionDevices()
475 topology_attr, device_assignment_attr)); in GetTPUCompilationAndExecutionDevices()
502 auto topology_attr = in GetHostDeviceOutsideComputation() local
504 if (!topology_attr) in GetHostDeviceOutsideComputation()
534 num_cores_per_replica_attr.getInt(), topology_attr.getValue(), in GetHostDeviceOutsideComputation()
|
| D | tpu_rewrite_device_util_test.cc | 58 std::string topology_attr; in TEST_P() local
62 devices, /*num_replicas=*/1, /*num_cores_per_replica=*/1, topology_attr, in TEST_P()
299 std::string topology_attr = topology_proto.SerializeAsString(); in TEST() local
308 devices, /*num_replicas=*/1, /*num_cores_per_replica=*/1, topology_attr, in TEST()
322 std::string topology_attr; in TEST() local
326 devices, /*num_replicas=*/8, /*num_cores_per_replica=*/1, topology_attr, in TEST()
418 std::string topology_attr = topology_proto.SerializeAsString(); in TEST() local
429 devices, /*num_replicas=*/4, /*num_cores_per_replica=*/2, topology_attr, in TEST()
532 std::string topology_attr = topology_proto.SerializeAsString(); in TEST() local
542 devices, /*num_replicas=*/2, /*num_cores_per_replica=*/3, topology_attr, in TEST()
|
| D | tpu_rewrite_device_util.h | 238 int num_cores_per_replica, llvm::StringRef topology_attr,
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/ |
| D | tpu_rewrite_pass.cc | 732 auto topology_attr = in Rewrite() local
734 if (!topology_attr) in Rewrite()
757 topology_attr.getValue(), status_or_device_coodinates.value()); in Rewrite()
|
| /external/tensorflow/tensorflow/core/tpu/graph_rewrite/ |
| D | distributed_tpu_rewrite_pass.h | 277 int num_cores_per_replica, const string& topology_attr,
|
| D | distributed_tpu_rewrite_pass.cc | 1593 static Status ParseTopologyAttr(const string& topology_attr, in ParseTopologyAttr() argument
1599 proto.ParseFromString(topology_attr); in ParseTopologyAttr()
1778 int num_cores_per_replica, const string& topology_attr, in BuildDeviceAssignment() argument
1815 if (topology_attr.empty()) { in BuildDeviceAssignment()
1860 TF_RETURN_IF_ERROR(ParseTopologyAttr(topology_attr, tpu_topology, num_tasks, in BuildDeviceAssignment()
|