feat(ir/codegen): add MX matmul_mx family and tget_scale_addr - #2237
feat(ir/codegen): add MX matmul_mx family and tget_scale_addr#2237yanghaoran29 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:
📝 WalkthroughWalkthroughThis PR adds MXFP8 block-scale matmul operations and scale-address binding to the IR, Python APIs, DSL exports, and PTO backend. It adds validation, accumulator and bias variants, code generation, unit tests, and English and Chinese documentation. ChangesMXFP8 matmul implementation
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related issues
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: 6ceccc0f0d
ℹ️ 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.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/backend/common/pto_ops_elementwise.cpp (1)
727-742: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winEnforce the all-or-none type-annotation invariant in
make_acc_codegen.
ins_type_partsonly pushesdst_typewhen it is non-empty, and only pushes eachoperand_typewhen it is non-empty. Operand SSA values are always emitted, so a partial type set desyncs the position of each entry in the emitted: type1, type2, ...clause from its actual operand.The same file's
MakeScatterCodegenPTOinpto_ops_datamove.cpprejects this exact situation with an explicitINTERNAL_CHECK_SPANthat requires all type annotations to be present or all absent. Apply the same rule here, so a future operand whose type does not resolve fails loudly instead of emitting a misaligned type list.🛡️ Proposed fix to enforce all-or-none type annotations
std::ostringstream acc_inst; acc_inst << pto_op << " ins(" << dst; - std::vector<std::string> ins_type_parts; - if (!dst_type.empty()) ins_type_parts.push_back(dst_type); - for (size_t i = 1; i < op->args_.size(); ++i) { - acc_inst << ", " << codegen.GetExprAsCode(op->args_[i]); - std::string operand_type = codegen.GetExprTypeAnnotation(op->args_[i]); - if (!operand_type.empty()) ins_type_parts.push_back(operand_type); - } + std::vector<std::string> operand_types = {dst_type}; + for (size_t i = 1; i < op->args_.size(); ++i) { + acc_inst << ", " << codegen.GetExprAsCode(op->args_[i]); + operand_types.push_back(codegen.GetExprTypeAnnotation(op->args_[i])); + } + bool any_type_present = std::any_of(operand_types.begin(), operand_types.end(), + [](const std::string& t) { return !t.empty(); }); + bool all_types_present = std::all_of(operand_types.begin(), operand_types.end(), + [](const std::string& t) { return !t.empty(); }); + INTERNAL_CHECK(!any_type_present || all_types_present) + << "Internal error: " << pto_op + << " operand type annotations must all be present or all absent, got a partial set"; + std::vector<std::string> ins_type_parts = all_types_present ? operand_types + : std::vector<std::string>{};🤖 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/backend/common/pto_ops_elementwise.cpp` around lines 727 - 742, Update make_acc_codegen so type annotations are emitted only when dst_type and every operand type are present; otherwise emit no type clause. Add an INTERNAL_CHECK_SPAN matching MakeScatterCodegenPTO to reject partial annotations before constructing acc_inst, while preserving the existing operand emission and annotation ordering.
🧹 Nitpick comments (1)
python/pypto/language/op/tile_ops.py (1)
1220-1245: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocstrings for the new MX ops omit Args/Returns sections.
matmul_mx,matmul_mx_acc,matmul_mx_bias, andtget_scale_addreach use a single-line docstring. The adjacent functionsmatmul,matmul_acc, andmatmul_biasdocument each parameter with "Args:" and "Returns:" sections. Add the same sections to the new functions for consistency.📝 Proposed docstring expansion (example for `matmul_mx`)
def matmul_mx(lhs: Tile, lhs_scale: Tile, rhs: Tile, rhs_scale: Tile) -> Tile: - """MX block-scale matrix multiplication.""" + """MX block-scale matrix multiplication. + + Args: + lhs: Left-hand side data tile (FP8E4M3FN) + lhs_scale: Left-hand side scale tile (FP8E8M0) + rhs: Right-hand side data tile (FP8E4M3FN) + rhs_scale: Right-hand side scale tile (FP8E8M0) + + Returns: + Tile wrapping the matmul_mx operation + """ call_expr = _ir_ops.matmul_mx(lhs.unwrap(), lhs_scale.unwrap(), rhs.unwrap(), rhs_scale.unwrap()) return Tile(expr=call_expr)🤖 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 1220 - 1245, Expand the docstrings for matmul_mx, matmul_mx_acc, matmul_mx_bias, and tget_scale_addr with Args: entries describing every parameter and a Returns: entry describing the returned Tile, matching the format and level of detail used by the adjacent matmul, matmul_acc, and matmul_bias functions.
🤖 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.
Outside diff comments:
In `@src/backend/common/pto_ops_elementwise.cpp`:
- Around line 727-742: Update make_acc_codegen so type annotations are emitted
only when dst_type and every operand type are present; otherwise emit no type
clause. Add an INTERNAL_CHECK_SPAN matching MakeScatterCodegenPTO to reject
partial annotations before constructing acc_inst, while preserving the existing
operand emission and annotation ordering.
---
Nitpick comments:
In `@python/pypto/language/op/tile_ops.py`:
- Around line 1220-1245: Expand the docstrings for matmul_mx, matmul_mx_acc,
matmul_mx_bias, and tget_scale_addr with Args: entries describing every
parameter and a Returns: entry describing the returned Tile, matching the format
and level of detail used by the adjacent matmul, matmul_acc, and matmul_bias
functions.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 1d9800fb-d8bf-418c-b153-d7f72126ecc9
📒 Files selected for processing (18)
CMakeLists.txtdocs/en/dev/ir/05-operators.mddocs/en/dev/passes/17-infer_tile_memory_space.mddocs/en/dev/ptoas-op-status.mddocs/zh/dev/ir/05-operators.mddocs/zh/dev/passes/17-infer_tile_memory_space.mddocs/zh/dev/ptoas-op-status.mdpython/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_datamove.cppsrc/backend/common/pto_ops_elementwise.cppsrc/codegen/pto/pto_codegen.cppsrc/ir/op/tile_ops/matmul_mx.cppsrc/ir/transforms/infer_tile_memory_space_pass.cpptests/ut/codegen/test_mx_ops_codegen.pytests/ut/ir/operators/test_mx_ops.py
💤 Files with no reviewable changes (1)
- src/ir/transforms/infer_tile_memory_space_pass.cpp
6ceccc0 to
8f52b07
Compare
|
Addressed AI review feedback in the amended commit:
Also fixed clang-format failures that blocked CI pre-commit. |
1ccb315 to
0cff7de
Compare
31fb37a to
1b33b35
Compare
Add tile.matmul_mx / matmul_mx_acc / matmul_mx_bias with shared MX alignment and type checks, plus compiler-only tile.tget_scale_addr and PTO lowering for the MX path. InferTileMemorySpace resolves Left/LeftScale and Right/RightScale and inserts tile.move. InsertMxScaleAddr then emits a fresh tget_scale_addr for every MX matmul consumer and rewrites scale operands to the bound SSA. It deliberately avoids cross-consumer CSE because SSA aliases, views, and bound results may share the same stateful physical scale buffer. Operands use AsVarLike so loop IterArgs are accepted. Also validate constant rhs physical/valid K and scale geometry independently of lhs, streamline the MX regression tests, and document the new pass as pass 18 in en/zh with subsequent pass documents renumbered.
1b33b35 to
a8b7913
Compare
1 participant
Summary
Testing
Dependency