| # Copyright 2026 The IREE Authors |
| # |
| # Licensed under the Apache License v2.0 with LLVM Exceptions. |
| # See https://llvm.org/LICENSE.txt for license information. |
| # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| |
| import argparse |
| import json |
| import os |
| import subprocess |
| import sys |
| import tempfile |
| import unittest |
| |
| _TOOL_PATHS = argparse.Namespace(fixture_generator=None, iree_profile=None) |
| |
| |
| def _runfile_path(relative_path): |
| workspace = os.environ.get("TEST_WORKSPACE") |
| roots = [ |
| os.environ.get("RUNFILES_DIR"), |
| os.environ.get("TEST_SRCDIR"), |
| os.getcwd(), |
| ] |
| suffixes = ["", ".exe"] if os.name == "nt" else [""] |
| for root in roots: |
| if not root: |
| continue |
| prefixes = [root] |
| if workspace: |
| prefixes.insert(0, os.path.join(root, workspace)) |
| for prefix in prefixes: |
| for suffix in suffixes: |
| path = os.path.join(prefix, relative_path) + suffix |
| if os.path.exists(path): |
| return path |
| raise FileNotFoundError(relative_path) |
| |
| |
| def _run_checked(argv): |
| completed = subprocess.run( |
| argv, |
| check=False, |
| stdout=subprocess.PIPE, |
| stderr=subprocess.PIPE, |
| text=True, |
| ) |
| if completed.returncode != 0: |
| command = " ".join(argv) |
| raise AssertionError( |
| f"command failed: {command}\n" |
| f"exit_code={completed.returncode}\n" |
| f"stdout:\n{completed.stdout}\n" |
| f"stderr:\n{completed.stderr}" |
| ) |
| return completed |
| |
| |
| def _run_jsonl(argv): |
| completed = _run_checked(argv) |
| return [json.loads(line) for line in completed.stdout.splitlines() if line.strip()] |
| |
| |
| def _find_row(rows, row_type): |
| for row in rows: |
| if row.get("type") == row_type or row.get("record_type") == row_type: |
| return row |
| raise AssertionError(f"missing row type {row_type!r}: {rows!r}") |
| |
| |
| def _row_type_set(rows): |
| return {row.get("type", row.get("record_type")) for row in rows} |
| |
| |
| def _parse_tool_paths(argv): |
| parser = argparse.ArgumentParser(add_help=False) |
| parser.add_argument("--fixture-generator") |
| parser.add_argument("--iree-profile") |
| return parser.parse_known_args(argv) |
| |
| |
| class IreeProfileCliTest(unittest.TestCase): |
| |
| @classmethod |
| def setUpClass(cls): |
| cls.fixture_generator = ( |
| _TOOL_PATHS.fixture_generator |
| if _TOOL_PATHS.fixture_generator |
| else _runfile_path( |
| "runtime/src/iree/tooling/profile/test_fixture_generator" |
| ) |
| ) |
| cls.iree_profile = ( |
| _TOOL_PATHS.iree_profile |
| if _TOOL_PATHS.iree_profile |
| else _runfile_path("tools/iree-profile") |
| ) |
| |
| def _create_profile(self, directory, *fixture_flags): |
| profile_path = os.path.join(directory, "smoke.ireeprof") |
| _run_checked( |
| [self.fixture_generator, *fixture_flags, profile_path], |
| ) |
| return profile_path |
| |
| def _profile_jsonl(self, profile_path, command, *flags): |
| return _run_jsonl( |
| [self.iree_profile, "--format=jsonl", *flags, command, profile_path] |
| ) |
| |
| def test_report_commands_load_production_sink_bundle(self): |
| with tempfile.TemporaryDirectory() as directory: |
| profile_path = self._create_profile(directory) |
| summary_rows = self._profile_jsonl(profile_path, "summary") |
| queue_rows = self._profile_jsonl(profile_path, "queue") |
| command_rows = self._profile_jsonl(profile_path, "command") |
| executable_rows = self._profile_jsonl(profile_path, "executable") |
| dispatch_rows = self._profile_jsonl(profile_path, "dispatch") |
| statistics_rows = self._profile_jsonl(profile_path, "statistics") |
| |
| self.assertIn("summary", _row_type_set(summary_rows)) |
| self.assertIn("device_summary", _row_type_set(summary_rows)) |
| summary = _find_row(summary_rows, "summary") |
| self.assertGreater(summary["device_metric_sample_records"], 0) |
| self.assertIn("queue_event", _row_type_set(queue_rows)) |
| self.assertIn("queue_device_event", _row_type_set(queue_rows)) |
| self.assertIn("host_execution_event", _row_type_set(queue_rows)) |
| self.assertIn("command_operation", _row_type_set(command_rows)) |
| self.assertIn("command_execution", _row_type_set(command_rows)) |
| self.assertIn("command_host_execution", _row_type_set(command_rows)) |
| self.assertIn( |
| "executable_function_host_dispatch_group", _row_type_set(executable_rows) |
| ) |
| self.assertIn("host_dispatch_group", _row_type_set(dispatch_rows)) |
| self.assertIn("statistics_summary", _row_type_set(statistics_rows)) |
| self.assertIn("statistics_row", _row_type_set(statistics_rows)) |
| |
| def test_statistics_reports_aggregate_row_families(self): |
| with tempfile.TemporaryDirectory() as directory: |
| rows = self._profile_jsonl(self._create_profile(directory), "statistics") |
| |
| summary = _find_row(rows, "statistics_summary") |
| self.assertGreater(summary["row_count"], 0) |
| self.assertEqual(summary["dropped_record_count"], 0) |
| |
| row_types = { |
| row["row_type"] for row in rows if row.get("type") == "statistics_row" |
| } |
| self.assertIn("dispatch_function", row_types) |
| self.assertIn("queue_device_operation", row_types) |
| self.assertIn("queue_host_operation", row_types) |
| self.assertIn("host_execution_function", row_types) |
| self.assertIn("memory_lifecycle", row_types) |
| |
| function_row = next( |
| row |
| for row in rows |
| if row.get("type") == "statistics_row" |
| and row.get("row_type") == "dispatch_function" |
| ) |
| self.assertEqual(function_row["function_name"], "smoke_function") |
| self.assertGreater(function_row["total_duration_ns"], 0) |
| |
| def test_queue_projection_separates_host_and_device_time(self): |
| with tempfile.TemporaryDirectory() as directory: |
| rows = self._profile_jsonl(self._create_profile(directory), "queue") |
| |
| queue_event = _find_row(rows, "queue_event") |
| self.assertEqual(queue_event["host_time_domain"], "iree_host_time_ns") |
| self.assertLessEqual( |
| queue_event["host_time_ns"], queue_event["ready_host_time_ns"] |
| ) |
| |
| device_event = _find_row(rows, "queue_device_event") |
| self.assertIn("start_tick", device_event) |
| self.assertIn("end_tick", device_event) |
| self.assertIn("duration_ns", device_event) |
| |
| host_event = _find_row(rows, "host_execution_event") |
| self.assertEqual(host_event["host_time_domain"], "iree_host_time_ns") |
| self.assertEqual(queue_event["submission_id"], device_event["submission_id"]) |
| self.assertEqual(queue_event["submission_id"], host_event["submission_id"]) |
| self.assertEqual(queue_event["op"], device_event["op"]) |
| self.assertEqual(queue_event["op"], host_event["op"]) |
| |
| def test_command_and_dispatch_projections_join_ids(self): |
| with tempfile.TemporaryDirectory() as directory: |
| profile_path = self._create_profile(directory) |
| command_rows = self._profile_jsonl(profile_path, "command") |
| dispatch_rows = self._profile_jsonl( |
| profile_path, "dispatch", "--dispatch_events" |
| ) |
| |
| operation = _find_row(command_rows, "command_operation") |
| self.assertEqual(operation["op"], "dispatch") |
| |
| execution = _find_row(command_rows, "command_execution") |
| self.assertEqual(operation["command_buffer_id"], execution["command_buffer_id"]) |
| |
| dispatch_event = _find_row(dispatch_rows, "dispatch_event") |
| self.assertEqual( |
| operation["command_buffer_id"], dispatch_event["command_buffer_id"] |
| ) |
| self.assertEqual(operation["command_index"], dispatch_event["command_index"]) |
| self.assertEqual(dispatch_event["device_tick_domain"], "device_tick") |
| self.assertEqual( |
| dispatch_event["duration_time_domain"], "device_tick_duration_ns" |
| ) |
| |
| host_dispatch_event = _find_row(dispatch_rows, "host_dispatch_event") |
| self.assertEqual( |
| operation["command_buffer_id"], host_dispatch_event["command_buffer_id"] |
| ) |
| self.assertEqual( |
| operation["command_index"], host_dispatch_event["command_index"] |
| ) |
| self.assertEqual(host_dispatch_event["timing_source"], "host_execution_event") |
| self.assertEqual(host_dispatch_event["time_domain"], "iree_host_time_ns") |
| self.assertEqual( |
| host_dispatch_event["duration_time_domain"], "iree_host_duration_ns" |
| ) |
| |
| def test_explain_uses_host_dispatch_spans_without_device_dispatches(self): |
| with tempfile.TemporaryDirectory() as directory: |
| profile_path = self._create_profile(directory, "--omit-dispatch-events") |
| rows = self._profile_jsonl(profile_path, "explain") |
| |
| summary = _find_row(rows, "explain_summary") |
| self.assertEqual(summary["valid_dispatches"], 0) |
| self.assertGreater(summary["valid_host_dispatches"], 0) |
| |
| top_function = _find_row(rows, "explain_top_function") |
| self.assertEqual(top_function["timing_source"], "host_execution_event") |
| self.assertGreater(top_function["total_ns"], 0) |
| self.assertGreater(top_function["total_tile_count"], 0) |
| |
| top_dispatch = _find_row(rows, "explain_top_dispatch") |
| self.assertEqual(top_dispatch["timing_source"], "host_execution_event") |
| self.assertGreater(top_dispatch["duration_ns"], 0) |
| self.assertGreater(top_dispatch["tile_count"], 0) |
| |
| def test_export_ireeperf_jsonl_decodes_record_families(self): |
| with tempfile.TemporaryDirectory() as directory: |
| profile_path = self._create_profile(directory) |
| output_path = os.path.join(directory, "smoke.ireeperf.jsonl") |
| _run_checked( |
| [ |
| self.iree_profile, |
| "--format=ireeperf-jsonl", |
| f"--output={output_path}", |
| "export", |
| profile_path, |
| ] |
| ) |
| with open(output_path, "r", encoding="utf-8") as file: |
| rows = [json.loads(line) for line in file if line.strip()] |
| |
| schema = _find_row(rows, "schema") |
| self.assertEqual(schema["format"], "ireeperf-jsonl") |
| self.assertEqual(schema["row_key"], "record_type") |
| |
| queue_event = _find_row(rows, "queue_event") |
| self.assertEqual(queue_event["host_time_domain"], "iree_host_time_ns") |
| self.assertLessEqual( |
| queue_event["host_time_ns"], queue_event["ready_host_time_ns"] |
| ) |
| |
| device_event = _find_row(rows, "queue_device_event") |
| self.assertIn("derived_time_available", device_event) |
| self.assertIn("duration_ns", device_event) |
| |
| relationship = _find_row(rows, "event_relationship") |
| self.assertEqual(relationship["kind"], "queue_event_host_execution_event") |
| self.assertEqual(relationship["source_type"], "queue_event") |
| self.assertEqual(relationship["target_type"], "host_execution_event") |
| |
| metric_source = _find_row(rows, "device_metric_source") |
| self.assertEqual(metric_source["device_class"], "gpu") |
| self.assertEqual(metric_source["name"], "smoke.metrics") |
| |
| metric_descriptor = _find_row(rows, "device_metric_descriptor") |
| self.assertIn(metric_descriptor["unit"], {"count", "millipercent"}) |
| |
| metric_sample = _find_row(rows, "device_metric_sample") |
| self.assertEqual(metric_sample["host_time_domain"], "iree_host_time_ns") |
| self.assertEqual(len(metric_sample["values"]), 2) |
| value_names = {value["name"] for value in metric_sample["values"]} |
| self.assertIn("activity.compute", value_names) |
| self.assertIn("smoke.source_specific", value_names) |
| self.assertTrue(all("value" in value for value in metric_sample["values"])) |
| |
| |
| if __name__ == "__main__": |
| _TOOL_PATHS, remaining_argv = _parse_tool_paths(sys.argv[1:]) |
| unittest.main(argv=[sys.argv[0], *remaining_argv]) |
