mindspore.runtime.StreamLimitCtx

View Source On Gitee
class mindspore.runtime.StreamLimitCtx(stream, cube_num=- 1, vector_num=- 1)[source]

Context-manager that selects a given stream core number limit.

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

Parameters

  • stream (mindspore.runtime.Stream) – selected stream.

  • cube_num (int, optional) – set cube num for stream. Default is -1, indicating that it is not set.

  • vector_num (int, optional) – set vector num for stream. Default is -1, indicating that it is not set.

Raises

TypeError – If stream is not mindspore.runtime.Stream.

Supported Platforms:

Ascend

Examples

>>> import mindspore as ms
>>> import numpy as np
>>> from mindspore import Tensor, ops
>>> ms.set_device("Ascend", 0)
>>> a = Tensor(np.ones([1024, 2048]), ms.float32)
>>> b = Tensor(np.ones([2048, 4096]), ms.float32)
>>> s1 = ms.runtime.Stream()
>>> with ms.runtime.StreamLimitCtx(s1, 8, 8):
...     c = ops.matmul(a, b)
>>> ms.runtime.synchronize()