1 /* Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 // This proto file defines messages which store the state of XLA 17 // computations within the XLA service. A computation is stored as a record 18 // of the operation requests used to build it. 19 syntax = "proto3"; 20 21 import "tensorflow/compiler/xla/xla_data.proto"; 22 23 package xla; 24 25 // Describes a single operation request. 26 message OperationRequest { 27 ComputationDataHandle output_handle = 1; 28 Shape output_shape = 2; 29 30 // For operations which call embedded computations such as "Map", these are 31 // the version(s) that the embedded computation should be called at. A version 32 // value of a computation is the ComputationDataHandle of the root of the 33 // computation at the point in time. 34 // 35 // "Call", "Map", "Reduce", and "ReduceWindow" operations take a single 36 // embedded computation so this field will have a single value for those 37 // operations. 38 // 39 // "While" operation takes two; index 0 is the "condition" version and index 1 40 // is the "body" version. 41 repeated int64 embedded_computation_versions = 3; 42 43 // The actual request, which in itself is a tagged union of all possible 44 // operation request types. 45 OpRequest request = 4; 46 } 47 48 // Describes a sequence of operation requests which define an XLA 49 // computation. 50 message SessionComputation { 51 string name = 1; 52 53 // The ComputationHandle used to refer to this computation in the XLA 54 // service. 55 ComputationHandle computation_handle = 2; 56 57 // Map from ComputationDataHandle value to operation request. The highest 58 // ComputationDataHandle value corresponds to the root of the computation. 59 map<int64, OperationRequest> requests = 3; 60 } 61 62 // Describes a group of SessionComputations with an "entry point" computation 63 // that may refer to the other non-entry (AKA embedded) computations. 64 // 65 // This message is used to serialize a computation that has been built via the 66 // XLA service API, along with its dependencies, for purposes such as 67 // analysis/replay/file-storage. 68 message SessionModule { 69 // The entry computation, which was requested for serialization. This may have 70 // referred to embedded computations, which are reflected below. 71 SessionComputation entry = 1; 72 73 // Embedded computations that are transitively referred to by the entry 74 // computation. 75 repeated SessionComputation embedded_computations = 2; 76 77 // The arguments passed to the computation. 78 repeated LiteralProto arguments = 3; 79 80 // The result of the computation. 81 LiteralProto result = 4; 82 83 // The name of the platform used to run the computation. 84 string execution_platform = 5; 85 } 86
