chore(toolchain): Update PTOAS to v0.52 and adapt ptoas discovery - #2135
chore(toolchain): Update PTOAS to v0.52 and adapt ptoas discovery#2135luohuan19 wants to merge 6 commits into
Conversation
📝 WalkthroughWalkthroughPTOAS v0.51 support adds centralized executable discovery for multiple release layouts, updates backend consumers and tests, and adds two DeepSeek-V4-Flash examples to model CI coverage. ChangesPTOAS discovery and integration
Model CI coverage
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
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: 45edc35fb5
ℹ️ 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".
b66bb0b to
b01223d
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 @.github/workflows/ci.yml:
- Around line 753-759: Remove the duplicate “Run pypto-lib DeepSeek-V4-Flash
prefill_attention_csa example” step from the workflow, while preserving the
existing identical prefill_attention_csa.py step elsewhere in the CI job.
🪄 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: a48c26e3-98c1-4c4e-a2b0-c51831aa892d
📒 Files selected for processing (7)
.github/workflows/ci.ymlpython/pypto/backend/_ptoas_locate.pypython/pypto/backend/pto_backend.pypython/pypto/jit/decorator.pypython/pypto/runtime/debug/pto_rebuild.pytests/ut/backend/test_ptoas_locate.pytoolchain/versions.env
…epSeek-V4-Flash moe to CI job
PTOAS v0.51 reorganised the release tarball: `<root>/ptoas` changed from a shell launcher into a Python package *directory*, leaving `<root>/bin/ptoas` as the only executable. Every codegen call resolved `$PTOAS_ROOT/ptoas` and so failed with "does not exist or is not executable", taking down the system-tests, system-tests-direct, system-tests-a5sim, dist-system-tests and pypto-lib-model jobs. Probe `<root>/ptoas` first and fall back to `<root>/bin/ptoas`. The order matters and is not cosmetic: up to v0.50 the top-level entry is a launcher that exports LD_LIBRARY_PATH=<root>/lib before exec'ing the bare binary, which carries no RUNPATH and otherwise dies on its bundled MLIR shared objects. From v0.51 the top-level entry is a directory, so the probe falls through to the now self-sufficient bin/ptoas. Launcher-first is correct on both layouts; bin/ptoas-first silently breaks every pre-v0.51 toolchain. The same discovery logic was duplicated in three modules, two of which degrade silently rather than raising: jit's `_ptoas_available()` flips `skip_ptoas` on, and `pto_rebuild` prints a skip line and rebuilds the stale kernel cpp. Both would have gone quiet under v0.51, so the probe is shared rather than patched at the one site CI happens to exercise. Also drop a trailing-whitespace-only line in the new pypto-lib CI steps that failed the pre-commit hook.
`pypto-lib-model` already runs `models/deepseek/v4-flash/prefill_attention_csa.py`, so the step added here was a byte-identical second copy: the example compiled and ran twice and burned an extra device allocation for no added coverage. The `moe` step is genuinely new and stays.
Placed with the other decode_attention_* steps in `pypto-lib-model`. This is the step the branch was meant to add: the earlier revision added a second `prefill_attention_csa`, which the job already ran.
Sharing the ptoas probe removed the module-level `_ptoas_binary`, which `tests/ut/runtime/test_pto_rebuild.py` monkeypatches in 15 tests — they all died with AttributeError. Bind the shared helper to that name instead of importing it under a new one, so the seam the tests patch still exists.
Both sha256 values verified against the release asset digests of hw-native-sys/PTOAS v0.52 (ptoas-bin-aarch64.tar.gz / ptoas-bin-x86_64.tar.gz). The tarball layout probe added for v0.51 (python/pypto/backend/ _ptoas_locate.py) is version-agnostic — it tests each candidate for being an executable file rather than keying off the release — so no discovery change is needed for v0.52.
…-level CommRemoteOffset helper (#2168) > **Rebased onto `main` (`a55399d4`); #2135 is closed, so its 5 toolchain commits (`update PTOAS to v0.51` … `fix(runtime): Keep the _ptoas_binary seam`) now land here.** The PTOAS pin moved on to **v0.54** in the process. ## What 1. Pin PTOAS **v0.54**. 2. Reland the `InsertCommFence` pass (reverts #2138). 3. Collapse the `system.cacheinvalid` codegen down to a single path. 4. Restore the module-level `@CommRemoteOffset_<dtype>` helper that #2161 inlined. ## 1. PTOAS v0.54 `toolchain/versions.env` is the single source of truth — every workflow reads it through the `toolchain` lead job, so the bump needs no workflow change. Both digests were verified by downloading the release assets and hashing them locally: | Arch | sha256 | | ---- | ------ | | aarch64 | `011e980dbc46c796e31a1b213051d943ba8eb4c67d356ae6bda2148fe512964a` | | x86_64 | `4e3acb9623384c18fe264610525777210095f2ba24f6c52b9823bf1cb81d7a99` | The release tarball layout changed at v0.51 (`<root>/ptoas` went from an executable launcher script to a Python package directory). `python/pypto/backend/_ptoas_locate.py` probes each candidate for being an executable file rather than keying off the release version, so it is version-agnostic and needed no change for v0.52 or v0.54. ## 2. Reland InsertCommFence (reverts #2138) #2138 backed out #2076 because ptoas could not lower the publish-side region `cacheinvalid`: | ptoas | `pto.cmo.cacheinvalid <partition_tensor_view> single_cache_line` | | ----- | --------------------------------------------------------------- | | 0.50 | parsed, but **no call emitted** — the marker never reached the device | | 0.51 | emitted `PTOAS__DCCI_SINGLE_CACHE_LINE(<GlobalTensor>)`, whose body casts to `__gm__ void*` — a conversion `GlobalTensor` does not have, so every kernel carrying the op failed to compile | hw-native-sys/PTOAS#1001 fixes this with a `GlobalTensor` overload that takes the address via `tensor.data()`, shipped in v0.52 and carried by the v0.54 pin. Verified against the real 0.52 binary — the emitted C++ now carries both overloads and binds the object one: ```cpp static AICORE inline void PTOAS__DCCI_SINGLE_CACHE_LINE( pto::GlobalTensor<Element, Shape, Stride, TensorLayout> &tensor) { dcci((__gm__ void*)tensor.data(), cache_line_t::SINGLE_CACHE_LINE); } ``` This is a plain revert of #2138 except for three deliberate deviations: - `toolchain/versions.env` stays on this branch's pin (now v0.54); the revert's restore of the v0.50 pin is dropped. - The pass doc is renumbered **43 → 44**. #2141 landed `LegalizeTileCast` at slot 14 and shifted everything below it, so `classify_iter_arg_carry` now owns 43 and `InsertCommFence` — still dead last in the pipeline — takes 44. - The revert's CommRemoteOffset inlining is **not** relanded, because #2161 landed the same inlining on `main` independently. See §4 — this PR takes that emission the other way. `44-insert_comm_fence.md` is wired into the mkdocs nav, `passes/index.md` and `00-pass_manager.md` (en + zh). The docs became an MkDocs site in #2193, and `mkdocs build --strict` fails on a page absent from the nav; `docs/zh-cn/` was also renamed to `docs/zh/` there, which the rebase picked up. ## 3. Route every `cacheinvalid` region through `partition_view` The scalar-write branch of `system.cacheinvalid` codegen emitted a bare pointer operand. Measured against ptoas 0.52: | operand | result | | ------- | ------ | | `!pto.ptr`, no type annotation ← **what we emitted** | parse error: `expected ':'` | | `!pto.ptr`, with type annotation | lowering: `addptr must feed make_tensor_view, ...` | | `!pto.tensor_view<1xf32>` | compiles | | `!pto.partition_tensor_view<1x1xf32>` | compiles | So that branch has never produced working code. It is reachable from the DSL (`pl.system.cacheinvalid(t, [1, 1], off)`) and, with this reland, from `InsertCommFence` too — `MakeCacheInvalid` uses the target's full shape, so any published tensor that is itself 1x1 lands there. The special case only existed because the region path was broken on ptoas <= 0.51. Now that it lowers correctly, a 1x1 `partition_view` is right: verified on the real binary, a `[1, 1]` region at offsets `[0, 8]` over a `[16, 16]` f32 tensor emits `GlobalTensor<float, Shape<1,1,1,1,1>, Stride<16,16,16,16,1>>` at `v1 + 8`. End-to-end, pypto's own generated `.pto` for that case now compiles where it previously hit the parse error. The branch is therefore deleted — one construction, one emit site. `GetFlatOffsetSSA` and `GetTensorBasePtr` remain in use elsewhere, so nothing is orphaned. **This supersedes #2137**, which fixes the same bug by routing the scalar case through `tensor_view<1xT>` and converging the two branches on a shared emit. Both work on 0.52; this one removes the branch entirely. Closing one of the two is a call for the authors — see "Open questions". `test_cacheinvalid_scalar_write_emits_ptr` asserted the broken form (`"partition_tensor_view" not in cmo_line`) and could not survive the fix as written. It is folded into a parametrized `test_cacheinvalid_region_emits_partition_view` covering both sizes; the dynamic-offset test now asserts the partition-view operand instead of `pto.addptr`. ## 4. Restore the module-level CommRemoteOffset helper #2161 ("Complete arbitrary-length allreduce support") inlined the distributed peer-address calculation at every remote-op call site and deleted the per-dtype `@CommRemoteOffset_<dtype>` helper. Its stated reason was that ptoas' `pto-memory-consistency` pass rejected a `func.call` to a callee holding the CommContext `pto.load_scalar` reads. **ptoas no longer performs that check**, so the helper form is viable again, and this PR restores it: - Op lowering registers the dtype via `PTOCodegen::RegisterCommRemoteOffsetHelper` and emits one `func.call @CommRemoteOffset_<dtype>(ctx, peer) -> index`. - `EmitCommRemoteOffsetHelpers` flushes one `func.func private` body per registered dtype at module end; MLIR resolves the forward references whole-module. - Sharing the CommContext field reads and the byte→element division across call sites keeps the emitted kernels smaller than the inlined form. `pto.addptr` and `pto.make_tensor_view` stay at the call site, as they did both before and after #2161 — PTOAS verifies per-function that `addptr` feeds `make_tensor_view`, and a tensor view cannot cross a func boundary because its lowered memref is strided. Returning the element offset is the only shape that satisfies both constraints. **Everything else from #2161 is preserved**: ragged-tail handling, DN / column-vector stride derivation, the tightened valid-shape inference, and the `allow_physical_tail_padding` attr on `pld.tile.remote_load`. Docs (en/zh), op descriptions, DSL docstrings, binding comments and the codegen unit-test assertions move back to the helper wording in lockstep. ## 5. Document that ptoas does not gate the markers The pass doc claimed that removing the wait-side `cacheinvalid all` from the ring-allreduce `.pto` "is rejected". That was measured on 0.50 and no longer holds. Re-run on 0.52 against the same `ring_step` kernel from `tests/st/distributed/collectives/test_l3_allreduce_ring.py`: | variant | ptoas 0.52 | | ------- | ---------- | | unmodified | accepted | | wait-side `cacheinvalid all` removed | accepted | | every `cacheinvalid` **and** `system.fence` removed | accepted | All three compile with no diagnostic; the instructions are simply absent from the generated C++. The doc now qualifies the original claim with the version it was measured on and adds a section stating there is no compile-time gate — a missing marker is a data race, not a build error — with an explicit warning against reading a green test run as evidence a marker is unnecessary. ## Testing - [x] Full unit suite on the rebased tree: `pytest tests/ut/ -n auto` → **8116 passed, 2 skipped** - [x] One pre-existing failure unrelated to this branch: `test_benchmark.py::test_benchmark_l3_ignores_prepare_setup_groups` — the test double's `prepare()` predates the `persistent=` kwarg added by #2163; `bench.py` and that test file are byte-identical to `main` here - [x] `tests/ut/ir/transforms/test_insert_comm_fence.py` + `tests/ut/codegen/distributed/` → 95 passed - [x] `ruff check` / `ruff format --check` / `clang-format --dry-run --Werror` clean on the diff - [x] Docs updated en + zh; new pass doc wired into the nav and both indices - [x] The ptoas claims in §2/§3/§5 were measured against the real v0.52 binary, not inferred - [ ] **`dist-system-tests` not run** — needs a 2-device host. This is the meaningful acceptance gate: those 36 cases are what #2138 backed the pass out for, and they are also what would exercise the restored `func.call` form from §4 on device. ## Open questions 1. **Overlap with #2137.** Same bug, same file, two fixes. One should close. 2. **The pass's benefit is inferred, not observed.** Under the 0.50 pin the publish-side region `cacheinvalid` emitted no call at all (hw-native-sys/PTOAS#995), so it never reached the device — and the distributed suite was green throughout. Combined with §5 (ptoas does not check for the markers), no existing test demonstrates what this pass fixes. Confirming it needs a case built to expose the race — large transfers, multiple ranks, repeated runs. Relates to #2076, #2138, #2161, hw-native-sys/PTOAS#995, hw-native-sys/PTOAS#1001
What
Bump the pinned PTOAS toolchain to v0.52, adapt PyPTO's
ptoasdiscovery tothe tarball layout change that came with v0.51, and add two DeepSeek-V4-Flash
examples to the
pypto-lib-modelCI job.Changes
1. PTOAS v0.52 pin
toolchain/versions.env—PTOAS_VERSIONv0.48 → v0.52, with refreshedaarch64 / x86_64 SHA256s. Both hashes were checked against the release asset
digests of
hw-native-sys/PTOASv0.52 (ptoas-bin-aarch64.tar.gz/ptoas-bin-x86_64.tar.gz).2. Adapt
ptoasdiscovery to the post-v0.50 tarball layoutThe v0.51 release reorganised the tarball.
<root>/ptoaschanged from a shelllauncher into a Python package directory, leaving
<root>/bin/ptoasas theonly executable. PyPTO resolved
$PTOAS_ROOT/ptoas, so from v0.51 on itresolved a directory and every codegen call failed with:
That took down
system-tests,system-tests-direct,system-tests-a5sim,dist-system-testsandpypto-lib-model.Discovery now probes
<root>/ptoasfirst and falls back to<root>/bin/ptoas.The order is load-bearing, not cosmetic:
<root>/ptoas<root>/bin/ptoasLD_LIBRARY_PATH=<root>/lib, then execs the binarylddreports 337 missing objectslddclean, runs standaloneLauncher-first is correct on both layouts: pre-v0.51 it picks the launcher, and
from v0.51 the top-level entry is a directory so the probe falls through to
bin/ptoas. Probingbin/ptoasfirst would turn CI green while silentlybreaking every pre-v0.51 toolchain on its bundled MLIR shared objects.
The probe is version-agnostic — it tests each candidate for being an executable
file rather than keying off the release — so the v0.52 bump needs no discovery
change on top of it.
The same logic was duplicated across three modules, so it moves to a shared
pypto/backend/_ptoas_locate.py. Two of the three call sites degradesilently rather than raising, which is why patching only the site CI exercises
was not enough:
pypto/jit/decorator.py—_ptoas_available()returningFalseflipsskip_ptoason, so@pl.jit"compiles successfully" but emits only raw.pto, no kernel cpp/so, and no error.pypto/runtime/debug/pto_rebuild.py— prints a one-line skip and rebuilds thestale kernel cpp, so an edited
.ptoappears not to take effect on replay.3. DeepSeek-V4-Flash CI coverage
Two examples added to the
pypto-lib-modeljob:decode_attention_csa(single device)moe(--device-num 2 --ep 2)A duplicated
prefill_attention_csastep is dropped in the same pass.Note on
InsertCommFenceAn earlier revision of this PR also suspended the publish-side region
system.cacheinvalid, because ptoas 0.51 lowered it toPTOAS__DCCI_SINGLE_CACHE_LINE(<GlobalTensor>)— whose body callsdcci((__gm__ void*)ptr, ...), andGlobalTensorhas no conversion to__gm__ void*in any pto-isa revision — which broke 36dist-system-testsandthe
moeexample.That is no longer part of this PR: #2138 reverted the whole
InsertCommFencepass (#2076), so the offending emission no longer exists on
mainand thisbranch has been rebased on top of that revert. The underlying ptoas gap is
tracked at
hw-native-sys/PTOAS#995 and
must be resolved before
InsertCommFenceis re-landed.Testing
tests/ut/backend/test_ptoas_locate.pycovers both tarball layouts and thelauncher-first ordering