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 receives a new stream from numba.cuda.stream() and the “default” group is backed by numba.cuda.default_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.

Returns:

None.

Return type:

None

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.

Returns:

None.

Return type:

None

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_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.

Returns:

None.

Return type:

None

Notes

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

streams: dict[str, FakeStream | int]