feat: add B09 partial arg and histogram ops - #2190
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 PyPTO
participant TileIR
participant PTOCodegen
participant PTOISA
PyPTO->>TileIR: Emit tile operation
TileIR->>PTOCodegen: Pass validated operation and types
PTOCodegen->>PTOISA: Generate corresponding PTO instruction
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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (3)
tests/st/runtime/ops/test_partial_arg.py (1)
156-162: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winPartial-shape coverage only varies rows, never columns.
valid_shapeis parametrized as[(M, N), (11, N)](Line 159), sov_colsis alwaysN(full) in both cases. The column-partial masking branch incompute_expected(choose0[:, self._v_cols :] = True, Line 151) and the corresponding column-boundary behavior ofpart_argmax/part_argminare never actually exercised.Add a case that also shrinks
v_cols(e.g.(M, 11)or(11, 11)) to cover the 2D partial-validity boundary.🤖 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/runtime/ops/test_partial_arg.py` around lines 156 - 162, Expand the valid_shape parametrization in test_partial_arg to include a case with partial columns, such as (M, 11) or (11, 11). Preserve the existing aligned and row-partial cases while ensuring compute_expected and both part_argmax/part_argmin exercise the column-boundary masking path.tests/st/runtime/ops/test_histogram.py (1)
23-41: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winHistogram filtering path is never actually exercised.
idxis always constructed to exactly equal the real higher-byte values already present insrc(e.g._idx16== row index == the upper byte baked into_src16;_idx32== the fixed0x12/0x34/0x56bytes baked into_src32). As a result, every filter predicate incompute_expected(Lines 141-143, 147-154) evaluates to "all elements pass" in every parametrized case. The test only proves "matching idx includes the element," never "mismatched idx excludes it" — the core selection behavior of the newtile.histogramop is untested.Consider adding at least one case per dtype where some rows/columns have an idx value that does not match the corresponding higher byte, and asserting those elements are excluded from the cumulative counts.
Also applies to: 136-155
🤖 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/runtime/ops/test_histogram.py` around lines 23 - 41, Update the histogram test data helpers _idx16 and _idx32, or add dedicated parametrized cases, so each supported dtype includes at least one idx value differing from the corresponding higher byte in src. Ensure compute_expected’s filtering logic and the cumulative-count assertions verify mismatched elements are excluded while matching elements remain included.src/ir/op/tile_ops/elementwise.cpp (1)
569-573: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse the existing
GetKwargOrhelper instead of a manual loop forbyte.This file already uses
GetKwargOr<bool>(kwargs, "high_precision", false)(line 337) for the same "extract an attribute with a default" pattern. The newbyteextraction reimplements the linear-search loop manually instead of reusing that helper.Based on learnings, kwargs should be looked up via the standard `GetKwarg` helper pattern rather than reimplemented per call site.♻️ Proposed refactor
- int byte = 1; - for (const auto& [key, value] : kwargs) { - if (key == "byte") byte = AnyCast<int>(value, "kwarg key: byte"); - } + int byte = GetKwargOr<int>(kwargs, "byte", 1);🤖 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/elementwise.cpp` around lines 569 - 573, Replace the manual kwargs loop initializing byte with the existing GetKwargOr helper, using the current default value of 1 and preserving the AnyCast<int> conversion and subsequent byte range validation in the histogram operator.Source: Learnings
🤖 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 `@src/backend/common/pto_ops_elementwise.cpp`:
- Around line 353-412: Update MakePartArgCodegenPTO and MakeHistogramCodegenPTO
to omit operand type annotations when the corresponding GetExprTypeAnnotation
result is empty, using EmitInsOuts or the existing conditional annotation
pattern. Preserve annotations for non-empty types and keep the generated
ins/outs structure valid for all operands.
- Around line 360-365: Ensure both tuple outputs from ResolveTupleResultElements
remain materialized before MakePartArgCodegenPTO processes them, even when one
output is unused by the kernel. Preserve both TupleGetItemExpr consumer bindings
through DCE, or update the codegen path around GetCurrentResultVar and
ResolveTupleResultElements to support a single-output consumer without
triggering the two-element internal check.
In `@src/ir/op/tile_ops/elementwise.cpp`:
- Around line 453-476: Update DeduceTilePartArgType to validate the required
layouts for tile.part_argmax/tile.part_argmin operands: enforce the value layout
contract on src0 and src1 and the index layout contract on src0_idx and
src1_idx, matching the layouts used by pto.tpartargmax/pto.tpartargmin. Ensure
invalid layouts are rejected during IR type deduction before backend codegen.
- Around line 495-515: Update the result-type construction after the dominance
comparison to inherit each result tile’s layout from the dominating source
rather than always from src0. Select the value source corresponding to
result_valid (types[0] when src0_dominates, otherwise types[1]) and the matching
index source (types[2] or types[3]), then pass those sources to
InheritTileViewLayout while preserving the existing valid_shape and type
metadata.
---
Nitpick comments:
In `@src/ir/op/tile_ops/elementwise.cpp`:
- Around line 569-573: Replace the manual kwargs loop initializing byte with the
existing GetKwargOr helper, using the current default value of 1 and preserving
the AnyCast<int> conversion and subsequent byte range validation in the
histogram operator.
In `@tests/st/runtime/ops/test_histogram.py`:
- Around line 23-41: Update the histogram test data helpers _idx16 and _idx32,
or add dedicated parametrized cases, so each supported dtype includes at least
one idx value differing from the corresponding higher byte in src. Ensure
compute_expected’s filtering logic and the cumulative-count assertions verify
mismatched elements are excluded while matching elements remain included.
In `@tests/st/runtime/ops/test_partial_arg.py`:
- Around line 156-162: Expand the valid_shape parametrization in
test_partial_arg to include a case with partial columns, such as (M, 11) or (11,
11). Preserve the existing aligned and row-partial cases while ensuring
compute_expected and both part_argmax/part_argmin exercise the column-boundary
masking path.
🪄 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: d5e61347-c1d1-4df4-b961-57de2f87a2f5
📒 Files selected for processing (13)
docs/en/dev/codegen/00-pto_codegen.mddocs/en/dev/ptoas-op-status.mddocs/zh-cn/dev/codegen/00-pto_codegen.mddocs/zh-cn/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_elementwise.cppsrc/ir/op/tile_ops/elementwise.cpptests/st/runtime/ops/test_histogram.pytests/st/runtime/ops/test_partial_arg.pytests/ut/ir/operators/test_partial_arg_histogram.py
3ab0ce1 to
e44e3bd
Compare
|
Review follow-up: e44e3bd expands partial-shape runtime coverage to rows, columns, and both axes; the UINT32 histogram fixture now contains deliberately mismatching high bytes so its filter path removes data; and the byte kwarg parsing now uses GetKwargOr. |
e44e3bd to
5730b88
Compare
5730b88 to
3a6d8f6
Compare
1 participant
Summary
tile.part_argmaxandtile.part_argminacross C++ IR, Python IR/DSL, tuple codegen, tests, and docstile.histogramwith UINT16/UINT32 contracts, PTOAS lowering, runtime coverage, and docsValidation
/data/chenshenai/test2pto.tpartargmax/pto.tpartargmin--precompile-workers 1, device 1)pto.thistogramA5 simulator execution was attempted but is unavailable on this host because its
libstdc++.so.6lacksGLIBCXX_3.4.32; the status docs therefore describe compilation coverage and leave A5 execution pending.