feat(codegen): support single-row GEMV on A2/A3 - #2185
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:
📝 WalkthroughWalkthroughGEMV tile inference now models padded physical shapes separately from logical valid regions. Mat loads select the required single-row layout, runtime and codegen tests cover the GEMV family, and English/Chinese documentation records the updated shape and validation contracts. ChangesGEMV shape and execution flow
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant TestGemv
participant tile.gemv
participant PTOAS
participant Output
TestGemv->>tile.gemv: construct GEMV tiles
tile.gemv->>PTOAS: lower to pto.tgemv
PTOAS->>Output: write padded physical result
Output-->>TestGemv: validate logical valid region
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 37b5704323
ℹ️ 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".
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
tests/st/runtime/ops/test_gemv.py (1)
457-468: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueNit: hoist the loop-invariant
valid_k.
valid_kdoes not depend onchunk; compute it once above the loop (asget_programalready does).🤖 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_gemv.py` around lines 457 - 468, In compute_expected, move the loop-invariant valid_k assignment before the for chunk loop, matching the existing calculation in get_program, and keep the loop’s matrix multiplication behavior unchanged.src/ir/op/tile_ops/matmul.cpp (1)
299-316: 🗄️ Data Integrity & Integration | 🔵 Trivial | 💤 Low value
gemv_accresult drops the accumulator's own valid extent.The returned type takes its
valid_shapepurely fromlhs/rhs. If a later chunk narrows the rhs valid N, the result claims a smaller valid region than the accumulator actually holds from earlier chunks. Consider taking the element-wise max ofacc's valid shape and the product's, or explicitly checking they agree.🤖 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/matmul.cpp` around lines 299 - 316, Update the gemv_acc result-type logic around BuildGemvResultType so the returned valid_shape preserves the accumulator’s existing valid extent instead of relying solely on the lhs/rhs product. Combine acc_type’s valid shape with the product’s using the intended element-wise maximum, or reject mismatches explicitly if equal extents are required, while preserving the existing physical-shape and dtype validation.
🤖 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/ir/op/tile_ops/matmul.cpp`:
- Around line 263-279: Validate in BuildGemvResultType that lhs_valid[0]
represents exactly one logical row before constructing the result, rejecting
non-single-row lhs inputs. Preserve the physical output shape [16, N] and
existing rhs validity handling, and do not propagate a larger lhs row extent
into tile_view.valid_shape.
---
Nitpick comments:
In `@src/ir/op/tile_ops/matmul.cpp`:
- Around line 299-316: Update the gemv_acc result-type logic around
BuildGemvResultType so the returned valid_shape preserves the accumulator’s
existing valid extent instead of relying solely on the lhs/rhs product. Combine
acc_type’s valid shape with the product’s using the intended element-wise
maximum, or reject mismatches explicitly if equal extents are required, while
preserving the existing physical-shape and dtype validation.
In `@tests/st/runtime/ops/test_gemv.py`:
- Around line 457-468: In compute_expected, move the loop-invariant valid_k
assignment before the for chunk loop, matching the existing calculation in
get_program, and keep the loop’s matrix multiplication behavior unchanged.
🪄 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: f660f89e-3544-4493-8a05-a8a3b1d8486c
📒 Files selected for processing (9)
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.mdsrc/ir/op/tile_ops/matmul.cppsrc/ir/op/tile_ops/memory.cpptests/st/runtime/ops/test_gemv.pytests/ut/codegen/test_pto_codegen.pytests/ut/ir/operators/test_tile_ops.py
cd5c9f5 to
eb3558f
Compare
|
Review follow-up: 97230be additionally requires gemv_acc accumulator valid_shape to match the computed GEMV result, preventing accumulation into a physically compatible but logically narrower tile. The runtime loop-invariant suggestion is test-only cleanup and does not affect the operator contract, so it was left unchanged. |
Summary
[1, N]validitypto.tgemv,pto.tgemv.acc, andpto.tgemv.bias, including FP16/BF16, narrow-K, bias, and accumulation casesValidation
ruff checkandruff format --checkon touched Python testsclang-format --dry-run --Werroron touched C++ filesclang-tidy --jobs=1passed (remote host LLVM 18.1.8; project pin warning expects 21.1.0).ptofiles with 16-row Acc tiles and exact base/acc/bias ops