// Copyright (c) Meta Platforms, Inc. and affiliates.

namespace fb_xnnpack;

// datatype for xnn-values
enum XNNDatatype : short {
  /// Invalid data type. Valid Values never have this datatype.
  xnn_datatype_invalid = 0,
  /// IEEE754 single-precision floating-point.
  xnn_datatype_fp32 = 1,
  /// IEEE754 half-precision floating-point.
  xnn_datatype_fp16 = 2,
  /// Quantized 8-bit signed integer with shared per-Value quantization parameters.
  xnn_datatype_qint8 = 3,
  /// Quantized 32-bit signed integer with shared per-Value quantization parameters.
  xnn_datatype_qint32 = 4,
}

// taken from executorch
// Data buffer abstraction.
table Buffer {
  storage:[ubyte] (force_align: 16);
}

table XNNTensorValue {
  // type of the tensor elements.
  datatype:XNNDatatype;
  // number of dimensions in the shape.
  num_dims:uint;
  // pointer to an array of @a num_dims shape dimensions. If num_dims is 0, this pointer can be NULL.
  // XNNPACK does not keep any pointers to this array after the function returns.
  dims:[uint];
  // Index to the program's constant buffer table, value 0 is reserved to indicate non constant
  constant_buffer_idx:uint;
  // external ID for the Value. The ID must be within the range of reserved Value IDs specified on
  // the Subgraph creation. If the external ID is XNN_INVALID_VALUE_ID, an internal ID will be
  // created for the Value.
  external_id:uint;
  // binary features of the Value. Supported values are any combination of XNN_VALUE_FLAG_EXTERNAL_INPUT
  // and XNN_VALUE_FLAG_EXTERNAL_OUTPUT.
  flags:uint;
  // pointer to the variable that will be initialized with the Value ID upon successful return. If a
  // valid @a external_id was provided, the variable will be initialized with the @a external_id value.
  id_out:uint;
}

union XNodeUnion {
  XNNAdd,
}

union XValueUnion {
  XNNTensorValue,
}

table XNode {
  xnode:XNodeUnion;
  // An int which can be linked back to the node in the origin graph
  debug_handle:uint;
}

table XValue {
  xvalue:XValueUnion;
}

table XNNAdd {
  input1_id:uint;
  input2_id:uint;
  output_id:uint;
  flags:uint;
}

table XNNGraph {
  // Schema version.
  version:string;
  xnodes:[XNode];
  xvalues:[XValue];

  // Number of external inputs/outputs
  num_externs:uint;

  // Ids of external inputs
  input_ids:[uint];

  // Ids of external outputs
  output_ids:[uint];

  // Tables of constant data, used for constant Values (e.g.
  // data field of weight tensors). Each constant is assigned an index into the table
  // which are each individually aligned. 0 index is reserved to be pointed to by non-constant
  // Tensors
  constant_buffer:[Buffer];

  // the list index is memory buffer id, the value is the memory buffer size.
  mem_buffer_sizes: [uint];
}

root_type XNNGraph;