Tree traversal refactor

[kyle-hoffmeyer]

NOTE: CSV changes are due to changes in synthetic op classifications and UID shifts for synthetic ops.

Summary

This PR refactors the tree traversal in tree_perf to be simpler and more comprehensive.

The existing traversal performed two passes over the tree:

1. A primary top -> bottom DFS search that started from the roots, collecting perf-modeled ops and leaf cpu_ops.
2. A separate pass that gathers kernels not traversed in the primary pass. For these kernels, it does a bottom -> top traversal, walking each kernel to find it's closest cpu op. These became synthetic operations.

The second pass re-derives information the first pass already has: "what is this kernel's nearest cpu_op, and what are its shape args" leading to unncessary traversals. The second pass is also detached from the live call stack, so synthetic ops don't have their call stacks populated. Generating the call stacks for these synthetic ops in this second pass leads to redundant traversals.

Solution

This PR modifies the tree traversal to perform the entire tree traversal in the primary pass and remove the second pass while still generating synthetic operations and complete call stacks for all operations. Synthetic ops now have call stack information and redundant tree traversals are avoided.

Implementation

All changes are in TraceLens/TreePerf/tree_perf.py, in and around collect_unified_perf_events.

Single unified set of roots. The old code seeded the traversal from two separate root loops — parentless python_function roots (only when add_python_func was on), then cpu_root_nodes:

# When python_function events are in the tree, start from
# parentless python_function roots that have GPU work
if self.add_python_func:
    for evt in self.tree.events:
        if (
            self.event_to_category(evt) == "python_function"
            and evt.get("parent") is None
            and evt.get("gpu_events")
        ):
            traverse(evt["UID"])

# Start from cpu_root_nodes (any not already visited via python roots)
for root_uid in self.tree.cpu_root_nodes:
    traverse(root_uid)

Neither loop enumerated a bare kernel launchers that are itself parentless, so parentless launchers were never entered by the main pass and were left to the post-pass. Both loops are replaced by one:

for evt in self.tree.events:
    if evt.get("parent") is None and evt.get("gpu_events"):
        traverse(evt["UID"])

Starting from every parentless event that has GPU work covers cpu_op roots, python_function roots, and bare cuda_runtime launchers with no CPU op above them. Combined with removing the early return that skipped non-cpu_op events when add_python_func was off, this is what lets orphan launchers be reached during the main pass regardless of add_python_func.

nearest_cpu_op threaded down the recursion. traverse now carries a third argument, the innermost cpu_op ancestor thus far in the traversal. This single value is the discriminator between the two kinds of synthetic ops:

• A synthetic op under a cpu_op: Here cpu_op_E is not a leaf operation (it contains a finer kernel-launching cpu_op), so it recurses; cpu_op_C/k_C gets its own normal row, and cpu_op_E's own directly-launched k_own is left over. When the recursion reaches that launcher, nearest_cpu_op is cpu_op_E, so a cpu_op_E->k_own (Synthetic Op) is created:

  cpu_op_E                          (nearest_cpu_op for its subtree = cpu_op_E)
  ├── cpu_op_C -> runtime -> k_C    (finer cpu_op → its own normal row)
  └── runtime -> k_own              (E's own kernel, no cpu_op between)
  

  → cpu_op_E->k_own (Synthetic Op), args from cpu_op_E.

• An orphan (no cpu_op anywhere on the path): Here there is no cpu_op above the kernel at all, so nearest_cpu_op is None when the launcher is reached (Exit 4):

  py_func_A                         (python_function root, no cpu_op anywhere)
  └── hipLaunchKernel -> kernel_A   (nearest_cpu_op = None)
  

  → hipLaunchKernel->kernel_A (Synthetic Op), args from the launcher.

Strict definition of a "leaf" cpu_op. Previously _is_leaf_cpu_op would return True if it directly launched a kernel even if it launched other kernel-launching cpu ops. Now, if the cpu_op has calls kernel-launching cpu ops, it returns False. This change was made, because previously a leaf cpu_op would collect its entire subtree as one row, subsuming any nested kernel-launching cpu_ops. Now such an op is no longer a leaf, it recurses, the nested cpu_op gets its own row, and the parent's own directly-launched kernels become synthetic ops.
This is the one behavioral change — nested kernel-launching cpu_ops now always surface at their own granularity, independent of whether anything in the subtree is perf-modeled.

Renamed and broadened the descendant check. _has_descendant_cpu_op_with_own_perf_model → _has_descendant_cpu_op_with_kernels. The old predicate only fired when a descendant cpu_op had both a perf model and GPU work; the new one fires for any descendant cpu_op that launches kernels:

-            if self._has_perf_model(node) and self._launches_gpu_kernels(node):
+            if self._launches_gpu_kernels(node):
                 return True

The rationale is granularity consistency: the innermost cpu_op should own its kernels whether or not some descendant happens to be perf-modeled, so recursion shouldn't hinge on a perf-model flag. This predicate now drives both the strict-leaf rule above and the backward-op exit (Exit 2), which previously used the perf-model-only variant.

Three small helpers factor out the inline synthesis:

• _create_synthetic_op(prefix_name, kernel, cpu_op) — appends one "<prefix>-><kernel> (Synthetic Op)" row owning exactly that kernel, as a copy of cpu_op (so it inherits its shape args) with a fresh UID. cpu_op is the enclosing cpu_op, or the launcher itself when there is no cpu_op on the path.
• _direct_kernels(event) — the kernels whose innermost cpu_op is event: descends through non-cpu_op children and stops at any nested cpu_op. No depth cap, so it still reaches very deep launchers.
• _recurse(event, call_stack, child_nearest_cpu_op, is_cpu_op) — recurses into all children and, if event is a cpu_op, emits a synthetic op per direct kernel. Called from Exit 2 and the generic recurse.

New parentless launcher colelction. When there are no cpu ops in the subtree, the traversal emits one synthetic row per directly-launched kernel child of the launcher, then recurses to reach deeper launchers — replacing the deleted parent_to_orphans grouping in the old post-pass.

Tests

tests/test_synthetic_op.py builds small hand-crafted trees and asserts the row classification and the every-kernel-once invariant directly against build_df_unified_perf_table.

Row-classification matrix

To make the traversal's behavior concrete, consider this one fixed tree shape and vary only whether each cpu_op has a perf model:

cpu_op_A
├── cpu_op_B -> kernel_B
├── cpu_op_C -> kernel_C
└── kernel_A            (A's own kernel, launched directly by A)

cpu_op_A is never a leaf here — it contains nested kernel-launching cpu_ops (B, C), so under the strict-leaf rule it can only be collected via a perf model; otherwise it recurses. That single fact drives the whole table:

A perf model?	B perf model?	C perf model?	Rows produced (row → kernels owned)	Why
true	true/false	true/false	cpu_op_A → {kernel_A, kernel_B, kernel_C}	A perf-modeled op owns its entire subtree and the traversal stops — it never descends, so B/C are subsumed and kernel_A yields no synthetic op. B/C's own perf-model status is irrelevant.
false	true	true	cpu_op_B → {kernel_B} cpu_op_C → {kernel_C} cpu_op_A->kernel_A (Synthetic Op) → {kernel_A}	A has no perf model and isn't a leaf → recurses. B and C each own their kernel. A's own kernel_A is left over → synthetic op sourced from A.
false	true	false	cpu_op_B → {kernel_B} cpu_op_C → {kernel_C} cpu_op_A->kernel_A (Synthetic Op) → {kernel_A}	Same shape. Both B and C are leaf cpu_ops.
false	false	true	cpu_op_B → {kernel_B} cpu_op_C → {kernel_C} cpu_op_A->kernel_A (Synthetic Op) → {kernel_A}	Same shape. Both B and C are leaf cpu_ops.
false	false	false	cpu_op_B → {kernel_B} cpu_op_C → {kernel_C} cpu_op_A->kernel_A (Synthetic Op) → {kernel_A}	Same shape. Both B and C are leaf cpu_ops.

Tests

tests/test_synthetic_op.py builds small hand-crafted trees and asserts the row classification and the every-kernel-once invariant directly against build_df_unified_perf_table.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 98.21429% with 1 line in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
TraceLens/TreePerf/tree_perf.py	98.21%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!