
// Copyright (c) 2020-2023, ARM Limited.
//
//    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.

namespace tosa;

// This corresponds to the version.
file_identifier "TOSA";
// File extension of any written files.
file_extension "tosa";

enum DType:uint32 {
  UNKNOWN = 0,
  BOOL,
  UINT8,
  INT4,
  INT8,
  INT16,
  INT32,
  INT48,
  FP32,
  UINT16,
  FP16,
  BF16,
}

enum ResizeMode:uint32 {
  UNKNOWN = 0,
  NEAREST,
  BILINEAR,
}

enum Op:uint32 {
  UNKNOWN = 0,

  // Tensor Operator
  ARGMAX,
  AVG_POOL2D,
  CONV2D,
  CONV3D,
  DEPTHWISE_CONV2D,
  FULLY_CONNECTED,
  MATMUL,
  MAX_POOL2D,
  TRANSPOSE_CONV2D,

  // Activation
  CLAMP,
  RESERVED,
  SIGMOID,
  TANH,

  // Elementwise-Binary
  ADD,
  ARITHMETIC_RIGHT_SHIFT,
  BITWISE_AND,
  BITWISE_OR,
  BITWISE_XOR,
  INTDIV,
  LOGICAL_AND,
  LOGICAL_LEFT_SHIFT,
  LOGICAL_RIGHT_SHIFT,
  LOGICAL_OR,
  LOGICAL_XOR,
  MAXIMUM,
  MINIMUM,
  MUL,
  POW,
  SUB,
  TABLE,

  // Elementwise-Unary
  ABS,
  BITWISE_NOT,
  CEIL,
  CLZ,
  EXP,
  FLOOR,
  LOG,
  LOGICAL_NOT,
  NEGATE,
  RECIPROCAL,
  RSQRT,

  // Elementwise-Ternary
  SELECT,

  // Logical
  EQUAL,
  GREATER,
  GREATER_EQUAL,

  // Reduction
  REDUCE_ANY,
  REDUCE_ALL,
  REDUCE_MAX,
  REDUCE_MIN,
  REDUCE_PRODUCT,
  REDUCE_SUM,

  // Data layout operation
  CONCAT,
  PAD,
  RESHAPE,
  REVERSE,
  SLICE,
  TILE,
  TRANSPOSE,

  // Gather/scatter operation
  GATHER,
  SCATTER,

  // Image
  RESIZE,

  // Type conversion
  CAST,
  RESCALE,

  // Data Nodes
  CONST,
  IDENTITY,

  // Custom operations
  CUSTOM,

  // Control flow operators
  COND_IF,
  WHILE_LOOP,
}

union Attribute {
  PoolAttribute,
  ConvAttribute,
  TransposeConvAttribute,
  PadAttribute,
  AxisAttribute,
  ReshapeAttribute,
  SliceAttribute,
  TileAttribute,
  ResizeAttribute,
  ClampAttribute,
  RescaleAttribute,
  MulAttribute,
  ArithmeticRightShiftAttribute,
  CondIfAttribute,
  WhileLoopAttribute,
  TransposeAttribute,
  TableAttribute,
  MatMulAttribute,
  FullyConnectedAttribute,
  NegateAttribute,
  CustomAttribute
}

table PoolAttribute {
  pad: [int32];
  kernel: [int32];
  stride: [int32];
  input_zp: int32;
  output_zp: int32;
  accum_dtype: DType;
}

table ConvAttribute {
  pad: [int32];
  stride: [int32];
  dilation: [int32];
  input_zp: int32;
  weight_zp: int32;
  accum_dtype: DType;
}

table TransposeConvAttribute {
  out_pad: [int32];
  stride: [int32];
  output_shape: [int32];
  input_zp: int32;
  weight_zp: int32;
  accum_dtype: DType;
}

table PadAttribute {
  padding: [int32];
  pad_const_int: int32;
  pad_const_fp: float;
}

table AxisAttribute {
  axis: int32;
}

table ReshapeAttribute {
  new_shape: [int32];
}

table SliceAttribute {
  start: [int32];
  size: [int32];
}

table TileAttribute {
  multiples: [int32];
}

table ResizeAttribute {
  scale: [int16];
  offset: [int16];
  border: [int16];
  mode: ResizeMode;
}

table ClampAttribute {
  min_int: int32;
  max_int: int32;
  min_fp: float;
  max_fp: float;
}

table RescaleAttribute {
  input_zp: int32;
  output_zp: int32;
  multiplier: [int32];
  shift: [int32];
  scale32: bool;
  double_round: bool;
  per_channel: bool;
}

table MulAttribute {
  shift: int32;
}

table ArithmeticRightShiftAttribute {
  round: bool;
}

table CondIfAttribute {
  then_branch: string;
  else_branch: string;
}

table WhileLoopAttribute {
  cond_branch: string;
  body_branch: string;
}

table TransposeAttribute {
  perms: [int32];
}

table TableAttribute {
  table: [int16];
}

table MatMulAttribute {
  a_zp: int32;
  b_zp: int32;
  accum_dtype: DType;
}

table FullyConnectedAttribute {
  input_zp: int32;
  weight_zp: int32;
  accum_dtype: DType;
}

table NegateAttribute {
  input1_zp: int32;
  output_zp: int32;
}

table CustomAttribute {
  identifier:string;
  config:string;
  implementation_attrs:[ubyte];
}

table Version {
  _major: int32 = 0;
  _minor: int32 = 51;
  _patch: int32 = 0;
  _draft: bool = true;
}

table TosaTensor {
  name:string;                      // name of the tensor, used for solving dependency
  shape:[int32];                    // shape of the tensor
  type:DType;                       // data type of the tensor
  data: [ubyte] (force_align: 8);   // raw data array if it's a constant tensor.
}

table TosaOperator {
  op:Op;                    // operator enum
  attribute:Attribute;      // union structure. operator attribute
  inputs:[string];          // list of input tensor names
  outputs:[string];         // list of output tensor names
}

table TosaBasicBlock {
  name:string;              // basic block name
  operators:[TosaOperator]; // operators array
  tensors:[TosaTensor];     // tensors array
  inputs:[string];          // name of graph inputs
  outputs:[string];         // name of graph outputs
}

table TosaRegion {
  name:string;             // name of region
  blocks:[TosaBasicBlock]; // basic blocks array
}

table TosaGraph {
  version:Version;
  regions:[TosaRegion];       // regions array
}

root_type TosaGraph;