1/* Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 3Licensed under the Apache License, Version 2.0 (the "License"); 4you may not use this file except in compliance with the License. 5You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9Unless required by applicable law or agreed to in writing, software 10distributed under the License is distributed on an "AS IS" BASIS, 11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12See the License for the specific language governing permissions and 13limitations under the License. 14==============================================================================*/ 15syntax = "proto3"; 16 17package tensorflow.tpu; 18 19import "tensorflow/core/framework/attr_value.proto"; 20import "tensorflow/core/framework/function.proto"; 21import "tensorflow/core/framework/tensor.proto"; 22import "tensorflow/core/framework/tensor_shape.proto"; 23import "tensorflow/core/protobuf/tpu/compile_metadata.proto"; 24 25// TPU compilation request for compiling computations into XLA HLO IR and build 26// TPU programs. 27message TpuCompilationRequestProto { 28 // A flag reserved for using experimental version of the compilation. By 29 // default the value should be false. 30 bool use_experimental = 1; 31 32 // Use mlir to lower computation(s) to Hlo. 33 bool use_mlir = 2; 34 35 // If true, returns hlo metadatas. 36 bool return_hlo_protos = 3; 37 38 // If true, unloads cache on session close. 39 bool unload_cache_on_session_close = 4; 40 41 // Compilation metadata. 42 TPUCompileMetadataProto metadata = 5; 43 44 // Computation argument shapes. 45 repeated TensorShapeProto arg_shapes = 6; 46 47 // Input tensor that gives const guarantee to the TF runtime. 48 repeated TensorProto guaranteed_constants = 7; 49 50 // MLIR module definition. 51 string mlir_module = 8; 52 53 // A set of named functions used as the input to lowering to Hlo when mlir is 54 // not used. 55 FunctionDefLibrary fdef_lib = 9; 56 57 // The version of the graph definition used to lower TF function to Hlo. 58 int32 graph_def_version = 10; 59 60 // Function containing the computation to compile. 61 NameAttrList function = 11; 62} 63