/* Copyright 2020 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 = "proto3";

package tensorflow.tpu;

import "tensorflow/core/framework/attr_value.proto";
import "tensorflow/core/framework/function.proto";
import "tensorflow/core/framework/tensor.proto";
import "tensorflow/core/framework/tensor_shape.proto";
import "tensorflow/core/protobuf/tpu/compile_metadata.proto";

// TPU compilation request for compiling computations into XLA HLO IR and build
// TPU programs.
message TpuCompilationRequestProto {
  // A flag reserved for using experimental version of the compilation. By
  // default the value should be false.
  bool use_experimental = 1;

  // Use mlir to lower computation(s) to Hlo.
  bool use_mlir = 2;

  // If true, returns hlo metadatas.
  bool return_hlo_protos = 3;

  // If true, unloads cache on session close.
  bool unload_cache_on_session_close = 4;

  // Compilation metadata.
  TPUCompileMetadataProto metadata = 5;

  // Computation argument shapes.
  repeated TensorShapeProto arg_shapes = 6;

  // Input tensor that gives const guarantee to the TF runtime.
  repeated TensorProto guaranteed_constants = 7;

  // MLIR module definition.
  string mlir_module = 8;

  // A set of named functions used as the input to lowering to Hlo when mlir is
  // not used.
  FunctionDefLibrary fdef_lib = 9;

  // The version of the graph definition used to lower TF function to Hlo.
  int32 graph_def_version = 10;

  // Function containing the computation to compile.
  NameAttrList function = 11;
}
