/* Copyright (c) 2016 PaddlePaddle 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 cinn.frontend.paddle.proto;

// Any incompatible changes to ProgramDesc and its dependencies should
// raise the version defined version.h.
//
// Serailization and Deserialization codes should be modified in a way
// that supports old versions following the version and compatibility policy.
message Version { optional int64 version = 1 [ default = 0 ]; }

enum AttrType {
  INT = 0;
  FLOAT = 1;
  STRING = 2;
  INTS = 3;
  FLOATS = 4;
  STRINGS = 5;
  BOOLEAN = 6;
  BOOLEANS = 7;
  BLOCK = 8;
  LONG = 9;
  BLOCKS = 10;
  LONGS = 11;
}

// OpDesc describes an instance of a C++ framework::OperatorBase
// derived class type.
message OpDesc {

  message Attr {
    required string name = 1;
    required AttrType type = 2;
    optional int32 i = 3;
    optional float f = 4;
    optional string s = 5;
    repeated int32 ints = 6;
    repeated float floats = 7;
    repeated string strings = 8;
    optional bool b = 10;
    repeated bool bools = 11;
    optional int32 block_idx = 12;
    optional int64 l = 13;
    repeated int32 blocks_idx = 14;
    repeated int64 longs = 15;
  };

  message Var {
    required string parameter = 1;
    repeated string arguments = 2;
  };

  required string type = 3;
  repeated Var inputs = 1;
  repeated Var outputs = 2;
  repeated Attr attrs = 4;
  optional bool is_target = 5 [ default = false ];
};

// OpProto describes a C++ framework::OperatorBase derived class.
message OpProto {

  // VarProto describes the C++ type framework::Variable.
  message Var {
    required string name = 1;
    required string comment = 2;

    optional bool duplicable = 3 [ default = false ];
    optional bool intermediate = 4 [ default = false ];
    optional bool dispensable = 5 [ default = false ];
  }

  // AttrProto describes the C++ type Attribute.
  message Attr {
    required string name = 1;
    required AttrType type = 2;
    required string comment = 3;
    // If that attribute is generated, it means the Paddle third
    // language binding has responsibility to fill that
    // attribute. End-User should not set that attribute.
    optional bool generated = 4 [ default = false ];
  }

  required string type = 1;
  repeated Var inputs = 2;
  repeated Var outputs = 3;
  repeated Attr attrs = 4;
  required string comment = 5;
}

message VarType {
  enum Type {
    // Pod Types
    BOOL = 0;
    INT16 = 1;
    INT32 = 2;
    INT64 = 3;
    FP16 = 4;
    FP32 = 5;
    FP64 = 6;
    // Tensor<size_t> is used in C++.
    SIZE_T = 19;
    UINT8 = 20;
    INT8 = 21;
    BF16 = 22;

    // Other types that may need additional descriptions
    LOD_TENSOR = 7;
    SELECTED_ROWS = 8;
    FEED_MINIBATCH = 9;
    FETCH_LIST = 10;
    STEP_SCOPES = 11;
    LOD_RANK_TABLE = 12;
    LOD_TENSOR_ARRAY = 13;
    PLACE_LIST = 14;
    READER = 15;
    // Any runtime decided variable type is raw
    // raw variables should manage their own allocations
    // in operators like nccl_op
    RAW = 17;
    TUPLE = 18;
  }

  required Type type = 1;

  message TensorDesc {
    // Should only be PODType. Is enforced in C++
    required Type data_type = 1;
    repeated int64 dims = 2; // [UNK, 640, 480] is saved as [-1, 640, 480]
  }
  optional TensorDesc selected_rows = 2;

  message LoDTensorDesc {
    required TensorDesc tensor = 1;
    optional int32 lod_level = 2 [ default = 0 ];
  }
  optional LoDTensorDesc lod_tensor = 3;

  message LoDTensorArrayDesc {
    required TensorDesc tensor = 1;
    optional int32 lod_level = 2 [ default = 0 ];
  }
  optional LoDTensorArrayDesc tensor_array = 4;

  message ReaderDesc { repeated LoDTensorDesc lod_tensor = 1; }
  optional ReaderDesc reader = 5;

  message Tuple { repeated Type element_type = 1; }
  optional Tuple tuple = 7;
}

message VarDesc {
  required string name = 1;
  required VarType type = 2;
  optional bool persistable = 3 [ default = false ];
  // True if the variable is an input data and
  // have to check the feed data shape and dtype
  optional bool need_check_feed = 4 [ default = false ];
}

message BlockDesc {
  required int32 idx = 1;
  required int32 parent_idx = 2;
  repeated VarDesc vars = 3;
  repeated OpDesc ops = 4;
  optional int32 forward_block_idx = 5 [ default = -1 ];
}

// CompatibleInfo is used to determine if a feature is compatible and
// provides the information.
message CompatibleInfo {
  enum Type {
    COMPATIBLE = 0;
    DEFINITELY_NOT = 1;
    POSSIBLE = 2;
    BUG_FIX = 3;
    PRECISION_CHANGE = 4;
  }
  required string version = 1;
  required Type type = 2;
}

// In some cases, Paddle Fluid may perform operator definition iterations,
// and the operator uses OpCompatibleMap for compatibility testing.
message OpCompatibleMap {
  message OpCompatiblePair {
    required string op_name = 1;
    required CompatibleInfo compatible_info = 2;
  }
  repeated OpCompatiblePair pair = 1;
  optional string default_required_version = 2;
}

// Please refer to
// https://github.com/PaddlePaddle/Paddle/blob/develop/doc/design/program.md
// for more details.
// TODO(panyx0718): A model can have multiple programs. Need a
// way to distinguish them. Maybe ID or name?
message ProgramDesc {
  reserved 2; // For backward compatibility.
  repeated BlockDesc blocks = 1;
  optional Version version = 4;
  optional OpCompatibleMap op_compatible_map = 3;
}
