xref: /aosp_15_r20/external/tensorflow/tensorflow/core/profiler/protobuf/memory_viewer_preprocess.proto (revision b6fb3261f9314811a0f4371741dbb8839866f948) 
1// Protobuf definitions for communicating the results of the memory
2// visualization analysis subprocess (written in C++) to the outer script which
3// renders HTML from Python.
4
5syntax = "proto3";
6
7package tensorflow.profiler;
8
9// Describes a heap object that is displayed in a plot in the memory
10// visualization HTML.
11message HeapObject {
12  oneof color {
13    int32 numbered = 1;
14    string named = 2;
15  }
16  string label = 3;
17  int32 logical_buffer_id = 4;
18  double logical_buffer_size_mib = 5;
19  double unpadded_shape_mib = 6;
20  string instruction_name = 7;
21  string shape_string = 8;
22  string tf_op_name = 9;
23  string group_name = 10;
24  string op_code = 11;
25}
26
27// Describes the start / exclusive limit HLO program points for a given buffer
28// lifetime, used for rendering a box on the plot.
29message BufferSpan {
30  int32 start = 1;
31  int32 limit = 2;
32}
33
34message LogicalBuffer {
35  int64 id = 1;
36  string shape = 2;
37  double size_mib = 3;
38  string hlo_name = 4;
39  repeated int64 shape_index = 5;
40}
41
42message BufferAllocation {
43  int64 id = 1;
44  double size_mib = 2;
45  repeated string attributes = 3;
46  repeated LogicalBuffer logical_buffers = 4;
47  string common_shape = 5;
48}
49
50// Groups together all results from the preprocessing C++ step.
51message PreprocessResult {
52  // Heap sizes at each HLO program point (the HLO sequential order).
53  repeated double heap_sizes = 1;
54
55  // Unpadded heap sizes (calculated as the minimal sizes based on the data type
56  // and dimensionality) at each HLO program point (the HLO sequential order).
57  repeated double unpadded_heap_sizes = 2;
58
59  // Heap objects at the peak memory usage point ordered by HLO program "birth"
60  // time.
61  repeated HeapObject max_heap = 3;
62
63  // Heap objects at the peak memory usage point ordered by size, descending.
64  repeated HeapObject max_heap_by_size = 4;
65
66  // Mapping from logical buffer ID to the HLO sequential order span in which it
67  // is alive.
68  map<int32, BufferSpan> logical_buffer_spans = 5;
69
70  // Indexes to get back and forth from the by-size and by-program-order
71  // sequences.
72  repeated int32 max_heap_to_by_size = 6;
73  repeated int32 by_size_to_max_heap = 7;
74
75  string module_name = 8;
76  string entry_computation_name = 9;
77
78  // Peak heap size for the HLO program.
79  double peak_heap_mib = 10;
80
81  // Peak unpadded heap size for the HLO program.
82  double peak_unpadded_heap_mib = 11;
83
84  // HLO program point number at which the peak heap size occurs.
85  int32 peak_heap_size_position = 12;
86
87  // Size of the entry computation parameters in MiB.
88  //
89  // This does not reflect whether those MiB are reusable during the computation
90  // or not, it is simply a size value.
91  double entry_computation_parameters_mib = 13;
92
93  double non_reusable_mib = 14;
94
95  double maybe_live_out_mib = 15;
96
97  repeated BufferAllocation indefinite_lifetimes = 16;
98}
99