Add perfetto support for counter collection

Fix endtimestamp for counter tracks Add fix for rocprofv3 counter collection tests Fix formats and refactors Added docs and addressed review comments Address more review comments.
2025-02-20 10:14:37 +00:00
parent 4dcb239872
commit c9ca876b79
9 changed files with 234 additions and 10 deletions
@@ -276,10 +276,73 @@ class PerfettoReader:
        """Extracts all the necessary data from the trace processor"""
        self.configure(**kwargs)

+        # generate empty dictionaries for each trace processor
+        self.track_ids = [{} for _ in range(len(self.trace_processor))]
+
        self.dataframe = self.query_tp(
            "SELECT slice_id, track_id, category, depth, stack_id, parent_stack_id, ts, dur, name FROM slice"
        )

+        counter_df = self.query_tp(
+            """SELECT
+                counter_track.id as slice_id,
+                counter.track_id,
+                counter_track.name as track_name,
+                'counter_collection' as category,
+                0 as depth,
+                0 as stack_id,
+                0 as parent_stack_id,
+                MIN(CASE WHEN counter.value > 0 THEN counter.ts ELSE NULL END) as ts,
+                0 as dur,
+                counter_track.name as name
+            FROM counter_track
+            JOIN counter ON counter.track_id = counter_track.id
+            WHERE counter_track.name LIKE 'AGENT%'
+            AND counter.value > 0
+            GROUP BY counter.track_id"""
+        )
+
+        # Transform counter data to match the main dataframe schema
+        if not counter_df.empty:
+            # Register counter track IDs in self.track_ids before adding to dataframe
+            for row in counter_df.itertuples():
+                if (
+                    row.tp_index < len(self.track_ids)
+                    and row.track_id not in self.track_ids[row.tp_index]
+                ):
+                    # Add the counter track to track_ids with reasonable default values
+                    self.track_ids[row.tp_index][row.track_id] = {
+                        "tp_index": row.tp_index,
+                        "pid": 0,
+                        "tid": 0,
+                        "rank": 0,
+                        "thread": 0,
+                        "prio": 2,
+                        "process_name": "counter_process",
+                        "thread_name": f"counter_track_{row.category}",
+                    }
+
+            # Create a new dataframe with the right columns
+            counter_collection_df = pd.DataFrame(
+                {
+                    "tp_index": counter_df["tp_index"],
+                    "slice_id": counter_df["slice_id"],
+                    "track_id": counter_df["track_id"],
+                    "category": "counter_collection",
+                    "depth": 0,
+                    "stack_id": 0,
+                    "parent_stack_id": 0,
+                    "ts": counter_df["ts"],
+                    "dur": 0,
+                    "name": counter_df["name"].astype(str),
+                }
+            )
+
+            # Concatenate with main dataframe
+            self.dataframe = pd.concat(
+                [self.dataframe, counter_collection_df], ignore_index=True
+            )
+
        self.df_categories = sorted(list(self.dataframe["category"].unique()))

        # check for update to include/exclude category
@@ -349,9 +412,6 @@ class PerfettoReader:
            "SELECT thread.utid AS thread_utid, thread.id AS thread_id, thread.tid, thread.name as thread_name, thread.is_main_thread, thread_track.id AS track_id, thread_track.parent_id AS track_parent_id, thread_track.name AS track_name from thread JOIN thread_track ON thread_track.utid = thread.utid"
        )

-        # generate empty dictionaries for each trace processor
-        self.track_ids = [{} for _ in range(len(self.trace_processor))]
-
        # generate mapping from track IDs to process and thread info.
        # the "pid" and "tid" fields are the system value. we want to
        # assign a "rank" and "thread" value for "pid" and "tid",
@@ -35,6 +35,7 @@ def test_perfetto_data(
        "memory_allocation",
        "rocdecode_api",
        "rocjpeg_api",
+        "counter_collection",
    ),
 ):

@@ -47,6 +48,7 @@ def test_perfetto_data(
        "memory_allocation": ("memory_allocation", "memory_allocation"),
        "rocdecode_api": ("rocdecode_api", "rocdecode_api"),
        "rocjpeg_api": ("rocjpeg_api", "rocjpeg_api"),
+        "counter_collection": ("counter_collection", "counter_collection"),
    }

    # make sure they specified valid categories
@@ -57,7 +59,23 @@ def test_perfetto_data(
        itr for key, itr in mapping.items() if key in categories
    ]:
        _pf_data = pftrace_data.loc[pftrace_data["category"] == pf_category]
-        _js_data = json_data["rocprofiler-sdk-tool"]["buffer_records"][js_category]
+
+        _js_data = []
+        if js_category != "counter_collection":
+            _js_data = json_data["rocprofiler-sdk-tool"]["buffer_records"][js_category]
+        else:
+            unique_counter_ids = set()
+
+            for dispatch_entry in json_data["rocprofiler-sdk-tool"]["callback_records"][
+                js_category
+            ]:
+                counter_records = dispatch_entry["records"]
+
+                for record in counter_records:
+                    counter_id = record["counter_id"]["handle"]
+                    unique_counter_ids.add(counter_id)
+
+            _js_data = [{"counter_id": id} for id in unique_counter_ids]

        assert len(_pf_data) == len(
            _js_data