StreamGroups

class cubie.memory.stream_groups.StreamGroups(groups: dict[str, list[int]] | None = NOTHING, streams: dict[str, FakeStream | int] = NOTHING)

Bases: object

Container for organizing instances into groups with shared streams.

Parameters:

• groups (dict[str, list[int]] | None) – Dictionary mapping group names to lists of instance identifiers. When omitted, an empty mapping is created and populated with the “default” group.

• streams (dict[str, cubie.cuda_simsafe.FakeStream | int]) – Dictionary mapping group names to CUDA streams. When omitted, each group, including “default”, receives a dedicated stream from numba.cuda.stream() on first use. No group is ever backed by the device-wide default stream, so work in one process never orders against the CUDA null stream.

groups

Dictionary mapping group names to lists of instance identifiers.

Type:

dict[str, list[int]] | None

streams

Dictionary mapping group names to CUDA streams.

Type:

dict[str, cubie.cuda_simsafe.FakeStream | int]

Notes

Each group has an associated CUDA stream that all instances in the group share for coordinated operations. The “default” group is created automatically.

add_instance(instance: Any, group: str) → None

Add an instance to a stream group.

Parameters:

• instance – Host object or integer identifier to register with a stream group.

• group – Name of the destination group.

Raises:

ValueError – If the instance is already in a stream group.

Notes

If the group does not exist, it is created with a new CUDA stream.

change_group(instance: Any, new_group: str) → None

Move an instance to another stream group.

Parameters:

• instance – Host object or integer identifier to move.

• new_group – Name of the destination group.

Notes

If the new group does not exist, it is created with a new CUDA stream.

get_group(instance: Any) → str

Get the stream group associated with an instance.

Parameters:

instance – Host object or integer identifier whose group is requested.

Returns:

Name of the group containing the instance.

Return type:

str

Raises:

ValueError – If the instance is not in any stream groups.

get_group_stream(group: str = 'default') → FakeStream | int

Get the dedicated stream for a named group, creating it if needed.

Parameters:

group – Name of the stream group.

Returns:

The group’s dedicated CUDA stream. A missing group (or a group created without a stream) receives a fresh stream from numba.cuda.stream(), never the device-wide default stream.

Return type:

Stream

get_instances_in_group(group: str) → list[int]

Get all instances in a stream group.

Parameters:

group – Name of the group to inspect.

Returns:

List of instance identifiers associated with the group, or an empty list when the group has not been created.

Return type:

list of int

get_stream(instance: Any) → FakeStream | int

Get the CUDA stream associated with an instance.

Parameters:

instance – Host object or integer identifier whose stream is requested.

Returns:

CUDA stream associated with the instance’s group.

Return type:

Stream or int

groups: dict[str, list[int]] | None

reinit_streams() → None

Reinitialize all streams after a context reset.

Notes

Called after CUDA context reset to create fresh streams for all groups.

remove_instance(instance: Any) → None

Remove an instance from its stream group, if it has one.

Parameters:

instance – Host object or integer identifier to remove.

Notes

Removing an instance that is in no group is a no-op; the group and its stream are kept for the remaining members.

streams: dict[str, FakeStream | int]