syntax = "proto3";

package tensorflow.profiler.op_profile;

import "plugin/xprof/protobuf/source_info.proto";

// Profile is the top-level data that summarizes a program.
message Profile {
  reserved 2;
  reserved "by_program_structure";
  reserved 3;
  reserved "per_program";
  // Root of a profile broken down by instruction category.
  Node by_category = 1;
  // Root of a profile broken down by program.
  Node by_program = 4;

  // Device type.
  string device_type = 5;

  // DVFS time scale multiplier weighted aggregated across all ops.
  // Multiplied it with the raw time to get estimated time as if running on
  // default pstate.
  double agg_dvfs_time_scale_multiplier = 10;

  // Exclude idle ops.
  Node by_category_exclude_idle = 6;
  Node by_program_exclude_idle = 7;
  // Root of a profile broken down by program and provenance.
  Node by_provenance = 8;
  Node by_provenance_exclude_idle = 9;
}

// An entry in the profile tree. (An instruction, or set of instructions).
message Node {
  string name = 1;             // Semantics depend on contents.
  Metrics metrics = 2;         // May be omitted e.g. for fused instructions.
  repeated Node children = 3;  // Subjected to pruning.

  // Details about what this node represents.
  oneof contents {
    InstructionCategory category = 4;
    XLAInstruction xla = 5;
  }

  int32 num_children = 6;  // Total number of children before pruning.
  // A category of XLA instructions.
  // name is a descriptive string, like "data formatting".
  message InstructionCategory {}
  // A single XLA instruction.
  // name is the unique instruction id, like "%multiply.5".
  message XLAInstruction {
    string op = 1;          // Opcode like %multiply
    string expression = 2;  // %multiply = [shape]multiply(operand1, operand2)
    string provenance = 3;  // Provenance op name, eg. TF Op name, JAX Op name
    string category = 4;
    // Describes the physical memory layout of the instruction's primary input.
    // e.g. for a convolution, this analyzes the image and ignores the kernel.
    LayoutAnalysis layout = 5;
    uint32 computation_primitive_size = 6;
    uint64 fingerprint = 7;
    uint64 program_id = 8;
    optional SourceInfo source_info = 9;
    message LayoutAnalysis {
      // The physical data layout, from most-minor to most-major dimensions.
      repeated Dimension dimensions = 1;
      message Dimension {
        int32 size = 1;       // Size of the data in this dimension.
        int32 alignment = 2;  // Data must be padded to a multiple of alignment.
        string semantics = 3;  // What the dimension represents, e.g. "spatial".
      }
    }
    // Xprof kernel metadata extracted from the HLO expression's frontend
    // attributes.
    optional string xprof_kernel_metadata = 10;
  }
}

// Measurements of an operation (or aggregated set of operations).
// Metrics are always "total" rather than "self".
// Next ID: 21
message Metrics {
  // Floating point computations performed by this operation, as a fraction of
  // peak core FLOPS * program time. This representation has useful properties:
  //  - it is proportional to the number of floating point operations performed
  //  - utilization is flops/time
  //  - wasted potential flops is proportional to time - flops
  //  - it does not reveal the peak core FLOPS of the hardware
  //  - it is capped at 1.0
  double flops = 2;

  // Uncapped FLOPs.
  double uncapped_flops = 19;

  // The memory bandwidth used to load operands, as a fraction of
  // thereotical memory bandwidth on the specific hardware.
  // Index into array using MemBwType enum.
  repeated double bandwidth_utils = 5;

  // The raw stats below are aggregated across all occurrences.
  double raw_time = 11;   // Elapsed core-time in picoseconds.
  double raw_flops = 12;  // Total floating-point operations performed.
  // Total FLOPs Normalized to the bf16 (default) devices peak bandwidth.
  double bf16_flops = 18;
  // Total time in picoseconds normalized.
  double normalized_time_ps = 20;
  // Total bytes accessed for each memory type.
  // Index into array using MemBwType enum.
  repeated double raw_bytes_accessed_array = 15;
  // Number of executions.
  uint32 occurrences = 16;
  // Average "accumlated" time in picoseconds that the operation took.
  double avg_time_ps = 17;

  reserved 1, 3, 4, 13, 14;
}
