mindspore.runtime.use_mem_pool

mindspore.runtime.use_mem_pool(pool: MemPool)

将内存分配和释放操作路由到指定内存池的上下文管理器。

说明

• 该上下文管理器仅会使当前线程的内存分配操作路由到指定内存池。

• 若在上下文管理器内部创建新线程，该线程的内存分配将不会路由到指定内存池。

• 只有在上下文管理器内部分配Device内存，才能将分配操作路由到指定内存池。

参数：

• pool (mindspore.runtime.MemPool) - 封装了PluggableAllocator的MemPool对象。

支持平台：

Ascend

样例：

>>> import mindspore as ms
>>> path = "/path/to/allocator.so"
>>> allocator = ms.runtime.PluggableAllocator(path, "Alloc", "Free")
>>> mem_pool = ms.runtime.MemPool(allocator)
>>> shape = (1024, 1024)
>>> x = ms.ops.Ones()(shape, ms.float32)
>>> with ms.runtime.use_mem_pool(mem_pool):
>>>     y = ms.ops.Ones()(shape, ms.float32)
>>> output = x + y