syntax = "proto3";

package tensorflow.profiler;

// This proto is based on MegaScaleInfoProto and should be consistent with it.
message DcnCollectiveInfoProto {
  enum TransferType {
    UNKNOWN_TRANSFER_TYPE = 0;

    // XLA AllToAll transfer.
    // Needs `endpoint_groups`.
    ALL_TO_ALL = 1;

    // Peer-To-Peer DCN transfer from source to one destination.
    // Needs one_to_one_groups.
    ONE_TO_ONE = 2;

    // XLA reduce-scatter transfer.
    // Needs `endpoint_groups`.
    REDUCE_SCATTER = 3;

    // XLA AllGather transfer.
    // Needs `endpoint_groups`.
    ALL_GATHER = 4;

    // XLA all-reduce transfer.
    // Needs `endpoint_groups`.
    ALL_REDUCE = 5;

    // XLA ragged all-to-all transfer.
    // Needs `endpoint_groups`.
    RAGGED_ALL_TO_ALL = 6;
  }

  message Endpoint {
    int32 slice_id = 1;
    int32 device_id = 2;
  }

  message EndpointGroup {
    repeated Endpoint endpoints = 1;
  }

  message OneToOneGroup {
    Endpoint source = 1;
    Endpoint destination = 2;
  }

  // The type of DCN transfer.
  TransferType transfer_type = 1;

  // Groups of endpoints (in the form of slice id and device id) involved in
  // `ALL_TO_ALL`, `REDUCE_SCATTER`, `ALL_REDUCE` and `ALL_GATHER` transfer.
  repeated EndpointGroup endpoint_groups = 2;

  // Groups of endpoints (in the form of slice id and device id) involved in
  // `ONE_TO_ONE` transfer.
  repeated OneToOneGroup one_to_one_groups = 3;
}
