
// Copyright (c) 2020-2024, 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";

// NOTE: New values added to the schema should be placed
// at the end of the list in order to keep schema stable.

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

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

enum Op:uint32 {
  UNKNOWN = 0,
  ARGMAX,
  AVG_POOL2D,
  CONV2D,
  CONV3D,
  DEPTHWISE_CONV2D,
  FULLY_CONNECTED,
  MATMUL,
  MAX_POOL2D,
  TRANSPOSE_CONV2D,
  CLAMP,
  RESERVED,
  SIGMOID,
  TANH,
  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,
  ABS,
  BITWISE_NOT,
  CEIL,
  CLZ,
  EXP,
  FLOOR,
  LOG,
  LOGICAL_NOT,
  NEGATE,
  RECIPROCAL,
  RSQRT,
  SELECT,
  EQUAL,
  GREATER,
  GREATER_EQUAL,
  REDUCE_ANY,
  REDUCE_ALL,
  REDUCE_MAX,
  REDUCE_MIN,
  REDUCE_PRODUCT,
  REDUCE_SUM,
  CONCAT,
  PAD,
  RESHAPE,
  REVERSE,
  SLICE,
  TILE,
  TRANSPOSE,
  GATHER,
  SCATTER,
  RESIZE,
  CAST,
  RESCALE,
  CONST,
  IDENTITY,
  CUSTOM,
  COND_IF,
  WHILE_LOOP,
  FFT2D,
  RFFT2D,
  ERF,
  DIM,
  CONST_SHAPE,
  CONCAT_SHAPE,
  ADD_SHAPE,
  SUB_SHAPE,
  MUL_SHAPE,
  DIV_SHAPE,
}

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

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;
  local_bound: bool;
}

table TransposeConvAttribute {
  out_pad: [int32];
  stride: [int32];
  output_shape: [int32];
  input_zp: int32;
  weight_zp: int32;
  local_bound: bool;
}

table PadAttribute {
  padding: [int32];
  pad_const_int: int32;
  pad_const_fp: [ubyte] (force_align: 8);
}

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: [ubyte] (force_align: 8);
  max_fp: [ubyte] (force_align: 8);
}

table RescaleAttribute {
  input_zp: int32;
  output_zp: int32;
  multiplier: [int32];
  shift: [int32];
  scale32: bool;
  double_round: bool;
  per_channel: bool;
  input_unsigned: bool;
  output_unsigned: 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;
}

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

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

table CustomAttribute {
  operator_name:string;
  domain_name:string;
  implementation_attrs:[ubyte];
}

table FFTAttribute {
  inverse: bool;
  local_bound: bool;
}

table RFFTAttribute {
  local_bound: bool;
}

table Version {
  _major: int32 = -1;
  _minor: int32 = -1;
  _patch: int32 = -1;
  _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.
  variable: bool;                   // is this a variable tensor
  is_unranked: bool;                // whether this is an unranked tensor
  variable_name:string;             // name for variable attribute
}

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 (required);
  regions:[TosaRegion];       // regions array
}

root_type TosaGraph;