Skip to content

[FIX] Add __call__ method to Trace class for nested JIT function calls - #124

Merged
mark14wu merged 5 commits into
mainfrom
fix_trace_callable
Sep 19, 2025
Merged

mark14wu merged 5 commits into
mainfrom
fix_trace_callable

Conversation

@mark14wu

@mark14wu mark14wu commented Sep 12, 2025 •

Copy link
Copy Markdown
Collaborator

When a traced JIT function is called from within another JIT function, the Trace wrapper needs to properly delegate to the underlying function. See #135.

mark14wu and others added 2 commits September 12, 2025 08:02
When a traced JIT function is called from within another JIT function,
the Trace wrapper needs to properly delegate to the underlying function.
This fix ensures compatibility with nested Triton JIT function calls.

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <noreply@anthropic.com>
Comment thread triton_viz/core/trace.py
self.finalize()
return ret

def __call__(self, *args, **kwargs):

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

We should interpret this function as well using its interpreter_fn

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Otherwise you cannot intercept any loads and stores in this function.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Please also add a test

@mark14wu
mark14wu merged commit f70956b into main Sep 19, 2025
1 check passed
@mark14wu
mark14wu deleted the fix_trace_callable branch September 19, 2025 13:22
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants