// LINT: LEGACY_NAMES

syntax = "proto3";

package tensorflow.profiler;

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

// Predicted computational cost of the instruction associated with the symbol.
// Estimated by traversing the HLO graph.
message PerformanceInfo {
  // The number of floating-point operations computed.
  optional int64 flops = 1;
  // The sum of bytes read and bytes written.
  optional int64 bytes_accessed = 2;
  // Breakdown of memory accessed by read/write and memory space.
  message MemoryAccessed {
    optional bool is_read = 1;
    enum MemorySpace {
      UNKNOWN = 0;
      HBM = 1;
      CMEM = 2;
      VMEM = 3;
    }
    optional MemorySpace memory_space = 2;
    optional int64 bytes_accessed = 3;
  }
  repeated MemoryAccessed memory_accessed_breakdown = 3;

  reserved 4;
}

// Types of memory bandwidth we track in the system.
enum MemBwType {
  option allow_alias = true;

  // We use FIRST and LAST enum values to be able to iterate over this enum
  // in TypeScript, since the _MIN and _MAX values are not automatically
  // available as in C++.
  MEM_BW_TYPE_FIRST = 0;
  // Aggregated BW across on-chip and off-chip memory.
  // For GPU, 1/2 is shared memory bandwisth.
  MEM_BW_TYPE_HBM_RW = 0;
  // On-chip memory read bw.
  MEM_BW_TYPE_SRAM_RD = 1;
  // On-chip memory write bw.
  MEM_BW_TYPE_SRAM_WR = 2;
  MEM_BW_TYPE_CMEM_RD = 3;
  MEM_BW_TYPE_CMEM_WR = 4;
  MEM_BW_TYPE_VMEM_RD = 5;
  MEM_BW_TYPE_VMEM_WR = 6;

  // Leave last.
  // Leave this MAX unchanged now to avoid op profile changes.
  // TODO(b/359279074) Revisit the memory breakdown in op profile since we have
  // more memory types now.
  MEM_BW_TYPE_MAX = 2;
}

// Tensorflow generic memory space names.
// These space names are used in analysis code to get memory bandwidth per core.
enum MemorySpace {
  MEMORY_SPACE_UNDEFINED = 0;
  // Off-chip memory.
  // Assume all backends use 1 for HBM/off-chip memory.
  MEMORY_SPACE_HBM = 1;
  // On-chip memory.
  MEMORY_SPACE_ON_CHIP = 0x7FFFFFFE;
  // Any memory.
  MEMORY_SPACE_ALL = 0x7FFFFFFF;
}

// What the dimension represents, e.g. spatial, feature or batch.
enum LayoutDimensionSemantics {
  UNKNOWN_SEMANTICS = 0;
  FEATURE = 1;
  BATCH = 2;
  SPATIAL = 3;
}

// Data layout of an op.
message LayoutAnalysis {
  // Physical data layout in each tensor dimension.
  message Dimension {
    // Size of the data in this dimension.
    int32 size = 1;
    // Data must be padded to a multiple of alignment.
    int32 alignment = 2;
    // What the dimension represents.
    LayoutDimensionSemantics semantics = 3;
  }
  // The physical data layout, from most-minor to most-major dimensions.
  repeated Dimension dimensions = 1;
}

// A container to serialize this repeated field in "symbolized xplane."
message MemoryAccessBreakdown {
  repeated OpMetrics.MemoryAccessed memory_accessed = 1;
}

// Metrics for an operation (accumulated over all occurrences).
// Next ID: 33
message OpMetrics {
  // HLO module id. 0 for Framework ops.
  uint64 hlo_module_id = 13;
  // A deterministic and unique identifier of the op within the HLO module.
  // Unlike `name`, which is only useful as an ID within the compiling task,
  // this ID is consistent across compilations because the module is
  // canonicalized during the compilation process.
  optional int64 stable_op_id = 32;
  // Name of this op.
  string name = 6;
  // Long name of this op (e.g., HLO expression).
  string long_name = 20;
  // Category of this op. (e.g. Hlo op category, Framework op type, input
  // pipeline stage category) Could be parsed from provenance if it is a
  // framework op.
  string category = 11;
  // Provenance of this op if it is an HLO Op. (e.g. TF Op name, JAX Op name)
  // TODO(b/310434797) Extends this for JAX as now only TF Op is populated.
  string provenance = 12;
  // Whether it is executed eagerly.
  bool is_eager = 18;
  // Number of executions.
  uint32 occurrences = 3;
  // Total time (self + children) in picoseconds.
  uint64 time_ps = 7;
  // Total time (self + children) in picoseconds needs to be spent if running on
  // the default pstate.
  uint64 normalized_time_ps = 27;
  // Minimum time (self + children) among all occurrences.
  uint64 min_time_ps = 17;
  // Total self time in picoseconds.
  uint64 self_time_ps = 1;
  // Deprecated. Use flops_v2 instead.
  uint64 flops = 2 [deprecated = true];
  // Deprecated. Use model_flops_v2 instead.
  uint64 model_flops = 24 [deprecated = true];
  // Total FLOPs. Normalized to the devices peak bandwidth.
  double flops_v2 = 29;
  // Total FLOPs for the model. Can be 0, in which case assume it's same as
  // flops
  double model_flops_v2 = 30;
  // Total VDD energy in Joules.
  optional double vdd_energy_j = 31;
  // Fingerprint of the symbol (cs/xla::HloPrintOptions::Fingerprint),
  // if 0, the fingerprint is not set.
  uint64 fingerprint = 25;
  // Total bytes accessed.
  uint64 bytes_accessed = 5;
  // Breakdown of memory accessed by operation type and memory space.
  message MemoryAccessed {
    enum OperationType {
      UNKNOWN = 0;
      READ = 1;
      WRITE = 2;
    }
    OperationType operation_type = 1;
    // Device-specific id of memory space.
    uint64 memory_space = 2;
    uint64 bytes_accessed = 3;
  }
  repeated MemoryAccessed memory_accessed_breakdown = 19;
  // Total dma stall time in picoseconds.
  uint64 dma_stall_ps = 10;
  // The data layout for this op. Only set for convolution ops for now.
  LayoutAnalysis layout = 14;
  // Deduplicated HLO name for this op. Not set for TF ops.
  string deduplicated_name = 15;
  // Children of the op. e.g. fused ops if this op is fusion.
  OpMetricsDb children = 16;
  // Number of cores this op occurs.
  uint32 num_cores = 21;
  // Computation primitive size in BITS. This is the size of the type of the
  // hardware computation. In the future this may be extended to include info
  // such as signed/unsigned, int/fp, etc. Currently only the size is needed.
  uint32 computation_primitive_size = 22;
  // Whether the op is autotuned.
  bool autotuned = 23;
  optional SourceInfo source_info = 26;
  enum TpuCoreType {
    UNKNOWN = 0;
    TENSOR_CORE = 1;
    SPARSE_CORE = 2;
  }
  TpuCoreType core_type = 28;
  reserved 4, 8, 9;
}

// Statistics about the various precision used in computation.
message PrecisionStats {
  // Amount of time spent on 16-bit computation (in ps).
  uint64 compute_16bit_ps = 1;
  // Amount of time spent on 32-bit computation (in ps).
  uint64 compute_32bit_ps = 2;
}

// A database for OpMetrics.
// Next ID: 19
message OpMetricsDb {
  // A bunch of OpMetrics.
  repeated OpMetrics metrics_db = 10;
  // The total host infeed-enqueue duration in picoseconds.
  uint64 total_host_infeed_enq_duration_ps = 2;
  // The total of the difference between the start times of two
  // consecutive infeed-enqueues (per host) in picoseconds.
  uint64 total_host_infeed_enq_start_timestamp_ps_diff = 3;
  // The total time in picoseconds.
  uint64 total_time_ps = 11;
  // The total time incurred by OPs in picoseconds.
  uint64 total_op_time_ps = 12;
  // The weighted average total op time in picoseconds if running all ops on the
  // default pstate.
  uint64 normalized_total_op_time_ps = 16;
  // Precision-related stats.
  PrecisionStats precision_stats = 13;
  // The below two stats will be different from the total time ps and total op
  // time ps because they are unioned all cores (and not summed).
  // For duty cycle, a device is idle if all the cores are idle.
  uint64 idle_time_ps = 14;
  // For duty cycle, a device is busy if any of the cores is busy.
  uint64 busy_time_ps = 15;
  // For duty cycle high confidence, including all custom calls.
  uint64 busy_time_high_confidence_ps = 17;
  // Idle time for high confidence duty cycle.
  uint64 idle_time_high_confidence_ps = 18;
  reserved 1, 4, 5, 6, 7, 8, 9;
}
