// Copyright 2019 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================

syntax = "proto2";

package tpu_driver;

import "tensorflow/compiler/xla/python/tpu_driver/tpu_driver.proto";
import "tensorflow/compiler/xla/service/hlo.proto";
import "tensorflow/compiler/xla/xla_data.proto";

option optimize_for = SPEED;

message StatusMessage {
  required int32 code = 1;
  optional string message = 2;
}

message AllocateRequest {
  required int32 core_id = 1;
  required MemoryRegion region = 2;
  oneof size {
    int64 num_bytes = 3;
    xla.ShapeProto shape = 4;
  }
}

message AllocateTupleRequest {
  required int32 core_id = 1;
  required MemoryRegion region = 2;
  repeated int64 children = 3;
}

message DeallocateRequest {
  required int64 handle = 1;
}

message TransferToDeviceRequest {
  required int64 target_handle = 1;
  required bytes data = 2;
}

message TransferFromDeviceRequest {
  required int64 source_handle = 1;
}

message TransferFromDeviceResponse {
  required bytes data = 2;
}

message TransferFromDeviceToDeviceRequest {
  required int64 source_handle = 1;
  required int64 target_handle = 2;
}

message CompileRequest {
  required xla.HloProto hlo_program = 1;
  optional int64 num_replicas = 2;
}

message CompiledProgramMetadata {
  required xla.ProgramShapeProto program_shape = 1;
}

message CompileResponse {
  required CompiledProgramMetadata metadata = 1;
}

message LoadProgramRequest {
  required int32 core_id = 1;
  required int64 compiled_program_handle = 2;
}

message UnloadProgramRequest {
  required int64 loaded_program_handle = 1;
}

message ExecuteRequest {
  required int64 loaded_program_handle = 1;
  repeated int64 input_handle = 2;
  repeated int64 output_handle = 3;
  optional xla.DeviceAssignmentProto device_assignment = 4;
}

message StreamRequest {
  message Entry {
    oneof request {
      AllocateRequest alloc = 1;
      AllocateTupleRequest alloc_tuple = 2;
      DeallocateRequest dealloc = 3;
      TransferToDeviceRequest transfer_to = 4;
      TransferFromDeviceRequest transfer_from = 5;
      TransferFromDeviceToDeviceRequest transfer_from_to = 10;
      CompileRequest compile = 6;
      LoadProgramRequest load = 7;
      UnloadProgramRequest unload = 8;
      ExecuteRequest execute = 9;
    }
    // If specified, a list of encoded EventId values.
    repeated int64 wait_for_id = 20;
    // A unique, encoded EventId value.
    // For Allocate, Compile, and Load, this also defines the result handle.
    required int64 operation_id = 21;

    // A unique identifier for the thread that issued this request. Currently
    // for debugging purposes only.
    optional int64 thread_id = 22;
  }

  repeated Entry entry = 30;
}

message StreamResponse {
  message Entry {
    oneof response {
      TransferFromDeviceResponse transfer_from = 3;
      CompileResponse compile = 4;
    }
    required StatusMessage status = 10;
    // Echos the given encoded EventId value.
    required int64 operation_id = 11;
  }

  repeated Entry entry = 20;
}

message OpenRequest {
  // The version number for this client. Versions are bumped in case of
  // backwards incompatible client-server protocol changes. Servers will reject
  // clients with an unsupported version.
  optional int32 client_version = 1 [default = 0];
}

message OpenResponse {
  required uint32 client_id = 1;

  // Maximum time this client can be idle before it is GC'ed and all resources
  // released.
  optional int32 max_idle_time_seconds = 2 [default = 3600];
}

message CloseRequest {
  required fixed32 client_id = 1;
}

message CloseResponse {}

message ResetRequest {}

message ResetResponse {}

message QuerySystemInfoRequest {}

message QuerySystemInfoResponse {
  required SystemInfo system_info = 1;
}

service CloudTpuDriver {
  // Open the driver. If the driver is already open, return an error.
  rpc Open(OpenRequest) returns (OpenResponse);

  // Close the driver. Any outstanding requests will be terminated.
  rpc Close(CloseRequest) returns (CloseResponse);

  // Reset the driver. All connected clients will be disconnected.
  rpc Reset(ResetRequest) returns (ResetResponse);

  // Query the driver for current system performance information.
  rpc QuerySystemInfo(QuerySystemInfoRequest) returns (QuerySystemInfoResponse);

  // Enqueue an operation to be executed when its dependencies are satisfied.
  rpc StreamExecute(stream StreamRequest) returns (stream StreamResponse);
}
