Add counter extraction to get counter names from files.

plugin/build_pip_package.sh

@@ -93,6 +93,10 @@ fi
 
 # Copy static files.
 cd xprof
+mkdir -p utils
+cp "${ROOT_RUNFILE_DIR}/xprof/utils/tpu_counter_ids_v7x.h" utils/
+cp "${ROOT_RUNFILE_DIR}/xprof/utils/tpu_counter_ids_v6e.h" utils/
+
 mkdir -p static
 cd static
 cp "$PLUGIN_RUNFILE_DIR/xprof/static/index.html" .

plugin/setup.py

@@ -83,6 +83,7 @@ def get_readme():
     package_data={
         'xprof': [
             'static/**',
+            'utils/*.h',
             'convert/profiler_plugin_c_api.so',
             'convert/profiler_plugin_c_api.pyd',
             'convert/profiler_plugin_c_api.dylib',

plugin/xprof/convert/BUILD

@@ -14,6 +14,7 @@ py_library(
     name = "all_libs",
     # strict_deps = False,
     deps = [
+        ":counter_extractor",
         ":raw_to_tool_data",
         ":trace_events_json",
     ],
@@ -59,6 +60,23 @@ py_library(
     ],
 )
 
+py_library(
+    name = "counter_extractor",
+    srcs = ["counter_extractor.py"],
+    data = ["@org_xprof//xprof/utils:tpu_counter_ids_headers"],
+    # strict_deps = False,
+)
+
+py_test(
+    name = "counter_extractor_test",
+    srcs = ["counter_extractor_test.py"],
+    # strict_deps = False,
+    deps = [
+        ":counter_extractor",
+        "@absl_py//absl/testing:absltest",
+    ],
+)
+
 py_test(
     name = "raw_to_tool_data_test",
     srcs = ["raw_to_tool_data_test.py"],

plugin/xprof/convert/counter_extractor.py

@@ -0,0 +1,49 @@
+"""Defines the counter extractor routine for TPU counters."""
+
+import functools
+import importlib.resources
+import re
+
+VALID_DEVICE_TYPES = frozenset({'v6e', 'v7x'})
+
+
+@functools.lru_cache()
+def get_all_counters(device_type: str) -> list[dict[str, int]]:
+  """Parses tpu_counter_ids_{device_type}.h and returns a list of counters with names and indices."""
+  if device_type not in VALID_DEVICE_TYPES:
+    raise ValueError(f'Unsupported device_type: {device_type}')
+  filename = f'tpu_counter_ids_{device_type}.h'
+
+  try:
+    ref = importlib.resources.files('xprof') / 'utils' / filename
+  except ImportError:
+    pass
+  content = ref.read_text(encoding='utf-8')
+  # Remove C++ style line comments from content.
+  content = re.sub(r'//.*', '', content)
+
+  # The counter file is a C++ header file with an enum.
+  # Example:
+  # enum TpuCounterIdsTpu{version} : uint64_t {
+  #   TPU_COUNTER_ID_FOO = 12345,
+  #   TPU_COUNTER_ID_BAR = 67890,
+  #   ...
+  # };
+  match = re.search(r'enum\s+\w+\s*:\s*uint64_t\s*\{([^}]+)\}', content)
+  if not match:
+    return []
+
+  enum_body = match.group(1)
+  counters = []
+  for line in enum_body.split(','):
+    line = line.strip()
+    if '=' in line:
+      name, val_str = line.split('=', 1)
+      try:
+        name = name.strip()
+        val = int(val_str.strip(), 0)
+        counters.append({'name': name.lower(), 'val': val})
+      except ValueError:
+        continue
+
+  return counters

plugin/xprof/convert/counter_extractor_test.py

@@ -0,0 +1,167 @@
+# Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Tests for the counter_extractor module."""
+
+from unittest import mock
+
+from absl.testing import absltest
+
+from xprof.convert import counter_extractor
+
+
+class CounterExtractorTest(absltest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    counter_extractor.get_all_counters.cache_clear()
+
+  @mock.patch("importlib.resources.files")
+  def test_get_all_counters_success(self, mock_files):
+    mock_content = """
+    #ifndef THIRD_PARTY_XPROF_UTILS_TPU_COUNTER_IDS_H_
+    #define THIRD_PARTY_XPROF_UTILS_TPU_COUNTER_IDS_H_
+    namespace xprof {
+    enum TpuCounterIdsTpu7x : uint64_t {
+      // NOLINTBEGIN
+      TPU_COUNTER_ID_FOO = 12345,
+      // Some comment
+      TPU_COUNTER_ID_BAR =
+          67890,
+      TPU_COUNTER_ID_HEX = 0x1234,
+      TPU_COUNTER_ID_MULTI_EQ = 1 = 2,
+      // NOLINTEND
+    };
+    }
+    #endif
+    """
+
+    mock_file = mock.MagicMock()
+    mock_file.read_text.return_value = mock_content
+    # Chain: files() -> / 'utils' -> / filename
+    mock_pkg = mock.MagicMock()
+    mock_utils = mock.MagicMock()
+    mock_pkg.__truediv__.return_value = mock_utils
+    mock_utils.__truediv__.return_value = mock_file
+    mock_files.return_value = mock_pkg
+
+    counters = counter_extractor.get_all_counters("v7x")
+    self.assertEqual(
+        counters,
+        [
+            {"name": "tpu_counter_id_foo", "val": 12345},
+            {"name": "tpu_counter_id_bar", "val": 67890},
+            {"name": "tpu_counter_id_hex", "val": 4660},
+        ],
+    )
+
+  @mock.patch("importlib.resources.files")
+  def test_get_all_counters_fallback_success(self, mock_files):
+    mock_content = """
+    #ifndef THIRD_PARTY_XPROF_UTILS_TPU_COUNTER_IDS_H_
+    #define THIRD_PARTY_XPROF_UTILS_TPU_COUNTER_IDS_H_
+    namespace xprof {
+    enum TpuCounterIdsTpu7x : uint64_t {
+      TPU_COUNTER_ID_FOO = 12345,
+    };
+    }
+    #endif
+    """
+
+    mock_file = mock.MagicMock()
+    mock_file.read_text.return_value = mock_content
+    mock_pkg = mock.MagicMock()
+    mock_utils = mock.MagicMock()
+    mock_pkg.__truediv__.return_value = mock_utils
+    mock_utils.__truediv__.return_value = mock_file
+
+    def side_effect(package_name):
+      if package_name == "google3.third_party.xprof":
+        raise ImportError("No module named google3")
+      elif package_name == "xprof":
+        return mock_pkg
+      raise ValueError(f"Unexpected package: {package_name}")
+
+    mock_files.side_effect = side_effect
+
+    counters = counter_extractor.get_all_counters("v7x")
+    self.assertEqual(
+        counters,
+        [
+            {"name": "tpu_counter_id_foo", "val": 12345},
+        ],
+    )
+
+  @mock.patch("importlib.resources.files")
+  def test_get_all_counters_missing_file(self, mock_files):
+    mock_utils = mock.MagicMock()
+    # Simulate FileNotFound when reading text
+    mock_file = mock.MagicMock()
+    mock_file.read_text.side_effect = FileNotFoundError("File not found")
+    mock_pkg = mock.MagicMock()
+    mock_pkg.__truediv__.return_value = mock_utils
+    mock_utils.__truediv__.return_value = mock_file
+    mock_files.return_value = mock_pkg
+
+    with self.assertRaises(FileNotFoundError):
+      counter_extractor.get_all_counters("v7x")
+
+  @mock.patch("importlib.resources.files")
+  def test_get_all_counters_no_match(self, mock_files):
+    mock_content = "some file without enum"
+    mock_file = mock.MagicMock()
+    mock_file.read_text.return_value = mock_content
+    mock_pkg = mock.MagicMock()
+    mock_utils = mock.MagicMock()
+    mock_pkg.__truediv__.return_value = mock_utils
+    mock_utils.__truediv__.return_value = mock_file
+    mock_files.return_value = mock_pkg
+
+    counters = counter_extractor.get_all_counters("v7x")
+    self.assertEqual(counters, [])
+
+  @mock.patch("importlib.resources.files")
+  def test_get_all_counters_empty_enum(self, mock_files):
+    mock_content_empty = "enum TpuCounterIdsTpu7x : uint64_t {};"
+    mock_content_whitespace = "enum TpuCounterIdsTpu7x : uint64_t { \n };"
+
+    mock_file = mock.MagicMock()
+    mock_pkg = mock.MagicMock()
+    mock_utils = mock.MagicMock()
+    mock_pkg.__truediv__.return_value = mock_utils
+    mock_utils.__truediv__.return_value = mock_file
+    mock_files.return_value = mock_pkg
+
+    # Test completely empty
+    mock_file.read_text.return_value = mock_content_empty
+    counters = counter_extractor.get_all_counters("v7x")
+    self.assertEqual(counters, [])
+
+    # Test whitespace only
+    mock_file.read_text.return_value = mock_content_whitespace
+    counters = counter_extractor.get_all_counters("v7x")
+    self.assertEqual(counters, [])
+
+  def test_get_all_counters_invalid_device_type(self):
+    with self.assertRaises(ValueError):
+      counter_extractor.get_all_counters("../invalid")
+
+  def test_get_all_counters_security_path_separators(self):
+    with self.assertRaises(ValueError):
+      counter_extractor.get_all_counters("v7x/../etc/passwd")
+
+
+if __name__ == "__main__":
+  absltest.main()

plugin/xprof/convert/scratch/test_extractor.py

@@ -0,0 +1,8 @@
+"""Test reading TPU counter IDs resources."""
+from xprof.convert import counter_extractor
+
+try:
+  counters = counter_extractor.get_all_counters("v7x")
+  print(f"Success! Found {len(counters)} counters.")
+except (ImportError, ValueError) as e:
+  print(f"Failed: {e}")

plugin/xprof/convert/scratch/test_resources.py

@@ -0,0 +1,24 @@
+"""Test reading TPU counter IDs resources."""
+import importlib.resources
+
+try:
+  # Try to read using importlib.resources
+  content = importlib.resources.read_text(
+      "google3.third_party.xprof.utils", "tpu_counter_ids_v7x.h"
+  )
+  print("Success with importlib.resources!")
+  print(content[:100])
+except ImportError as e:
+  print(f"Failed with importlib.resources: {e}")
+
+try:
+  # Try with files()
+  ref = (
+      importlib.resources.files("google3.third_party.xprof.utils")
+      / "tpu_counter_ids_v7x.h"
+  )
+  content = ref.read_text()
+  print("Success with importlib.resources.files!")
+  print(content[:100])
+except ImportError as e:
+  print(f"Failed with importlib.resources.files: {e}")

plugin/xprof/profile_plugin.py

@@ -37,6 +37,7 @@
 
 from xprof import profile_io
 from xprof import version
+from xprof.convert import counter_extractor
 from xprof.convert import raw_to_tool_data as convert
 from xprof.standalone.tensorboard_shim import base_plugin
 from xprof.standalone.tensorboard_shim import plugin_asset_util
@@ -1211,6 +1212,22 @@ def data_impl(
     """
     run = request.args.get('run')
     tool = request.args.get('tag')
+
+    if tool == 'perf_counters' and request.args.get('names_only') == '1':
+      device_type = request.args.get('device_type')
+      if not device_type:
+        raise ValueError(
+            'device_type is required for perf_counters with names_only'
+        )
+      try:
+        names = counter_extractor.get_all_counters(device_type)
+      except FileNotFoundError:
+        logger.warning(
+            'Failed to get counter names for device type: %s', device_type
+        )
+        names = []
+      return json.dumps(names), 'application/json', None
+
     hosts_param = request.args.get('hosts')
     host = request.args.get('host')
     module_name = request.args.get('module_name')

plugin/xprof/profile_plugin_test.py

@@ -377,6 +377,47 @@ def testData(self):
           )
       )
 
+  def testDataNamesOnly_HappyPath(self):
+    data, content_type, _ = self.plugin.data_impl(
+        utils.make_data_request(
+            utils.DataRequestOptions(
+                tool='perf_counters', names_only='1', device_type='v7x'
+            )
+        )
+    )
+    self.assertEqual(content_type, 'application/json')
+    names = json.loads(data)
+    self.assertIsInstance(names, list)
+    self.assertNotEmpty(names)
+    self.assertIn('name', names[0])
+    self.assertIn('val', names[0])
+
+  def testDataNamesOnly_MissingDeviceType(self):
+    with self.assertRaisesRegex(ValueError, 'device_type is required'):
+      self.plugin.data_impl(
+          utils.make_data_request(
+              utils.DataRequestOptions(tool='perf_counters', names_only='1')
+          )
+      )
+
+  def testDataNamesOnly_FileNotFoundFallback(self):
+    with mock.patch.object(
+        profile_plugin.counter_extractor, 'get_all_counters'
+    ) as mock_get_counters:
+      mock_get_counters.side_effect = FileNotFoundError(
+          'Simulated missing file'
+      )
+      data, content_type, _ = self.plugin.data_impl(
+          utils.make_data_request(
+              utils.DataRequestOptions(
+                  tool='perf_counters', names_only='1', device_type='v7x'
+              )
+          )
+      )
+      self.assertEqual(content_type, 'application/json')
+      names = json.loads(data)
+      self.assertEqual(names, [])
+
   def testDataWithCache(self):
     generate_testdata(self.logdir)
     self.multiplexer.AddRunsFromDirectory(self.logdir)