tensorboard-plugin-profile 2.18.0__tar.gz → 2.19.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (53) hide show
  1. tensorboard_plugin_profile-2.19.0/PKG-INFO +110 -0
  2. tensorboard_plugin_profile-2.19.0/README.md +89 -0
  3. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/setup.py +5 -2
  4. tensorboard_plugin_profile-2.19.0/tensorboard_plugin_profile/convert/hlo_stats_proto_to_gviz.py +162 -0
  5. tensorboard_plugin_profile-2.19.0/tensorboard_plugin_profile/convert/inference_stats_proto_to_gviz.py +327 -0
  6. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/convert/input_pipeline_proto_to_gviz.py +238 -10
  7. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/convert/overview_page_proto_to_gviz.py +44 -3
  8. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/convert/raw_to_tool_data.py +24 -4
  9. tensorboard_plugin_profile-2.19.0/tensorboard_plugin_profile/convert/roofline_model_proto_to_gviz.py +392 -0
  10. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/convert/trace_events_json.py +0 -1
  11. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/integration_tests/tpu/tensorflow/tpu_tf2_keras_test.py +3 -0
  12. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/profile_plugin.py +60 -26
  13. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/profile_plugin_loader.py +8 -3
  14. tensorboard_plugin_profile-2.19.0/tensorboard_plugin_profile/protobuf/hlo_stats_pb2.py +27 -0
  15. tensorboard_plugin_profile-2.19.0/tensorboard_plugin_profile/protobuf/inference_stats_pb2.py +75 -0
  16. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/protobuf/overview_page_pb2.py +7 -3
  17. tensorboard_plugin_profile-2.19.0/tensorboard_plugin_profile/protobuf/roofline_model_pb2.py +30 -0
  18. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/protobuf/tf_stats_pb2.py +2 -2
  19. tensorboard_plugin_profile-2.19.0/tensorboard_plugin_profile/protobuf/tpu_input_pipeline_pb2.py +32 -0
  20. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/static/bundle.js +3816 -682
  21. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/static/trace_viewer_index.html +46 -1
  22. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/static/trace_viewer_index.js +258 -26
  23. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/version.py +1 -1
  24. tensorboard_plugin_profile-2.19.0/tensorboard_plugin_profile.egg-info/PKG-INFO +110 -0
  25. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile.egg-info/SOURCES.txt +8 -1
  26. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile.egg-info/requires.txt +1 -0
  27. tensorboard_plugin_profile-2.18.0/PKG-INFO +0 -21
  28. tensorboard_plugin_profile-2.18.0/README.rst +0 -1
  29. tensorboard_plugin_profile-2.18.0/tensorboard_plugin_profile.egg-info/PKG-INFO +0 -21
  30. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/setup.cfg +0 -0
  31. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/__init__.py +0 -0
  32. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/convert/dcn_collective_stats_proto_to_gviz.py +0 -0
  33. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/convert/diagnostics.py +0 -0
  34. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/convert/kernel_stats_proto_to_gviz.py +0 -0
  35. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/convert/tf_data_stats_proto_to_gviz.py +0 -0
  36. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/convert/tf_stats_proto_to_gviz.py +0 -0
  37. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/integration_tests/tf_mnist.py +0 -0
  38. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/integration_tests/tf_profiler_session.py +0 -0
  39. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/protobuf/dcn_slack_analysis_pb2.py +0 -0
  40. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/protobuf/diagnostics_pb2.py +0 -0
  41. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/protobuf/input_pipeline_pb2.py +0 -0
  42. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/protobuf/kernel_stats_pb2.py +0 -0
  43. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/protobuf/power_metrics_pb2.py +0 -0
  44. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/protobuf/tf_data_stats_pb2.py +0 -0
  45. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/protobuf/trace_events_pb2.py +0 -0
  46. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/static/index.html +0 -0
  47. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/static/index.js +0 -0
  48. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/static/materialicons.woff2 +0 -0
  49. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/static/styles.css +0 -0
  50. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile/static/zone.js +0 -0
  51. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile.egg-info/dependency_links.txt +0 -0
  52. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile.egg-info/entry_points.txt +0 -0
  53. {tensorboard_plugin_profile-2.18.0 → tensorboard_plugin_profile-2.19.0}/tensorboard_plugin_profile.egg-info/top_level.txt +0 -0
@@ -0,0 +1,110 @@
1
+ Metadata-Version: 2.1
2
+ Name: tensorboard_plugin_profile
3
+ Version: 2.19.0
4
+ Summary: Profile Tensorboard Plugin
5
+ Home-page: https://github.com/tensorflow/profiler
6
+ Author: Google Inc.
7
+ Author-email: packages@tensorflow.org
8
+ License: Apache 2.0
9
+ Keywords: jax pytorch xla tensorflow tensorboard xprof profile plugin
10
+ Classifier: Intended Audience :: Developers
11
+ Classifier: Intended Audience :: Education
12
+ Classifier: Intended Audience :: Science/Research
13
+ Classifier: License :: OSI Approved :: Apache Software License
14
+ Classifier: Programming Language :: Python :: 2
15
+ Classifier: Programming Language :: Python :: 3
16
+ Classifier: Topic :: Scientific/Engineering :: Mathematics
17
+ Classifier: Topic :: Software Development :: Libraries :: Python Modules
18
+ Classifier: Topic :: Software Development :: Libraries
19
+ Requires-Python: >= 2.7, != 3.0.*, != 3.1.*
20
+ Description-Content-Type: text/markdown
21
+
22
+ # Tensorboard Profiler Plugin
23
+ The profiler includes a suite of tools for [JAX](https://jax.readthedocs.io/), [TensorFlow](https://www.tensorflow.org/), and [PyTorch/XLA](https://github.com/pytorch/xla). These tools help you understand, debug and optimize programs to run on CPUs, GPUs and TPUs.
24
+
25
+ The profiler plugin offers a number of tools to analyse and visualize the
26
+ performance of your model across multiple devices. Some of the tools include:
27
+
28
+ * **Overview**: A high-level overview of the performance of your model. This
29
+ is an aggregated overview for your host and all devices. It includes:
30
+ * Performance summary and breakdown of step times.
31
+ * A graph of individual step times.
32
+ * A table of the top 10 most expensive operations.
33
+ * **Trace Viewer**: Displays a timeline of the execution of your model that shows:
34
+ * The duration of each op.
35
+ * Which part of the system (host or device) executed an op.
36
+ * The communication between devices.
37
+ * **Memory Profile Viewer**: Monitors the memory usage of your model.
38
+ * **Graph Viewer**: A visualization of the graph structure of HLOs of your model.
39
+
40
+ ## Demo
41
+ First time user? Come and check out this [Colab Demo](https://www.tensorflow.org/tensorboard/tensorboard_profiling_keras).
42
+
43
+ ## Prerequisites
44
+ * TensorFlow >= 2.18.0
45
+ * TensorBoard >= 2.18.0
46
+ * tensorboard-plugin-profile >= 2.18.0
47
+
48
+ Note: The Tensorboard Profiler Plugin requires access to the Internet to load the [Google Chart library](https://developers.google.com/chart/interactive/docs/basic_load_libs#basic-library-loading).
49
+ Some charts and tables may be missing if you run TensorBoard entirely offline on
50
+ your local machine, behind a corporate firewall, or in a datacenter.
51
+
52
+ To profile on a **single GPU** system, the following NVIDIA software must be installed on your system:
53
+
54
+ 1. NVIDIA GPU drivers and CUDA Toolkit:
55
+ * CUDA 12.5 requires 525.60.13 and higher.
56
+ 2. Ensure that CUPTI 10.1 exists on the path.
57
+
58
+ ```shell
59
+ $ /sbin/ldconfig -N -v $(sed 's/:/ /g' <<< $LD_LIBRARY_PATH) | grep libcupti
60
+ ```
61
+
62
+ If you don't see `libcupti.so.12.5` on the path, prepend its installation directory to the $LD_LIBRARY_PATH environmental variable:
63
+
64
+ ```shell
65
+ $ export LD_LIBRARY_PATH=/usr/local/cuda/extras/CUPTI/lib64:$LD_LIBRARY_PATH
66
+ ```
67
+ Run the ldconfig command above again to verify that the CUPTI 12.5 library is found.
68
+
69
+ If this doesn't work, try:
70
+ ```shell
71
+ $ sudo apt-get install libcupti-dev
72
+ ```
73
+
74
+ To profile a system with **multiple GPUs**, see this [guide](docs/profile_multi_gpu.md) for details.
75
+
76
+ To profile multi-worker GPU configurations, profile individual workers independently.
77
+
78
+ To profile cloud TPUs, you must have access to Google Cloud TPUs.
79
+
80
+ ## Quick Start
81
+ The profiler plugin follows the TensorFlow versioning scheme. As a result, the
82
+ `tensorboard-plugin-profile` PyPI package can be behind the `tbp-nightly` PyPI
83
+ package. In order to get the latest version of the profiler plugin, you can
84
+ install the nightly package..
85
+
86
+ To install the nightly version of profiler:
87
+
88
+ ```
89
+ $ pip uninstall tensorboard-plugin-profile
90
+ $ pip install tbp-nightly
91
+ ```
92
+
93
+ Run TensorBoard:
94
+
95
+ ```
96
+ $ tensorboard --logdir=profiler/demo
97
+ ```
98
+ If you are behind a corporate firewall, you may need to include the `--bind_all`
99
+ tensorboard flag.
100
+
101
+ Go to `localhost:6006/#profile` of your browser, you should now see the demo overview page show up.
102
+ ![Overview Page](docs/images/overview_page.png)
103
+ Congratulations! You're now ready to capture a profile.
104
+
105
+ ## Next Steps
106
+ * JAX Profiling Guide: https://jax.readthedocs.io/en/latest/profiling.html
107
+ * TensorFlow Profiling Guide: https://tensorflow.org/guide/profiler
108
+ * Cloud TPU Profiling Guide: https://cloud.google.com/tpu/docs/cloud-tpu-tools
109
+ * Colab Tutorial: https://www.tensorflow.org/tensorboard/tensorboard_profiling_keras
110
+ * MiniGPT Example: https://docs.jaxstack.ai/en/latest/JAX_for_LLM_pretraining.html
@@ -0,0 +1,89 @@
1
+ # Tensorboard Profiler Plugin
2
+ The profiler includes a suite of tools for [JAX](https://jax.readthedocs.io/), [TensorFlow](https://www.tensorflow.org/), and [PyTorch/XLA](https://github.com/pytorch/xla). These tools help you understand, debug and optimize programs to run on CPUs, GPUs and TPUs.
3
+
4
+ The profiler plugin offers a number of tools to analyse and visualize the
5
+ performance of your model across multiple devices. Some of the tools include:
6
+
7
+ * **Overview**: A high-level overview of the performance of your model. This
8
+ is an aggregated overview for your host and all devices. It includes:
9
+ * Performance summary and breakdown of step times.
10
+ * A graph of individual step times.
11
+ * A table of the top 10 most expensive operations.
12
+ * **Trace Viewer**: Displays a timeline of the execution of your model that shows:
13
+ * The duration of each op.
14
+ * Which part of the system (host or device) executed an op.
15
+ * The communication between devices.
16
+ * **Memory Profile Viewer**: Monitors the memory usage of your model.
17
+ * **Graph Viewer**: A visualization of the graph structure of HLOs of your model.
18
+
19
+ ## Demo
20
+ First time user? Come and check out this [Colab Demo](https://www.tensorflow.org/tensorboard/tensorboard_profiling_keras).
21
+
22
+ ## Prerequisites
23
+ * TensorFlow >= 2.18.0
24
+ * TensorBoard >= 2.18.0
25
+ * tensorboard-plugin-profile >= 2.18.0
26
+
27
+ Note: The Tensorboard Profiler Plugin requires access to the Internet to load the [Google Chart library](https://developers.google.com/chart/interactive/docs/basic_load_libs#basic-library-loading).
28
+ Some charts and tables may be missing if you run TensorBoard entirely offline on
29
+ your local machine, behind a corporate firewall, or in a datacenter.
30
+
31
+ To profile on a **single GPU** system, the following NVIDIA software must be installed on your system:
32
+
33
+ 1. NVIDIA GPU drivers and CUDA Toolkit:
34
+ * CUDA 12.5 requires 525.60.13 and higher.
35
+ 2. Ensure that CUPTI 10.1 exists on the path.
36
+
37
+ ```shell
38
+ $ /sbin/ldconfig -N -v $(sed 's/:/ /g' <<< $LD_LIBRARY_PATH) | grep libcupti
39
+ ```
40
+
41
+ If you don't see `libcupti.so.12.5` on the path, prepend its installation directory to the $LD_LIBRARY_PATH environmental variable:
42
+
43
+ ```shell
44
+ $ export LD_LIBRARY_PATH=/usr/local/cuda/extras/CUPTI/lib64:$LD_LIBRARY_PATH
45
+ ```
46
+ Run the ldconfig command above again to verify that the CUPTI 12.5 library is found.
47
+
48
+ If this doesn't work, try:
49
+ ```shell
50
+ $ sudo apt-get install libcupti-dev
51
+ ```
52
+
53
+ To profile a system with **multiple GPUs**, see this [guide](docs/profile_multi_gpu.md) for details.
54
+
55
+ To profile multi-worker GPU configurations, profile individual workers independently.
56
+
57
+ To profile cloud TPUs, you must have access to Google Cloud TPUs.
58
+
59
+ ## Quick Start
60
+ The profiler plugin follows the TensorFlow versioning scheme. As a result, the
61
+ `tensorboard-plugin-profile` PyPI package can be behind the `tbp-nightly` PyPI
62
+ package. In order to get the latest version of the profiler plugin, you can
63
+ install the nightly package..
64
+
65
+ To install the nightly version of profiler:
66
+
67
+ ```
68
+ $ pip uninstall tensorboard-plugin-profile
69
+ $ pip install tbp-nightly
70
+ ```
71
+
72
+ Run TensorBoard:
73
+
74
+ ```
75
+ $ tensorboard --logdir=profiler/demo
76
+ ```
77
+ If you are behind a corporate firewall, you may need to include the `--bind_all`
78
+ tensorboard flag.
79
+
80
+ Go to `localhost:6006/#profile` of your browser, you should now see the demo overview page show up.
81
+ ![Overview Page](docs/images/overview_page.png)
82
+ Congratulations! You're now ready to capture a profile.
83
+
84
+ ## Next Steps
85
+ * JAX Profiling Guide: https://jax.readthedocs.io/en/latest/profiling.html
86
+ * TensorFlow Profiling Guide: https://tensorflow.org/guide/profiler
87
+ * Cloud TPU Profiling Guide: https://cloud.google.com/tpu/docs/cloud-tpu-tools
88
+ * Colab Tutorial: https://www.tensorflow.org/tensorboard/tensorboard_profiling_keras
89
+ * MiniGPT Example: https://docs.jaxstack.ai/en/latest/JAX_for_LLM_pretraining.html
@@ -30,11 +30,12 @@ REQUIRED_PACKAGES = [
30
30
  'setuptools >= 41.0.0',
31
31
  'six >= 1.10.0',
32
32
  'werkzeug >= 0.11.15',
33
+ 'etils[epath] >= 1.0.0'
33
34
  ]
34
35
 
35
36
 
36
37
  def get_readme():
37
- with open('README.rst') as f:
38
+ with open('README.md') as f:
38
39
  return f.read()
39
40
 
40
41
 
@@ -43,6 +44,7 @@ setuptools.setup(
43
44
  version=VERSION,
44
45
  description='Profile Tensorboard Plugin',
45
46
  long_description=get_readme(),
47
+ long_description_content_type='text/markdown',
46
48
  author='Google Inc.',
47
49
  author_email='packages@tensorflow.org',
48
50
  url='https://github.com/tensorflow/profiler',
@@ -53,6 +55,7 @@ setuptools.setup(
53
55
  ),
54
56
  package_data={
55
57
  'tensorboard_plugin_profile': ['static/**'],
58
+ '': ['_pywrap_profiler_plugin.so'],
56
59
  },
57
60
  entry_points={
58
61
  'tensorboard_plugins': [
@@ -78,5 +81,5 @@ setuptools.setup(
78
81
  'Topic :: Software Development :: Libraries',
79
82
  ],
80
83
  license='Apache 2.0',
81
- keywords='tensorflow tensorboard xprof profile plugin',
84
+ keywords='jax pytorch xla tensorflow tensorboard xprof profile plugin',
82
85
  )
@@ -0,0 +1,162 @@
1
+ # Copyright 2024 The TensorFlow Authors. All Rights Reserved.
2
+ #
3
+ # Licensed under the Apache License, Version 2.0 (the "License");
4
+ # you may not use this file except in compliance with the License.
5
+ # You may obtain a copy of the License at
6
+ #
7
+ # http://www.apache.org/licenses/LICENSE-2.0
8
+ #
9
+ # Unless required by applicable law or agreed to in writing, software
10
+ # distributed under the License is distributed on an "AS IS" BASIS,
11
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
+ # See the License for the specific language governing permissions and
13
+ # limitations under the License.
14
+ # ==============================================================================
15
+ """For conversion of HloStats protos to GViz DataTables.
16
+
17
+ Usage:
18
+ gviz_data_tables = generate_hlo_stats_table(hlo_stats_db)
19
+ """
20
+
21
+ from __future__ import absolute_import
22
+ from __future__ import division
23
+ from __future__ import print_function
24
+
25
+ from typing import Any
26
+
27
+ import gviz_api
28
+
29
+ from tensorboard_plugin_profile.protobuf import hlo_stats_pb2
30
+
31
+
32
+ def get_hlo_op_name_from_expression(hlo_op_expression: str) -> str:
33
+ """Returns the hlo op name from the hlo op expression.
34
+
35
+ Args:
36
+ hlo_op_expression: A string of hlo op expression.
37
+
38
+ Returns:
39
+ Returns the hlo op name extracted from the hlo op expression.
40
+ """
41
+ # The parse logic based on the assumption that the hlo op text is in format of
42
+ # '%op_name = <long name>'
43
+ parts = hlo_op_expression.split(" = ")
44
+ hlo_op_name = parts[0]
45
+ if hlo_op_name and hlo_op_name[0] == "%":
46
+ hlo_op_name = hlo_op_name[1:]
47
+ return hlo_op_name
48
+
49
+
50
+ TableColumnDescription = tuple[str, str, str]
51
+ TableRow = list[Any]
52
+ TableProperties = dict[str, Any]
53
+
54
+
55
+ def get_hlo_stats_table_args(
56
+ hlo_stats_db: hlo_stats_pb2.HloStatsDatabase,
57
+ ) -> tuple[list[TableColumnDescription], list[TableRow], TableProperties]:
58
+ """Creates hlo op stats table from a hlo stats proto.
59
+
60
+ Args:
61
+ hlo_stats_db: A HloStatsDatabase proto.
62
+
63
+ Returns:
64
+ Returns table description(column defubutuibs), data(rows data) and custom
65
+ properties for preparing the HloStats gviz table.
66
+ """
67
+
68
+ table_description = [
69
+ ("rank", "number", "Rank"),
70
+ ("program_id", "string", "Program id"),
71
+ ("category", "string", "HLO op category"),
72
+ ("hlo_op_name", "string", "HLO op name"),
73
+ ("hlo_op_expression", "string", "HLO op text"),
74
+ ("tf_op_name", "string", "Framework op name"),
75
+ ("occurrences", "number", "#Occurrences"),
76
+ ("total_time", "number", "Total time (us)"),
77
+ ("avg_time", "number", "Avg. time (us)"),
78
+ ("total_self_time", "number", "Total self time (us)"),
79
+ ("avg_self_time", "number", "Avg. self time (us)"),
80
+ ("total_self_time_percent", "number", "Total self time (%)"),
81
+ (
82
+ "cumulative_total_self_time_percent",
83
+ "number",
84
+ "Cumulative total self time (%)",
85
+ ),
86
+ ("dma_stall_percent", "number", "%time stalled by DMA"),
87
+ ("model_flop_rate", "number", "Model GFLOP/s"),
88
+ ("normalized_flop_rate", "number", "Normalized GFLOP/s"),
89
+ ("measured_memory_bw", "number", "Measured memory BW (GiB/s)"),
90
+ ("hbm_bw", "number", "HBM BW (GiB/s)"),
91
+ ("cmem_read_bw", "number", "CMEM Read BW (GiB/s)"),
92
+ ("cmem_write_bw", "number", "CMEM Write BW (GiB/s)"),
93
+ ("operational_intensity", "number", "Operational intensity (FLOPS/Byte)"),
94
+ ("bound_by", "string", "Bound by"),
95
+ ("hlo_rematerialization", "string", "Rematerialization"),
96
+ ("outside_compilation", "string", "Outside Compilation"),
97
+ ("autotuned", "string", "Autotuned"),
98
+ ]
99
+
100
+ data = []
101
+ for record in hlo_stats_db.hlo_stats_record:
102
+ row = [
103
+ record.rank,
104
+ str(record.program_id),
105
+ record.hlo_category,
106
+ get_hlo_op_name_from_expression(record.hlo_expression),
107
+ record.hlo_expression,
108
+ record.tf_op_name,
109
+ record.occurrences,
110
+ record.total_time_in_us,
111
+ record.avg_time_in_us,
112
+ record.total_self_time_in_us,
113
+ record.avg_self_time_in_us,
114
+ record.total_self_time_as_fraction,
115
+ record.cumulative_total_self_time_as_fraction,
116
+ record.dma_stall_fraction,
117
+ record.model_flop_rate,
118
+ record.measured_flop_rate,
119
+ record.measured_memory_bw,
120
+ record.hbm_bw,
121
+ record.cmem_read_bw,
122
+ record.cmem_write_bw,
123
+ record.operational_intensity,
124
+ record.bound_by,
125
+ "Yes" if record.rematerialization else "No",
126
+ "Yes" if record.outside_compilation else "No",
127
+ "Yes" if record.autotuned else "No",
128
+ ]
129
+ data.append(row)
130
+
131
+ return (table_description, data, {})
132
+
133
+
134
+ def generate_hlo_stats_table(
135
+ hlo_stats_db: hlo_stats_pb2.HloStatsDatabase,
136
+ ) -> gviz_api.DataTable:
137
+ """Converts a HloStatsDb proto to a GViz DataTable.
138
+
139
+ Args:
140
+ hlo_stats_db: A HloStatsDatabase proto.
141
+
142
+ Returns:
143
+ Returns a GViz DataTable for HloStats.
144
+ """
145
+ (table_description, data, custom_properties) = get_hlo_stats_table_args(
146
+ hlo_stats_db
147
+ )
148
+ return gviz_api.DataTable(table_description, data, custom_properties)
149
+
150
+
151
+ def to_json(raw_data: bytes) -> str:
152
+ """Converts a serialized HloStatsDb bytes to json.
153
+
154
+ Args:
155
+ raw_data: Bytes of serialized HloStatsDb.
156
+
157
+ Returns:
158
+ Returns a json string of HloStats Table.
159
+ """
160
+ hlo_stats_db = hlo_stats_pb2.HloStatsDatabase()
161
+ hlo_stats_db.ParseFromString(raw_data)
162
+ return generate_hlo_stats_table(hlo_stats_db).ToJSon()