// This proto describes the format of the output profile file from
// the HLO stats tool.
syntax = "proto2";

package tensorflow.profiler.hlo_stats;

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

// A database of HloStats records.
message HloStatsDatabase {
  // All HloStats records, one for each HLO operation.
  repeated HloStatsRecord hlo_stats_record = 1;
}

// There is one HloStatsRecord for each HLO operation profiled.
// Next ID: 44
message HloStatsRecord {
  // The rank by self time
  optional uint64 rank = 1;

  // program_id for this op
  optional uint64 program_id = 30;

  // The HLO category name.
  optional string hlo_category = 17;

  // The HLO expression.
  optional string hlo_expression = 2;

  // The framework op name (TF Op, JAX Op)
  optional string tf_op_name = 21;

  // Number of occurrences of the operation.
  optional int64 occurrences = 3;

  // Total "accumulated" time in micro-seconds that the operation
  // took. If this operation has any children operations,
  // the "accumulated" time includes the time spent inside children.
  optional double total_time_in_us = 4;

  // Average "accumulated" time in micro-seconds that each
  // occurrence of the operation took.
  optional double avg_time_in_us = 5;

  // Total "self" time in micro-seconds that the operation took.
  // If this operation has any children operations, the "self" time
  // doesn't include the time spent inside children.
  optional double total_self_time_in_us = 6;

  // Average "self" time in micro-seconds that the operation took.
  optional double avg_self_time_in_us = 7;
  optional double total_self_time_as_fraction = 8;
  optional double cumulative_total_self_time_as_fraction = 9;

  // Percentage of the total "accumulated" time that was caused by
  // DMA stall.
  optional double dma_stall_fraction = 10;

  // Total floating-point operations (FLOPs) performed per second normalized to
  // the bf16 peak capacity.
  optional double measured_flop_rate = 13;

  // Total Floating-point operations for the op per second.
  optional double model_flop_rate = 34;

  // Number of total bytes (including both read and write) accessed per second.
  optional double measured_memory_bw = 14;

  // Number of bytes accessed from HBM (including both read and write) per
  // second.
  optional double hbm_bw = 22;

  // Number of bytes read from CMEM per second.
  optional double cmem_read_bw = 23;

  // Number of bytes written to CMEM per second.
  optional double cmem_write_bw = 24;

  // Number of bytes read from VMEM per second.
  optional double vmem_read_bw = 35;

  // Number of bytes written to VMEM per second.
  optional double vmem_write_bw = 36;

  // Overall operational intensity in FLOP/Byte.
  optional double operational_intensity = 15;

  // Operational intensity based on HBM in FLOP/Byte.
  optional double hbm_operational_intensity = 26;

  // Operational intensity based on CMEM read in FLOP/Byte.
  optional double cmem_read_operational_intensity = 27;

  // Operational intensity based on CMEM write in FLOP/Byte.
  optional double cmem_write_operational_intensity = 28;

  // Operational intensity based on VMEM read in FLOP/Byte.
  optional double vmem_read_operational_intensity = 37;

  // Operational intensity based on VMEM write in FLOP/Byte.
  optional double vmem_write_operational_intensity = 38;

  // Operational intensity based on the bottleneck resource in FLOP/Byte.
  optional double bottleneck_operational_intensity = 29;

  // Whether this operation is "Compute", "HBM", "CMEM Read", "CMEM Write"
  // bound, according to the Roofline Model.
  optional string bound_by = 16;

  // Whether this operation is for HLO or Framework rematerialization.
  optional bool rematerialization = 20;

  // Whether this op is for outside compilation.
  optional bool outside_compilation = 25;

  // Whether this op is autotuned.
  optional bool autotuned = 31;

  // Deprecated. Use flops_v2 instead.
  optional uint64 flops = 32 [deprecated = true];

  // Flops for the record
  optional double flops_v2 = 40;

  // Bytes accessed for the record
  optional uint64 bytes_accessed = 33;

  // Infrmation about the corresponding source code.
  optional SourceInfo source_info = 39;

  // The TPU core type of the op.
  optional string core_type = 41;

  // The op name of the parent op if this op is a child op. Ops without a parent
  // op are assumed to be at the top level.
  optional string parent_op_name = 42;

  // VDD energy in Joules.
  optional double vdd_energy_j = 43;

  reserved 11, 12, 18, 19;
}
