[FIX] Derive constexpr names from JITFunction to fix string annotations#330
Open
mark14wu wants to merge 6 commits into
Open
[FIX] Derive constexpr names from JITFunction to fix string annotations#330mark14wu wants to merge 6 commits into
mark14wu wants to merge 6 commits into
Conversation
Verify the sanitizer correctly handles tl.constexpr parameters passed as positional (non-keyword) arguments to a kernel using tl.arange.
Sanitizer Performance Benchmark
Iterations: 1 warmup + 20 measured |
On Python 3.14+ the AST rewriter double-quotes string annotations (e.g. 'tl.constexpr' becomes "'tl.constexpr'"), causing GridExecutor.constexprs to be empty. This makes constexpr parameters go through _implicit_cvt, turning them into tensors, which then causes tl.arange to raise ValueError. Fix: derive constexpr names from the original JITFunction.params (passed via jit_fn kwarg) instead of relying on self.constexprs from the rewritten function.
mark14wu
changed the title
Jokeren
reviewed
Mar 16, 2026
| # reliable annotation info. The rewritten function's __annotations__ | ||
| # can be corrupted on Python 3.14+ (string annotations get | ||
| # double-quoted), making self.constexprs unreliable. | ||
| if jit_fn is not None and hasattr(jit_fn, "params"): |
Member
There was a problem hiding this comment.
Why not just keep the code in the if path? That is, only checking jit_fn.params
Collaborator
Author
There was a problem hiding this comment.
Good call — removed the fallback. jit_fn is always passed through kwargs from trace.py:run, so the else branch was dead code. Simplified to just {p.name for p in jit_fn.params if p.is_constexpr}.
Always derive constexpr names from jit_fn.params since jit_fn is always available in this code path. Addresses review feedback.
jit_fn is None for nested JIT calls and some Autotune paths. Fall back to self.constexprs when jit_fn is unavailable.
Jokeren
approved these changes
Mar 18, 2026
Collaborator
Author
|
@codex review |
|
Codex Review: Didn't find any major issues. Chef's kiss. ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
Member
|
Closing as stale during PR cleanup. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
'tl.constexpr'becomes"'tl.constexpr'"), causingGridExecutor.constexprsto be empty. This makes constexpr parameters go through_implicit_cvt, turning them into tensors, which then causestl.arangeto raiseValueError: arange's arguments must be of type tl.constexpr.JITFunction.params(passed viajit_fnkwarg) instead of relying onself.constexprsfrom the rewritten function."tl.constexpr"passed as a positional argument.Test plan
test_constexpr_string_annotation_positional_argpasses with fix, fails without