Skip to content

Runtime improvement for pseudo op extension #617

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
devalshahamd wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Runtime improvement for pseudo op extension #617

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e6265b5
Runtime improvement for pseudo op extension
devalshah-amd Apr 30, 2026
77151c8
reformatting
devalshah-amd Apr 30, 2026
File filter

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 39 additions & 5 deletions TraceLens/Trace2Tree/extensions/pseudo_ops_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -159,18 +159,48 @@ def inject_pseudo_op_wrap_children(
parent_evt["children"] = [pseudo_evt["UID"]]
pseudo_evt["parent"] = parent_evt["UID"]

# Descendants that were cpu_root_nodes are no longer roots since they
# now live under the pseudo op. Remove them and promote the pseudo op.
root_set = set(tree.cpu_root_nodes)
# Descendants that were cpu_root_nodes are no longer roots since they now
# live under the pseudo op. Removals/additions are accumulated into per-tree
# pending sets and callers MUST invoke ``finalize_pseudo_op_mutations(tree)`` after the
# last call to commit the accumulated changes back to ``tree.cpu_root_nodes``.
pending_removals = getattr(tree, "_cpu_root_nodes_pending_removals", None)
if pending_removals is None:
pending_removals = set()
tree._cpu_root_nodes_pending_removals = pending_removals
tree._cpu_root_nodes_pending_additions = []
tree._cpu_root_nodes_set = set(tree.cpu_root_nodes)
pending_additions = tree._cpu_root_nodes_pending_additions
root_set = tree._cpu_root_nodes_set

stack = list(children_uids)
while stack:
uid = stack.pop()
if uid in root_set:
tree.cpu_root_nodes.remove(uid)
pending_removals.add(uid)
root_set.discard(uid)
evt = tree.get_UID2event(uid)
stack.extend(evt.get("children", []))
tree.cpu_root_nodes.append(pseudo_evt["UID"])
pending_additions.append(pseudo_evt["UID"])
Comment on lines +179 to +183

Copilot AI Apr 30, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The pending-root tracking doesn’t account for pseudo-ops added during the current pending-mutation window. Because root_set is only initialized from the original tree.cpu_root_nodes and never updated with pseudo_evt['UID'], a subsequent inject_pseudo_op_wrap_children that wraps a subtree containing an earlier pseudo-op will fail to mark that earlier pseudo-op for removal (the old eager implementation would remove it). Additionally, if you do add pseudo UIDs into root_set, ensure removals take precedence over additions (e.g., filter pending_additions against pending_removals in finalize_pseudo_op_mutations).
Copilot uses AI. Check for mistakes.


def finalize_pseudo_op_mutations(tree):
"""Commit any pending cpu_root_nodes mutations from
``inject_pseudo_op_wrap_children`` back to ``tree.cpu_root_nodes``
"""
pending_removals = getattr(tree, "_cpu_root_nodes_pending_removals", None)
if pending_removals is None:
return
pending_additions = tree._cpu_root_nodes_pending_additions
if pending_removals:
tree.cpu_root_nodes = [
u for u in tree.cpu_root_nodes if u not in pending_removals
]
if pending_additions:
tree.cpu_root_nodes.extend(pending_additions)
del tree._cpu_root_nodes_pending_removals
del tree._cpu_root_nodes_pending_additions
Comment on lines +193 to +201

Copilot AI Apr 30, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

finalize_pseudo_op_mutations introduces new deferred-mutation semantics for cpu_root_nodes, but there’s no test coverage exercising the commit behavior (including edge cases like multiple injections and nested wrapping where a previously-added pseudo-op root should be removed). Adding a focused unit test would help prevent regressions in root-node bookkeeping.

Suggested change
pending_additions = tree._cpu_root_nodes_pending_additions
if pending_removals:
tree.cpu_root_nodes = [
u for u in tree.cpu_root_nodes if u not in pending_removals
]
if pending_additions:
tree.cpu_root_nodes.extend(pending_additions)
del tree._cpu_root_nodes_pending_removals
del tree._cpu_root_nodes_pending_additions
pending_additions = getattr(tree, "_cpu_root_nodes_pending_additions", [])
if pending_removals:
tree.cpu_root_nodes = [
u for u in tree.cpu_root_nodes if u not in pending_removals
]
if pending_additions:
tree.cpu_root_nodes.extend(pending_additions)
if hasattr(tree, "_cpu_root_nodes_pending_removals"):
del tree._cpu_root_nodes_pending_removals
if hasattr(tree, "_cpu_root_nodes_pending_additions"):
del tree._cpu_root_nodes_pending_additions
Copilot uses AI. Check for mistakes.
if hasattr(tree, "_cpu_root_nodes_set"):
del tree._cpu_root_nodes_set


def apply_pseudo_op_extensions(tree, verbose: bool = False):
Expand Down Expand Up @@ -256,3 +286,7 @@ def apply_pseudo_op_extensions(tree, verbose: bool = False):
ext_func(tree)
except Exception as e:
logger.warning(f"Failed to apply pseudo-op extension {ext_name}: {e}")

# Commit any pending cpu_root_nodes mutations accumulated by
# inject_pseudo_op_wrap_children (no-op if none were made).
finalize_pseudo_op_mutations(tree)
Loading