//
//  MNN.fbs
//  MNN
//
//  Created by jiangxiaotang on 2019/1/4.
//  Copyright © 2018, Alibaba Group Holding Limited
//

include "CaffeOp.fbs";
include "TensorflowOp.fbs";
include "TFQuantizeOp.fbs";
include "ExtraInfo.fbs";
include "UserDefine.fbs";
namespace MNN;
attribute "priority";

enum OpType : int {
    AbsVal,
    QuantizedAdd,
    ArgMax,
    AsString,
    InstanceNorm,
    BatchToSpaceND,
    Copy,
    BinaryOp,
    Bnll,
    Cast,
    Concat,
    Const,
    Convolution,
    ConvolutionDepthwise,
    Crop,
    CropAndResize,
    ImageProcess,
    Deconvolution,
    DeconvolutionDepthwise,
    Dequantize,
    DetectionOutput,
    Dropout,
    Eltwise,
    ELU,
    Unique,
    Exp,
    ExpandDims,
    Fill,
    Flatten,
    Im2Col,
    Gather,
    GatherV2,
    Im2Seq,
    InnerProduct,
    Input,
    Interp,
    Log,
    LRN,
    LSTM,
    MatMul,
    MoE,
    NonMaxSuppression,
    NonMaxSuppressionV2,
    Normalize,
    Pack,
    Padding,
    Permute,
    Pooling,
    Power,
    PReLU,
    PriorBox,
    Proposal,
    QuantizedAvgPool,
    QuantizedBiasAdd,
    QuantizedConcat,
    QuantizedDepthwiseConv2D,
    QuantizedLogistic,
    RasterAndInterpolate,
    QuantizedMaxPool,
    Texture,
    RasterDiff,
    QuantizedReshape,
    QuantizedSoftmax,
    QuantizeMaxMin,
    QuantizeV2,
    Range,
    Rank,
    ReduceJoin,
    Reduction,
    ReLU,
    ReLU6, // Use as Clip
    RequantizationRange,
    Requantize,
    Reshape,
    Resize,
    RNN,
    ROIPooling,
    Scale,
    Selu,
    Seq2Out,
    Shape,
    Sigmoid,
    Size,
    Slice,
    SliceTf,
    Softmax,
    SpaceToBatchND,
    SpatialProduct,
    Col2Im,
    Segment,
    Squeeze,
    StridedSlice,
    CastLike,
    StringSplit,
    StringToNumber,
    TanH,
    TfQuantizedConv2D,
    Threshold,
    Tile,
    TopKV2,
    Transpose,
    UnaryOp,
    Unpack,
    Where,

    Moments,
    RNNSequenceGRU,
    BatchMatMul,
    Unsqueeze,
    CosineSimilarity,
    DepthToSpace,
    SpaceToDepth,
    ReverseSequence,
    Pooling3D,
    Convolution3D,
    MatrixBandPart,
    GatherND,
    DetectionPostProcess,
    UnravelIndex,
    ScatterNd,
    OneHot,
    BroadcastTo,
    Dilation2D,
    Interp3D,
    Raster = 128,
    ConvertTensor = 129,
    ArgMin = 130,
    LinSpace = 131,
    RandomUniform = 132,
    // TensorArray Ops
    TensorArray = 133,
    TensorArraySize = 134,
    TensorArrayRead = 135,
    TensorArrayWrite = 136,
    TensorArrayGather = 137,
    TensorArrayScatter = 138,
    TensorArraySplit = 139,
    TensorArrayConcat = 140,
    LSTMBlockCell = 141,
    Reverse = 142,
    ROIAlign = 143,
    RandomNormal = 144,
    TensorArrayInsert = 145,
    TensorArrayErase = 146,
    EyeLike = 147,
    CumSum = 148,
    Det = 149,
    CumProd = 150,
    ScatterElements = 151,
    GatherElements = 152,
    Svd = 153,
    Histogram = 154,
    DynamicQuant = 155,
    Stft = 156,

    Plugin = 256, //The Type load from plugin
    //Training Op Start from 257
    Select = 257,
    ZerosLike,
    Broastcast,
    SetDiff1D,
    ReluGrad,
    Identity,
    PoolGrad,
    SoftmaxGrad,
    Conv2DBackPropFilter,
    TrainableParam,
    BatchNorm,
    ConvTranspose3D,
    // Use for self defined grad
    ZeroGrad,

    // User define op
    Attention = 299,
    FmhaV2 = 300,
    Fmhca = 301,
    SeqLen2Spatial = 302,
    SplitGeLU = 303,
    GroupNorm = 304,
    LinearAttention = 305,
    RoPE = 306,

    Extra = 512,
    // quantization
    ConvInt8 = 513,
    Int8ToFloat = 514,
    DepthwiseConvInt8 = 515,
    FloatToInt8 = 517,
    EltwiseInt8 = 518,

    While = 600,
    If    = 601,
    LayerNorm = 603,
    GridSample = 604,
}

table Plugin {
    type: string;
    attr: [Attribute];
}

table Extra {
    type: string;
    engine: string;
    info: [byte];
    attr:[Attribute];
    // The Extra Op can be vectorized for execution
    vector: bool;
}

table StringVec {
    data: [string];
}

table AttentionParam {
    kv_cache: bool = true;
    kv_shared_layer: string;
    layer_index: int = -1;
    kv_shared_layer_index: int = -1;
    mhq_quant:[TensorQuantInfo]; // qk_scale_q, qk_scale_k, sv_scale_s, sv_scale_v
    output_c4:bool = false;
    attnScale: float = 0.0;
}

table LinearAttentionParam {
    attn_type: string;
    num_k_heads: int;
    num_v_heads: int;
    head_k_dim: int;
    head_v_dim: int;
    use_qk_l2norm: bool;
}

table RoPEParam {
    rope_cut_head_dim: int = 0;
    num_head: int = 0;
    kv_num_head: int = 0;
    head_dim: int = 0;
    q_norm: LayerNorm;
    k_norm: LayerNorm;
}

table FmhaV2Param {
    heads: int;
}

table FmhcaParam {
    heads: int;
}

table StftParam {
    n_fft: int;
    hop_length: int;
    abs: bool = true;
}

table ShapeParam {
    hasStart: bool = false;
    start: int = 0;
    hasEnd: bool = false;
    end: int = 0;
}

table WhileParam {
    // The name of condition subgraph.
    cond_graph: string;

    // The name of body subgraph.
    body_graph: string;

    // Aliases the forigen inputs and subgraph inputs.
    // One forign input maybe correspond to multiple subgraph inputs.
    aliases_inputs: [StringVec];

    // Aliases the forigen outputs and subgraph outputs.
    aliases_outputs: [string];

    // Aliases body subgraph outputs and inputs.
    // One input should be updated by only one output.
    aliases_updates: [StringVec];
}

table IfParam {
    // The name of then subgraph.
    then_graph: string;

    // The name of else subgraph.
    else_graph: string;

    // Aliases the forigen inputs and subgraph inputs.
    // One forign input maybe correspond to multiple subgraph inputs.
    aliases_inputs: [StringVec];

    // Aliases the forigen outputs and subgraph outputs.
    // Each output should has two possible value from then branch and
    // else branch.
    aliases_outputs: [StringVec];
}

table RegionCommand {
    op: Op;
    steps:[int];
    size:[int];
    indexes:[int];
    view:[View];

    // How to treat output's old value for the region position
    // If fuse < 0, it means out_new = compute
    // -3 means output outside region should be set as zero
    // For backend memory layout is not linear:
    // -2 means output is totally invalid, we can rewrite output outside region
    // -1 means the output has valid value outside compute region, we can only rewrite the value in region

    // If fuse >= 0, use as a BinaryOpOperation op
    // For example, BinaryOpOperation_DIV means: out_new = out_old / compute
    fuse: int = -1;

    // If iterIndex == -1, use 0 -> loopNumber - 1
    // If iterIndex >= 0, for i in (inputTensor[loopIndexes])
    iterIndexes:[int];
}

table LoopParam {
    tensorNumber:int;
    outputIndexes:[int];
    inputIndexes:[int];
    extraTensorInfos:[TensorDescribe];
    // The loop can be parallel executed
    parallel:bool = true;
    // The number of loop (for i = 0; i<loopNumber; ++i)
    loopNumber:int;

    commands:[RegionCommand];
    // init output Tensor before loop:
    // 1. nullptr mean no init;
    // 2. RegionCommand whiout any param mean all zero init;
    // 3. Execute RegionCommand to init output;
    initCommand:[RegionCommand];
}

union OpParameter {
    QuantizedAdd,
    ArgMax,
    AsString,
    Axis,
    BatchNorm,
    BinaryOp,
    Blob,
    CastParam,
    Convolution2D,
    Crop,
    CropAndResize,
    Dequantize,
    DetectionOutput,
    Eltwise,
    ExpandDims,
    Fill,
    Flatten,
    Gather,
    GatherV2,
    InnerProduct,
    Input,
    Interp,
    LRN,
    LSTM,
    MatMul,
    NonMaxSuppressionV2,
    Normalize,
    PackParam,
    Permute,
    Plugin,
    Pool,
    PRelu,
    PriorBox,
    Proposal,
    QuantizedAvgPool,
    QuantizedBiasAdd,
    QuantizedConcat,
    QuantizedLogistic,
    QuantizedMatMul,
    QuantizedMaxPool,
    QuantizedRelu,
    QuantizedRelu6,
    QuantizedReshape,
    QuantizedSoftmax,
    QuantizeMaxMin,
    QuantizeV2,
    Range,
    Rank,
    ReduceJoin,
    ReductionParam,
    Relu,
    Relu6,
    RequantizationRange,
    Requantize,
    Reshape,
    Resize,
    RoiParameters,
    Scale,
    Selu,
    Size,
    Slice,
    SliceTf,
    SpaceBatch,
    SqueezeParam,
    StridedSliceParam,
    TensorConvertInfo,
    TfQuantizedConv2D,
    TopKV2,
    Transpose,
    UnaryOp,
    MomentsParam,
    RNNParam,
    BatchMatMulParam,
    QuantizedFloatParam,
    DepthSpaceParam, // DepthToSpace and SpaceToDepth using the same parameter
    EltwiseInt8,
    ReverseSequenceParam,
    Extra,
    Pool3D,
    Convolution3D,
    ELU,
    DetectionPostProcessParam,
    OneHotParam,
    PadParam,
    WhileParam,
    IfParam,
    RandomUniform,
    LayerNorm,
    TensorArray,
    LSTMBlockCell,
    GridSample,
    LoopParam,
    ImageProcessParam,
    CumSum,
    GroupNorm,
    FmhaV2Param,
    FmhcaParam,
    AttentionParam,
    StftParam,
    LinearAttentionParam,
    ShapeParam,
    RoPEParam
}

table Op {
    inputIndexes: [int];
    main: OpParameter;
    name: string;
    outputIndexes: [int];
    type: OpType;
    defaultDimentionFormat : MNN_DATA_FORMAT = NHWC;
    externalPath:string;
}

table View {
    offset:int;
    stride:[int];
}

table Region {
    src:View;
    dst:View;
    size:[int];
    origin:int;
}

table TensorDescribe {
    blob: Blob;
    index: int;
    name: string;
    regions:[Region];
    quantInfo:TensorQuantInfo;
}

enum ForwardType : byte {
    CPU = 0,
    METAL,
    CUDA,
    OPENCL,
    AUTO,
    NNAPI,
    OPENGLES,
    VULKAN,
}
enum Usage : byte {
    INFERENCE = 0,
    TRAIN = 1,
    INFERENCE_STATIC = 2
}

table SubGraphProto {
    // Subgraph unique name.
    name: string;

    // The ids of input tensors.
    inputs: [int];

    // The ids of output tensors.
    outputs: [int];

    // All tensor names.
    // The id of each tensor is the index in the vector names.
    tensors: [string];

    // Nodes of the subgraph.
    nodes: [Op];

    // Tensor describe info
    extraTensorDescribe:[TensorDescribe];
}

table TensorQuantInfo {
    scale:float;
    zero:float = 0;
    min:float = -128;
    max:float = 127;
    type:DataType;
}

table Net {
    bizCode: string;
    extraTensorDescribe: [TensorDescribe];
    extraInfo: ExtraInfo;
    oplists: [Op];
    outputName: [string];
    preferForwardType: ForwardType = CPU;
    sourceType: NetSource = CAFFE;
    tensorName: [string];
    tensorNumber: int = 0;
    usage:Usage = INFERENCE;  // used to more compatibility in future

    // Subgraphs of the Net.
    subgraphs: [SubGraphProto];
    mnn_uuid: string;
}

root_type Net;