| 530 | } |
| 531 | |
| 532 | Status XrtContext::InitializeTPU() { |
| 533 | LOG(INFO) << "Initializing TPU devices."; |
| 534 | TF_RETURN_IF_ERROR(RegisterTPUInitializeFunction(tf_context_.get())); |
| 535 | |
| 536 | TensorProto index_proto; |
| 537 | index_proto.set_dtype(DT_INT32); |
| 538 | index_proto.add_int_val(0); |
| 539 | XrtTensorHandle device_ordinal = EnqueueConst( |
| 540 | tf_context_.get(), /*device_id=*/tf_context_->cpu_device_id(), |
| 541 | index_proto, /*host_memory=*/false); |
| 542 | |
| 543 | protobuf::Map<string, AttrValue> attrs; |
| 544 | attrs["f"].mutable_func()->set_name("TPUInitFunc"); |
| 545 | attrs["Tin"].mutable_list(); |
| 546 | attrs["Tout"].mutable_list()->add_type(DT_STRING); |
| 547 | XrtTensorHandle t = std::move( |
| 548 | tf_context_->EnqueueOp("TPUPartitionedCall", {&device_ordinal}, |
| 549 | /*output_arity=*/1, |
| 550 | /*attrs=*/attrs, tf_context_->cpu_device_id())[0]); |
| 551 | |
| 552 | auto result = tf_context_->RecvTensor(t, DT_STRING, /*host_memory=*/false); |
| 553 | TF_ASSIGN_OR_RETURN(RecvTensorResponse * response, result->Get()); |
| 554 | VLOG(10) << "TPU topology " << response->DebugString(); |
| 555 | |
| 556 | TF_ASSIGN_OR_RETURN(std::string data, |
| 557 | DeserializeTensorProtoAsString(response->tensor())); |
| 558 | |
| 559 | tpu::TopologyProto tpu_topology; |
| 560 | tpu_topology.ParsePartialFromString(data); |
| 561 | VLOG(4) << "TPU topology:\n" << tpu_topology.DebugString(); |
| 562 | |
| 563 | TF_RET_CHECK(tpu_topology.num_tasks() == 1) << tpu_topology.DebugString(); |
| 564 | TF_RET_CHECK(tpu_topology.num_tpu_devices_per_task() == tf_device_ids_.size()) |
| 565 | << tpu_topology.DebugString() << " " << tf_device_ids_.size(); |
| 566 | |
| 567 | const int mesh_rank = tpu_topology.mesh_shape_size(); |
| 568 | TF_RET_CHECK(tpu_topology.device_coordinates_size() == |
| 569 | tf_device_ids_.size() * mesh_rank); |
| 570 | |
| 571 | for (int i = 0; i < tf_device_ids_.size(); ++i) { |
| 572 | device_mesh_coordinates_.push_back({}); |
| 573 | auto& coords = device_mesh_coordinates_.back(); |
| 574 | for (int j = 0; j < mesh_rank; ++j) { |
| 575 | coords.add_value(tpu_topology.device_coordinates(i * mesh_rank + j)); |
| 576 | } |
| 577 | } |
| 578 | |
| 579 | LOG(INFO) << "TPU initialization succeeded."; |
| 580 | return Status::OK(); |
| 581 | } |
| 582 | |
| 583 | XrtContext::ExecuteReplicatedKey::ExecuteReplicatedKey( |
| 584 | absl::Span<const int> input_arity, xla::DeviceAssignment device_assignment) |
no test coverage detected