feat(ops): add B03 tri and gather operators - #2184
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughAdds ChangesTile operation support
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant PythonAPI
participant TileIR
participant PTOCodegen
participant RuntimeTests
PythonAPI->>TileIR: create tri, gatherb, or mgather
TileIR->>PTOCodegen: validate and infer tile result
PTOCodegen->>RuntimeTests: emit canonical PTO operation
RuntimeTests->>RuntimeTests: execute and compare expected output
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 |
![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: d37d427536
ℹ️ 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".
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (3)
python/pypto/language/op/tile_ops.py (1)
2689-2698: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
mem: Tensortype hint is narrower than the IR contract.
DeduceTileMgatherTypeacceptsTensorTypeorDistributedTensorType(viaAsTensorTypeLike), same astile.load. Sibling opmscatterin this file already uses a generic_TensorTfor this reason. Typingmemas plainTensorhere would make static type checkers flag legitimateDistributedTensorcall sites even though the call itself works at runtime.♻️ Proposed fix
-def mgather(mem: Tensor, idx: Tile, coalesce: str | int = "row") -> Tile: +def mgather(mem: _TensorT, idx: Tile, coalesce: str | int = "row") -> Tile:🤖 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/language/op/tile_ops.py` around lines 2689 - 2698, Update the mgather function’s mem parameter annotation from Tensor to the generic _TensorT type used by mscatter, while preserving the existing unwrap and IR-call behavior so both Tensor and DistributedTensor inputs are accepted by static type checkers.python/pypto/backend/_ptoas_preprocess.py (1)
33-50: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winSkip the whole-file identifier/alias scan when no MGATHER call is present.
_restore_mgather_wrapper_operandsunconditionally buildsaliases/identifier_occurrencesby scanning every line of every PTOAS output file, even though most generated kernels don't usetile.mgatherat all. A cheap early-exit avoids this cost on the common path.⚡ Proposed fix
def _restore_mgather_wrapper_operands(content: str) -> str: """...""" + if "MGATHER" not in content: + return content lines = content.splitlines(keepends=True)🤖 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/backend/_ptoas_preprocess.py` around lines 33 - 50, Update _restore_mgather_wrapper_operands to detect the absence of an MGATHER call and return content before building aliases or identifier_occurrences. Perform this lightweight check before the whole-file line scan, while preserving the existing rewrite behavior when an MGATHER call is present.src/ir/op/tile_ops/memory.cpp (1)
1053-1062: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCoalesce-mode 0/1 mapping is duplicated across IR and backend with no shared enum. Both sites independently encode "0=row, 1=elem" as raw literals; the codebase's own
ir::AtomicType(shared betweenDeduceTileStoreTypeand its store codegen) is the established pattern for this exact situation.
src/ir/op/tile_ops/memory.cpp#L1053-L1062: movekMgatherCoalesceRow/kMgatherCoalesceElemout of the anonymous namespace into a small shared header enum (e.g.enum class CoalesceMode { kRow = 0, kElem = 1 }), mirroringAtomicType.src/backend/common/pto_ops_memory.cpp#L350-L353: replace the rawop->GetKwarg<int>("coalesce", 0)+coalesce == 1 ? "elem" : "row"ternary with the shared enum from the header above, instead of re-hardcoding the mapping.🤖 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 `@src/ir/op/tile_ops/memory.cpp` around lines 1053 - 1062, Replace the duplicated coalesce-mode literals with a shared enum, following the existing ir::AtomicType pattern. In src/ir/op/tile_ops/memory.cpp lines 1053-1062, move the row/element constants from the anonymous namespace into a shared header enum such as CoalesceMode; in src/backend/common/pto_ops_memory.cpp lines 350-353, use that enum for the default and comparison instead of raw integer mapping.
🤖 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.
Nitpick comments:
In `@python/pypto/backend/_ptoas_preprocess.py`:
- Around line 33-50: Update _restore_mgather_wrapper_operands to detect the
absence of an MGATHER call and return content before building aliases or
identifier_occurrences. Perform this lightweight check before the whole-file
line scan, while preserving the existing rewrite behavior when an MGATHER call
is present.
In `@python/pypto/language/op/tile_ops.py`:
- Around line 2689-2698: Update the mgather function’s mem parameter annotation
from Tensor to the generic _TensorT type used by mscatter, while preserving the
existing unwrap and IR-call behavior so both Tensor and DistributedTensor inputs
are accepted by static type checkers.
In `@src/ir/op/tile_ops/memory.cpp`:
- Around line 1053-1062: Replace the duplicated coalesce-mode literals with a
shared enum, following the existing ir::AtomicType pattern. In
src/ir/op/tile_ops/memory.cpp lines 1053-1062, move the row/element constants
from the anonymous namespace into a shared header enum such as CoalesceMode; in
src/backend/common/pto_ops_memory.cpp lines 350-353, use that enum for the
default and comparison instead of raw integer mapping.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 7022cf95-b628-4ead-8a32-d5087805f99e
📒 Files selected for processing (19)
docs/en/dev/ir/05-operators.mddocs/en/dev/ptoas-op-status.mddocs/zh-cn/dev/ir/05-operators.mddocs/zh-cn/dev/ptoas-op-status.mdpython/pypto/backend/_ptoas_preprocess.pypython/pypto/ir/op/tile_ops.pypython/pypto/language/__init__.pypython/pypto/language/op/__init__.pypython/pypto/language/op/tile_ops.pysrc/backend/common/pto_ops_elementwise.cppsrc/backend/common/pto_ops_memory.cppsrc/ir/op/tile_ops/gather.cppsrc/ir/op/tile_ops/memory.cpptests/st/runtime/ops/test_gatherb.pytests/st/runtime/ops/test_mgather.pytests/st/runtime/ops/test_tri.pytests/ut/codegen/test_pto_codegen.pytests/ut/codegen/test_pto_codegen_ops.pytests/ut/ir/operators/test_tile_ops.py
|
Review follow-up: 6137e4f adds the requested no-MGATHER preprocessing fast path. I did not change mgather(mem: Tensor, ...) to _TensorT: DistributedTensor is already a Tensor subclass and this API returns a Tile, so a TypeVar would not preserve any input/output type relationship. I also kept the internal integer coalesce encoding local; the public API remains the string contract, and introducing a cross-language enum/binding solely for this serialized attribute would expand this batch without changing safety. |
1 participant
Summary
pto.ttri,pto.tgatherb, and canonicalpto.mgathertgatherbwith the current 32-byte block-offset semantics and row/column valid-shape tailsmgatherrow/elem coalesce, valid shapes, signed/unsigned payload dtypes, and printed-IR integer enum round tripsMGATHERto the legacy raw-pointer ABIee7b686a,64c6b160,711e45c7,7ec0e5f2) and remove the stalepto.tmgatherspellingThis branch is based directly on current
mainand contains no B02 changes, so B02 and B03 can be reviewed and squash-merged independently.Validation
task_20260728_073610_14453233363, device 1)ttri: 10/10 pytest items, 9 unique device programstgatherb: 8/8 pytest items, 7 unique device programsmgather: 9/9 pytest items, 8 unique device programspto.ttri,pto.tgatherb, andpto.mgather; legacypto.tmgatheris absent