[metal] Add @metal_jit decorator for AST-to-MSL compilation#1991
Merged
[metal] Add @metal_jit decorator for AST-to-MSL compilation#1991
Conversation
aditvenk
force-pushed
the
aditvenk/stack/25
branch
from
276e63c to
951fddc
Compare
This was referenced Apr 9, 2026
meta-cla
bot
added
the
CLA Signed
aditvenk
force-pushed
the
aditvenk/stack/25
branch
from
951fddc to
67b7d10
Compare
aditvenk
force-pushed
the
aditvenk/stack/25
branch
from
67b7d10 to
26519e4
Compare
aditvenk
force-pushed
the
aditvenk/stack/25
branch
from
26519e4 to
0584fad
Compare
jansel
approved these changes
Apr 13, 2026
Add metal_jit decorator that JIT-compiles a Python function to an MSL Metal shader. On each call the decorator: 1. Parses the decorated function's source to recover the Python AST 2. Infers tensor dtypes and scalar args from the actual call arguments 3. Reads block sizes from module-level globals (_BLOCK_SIZE_N) 4. Translates the AST body to MSL C++ via _generate_msl 5. Compiles the MSL via torch.mps.compile_shader No decorator arguments needed — all metadata is inferred at runtime. stack-info: PR: #1991, branch: aditvenk/stack/25
aditvenk
force-pushed
the
aditvenk/stack/25
branch
from
0584fad to
713507e
Compare
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.
Stacked PRs:
[metal] Add @metal_jit decorator for AST-to-MSL compilation
Add metal_jit decorator that JIT-compiles a Python function to an MSL
Metal shader. On each call the decorator:
No decorator arguments needed — all metadata is inferred at runtime.