// Copyright 2022 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.

// prototype for stablehlo schema, WIP
// WARNING: converting to stablehlo file is experimental feature, and no runtime
// support is provided

namespace stablehlo.flatbuf;

enum DataType: byte{
  FLOAT16 = 0,
  FLOAT32 = 1,
  FLOAT64 = 2,
  INT4 = 3,
  INT8 = 4,
  INT16 = 5,
  INT32 = 6,
  UINT4 = 7,
  UINT8 = 8,
  UINT16 = 9,
  UINT32 = 10,
  UINT64 = 11,
  INT64 = 12,
}

table Tensor {
  // The tensor shape.
  shape:[int];
  type:DataType;
  // An index that refers to the buffers table at the root of the model. Or,
  // if there is no data buffer associated (i.e. intermediate results), then
  // this is 0 (which refers to an always existent empty buffer).
  //
  // The data_buffer itself is an opaque container, with the assumption that the
  // target device is little-endian. In addition, all builtin operators assume
  // the memory is ordered such that if `shape` is [4, 3, 2], then index
  // [i, j, k] maps to data_buffer[i*3*2 + j*2 + k].
  buffer:uint;
  name:string;  // For debugging and importing back into tensorflow.

  //for now assuming the tensor always have rank
}

enum OperatorCode : int32 {
  DOT = 0,
  ADD = 1,
  CONVOLUTION = 2,
  MAXIMUM = 3,
  MINIMUM = 4,
  RESHAPE = 5,
  DIVIDE = 6,
  MULTIPLY = 7,
  REDUCE = 8,
  REDUCE_WINDOW = 9,
  BROADCAST_IN_DIM = 10,
  LOGISTIC = 11,
  CUSTOM_CALL = 12,
  BATCH_NORM_INFERENCE = 13,
  CLAMP = 14,
  SLICE = 15,
  CONCATENATE = 16,
  IOTA = 17,
  SUBTRACT = 18,
  CEIL = 19,
  CONVERT = 20,
  GATHER = 21,
  ABS = 22,
  DOT_GENERAL = 23,
  RESIZE_BILINEAR = 24,
}

// Options for stablehlo operators.
union OperatorOptions {
  DotOptions,
  AddOptions,
  ConvolutionOptions,
  MaximumOptions,
  MinimumOptions,
  ReshapeOptions,
  DivideOptions,
  MultiplyOptions,
  ReduceOptions,
  ReduceWindowOptions,
  BroadcastInDimOptions,
  LogisticOptions,
  CustomCallOptions,
  BatchNormInferenceOptions,
  ClampOptions,
  SliceOptions,
  ConcatenateOptions,
  IotaOptions,
  SubtractOptions,
  CeilOptions,
  ConvertOptions,
  GatherOptions,
  AbsOptions,
  DotGeneralOptions,
  ResizeBilinearOptions,
}

table DotOptions {
}

table AddOptions {
}

table ConvolutionOptions {
  window_strides:[long];
  padding:[long];
  lhs_dilation:[long];
  rhs_dilation:[long];
  window_reversal:[bool];

  //following is expanded ConvDimensionNumbersAttr
  input_batch_dimension:long;
  input_feature_dimention:long;
  input_spatial_dimensions:[long];
  kernel_input_feature_dimension:long;
  kernel_output_feature_dimension:long;
  kernel_spatial_dimensions:[long];
  output_batch_dimension:long;
  output_feature_dimension:long;
  output_spatial_dimensions:[long];

  feature_group_count:long;
  batch_group_count:long;
}

table MaximumOptions {
}

table MinimumOptions {
}

table ReshapeOptions {
}

table DivideOptions {
}

table MultiplyOptions {
}

table ReduceOptions {
  dimensions:[long];
  // computation points to another subgraph in the model
  computation:int;
}

table ReduceWindowOptions {
  window_dimension:[long];
  window_strides:[long];
  base_dilations:[long];
  window_dilations:[long];
  padding:[long];
  // computation points to another subgraph in the model
  computation:int;
}

table BroadcastInDimOptions {
  broadcast_dimensions:[long];
}

table LogisticOptions {
}

table CustomCallOptions {
  call_target_name:string;
  backend_config:[ubyte];
}

table BatchNormInferenceOptions {
  epsilon:float;
  feature_index:long;
}

table ClampOptions {
}

table SliceOptions {
  start_indices:[long];
  limit_indices:[long];
  strides:[long];
}

table ConcatenateOptions {
  dimension:long;
}

table IotaOptions {
  iota_dimension:long;
}

table SubtractOptions {
}

table CeilOptions {
}

table ConvertOptions {
}

table GatherOptions {
  slice_sizes:[long];
  indices_are_sorted:bool;

  //following is expanded GatherDimensionNumbersAttr
  offset_dims:[long];
  collapsed_slice_dims:[long];
  start_index_map:[long];
  index_vector_dim:long;
}

table AbsOptions {
}

table DotGeneralOptions {
  //following is expanded DotDimensionNumbersAttr
  lhs_batching_dimensions:[long];
  rhs_batching_dimensions:[long];
  lhs_contracting_dimensions:[long];
  rhs_contracting_dimensions:[long];
}

table ResizeBilinearOptions {
  align_corners: bool;
  half_pixel_centers: bool;
}

table Operator {
  opcode_index:uint;

  // Optional inputs/outputs are indicated by -1.
  inputs:[int];
  outputs:[int];

  operator_options:OperatorOptions;
}

// The root type, defining a subgraph, which typically represents an entire
// model.
table SubGraph {
  // A list of all tensors used in this subgraph.
  tensors:[Tensor];

  // Indices of the tensors that are inputs into this subgraph. Note this is
  // the list of non-static tensors that feed into the subgraph for inference.
  inputs:[int];

  // Indices of the tensors that are outputs out of this subgraph. Note this is
  // the list of output tensors that are considered the product of the
  // subgraph's inference.
  outputs:[int];

  // All operators, in execution order.
  operators:[Operator];

  // Name of this subgraph (used for debugging).
  name:string;
}

// Table of raw data buffers (used for constant tensors). Referenced by tensors
// by index. The generous alignment accommodates mmap-friendly data structures.
table Buffer {
  data:[ubyte] (force_align: 16);
}


table Model {
  // Version of the schema.
  version:uint;

  // A list of all operator codes used in this model. This is
  // kept in order because operators carry an index into this
  // vector.
  operator_codes:[OperatorCode];

  // All the subgraphs of the model. The 0th is assumed to be the main
  // model.
  subgraphs:[SubGraph];

  // Buffers of the model.
  // Note the 0th entry of this array must be an empty buffer (sentinel).
  // This is a convention so that tensors without a buffer can provide 0 as
  // their buffer.
  buffers:[Buffer];

}

root_type Model;