Skip to content

fix(store): propagate phase for phased accumulators - #2396

Draft
Little-oil wants to merge 6 commits into
hw-native-sys:mainfrom
Little-oil:fix/store-st-phase
Draft

fix(store): propagate phase for phased accumulators#2396
Little-oil wants to merge 6 commits into
hw-native-sys:mainfrom
Little-oil:fix/store-st-phase

Conversation

@Little-oil

@Little-oil Little-oil commented Aug 17, 2026

Copy link
Copy Markdown
Contributor

Summary

How it was found

This bug was exposed while investigating repeated system-tests-direct failures on the unrelated reciprocal-precision PR #2376. That PR did not change GEMV, tile.store, the runtime, or the failing test; it only caused the full device suite to run. The first genuine failure repeatedly occurred in TestGemvAcc::test_tile_gemv_acc_partial_final_phases as a scheduler S1:running-stalled timeout, while later failures were cascades after the shared device runner became unusable.

The device shard executes up to 64 compiled artifacts in one hot process and reuses one ChipWorker. This made the existing cross-invocation unit-flag leak observable: an individual final-phase GEMV could return the correct value and pass, yet leave state that stalled a later phased GEMV on the same worker. Repeated failures on different runners, together with clean-main/control passes and the lack of any GEMV path in the #2376 diff, isolated the issue from the reciprocal change.

Root cause

On the A2/A3 phase-aware path, a final accumulator producer performs producer-side check-and-set. PyPTO then emitted a plain pto.tstore, because pl.store had no way to forward a store phase and does not inherit the producer phase from its input. The consumer therefore did not perform the matching check-and-clear operation. When the worker was reused, the uncleared unit flag could leave a later phased GEMV waiting indefinitely.

Before this change:

acc = pl.tile.gemv(lhs0, rhs0, acc_phase="partial")
acc = pl.tile.gemv_acc(acc, lhs1, rhs1, acc_phase="final")
out = pl.store(acc, [0, 0], out)  # no matching check-and-clear

Fix

  • expose producer and consumer phases as distinct typed pl.AccPhase and pl.STPhase enums; AccPhase supports Unspecified, Partial, and Final, while STPhase intentionally supports only Unspecified and Final and rejects PTO-ISA's lower-level check-only store phase
  • add keyword-only st_phase forwarding to pl.store and the direct tile.store IR API
  • lower explicit enum phases to the exact PTO accPhase and stPhase attributes while keeping default stores attribute-free and combining stPhase cleanly with atomicType
  • add AccStorePhaseValid verification immediately after InlineFunctions: track the exact producer SSA value (including plain aliases), require exactly one matching final store in the same straight-line region, and reject missing, mismatched, duplicate, or cross-control-flow pairs before code generation
  • preserve store kwargs through fused batch-matmul flattening
  • teach the Python printer to restore typed enum spellings and document the protocol in English and Chinese
  • add IR, printer, transform, codegen, verifier, enum, and system-test coverage

Usage

The corrected producer/consumer pairing is:

acc = pl.tile.gemv(lhs0, rhs0, acc_phase=pl.AccPhase.Partial)
acc = pl.tile.gemv_acc(acc, lhs1, rhs1, acc_phase=pl.AccPhase.Final)
out = pl.store(acc, [0, 0], out, st_phase=pl.STPhase.Final)

Existing non-phased calls remain unchanged:

out = pl.store(tile, [0, 0], out)

The default pl.STPhase.Unspecified is omitted from IR and PTO output, preserving previous store codegen exactly. A final GEMV-family producer without its matching final store—or a final store without its producer—is now rejected at compile time instead of being allowed to stall the device silently.

Testing

  • cmake --build build --parallel
  • PYTHONPATH=python pytest tests/ut -n 2 — 9,724 passed, 8 skipped
  • focused enum, AccStorePhaseValid, store phase/atomic, GEMV codegen, printer, and IR tests — 27 passed
  • targeted Pyright — 0 errors, 0 warnings
  • full pre-commit suite, including docs parity/coverage, C++ style, Ruff, and full Pyright — passed

Related to #2376; #2376 exposed but did not introduce this bug.

@coderabbitai

coderabbitai Bot commented Aug 17, 2026

Copy link
Copy Markdown

Review Change Stack

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: c4dbf6ba-cbf3-4843-9b49-660e596836f4

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

  • 🔍 Trigger review

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 7b19665b-7902-4081-8a29-6c1fd502f0d2

📥 Commits

Reviewing files that changed from the base of the PR and between b10cae3 and d11b70f.

📒 Files selected for processing (13)
  • docs/en/dev/ir/05-operators.md
  • docs/zh/dev/ir/05-operators.md
  • python/pypto/ir/op/tile_ops.py
  • python/pypto/language/op/tile_ops.py
  • src/backend/common/pto_ops_memory.cpp
  • src/ir/op/tile_ops/memory.cpp
  • src/ir/transforms/flatten_tile_nd_to_2d/batch_matmul.cpp
  • tests/st/runtime/ops/test_gemv.py
  • tests/ut/codegen/test_pto_codegen.py
  • tests/ut/codegen/test_pto_codegen_ops.py
  • tests/ut/ir/operators/test_tile_ops.py
  • tests/ut/ir/printing/test_python_printer.py
  • tests/ut/ir/transforms/test_flatten_tile_nd_to_2d.py

Included review availability: Your plan includes up to 2 reviews per rolling hour; 1 remains after this review.


📝 Walkthrough

Walkthrough

tile.store now supports st_phase values "unspecified", "partial", and "final". The phase flows through Python APIs, IR validation, PTO emission, lowering, GEMV programs, tests, and documentation.

Changes

Tile Store Phase Protocol

Layer / File(s) Summary
API and IR store-phase contract
python/pypto/ir/op/tile_ops.py, python/pypto/language/op/tile_ops.py, src/ir/op/tile_ops/memory.cpp
Both Python store APIs accept st_phase. The IR operation declares and validates the supported phase values.
PTO store attribute emission
src/backend/common/pto_ops_memory.cpp
The backend emits non-default stPhase values and combines them with atomicType when required.
Lowering and GEMV store integration
src/ir/transforms/flatten_tile_nd_to_2d/batch_matmul.cpp, tests/st/runtime/ops/test_gemv.py
Batch-matmul lowering preserves store keyword arguments. GEMV programs pass final or unspecified phases according to accumulation state.
Protocol validation and documentation
tests/ut/codegen/*, tests/ut/ir/*, docs/en/dev/ir/05-operators.md, docs/zh/dev/ir/05-operators.md
Tests cover serialization, validation, printing, code generation, and lowering. Documentation describes final, partial, and unspecified store phases.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Merge Risk: ⚪ Minimal · up to d11b7

The change adds phased accumulator-store handling while preserving default store behavior, with matching producer and final-store phases and extensive validation. No actionable merge-blocking risk remains beyond normal checks and review.

Sequence Diagram(s)

sequenceDiagram
  participant LanguageTileStore
  participant TileStoreIR
  participant PTOStoreEmitter
  LanguageTileStore->>TileStoreIR: forward st_phase
  TileStoreIR->>TileStoreIR: validate st_phase
  TileStoreIR->>PTOStoreEmitter: provide validated attributes
  PTOStoreEmitter->>PTOStoreEmitter: emit stPhase and atomicType
Loading

Poem

A rabbit stores with phases bright,
Final flags clear at ending light.
Partial checks hop through the flow,
Plain stores keep their defaults below.
Tests stamp each path just right.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 64.29% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
Title check ✅ Passed The title clearly summarizes the main change: propagating store phases for phased accumulators.
Description check ✅ Passed The description directly explains the root cause, fix, affected APIs, verification steps, and compatibility behavior.

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

Youhezhen added 4 commits August 17, 2026 05:39
- replace string phase arguments with typed AccPhase/STPhase enums
- keep verifier, codegen, printer, docs, and tests synchronized
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

Status: No status

Development

Successfully merging this pull request may close these issues.

1 participant