[Lang] Cut per-compilation Python overhead (sys.modules scan + per-node source formatting)#804
[Lang] Cut per-compilation Python overhead (sys.modules scan + per-node source formatting)#804duburcqa wants to merge 2 commits into
Conversation
There was a problem hiding this comment.
💡 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() |
There was a problem hiding this comment.
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 👍 / 👎.
There was a problem hiding this comment.
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.
|
@codex review |
|
Interesting |
There was a problem hiding this comment.
💡 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() |
There was a problem hiding this comment.
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 👍 / 👎.
There was a problem hiding this comment.
again, this seems like not something we do typically, but silently being incorrect seems not great?

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.modulesscan_inside_classruns once per kernel creation and usedinspect.getframeinfo(frame), which resolves theframe's module through
inspect.getmodule— anO(len(sys.modules))scan over every loaded module. Inapplications 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 touchingsys.modules.inspectis nolonger needed here.
2. Kernel build: attach only a cheap source header per AST node
ASTTransformerattached full source-position info to every generated IR node viaget_pos_info, whichbuilds a
TextWrapper-formatted source-line-plus-caret hint for each node. That formatting ran for everyAST node of every kernel compilation. Now the eager per-node path attaches only the cheap
File "...", line N, in fnheader (get_pos_header); the full hint is still produced on the actualcompile-error path (the
excepthandler callsget_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.buildinstrumentation for both):Scene.build: 58,094s -> 50,808s (-12.5%)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:
Scene.buildsuite-wide.sys.modulesscan): ~-2% ofScene.buildsuite-wide, but ~+20% in isolation on small,build-light tests (it removes an
O(sys.modules)scan paid per kernel creation, which dominates a tinybuild but is a small slice of a large one).
Isolation A/B on a single kernel-build-heavy test (interleaved, warm caches): median wall ~29.6s -> ~25.7s.