feat(ir): add MX LeftScale/RightScale memspaces and scale load/move - #2195
feat(ir): add MX LeftScale/RightScale memspaces and scale load/move#2195yanghaoran29 wants to merge 1 commit into
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 Ascend950 MX scale support through new LeftScale and RightScale memory spaces, MX-aware ChangesMX scale memory support
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant PythonTileOps
participant OpRegistry
participant TileTypeDeduction
participant InferTileMemorySpace
participant PTOCodegen
PythonTileOps->>OpRegistry: create tile.load with mx_layout
OpRegistry->>OpRegistry: default target_memory to Mat
OpRegistry->>TileTypeDeduction: validate layout, dtype, rank, and derive TileView
TileTypeDeduction-->>OpRegistry: return Mat tile with fractal=32
OpRegistry->>InferTileMemorySpace: infer memory space
InferTileMemorySpace-->>OpRegistry: preserve MX TileView
OpRegistry->>PTOCodegen: lower MX tile.load
PTOCodegen->>PTOCodegen: emit MX tensor view and pto.tload
Possibly related PRs
Suggested labels: 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: 99aa17eaf9
ℹ️ 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.
Actionable comments posted: 1
🧹 Nitpick comments (1)
src/ir/op/tile_ops/memory.cpp (1)
395-465: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winMissing 2D-shape validation for tile.move into LeftScale/RightScale.
DeduceTileLoadType's MX path enforcestensor_type->shape_.size() == 2(line 162-164), butDeduceTileMoveType's LeftScale/RightScale path only validates dtype, not shape/rank — even thoughfractal=32is only meaningful for a 2D scale tile, and the project's own docs (docs/zh-cn/dev/ir/05-operators.md, PTOAS constraints table) state a shape-matching precondition for Mat→Scaletmov. A malformed-rank tile moved into LeftScale/RightScale would pass this type deduction and only fail deep in codegen or on hardware.♻️ Proposed check
DataType out_dtype = tile_type->dtype_; if (space == MemorySpace::LeftScale || space == MemorySpace::RightScale) { + CHECK(input_shape.size() == 2) + << "The operator " << op_name << " into LeftScale/RightScale requires a 2D tile, got rank " + << input_shape.size(); CHECK(out_dtype == DataType::UINT8 || out_dtype == DataType::FP8E8M0) << "The operator " << op_name << " into LeftScale/RightScale requires UINT8 or FP8E8M0 dtype, but got " << out_dtype.ToString();🤖 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 395 - 465, Update DeduceTileMoveType so the LeftScale and RightScale branches validate that the input tile shape is 2D before constructing the destination TileType. Preserve the existing UINT8/FP8E8M0 dtype validation and reject invalid ranks with an operator-context error consistent with DeduceTileLoadType.
🤖 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/ut/ir/operators/test_mx_scale_memspace.py`:
- Around line 83-96: Update test_mx_layout_load_defaults_to_mat so tile.load is
called with mx_layout="mx_a_zz" but without the explicit target_memory argument.
Keep the existing assertions, ensuring the test validates that the default
memory space is ir.MemorySpace.Mat.
---
Nitpick comments:
In `@src/ir/op/tile_ops/memory.cpp`:
- Around line 395-465: Update DeduceTileMoveType so the LeftScale and RightScale
branches validate that the input tile shape is 2D before constructing the
destination TileType. Preserve the existing UINT8/FP8E8M0 dtype validation and
reject invalid ranks with an operator-context error consistent with
DeduceTileLoadType.
🪄 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: 9bd5fe4e-15d1-4b56-96d1-4dcb8898817d
📒 Files selected for processing (31)
docs/en/dev/07-memory-map.mddocs/en/dev/ir/05-operators.mddocs/en/dev/passes/17-infer_tile_memory_space.mddocs/zh-cn/dev/07-memory-map.mddocs/zh-cn/dev/ir/05-operators.mddocs/zh-cn/dev/passes/17-infer_tile_memory_space.mdinclude/pypto/ir/memory_space.hinclude/pypto/ir/pipe.hinclude/pypto/ir/tile_view_semantics.hinclude/pypto/ir/transforms/utils/attrs.hpython/bindings/modules/ir.cpppython/pypto/ir/op/tile_ops.pypython/pypto/ir/type.pypython/pypto/language/op/tile_ops.pypython/pypto/language/parser/type_resolver.pypython/pypto/pypto_core/ir.pyipython/pypto/tools/memory_map.pypython/pypto/tools/templates/memory_map.htmlsrc/backend/common/pto_ops_elementwise.cppsrc/backend/common/pto_ops_memory.cppsrc/backend/common/soc.cppsrc/codegen/pto/pto_type_utils.cppsrc/ir/memref.cppsrc/ir/op/tile_ops/memory.cppsrc/ir/op_registry.cppsrc/ir/transforms/infer_tile_memory_space_pass.cppsrc/ir/transforms/memory_reuse_pass.cpptests/ut/backend/test_backend_950.pytests/ut/ir/operators/test_mx_ops.pytests/ut/ir/operators/test_mx_scale_memspace.pytests/ut/tools/test_memory_map.py
- Validate 2D MX load windows and scale tiles - Require Mat inputs for LeftScale and RightScale moves - Cover default Mat inference and invalid scale paths
5668158 to
b1fb419
Compare
Reject unsafe MX loads from strided or slice-derived tensors. Enforce fixed scale move layouts and reject scale tile.create targets. Add PTO codegen and orchestration provenance regressions for PR hw-native-sys#2195.
Add LeftScale and RightScale memory spaces with MX load and move lowering. Preserve hardware-fixed scale layouts, reject unsupported create and unsafe tensor views, and cover PTO codegen and orchestration slice provenance.
4232652 to
72caa6f
Compare
1 participant
Summary
LeftScale/RightScale) with matching SoC, reuse, bindings, and memory-map support.tile.load(..., mx_layout=mx_a_zz|mx_b_nn)(Mat-only, stamptarget_memory) and Mat→scaletile.movewith UINT8/FP8E8M0 checks and UINT8→FP8E8M0 promotion; preserve fractal=32 through Infer.Test plan
SKIP=pyright pre-commit run --from-ref upstream/main --to-ref HEAD(pyright skipped locally due to env install/network)tests/ut/ir/operators/test_mx_scale_memspace.py,tests/ut/ir/operators/test_mx_ops.py(mx load / dtype-space)tests/ut/backend/test_backend_950.py,tests/ut/tools/test_memory_map.pyCo-authored-by, and docs do not claim matmul/tgetlanding