fix(jit): Carry a tensor annotation's layout slot through specialization - #2258
Open
lyfne123 wants to merge 2 commits into
Open
fix(jit): Carry a tensor annotation's layout slot through specialization#2258lyfne123 wants to merge 2 commits into
lyfne123 wants to merge 2 commits into
Conversation
📝 WalkthroughWalkthroughJIT tensor annotations now preserve supported layouts through metadata extraction, specialization, and cache keys. Unsupported layouts raise explicit errors. Source maps include signature locations, and tests and documentation cover layout behavior. ChangesJIT Tensor Layout Handling
Estimated code review effort: 3 (Moderate) | ~30 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant JITDecorator
participant TensorSpecializer
participant CacheKey
Caller->>JITDecorator: invoke JIT function with tensor annotations
JITDecorator->>TensorSpecializer: pass tensor shape, dtype, and layout metadata
TensorSpecializer->>CacheKey: include entry and dependency layouts
CacheKey-->>JITDecorator: return layout-specific cache key
JITDecorator-->>Caller: compile or reuse specialized function
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Warning Review ran into problems🔥 ProblemsGit: Failed to clone repository. Please run the Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 445d816da9
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
lyfne123
force-pushed
the
fix/jit-propagate-tensor-layout
branch
from
August 3, 2026 08:48
445d816 to
8096d62
Compare
@pl.jit does not resolve the user's annotations — it records shape and dtype in TensorMeta and regenerates the annotation from that record. TensorMeta had no layout field and _build_tensor_annotation emitted only two slots, so the third one never survived: pl.NZ silently compiled as ND, and pl.DN never reached the type resolver, so it neither errored nor took effect. Give TensorMeta a layout, read it from the annotation on both binding paths (a torch tensor carries no PyPTO layout, so the annotation is the only source even there), and emit the slot again. DN's rejection then falls out of the resolver rather than being re-implemented here. Layout rides along wherever the tensor's identity survives — DynDim overlay, per-rank x[r], pl.slice, subscript — and pl.reshape now declines the meta instead of claiming ND, since it re-groups the dims a layout describes. A dep's own declaration gets the same overlay: no caller argument carries it, so a dep declaring pl.NZ under a caller that declares none was compiling as ND. pl.TensorView in the slot now raises instead of being dropped. Tensor's subscript routes any non-MemRef third element into `layout`, so a view landed where TensorMeta has no field for it and its stride vanished — reachable straight from the DN rejection's own migration hint. TensorCacheInfo gains the layout because it can enter the annotation through a closure variable, leaving the source text (and source_hash) unchanged while the artifact differs. The specializer source map now covers the signature line, so an annotation-level error reports against the user's file instead of the generated <jit:...> source.
- Include dep-declared layouts in the JIT cache key. _overlay_dep_declared_ layouts reads a dep's own annotation, but the key only serialized the entry function's tensor_meta. A postponed annotation (pl.Tensor[..., L] with a module-level L) keeps the source text — and so source_hash — identical when L is rebound, so the second compile() hit L1 and returned the artifact built with the old dep layout. - Stop implying pl.NZ works on a tensor annotation. NZ in that slot parses but cannot compile: ptoas fails it with "layout mismatch: user-specified layout=nz but inferred=nd". That is true of @pl.function too, so carrying the layout through is parity with the non-JIT path, not a claim that NZ tensors lower. The positive tests move to MX_A_ZZ (a layout the pipeline genuinely accepts, with target_memory=Mat and a UINT8/FP8E8M0 dtype), and TestNzOnTensorIsNotJitSpecific pins the shared behaviour on both paths. - Correct the language guide, which listed pl.Tensor[[64, 128], pl.FP16, pl.NZ] as a valid annotation four lines above stating NZ is never a TensorType annotation. The example cannot compile; the prohibition is right. Note the ptoas failure explicitly so the limitation is discoverable. test_basic_key_structure asserts compile_opts exactly, so it gains the new dep_layouts entry; test_different_tensor_layouts_cause_miss and test_different_dep_layouts_cause_miss pin that both components split the key.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/en/user/01-language_guide.md`:
- Line 34: Reconcile the Tensor layout guidance by explicitly identifying
MX_A_ZZ as the supported exception or revising the absolute no-layout
instruction. Apply the equivalent clarification at
docs/en/user/01-language_guide.md lines 34-34 and
docs/zh/user/01-language_guide.md lines 34-34, preserving MX_A_ZZ as the valid
positive layout example.
In `@python/pypto/jit/decorator.py`:
- Around line 264-297: Update the rejection message in _annotation_layout to use
pl.MX_A_ZZ instead of pl.NZ as the plain-layout example, matching the documented
valid layout while leaving the surrounding error guidance unchanged.
- Around line 1579-1602: Cache the computed result of
JITFunction._dep_declared_layouts on the instance so _param_layouts and
dependency traversal run only once. Reuse the cached tuple on subsequent calls,
while preserving the existing sorted (dependency name, parameter, layout)
contents and cache-key behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 3cbfdad7-d296-4ed7-99a7-a7e782aa1c85
📒 Files selected for processing (10)
docs/en/dev/language/00-python_syntax.mddocs/en/user/01-language_guide.mddocs/zh/dev/language/00-python_syntax.mddocs/zh/user/01-language_guide.mdpython/pypto/jit/cache.pypython/pypto/jit/decorator.pypython/pypto/jit/specializer.pytests/ut/jit/test_cache.pytests/ut/jit/test_jit_compile_extraction.pytests/ut/jit/test_specializer.py
| x: pl.Tensor[[64, 128], pl.FP32] # 2D, 64×128, float32 | ||
| y: pl.Tensor[[256], pl.FP16] # 1D, 256 elements, float16 | ||
| z: pl.Tensor[[64, 128], pl.FP16, pl.NZ] # With NZ layout | ||
| s: pl.Tensor[[64, 128], pl.UINT8, pl.MX_A_ZZ] # With a layout (see below) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Reconcile the layout annotation guidance.
Line 34 shows pl.MX_A_ZZ as a valid Tensor annotation. Nearby text instructs users to write Tensor annotations without layout markers. State that this example is a supported exception, or revise the absolute guidance.
docs/en/user/01-language_guide.md#L34-L34: Clarify when a Tensor layout annotation is valid.docs/zh/user/01-language_guide.md#L34-L34: Apply the equivalent clarification.
Based on PR objectives, MX_A_ZZ is the valid positive layout example.
📍 Affects 2 files
docs/en/user/01-language_guide.md#L34-L34(this comment)docs/zh/user/01-language_guide.md#L34-L34
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/en/user/01-language_guide.md` at line 34, Reconcile the Tensor layout
guidance by explicitly identifying MX_A_ZZ as the supported exception or
revising the absolute no-layout instruction. Apply the equivalent clarification
at docs/en/user/01-language_guide.md lines 34-34 and
docs/zh/user/01-language_guide.md lines 34-34, preserving MX_A_ZZ as the valid
positive layout example.
Comment on lines
+264
to
+297
| def _annotation_layout(annotation: Any, param_name: str, func_name: str) -> _ir.TensorLayout | None: | ||
| """Read the layout slot off an already-resolved tensor annotation. | ||
|
|
||
| ``pl.Tensor[[...], dtype, pl.NZ]`` evaluates to a ``Tensor`` instance whose | ||
| ``layout`` holds the third slot; the two-slot form leaves it None. | ||
|
|
||
| Args: | ||
| annotation: Resolved parameter annotation (any object — non-tensor | ||
| annotations simply carry no layout) | ||
| param_name: Parameter the annotation belongs to, for diagnostics | ||
| func_name: Enclosing kernel name, for diagnostics | ||
|
|
||
| Returns: | ||
| The annotated layout, or None when the annotation declares none | ||
|
|
||
| Raises: | ||
| TypeError: If the slot holds a ``pl.TensorView`` — specialization has | ||
| nowhere to carry it, and silently dropping it would mis-declare the | ||
| parameter | ||
| """ | ||
| layout = getattr(annotation, "layout", None) | ||
| if layout is None or isinstance(layout, _ir.TensorLayout): | ||
| return layout | ||
| # ``Tensor.__getitem__`` routes any non-MemRef third element into ``layout``, | ||
| # so a pl.TensorView(...) lands here. TensorMeta has no field for it, and a | ||
| # dropped stride is silently wrong code — refuse instead. This is reachable | ||
| # from the DN rejection's own migration hint, so the message has to be plain. | ||
| raise TypeError( | ||
| f"@pl.jit function {func_name!r}: parameter {param_name!r} annotates a " | ||
| f"{type(layout).__name__} in its layout slot, which @pl.jit does not yet " | ||
| f"support — it would be dropped and the parameter compiled as ND. Use a " | ||
| f"plain layout (e.g. pl.NZ), or declare the kernel with @pl.function, " | ||
| f"which resolves the annotation directly." | ||
| ) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Fix the misleading pl.NZ example in the TensorView rejection message.
The message at Line 295 suggests a TensorView in its layout slot, which @pl.jit does not yet support — it would be dropped and the parameter compiled as ND. Use a plain layout (e.g. pl.NZ), or declare the kernel with @pl.function, which resolves the annotation directly. The PR's own documentation states pl.NZ is rejected by the compilation pipeline on both JIT and non-JIT paths. Recommending pl.NZ here trades one error for another, unresolved failure.
Use pl.MX_A_ZZ in the example instead, since the PR documents it as the valid positive-coverage layout.
📝 Proposed fix
raise TypeError(
f"`@pl.jit` function {func_name!r}: parameter {param_name!r} annotates a "
f"{type(layout).__name__} in its layout slot, which `@pl.jit` does not yet "
f"support — it would be dropped and the parameter compiled as ND. Use a "
- f"plain layout (e.g. pl.NZ), or declare the kernel with `@pl.function`, "
+ f"plain layout (e.g. pl.MX_A_ZZ), or declare the kernel with `@pl.function`, "
f"which resolves the annotation directly."
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| def _annotation_layout(annotation: Any, param_name: str, func_name: str) -> _ir.TensorLayout | None: | |
| """Read the layout slot off an already-resolved tensor annotation. | |
| ``pl.Tensor[[...], dtype, pl.NZ]`` evaluates to a ``Tensor`` instance whose | |
| ``layout`` holds the third slot; the two-slot form leaves it None. | |
| Args: | |
| annotation: Resolved parameter annotation (any object — non-tensor | |
| annotations simply carry no layout) | |
| param_name: Parameter the annotation belongs to, for diagnostics | |
| func_name: Enclosing kernel name, for diagnostics | |
| Returns: | |
| The annotated layout, or None when the annotation declares none | |
| Raises: | |
| TypeError: If the slot holds a ``pl.TensorView`` — specialization has | |
| nowhere to carry it, and silently dropping it would mis-declare the | |
| parameter | |
| """ | |
| layout = getattr(annotation, "layout", None) | |
| if layout is None or isinstance(layout, _ir.TensorLayout): | |
| return layout | |
| # ``Tensor.__getitem__`` routes any non-MemRef third element into ``layout``, | |
| # so a pl.TensorView(...) lands here. TensorMeta has no field for it, and a | |
| # dropped stride is silently wrong code — refuse instead. This is reachable | |
| # from the DN rejection's own migration hint, so the message has to be plain. | |
| raise TypeError( | |
| f"@pl.jit function {func_name!r}: parameter {param_name!r} annotates a " | |
| f"{type(layout).__name__} in its layout slot, which @pl.jit does not yet " | |
| f"support — it would be dropped and the parameter compiled as ND. Use a " | |
| f"plain layout (e.g. pl.NZ), or declare the kernel with @pl.function, " | |
| f"which resolves the annotation directly." | |
| ) | |
| def _annotation_layout(annotation: Any, param_name: str, func_name: str) -> _ir.TensorLayout | None: | |
| """Read the layout slot off an already-resolved tensor annotation. | |
| ``pl.Tensor[[...], dtype, pl.NZ]`` evaluates to a ``Tensor`` instance whose | |
| ``layout`` holds the third slot; the two-slot form leaves it None. | |
| Args: | |
| annotation: Resolved parameter annotation (any object — non-tensor | |
| annotations simply carry no layout) | |
| param_name: Parameter the annotation belongs to, for diagnostics | |
| func_name: Enclosing kernel name, for diagnostics | |
| Returns: | |
| The annotated layout, or None when the annotation declares none | |
| Raises: | |
| TypeError: If the slot holds a ``pl.TensorView`` — specialization has | |
| nowhere to carry it, and silently dropping it would mis-declare the | |
| parameter | |
| """ | |
| layout = getattr(annotation, "layout", None) | |
| if layout is None or isinstance(layout, _ir.TensorLayout): | |
| return layout | |
| # ``Tensor.__getitem__`` routes any non-MemRef third element into ``layout``, | |
| # so a pl.TensorView(...) lands here. TensorMeta has no field for it, and a | |
| # dropped stride is silently wrong code — refuse instead. This is reachable | |
| # from the DN rejection's own migration hint, so the message has to be plain. | |
| raise TypeError( | |
| f"`@pl.jit` function {func_name!r}: parameter {param_name!r} annotates a " | |
| f"{type(layout).__name__} in its layout slot, which `@pl.jit` does not yet " | |
| f"support — it would be dropped and the parameter compiled as ND. Use a " | |
| f"plain layout (e.g. pl.MX_A_ZZ), or declare the kernel with `@pl.function`, " | |
| f"which resolves the annotation directly." | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/pypto/jit/decorator.py` around lines 264 - 297, Update the rejection
message in _annotation_layout to use pl.MX_A_ZZ instead of pl.NZ as the
plain-layout example, matching the documented valid layout while leaving the
surrounding error guidance unchanged.
Comment on lines
+1579
to
+1602
| def _dep_declared_layouts(self) -> tuple[tuple[str, str, str], ...]: | ||
| """Layouts every reachable dep declares on its own parameters. | ||
|
|
||
| ``_overlay_dep_declared_layouts`` folds these into the generated dep | ||
| signatures, so they change the artifact — but they live outside the | ||
| entry's ``tensor_meta``, and a postponed annotation | ||
| (``pl.Tensor[..., L]`` with a module-level ``L``) keeps the source text, | ||
| and therefore ``source_hash``, identical when ``L`` is rebound. Without | ||
| them in the key, rebinding ``L`` would hand the second call the first | ||
| one's artifact. | ||
|
|
||
| Returns: | ||
| Sorted ``(dep name, parameter, layout)`` triples — a stable, | ||
| hashable component for the cache key | ||
| """ | ||
| deps, _, _, _ = self._get_dep_graph() | ||
| return tuple( | ||
| sorted( | ||
| (dep.__name__, param, str(layout)) | ||
| for dep in deps | ||
| for param, layout in _param_layouts(dep._func, dep.__name__).items() | ||
| ) | ||
| ) | ||
|
|
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win
Cache _dep_declared_layouts() so it doesn't recompute on every call.
_resolve_compiled() calls self._dep_declared_layouts() at Line 2016 to build the cache key, on every invocation — including cache hits. This method (Lines 1579-1602) recomputes _param_layouts(dep._func, ...) for every reachable dep every time, which runs inspect.signature() and, for string annotations, eval(), on each call.
A dep's own layout annotations never change during the JITFunction's lifetime. _get_dep_graph() and _get_source_hash() already cache their results as instance attributes for exactly this reason. _dep_declared_layouts() should do the same, so its cost is paid once instead of on every kernel invocation.
⚡ Proposed fix: cache the computed layouts
self._cache: dict[CacheKey, Any] = {} # CacheKey → CompiledProgram
self._source_hash: str | None = None
+ self._dep_layouts_cache: tuple[tuple[str, str, str], ...] | None = None def _dep_declared_layouts(self) -> tuple[tuple[str, str, str], ...]:
"""Layouts every reachable dep declares on its own parameters.
...
"""
- deps, _, _, _ = self._get_dep_graph()
- return tuple(
- sorted(
- (dep.__name__, param, str(layout))
- for dep in deps
- for param, layout in _param_layouts(dep._func, dep.__name__).items()
- )
- )
+ if self._dep_layouts_cache is None:
+ deps, _, _, _ = self._get_dep_graph()
+ self._dep_layouts_cache = tuple(
+ sorted(
+ (dep.__name__, param, str(layout))
+ for dep in deps
+ for param, layout in _param_layouts(dep._func, dep.__name__).items()
+ )
+ )
+ return self._dep_layouts_cache🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/pypto/jit/decorator.py` around lines 1579 - 1602, Cache the computed
result of JITFunction._dep_declared_layouts on the instance so _param_layouts
and dependency traversal run only once. Reuse the cached tuple on subsequent
calls, while preserving the existing sorted (dependency name, parameter, layout)
contents and cache-key behavior.
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
1 participant
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
@pl.jitdoes not resolve the user's annotations. It records shape and dtype inTensorMetaand regenerates the annotation from that record, so the third subscript slot never survived:_build_tensor_annotationemitted onlypl.Tensor[[dims], dtype], andgrep -r layout python/pypto/jit/matched nothing at all. Because the slot never reached the type resolver, thepl.DNrejection added in #2250 was unreachable from this path — writingpl.DNunder@pl.jitneither errored nor produced DN.What changed
Layout reaches the generated annotation.
TensorMetagains alayout; both binding paths read it from the annotation (a torch tensor carries no PyPTO layout, so the annotation is the only source even on the tensor-argument path);_build_tensor_annotationemits the slot again. The DN rejection then comes from the resolver rather than a second copy of the policy here.Propagation follows tensor identity. Layout rides along through the DynDim overlay, per-rank
x[r]dispatch,pl.slice, and subscript views.pl.reshapedeclines the meta instead — it re-groups the dims a layout describes, and claiming ND would be the same silent mis-declaration.A dep's own declaration is honored. No caller argument carries it, so a dep declaring
pl.Tensor[[...], pl.NZ]under a caller that declares none was compiling as ND. Same overlay shape as the existingdep_dyn_mapfill, caller wins on conflict.pl.TensorViewin the slot now raises.Tensor.__getitem__routes any non-MemRefthird element intolayout, so a view landed whereTensorMetahas no field for it and its stride vanished. This is reachable directly from the DN rejection's own hint (which recommendspl.TensorView(stride=..., layout=DN)), so a user following that advice would have walked from a clear error into silent wrong code. It now fails with a message naming the parameter and pointing at@pl.function.Cache key. A layout can enter the annotation through a closure variable (
L = pl.NZ; ... pl.Tensor[[...], pl.FP32, L]), leaving the source text — and thussource_hash— unchanged while the artifact differs, soTensorCacheInfocarries it.Diagnostics. The specializer source map now covers the signature line, so an annotation-level error reports against the user's
.pyinstead of the generated<jit:...>source.Testing
tests/ut— 8596 passed, 2 skippedpyright python/pypto/jit/— 0 errorsruff check/ruff format --check(pinned 0.14.8) — clean, no new findings vs. baseNew coverage in
tests/ut/jit/: layout rendering and the byte-identical two-slot path, layout withpl.Outand with aDynDimshape, end-to-end NZ reachingtensor_view.layout, DN rejected with its span pointing at the user's file, a dep's own declaration surviving, and thepl.TensorViewrefusal.Two existing
TestSpecializerSourceMapassertions were exact-equality over the whole map, which the new signature entry breaks. Both were narrowed to name the signature and statement entries explicitly — the multiset assertion is now stricter than before (it also pins columns, which it previously ignored).Related
Follows #2250, which added the
pl.Tensor[..., pl.DN]rejection on the@pl.functionpath.