// This proto describes the format of the output profile file from
// the Pod Viewer tool.
syntax = "proto3";

package tensorflow.profiler;

import "plugin/xprof/protobuf/diagnostics.proto";
import "plugin/xprof/protobuf/pod_stats.proto";
import "plugin/xprof/protobuf/topology.proto";

// Describes the replica groups in a cross replica op (e.g., all-reduce and
// all-to-all).
message ReplicaGroup {
  // The ids of the replicas that belongs to the same group. The ordering of the
  // ids matters in some ops (e.g., all-to-all).
  repeated int64 replica_ids = 1;
}

message AllReduceOpInfo {
  // Name of this OP.
  string name = 1;
  // Number of instances that this OP occurred.
  uint32 occurrences = 2;
  // The time in microseconds spent in this OP (averaged across all of its
  // occurrences).
  double duration_us = 3;
  // Byte size of data transferred.
  uint64 data_size = 4;
  // Replica groups.
  repeated ReplicaGroup replica_groups = 5;
  // Description (e.g. XLA expression).
  string description = 6;
}

// Result proto for information in a step across all cores.
message PodStatsMap {
  // The (micro) step number.
  uint32 step_num = 1;
  // A map from core_id to PodStatsRecord.
  map<uint32, PodStatsRecord> pod_stats_per_core = 2;
  // A database of channel info.
  repeated ChannelInfo channel_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 = 4;
  // A database of all reduce ops.
  repeated AllReduceOpInfo all_reduce_op_db = 5;
}

// A sequence of PodStatsMap for each step.
message PodStatsSequence {
  repeated PodStatsMap pod_stats_map = 1;
}

// Next ID: 14
// Information about a send and recv channel.
message ChannelInfo {
  // Id of the channel.
  int64 channel_id = 1;
  // Core ids of send ops.
  repeated uint32 src_core_ids = 11;
  // Core ids of recv ops.
  repeated uint32 dst_core_ids = 12;
  // Byte size of the data transferred.
  uint64 data_size = 4;
  // Duration from the beginning of send to the end of recv-done in
  // microseconds.
  double duration_us = 5;
  // Number of occurrences of a channel.
  uint32 occurrences = 6;
  // Percentage of the link BW utilized over the peak link BW.
  double utilization = 7;
  // A list of hlo names associated with this channel id.
  repeated string hlo_names = 8;
  // Duration from the beginning of the recv-done to the beginning of send in
  // microseconds. If the recv-done op starts after the beginning of the send
  // op, the delay is zero.
  double send_delay_us = 9;
  // Description (e.g. XLA expression).
  string description = 13;

  reserved 2, 3, 10;
}

message PodViewerSummary {
  repeated string warnings = 1;
}

// Next ID: 9
// Topology graph draws all the cores in the system in a 2-D rectangle or
// 3-D cube. It is hierarchically grouped by host, chip and core.
message PodViewerTopology {
  // Number of chips in the x dimension of the rectangle/cube.
  int32 x_dimension = 1;
  // Number of chips in the y dimension of the rectangle/cube.
  int32 y_dimension = 2;
  // Number of chips in the z dimension of the cube.
  int32 z_dimension = 3;
  // Number of chips in the x dimension of each host.
  int32 host_x_stride = 4;
  // Number of chips in the y dimension of each host.
  int32 host_y_stride = 5;
  // Number of chips in the z dimension of each host.
  int32 host_z_stride = 6;
  // Number of cores per chip.
  int32 num_cores_per_chip = 7;
  // Core locations.
  repeated TopologyLocation cores = 8;
}

// Next ID: 12
// A database of pod viewer records.
message PodViewerDatabase {
  // The type of device used.
  string device_type = 10;
  // Pod level stats for each step.
  PodStatsSequence pod_stats_sequence = 3;
  // Top level summary of pod viewer.
  PodViewerSummary summary = 7;
  // Error and warning messages for diagnosing profiling issues.
  Diagnostics diagnostics = 8;
  // A map from event type number to event name string for step breakdown.
  repeated StepBreakdownEvents step_breakdown_events = 9;
  // Info to draw the topology graph.
  PodViewerTopology topology = 11;

  reserved 1, 2, 4, 5, 6;
}
