/* Copyright 2017 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.
==============================================================================*/

// This proto file defines messages which store the state of XLA
// computations within the XLA service. A computation is stored as a record
// of the operation requests used to build it.
syntax = "proto3";

import "tensorflow/compiler/xla/xla_data.proto";

package xla;

// Describes a single operation request.
message OperationRequest {
  ComputationDataHandle output_handle = 1;
  Shape output_shape = 2;

  // For operations which call embedded computations such as "Map", these are
  // the version(s) that the embedded computation should be called at. A version
  // value of a computation is the ComputationDataHandle of the root of the
  // computation at the point in time.
  //
  // "Call", "Map", "Reduce", and "ReduceWindow" operations take a single
  // embedded computation so this field will have a single value for those
  // operations.
  //
  // "While" operation takes two; index 0 is the "condition" version and index 1
  // is the "body" version.
  repeated int64 embedded_computation_versions = 3;

  // The actual request, which in itself is a tagged union of all possible
  // operation request types.
  OpRequest request = 4;
}

// Describes a sequence of operation requests which define an XLA
// computation.
message SessionComputation {
  string name = 1;

  // The ComputationHandle used to refer to this computation in the XLA
  // service.
  ComputationHandle computation_handle = 2;

  // Map from ComputationDataHandle value to operation request. The highest
  // ComputationDataHandle value corresponds to the root of the computation.
  map<int64, OperationRequest> requests = 3;
}

// Describes a group of SessionComputations with an "entry point" computation
// that may refer to the other non-entry (AKA embedded) computations.
//
// This message is used to serialize a computation that has been built via the
// XLA service API, along with its dependencies, for purposes such as
// analysis/replay/file-storage.
message SessionModule {
  // The entry computation, which was requested for serialization. This may have
  // referred to embedded computations, which are reflected below.
  SessionComputation entry = 1;

  // Embedded computations that are transitively referred to by the entry
  // computation.
  repeated SessionComputation embedded_computations = 2;

  // The arguments passed to the computation.
  repeated LiteralProto arguments = 3;

  // The result of the computation.
  LiteralProto result = 4;

  // The name of the platform used to run the computation.
  string execution_platform = 5;
}
