Skip to content

simpler.task_interface

Generated from the source. For the curated view — which **config keys Worker accepts, CallConfig defaults, and the argument-order footguns — see Python API.

simpler.task_interface

Public Python API for task_interface nanobind bindings.

Re-exports the canonical C++ types (DataType, Tensor, ChipStorageTaskArgs, TaskArgs, TensorArgType) plus scalar_to_uint64. Torch-aware helpers (make_tensor_arg, torch_dtype_to_datatype) live in simpler_setup.torch_interop — this module has no torch dependency.

Usage

from simpler.task_interface import DataType, Tensor, ChipStorageTaskArgs from simpler_setup.torch_interop import make_tensor_arg

RemoteAddressSpace

Bases: IntEnum

How a remote buffer's bytes are reached.

HOST_INLINE carries the payload in the message itself rather than naming remote memory. REMOTE_WINDOW and UB_LDST are protocol placeholders: the shipped transport is simulation-backed.

RemoteBufferHandle

A reference to memory on a remote L3 worker.

Returned by Worker.remote_malloc (an owner handle) or by Worker.remote_import (an imported handle, told apart by is_imported). The two are not interchangeable: owner handles are freed with remote_free, imported ones with remote_release_import.

RemoteTensorRef.host_inline produces a third form, with address_space of HOST_INLINE: it carries its bytes in the message and names no remote allocation, so neither release call applies — remote_free rejects it outright and it is never is_imported.

Construct only through Worker or RemoteTensorRef.host_inline; the constructor is token-guarded.

worker_id property

worker_id: int

Worker holding this reference — the importer, for an imported handle.

owner_worker_id property

owner_worker_id: int

Worker that owns the underlying allocation.

import_id property

import_id: int

Nonzero on an imported handle; 0 on an owner handle.

address_space property

address_space: RemoteAddressSpace

How these bytes are reached; see RemoteAddressSpace.

nbytes property

nbytes: int

Size of the allocation in bytes, or of the payload for HOST_INLINE.

released property

released: bool

Whether the handle has been freed or released.

access_flags property

access_flags: int

Permitted access as a read/write bitmask; an export may only narrow it.

is_imported property

is_imported: bool

Whether this came from remote_import rather than remote_malloc.

RemoteBufferExport

Opaque descriptor returned by Worker.remote_export.

The transport fields are intentionally kept private so callers cannot forge or log remote keys by accidentally treating the export as a plain dataclass.

owner_worker_id property

owner_worker_id: int

Worker that owns the exported allocation.

address_space property

address_space: RemoteAddressSpace

How the exported bytes are reached.

offset property

offset: int

Start of the exported range within the owner buffer.

nbytes property

nbytes: int

Length of the exported range in bytes.

access_flags property

access_flags: int

Access granted here; a subset of the owner handle's flags.

transport_profile property

transport_profile: str

Transport this export was minted for.

RemoteTensorRef dataclass

A tensor argument that lives on, or travels to, a remote worker.

host_inline classmethod

host_inline(payload: bytes, *, shape: tuple[int, ...], dtype: DataType) -> RemoteTensorRef

Build a reference whose payload travels inline, naming no remote memory.

payload length must equal the byte size implied by shape and dtype, and shape entries must be non-negative.

CommBufferSpec dataclass

A named slice of the per-rank communicator window.

Buffers are placed sequentially inside the window in declaration order — Buffers are placed sequentially inside the window in declaration order. The CommDomainHandle.contexts[chip_idx].buffer_ptrs dict returned by Orchestrator.allocate_domain is keyed by CommBufferSpec.name.

ChipDomainContext dataclass

Per-domain view handed to a chip worker: its rank within the domain and the local slice of the symmetric window.

CommDomainHandle

User-facing handle for one dynamically-allocated CommDomain.

Returned by Orchestrator.allocate_domain(...). Acts as a context manager: with exit marks the handle for release and prevents further use; the actual backend free runs after Worker.run has drained any tasks the orch function submitted using this domain. This is required because submit_* only enqueues to the DAG — freeing before drain would create a use-after-free on the chip side.

Lifecycle states::

live           — allocated, indexable, can be passed to submit_*
released       — release() called; further indexing raises;
                  backend memory still alive until Worker.run drain
freed          — backend release_domain has executed, memory gone

Most users only see released; the live → released transition happens at with exit (or explicit release()), and the released → freed transition is the runtime's job at end-of-run.

released property

released: bool

True once release() (or with exit) has been called.

Backend memory may still be alive — it is freed by the Worker after DAG drain at end-of-run. Use this to gate further indexing / submission, not to assert physical teardown (use freed for that).

freed property

freed: bool

True once the backend comm_release_domain_windows has executed.

Only flips after the owning Worker.run completes and processes the pending-release queue. An orch_fn will never observe True for a handle it released within the same run call.

release

release() -> None

Mark this handle for collective release. Idempotent.

Inside an orch function, this is a non-blocking mark — the actual backend comm_release_domain_windows runs after the owning run's completion wait so that tasks already submitted with this domain's device_ctx see live memory through execution.

After this returns, the handle is treated as released for the user's purposes: __getitem__ raises, repeated release() is a no-op, and the orch function must not pass it to further submit_* calls.

ChipWorker

Unified execution interface wrapping the host runtime C API.

The runtime library and target device are bound once via init() and cannot be changed. Public dispatch uses opaque CallableHandle values. Integer execution slots are private to this wrapper and the runtime ABI.

Usage::

worker = ChipWorker()
worker.init(device_id=0, bins=bins)
handle = worker.register_callable(chip_callable)
worker.run(handle, args=orch_args, config=CallConfig())
worker.unregister_callable(handle)
worker.finalize()

aicpu_dlopen_count property

aicpu_dlopen_count

Number of distinct callable identities the AICPU has dlopened for.

host_dlopen_count property

host_dlopen_count

Number of host-side orch SO dlopens (host_build_graph variants).

run_stream_set_create_count property

run_stream_set_create_count

Number of AICore run streams the bound runner has created.

runtime_buffer_addrs property

runtime_buffer_addrs

Address of each host Runtime staging buffer, in slot order.

device_id property

device_id

ACL device ordinal this worker is bound to.

initialized property

initialized

Whether the underlying native worker has completed init.

init

init(device_id: int, bins: Any, log_level: int | None = None, prewarm_config: CallConfig | None = None, enable_sdma: bool = False)

Attach the calling thread to device_id, load the host runtime library, and cache platform binaries.

Can only be called once — the runtime and device cannot be changed after init.

Performs the process-wide RTLD_GLOBAL bootstrap (libsimpler_log.so, plus libcpu_sim_context.so on sim platforms) and seeds the HostLogger via simpler_log_init before the C++ _ChipWorker.init dlopens host_runtime.so — host_runtime.so resolves its undefined HostLogger / unified_log_ (and, on sim, sim_context_) symbols against those globals, and any LOG_* macro firing during its dlopen-time constructors must already see the right filter.

Parameters:

Name Type Description Default
device_id int

NPU device ID to attach the calling thread to.

required
bins Any

A simpler_setup.runtime_builder.RuntimeBinaries (or any object exposing host_path / aicpu_path / aicore_path / simpler_log_path / sim_context_path / dispatcher_path). dispatcher_path is required for onboard platforms and ignored on sim (set to None).

required
log_level int | None

Threshold (10=DEBUG, 20=INFO, 25=TIMING, 30=WARN, 40=ERROR, 60=NUL). Defaults to a snapshot of the simpler logger via _log.get_current_config().

None

For tests that need to drive the binding directly with arbitrary path strings (e.g. to assert dlopen failure on /nonexistent/foo.so), call _ChipWorker.init(...) from _task_interface instead of going through this wrapper.

finalize

finalize()

Tear down everything: device resources and runtime library.

Terminal operation — the object cannot be reused after this.

register_callable

register_callable(callable)

Prepare a ChipCallable and return an opaque handle.

The runtime still uses an integer slot internally, but the caller never chooses or observes it.

run

run(handle: CallableHandle, args: ChipStorageTaskArgs, config: CallConfig | None = None, **kwargs: Any)

Launch a callable previously returned by register_callable.

Parameters:

Name Type Description Default
handle CallableHandle

CallableHandle returned by register_callable.

required
args ChipStorageTaskArgs

ChipStorageTaskArgs for this invocation.

required
config CallConfig | None

Optional CallConfig. If None, a default is created.

None
**kwargs Any

Overrides applied to config (e.g. aicpu_thread_num=4). A run always takes the whole device; orchestration reads the resulting width back through rt_available_cluster_count().

{}

Returns None. Per-stage run timing is emitted as [STRACE] log markers by the platform — see docs/dfx/host-trace.md.

unregister_callable

unregister_callable(handle) -> None

Drop one live callable handle and release its private resources when final.

arena_bank_gm_heap_base

arena_bank_gm_heap_base(bank_id)

Committed GM heap base of one arena bank, or 0 when uncommitted.

retained_temp_addr

retained_temp_addr(slot_id)

Retained temporary-buffer address for one slot, or 0 when unheld.

malloc

malloc(size)

Allocate memory. Returns a pointer (uint64).

free

free(ptr)

Free memory allocated by malloc().

copy_to

copy_to(dst, src, size)

Copy size bytes from host src to worker dst.

copy_from

copy_from(dst, src, size)

Copy size bytes from worker src to host dst.

comm_init

comm_init(rank: int, nranks: int, rootinfo_path: str) -> int

Initialize a distributed communicator for this rank.

ChipWorker owns ACL bring-up and the aclrtStream internally, so callers never touch aclInit / aclrtSetDevice / stream lifetimes. On sim, ACL / stream are not used. Pair with comm_destroy for teardown.

Parameters:

Name Type Description Default
rank int

This process's rank (0-based).

required
nranks int

Total number of ranks.

required
rootinfo_path str

Filesystem path used for rank handshake.

required

Returns:

Type Description
int

Opaque communicator handle (uint64) for the other comm_* calls.

comm_alloc_windows

comm_alloc_windows(comm_handle: int, win_size: int) -> int

Allocate per-rank windows. Returns a device CommContext pointer (uint64).

comm_get_local_window_base

comm_get_local_window_base(comm_handle: int) -> int

Return this rank's local window base address (uint64).

comm_get_window_size

comm_get_window_size(comm_handle: int) -> int

Return the actual per-rank window size in bytes.

comm_derive_context

comm_derive_context(comm_handle: int, rank_ids: list[int], domain_rank: int, window_offset: int, window_size: int) -> int

Derive a domain-local device CommContext from an allocated base communicator.

comm_barrier

comm_barrier(comm_handle: int) -> None

Synchronize all ranks.

comm_destroy

comm_destroy(comm_handle: int) -> None

Destroy the communicator and release its resources.

comm_destroy_all

comm_destroy_all() -> None

Destroy all communicators owned by this worker.

scalar_to_uint64

scalar_to_uint64(value) -> int

Convert a scalar value to uint64.

value can be a Python int, float, a ctypes scalar (c_int64, c_float, etc.), or any object convertible to int.

Python float values are converted to IEEE 754 single precision (32-bit) and their bit pattern is zero-extended to uint64. This may cause a loss of precision. For double precision, use ctypes.c_double.