syntax = "proto3";

package tensorflow.data.model;

option cc_enable_arenas = true;

// Class of a node in the performance model.
enum NodeClass {
  UNKNOWN = 0;
  INTERLEAVE_MANY = 1;
  ASYNC_INTERLEAVE_MANY = 2;
  KNOWN_RATIO = 3;
  ASYNC_KNOWN_RATIO = 4;
  UNKNOWN_RATIO = 5;
}

// Algorithm used for model autotuning optimization.
enum AutotuneAlgorithm {
  HILL_CLIMB = 0;
  GRADIENT_DESCENT = 1;
}

// Protocol buffer representing the data used by the autotuning modeling
// framework.
message ModelProto {
  // General representation of a node in the model.
  message Node {
    // Unique node ID.
    int64 id = 1;

    // Human-readable name of the node.
    string name = 2;

    // An indication whether autotuning is enabled for this node.
    bool autotune = 3;

    // The number of bytes stored in this node's buffer.
    int64 buffered_bytes = 4;

    // The number of elements stored in this node's buffer.
    int64 buffered_elements = 5;

    // The number of bytes consumed by the node.
    int64 bytes_consumed = 6;

    // The number of bytes produced by the node.
    int64 bytes_produced = 7;

    // The number of elements produced by the node.
    int64 num_elements = 8;

    // The aggregate processing time spent in this node.
    int64 processing_time = 9;

    // An indication whether this node records metrics about produced and
    // consumed elements.
    bool record_metrics = 10;

    // Represents a node parameter.
    message Parameter {
      // Human-readable name of the parameter.
      string name = 1;

      // Identifies the model value of the parameter. This can be different from
      // the actual value (e.g. during optimization search).
      double value = 2;

      // The actual value of the parameter.
      double state_value = 3;

      // Minimum value of the parameter.
      double min = 4;

      // Maximum value of the parameter.
      double max = 5;

      // Identifies whether the parameter should participate in autotuning.
      bool tunable = 6;
    }

    // Parameters of this node.
    repeated Parameter parameters = 11;

    // Statistic of inputs processing time history.
    double input_processing_time_sum = 12;
    int64 input_processing_time_count = 13;

    // Inputs of this node.
    repeated Node inputs = 14;

    // Class of this node.
    NodeClass node_class = 15;

    // Ratio of input to output elements. This is only used by KNOWN_RATIO and
    // ASYNC_KNOWN_RATIO nodes.
    double ratio = 16;

    // Ratio identifies how many parallelism calls are introduced by one
    // buffered element. This is only used by ASYNC_KNOWN_RATIO nodes.
    double memory_ratio = 17;
  }

  // Output node of this model.
  Node output = 1;

  // Counter for node IDs of this model.
  int64 id_counter = 2;

  // Indicates whether the modeling framework should collect resource usage,
  // e.g. CPU, memory.
  bool collect_resource_usage = 3;

  // Contains parameters of the model autotuning optimization.
  message OptimizationParams {
    // Algorithm used for autotuning optimization.
    AutotuneAlgorithm algorithm = 1;

    // Number of available logical threads.
    int64 cpu_budget = 2;

    // Amount of available memory in bytes.
    int64 ram_budget = 3;

    // Time between two consecutive `GetNext` calls to the iterator represented
    // by the output node.
    double model_input_time = 4;
  }

  OptimizationParams optimization_params = 4;
}
