mindspore.hal.StreamCtx

class mindspore.hal.StreamCtx(ctx_stream)[source]

Context-manager that selects a given stream.

All kernels queued within its context will be enqueued on a selected stream.

Parameters: ctx_stream (Stream) – selected stream. This manager is a no-op if it's None.
Raises: TypeError – If 'stream' is neither a mindspore.hal.Stream nor a None.

Examples

>>> import mindspore as ms
>>> import numpy as np
>>> from mindspore import Tensor, ops
>>> a = Tensor(np.ones([1024, 2048]), ms.float32)
>>> b = Tensor(np.ones([2048, 4096]), ms.float32)
>>> s1 = ms.hal.Stream()
>>> with ms.hal.StreamCtx(s1):
...     c = ops.matmul(a, b)
>>> ms.hal.synchronize()
>>> assert s1.query()