static*/
| 215 | } |
| 216 | |
| 217 | /*static*/ xla::StatusOr<std::shared_ptr<XrtExecutable>> XrtExecutable::Compile( |
| 218 | std::shared_ptr<XrtContext> context, |
| 219 | const xla::HloModuleProto& hlo_module_proto, |
| 220 | const std::vector<xla::Shape>& argument_shapes, |
| 221 | const xla::Shape& result_shape, xla::DeviceAssignment device_assignment) { |
| 222 | if (device_assignment.replica_count() <= 0 || |
| 223 | device_assignment.computation_count() <= 0) { |
| 224 | return errors::InvalidArgument( |
| 225 | "Device assignment must be non-empty; got ", |
| 226 | device_assignment.replica_count(), " replicas and ", |
| 227 | device_assignment.computation_count(), " computations per replica."); |
| 228 | } |
| 229 | |
| 230 | // TODO(phawkins): add support for per-core argument and return shapes. |
| 231 | TF_RET_CHECK(device_assignment.computation_count() == 1) |
| 232 | << "Computation count != 1 not implemented"; |
| 233 | |
| 234 | xrt::XLAComputation computation; |
| 235 | computation.mutable_config()->set_num_replicas( |
| 236 | device_assignment.replica_count()); |
| 237 | computation.mutable_config()->set_num_cores_per_replica( |
| 238 | device_assignment.computation_count()); |
| 239 | |
| 240 | xrt::DeviceAssignment* xrt_assignment = |
| 241 | computation.mutable_config()->mutable_device_assignment(); |
| 242 | for (int computation = 0; computation < device_assignment.computation_count(); |
| 243 | ++computation) { |
| 244 | xrt::DeviceAssignment::ComputationDevice* xrt_devices = |
| 245 | xrt_assignment->add_computation_devices(); |
| 246 | for (int replica = 0; replica < device_assignment.replica_count(); |
| 247 | ++replica) { |
| 248 | int xrt_device_ordinal = device_assignment(replica, computation); |
| 249 | if (xrt_device_ordinal < 0 || |
| 250 | xrt_device_ordinal >= context->tf_device_ids().size()) { |
| 251 | return errors::InvalidArgument("Invalid device ordinal in device ", |
| 252 | "assignment: ", xrt_device_ordinal); |
| 253 | } |
| 254 | *xrt_devices->add_replica_devices() = |
| 255 | context->device_mesh_coordinates().at(xrt_device_ordinal); |
| 256 | } |
| 257 | } |
| 258 | |
| 259 | xla::ProgramShape program_shape; |
| 260 | for (const xla::Shape& shape : argument_shapes) { |
| 261 | xla::Shape* param_shape = program_shape.add_parameters(); |
| 262 | *param_shape = shape; |
| 263 | if (!xla::LayoutUtil::HasLayout(shape)) { |
| 264 | xla::LayoutUtil::SetToDefaultLayout(param_shape); |
| 265 | } |
| 266 | } |
| 267 | *program_shape.mutable_result() = result_shape; |
| 268 | if (!xla::LayoutUtil::HasLayout(result_shape)) { |
| 269 | xla::LayoutUtil::SetToDefaultLayout(program_shape.mutable_result()); |
| 270 | } |
| 271 | *computation.mutable_config()->mutable_program_shape() = |
| 272 | program_shape.ToProto(); |
| 273 | *computation.mutable_hlo_snapshot()->mutable_hlo()->mutable_hlo_module() = |
| 274 | hlo_module_proto; |