Method Compile

tensorflow/compiler/xrt/kernels/xrt_compile_ops.cc:82–138

XRTCompileOp::XRTCompileOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}

Status XRTCompileOp::Compile(OpKernelContext* ctx,
                             const xrt::XLAComputation& computation_proto,
                             std::unique_ptr<xla::LocalExecutable>* program) {
  const xrt::XLAComputationConfig& config = computation_proto.config();

  // The default config value is 0; treat it as 1 for convenience.
  int num_replicas = config.num_replicas() ? config.num_replicas() : 1;
  TF_RET_CHECK(num_replicas == 1);
  int num_cores_per_replica =
      config.num_cores_per_replica() ? config.num_cores_per_replica() : 1;
  TF_RET_CHECK(num_cores_per_replica == 1);
  TF_RET_CHECK(config.per_core_program_shape_size() == 0);

  // We are guaranteed that the underlying device object won't be deleted out
  // from under us, while the ScopedRef is live.
  class XRTGenericDeviceAccessor::ScopedRef device_ref;
  TF_RETURN_IF_ERROR(
      XRTGenericDeviceAccessor::InitScopedRef(ctx, 0, &device_ref));

  xla::LocalClient* client = device_ref.client();

  // There is officially no way to use XLA in a client/server architecture where
  // client and server are built from different revisions, because the XLA team
  // does not want to give any guarantees about the stability of the Hlo
  // proto. For cloud TPU this is fine because server and client versions can be
  // assumed to be synced to the same version. For general use the mechanism
  // here (using a snapshot from XlaComputation) works as well as the "official"
  // XLA client/server design, which serializes the same proto between client
  // and server, so in reality is probably fine.
  TF_ASSIGN_OR_RETURN(xla::XlaComputation computation,
                      client->LoadSnapshot(computation_proto.hlo_snapshot()));

  std::vector<xla::Shape> argument_layouts(
      config.program_shape().parameters_size());
  std::vector<const xla::Shape*> argument_layout_ptrs(
      config.program_shape().parameters_size());
  for (int i = 0; i < config.program_shape().parameters_size(); ++i) {
    argument_layouts[i] = xla::Shape(config.program_shape().parameters(i));
    argument_layout_ptrs[i] = &argument_layouts[i];
  }
  xla::ExecutableBuildOptions build_options;
  build_options.set_device_ordinal(client->default_device_ordinal());
  build_options.set_result_layout(xla::Shape(config.program_shape().result()));
  build_options.set_device_allocator(device_ref.backend()->memory_allocator());
  if (config.has_debug_options()) {
    *build_options.mutable_debug_options() =
        BuildXlaDebugOptions(config.debug_options());
  }

  VLOG(1) << "Building executable";
  TF_ASSIGN_OR_RETURN(
      auto executables,
      client->Compile(computation, argument_layout_ptrs, build_options));
  TF_RET_CHECK(executables.size() == 1);
  *program = std::move(executables[0]);
  return Status::OK();
}