feat(distributed): pld.alloc_window_buffer / pld.window ops; remodel WindowBuffer as a Var#1346
Merged
Merged
Conversation
…emodel WindowBuffer as a Var WindowBuffer is now a Var subclass (mirroring MemRef) typed by the new singleton WindowBufferType. The buffer's runtime-unique identifier flows through the inherited name_hint; per-rank size is stored in bytes, making it dtype-agnostic and aligning ChipBufferSpec.nbytes 1:1 with the manifest. DistributedTensorType gains an optional window_buffer back-reference so two same-shape/dtype slices of distinct allocations stay structurally distinct. Adds pld.alloc_window_buffer (pure address-space alloc returning a Ptr) and pld.window (materialise a Ptr as a DistributedTensorType view) at the C++, binding, parser, and DSL layers. The parser intercepts the alloc assignment to derive a program-global unique name from the LHS and to reject tuple unpacking / user-supplied kwargs. Updates serialization, structural_equal, structural_hash, and the python_printer to handle the new type and the optional back-reference. The comm-manifest schema drops dtype/bits-per-element and carries only name + nbytes; the runtime passes an opaque dtype placeholder to ChipBufferSpec, which does not consume that field.
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThis PR introduces first-class distributed window-buffer allocation and view-materialization operations ( ChangesDistributed Window-Buffer Support
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. 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 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request redesigns the WindowBuffer and DistributedTensorType infrastructure to mirror the MemRef pattern, making WindowBuffer a Var subclass that is dtype-agnostic and carries size in bytes. Key changes include the introduction of pld.alloc_window_buffer and pld.window ops, updated Python AST parsing to enforce global name uniqueness, and the addition of a window_buffer back-reference in DistributedTensorType to maintain structural identity. Feedback identifies a missing field in the DistributedTensorType deserialization logic, corrects an incorrect hint in a parser error message, and recommends using INTERNAL_CHECK for defensive null checks in C++ type deduction logic.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
python/bindings/modules/ir.cpp (1)
1518-1523: 💤 Low valueConsider clarifying the base parameter requirement in the docstring.
The constructor accepts a
VarPtr basebut doesn't validate that it's actually typed asPtrType. While this mirrors theMemRefpattern (which similarly doesn't validate the base), the docstring could explicitly state the expected type constraint to guide users.📝 Suggested docstring clarification
nb::arg("span") = Span::unknown(), - "Create a WindowBuffer wrapping the given Ptr Var. The buffer's " + "Create a WindowBuffer wrapping the given Var (which must be of type " + "PtrType, produced by pld.alloc_window_buffer). The buffer's " "runtime-unique identifier flows through the inherited " "Var.name_hint (taken from base.name_hint).");🤖 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/bindings/modules/ir.cpp` around lines 1518 - 1523, The docstring for the WindowBuffer constructor (the nb::init<VarPtr, ExprPtr, bool, bool, Span>() bound on window_buffer_class) should state that the base parameter must be a pointer-typed Var (PtrType) even though no runtime check is performed; update the string literal passed to window_buffer_class.def to explicitly mention that base is expected to be a Var with PtrType (or equivalent pointer semantics) and any consequences if a non-PtrType is supplied so users know the requirement when calling the constructor.include/pypto/ir/program.h (1)
55-61: ⚡ Quick winConsider adding a null-check for the base parameter.
The constructor dereferences
base->name_hint_without validating thatbaseis non-null. While call sites likely ensure a valid base, adding a defensive check or an assertion would prevent undefined behavior if misused.🛡️ Suggested defensive check
WindowBuffer(VarPtr base, ExprPtr size, bool load_from_host = false, bool store_to_host = false, Span span = Span::unknown()) - : Var(base->name_hint_, GetWindowBufferType(), std::move(span)), + : Var(base ? base->name_hint_ : "", GetWindowBufferType(), std::move(span)), base_(std::move(base)), size_(std::move(size)), load_from_host_(load_from_host), - store_to_host_(store_to_host) {} + store_to_host_(store_to_host) { + if (!base_) { + throw pypto::ValueError("WindowBuffer requires a non-null base Ptr Var"); + } + }Alternatively, if null base is genuinely impossible (guaranteed by parser), a debug assertion (
PYPTO_CHECK(base != nullptr)) would document the precondition without runtime cost in release builds.🤖 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 `@include/pypto/ir/program.h` around lines 55 - 61, The WindowBuffer constructor currently dereferences base->name_hint_ without checking base; add a defensive check or debug assertion at the start of the WindowBuffer(VarPtr base, ExprPtr size, ...) constructor to ensure base is non-null (e.g., PYPTO_CHECK(base != nullptr) or an explicit if-check throwing/logging), then proceed to use base->name_hint_ when constructing the Var; this documents the precondition and prevents undefined behavior in WindowBuffer and related uses.
🤖 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 `@python/pypto/language/parser/ast_parser.py`:
- Around line 4451-4455: The diagnostic message for pld.alloc_window_buffer is
giving a hint that includes a dtype keyword, but the parser path rejects kwargs
for pld.alloc_window_buffer; update the hint string in the error construction
(the block that emits "pld.alloc_window_buffer must appear as the RHS..." near
the pld.alloc_window_buffer handling in ast_parser.py) to show the correct
positional-only call signature (e.g., "Write 'buf = pld.alloc_window_buffer(N,
pl.FP32)'" or similar) so the hint matches the enforced signature and does not
suggest kwargs.
---
Nitpick comments:
In `@include/pypto/ir/program.h`:
- Around line 55-61: The WindowBuffer constructor currently dereferences
base->name_hint_ without checking base; add a defensive check or debug assertion
at the start of the WindowBuffer(VarPtr base, ExprPtr size, ...) constructor to
ensure base is non-null (e.g., PYPTO_CHECK(base != nullptr) or an explicit
if-check throwing/logging), then proceed to use base->name_hint_ when
constructing the Var; this documents the precondition and prevents undefined
behavior in WindowBuffer and related uses.
In `@python/bindings/modules/ir.cpp`:
- Around line 1518-1523: The docstring for the WindowBuffer constructor (the
nb::init<VarPtr, ExprPtr, bool, bool, Span>() bound on window_buffer_class)
should state that the base parameter must be a pointer-typed Var (PtrType) even
though no runtime check is performed; update the string literal passed to
window_buffer_class.def to explicitly mention that base is expected to be a Var
with PtrType (or equivalent pointer semantics) and any consequences if a
non-PtrType is supplied so users know the requirement when calling the
constructor.
🪄 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: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: f5ac3c1b-5622-401f-8be0-21d33a30294a
📒 Files selected for processing (27)
CMakeLists.txtdocs/en/dev/ir/02-types.mddocs/zh-cn/dev/ir/02-types.mdinclude/pypto/ir/core.hinclude/pypto/ir/kind_traits.hinclude/pypto/ir/program.hinclude/pypto/ir/type.hpython/bindings/modules/ir.cpppython/pypto/ir/comm_manifest.pypython/pypto/language/distributed/__init__.pypython/pypto/language/distributed/alloc.pypython/pypto/language/parser/ast_parser.pypython/pypto/language/parser/decorator.pypython/pypto/pypto_core/ir.pyipython/pypto/runtime/distributed_runner.pysrc/ir/op/distributed/memory.cppsrc/ir/serialization/deserializer.cppsrc/ir/serialization/serializer.cppsrc/ir/serialization/type_deserializers.cppsrc/ir/transforms/python_printer.cppsrc/ir/transforms/structural_equal.cppsrc/ir/transforms/structural_hash.cpptests/ut/ir/core/test_comm_group_schema.pytests/ut/ir/parser/test_alloc_window_buffer.pytests/ut/ir/parser/test_window_op.pytests/ut/ir/test_distributed_ops.pytests/ut/runtime/test_chip_bootstrap_configs.py
YunjiQin
added a commit
to YunjiQin/pypto
that referenced
this pull request
May 12, 2026
- Round-trip the DistributedTensorType.window_buffer back-reference through serialization. The serializer now writes window_buffer as a shared node reference (WindowBuffer is an IRNode, so the ref-table handles identity); the deserializer reads it and uses a new full- fields DistributedTensorType constructor to populate every optional field (memref, tensor_view, window_buffer) in one shot. - Correct the pld.alloc_window_buffer error hint: the op rejects all user kwargs, so the previous "dtype=pl.FP32" suggestion was contradictory.
- Round-trip the DistributedTensorType.window_buffer back-reference through serialization. The serializer now writes window_buffer as a shared node reference (WindowBuffer is an IRNode, so the ref-table handles identity); the deserializer reads it and uses a new full- fields DistributedTensorType constructor to populate every optional field (memref, tensor_view, window_buffer) in one shot. - Correct the pld.alloc_window_buffer error hint: the op rejects all user kwargs, so the previous "dtype=pl.FP32" suggestion was contradictory.
YunjiQin
force-pushed
the
feat/distributed-op
branch
from
24a730d to
3b42458
Compare
lyfne123
approved these changes
May 13, 2026
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
WindowBufferas aVarsubclass (mirroringMemRef) typed by a new singletonWindowBufferType. The buffer's runtime-unique identifier flows through the inheritedname_hint; per-rank size is stored in bytes, making the buffer dtype-agnostic and aligningChipBufferSpec.nbytes1:1 with the manifest.DistributedTensorType.window_buffer— an optional back-reference to the sourceWindowBuffer— so two same-shape/dtype slices of distinct allocations stay structurally distinct.Nonefor user-declared parameter annotations.pld.alloc_window_buffer(size_bytes)— pure address-space alloc returning aPtr. The parser intercepts the assignment to derive a program-global unique name from the LHS and rejects tuple unpacking / user-supplied kwargs.pld.window(buf, [shape], dtype=...)— materialise aPtrhandle as aDistributedTensorTypeview.structural_equal,structural_hash, and the python printer for the new type and optional back-reference.name+nbytes. The runtime passes an opaque dtype placeholder toChipBufferSpec, whosedtypefield is not consumed.docs/en/dev/ir/02-types.mdand the zh-cn mirror for the newwindow_bufferback-reference and the redesignedWindowBufferschema.Relative Issue
#1127 Communication Interface impl N2
Test plan
test_split_vector_kernel.pyregex, untouched by this PR)pld.alloc_window_bufferparser behaviour (tests/ut/ir/parser/test_alloc_window_buffer.py),pld.windowparser behaviour (tests/ut/ir/parser/test_window_op.py), distributed-op type round-trip (tests/ut/ir/test_distributed_ops.py), structural equality with shared/independent slot Vars (tests/ut/ir/core/test_comm_group_schema.py), and the manifest pipeline end-to-end (tests/ut/runtime/test_chip_bootstrap_configs.py)