syntax = "proto3";

package tensorflow.profiler;

import "google/protobuf/any.proto";
import "plugin/xprof/protobuf/op_metrics.proto";
import "plugin/xprof/protobuf/flat_op_metrics.proto";

// Breakdown of step-time on generic hardware. Note that these components are
// mutually exclusive so that adding them together is equal to the step time. If
// an execution time interval has multiple types of event happening, we need to
// pick one of the event type to attribute the time interval to.
message GenericStepBreakdown {
  // Map event type to the accumulated duration in
  // picoseconds of that type.
  map<int32, uint64> type_ps = 1;

  // Map of string category to accumulated duration in picoseconds for
  // that category.
  map<string, uint64> category_ps = 2;
}

// Breakdown of step-time on TPU.
// Next ID: 20
message TpuStepBreakdown {
  // The infeed duration (host to TensorCore) in picoseconds.
  uint64 infeed_duration_ps = 1;

  // The outfeed duration (TensorCore to host) in picoseconds.
  uint64 host_outfeed_ps = 2;

  // The TensorCore time that is waiting for SparseCoreV0 in picoseconds.
  uint64 wait_for_scv0_duration_ps = 3;

  // The TensorCore time spent transforming activations in SparseCoreV0 layout
  // into XLA layout.
  uint64 scv0_infeed_transform_ps = 4;

  // The outfeed duration (TensorCore to SparseCoreV0) in picoseconds.
  uint64 scv0_outfeed_ps = 5;

  // The time spent on all-reduce (used to be cross-replica-sum) in picoseconds.
  uint64 crs_duration_ps = 6;

  // The percentage of the SparseCoreV0 time that spends on infeed from host
  // (including both data and instruction).
  double scv0_infeed_percent = 7;

  // The time spent on send operation.
  uint64 send_duration_ps = 8;

  // The time spent on recv operation.
  uint64 recv_duration_ps = 9;

  // The time spent on host send operation.
  uint64 host_send_duration_ps = 15;

  // The time spent on host recv operation.
  uint64 host_recv_duration_ps = 16;

  // Megacore fusion runs different operations on each core, e.g., a convolution
  // on one core and an all-reduce on the other core. This is the time that the
  // core executing the faster operation waits for the core executing the slower
  // operation to reach the synchronization point.
  uint64 wait_for_megacore_fusion_peer_duration_ps = 14;

  // The time waiting for overlay DMAs in picoseconds.
  uint64 overlay_wait_duration_ps = 11;

  // The time spent running high flops ops, such as convolution and output
  // fusion.
  uint64 high_flops_compute_ps = 12;

  // The time that the Tensorcore is idle but not waiting for input or
  // SparseCoreV0.
  uint64 tc_idle_ps = 13;

  // The TensorCore time that is busy in picoseconds.
  uint64 tc_busy_ps = 17;

  // The SparseCoreV0 time that is busy in picoseconds (equal to
  // SparseCoreV0 time - HOST_INSTRUCTION_STALL - HOST_DATA_STALL -
  // TENSOR_CORE_STALL).
  uint64 scv0_busy_ps = 18;

  // SparseCoreV0 step time in picoseconds (equal to SparseCoreV0 time -
  // TENSOR_CORE_STALL).
  uint64 scv0_step_ps = 19;

  reserved 10;
}

// Breakdown of step-time on SparseCore.
message SparseCoreStepBreakdown {
  // SparseCore step time in picoseconds (equal to SparseCore time - sc_idle -
  // sc_wait_time).
  uint64 sc_compute_ps = 1;

  // Host to sparse core time in picoseconds.
  uint64 sc_infeed_ps = 2;

  // SparseCore to host time in picoseconds.
  uint64 sc_outfeed_ps = 3;

  // Idle time but not waiting for input in picoseconds.
  uint64 sc_idle_ps = 4;

  // SparseCore busy time in picoseconds.
  uint64 sc_busy_ps = 5;
}

// Information about memory transfer to/from device memory.
message DeviceMemoryTransfer {
  uint64 occurrence = 1;
  double time_us = 2;
  uint64 bytes_transferred = 3;
}

// Next ID: 7
// Result proto for StepInfo.
message StepInfoResult {
  // The step number.
  uint32 step_num = 1;
  // The step name.
  string step_name = 5;
  // The step duration in picoseconds.
  uint64 duration_ps = 2;
  // The start time of this step in picoseconds.
  uint64 begin_ps = 3;
  // Breakdown of the step-time. Can be unpacked into a GenericStepBreakdown.
  google.protobuf.Any step_breakdown = 4;
  // Total time/bytes/occurrences for collectives. (All-Reduce, All-to-All etc)
  DeviceMemoryTransfer collectives = 6;
}

// Result proto for all -educe ops.
message AllReduceInfo {
  // Unique id for all-reduce ops.
  uint64 id = 1;
  // The name of the hlo op. This field is no longer set by the profiler.
  string name = 2 [deprecated = true];
  // For all-reduce nodes from different modules, if they have the same
  // all_reduce_id, they will be 'Allreduce'd'. If empty, AllReduce will not be
  // applied across modules.
  uint64 all_reduce_id = 3;
  // The start time in picoseconds of the op event.
  uint64 start_time_ps = 4;
  // The end time in picoseconds of the op event.
  uint64 end_time_ps = 5;
  // The size of the op in bytes.
  uint64 byte_size = 6;
}

// Result database for all-reduce ops.
message AllReduceDbResult {
  repeated AllReduceInfo all_reduce_info = 1;
}

// Result proto for information in a step across all cores.
message PerCoreStepInfo {
  // The step number.
  uint32 step_num = 1;
  // A map from core_id to StepInfo.
  map<uint32, StepInfoResult> step_info_per_core = 2;
  // The result for the per-step HLO-metric database.
  OpMetricsDb hlo_metrics_db = 3;
  // A map from core ID to program replica id. Replica id map could change
  // during a profile session, but should stay stable within a step.
  map<uint32, uint32> core_id_to_replica_id_map = 5;
  // A map from core_id to all-reduce ops.
  map<uint32, AllReduceDbResult> all_reduce_db_per_core = 6;
  // Information about device memory transfers, categoried by source and
  // destination. Ordered by following categories:
  // 1. HostToDevice
  // 2. DeviceToHost
  // 3. DeviceToDevice
  // Cores are normally sharing host interfaces (i.e. PCIe).
  repeated DeviceMemoryTransfer device_memory_transfers = 7;
  // The new path for per-step HLO metrics. hlo_metrics_db will be deprecated
  // soon.
  tensorflow.profiler.FlatOpMetricsDb flat_hlo_metrics_db = 8;

  reserved 4;
}

// Result proto for a StepDatabase.
message StepDatabaseResult {
  // A sequence of PerCoreStepInfo.
  repeated PerCoreStepInfo step_sequence = 1;
  // Whether the step db uses incomplete step information.
  // This flag is set to true when:
  // 1) no step marker or annotation present.
  // 2) profiling duration is too short to cover a full step.
  // If this flag is false, we will group and breakdown the
  // profile by complete steps only and ignore incomplete steps.
  // If this flag is true, we will simply aggregate and breakdown over the total
  // profile as a single step.
  bool use_incomplete_step = 2;
  // Number of steps dropped during post processing.
  uint32 num_steps_dropped = 3;
  // If the step_sequence is empty because:
  //   * there is no step profiled on any host, then empty_intersect is false.
  //   * there are steps profiled on some host, but the intersection of steps
  //     over all hosts is empty, then empty_intersect is true.
  bool empty_intersect = 4;
}
