| 86 | } |
| 87 | |
| 88 | void XlaCompilationDevice::Compute(OpKernel* op_kernel, |
| 89 | OpKernelContext* context) { |
| 90 | VLOG(4) << "XlaCompilationDevice::Compute " |
| 91 | << FormatNodeDefForError(op_kernel->def()); |
| 92 | auto* b = XlaContext::Get(context).builder(); |
| 93 | xla::OpMetadata metadata; |
| 94 | metadata.set_op_type(op_kernel->type_string()); |
| 95 | metadata.set_op_name(op_kernel->name()); |
| 96 | b->SetOpMetadata(metadata); |
| 97 | |
| 98 | auto sharding_parse_result = ParseShardingFromDevice( |
| 99 | op_kernel->def(), std::numeric_limits<int>::max()); |
| 100 | OP_REQUIRES_OK(context, sharding_parse_result.status()); |
| 101 | absl::optional<xla::OpSharding> op_sharding = |
| 102 | sharding_parse_result.ValueOrDie(); |
| 103 | |
| 104 | // If no sharding metadata is found, XLA is free to use whatever device it |
| 105 | // wants. In practice this usually has the effect of placing things on device |
| 106 | // 0. |
| 107 | xla::XlaScopedShardingAssignment assign_sharding(b, op_sharding); |
| 108 | op_kernel->Compute(context); |
| 109 | |
| 110 | b->ClearOpMetadata(); |
| 111 | VLOG(4) << "Done"; |
| 112 | } |
| 113 | |
| 114 | Status XlaCompilationDevice::Sync() { return Status::OK(); } |
| 115 |
nothing calls this directly
no test coverage detected