mstar.engine.accelerator_graph_backend

mstar.engine.accelerator_graph_backend#

Functions

get_accelerator_graph_backend(device)

Classes

AcceleratorGraphBackend(device)

Graph and stream operations required by the accelerator graph runners.

AcceleratorGraphPool(handle)

CUDAGraphBackend(device)

CapturedGraph(*args, **kwargs)

XPUGraphBackend(device)

XPUGraphPool(handle[, failed_graph])

class mstar.engine.accelerator_graph_backend.AcceleratorGraphBackend(device)[source]#

Bases: ABC

Graph and stream operations required by the accelerator graph runners.

Implementations own their native graph API and capture recovery. Runners do not need to know which runtime supplies those operations.

Parameters:

device (torch.device)

abstractmethod capture(graph, *, pool=None, stream=None)[source]#
Parameters:
Return type:

AbstractContextManager

abstractmethod create_graph(keep_graph=False)[source]#
Parameters:

keep_graph (bool)

Return type:

CapturedGraph

abstractmethod current_stream()[source]#
Return type:

Any

property device_type: str#
abstractmethod graph_pool_handle()[source]#
Return type:

AcceleratorGraphPool

abstractmethod is_available()[source]#
Return type:

bool

abstractmethod is_current_stream_capturing()[source]#
Return type:

bool

abstractmethod memory_allocated()[source]#
Return type:

int

abstractmethod new_stream()[source]#
Return type:

Any

abstractmethod recover_failed_capture(pool, stream)[source]#
Parameters:
Return type:

None

abstractmethod set_device()[source]#
Return type:

None

abstractmethod stream_context(stream)[source]#
Parameters:

stream (Any)

Return type:

AbstractContextManager

abstractmethod synchronize()[source]#
Return type:

None

class mstar.engine.accelerator_graph_backend.AcceleratorGraphPool(handle: 'Any')[source]#

Bases: object

Parameters:

handle (Any)

handle: Any#
on_failure(graph)[source]#
Parameters:

graph (CapturedGraph)

Return type:

None

on_success()[source]#
Return type:

None

class mstar.engine.accelerator_graph_backend.CUDAGraphBackend(device)[source]#

Bases: _TorchGraphBackend

Parameters:

device (torch.device)

create_graph(keep_graph=False)[source]#
Parameters:

keep_graph (bool)

Return type:

CapturedGraph

class mstar.engine.accelerator_graph_backend.CapturedGraph(*args, **kwargs)[source]#

Bases: Protocol

pool()[source]#
Return type:

Any

replay()[source]#
Return type:

None

reset()[source]#
Return type:

None

class mstar.engine.accelerator_graph_backend.XPUGraphBackend(device)[source]#

Bases: _TorchGraphBackend

Parameters:

device (torch.device)

create_graph(keep_graph=False)[source]#
Parameters:

keep_graph (bool)

Return type:

CapturedGraph

graph_pool_handle()[source]#
Return type:

AcceleratorGraphPool

class mstar.engine.accelerator_graph_backend.XPUGraphPool(handle: 'Any', failed_graph: 'CapturedGraph | None' = None)[source]#

Bases: AcceleratorGraphPool

Parameters:
failed_graph: CapturedGraph | None = None#
on_failure(graph)[source]#
Parameters:

graph (CapturedGraph)

Return type:

None

on_success()[source]#
Return type:

None

mstar.engine.accelerator_graph_backend.get_accelerator_graph_backend(device)[source]#
Parameters:

device (device)

Return type:

AcceleratorGraphBackend