syntax = "proto3";

package tensorflow;

option go_package = "github.com/google/tsl/tsl/go/protobuf/for_core_protos_go_proto";

// Some of the data from AllocatorStats
message MemAllocatorStats {
  int64 num_allocs = 1;
  int64 bytes_in_use = 2;
  int64 peak_bytes_in_use = 3;
  int64 largest_alloc_size = 4;
  float fragmentation_metric = 5;
}

message MemChunk {
  uint64 address = 1;
  int64 size = 2;
  int64 requested_size = 3;
  int32 bin = 4;
  string op_name = 5;
  uint64 freed_at_count = 6;
  uint64 action_count = 7;
  bool in_use = 8;
  uint64 step_id = 9;
}

message BinSummary {
  int32 bin = 1;
  int64 total_bytes_in_use = 2;
  int64 total_bytes_in_bin = 3;
  int64 total_chunks_in_use = 4;
  int64 total_chunks_in_bin = 5;
}

message SnapShot {
  uint64 action_count = 1;
  int64 size = 2;
}

message MemoryDump {
  string allocator_name = 1;
  repeated BinSummary bin_summary = 2;
  repeated MemChunk chunk = 3;
  repeated SnapShot snap_shot = 4;
  MemAllocatorStats stats = 5;
}