diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 8e1740ee76..ab9e985df0 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -743,23 +743,25 @@ jobs: - name: Run Python hardware unit tests (a5) run: | source .venv/bin/activate - DEVICE_LIST=$(python -c "p='${DEVICE_RANGE}'.split('-'); s,e=p[0],p[-1]; print(','.join(str(i) for i in range(int(s),int(e)+1)))") - task-submit --timeout 1800 --max-time 1800 --device "$DEVICE_LIST" --run "python -m pytest tests -m requires_hardware --platform a5 --device ${DEVICE_RANGE} -v" + DEVICE_NUM=$(python -c "p='${DEVICE_RANGE}'.split('-'); print(int(p[-1])-int(p[0])+1)") + PYTEST_BASETEMP="${RUNNER_TEMP}/pytest-${GITHUB_JOB}-${GITHUB_RUN_ID}-${GITHUB_RUN_ATTEMPT}" + task-submit --timeout 1800 --max-time 1800 --device auto --device-num "$DEVICE_NUM" \ + --run "python -m pytest tests -m requires_hardware --platform a5 \ + --device \$TASK_DEVICE --basetemp $PYTEST_BASETEMP -v" - name: Build and run C++ hardware unit tests (a5) run: | source .venv/bin/activate cmake -B tests/ut/cpp/build -S tests/ut/cpp -DSIMPLER_ENABLE_HARDWARE_TESTS=ON cmake --build tests/ut/cpp/build - python3 -c " - import json, os - p = os.environ['DEVICE_RANGE'].split('-'); s, e = p[0], p[-1] # tolerate a single id (no hyphen) - npus = [{'id': str(i), 'slots': 1} for i in range(int(s), int(e)+1)] - json.dump({'version': {'major': 1, 'minor': 0}, 'local': [{'npus': npus}]}, - open('tests/ut/cpp/build/resources.json', 'w')) - " - DEVICE_LIST=$(python -c "p='${DEVICE_RANGE}'.split('-'); s,e=p[0],p[-1]; print(','.join(str(i) for i in range(int(s),int(e)+1)))") - task-submit --timeout 1800 --max-time 1800 --device "$DEVICE_LIST" --run "ctest --test-dir tests/ut/cpp/build -L '^requires_hardware(_a5)?\$' --resource-spec-file $PWD/tests/ut/cpp/build/resources.json -j$(nproc) --output-on-failure" + DEVICE_NUM=$(python -c "p='${DEVICE_RANGE}'.split('-'); print(int(p[-1])-int(p[0])+1)") + RESOURCE_SPEC="$PWD/tests/ut/cpp/build/resources.json" + task-submit --timeout 1800 --max-time 1800 --device auto --device-num "$DEVICE_NUM" \ + --run "python3 -c 'import json, os; ids=os.environ[\"TASK_DEVICE\"].split(\",\"); \ + spec={\"version\":{\"major\":1,\"minor\":0},\"local\":[{\"npus\":[{\"id\":i,\"slots\":1} for i in ids]}]}; \ + json.dump(spec, open(\"$RESOURCE_SPEC\", \"w\"))' && \ + ctest --test-dir tests/ut/cpp/build -L '^requires_hardware(_a5)?\$' \ + --resource-spec-file $RESOURCE_SPEC -j$(nproc) --output-on-failure" - name: Build cann-examples/query (CANN host-API smoke) run: | @@ -813,9 +815,12 @@ jobs: - name: Run pytest scene tests (a5) run: | source .venv/bin/activate - DEVICE_LIST=$(python -c "p='${DEVICE_RANGE}'.split('-'); s,e=p[0],p[-1]; print(','.join(str(i) for i in range(int(s),int(e)+1)))") - PYTEST="python -m pytest examples tests/st --platform a5 --device ${DEVICE_RANGE} -v --require-pto-isa" - task-submit --timeout 1800 --max-time 1800 --device "$DEVICE_LIST" --run "$PYTEST --pto-session-timeout 1200" + DEVICE_NUM=$(python -c "p='${DEVICE_RANGE}'.split('-'); print(int(p[-1])-int(p[0])+1)") + PYTEST_BASETEMP="${RUNNER_TEMP}/pytest-${GITHUB_JOB}-${GITHUB_RUN_ID}-${GITHUB_RUN_ATTEMPT}" + PYTEST="python -m pytest examples tests/st --platform a5 --device \$TASK_DEVICE \ + --basetemp $PYTEST_BASETEMP -v --require-pto-isa" + task-submit --timeout 1800 --max-time 1800 --device auto --device-num "$DEVICE_NUM" \ + --run "$PYTEST --pto-session-timeout 1200" # DFX per-feature smokes — hardware mirror of the st-sim-a5 set. The # race window in dep_gen only fires on real timing, so this row is @@ -824,35 +829,43 @@ jobs: - name: dep_gen smoke run: | source .venv/bin/activate - DEVICE_LIST=$(python -c "p='${DEVICE_RANGE}'.split('-'); s,e=p[0],p[-1]; print(','.join(str(i) for i in range(int(s),int(e)+1)))") + DEVICE_NUM=$(python -c "p='${DEVICE_RANGE}'.split('-'); print(int(p[-1])-int(p[0])+1)") + PYTEST_BASETEMP="${RUNNER_TEMP}/pytest-${GITHUB_JOB}-${GITHUB_RUN_ID}-${GITHUB_RUN_ATTEMPT}" PYTEST="python -m pytest tests/st/a5/tensormap_and_ringbuffer/dfx/dep_gen/test_dep_gen.py \ - --platform a5 --device ${DEVICE_RANGE} -p no:xdist --pto-session-timeout 600 \ + --platform a5 --device \$TASK_DEVICE --basetemp $PYTEST_BASETEMP \ + -p no:xdist --pto-session-timeout 600 \ --require-pto-isa --enable-dep-gen" - task-submit --timeout 1800 --max-time 1800 --device "$DEVICE_LIST" --run "$PYTEST" + task-submit --timeout 1800 --max-time 1800 --device auto --device-num "$DEVICE_NUM" --run "$PYTEST" - name: l2_swimlane smoke run: | source .venv/bin/activate - DEVICE_LIST=$(python -c "p='${DEVICE_RANGE}'.split('-'); s,e=p[0],p[-1]; print(','.join(str(i) for i in range(int(s),int(e)+1)))") + DEVICE_NUM=$(python -c "p='${DEVICE_RANGE}'.split('-'); print(int(p[-1])-int(p[0])+1)") + PYTEST_BASETEMP="${RUNNER_TEMP}/pytest-${GITHUB_JOB}-${GITHUB_RUN_ID}-${GITHUB_RUN_ATTEMPT}" PYTEST="python -m pytest tests/st/a5/tensormap_and_ringbuffer/dfx/l2_swimlane/ \ - --platform a5 --device ${DEVICE_RANGE} -p no:xdist --pto-session-timeout 600 \ + --platform a5 --device \$TASK_DEVICE --basetemp $PYTEST_BASETEMP \ + -p no:xdist --pto-session-timeout 600 \ --require-pto-isa --enable-l2-swimlane --enable-dep-gen" - task-submit --timeout 1800 --max-time 1800 --device "$DEVICE_LIST" --run "$PYTEST" + task-submit --timeout 1800 --max-time 1800 --device auto --device-num "$DEVICE_NUM" --run "$PYTEST" - name: PMU smoke run: | source .venv/bin/activate - DEVICE_LIST=$(python -c "p='${DEVICE_RANGE}'.split('-'); s,e=p[0],p[-1]; print(','.join(str(i) for i in range(int(s),int(e)+1)))") + DEVICE_NUM=$(python -c "p='${DEVICE_RANGE}'.split('-'); print(int(p[-1])-int(p[0])+1)") + PYTEST_BASETEMP="${RUNNER_TEMP}/pytest-${GITHUB_JOB}-${GITHUB_RUN_ID}-${GITHUB_RUN_ATTEMPT}" PYTEST="python -m pytest tests/st/a5/tensormap_and_ringbuffer/dfx/pmu/test_pmu.py \ - --platform a5 --device ${DEVICE_RANGE} -p no:xdist --pto-session-timeout 600 \ + --platform a5 --device \$TASK_DEVICE --basetemp $PYTEST_BASETEMP \ + -p no:xdist --pto-session-timeout 600 \ --require-pto-isa --enable-pmu 2" - task-submit --timeout 1800 --max-time 1800 --device "$DEVICE_LIST" --run "$PYTEST" + task-submit --timeout 1800 --max-time 1800 --device auto --device-num "$DEVICE_NUM" --run "$PYTEST" - name: args_dump smoke run: | source .venv/bin/activate - DEVICE_LIST=$(python -c "p='${DEVICE_RANGE}'.split('-'); s,e=p[0],p[-1]; print(','.join(str(i) for i in range(int(s),int(e)+1)))") + DEVICE_NUM=$(python -c "p='${DEVICE_RANGE}'.split('-'); print(int(p[-1])-int(p[0])+1)") + PYTEST_BASETEMP="${RUNNER_TEMP}/pytest-${GITHUB_JOB}-${GITHUB_RUN_ID}-${GITHUB_RUN_ATTEMPT}" PYTEST="python -m pytest tests/st/a5/tensormap_and_ringbuffer/dfx/args_dump/test_args_dump.py \ - --platform a5 --device ${DEVICE_RANGE} -p no:xdist --pto-session-timeout 600 \ + --platform a5 --device \$TASK_DEVICE --basetemp $PYTEST_BASETEMP \ + -p no:xdist --pto-session-timeout 600 \ --require-pto-isa --dump-args" - task-submit --timeout 1800 --max-time 1800 --device "$DEVICE_LIST" --run "$PYTEST" + task-submit --timeout 1800 --max-time 1800 --device auto --device-num "$DEVICE_NUM" --run "$PYTEST" diff --git a/README.md b/README.md index e0f8f27959..bb45136a93 100644 --- a/README.md +++ b/README.md @@ -73,6 +73,7 @@ export ASCEND_HOME_PATH=/usr/local/Ascend/ascend-toolkit/latest | [Orchestrator](docs/orchestrator.md) | DAG submission internals: submit flow, TensorMap, Scope, Ring, task state machine | | [Scheduler](docs/scheduler.md) | DAG dispatch internals: wiring/ready/completion queues, dispatch loop | | [Worker Manager](docs/worker-manager.md) | Worker pool, WorkerThread, THREAD/PROCESS modes, fork + mailbox mechanics | +| [Streaming Request Sessions](docs/request-session.md) | Bounded multi-flight runs, per-request token streams, and two-slot HostGraph execution | | [Getting Started](docs/getting-started.md) | Setup, prerequisites, build process, configuration | | [Developer Guide](docs/developer-guide.md) | Directory structure, role ownership, conventions | | [Testing Guide](docs/testing.md) | CI pipeline, test types, writing new tests | diff --git a/docs/dynamic-linking.md b/docs/dynamic-linking.md index 370db402c9..1cd1541a2d 100644 --- a/docs/dynamic-linking.md +++ b/docs/dynamic-linking.md @@ -179,6 +179,7 @@ uses `pthread_key_t` (POSIX TLS) for per-thread state in framework SOs. | `s_orch_thread_idx` | `__thread int` | AICPU SO | Profiling thread index (profiling off by default) | | `g_platform_phase_base` | plain global + `extern "C"` setter | AICPU SO | Device-phase buffer base; published by host (onboard kernel.cpp / sim dlsym), read by the `[STRACE]` phase stamps. Per-thread slotting via the affinity pthread-key index, not TLS. | | strace `inv`/`depth`/`hid` | `pthread_key_t` (`ThreadState`) | host runtime SO | Per-thread `[STRACE]` host-trace state (was C++ `thread_local`, converted per this rule). | +| active arena bank | `pthread_key_t` | onboard host platform SO | Per-thread Host O arena selection for overlapping requests. | | `execution_context` | `thread_local` | Kernel SO (PTO ISA) | Per-thread execution context (fallback, cached values only) | | `NPUMemoryModel::instance` | `thread_local` | Kernel SO (PTO ISA) | Per-core memory model simulation | diff --git a/docs/hierarchical_level_runtime.md b/docs/hierarchical_level_runtime.md index adbd5823bc..8fb8bd8718 100644 --- a/docs/hierarchical_level_runtime.md +++ b/docs/hierarchical_level_runtime.md @@ -310,3 +310,6 @@ device allocation algorithm. | `src/common/worker/chip_worker.{h,cpp}` | L2 `ChipWorker` (kernel-running leaf, runs in the forked chip child) | | `python/bindings/` | nanobind exposure of C++ engine to Python | | `python/simpler/worker.py` | Python `Worker` factory + lifecycle wrapper | + +For bounded multi-flight request execution and per-request token streams on one +L3 Worker, see [Streaming Request Sessions](request-session.md). diff --git a/docs/request-session.md b/docs/request-session.md new file mode 100644 index 0000000000..92015d464d --- /dev/null +++ b/docs/request-session.md @@ -0,0 +1,130 @@ +# Streaming Request Sessions + +`RequestSession` lets an initialized hierarchical `Worker` keep multiple +`Worker.run` calls in flight. It also provides a per-request output stream, +bounded submission, cancellation, and failure isolation. + +`max_active_runs` controls the number of dispatcher threads. Each dispatcher +enters a real `Worker.run` call with its own outer scope. Submission mutates the +shared orchestrator under one lock, while callbacks may wait for token messages +independently. Onboard a2a3 HostGraph chip workers expose two mailbox credits, +so two session dispatchers can execute complete runs concurrently. Other +endpoint types remain single-credit. + +## PR1379 Multi-Flight Architecture + +The bounded multi-flight mailbox has these invariants: + +- one resident child message thread owns mailbox state transitions; +- two long-lived run threads execute complete HostGraph runs, each through an + isolated `ChipWorker`/`DeviceRunner`; +- each run thread uses bank 0 of its private runner-owned arena; +- each run thread owns its Runtime storage and mutable Host run context; +- Host O may overlap an earlier Device S, while a process-level device gate + serializes Device S because the resident AICPU runtime owns process-global + executor, affinity, register, and diagnostic state; +- task completion is correlated by slot generation, so completion order may + differ from submission order; +- control and shutdown operations are exclusive and wait for active task slots + to drain; +- cancellation never interrupts active device work, but it suppresses output + and releases the slot after completion; +- onboard HostGraph endpoints advertise two credits; TRB, simulation, SUB, and + remote endpoints initially remain single-credit. + +The target request path is: + +```text +parent scheduler -> task slot 0 -> run thread 0 -> runner 0 arena -> completion + -> task slot 1 -> run thread 1 -> runner 1 arena -> completion +``` + +The parent-side endpoint publishes a ready slot without waiting for Device S. +A completion pump reports the matching scheduler completion only after the +child marks that slot done. This is the non-blocking boundary; changing only +`submit_next_level` would leave the current synchronous endpoint bottleneck in +place. + +HostApi thread context, scheduler completion latches, the HostGraph epoch/token +pipeline, runner-owned arenas, and concurrent top-level `Worker.run` batching +remain. +There is no cross-request admission lane or split prepare/execute ABI: every +mailbox task is a self-contained run. + +## API + +Open one session on an initialized L3-or-higher Worker: + +```python +def request_orch(orch, request, request_id, emitter, config): + token = build_and_submit(orch, request, request_id, config) + emitter.emit(token, final=token.is_final) + +with worker.open_request_session( + request_orch, max_pending=8, max_active_runs=2 +) as session: + stream_a = session.submit(request_a, config=config, request_id=100) + stream_b = session.submit(request_b, config=config, request_id=101) + + forward_concurrently(stream_a, stream_b) +``` + +`submit` returns immediately after enqueueing. Request IDs are unsigned 64-bit +integers and must be unique among live requests in the session. The session +rejects a full pending queue with `RequestBackpressureError`; `max_pending` +bounds queued work and does not count requests already owned by dispatchers. + +`RequestStream.next(timeout)` returns one item, raises `StopIteration` after a +successful terminal item, and re-raises the request's terminal exception on +failure. `RequestEmitter.emit()` does not return until `next()` has received +that item. This one-item handoff prevents L3 from accumulating several tokens +before returning them to the user. Consumers must drain the stream; +`wait(timeout)` alone cannot complete while a token is waiting for delivery. +A callback exception terminates its request stream, then the dispatcher +continues with later queued requests. + +Only one request session may own a Worker. Its dispatcher threads are the only +threads allowed to call `Worker.run` until the session closes. Closing a +session stops new admission, drains every dispatcher, and releases ownership. + +## Cancellation + +`RequestStream.cancel()` is a soft cancellation: + +- a queued request is skipped before `Worker.run`; +- active device work is drained to preserve arena and scope lifetimes; +- later output items are suppressed; +- consumers observe `RequestCancelledError`. + +Cancellation does not interrupt an AICPU or AICore operation. Reusing its +buffers before Device S finishes would allow the next Host O to overwrite live +state, so active cancellation deliberately favors memory safety over immediate +device preemption. + +Each run thread has a dedicated `ChipWorker`, `DeviceRunner`, runtime buffer, +and device arena. Both select bank 0 because HostGraph async threads do not +inherit the caller's thread-local bank selection. Host O can overlap an +earlier Device S without sharing mutable Host runtime state. Device S remains +bounded to one active launch per child process. Control operations are still +exclusive: create token queues and other L3-L2 resources before publishing the +first task when later requests will reuse them. + +The HostGraph token trailer carries `request_id`, sequence, token value, final +state, and status. L3 decodes it into `HostGraphToken`. A depth-one L2-to-L3 +token queue plus the acknowledged `emit()` handoff forms this flow-control +chain: + +```text +L2 publish → L3 read → user next() → L3 release → L2 may publish next token +``` + +## Implementation + +| Path | Role | +| ---- | ---- | +| `python/simpler/request_session.py` | Session and protocol | +| `python/simpler/worker.py` | Worker ownership and two-slot child execution | +| `src/common/hierarchical/scope.cpp` | Per-run thread-local scope frames | +| `src/common/hierarchical/worker_manager.cpp` | Credit and mailbox-slot protocol | +| `src/common/worker/chip_worker.cpp` | Per-run runtime ownership and arena selection | +| `src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp` | HostGraph epoch/token pipeline | diff --git a/docs/worker-manager.md b/docs/worker-manager.md index 976c137116..4246b04e24 100644 --- a/docs/worker-manager.md +++ b/docs/worker-manager.md @@ -14,6 +14,12 @@ drives a shared-memory mailbox consumed by a forked Python child. The child runs the real worker (a `ChipWorker` for NEXT_LEVEL, a Python callable for SUB) in its own address space. +`WorkerEndpointCaps::max_in_flight` defines the dispatch credits of one +endpoint. `WorkerThread` owns that many executor threads and reports itself +available while at least one credit remains. Onboard a2a3 HostGraph chip +mailboxes advertise two credits; TRB, simulation, SUB, nested Worker, and +remote endpoints advertise one. + The remote L3 design keeps this local fork/shm path behind `LocalMailboxEndpoint` and reserves the same `WorkerEndpoint` boundary for a framed `RemoteL3Endpoint` for cross-host NEXT_LEVEL children. A remote endpoint @@ -79,8 +85,8 @@ the public worker id from the local worker vector index. - **Directed NEXT_LEVEL lookup**: `get_worker_by_id` resolves the exact stable target selected by the user; the Scheduler never asks the manager to choose another NEXT_LEVEL worker -- **SUB-only idle selection**: `pick_idle_sub_excluding` chooses an idle SUB - worker not already used by the same SUB group +- **SUB-only credit selection**: `pick_idle_sub_excluding` chooses a SUB worker + with available endpoint credit that is not already used by the same SUB group Callable and remote-buffer eligibility is validated against the exact target during Orchestrator submission. It is not scheduling metadata and is not @@ -112,19 +118,21 @@ public: private: Ring *ring_; // reads slot state via ring->slot_state(id) std::unique_ptr endpoint_; - std::thread thread_; + std::vector threads_; std::queue queue_; std::mutex mu_; std::condition_variable cv_; + uint32_t capacity_; + std::atomic available_; void loop(); WorkerCompletion dispatch_process(WorkerDispatch d); }; ``` -The WorkerThread's `std::thread` pumps the internal queue and calls -`endpoint->run(...)` once per dispatch. `LocalMailboxEndpoint::run` drives the -shm handshake — one mailbox round trip per dispatch. The forked child loop +The WorkerThread executor threads pump one shared queue and call +`endpoint->run(...)` once per reserved credit. `LocalMailboxEndpoint::run` +drives one mailbox frame round trip per dispatch. The forked child loop that consumes the mailbox lives in Python (`_chip_process_loop` / `_sub_worker_loop` in `python/simpler/worker.py`); the parent does not fork children. @@ -143,7 +151,13 @@ group sub-index. ## 3. Dispatch via shm mailbox -Each `LocalMailboxEndpoint` drives a `MAILBOX_SIZE`-byte `MAP_SHARED` region. +Each `LocalMailboxEndpoint` drives a `MAILBOX_SIZE`-byte `MAP_SHARED` region +containing two fixed-size task frames. Slot generation identifies frame reuse; +the protocol magic/version rejects mismatched parent and child layouts. An +onboard a2a3 HostGraph child owns two long-lived task threads. Each thread has +its own `ChipWorker`/`DeviceRunner`, Runtime storage, streams, kernel arguments, +diagnostic collectors, and private device arena. This isolates Host run state +instead of making `DeviceRunnerBase`'s mutable launch state concurrent. The Python facade forks one child per mailbox **before** `WorkerManager::start()` (so the parent has only the Python main thread when fork runs, avoiding the classical "fork in a multi-threaded process" hazard) @@ -154,37 +168,42 @@ and the child polls the mailbox for the lifetime of the worker. ```cpp WorkerCompletion LocalMailboxEndpoint::run(Ring *ring, WorkerDispatch d) { TaskSlotState &s = *ring->slot_state(d.task_slot); - char *m = static_cast(mailbox_); + size_t slot = acquire_task_slot(); + char *m = static_cast(mailbox_) + slot * MAILBOX_TASK_SLOT_SIZE; - // Write task data: reserved callable field, config, digest prefix, then + // Write task data: generation, protocol, config, digest prefix, then // length-prefixed TaskArgs blob. Tags are stripped; only // [digest][T][S][tensors][scalars] crosses the fork boundary. - uint64_t reserved = 0; - memcpy(m + MAILBOX_OFF_CALLABLE, &reserved, sizeof(reserved)); + uint64_t generation = next_generation(slot); + memcpy(m + MAILBOX_OFF_GENERATION, &generation, sizeof(generation)); + write_protocol(m, MAILBOX_PROTOCOL_MAGIC_VERSION); memcpy(m + MAILBOX_OFF_CONFIG, &s.config, sizeof(CallConfig)); memcpy(m + MAILBOX_OFF_TASK_CALLABLE_HASH, s.callable.digest.data(), 32); const TaskArgs &args = s.is_group() ? s.task_args_list[d.group_index] : s.task_args; write_blob(m + MAILBOX_OFF_TASK_ARGS_BLOB, args); // Signal child - write_state(mailbox_, MailboxState::TASK_READY); + write_state(m, MailboxState::TASK_READY); // Poll for completion - while (read_state(mailbox_) != MailboxState::TASK_DONE) + while (read_state(m) != MailboxState::TASK_DONE) std::this_thread::sleep_for(std::chrono::microseconds(50)); - int err = read_error(mailbox_); - write_state(mailbox_, MailboxState::IDLE); + check_generation(m, generation); + int err = read_error(m); + write_state(m, MailboxState::IDLE); + release_task_slot(slot); return err == 0 ? WorkerCompletion{d.task_slot, d.group_index, EndpointOutcome::SUCCESS, ""} : WorkerCompletion{d.task_slot, d.group_index, EndpointOutcome::TASK_FAILURE, - read_error_msg(mailbox_)}; + read_error_msg(m)}; } ``` Parent-side cost per dispatch: -- One reserved `uint64`, one `CallConfig`, one 32-byte digest, and one +- One generation `uint64`, one protocol word, one `CallConfig`, one 32-byte + digest, and one TaskArgs blob - One signal (`write_state`) - Poll loop with `sleep_for(50us)` (not busy-wait) @@ -205,23 +224,30 @@ The child inherits the parent's full address space at fork time, so: - The Python callable registry is COW-visible - Tensor data in `torch.share_memory_()` regions is fully shared (MAP_SHARED) +Control commands remain exclusive with task frames. Request-session resources +that require control operations, such as L3-L2 message regions, must therefore +be created before the first device task is published. The two-flight hardware +test creates both token queues before request A submits its task; request B then +reuses its existing queue without draining A. + ### 3.3 Mailbox layout ```text offset 0: int32 state offset 4: int32 error -offset 8: uint64 reserved task callable field +offset 8: uint64 task generation or control sub-command offset 16: CallConfig config MAILBOX_OFF_TASK_CALLABLE_HASH: uint8[32] callable digest MAILBOX_OFF_TASK_ARGS_BLOB: bytes [int32 T][int32 S] [Tensor x T][uint64_t x S] -tail: fixed-size NUL-terminated error message +MAILBOX_OFF_PROTOCOL: uint64 protocol magic/version +frame tail: fixed-size NUL-terminated error message ``` -The current mailbox size is the C++ `MAILBOX_SIZE` constant exported through -the nanobind module; Python derives its offsets from the same binding where -possible so the two sides cannot drift silently. +`MAILBOX_SIZE` is two `MAILBOX_TASK_SLOT_SIZE` frames. The constants are +exported through the nanobind module; Python derives its protocol, offsets, +and allocation size from that binding so the two sides cannot drift silently. ### 3.4 Shutdown @@ -229,7 +255,7 @@ possible so the two sides cannot drift silently. endpoint; for `LocalMailboxEndpoint` this writes `SHUTDOWN` to the mailbox. Each child loop sees it on its next poll and exits. The Python facade owns the child PIDs and calls `waitpid()` after writing `SHUTDOWN` to its own mailbox -copy. The parent's `WorkerThread::stop()` only joins the C++ dispatcher thread +copy. The parent's `WorkerThread::stop()` only joins the C++ dispatcher threads — it does not own the child process. --- diff --git a/python/bindings/task_interface.cpp b/python/bindings/task_interface.cpp index 15a20ca6c7..45aec2e815 100644 --- a/python/bindings/task_interface.cpp +++ b/python/bindings/task_interface.cpp @@ -876,6 +876,22 @@ NB_MODULE(_task_interface, m) { .def("tensor_count", &TaskArgs::tensor_count) .def("scalar_count", &TaskArgs::scalar_count) + .def( + "_write_blob", + [](const TaskArgs &self, uint64_t dst_ptr, size_t capacity) { + const size_t needed = task_args_blob_size(self); + if (dst_ptr == 0 || needed > capacity) { + throw std::runtime_error( + "TaskArgs._write_blob: need " + std::to_string(needed) + ", capacity " + + std::to_string(capacity) + ); + } + write_blob(reinterpret_cast(dst_ptr), self); + return needed; + }, + nb::arg("dst_ptr"), nb::arg("capacity"), "Serialize this TaskArgs into caller-owned shared memory." + ) + .def("clear", &TaskArgs::clear) .def( @@ -1235,6 +1251,23 @@ NB_MODULE(_task_interface, m) { // --- CallConfig --- nb::class_(m, "CallConfig") .def(nb::init<>()) + .def_prop_ro_static( + "_blob_size", + [](nb::handle) { + return sizeof(CallConfig); + } + ) + .def( + "_write_blob", + [](const CallConfig &self, uint64_t dst_ptr, size_t capacity) { + if (dst_ptr == 0 || capacity < sizeof(CallConfig)) { + throw std::runtime_error("CallConfig._write_blob: destination is null or too small"); + } + std::memcpy(reinterpret_cast(dst_ptr), &self, sizeof(CallConfig)); + return sizeof(CallConfig); + }, + nb::arg("dst_ptr"), nb::arg("capacity"), "Serialize the CallConfig POD into shared memory." + ) .def_rw("block_dim", &CallConfig::block_dim) .def_rw("aicpu_thread_num", &CallConfig::aicpu_thread_num) // runtime_env returns an internal reference so `cfg.runtime_env.ring_heap = X` @@ -1424,7 +1457,7 @@ NB_MODULE(_task_interface, m) { .def( "run_from_blob", [](ChipWorker &self, int32_t callable_id, uint64_t args_blob_ptr, size_t blob_capacity, - const CallConfig &config) { + const CallConfig &config, unsigned pipeline_slot, uint64_t accepted_state_addr) { // The mailbox region is the on-wire format `write_blob` produced; // `read_blob` is the matching reader that returns a zero-copy // TaskArgsView into the caller-owned bytes. Forwards to the @@ -1432,9 +1465,14 @@ NB_MODULE(_task_interface, m) { // loops never re-implement the tensor/scalar layout in Python // (where it has historically dropped fields like child_memory). TaskArgsView view = read_blob(reinterpret_cast(args_blob_ptr), blob_capacity); - self.run(callable_id, view, config); + ChipStorageTaskArgs storage = view_to_chip_storage(view); + self.run( + callable_id, &storage, config, pipeline_slot, + reinterpret_cast(accepted_state_addr) + ); }, nb::arg("callable_id"), nb::arg("args_blob_ptr"), nb::arg("blob_capacity"), nb::arg("config"), + nb::arg("pipeline_slot") = 0, nb::arg("accepted_state_addr") = 0, nb::call_guard(), "Launch a callable_id from a raw mailbox-blob pointer + capacity " "(used by chip-child mailbox loops to avoid Python-side re-deserialisation " "of the per-task tensor/scalar layout). The blob must be in the format " @@ -1452,6 +1490,14 @@ NB_MODULE(_task_interface, m) { ) .def_prop_ro("device_id", &ChipWorker::device_id) .def_prop_ro("initialized", &ChipWorker::initialized) + .def_prop_ro( + "pipeline_slot_count", &ChipWorker::pipeline_slot_count, + "Pipeline depth declared by the bound runtime's resource contract." + ) + .def_prop_ro( + "arena_bank_count", &ChipWorker::arena_bank_count, + "Number of independently filled arena banks declared by the bound runtime." + ) .def_prop_ro( "aicpu_dlopen_count", &ChipWorker::aicpu_dlopen_count, "Number of distinct callable entries the AICPU has dlopened for on the " diff --git a/python/bindings/worker_bind.h b/python/bindings/worker_bind.h index a74ee3dd84..b470f14ff9 100644 --- a/python/bindings/worker_bind.h +++ b/python/bindings/worker_bind.h @@ -372,20 +372,21 @@ inline void bind_worker(nb::module_ &m) { .def( "add_next_level_worker", - [](Worker &self, uint64_t mailbox_ptr) { - self.add_worker(WorkerType::NEXT_LEVEL, reinterpret_cast(mailbox_ptr)); + [](Worker &self, uint64_t mailbox_ptr, uint32_t max_in_flight) { + self.add_worker(WorkerType::NEXT_LEVEL, reinterpret_cast(mailbox_ptr), max_in_flight); }, - nb::arg("mailbox_ptr"), + nb::arg("mailbox_ptr"), nb::arg("max_in_flight") = 1, "Add a NEXT_LEVEL sub-worker. `mailbox_ptr` is the address of a " "MAILBOX_SIZE-byte MAP_SHARED region; the child process loop is " "Python-managed (fork + _chip_process_loop)." ) .def( "add_next_level_worker_at", - [](Worker &self, int32_t worker_id, uint64_t mailbox_ptr) { - self.add_next_level_worker(worker_id, reinterpret_cast(mailbox_ptr)); + [](Worker &self, int32_t worker_id, uint64_t mailbox_ptr, uint32_t max_in_flight) { + self.add_next_level_worker(worker_id, reinterpret_cast(mailbox_ptr), max_in_flight); }, - nb::arg("worker_id"), nb::arg("mailbox_ptr"), "Add a NEXT_LEVEL sub-worker with an explicit worker id." + nb::arg("worker_id"), nb::arg("mailbox_ptr"), nb::arg("max_in_flight") = 1, + "Add a NEXT_LEVEL sub-worker with an explicit worker id." ) .def( "add_sub_worker", @@ -740,7 +741,11 @@ inline void bind_worker(nb::module_ &m) { m.attr("DEFAULT_HEAP_RING_SIZE") = static_cast(DEFAULT_HEAP_RING_SIZE); m.attr("MAILBOX_SIZE") = static_cast(MAILBOX_SIZE); + m.attr("MAILBOX_TASK_SLOT_SIZE") = static_cast(MAILBOX_TASK_SLOT_SIZE); + m.attr("MAILBOX_TASK_SLOT_COUNT") = static_cast(MAILBOX_TASK_SLOT_COUNT); m.attr("MAILBOX_OFF_ERROR_MSG") = static_cast(MAILBOX_OFF_ERROR_MSG); + m.attr("MAILBOX_OFF_PROTOCOL") = static_cast(MAILBOX_OFF_PROTOCOL); + m.attr("MAILBOX_PROTOCOL_MAGIC_VERSION") = MAILBOX_PROTOCOL_MAGIC_VERSION; m.attr("MAILBOX_ERROR_MSG_SIZE") = static_cast(MAILBOX_ERROR_MSG_SIZE); m.attr("MAX_RING_DEPTH") = static_cast(MAX_RING_DEPTH); m.attr("MAX_SCOPE_DEPTH") = static_cast(MAX_SCOPE_DEPTH); diff --git a/python/simpler/l3_l2_orch_comm.py b/python/simpler/l3_l2_orch_comm.py index 4cb5566066..11ed90e775 100644 --- a/python/simpler/l3_l2_orch_comm.py +++ b/python/simpler/l3_l2_orch_comm.py @@ -12,8 +12,11 @@ import ctypes import struct +import threading +import time from dataclasses import dataclass from enum import IntEnum +from multiprocessing.shared_memory import SharedMemory from typing import Any from _task_interface import ( # pyright: ignore[reportMissingImports] @@ -23,11 +26,23 @@ _l3_host_mapped_payload_read, _l3_host_mapped_payload_write, _l3_host_mapped_region_close, + _mailbox_load_i32, + _mailbox_store_i32, ) from .task_interface import Tensor +class L3L2OrchCommCmd(IntEnum): + ALLOC_REGION = 1 + FREE_REGION = 2 + PAYLOAD_WRITE = 3 + PAYLOAD_READ = 4 + SIGNAL_NOTIFY = 5 + SIGNAL_WAIT = 6 + SIGNAL_TEST = 7 + + class NotifyOp(IntEnum): Set = 0 Add = 1 @@ -48,6 +63,17 @@ class L3L2RegionAccessProfile(IntEnum): SIM_POSIX_SHM = 2 +class _ServiceError(IntEnum): + COPY_FAILED = 6 + SIGNAL_TIMEOUT = 7 + + +_STATE_IDLE = 0 +_STATE_READY = 1 +_STATE_DONE = 3 +_POLL_INTERVAL_S = 0.00005 +_DEFAULT_SUBMIT_TIMEOUT_S = 5.0 +_SIGNAL_TIMEOUT_MARGIN_S = 1.0 _MAX_SIGNED_CHRONO_TIMEOUT_NS = 2**63 - 1 # Wire values returned by _l3_host_mapped_counter_wait in _task_interface; must @@ -86,6 +112,16 @@ def _checked_add_u64(lhs: int, rhs: int) -> int: return result +_REQUEST = struct.Struct(" None: + self._shm = shm + buf = shm.buf + assert buf is not None + self._buf = buf + self._state_addr = ctypes.addressof(ctypes.c_char.from_buffer(self._buf)) + _CONTROL_OFF_STATE + self._mu = threading.Lock() + _mailbox_store_i32(self._state_addr, _STATE_IDLE) + + def submit( + self, request: L3L2OrchCommRequest, timeout_s: float = _DEFAULT_SUBMIT_TIMEOUT_S + ) -> L3L2OrchCommResponse: + deadline = time.monotonic() + float(timeout_s) + with self._mu: + while _mailbox_load_i32(self._state_addr) != _STATE_IDLE: + if time.monotonic() >= deadline: + raise TimeoutError("L3-L2 orch comm client timed out waiting for IDLE") + time.sleep(_POLL_INTERVAL_S) + + _REQUEST.pack_into( + self._buf, + _CONTROL_OFF_REQUEST, + int(request.cmd), + int(request.op), + int(request.region_id), + int(request.payload_offset), + int(request.host_ptr), + int(request.payload_bytes), + int(request.counter_addr), + int(request.counter_bytes), + int(request.counter_operand), + 0, + int(request.timeout_ns), + ) + self._buf[_CONTROL_OFF_RESPONSE : _CONTROL_OFF_RESPONSE + _RESPONSE.size] = b"\x00" * _RESPONSE.size + _mailbox_store_i32(self._state_addr, _STATE_READY) + + while _mailbox_load_i32(self._state_addr) != _STATE_DONE: + if time.monotonic() >= deadline: + raise TimeoutError("L3-L2 orch comm client timed out waiting for DONE") + time.sleep(_POLL_INTERVAL_S) + + response = self._read_response() + _mailbox_store_i32(self._state_addr, _STATE_IDLE) + return response + + def _read_response(self) -> L3L2OrchCommResponse: + fields = _RESPONSE.unpack_from(self._buf, _CONTROL_OFF_RESPONSE) + status, error_kind, region_id, observed_counter, matched = fields[:5] + desc_values = fields[5:11] + raw_message = fields[11] + desc = None + if any(int(v) != 0 for v in desc_values): + desc = L3L2OrchRegionDesc(*[int(v) for v in desc_values]) + message = raw_message.split(b"\x00", 1)[0].decode("utf-8", "replace") + return L3L2OrchCommResponse( + status=int(status), + error_kind=int(error_kind), + region_id=int(region_id), + observed_counter=int(observed_counter), + matched=bool(matched), + desc=desc, + message=message, + ) + + @dataclass(frozen=True) class L3L2RegionCreateRequest: magic_version: int diff --git a/python/simpler/request_session.py b/python/simpler/request_session.py new file mode 100644 index 0000000000..11ee6d4721 --- /dev/null +++ b/python/simpler/request_session.py @@ -0,0 +1,336 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Non-blocking request submission for a hierarchical Worker.""" + +from __future__ import annotations + +import queue +import struct +import threading +from dataclasses import dataclass +from typing import Any, Callable + +HOST_GRAPH_TOKEN_STREAM_MAGIC_VERSION = 0x4847535400010000 +HOST_GRAPH_TOKEN_STREAM_TRAILER_SCALARS = 14 +HOST_GRAPH_TOKEN_FINAL = 1 << 0 +HOST_GRAPH_TOKEN_SYNTHETIC = 1 << 1 +_HOST_GRAPH_TOKEN_PACKET = struct.Struct(" None: + """Append the optional HostGraph token-stream trailer to L2 task args.""" + request_id = int(request_id) + if request_id < 0 or request_id > (1 << 64) - 1: + raise ValueError("HostGraph request_id must fit uint64") + queue_scalars = queue_handle.l2_task_arg_scalars() + if len(queue_scalars) + 2 != HOST_GRAPH_TOKEN_STREAM_TRAILER_SCALARS: + raise RuntimeError("HostGraph token-stream trailer layout has changed") + task_args.add_scalar(HOST_GRAPH_TOKEN_STREAM_MAGIC_VERSION) + task_args.add_scalar(request_id) + for scalar in queue_scalars: + task_args.add_scalar(int(scalar)) + + +def decode_host_graph_token(payload: bytes | bytearray | memoryview) -> HostGraphToken: + """Decode one HostGraph token packet read from an L3-L2 queue.""" + view = memoryview(payload) + if view.nbytes != _HOST_GRAPH_TOKEN_PACKET.size: + raise ValueError(f"HostGraph token packet has {view.nbytes} bytes, expected {_HOST_GRAPH_TOKEN_PACKET.size}") + magic, request_id, token_seq, token_id, flags, status = _HOST_GRAPH_TOKEN_PACKET.unpack(view) + if magic != HOST_GRAPH_TOKEN_STREAM_MAGIC_VERSION: + raise ValueError(f"HostGraph token packet has unsupported magic/version 0x{magic:016x}") + return HostGraphToken( + request_id=request_id, + token_seq=token_seq, + token_id=token_id, + is_final=bool(flags & HOST_GRAPH_TOKEN_FINAL), + status=status, + synthetic=bool(flags & HOST_GRAPH_TOKEN_SYNTHETIC), + ) + + +@dataclass(frozen=True) +class _Terminal: + error: BaseException | None + + +@dataclass(frozen=True) +class _StreamItem: + value: Any + consumed: threading.Event + + +class RequestBackpressureError(RuntimeError): + """Raised when the session's bounded pending queue is full.""" + + +class RequestCancelledError(RuntimeError): + """Raised by a stream whose request was cancelled.""" + + +class RequestStream: + """Thread-safe stream returned immediately by ``RequestSession.submit``.""" + + def __init__(self, request_id: int, session: RequestSession) -> None: + self.request_id = int(request_id) + self._session = session + self._items: queue.Queue[Any] = queue.Queue() + self._done = threading.Event() + self._terminal_seen = False + self._terminal_error: BaseException | None = None + self._finish_lock = threading.Lock() + self._pending_consumed: threading.Event | None = None + + @property + def done(self) -> bool: + return self._done.is_set() + + def next(self, timeout: float | None = None) -> Any: + if self._terminal_seen: + if self._terminal_error is not None: + raise self._terminal_error + raise StopIteration + try: + item = self._items.get(timeout=timeout) + except queue.Empty as exc: + raise TimeoutError(f"request {self.request_id} stream timed out") from exc + if isinstance(item, _Terminal): + self._terminal_seen = True + self._terminal_error = item.error + if item.error is not None: + raise item.error + raise StopIteration + assert isinstance(item, _StreamItem) + item.consumed.set() + return item.value + + def wait(self, timeout: float | None = None) -> None: + if not self._done.wait(timeout): + raise TimeoutError(f"request {self.request_id} did not finish") + if self._terminal_error is not None: + raise self._terminal_error + + def __iter__(self): + return self + + def __next__(self): + return self.next() + + def cancel(self) -> bool: + """Request cancellation; active device work is drained but no longer emitted.""" + return self._session.cancel(self.request_id) + + def _emit(self, item: Any, *, final: bool = False) -> None: + consumed = threading.Event() + with self._finish_lock: + if self._done.is_set(): + if isinstance(self._terminal_error, RequestCancelledError): + return + raise RuntimeError(f"request {self.request_id} stream is already complete") + if self._pending_consumed is not None: + raise RuntimeError(f"request {self.request_id} already has an outstanding item") + self._pending_consumed = consumed + self._items.put(_StreamItem(item, consumed)) + consumed.wait() + with self._finish_lock: + if self._pending_consumed is consumed: + self._pending_consumed = None + if self._done.is_set(): + if isinstance(self._terminal_error, RequestCancelledError): + return + if self._terminal_error is not None: + raise self._terminal_error + if final: + self._finish(None) + + def _finish(self, error: BaseException | None) -> None: + with self._finish_lock: + if self._done.is_set(): + return + self._terminal_error = error + self._items.put(_Terminal(error)) + self._done.set() + if self._pending_consumed is not None: + self._pending_consumed.set() + + +class RequestEmitter: + """Producer handle passed to the session orchestration callback.""" + + def __init__(self, stream: RequestStream, session: RequestSession | None = None) -> None: + self._stream = stream + self._session = session + + def emit(self, item: Any, *, final: bool = False) -> None: + self._stream._emit(item, final=final) + + +@dataclass +class _RequestWork: + request: Any + request_id: int + config: Any + stream: RequestStream + cancelled: threading.Event + + +RequestOrchestration = Callable[[Any, Any, int, RequestEmitter, Any], None] + + +class RequestSession: + """Dispatch bounded concurrent ``Worker.run`` calls with per-request streams.""" + + def __init__( + self, + worker: Any, + orchestration: RequestOrchestration, + *, + max_pending: int = 8, + max_active_runs: int = 1, + ) -> None: + if not callable(orchestration): + raise TypeError("RequestSession orchestration must be callable") + if int(max_pending) <= 0: + raise ValueError("RequestSession max_pending must be positive") + if int(max_active_runs) <= 0: + raise ValueError("RequestSession max_active_runs must be positive") + self._worker = worker + self._orchestration = orchestration + self._requests: queue.Queue[_RequestWork | object] = queue.Queue(maxsize=int(max_pending)) + self._state_lock = threading.Lock() + self._accepting = True + self._next_request_id = 1 + self._live_work: dict[int, _RequestWork] = {} + self._dispatcher_idents: set[int] = set() + self._close_lock = threading.Lock() + self._closed = threading.Event() + self._threads = [ + threading.Thread( + target=self._dispatch_loop, + name=f"simpler-request-session-{index}", + daemon=True, + ) + for index in range(int(max_active_runs)) + ] + + def _start(self) -> None: + for thread in self._threads: + thread.start() + + def submit(self, request: Any, *, config: Any = None, request_id: int | None = None) -> RequestStream: + with self._state_lock: + if not self._accepting: + raise RuntimeError("RequestSession is closed") + if request_id is None: + while self._next_request_id in self._live_work: + self._next_request_id += 1 + request_id = self._next_request_id + self._next_request_id += 1 + request_id = int(request_id) + if request_id < 0 or request_id > (1 << 64) - 1: + raise ValueError("request_id must fit uint64") + if request_id in self._live_work: + raise ValueError(f"request_id {request_id} is already live in this session") + stream = RequestStream(request_id, self) + work = _RequestWork(request, request_id, config, stream, threading.Event()) + try: + self._requests.put_nowait(work) + except queue.Full as exc: + raise RequestBackpressureError( + f"RequestSession pending queue is full (request_id={request_id})" + ) from exc + self._live_work[request_id] = work + return stream + + def cancel(self, request_id: int) -> bool: + request_id = int(request_id) + with self._state_lock: + work = self._live_work.get(request_id) + if work is None or work.cancelled.is_set() or work.stream.done: + return False + work.cancelled.set() + work.stream._finish(RequestCancelledError(f"request {request_id} was cancelled")) + return True + + def close(self) -> None: + with self._close_lock: + if self._closed.is_set(): + return + with self._state_lock: + self._accepting = False + live_work = list(self._live_work.values()) + for work in live_work: + work.cancelled.set() + + for work in live_work: + work.stream._finish(RequestCancelledError(f"request {work.request_id} was cancelled by session close")) + for _thread in self._threads: + self._requests.put(_STOP) + current_ident = threading.get_ident() + for thread in self._threads: + if thread.ident != current_ident: + thread.join() + self._worker._release_request_session(self) + self._closed.set() + + def _owns_current_thread(self) -> bool: + with self._state_lock: + return threading.get_ident() in self._dispatcher_idents + + def _retire(self, work: _RequestWork) -> None: + with self._state_lock: + if self._live_work.get(work.request_id) is work: + self._live_work.pop(work.request_id, None) + + def _dispatch_loop(self) -> None: + ident = threading.get_ident() + with self._state_lock: + self._dispatcher_idents.add(ident) + try: + while True: + work = self._requests.get() + if work is _STOP: + return + assert isinstance(work, _RequestWork) + if work.cancelled.is_set(): + self._retire(work) + continue + emitter = RequestEmitter(work.stream, self) + + def task_orch(orch, _args, cfg, _work=work, _emitter=emitter): + self._orchestration(orch, _work.request, _work.request_id, _emitter, cfg) + + try: + self._worker.run(task_orch, args=None, config=work.config) + work.stream._finish(None) + except BaseException as exc: # noqa: BLE001 + work.stream._finish(exc) + finally: + self._retire(work) + finally: + with self._state_lock: + self._dispatcher_idents.discard(ident) + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + self.close() + return False diff --git a/python/simpler/task_interface.py b/python/simpler/task_interface.py index 5c2190471a..9a72fb31bd 100644 --- a/python/simpler/task_interface.py +++ b/python/simpler/task_interface.py @@ -34,7 +34,11 @@ from _task_interface import ( # pyright: ignore[reportMissingImports] MAILBOX_ERROR_MSG_SIZE, MAILBOX_OFF_ERROR_MSG, + MAILBOX_OFF_PROTOCOL, + MAILBOX_PROTOCOL_MAGIC_VERSION, MAILBOX_SIZE, + MAILBOX_TASK_SLOT_COUNT, + MAILBOX_TASK_SLOT_SIZE, MAX_REGISTERED_CALLABLE_IDS, MAX_TENSOR_DIMS, ArgDirection, @@ -83,7 +87,11 @@ "TaskState", "_Worker", "MAILBOX_SIZE", + "MAILBOX_TASK_SLOT_SIZE", + "MAILBOX_TASK_SLOT_COUNT", "MAILBOX_OFF_ERROR_MSG", + "MAILBOX_OFF_PROTOCOL", + "MAILBOX_PROTOCOL_MAGIC_VERSION", "MAILBOX_ERROR_MSG_SIZE", "read_args_from_blob", # Dynamic CommDomain allocation (orch-only API) @@ -1029,6 +1037,16 @@ def finalize(self): self._identity_registry.clear() self._live_handles.clear() + @property + def pipeline_slot_count(self) -> int: + """Pipeline depth declared by the initialized runtime.""" + return int(self._impl.pipeline_slot_count) + + @property + def arena_bank_count(self) -> int: + """Number of independently filled runtime arena banks.""" + return int(self._impl.arena_bank_count) + def _allocate_slot_locked(self) -> int: for slot_id in range(MAX_REGISTERED_CALLABLE_IDS): if slot_id not in self._callable_registry: diff --git a/python/simpler/worker.py b/python/simpler/worker.py index 3236185cd7..c4be6b4b5b 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -129,7 +129,10 @@ def my_l4_orch(orch, args, config): from .task_interface import ( MAILBOX_ERROR_MSG_SIZE, MAILBOX_OFF_ERROR_MSG, + MAILBOX_OFF_PROTOCOL, + MAILBOX_PROTOCOL_MAGIC_VERSION, MAILBOX_SIZE, + MAILBOX_TASK_SLOT_SIZE, CallConfig, ChipCallable, ChipDomainContext, @@ -156,6 +159,13 @@ def my_l4_orch(orch, args, config): _L3_L2_ENDPOINT_ERROR_REGION_RE = re.compile(r"\bL3-L2 endpoint error\b[^\n]*\bregion=(\d+)\b") +class _ConcurrentRunBatch: + def __init__(self) -> None: + self.active = 0 + self.done = threading.Event() + self.error: BaseException | None = None + + # --------------------------------------------------------------------------- # Unified mailbox layout (must match worker_manager.h MAILBOX_OFF_*) # --------------------------------------------------------------------------- @@ -186,7 +196,7 @@ def my_l4_orch(orch, args, config): # MAILBOX_ARGS_CAPACITY mirrors the C++ constexpr in worker_manager.h so the # Python reader can bounds-check incoming args blobs. Source-of-truth for the # constants on the right is the nanobind binding (cannot drift). -_MAILBOX_ARGS_CAPACITY = MAILBOX_SIZE - _OFF_TASK_ARGS_BLOB - MAILBOX_ERROR_MSG_SIZE +_MAILBOX_ARGS_CAPACITY = MAILBOX_OFF_PROTOCOL - _OFF_TASK_ARGS_BLOB _OFF_CONTROL_CALLABLE_HASH = _OFF_ARGS + 32 # MAILBOX_OFF_ERROR_MSG / MAILBOX_ERROR_MSG_SIZE come from the C++ # nanobind module so the two sides cannot drift. @@ -194,6 +204,7 @@ def my_l4_orch(orch, args, config): _IDLE = 0 _TASK_READY = 1 _TASK_DONE = 2 +_TASK_ACCEPTED = 8 _SHUTDOWN = 3 _CONTROL_REQUEST = 4 _CONTROL_DONE = 5 @@ -222,6 +233,11 @@ def my_l4_orch(orch, args, config): # to close gracefully (so it unlinks the nested mailbox shms only it knows the # names of) before being SIGKILLed. This bounds that graceful wait. _ROLLBACK_GRACEFUL_TIMEOUT_S = 10.0 +# Communicator teardown can legitimately remain inside the native HCCL +# destroy barrier for up to 120 seconds. Give chip children a matching grace +# window once a base communicator exists; ordinary child teardown keeps the +# shorter rollback bound above. +_COMM_REAP_GRACEFUL_TIMEOUT_S = 130.0 # Bounded re-check interval for a close() joiner waiting on an in-flight # _CloseAttempt. A joiner normally wakes immediately on the completing thread's # notify_all(); the timeout is a backstop so that if that notify is skipped (an @@ -263,6 +279,7 @@ def my_l4_orch(orch, args, config): _CTRL_PY_REGISTER = 10 _CTRL_PY_UNREGISTER = 11 _CTRL_PY_IMPORT_REGISTER = 12 +_CTRL_L3_L2_ORCH_COMM_INIT = 13 # Host-buffer registration. MAP_HOST maps a named host-buffer shm # into every local L3 child *post-fork* and keeps it mapped so later runs can copy # through it; UNMAP_HOST drops one. The child also records the parent VA range @@ -862,8 +879,9 @@ def _read_control_digest(buf) -> bytes: return bytes(buf[_OFF_CONTROL_CALLABLE_HASH : _OFF_CONTROL_CALLABLE_HASH + CALLABLE_HASH_DIGEST_BYTES]) -def _read_task_digest(buf) -> bytes: - return bytes(buf[_OFF_TASK_CALLABLE_HASH : _OFF_TASK_CALLABLE_HASH + CALLABLE_HASH_DIGEST_BYTES]) +def _read_task_digest(buf, frame_offset: int = 0) -> bytes: + start = frame_offset + _OFF_TASK_CALLABLE_HASH + return bytes(buf[start : start + CALLABLE_HASH_DIGEST_BYTES]) def _format_digest(digest: bytes) -> str: @@ -935,7 +953,7 @@ def _buffer_field_addr(buf, offset: int) -> int: return ctypes.addressof(ctypes.c_char.from_buffer(buf)) + offset -def _write_error(buf, code: int, msg: str = "") -> None: +def _write_error(buf, code: int, msg: str = "", frame_offset: int = 0) -> None: """Write an (error code, message) tuple into the mailbox error region. The message is UTF-8-encoded and truncated to ``MAILBOX_ERROR_MSG_SIZE - 1`` @@ -943,10 +961,10 @@ def _write_error(buf, code: int, msg: str = "") -> None: NUL-terminated content. On success (code=0) callers may pass an empty message; the region is zero-padded. """ - struct.pack_into("i", buf, _OFF_ERROR, code) + struct.pack_into("i", buf, frame_offset + _OFF_ERROR, code) encoded = msg.encode("utf-8", "replace") n = min(len(encoded), MAILBOX_ERROR_MSG_SIZE - 1) - start = MAILBOX_OFF_ERROR_MSG + start = frame_offset + MAILBOX_OFF_ERROR_MSG buf[start : start + n] = encoded[:n] # Zero-pad the remaining bytes so stale content from a previous dispatch # never leaks into the current error report. @@ -966,7 +984,15 @@ def _format_exc(prefix: str, exc: BaseException) -> str: return f"{prefix}: {type(exc).__name__}: {exc}" -def _read_args_from_mailbox(buf) -> TaskArgs: +def _validate_task_protocol(buf, frame_offset: int = 0) -> None: + protocol = struct.unpack_from("Q", buf, frame_offset + MAILBOX_OFF_PROTOCOL)[0] + if protocol != MAILBOX_PROTOCOL_MAGIC_VERSION: + raise RuntimeError( + f"mailbox protocol mismatch: got 0x{protocol:016x}, expected 0x{MAILBOX_PROTOCOL_MAGIC_VERSION:016x}" + ) + + +def _read_args_from_mailbox(buf, frame_offset: int = 0) -> TaskArgs: """Decode the TaskArgs blob written by C++ write_blob from the mailbox. Used by the Python-targeted child loops (sub_worker, nested L4+ child) @@ -983,10 +1009,10 @@ def _read_args_from_mailbox(buf) -> TaskArgs: (HCCL window slots etc.) — now structurally impossible. """ mailbox_addr = ctypes.addressof(ctypes.c_char.from_buffer(buf)) - return read_args_from_blob(mailbox_addr + _OFF_TASK_ARGS_BLOB) + return read_args_from_blob(mailbox_addr + frame_offset + _OFF_TASK_ARGS_BLOB) -def _sub_worker_loop( +def _sub_worker_loop( # noqa: PLR0912 -- unified task/control mailbox state machine buf, registry: dict[int, Any], identity_table: dict[bytes, int], @@ -1006,6 +1032,12 @@ def _sub_worker_loop( while True: state = _mailbox_load_i32(state_addr) if state == _TASK_READY: + try: + _validate_task_protocol(buf) + except Exception as e: # noqa: BLE001 + _write_error(buf, 1, _format_exc("sub_worker", e)) + _mailbox_store_i32(state_addr, _TASK_DONE) + continue digest = _read_task_digest(buf) cid = identity_table.get(digest) fn = registry.get(int(cid)) if cid is not None else None @@ -1230,6 +1262,25 @@ def _create_sim_l3_l2_region( return region, meta +def _handle_ctrl_l3_l2_orch_comm_init(cw: ChipWorker, buf: memoryview) -> SharedMemory: + control_shm_name = _read_shm_name(buf, _OFF_ARGS) + control_shm = SharedMemory(name=control_shm_name) + control_buf = control_shm.buf + assert control_buf is not None + exported = ctypes.c_char.from_buffer(control_buf) + success = False + try: + control_block_addr = ctypes.addressof(exported) + cw.l3_l2_orch_comm_init_from_addr(control_block_addr, control_shm.size) + success = True + finally: + del exported + del control_buf + if not success: + control_shm.close() + return control_shm + + def _create_onboard_l3_l2_region( cw: ChipWorker, request: L3L2RegionCreateRequest, region_id: int, counter_offset: int, total_bytes: int ) -> tuple[_L2HostL3L2Region, _L2HostL3L2RegionReplyMeta]: @@ -1369,14 +1420,27 @@ def _comm_base_handle(cw: ChipWorker) -> int: return int(handle) -def _ensure_prepared(cw, registry, prepared, cid: int, *, device_id: int) -> None: - if cid in prepared: +def _ensure_registered(run_workers, registry, registered_callables, cid: int, *, device_id: int) -> None: + if cid in registered_callables: return + if not isinstance(run_workers, (list, tuple)): + run_workers = [run_workers] callable_obj = registry.get(cid) if callable_obj is None: raise RuntimeError(f"chip_process dev={device_id}: cid {cid} not in registry") - cw._register_callable_at_slot(cid, callable_obj) - prepared.add(cid) + registered = [] + try: + for worker in run_workers: + worker._register_callable_at_slot(cid, callable_obj) + registered.append(worker) + except Exception: + for worker in reversed(registered): + try: + worker._unregister_slot(cid) + except Exception: # noqa: BLE001 + pass + raise + registered_callables.add(cid) def _run_chip_main_loop( # noqa: PLR0912, PLR0913, PLR0915 -- unified TASK_READY / CONTROL_REQUEST state machine @@ -1393,6 +1457,8 @@ def _run_chip_main_loop( # noqa: PLR0912, PLR0913, PLR0915 -- unified TASK_READ chip_runtime: str = "", on_task_done_success=None, prepared: set[int] | None = None, + run_workers=None, + pipeline_slots: int = 1, ) -> None: """Unified TASK_READY / CONTROL_REQUEST / SHUTDOWN state machine. @@ -1403,68 +1469,93 @@ def _run_chip_main_loop( # noqa: PLR0912, PLR0913, PLR0915 -- unified TASK_READ Returning a non-zero code overrides the kernel's success. TASK_READY carries a callable digest. The child resolves it to a - target-local slot and runs it. The slot must already be prepared: initial - startup-snapshot ChipCallables are prepared before INIT_READY (carried in via - ``prepared``), and callables registered dynamically after startup arrive via - ``_CTRL_PREPARE``. A TASK_READY for an unprepared slot is a control-flow - error and fails the task rather than lazily preparing it. + target-local slot and runs it. Startup-snapshot ChipCallables are registered + before INIT_READY (carried in via ``prepared``), while dynamic registrations + arrive through ``_CTRL_PREPARE``. A TASK_READY for an unregistered slot is a + control-flow error and fails the task rather than registering it lazily. """ - prepared = prepared if prepared is not None else set() + registered_callables = prepared if prepared is not None else set() l3_l2_region_store = _L2HostL3L2RegionStore() + l3_l2_control_shms: list[SharedMemory] = [] # Post-fork host buffers mapped into this child. `host_buf_table` # owns the mmap per token (for unmap + teardown); `host_buf_ranges` is the # parent-VA → child-VA translation table the per-task blob rewrite consults, # rebuilt from the table on every map/unmap. host_buf_table: dict[int, tuple[SharedMemory, int, int, int]] = {} # token -> (shm, lo, hi, child_base) host_buf_ranges: list[tuple[int, int, int]] = [] # (parent_lo, parent_hi, child_base) + control_workers = list(run_workers) if run_workers is not None else [cw] + if not control_workers or control_workers[0] is not cw: + raise RuntimeError("chip run worker 0 must be the control ChipWorker") + if pipeline_slots not in (1, 2): + raise RuntimeError(f"chip pipeline_slots must be 1 or 2, got {pipeline_slots}") + slot_workers = [cw] * pipeline_slots + stop_run_workers = threading.Event() + + def run_task(slot_index: int, slot_worker: ChipWorker) -> None: + frame_offset = slot_index * MAILBOX_TASK_SLOT_SIZE + slot_state_addr = state_addr if slot_index == 0 else mailbox_addr + frame_offset + _OFF_STATE + code = 0 + msg = "" + try: + _validate_task_protocol(buf, frame_offset) + digest = _read_task_digest(buf, frame_offset) + cid = identity_table.get(digest) + if cid is None: + raise RuntimeError(f"callable hash {_format_digest(digest)} not registered") + if cid not in registered_callables: + raise RuntimeError( + f"chip_process dev={device_id}: cid {cid} not registered before TASK_READY " + f"(register via _CTRL_PREPARE first)" + ) + blob_offset = frame_offset + _OFF_TASK_ARGS_BLOB + if host_buf_ranges: + _rewrite_blob_host_addrs(buf, blob_offset, host_buf_ranges) + cfg = _read_config_from_mailbox(buf, frame_offset) + # The shared ChipWorker owns one DeviceRunner and two Runtime + # buffers. The slot index selects the matching Runtime/image bank; + # arena selection is thread-local, so Host O and Device S for this + # request stay on the same bank. + slot_worker._impl.run_from_blob( + int(cid), mailbox_addr + blob_offset, _MAILBOX_ARGS_CAPACITY, cfg, slot_index, slot_state_addr + ) + except Exception as e: # noqa: BLE001 + code = 1 + msg = _format_exc(f"chip_process dev={device_id} slot={slot_index}", e) + + if code == 0 and on_task_done_success is not None: + code, msg = on_task_done_success() + _write_error(buf, code, msg, frame_offset) + _mailbox_store_i32(slot_state_addr, _TASK_DONE) + + def run_task_loop(slot_index: int, slot_worker: ChipWorker) -> None: + frame_offset = slot_index * MAILBOX_TASK_SLOT_SIZE + slot_state_addr = state_addr if slot_index == 0 else mailbox_addr + frame_offset + _OFF_STATE + while not stop_run_workers.is_set(): + if _mailbox_load_i32(slot_state_addr) != _TASK_READY: + time.sleep(0.00005) + continue + run_task(slot_index, slot_worker) + + # Preserve the native runtime's thread affinity on every single-slot + # platform. Only the two-slot a2a3 HostGraph path needs background run + # threads so its control loop can admit another request concurrently. + task_threads = [] + if len(slot_workers) > 1: + task_threads = [ + threading.Thread( + target=run_task_loop, + args=(slot_index, slot_worker), + name=f"simpler-chip-run-{slot_index}", + ) + for slot_index, slot_worker in enumerate(slot_workers) + ] + for thread in task_threads: + thread.start() try: while True: state = _mailbox_load_i32(state_addr) - if state == _TASK_READY: - digest = _read_task_digest(buf) - cid = identity_table.get(digest) - cfg = _read_config_from_mailbox(buf) - - code = 0 - msg = "" - try: - if cid is None: - raise RuntimeError(f"callable hash {_format_digest(digest)} not registered") - # Run only consumes a prepared slot — it never lazily - # prepares. The callable must have been staged via - # _CTRL_PREPARE first; reaching TASK_READY without it is a - # control-flow bug, so fail loudly instead of masking the - # missing-prepare with a first-task latency spike. - if cid not in prepared: - raise RuntimeError( - f"chip_process dev={device_id}: cid {cid} not prepared before TASK_READY " - f"(register via _CTRL_PREPARE first)" - ) - # Redirect any registered host pointer (a parent VA) in the - # blob to this child's own mapping before the runtime reads it. - # No-op when nothing is registered. - if host_buf_ranges: - _rewrite_blob_host_addrs(buf, _OFF_TASK_ARGS_BLOB, host_buf_ranges) - # Hand the mailbox bytes straight to C++ (zero-copy zero-decode): - # the blob layout is what `write_blob` already wrote, so re-parsing - # it in Python is N×40B of avoidable work and a permanent - # opportunity to drop a field. C++ reinterpret_cast - # is the source of truth. - cw._impl.run_from_blob(cid, mailbox_addr + _OFF_TASK_ARGS_BLOB, _MAILBOX_ARGS_CAPACITY, cfg) - except Exception as e: # noqa: BLE001 - code = 1 - msg = _format_exc(f"chip_process dev={device_id}", e) - - # On a successful kernel run, give the caller a chance to do - # post-run work (e.g. store_to_host D2H staging) before the - # parent sees TASK_DONE. The kernel's failure path skips the - # hook because the device output region is undefined and - # staging garbage would mask the real error in post-mortems. - if code == 0 and on_task_done_success is not None: - code, msg = on_task_done_success() - - _write_error(buf, code, msg) - _mailbox_store_i32(state_addr, _TASK_DONE) + if len(slot_workers) == 1 and state == _TASK_READY: + run_task(0, slot_workers[0]) elif state == _CONTROL_REQUEST: sub_cmd = struct.unpack_from("Q", buf, _OFF_CALLABLE)[0] code = 0 @@ -1494,7 +1585,9 @@ def _run_chip_main_loop( # noqa: PLR0912, PLR0913, PLR0915 -- unified TASK_READ raise RuntimeError( f"prepare chip={device_id}: callable hash {_format_digest(digest)} not registered" ) - _ensure_prepared(cw, registry, prepared, int(cid), device_id=device_id) + _ensure_registered( + control_workers, registry, registered_callables, int(cid), device_id=device_id + ) elif sub_cmd == _CTRL_REGISTER: digest = _read_control_digest(buf) payload_size = struct.unpack_from("Q", buf, _CTRL_OFF_ARG0)[0] @@ -1523,20 +1616,32 @@ def _run_chip_main_loop( # noqa: PLR0912, PLR0913, PLR0915 -- unified TASK_READ ) # Self-heal when a prior unregister popped the local # identity table but failed before clearing device - # prepared state for the reusable private slot. - if int(cid) in prepared: + # registered state for the reusable private slot. + if int(cid) in registered_callables: try: - cw._unregister_slot(int(cid)) + for slot_worker in control_workers: + slot_worker._unregister_slot(int(cid)) except Exception: # noqa: BLE001 pass - prepared.discard(int(cid)) + registered_callables.discard(int(cid)) exported = ctypes.c_char.from_buffer(shm_buf) try: addr = ctypes.addressof(exported) - cw._impl.register_callable_from_blob(int(cid), addr) + registered_workers = [] + try: + for slot_worker in control_workers: + slot_worker._impl.register_callable_from_blob(int(cid), addr) + registered_workers.append(slot_worker) + except Exception: + for slot_worker in reversed(registered_workers): + try: + slot_worker._unregister_slot(int(cid)) + except Exception: # noqa: BLE001 + pass + raise finally: del exported - prepared.add(int(cid)) + registered_callables.add(int(cid)) finally: shm_buf.release() # Release the local mmap as soon as prepare returns; @@ -1547,14 +1652,18 @@ def _run_chip_main_loop( # noqa: PLR0912, PLR0913, PLR0915 -- unified TASK_READ digest = _read_control_digest(buf) cid, removed = _remove_local_identity(registry, identity_table, identity_refs, digest) if removed and cid is not None: - cw._unregister_slot(int(cid)) - prepared.discard(int(cid)) + for slot_worker in control_workers: + slot_worker._unregister_slot(int(cid)) + registered_callables.discard(int(cid)) elif sub_cmd == _CTRL_ALLOC_DOMAIN: _handle_ctrl_alloc_domain(cw, buf) elif sub_cmd == _CTRL_RELEASE_DOMAIN: _handle_ctrl_release_domain(cw, buf) elif sub_cmd == _CTRL_COMM_INIT: _handle_ctrl_comm_init(cw, buf) + elif sub_cmd == _CTRL_L3_L2_ORCH_COMM_INIT: + control_shm = _handle_ctrl_l3_l2_orch_comm_init(cw, buf) + l3_l2_control_shms.append(control_shm) elif sub_cmd == _CTRL_MAP_HOST: _handle_ctrl_map_host(buf, host_buf_table, host_buf_ranges) elif sub_cmd == _CTRL_UNMAP_HOST: @@ -1576,8 +1685,27 @@ def _run_chip_main_loop( # noqa: PLR0912, PLR0913, PLR0915 -- unified TASK_READ _mailbox_store_i32(state_addr, _CONTROL_DONE) elif state == _SHUTDOWN: break + else: + time.sleep(0.00005) finally: + stop_run_workers.set() + for thread in task_threads: + thread.join() _sweep_l2_host_l3_l2_regions(l3_l2_region_store) + if l3_l2_control_shms: + try: + cw.l3_l2_orch_comm_shutdown() + except Exception as e: # noqa: BLE001 + sys.stderr.write( + f"[chip_process pid={os.getpid()} dev={device_id}] " + f"WARN: l3_l2_orch_comm_shutdown failed: {type(e).__name__}: {e}\n" + ) + sys.stderr.flush() + for control_shm in reversed(l3_l2_control_shms): + try: + control_shm.close() + except Exception: # noqa: BLE001 + pass for host_shm, _lo, _hi, _base in host_buf_table.values(): try: host_shm.close() @@ -1610,6 +1738,7 @@ def _chip_process_loop( # noqa: PLR0913 -- fork-child entry: all context (bins, """ import traceback as _tb # noqa: PLC0415 + run_workers: list[ChipWorker] = [] try: cw = ChipWorker() cw.init( @@ -1620,6 +1749,7 @@ def _chip_process_loop( # noqa: PLR0913 -- fork-child entry: all context (bins, prewarm_config=prewarm_config, enable_sdma=enable_sdma, ) + run_workers.append(cw) except Exception as e: _tb.print_exc() # Publish the cause into the mailbox and flag INIT_FAILED so the @@ -1627,28 +1757,35 @@ def _chip_process_loop( # noqa: PLR0913 -- fork-child entry: all context (bins, # spinning forever on a child that will never reach INIT_READY. _write_error(buf, 1, _format_exc(f"chip_process dev={device_id} init", e)) _mailbox_store_i32(_buffer_field_addr(buf, _OFF_STATE), _INIT_FAILED) + for worker in reversed(run_workers): + worker.finalize() return # Prepare every ChipCallable in the startup snapshot before publishing # INIT_READY, so the H2D upload + device-orch load is charged inside the # readiness barrier and the first task dispatch pays no upload. The set of - # prepared cids carries into the main loop, which requires a cid be prepared + # registered cids carry into the main loop, which requires a cid be registered # before it dispatches. The parent therefore issues no post-READY # control_prepare for the initial snapshot. prepared: set[int] = set() try: for cid, target in registry.items(): if isinstance(target, ChipCallable): - _ensure_prepared(cw, registry, prepared, int(cid), device_id=device_id) + _ensure_registered(run_workers, registry, prepared, int(cid), device_id=device_id) except Exception as e: _tb.print_exc() _write_error(buf, 1, _format_exc(f"chip_process dev={device_id} prepare", e)) _mailbox_store_i32(_buffer_field_addr(buf, _OFF_STATE), _INIT_FAILED) - cw.finalize() + for worker in reversed(run_workers): + worker.finalize() return mailbox_addr = ctypes.addressof(ctypes.c_char.from_buffer(buf)) state_addr = mailbox_addr + _OFF_STATE + # Publish the runtime's resource-contract depth before INIT_READY. The + # parent reads this immutable bootstrap field when assigning endpoint + # credits, so admission is contract-driven rather than runtime-name based. + struct.pack_into("Q", buf, _CTRL_OFF_RESULT, cw.pipeline_slot_count) # Signal init success. The parent's readiness barrier waits for every chip # child to reach _INIT_READY before dispatching the first task, so the # per-rank host-side stream sync budget only covers actual op execution @@ -1670,12 +1807,15 @@ def _chip_process_loop( # noqa: PLR0913 -- fork-child entry: all context (bins, chip_platform=platform, chip_runtime=runtime, prepared=prepared, + run_workers=run_workers, + pipeline_slots=cw.pipeline_slot_count, ) finally: - cw.finalize() + for worker in reversed(run_workers): + worker.finalize() -def _read_config_from_mailbox(buf: memoryview) -> CallConfig: +def _read_config_from_mailbox(buf: memoryview, frame_offset: int = 0) -> CallConfig: """Reconstruct a CallConfig from the unified mailbox layout.""" ( block_dim, @@ -1687,7 +1827,7 @@ def _read_config_from_mailbox(buf: memoryview) -> CallConfig: scope_stats, *ring_values, prefix_bytes, - ) = _CFG_FMT.unpack_from(buf, _OFF_CONFIG) + ) = _CFG_FMT.unpack_from(buf, frame_offset + _OFF_CONFIG) ring_task_window = list(ring_values[:RUNTIME_ENV_RING_COUNT]) ring_heap = list(ring_values[RUNTIME_ENV_RING_COUNT : 2 * RUNTIME_ENV_RING_COUNT]) ring_dep_pool = list(ring_values[2 * RUNTIME_ENV_RING_COUNT : 3 * RUNTIME_ENV_RING_COUNT]) @@ -1727,6 +1867,12 @@ def _child_worker_loop( while True: state = _mailbox_load_i32(state_addr) if state == _TASK_READY: + try: + _validate_task_protocol(buf) + except Exception as e: # noqa: BLE001 + _write_error(buf, 1, _format_exc(f"child_worker level={inner_worker.level}", e)) + _mailbox_store_i32(state_addr, _TASK_DONE) + continue digest = _read_task_digest(buf) cid = identity_table.get(digest) orch_fn = registry.get(int(cid)) if cid is not None else None @@ -1991,6 +2137,11 @@ def __init__( # fast (this worker does not cancel an in-progress init); any thread may # join an in-flight close(). self._init_owner_thread: threading.Thread | None = None + self._request_session: Any | None = None + self._request_session_lock = threading.Lock() + self._run_cv = threading.Condition() + self._run_batch: _ConcurrentRunBatch | None = None + self._run_draining = False # Narrow lock around `_callable_registry` mutation so concurrent # register / unregister calls don't trip CPython's non-atomic @@ -2084,6 +2235,9 @@ def __init__( # starts the C++ scheduler; no comm work happens there. self._comm_base_ready: bool = False + self._l3_l2_orch_comm_ready: set[int] = set() + self._l3_l2_orch_comm_shms: dict[int, SharedMemory] = {} + self._l3_l2_orch_comm_clients: dict[int, Any] = {} self._live_l3_l2_regions: list[Any] = [] self._l3_l2_orch_comm_host_buffers: dict[int, int] = {} @@ -4279,7 +4433,11 @@ def _setup(inner=inner_worker): # Register chip workers as NEXT_LEVEL (L3) if device_ids: for shm in self._chip_shms: - dw.add_next_level_worker(_mailbox_addr(shm)) + assert shm.buf is not None + max_in_flight = int(struct.unpack_from("Q", shm.buf, _CTRL_OFF_RESULT)[0]) + if max_in_flight not in (1, 2): + raise RuntimeError(f"chip child published invalid pipeline slot count {max_in_flight}") + dw.add_next_level_worker(_mailbox_addr(shm), max_in_flight=max_in_flight) # Register Worker children as NEXT_LEVEL (L4+) if self._next_level_shms and not hasattr(dw, "add_next_level_worker_at"): @@ -4564,6 +4722,50 @@ def live_domains(self) -> dict[str, CommDomainHandle]: """ return dict(self._live_domains) + def _make_l3_l2_orch_comm_client(self, shm: SharedMemory): + from .l3_l2_orch_comm import L3L2OrchCommClient # noqa: PLC0415 + + return L3L2OrchCommClient(shm) + + def _ensure_l3_l2_orch_comm(self, worker_id: int): + from .l3_l2_orch_comm import CONTROL_SHM_SIZE # noqa: PLC0415 + + if self.level < 3: + raise RuntimeError("create_l3_l2_region requires a hierarchical Worker") + if self._worker is None: + raise RuntimeError("create_l3_l2_region requires Worker.init()") + device_ids = self._config.get("device_ids", []) + if worker_id < 0 or worker_id >= len(device_ids): + raise ValueError(f"create_l3_l2_region: worker_id {worker_id} outside [0, {len(device_ids)})") + if worker_id in self._l3_l2_orch_comm_ready: + return self._l3_l2_orch_comm_clients[worker_id] + + chip_shm = self._chip_shms[worker_id] + assert chip_shm.buf is not None + state = _mailbox_load_i32(_buffer_field_addr(chip_shm.buf, _OFF_STATE)) + if state != _IDLE: + raise RuntimeError( + f"create_l3_l2_region bootstrap failed: target worker {worker_id} is busy and " + "the L3-L2 service is not ready" + ) + + control_shm = SharedMemory(create=True, size=CONTROL_SHM_SIZE) + try: + client = self._make_l3_l2_orch_comm_client(control_shm) + self._worker.control_l3_l2_orch_comm_init(worker_id, control_shm.name) + except Exception: + try: + control_shm.close() + control_shm.unlink() + except Exception: # noqa: BLE001 + pass + raise + + self._l3_l2_orch_comm_shms[worker_id] = control_shm + self._l3_l2_orch_comm_clients[worker_id] = client + self._l3_l2_orch_comm_ready.add(worker_id) + return client + def _validate_l3_l2_worker_id(self, worker_id: int) -> None: if self.level < 3: raise RuntimeError("create_l3_l2_region requires a hierarchical Worker") @@ -4573,6 +4775,10 @@ def _validate_l3_l2_worker_id(self, worker_id: int) -> None: if worker_id < 0 or worker_id >= len(device_ids): raise ValueError(f"create_l3_l2_region: worker_id {worker_id} outside [0, {len(device_ids)})") + def _l3_l2_orch_comm_submit(self, worker_id: int, request, timeout_s: float): + client = self._ensure_l3_l2_orch_comm(int(worker_id)) + return client.submit(request, timeout_s) + def _poison_l3_l2_region_from_endpoint_error(self, exc: BaseException) -> bool: match = _L3_L2_ENDPOINT_ERROR_REGION_RE.search(str(exc)) if match is None: @@ -4729,7 +4935,16 @@ def _close_l3_l2_orch_comm(self) -> None: except RuntimeError: pass self._live_l3_l2_regions.clear() + self._l3_l2_orch_comm_clients.clear() + self._l3_l2_orch_comm_ready.clear() self._l3_l2_orch_comm_host_buffers.clear() + for shm in self._l3_l2_orch_comm_shms.values(): + try: + shm.close() + shm.unlink() + except Exception: # noqa: BLE001 + pass + self._l3_l2_orch_comm_shms.clear() # ------------------------------------------------------------------ # Dynamic CommDomain allocation (driven by Orchestrator.allocate_domain; @@ -5614,6 +5829,59 @@ def _find_host_buf_entry(self, addr: int, nbytes: int) -> _HostBufEntry | None: # run — uniform entry point # ------------------------------------------------------------------ + def _enter_concurrent_run(self) -> _ConcurrentRunBatch: + assert self._orch is not None + with self._run_cv: + while self._run_draining: + self._run_cv.wait() + batch = self._run_batch + if batch is None: + batch = _ConcurrentRunBatch() + self._run_batch = batch + self._orch._clear_error() + batch.active += 1 + return batch + + def _leave_concurrent_run(self, batch: _ConcurrentRunBatch) -> bool: + with self._run_cv: + if self._run_batch is not batch or batch.active <= 0: + raise RuntimeError("Worker concurrent run batch state is inconsistent") + batch.active -= 1 + if batch.active != 0: + return False + self._run_batch = None + self._run_draining = True + return True + + def _drain_concurrent_run_batch(self, batch: _ConcurrentRunBatch) -> None: + assert self._orch is not None + try: + try: + self._orch._drain() + except Exception as exc: + self._poison_l3_l2_region_from_endpoint_error(exc) + raise + except BaseException as exc: # noqa: BLE001 + batch.error = exc + finally: + try: + self._release_active_remote_slot_refs() + self._flush_pending_remote_frees() + try: + self._cleanup_l3_l2_regions() + finally: + self._l3_l2_orch_comm_host_buffers.clear() + self._execute_pending_domain_releases() + if self._live_domains: + self._release_all_live_domains() + except BaseException as exc: # noqa: BLE001 + if batch.error is None: + batch.error = exc + with self._run_cv: + self._run_draining = False + batch.done.set() + self._run_cv.notify_all() + def run(self, callable, args=None, config=None) -> None: """Execute one task (L2) or one DAG (L3+) synchronously. @@ -5638,6 +5906,10 @@ def run(self, callable, args=None, config=None) -> None: self._run_locked(callable, args, config) def _run_locked(self, callable, args, config) -> None: + with self._request_session_lock: + request_session = self._request_session + if request_session is not None and not request_session._owns_current_thread(): + raise RuntimeError("Worker.run is owned by an active RequestSession") cfg = config if config is not None else CallConfig() if self.level == 2: @@ -5648,52 +5920,67 @@ def _run_locked(self, callable, args, config) -> None: assert self._orch is not None assert self._worker is not None - # Drop any error stashed by a previous run() so this call starts - # clean. drain() rethrows on the way out; every successful run() - # leaves the error slot empty, but an unrelated caller may have - # poked it. - self._orch._clear_error() - self._orch._scope_begin() + batch = self._enter_concurrent_run() + own_error: BaseException | None = None + own_traceback = None + scope_open = False try: + self._orch._scope_begin() + scope_open = True callable(self._orch, args, cfg) + except BaseException as exc: # noqa: BLE001 + own_error = exc + own_traceback = exc.__traceback__ finally: - # Always release scope refs and drain so ring slots aren't - # stranded when the orch fn raises mid-DAG. drain() also - # rethrows the first dispatch failure for this run — that - # is how child-task exceptions surface to the caller of - # Worker.run(). scope_end deliberately does NOT throw: if - # it did, released refs would be incomplete and drain - # would hang on in-flight tasks. - self._orch._scope_end() - # ORDER MATTERS: drain() must complete first so any in-flight - # task that captured a now-pending handle's device_ctx / - # buffer_ptrs sees live memory. THEN execute the pending - # backend releases. Last, sweep any handles that the orch - # function neither released nor passed out (covers exception - # unwind and "forgot to release" — auto-release in LIFO). - # drain() rethrows the first chip-task/dispatch failure, so the - # cleanup lives in a finally: a failed task must not strand - # backend domain allocations into the next run. - try: + if scope_open: try: - self._orch._drain() - except Exception as e: - self._poison_l3_l2_region_from_endpoint_error(e) - raise - finally: - self._release_active_remote_slot_refs() - self._flush_pending_remote_frees() - try: - self._cleanup_l3_l2_regions() - finally: - self._l3_l2_orch_comm_host_buffers.clear() - self._execute_pending_domain_releases() - if self._live_domains: - self._release_all_live_domains() + self._orch._scope_end() + except BaseException as exc: # noqa: BLE001 + if own_error is None: + own_error = exc + own_traceback = exc.__traceback__ + is_last = self._leave_concurrent_run(batch) + if is_last: + self._drain_concurrent_run_batch(batch) + else: + batch.done.wait() + if own_error is not None: + raise own_error.with_traceback(own_traceback) + if batch.error is not None: + raise batch.error # L3+ returns None like every other worker level; per-L2-child timing # is emitted as `[STRACE]` markers from each simpler_run. return None + def open_request_session(self, orchestration, *, max_pending: int = 8, max_active_runs: int = 1): + """Open a non-blocking request stream over this hierarchical Worker.""" + if self.level < 3: + raise RuntimeError("Worker.open_request_session requires a hierarchical Worker") + from .request_session import RequestSession # noqa: PLC0415 + + with self._request_session_lock: + # Re-check under the session lock. Worker.close() publishes CLOSED + # before taking this lock, so a racing open is either observed and + # closed by that teardown or rejected here. + if not self._initialized: + raise RuntimeError("Worker.open_request_session requires Worker.init()") + if self._request_session is not None: + raise RuntimeError("Worker already has an active RequestSession") + session = RequestSession( + self, + orchestration, + max_pending=int(max_pending), + max_active_runs=int(max_active_runs), + ) + self._request_session = session + session._start() + return session + + def _release_request_session(self, session) -> None: + with self._request_session_lock: + if self._request_session is session: + self._request_session = None + @property def aicpu_dlopen_count(self) -> int: """L2 only: number of distinct callable identities the AICPU has dlopened for. @@ -5869,6 +6156,14 @@ def close(self) -> None: # noqa: PLR0912, PLR0915 -- lifecycle linearization: r if teardown_tree: self._teardown_attempted = True if teardown_tree: + # CLOSED already rejects new run leases. Closing the session now + # stops its idle dispatchers and cancels queued requests before + # their backing native tree is dismantled. An active dispatcher + # was covered by the lease drain above. + with self._request_session_lock: + request_session = self._request_session + if request_session is not None: + request_session.close() self._teardown_ready_tree() except BaseException as exc: # noqa: BLE001 if result is None: @@ -6098,7 +6393,8 @@ def _close_worker() -> None: # teardown entry — so the (blocking) pre-child cleanup above cannot # eat it. Reap removes reclaimed pids/shms in place; a surviving child # is left in place and reported as an error (terminal, not retried). - reap_deadline = time.monotonic() + _ROLLBACK_GRACEFUL_TIMEOUT_S + reap_timeout_s = _COMM_REAP_GRACEFUL_TIMEOUT_S if self._comm_base_ready else _ROLLBACK_GRACEFUL_TIMEOUT_S + reap_deadline = time.monotonic() + reap_timeout_s _step(lambda: self._reap_child_groups(groups, reap_deadline)) _step(self._close_l3_l2_orch_comm) # Drop next-level worker refs only once their pids/shms are reclaimed. diff --git a/src/a2a3/platform/onboard/host/device_runner.cpp b/src/a2a3/platform/onboard/host/device_runner.cpp index b58ddaf875..9c1d84a651 100644 --- a/src/a2a3/platform/onboard/host/device_runner.cpp +++ b/src/a2a3/platform/onboard/host/device_runner.cpp @@ -26,6 +26,7 @@ #include #include #include +#include #include #include #include "acl/acl.h" @@ -60,6 +61,20 @@ extern "C" __attribute__((weak, visibility("hidden"))) int dep_gen_replay_emit_d return -1; } +// host_build_graph overrides this hook to release its Host O worker only after +// per-run profiling collectors are live. Other runtimes keep the no-op. +extern "C" __attribute__((weak, visibility("hidden"))) int release_async_host_graph_pipeline( + Runtime * /*runtime*/ +) { + return 0; +} + +// host_build_graph opens its Host Stage1 worker before Gate B so the build can +// overlap the prior request's Device S. Other runtimes keep the no-op. +extern "C" __attribute__((weak, visibility("hidden"))) int open_async_host_graph_pipeline(Runtime * /*runtime*/) { + return 0; +} + // ============================================================================= // Lazy-loaded HAL (ascend_hal) for halHostRegister / halHostUnregister // ============================================================================= @@ -192,6 +207,17 @@ int DeviceRunner::destroy_comm_stream(void *stream) { // `src/common/platform/onboard/host/device_runner_base.cpp`. int DeviceRunner::run(Runtime &runtime, const CallConfig &config) { + int rc = open_async_host_graph_pipeline(&runtime); + if (rc != 0) { + LOG_ERROR("open_async_host_graph_pipeline failed: %d", rc); + return rc; + } + + // Gate B covers every runner-owned per-run object, not only KernelLaunch. + // The HBG Host Stage1 worker was opened above and uses its banked image, so + // it can continue while this request waits for the prior Device S + reap. + std::unique_lock device_run_lock(device_run_mutex_); + // Latch this run's diagnostic enables onto the runner before the collector // paths below read them; block_dim/aicpu_thread_num are consumed locally. apply_call_config(config); @@ -216,7 +242,7 @@ int DeviceRunner::run(Runtime &runtime, const CallConfig &config) { } if (validate_launch_aicpu_num(launch_aicpu_num) != 0) return -1; - int rc = ensure_device_initialized(); + rc = ensure_device_initialized(); if (rc != 0) { LOG_ERROR("ensure_device_initialized failed: %d", rc); return rc; @@ -428,6 +454,30 @@ int DeviceRunner::run(Runtime &runtime, const CallConfig &config) { l2_swimlane_collector_.set_core_types(core_types.data(), num_aicore); } + rc = release_async_host_graph_pipeline(&runtime); + if (rc != 0) { + LOG_ERROR("release_async_host_graph_pipeline failed: %d", rc); + return rc; + } + + rc = launch_run(runtime, num_aicore, launch_aicpu_num); + if (rc != 0) return rc; + + rc = reap_run(); + if (rc != 0) return rc; + + // Print handshake results (reads from device memory, must be before free) + print_handshake_results(); + + return 0; +} + +int DeviceRunner::launch_run(Runtime &runtime, int num_aicore, int launch_aicpu_num) { + // KernelLaunch is the pipeline boundary: this method clears the handshake + // consumed by the launch and submits exactly the AICore and AICPU kernels. + // It intentionally performs no stream synchronization or per-run cleanup. + int rc = 0; + // Launch the AICore worker BEFORE the AICPU Run task — mirrors the a5 path // so the two arches stay symmetric. First-launch latency optimization + // op-timeout-family defense-in-depth: with the AICPU Run task launched first @@ -453,7 +503,7 @@ int DeviceRunner::run(Runtime &runtime, const CallConfig &config) { LOG_INFO_V0("=== launch_aicore_kernel ==="); // Launch AICore kernel (pass device copy of KernelArgs) - rc = launch_aicore_kernel(stream_aicore_, kernel_args_.device_k_args_); + rc = launch_aicore_kernel(run_stream_aicore(), kernel_args_.device_k_args_); if (rc != 0) { LOG_ERROR("launch_aicore_kernel failed: %d", rc); recover_device_or_mark_unusable(rc); @@ -462,7 +512,7 @@ int DeviceRunner::run(Runtime &runtime, const CallConfig &config) { LOG_INFO_V0("=== launch_aicpu_kernel %s ===", host::KernelNames::RunName); int aicpu_launch_n = (runtime.get_aicpu_launch_count() > 0) ? runtime.get_aicpu_launch_count() : launch_aicpu_num; - rc = launch_aicpu_kernel(stream_aicpu_, &kernel_args_.args, host::KernelNames::RunName, aicpu_launch_n); + rc = launch_aicpu_kernel(run_stream_aicpu(), &kernel_args_.args, host::KernelNames::RunName, aicpu_launch_n); if (rc != 0) { LOG_ERROR("launch_aicpu_kernel (main) failed: %d", rc); // The AICore worker was already launched above and is now spinning in @@ -474,7 +524,16 @@ int DeviceRunner::run(Runtime &runtime, const CallConfig &config) { return rc; } - rc = sync_run_streams(); + // Both launches are now accepted by the runtime. Publish the flight ACK + // before reap blocks on stream completion; TASK_DONE is written later by + // the child loop after validation. + publish_task_accepted(); + + return 0; +} + +int DeviceRunner::reap_run() { + int rc = sync_run_streams(); if (rc != 0) { // sync_run_streams surfaces the AICore op-timeout (STARS-reaped op -> // 507000/507018/507046 at AICPU/AICore stream sync). The op-timeout @@ -513,9 +572,6 @@ int DeviceRunner::run(Runtime &runtime, const CallConfig &config) { } } - // Print handshake results (reads from device memory, must be before free) - print_handshake_results(); - return 0; } diff --git a/src/a2a3/platform/onboard/host/device_runner.h b/src/a2a3/platform/onboard/host/device_runner.h index 62904e2939..a8db15bf92 100644 --- a/src/a2a3/platform/onboard/host/device_runner.h +++ b/src/a2a3/platform/onboard/host/device_runner.h @@ -29,6 +29,7 @@ #include #include #include +#include #include #include #include @@ -219,6 +220,13 @@ class DeviceRunner : public DeviceRunnerBase { // recovery. See run() and recover_device_or_mark_unusable(). bool device_unusable_{false}; + // Gate B: one Device S phase at a time for this runner. Host Stage1 uses + // independent resource banks and remains outside this gate. + std::mutex device_run_mutex_; + + int launch_run(Runtime &runtime, int num_aicore, int launch_aicpu_num); + int reap_run(); + // On an AICore launch/sync error, best-effort drain the device so a later // run() on the same DeviceRunner can recover in place; if the drain itself // errors the context is unrecoverable without a full reset, so flip diff --git a/src/a2a3/platform/sim/host/device_runner.cpp b/src/a2a3/platform/sim/host/device_runner.cpp index 614c714b77..e721c47927 100644 --- a/src/a2a3/platform/sim/host/device_runner.cpp +++ b/src/a2a3/platform/sim/host/device_runner.cpp @@ -57,6 +57,12 @@ extern "C" __attribute__((weak, visibility("hidden"))) int dep_gen_replay_emit_d return -1; } +extern "C" __attribute__((weak, visibility("hidden"))) int release_async_host_graph_pipeline( + Runtime * /*runtime*/ +) { + return 0; +} + DeviceRunner::~DeviceRunner() { finalize(); } int DeviceRunner::ensure_binaries_loaded() { @@ -207,6 +213,10 @@ int DeviceRunner::invoke_device_register(const RegisterCallableArgs ®_args) { } int DeviceRunner::run(Runtime &runtime, const CallConfig &config) { + // Match onboard Gate B: two pipeline slots may finish Host O concurrently, + // but one shared runner cannot mutate its per-run state from both slots. + std::unique_lock device_run_lock(device_run_mutex_); + apply_call_config(config); int block_dim = config.block_dim; const int launch_aicpu_num = config.aicpu_thread_num; @@ -474,6 +484,12 @@ int DeviceRunner::run(Runtime &runtime, const CallConfig &config) { scope_stats_collector_.start(thread_factory); } + rc = release_async_host_graph_pipeline(&runtime); + if (rc != 0) { + LOG_ERROR("release_async_host_graph_pipeline failed: %d", rc); + return rc; + } + constexpr int over_launch = PLATFORM_MAX_AICPU_THREADS_JUST_FOR_LAUNCH; LOG_INFO_V0("Launching %d AICPU threads (logical=%d)", over_launch, launch_aicpu_num); std::vector aicpu_threads; diff --git a/src/a2a3/platform/sim/host/device_runner.h b/src/a2a3/platform/sim/host/device_runner.h index 0e762a78cb..a5d976347b 100644 --- a/src/a2a3/platform/sim/host/device_runner.h +++ b/src/a2a3/platform/sim/host/device_runner.h @@ -21,6 +21,7 @@ #define SRC_A2A3_PLATFORM_SIM_HOST_DEVICE_RUNNER_H_ #include +#include #include #include @@ -80,6 +81,11 @@ class DeviceRunner : public SimDeviceRunnerBase { // a2a3-only; a5 has no dep_gen. DepGenCollector dep_gen_collector_; bool enable_dep_gen_{false}; + + // Gate B: the sim runner has process-shared per-run state (KernelArgs, + // executor symbols, collectors, and temporary SO handles). Host binding + // happens before DeviceRunner::run(), so this only serializes Device S. + std::mutex device_run_mutex_; }; #endif // SRC_A2A3_PLATFORM_SIM_HOST_DEVICE_RUNNER_H_ diff --git a/src/a2a3/runtime/host_build_graph/aicpu/aicpu_executor.cpp b/src/a2a3/runtime/host_build_graph/aicpu/aicpu_executor.cpp index b2f7df3572..e4ba17ee6b 100644 --- a/src/a2a3/runtime/host_build_graph/aicpu/aicpu_executor.cpp +++ b/src/a2a3/runtime/host_build_graph/aicpu/aicpu_executor.cpp @@ -16,12 +16,15 @@ #include #include #include +#include #ifdef __linux__ #include #endif +#include "aicpu/l3_l2_message_queue.h" #include "aicpu/device_time.h" #include "callable_protocol.h" +#include "host_graph_token_stream.h" #include "pto2_dispatch_payload.h" #include "runtime.h" #include "spin_hint.h" @@ -78,6 +81,75 @@ static int32_t read_pto2_runtime_status(Runtime *runtime) { static PTO2Runtime *rt{nullptr}; +namespace { + +constexpr uint64_t kHostGraphTokenQueueTimeoutNs = 5000000000ULL; +using HostGraphTokenQueue = L3L2QueueEndpoint<1>; + +class HostGraphTokenPublisher { +public: + explicit HostGraphTokenPublisher(Runtime *runtime) { + const ChipStorageTaskArgs &args = runtime->get_orch_args(); + int32_t scalar_count = args.scalar_count(); + if (scalar_count < PTO2_HOST_GRAPH_TOKEN_STREAM_TRAILER_SCALARS) return; + + int32_t trailer = scalar_count - PTO2_HOST_GRAPH_TOKEN_STREAM_TRAILER_SCALARS; + if (args.scalar(trailer) != PTO2_HOST_GRAPH_TOKEN_STREAM_MAGIC_VERSION) return; + + request_id_ = args.scalar(trailer + 1); + L3L2OrchRegionDesc desc{ + args.scalar(trailer + 2), args.scalar(trailer + 3), args.scalar(trailer + 4), + args.scalar(trailer + 5), args.scalar(trailer + 6), args.scalar(trailer + 7), + }; + L3L2QueueArgs queue_args{ + args.scalar(trailer + 8), args.scalar(trailer + 9), args.scalar(trailer + 10), + args.scalar(trailer + 11), args.scalar(trailer + 12), args.scalar(trailer + 13), + }; + queue_ = new (queue_storage_) HostGraphTokenQueue(desc, queue_args); + enabled_ = true; + } + + ~HostGraphTokenPublisher() { + if (queue_ != nullptr) queue_->~HostGraphTokenQueue(); + } + + HostGraphTokenPublisher(const HostGraphTokenPublisher &) = delete; + HostGraphTokenPublisher &operator=(const HostGraphTokenPublisher &) = delete; + + bool valid() const { return !enabled_ || (queue_ != nullptr && queue_->error().kind == L3L2QueueErrorKind::NONE); } + + bool publish(uint64_t epoch, bool final_epoch) { + if (!enabled_) return true; + if (!valid()) return false; + + L3L2QueueOutputReservation output{}; + if (!queue_->output().reserve(sizeof(PTO2HostGraphTokenPacket), kHostGraphTokenQueueTimeoutNs, output)) { + return false; + } + + PTO2HostGraphTokenPacket packet{ + PTO2_HOST_GRAPH_TOKEN_STREAM_MAGIC_VERSION, + request_id_, + epoch, + static_cast(epoch), + PTO2_HOST_GRAPH_TOKEN_SYNTHETIC | (final_epoch ? PTO2_HOST_GRAPH_TOKEN_FINAL : 0U), + 0, + }; + void *payload = reinterpret_cast(static_cast(output.payload.gm_addr)); + memcpy(payload, &packet, sizeof(packet)); + cache_flush_range(payload, sizeof(packet)); + return queue_->output().publish(output, L3L2QueueOpcode::DATA); + } + +private: + alignas(HostGraphTokenQueue) uint8_t queue_storage_[sizeof(HostGraphTokenQueue)]{}; + HostGraphTokenQueue *queue_{nullptr}; + uint64_t request_id_{0}; + bool enabled_{false}; +}; + +} // namespace + struct AicpuExecutor { int32_t sched_thread_num_; bool orch_to_sched_{false}; @@ -225,15 +297,36 @@ int32_t AicpuExecutor::run(Runtime *runtime) { int32_t run_rc = 0; LOG_INFO_V0("Thread %d: Start (exec_idx=%d)", thread_idx, affinity_exec_idx); - // Boot thread (thread N-1): host_build_graph host-orch boot. The - // orchestrator already ran on the host, which also relocated every - // cross-task pointer to its final device address before H2D — so the - // SM/arena this thread sees are already fully device-addressed. This thread - // attaches the prebuilt arena, points the SM handle's ring-header pointers - // at the device SM WITHOUT resetting the host-populated data, releases the - // scheduler threads, and hands the host-computed task count to the - // scheduler. It owns no AICore cores, so it does not dispatch. + // Boot thread (thread N-1): attach the HostGraph image without resetting + // its host-populated SM. Whole-graph runs consume one prebuilt image; + // streaming runs observe and release epoch images as the host publishes + // them. This thread owns no AICore cores and never dispatches tasks itself. if (thread_idx >= sched_thread_num_) { + void *sm_ptr = runtime->get_gm_sm_ptr(); + auto *pending_header = static_cast(sm_ptr); + const bool async_host_orch = runtime->host_total_tasks < 0; + if (async_host_orch) { + uint64_t publish_epoch = + pending_header->host_graph_epochs.host_publish_epoch.load(std::memory_order_acquire); + const uint64_t wait_start = get_sys_cnt_aicpu(); + while (publish_epoch == 0) { + if (pending_header->orch_error_code.load(std::memory_order_acquire) != PTO2_ERROR_NONE || + pending_header->host_graph_epochs.failed_epoch.load(std::memory_order_acquire) != 0) { + LOG_ERROR("Thread %d: HostGraph failed before final publication", thread_idx); + runtime_init_ready_.store(true, std::memory_order_release); + return -1; + } + if (get_sys_cnt_aicpu() - wait_start > SCHEDULER_TIMEOUT_CYCLES) { + LOG_ERROR("Thread %d: timed out waiting for final HostGraph publication", thread_idx); + pending_header->orch_error_code.store(PTO2_ERROR_FLOW_CONTROL_DEADLOCK, std::memory_order_release); + runtime_init_ready_.store(true, std::memory_order_release); + return -1; + } + SPIN_WAIT_HINT(); + publish_epoch = pending_header->host_graph_epochs.host_publish_epoch.load(std::memory_order_acquire); + } + } + void *prebuilt_arena = runtime->get_prebuilt_arena_base(); size_t off_runtime = runtime->get_prebuilt_runtime_offset(); if (prebuilt_arena == nullptr) { @@ -245,7 +338,6 @@ int32_t AicpuExecutor::run(Runtime *runtime) { rt = reinterpret_cast(static_cast(prebuilt_arena) + off_runtime); runtime_wire_arena_pointers(runtime_arena_, rt->prebuilt_layout, rt); - void *sm_ptr = runtime->get_gm_sm_ptr(); uint64_t sm_size = PTO2SharedMemoryHandle::calculate_size_per_ring(rt->prebuilt_layout.task_window_sizes); memset(rt->sm_handle, 0, sizeof(*rt->sm_handle)); if (!rt->sm_handle->attach_populated(sm_ptr, sm_size, rt->prebuilt_layout.task_window_sizes)) { @@ -261,25 +353,148 @@ int32_t AicpuExecutor::run(Runtime *runtime) { sched_ctx_.bind_runtime(rt); - // Latch the host-built task count (on_orchestration_done sets total_tasks_) - // BEFORE the runtime_init_ready_ release below — that store is the barrier - // that unblocks the scheduler threads. Otherwise they would acquire - // runtime_init_ready_ with total_tasks_=0 and race to an early exit before - // the host task count is visible (host-orch has no concurrent orchestrator - // to keep them alive). - // NOTE: do NOT call rt_orchestration_done(rt) here. The HOST already - // called it in run_host_orchestration; the orchestrator's own - // task-allocator pointers are intentionally NOT relocated (only the - // SM cross-task pointers and the host-built fanout adjacency — - // dep_pool / ready queues / fanout_head — were), so they still hold - // host addresses and mark_done()'s active_count() read would - // dereference host memory and fault the AICPU. on_orchestration_done - // only needs total_tasks and the scalar - // orchestrator.inline_completed_tasks, both already valid. - sched_ctx_.on_orchestration_done(runtime, rt, thread_idx, runtime->host_total_tasks); - - runtime_init_ready_.store(true, std::memory_order_release); - LOG_INFO_V0("Thread %d: host-orch boot complete (%d tasks)", thread_idx, runtime->host_total_tasks); + if (async_host_orch) { + PTO2SharedMemoryHeader *header = rt->sm_handle->header; + PTO2HostGraphEpochControl &control = header->host_graph_epochs; + uint64_t released_epoch = 0; + uint64_t completed_epoch = 0; + int32_t released_task_end = 0; + bool schedulers_started = false; + HostGraphTokenPublisher token_publisher(runtime); + if (!token_publisher.valid()) { + LOG_ERROR("Thread %d: invalid HostGraph token stream descriptor", thread_idx); + header->orch_error_code.store(PTO2_ERROR_INVALID_ARGS, std::memory_order_release); + runtime_init_ready_.store(true, std::memory_order_release); + return -1; + } + + auto mark_epoch_complete = [&]() { + if (released_epoch == 0 || sched_ctx_.completed_tasks_count() < released_task_end) return false; + if (completed_epoch >= released_epoch) return true; + size_t slot_index = + static_cast((released_epoch - 1) % static_cast(PTO2_HOST_GRAPH_EPOCH_SLOT_COUNT)); + PTO2HostGraphEpochSlot &completed_slot = control.slots[slot_index]; + if (!token_publisher.publish(released_epoch, completed_slot.range.final_epoch != 0)) { + header->sched_error_code.store(PTO2_ERROR_EXPLICIT_ORCH_FATAL, std::memory_order_release); + return false; + } + control.device_exec_done_epoch.store(released_epoch, std::memory_order_release); + control.device_buffer_free_epoch.store(released_epoch, std::memory_order_release); + completed_epoch = released_epoch; + return true; + }; + + while (true) { + uint64_t publish_epoch = control.host_publish_epoch.load(std::memory_order_acquire); + uint64_t wait_start = get_sys_cnt_aicpu(); + while (publish_epoch <= released_epoch) { + (void)mark_epoch_complete(); + if (header->orch_error_code.load(std::memory_order_acquire) != PTO2_ERROR_NONE || + header->sched_error_code.load(std::memory_order_acquire) != PTO2_ERROR_NONE || + control.failed_epoch.load(std::memory_order_acquire) != 0) { + LOG_ERROR( + "Thread %d: HostGraph failed while waiting for epoch=%" PRIu64, thread_idx, + released_epoch + 1 + ); + runtime_init_ready_.store(true, std::memory_order_release); + return -1; + } + if (get_sys_cnt_aicpu() - wait_start > SCHEDULER_TIMEOUT_CYCLES) { + LOG_ERROR( + "Thread %d: timed out waiting for HostGraph epoch=%" PRIu64, thread_idx, released_epoch + 1 + ); + header->orch_error_code.store(PTO2_ERROR_FLOW_CONTROL_DEADLOCK, std::memory_order_release); + runtime_init_ready_.store(true, std::memory_order_release); + return -1; + } + SPIN_WAIT_HINT(); + publish_epoch = control.host_publish_epoch.load(std::memory_order_acquire); + } + if (publish_epoch != released_epoch + 1) { + LOG_ERROR( + "Thread %d: HostGraph skipped publication epoch (released=%" PRIu64 " published=%" PRIu64 ")", + thread_idx, released_epoch, publish_epoch + ); + header->orch_error_code.store(PTO2_ERROR_INVALID_ARGS, std::memory_order_release); + runtime_init_ready_.store(true, std::memory_order_release); + return -1; + } + + size_t slot_index = + static_cast((publish_epoch - 1) % static_cast(PTO2_HOST_GRAPH_EPOCH_SLOT_COUNT)); + PTO2HostGraphEpochSlot &slot = control.slots[slot_index]; + if (slot.owner_epoch.load(std::memory_order_acquire) != publish_epoch || + slot.range.task_begin != released_task_end || slot.range.task_end <= slot.range.task_begin) { + LOG_ERROR( + "Thread %d: invalid HostGraph epoch=%" PRIu64 " range=[%d,%d) expected_begin=%d", thread_idx, + publish_epoch, slot.range.task_begin, slot.range.task_end, released_task_end + ); + header->orch_error_code.store(PTO2_ERROR_INVALID_ARGS, std::memory_order_release); + runtime_init_ready_.store(true, std::memory_order_release); + return -1; + } + + sched_ctx_.on_host_graph_published( + runtime, rt, thread_idx, slot.range.task_end, slot.range.inline_completed, + slot.range.final_epoch != 0 + ); + control.device_release_epoch.store(publish_epoch, std::memory_order_release); + released_epoch = publish_epoch; + released_task_end = slot.range.task_end; + LOG_INFO_V9( + "Thread %d: observed HostGraph epoch=%" PRIu64 " tasks=[%d,%d) inline=%d final=%d", thread_idx, + publish_epoch, slot.range.task_begin, slot.range.task_end, slot.range.inline_completed, + slot.range.final_epoch + ); + + if (!schedulers_started) { + runtime_init_ready_.store(true, std::memory_order_release); + schedulers_started = true; + } + + // Classify the newly published range against the monotonic + // completion flags. Cross-epoch consumers wait on the first + // incomplete producer and are reclassified when it completes. + for (int32_t task_id = slot.range.task_begin; task_id < slot.range.task_end; ++task_id) { + PTO2TaskSlotState &task_slot = header->ring.get_slot_state_by_task_id(task_id); + if (task_slot.task_state.load(std::memory_order_acquire) >= PTO2_TASK_COMPLETED) continue; + int32_t fanin_state = rt->scheduler.classify_fanin_state(&task_slot); + if (fanin_state < 0) { + rt->scheduler.push_ready_routed(&task_slot); + } else { + rt->scheduler.register_wake( + &header->ring.get_slot_state_by_task_id(task_slot.payload->fanin_local_ids[fanin_state]), + &task_slot + ); + } + } + if (slot.range.final_epoch == 0) continue; + + uint64_t complete_wait_start = get_sys_cnt_aicpu(); + while (!mark_epoch_complete()) { + if (header->sched_error_code.load(std::memory_order_acquire) != PTO2_ERROR_NONE) { + LOG_ERROR("Thread %d: failed to publish final HostGraph token", thread_idx); + return -1; + } + if (get_sys_cnt_aicpu() - complete_wait_start > SCHEDULER_TIMEOUT_CYCLES) { + LOG_ERROR("Thread %d: timed out waiting for final HostGraph completion", thread_idx); + header->sched_error_code.store(PTO2_ERROR_SCHEDULER_TIMEOUT, std::memory_order_release); + return -1; + } + SPIN_WAIT_HINT(); + } + break; + } + LOG_INFO_V9( + "Thread %d: streaming host-orch boot complete (%d tasks, epochs=%" PRIu64 ")", thread_idx, + released_task_end, released_epoch + ); + } else { + int32_t total_tasks = runtime->host_total_tasks; + sched_ctx_.on_orchestration_done(runtime, rt, thread_idx, total_tasks); + runtime_init_ready_.store(true, std::memory_order_release); + LOG_INFO_V9("Thread %d: host-orch boot complete (%d tasks)", thread_idx, total_tasks); + } } // Scheduler thread (orchestrator threads skip dispatch when orch_to_sched_ is false) diff --git a/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp b/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp index 59f791686a..bd58f07852 100644 --- a/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp +++ b/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp @@ -35,16 +35,24 @@ #include #include +#include #include #include #include #include #include +#include #include #include +#include #include +#include +#include +#include #include +#include #include +#include #include #include "../common/pto_runtime_status.h" @@ -56,12 +64,32 @@ #include "../runtime/runtime.h" #include "../../../../common/runtime_status/error_log.h" #include "../../../../common/task_interface/call_config.h" +#include "../../../../common/worker/pto_runtime_c_api.h" #include "callable.h" #include "common/platform_config.h" +#include "common/strace.h" #include "common/unified_log.h" +#include "host/raii_scope_guard.h" #include "utils/device_arena.h" #include "prepare_callable_common.h" +extern "C" const PipelineContract *get_pipeline_contract(void) { + static const PipelineContract contract = { + PTO_PIPELINE_CONTRACT_ABI_VERSION, + 5, + 2, + 2, + { + {PTO_PIPELINE_GM_HEAP, PTO_PIPELINE_FILL_MEM, 0}, + {PTO_PIPELINE_GM_SM, PTO_PIPELINE_FILL_MEM, 0}, + {PTO_PIPELINE_RUNTIME_IMAGE, PTO_PIPELINE_FILL_MEM, 0}, + {PTO_PIPELINE_AICPU_STREAM, PTO_PIPELINE_EXEC_HANDLE, 0}, + {PTO_PIPELINE_AICORE_STREAM, PTO_PIPELINE_EXEC_HANDLE, 0}, + }, + }; + return &contract; +} + // RuntimeEnv (call_config.h) is the cross-runtime ABI for per-ring config and // carries RUNTIME_ENV_RING_COUNT slots, shared with tensormap_and_ringbuffer. // host_build_graph is single-ring (PTO2_MAX_RING_DEPTH == 1) and reads only the @@ -419,28 +447,246 @@ static bool relocate_host_orch_image( return ok; } +struct HostGraphEpochCapture { + int32_t task_begin{0}; + int64_t build_start_ns{0}; + std::vector ranges; + bool (*boundary_handler)(PTO2Runtime *, PTO2HostGraphEpochRange *, size_t, void *){nullptr}; + void *boundary_context{nullptr}; +}; + +static bool capture_host_graph_epoch(PTO2Runtime *rt, bool final_epoch, void *opaque) { + auto *capture = static_cast(opaque); + if (capture == nullptr || rt == nullptr || rt->orchestrator.sm_header == nullptr) { + LOG_ERROR("host-orch: graph boundary has no active capture context"); + return false; + } + + int32_t task_end = rt->orchestrator.sm_header->ring.fc.current_task_index.load(std::memory_order_acquire); + if (task_end <= capture->task_begin) { + LOG_ERROR("host-orch: empty or reversed graph range [%d,%d)", capture->task_begin, task_end); + return false; + } + + PTO2HostGraphEpochRange range; + range.task_begin = capture->task_begin; + range.task_end = task_end; + range.final_epoch = final_epoch ? 1 : 0; + size_t epoch_index = capture->ranges.size(); + auto boundary_time = std::chrono::steady_clock::now(); + int64_t boundary_ns = + std::chrono::duration_cast(boundary_time.time_since_epoch()).count(); + if (capture->build_start_ns > 0 && boundary_ns >= capture->build_start_ns) { + char attrs[160]; + std::snprintf( + attrs, sizeof(attrs), "layer=%zu tasks=%d task_begin=%d task_end=%d final=%d", epoch_index, + range.task_end - range.task_begin, range.task_begin, range.task_end, range.final_epoch + ); + STRACE_HOST_SPAN_AT( + "simpler_run.host_orch.layer_build", capture->build_start_ns, boundary_ns - capture->build_start_ns, 1, + attrs + ); + } + if (capture->boundary_handler != nullptr && + !capture->boundary_handler(rt, &range, epoch_index, capture->boundary_context)) { + return false; + } + capture->ranges.push_back(range); + capture->task_begin = task_end; + auto next_build_time = std::chrono::steady_clock::now(); + capture->build_start_ns = + std::chrono::duration_cast(next_build_time.time_since_epoch()).count(); + return true; +} + +static bool +materialize_streaming_host_graph_range(PTO2Runtime *source, PTO2Runtime *target, PTO2HostGraphEpochRange *range) { + if (source == nullptr || target == nullptr || range == nullptr || source->orchestrator.sm_header == nullptr || + target->orchestrator.sm_header == nullptr) { + return false; + } + + PTO2SharedMemoryRingHeader &source_ring = source->orchestrator.sm_header->ring; + PTO2SharedMemoryRingHeader &target_ring = target->orchestrator.sm_header->ring; + int32_t source_total = source_ring.fc.current_task_index.load(std::memory_order_acquire); + if (range->task_begin < 0 || range->task_end <= range->task_begin || range->task_end > source_total || + static_cast(range->task_end) >= target_ring.task_window_size || + target->orchestrator.ring.task_allocator.task_head() != range->task_begin) { + LOG_ERROR( + "host-orch: invalid streaming range [%d,%d), source_total=%d target_head=%d window=%" PRIu64, + range->task_begin, range->task_end, source_total, target->orchestrator.ring.task_allocator.task_head(), + target_ring.task_window_size + ); + return false; + } + + range->inline_completed = 0; + + for (int32_t task_id = range->task_begin; task_id < range->task_end; ++task_id) { + PTO2TaskSlotState &source_slot = source_ring.get_slot_state_by_task_id(task_id); + if (source_slot.task == nullptr || source_slot.payload == nullptr) { + LOG_ERROR("host-orch: source task %d has incomplete slot bindings", task_id); + return false; + } + + PTO2TaskAllocResult alloc = target->orchestrator.ring.task_allocator.alloc(0); + if (alloc.failed() || alloc.task_id != task_id) { + LOG_ERROR("host-orch: target task allocation diverged at task %d", task_id); + return false; + } + PTO2TaskDescriptor &target_task = target_ring.get_task_by_slot(alloc.slot); + PTO2TaskPayload &target_payload = target_ring.get_payload_by_slot(alloc.slot); + PTO2TaskSlotState &target_slot = target_ring.get_slot_state_by_slot(alloc.slot); + std::memcpy(&target_task, source_slot.task, sizeof(target_task)); + std::memcpy(&target_payload, source_slot.payload, sizeof(target_payload)); + + if (target_payload.fanin_count < 0 || target_payload.fanin_count > PTO2_MAX_FANIN) { + LOG_ERROR("host-orch: streaming task %d has invalid fanin count %d", task_id, target_payload.fanin_count); + return false; + } + for (int32_t i = 0; i < target_payload.fanin_count; ++i) { + int32_t producer_id = target_payload.fanin_local_ids[i]; + if (producer_id < 0 || producer_id >= task_id) { + LOG_ERROR("host-orch: task %d has invalid producer id %d", task_id, producer_id); + return false; + } + } + target_payload.dispatch_fanin.store(0, std::memory_order_relaxed); + + PTO2TaskState task_state = source_slot.task_state.load(std::memory_order_acquire); + target_slot.reset_for_reuse(); + target_slot.bind_buffers(&target_payload, &target_task); + target_slot.last_consumer_local_id = source_slot.last_consumer_local_id; + target_slot.task_state.store(task_state, std::memory_order_relaxed); + target_slot.active_mask = source_slot.active_mask; + target_slot.task_attrs = source_slot.task_attrs; + target_slot.task_attrs.set_early_resolve(false); + target_slot.total_required_subtasks = source_slot.total_required_subtasks; + target_slot.logical_block_num = source_slot.logical_block_num; + uint8_t completed = + source_ring.completion_flags[source_ring.get_slot_by_task_id(task_id)].load(std::memory_order_acquire); + target_ring.completion_flags[alloc.slot].store(completed, std::memory_order_relaxed); + if (completed != 0 || task_state >= PTO2_TASK_COMPLETED) range->inline_completed++; + } + + target->orchestrator.inline_completed_tasks += static_cast(range->inline_completed); + return true; +} + +static bool materialize_host_graph_ranges( + PTO2SharedMemoryHandle &source_sm_handle, PTO2SharedMemoryHandle &target_sm_handle, + const std::vector &ranges, int32_t total_tasks +) { + if (source_sm_handle.header == nullptr || target_sm_handle.header == nullptr) { + LOG_ERROR("host-orch: cannot materialize graph ranges without source and target SM headers"); + return false; + } + + PTO2SharedMemoryRingHeader &source_ring = source_sm_handle.header->ring; + PTO2SharedMemoryRingHeader &target_ring = target_sm_handle.header->ring; + if (total_tasks < 0 || static_cast(total_tasks) > source_ring.task_window_size || + source_ring.task_window_size != target_ring.task_window_size) { + LOG_ERROR( + "host-orch: invalid materialization size tasks=%d source_window=%" PRIu64 " target_window=%" PRIu64, + total_tasks, source_ring.task_window_size, target_ring.task_window_size + ); + return false; + } + + int32_t expected_begin = 0; + for (size_t index = 0; index < ranges.size(); ++index) { + const PTO2HostGraphEpochRange &range = ranges[index]; + if (range.task_begin != expected_begin || range.task_end <= range.task_begin || range.task_end > total_tasks) { + LOG_ERROR( + "host-orch: epoch %zu does not form a contiguous graph partition: expected_begin=%d range=[%d,%d) " + "total=%d", + index, expected_begin, range.task_begin, range.task_end, total_tasks + ); + return false; + } + + char attrs[160]; + std::snprintf( + attrs, sizeof(attrs), "epoch=%zu slot=%zu task_begin=%d task_end=%d tasks=%d", index, + index % static_cast(PTO2_HOST_GRAPH_EPOCH_SLOT_COUNT), range.task_begin, range.task_end, + range.task_end - range.task_begin + ); + { + STRACE_A("simpler_run.host_orch.epoch.materialize", attrs); + for (int32_t task_id = range.task_begin; task_id < range.task_end; ++task_id) { + int32_t source_slot = source_ring.get_slot_by_task_id(task_id); + int32_t target_slot = target_ring.get_slot_by_task_id(task_id); + std::memcpy( + &target_ring.task_descriptors[target_slot], &source_ring.task_descriptors[source_slot], + sizeof(PTO2TaskDescriptor) + ); + std::memcpy( + &target_ring.task_payloads[target_slot], &source_ring.task_payloads[source_slot], + sizeof(PTO2TaskPayload) + ); + std::memcpy( + &target_ring.slot_states[target_slot], &source_ring.slot_states[source_slot], + sizeof(PTO2TaskSlotState) + ); + target_ring.completion_flags[target_slot].store( + source_ring.completion_flags[source_slot].load(std::memory_order_acquire), std::memory_order_relaxed + ); + } + } + expected_begin = range.task_end; + } + + if (ranges.empty() || expected_begin != total_tasks) { + LOG_ERROR("host-orch: captured graph ranges cover [0,%d), expected [0,%d)", expected_begin, total_tasks); + return false; + } + return true; +} + int32_t run_host_orchestration( Runtime *runtime, const HostApi *api, PTO2Runtime *rt, DeviceArena &host_arena, const PTO2RuntimeArenaLayout &layout, void *device_sm, uint64_t sm_size, void *device_arena, void *gm_heap, const uint64_t eff_heap_sizes[PTO2_MAX_RING_DEPTH], const uint64_t eff_task_window_sizes[PTO2_MAX_RING_DEPTH], - void *host_orch_func_ptr, const L2TaskArgs &orch_l2 + void *host_orch_func_ptr, const L2TaskArgs &orch_l2, + bool (*boundary_handler)(PTO2Runtime *, PTO2HostGraphEpochRange *, size_t, void *) = nullptr, + void *boundary_context = nullptr ) { - std::vector host_sm_buf(sm_size, 0); - void *host_sm = host_sm_buf.data(); + constexpr char epoch_attrs[] = "epoch=0 slot=0 final=1"; + std::vector source_sm_buf(sm_size, 0); + void *source_sm = source_sm_buf.data(); + + DeviceArena source_arena; + PTO2RuntimeArenaLayout source_layout = runtime_reserve_layout(source_arena, eff_task_window_sizes, eff_heap_sizes); + if (source_layout.arena_size != layout.arena_size || source_layout.off_runtime != layout.off_runtime || + source_arena.commit(DeviceArena::kDefaultBaseAlign) == nullptr) { + LOG_ERROR("host-orch: failed to create an equivalent source arena for epoch capture"); + return -1; + } + + PTO2Runtime *source_rt = runtime_init_data_from_layout( + source_arena, source_layout, PTO2_MODE_EXECUTE, device_sm, sm_size, gm_heap, eff_heap_sizes + ); + if (source_rt == nullptr) { + LOG_ERROR("host-orch: source runtime init failed"); + return -1; + } + runtime_wire_arena_pointers(source_arena, source_layout, source_rt); // Re-point the orchestrator half at the host SM (scheduler keeps device SM). // init_data_from_layout resets the orchestrator state, so this is safe. - if (!rt->orchestrator.init_data_from_layout( - layout.orch, host_arena, host_sm, gm_heap, eff_heap_sizes[0], eff_task_window_sizes[0] + if (!source_rt->orchestrator.init_data_from_layout( + source_layout.orch, source_arena, source_sm, gm_heap, eff_heap_sizes[0], eff_task_window_sizes[0] )) { LOG_ERROR("host-orch: orchestrator re-init against host SM failed"); return -1; } - rt->orchestrator.wire_arena_pointers(layout.orch, host_arena, &rt->scheduler); + source_rt->orchestrator.wire_arena_pointers(source_layout.orch, source_arena, &source_rt->scheduler); + source_rt->scheduler.sm_header = source_rt->orchestrator.sm_header; + source_rt->scheduler.ring_sched_state.ring = &source_rt->orchestrator.sm_header->ring; // Initialize the host SM header (ring flow control) so submit_task can run. - PTO2SharedMemoryHandle host_sm_handle; - if (!host_sm_handle.init_per_ring(host_sm, sm_size, eff_task_window_sizes, eff_heap_sizes)) { + PTO2SharedMemoryHandle source_sm_handle; + if (!source_sm_handle.init_per_ring(source_sm, sm_size, eff_task_window_sizes, eff_heap_sizes)) { LOG_ERROR("host-orch: host SM init_per_ring failed"); return -1; } @@ -449,8 +695,8 @@ int32_t run_host_orchestration( // re-applied with the real device values on the AICPU at boot; the values // here only feed cluster spreading during this host submit and are unused // by the migrated non-cluster examples. - runtime_finalize_after_wire(rt, /*aic*/ 24, /*aiv*/ 48); - rt->mode = PTO2_MODE_EXECUTE; + runtime_finalize_after_wire(source_rt, /*aic*/ 24, /*aiv*/ 48); + source_rt->mode = PTO2_MODE_EXECUTE; // get_tensor_data/set_tensor_data dereference buffer.addr directly: the // input tensors were mapped into host address space at staging time // (HostApi::register_device_memory_to_host), so the host orchestrator can @@ -461,20 +707,102 @@ int32_t run_host_orchestration( // rt_scope_* / rt_orchestration_done) and the orch .so's own copy (used by // its inline rt_submit_* -> current_runtime()). const HostOrchEntryPoints *eps = reinterpret_cast(host_orch_func_ptr); - framework_bind_runtime(rt); + framework_bind_runtime(source_rt); if (eps->bind != nullptr) { - eps->bind(rt); + eps->bind(source_rt); } else { LOG_ERROR("host-orch: orch .so framework_bind_runtime was not resolved"); return -1; } - rt_scope_begin(rt); - eps->entry(orch_l2); - rt_scope_end(rt); - rt_orchestration_done(rt); + HostGraphEpochCapture epoch_capture; + epoch_capture.boundary_handler = boundary_handler; + epoch_capture.boundary_context = boundary_context; + { + runtime_set_graph_boundary_callback(source_rt, capture_host_graph_epoch, &epoch_capture); + auto capture_guard = RAIIScopeGuard([source_rt]() { + runtime_set_graph_boundary_callback(source_rt, nullptr, nullptr); + }); + { + STRACE_A("simpler_run.host_orch.epoch.build", epoch_attrs); + rt_scope_begin(source_rt); + auto build_start = std::chrono::steady_clock::now(); + epoch_capture.build_start_ns = + std::chrono::duration_cast(build_start.time_since_epoch()).count(); + eps->entry(orch_l2); + rt_scope_end(source_rt); + // An orchestration without explicit boundaries is one final epoch. + if (epoch_capture.ranges.empty()) rt_graph_boundary(source_rt, true); + rt_orchestration_done(source_rt); + } + } + + int32_t total_tasks = pto2_sm_layout::ring_current_task_index_addr(source_sm)->load(std::memory_order_acquire); + uint64_t captured_epochs = epoch_capture.ranges.size(); + char capture_attrs[96]; + std::snprintf( + capture_attrs, sizeof(capture_attrs), "captured_epochs=%" PRIu64 " tasks=%d", captured_epochs, total_tasks + ); + { STRACE_A("simpler_run.host_orch.boundary_capture", capture_attrs); } + for (size_t index = 0; index < epoch_capture.ranges.size(); ++index) { + const PTO2HostGraphEpochRange &range = epoch_capture.ranges[index]; + char range_attrs[160]; + std::snprintf( + range_attrs, sizeof(range_attrs), "epoch=%zu slot=%zu task_begin=%d task_end=%d tasks=%d", index, + index % static_cast(PTO2_HOST_GRAPH_EPOCH_SLOT_COUNT), range.task_begin, range.task_end, + range.task_end - range.task_begin + ); + { STRACE_A("simpler_run.host_orch.epoch.capture", range_attrs); } + } + LOG_INFO_V0("host-orch: captured epochs=%" PRIu64 " tasks=%d", captured_epochs, total_tasks); + + if (boundary_handler != nullptr) { + if (epoch_capture.ranges.empty() || epoch_capture.ranges.back().final_epoch == 0 || + epoch_capture.ranges.back().task_end != total_tasks) { + LOG_ERROR("host-orch: streaming orchestration did not publish one final range covering all tasks"); + return -1; + } + return total_tasks; + } + + std::memcpy(host_arena.base(), source_arena.base(), layout.arena_size); + runtime_wire_arena_pointers(host_arena, layout, rt); + + std::vector target_sm_buf(sm_size, 0); + void *target_sm = target_sm_buf.data(); + std::memcpy(target_sm, source_sm, sizeof(PTO2SharedMemoryHeader)); + PTO2SharedMemoryHandle target_sm_handle; + if (!target_sm_handle.attach_populated(target_sm, sm_size, eff_task_window_sizes)) { + LOG_ERROR("host-orch: target SM attach failed during epoch materialization"); + return -1; + } + if (!materialize_host_graph_ranges(source_sm_handle, target_sm_handle, epoch_capture.ranges, total_tasks)) { + return -1; + } - int32_t total_tasks = pto2_sm_layout::ring_current_task_index_addr(host_sm)->load(std::memory_order_acquire); + PTO2HostGraphEpochControl &epoch_control = target_sm_handle.header->host_graph_epochs; + PTO2HostGraphEpochSlot &published_slot = epoch_control.slots[0]; + published_slot.owner_epoch.store(1, std::memory_order_relaxed); + published_slot.range.task_begin = 0; + published_slot.range.task_end = total_tasks; + published_slot.range.final_epoch = 1; + epoch_control.host_publish_epoch.store(0, std::memory_order_relaxed); + + // The source arena is a byte-for-byte execution-state image. Rewiring above + // makes all arena-private bases point at the target before relocation walks + // them. The scheduler's SM anchors were device addresses before the copy and + // remain so; only cross-task pointers move through source -> target -> device. + const int64_t source_to_target_sm_delta = static_cast(reinterpret_cast(target_sm)) - + static_cast(reinterpret_cast(source_sm)); + const int64_t source_to_target_arena_delta = static_cast(reinterpret_cast(host_arena.base())) - + static_cast(reinterpret_cast(source_arena.base())); + if (!relocate_host_orch_image( + target_sm_handle, rt, reinterpret_cast(source_sm), sm_size, source_to_target_sm_delta, + reinterpret_cast(source_arena.base()), layout.arena_size, source_to_target_arena_delta + )) { + LOG_ERROR("host-orch: source-to-target epoch relocation failed"); + return -1; + } // Relocate the host-DDR cross-task pointers to their final DEVICE addresses // on the host, before the SM and arena leave for the device. Pointers into @@ -482,26 +810,582 @@ int32_t run_host_orchestration( // queue) shift by arena_delta. After this both the SM and arena carry device // addresses, so the device boots scheduler-only. const int64_t sm_delta = static_cast(reinterpret_cast(device_sm)) - - static_cast(reinterpret_cast(host_sm)); + static_cast(reinterpret_cast(target_sm)); const int64_t arena_delta = static_cast(reinterpret_cast(device_arena)) - static_cast(reinterpret_cast(host_arena.base())); - if (!relocate_host_orch_image( - host_sm_handle, rt, reinterpret_cast(host_sm), sm_size, sm_delta, - reinterpret_cast(host_arena.base()), layout.arena_size, arena_delta - )) { - LOG_ERROR("host-orch: relocation failed; refusing to H2D an image with unrelocated host pointers"); - return -1; + { + STRACE_A("simpler_run.host_orch.image.relocate", epoch_attrs); + if (!relocate_host_orch_image( + target_sm_handle, rt, reinterpret_cast(target_sm), sm_size, sm_delta, + reinterpret_cast(host_arena.base()), layout.arena_size, arena_delta + )) { + LOG_ERROR("host-orch: relocation failed; refusing to H2D an image with unrelocated host pointers"); + return -1; + } } - if (api->copy_to_device(device_sm, host_sm, sm_size) != 0) { - LOG_ERROR("host-orch: H2D of populated SM failed"); - return -1; + { + STRACE_A("simpler_run.host_orch.epoch.stage_sm", epoch_attrs); + if (api->copy_to_device(device_sm, target_sm, sm_size) != 0) { + LOG_ERROR("host-orch: H2D of populated SM failed"); + return -1; + } } return total_tasks; } +std::shared_ptr stage1_gate_for_runner(void *runner_context) { + static std::mutex registry_mutex; + static std::unordered_map> gates; + + std::lock_guard lock(registry_mutex); + auto gate = gates[runner_context].lock(); + if (gate == nullptr) { + gate = std::make_shared(); + gates[runner_context] = gate; + } + return gate; +} + +class HostGraphAsyncJob { +public: + ~HostGraphAsyncJob() { (void)join(); } + + bool prepare_and_start( + Runtime *runtime, const HostApi *api, const PTO2RuntimeArenaLayout &layout, void *device_sm, uint64_t sm_size, + void *device_arena, void *gm_heap, const uint64_t heap_sizes[PTO2_MAX_RING_DEPTH], + const uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH], void *host_orch_func_ptr, + const ChipStorageTaskArgs &device_args + ) { + runtime_ = runtime; + api_ = api; + layout_ = layout; + device_sm_ = device_sm; + sm_size_ = sm_size; + device_arena_ = device_arena; + gm_heap_ = gm_heap; + host_orch_func_ptr_ = host_orch_func_ptr; + device_args_ = device_args; + std::memcpy(heap_sizes_, heap_sizes, sizeof(heap_sizes_)); + std::memcpy(task_window_sizes_, task_window_sizes, sizeof(task_window_sizes_)); +#if SIMPLER_HOST_STRACE + strace_inv_ = simpler::strace::StraceScope::current_inv(); + strace_hid_ = simpler::strace::StraceScope::current_hid(); +#endif + + if (runtime_ == nullptr || api_ == nullptr || device_sm_ == nullptr || device_arena_ == nullptr || + gm_heap_ == nullptr || host_orch_func_ptr_ == nullptr || api_->capture_thread_context == nullptr || + api_->bind_thread_context == nullptr || api_->unbind_thread_context == nullptr || + api_->store_u64_release_to_device == nullptr || api_->load_u64_acquire_from_device == nullptr) { + LOG_ERROR("host-orch: async job has incomplete inputs"); + return false; + } + + PTO2RuntimeArenaLayout target_layout = runtime_reserve_layout(target_arena_, task_window_sizes_, heap_sizes_); + if (target_layout.arena_size != layout_.arena_size || target_layout.off_runtime != layout_.off_runtime || + target_arena_.commit(DeviceArena::kDefaultBaseAlign) == nullptr) { + LOG_ERROR("host-orch: async target arena layout mismatch or allocation failure"); + return false; + } + target_rt_ = runtime_init_data_from_layout( + target_arena_, layout_, PTO2_MODE_EXECUTE, device_sm_, sm_size_, gm_heap_, heap_sizes_ + ); + if (target_rt_ == nullptr) { + LOG_ERROR("host-orch: async target runtime init failed"); + return false; + } + runtime_wire_arena_pointers(target_arena_, layout_, target_rt_); + target_rt_->prebuilt_layout = layout_; + + target_sm_.resize(sm_size_, 0); + if (!target_sm_handle_.init_per_ring(target_sm_.data(), sm_size_, task_window_sizes_, heap_sizes_) || + !target_rt_->orchestrator.init_data_from_layout( + layout_.orch, target_arena_, target_sm_.data(), gm_heap_, heap_sizes_[0], task_window_sizes_[0] + )) { + LOG_ERROR("host-orch: failed to initialize the streaming target image"); + return false; + } + target_rt_->orchestrator.wire_arena_pointers(layout_.orch, target_arena_, &target_rt_->scheduler); + + // Publish the empty runtime arena once. Device-side wire restores every + // arena-private pointer; the orchestrator's host SM anchor is unused by + // scheduler-only execution and is temporarily replaced for this image. + PTO2SharedMemoryHeader *target_host_header = target_rt_->orchestrator.sm_header; + target_rt_->orchestrator.sm_header = static_cast(device_sm_); + int arena_rc = api_->copy_to_device(device_arena_, target_arena_.base(), layout_.arena_size); + target_rt_->orchestrator.sm_header = target_host_header; + if (arena_rc != 0 || api_->copy_to_device(device_sm_, target_sm_.data(), sm_size_) != 0) { + LOG_ERROR("host-orch: failed to reset the async publication SM"); + return false; + } + + runtime_->set_prebuilt_arena(device_arena_, layout_.off_runtime); + runtime_->host_total_tasks = -1; + thread_context_ = api_->capture_thread_context(); + if (thread_context_ == nullptr) { + LOG_ERROR("host-orch: failed to capture the runner thread context"); + return false; + } + stage1_gate_ = stage1_gate_for_runner(thread_context_); + + try { + worker_ = std::thread([this]() { +#if SIMPLER_HOST_STRACE + STRACE_SET_CONTEXT(strace_inv_, strace_hid_); +#endif + if (api_->bind_thread_context(thread_context_) != 0) { + result_ = -1; + worker_start_state_.store(-1, std::memory_order_release); + return; + } + worker_start_state_.store(1, std::memory_order_release); + { + std::unique_lock lock(run_gate_mutex_); + run_gate_cv_.wait(lock, [this]() { + return run_gate_state_.load(std::memory_order_acquire) != 0; + }); + } + if (run_gate_state_.load(std::memory_order_acquire) < 0) { + result_ = -1; + api_->unbind_thread_context(); + return; + } + // Host orchestration uses one control flow per runner. A later + // request may prepare while the prior request executes on the + // device, but two Host Stage1 sections must never mutate the + // orchestration runtime at the same time. + std::unique_lock stage1_lock(*stage1_gate_); + try { + run(); + } catch (const std::exception &e) { + LOG_ERROR("host-orch: async worker exception: %s", e.what()); + publish_failure(PTO2_ERROR_EXPLICIT_ORCH_FATAL); + result_ = -1; + } catch (...) { + LOG_ERROR("host-orch: async worker unknown exception"); + publish_failure(PTO2_ERROR_EXPLICIT_ORCH_FATAL); + result_ = -1; + } + api_->unbind_thread_context(); + }); + } catch (const std::exception &e) { + LOG_ERROR("host-orch: failed to start async worker: %s", e.what()); + return false; + } + + while (worker_start_state_.load(std::memory_order_acquire) == 0) + std::this_thread::yield(); + if (worker_start_state_.load(std::memory_order_acquire) < 0) { + worker_.join(); + return false; + } + return true; + } + + bool open_run_gate() { + int32_t expected = 0; + if (!run_gate_state_.compare_exchange_strong( + expected, 1, std::memory_order_release, std::memory_order_acquire + )) { + return expected == 1; + } + run_gate_cv_.notify_one(); + LOG_INFO_V9("host-orch: async worker run gate opened"); + return true; + } + + bool release_for_run() { + if (!open_run_gate()) return false; + int32_t expected = 0; + if (!publish_gate_state_.compare_exchange_strong( + expected, 1, std::memory_order_release, std::memory_order_acquire + )) { + return expected == 1; + } + run_gate_cv_.notify_all(); + LOG_INFO_V9("host-orch: publication gate released after runner profiling setup"); + return true; + } + + int join() { + int32_t expected = 0; + if (run_gate_state_.compare_exchange_strong( + expected, -1, std::memory_order_release, std::memory_order_acquire + )) { + run_gate_cv_.notify_all(); + } + expected = 0; + (void)publish_gate_state_.compare_exchange_strong( + expected, -1, std::memory_order_release, std::memory_order_acquire + ); + run_gate_cv_.notify_all(); + if (worker_.joinable()) worker_.join(); + return result_; + } + +private: + bool wait_for_publish_gate() { + std::unique_lock lock(run_gate_mutex_); + run_gate_cv_.wait(lock, [this]() { + return publish_gate_state_.load(std::memory_order_acquire) != 0; + }); + return publish_gate_state_.load(std::memory_order_acquire) > 0; + } + + void publish_failure(int32_t error_code) const { + if (api_ == nullptr || device_sm_ == nullptr) return; + uint64_t failed_epoch = 1; + (void)api_->copy_to_device( + static_cast(device_sm_) + offsetof(PTO2SharedMemoryHeader, orch_error_code), &error_code, + sizeof(error_code) + ); + (void)api_->store_u64_release_to_device( + static_cast(device_sm_) + offsetof(PTO2SharedMemoryHeader, host_graph_epochs) + + offsetof(PTO2HostGraphEpochControl, failed_epoch), + failed_epoch + ); + } + + bool wait_for_device_epoch(size_t field_offset, uint64_t epoch, const char *kind) const { + char attrs[96]; + std::snprintf(attrs, sizeof(attrs), "kind=%s target_epoch=%" PRIu64, kind, epoch); + STRACE_A("simpler_run.host_orch.epoch.wait_device", attrs); + auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(30); + uint64_t observed = 0; + while (observed < epoch) { + if (api_->load_u64_acquire_from_device(&observed, static_cast(device_sm_) + field_offset) != + 0) { + LOG_ERROR("host-orch: failed to read Device %s epoch", kind); + return false; + } + if (observed >= epoch) return true; + if (std::chrono::steady_clock::now() >= deadline) { + LOG_ERROR( + "host-orch: timed out waiting for Device %s epoch=%" PRIu64 " (observed=%" PRIu64 ")", kind, epoch, + observed + ); + return false; + } + std::this_thread::sleep_for(std::chrono::microseconds(50)); + } + return true; + } + + bool reclaim_target_through(uint64_t epoch) { + if (epoch <= target_reclaimed_epoch_) return true; + if (epoch > published_ranges_.size()) { + LOG_ERROR( + "host-orch: cannot reclaim unpublished target epoch=%" PRIu64 " (published=%zu)", epoch, + published_ranges_.size() + ); + return false; + } + + target_reclaimed_epoch_ = epoch; + LOG_INFO_V9("host-orch: epoch metadata slot is reusable through epoch=%" PRIu64, epoch); + return true; + } + + bool wait_and_reclaim_target(uint64_t epoch, const char *kind) { + if (epoch <= target_reclaimed_epoch_) return true; + size_t field_offset = offsetof(PTO2SharedMemoryHeader, host_graph_epochs) + + offsetof(PTO2HostGraphEpochControl, device_buffer_free_epoch); + return wait_for_device_epoch(field_offset, epoch, kind) && reclaim_target_through(epoch); + } + + bool ensure_target_capacity(const PTO2HostGraphEpochRange &range) { + PTO2TaskAllocator &allocator = target_rt_->orchestrator.ring.task_allocator; + int32_t task_count = range.task_end - range.task_begin; + if (task_count <= 0 || range.task_begin != allocator.task_head() || range.task_end >= allocator.window_size()) { + LOG_ERROR( + "host-orch: epoch cannot fit whole-graph target storage " + "(range=[%d,%d) target_head=%d window=%d)", + range.task_begin, range.task_end, allocator.task_head(), allocator.window_size() + ); + return false; + } + return true; + } + + bool stage_epoch(const PTO2HostGraphEpochRange &range, size_t epoch_index) { + const int32_t task_count = range.task_end - range.task_begin; + PTO2SharedMemoryRingHeader &ring = target_sm_handle_.header->ring; + if (task_count <= 0 || range.task_begin < 0 || static_cast(range.task_end) >= ring.task_window_size) { + LOG_ERROR( + "host-orch: invalid streaming task range size=%d window=%" PRIu64, task_count, ring.task_window_size + ); + return false; + } + + char attrs[192]; + std::snprintf( + attrs, sizeof(attrs), "epoch=%zu slot=%zu task_begin=%d task_end=%d tasks=%d final=%d", epoch_index + 1, + epoch_index % static_cast(PTO2_HOST_GRAPH_EPOCH_SLOT_COUNT), range.task_begin, range.task_end, + task_count, range.final_epoch + ); + STRACE_A("simpler_run.host_orch.epoch.stage_upload", attrs); + + auto relocate_target_pointer = [&](auto *&pointer) -> bool { + using Pointer = std::remove_reference_t; + uint64_t value = reinterpret_cast(pointer); + if (value == 0) return true; + uint64_t host_sm = reinterpret_cast(target_sm_.data()); + uint64_t host_arena = reinterpret_cast(target_arena_.base()); + if (value >= host_sm && value < host_sm + sm_size_) { + pointer = reinterpret_cast( + reinterpret_cast(device_sm_) + static_cast(value - host_sm) + ); + return true; + } + if (value >= host_arena && value < host_arena + layout_.arena_size) { + pointer = reinterpret_cast( + reinterpret_cast(device_arena_) + static_cast(value - host_arena) + ); + return true; + } + LOG_ERROR("host-orch: streaming pointer %#lx is outside target SM/arena", value); + return false; + }; + + auto sm_offsets = pto2_sm_layout::ring_segment_offsets(ring.task_window_size); + char *device_sm_bytes = static_cast(device_sm_); + auto stage_task_run = [&](int32_t first_slot, int32_t run_count) { + std::vector payload_image(static_cast(run_count) * sizeof(PTO2TaskPayload)); + std::vector slot_image(static_cast(run_count) * sizeof(PTO2TaskSlotState)); + std::vector completion_image(static_cast(run_count)); + for (int32_t i = 0; i < run_count; ++i) { + PTO2TaskPayload payload_copy; + PTO2TaskSlotState slot_copy; + std::memcpy(&payload_copy, &ring.task_payloads[first_slot + i], sizeof(payload_copy)); + std::memcpy(&slot_copy, &ring.slot_states[first_slot + i], sizeof(slot_copy)); + if (payload_copy.fanin_count < 0 || payload_copy.fanin_count > PTO2_MAX_FANIN) return false; + if (!relocate_target_pointer(slot_copy.task) || !relocate_target_pointer(slot_copy.payload)) { + return false; + } + std::memcpy( + payload_image.data() + static_cast(i) * sizeof(payload_copy), &payload_copy, + sizeof(payload_copy) + ); + std::memcpy( + slot_image.data() + static_cast(i) * sizeof(slot_copy), &slot_copy, sizeof(slot_copy) + ); + completion_image[static_cast(i)] = + ring.completion_flags[first_slot + i].load(std::memory_order_acquire); + } + + size_t descriptor_bytes = static_cast(run_count) * sizeof(PTO2TaskDescriptor); + return api_->copy_to_device( + device_sm_bytes + sm_offsets.descriptors + + static_cast(first_slot) * sizeof(PTO2TaskDescriptor), + &ring.task_descriptors[first_slot], descriptor_bytes + ) == 0 && + api_->copy_to_device( + device_sm_bytes + sm_offsets.payloads + + static_cast(first_slot) * sizeof(PTO2TaskPayload), + payload_image.data(), payload_image.size() + ) == 0 && + api_->copy_to_device( + device_sm_bytes + sm_offsets.slot_states + + static_cast(first_slot) * sizeof(PTO2TaskSlotState), + slot_image.data(), slot_image.size() + ) == 0 && + api_->copy_to_device( + device_sm_bytes + sm_offsets.completion_flags + static_cast(first_slot), + completion_image.data(), completion_image.size() + ) == 0; + }; + int32_t staged_tasks = 0; + while (staged_tasks < task_count) { + int32_t task_id = range.task_begin + staged_tasks; + int32_t first_slot = ring.get_slot_by_task_id(task_id); + int32_t run_count = + std::min(task_count - staged_tasks, static_cast(ring.task_window_size) - first_slot); + if (!stage_task_run(first_slot, run_count)) { + LOG_ERROR("host-orch: failed to stage task range [%d,%d)", range.task_begin, range.task_end); + return false; + } + staged_tasks += run_count; + } + + PTO2HostGraphEpochControl &control = target_sm_handle_.header->host_graph_epochs; + size_t control_slot = epoch_index % static_cast(PTO2_HOST_GRAPH_EPOCH_SLOT_COUNT); + PTO2HostGraphEpochSlot &epoch_slot = control.slots[control_slot]; + epoch_slot.owner_epoch.store(epoch_index + 1, std::memory_order_relaxed); + epoch_slot.range = range; + + int32_t current_task_index = range.task_end; + size_t epoch_slot_offset = offsetof(PTO2SharedMemoryHeader, host_graph_epochs) + + offsetof(PTO2HostGraphEpochControl, slots) + + control_slot * sizeof(PTO2HostGraphEpochSlot); + if (api_->copy_to_device( + pto2_sm_layout::ring_current_task_index_addr(device_sm_), ¤t_task_index, + sizeof(current_task_index) + ) != 0 || + api_->copy_to_device(device_sm_bytes + epoch_slot_offset, &epoch_slot, sizeof(epoch_slot)) != 0) { + LOG_ERROR("host-orch: failed to stage epoch metadata"); + return false; + } + return true; + } + + bool commit_publication(const PTO2HostGraphEpochRange &range, size_t epoch_index) const { + uint64_t publish_epoch = epoch_index + 1; + size_t control_slot = epoch_index % static_cast(PTO2_HOST_GRAPH_EPOCH_SLOT_COUNT); + char attrs[160]; + std::snprintf( + attrs, sizeof(attrs), "epoch=%" PRIu64 " slot=%zu task_begin=%d task_end=%d final=%d", publish_epoch, + control_slot, range.task_begin, range.task_end, range.final_epoch + ); + STRACE_A("simpler_run.host_orch.epoch.commit", attrs); + size_t epoch_offset = offsetof(PTO2SharedMemoryHeader, host_graph_epochs) + + offsetof(PTO2HostGraphEpochControl, host_publish_epoch); + return api_->store_u64_release_to_device(static_cast(device_sm_) + epoch_offset, publish_epoch) == 0; + } + + bool publish_boundary(PTO2Runtime *source, PTO2HostGraphEpochRange *range, size_t epoch_index) { + uint64_t publish_epoch = epoch_index + 1; + if (epoch_index >= static_cast(PTO2_HOST_GRAPH_EPOCH_SLOT_COUNT) && + !wait_and_reclaim_target(publish_epoch - PTO2_HOST_GRAPH_EPOCH_SLOT_COUNT, "slot-free")) { + return false; + } + if (!ensure_target_capacity(*range)) return false; + + char attrs[160]; + std::snprintf( + attrs, sizeof(attrs), "epoch=%" PRIu64 " task_begin=%d task_end=%d final=%d", publish_epoch, + range->task_begin, range->task_end, range->final_epoch + ); + { + STRACE_A("simpler_run.host_orch.epoch.materialize", attrs); + if (!materialize_streaming_host_graph_range(source, target_rt_, range)) return false; + } + if (!stage_epoch(*range, epoch_index)) return false; + if (epoch_index == 0) { + char attrs[96]; + std::snprintf(attrs, sizeof(attrs), "epoch=%" PRIu64 " kind=publish-gate", publish_epoch); + STRACE_A("simpler_run.host_orch.epoch.wait_publish", attrs); + if (!wait_for_publish_gate()) return false; + } + + size_t control_base = offsetof(PTO2SharedMemoryHeader, host_graph_epochs); + if (epoch_index > 0 && !wait_for_device_epoch( + control_base + offsetof(PTO2HostGraphEpochControl, device_exec_done_epoch), + publish_epoch - 1, "exec-done" + )) { + return false; + } + if (!commit_publication(*range, epoch_index)) return false; + published_ranges_.push_back(*range); + if (!wait_for_device_epoch( + control_base + offsetof(PTO2HostGraphEpochControl, device_release_epoch), publish_epoch, "release" + )) { + return false; + } + source->orchestrator.sm_header->ring.fc.last_task_alive.store(range->task_end, std::memory_order_release); + return true; + } + + static bool + boundary_handler(PTO2Runtime *source, PTO2HostGraphEpochRange *range, size_t epoch_index, void *context) { + auto *job = static_cast(context); + return job != nullptr && job->publish_boundary(source, range, epoch_index); + } + + void run() { + L2TaskArgs orch_l2; + orch_l2.create_from_chip_args(device_args_); + int32_t total_tasks = run_host_orchestration( + runtime_, api_, target_rt_, target_arena_, layout_, device_sm_, sm_size_, device_arena_, gm_heap_, + heap_sizes_, task_window_sizes_, host_orch_func_ptr_, orch_l2, boundary_handler, this + ); + if (total_tasks < 0) { + publish_failure(PTO2_ERROR_EXPLICIT_ORCH_FATAL); + result_ = -1; + return; + } + + runtime_->host_total_tasks = total_tasks; + LOG_INFO_V9("host-orch: async streaming graph completed with %d tasks", total_tasks); + result_ = 0; + } + + Runtime *runtime_{nullptr}; + const HostApi *api_{nullptr}; + void *thread_context_{nullptr}; + std::shared_ptr stage1_gate_; + PTO2RuntimeArenaLayout layout_{}; + DeviceArena target_arena_; + PTO2Runtime *target_rt_{nullptr}; + std::vector target_sm_; + PTO2SharedMemoryHandle target_sm_handle_{}; + void *device_sm_{nullptr}; + uint64_t sm_size_{0}; + void *device_arena_{nullptr}; + void *gm_heap_{nullptr}; + uint64_t heap_sizes_[PTO2_MAX_RING_DEPTH]{}; + uint64_t task_window_sizes_[PTO2_MAX_RING_DEPTH]{}; + std::vector published_ranges_; + uint64_t target_reclaimed_epoch_{0}; + void *host_orch_func_ptr_{nullptr}; + ChipStorageTaskArgs device_args_; + std::thread worker_; + std::mutex run_gate_mutex_; + std::condition_variable run_gate_cv_; + std::atomic run_gate_state_{0}; + std::atomic publish_gate_state_{0}; + std::atomic worker_start_state_{0}; + int result_{-1}; +#if SIMPLER_HOST_STRACE + unsigned strace_inv_{0}; + uint64_t strace_hid_{0}; +#endif +}; + +static bool start_async_host_graph_pipeline( + Runtime *runtime, const HostApi *api, const PTO2RuntimeArenaLayout &layout, void *device_sm, uint64_t sm_size, + void *device_arena, void *gm_heap, const uint64_t heap_sizes[PTO2_MAX_RING_DEPTH], + const uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH], void *host_orch_func_ptr, + const ChipStorageTaskArgs &device_args +) { + auto *job = new (std::nothrow) HostGraphAsyncJob{}; + if (job == nullptr) return false; + if (!job->prepare_and_start( + runtime, api, layout, device_sm, sm_size, device_arena, gm_heap, heap_sizes, task_window_sizes, + host_orch_func_ptr, device_args + )) { + delete job; + return false; + } + runtime->set_host_orch_job(job); + return true; +} + +static int join_async_host_graph_pipeline(Runtime *runtime) { + auto *job = static_cast(runtime->take_host_orch_job()); + if (job == nullptr) return 0; + int rc = job->join(); + delete job; + return rc; +} + } // namespace +extern "C" int open_async_host_graph_pipeline(Runtime *runtime) { + if (runtime == nullptr) return -1; + auto *job = static_cast(runtime->take_host_orch_job()); + if (job == nullptr) return 0; + runtime->set_host_orch_job(job); + return job->open_run_gate() ? 0 : -1; +} + +extern "C" int release_async_host_graph_pipeline(Runtime *runtime) { + if (runtime == nullptr) return -1; + auto *job = static_cast(runtime->take_host_orch_job()); + if (job == nullptr) return 0; + runtime->set_host_orch_job(job); + return job->release_for_run() ? 0 : -1; +} + /** * Stage the per-callable resources (kernel binaries + orchestration SO) into * the supplied runtime so a subsequent bind_callable_to_runtime_impl can use @@ -753,12 +1637,8 @@ extern "C" int bind_callable_to_runtime_impl( uint64_t sm_size = PTO2SharedMemoryHandle::calculate_size_per_ring(eff_task_window_sizes); int64_t t_prebuilt_start = _now_ms(); - DeviceArena host_arena; // libc malloc backend by default - PTO2RuntimeArenaLayout layout = runtime_reserve_layout(host_arena, eff_task_window_sizes, eff_heap_sizes); - if (host_arena.commit(DeviceArena::kDefaultBaseAlign) == nullptr) { - LOG_ERROR("Failed to commit host arena for prebuilt runtime image"); - return -1; - } + DeviceArena layout_probe; + PTO2RuntimeArenaLayout layout = runtime_reserve_layout(layout_probe, eff_task_window_sizes, eff_heap_sizes); int64_t t_setup_start = _now_ms(); if (api->setup_static_arena(total_heap_size, sm_size, layout.arena_size) != 0) { @@ -794,62 +1674,17 @@ extern "C" int bind_callable_to_runtime_impl( // Set up orchestration state (consumed by the host orchestrator below) runtime->set_orch_args(device_args); - // ------------------------------------------------------------------------- - // Build the prebuilt runtime-arena image on host. - // - // We pre-compute every byte the AICPU's runtime arena would otherwise have - // to write at boot: layout offsets, sub-structure init data, and pointers - // back to the SM / GM heap. Then we rtMemcpy the image into the pooled - // runtime-arena region that DeviceRunner keeps alive across runs. AICPU - // boot becomes attach + wire (cheap pointer fixup) + sm_handle->init (SM - // reset) + a handful of device-only field fixups. - // ------------------------------------------------------------------------- - PTO2Runtime *rt = - runtime_init_data_from_layout(host_arena, layout, PTO2_MODE_EXECUTE, sm_ptr, sm_size, gm_heap, eff_heap_sizes); - if (rt == nullptr) { - LOG_ERROR("runtime_init_data_from_layout failed"); - return -1; - } - runtime_wire_arena_pointers(host_arena, layout, rt); - - // host_build_graph host-orch: run the orchestrator on the host now, against - // a host SM mirror, and ship the populated SM to the device. The arena - // (copied to the device below) carries the resulting orchestrator/scheduler - // state; the device boots scheduler-only. register_callable_impl guarantees - // host_orch_func_ptr is non-null on success (it fails the whole prepare - // otherwise), so this is an assertion-style guard, not a fallback path. if (host_orch_func_ptr == nullptr) { LOG_ERROR("host-orch: orchestration entry points were not resolved"); return -1; } - { - L2TaskArgs orch_l2; - orch_l2.create_from_chip_args(device_args); - int32_t total_tasks = run_host_orchestration( - runtime, api, rt, host_arena, layout, sm_ptr, sm_size, runtime_arena_dev, gm_heap, eff_heap_sizes, - eff_task_window_sizes, host_orch_func_ptr, orch_l2 - ); - if (total_tasks < 0) { - LOG_ERROR("host-orch: orchestration run failed"); - return -1; - } - runtime->host_total_tasks = total_tasks; - LOG_INFO_V0("host-orch: submitted %d tasks on host", total_tasks); - } - - // Stash the layout inside the PTO2Runtime image so the AICPU can recover - // every arena-internal offset after rtMemcpy. The runtime arena's device - // base does NOT travel in this image — it's on the host Runtime - // (set_prebuilt_arena below), since the AICPU needs that pointer - // *before* it can dereference the image. - rt->prebuilt_layout = layout; - - int rc_upload = api->copy_to_device(runtime_arena_dev, host_arena.base(), layout.arena_size); - if (rc_upload != 0) { - LOG_ERROR("Failed to rtMemcpy prebuilt runtime arena to device (rc=%d)", rc_upload); + if (!start_async_host_graph_pipeline( + runtime, api, layout, sm_ptr, sm_size, runtime_arena_dev, gm_heap, eff_heap_sizes, eff_task_window_sizes, + host_orch_func_ptr, device_args + )) { + LOG_ERROR("host-orch: failed to prepare the async whole-graph worker"); return -1; } - runtime->set_prebuilt_arena(runtime_arena_dev, layout.off_runtime); int64_t t_prebuilt_end = _now_ms(); LOG_INFO_V0("Device orchestration ready: %d tensors + %d scalars", tensor_count, scalar_count); @@ -889,6 +1724,11 @@ extern "C" int validate_runtime_impl(Runtime *runtime, const HostApi *api, int e int rc = 0; + if (join_async_host_graph_pipeline(runtime) != 0) { + LOG_ERROR("host-orch: async orchestration failed"); + rc = -1; + } + LOG_INFO_V0("=== Copying Results Back to Host ==="); // Copy all recorded tensors from device back to host diff --git a/src/a2a3/runtime/host_build_graph/orchestration/common.cpp b/src/a2a3/runtime/host_build_graph/orchestration/common.cpp index c4878a1c26..9e37d78522 100644 --- a/src/a2a3/runtime/host_build_graph/orchestration/common.cpp +++ b/src/a2a3/runtime/host_build_graph/orchestration/common.cpp @@ -10,6 +10,8 @@ */ #include "common.h" +#include + #ifdef __linux__ #include #include @@ -32,20 +34,39 @@ struct PTO2Runtime; extern "C" void unified_log_error(const char *func, const char *fmt, ...); namespace { -// Plain global (not thread_local) to avoid glibc TLSDESC stale-resolution -// crash (BZ #32412) when the orchestration SO is dlclose'd/re-dlopen'd -// between execution rounds. All orchestrator threads bind the same rt -// value, so per-thread storage is unnecessary. -PTO2Runtime *g_current_runtime = nullptr; +// HostGraph requests can orchestrate concurrently. POSIX TLS keeps each Host +// job on its own Runtime without using C++ thread_local, whose TLSDESC entries +// can stay stale when an orchestration SO is dlclose'd/re-dlopen'd. +pthread_key_t g_current_runtime_key; +pthread_once_t g_current_runtime_once = PTHREAD_ONCE_INIT; +bool g_current_runtime_key_ready = false; + +void create_current_runtime_key() { + if (pthread_key_create(&g_current_runtime_key, nullptr) == 0) { + g_current_runtime_key_ready = true; + } +} + +__attribute__((destructor)) void destroy_current_runtime_key() { + if (g_current_runtime_key_ready) { + pthread_key_delete(g_current_runtime_key); + g_current_runtime_key_ready = false; + } +} } // namespace extern "C" __attribute__((visibility("default"))) void framework_bind_runtime(PTO2Runtime *rt) { - g_current_runtime = rt; + pthread_once(&g_current_runtime_once, create_current_runtime_key); + if (g_current_runtime_key_ready) pthread_setspecific(g_current_runtime_key, rt); } // Keep current_runtime local to this .so so orchestration helpers do not // accidentally bind to the AICPU binary's same-named symbol. -extern "C" __attribute__((visibility("hidden"))) PTO2Runtime *framework_current_runtime() { return g_current_runtime; } +extern "C" __attribute__((visibility("hidden"))) PTO2Runtime *framework_current_runtime() { + pthread_once(&g_current_runtime_once, create_current_runtime_key); + return g_current_runtime_key_ready ? static_cast(pthread_getspecific(g_current_runtime_key)) : + nullptr; +} /** * Use addr2line to convert an address to file:line information. diff --git a/src/a2a3/runtime/host_build_graph/orchestration/pto_orchestration_api.h b/src/a2a3/runtime/host_build_graph/orchestration/pto_orchestration_api.h index f8321ece0f..fac6188f1c 100644 --- a/src/a2a3/runtime/host_build_graph/orchestration/pto_orchestration_api.h +++ b/src/a2a3/runtime/host_build_graph/orchestration/pto_orchestration_api.h @@ -91,6 +91,7 @@ typedef struct PTO2RuntimeOps { // collector can log it. Always present to keep ops-table layout stable // across SIMPLER_DFX settings; set to nullptr at SIMPLER_DFX=0. void (*scope_set_site)(const char *file, int line); + void (*graph_boundary)(PTO2Runtime *rt, bool final_epoch); } PTO2RuntimeOps; /** @@ -225,6 +226,14 @@ static inline void rt_orchestration_done() { rt->ops->orchestration_done(rt); } +static inline void rt_graph_boundary(bool final_epoch = false) { + PTO2Runtime *rt = current_runtime(); + if (rt->ops->is_fatal(rt) || rt->ops->graph_boundary == nullptr) { + return; + } + rt->ops->graph_boundary(rt, final_epoch); +} + static inline bool rt_is_fatal() { PTO2Runtime *rt = current_runtime(); return rt->ops->is_fatal(rt); diff --git a/src/a2a3/runtime/host_build_graph/runtime/host_graph_epoch.h b/src/a2a3/runtime/host_build_graph/runtime/host_graph_epoch.h new file mode 100644 index 0000000000..15c0fe331d --- /dev/null +++ b/src/a2a3/runtime/host_build_graph/runtime/host_graph_epoch.h @@ -0,0 +1,63 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +#pragma once + +#include +#include +#include + +inline constexpr int32_t PTO2_HOST_GRAPH_EPOCH_SLOT_COUNT = 2; + +struct PTO2HostGraphEpochRange { + int32_t task_begin{0}; + int32_t task_end{0}; + int32_t dep_begin{0}; + int32_t dep_end{0}; + int32_t inline_completed{0}; + int32_t final_epoch{0}; +}; + +struct alignas(64) PTO2HostGraphEpochSlot { + std::atomic owner_epoch; + PTO2HostGraphEpochRange range; + + void init() { + owner_epoch.store(0, std::memory_order_relaxed); + range = PTO2HostGraphEpochRange{}; + } +}; + +struct alignas(64) PTO2HostGraphEpochControl { + std::atomic host_publish_epoch; + std::atomic device_release_epoch; + std::atomic device_exec_done_epoch; + std::atomic device_buffer_free_epoch; + std::atomic failed_epoch; + PTO2HostGraphEpochSlot slots[PTO2_HOST_GRAPH_EPOCH_SLOT_COUNT]; + + void init() { + host_publish_epoch.store(0, std::memory_order_relaxed); + device_release_epoch.store(0, std::memory_order_relaxed); + device_exec_done_epoch.store(0, std::memory_order_relaxed); + device_buffer_free_epoch.store(0, std::memory_order_relaxed); + failed_epoch.store(0, std::memory_order_relaxed); + for (int32_t slot = 0; slot < PTO2_HOST_GRAPH_EPOCH_SLOT_COUNT; ++slot) { + slots[slot].init(); + } + } +}; + +static_assert(std::is_standard_layout_v); +static_assert(std::is_standard_layout_v); +static_assert(std::is_standard_layout_v); +static_assert(sizeof(PTO2HostGraphEpochSlot) % 64 == 0); +static_assert(sizeof(PTO2HostGraphEpochControl) % 64 == 0); diff --git a/src/a2a3/runtime/host_build_graph/runtime/host_graph_token_stream.h b/src/a2a3/runtime/host_build_graph/runtime/host_graph_token_stream.h new file mode 100644 index 0000000000..dc9418f4e9 --- /dev/null +++ b/src/a2a3/runtime/host_build_graph/runtime/host_graph_token_stream.h @@ -0,0 +1,36 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +#ifndef SRC_A2A3_RUNTIME_HOST_BUILD_GRAPH_RUNTIME_HOST_GRAPH_TOKEN_STREAM_H_ +#define SRC_A2A3_RUNTIME_HOST_BUILD_GRAPH_RUNTIME_HOST_GRAPH_TOKEN_STREAM_H_ + +#include + +constexpr uint64_t PTO2_HOST_GRAPH_TOKEN_STREAM_MAGIC_VERSION = 0x4847535400010000ULL; +constexpr int32_t PTO2_HOST_GRAPH_TOKEN_STREAM_TRAILER_SCALARS = 14; + +enum PTO2HostGraphTokenFlags : uint32_t { + PTO2_HOST_GRAPH_TOKEN_FINAL = 1U << 0, + PTO2_HOST_GRAPH_TOKEN_SYNTHETIC = 1U << 1, +}; + +struct PTO2HostGraphTokenPacket { + uint64_t magic_version; + uint64_t request_id; + uint64_t token_seq; + int64_t token_id; + uint32_t flags; + int32_t status; +}; + +static_assert(sizeof(PTO2HostGraphTokenPacket) == 40, "HostGraph token packet ABI drift"); + +#endif // SRC_A2A3_RUNTIME_HOST_BUILD_GRAPH_RUNTIME_HOST_GRAPH_TOKEN_STREAM_H_ diff --git a/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp b/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp index df383b8640..27e08eeeaf 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp @@ -405,6 +405,14 @@ static bool prepare_task( out->task = &orch->sm_header->ring.task_descriptors[out->alloc_result.slot]; out->payload = &orch->sm_header->ring.task_payloads[out->alloc_result.slot]; + // Streaming Host O advances last_task_alive at graph boundaries, so task IDs + // beyond the first physical window can reuse a slot. The original whole-graph + // HBG path initialized each slot only once and therefore never needed this + // reset on submit. + if (out->alloc_result.task_id >= allocator.window_size()) { + out->slot_state->reset_for_reuse(); + } + out->payload->prefetch(args.tensor_count(), args.scalar_count()); // Re-bind payload/task pointers each submit. Value is per-slot constant @@ -421,14 +429,12 @@ static bool prepare_task( // early-dispatch fields) is initialized in PTO2TaskPayload::init, the // single payload-init point, which runs before Orch-side wiring publish. - // Fields already zeroed by reset_for_reuse() at slot init: + // Fields already zeroed by reset_for_reuse() at slot init or physical reuse: // fanout_lock=0, fanout_count=PTO2_FANOUT_SCOPE_BIT, fanout_head=nullptr, // fanin_refcount=0, fanout_refcount=0, completed_subtasks=0, next_block_idx=0 // Fields immutable after RingSchedState::init(): // ring_id - // task_state is set to PENDING here as the orchestrator populates the slot - // (host_build_graph does not recycle slots at runtime, so there is no - // post-CONSUMED reset path). + // task_state is set to PENDING here as the orchestrator populates the slot. out->slot_state->task_state.store(PTO2_TASK_PENDING, std::memory_order_relaxed); int16_t block_num = args.launch_spec.block_num(); out->slot_state->total_required_subtasks = diff --git a/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_runtime2.cpp b/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_runtime2.cpp index 0ac79ba90a..8bc068cd67 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_runtime2.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_runtime2.cpp @@ -88,6 +88,19 @@ void rt_scope_end(PTO2Runtime *rt) { rt->orchestrator.end_scope(); } void rt_orchestration_done(PTO2Runtime *rt) { rt->orchestrator.mark_done(); } +void rt_graph_boundary(PTO2Runtime *rt, bool final_epoch) { + PTO2SharedMemoryHeader *header = rt->orchestrator.sm_header; + if (header == nullptr) { + rt->orchestrator.report_fatal(PTO2_ERROR_INVALID_ARGS, __FUNCTION__, "shared-memory header is null"); + return; + } + if (rt->graph_boundary_callback != nullptr && + !rt->graph_boundary_callback(rt, final_epoch, rt->graph_boundary_context)) { + rt->orchestrator.report_fatal(PTO2_ERROR_EXPLICIT_ORCH_FATAL, __FUNCTION__, "Host graph capture failed"); + return; + } +} + static bool is_fatal_impl(PTO2Runtime *rt) { return rt->orchestrator.fatal; } void rt_report_fatal(PTO2Runtime *rt, int32_t error_code, const char *func, const char *fmt, ...) { @@ -304,6 +317,7 @@ static const PTO2RuntimeOps s_runtime_ops = { #else .scope_set_site = nullptr, #endif + .graph_boundary = rt_graph_boundary, }; // ============================================================================= @@ -327,3 +341,11 @@ void runtime_set_mode(PTO2Runtime *rt, PTO2RuntimeMode mode) { rt->mode = mode; } } + +void runtime_set_graph_boundary_callback(PTO2Runtime *rt, PTO2GraphBoundaryCallback callback, void *context) { + if (rt == nullptr) { + return; + } + rt->graph_boundary_callback = callback; + rt->graph_boundary_context = context; +} diff --git a/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2.h b/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2.h index ad7989e2fa..061376203e 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2.h +++ b/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2.h @@ -65,6 +65,7 @@ enum PTO2RuntimeMode { * dependencies on runtime .cpp files. */ typedef struct PTO2Runtime PTO2Runtime; // forward declare for ops signatures +using PTO2GraphBoundaryCallback = bool (*)(PTO2Runtime *rt, bool final_epoch, void *context); struct PTO2RuntimeOps { TaskOutputTensors (*submit_task)(PTO2Runtime *rt, const MixedKernels &mixed_kernels, const L0TaskArgs &args); @@ -93,6 +94,7 @@ struct PTO2RuntimeOps { // collector. Always present in the struct to keep ops-table layout stable // across SIMPLER_DFX settings; set to nullptr at SIMPLER_DFX=0. void (*scope_set_site)(const char *file, int line); + void (*graph_boundary)(PTO2Runtime *rt, bool final_epoch); }; /** @@ -146,6 +148,11 @@ struct PTO2Runtime { // Statistics int64_t total_cycles; + // Host-only synchronous capture hook. The host clears both fields before + // the runtime-arena image is copied to the device. + PTO2GraphBoundaryCallback graph_boundary_callback; + void *graph_boundary_context; + // Prebuilt-arena fast path metadata. Carries every offset // wire_arena_pointers needs at AICPU boot so the AICPU can reconstruct // all arena-internal pointer fields without re-running init_data. The @@ -231,6 +238,8 @@ void runtime_destroy(PTO2Runtime *rt, DeviceArena &arena); */ void runtime_set_mode(PTO2Runtime *rt, PTO2RuntimeMode mode); +void runtime_set_graph_boundary_callback(PTO2Runtime *rt, PTO2GraphBoundaryCallback callback, void *context); + // ============================================================================= // Orchestration API (called by orchestration function) // ============================================================================= @@ -259,6 +268,8 @@ void rt_scope_end(PTO2Runtime *rt); */ void rt_orchestration_done(PTO2Runtime *rt); +void rt_graph_boundary(PTO2Runtime *rt, bool final_epoch = false); + /** * Enter fatal state explicitly from orchestration. */ diff --git a/src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h b/src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h index 1d861b6cf5..739ef61ed9 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h +++ b/src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h @@ -34,6 +34,7 @@ #include +#include "host_graph_epoch.h" #include "utils/device_arena.h" #include "pto_runtime2_types.h" @@ -156,13 +157,19 @@ struct alignas(PTO2_ALIGN_SIZE) PTO2SharedMemoryHeader { std::atomic sched_error_bitmap; // Bit X set = thread X had error std::atomic sched_error_code; // Last scheduler error code (last-writer-wins) std::atomic sched_error_thread; // Thread index of last error writer + + PTO2HostGraphEpochControl host_graph_epochs; }; -static_assert(sizeof(PTO2SharedMemoryHeader) == 320, "PTO2SharedMemoryHeader layout drift"); static_assert(offsetof(PTO2SharedMemoryHeader, total_size) == 264, "PTO2SharedMemoryHeader total_size layout drift"); static_assert( offsetof(PTO2SharedMemoryHeader, orch_error_code) == 272, "PTO2SharedMemoryHeader orch_error_code layout drift" ); +static_assert(offsetof(PTO2SharedMemoryHeader, host_graph_epochs) == 320, "HostGraph epoch control layout drift"); +static_assert( + (sizeof(PTO2SharedMemoryHeader) % PTO2_ALIGN_SIZE == 0) && (sizeof(PTO2SharedMemoryHeader) < 4096), + "PTO2SharedMemoryHeader should be aligned and reasonably sized" +); // ============================================================================= // Shared Memory Handle diff --git a/src/a2a3/runtime/host_build_graph/runtime/runtime.h b/src/a2a3/runtime/host_build_graph/runtime/runtime.h index d49e6d7242..6791ee7fb5 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/runtime.h +++ b/src/a2a3/runtime/host_build_graph/runtime/runtime.h @@ -270,6 +270,8 @@ class Runtime { void set_prebuilt_arena(void *arena_base, size_t runtime_off); void *get_prebuilt_arena_base() const; size_t get_prebuilt_runtime_offset() const; + void set_host_orch_job(void *job); + void *take_host_orch_job(); // Orchestration metadata written by the platform host (DeviceRunner) at // callable registration. Shared ABI with tensormap_and_ringbuffer; the @@ -314,6 +316,7 @@ class Runtime { // garbage, identical to host_api above. No fixed cap — grows with the // chip-level entry-tensor count. std::vector tensor_pairs_; + void *host_orch_job_{nullptr}; }; // Number of bytes of the Runtime image that must be copied to the device. diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp index c7bc14c22c..16af627105 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp @@ -57,7 +57,8 @@ LoopAction SchedulerContext::handle_orchestrator_exit( LOG_ERROR( "Thread %d: Fatal error (code=%d), sending EXIT_SIGNAL to all cores. " "completed_tasks=%d, total_tasks=%d", - thread_idx, orch_err, completed_tasks_.load(std::memory_order_relaxed), total_tasks_ + thread_idx, orch_err, completed_tasks_.load(std::memory_order_relaxed), + total_tasks_.load(std::memory_order_relaxed) ); if (!completed_.exchange(true, std::memory_order_acq_rel)) { emergency_shutdown(runtime); @@ -73,8 +74,9 @@ LoopAction SchedulerContext::handle_orchestrator_exit( return LoopAction::BREAK_LOOP; } - task_count = total_tasks_; - if (task_count > 0 && completed_tasks_.load(std::memory_order_relaxed) >= task_count) { + task_count = total_tasks_.load(std::memory_order_acquire); + if (orchestration_done_.load(std::memory_order_acquire) && task_count > 0 && + completed_tasks_.load(std::memory_order_relaxed) >= task_count) { completed_.store(true, std::memory_order_release); LOG_INFO_V0( "Thread %d: PTO2 completed tasks %d/%d", thread_idx, completed_tasks_.load(std::memory_order_relaxed), @@ -359,7 +361,9 @@ void SchedulerContext::log_shutdown_stall_snapshot( thread_count = thread_count < 0 ? 0 : MAX_AICPU_THREADS; } for (int32_t t = 0; t < thread_count; t++) { - log_stall_diagnostics(t, total_tasks_, trigger_idle_iterations, trigger_last_progress_count); + log_stall_diagnostics( + t, total_tasks_.load(std::memory_order_relaxed), trigger_idle_iterations, trigger_last_progress_count + ); } } @@ -901,7 +905,11 @@ int32_t SchedulerContext::post_handshake_init(Runtime *runtime) { #endif // Initialize task counters. Task count comes from PTO2 shared memory. - if (runtime->get_gm_sm_ptr()) { + if (runtime->host_total_tasks < 0) { + // Async Host O publishes the first range after AICPU init. Do not + // mistake a concurrently staged ring head for a final task count. + total_tasks_.store(0, std::memory_order_relaxed); + } else if (runtime->get_gm_sm_ptr()) { auto *header = static_cast(runtime->get_gm_sm_ptr()); // Read at one-time boot init, before the SM is reset for the run, so a // ring not yet written holds uninitialized memory (0xbe... under ASAN's @@ -915,11 +923,12 @@ int32_t SchedulerContext::post_handshake_init(Runtime *runtime) { int32_t ring_tasks = header->ring.fc.current_task_index.load(std::memory_order_acquire); if (ring_tasks > 0 && ring_tasks <= PTO2_SCOPE_TASKS_CAP) pto2_count += ring_tasks; } - total_tasks_ = static_cast(pto2_count); + total_tasks_.store(static_cast(pto2_count), std::memory_order_relaxed); } else { - total_tasks_ = 0; + total_tasks_.store(0, std::memory_order_relaxed); } completed_tasks_.store(0, std::memory_order_release); + orchestration_done_.store(false, std::memory_order_relaxed); // prepare_subtask_to_core fully writes a per-core payload / deferred-slab slot // before the AICore is told to read it: build_payload sets @@ -1011,7 +1020,8 @@ void SchedulerContext::deinit() { // Reset task counters and orchestrator state completed_tasks_.store(0, std::memory_order_release); - total_tasks_ = 0; + total_tasks_.store(0, std::memory_order_relaxed); + orchestration_done_.store(false, std::memory_order_relaxed); completed_.store(false, std::memory_order_release); // Reset core discovery and assignment state @@ -1056,7 +1066,7 @@ void SchedulerContext::on_orchestration_done( } #endif - total_tasks_ = total_tasks; + total_tasks_.store(total_tasks, std::memory_order_release); // Fold tasks completed inline during orchestration int32_t inline_completed = static_cast(rt->orchestrator.inline_completed_tasks); @@ -1106,6 +1116,8 @@ void SchedulerContext::on_orchestration_done( } } + orchestration_done_.store(true, std::memory_order_release); + #if SIMPLER_DFX // Write the core-to-thread mapping so the profiling data reflects the // scheduler threads' final core distribution. @@ -1119,3 +1131,46 @@ void SchedulerContext::on_orchestration_done( } #endif } + +void SchedulerContext::on_host_graph_published( + Runtime *runtime, [[maybe_unused]] PTO2Runtime *rt, [[maybe_unused]] int32_t thread_idx, int32_t total_tasks, + int32_t inline_completed, bool final_epoch +) { + total_tasks_.store(total_tasks, std::memory_order_release); + if (inline_completed > 0) { + completed_tasks_.fetch_add(inline_completed, std::memory_order_acq_rel); +#if SIMPLER_SCHED_PROFILING + rt->scheduler.tasks_completed.fetch_add(inline_completed, std::memory_order_relaxed); +#endif + } + + if (!final_epoch) return; + + int32_t orch_err = PTO2_ERROR_NONE; + if (sched_->sm_header != nullptr) { + orch_err = sched_->sm_header->orch_error_code.load(std::memory_order_acquire); + } + if (orch_err != PTO2_ERROR_NONE && !completed_.exchange(true, std::memory_order_acq_rel)) { + emergency_shutdown(runtime); + } + +#if SIMPLER_DFX + if (l2_swimlane_level_ >= L2SwimlaneLevel::ORCH_PHASES) { + l2_swimlane_aicpu_flush_orch_phase_buffer(thread_idx); + } + if (l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES) { + l2_swimlane_aicpu_init_core_assignments(cores_total_num_); + for (int32_t t = 0; t < active_sched_threads_; t++) { + l2_swimlane_aicpu_write_core_assignments_for_thread( + t, core_trackers_[t].core_ids(), core_trackers_[t].core_num() + ); + } + } +#endif + + orchestration_done_.store(true, std::memory_order_release); + LOG_INFO_V9( + "HostGraph final publication latched: total=%d completed=%d", total_tasks, + completed_tasks_.load(std::memory_order_relaxed) + ); +} diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h index 6c5e3200bb..3389288507 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h @@ -100,6 +100,14 @@ class SchedulerContext { // step belongs to the orchestrator lifecycle, not the scheduler. void on_orchestration_done(Runtime *runtime, PTO2Runtime *rt, int32_t thread_idx, int32_t total_tasks); + // Publish one Host-built epoch while scheduler threads are already live. + // `inline_completed` is the count within this epoch only; `final_epoch` + // latches the global completion condition after the last range arrives. + void on_host_graph_published( + Runtime *runtime, PTO2Runtime *rt, int32_t thread_idx, int32_t total_tasks, int32_t inline_completed, + bool final_epoch + ); + // Bind the PTO2Runtime scheduler pointer. void bind_runtime(PTO2Runtime *rt); @@ -149,7 +157,8 @@ class SchedulerContext { // --- Task-execution tracking --- std::atomic completed_tasks_{0}; - int32_t total_tasks_{0}; + std::atomic total_tasks_{0}; + std::atomic orchestration_done_{false}; std::atomic completed_{false}; uint64_t *func_id_to_addr_{nullptr}; diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp index 75b83958e2..c73475d5cc 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp @@ -1313,7 +1313,9 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ } if (idle_iterations % STALL_LOG_INTERVAL == 0) { - log_stall_diagnostics(thread_idx, total_tasks_, idle_iterations, last_progress_count); + log_stall_diagnostics( + thread_idx, total_tasks_.load(std::memory_order_relaxed), idle_iterations, last_progress_count + ); } // Wall-clock budget gate, with two fatal-latch branches: // @@ -1333,8 +1335,9 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ // observability is preserved. if (get_sys_cnt_aicpu() - last_progress_ts > scheduler_timeout_cycles) { bool self_owns = self_owns_running_task(thread_idx); - bool global_stuck = !self_owns && total_tasks_ > 0 && - completed_tasks_.load(std::memory_order_relaxed) < total_tasks_ && + int32_t published_tasks = total_tasks_.load(std::memory_order_relaxed); + bool global_stuck = !self_owns && published_tasks > 0 && + completed_tasks_.load(std::memory_order_relaxed) < published_tasks && no_thread_owns_running_task(); if (self_owns || global_stuck) { // Latch the error + emergency_shutdown, then break to the diff --git a/src/a2a3/runtime/host_build_graph/runtime/shared/pto_shared_memory.cpp b/src/a2a3/runtime/host_build_graph/runtime/shared/pto_shared_memory.cpp index 9b33075a72..135f953c75 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/shared/pto_shared_memory.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/shared/pto_shared_memory.cpp @@ -177,6 +177,7 @@ void PTO2SharedMemoryHandle::init_header_per_ring( header->sched_error_bitmap.store(0, std::memory_order_relaxed); header->sched_error_code.store(PTO2_ERROR_NONE, std::memory_order_relaxed); header->sched_error_thread.store(-1, std::memory_order_relaxed); + header->host_graph_epochs.init(); // Per-ring slot_states reset. Previously lived in // PTO2SchedulerState::RingSchedState::init(), but it writes into diff --git a/src/a2a3/runtime/host_build_graph/runtime/shared/runtime.cpp b/src/a2a3/runtime/host_build_graph/runtime/shared/runtime.cpp index 79a8001815..23c63b03be 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/shared/runtime.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/shared/runtime.cpp @@ -47,6 +47,7 @@ Runtime::Runtime() { orch_args_storage_.clear(); prebuilt_arena_base_ = nullptr; prebuilt_runtime_offset_ = 0; + host_orch_job_ = nullptr; active_callable_id_ = -1; dev_orch_so_addr_ = 0; @@ -82,6 +83,14 @@ void Runtime::set_prebuilt_arena(void *arena_base, size_t runtime_off) { void *Runtime::get_prebuilt_arena_base() const { return prebuilt_arena_base_; } size_t Runtime::get_prebuilt_runtime_offset() const { return prebuilt_runtime_offset_; } +void Runtime::set_host_orch_job(void *job) { host_orch_job_ = job; } + +void *Runtime::take_host_orch_job() { + void *job = host_orch_job_; + host_orch_job_ = nullptr; + return job; +} + // Orchestration metadata written by the platform host (DeviceRunner) at // callable registration. host_build_graph runs the orchestrator on the host so // the device side never reads these back, but the platform registration path is diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp index 1fef4668b9..d157dccbf6 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp @@ -258,14 +258,23 @@ int32_t AicpuExecutor::init(Runtime *runtime) { } // The orchestrator (top thread, tidx == nthreads-1) does not dispatch to - // cores, so it skips the handshake entirely and returns to - // run() immediately to build the graph — overlapping the schedulers' handshake. + // cores, so it skips the handshake entirely. The non-DFX path returns to + // run() immediately to build the graph, overlapping the schedulers' handshake. + // DFX must wait for the scheduler leader to initialize the shared collectors; + // otherwise the first submit can race dep_gen_aicpu_init() and disappear from + // tasks[] while later tensor-map edges still reference it. // Core counts it needs are derived from cores_total_num_ (fixed 1:2 ratio) in // run(). The remaining (nthreads-1) threads re-partition ALL cores among // themselves, so every core still gets its register window opened. const bool decouple_orch = (nthreads > 1) && !serial_orch_sched_; const bool is_orchestrator = (tidx == nthreads - 1); if (decouple_orch && is_orchestrator) { +#if SIMPLER_DFX + while (!init_done_.load(std::memory_order_acquire)) { + if (init_failed_.load(std::memory_order_acquire)) return -1; + } + if (init_failed_.load(std::memory_order_acquire)) return -1; +#endif return 0; // do NOT touch the handshake or hs_arrived_ barrier } const int32_t hs_nthreads = decouple_orch ? (nthreads - 1) : nthreads; diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp index c3f0092781..d8b72166b3 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cpp @@ -38,6 +38,7 @@ #include #include #include +#include #include #include @@ -47,6 +48,7 @@ #include "../runtime/runtime.h" #include "../../../../common/runtime_status/error_log.h" #include "../../../../common/task_interface/call_config.h" +#include "../../../../common/worker/pto_runtime_c_api.h" #include "callable.h" #include "common/platform_config.h" #include "common/strace.h" @@ -57,6 +59,24 @@ #include "utils/device_arena.h" #include "prepare_callable_common.h" +extern "C" const PipelineContract *get_pipeline_contract(void) { + static const PipelineContract contract = { + PTO_PIPELINE_CONTRACT_ABI_VERSION, + 4, + 2, + 1, + { + {PTO_PIPELINE_TASK_ARGS, PTO_PIPELINE_FILL_MEM, 0}, + {PTO_PIPELINE_RUNTIME_IMAGE, PTO_PIPELINE_REUSE_MEM, 0}, + {PTO_PIPELINE_AICPU_STREAM, PTO_PIPELINE_EXEC_HANDLE, 0}, + {PTO_PIPELINE_AICORE_STREAM, PTO_PIPELINE_EXEC_HANDLE, 0}, + }, + }; + return &contract; +} + +static std::mutex g_stage1_mutex; + static_assert( RUNTIME_ENV_RING_COUNT == PTO2_MAX_RING_DEPTH, "RuntimeEnv ring count must match PTO2 runtime ring depth" ); @@ -833,6 +853,11 @@ extern "C" int bind_callable_to_runtime_impl( return -1; } + // Gate A: one host Stage1 at a time. Releasing this lock at bind return + // lets the next request fill its independent TaskArg slot while this + // request waits in or executes the runner's Gate-B-protected Device S. + std::lock_guard stage1_lock(g_stage1_mutex); + int tensor_count = orch_args->tensor_count(); int scalar_count = orch_args->scalar_count(); LOG_INFO_V0("RT2 bind: %d tensors + %d scalars, device orchestration mode", tensor_count, scalar_count); @@ -873,6 +898,8 @@ extern "C" int bind_callable_to_runtime_impl( int64_t t_prebuilt_start = _now_ms(); { STRACE("simpler_run.bind.prebuilt"); + // Gate A also serializes a cold miss, so one request builds and uploads + // the single REUSE_MEM arena and the next observes the immutable cache. PrebuiltRuntimeArenaCacheProbe cache_probe = make_prebuilt_runtime_arena_cache_probe(sizing); int cache_rc = bind_cached_runtime_image(runtime, api, cache_probe, device_args); if (cache_rc < 0) { diff --git a/src/common/hierarchical/orchestrator.cpp b/src/common/hierarchical/orchestrator.cpp index dda1db8107..991219f9fe 100644 --- a/src/common/hierarchical/orchestrator.cpp +++ b/src/common/hierarchical/orchestrator.cpp @@ -68,6 +68,7 @@ TaskSlotState &Orchestrator::slot_state(TaskSlot s) { uint64_t Orchestrator::output_alloc_bytes(const Tensor &t) { return align_up(t.nbytes(), HEAP_ALIGN); } Tensor Orchestrator::alloc(const std::vector &shape, DataType dtype) { + std::lock_guard submit_lk(submit_mu_); if (shape.empty()) { // Rank-0 tensors are not supported across the ABI (Tensor enforces // ndims > 0). Reject here so we never allocate + register a buffer in @@ -189,6 +190,7 @@ SubmitResult Orchestrator::submit_impl( std::vector target_worker_ids, std::vector> eligible_worker_ids, std::vector remote_sidecars ) { + std::lock_guard submit_lk(submit_mu_); if (args_list.empty()) throw std::invalid_argument("Orchestrator: args_list must not be empty"); config.validate(); validate_worker_eligibility(worker_type, args_list.size(), target_worker_ids, eligible_worker_ids); @@ -629,9 +631,13 @@ void Orchestrator::infer_deps( // Scope // ============================================================================= -void Orchestrator::scope_begin() { scope_->scope_begin(); } +void Orchestrator::scope_begin() { + std::lock_guard submit_lk(submit_mu_); + scope_->scope_begin(); +} void Orchestrator::scope_end() { + std::lock_guard submit_lk(submit_mu_); scope_->scope_end([this](TaskSlot slot) { release_ref(slot); }); diff --git a/src/common/hierarchical/orchestrator.h b/src/common/hierarchical/orchestrator.h index cd6d266065..a0c1dc0c6d 100644 --- a/src/common/hierarchical/orchestrator.h +++ b/src/common/hierarchical/orchestrator.h @@ -172,6 +172,7 @@ class Orchestrator { std::atomic active_tasks_{0}; std::mutex drain_mu_; std::condition_variable drain_cv_; + std::mutex submit_mu_; // Scheduler's loop mutex (not owned). Held across reset_to_empty() in // drain() so the scheduler can't be mid-on_task_complete during teardown. std::mutex *sched_loop_mu_{nullptr}; diff --git a/src/common/hierarchical/scope.cpp b/src/common/hierarchical/scope.cpp index 055074a17c..402bfa1a20 100644 --- a/src/common/hierarchical/scope.cpp +++ b/src/common/hierarchical/scope.cpp @@ -12,19 +12,45 @@ #include "scope.h" void Scope::scope_begin() { - if (depth() >= MAX_SCOPE_DEPTH) throw std::runtime_error("Scope: maximum nesting depth exceeded"); - stack_.push_back(ScopeFrame{}); + std::lock_guard lk(mu_); + auto &stack = stacks_[std::this_thread::get_id()]; + if (stack.size() >= static_cast(MAX_SCOPE_DEPTH)) { + throw std::runtime_error("Scope: maximum nesting depth exceeded"); + } + stack.push_back(ScopeFrame{}); } void Scope::scope_end(const std::function &release_fn) { - if (stack_.empty()) throw std::runtime_error("Scope: scope_end without scope_begin"); - ScopeFrame &frame = stack_.back(); - for (TaskSlot slot : frame.tasks) + ScopeFrame frame; + { + std::lock_guard lk(mu_); + auto it = stacks_.find(std::this_thread::get_id()); + if (it == stacks_.end() || it->second.empty()) { + throw std::runtime_error("Scope: scope_end without scope_begin"); + } + frame = std::move(it->second.back()); + it->second.pop_back(); + if (it->second.empty()) stacks_.erase(it); + } + for (TaskSlot slot : frame.tasks) { release_fn(slot); - stack_.pop_back(); + } } void Scope::register_task(TaskSlot slot) { - if (stack_.empty()) return; // no open scope — task has no scope ref - stack_.back().tasks.push_back(slot); + std::lock_guard lk(mu_); + auto it = stacks_.find(std::this_thread::get_id()); + if (it == stacks_.end() || it->second.empty()) return; + it->second.back().tasks.push_back(slot); +} + +int32_t Scope::depth() const { + std::lock_guard lk(mu_); + auto it = stacks_.find(std::this_thread::get_id()); + return it == stacks_.end() ? 0 : static_cast(it->second.size()); +} + +int32_t Scope::current_depth() const { + int32_t d = depth(); + return d > 0 ? d - 1 : 0; } diff --git a/src/common/hierarchical/scope.h b/src/common/hierarchical/scope.h index f430c640b6..33112e126a 100644 --- a/src/common/hierarchical/scope.h +++ b/src/common/hierarchical/scope.h @@ -17,7 +17,8 @@ * scope_end() is called, that reference is released for every task in the scope, * allowing tasks that have no downstream consumers to reach CONSUMED. * - * Orch-owned: single-threaded, no locking required. + * Frames are owned by the calling orchestration thread. Multiple concurrent + * Worker::run calls may open independent outer scopes on one Orchestrator. * * Mirrors L2 scope_begin / scope_end semantics. */ @@ -25,7 +26,10 @@ #pragma once #include +#include #include +#include +#include #include #include "types.h" @@ -44,7 +48,7 @@ class Scope { void register_task(TaskSlot slot); // Current nesting depth (0 = no open scope). - int32_t depth() const { return static_cast(stack_.size()); } + int32_t depth() const; // L2-style 0-based scope index: the innermost open scope, or 0 when // none is open. Used by Ring::alloc to choose a heap ring: @@ -54,14 +58,12 @@ class Scope { // scope maps to ring 1, and so on. Returns 0 when no scope is open // so tasks submitted outside `Worker::run` still have a deterministic // ring assignment. - int32_t current_depth() const { - int32_t d = depth(); - return d > 0 ? d - 1 : 0; - } + int32_t current_depth() const; private: struct ScopeFrame { std::vector tasks; }; - std::vector stack_; + mutable std::mutex mu_; + std::unordered_map> stacks_; }; diff --git a/src/common/hierarchical/tensormap.cpp b/src/common/hierarchical/tensormap.cpp index 3e7fa9e4e8..d846c81308 100644 --- a/src/common/hierarchical/tensormap.cpp +++ b/src/common/hierarchical/tensormap.cpp @@ -12,16 +12,24 @@ #include "tensormap.h" TaskSlot TensorMap::lookup(TensorKey key) const { + std::lock_guard lk(mu_); auto it = map_.find(key); if (it == map_.end()) return INVALID_SLOT; return it->second; } -void TensorMap::insert(TensorKey key, TaskSlot producer) { map_[key] = producer; } +void TensorMap::insert(TensorKey key, TaskSlot producer) { + std::lock_guard lk(mu_); + map_[key] = producer; +} void TensorMap::erase_task_outputs(const std::vector &keys) { + std::lock_guard lk(mu_); for (const auto &key : keys) map_.erase(key); } -int32_t TensorMap::size() const { return static_cast(map_.size()); } +int32_t TensorMap::size() const { + std::lock_guard lk(mu_); + return static_cast(map_.size()); +} diff --git a/src/common/hierarchical/tensormap.h b/src/common/hierarchical/tensormap.h index 326970e269..0184db3048 100644 --- a/src/common/hierarchical/tensormap.h +++ b/src/common/hierarchical/tensormap.h @@ -23,11 +23,13 @@ * - Does not perform overlap detection (each key maps to one producer) * - Cleans up entries actively when a task is CONSUMED * - * Owned exclusively by the Orchestrator (main thread); no locking required. + * Concurrent orchestration threads may submit while the scheduler consumes + * completed tasks, so every map operation is internally synchronized. */ #pragma once +#include #include #include @@ -51,5 +53,6 @@ class TensorMap { int32_t size() const; private: + mutable std::mutex mu_; std::unordered_map map_; }; diff --git a/src/common/hierarchical/worker.cpp b/src/common/hierarchical/worker.cpp index e061248eea..2f0bb0fba6 100644 --- a/src/common/hierarchical/worker.cpp +++ b/src/common/hierarchical/worker.cpp @@ -68,15 +68,18 @@ Worker::~Worker() { if (initialized_) close(); } -void Worker::add_worker(WorkerType type, void *mailbox) { +void Worker::add_worker(WorkerType type, void *mailbox, uint32_t max_in_flight) { if (initialized_) throw std::runtime_error("Worker: add_worker after init"); - if (type == WorkerType::NEXT_LEVEL) manager_.add_next_level(mailbox); - else manager_.add_sub(mailbox); + if (type == WorkerType::NEXT_LEVEL) manager_.add_next_level(mailbox, max_in_flight); + else { + if (max_in_flight != 1) throw std::invalid_argument("Worker: SUB worker max_in_flight must be 1"); + manager_.add_sub(mailbox); + } } -void Worker::add_next_level_worker(int32_t worker_id, void *mailbox) { +void Worker::add_next_level_worker(int32_t worker_id, void *mailbox, uint32_t max_in_flight) { if (initialized_) throw std::runtime_error("Worker: add_next_level_worker after init"); - manager_.add_next_level_at(worker_id, mailbox); + manager_.add_next_level_at(worker_id, mailbox, max_in_flight); } void Worker::add_remote_l3_socket( diff --git a/src/common/hierarchical/worker.h b/src/common/hierarchical/worker.h index 1ced949bd5..29750e0969 100644 --- a/src/common/hierarchical/worker.h +++ b/src/common/hierarchical/worker.h @@ -74,8 +74,8 @@ class Worker { // MAILBOX_SIZE-byte MAP_SHARED region; the real worker (a `ChipWorker` // for NEXT_LEVEL, a Python callable for SUB) lives in the forked // child and consumes the mailbox via the Python child loop. - void add_worker(WorkerType type, void *mailbox); - void add_next_level_worker(int32_t worker_id, void *mailbox); + void add_worker(WorkerType type, void *mailbox, uint32_t max_in_flight = 1); + void add_next_level_worker(int32_t worker_id, void *mailbox, uint32_t max_in_flight = 1); // Register a REMOTE_L3 endpoint only after its session runner completed // prestart and reported HELLO READY on the command lane. diff --git a/src/common/hierarchical/worker_manager.cpp b/src/common/hierarchical/worker_manager.cpp index 5044b37f6a..246c6865cd 100644 --- a/src/common/hierarchical/worker_manager.cpp +++ b/src/common/hierarchical/worker_manager.cpp @@ -126,14 +126,22 @@ void WorkerEndpoint::control_l3_l2_region_release(uint64_t) { // LocalMailboxEndpoint — mailbox helpers // ============================================================================= -LocalMailboxEndpoint::LocalMailboxEndpoint(int32_t worker_id, void *mailbox) : +LocalMailboxEndpoint::LocalMailboxEndpoint(int32_t worker_id, void *mailbox, uint32_t task_slot_count) : mailbox_(mailbox) { if (mailbox == nullptr) throw std::invalid_argument("LocalMailboxEndpoint: null mailbox"); + if (task_slot_count == 0 || task_slot_count > MAILBOX_TASK_SLOT_COUNT) { + throw std::invalid_argument("LocalMailboxEndpoint: invalid task_slot_count"); + } caps_.worker_id = worker_id; + caps_.max_in_flight = task_slot_count; + for (size_t slot = 0; slot < task_slot_count; ++slot) { + free_task_slots_.push(slot); + generations_[slot].store(0, std::memory_order_relaxed); + } } -MailboxState LocalMailboxEndpoint::read_mailbox_state() const { - volatile int32_t *ptr = reinterpret_cast(mbox() + MAILBOX_OFF_STATE); +MailboxState LocalMailboxEndpoint::read_mailbox_state(const char *frame) const { + const volatile int32_t *ptr = reinterpret_cast(frame + MAILBOX_OFF_STATE); int32_t v; #if defined(__aarch64__) __asm__ volatile("ldar %w0, [%1]" : "=r"(v) : "r"(ptr) : "memory"); @@ -146,8 +154,8 @@ MailboxState LocalMailboxEndpoint::read_mailbox_state() const { return static_cast(v); } -void LocalMailboxEndpoint::write_mailbox_state(MailboxState s) { - volatile int32_t *ptr = reinterpret_cast(mbox() + MAILBOX_OFF_STATE); +void LocalMailboxEndpoint::write_mailbox_state(char *frame, MailboxState s) { + volatile int32_t *ptr = reinterpret_cast(frame + MAILBOX_OFF_STATE); int32_t v = static_cast(s); #if defined(__aarch64__) __asm__ volatile("stlr %w0, [%1]" : : "r"(v), "r"(ptr) : "memory"); @@ -159,7 +167,28 @@ void LocalMailboxEndpoint::write_mailbox_state(MailboxState s) { #endif } -void LocalMailboxEndpoint::shutdown_child() { write_mailbox_state(MailboxState::SHUTDOWN); } +void LocalMailboxEndpoint::shutdown_child() { + std::unique_lock lk(mailbox_mu_); + write_mailbox_state(mbox(), MailboxState::SHUTDOWN); +} + +size_t LocalMailboxEndpoint::acquire_task_slot() { + std::unique_lock lk(free_slot_mu_); + free_slot_cv_.wait(lk, [this] { + return !free_task_slots_.empty(); + }); + size_t task_slot = free_task_slots_.front(); + free_task_slots_.pop(); + return task_slot; +} + +void LocalMailboxEndpoint::release_task_slot(size_t task_slot) { + { + std::lock_guard lk(free_slot_mu_); + free_task_slots_.push(task_slot); + } + free_slot_cv_.notify_one(); +} // ============================================================================= // WorkerThread — lifecycle @@ -175,12 +204,23 @@ void WorkerThread::start( on_complete_ = on_complete; endpoint_ = std::move(endpoint); shutdown_ = false; - idle_.store(true, std::memory_order_relaxed); - thread_ = std::thread(&WorkerThread::loop, this); + capacity_ = std::max(1, endpoint_->caps().max_in_flight); + available_.store(capacity_, std::memory_order_relaxed); + threads_.reserve(capacity_); + for (uint32_t i = 0; i < capacity_; ++i) { + threads_.emplace_back(&WorkerThread::loop, this); + } } void WorkerThread::dispatch(WorkerDispatch d) { - idle_.store(false, std::memory_order_release); + uint32_t available = available_.load(std::memory_order_acquire); + while (true) { + if (available == 0) throw std::runtime_error("WorkerThread::dispatch: no available endpoint credit"); + if (available_.compare_exchange_weak( + available, available - 1, std::memory_order_acq_rel, std::memory_order_acquire + )) + break; + } std::lock_guard lk(mu_); queue_.push(d); cv_.notify_one(); @@ -192,7 +232,10 @@ void WorkerThread::stop() { shutdown_ = true; } cv_.notify_all(); - if (thread_.joinable()) thread_.join(); + for (auto &thread : threads_) { + if (thread.joinable()) thread.join(); + } + threads_.clear(); } void WorkerThread::shutdown_child() { @@ -242,7 +285,7 @@ void WorkerThread::loop() { manager_->report_error(std::make_exception_ptr(std::runtime_error(completion.error_message))); } - idle_.store(true, std::memory_order_release); + available_.fetch_add(1, std::memory_order_release); on_complete_(std::move(completion)); } } @@ -253,7 +296,21 @@ WorkerCompletion WorkerThread::dispatch_process(WorkerDispatch d) { } WorkerCompletion LocalMailboxEndpoint::run(Ring *ring, const WorkerDispatch &dispatch) { + std::shared_lock mailbox_lk(mailbox_mu_); + size_t task_slot = acquire_task_slot(); + try { + WorkerCompletion completion = run_task_slot(ring, dispatch, task_slot); + release_task_slot(task_slot); + return completion; + } catch (...) { + release_task_slot(task_slot); + throw; + } +} + +WorkerCompletion LocalMailboxEndpoint::run_task_slot(Ring *ring, const WorkerDispatch &dispatch, size_t task_slot) { if (ring == nullptr) throw std::invalid_argument("LocalMailboxEndpoint::run: null ring"); + char *frame = mbox(task_slot); TaskSlotState &s = *ring->slot_state(dispatch.task_slot); int32_t group_index = dispatch.group_index; TaskArgsView view = s.args_view(group_index); @@ -269,11 +326,6 @@ WorkerCompletion LocalMailboxEndpoint::run(Ring *ring, const WorkerDispatch &dis completion.task_slot = dispatch.task_slot; completion.group_index = group_index; - // Hold mailbox_mu_ for the entire round trip (write payload + state + - // spin-poll TASK_DONE + reset to IDLE). Any control_* request from the - // orch thread waits for the dispatch to finish before claiming the - // mailbox; without this they would race on MAILBOX_OFF_STATE. - std::lock_guard lk(mailbox_mu_); if (mailbox_control_timed_out_) { completion.outcome = EndpointOutcome::ENDPOINT_FAILURE; completion.error_message = "LocalMailboxEndpoint::run: mailbox has an unresolved timed-out control command"; @@ -283,14 +335,16 @@ WorkerCompletion LocalMailboxEndpoint::run(Ring *ring, const WorkerDispatch &dis // Clear the child-writable error fields so stale bytes from a prior // dispatch cannot masquerade as a fresh failure. int32_t zero_err = 0; - std::memcpy(mbox() + MAILBOX_OFF_ERROR, &zero_err, sizeof(int32_t)); - std::memset(mbox() + MAILBOX_OFF_ERROR_MSG, 0, MAILBOX_ERROR_MSG_SIZE); + std::memcpy(frame + MAILBOX_OFF_ERROR, &zero_err, sizeof(int32_t)); + std::memset(frame + MAILBOX_OFF_ERROR_MSG, 0, MAILBOX_ERROR_MSG_SIZE); - uint64_t reserved_callable = 0; - std::memcpy(mbox() + MAILBOX_OFF_CALLABLE, &reserved_callable, sizeof(uint64_t)); + uint64_t generation = generations_[task_slot].fetch_add(1, std::memory_order_relaxed) + 1; + std::memcpy(frame + MAILBOX_OFF_GENERATION, &generation, sizeof(uint64_t)); + uint64_t protocol = MAILBOX_PROTOCOL_MAGIC_VERSION; + std::memcpy(frame + MAILBOX_OFF_PROTOCOL, &protocol, sizeof(uint64_t)); // Write config as a single packed POD block (see call_config.h). - std::memcpy(mbox() + MAILBOX_OFF_CONFIG, &s.config, sizeof(CallConfig)); + std::memcpy(frame + MAILBOX_OFF_CONFIG, &s.config, sizeof(CallConfig)); // Write length-prefixed TaskArgs blob: [T][S][tensors][scalars]. size_t blob_bytes = TASK_ARGS_BLOB_HEADER_SIZE + static_cast(view.tensor_count) * sizeof(Tensor) + @@ -303,10 +357,10 @@ WorkerCompletion LocalMailboxEndpoint::run(Ring *ring, const WorkerDispatch &dis " bytes, tensors=" + std::to_string(view.tensor_count) + ", scalars=" + std::to_string(view.scalar_count); return completion; } - uint8_t *hash_dst = reinterpret_cast(mbox() + MAILBOX_OFF_TASK_CALLABLE_HASH); + uint8_t *hash_dst = reinterpret_cast(frame + MAILBOX_OFF_TASK_CALLABLE_HASH); std::memcpy(hash_dst, s.callable.digest.data(), CALLABLE_HASH_DIGEST_SIZE); - uint8_t *d = reinterpret_cast(mbox() + MAILBOX_OFF_TASK_ARGS_BLOB); + uint8_t *d = reinterpret_cast(frame + MAILBOX_OFF_TASK_ARGS_BLOB); std::memcpy(d + 0, &view.tensor_count, sizeof(int32_t)); std::memcpy(d + 4, &view.scalar_count, sizeof(int32_t)); if (view.tensor_count > 0) { @@ -322,10 +376,12 @@ WorkerCompletion LocalMailboxEndpoint::run(Ring *ring, const WorkerDispatch &dis } // Signal child process. - write_mailbox_state(MailboxState::TASK_READY); + write_mailbox_state(frame, MailboxState::TASK_READY); - // Spin-poll until child signals TASK_DONE. - while (read_mailbox_state() != MailboxState::TASK_DONE) { + // TASK_ACCEPTED is the post-KernelLaunch flight ACK. This endpoint keeps + // ownership of the slot until TASK_DONE so generation/error validation and + // Gate-A bank reuse remain completion-driven. + while (read_mailbox_state(frame) != MailboxState::TASK_DONE) { std::this_thread::sleep_for(std::chrono::microseconds(50)); } @@ -334,10 +390,18 @@ WorkerCompletion LocalMailboxEndpoint::run(Ring *ring, const WorkerDispatch &dis // caught an exception and filled OFF_ERROR_MSG with // `f"{type(e).__name__}: {e}"` (truncated to MAILBOX_ERROR_MSG_SIZE). int32_t error_code = 0; - std::memcpy(&error_code, mbox() + MAILBOX_OFF_ERROR, sizeof(int32_t)); + uint64_t completed_generation = 0; + std::memcpy(&completed_generation, frame + MAILBOX_OFF_GENERATION, sizeof(uint64_t)); + if (completed_generation != generation) { + write_mailbox_state(frame, MailboxState::IDLE); + completion.outcome = EndpointOutcome::ENDPOINT_FAILURE; + completion.error_message = "LocalMailboxEndpoint::run: task generation changed before completion"; + return completion; + } + std::memcpy(&error_code, frame + MAILBOX_OFF_ERROR, sizeof(int32_t)); if (error_code != 0) { - std::string msg = read_error_msg(mbox()); - write_mailbox_state(MailboxState::IDLE); + std::string msg = read_error_msg(frame); + write_mailbox_state(frame, MailboxState::IDLE); completion.outcome = EndpointOutcome::TASK_FAILURE; completion.error_message = "LocalMailboxEndpoint::run: child failed (worker_id=" + std::to_string(caps_.worker_id) + @@ -345,7 +409,7 @@ WorkerCompletion LocalMailboxEndpoint::run(Ring *ring, const WorkerDispatch &dis return completion; } - write_mailbox_state(MailboxState::IDLE); + write_mailbox_state(frame, MailboxState::IDLE); completion.outcome = EndpointOutcome::SUCCESS; return completion; } @@ -354,13 +418,16 @@ WorkerCompletion LocalMailboxEndpoint::run(Ring *ring, const WorkerDispatch &dis // WorkerManager // ============================================================================= -void WorkerManager::add_next_level(void *mailbox) { - add_next_level_at(static_cast(next_level_entries_.size()), mailbox); +void WorkerManager::add_next_level(void *mailbox, uint32_t max_in_flight) { + add_next_level_at(static_cast(next_level_entries_.size()), mailbox, max_in_flight); } -void WorkerManager::add_next_level_at(int32_t worker_id, void *mailbox) { +void WorkerManager::add_next_level_at(int32_t worker_id, void *mailbox, uint32_t max_in_flight) { if (worker_id < 0) throw std::invalid_argument("WorkerManager::add_next_level_at: negative worker_id"); - next_level_entries_.push_back(LocalNextLevelEntry{worker_id, mailbox}); + if (max_in_flight == 0 || max_in_flight > MAILBOX_TASK_SLOT_COUNT) { + throw std::invalid_argument("WorkerManager::add_next_level_at: invalid max_in_flight"); + } + next_level_entries_.push_back(LocalNextLevelEntry{worker_id, mailbox, max_in_flight}); } void WorkerManager::add_next_level_endpoint(std::unique_ptr endpoint) { @@ -397,7 +464,7 @@ void WorkerManager::start(Ring *ring, const OnCompleteFn &on_complete) { auto make_next_level_threads = [&]() { for (const auto &entry : next_level_entries_) { auto wt = std::make_unique(); - auto endpoint = std::make_unique(entry.worker_id, entry.mailbox); + auto endpoint = std::make_unique(entry.worker_id, entry.mailbox, entry.max_in_flight); wt->start(ring, this, on_complete, std::move(endpoint)); next_level_threads_.push_back(std::move(wt)); } @@ -518,14 +585,14 @@ void LocalMailboxEndpoint::run_control_command(const char *op_name, double timeo int32_t zero_err = 0; std::memcpy(mbox() + MAILBOX_OFF_ERROR, &zero_err, sizeof(int32_t)); std::memset(mbox() + MAILBOX_OFF_ERROR_MSG, 0, MAILBOX_ERROR_MSG_SIZE); - write_mailbox_state(MailboxState::CONTROL_REQUEST); + write_mailbox_state(mbox(), MailboxState::CONTROL_REQUEST); auto deadline = std::chrono::steady_clock::time_point::max(); if (timeout_s >= 0.0) { deadline = std::chrono::steady_clock::now() + std::chrono::duration_cast(std::chrono::duration(timeout_s)); } - while (read_mailbox_state() != MailboxState::CONTROL_DONE) { + while (read_mailbox_state(mbox()) != MailboxState::CONTROL_DONE) { if (std::chrono::steady_clock::now() >= deadline) { mailbox_control_timed_out_ = true; throw std::runtime_error(std::string(op_name) + " timed out waiting for CONTROL_DONE"); @@ -535,28 +602,28 @@ void LocalMailboxEndpoint::run_control_command(const char *op_name, double timeo std::memcpy(&err, mbox() + MAILBOX_OFF_ERROR, sizeof(int32_t)); if (err != 0) { std::string msg = read_error_msg(mbox()); - write_mailbox_state(MailboxState::IDLE); + write_mailbox_state(mbox(), MailboxState::IDLE); throw std::runtime_error(std::string(op_name) + " failed on child: " + msg); } - write_mailbox_state(MailboxState::IDLE); + write_mailbox_state(mbox(), MailboxState::IDLE); } uint64_t LocalMailboxEndpoint::control_malloc(size_t size) { - std::lock_guard lk(mailbox_mu_); + std::unique_lock lk(mailbox_mu_); write_control_args(mbox(), CTRL_MALLOC, static_cast(size)); run_control_command("control_malloc"); return read_control_result(mbox()); } void LocalMailboxEndpoint::control_prepare(const uint8_t *digest) { - std::lock_guard lk(mailbox_mu_); + std::unique_lock lk(mailbox_mu_); write_control_args(mbox(), CTRL_PREPARE); write_control_digest(mbox(), digest); run_control_command("control_prepare"); } void LocalMailboxEndpoint::control_register(const char *shm_name, size_t blob_size, const uint8_t *digest) { - std::lock_guard lk(mailbox_mu_); + std::unique_lock lk(mailbox_mu_); // OFF_ERROR / OFF_ERROR_MSG are cleared by run_control_command — no // prelude memset needed (matches the other control_* methods). uint64_t sub_cmd = CTRL_REGISTER; @@ -576,7 +643,7 @@ void LocalMailboxEndpoint::control_register(const char *shm_name, size_t blob_si } void LocalMailboxEndpoint::control_unregister(const uint8_t *digest) { - std::lock_guard lk(mailbox_mu_); + std::unique_lock lk(mailbox_mu_); write_control_args(mbox(), CTRL_UNREGISTER); write_control_digest(mbox(), digest); run_control_command("control_unregister"); @@ -637,7 +704,7 @@ void LocalMailboxEndpoint::control_remote_release_import(const RemoteBufferHandl void LocalMailboxEndpoint::control_generic( uint64_t sub_cmd, const char *shm_name, size_t staged_payload_size, double timeout_s, const uint8_t *digest ) { - std::lock_guard lk(mailbox_mu_); + std::unique_lock lk(mailbox_mu_); std::memcpy(mbox() + MAILBOX_OFF_CALLABLE, &sub_cmd, sizeof(uint64_t)); uint64_t payload_size = static_cast(staged_payload_size); std::memcpy(mbox() + CTRL_OFF_ARG0, &payload_size, sizeof(uint64_t)); @@ -653,19 +720,19 @@ void LocalMailboxEndpoint::control_generic( } void LocalMailboxEndpoint::control_free(uint64_t ptr) { - std::lock_guard lk(mailbox_mu_); + std::unique_lock lk(mailbox_mu_); write_control_args(mbox(), CTRL_FREE, ptr); run_control_command("control_free"); } void LocalMailboxEndpoint::control_copy_to(uint64_t dst, uint64_t src, size_t size) { - std::lock_guard lk(mailbox_mu_); + std::unique_lock lk(mailbox_mu_); write_control_args(mbox(), CTRL_COPY_TO, dst, src, static_cast(size)); run_control_command("control_copy_to"); } void LocalMailboxEndpoint::control_copy_from(uint64_t dst, uint64_t src, size_t size) { - std::lock_guard lk(mailbox_mu_); + std::unique_lock lk(mailbox_mu_); write_control_args(mbox(), CTRL_COPY_FROM, dst, src, static_cast(size)); run_control_command("control_copy_from"); } @@ -691,7 +758,7 @@ void LocalMailboxEndpoint::control_alloc_domain(const char *request_shm_name, co if (!request_shm_name || !*request_shm_name || !reply_shm_name || !*reply_shm_name) { throw std::runtime_error("control_alloc_domain: request and reply shm names must be non-empty"); } - std::lock_guard lk(mailbox_mu_); + std::unique_lock lk(mailbox_mu_); uint64_t sub_cmd = CTRL_ALLOC_DOMAIN; std::memcpy(mbox() + MAILBOX_OFF_CALLABLE, &sub_cmd, sizeof(uint64_t)); write_shm_name_pair(mbox(), request_shm_name, reply_shm_name); @@ -702,7 +769,7 @@ void LocalMailboxEndpoint::control_release_domain(const char *request_shm_name) if (!request_shm_name || !*request_shm_name) { throw std::runtime_error("control_release_domain: request shm name must be non-empty"); } - std::lock_guard lk(mailbox_mu_); + std::unique_lock lk(mailbox_mu_); uint64_t sub_cmd = CTRL_RELEASE_DOMAIN; std::memcpy(mbox() + MAILBOX_OFF_CALLABLE, &sub_cmd, sizeof(uint64_t)); write_shm_name_pair(mbox(), request_shm_name, ""); @@ -713,7 +780,7 @@ void LocalMailboxEndpoint::control_comm_init(const char *request_shm_name) { if (!request_shm_name || !*request_shm_name) { throw std::runtime_error("control_comm_init: request shm name must be non-empty"); } - std::lock_guard lk(mailbox_mu_); + std::unique_lock lk(mailbox_mu_); uint64_t sub_cmd = CTRL_COMM_INIT; std::memcpy(mbox() + MAILBOX_OFF_CALLABLE, &sub_cmd, sizeof(uint64_t)); write_shm_name_pair(mbox(), request_shm_name, ""); @@ -724,7 +791,7 @@ void LocalMailboxEndpoint::control_l3_l2_region_create(const char *request_shm_n if (!request_shm_name || !*request_shm_name || !reply_shm_name || !*reply_shm_name) { throw std::runtime_error("control_l3_l2_region_create: request and reply shm names must be non-empty"); } - std::lock_guard lk(mailbox_mu_); + std::unique_lock lk(mailbox_mu_); uint64_t sub_cmd = CTRL_L3_L2_REGION_CREATE; std::memcpy(mbox() + MAILBOX_OFF_CALLABLE, &sub_cmd, sizeof(uint64_t)); write_shm_name_pair(mbox(), request_shm_name, reply_shm_name); @@ -732,7 +799,7 @@ void LocalMailboxEndpoint::control_l3_l2_region_create(const char *request_shm_n } void LocalMailboxEndpoint::control_l3_l2_region_release(uint64_t region_id) { - std::lock_guard lk(mailbox_mu_); + std::unique_lock lk(mailbox_mu_); write_control_args(mbox(), CTRL_L3_L2_REGION_RELEASE, region_id); run_control_command("control_l3_l2_region_release"); } @@ -877,9 +944,9 @@ void WorkerThread::control_l3_l2_region_release(uint64_t region_id) { bool WorkerManager::any_busy() const { for (auto &wt : next_level_threads_) - if (!wt->idle()) return true; + if (wt->busy()) return true; for (auto &wt : sub_threads_) - if (!wt->idle()) return true; + if (wt->busy()) return true; return false; } diff --git a/src/common/hierarchical/worker_manager.h b/src/common/hierarchical/worker_manager.h index 7d4cfe1fd8..5198bce99a 100644 --- a/src/common/hierarchical/worker_manager.h +++ b/src/common/hierarchical/worker_manager.h @@ -18,7 +18,8 @@ * * Each WorkerThread encodes `(callable digest, config, args_blob)` into a * pre-forked child's shared-memory mailbox, signals TASK_READY, and - * spin-polls TASK_DONE. The child process loop (Python) reads the + * observes TASK_ACCEPTED after native KernelLaunch, then spin-polls TASK_DONE. + * The child process loop (Python) reads the * digest, resolves it to a child-local slot, and runs that slot on its * `ChipWorker` (NEXT_LEVEL) or registered Python callable (SUB) in its * own address space. @@ -27,6 +28,7 @@ #pragma once #include +#include #include #include #include @@ -36,6 +38,7 @@ #include #include #include +#include #include #include #include @@ -72,6 +75,9 @@ enum class MailboxState : int32_t { // PLATFORM_STREAM_SYNC_TIMEOUT_MS budget (issue #897). INIT_READY = 6, INIT_FAILED = 7, + // Native a2a3 publishes this only after both AICore and AICPU launches + // succeed. Completion and mailbox reuse still require TASK_DONE. + TASK_ACCEPTED = 8, }; // Sized so the args region can hold any TaskArgs the runtime itself accepts @@ -82,7 +88,10 @@ enum class MailboxState : int32_t { // to the unified 128 B Tensor: the worst-case blob (CHIP_MAX_TENSOR_ARGS tensors) // grew ~3x, and 128*128 B = 16 KB alone exceeded the old mailbox (see the // capacity static_assert after MAILBOX_ARGS_CAPACITY). -static constexpr size_t MAILBOX_SIZE = 32768; +static constexpr size_t MAILBOX_TASK_SLOT_SIZE = 32768; +static constexpr size_t MAILBOX_TASK_SLOT_COUNT = 2; +static constexpr size_t MAILBOX_SIZE = MAILBOX_TASK_SLOT_SIZE * MAILBOX_TASK_SLOT_COUNT; +static constexpr uint64_t MAILBOX_PROTOCOL_MAGIC_VERSION = 0x534D424F00020000ULL; // Error message region lives at the mailbox tail. 256 B of headroom is // enough for `: ` produced by the child-side @@ -110,15 +119,16 @@ static_assert( "CallConfig overflows reserved config region" ); static constexpr ptrdiff_t MAILBOX_OFF_ERROR_MSG = - static_cast(MAILBOX_SIZE) - static_cast(MAILBOX_ERROR_MSG_SIZE); + static_cast(MAILBOX_TASK_SLOT_SIZE) - static_cast(MAILBOX_ERROR_MSG_SIZE); +static constexpr ptrdiff_t MAILBOX_OFF_PROTOCOL = MAILBOX_OFF_ERROR_MSG - static_cast(sizeof(uint64_t)); +static constexpr ptrdiff_t MAILBOX_OFF_GENERATION = MAILBOX_OFF_CALLABLE; static constexpr ptrdiff_t MAILBOX_OFF_TASK_CALLABLE_HASH = MAILBOX_OFF_ARGS; static constexpr ptrdiff_t MAILBOX_OFF_TASK_ARGS_BLOB = MAILBOX_OFF_TASK_CALLABLE_HASH + static_cast(CALLABLE_HASH_DIGEST_SIZE); static constexpr size_t CTRL_SHM_NAME_BYTES = 32; static constexpr ptrdiff_t MAILBOX_OFF_CONTROL_CALLABLE_HASH = MAILBOX_OFF_ARGS + static_cast(CTRL_SHM_NAME_BYTES); -static constexpr size_t MAILBOX_ARGS_CAPACITY = - MAILBOX_SIZE - static_cast(MAILBOX_OFF_TASK_ARGS_BLOB) - MAILBOX_ERROR_MSG_SIZE; +static constexpr size_t MAILBOX_ARGS_CAPACITY = static_cast(MAILBOX_OFF_PROTOCOL - MAILBOX_OFF_TASK_ARGS_BLOB); static_assert( MAILBOX_ARGS_CAPACITY >= TASK_ARGS_BLOB_HEADER_SIZE + static_cast(CHIP_MAX_TENSOR_ARGS) * sizeof(Tensor) + static_cast(CHIP_MAX_SCALAR_ARGS) * sizeof(uint64_t), @@ -193,6 +203,7 @@ struct WorkerEndpointCaps { bool remote{false}; bool supports_task_dispatch{true}; bool supports_control{true}; + uint32_t max_in_flight{1}; std::string transport{"local-mailbox"}; }; @@ -249,7 +260,7 @@ class WorkerEndpoint { class LocalMailboxEndpoint : public WorkerEndpoint { public: - LocalMailboxEndpoint(int32_t worker_id, void *mailbox); + LocalMailboxEndpoint(int32_t worker_id, void *mailbox, uint32_t task_slot_count = 1); const WorkerEndpointCaps &caps() const override { return caps_; } WorkerCompletion run(Ring *ring, const WorkerDispatch &dispatch) override; @@ -300,12 +311,21 @@ class LocalMailboxEndpoint : public WorkerEndpoint { private: WorkerEndpointCaps caps_; void *mailbox_{nullptr}; - std::mutex mailbox_mu_; + std::shared_mutex mailbox_mu_; + std::mutex free_slot_mu_; + std::condition_variable free_slot_cv_; + std::queue free_task_slots_; + std::array, MAILBOX_TASK_SLOT_COUNT> generations_{}; bool mailbox_control_timed_out_{false}; - char *mbox() const { return static_cast(mailbox_); } - MailboxState read_mailbox_state() const; - void write_mailbox_state(MailboxState s); + char *mbox(size_t task_slot = 0) const { + return static_cast(mailbox_) + task_slot * MAILBOX_TASK_SLOT_SIZE; + } + MailboxState read_mailbox_state(const char *frame) const; + void write_mailbox_state(char *frame, MailboxState s); + size_t acquire_task_slot(); + void release_task_slot(size_t task_slot); + WorkerCompletion run_task_slot(Ring *ring, const WorkerDispatch &dispatch, size_t task_slot); void run_control_command(const char *op_name, double timeout_s = -1.0); }; @@ -356,7 +376,9 @@ class WorkerThread { void dispatch(WorkerDispatch d); // True if the worker has no active task. - bool idle() const { return idle_.load(std::memory_order_acquire); } + bool idle() const { return available_.load(std::memory_order_acquire) > 0; } + bool busy() const { return available_.load(std::memory_order_acquire) < capacity_; } + uint32_t capacity() const { return capacity_; } const WorkerEndpointCaps &caps() const; int32_t worker_id() const; @@ -441,12 +463,13 @@ class WorkerThread { std::unique_ptr endpoint_; std::function on_complete_; - std::thread thread_; + std::vector threads_; std::queue queue_; std::mutex mu_; std::condition_variable cv_; bool shutdown_{false}; - std::atomic idle_{true}; + uint32_t capacity_{1}; + std::atomic available_{1}; void loop(); WorkerCompletion dispatch_process(WorkerDispatch d); @@ -463,8 +486,8 @@ class WorkerManager { // Register a worker. `mailbox` is a MAILBOX_SIZE-byte MAP_SHARED // region; the real worker (a `ChipWorker` for NEXT_LEVEL, a Python // callable for SUB) lives in the forked child. - void add_next_level(void *mailbox); - void add_next_level_at(int32_t worker_id, void *mailbox); + void add_next_level(void *mailbox, uint32_t max_in_flight = 1); + void add_next_level_at(int32_t worker_id, void *mailbox, uint32_t max_in_flight = 1); void add_next_level_endpoint(std::unique_ptr endpoint); void add_sub(void *mailbox); @@ -554,6 +577,7 @@ class WorkerManager { struct LocalNextLevelEntry { int32_t worker_id{-1}; void *mailbox{nullptr}; + uint32_t max_in_flight{1}; }; std::vector next_level_entries_; std::vector sub_entries_; diff --git a/src/common/log/include/common/strace.h b/src/common/log/include/common/strace.h index acfa96a157..58b9fe15bf 100644 --- a/src/common/log/include/common/strace.h +++ b/src/common/log/include/common/strace.h @@ -155,6 +155,12 @@ class StraceScope { } /** Set the callable hash for spans emitted on this thread. */ static void set_hid(uint64_t h) { hid() = h; } + /** Propagate an invocation onto an auxiliary worker thread. */ + static void set_context(unsigned invocation, uint64_t h) { + inv() = invocation; + hid() = h; + depth() = 0; + } /** Current invocation id / callable hash for this thread (for emit_span_at). */ static unsigned current_inv() { return inv(); } static uint64_t current_hid() { return hid(); } @@ -202,9 +208,14 @@ emit_span_at(const char *name, long long ts_ns, long long dur_ns, int depth, con #define STRACE_NEW_INV() ::simpler::strace::StraceScope::next_inv() /** Set the callable hash for subsequent spans on this thread. */ #define STRACE_SET_HID(h) ::simpler::strace::StraceScope::set_hid(h) +/** Propagate the active invocation/hash onto an auxiliary thread. */ +#define STRACE_SET_CONTEXT(inv, h) ::simpler::strace::StraceScope::set_context((inv), (h)) /** Emit a device-domain span (device-clock start `ts_ns` + measured `dur_ns`). */ #define STRACE_DEV_SPAN_AT(name, ts_ns, dur_ns, depth) \ ::simpler::strace::emit_span_at((name), (ts_ns), (dur_ns), (depth)) +/** Emit a host steady-clock span measured outside an RAII scope. */ +#define STRACE_HOST_SPAN_AT(name, ts_ns, dur_ns, depth, attrs) \ + ::simpler::strace::emit_span_at((name), (ts_ns), (dur_ns), (depth), (attrs)) #else // !SIMPLER_HOST_STRACE @@ -212,7 +223,9 @@ emit_span_at(const char *name, long long ts_ns, long long dur_ns, int depth, con #define STRACE_A(name, attrs) ((void)0) #define STRACE_NEW_INV() ((void)0) #define STRACE_SET_HID(h) ((void)0) +#define STRACE_SET_CONTEXT(inv, h) ((void)0) #define STRACE_DEV_SPAN_AT(name, ts_ns, dur_ns, depth) ((void)0) +#define STRACE_HOST_SPAN_AT(name, ts_ns, dur_ns, depth, attrs) ((void)0) #endif // SIMPLER_HOST_STRACE diff --git a/src/common/platform/include/common/host_api.h b/src/common/platform/include/common/host_api.h index e4ff984849..c42c54ab98 100644 --- a/src/common/platform/include/common/host_api.h +++ b/src/common/platform/include/common/host_api.h @@ -28,10 +28,22 @@ * the field set stays defined in exactly one place. */ struct HostApi { + // Async Host orchestration must inherit both the runner TLS and the + // backend's per-thread device binding before using callbacks below. + void *(*capture_thread_context)(); + int (*bind_thread_context)(void *context); + void (*unbind_thread_context)(); void *(*device_malloc)(size_t size); void (*device_free)(void *dev_ptr); int (*copy_to_device)(void *dev_ptr, const void *host_ptr, size_t size); int (*copy_from_device)(void *host_ptr, const void *dev_ptr, size_t size); + // Publish and observe 64-bit device control words with release/acquire + // semantics. Onboard backends use synchronous device copies. Sim backends + // must use atomic operations because Host and Device are concurrent threads + // sharing one address space; a plain memcpy does not form a publication + // edge on weakly ordered hosts. + int (*store_u64_release_to_device)(void *dev_ptr, uint64_t value); + int (*load_u64_acquire_from_device)(uint64_t *value, const void *dev_ptr); // Map a device buffer into host address space and return a host-readable VA // (nullptr on failure); the paired unregister releases it. The returned VA // may differ from dev_ptr, so callers must use it, not dev_ptr, for host diff --git a/src/common/platform/onboard/host/c_api_shared.cpp b/src/common/platform/onboard/host/c_api_shared.cpp index 7fe5334066..56830eb309 100644 --- a/src/common/platform/onboard/host/c_api_shared.cpp +++ b/src/common/platform/onboard/host/c_api_shared.cpp @@ -58,6 +58,7 @@ extern "C" { * =========================================================================== */ int register_callable_impl(const ChipCallable *callable, uint64_t (*upload_fn)(const void *), CallableArtifacts *out); int validate_runtime_impl(Runtime *runtime, const HostApi *api, int execution_rc); +int release_async_host_graph_pipeline(Runtime *runtime); /* =========================================================================== * Per-thread DeviceRunnerBase binding (set by simpler_register_callable / simpler_run) @@ -69,6 +70,20 @@ static void create_runner_key() { pthread_key_create(&g_runner_key, nullptr); } static DeviceRunnerBase *current_runner() { return static_cast(pthread_getspecific(g_runner_key)); } +static void *capture_thread_context() { return current_runner(); } + +static int bind_thread_context(void *context) { + auto *runner = static_cast(context); + if (runner == nullptr) return -1; + pthread_once(&g_runner_key_once, create_runner_key); + pthread_setspecific(g_runner_key, runner); + int rc = runner->attach_current_thread(runner->device_id()); + if (rc != 0) pthread_setspecific(g_runner_key, nullptr); + return rc; +} + +static void unbind_thread_context() { pthread_setspecific(g_runner_key, nullptr); } + /* =========================================================================== * Internal device-memory functions (wired into a HostApi and passed to the * runtime impls, NOT dlsym'd) @@ -107,6 +122,24 @@ static int copy_from_device(void *host_ptr, const void *dev_ptr, size_t size) { } } +static int store_u64_release_to_device(void *dev_ptr, uint64_t value) { + if (dev_ptr == nullptr) return -1; + try { + return current_runner()->copy_to_device(dev_ptr, &value, sizeof(value)); + } catch (...) { + return -1; + } +} + +static int load_u64_acquire_from_device(uint64_t *value, const void *dev_ptr) { + if (value == nullptr || dev_ptr == nullptr) return -1; + try { + return current_runner()->copy_from_device(value, dev_ptr, sizeof(*value)); + } catch (...) { + return -1; + } +} + static void *register_device_memory_to_host(void *dev_ptr, size_t bytes) { try { return current_runner()->register_device_memory_to_host(dev_ptr, bytes); @@ -222,10 +255,15 @@ extern "C" int prewarm_config_impl( // time rather than reassembling the pointer table on each simpler_run. Passed by // address into bind_callable_to_runtime_impl / validate_runtime_impl. static const HostApi g_host_api = { + .capture_thread_context = capture_thread_context, + .bind_thread_context = bind_thread_context, + .unbind_thread_context = unbind_thread_context, .device_malloc = device_malloc, .device_free = device_free, .copy_to_device = copy_to_device, .copy_from_device = copy_from_device, + .store_u64_release_to_device = store_u64_release_to_device, + .load_u64_acquire_from_device = load_u64_acquire_from_device, .register_device_memory_to_host = register_device_memory_to_host, .unregister_device_memory_from_host = unregister_device_memory_from_host, .device_memset = device_memset, @@ -633,6 +671,33 @@ int simpler_run( } } +int select_arena_bank_ctx(DeviceContextHandle ctx, unsigned arena_bank) { + if (ctx == NULL) return -1; + try { + return static_cast(ctx)->select_arena_bank(arena_bank); + } catch (...) { + return -1; + } +} + +int select_pipeline_slot_ctx(DeviceContextHandle ctx, unsigned pipeline_slot) { + if (ctx == NULL) return -1; + try { + return static_cast(ctx)->select_pipeline_slot(pipeline_slot); + } catch (...) { + return -1; + } +} + +int set_task_accepted_state_ctx(DeviceContextHandle ctx, volatile int32_t *state, int32_t accepted_value) { + if (ctx == NULL) return -1; + try { + return static_cast(ctx)->set_task_accepted_state(state, accepted_value); + } catch (...) { + return -1; + } +} + int simpler_unregister_callable(DeviceContextHandle ctx, int32_t callable_id) { if (ctx == NULL) return -1; try { diff --git a/src/common/platform/onboard/host/device_runner_base.cpp b/src/common/platform/onboard/host/device_runner_base.cpp index dc89003253..042e4c0d54 100644 --- a/src/common/platform/onboard/host/device_runner_base.cpp +++ b/src/common/platform/onboard/host/device_runner_base.cpp @@ -27,6 +27,7 @@ #include #include #include +#include #include #include @@ -63,6 +64,10 @@ extern "C" const char *const *runtime_extra_aicpu_symbols(size_t *count); namespace { +thread_local unsigned g_pipeline_slot = 0; +thread_local volatile int32_t *g_task_accepted_state = nullptr; +thread_local int32_t g_task_accepted_value = 0; + HostRuntimeTimeoutConfig resolve_onboard_timeout_config() { RuntimeTimeoutConfig order_defaults{ PLATFORM_OP_EXECUTE_TIMEOUT_US, PLATFORM_STREAM_SYNC_TIMEOUT_MS, PLATFORM_ONBOARD_SCHEDULER_TIMEOUT_MS @@ -118,12 +123,72 @@ HostRuntimeTimeoutConfig resolve_onboard_timeout_config() { return HostRuntimeTimeoutConfig{cfg.op_execute_timeout_us, cfg.stream_sync_timeout_ms, scheduler_override}; } +pthread_key_t g_arena_bank_key; +pthread_once_t g_arena_bank_once = PTHREAD_ONCE_INIT; +bool g_arena_bank_key_ready = false; + +void create_arena_bank_key() { + if (pthread_key_create(&g_arena_bank_key, nullptr) == 0) { + g_arena_bank_key_ready = true; + } +} + +unsigned current_arena_bank() { + pthread_once(&g_arena_bank_once, create_arena_bank_key); + if (!g_arena_bank_key_ready) return 0; + return static_cast(reinterpret_cast(pthread_getspecific(g_arena_bank_key))); +} + +__attribute__((destructor)) void destroy_arena_bank_key() { + if (g_arena_bank_key_ready) { + pthread_key_delete(g_arena_bank_key); + g_arena_bank_key_ready = false; + } +} + } // namespace DeviceRunnerBase::DeviceRunnerBase() : gm_heap_arena_(&arena_alloc_trampoline, &arena_free_trampoline, &mem_alloc_), gm_sm_arena_(&arena_alloc_trampoline, &arena_free_trampoline, &mem_alloc_), - runtime_arena_pool_(&arena_alloc_trampoline, &arena_free_trampoline, &mem_alloc_) {} + runtime_arena_pool_(&arena_alloc_trampoline, &arena_free_trampoline, &mem_alloc_), + gm_heap_arena_bank1_(&arena_alloc_trampoline, &arena_free_trampoline, &mem_alloc_), + gm_sm_arena_bank1_(&arena_alloc_trampoline, &arena_free_trampoline, &mem_alloc_), + runtime_arena_pool_bank1_(&arena_alloc_trampoline, &arena_free_trampoline, &mem_alloc_) {} + +int DeviceRunnerBase::select_arena_bank(unsigned bank) { + if (bank > 1) { + LOG_ERROR("arena bank %u is outside [0, 2)", bank); + return -1; + } + pthread_once(&g_arena_bank_once, create_arena_bank_key); + if (!g_arena_bank_key_ready || pthread_setspecific(g_arena_bank_key, reinterpret_cast(bank)) != 0) { + LOG_ERROR("failed to select arena bank %u for the current thread", bank); + return -1; + } + return 0; +} + +int DeviceRunnerBase::select_pipeline_slot(unsigned slot) { + if (slot > 1) { + LOG_ERROR("pipeline slot %u is outside [0, 2)", slot); + return -1; + } + g_pipeline_slot = slot; + return 0; +} + +int DeviceRunnerBase::set_task_accepted_state(volatile int32_t *state, int32_t accepted_value) { + g_task_accepted_state = state; + g_task_accepted_value = accepted_value; + return 0; +} + +void DeviceRunnerBase::publish_task_accepted() { + if (g_task_accepted_state != nullptr) { + __atomic_store_n(g_task_accepted_state, g_task_accepted_value, __ATOMIC_RELEASE); + } +} void *DeviceRunnerBase::allocate_tensor(std::size_t bytes) { return mem_alloc_.alloc(bytes); } @@ -146,44 +211,50 @@ int DeviceRunnerBase::device_memset(void *dev_ptr, int value, std::size_t bytes) } void DeviceRunnerBase::get_retained_temp_buffer(void **addr, size_t *size) { - if (addr != nullptr) *addr = retained_temp_addr_; - if (size != nullptr) *size = retained_temp_size_; + if (addr != nullptr) *addr = retained_temp_addrs_[g_pipeline_slot]; + if (size != nullptr) *size = retained_temp_sizes_[g_pipeline_slot]; } void DeviceRunnerBase::set_retained_temp_buffer(void *addr, size_t size) { - retained_temp_addr_ = addr; - retained_temp_size_ = size; + retained_temp_addrs_[g_pipeline_slot] = addr; + retained_temp_sizes_[g_pipeline_slot] = size; } void DeviceRunnerBase::clear_temporary_buffer() { - if (retained_temp_addr_ != nullptr) { - mem_alloc_.free(retained_temp_addr_); - retained_temp_addr_ = nullptr; - retained_temp_size_ = 0; + for (size_t slot = 0; slot < retained_temp_addrs_.size(); ++slot) { + if (retained_temp_addrs_[slot] != nullptr) { + mem_alloc_.free(retained_temp_addrs_[slot]); + retained_temp_addrs_[slot] = nullptr; + retained_temp_sizes_[slot] = 0; + } } } void *DeviceRunnerBase::acquire_pooled_gm_heap() { - if (!gm_heap_arena_.is_committed()) return nullptr; - return gm_heap_arena_.base(); + DeviceArena &arena = current_arena_bank() == 0 ? gm_heap_arena_ : gm_heap_arena_bank1_; + if (!arena.is_committed()) return nullptr; + return arena.base(); } void *DeviceRunnerBase::acquire_pooled_gm_sm() { - if (!gm_sm_arena_.is_committed()) return nullptr; - return gm_sm_arena_.base(); + DeviceArena &arena = current_arena_bank() == 0 ? gm_sm_arena_ : gm_sm_arena_bank1_; + if (!arena.is_committed()) return nullptr; + return arena.base(); } void *DeviceRunnerBase::acquire_pooled_runtime_arena() { // hbg calls setup_static_arena(...,0) and leaves runtime_arena_pool_ // uncommitted — fail loudly if a caller asks for it anyway. - if (!runtime_arena_pool_.is_committed()) return nullptr; - return runtime_arena_pool_.base(); + DeviceArena &arena = current_arena_bank() == 0 ? runtime_arena_pool_ : runtime_arena_pool_bank1_; + if (!arena.is_committed()) return nullptr; + return arena.base(); } bool DeviceRunnerBase::lookup_prebuilt_runtime_arena_cache( uint64_t hash, const void *key_data, size_t key_size, void **gm_heap_base, void **sm_base, void **runtime_arena_base, size_t *runtime_off, const void **image_data, size_t *image_size ) const { + if (current_arena_bank() != 0) return false; if (!prebuilt_runtime_arena_cache_valid_ || prebuilt_runtime_arena_cache_hash_ != hash || prebuilt_runtime_arena_cache_key_.size() != key_size || key_data == nullptr || gm_heap_base == nullptr || sm_base == nullptr || runtime_arena_base == nullptr || runtime_off == nullptr || image_data == nullptr || @@ -206,6 +277,7 @@ void DeviceRunnerBase::mark_prebuilt_runtime_arena_cached( uint64_t hash, const void *key_data, size_t key_size, void *gm_heap_base, void *sm_base, void *runtime_arena_base, size_t runtime_off, const void *image_data, size_t image_size ) { + if (current_arena_bank() != 0) return; prebuilt_runtime_arena_cache_valid_ = false; prebuilt_runtime_arena_cache_hash_ = hash; prebuilt_runtime_arena_cache_key_.assign( @@ -232,6 +304,14 @@ int DeviceRunnerBase::setup_static_arena(size_t gm_heap_size, size_t gm_sm_size, // worker's lifetime). If a caller asks for a larger layout on any // region, redo just that region — already-committed peers stay alive // so their callers don't have to re-acquire. + const unsigned arena_bank = current_arena_bank(); + DeviceArena &gm_heap_arena = arena_bank == 0 ? gm_heap_arena_ : gm_heap_arena_bank1_; + DeviceArena &gm_sm_arena = arena_bank == 0 ? gm_sm_arena_ : gm_sm_arena_bank1_; + DeviceArena &runtime_arena_pool = arena_bank == 0 ? runtime_arena_pool_ : runtime_arena_pool_bank1_; + size_t &cached_gm_heap_size = arena_bank == 0 ? cached_gm_heap_size_ : cached_gm_heap_size_bank1_; + size_t &cached_gm_sm_size = arena_bank == 0 ? cached_gm_sm_size_ : cached_gm_sm_size_bank1_; + size_t &cached_runtime_arena_size = arena_bank == 0 ? cached_runtime_arena_size_ : cached_runtime_arena_size_bank1_; + bool arena_changed = false; auto commit_region = [&arena_changed](DeviceArena &arena, size_t &cached_size, size_t requested_size) -> int { if (requested_size == 0) { @@ -269,25 +349,27 @@ int DeviceRunnerBase::setup_static_arena(size_t gm_heap_size, size_t gm_sm_size, // asking for a larger layout) fails midway, defeating the // "failure means failure" guarantee. Reset everything to the // post-construction state so the caller can retry with a new layout. - bool ok = commit_region(gm_heap_arena_, cached_gm_heap_size_, gm_heap_size) == 0; - ok = ok && commit_region(gm_sm_arena_, cached_gm_sm_size_, gm_sm_size) == 0; - ok = ok && commit_region(runtime_arena_pool_, cached_runtime_arena_size_, runtime_arena_size) == 0; + bool ok = commit_region(gm_heap_arena, cached_gm_heap_size, gm_heap_size) == 0; + ok = ok && commit_region(gm_sm_arena, cached_gm_sm_size, gm_sm_size) == 0; + ok = ok && commit_region(runtime_arena_pool, cached_runtime_arena_size, runtime_arena_size) == 0; if (!ok) { - gm_heap_arena_.release(); - gm_sm_arena_.release(); - runtime_arena_pool_.release(); - cached_gm_heap_size_ = 0; - cached_gm_sm_size_ = 0; - cached_runtime_arena_size_ = 0; - prebuilt_runtime_arena_cache_valid_ = false; - prebuilt_runtime_arena_cache_key_.clear(); - prebuilt_runtime_arena_cache_gm_heap_base_ = nullptr; - prebuilt_runtime_arena_cache_sm_base_ = nullptr; - prebuilt_runtime_arena_cache_runtime_arena_base_ = nullptr; - prebuilt_runtime_arena_cache_image_.clear(); + gm_heap_arena.release(); + gm_sm_arena.release(); + runtime_arena_pool.release(); + cached_gm_heap_size = 0; + cached_gm_sm_size = 0; + cached_runtime_arena_size = 0; + if (arena_bank == 0) { + prebuilt_runtime_arena_cache_valid_ = false; + prebuilt_runtime_arena_cache_key_.clear(); + prebuilt_runtime_arena_cache_gm_heap_base_ = nullptr; + prebuilt_runtime_arena_cache_sm_base_ = nullptr; + prebuilt_runtime_arena_cache_runtime_arena_base_ = nullptr; + prebuilt_runtime_arena_cache_image_.clear(); + } return -1; } - if (arena_changed) { + if (arena_changed && arena_bank == 0) { prebuilt_runtime_arena_cache_valid_ = false; prebuilt_runtime_arena_cache_key_.clear(); prebuilt_runtime_arena_cache_gm_heap_base_ = nullptr; @@ -362,34 +444,30 @@ int DeviceRunnerBase::ensure_device_initialized() { return rc; } - bool aicpu_created_here = false; - bool aicore_created_here = false; - if (stream_aicpu_ == nullptr) { - rc = rtStreamCreate(&stream_aicpu_, 0); - if (rc != 0) { - LOG_ERROR("rtStreamCreate (AICPU) failed: %d", rc); - ACL_LOG_ERROR_DETAIL(rc); - return rc; + std::vector created_here; + auto ensure_stream = [&](rtStream_t *stream, const char *name) -> int { + if (*stream != nullptr) return 0; + int create_rc = rtStreamCreate(stream, 0); + if (create_rc != 0) { + LOG_ERROR("rtStreamCreate (%s) failed: %d", name, create_rc); + ACL_LOG_ERROR_DETAIL(create_rc); + return create_rc; } - aicpu_created_here = true; - } - if (stream_aicore_ == nullptr) { - rc = rtStreamCreate(&stream_aicore_, 0); - if (rc != 0) { - LOG_ERROR("rtStreamCreate (AICore) failed: %d", rc); - ACL_LOG_ERROR_DETAIL(rc); - // Roll back only the AICPU stream we just created, not a - // pre-existing persistent one. - if (aicpu_created_here) { - rtStreamDestroy(stream_aicpu_); - stream_aicpu_ = nullptr; - } - return rc; + created_here.push_back(stream); + return 0; + }; + if ((rc = ensure_stream(&stream_aicpu_, "AICPU slot 0")) != 0 || + (rc = ensure_stream(&stream_aicore_, "AICore slot 0")) != 0 || + (rc = ensure_stream(&stream_aicpu_bank1_, "AICPU slot 1")) != 0 || + (rc = ensure_stream(&stream_aicore_bank1_, "AICore slot 1")) != 0) { + for (auto it = created_here.rbegin(); it != created_here.rend(); ++it) { + rtStreamDestroy(**it); + **it = nullptr; } - aicore_created_here = true; + return rc; } - if (aicpu_created_here || aicore_created_here) { - LOG_INFO_V0("DeviceRunner: device=%d set, streams created", device_id_); + if (!created_here.empty()) { + LOG_INFO_V0("DeviceRunner: device=%d set, two stream sets created", device_id_); } rc = ensure_binaries_loaded(); @@ -1023,6 +1101,14 @@ int DeviceRunnerBase::finalize_common() { capture(rtStreamDestroy(stream_aicore_)); stream_aicore_ = nullptr; } + if (stream_aicpu_bank1_ != nullptr) { + capture(rtStreamDestroy(stream_aicpu_bank1_)); + stream_aicpu_bank1_ = nullptr; + } + if (stream_aicore_bank1_ != nullptr) { + capture(rtStreamDestroy(stream_aicore_bank1_)); + stream_aicore_bank1_ = nullptr; + } // Release the async-DMA provider (SDMA STARS streams + workspace) while RTS // is live, before the subclass device reset. Null unless the Worker was @@ -1077,6 +1163,9 @@ int DeviceRunnerBase::finalize_common() { gm_heap_arena_.release(); gm_sm_arena_.release(); runtime_arena_pool_.release(); + gm_heap_arena_bank1_.release(); + gm_sm_arena_bank1_.release(); + runtime_arena_pool_bank1_.release(); prebuilt_runtime_arena_cache_valid_ = false; prebuilt_runtime_arena_cache_key_.clear(); prebuilt_runtime_arena_cache_gm_heap_base_ = nullptr; @@ -1104,6 +1193,9 @@ int DeviceRunnerBase::finalize_common() { cached_gm_heap_size_ = 0; cached_gm_sm_size_ = 0; cached_runtime_arena_size_ = 0; + cached_gm_heap_size_bank1_ = 0; + cached_gm_sm_size_bank1_ = 0; + cached_runtime_arena_size_bank1_ = 0; return rc; } @@ -1213,14 +1305,14 @@ int DeviceRunnerBase::resolve_block_dim(int requested_block_dim) { // auto branch skips validate so we don't pay the ACL syscalls twice. int resolved = requested_block_dim; if (resolved == 0) { - resolved = query_max_block_dim(stream_aicore_); + resolved = query_max_block_dim(run_stream_aicore()); LOG_INFO_V0("block_dim auto-resolved to %d", resolved); if (resolved < 1) { LOG_ERROR("block_dim auto-resolved to invalid value %d", resolved); return -1; } } else { - int rc = validate_block_dim(stream_aicore_, resolved); + int rc = validate_block_dim(run_stream_aicore(), resolved); if (rc != 0) { return -1; } @@ -1268,8 +1360,10 @@ int DeviceRunnerBase::prepare_runtime_for_launch(Runtime &runtime, int block_dim } int DeviceRunnerBase::sync_run_streams() { - LOG_INFO_V0("=== aclrtSynchronizeStreamWithTimeout stream_aicpu_ ==="); - int rc = aclrtSynchronizeStreamWithTimeout(stream_aicpu_, timeout_config_.stream_sync_timeout_ms); + rtStream_t aicpu_stream = run_stream_aicpu(); + rtStream_t aicore_stream = run_stream_aicore(); + LOG_INFO_V0("=== aclrtSynchronizeStreamWithTimeout AICPU pipeline slot %u ===", g_pipeline_slot); + int rc = aclrtSynchronizeStreamWithTimeout(aicpu_stream, timeout_config_.stream_sync_timeout_ms); if (rc == ACL_ERROR_RT_STREAM_SYNC_TIMEOUT) { LOG_ERROR( "Stream sync timeout: stream=AICPU timeout_ms=%d device_id=%d block_dim=%d", @@ -1284,8 +1378,8 @@ int DeviceRunnerBase::sync_run_streams() { return rc; } - LOG_INFO_V0("=== aclrtSynchronizeStreamWithTimeout stream_aicore_ ==="); - rc = aclrtSynchronizeStreamWithTimeout(stream_aicore_, timeout_config_.stream_sync_timeout_ms); + LOG_INFO_V0("=== aclrtSynchronizeStreamWithTimeout AICore pipeline slot %u ===", g_pipeline_slot); + rc = aclrtSynchronizeStreamWithTimeout(aicore_stream, timeout_config_.stream_sync_timeout_ms); if (rc == ACL_ERROR_RT_STREAM_SYNC_TIMEOUT) { LOG_ERROR( "Stream sync timeout: stream=AICore timeout_ms=%d device_id=%d block_dim=%d", @@ -1302,6 +1396,14 @@ int DeviceRunnerBase::sync_run_streams() { return 0; } +rtStream_t DeviceRunnerBase::run_stream_aicpu() const { + return g_pipeline_slot == 0 ? stream_aicpu_ : stream_aicpu_bank1_; +} + +rtStream_t DeviceRunnerBase::run_stream_aicore() const { + return g_pipeline_slot == 0 ? stream_aicore_ : stream_aicore_bank1_; +} + void DeviceRunnerBase::read_device_wall_ns() { // Pull the per-thread AICPU phase records back from the device buffer that // AICPU writes through via KernelArgs::device_wall_data_base. (We can't use diff --git a/src/common/platform/onboard/host/device_runner_base.h b/src/common/platform/onboard/host/device_runner_base.h index 9260f9a18c..50c4ca5033 100644 --- a/src/common/platform/onboard/host/device_runner_base.h +++ b/src/common/platform/onboard/host/device_runner_base.h @@ -38,6 +38,7 @@ #include +#include #include #include #include @@ -134,6 +135,16 @@ class DeviceRunnerBase { */ int setup_static_arena(size_t gm_heap_size, size_t gm_sm_size, size_t runtime_arena_size); + /** Select the pooled-arena bank used by HostApi calls on this thread. */ + int select_arena_bank(unsigned bank); + + /** Select this caller thread's per-run pipeline resource slot. */ + int select_pipeline_slot(unsigned slot); + + /** Bind and publish the mailbox ACK associated with this caller thread. */ + int set_task_accepted_state(volatile int32_t *state, int32_t accepted_value); + void publish_task_accepted(); + /** * Return the pooled GM heap / PTO2 SM / runtime arena base pointer. * `setup_static_arena` (arch subclass) must have already committed @@ -676,6 +687,9 @@ class DeviceRunnerBase { */ int sync_run_streams(); + rtStream_t run_stream_aicpu() const; + rtStream_t run_stream_aicore() const; + /** * Pull the device wall (ns) back from `device_wall_dev_ptr_` and * cache it on `device_wall_ns_`. D2H copy failure is a soft warn — @@ -870,14 +884,16 @@ class DeviceRunnerBase { host::LoadAicpuOp load_aicpu_op_; MemoryAllocator mem_alloc_; - // Retained temporary buffer slot for TRB device-arg staging (see HostApi - // get/set_retained_temp_buffer). Just a remembered {addr, size} reused - // across runs and freed in finalize; the grow/pack logic lives in trb bind. - void *retained_temp_addr_ = nullptr; - std::size_t retained_temp_size_ = 0; + // Two retained TaskArg staging buffers. TMR selects one by pipeline slot; + // its large arena remains independently single-banked REUSE_MEM. + std::array retained_temp_addrs_{{nullptr, nullptr}}; + std::array retained_temp_sizes_{{0, 0}}; DeviceArena gm_heap_arena_; DeviceArena gm_sm_arena_; DeviceArena runtime_arena_pool_; + DeviceArena gm_heap_arena_bank1_; + DeviceArena gm_sm_arena_bank1_; + DeviceArena runtime_arena_pool_bank1_; // Cached arena sizes for `setup_static_arena`'s "fits" check — avoids // re-allocating the same buffer when a later worker init asks for an @@ -886,6 +902,9 @@ class DeviceRunnerBase { size_t cached_gm_heap_size_{0}; size_t cached_gm_sm_size_{0}; size_t cached_runtime_arena_size_{0}; + size_t cached_gm_heap_size_bank1_{0}; + size_t cached_gm_sm_size_bank1_{0}; + size_t cached_runtime_arena_size_bank1_{0}; bool prebuilt_runtime_arena_cache_valid_{false}; uint64_t prebuilt_runtime_arena_cache_hash_{0}; std::vector prebuilt_runtime_arena_cache_key_; @@ -900,6 +919,8 @@ class DeviceRunnerBase { // `finalize()`. `nullptr` before init. rtStream_t stream_aicpu_{nullptr}; rtStream_t stream_aicore_{nullptr}; + rtStream_t stream_aicpu_bank1_{nullptr}; + rtStream_t stream_aicore_bank1_{nullptr}; KernelArgsHelper kernel_args_; // Platform-level device phase buffer: device-resident diff --git a/src/common/platform/sim/host/c_api_shared.cpp b/src/common/platform/sim/host/c_api_shared.cpp index 4704006882..4621982d72 100644 --- a/src/common/platform/sim/host/c_api_shared.cpp +++ b/src/common/platform/sim/host/c_api_shared.cpp @@ -32,6 +32,7 @@ #include #include +#include #include #include #include @@ -65,6 +66,23 @@ static SimDeviceRunnerBase *current_runner() { return static_cast(pthread_getspecific(g_runner_key)); } +static void *capture_thread_context() { return current_runner(); } + +static int bind_thread_context(void *context) { + auto *runner = static_cast(context); + if (runner == nullptr) return -1; + pthread_once(&g_runner_key_once, create_runner_key); + pthread_setspecific(g_runner_key, runner); + int rc = runner->attach_current_thread(runner->device_id()); + if (rc != 0) pthread_setspecific(g_runner_key, nullptr); + return rc; +} + +static void unbind_thread_context() { + pto_cpu_sim_bind_device(-1); + pthread_setspecific(g_runner_key, nullptr); +} + /* =========================================================================== * Internal device-memory functions (wired into a HostApi and passed to the * runtime impls, NOT dlsym'd) @@ -103,6 +121,18 @@ static int copy_from_device(void *host_ptr, const void *dev_ptr, size_t size) { } } +static int store_u64_release_to_device(void *dev_ptr, uint64_t value) { + if (dev_ptr == nullptr) return -1; + reinterpret_cast *>(dev_ptr)->store(value, std::memory_order_release); + return 0; +} + +static int load_u64_acquire_from_device(uint64_t *value, const void *dev_ptr) { + if (value == nullptr || dev_ptr == nullptr) return -1; + *value = reinterpret_cast *>(dev_ptr)->load(std::memory_order_acquire); + return 0; +} + static void *register_device_memory_to_host(void *dev_ptr, size_t bytes) { try { return current_runner()->register_device_memory_to_host(dev_ptr, bytes); @@ -219,10 +249,15 @@ extern "C" int prewarm_config_impl( ); static const HostApi g_host_api = { + .capture_thread_context = capture_thread_context, + .bind_thread_context = bind_thread_context, + .unbind_thread_context = unbind_thread_context, .device_malloc = device_malloc, .device_free = device_free, .copy_to_device = copy_to_device, .copy_from_device = copy_from_device, + .store_u64_release_to_device = store_u64_release_to_device, + .load_u64_acquire_from_device = load_u64_acquire_from_device, .register_device_memory_to_host = register_device_memory_to_host, .unregister_device_memory_from_host = unregister_device_memory_from_host, .device_memset = device_memset, @@ -279,6 +314,28 @@ int copy_from_device_ctx(DeviceContextHandle ctx, void *host_ptr, const void *de } } +int select_pipeline_slot_ctx(DeviceContextHandle ctx, unsigned pipeline_slot) { + if (ctx == NULL) return -1; + try { + return static_cast(ctx)->select_pipeline_slot(pipeline_slot); + } catch (...) { + return -1; + } +} + +int select_arena_bank_ctx(DeviceContextHandle ctx, unsigned arena_bank) { + if (ctx == NULL || arena_bank > 1) return -1; + return 0; +} + +int set_task_accepted_state_ctx(DeviceContextHandle ctx, volatile int32_t *state, int32_t accepted_value) { + // Sim has no asynchronous KernelLaunch boundary. Keep the ABI uniform; + // TASK_DONE remains its completion signal. + (void)state; + (void)accepted_value; + return ctx == NULL ? -1 : 0; +} + int finalize_device(DeviceContextHandle ctx) { if (ctx == NULL) return -1; try { diff --git a/src/common/platform/sim/host/device_runner_base.cpp b/src/common/platform/sim/host/device_runner_base.cpp index 023afc0113..87fe1ffac1 100644 --- a/src/common/platform/sim/host/device_runner_base.cpp +++ b/src/common/platform/sim/host/device_runner_base.cpp @@ -31,6 +31,10 @@ #include "task_args.h" #include "utils/elf_build_id.h" +namespace { +thread_local unsigned g_pipeline_slot = 0; +} + namespace simpler::common::sim_host { namespace { @@ -267,23 +271,31 @@ int SimDeviceRunnerBase::device_memset(void *dev_ptr, int value, size_t bytes) { } void SimDeviceRunnerBase::get_retained_temp_buffer(void **addr, size_t *size) { - if (addr != nullptr) *addr = retained_temp_addr_; - if (size != nullptr) *size = retained_temp_size_; + if (addr != nullptr) *addr = retained_temp_addrs_[g_pipeline_slot]; + if (size != nullptr) *size = retained_temp_sizes_[g_pipeline_slot]; } void SimDeviceRunnerBase::set_retained_temp_buffer(void *addr, size_t size) { - retained_temp_addr_ = addr; - retained_temp_size_ = size; + retained_temp_addrs_[g_pipeline_slot] = addr; + retained_temp_sizes_[g_pipeline_slot] = size; } void SimDeviceRunnerBase::clear_temporary_buffer() { - if (retained_temp_addr_ != nullptr) { - mem_alloc_.free(retained_temp_addr_); - retained_temp_addr_ = nullptr; - retained_temp_size_ = 0; + for (size_t slot = 0; slot < retained_temp_addrs_.size(); ++slot) { + if (retained_temp_addrs_[slot] != nullptr) { + mem_alloc_.free(retained_temp_addrs_[slot]); + retained_temp_addrs_[slot] = nullptr; + retained_temp_sizes_[slot] = 0; + } } } +int SimDeviceRunnerBase::select_pipeline_slot(unsigned slot) { + if (slot > 1) return -1; + g_pipeline_slot = slot; + return 0; +} + int SimDeviceRunnerBase::stamp_orch_so(Runtime &runtime, int32_t cid) { // Registered-callable flow only: the orch SO was already delivered to the // sim AICPU at launch_device_register time. A run just needs the active diff --git a/src/common/platform/sim/host/device_runner_base.h b/src/common/platform/sim/host/device_runner_base.h index b1ecca0ad5..5857fab671 100644 --- a/src/common/platform/sim/host/device_runner_base.h +++ b/src/common/platform/sim/host/device_runner_base.h @@ -47,6 +47,8 @@ #include "common/unified_log.h" #include "host/memory_allocator.h" #include "host/l2_swimlane_collector.h" +#include + #include "host/args_dump_collector.h" #include "host/pmu_collector.h" #include "host/scope_stats_collector.h" @@ -99,6 +101,7 @@ class SimDeviceRunnerBase { void get_retained_temp_buffer(void **addr, size_t *size); void set_retained_temp_buffer(void *addr, size_t size); void clear_temporary_buffer(); + int select_pipeline_slot(unsigned slot); // On sim, allocate_tensor returns a plain host pointer, so the "device" // address is already host-readable — register is identity, unregister a @@ -216,8 +219,8 @@ class SimDeviceRunnerBase { std::vector aicore_kernel_binary_; MemoryAllocator mem_alloc_; - void *retained_temp_addr_ = nullptr; - size_t retained_temp_size_ = 0; + std::array retained_temp_addrs_{{nullptr, nullptr}}; + std::array retained_temp_sizes_{{0, 0}}; // Three independent per-Worker arenas, each backing a single pooled // region (PTO2 GM heap / PTO2 shared memory / trb prebuilt runtime diff --git a/src/common/worker/chip_worker.cpp b/src/common/worker/chip_worker.cpp index 567eb43d6e..cab68205a9 100644 --- a/src/common/worker/chip_worker.cpp +++ b/src/common/worker/chip_worker.cpp @@ -37,6 +37,14 @@ T load_symbol(void *handle, const char *name) { return reinterpret_cast(sym); } +template +T load_optional_symbol(void *handle, const char *name) { + dlerror(); + void *sym = dlsym(handle, name); + if (dlerror() != nullptr) return nullptr; + return reinterpret_cast(sym); +} + std::vector read_binary_file(const std::string &path) { std::ifstream f(path, std::ios::binary | std::ios::ate); if (!f) { @@ -105,6 +113,11 @@ void ChipWorker::init( simpler_init_fn_ = load_symbol(handle, "simpler_init"); register_callable_fn_ = load_symbol(handle, "simpler_register_callable"); run_fn_ = load_symbol(handle, "simpler_run"); + select_arena_bank_fn_ = load_optional_symbol(handle, "select_arena_bank_ctx"); + select_pipeline_slot_fn_ = load_optional_symbol(handle, "select_pipeline_slot_ctx"); + set_task_accepted_state_fn_ = + load_optional_symbol(handle, "set_task_accepted_state_ctx"); + get_pipeline_contract_fn_ = load_optional_symbol(handle, "get_pipeline_contract"); unregister_callable_fn_ = load_symbol(handle, "simpler_unregister_callable"); get_aicpu_dlopen_count_fn_ = load_symbol(handle, "get_aicpu_dlopen_count"); get_host_dlopen_count_fn_ = load_symbol(handle, "get_host_dlopen_count"); @@ -136,6 +149,19 @@ void ChipWorker::init( lib_handle_ = handle; + if (get_pipeline_contract_fn_ != nullptr) { + const PipelineContract *contract = get_pipeline_contract_fn_(); + if (contract == nullptr || contract->abi_version != PTO_PIPELINE_CONTRACT_ABI_VERSION || + contract->resource_count > PTO_PIPELINE_MAX_RESOURCES || contract->pipeline_slots == 0 || + contract->pipeline_slots > PTO_PIPELINE_MAX_SLOTS || contract->arena_banks == 0 || + contract->arena_banks > contract->pipeline_slots) { + dlclose(handle); + lib_handle_ = nullptr; + throw std::runtime_error("host runtime returned an invalid PipelineContract"); + } + pipeline_contract_ = *contract; + } + device_ctx_ = create_device_context_fn_(); if (device_ctx_ == nullptr) { dlclose(handle); @@ -143,7 +169,10 @@ void ChipWorker::init( throw std::runtime_error("create_device_context returned null"); } - runtime_buf_.resize(get_runtime_size_fn_()); + const size_t runtime_size = get_runtime_size_fn_(); + for (auto &runtime_buf : runtime_bufs_) { + runtime_buf.resize(runtime_size); + } // One-shot platform-side init: attach the calling thread to `device_id` // (rtSetDevice on onboard, sim bind+acquire on sim), transfer ownership @@ -191,6 +220,10 @@ void ChipWorker::init( simpler_init_fn_ = nullptr; register_callable_fn_ = nullptr; run_fn_ = nullptr; + select_arena_bank_fn_ = nullptr; + select_pipeline_slot_fn_ = nullptr; + set_task_accepted_state_fn_ = nullptr; + get_pipeline_contract_fn_ = nullptr; unregister_callable_fn_ = nullptr; get_aicpu_dlopen_count_fn_ = nullptr; get_host_dlopen_count_fn_ = nullptr; @@ -207,7 +240,8 @@ void ChipWorker::init( comm_release_domain_windows_fn_ = nullptr; comm_barrier_fn_ = nullptr; comm_destroy_fn_ = nullptr; - runtime_buf_.clear(); + for (auto &runtime_buf : runtime_bufs_) + runtime_buf.clear(); throw; } if (init_rc != 0) { @@ -230,6 +264,10 @@ void ChipWorker::init( simpler_init_fn_ = nullptr; register_callable_fn_ = nullptr; run_fn_ = nullptr; + select_arena_bank_fn_ = nullptr; + select_pipeline_slot_fn_ = nullptr; + set_task_accepted_state_fn_ = nullptr; + get_pipeline_contract_fn_ = nullptr; unregister_callable_fn_ = nullptr; get_aicpu_dlopen_count_fn_ = nullptr; get_host_dlopen_count_fn_ = nullptr; @@ -247,7 +285,8 @@ void ChipWorker::init( comm_release_domain_windows_fn_ = nullptr; comm_barrier_fn_ = nullptr; comm_destroy_fn_ = nullptr; - runtime_buf_.clear(); + for (auto &runtime_buf : runtime_bufs_) + runtime_buf.clear(); throw std::runtime_error("simpler_init failed with code " + std::to_string(init_rc)); } @@ -295,6 +334,10 @@ void ChipWorker::finalize() { get_runtime_size_fn_ = nullptr; register_callable_fn_ = nullptr; run_fn_ = nullptr; + select_arena_bank_fn_ = nullptr; + select_pipeline_slot_fn_ = nullptr; + set_task_accepted_state_fn_ = nullptr; + get_pipeline_contract_fn_ = nullptr; unregister_callable_fn_ = nullptr; get_aicpu_dlopen_count_fn_ = nullptr; get_host_dlopen_count_fn_ = nullptr; @@ -312,7 +355,9 @@ void ChipWorker::finalize() { comm_release_domain_windows_fn_ = nullptr; comm_barrier_fn_ = nullptr; comm_destroy_fn_ = nullptr; - runtime_buf_.clear(); + for (auto &runtime_buf : runtime_bufs_) + runtime_buf.clear(); + pipeline_contract_ = {PTO_PIPELINE_CONTRACT_ABI_VERSION, 0, 1, 1, {}}; initialized_ = false; device_id_ = -1; finalized_ = true; @@ -332,26 +377,115 @@ void ChipWorker::register_callable(int32_t callable_id, const void *callable) { } void ChipWorker::run(int32_t callable_id, TaskArgsView args, const CallConfig &config) { + run(callable_id, args, config, 0); +} + +void ChipWorker::run(int32_t callable_id, TaskArgsView args, const CallConfig &config, unsigned pipeline_slot) { ChipStorageTaskArgs chip_storage = view_to_chip_storage(args); - run(callable_id, &chip_storage, config); + run(callable_id, &chip_storage, config, pipeline_slot); } void ChipWorker::run(int32_t callable_id, const ChipStorageTaskArgs *args, const CallConfig &config) { + run(callable_id, args, config, 0); +} + +void ChipWorker::run( + int32_t callable_id, const ChipStorageTaskArgs *args, const CallConfig &config, unsigned pipeline_slot +) { + run(callable_id, args, config, pipeline_slot, nullptr); +} + +void ChipWorker::run( + int32_t callable_id, const ChipStorageTaskArgs *args, const CallConfig &config, unsigned pipeline_slot, + volatile int32_t *accepted_state +) { config.validate(); if (!initialized_) { throw std::runtime_error("ChipWorker not initialized; call init() first"); } - void *rt = runtime_buf_.data(); + if (pipeline_slot >= pipeline_contract_.pipeline_slots) { + throw std::runtime_error("run pipeline slot is outside the runtime PipelineContract"); + } + const unsigned arena_bank = pipeline_slot % pipeline_contract_.arena_banks; + if (pipeline_slot != 0 && select_pipeline_slot_fn_ == nullptr) { + throw std::runtime_error("bound runtime does not support nonzero pipeline slots"); + } + if (arena_bank != 0 && select_arena_bank_fn_ == nullptr) { + throw std::runtime_error("bound runtime does not support nonzero arena banks"); + } + void *rt = runtime_bufs_[pipeline_slot].data(); // Per-stage timing is emitted by the platform as `[STRACE]` log markers, not // returned (see chip_worker.h::run). CallConfig is threaded through to the C // ABI as a single pointer rather than unpacked into per-field args. - int rc = run_fn_(device_ctx_, rt, callable_id, args, &config); + if (select_pipeline_slot_fn_ != nullptr && select_pipeline_slot_fn_(device_ctx_, pipeline_slot) != 0) { + throw std::runtime_error("select_pipeline_slot_ctx failed"); + } + if (select_arena_bank_fn_ != nullptr && select_arena_bank_fn_(device_ctx_, arena_bank) != 0) { + throw std::runtime_error("select_arena_bank_ctx failed"); + } + if (accepted_state != nullptr) { + if (set_task_accepted_state_fn_ == nullptr || + set_task_accepted_state_fn_(device_ctx_, accepted_state, 8) != 0) { + throw std::runtime_error("set_task_accepted_state_ctx failed"); + } + } + + const bool reuse_admitted = has_shared_reuse_resource(); + if (reuse_admitted) { + acquire_reuse_resource(config.runtime_env); + } + + int rc = 0; + try { + rc = run_fn_(device_ctx_, rt, callable_id, args, &config); + if (set_task_accepted_state_fn_ != nullptr) { + (void)set_task_accepted_state_fn_(device_ctx_, nullptr, 0); + } + } catch (...) { + if (set_task_accepted_state_fn_ != nullptr) { + (void)set_task_accepted_state_fn_(device_ctx_, nullptr, 0); + } + if (reuse_admitted) { + release_reuse_resource(); + } + throw; + } + if (reuse_admitted) { + release_reuse_resource(); + } if (rc != 0) { throw std::runtime_error("run failed with code " + std::to_string(rc)); } } +bool ChipWorker::has_shared_reuse_resource() const { + if (pipeline_contract_.arena_banks >= pipeline_contract_.pipeline_slots) return false; + for (uint32_t i = 0; i < pipeline_contract_.resource_count; ++i) { + if (pipeline_contract_.resources[i].resource_class == PTO_PIPELINE_REUSE_MEM) return true; + } + return false; +} + +void ChipWorker::acquire_reuse_resource(const RuntimeEnv &runtime_env) { + std::unique_lock lock(reuse_resource_mutex_); + reuse_resource_cv_.wait(lock, [&]() { + return active_reuse_runs_ == 0 || + std::memcmp(&active_reuse_runtime_env_, &runtime_env, sizeof(RuntimeEnv)) == 0; + }); + if (active_reuse_runs_ == 0) { + active_reuse_runtime_env_ = runtime_env; + } + ++active_reuse_runs_; +} + +void ChipWorker::release_reuse_resource() { + std::lock_guard lock(reuse_resource_mutex_); + if (--active_reuse_runs_ == 0) { + reuse_resource_cv_.notify_all(); + } +} + void ChipWorker::unregister_callable(int32_t callable_id) { if (!initialized_) { throw std::runtime_error("ChipWorker not initialized; call init() first"); diff --git a/src/common/worker/chip_worker.h b/src/common/worker/chip_worker.h index a09108143b..8d82479db3 100644 --- a/src/common/worker/chip_worker.h +++ b/src/common/worker/chip_worker.h @@ -12,7 +12,10 @@ #ifndef SRC_COMMON_WORKER_CHIP_WORKER_H_ #define SRC_COMMON_WORKER_CHIP_WORKER_H_ +#include +#include #include +#include #include #include #include @@ -69,12 +72,17 @@ class ChipWorker { // platform as `[STRACE]` log markers — see src/common/log/.../strace.h — not // returned, so the L3 dispatcher and L2 child are observed uniformly. void run(int32_t callable_id, TaskArgsView args, const CallConfig &config); + void run(int32_t callable_id, TaskArgsView args, const CallConfig &config, unsigned pipeline_slot); // Same launch, but the caller already holds the runtime.so-ABI POD — // skip the view→storage memcpy and hand the pointer straight to the C ABI. // Used by the ChipStorageTaskArgs path in the nanobind binding. void run(int32_t callable_id, const ChipStorageTaskArgs *args, const CallConfig &config); + void run(int32_t callable_id, const ChipStorageTaskArgs *args, const CallConfig &config, unsigned pipeline_slot); + void + run(int32_t callable_id, const ChipStorageTaskArgs *args, const CallConfig &config, unsigned pipeline_slot, + volatile int32_t *accepted_state); - // Per-callable_id preparation. Requires init() first and a callable_id + // Per-callable_id registration. Requires init() first and a callable_id // in [0, MAX_REGISTERED_CALLABLE_IDS) (cap 64). void register_callable(int32_t callable_id, const void *callable); void unregister_callable(int32_t callable_id); @@ -90,6 +98,9 @@ class ChipWorker { /// `aicpu_dlopen_count` for the trb path; returns 0 on device-orch variants. size_t host_dlopen_count() const; + unsigned pipeline_slot_count() const { return pipeline_contract_.pipeline_slots; } + unsigned arena_bank_count() const { return pipeline_contract_.arena_banks; } + uint64_t malloc(size_t size); void free(uint64_t ptr); void copy_to(uint64_t dst, uint64_t src, size_t size); @@ -157,6 +168,10 @@ class ChipWorker { ); using SimplerRegisterCallableFn = int (*)(void *, int32_t, const void *); using SimplerRunFn = int (*)(void *, void *, int32_t, const void *, const CallConfig *); + using SelectArenaBankFn = int (*)(void *, unsigned); + using SelectPipelineSlotFn = int (*)(void *, unsigned); + using SetTaskAcceptedStateFn = int (*)(void *, volatile int32_t *, int32_t); + using GetPipelineContractFn = const PipelineContract *(*)(); using SimplerUnregisterCallableFn = int (*)(void *, int32_t); using GetAicpuDlopenCountFn = size_t (*)(void *); using SimplerProvisionDmaWorkspaceFn = int (*)(void *, uint32_t); @@ -186,6 +201,9 @@ class ChipWorker { void *create_comm_stream_checked(const char *op_name); void destroy_comm_stream_best_effort(void *stream, int *rc); + bool has_shared_reuse_resource() const; + void acquire_reuse_resource(const RuntimeEnv &runtime_env); + void release_reuse_resource(); CommSession *find_comm_session(uint64_t comm_handle); CommSession *create_comm_session(void *handle, void *stream, bool is_base); int destroy_comm_session(CommSession &session); @@ -203,6 +221,10 @@ class ChipWorker { SimplerInitFn simpler_init_fn_ = nullptr; SimplerRegisterCallableFn register_callable_fn_ = nullptr; SimplerRunFn run_fn_ = nullptr; + SelectArenaBankFn select_arena_bank_fn_ = nullptr; + SelectPipelineSlotFn select_pipeline_slot_fn_ = nullptr; + SetTaskAcceptedStateFn set_task_accepted_state_fn_ = nullptr; + GetPipelineContractFn get_pipeline_contract_fn_ = nullptr; SimplerUnregisterCallableFn unregister_callable_fn_ = nullptr; GetAicpuDlopenCountFn get_aicpu_dlopen_count_fn_ = nullptr; GetAicpuDlopenCountFn get_host_dlopen_count_fn_ = nullptr; @@ -225,12 +247,12 @@ class ChipWorker { std::unordered_map comm_session_index_; uint64_t base_comm_handle_ = 0; - std::vector runtime_buf_; - // device_id_ is set once in init() and never modified afterward. All - // ChipWorker callers run on the thread that called init() (the same - // thread is the only one that subsequently calls malloc / copy_to / - // run / finalize), so plain `int` is sufficient — no cross-thread - // synchronization required. + std::array, 2> runtime_bufs_; + PipelineContract pipeline_contract_{PTO_PIPELINE_CONTRACT_ABI_VERSION, 0, 1, 1, {}}; + std::mutex reuse_resource_mutex_; + std::condition_variable reuse_resource_cv_; + RuntimeEnv active_reuse_runtime_env_{}; + unsigned active_reuse_runs_{0}; int device_id_ = -1; bool initialized_ = false; bool finalized_ = false; diff --git a/src/common/worker/pto_runtime_c_api.h b/src/common/worker/pto_runtime_c_api.h index 38035da609..0b90de04f7 100644 --- a/src/common/worker/pto_runtime_c_api.h +++ b/src/common/worker/pto_runtime_c_api.h @@ -15,16 +15,15 @@ * Both the ChipWorker (consumer, resolves public symbols via dlsym) and the * platform implementations (producers, define all symbols) include this file. * - * Public API — resolved by ChipWorker via dlsym (every host_runtime.so must - * export ALL of these; runtimes without a real backend ship not-supported - * stubs rather than omitting symbols, so ChipWorker can dlsym the full set - * unconditionally without per-symbol probing): + * Public API — resolved by ChipWorker via dlsym. Core lifecycle, memory, + * registration, run, communication, and stream symbols are required from + * every host_runtime.so. * - lifecycle: create_device_context, destroy_device_context, * simpler_init, finalize_device * - sizing: get_runtime_size * - device-mem: device_malloc_ctx, device_free_ctx, * copy_to_device_ctx, copy_from_device_ctx - * - prepared run: simpler_register_callable, simpler_run, unregister_callable, + * - callable/run: simpler_register_callable, simpler_run, unregister_callable, * get_aicpu_dlopen_count, get_host_dlopen_count, * simpler_provision_dma_workspace * - ACL/stream: ensure_acl_ready_ctx, create_comm_stream_ctx, @@ -61,6 +60,42 @@ enum { PTO_RUNTIME_ERR_UNSUPPORTED = -2, }; +enum { + PTO_PIPELINE_CONTRACT_ABI_VERSION = 1, + PTO_PIPELINE_MAX_RESOURCES = 8, + PTO_PIPELINE_MAX_SLOTS = 2, +}; + +typedef enum PipelineResourceClass { + PTO_PIPELINE_FILL_MEM = 0, + PTO_PIPELINE_REUSE_MEM = 1, + PTO_PIPELINE_EXEC_HANDLE = 2, +} PipelineResourceClass; + +typedef enum PipelineResourceKind { + PTO_PIPELINE_GM_HEAP = 0, + PTO_PIPELINE_GM_SM = 1, + PTO_PIPELINE_RUNTIME_IMAGE = 2, + PTO_PIPELINE_TASK_ARGS = 3, + PTO_PIPELINE_AICPU_STREAM = 4, + PTO_PIPELINE_AICORE_STREAM = 5, +} PipelineResourceKind; + +typedef struct PipelineResource { + uint32_t kind; + uint32_t resource_class; + size_t bytes_per_slot; +} PipelineResource; + +/** Runtime-owned declaration of resources crossing KernelLaunch. */ +typedef struct PipelineContract { + uint32_t abi_version; + uint32_t resource_count; + uint32_t pipeline_slots; + uint32_t arena_banks; + PipelineResource resources[PTO_PIPELINE_MAX_RESOURCES]; +} PipelineContract; + /* Per-stage run timing is no longer returned. The platform emits it as * `[STRACE]` log markers (host stages + the AICPU device-phase breakdown, * gated on SIMPLER_HOST_STRACE) — parse with simpler_setup.tools.strace_timing. @@ -70,6 +105,9 @@ enum { * Public API (resolved by ChipWorker via dlsym) * =========================================================================== */ +/** Return this runtime's immutable pipeline resource declaration. Optional. */ +const PipelineContract *get_pipeline_contract(void); + /** * Create a new device context (heap-allocated DeviceRunner). * Each ChipWorker should own one context for the lifetime of its init→finalize cycle. @@ -206,6 +244,15 @@ int simpler_run( DeviceContextHandle ctx, RuntimeHandle runtime, int32_t callable_id, const void *args, const CallConfig *config ); +/** Select the arena bank used by subsequent calls on the current thread. */ +int select_arena_bank_ctx(DeviceContextHandle ctx, unsigned arena_bank); + +/** Select the per-thread pipeline slot independently from arena banking. */ +int select_pipeline_slot_ctx(DeviceContextHandle ctx, unsigned pipeline_slot); + +/** Register the mailbox state word that receives the post-KernelLaunch ACK. */ +int set_task_accepted_state_ctx(DeviceContextHandle ctx, volatile int32_t *state, int32_t accepted_value); + /** * Drop the prepared state for `callable_id` and release the per-id share of * the device orch SO buffer. The buffer itself is freed only when its diff --git a/tests/st/a2a3/host_build_graph/batch_paged_attention_three_layer/kernels/aic/aic_pv_matmul.cpp b/tests/st/a2a3/host_build_graph/batch_paged_attention_three_layer/kernels/aic/aic_pv_matmul.cpp new file mode 100644 index 0000000000..825665b708 --- /dev/null +++ b/tests/st/a2a3/host_build_graph/batch_paged_attention_three_layer/kernels/aic/aic_pv_matmul.cpp @@ -0,0 +1,137 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +// Batched PV Matmul Kernel: for each batch b, pij(M, K) @ vj(K, N) -> oi_new(M, N) +// +// Processes batch_count batches in a single kernel invocation. +// Per-batch addresses are computed from global tensor bases + block_table lookup. +// +// Supports two tile configurations via runtime dispatch: +// Case1: (16, 128) @ (128, 128) -> (16, 128) +// Case2: (64, 64) @ ( 64, 128) -> (64, 128) +// +// Template: M=q_tile, K=block_size, N=head_dim + +#include +#include + +#include "tensor.h" + +// NOLINTNEXTLINE(build/namespaces) +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +template +static __aicore__ void pv_matmul_batch_impl( + __gm__ Tensor *pij_batch, __gm__ Tensor *value_cache, __gm__ Tensor *block_table_t, __gm__ Tensor *oi_new_batch, + uint64_t batch_count, uint64_t block_idx, uint64_t block_num, uint64_t batch_start +) { + __gm__ bfloat16_t *pij_base = reinterpret_cast<__gm__ bfloat16_t *>(pij_batch->buffer.addr); + __gm__ bfloat16_t *val_base = reinterpret_cast<__gm__ bfloat16_t *>(value_cache->buffer.addr); + __gm__ float *oi_base = reinterpret_cast<__gm__ float *>(oi_new_batch->buffer.addr); + __gm__ int32_t *bt = reinterpret_cast<__gm__ int32_t *>(block_table_t->buffer.addr); + + using GlobalA = GlobalTensor, Stride>; + using GlobalB = GlobalTensor, Stride>; + using GlobalOut = GlobalTensor, Stride>; + + using TileMatA = Tile; + using TileMatB = Tile; + + using LeftTile = TileLeft; + using RightTile = TileRight; + using AccTile = TileAcc; + + TileMatA aMatTile; + TileMatB bMatTile; + TASSIGN(aMatTile, 0x0); + TASSIGN(bMatTile, 0x20000); + + LeftTile aTile; + RightTile bTile; + AccTile cTile; + TASSIGN(aTile, 0x0); + TASSIGN(bTile, 0x0); + TASSIGN(cTile, 0x0); + + for (uint64_t b = 0; b < batch_count; b++) { + __gm__ bfloat16_t *pij_addr = pij_base + b * M * K; + int32_t phys_block = bt[(batch_start + b) * block_num + block_idx]; + __gm__ bfloat16_t *vj_addr = val_base + static_cast(phys_block) * K * N; + __gm__ float *oi_addr = oi_base + b * M * N; + + GlobalA pijGlobal(pij_addr); + GlobalB vjGlobal(vj_addr); + GlobalOut oiGlobal(oi_addr); + + TLOAD(aMatTile, pijGlobal); + TLOAD(bMatTile, vjGlobal); + + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + + TMOV(aTile, aMatTile); + TMOV(bTile, bMatTile); + + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + + TMATMUL(cTile, aTile, bTile); + + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + + TSTORE(oiGlobal, cTile); + + if (b + 1 < batch_count) { + pipe_barrier(PIPE_ALL); + } + } + + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *pij_batch = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *value_cache = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *block_table_t = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ Tensor *oi_new_batch = reinterpret_cast<__gm__ Tensor *>(args[3]); + uint64_t batch_count = static_cast(args[4]); + uint64_t block_idx = static_cast(args[5]); + uint64_t block_num = static_cast(args[6]); + uint64_t batch_start = static_cast(args[7]); + + uint64_t q_tile_size = static_cast(pij_batch->shapes[0] / batch_count); + uint64_t block_size = static_cast(pij_batch->shapes[1]); + + if (q_tile_size == 16 && block_size <= 16) { + pv_matmul_batch_impl<16, 16, 16>( + pij_batch, value_cache, block_table_t, oi_new_batch, batch_count, block_idx, block_num, batch_start + ); + } else if (q_tile_size == 16) { + pv_matmul_batch_impl<16, 128, 128>( + pij_batch, value_cache, block_table_t, oi_new_batch, batch_count, block_idx, block_num, batch_start + ); + } else { + pv_matmul_batch_impl<64, 64, 128>( + pij_batch, value_cache, block_table_t, oi_new_batch, batch_count, block_idx, block_num, batch_start + ); + } +} diff --git a/tests/st/a2a3/host_build_graph/batch_paged_attention_three_layer/kernels/aic/aic_qk_matmul.cpp b/tests/st/a2a3/host_build_graph/batch_paged_attention_three_layer/kernels/aic/aic_qk_matmul.cpp new file mode 100644 index 0000000000..0bf394f93b --- /dev/null +++ b/tests/st/a2a3/host_build_graph/batch_paged_attention_three_layer/kernels/aic/aic_qk_matmul.cpp @@ -0,0 +1,144 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +// Batched QK Matmul Kernel: for each batch b, qi(M, K) @ kj.T(K, N) -> sij(M, N) +// +// Processes batch_count batches in a single kernel invocation. +// Per-batch addresses are computed from global tensor bases + block_table lookup. +// +// Supports two tile configurations via runtime dispatch: +// Case1: (16, 128) @ (128, 128).T -> (16, 128) +// Case2: (64, 128) @ (128, 64).T -> (64, 64) +// +// Template: M=q_tile, K=head_dim, N=block_size + +#include +#include + +#include "tensor.h" + +// NOLINTNEXTLINE(build/namespaces) +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +template +static __aicore__ void qk_matmul_batch_impl( + __gm__ Tensor *query, __gm__ Tensor *key_cache, __gm__ Tensor *block_table_t, __gm__ Tensor *sij_batch, + uint64_t batch_count, uint64_t block_idx, uint64_t q_offset, uint64_t block_num, uint64_t num_heads, + uint64_t batch_start +) { + __gm__ bfloat16_t *query_base = reinterpret_cast<__gm__ bfloat16_t *>(query->buffer.addr); + __gm__ bfloat16_t *key_base = reinterpret_cast<__gm__ bfloat16_t *>(key_cache->buffer.addr); + __gm__ float *sij_base = reinterpret_cast<__gm__ float *>(sij_batch->buffer.addr); + __gm__ int32_t *bt = reinterpret_cast<__gm__ int32_t *>(block_table_t->buffer.addr); + + using GlobalA = GlobalTensor, Stride>; + using GlobalB = GlobalTensor, Stride, Layout::DN>; + using GlobalOut = GlobalTensor, Stride>; + + using TileMatA = Tile; + using TileMatB = Tile; + + using LeftTile = TileLeft; + using RightTile = TileRight; + using AccTile = TileAcc; + + TileMatA aMatTile; + TileMatB bMatTile; + TASSIGN(aMatTile, 0x0); + TASSIGN(bMatTile, 0x20000); + + LeftTile aTile; + RightTile bTile; + AccTile cTile; + TASSIGN(aTile, 0x0); + TASSIGN(bTile, 0x0); + TASSIGN(cTile, 0x0); + + for (uint64_t b = 0; b < batch_count; b++) { + __gm__ bfloat16_t *qi_addr = query_base + ((batch_start + b) * num_heads + q_offset) * K; + int32_t phys_block = bt[(batch_start + b) * block_num + block_idx]; + __gm__ bfloat16_t *kj_addr = key_base + static_cast(phys_block) * N * K; + __gm__ float *sij_addr = sij_base + b * M * N; + + GlobalA qiGlobal(qi_addr); + GlobalB kjGlobal(kj_addr); + GlobalOut sijGlobal(sij_addr); + + TLOAD(aMatTile, qiGlobal); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + TLOAD(bMatTile, kjGlobal); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + TMOV(aTile, aMatTile); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + TMOV(bTile, bMatTile); + + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + + TMATMUL(cTile, aTile, bTile); + + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + + TSTORE(sijGlobal, cTile); + + if (b + 1 < batch_count) { + pipe_barrier(PIPE_ALL); + } + } + + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *query = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *key_cache = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *block_table_t = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ Tensor *sij_batch = reinterpret_cast<__gm__ Tensor *>(args[3]); + uint64_t batch_count = static_cast(args[4]); + uint64_t block_idx = static_cast(args[5]); + uint64_t q_offset = static_cast(args[6]); + uint64_t block_num = static_cast(args[7]); + uint64_t num_heads = static_cast(args[8]); + uint64_t batch_start = static_cast(args[9]); + + uint64_t q_tile_size = static_cast(sij_batch->shapes[0] / batch_count); + uint64_t block_size = static_cast(sij_batch->shapes[1]); + + if (q_tile_size == 16 && block_size <= 16) { + qk_matmul_batch_impl<16, 16, 16>( + query, key_cache, block_table_t, sij_batch, batch_count, block_idx, q_offset, block_num, num_heads, + batch_start + ); + } else if (q_tile_size == 16) { + qk_matmul_batch_impl<16, 128, 128>( + query, key_cache, block_table_t, sij_batch, batch_count, block_idx, q_offset, block_num, num_heads, + batch_start + ); + } else { + qk_matmul_batch_impl<64, 128, 64>( + query, key_cache, block_table_t, sij_batch, batch_count, block_idx, q_offset, block_num, num_heads, + batch_start + ); + } +} diff --git a/tests/st/a2a3/host_build_graph/batch_paged_attention_three_layer/kernels/aiv/aiv_online_update.cpp b/tests/st/a2a3/host_build_graph/batch_paged_attention_three_layer/kernels/aiv/aiv_online_update.cpp new file mode 100644 index 0000000000..b8955c3b51 --- /dev/null +++ b/tests/st/a2a3/host_build_graph/batch_paged_attention_three_layer/kernels/aiv/aiv_online_update.cpp @@ -0,0 +1,230 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ +// Batched Online Softmax Update + Normalize Kernel (AIV) +// +// Processes batch_count batches in a single kernel invocation. +// For each batch b, updates accumulators mi/li/oi with new block's mij/lij/oi_new. +// On is_last, normalizes and writes to the output tensor at the correct batch offset. +// +// Supports two tile configurations via runtime dispatch: +// Case1: (16, 128) -- q_tile=16, head_dim=128 +// Case2: (64, 128) -- q_tile=64, head_dim=128 +// +// Scalar layout strategy: +// M scalar floats stored contiguously in GM can be loaded as either: +// - ND (kScalarRows, kScalarCols) RowMajor for element-wise ops +// - DN (kAlignedRows, 1) ColMajor for row-broadcast ops +// Conversion between layouts uses TRESHAPE (UB-internal, zero GM access). + +#include +#include + +#include "tensor.h" + +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +template +static __aicore__ void online_update_batch_impl( + __gm__ Tensor *mij_batch, __gm__ Tensor *lij_batch, __gm__ Tensor *oi_new_batch, __gm__ Tensor *mi_batch, + __gm__ Tensor *li_batch, __gm__ Tensor *oi_batch, __gm__ Tensor *out, uint64_t is_first, uint64_t is_last, + uint64_t batch_count, uint64_t q_offset, uint64_t num_heads, uint64_t batch_start +) { + __gm__ float *mij_base = reinterpret_cast<__gm__ float *>(mij_batch->buffer.addr); + __gm__ float *lij_base = reinterpret_cast<__gm__ float *>(lij_batch->buffer.addr); + __gm__ float *oi_new_base = reinterpret_cast<__gm__ float *>(oi_new_batch->buffer.addr); + __gm__ float *mi_base = reinterpret_cast<__gm__ float *>(mi_batch->buffer.addr); + __gm__ float *li_base = reinterpret_cast<__gm__ float *>(li_batch->buffer.addr); + __gm__ float *oi_base = reinterpret_cast<__gm__ float *>(oi_batch->buffer.addr); + __gm__ float *out_base = reinterpret_cast<__gm__ float *>(out->buffer.addr); + + constexpr int kScalarCols = 32 / sizeof(float); + constexpr int kScalarRows = M / kScalarCols; + constexpr int kAlignedRows = ((M * sizeof(float) + 31) / 32) * (32 / sizeof(float)); + + using GlobalDataMxN = GlobalTensor, Stride<1, 1, 1, N, 1>>; + using GlobalScalarND = + GlobalTensor, Stride<1, 1, 1, kScalarCols, 1>>; + + using TileDataMxN = Tile; + using TileScalarND = + Tile; + using TileScalarDN = Tile; + + constexpr int kDataBytes = M * N * sizeof(float); + constexpr int kScalarNDBytes = kScalarRows * kScalarCols * sizeof(float); + + TileDataMxN oiNewTile; + TileDataMxN oiTile; + + TileScalarND mijND, lijND, miND, liND; + TileScalarND miNewND, alphaND, betaND, tmpND; + + TileScalarDN alphaDN, betaDN, liDN; + + TASSIGN(oiNewTile, 0); + TASSIGN(oiTile, kDataBytes); + TASSIGN(mijND, 2 * kDataBytes); + TASSIGN(lijND, 2 * kDataBytes + kScalarNDBytes); + TASSIGN(miND, 2 * kDataBytes + 2 * kScalarNDBytes); + TASSIGN(liND, 2 * kDataBytes + 3 * kScalarNDBytes); + TASSIGN(miNewND, 2 * kDataBytes + 4 * kScalarNDBytes); + TASSIGN(alphaND, 2 * kDataBytes + 5 * kScalarNDBytes); + TASSIGN(betaND, 2 * kDataBytes + 6 * kScalarNDBytes); + TASSIGN(tmpND, 2 * kDataBytes + 7 * kScalarNDBytes); + + for (uint64_t b = 0; b < batch_count; b++) { + __gm__ float *mij_ptr = mij_base + b * M; + __gm__ float *lij_ptr = lij_base + b * M; + __gm__ float *oi_new_ptr = oi_new_base + b * M * N; + __gm__ float *mi_ptr = mi_base + b * M; + __gm__ float *li_ptr = li_base + b * M; + __gm__ float *oi_ptr = oi_base + b * M * N; + __gm__ float *dst_ptr = out_base + ((batch_start + b) * num_heads + q_offset) * N; + + GlobalDataMxN oiNewGlobal(oi_new_ptr); + GlobalDataMxN oiGlobal(oi_ptr); + GlobalDataMxN dstGlobal(dst_ptr); + + GlobalScalarND mijGlobalND(mij_ptr); + GlobalScalarND lijGlobalND(lij_ptr); + GlobalScalarND miGlobalND(mi_ptr); + GlobalScalarND liGlobalND(li_ptr); + + if (is_first) { + TLOAD(oiNewTile, oiNewGlobal); + TLOAD(mijND, mijGlobalND); + TLOAD(lijND, lijGlobalND); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(miGlobalND, mijND); + TSTORE(liGlobalND, lijND); + TSTORE(oiGlobal, oiNewTile); + + if (is_last) { + TRESHAPE(liDN, lijND); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); + TROWEXPANDDIV(oiNewTile, oiNewTile, liDN); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + TSTORE(dstGlobal, oiNewTile); + } + } else { + TLOAD(oiNewTile, oiNewGlobal); + TLOAD(oiTile, oiGlobal); + TLOAD(mijND, mijGlobalND); + TLOAD(lijND, lijGlobalND); + TLOAD(miND, miGlobalND); + TLOAD(liND, liGlobalND); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + TMAX(miNewND, miND, mijND); + pipe_barrier(PIPE_V); + TSUB(alphaND, miND, miNewND); + pipe_barrier(PIPE_V); + TEXP(alphaND, alphaND); + pipe_barrier(PIPE_V); + TSUB(betaND, mijND, miNewND); + pipe_barrier(PIPE_V); + TEXP(betaND, betaND); + pipe_barrier(PIPE_V); + TMUL(liND, alphaND, liND); + pipe_barrier(PIPE_V); + TMUL(tmpND, betaND, lijND); + pipe_barrier(PIPE_V); + TADD(liND, liND, tmpND); + + TRESHAPE(alphaDN, alphaND); + TRESHAPE(betaDN, betaND); + if (is_last) { + TRESHAPE(liDN, liND); + } + + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(miGlobalND, miNewND); + TSTORE(liGlobalND, liND); + + TROWEXPANDMUL(oiTile, oiTile, alphaDN); + TROWEXPANDMUL(oiNewTile, oiNewTile, betaDN); + pipe_barrier(PIPE_V); + TADD(oiTile, oiTile, oiNewTile); + + if (is_last) { + pipe_barrier(PIPE_V); + TROWEXPANDDIV(oiTile, oiTile, liDN); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + TSTORE(dstGlobal, oiTile); + } else { + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + TSTORE(oiGlobal, oiTile); + } + } + + if (b + 1 < batch_count) { + pipe_barrier(PIPE_ALL); + } + } + + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *mij_batch = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *lij_batch = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *oi_new_batch = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ Tensor *mi_batch = reinterpret_cast<__gm__ Tensor *>(args[3]); + __gm__ Tensor *li_batch = reinterpret_cast<__gm__ Tensor *>(args[4]); + __gm__ Tensor *oi_batch = reinterpret_cast<__gm__ Tensor *>(args[5]); + __gm__ Tensor *out = reinterpret_cast<__gm__ Tensor *>(args[6]); + uint64_t is_first = static_cast(args[7]); + uint64_t is_last = static_cast(args[8]); + uint64_t batch_count = static_cast(args[9]); + uint64_t q_offset = static_cast(args[10]); + uint64_t num_heads = static_cast(args[11]); + uint64_t batch_start = static_cast(args[12]); + + uint64_t q_tile_size = static_cast(mij_batch->shapes[0] / batch_count); + uint64_t head_dim = static_cast(oi_new_batch->shapes[1]); + + if (q_tile_size == 16 && head_dim <= 16) { + online_update_batch_impl<16, 16>( + mij_batch, lij_batch, oi_new_batch, mi_batch, li_batch, oi_batch, out, is_first, is_last, batch_count, + q_offset, num_heads, batch_start + ); + } else if (q_tile_size == 16) { + online_update_batch_impl<16, 128>( + mij_batch, lij_batch, oi_new_batch, mi_batch, li_batch, oi_batch, out, is_first, is_last, batch_count, + q_offset, num_heads, batch_start + ); + } else { + online_update_batch_impl<64, 128>( + mij_batch, lij_batch, oi_new_batch, mi_batch, li_batch, oi_batch, out, is_first, is_last, batch_count, + q_offset, num_heads, batch_start + ); + } +} diff --git a/tests/st/a2a3/host_build_graph/batch_paged_attention_three_layer/kernels/aiv/aiv_softmax_prepare.cpp b/tests/st/a2a3/host_build_graph/batch_paged_attention_three_layer/kernels/aiv/aiv_softmax_prepare.cpp new file mode 100644 index 0000000000..3ce77eaa09 --- /dev/null +++ b/tests/st/a2a3/host_build_graph/batch_paged_attention_three_layer/kernels/aiv/aiv_softmax_prepare.cpp @@ -0,0 +1,200 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +// Batched Softmax Preparation Kernel (AIV) +// +// Processes batch_count batches in a single kernel invocation. +// For each batch b at block_idx bn: +// valid_len = min(N, context_lens[b] - bn * N) +// sij_masked = pad(sij[b], valid_len, -inf) +// sij_scale = sij_masked * scale +// mij[b] = row_max(sij_scale) +// pij[b] = exp(sij_scale - mij[b]) (truncated to bf16 then back) +// lij[b] = row_sum(pij[b]) +// +// Supports two tile configurations via runtime dispatch: +// Case1: (16, 128) -- q_tile=16, block_size=128 +// Case2: (64, 64) -- q_tile=64, block_size=64 + +#include +#include + +#include "tensor.h" + +// NOLINTNEXTLINE(build/namespaces) +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +template +static __aicore__ void softmax_prepare_batch_impl( + __gm__ Tensor *sij_batch, __gm__ Tensor *context_lens_t, __gm__ Tensor *pij_batch, __gm__ Tensor *mij_batch, + __gm__ Tensor *lij_batch, float scale_value, uint64_t batch_count, uint64_t block_idx, uint64_t batch_start +) { + __gm__ float *sij_base = reinterpret_cast<__gm__ float *>(sij_batch->buffer.addr); + __gm__ bfloat16_t *pij_base = reinterpret_cast<__gm__ bfloat16_t *>(pij_batch->buffer.addr); + __gm__ float *mij_base = reinterpret_cast<__gm__ float *>(mij_batch->buffer.addr); + __gm__ float *lij_base = reinterpret_cast<__gm__ float *>(lij_batch->buffer.addr); + __gm__ int32_t *ctx_lens = reinterpret_cast<__gm__ int32_t *>(context_lens_t->buffer.addr); + + constexpr int kAlignedRows = ((M * sizeof(float) + 31) / 32) * (32 / sizeof(float)); + + using GlobalDataMxN = GlobalTensor, Stride<1, 1, 1, N, 1>>; + using GlobalDataMxN_bf16 = GlobalTensor, Stride<1, 1, 1, N, 1>>; + using GlobalScalarDN = GlobalTensor, Stride<1, 1, 1, 1, 1>, Layout::DN>; + + using TileSijDyn = Tile; + using TileSijPad = Tile; + + using TileVecMxN = Tile; + using TileVecMxN_bf16 = Tile; + using TileScalarDN = Tile; + + TileVecMxN sijTile; + TileSijPad sijPadTile; + TileVecMxN pijTile; + TileVecMxN tmpTile; + TileScalarDN maxTile; + TileScalarDN sumTile; + TileVecMxN_bf16 pijBf16Tile; + + TASSIGN(sijTile, 0x0); + TASSIGN(sijPadTile, 0x0); + TASSIGN(pijTile, M * N * sizeof(float)); + TASSIGN(tmpTile, 2 * M * N * sizeof(float)); + TASSIGN(maxTile, 3 * M * N * sizeof(float)); + TASSIGN(sumTile, 3 * M * N * sizeof(float) + kAlignedRows * sizeof(float)); + TASSIGN(pijBf16Tile, 3 * M * N * sizeof(float) + 2 * kAlignedRows * sizeof(float)); + + for (uint64_t b = 0; b < batch_count; b++) { + int32_t cur_seq = ctx_lens[batch_start + b]; + uint64_t start = block_idx * N; + uint64_t valid_len = 0; + if (start < static_cast(cur_seq)) { + uint64_t remaining = static_cast(cur_seq) - start; + valid_len = (remaining < N) ? remaining : N; + } + + __gm__ float *sij_addr = sij_base + b * M * N; + __gm__ bfloat16_t *pij_addr = pij_base + b * M * N; + __gm__ float *mij_addr = mij_base + b * M; + __gm__ float *lij_addr = lij_base + b * M; + + GlobalDataMxN sijGlobal(sij_addr); + GlobalDataMxN_bf16 pijGlobal(pij_addr); + GlobalScalarDN mijGlobal(mij_addr); + GlobalScalarDN lijGlobal(lij_addr); + + if (valid_len == 0) { + // Block entirely beyond sequence: write mij=-1e30, lij=0, pij=0 + // Use -1e30 instead of -inf to avoid NaN in online_update (exp(-inf - (-inf)) = NaN) + constexpr float NEG_LARGE = -1e30f; + for (int i = 0; i < kAlignedRows; i++) { + maxTile.SetValue(i, NEG_LARGE); + sumTile.SetValue(i, 0.0f); + } + for (int i = 0; i < M * N; i++) { + pijBf16Tile.SetValue(i, static_cast(0.0f)); + } + + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(mijGlobal, maxTile); + TSTORE(lijGlobal, sumTile); + TSTORE(pijGlobal, pijBf16Tile); + + if (b + 1 < batch_count) { + pipe_barrier(PIPE_ALL); + } + continue; + } + + TLOAD(sijTile, sijGlobal); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + TileSijDyn sijDynTile(static_cast(valid_len)); + TASSIGN(sijDynTile, 0x0); + TFILLPAD_INPLACE(sijPadTile, sijDynTile); + pipe_barrier(PIPE_V); + + TMULS(sijTile, sijTile, scale_value); + pipe_barrier(PIPE_V); + TROWMAX(maxTile, sijTile, tmpTile); + pipe_barrier(PIPE_V); + TROWEXPANDSUB(pijTile, sijTile, maxTile); + pipe_barrier(PIPE_V); + TEXP(pijTile, pijTile); + pipe_barrier(PIPE_V); + // Truncate pij to bf16 first, then compute lij from truncated values (matches golden) + TCVT(pijBf16Tile, pijTile, RoundMode::CAST_ROUND); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + pipe_barrier(PIPE_V); + TCVT(pijTile, pijBf16Tile, RoundMode::CAST_ROUND); + pipe_barrier(PIPE_V); + TROWSUM(sumTile, pijTile, tmpTile); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(pijGlobal, pijBf16Tile); + TSTORE(mijGlobal, maxTile); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + TSTORE(lijGlobal, sumTile); + + if (b + 1 < batch_count) { + pipe_barrier(PIPE_ALL); + } + } + + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *sij_batch = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *context_lens_t = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *pij_batch = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ Tensor *mij_batch = reinterpret_cast<__gm__ Tensor *>(args[3]); + __gm__ Tensor *lij_batch = reinterpret_cast<__gm__ Tensor *>(args[4]); + union { + uint64_t u; + float f; + } scale_conv; + scale_conv.u = static_cast(args[5]); + float scale_value = scale_conv.f; + uint64_t batch_count = static_cast(args[6]); + uint64_t block_idx = static_cast(args[7]); + uint64_t batch_start = static_cast(args[8]); + + uint64_t q_tile_size = static_cast(sij_batch->shapes[0] / batch_count); + uint64_t block_size = static_cast(pij_batch->shapes[1]); + + if (q_tile_size == 16 && block_size <= 16) { + softmax_prepare_batch_impl<16, 16>( + sij_batch, context_lens_t, pij_batch, mij_batch, lij_batch, scale_value, batch_count, block_idx, batch_start + ); + } else if (q_tile_size == 16) { + softmax_prepare_batch_impl<16, 128>( + sij_batch, context_lens_t, pij_batch, mij_batch, lij_batch, scale_value, batch_count, block_idx, batch_start + ); + } else { + softmax_prepare_batch_impl<64, 64>( + sij_batch, context_lens_t, pij_batch, mij_batch, lij_batch, scale_value, batch_count, block_idx, batch_start + ); + } +} diff --git a/tests/st/a2a3/host_build_graph/batch_paged_attention_three_layer/kernels/orchestration/paged_attention_three_layer_orch.cpp b/tests/st/a2a3/host_build_graph/batch_paged_attention_three_layer/kernels/orchestration/paged_attention_three_layer_orch.cpp new file mode 100644 index 0000000000..fd8c41ba96 --- /dev/null +++ b/tests/st/a2a3/host_build_graph/batch_paged_attention_three_layer/kernels/orchestration/paged_attention_three_layer_orch.cpp @@ -0,0 +1,214 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include +#include +#include + +#include "pto_orchestration_api.h" + +#define FUNC_QK_MATMUL 0 +#define FUNC_SOFTMAX_PREPARE 1 +#define FUNC_PV_MATMUL 2 +#define FUNC_ONLINE_UPDATE 3 + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; + return PTO2OrchestrationConfig{ + .expected_arg_count = 9, + }; +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) { + uint64_t batch = orch_args.tensor(0).ref().shapes[0]; + uint64_t num_heads = orch_args.tensor(0).ref().shapes[1]; + uint64_t head_dim = orch_args.tensor(0).ref().shapes[2]; + DataType data_type = orch_args.tensor(0).ref().dtype; + + uint64_t block_size = orch_args.tensor(1).ref().shapes[1]; + uint64_t block_num = orch_args.tensor(3).ref().shapes[1]; + uint64_t scale_value = orch_args.scalar(0); + uint64_t layer_count = orch_args.scalar(1); + uint64_t layers_per_epoch = orch_args.scalar(2); + if (layer_count == 0 || layers_per_epoch == 0) { + LOG_ERROR("batch_paged_attention_three_layer: layer and epoch sizes must be non-zero"); + return; + } + + uint64_t q_tile = std::min(num_heads, static_cast(128)); + uint64_t q_loop = (num_heads + q_tile - 1) / q_tile; + + void *query_ptr = orch_args.tensor(0).ref().data_as(); + void *kc_ptr = orch_args.tensor(1).ref().data_as(); + void *vc_ptr = orch_args.tensor(2).ref().data_as(); + void *out_ptr = orch_args.tensor(5).ref().data_as(); + + uint32_t bt_shapes[2] = {static_cast(batch), static_cast(block_num)}; + Tensor block_table = + make_tensor_external(orch_args.tensor(3).ref().data_as(), bt_shapes, 2, DataType::INT32, false); + + uint32_t cl_shapes[1] = {static_cast(batch)}; + Tensor context_lens = + make_tensor_external(orch_args.tensor(4).ref().data_as(), cl_shapes, 1, DataType::INT32, false); + + uint64_t max_bn = 0; + for (uint64_t b = 0; b < batch; b++) { + uint32_t cl_idx[1] = {static_cast(b)}; + uint64_t cur_seq = static_cast(get_tensor_data(context_lens, 1, cl_idx)); + max_bn = std::max(max_bn, (cur_seq + block_size - 1) / block_size); + } + + uint32_t query_shapes[2] = {static_cast(batch * num_heads), static_cast(head_dim)}; + uint64_t total_blocks_count = orch_args.tensor(1).ref().shapes[0]; + uint64_t kv_total_rows = total_blocks_count * block_size; + uint32_t key_cache_shapes[2] = {static_cast(kv_total_rows), static_cast(head_dim)}; + uint32_t value_cache_shapes[2] = {static_cast(kv_total_rows), static_cast(head_dim)}; + uint32_t out_shapes[2] = {static_cast(batch * num_heads), static_cast(head_dim)}; + + Tensor query = make_tensor_external(query_ptr, query_shapes, 2, data_type); + Tensor key_cache = make_tensor_external(kc_ptr, key_cache_shapes, 2, data_type); + Tensor value_cache = make_tensor_external(vc_ptr, value_cache_shapes, 2, data_type); + Tensor out = make_tensor_external(out_ptr, out_shapes, 2, DataType::FLOAT32, true); + + constexpr uint64_t IN_CORE_BATCH = 16; + constexpr uint32_t MAX_LAYER_DONE_DEPS = 64; + uint64_t num_chunks = (batch + IN_CORE_BATCH - 1) / IN_CORE_BATCH; + PTO2TaskId prev_layer_done = PTO2TaskId::invalid(); + + for (uint64_t layer = 0; layer < layer_count; layer++) { + PTO2TaskId layer_done_deps[MAX_LAYER_DONE_DEPS]; + uint32_t layer_done_dep_count = 0; + + for (uint64_t q_idx = 0; q_idx < q_loop; q_idx++) { + uint64_t q_offset = q_idx * q_tile; + + for (uint64_t chunk_idx = 0; chunk_idx < num_chunks; chunk_idx++) { + uint64_t chunk_bc = std::min(batch - chunk_idx * IN_CORE_BATCH, IN_CORE_BATCH); + uint64_t batch_start = chunk_idx * IN_CORE_BATCH; + PTO2TaskId chunk_done = PTO2TaskId::invalid(); + + PTO2_SCOPE() { + uint32_t oi_acc_shapes[2] = { + static_cast(chunk_bc * q_tile), static_cast(head_dim) + }; + uint32_t scalar_acc_shapes[1] = {static_cast(chunk_bc * q_tile)}; + TensorCreateInfo oi_batch_ci(oi_acc_shapes, 2, DataType::FLOAT32); + TensorCreateInfo scalar_acc_ci(scalar_acc_shapes, 1, DataType::FLOAT32); + TaskOutputTensors alloc_outs = alloc_tensors(oi_batch_ci, scalar_acc_ci, scalar_acc_ci); + const Tensor &oi_batch = alloc_outs.get_ref(0); + const Tensor &li_batch = alloc_outs.get_ref(1); + const Tensor &mi_batch = alloc_outs.get_ref(2); + + uint32_t sij_shapes[2] = { + static_cast(chunk_bc * q_tile), static_cast(block_size) + }; + uint32_t vec_shapes[1] = {static_cast(chunk_bc * q_tile)}; + uint32_t oi_new_shapes[2] = { + static_cast(chunk_bc * q_tile), static_cast(head_dim) + }; + TensorCreateInfo sij_ci(sij_shapes, 2, DataType::FLOAT32); + TensorCreateInfo pij_ci(sij_shapes, 2, data_type); + TensorCreateInfo vec_ci(vec_shapes, 1, DataType::FLOAT32); + TensorCreateInfo oi_new_ci(oi_new_shapes, 2, DataType::FLOAT32); + + for (uint64_t bn = 0; bn < max_bn; bn++) { + PTO2_SCOPE() { + L0TaskArgs params_qk; + params_qk.add_input(query); + params_qk.add_input(key_cache); + params_qk.add_input(block_table); + params_qk.add_output(sij_ci); + if (prev_layer_done.is_valid()) { + params_qk.set_dependencies(&prev_layer_done, 1); + } + params_qk.add_scalar(chunk_bc); + params_qk.add_scalar(bn); + params_qk.add_scalar(q_offset); + params_qk.add_scalar(block_num); + params_qk.add_scalar(num_heads); + params_qk.add_scalar(batch_start); + TaskOutputTensors qk_outs = rt_submit_aic_task(FUNC_QK_MATMUL, params_qk); + const Tensor &sij_b = qk_outs.get_ref(0); + + L0TaskArgs params_sf; + params_sf.add_input(sij_b); + params_sf.add_input(context_lens); + params_sf.add_output(pij_ci); + params_sf.add_output(vec_ci); + params_sf.add_output(vec_ci); + params_sf.add_scalar(scale_value); + params_sf.add_scalar(chunk_bc); + params_sf.add_scalar(bn); + params_sf.add_scalar(batch_start); + TaskOutputTensors sf_outs = rt_submit_aiv_task(FUNC_SOFTMAX_PREPARE, params_sf); + const Tensor &pij_b = sf_outs.get_ref(0); + const Tensor &mij_b = sf_outs.get_ref(1); + const Tensor &lij_b = sf_outs.get_ref(2); + + L0TaskArgs params_pv; + params_pv.add_input(pij_b); + params_pv.add_input(value_cache); + params_pv.add_input(block_table); + params_pv.add_output(oi_new_ci); + params_pv.add_scalar(chunk_bc); + params_pv.add_scalar(bn); + params_pv.add_scalar(block_num); + params_pv.add_scalar(batch_start); + TaskOutputTensors pv_outs = rt_submit_aic_task(FUNC_PV_MATMUL, params_pv); + const Tensor &oi_new_b = pv_outs.get_ref(0); + + L0TaskArgs params_up; + params_up.add_input(mij_b); + params_up.add_input(lij_b); + params_up.add_input(oi_new_b); + params_up.add_inout(mi_batch); + params_up.add_inout(li_batch); + params_up.add_inout(oi_batch); + params_up.add_inout(out); + params_up.add_scalar(bn == 0 ? 1 : 0); + params_up.add_scalar(bn == max_bn - 1 ? 1 : 0); + params_up.add_scalar(chunk_bc); + params_up.add_scalar(q_offset); + params_up.add_scalar(num_heads); + params_up.add_scalar(batch_start); + chunk_done = rt_submit_aiv_task(FUNC_ONLINE_UPDATE, params_up).task_id(); + } + } + } + + if (!chunk_done.is_valid() || layer_done_dep_count >= MAX_LAYER_DONE_DEPS) { + LOG_ERROR("batch_paged_attention_three_layer: layer barrier capacity exceeded"); + return; + } + layer_done_deps[layer_done_dep_count++] = chunk_done; + } + } + + L0TaskArgs layer_done_args; + layer_done_args.set_dependencies(layer_done_deps, layer_done_dep_count); + prev_layer_done = rt_submit_dummy_task(layer_done_args).task_id(); + bool final_layer = layer + 1 == layer_count; + if (final_layer || (layer + 1) % layers_per_epoch == 0) { + rt_graph_boundary(final_layer); + } + } + + uint64_t tasks_per_layer = q_loop * num_chunks * (1 + max_bn * 4) + 1; + LOG_INFO_V0( + "batch_paged_attention_three_layer: layers=%" PRIu64 ", layers_per_epoch=%" PRIu64 ", tasks=%" PRIu64 + ", tasks_per_layer=%" PRIu64, + layer_count, layers_per_epoch, layer_count * tasks_per_layer, tasks_per_layer + ); +} + +} // extern "C" diff --git a/tests/st/a2a3/host_build_graph/batch_paged_attention_three_layer/test_batch_paged_attention_three_layer.py b/tests/st/a2a3/host_build_graph/batch_paged_attention_three_layer/test_batch_paged_attention_three_layer.py new file mode 100644 index 0000000000..abc4cc1ab7 --- /dev/null +++ b/tests/st/a2a3/host_build_graph/batch_paged_attention_three_layer/test_batch_paged_attention_three_layer.py @@ -0,0 +1,171 @@ +#!/usr/bin/env python3 +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Three-layer batch paged attention baseline for host_build_graph.""" + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import Scalar, SceneTestCase, TaskArgsBuilder, Tensor, scene_test +from simpler_setup.goldens.paged_attention import compute_golden as _pa_compute_golden +from simpler_setup.goldens.paged_attention import generate_inputs as _pa_generate_inputs + + +@scene_test(level=2, runtime="host_build_graph") +class TestBatchPagedAttentionThreeLayer(SceneTestCase): + RTOL = 1e-3 + ATOL = 1e-3 + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/paged_attention_three_layer_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.IN, D.IN, D.IN, D.IN, D.IN, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "name": "QK", + "source": "kernels/aic/aic_qk_matmul.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 1, + "name": "SF", + "source": "kernels/aiv/aiv_softmax_prepare.cpp", + "core_type": "aiv", + "signature": [D.IN, D.OUT, D.OUT, D.OUT], + }, + { + "func_id": 2, + "name": "PV", + "source": "kernels/aic/aic_pv_matmul.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 3, + "name": "UP", + "source": "kernels/aiv/aiv_online_update.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.IN, D.INOUT, D.INOUT, D.INOUT, D.INOUT], + }, + ], + } + + CASES = [ + { + "name": "SmallThreeLayer", + "platforms": ["a2a3sim", "a2a3"], + "config": {"aicpu_thread_num": 4, "block_dim": 9}, + "params": { + "batch": 1, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 33, + "max_model_len": 256, + "dtype": "bfloat16", + "layer_count": 3, + }, + }, + { + "name": "PipelineStressThreeLayer", + "platforms": ["a2a3"], + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "manual": True, + "params": { + "batch": 16, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 128, + "block_size": 128, + "context_len": 8192, + "max_model_len": 8192, + "dtype": "bfloat16", + "layer_count": 3, + }, + }, + { + "name": "PipelineStressTwoTokenFortyLayer", + "platforms": ["a2a3"], + "config": { + "aicpu_thread_num": 4, + "block_dim": 24, + "runtime_env": { + "ring_task_window": 32768, + "ring_heap": 1024 * 1024 * 1024, + "ring_dep_pool": 32768, + }, + }, + "manual": True, + "params": { + "batch": 16, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 128, + "block_size": 128, + "context_len": 8192, + "max_model_len": 8192, + "dtype": "bfloat16", + "layer_count": 80, + "layers_per_epoch": 40, + }, + }, + { + "name": "PipelineStressThreeTokenFortyLayer", + "platforms": ["a2a3"], + "config": { + "aicpu_thread_num": 4, + "block_dim": 24, + "runtime_env": { + "ring_task_window": 32768, + "ring_heap": 1024 * 1024 * 1024, + "ring_dep_pool": 32768, + }, + }, + "manual": True, + "params": { + "batch": 16, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 128, + "block_size": 128, + "context_len": 8192, + "max_model_len": 8192, + "dtype": "bfloat16", + "layer_count": 120, + "layers_per_epoch": 40, + }, + }, + ] + + def generate_args(self, params): + specs = [] + for name, value in _pa_generate_inputs(params): + if isinstance(value, torch.Tensor): + specs.append(Tensor(name, value)) + else: + specs.append(Scalar(name, value)) + specs.append(Scalar("layer_count", params["layer_count"])) + specs.append(Scalar("layers_per_epoch", params.get("layers_per_epoch", 1))) + return TaskArgsBuilder(*specs) + + def compute_golden(self, args, params): + tensors = {spec.name: spec.value for spec in args.specs if isinstance(spec, Tensor)} + _pa_compute_golden(tensors, params) + for spec in args.specs: + if isinstance(spec, Tensor) and spec.name in tensors: + getattr(args, spec.name)[:] = tensors[spec.name] + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a2a3/host_build_graph/batch_paged_attention_three_layer/test_batch_paged_attention_three_layer_l3.py b/tests/st/a2a3/host_build_graph/batch_paged_attention_three_layer/test_batch_paged_attention_three_layer_l3.py new file mode 100644 index 0000000000..6f3db947fd --- /dev/null +++ b/tests/st/a2a3/host_build_graph/batch_paged_attention_three_layer/test_batch_paged_attention_three_layer_l3.py @@ -0,0 +1,298 @@ +#!/usr/bin/env python3 +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Level-3 coverage for streaming HostGraph paged attention.""" + +import threading +import time + +import torch +from simpler.l3_l2_message_queue import L3L2QueueOpcode +from simpler.request_session import ( + HostGraphToken, + append_host_graph_token_stream_args, + decode_host_graph_token, +) +from simpler.task_interface import ArgDirection as D + +from simpler_setup import Scalar, SceneTestCase, TaskArgsBuilder, Tensor, scene_test +from simpler_setup.goldens.paged_attention import compute_golden as _pa_compute_golden +from simpler_setup.goldens.paged_attention import generate_inputs as _pa_generate_inputs +from simpler_setup.scene_test import CallableNamespace, _build_l3_task_args, _compare_outputs, _RehostedTaskArgs + +_STREAM_TIMEOUT_S = 30.0 +_CROSS_REQUEST_STAGGER_S = 0.02 +_REQUEST_A = 1 +_REQUEST_B = 2 + +_SMALL_PARAMS = { + "batch": 1, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 33, + "max_model_len": 256, + "dtype": "bfloat16", + "layer_count": 3, + "layers_per_epoch": 1, +} +_STRESS_PARAMS = { + "batch": 16, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 128, + "block_size": 128, + "context_len": 8192, + "max_model_len": 8192, + "dtype": "bfloat16", + "layer_count": 120, + "layers_per_epoch": 40, +} +_STRESS_CONFIG = { + "device_count": 1, + "num_sub_workers": 0, + "aicpu_thread_num": 4, + "block_dim": 24, + "runtime_env": { + "ring_task_window": 32768, + "ring_heap": 1024 * 1024 * 1024, + "ring_dep_pool": 32768, + }, +} + + +def _read_token(queue, request_id: int, expected_seq: int, expected_epochs: int) -> tuple[HostGraphToken, object]: + message = queue.output.peek(timeout=_STREAM_TIMEOUT_S) + assert message.opcode == L3L2QueueOpcode.DATA + payload = bytearray(message.payload_nbytes) + queue.output.read_into(message, payload) + token = decode_host_graph_token(payload) + assert token.request_id == request_id + assert token.token_seq == expected_seq + assert token.token_id == expected_seq + assert token.is_final == (expected_seq == expected_epochs) + assert token.synthetic + assert token.status == 0 + return token, message + + +def run_dag(orch, callables, task_args, config, *, request_id=_REQUEST_A, token_sink=None): + chip_args, _ = _build_l3_task_args(task_args, callables.paged_attention_sig) + layer_count = int(chip_args.scalar(1)) + layers_per_epoch = int(chip_args.scalar(2)) + expected_epochs = (layer_count + layers_per_epoch - 1) // layers_per_epoch + queue = orch.create_l3_l2_queue(worker_id=0, depth=1, input_arena_bytes=64, output_arena_bytes=4096) + append_host_graph_token_stream_args(chip_args, queue, request_id) + callables.keep(chip_args) + orch.submit_next_level(callables.paged_attention, chip_args, config, worker=0) + + for seq in range(1, expected_epochs + 1): + token, message = _read_token(queue, request_id, seq, expected_epochs) + if token_sink is not None: + token_sink(token) + queue.output.release(message) + + +@scene_test(level=3, runtime="host_build_graph") +class TestBatchPagedAttentionThreeLayerL3(SceneTestCase): + RTOL = 1e-3 + ATOL = 1e-3 + + CALLABLE = { + "orchestration": run_dag, + "callables": [ + { + "name": "paged_attention", + "orchestration": { + "source": "kernels/orchestration/paged_attention_three_layer_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.IN, D.IN, D.IN, D.IN, D.IN, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "name": "QK", + "source": "kernels/aic/aic_qk_matmul.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 1, + "name": "SF", + "source": "kernels/aiv/aiv_softmax_prepare.cpp", + "core_type": "aiv", + "signature": [D.IN, D.OUT, D.OUT, D.OUT], + }, + { + "func_id": 2, + "name": "PV", + "source": "kernels/aic/aic_pv_matmul.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 3, + "name": "UP", + "source": "kernels/aiv/aiv_online_update.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.IN, D.INOUT, D.INOUT, D.INOUT, D.INOUT], + }, + ], + } + ], + } + + CASES = [ + { + "name": "SmallThreeLayerStreaming", + "platforms": ["a2a3sim", "a2a3"], + "config": {"device_count": 1, "num_sub_workers": 0, "aicpu_thread_num": 4, "block_dim": 9}, + "params": _SMALL_PARAMS, + }, + { + "name": "ConcurrentTwoRequestOverlapProfile", + "platforms": ["a2a3"], + "config": _STRESS_CONFIG, + "manual": True, + "params": _STRESS_PARAMS, + }, + ] + + def _run_and_validate_l3(self, worker, compiled_callables, sub_handles, case, **kwargs): + if case["name"] != "ConcurrentTwoRequestOverlapProfile": + return super()._run_and_validate_l3(worker, compiled_callables, sub_handles, case, **kwargs) + return self._run_concurrent_requests(worker, compiled_callables, sub_handles, case) + + def _build_submission(self, callables, request, request_id: int, queue): + chip_args, _ = _build_l3_task_args(request, callables.paged_attention_sig) + layer_count = int(chip_args.scalar(1)) + layers_per_epoch = int(chip_args.scalar(2)) + expected_epochs = (layer_count + layers_per_epoch - 1) // layers_per_epoch + append_host_graph_token_stream_args(chip_args, queue, request_id) + callables.keep(chip_args) + return chip_args, queue, expected_epochs + + def _relay_tokens(self, queue, request_id: int, expected_epochs: int, emitter) -> None: + for seq in range(1, expected_epochs + 1): + token, message = _read_token(queue, request_id, seq, expected_epochs) + print( + f"L3_TOKEN_RECEIVED request_id={request_id} token_seq={seq} " + f"received_ns={time.monotonic_ns()} final={token.is_final}" + ) + try: + emitter.emit(token, final=token.is_final) + print( + f"USER_TOKEN_DELIVERED request_id={request_id} token_seq={seq} " + f"delivered_ns={time.monotonic_ns()} final={token.is_final}" + ) + finally: + queue.output.release(message) + + def _run_concurrent_requests(self, worker, compiled_callables, sub_handles, case): + params = case["params"] + requests = {request_id: self.generate_args(params) for request_id in (_REQUEST_A, _REQUEST_B)} + goldens = {request_id: request.clone() for request_id, request in requests.items()} + for golden in goldens.values(): + self.compute_golden(golden, params) + rehosted = [] + try: + for request_id in (_REQUEST_A, _REQUEST_B): + rehosted.append(_RehostedTaskArgs(worker, requests[request_id])) + + config = self._build_config(case["config"]) + callables = CallableNamespace({**compiled_callables, **sub_handles}) + request_a_submitted = threading.Event() + queues_ready = threading.Event() + queues = {} + + def request_orch(orch, request, request_id, emitter, cfg): + if request_id == _REQUEST_A: + for queued_request_id in (_REQUEST_A, _REQUEST_B): + queues[queued_request_id] = orch.create_l3_l2_queue( + worker_id=0, depth=1, input_arena_bytes=64, output_arena_bytes=4096 + ) + queues_ready.set() + else: + assert queues_ready.wait(_STREAM_TIMEOUT_S) + queue = queues[request_id] + chip_args, queue, expected_epochs = self._build_submission(callables, request, request_id, queue) + orch.submit_next_level(callables.paged_attention, chip_args, cfg, worker=0) + if request_id == _REQUEST_A: + request_a_submitted.set() + self._relay_tokens(queue, request_id, expected_epochs, emitter) + + def consume(stream, output, errors): + try: + output.extend(stream) + except BaseException as exc: # noqa: BLE001 + errors.append(exc) + + with worker.open_request_session(request_orch, max_active_runs=2) as session: + streams = { + _REQUEST_A: session.submit(requests[_REQUEST_A], config=config, request_id=_REQUEST_A), + } + assert request_a_submitted.wait(_STREAM_TIMEOUT_S) + # Place B's Host O behind A's first publication instead of profiling + # two simultaneous cold starts. + time.sleep(_CROSS_REQUEST_STAGGER_S) + streams[_REQUEST_B] = session.submit(requests[_REQUEST_B], config=config, request_id=_REQUEST_B) + tokens = {request_id: [] for request_id in streams} + errors = [] + consumers = [ + threading.Thread(target=consume, args=(stream, tokens[request_id], errors)) + for request_id, stream in streams.items() + ] + for consumer in consumers: + consumer.start() + for consumer in consumers: + consumer.join(_STREAM_TIMEOUT_S) + assert not consumer.is_alive(), "request stream consumer timed out" + if errors: + raise errors[0] + + expected_epochs = (int(params["layer_count"]) + int(params["layers_per_epoch"]) - 1) // int( + params["layers_per_epoch"] + ) + for request_id in streams: + assert [token.token_seq for token in tokens[request_id]] == list(range(1, expected_epochs + 1)) + _compare_outputs( + requests[request_id], + goldens[request_id], + requests[request_id].tensor_names(), + self.RTOL, + self.ATOL, + ) + finally: + for request_buffers in reversed(rehosted): + request_buffers.release() + + def generate_args(self, params): + specs = [ + Tensor(name, value.share_memory_()) if isinstance(value, torch.Tensor) else Scalar(name, value) + for name, value in _pa_generate_inputs(params) + ] + specs.extend( + [ + Scalar("layer_count", params["layer_count"]), + Scalar("layers_per_epoch", params.get("layers_per_epoch", 1)), + ] + ) + return TaskArgsBuilder(*specs) + + def compute_golden(self, args, params): + tensors = {spec.name: spec.value for spec in args.specs if isinstance(spec, Tensor)} + _pa_compute_golden(tensors, params) + for spec in args.specs: + if isinstance(spec, Tensor) and spec.name in tensors: + getattr(args, spec.name)[:] = tensors[spec.name] + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_dependency.py b/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_dependency.py index ae14b9dd71..9913413003 100644 --- a/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_dependency.py +++ b/tests/st/a2a3/tensormap_and_ringbuffer/test_l3_dependency.py @@ -7,19 +7,21 @@ # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. # See LICENSE in the root of the software repository for the full text of the License. # ----------------------------------------------------------------------------------------------------------- -"""L3 ChipTask → SubTask dependency via TensorMap. +"""L3 ChipTask -> SubTask dependency via TensorMap. Worker(level=3) submits a ChipTask then a SubTask that depends on it. Verifies: TensorMap dependency inference, cross-fork data visibility, SubWorker reads result produced by ChipWorker. """ +import threading + import torch from simpler.task_interface import ArgDirection as D from simpler.task_interface import TaskArgs, TensorArgType from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, make_tensor_arg, scene_test -from simpler_setup.scene_test import _build_l3_task_args +from simpler_setup.scene_test import CallableNamespace, _build_l3_task_args, _compare_outputs, _RehostedTaskArgs KERNELS_BASE = "../../../../examples/a2a3/tensormap_and_ringbuffer/vector_example/kernels" @@ -88,8 +90,51 @@ class TestL3Dependency(SceneTestCase): "config": {"device_count": 1, "num_sub_workers": 1, "block_dim": 3, "aicpu_thread_num": 4}, "params": {}, }, + { + "name": "ConcurrentTwoRequestPipeline", + "platforms": ["a2a3"], + "config": {"device_count": 1, "num_sub_workers": 1, "block_dim": 3, "aicpu_thread_num": 4}, + "manual": True, + "params": {}, + }, ] + def _run_and_validate_l3(self, worker, compiled_callables, sub_handles, case, **kwargs): + if case["name"] != "ConcurrentTwoRequestPipeline": + return super()._run_and_validate_l3(worker, compiled_callables, sub_handles, case, **kwargs) + + requests = [self.generate_args(case["params"]) for _ in range(2)] + goldens = [request.clone() for request in requests] + for golden in goldens: + self.compute_golden(golden, case["params"]) + + rehosted = [] + try: + rehosted = [_RehostedTaskArgs(worker, request) for request in requests] + config = self._build_config(case["config"]) + callables = CallableNamespace({**compiled_callables, **sub_handles}) + ready = threading.Barrier(2, timeout=30.0) + + def request_orch(orch, request, request_id, emitter, cfg): + chip_args, _ = _build_l3_task_args(request, callables.vector_kernel_sig) + callables.keep(chip_args) + ready.wait() + orch.submit_next_level(callables.vector_kernel, chip_args, cfg, worker=0) + + with worker.open_request_session(request_orch, max_active_runs=2) as session: + streams = [ + session.submit(request, config=config, request_id=request_id) + for request_id, request in enumerate(requests, start=1) + ] + for stream in streams: + stream.wait(timeout=30.0) + + for request, golden in zip(requests, goldens): + _compare_outputs(request, golden, request.tensor_names(), self.RTOL, self.ATOL) + finally: + for request_buffers in reversed(rehosted): + request_buffers.release() + def generate_args(self, params): SIZE = 128 * 128 return TaskArgsBuilder( diff --git a/tests/ut/cpp/common/test_trb_runtime_temp_buffer.cpp b/tests/ut/cpp/common/test_trb_runtime_temp_buffer.cpp index 30dba36523..bc0308b538 100644 --- a/tests/ut/cpp/common/test_trb_runtime_temp_buffer.cpp +++ b/tests/ut/cpp/common/test_trb_runtime_temp_buffer.cpp @@ -178,24 +178,24 @@ uint64_t fake_upload_chip_callable_buffer(const void * /* callable */) { return // with_temporary_buffer=false leaves the retained-slot callbacks null, which // makes trb bind fall back to a per-tensor device_malloc for each tensor. HostApi make_host_api(bool with_temporary_buffer = true) { - return HostApi{ - fake_device_malloc, - fake_device_free, - fake_copy_to_device, - fake_copy_from_device, - fake_register_device_memory_to_host, - fake_unregister_device_memory_from_host, - fake_device_memset, - with_temporary_buffer ? fake_get_retained_temp_buffer : nullptr, - with_temporary_buffer ? fake_set_retained_temp_buffer : nullptr, - fake_setup_static_arena, - fake_acquire_pooled_gm_heap, - fake_acquire_pooled_gm_sm, - fake_acquire_pooled_runtime_arena, - fake_lookup_prebuilt_runtime_arena_cache, - fake_mark_prebuilt_runtime_arena_cached, - fake_upload_chip_callable_buffer, - }; + HostApi api{}; + api.device_malloc = fake_device_malloc; + api.device_free = fake_device_free; + api.copy_to_device = fake_copy_to_device; + api.copy_from_device = fake_copy_from_device; + api.register_device_memory_to_host = fake_register_device_memory_to_host; + api.unregister_device_memory_from_host = fake_unregister_device_memory_from_host; + api.device_memset = fake_device_memset; + api.get_retained_temp_buffer = with_temporary_buffer ? fake_get_retained_temp_buffer : nullptr; + api.set_retained_temp_buffer = with_temporary_buffer ? fake_set_retained_temp_buffer : nullptr; + api.setup_static_arena = fake_setup_static_arena; + api.acquire_pooled_gm_heap = fake_acquire_pooled_gm_heap; + api.acquire_pooled_gm_sm = fake_acquire_pooled_gm_sm; + api.acquire_pooled_runtime_arena = fake_acquire_pooled_runtime_arena; + api.lookup_prebuilt_runtime_arena_cache = fake_lookup_prebuilt_runtime_arena_cache; + api.mark_prebuilt_runtime_arena_cached = fake_mark_prebuilt_runtime_arena_cached; + api.upload_chip_callable_buffer = fake_upload_chip_callable_buffer; + return api; } Tensor make_tensor(std::vector &storage, bool child_memory = false) { diff --git a/tests/ut/cpp/hierarchical/test_scheduler.cpp b/tests/ut/cpp/hierarchical/test_scheduler.cpp index 7383744692..7e57080a1b 100644 --- a/tests/ut/cpp/hierarchical/test_scheduler.cpp +++ b/tests/ut/cpp/hierarchical/test_scheduler.cpp @@ -57,6 +57,7 @@ struct MockMailboxWorker { struct Record { uint8_t callable_hash0; uint64_t tensor_key; // first tensor's `data` field (unique per submit in tests) + uint64_t generation; }; alignas(8) std::array mailbox{}; @@ -70,6 +71,7 @@ struct MockMailboxWorker { std::string next_error_msg; std::atomic is_running{false}; std::atomic stop_flag{false}; + std::atomic protocol_valid{true}; std::thread loop_thread; void start() { @@ -146,6 +148,11 @@ struct MockMailboxWorker { if (stop_flag.load(std::memory_order_acquire)) return; MailboxState s = read_state(); if (s == MailboxState::TASK_READY) { + uint64_t protocol = 0; + std::memcpy(&protocol, mailbox.data() + MAILBOX_OFF_PROTOCOL, sizeof(uint64_t)); + protocol_valid.store(protocol == MAILBOX_PROTOCOL_MAGIC_VERSION, std::memory_order_release); + uint64_t generation = 0; + std::memcpy(&generation, mailbox.data() + MAILBOX_OFF_GENERATION, sizeof(uint64_t)); uint8_t callable_hash0 = static_cast(mailbox[MAILBOX_OFF_TASK_CALLABLE_HASH]); int32_t t_count = 0; std::memcpy(&t_count, mailbox.data() + MAILBOX_OFF_TASK_ARGS_BLOB, sizeof(int32_t)); @@ -159,7 +166,7 @@ struct MockMailboxWorker { } { std::lock_guard lk(dispatched_mu); - dispatched.push_back({callable_hash0, tensor_key}); + dispatched.push_back({callable_hash0, tensor_key, generation}); } is_running.store(true, std::memory_order_release); @@ -209,12 +216,13 @@ struct MockMailboxWorker { class FakeEndpoint final : public WorkerEndpoint { public: - explicit FakeEndpoint(int32_t worker_id, std::atomic *prepare_count = nullptr) : + explicit FakeEndpoint(int32_t worker_id, std::atomic *prepare_count = nullptr, uint32_t max_in_flight = 1) : prepare_count_(prepare_count) { caps_.kind = WorkerEndpointKind::REMOTE_L3; caps_.worker_id = worker_id; caps_.remote = true; caps_.transport = "test-remote"; + caps_.max_in_flight = max_in_flight; } const WorkerEndpointCaps &caps() const override { return caps_; } @@ -237,6 +245,127 @@ class FakeEndpoint final : public WorkerEndpoint { std::atomic *prepare_count_{nullptr}; }; +class BlockingEndpoint final : public WorkerEndpoint { +public: + BlockingEndpoint() { + caps_.kind = WorkerEndpointKind::REMOTE_L3; + caps_.worker_id = 9; + caps_.max_in_flight = 2; + } + + const WorkerEndpointCaps &caps() const override { return caps_; } + + WorkerCompletion run(Ring *, const WorkerDispatch &dispatch) override { + { + std::unique_lock lk(mu_); + ++entered_; + cv_.notify_all(); + cv_.wait(lk, [this] { + return released_; + }); + } + return WorkerCompletion{dispatch.task_slot, dispatch.group_index, EndpointOutcome::SUCCESS, ""}; + } + + bool wait_entered(int expected, int timeout_ms = 500) { + std::unique_lock lk(mu_); + return cv_.wait_for(lk, std::chrono::milliseconds(timeout_ms), [this, expected] { + return entered_ >= expected; + }); + } + + void release() { + std::lock_guard lk(mu_); + released_ = true; + cv_.notify_all(); + } + +private: + WorkerEndpointCaps caps_; + std::mutex mu_; + std::condition_variable cv_; + int entered_{0}; + bool released_{false}; +}; + +class TwoSlotMailboxChild { +public: + alignas(8) std::array mailbox{}; + + void start() { + for (size_t slot = 0; slot < MAILBOX_TASK_SLOT_COUNT; ++slot) { + write_state(slot, MailboxState::IDLE); + threads_.emplace_back(&TwoSlotMailboxChild::loop, this, slot); + } + } + + ~TwoSlotMailboxChild() { + stop_.store(true, std::memory_order_release); + release(); + for (auto &thread : threads_) { + if (thread.joinable()) thread.join(); + } + } + + void *mailbox_ptr() { return mailbox.data(); } + + bool wait_active(int expected, int timeout_ms = 500) { + auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(timeout_ms); + while (active_.load(std::memory_order_acquire) < expected && std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + return active_.load(std::memory_order_acquire) >= expected; + } + + void release() { released_.store(true, std::memory_order_release); } + uint32_t seen_mask() const { return seen_mask_.load(std::memory_order_acquire); } + bool protocol_valid() const { return protocol_valid_.load(std::memory_order_acquire); } + +private: + std::vector threads_; + std::atomic stop_{false}; + std::atomic released_{false}; + std::atomic active_{0}; + std::atomic seen_mask_{0}; + std::atomic protocol_valid_{true}; + + char *frame(size_t slot) { return mailbox.data() + slot * MAILBOX_TASK_SLOT_SIZE; } + + MailboxState read_state(size_t slot) const { + const char *base = mailbox.data() + slot * MAILBOX_TASK_SLOT_SIZE; + const auto *ptr = reinterpret_cast(base + MAILBOX_OFF_STATE); + return static_cast(__atomic_load_n(ptr, __ATOMIC_ACQUIRE)); + } + + void write_state(size_t slot, MailboxState state) { + auto *ptr = reinterpret_cast(frame(slot) + MAILBOX_OFF_STATE); + __atomic_store_n(ptr, static_cast(state), __ATOMIC_RELEASE); + } + + void loop(size_t slot) { + while (!stop_.load(std::memory_order_acquire)) { + if (read_state(slot) != MailboxState::TASK_READY) { + std::this_thread::sleep_for(std::chrono::microseconds(50)); + continue; + } + uint64_t protocol = 0; + std::memcpy(&protocol, frame(slot) + MAILBOX_OFF_PROTOCOL, sizeof(protocol)); + if (protocol != MAILBOX_PROTOCOL_MAGIC_VERSION) { + protocol_valid_.store(false, std::memory_order_release); + } + seen_mask_.fetch_or(1u << slot, std::memory_order_acq_rel); + active_.fetch_add(1, std::memory_order_acq_rel); + while (!released_.load(std::memory_order_acquire) && !stop_.load(std::memory_order_acquire)) { + std::this_thread::sleep_for(std::chrono::microseconds(50)); + } + int32_t zero_error = 0; + std::memcpy(frame(slot) + MAILBOX_OFF_ERROR, &zero_error, sizeof(zero_error)); + active_.fetch_sub(1, std::memory_order_acq_rel); + write_state(slot, MailboxState::TASK_DONE); + } + } +}; + // --------------------------------------------------------------------------- // Helper: build a TaskArgs whose only tensor has the given (data, tag). // --------------------------------------------------------------------------- @@ -374,6 +503,75 @@ TEST(WorkerManagerTest, ControlPrepareUsesStableNextLevelWorkerId) { EXPECT_EQ(worker3_prepares.load(std::memory_order_relaxed), 1); } +TEST(WorkerThreadTest, EndpointCreditsRunTwoDispatchesConcurrently) { + Ring allocator; + allocator.init(/*heap_bytes=*/0); + WorkerThread worker; + auto endpoint = std::make_unique(); + BlockingEndpoint *endpoint_ptr = endpoint.get(); + std::atomic completed{0}; + + worker.start( + &allocator, nullptr, + [&completed](WorkerCompletion) { + completed.fetch_add(1, std::memory_order_release); + }, + std::move(endpoint) + ); + EXPECT_EQ(worker.capacity(), 2u); + worker.dispatch(WorkerDispatch{1, 0}); + worker.dispatch(WorkerDispatch{2, 0}); + EXPECT_TRUE(endpoint_ptr->wait_entered(2)); + EXPECT_FALSE(worker.idle()); + + endpoint_ptr->release(); + auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(500); + while (completed.load(std::memory_order_acquire) != 2 && std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + EXPECT_EQ(completed.load(std::memory_order_acquire), 2); + EXPECT_TRUE(worker.idle()); + worker.stop(); + allocator.shutdown(); +} + +TEST(LocalMailboxEndpointTest, TwoTaskSlotsCompleteConcurrently) { + Ring allocator; + allocator.init(/*heap_bytes=*/0); + std::array slots{}; + for (size_t i = 0; i < slots.size(); ++i) { + AllocResult result = allocator.alloc(); + ASSERT_NE(result.slot, INVALID_SLOT); + slots[i] = result.slot; + TaskSlotState *state = allocator.slot_state(result.slot); + ASSERT_NE(state, nullptr); + state->callable = C(static_cast(40 + i)); + state->task_args = single_tensor_args(0x1000 + i, TensorArgType::OUTPUT); + } + + TwoSlotMailboxChild child; + child.start(); + LocalMailboxEndpoint endpoint(0, child.mailbox_ptr(), 2); + std::array completions{}; + std::array dispatchers; + for (size_t i = 0; i < dispatchers.size(); ++i) { + dispatchers[i] = std::thread([&, i]() { + completions[i] = endpoint.run(&allocator, WorkerDispatch{slots[i], 0}); + }); + } + + bool both_active = child.wait_active(2); + child.release(); + for (auto &thread : dispatchers) + thread.join(); + EXPECT_TRUE(both_active); + EXPECT_TRUE(child.protocol_valid()); + EXPECT_EQ(child.seen_mask(), 0x3u); + for (const auto &completion : completions) + EXPECT_EQ(completion.outcome, EndpointOutcome::SUCCESS); + allocator.shutdown(); +} + TEST_F(SchedulerFixture, IndependentTaskDispatchedAndConsumed) { auto args_a = single_tensor_args(0xCAFE, TensorArgType::OUTPUT); auto res = orch.submit_next_level(C(42), args_a, cfg, 0); @@ -383,6 +581,8 @@ TEST_F(SchedulerFixture, IndependentTaskDispatchedAndConsumed) { ASSERT_GE(mock_worker.dispatched_count(), 1); EXPECT_EQ(mock_worker.dispatched[0].tensor_key, 0xCAFEu); EXPECT_EQ(mock_worker.dispatched[0].callable_hash0, 42u); + EXPECT_TRUE(mock_worker.protocol_valid.load(std::memory_order_acquire)); + EXPECT_GT(mock_worker.dispatched[0].generation, 0u); mock_worker.complete(); wait_consumed(slot); @@ -406,6 +606,7 @@ TEST_F(SchedulerFixture, DependentTaskDispatchedAfterProducerCompletes) { } ASSERT_GE(mock_worker.dispatched_count(), 2); EXPECT_EQ(mock_worker.dispatched[1].callable_hash0, 11u); + EXPECT_EQ(mock_worker.dispatched[1].generation, mock_worker.dispatched[0].generation + 1); mock_worker.complete(); // B done wait_consumed(b.task_slot); diff --git a/tests/ut/cpp/hierarchical/test_scope.cpp b/tests/ut/cpp/hierarchical/test_scope.cpp index 273b33bfc8..9550cfc1f3 100644 --- a/tests/ut/cpp/hierarchical/test_scope.cpp +++ b/tests/ut/cpp/hierarchical/test_scope.cpp @@ -11,6 +11,10 @@ #include +#include +#include +#include + #include "scope.h" TEST(Scope, InitialDepthIsZero) { @@ -81,3 +85,61 @@ TEST(Scope, EmptyScopeReleasesNothing) { }); EXPECT_EQ(calls, 0); } + +TEST(Scope, ConcurrentThreadsOwnIndependentFrames) { + Scope sc; + std::mutex mu; + std::condition_variable cv; + int stage = 0; + std::vector released_a; + std::vector released_b; + + std::thread a([&]() { + sc.scope_begin(); + sc.register_task(10); + { + std::unique_lock lk(mu); + stage = 1; + cv.notify_all(); + cv.wait(lk, [&]() { + return stage >= 2; + }); + } + sc.scope_end([&](TaskSlot slot) { + released_a.push_back(slot); + }); + { + std::lock_guard lk(mu); + stage = 3; + cv.notify_all(); + } + }); + std::thread b([&]() { + { + std::unique_lock lk(mu); + cv.wait(lk, [&]() { + return stage >= 1; + }); + } + sc.scope_begin(); + sc.register_task(20); + { + std::unique_lock lk(mu); + stage = 2; + cv.notify_all(); + cv.wait(lk, [&]() { + return stage >= 3; + }); + } + sc.scope_end([&](TaskSlot slot) { + released_b.push_back(slot); + }); + }); + + a.join(); + b.join(); + ASSERT_EQ(released_a.size(), 1u); + ASSERT_EQ(released_b.size(), 1u); + EXPECT_EQ(released_a[0], 10); + EXPECT_EQ(released_b[0], 20); +} diff --git a/tests/ut/py/test_worker/test_ensure_prepared.py b/tests/ut/py/test_worker/test_ensure_prepared.py index 2d6e5d3ff8..63442f475c 100644 --- a/tests/ut/py/test_worker/test_ensure_prepared.py +++ b/tests/ut/py/test_worker/test_ensure_prepared.py @@ -6,18 +6,18 @@ # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. # See LICENSE in the root of the software repository for the full text of the License. # ----------------------------------------------------------------------------------------------------------- -"""Unit tests for the module-level ``_ensure_prepared`` helper. +"""Unit tests for the module-level ``_ensure_registered`` helper. -``_ensure_prepared`` is the ``_CTRL_PREPARE`` registration path: it stages a -callable's slot exactly once. There is no lazy/run-time preparation — the run -path (``_TASK_READY``) only consumes an already-prepared slot and raises if it +``_ensure_registered`` is the ``_CTRL_PREPARE`` registration path: it stages a +callable's slot exactly once. There is no lazy run-time registration: the run +path (``_TASK_READY``) only consumes an already-registered slot and raises if it is missing, so that behavior is covered by the chip-loop tests, not here. """ from unittest.mock import MagicMock import pytest -from simpler.worker import _ensure_prepared +from simpler.worker import _ensure_registered class TestEnsurePrepared: @@ -25,24 +25,24 @@ def test_prepares_slot(self, capsys): cw = MagicMock() callable_obj = object() registry = {2: callable_obj} - prepared: set[int] = set() + registered: set[int] = set() - _ensure_prepared(cw, registry, prepared, 2, device_id=0) + _ensure_registered(cw, registry, registered, 2, device_id=0) cw._register_callable_at_slot.assert_called_once_with(2, callable_obj) - assert prepared == {2} + assert registered == {2} # Preparation is a normal control-flow step now — no warning is emitted. assert capsys.readouterr().err == "" - def test_already_prepared_short_circuits(self): + def test_already_registered_short_circuits(self): cw = MagicMock() - # cid is already in `prepared`; helper must skip lookup and - # slot preparation entirely. Pass an empty registry to prove the + # cid is already in `registered`; helper must skip lookup and + # slot registration entirely. Pass an empty registry to prove the # lookup never happens (otherwise registry.get would return None # and the helper would raise). - _ensure_prepared(cw, {}, {5}, 5, device_id=0) + _ensure_registered(cw, {}, {5}, 5, device_id=0) cw._register_callable_at_slot.assert_not_called() def test_missing_cid_raises(self): with pytest.raises(RuntimeError, match="cid 9 not in registry"): - _ensure_prepared(MagicMock(), {}, set(), 9, device_id=0) + _ensure_registered(MagicMock(), {}, set(), 9, device_id=0) diff --git a/tests/ut/py/test_worker/test_host_buffer_registration.py b/tests/ut/py/test_worker/test_host_buffer_registration.py index 4027243ba2..664a782230 100644 --- a/tests/ut/py/test_worker/test_host_buffer_registration.py +++ b/tests/ut/py/test_worker/test_host_buffer_registration.py @@ -164,6 +164,7 @@ def test_l3_sub_worker_maps_rewrites_and_unmaps_host_buffer(monkeypatch): struct.pack_into(" SharedMemory: return shm -def test_chip_process_loop_inits_runs_and_finalizes(monkeypatch): +@pytest.mark.parametrize(("runtime", "expected_slots"), [("tensormap_and_ringbuffer", 2), ("host_build_graph", 2)]) +def test_chip_process_loop_inits_one_worker_with_runtime_slots(monkeypatch, runtime, expected_slots): events: list[tuple] = [] class FakeChipWorker: + pipeline_slot_count = expected_slots + def init(self, device_id, bins, *, log_level, log_info_v, prewarm_config=None, enable_sdma=False): events.append(("init", device_id, bins, log_level, log_info_v, prewarm_config, enable_sdma)) def finalize(self) -> None: events.append(("finalize",)) - def fake_run_chip_main_loop(cw, *_args, chip_platform, chip_runtime, prepared=None): - events.append(("main_loop", cw, chip_platform, chip_runtime)) + def fake_run_chip_main_loop( + cw, *_args, chip_platform, chip_runtime, prepared=None, run_workers=None, pipeline_slots=1 + ): + events.append(("main_loop", cw, chip_platform, chip_runtime, len(run_workers or [cw]), pipeline_slots)) monkeypatch.setattr(worker_mod, "ChipWorker", FakeChipWorker) monkeypatch.setattr(worker_mod, "_run_chip_main_loop", fake_run_chip_main_loop) @@ -150,16 +155,17 @@ def fake_run_chip_main_loop(cw, *_args, chip_platform, chip_runtime, prepared=No {}, {}, platform="a2a3", - runtime="tensormap_and_ringbuffer", + runtime=runtime, ) finally: shm.close() shm.unlink() - assert events[0] == ("init", 7, "bins", 1, 5, None, False) - assert events[1][0] == "main_loop" - assert events[1][2:] == ("a2a3", "tensormap_and_ringbuffer") - assert events[2] == ("finalize",) + assert [event[0] for event in events].count("init") == 1 + main_loop_index = 1 + assert events[main_loop_index][0] == "main_loop" + assert events[main_loop_index][2:] == ("a2a3", runtime, 1, expected_slots) + assert events[main_loop_index + 1 :] == [("finalize",)] def _chip_digest(callable_obj: ChipCallable, *, platform: str = "", runtime: str = "") -> bytes: @@ -1863,6 +1869,60 @@ def _spawn_loop(cw, buf, state_addr, registry=None, identity_table=None, identit t.start() return t + def test_single_slot_runs_on_chip_main_thread(self): + from unittest.mock import MagicMock # noqa: PLC0415 + + digest = b"\x05" * 32 + cw = MagicMock() + run_threads = [] + cw._impl.run_from_blob.side_effect = lambda *_args: run_threads.append(threading.get_ident()) + shm, buf, state_addr = self._build_mailbox() + monitor_errors = [] + + def shutdown_after_done(): + import time # noqa: PLC0415 + + deadline = time.monotonic() + 2.0 + try: + while _mailbox_load_i32(state_addr) != worker_mod._TASK_DONE: + if time.monotonic() >= deadline: + raise TimeoutError("single-slot run did not publish TASK_DONE") + time.sleep(0.001) + except BaseException as exc: # noqa: BLE001 + monitor_errors.append(exc) + finally: + _mailbox_store_i32(state_addr, worker_mod._SHUTDOWN) + + try: + struct.pack_into("Q", buf, worker_mod.MAILBOX_OFF_PROTOCOL, worker_mod.MAILBOX_PROTOCOL_MAGIC_VERSION) + start = worker_mod._OFF_TASK_CALLABLE_HASH + buf[start : start + len(digest)] = digest + monitor = threading.Thread(target=shutdown_after_done) + monitor.start() + caller_thread = threading.get_ident() + _mailbox_store_i32(state_addr, worker_mod._TASK_READY) + worker_mod._run_chip_main_loop( + cw, + buf, + _mailbox_addr(shm), + state_addr, + 0, + {0: object()}, + {digest: 0}, + {digest: 1}, + chip_platform="a5", + prepared={0}, + run_workers=[cw], + ) + monitor.join(timeout=2.0) + assert not monitor.is_alive() + assert not monitor_errors + assert run_threads == [caller_thread] + finally: + buf.release() + shm.close() + shm.unlink() + def test_register_uses_payload_size_and_allocates_local_slot(self): from unittest.mock import MagicMock # noqa: PLC0415 diff --git a/tests/ut/py/test_worker/test_request_session.py b/tests/ut/py/test_worker/test_request_session.py new file mode 100644 index 0000000000..c7e50142d8 --- /dev/null +++ b/tests/ut/py/test_worker/test_request_session.py @@ -0,0 +1,225 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +import threading +import time +from typing import cast + +import pytest +import simpler.worker as worker_mod +from simpler.request_session import RequestCancelledError, RequestEmitter, RequestSession, RequestStream +from simpler.worker import Worker + + +class _StreamSession: + def cancel(self, _request_id: int) -> bool: + return False + + +class _ConcurrentWorker: + def __init__(self) -> None: + self._lock = threading.Lock() + self.active = 0 + self.max_active = 0 + self.both_started = threading.Event() + self.release = threading.Event() + + def run(self, task_orch, args=None, config=None) -> None: + with self._lock: + self.active += 1 + self.max_active = max(self.max_active, self.active) + if self.active == 2: + self.both_started.set() + try: + task_orch(None, args, config) + assert self.release.wait(1.0) + finally: + with self._lock: + self.active -= 1 + + def _release_request_session(self, _session) -> None: + pass + + +class _ConcurrentOrchestrator: + def __init__(self) -> None: + self._lock = threading.Lock() + self.clear_count = 0 + self.drain_count = 0 + + def _clear_error(self) -> None: + with self._lock: + self.clear_count += 1 + + def _scope_begin(self) -> None: + pass + + def _scope_end(self) -> None: + pass + + def _drain(self) -> None: + with self._lock: + self.drain_count += 1 + + +class _InlineWorker: + def run(self, task_orch, args=None, config=None) -> None: + task_orch(None, args, config) + + def _release_request_session(self, _session) -> None: + pass + + +def test_emit_waits_until_consumer_receives_token() -> None: + stream = RequestStream(7, cast(RequestSession, _StreamSession())) + emitter = RequestEmitter(stream) + producer_done = threading.Event() + + def produce() -> None: + emitter.emit("token-1") + producer_done.set() + + thread = threading.Thread(target=produce) + thread.start() + time.sleep(0.02) + assert not producer_done.is_set() + assert stream.next(timeout=1.0) == "token-1" + assert producer_done.wait(1.0) + stream._finish(None) + thread.join() + + +def test_close_cancels_unconsumed_stream_and_joins_dispatcher() -> None: + callback_entered = threading.Event() + + def request_orch(_orch, _request, _request_id, emitter, _config) -> None: + callback_entered.set() + emitter.emit("unconsumed-token") + + session = RequestSession(_InlineWorker(), request_orch) + session._start() + stream = session.submit(object(), request_id=11) + assert callback_entered.wait(1.0) + + close_done = threading.Event() + + def close() -> None: + session.close() + close_done.set() + + closer = threading.Thread(target=close, daemon=True) + closer.start() + assert close_done.wait(1.0) + closer.join() + + with pytest.raises(RequestCancelledError, match="cancelled by session close"): + stream.wait(timeout=0.1) + assert all(not thread.is_alive() for thread in session._threads) + + +def test_cancel_unblocks_emit_and_allows_request_id_reuse() -> None: + callback_entered = threading.Event() + + def request_orch(_orch, request, _request_id, emitter, _config) -> None: + if request == "blocking": + callback_entered.set() + emitter.emit("unconsumed-token") + + session = RequestSession(_InlineWorker(), request_orch) + session._start() + try: + stream = session.submit("blocking", request_id=23) + assert callback_entered.wait(1.0) + assert stream.cancel() + with pytest.raises(RequestCancelledError, match="request 23 was cancelled"): + stream.wait(timeout=1.0) + + deadline = time.monotonic() + 1.0 + while 23 in session._live_work and time.monotonic() < deadline: + time.sleep(0.001) + assert 23 not in session._live_work + + replacement = session.submit("complete", request_id=23) + replacement.wait(timeout=1.0) + finally: + session.close() + + +def test_session_runs_two_submitted_requests_concurrently() -> None: + worker = _ConcurrentWorker() + session = RequestSession(worker, lambda *_args: None, max_active_runs=2) + session._start() + try: + stream_a = session.submit(object(), request_id=1) + stream_b = session.submit(object()) + assert stream_b.request_id == 2 + assert worker.both_started.wait(1.0) + assert worker.max_active == 2 + worker.release.set() + stream_a.wait(timeout=1.0) + stream_b.wait(timeout=1.0) + finally: + worker.release.set() + session.close() + + +def test_worker_close_closes_active_request_session() -> None: + worker = Worker(level=3, device_ids=[]) + worker.init() + session = worker.open_request_session(lambda *_args: None) + + worker.close() + + assert session._closed.is_set() + assert worker._request_session is None + assert all(not thread.is_alive() for thread in session._threads) + + +def test_worker_groups_overlapping_run_calls_into_one_drain() -> None: + worker = Worker(level=3, device_ids=[]) + orch = _ConcurrentOrchestrator() + worker._lifecycle = worker_mod._Lifecycle.READY + worker._worker = object() + worker._orch = orch + worker._start_hierarchical = lambda: None + worker._release_active_remote_slot_refs = lambda: None + worker._flush_pending_remote_frees = lambda: None + worker._cleanup_l3_l2_regions = lambda: None + worker._execute_pending_domain_releases = lambda: None + started = 0 + started_lock = threading.Lock() + both_started = threading.Event() + release = threading.Event() + errors = [] + + def task_orch(_orch, _args, _config) -> None: + nonlocal started + with started_lock: + started += 1 + if started == 2: + both_started.set() + assert release.wait(1.0) + + def run() -> None: + try: + worker.run(task_orch) + except BaseException as exc: # noqa: BLE001 + errors.append(exc) + + threads = [threading.Thread(target=run), threading.Thread(target=run)] + for thread in threads: + thread.start() + assert both_started.wait(1.0) + release.set() + for thread in threads: + thread.join() + + assert not errors + assert orch.clear_count == 1 + assert orch.drain_count == 1 diff --git a/tests/ut/py/test_worker/test_startup_readiness.py b/tests/ut/py/test_worker/test_startup_readiness.py index abaa718c69..813d78def5 100644 --- a/tests/ut/py/test_worker/test_startup_readiness.py +++ b/tests/ut/py/test_worker/test_startup_readiness.py @@ -934,6 +934,26 @@ def capture_reap(groups, deadline): # fixed at teardown entry. assert captured["remaining"] > 0.7 + def test_comm_reap_deadline_covers_native_destroy_budget(self, monkeypatch): + import types # noqa: PLC0415 + + import simpler.worker as worker_mod # noqa: PLC0415 + + monkeypatch.setattr(worker_mod, "_ROLLBACK_GRACEFUL_TIMEOUT_S", 1.0) + monkeypatch.setattr(worker_mod, "_COMM_REAP_GRACEFUL_TIMEOUT_S", 3.0) + w = Worker(level=3, num_sub_workers=0) + w._worker = types.SimpleNamespace(close=lambda: None) + w._comm_base_ready = True + captured = {} + + def capture_reap(groups, deadline): + captured["remaining"] = deadline - time.monotonic() + + monkeypatch.setattr(Worker, "_reap_child_groups", staticmethod(capture_reap)) + with _hard_timeout(_TEST_WALL_BUDGET_S): + w._teardown_ready_tree() + assert captured["remaining"] > 2.7 + def test_reap_child_groups_stuck_child_no_starvation(self, monkeypatch): # A stuck child in one group must not starve the reap of healthy children # in another. With every group SHUTDOWN up front, the interleaved reap diff --git a/tools/cross_request_host_device_perfetto.py b/tools/cross_request_host_device_perfetto.py new file mode 100644 index 0000000000..230ba96638 --- /dev/null +++ b/tools/cross_request_host_device_perfetto.py @@ -0,0 +1,358 @@ +#!/usr/bin/env python3 +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Build a minimal token-level O/S Perfetto trace for concurrent requests.""" + +from __future__ import annotations + +import argparse +import json +import re +from pathlib import Path + +STRACE_RE = re.compile( + r"\[STRACE\].*?\binv=(?P\d+).*?\bname=(?P\S+) " + r"ts=(?P\d+) dur=(?P\d+)(?: (?P.*))?$" +) +TOKEN_RE = re.compile( + r"L3_TOKEN_RECEIVED request_id=(?P\d+) token_seq=(?P\d+) " + r"received_ns=(?P\d+) final=(?PTrue|False)" +) +DELIVERY_RE = re.compile( + r"USER_TOKEN_DELIVERED request_id=(?P\d+) token_seq=(?P\d+) " + r"delivered_ns=(?P\d+) final=(?PTrue|False)" +) +ORCH_SUFFIXES = ("layer_build", "epoch.materialize", "epoch.stage_upload") +TRACE_PID = 1000 + + +def parse_attrs(raw: str | None) -> dict[str, str]: + result: dict[str, str] = {} + for token in (raw or "").split(): + if "=" in token: + key, value = token.split("=", 1) + result[key] = value + return result + + +def overlap_ms(first: tuple[int, int], second: tuple[int, int]) -> float: + return max(0, min(first[1], second[1]) - max(first[0], second[0])) / 1_000_000.0 + + +def token_epoch(span: dict) -> int | None: + attrs = span["attrs"] + if span["name"].endswith("layer_build") and "layer" in attrs: + return int(attrs["layer"]) + 1 + if "epoch" in attrs: + return int(attrs["epoch"]) + return None + + +def token_orch_windows(spans: list[dict]) -> dict[int, tuple[int, int]]: + grouped: dict[int, list[tuple[int, int]]] = {} + for span in spans: + if not span["name"].startswith("simpler_run.host_orch.") or not span["name"].endswith(ORCH_SUFFIXES): + continue + epoch = token_epoch(span) + if epoch is not None: + grouped.setdefault(epoch, []).append((span["ts"], span["ts"] + span["dur"])) + return { + epoch: (min(start for start, _end in values), max(end for _start, end in values)) + for epoch, values in grouped.items() + } + + +def token_schedule_windows(spans: list[dict], tokens: list[dict]) -> dict[int, tuple[int, int]]: + """Build device epoch windows from release and completion observations.""" + release_times = { + int(span["attrs"]["target_epoch"]): span["ts"] + span["dur"] + for span in spans + if span["name"].endswith("epoch.wait_device") + and span["attrs"].get("kind") == "release" + and "target_epoch" in span["attrs"] + } + completion_times = { + int(span["attrs"]["target_epoch"]): span["ts"] + span["dur"] + for span in spans + if span["name"].endswith("epoch.wait_device") + and span["attrs"].get("kind") == "exec-done" + and "target_epoch" in span["attrs"] + } + receive_times = {token["seq"]: token["ts"] for token in tokens} + final_epoch = max(receive_times, default=None) + windows = {} + for epoch, start in sorted(release_times.items()): + if epoch not in receive_times: + continue + if epoch in completion_times: + end = completion_times[epoch] + elif epoch == final_epoch: + end = receive_times[epoch] + else: + raise ValueError(f"non-final token {epoch} has no device exec-done observation") + if end < start: + raise ValueError(f"token {epoch} completed before its device release observation") + windows[epoch] = (start, end) + return windows + + +def parse_log(log: Path) -> tuple[dict[int, list[dict]], list[dict], list[dict], list[int]]: + spans: dict[int, list[dict]] = {} + tokens: list[dict] = [] + deliveries: list[dict] = [] + for line in log.read_text(encoding="utf-8", errors="replace").splitlines(): + if match := STRACE_RE.search(line): + spans.setdefault(int(match.group("inv")), []).append( + { + "name": match.group("name"), + "ts": int(match.group("ts")), + "dur": int(match.group("dur")), + "attrs": parse_attrs(match.group("attrs")), + } + ) + if match := TOKEN_RE.search(line): + tokens.append( + { + "request": int(match.group("request")), + "seq": int(match.group("seq")), + "ts": int(match.group("ts")), + } + ) + if match := DELIVERY_RE.search(line): + deliveries.append( + { + "request": int(match.group("request")), + "seq": int(match.group("seq")), + "ts": int(match.group("ts")), + } + ) + request_ids = list(dict.fromkeys(token["request"] for token in tokens)) + if not spans or len(request_ids) < 2: + raise ValueError("log must contain STRACE spans and tokens for at least two requests") + return spans, tokens, deliveries, request_ids + + +def map_host_invocations(spans: dict[int, list[dict]], request_ids: list[int]) -> dict[int, int]: + normal_invs = [inv for inv, values in spans.items() if any(span["name"] == "simpler_run" for span in values)] + normal_invs.sort(key=lambda inv: next(span["ts"] for span in spans[inv] if span["name"] == "simpler_run")) + if len(normal_invs) != len(request_ids): + raise ValueError("log must contain one complete simpler_run invocation per request") + return dict(zip(request_ids, normal_invs)) + + +def complete( + tid: int, + name: str, + category: str, + window: tuple[int, int], + base: int, + *, + request_id: int, + token: int, +) -> dict: + return { + "ph": "X", + "pid": TRACE_PID, + "tid": tid, + "name": name, + "cat": category, + "cname": "rail_animation" if category == "O" else "good", + "ts": (window[0] - base) / 1000.0, + "dur": (window[1] - window[0]) / 1000.0, + "args": {"request_id": request_id, "token": token}, + } + + +def build_summary( + request_ids: list[int], + orch_by_request: dict[int, dict[int, tuple[int, int]]], + schedule_by_request: dict[int, dict[int, tuple[int, int]]], + tokens: list[dict], + deliveries: list[dict], +) -> dict: + token_parallel = {} + for label_index, request_id in enumerate(request_ids): + label = chr(ord("A") + label_index) + overlaps = { + f"O_token_{token}_vs_S_token_{token - 1}": overlap_ms( + orch_by_request[request_id][token], schedule_by_request[request_id][token - 1] + ) + for token in sorted(orch_by_request[request_id]) + if token - 1 in schedule_by_request[request_id] + } + token_parallel[f"Request_{label}"] = overlaps + + request_parallel = {} + for previous_index, (previous_request, current_request) in enumerate(zip(request_ids, request_ids[1:])): + previous_label = chr(ord("A") + previous_index) + current_label = chr(ord("A") + previous_index + 1) + for current_token, current_o in sorted(orch_by_request[current_request].items()): + for previous_token, previous_s in sorted(schedule_by_request[previous_request].items()): + overlap = overlap_ms(current_o, previous_s) + if overlap > 0: + request_parallel[ + f"Request_{current_label}_O_token_{current_token}_vs_Request_{previous_label}_S_token_{previous_token}" + ] = overlap + received = {(token["request"], token["seq"]): token["ts"] for token in tokens} + delivered = {(token["request"], token["seq"]): token["ts"] for token in deliveries} + shared_keys = received.keys() & delivered.keys() + delivery_ms = {key: (delivered[key] - received[key]) / 1_000_000.0 for key in shared_keys} + one_to_one = received.keys() == delivered.keys() and all(delay >= 0 for delay in delivery_ms.values()) + no_batching = one_to_one and all( + delivered[(request_id, seq)] < received[(request_id, seq + 1)] + for request_id, seq in shared_keys + if (request_id, seq + 1) in received + ) + return { + "how_to_read": "Horizontal overlap means parallel execution.", + "token_parallel_ms": token_parallel, + "request_parallel_ms": request_parallel, + "stream_delivery": { + "strict_one_to_one": one_to_one, + "user_consumes_before_next_l3_token": no_batching, + "l3_to_user_ms": { + f"Request_{chr(ord('A') + label_index)}": { + str(seq): delivery_ms[(request_id, seq)] for req, seq in sorted(delivery_ms) if req == request_id + } + for label_index, request_id in enumerate(request_ids) + }, + "max_l3_to_user_ms": max(delivery_ms.values(), default=None), + }, + "schedule_window": ( + "Device-observed epoch release to exec-done; the final epoch ends at final L3 token receipt " + "because no later Host wait observes its exec-done boundary." + ), + } + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--log", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + + spans, tokens, deliveries, request_ids = parse_log(args.log) + host_inv_by_request = map_host_invocations(spans, request_ids) + tokens_by_request = { + request_id: [token for token in tokens if token["request"] == request_id] for request_id in request_ids + } + orch_by_request = { + request_id: token_orch_windows(spans[host_inv_by_request[request_id]]) for request_id in request_ids + } + schedule_by_request = { + request_id: token_schedule_windows(spans[host_inv_by_request[request_id]], tokens_by_request[request_id]) + for request_id in request_ids + } + all_windows = [ + window + for request_id in request_ids + for windows in (orch_by_request[request_id], schedule_by_request[request_id]) + for window in windows.values() + ] + if not all_windows: + raise ValueError("log contains no token-level O/S windows") + base = min(start for start, _end in all_windows) + + events = [ + { + "ph": "M", + "name": "process_name", + "pid": TRACE_PID, + "tid": 0, + "args": {"name": "O/S Pipeline + per-token stream"}, + } + ] + deliveries_by_key = {(delivery["request"], delivery["seq"]): delivery for delivery in deliveries} + for label_index, request_id in enumerate(request_ids): + label = chr(ord("A") + label_index) + orch_tid = label_index * 3 + 10 + schedule_tid = orch_tid + 1 + stream_tid = orch_tid + 2 + events.extend( + [ + { + "ph": "M", + "name": "thread_name", + "pid": TRACE_PID, + "tid": orch_tid, + "args": {"name": f"Request {label} · O (Orch)"}, + }, + { + "ph": "M", + "name": "thread_name", + "pid": TRACE_PID, + "tid": schedule_tid, + "args": {"name": f"Request {label} · S (device epoch)"}, + }, + { + "ph": "M", + "name": "thread_name", + "pid": TRACE_PID, + "tid": stream_tid, + "args": {"name": f"Request {label} · Stream (L3 → User)"}, + }, + ] + ) + for token, window in sorted(orch_by_request[request_id].items()): + events.append(complete(orch_tid, f"O token {token}", "O", window, base, request_id=request_id, token=token)) + for token, window in sorted(schedule_by_request[request_id].items()): + events.append( + complete(schedule_tid, f"S token {token}", "S", window, base, request_id=request_id, token=token) + ) + for receipt in tokens_by_request[request_id]: + delivery = deliveries_by_key.get((request_id, receipt["seq"])) + if delivery is None: + continue + flow_id = request_id * 100 + receipt["seq"] + delay_ms = (delivery["ts"] - receipt["ts"]) / 1_000_000.0 + events.extend( + [ + { + "ph": "s", + "pid": TRACE_PID, + "tid": schedule_tid, + "name": f"L3 token {receipt['seq']}", + "cat": "stream_delivery", + "id": flow_id, + "ts": (receipt["ts"] - base) / 1000.0, + }, + { + "ph": "f", + "bp": "e", + "pid": TRACE_PID, + "tid": stream_tid, + "name": f"User token {receipt['seq']}", + "cat": "stream_delivery", + "id": flow_id, + "ts": (delivery["ts"] - base) / 1000.0, + }, + { + "ph": "i", + "s": "t", + "pid": TRACE_PID, + "tid": stream_tid, + "name": f"User token {receipt['seq']}", + "cat": "stream_delivery", + "cname": "rail_response", + "ts": (delivery["ts"] - base) / 1000.0, + "args": {"request_id": request_id, "token": receipt["seq"], "l3_to_user_ms": delay_ms}, + }, + ] + ) + + summary = build_summary(request_ids, orch_by_request, schedule_by_request, tokens, deliveries) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps({"traceEvents": events, "summary": summary}, indent=2) + "\n", encoding="utf-8") + print(json.dumps(summary, indent=2)) + print(args.output.resolve()) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/host_orch_perfetto.py b/tools/host_orch_perfetto.py new file mode 100644 index 0000000000..e3a6627bcd --- /dev/null +++ b/tools/host_orch_perfetto.py @@ -0,0 +1,542 @@ +#!/usr/bin/env python3 +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Build a compact Host-O / Device-S Perfetto trace from one profiled run.""" + +from __future__ import annotations + +import argparse +import json +import re +from datetime import datetime +from pathlib import Path +from typing import TypedDict + + +class StraceSpan(TypedDict): + name: str + ts_ns: int + dur_ns: int + attrs: dict[str, str] + + +class SchedulerLayer(TypedDict): + layer: int + start_cycles: int + end_cycles: int + start_ns: int + dur_ns: int + records: int + unique_tasks: int + + +STRACE_RE = re.compile( + r"\[STRACE\].*?\binv=(?P\d+).*?\bname=(?P\S+) " + r"ts=(?P\d+) dur=(?P\d+)(?: (?P.*))?$" +) +WALL_TIME_RE = re.compile(r"^\[(?P\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}\.\d+)\]") + + +def parse_attrs(raw: str | None) -> dict[str, str]: + attrs: dict[str, str] = {} + for token in (raw or "").split(): + if "=" in token: + key, value = token.split("=", 1) + attrs[key] = value + return attrs + + +def load_strace_spans(log_path: Path) -> dict[int, list[StraceSpan]]: + rounds: dict[int, list[StraceSpan]] = {} + for line in log_path.read_text(encoding="utf-8", errors="replace").splitlines(): + match = STRACE_RE.search(line) + if match is None: + continue + invocation = int(match.group("inv")) + rounds.setdefault(invocation, []).append( + { + "name": match.group("name"), + "ts_ns": int(match.group("ts")), + "dur_ns": int(match.group("dur")), + "attrs": parse_attrs(match.group("attrs")), + } + ) + if not rounds: + raise ValueError(f"no STRACE spans found in {log_path}") + return rounds + + +def wall_time_ns(raw: str) -> int: + stamp = datetime.strptime(raw, "%Y-%m-%d %H:%M:%S.%f") + epoch = datetime(1970, 1, 1) + return round((stamp - epoch).total_seconds() * 1_000_000_000) + + +def load_whole_graph_anchors(log_path: Path) -> dict[str, int]: + anchors: dict[str, int] = {} + patterns = { + "orch_ready": "Device orchestration ready:", + "profile_ready": "Performance profiling initialized", + "aicpu_launch": "=== launch_aicpu_kernel simpler_aicpu_exec ===", + "device_complete": "=== aclrtSynchronizeStreamWithTimeout stream_aicore_ ===", + } + for line in log_path.read_text(encoding="utf-8", errors="replace").splitlines(): + match = WALL_TIME_RE.match(line) + if match is None: + continue + for name, marker in patterns.items(): + if marker in line: + anchors[name] = wall_time_ns(match.group("ts")) + required = {"orch_ready", "aicpu_launch", "device_complete"} + missing = sorted(required - anchors.keys()) + if missing: + raise ValueError(f"missing whole-graph Host anchors in {log_path}: {', '.join(missing)}") + return anchors + + +def layer_value(span: StraceSpan) -> int: + return int(span["attrs"].get("layer", "-1")) + + +def find_span(spans: list[StraceSpan], name: str, layer: int | None = None) -> StraceSpan: + matches = [span for span in spans if span["name"] == name and (layer is None or layer_value(span) == layer)] + if len(matches) != 1: + suffix = "" if layer is None else f" for layer {layer}" + raise ValueError(f"expected one {name} span{suffix}, found {len(matches)}") + return matches[0] + + +def find_optional_span( + spans: list[StraceSpan], + name: str, + layer: int, +) -> StraceSpan | None: + matches = [span for span in spans if span["name"] == name and layer_value(span) == layer] + if len(matches) > 1: + raise ValueError(f"expected at most one {name} span for layer {layer}, found {len(matches)}") + return matches[0] if matches else None + + +def metadata_event(pid: int, tid: int, kind: str, value: str) -> dict[str, object]: + return {"ph": "M", "name": kind, "pid": pid, "tid": tid, "args": {"name": value}} + + +def complete_event( + pid: int, + tid: int, + name: str, + category: str, + ts_ns: int, + dur_ns: int, + base_ns: int, + args: dict[str, object], +) -> dict[str, object]: + return { + "ph": "X", + "pid": pid, + "tid": tid, + "name": name, + "cat": category, + "ts": (ts_ns - base_ns) / 1000.0, + "dur": dur_ns / 1000.0, + "args": args, + } + + +def read_l2(l2_path: Path) -> tuple[int, list[list[int]]]: + data = json.loads(l2_path.read_text(encoding="utf-8")) + return int(data["metadata"]["clock_freq_hz"]), data["aicore_tasks"] + + +def build_pipeline_events( # noqa: PLR0912 -- translates independent trace event kinds + invocation: int, + round_index: int, + spans: list[StraceSpan], + l2_path: Path, + title: str, +) -> list[dict[str, object]]: + frequency, aicore_tasks = read_l2(l2_path) + task_layers = sorted( + { + layer_value(span) + for span in spans + if span["name"] == "simpler_run.host_orch.layer_build" and int(span["attrs"].get("tasks", "0")) > 0 # type: ignore[union-attr] + } + ) + if not task_layers: + raise ValueError(f"invocation {invocation} has no task-bearing Host-O layers") + + layer_ranges: dict[int, tuple[int, int]] = {} + for layer in task_layers: + build = find_span(spans, "simpler_run.host_orch.layer_build", layer) + attrs = build["attrs"] + assert isinstance(attrs, dict) + layer_ranges[layer] = (int(attrs["task_begin"]), int(attrs["task_end"])) + + s_layers: list[SchedulerLayer] = [] + for layer in task_layers: + begin, end = layer_ranges[layer] + records = [record for record in aicore_tasks if begin <= int(record[1]) < end] + if not records: + raise ValueError(f"{l2_path} has no AICore records for task range [{begin},{end})") + start_cycles = min(int(record[3]) for record in records) + end_cycles = max(int(record[4]) for record in records) + s_layers.append( + { + "layer": layer, + "start_cycles": start_cycles, + "end_cycles": end_cycles, + "start_ns": 0, + "dur_ns": 0, + "records": len(records), + "unique_tasks": len({int(record[1]) for record in records}), + } + ) + + completion_waits = [ + span + for span in spans + if span["name"] == "simpler_run.host_orch.epoch.wait_device" + and span["attrs"].get("kind") in {"exec-done", "storage-free", "slot-free"} # type: ignore[union-attr] + ] + if not completion_waits: + raise ValueError(f"invocation {invocation} has no Device completion acknowledgement") + completion_by_epoch: dict[int, StraceSpan] = {} + for wait in completion_waits: + epoch = int(wait["attrs"]["target_epoch"]) # type: ignore[index] + end_ns = int(wait["ts_ns"]) + int(wait["dur_ns"]) + current = completion_by_epoch.get(epoch) + if current is None or end_ns < int(current["ts_ns"]) + int(current["dur_ns"]): + completion_by_epoch[epoch] = wait + anchor_epoch = max(completion_by_epoch) + anchor_wait = completion_by_epoch[anchor_epoch] + anchor_layer = anchor_epoch - 1 + anchor_s = next((layer for layer in s_layers if layer["layer"] == anchor_layer), None) + if anchor_s is None: + raise ValueError(f"exec-done epoch {anchor_epoch} has no matching S layer") + anchor_host_ns = int(anchor_wait["ts_ns"]) + int(anchor_wait["dur_ns"]) + anchor_device_ns = round(int(anchor_s["end_cycles"]) * 1_000_000_000 / frequency) + device_to_host_offset_ns = anchor_host_ns - anchor_device_ns + for s_layer in s_layers: + start_ns = round(int(s_layer["start_cycles"]) * 1_000_000_000 / frequency) + device_to_host_offset_ns + end_ns = round(int(s_layer["end_cycles"]) * 1_000_000_000 / frequency) + device_to_host_offset_ns + s_layer["start_ns"] = start_ns + s_layer["dur_ns"] = end_ns - start_ns + + ack_residuals: dict[int, int] = {} + for epoch, wait in completion_by_epoch.items(): + layer = epoch - 1 + s_layer = next((item for item in s_layers if item["layer"] == layer), None) + if s_layer is None: + continue + mapped_end = int(s_layer["start_ns"]) + int(s_layer["dur_ns"]) + ack_end = int(wait["ts_ns"]) + int(wait["dur_ns"]) + ack_residuals[layer] = ack_end - mapped_end + + host_starts = [span["ts_ns"] for span in spans if span["name"].startswith("simpler_run.host_orch.")] + device_starts = [layer["start_ns"] for layer in s_layers] + base_ns = min(*host_starts, *device_starts) + pid = round_index + 1 + events: list[dict[str, object]] = [ + metadata_event(pid, 0, "process_name", f"{title} - round {round_index}"), + metadata_event(pid, 1, "thread_name", "O active (Host build/materialize/stage/commit)"), + metadata_event(pid, 2, "thread_name", "O synchronization (S release/completion)"), + metadata_event(pid, 3, "thread_name", "S execution (AICore envelope)"), + ] + + phase_names = { + "simpler_run.host_orch.layer_build": "build", + "simpler_run.host_orch.epoch.materialize": "materialize", + "simpler_run.host_orch.epoch.stage_upload": "stage upload", + "simpler_run.host_orch.epoch.commit": "commit", + } + for span in spans: + phase = phase_names.get(str(span["name"])) + if phase is None: + continue + attrs = dict(span["attrs"]) + layer = layer_value(span) if phase == "build" else int(attrs["epoch"]) - 1 + if layer not in task_layers: + continue + cache_mode = attrs.get("cache", "none") + events.append( + complete_event( + pid, + 1, + f"O L{layer} {phase} ({cache_mode})", + "host_orchestrator", + int(span["ts_ns"]), + int(span["dur_ns"]), + base_ns, + dict[str, object](layer=layer, phase=phase, **attrs), + ) + ) + + for wait in spans: + if wait["name"] != "simpler_run.host_orch.epoch.wait_device": + continue + attrs = dict(wait["attrs"]) + kind = attrs.get("kind", "device") + target_epoch = int(attrs.get("target_epoch", "0")) + events.append( + complete_event( + pid, + 2, + f"wait {kind} epoch {target_epoch}", + "host_backpressure", + int(wait["ts_ns"]), + int(wait["dur_ns"]), + base_ns, + dict(attrs), + ) + ) + + for s_layer in s_layers: + layer = int(s_layer["layer"]) + begin, end = layer_ranges[layer] + events.append( + complete_event( + pid, + 3, + f"S L{layer} ({s_layer['records']} AICore records)", + "device_scheduler", + int(s_layer["start_ns"]), + int(s_layer["dur_ns"]), + base_ns, + { + "layer": layer, + "task_range": f"[{begin},{end})", + "records": s_layer["records"], + "unique_task_ids": s_layer["unique_tasks"], + "timing": ( + "actual AICore cycles; clock offset anchored by earliest completion acknowledgement " + f"for epoch {anchor_epoch} ({anchor_wait['attrs']['kind']})" + ), + "ack_residual_ns": ack_residuals.get(layer), + }, + ) + ) + + for layer in task_layers[1:]: + build = find_span(spans, "simpler_run.host_orch.layer_build", layer) + prior_s = next(item for item in s_layers if item["layer"] == layer - 1) + build_start = int(build["ts_ns"]) + build_end = build_start + int(build["dur_ns"]) + s_start = int(prior_s["start_ns"]) + s_end = s_start + int(prior_s["dur_ns"]) + overlap_ns = max(0, min(build_end, s_end) - max(build_start, s_start)) + events.append( + { + "ph": "C", + "pid": pid, + "tid": 1, + "name": "O/S overlap ns", + "ts": (build_start - base_ns) / 1000.0, + "args": {f"O L{layer} with S L{layer - 1}": overlap_ns}, + } + ) + return events + + +def parse_layer_task_counts(raw: str) -> list[int]: + counts = [int(token.strip()) for token in raw.split(",") if token.strip()] + if not counts or any(count <= 0 for count in counts): + raise ValueError("--layer-task-counts must contain positive comma-separated integers") + return counts + + +def build_whole_graph_events( + spans: list[StraceSpan], + anchors: dict[str, int], + l2_path: Path, + layer_task_counts: list[int], + title: str, + alignment: str, +) -> list[dict[str, object]]: + bind = find_span(spans, "simpler_run.bind") + device_wall = find_span(spans, "simpler_run.runner_run.device_wall") + frequency, aicore_tasks = read_l2(l2_path) + if not aicore_tasks: + raise ValueError(f"{l2_path} has no AICore records") + + layer_ranges: list[tuple[int, int]] = [] + begin = 0 + for count in layer_task_counts: + layer_ranges.append((begin, begin + count)) + begin += count + + orch_end_ns = anchors["orch_ready"] + orch_dur_ns = int(bind["dur_ns"]) + orch_start_ns = orch_end_ns - orch_dur_ns + if alignment == "compact": + device_wall_start_ns = orch_end_ns + completion_ns = device_wall_start_ns + int(device_wall["dur_ns"]) + alignment_note = "device wall placed immediately after Host O; profiling setup gap omitted" + else: + device_wall_start_ns = anchors["device_complete"] - int(device_wall["dur_ns"]) + completion_ns = anchors["device_complete"] + alignment_note = "device completion aligned to the profiled Host wall-clock acknowledgement" + + global_end_cycles = max(int(record[4]) for record in aicore_tasks) + s_layers: list[dict[str, int]] = [] + for layer, (begin, end) in enumerate(layer_ranges): + records = [record for record in aicore_tasks if begin <= int(record[1]) < end] + if not records: + raise ValueError(f"{l2_path} has no AICore records for task range [{begin},{end})") + start_cycles = min(int(record[3]) for record in records) + end_cycles = max(int(record[4]) for record in records) + start_ns = completion_ns - round((global_end_cycles - start_cycles) * 1_000_000_000 / frequency) + end_ns = completion_ns - round((global_end_cycles - end_cycles) * 1_000_000_000 / frequency) + s_layers.append( + { + "layer": layer, + "start_ns": start_ns, + "dur_ns": end_ns - start_ns, + "records": len(records), + "unique_tasks": len({int(record[1]) for record in records}), + } + ) + + first_s_ns = min(layer["start_ns"] for layer in s_layers) + base_ns = min(orch_start_ns, first_s_ns) + pid = 1 + events: list[dict[str, object]] = [ + metadata_event(pid, 0, "process_name", title), + metadata_event(pid, 1, "thread_name", "O active (Host whole-graph bind/build/relocate/upload)"), + metadata_event(pid, 2, "thread_name", "Host/device setup and synchronization"), + metadata_event(pid, 3, "thread_name", "S execution (AICore envelope)"), + complete_event( + pid, + 1, + "O whole graph (Host bind/build/relocate/upload)", + "host_orchestrator", + orch_start_ns, + orch_dur_ns, + base_ns, + { + "layers": len(layer_task_counts), + "tasks": sum(layer_task_counts), + "timing": "Host bind STRACE duration; end aligned to Device orchestration ready log", + }, + ), + complete_event( + pid, + 2, + "Device run wall (AICPU + S + teardown)", + "device_wall", + device_wall_start_ns, + int(device_wall["dur_ns"]), + base_ns, + {"alignment": alignment_note}, + ), + ] + + if alignment == "profile-wall" and first_s_ns > orch_end_ns: + events.append( + complete_event( + pid, + 2, + "Host/device setup before first AICore task", + "host_device_setup", + orch_end_ns, + first_s_ns - orch_end_ns, + base_ns, + { + "profile_ready_seen": "profile_ready" in anchors, + "aicpu_launch_ns": anchors["aicpu_launch"], + }, + ) + ) + + for layer, s_layer in enumerate(s_layers): + begin, end = layer_ranges[layer] + events.append( + complete_event( + pid, + 3, + f"S L{layer} ({s_layer['records']} AICore records)", + "device_scheduler", + s_layer["start_ns"], + s_layer["dur_ns"], + base_ns, + { + "layer": layer, + "task_range": f"[{begin},{end})", + "records": s_layer["records"], + "unique_task_ids": s_layer["unique_tasks"], + "timing": f"actual AICore cycles; {alignment_note}", + }, + ) + ) + return events + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--log", type=Path, required=True, help="run.log containing Host STRACE spans") + parser.add_argument("--l2-dir", type=Path, required=True, help="test output containing L2 JSON files") + parser.add_argument("--output", type=Path, required=True, help="output Perfetto JSON") + parser.add_argument("--title", default="Host O / Device S", help="Perfetto process name") + parser.add_argument( + "--layer-task-counts", + help="whole-graph mode: comma-separated scheduler task counts for consecutive layers", + ) + parser.add_argument( + "--whole-graph-alignment", + choices=("compact", "profile-wall"), + default="compact", + help="whole-graph mode: omit profiling setup or retain its Host wall-clock gap", + ) + args = parser.parse_args() + + rounds = load_strace_spans(args.log) + has_pipeline_spans = any( + span["name"] == "simpler_run.host_orch.layer_build" for spans in rounds.values() for span in spans + ) + events: list[dict[str, object]] = [] + if has_pipeline_spans: + for round_index, invocation in enumerate(sorted(rounds)): + root_l2_path = args.l2_dir / "l2_swimlane_records.json" + l2_path = ( + root_l2_path + if len(rounds) == 1 and root_l2_path.is_file() + else args.l2_dir / f"round_{round_index:03d}" / "l2_swimlane_records.json" + ) + if not l2_path.is_file(): + raise FileNotFoundError(l2_path) + events.extend(build_pipeline_events(invocation, round_index, rounds[invocation], l2_path, args.title)) + mode = "pipeline" + else: + if args.layer_task_counts is None: + raise ValueError("whole-graph mode requires --layer-task-counts") + if len(rounds) != 1: + raise ValueError("whole-graph mode currently requires exactly one profiled invocation") + l2_path = args.l2_dir / "l2_swimlane_records.json" + if not l2_path.is_file(): + raise FileNotFoundError(l2_path) + invocation = next(iter(rounds)) + events = build_whole_graph_events( + rounds[invocation], + load_whole_graph_anchors(args.log), + l2_path, + parse_layer_task_counts(args.layer_task_counts), + args.title, + args.whole_graph_alignment, + ) + mode = "whole-graph" + + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps({"traceEvents": events}, indent=2) + "\n", encoding="utf-8") + print(f"mode={mode}") + print(args.output.resolve()) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())