// Copyright (c) 2019, Apple Inc. All rights reserved.
//
// Use of this source code is governed by a BSD-3-clause license that can be
// found in LICENSE.txt or at https://opensource.org/licenses/BSD-3-Clause

/*
 * - A Program is the container with following information
 *     - set of functions: Function defines a program block to be executed
 *     - A model can have multiple functions defined and will have a single
 * point of entry.
 * - A Function consists of
 *     - List of named inputs and output types
 *     - A block defining scope for a function - similar to a function in C/C++
 * - A Block consists of
 *     - List of named inputs and output names
 *     - Topologically sorted Ops
 * - A Op consists of
 *     - List of named inputs and outputs (name, type) pair
 *     - Optionally, blocks for Control-Flow
 *
 * Programs, functions, blocks, ops, and tensor types all can contain an
 * optional set of attributes.
 *
 * == Identifiers ==
 * Identifiers, generally used for names and keys, must match the
 * regular expression [A-Za-z\_][A-Za-z0-9\_@]*
 */

syntax = "proto3";
option optimize_for = LITE_RUNTIME;

package CoreML.Specification.MILSpec;

// The top level container.
message Program {
  int64 version = 1;

  // Must be unique within the containing program
  // Names must be valid identifiers as described above.
  map<string, Function> functions = 2;

  string docString = 3;

  // Any other attributes not described by other fields.
  // Keys must be valid identifiers as described above.
  map<string, Value> attributes = 4;
}

// A program-level function.
message Function {
  // Function inputs are unordered (name, ValueType) pairs.
  // Inputs intended to process images must be rank-4 Float32 tensors.
  // Dimensions are interpreted as NCHW, with N == 1 and C being 1 for grayscale
  // and 3 for RGB. Names must be valid identifiers as described above.
  repeated NamedValueType inputs = 1;

  // The active block is drawn from this named specialization.
  // This key must exist in `block_specializations`.
  string opset = 2;

  // Named specializations of this function.
  //
  // Specialization keys are the name of the opset that the
  // function specialization is written in. They must be valid
  // identifiers as described above.
  //
  // Outputs from all blocks must match. They define the outputs
  // of the function.
  // Each block inherits the lexical scope from the function.
  map<string, Block> block_specializations = 3;

  // Any other attributes not described by other fields.
  // Keys must be valid identifiers as described above.
  map<string, Value> attributes = 4;
}

// A basic block with a single entry and exit in SSA form.
message Block {
  // Infrequently used, these are for operators that may need to give
  // block-local names to input values (e.g. while_loop).
  repeated NamedValueType inputs = 1;

  // The names to give to values returned by this block. They must be
  // identifiers as described above.
  //
  // ValueType of outputs[i] is Operation[j].outputs[k].type where
  // i, j and k are indices of block output, block Operation and
  // jth operation's output respectively.
  // this is due to
  // 1. An operation can have more than one output
  // 2. Any one of operation's output could be potentially block's output
  repeated string outputs = 2;

  repeated Operation operations = 3;

  // Any other attributes not described by other fields.
  // Keys must be valid identifiers as described above.
  map<string, Value> attributes = 4;
}

// Argument is list of Binding to either name or value
message Argument {
  message Binding {
    oneof binding {
      // The name of a previously defined value.
      string name = 1;

      // A compile time constant.
      Value value = 2;
    }
  }

  repeated Binding arguments = 1;
};

// A single operation/node/layer.
message Operation {
  // Examples: "convolution", "cropResize". Operation type defines the
  // expected inputs and output.
  string type = 1;

  // Operator arguments
  //
  // Key: parameter name
  // Value: Argument (list of bindings)
  //
  // Value is list of argument binding to given parameter
  // Binding can be a string name (previous operation output or input given to
  // model/block/function)
  //             or a Value (known compile time value for given operation)
  // Argument can be of length 1 (general) or variable length (e.g. concat
  // layer) e.g. {'stride' : ['input_01']} e.g. {'x' : ['input_01', 'input_02',
  // 'input_03', false]}
  map<string, Argument> inputs = 2;

  // Names to which to bind values returned by this operation.
  // Names must be:
  //  (*) valid identifiers as described above; and
  //  (*) unique within the current scope.
  repeated NamedValueType outputs = 3;

  // Nested blocks for loops and conditionals. For example,
  // a conditional block will have two entries here.
  repeated Block blocks = 4;

  // Any other information not captured by other fields.
  // Keys must be valid identifiers as described above.
  map<string, Value> attributes = 5;
}

// Named Value parameters
// (name, type) pair
message NamedValueType {
  // The name of this parameter; must be a valid identifier as described above.
  string name = 1;

  // This parameter's required type.
  ValueType type = 2;
}

/* ========  Types ======= */

// Primer: Two fundamental representations of state:
//
// Variable: Variables are NEVER materialized at compile time and are only
// available at run time. Therefore, for Variables we only have ValueType,
// which may have unknown shapes in the IR. Variable encompasses familiar
// concepts such as placeholder, output of an Op.
//
// Value: Values are ALWAYS materialized at compile time, and MAY be modified
// at runtime (e.g., during on-device training). Value describes notions
// such as parameter, attributes of an op. Value is either stored inside
// proto (e.g., attributes) or outside of proto (e.g. parameters) and
// NEVER contains unknown shape in the IR.
//
// Comment(daviddai): A Variable with the potential to be materialized at
// compile time (e.g., through constant propagation) does NOT preclude it to
// be a Variable. Certain Ops such as LoadParameter and Const, their output
// has potential to be materialized at compile time but is still represented
// as Variable.

// A type of any kind
message ValueType {
  oneof type {
    TensorType tensorType = 1;
    ListType listType = 2;
    TupleType tupleType = 3;
    DictionaryType dictionaryType = 4;
    StateType stateType = 5;
  }
}

// Supported data types
enum DataType {
  // Comment: Two schemes of specifying field id: just start with 0
  // without reserving numbers, but keep track of the next field ID. The
  // other is assign blocks of ID to int / float / uint etc.

  // 0-10 reserved for special types
  UNUSED_TYPE = 0;  // not currently in use
  BOOL = 1;
  STRING = 2;  // arbitrary sequence of bytes

  // Floats
  FLOAT8E4M3FN = 40;
  FLOAT8E5M2 = 41;
  FLOAT16 = 10;
  FLOAT32 = 11;
  FLOAT64 = 12;
  BFLOAT16 = 13;

  // Ints
  INT8 = 21;
  INT16 = 22;
  INT32 = 23;
  INT64 = 24;
  INT4 = 25;

  // UInts
  UINT8 = 31;
  UINT16 = 32;
  UINT32 = 33;
  UINT64 = 34;

  UINT4 = 35;
  UINT2 = 36;
  UINT1 = 37;
  UINT6 = 38;
  UINT3 = 39;
}

message TensorType {
  // The data type stored in a tensor of this type
  DataType dataType = 1;

  // The number of dimensions in the tensor shape. rank == -1 implies
  // variable (not fixed) rank
  int64 rank = 2;

  // Tensor shape values; must be of length "rank"
  repeated Dimension dimensions = 3;

  // Any other tensor type attributes not described by other fields.
  // Keys must be valid identifiers in MIL text syntax.
  map<string, Value> attributes = 4;
}

message TupleType {
  // Recursively define TupleType from ValueType.
  repeated ValueType types = 1;
}

message ListType {
  // The type of element stored in a list of this type
  ValueType type = 1;

  // The number of elements in a list of this type. May be unknown (variable
  // length)
  Dimension length = 2;
}

// An unordered key-value mapping
message DictionaryType {
  ValueType keyType = 1;
  ValueType valueType = 2;
}

message StateType {
  ValueType wrappedType = 1;
}

message Dimension {
  oneof dimension {
    ConstantDimension constant = 1;
    UnknownDimension unknown = 2;
  }

  message ConstantDimension {
    uint64 size = 1;
  }

  message UnknownDimension {
    bool variadic = 1;
  }
}

/* ======== Values ======= */

// See Variable vs Value primer above.
message Value {
  string docString = 1;  // optional human-readable texts.
  ValueType type = 2;

  // An immediate value stored within the proto
  message ImmediateValue {
    oneof value {
      TensorValue tensor = 1;
      TupleValue tuple = 2;
      ListValue list = 3;
      DictionaryValue dictionary = 4;
    }
  }

  // Reference to a "blob v2" storage file
  message BlobFileValue {
    // name of file
    string fileName = 1;

    // byte offset to metadata
    uint64 offset = 2;
  }

  oneof value {
    ImmediateValue immediateValue = 3;
    BlobFileValue blobFileValue = 5;
  }
}

message TensorValue {
  oneof value {
    RepeatedFloats floats = 1;
    RepeatedInts ints = 2;
    RepeatedBools bools = 3;
    RepeatedStrings strings = 4;
    RepeatedLongInts longInts = 5;
    RepeatedDoubles doubles = 6;
    RepeatedBytes bytes = 7;
  }

  message RepeatedFloats {
    repeated float values = 1 [packed = true];
  }

  message RepeatedDoubles {
    repeated double values = 1 [packed = true];
  }

  message RepeatedInts {
    repeated int32 values = 1 [packed = true];
  }

  message RepeatedLongInts {
    repeated int64 values = 1 [packed = true];
  }

  message RepeatedBools {
    repeated bool values = 1 [packed = true];
  }

  message RepeatedStrings {
    repeated string values = 1;
  }

  message RepeatedBytes {
    bytes values = 1;
  }
}

message TupleValue {
  // Comment: TupleValue is recursively defined from Value.
  repeated Value values = 1;
}

message ListValue {
  repeated Value values = 1;
}

message DictionaryValue {
  message KeyValuePair {
    Value key = 1;
    Value value = 2;
  }
  repeated KeyValuePair values = 1;
}