Skip to content

[Lang] Cut per-compilation Python overhead (sys.modules scan + per-node source formatting)#804

Draft
duburcqa wants to merge 2 commits into
mainfrom
duburcqa/inside-class-linecache
Draft

[Lang] Cut per-compilation Python overhead (sys.modules scan + per-node source formatting)#804
duburcqa wants to merge 2 commits into
mainfrom
duburcqa/inside-class-linecache

Conversation

@duburcqa

@duburcqa duburcqa commented Jul 21, 2026

Copy link
Copy Markdown
Contributor

Two independent reductions of Python-side overhead paid on every kernel compilation. Both are pure
performance changes with no functional effect on generated kernels.

1. Class-kernel detection: drop the sys.modules scan

_inside_class runs once per kernel creation and used inspect.getframeinfo(frame), which resolves the
frame's module through inspect.getmodule — an O(len(sys.modules)) scan over every loaded module. In
applications importing many modules (a full physics/ML stack, several thousand modules) this runs on every
kernel materialization and dominates kernel build time. Replaced with linecache.getline(co_filename, f_lineno), which returns the same decoration-site line without touching sys.modules. inspect is no
longer needed here.

2. Kernel build: attach only a cheap source header per AST node

ASTTransformer attached full source-position info to every generated IR node via get_pos_info, which
builds a TextWrapper-formatted source-line-plus-caret hint for each node. That formatting ran for every
AST node of every kernel compilation. Now the eager per-node path attaches only the cheap
File "...", line N, in fn header (get_pos_header); the full hint is still produced on the actual
compile-error path (the except handler calls get_pos_info), so compile-error messages are unchanged.

Impact

Full-suite measurement on the Genesis test suite (1186 tests, 8 GPUs, forked; both changes vs the unpatched
baseline, same Scene.build instrumentation for both):

  • total Scene.build: 58,094s -> 50,808s (-12.5%)
  • total test call time: 74,225s -> 66,691s (-10.2%)
  • all tests pass (1186 passed)

Approximate per-lever split — the combined figure above is measured; the split below is inferred
(from the isolation A/B plus a gutted-formatting run), not independently measured:

Isolation A/B on a single kernel-build-heavy test (interleaved, warm caches): median wall ~29.6s -> ~25.7s.

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 53f3064db1

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

# resolves the frame's module through inspect.getmodule, an O(len(sys.modules)) scan run once per kernel
# creation. With thousands of modules loaded this dominates kernel build time; linecache.getline returns
# the same source line with no such scan.
first_statment = linecache.getline(maybe_class_frame.f_code.co_filename, maybe_class_frame.f_lineno).strip()

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P2 Badge Preserve loader-backed source lookup

When user kernels are defined in modules loaded from zipimport/zipapps/PEX-style archives, co_filename is a virtual path like /app.pyz/pkg/mod.py, and linecache.getline(filename, lineno) returns an empty string unless module globals/loader state are supplied. The previous inspect.getframeinfo path routed through inspect.findsource()/linecache.getlines(file, module.__dict__), so the loader could provide the source; with this change _inside_class returns False for @qd.kernel/@qd.func methods in those packaged modules, causing @qd.data_oriented not to wrap class kernels and self to be treated as a normal kernel argument.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Has this been addressed? If so, how?

I have mixed feelings about this codex comment. On the one hand, we dont use any zipimport style archives. On the other hand, being silently wrong seems not great. We should either explicitly forbid this and throw an exception, or somehow ensure this is correct (eg falling back on the old method), I feel.

@github-actions

Copy link
Copy Markdown

@duburcqa duburcqa changed the title [Lang] Speed up kernel build by avoiding a sys.modules scan in class-kernel detection [Lang] Cut per-compilation Python overhead (sys.modules scan + per-node source formatting) Jul 21, 2026
@github-actions

Copy link
Copy Markdown

@hughperkins

Copy link
Copy Markdown
Collaborator

@codex review

@hughperkins

Copy link
Copy Markdown
Collaborator

Interesting

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: b1817f7b02

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

# resolves the frame's module through inspect.getmodule, an O(len(sys.modules)) scan run once per kernel
# creation. With thousands of modules loaded this dominates kernel build time; linecache.getline returns
# the same source line with no such scan.
first_statment = linecache.getline(maybe_class_frame.f_code.co_filename, maybe_class_frame.f_lineno).strip()

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P2 Badge Refresh linecache before reading decorator lines

When a long-lived process redefines kernels from a filename that is already in linecache after the file has changed (for example, importlib.reload()/hot-reload after inserting lines above a @qd.data_oriented class), this can read the stale cached line. The old inspect.getframeinfo() path invalidated the cache via linecache.checkcache() before returning context, but linecache.getline() does not, so _inside_class() can match an old class/decorator line or miss the new one and misclassify kernels, breaking method binding through data_oriented. Refresh this filename before getline() or otherwise avoid stale cache reads.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

again, this seems like not something we do typically, but silently being incorrect seems not great?

@hughperkins

Copy link
Copy Markdown
Collaborator

For anyone else who didnt know, what is linecache:

Screenshot 2026-07-21 at 08 23 01

@duburcqa
duburcqa marked this pull request as draft July 22, 2026 12:48
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