Add: accept shared Tensor arguments in ChipWorker.run

[ChaoWao]

Problem and behavior

Worker accepts TaskArgs<Tensor>, but the public synchronous ChipWorker.run still requires callers to build the address-bearing ChipStorageTaskArgs POD. It now accepts shared TaskArgs through the existing snapshot, ImportRegistry, and native materializer, then calls the same native run entry.

This entry supports scalars and contiguous HOST/H2D views, including nonzero byte offsets. It rejects DEVICE descriptors before import: those use Worker(level=2).submit and its live-source registration/retention from #2482. HOST/NONE and noncontiguous H2D are also refused. The resolved POD input remains compatible; no device ABI or runtime execution mechanism changes.

Lifetime and validation

• Whole-call grant/overlap checks and chip-input restrictions run before import. Config errors also precede import. Caller descriptor/scalar mutation cannot change an accepted snapshot.
• Each call owns its imports through synchronous native return and output copies. A bind failure closes them. A native-call exception retains them, refuses further public runs, and releases them only after successful native teardown; a teardown exception preserves them for retry.
• Public run and finalize serialize around that lifetime. This owns imported mappings, not the caller's allocation: callers must retain their backing storage until return or successful teardown after an error.
• The new entry requires canonical row-major strides. Broader strided H2D support and standalone ChipWorker device-source attachment are outside this change; it does not complete all public POD migration.

Validation

Validated head cacc1415, based on main 8aae172e5:

• Full editable package/runtime build succeeds.

• Full Python UT: 2853 passed, 65 skipped, one existing torch readonly-buffer warning.

• HBG readiness suite: 7 passed on a2a3sim; 7 passed on a5sim.

• Ten focused argument tests cover offsets and shared-memory write-through, exact scalar bits, caller mutation, whole-call rejection, bind rollback, failed teardown, a real contending finalize thread, and invalid config. Existing ChipWorker tests preserve the POD compatibility path.

• Negative controls fail when using main's old entry, removing the snapshot, or unmapping on a native exception. The noncontiguous-input test also fails before adding its pre-import guard.

• The HBG scene calls real ChipWorker.run six times with shared HOST buffers, checks output/control values after H2D and device execution, and checks sentinels around offset views. Its suite also retains the six existing Worker-path cases.

• Applicable pre-commit checks pass. No onboard run: architecture precheck confirmed A2/A3, but all 16 devices had nonzero HBM usage even though the task-submit queue was empty. A5 hardware is unavailable on this host.

Only code, regression tests and API contract docs are included. Local design/roadmap materials are not committed.

[coderabbitai]

Warning

Review limit reached

Next included review available in 28 minutes.

Check out review usage here.

View limit details

Limit details: You’ve used the included review currently available.

You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository.

Learn how review limits work.

Review configuration:

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Advanced

Run ID: 332ca3cf-8cf3-478c-9e98-62237332e718

📥 Commits

Reviewing files that changed from the base of the PR and between 8aae172 and cacc141.

📒 Files selected for processing (5)

• docs/buffer-abi.md
• docs/user/reference/python-api.md
• python/simpler/task_interface.py
• tests/st/host_build_graph_readiness/test_host_readiness.py
• tests/ut/py/test_chip_worker_args.py

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

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