feat(ir): add DataType.get_byte() to eliminate magic dtype byte-size literals#2000
feat(ir): add DataType.get_byte() to eliminate magic dtype byte-size literals#2000georgebisbas wants to merge 1 commit into
Conversation
…literals
Add GetByte() C++ method on DataType (ceil(GetBit()/8)), expose as get_byte()
in Python bindings, and migrate all ~68 alloc_window_buffer call sites in
pypto tests (27 files) and docs (4 files) to use explicit dtype byte sizes.
Canonical API for non-traced code:
alloc_window_buffer(SIZE * pl.FP32.get_byte())
Traced-function equivalent (parser-compatible bare-name ints):
from tests.st.distributed._window_byte_constants import BYTE_FP32, ...
alloc_window_buffer(SIZE * BYTE_FP32)
Sub-byte types (INT4, BOOL) correctly round up to 1 byte via (GetBit()+7)/8,
matching the existing ad-hoc pattern at pto_codegen.cpp:1319.
RFC: hw-native-sys#1968 (resolved, Option 1)
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a GetByte() method to the DataType class (with corresponding Python bindings and stubs) to retrieve the byte size of data types, replacing magic numbers in window buffer allocations. It also adds a set of byte-size constants (BYTE_FP32, BYTE_INT32, etc.) for use in traced functions. The review feedback highlights several instances in the test files (test_l3_multi_group.py, test_l3_get.py, test_l3_put.py, and test_l3_remote_store.py) where signal buffer allocations were missed and still use the magic number 4 instead of the newly introduced BYTE_INT32 constant.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| # Per-group window buffers land in distinct CommGroups because | ||
| # MaterializeCommDomainScopes clusters by dispatch device subset. | ||
| scratch_a_buf = pld.alloc_window_buffer(SIZE * 4) | ||
| scratch_a_buf = pld.alloc_window_buffer(SIZE * BYTE_FP32) |
There was a problem hiding this comment.
The signal_a_buf allocation on line 126 and signal_b_buf on line 128 still use the magic number 4 instead of the newly introduced BYTE_INT32 constant. To fully eliminate magic dtype byte-size literals as intended by this PR, please update these allocations to use BYTE_INT32.
| src_buf = pld.alloc_window_buffer(SIZE * 4) # 1xSIZE x FP32 | ||
| dst_buf = pld.alloc_window_buffer(SIZE * 4) | ||
| src_buf = pld.alloc_window_buffer(SIZE * BYTE_FP32) | ||
| dst_buf = pld.alloc_window_buffer(SIZE * BYTE_FP32) |
There was a problem hiding this comment.
The signal_buf allocation on line 113 still uses the magic number 4 instead of the BYTE_INT32 constant. Please update it to use BYTE_INT32 to maintain consistency with the PR's goal of eliminating magic dtype byte-size literals.
| src_buf = pld.alloc_window_buffer(SIZE * 4) # 1×SIZE × FP32 (4 bytes) | ||
| dst_buf = pld.alloc_window_buffer(SIZE * 4) | ||
| src_buf = pld.alloc_window_buffer(SIZE * BYTE_FP32) | ||
| dst_buf = pld.alloc_window_buffer(SIZE * BYTE_FP32) |
There was a problem hiding this comment.
The signal_buf allocation on line 118 still uses the magic number 4 instead of the BYTE_INT32 constant. Please update it to use BYTE_INT32 to fully eliminate magic dtype byte-size literals.
| outputs: pl.Out[pl.Tensor[[2, PIPE_ROWS, PIPE_COLS], pl.FP32]], | ||
| ) -> pl.Tensor[[2, PIPE_ROWS, PIPE_COLS], pl.FP32]: | ||
| nbytes = PIPE_ROWS * PIPE_COLS * 4 | ||
| nbytes = PIPE_ROWS * PIPE_COLS * BYTE_FP32 |
There was a problem hiding this comment.
The signal_buf allocation on line 198 still uses the magic number 4 instead of the BYTE_INT32 constant. Please update it to use BYTE_INT32 to fully eliminate magic dtype byte-size literals.
| src_buf = pld.alloc_window_buffer(16 * 16 * 4) # 16×16 × INT32 | ||
| acc_buf = pld.alloc_window_buffer(16 * 16 * 4) | ||
| src_buf = pld.alloc_window_buffer(16 * 16 * BYTE_INT32) | ||
| acc_buf = pld.alloc_window_buffer(16 * 16 * BYTE_INT32) |
There was a problem hiding this comment.
The signal_buf allocation on line 289 still uses the magic number 4 instead of the BYTE_INT32 constant. Please update it to use BYTE_INT32 to fully eliminate magic dtype byte-size literals.
| outputs: pl.Out[pl.Tensor[[2, 1, SIZE], pl.FP32]], | ||
| ) -> pl.Tensor[[2, 1, SIZE], pl.FP32]: | ||
| dst_buf = pld.alloc_window_buffer(SIZE * 4) | ||
| dst_buf = pld.alloc_window_buffer(SIZE * BYTE_FP32) |
There was a problem hiding this comment.
The signal_buf allocation on line 109 still uses the magic number 4 instead of the BYTE_INT32 constant. Please update it to use BYTE_INT32 to fully eliminate magic dtype byte-size literals.
![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: 9e15ff3520
ℹ️ 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".
| outputs: pl.Out[pl.Tensor[[2, 1, 256], pl.FP32]], | ||
| ): | ||
| data_buf = pld.alloc_window_buffer(256 * 4) | ||
| data_buf = pld.alloc_window_buffer(256 * pl.FP32.get_byte()) |
There was a problem hiding this comment.
This line is inside a @pl.jit.host body, so it is parsed by the DSL AST parser instead of being executed as normal Python. alloc_window_buffer arguments go through ASTParser.parse_expression, where pl.FP32.get_byte() is treated as a pl.* operation call (FP32) and rejected as an unknown op; the changed STs avoid this by using closure constants like BYTE_FP32. Users copying this example will hit a parse error, so hoist pl.FP32.get_byte() to a module-level int or document the closure-constant pattern for traced bodies.
Useful? React with 👍 / 👎.
|
Caution Review failedAn error occurred during the review process. Please try again later. 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 |
📝 WalkthroughWalkthroughWindow allocation sizing now derives byte widths from dtypes. A ChangesWindow byte sizing
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
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 |
There was a problem hiding this comment.
Pull request overview
Adds a DataType.get_byte() / DataType::GetByte() accessor to centralize dtype byte-size computation (including correct rounding for sub-byte types), and migrates window-buffer allocation examples and test call sites away from hard-coded * 4 / * 2 byte multipliers.
Changes:
- Add C++
DataType::GetByte()and expose it to Python asDataType.get_byte()(with updated type stubs). - Update docs/examples to use
pl.<DT>.get_byte()instead of magic byte literals. - Refactor distributed system tests to use shared
BYTE_*constants for parser-compatible traced code.
Reviewed changes
Copilot reviewed 35 out of 35 changed files in this pull request and generated 27 comments.
Show a summary per file
| File | Description |
|---|---|
| include/pypto/core/dtype.h | Adds DataType::GetByte() implementing (GetBit() + 7) / 8. |
| python/bindings/modules/core.cpp | Binds DataType.get_byte() to the new C++ API. |
| python/pypto/pypto_core/init.pyi | Adds DataType.get_byte() stub and docstring. |
| python/pypto/language/distributed/op/system_ops.py | Updates docstring example to use pl.INT32.get_byte(). |
| docs/en/user/01-language_guide.md | Replaces * 4 with pl.FP32.get_byte() in user guide example. |
| docs/zh-cn/user/01-language_guide.md | Mirrors the same change in the Chinese user guide. |
| docs/en/dev/passes/37-synthesize_allreduce_signals.md | Replaces * 4 with pl.INT32.get_byte() in pass documentation example. |
| docs/zh-cn/dev/passes/37-synthesize_allreduce_signals.md | Mirrors the same change in the Chinese pass documentation. |
| tests/st/distributed/_window_byte_constants.py | Introduces shared BYTE_* constants for traced-function-compatible allocation sizing. |
| tests/st/distributed/test_l3_window_slice_incore.py | Replaces magic byte multipliers with BYTE_* constants. |
| tests/st/distributed/test_l3_remote_store.py | Replaces magic byte multipliers with BYTE_FP32. |
| tests/st/distributed/test_l3_remote_store_window_raw.py | Replaces magic byte multipliers with BYTE_BF16 / BYTE_INT32. |
| tests/st/distributed/test_l3_put.py | Replaces magic byte multipliers with BYTE_FP32 / BYTE_INT32. |
| tests/st/distributed/test_l3_multi_group.py | Replaces magic byte multipliers with BYTE_FP32. |
| tests/st/distributed/test_l3_host_tensor_reduce_scatter.py | Replaces magic byte multipliers with BYTE_FP32 / BYTE_INT32. |
| tests/st/distributed/test_l3_host_tensor_broadcast.py | Replaces magic byte multipliers with BYTE_FP32 / BYTE_INT32. |
| tests/st/distributed/test_l3_host_tensor_barrier.py | Replaces magic byte multipliers with BYTE_FP32 / BYTE_INT32. |
| tests/st/distributed/test_l3_host_tensor_allreduce.py | Replaces magic byte multipliers with BYTE_FP32 / BYTE_INT32. |
| tests/st/distributed/test_l3_host_tensor_allgather.py | Replaces magic byte multipliers with BYTE_FP32 / BYTE_INT32. |
| tests/st/distributed/test_l3_get.py | Replaces magic byte multipliers with BYTE_FP32. |
| tests/st/distributed/test_l3_gemm.py | Replaces magic byte multipliers with BYTE_INT32. |
| tests/st/distributed/test_l3_ep_dispatch_combine.py | Replaces magic byte multipliers with BYTE_FP32 / BYTE_BF16 / BYTE_INT32. |
| tests/st/distributed/collectives/test_l3_tensor_reduce_scatter_intrinsic.py | Replaces magic byte multipliers with BYTE_FP32 / BYTE_INT32. |
| tests/st/distributed/collectives/test_l3_tensor_broadcast_intrinsic.py | Replaces magic byte multipliers with BYTE_FP32 / BYTE_INT32. |
| tests/st/distributed/collectives/test_l3_tensor_barrier_intrinsic.py | Replaces magic byte multipliers with BYTE_FP32 / BYTE_INT32. |
| tests/st/distributed/collectives/test_l3_tensor_allreduce_ring_intrinsic.py | Replaces magic byte multipliers with BYTE_FP32 / BYTE_INT32. |
| tests/st/distributed/collectives/test_l3_tensor_allreduce_intrinsic.py | Replaces magic byte multipliers with BYTE_FP32 / BYTE_INT32. |
| tests/st/distributed/collectives/test_l3_tensor_allgather_intrinsic.py | Replaces magic byte multipliers with BYTE_FP32 / BYTE_INT32. |
| tests/st/distributed/collectives/test_l3_tensor_all_to_all_intrinsic.py | Replaces magic byte multipliers with BYTE_FP32 / BYTE_INT32. |
| tests/st/distributed/collectives/test_l3_reduce_scatter.py | Replaces magic byte multipliers with BYTE_FP32 / BYTE_INT32. |
| tests/st/distributed/collectives/test_l3_broadcast.py | Replaces magic byte multipliers with BYTE_FP32 / BYTE_INT32. |
| tests/st/distributed/collectives/test_l3_allreduce.py | Replaces magic byte multipliers with BYTE_FP32 / BYTE_INT32. |
| tests/st/distributed/collectives/test_l3_allreduce_ring.py | Replaces magic byte multipliers with BYTE_FP32 / BYTE_INT32. |
| tests/st/distributed/collectives/test_l3_allgather.py | Replaces magic byte multipliers with BYTE_FP32 / BYTE_INT32. |
| tests/st/distributed/collectives/test_l3_all_to_all.py | Replaces magic byte multipliers with BYTE_FP32 / BYTE_INT32. |
|
|
||
| import pypto.language as pl | ||
| import pypto.language.distributed as pld | ||
| from tests.st.distributed._window_byte_constants import BYTE_FP32, BYTE_INT32, BYTE_BF16, BYTE_FP16 |
|
|
||
| import pypto.language as pl | ||
| import pypto.language.distributed as pld | ||
| from tests.st.distributed._window_byte_constants import BYTE_FP32, BYTE_INT32, BYTE_BF16, BYTE_FP16 |
|
|
||
| import pypto.language as pl | ||
| import pypto.language.distributed as pld | ||
| from tests.st.distributed._window_byte_constants import BYTE_FP32, BYTE_INT32, BYTE_BF16, BYTE_FP16 |
|
|
||
| import pypto.language as pl | ||
| import pypto.language.distributed as pld | ||
| from tests.st.distributed._window_byte_constants import BYTE_FP32, BYTE_INT32, BYTE_BF16, BYTE_FP16 |
|
|
||
| import pypto.language as pl | ||
| import pypto.language.distributed as pld | ||
| from tests.st.distributed._window_byte_constants import BYTE_FP32, BYTE_INT32, BYTE_BF16, BYTE_FP16 |
|
|
||
| import pypto.language as pl | ||
| import pypto.language.distributed as pld | ||
| from tests.st.distributed._window_byte_constants import BYTE_FP32, BYTE_INT32, BYTE_BF16, BYTE_FP16 |
|
|
||
| import pypto.language as pl | ||
| import pypto.language.distributed as pld | ||
| from tests.st.distributed._window_byte_constants import BYTE_FP32, BYTE_INT32, BYTE_BF16, BYTE_FP16 |
|
|
||
| import pypto.language as pl | ||
| import pypto.language.distributed as pld | ||
| from tests.st.distributed._window_byte_constants import BYTE_FP32, BYTE_INT32, BYTE_BF16, BYTE_FP16 |
|
|
||
| import pypto.language as pl | ||
| import pypto.language.distributed as pld | ||
| from tests.st.distributed._window_byte_constants import BYTE_FP32, BYTE_INT32, BYTE_BF16, BYTE_FP16 |
| BYTE_FP32 = 4 | ||
| BYTE_FP16 = 2 | ||
| BYTE_BF16 = 2 | ||
| BYTE_INT32 = 4 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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 `@tests/st/distributed/test_l3_multi_group.py`:
- Around line 125-128: Replace the hardcoded size literal 4 in the signal_a_buf
and signal_b_buf allocations with the imported BYTE_INT32 constant, leaving the
scratch buffer allocations unchanged.
In `@tests/st/distributed/test_l3_remote_store.py`:
- Around line 108-109: Replace the literal allocation size in both `host_orch`
definitions with `BYTE_INT32`, updating each `signal_buf =
pld.alloc_window_buffer(4)` to use the imported constant consistent with the
`pl.INT32` signal tensor.
🪄 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
Run ID: 3e74ffee-c078-41cd-9884-73fbe5bd2d71
📒 Files selected for processing (35)
docs/en/dev/passes/37-synthesize_allreduce_signals.mddocs/en/user/01-language_guide.mddocs/zh-cn/dev/passes/37-synthesize_allreduce_signals.mddocs/zh-cn/user/01-language_guide.mdinclude/pypto/core/dtype.hpython/bindings/modules/core.cpppython/pypto/language/distributed/op/system_ops.pypython/pypto/pypto_core/__init__.pyitests/st/distributed/_window_byte_constants.pytests/st/distributed/collectives/test_l3_all_to_all.pytests/st/distributed/collectives/test_l3_allgather.pytests/st/distributed/collectives/test_l3_allreduce.pytests/st/distributed/collectives/test_l3_allreduce_ring.pytests/st/distributed/collectives/test_l3_broadcast.pytests/st/distributed/collectives/test_l3_reduce_scatter.pytests/st/distributed/collectives/test_l3_tensor_all_to_all_intrinsic.pytests/st/distributed/collectives/test_l3_tensor_allgather_intrinsic.pytests/st/distributed/collectives/test_l3_tensor_allreduce_intrinsic.pytests/st/distributed/collectives/test_l3_tensor_allreduce_ring_intrinsic.pytests/st/distributed/collectives/test_l3_tensor_barrier_intrinsic.pytests/st/distributed/collectives/test_l3_tensor_broadcast_intrinsic.pytests/st/distributed/collectives/test_l3_tensor_reduce_scatter_intrinsic.pytests/st/distributed/test_l3_ep_dispatch_combine.pytests/st/distributed/test_l3_gemm.pytests/st/distributed/test_l3_get.pytests/st/distributed/test_l3_host_tensor_allgather.pytests/st/distributed/test_l3_host_tensor_allreduce.pytests/st/distributed/test_l3_host_tensor_barrier.pytests/st/distributed/test_l3_host_tensor_broadcast.pytests/st/distributed/test_l3_host_tensor_reduce_scatter.pytests/st/distributed/test_l3_multi_group.pytests/st/distributed/test_l3_put.pytests/st/distributed/test_l3_remote_store.pytests/st/distributed/test_l3_remote_store_window_raw.pytests/st/distributed/test_l3_window_slice_incore.py
| scratch_a_buf = pld.alloc_window_buffer(SIZE * BYTE_FP32) | ||
| signal_a_buf = pld.alloc_window_buffer(4) | ||
| scratch_b_buf = pld.alloc_window_buffer(SIZE * 4) | ||
| scratch_b_buf = pld.alloc_window_buffer(SIZE * BYTE_FP32) | ||
| signal_b_buf = pld.alloc_window_buffer(4) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Replace remaining hardcoded 4 with BYTE_INT32 for signal buffers.
signal_a_buf and signal_b_buf still use the magic literal 4 while BYTE_INT32 is imported but unused. This is inconsistent with the PR's goal of eliminating magic byte-size literals. The signal windows are [1] with dtype=pl.INT32, so BYTE_INT32 is the correct constant.
♻️ Proposed fix
scratch_a_buf = pld.alloc_window_buffer(SIZE * BYTE_FP32)
- signal_a_buf = pld.alloc_window_buffer(4)
+ signal_a_buf = pld.alloc_window_buffer(BYTE_INT32)
scratch_b_buf = pld.alloc_window_buffer(SIZE * BYTE_FP32)
- signal_b_buf = pld.alloc_window_buffer(4)
+ signal_b_buf = pld.alloc_window_buffer(BYTE_INT32)📝 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.
| scratch_a_buf = pld.alloc_window_buffer(SIZE * BYTE_FP32) | |
| signal_a_buf = pld.alloc_window_buffer(4) | |
| scratch_b_buf = pld.alloc_window_buffer(SIZE * 4) | |
| scratch_b_buf = pld.alloc_window_buffer(SIZE * BYTE_FP32) | |
| signal_b_buf = pld.alloc_window_buffer(4) | |
| scratch_a_buf = pld.alloc_window_buffer(SIZE * BYTE_FP32) | |
| signal_a_buf = pld.alloc_window_buffer(BYTE_INT32) | |
| scratch_b_buf = pld.alloc_window_buffer(SIZE * BYTE_FP32) | |
| signal_b_buf = pld.alloc_window_buffer(BYTE_INT32) |
🤖 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 `@tests/st/distributed/test_l3_multi_group.py` around lines 125 - 128, Replace
the hardcoded size literal 4 in the signal_a_buf and signal_b_buf allocations
with the imported BYTE_INT32 constant, leaving the scratch buffer allocations
unchanged.
| dst_buf = pld.alloc_window_buffer(SIZE * BYTE_FP32) | ||
| signal_buf = pld.alloc_window_buffer(4) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Missed migration: signal_buf still uses literal 4.
dst_buf was migrated to SIZE * BYTE_FP32, but the adjacent signal_buf = pld.alloc_window_buffer(4) (both host_orch definitions) still uses a bare literal even though BYTE_INT32 is imported and the signal tensor is declared pl.INT32. This leaves the exact magic-constant pattern the PR intends to eliminate.
🔧 Proposed fix
- signal_buf = pld.alloc_window_buffer(4)
+ signal_buf = pld.alloc_window_buffer(BYTE_INT32)Apply to both occurrences (lines 109 and 194).
Also applies to: 193-194
🤖 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 `@tests/st/distributed/test_l3_remote_store.py` around lines 108 - 109, Replace
the literal allocation size in both `host_orch` definitions with `BYTE_INT32`,
updating each `signal_buf = pld.alloc_window_buffer(4)` to use the imported
constant consistent with the `pl.INT32` signal tensor.
Summary
Adds
DataType.get_byte()(C++GetByte()→(GetBit()+7)/8) and migrates all ~68 magic-multiplieralloc_window_buffercall sites across pypto tests (27 files) and docs (4 files) to use explicit dtype byte sizes.This implements Option 1 from the resolved RFC.
RFC: #1968 (discussed with @YunjiQin — 2026-07-09)
Changes
C++ + Bindings + Stubs (3 files)
include/pypto/core/dtype.h—GetByte()inline one-liner afterGetBit()python/bindings/modules/core.cpp—.def("get_byte", ...)bindingpython/pypto/pypto_core/__init__.pyi—get_byte()stubTest call sites (27 files)
All ~68
alloc_window_buffer(X * 4)/alloc_window_buffer(X * 2)replaced withalloc_window_buffer(X * BYTE_FP32)/alloc_window_buffer(X * BYTE_INT32)/alloc_window_buffer(X * BYTE_BF16)using parser-compatible bare-name ints from a shared_window_byte_constants.pytest utility. Inline comments removed.Documentation (4 files + 1 docstring)
docs/en/user/01-language_guide.md+zh-cnmirror — usepl.FP32.get_byte()patterndocs/en/dev/passes/37-synthesize_allreduce_signals.md+zh-cnmirror — samepython/pypto/language/distributed/op/system_ops.pydocstring — sameDesign note
Inside traced
@pl.functionbodies the AST parser cannot resolve method calls onDataTypeobjects (pl.FP32.get_byte()). The workaround: traced code uses bare-nameBYTE_FP32etc. from a shared test-only module, which the parser resolves as closure variables. Non-traced code (docs, standalone Python) usespl.FP32.get_byte()directly.Verification
FP32=4, FP16=2, INT32=4, INT4=1, BOOL=1, BF16=2, UINT32=4, INDEX=8✅pytest tests/st/distributed/ -v --platform a2a3sim)