mindspore.runtime.use_mem_pool

mindspore.runtime.use_mem_pool(pool: MemPool)

A context manager that routes allocations and deallocations to a given pool.

Note

• This context manager makes only current thread's allocations route to the given pool.

• If a new thread is spawned inside the context manager the allocations in that thread will not route to the given pool.

• Only by allocating Device memory inside the context manager, the allocation operation can be routed to the given pool.

Parameters

pool (mindspore.runtime.MemPool) – a MemPool object that warp a PluggableAllocator.

Supported Platforms:

Ascend

Examples

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