ENH: allow using aligned memory allocation, or exposing an API for memory management

[mattip]

Submitting this here since it is easier to link to all the issues and PRs dealing with memory management. Will submit it to the mailing list after a first-discussion here if needed.

There are three allocation to instantiate a PyArrayObject: the python allocation of the PyObject, the use of PyArray_malloc/PyArray_realloc/PyArray_free to allocate array dims and strides, and the PyDataMem_NEW/PyDataMem_NEWZEROED/PyDataMem_FREE calls for the array data.

I would like to make the case for an API for overloading internal use of PyDataMem_NEW, PyDataMem_NEWZEROED and PyDataMem_FREE. Work on this was begun in gh-5470, and referred to in issues gh-5312 (asking for an aligned allocator). In addition there is PR gh-5457, to provide a global interface for specifying aligned allocations.

As mentioned in the issue, alignment can be crucial for some applications, but in general is just extra overhead, so it should be configurable by the user/app.

I think I would prefer providing an API hook as in gh-5470. That PR actually had two parts: a configurable PyDataMem* overrides and a hook mechanism. The hook mechanism was merged (does anyone use that?) but the part I am interested in would requre an API change.

A fallback proposal would be to allow specifying a default allocation alignment aka gh-5457, perhaps using PyArray_malloc_aligned and friends as part of NumPy with the random API refactor.

[seberg]

I am not opposed, fo rme it was mainly something that other people were always more interested in (e.g. Julian). One issue I see glancing at the PR, is that currently we support for a user to provide the memory block (and NumPy taking ownership). And that usage is broken by gh-5470. So I wonder if we would need an allocation-strategy flag on the array, or something similar?

Other than that, the code seems simple enough that it is not an issue, unless there is some negative performance implication. One more question: You want such a higher-aligned memory, but you do not actually care for knowing whether an array provides that alignment? So this doesn't affect the question in gh-17359, right?

[mattip]

So this doesn't affect the question in gh-17359, right?

Right, it is orthogonal.

[mattip]

we support for a user to provide the memory block (and NumPy taking ownership). And that usage is broken by gh-5470

As envisioned in the PR, the functions are part of the PyArrayObject so each ndarray object would carry the information with it. The PR would have to make sure we use the correct functions consistently in all the right places: everywhere PyDataMem_NEW, PyDataMem_NEWZEROED and PyDataMem_FREE (and realloc would have to be carefully checked) are used.

[seberg]

Ah, I missed that it was stored on the array-object. That seems like a good solution to be able to change things safely (I am a bit worried that the cache must only be used if it is the current default allocator, maybe be part of the allocator struct, but that is a small thing).

So maybe this had also stalled because it required extending the array struct? I think I am in favor of being open to that in any case, it is also useful for cleaning up buffer handling (and who knows what else).

[leofang]

Speaking of memory management, it would be great if NumPy provides an API in either Python (preferred) or C (we can still make it work) for swapping in a memory manager implemented by 3rd parties. This need comes from GPU libraries: We use something called pinned/pagelocked memory to speed up host-device transfer, and in CuPy, for example, there is a custom memory manager to handle this. But it's not very convenient when we want every NumPy array to be backed by our memory pool, and it's better if NumPy offers an API to do so:

# very rough sketch just to illustrate the idea
import numpy as np
import cupy as cp

np.set_allocator(cp.cuda.alloc_pinned_memory)
a = np.empty(100)  # use pinned memory

This would make our life a lot easier (rel: cupy/cupy#3625).

cc possibly interested people @jakirkham @kmaehashi @gmarkall

[jakirkham]

cc @pentschev @hmaarrfk (also likely interested)

[mattip]

@leofang could you take a look at gh-17582? here is the rendered documentation. Do you require overriding memcpy as well?

[jakirkham]

Thanks Matti! 😄

There can be some value to overriding memcpy if we know we are going to/from pinned memory and/or device memory as this could be done asynchronously on a stream. Additionally there are tools like gdrcopy, which could speed up copies further if plugged in. Though I guess how useful this would be would depend on how complex this is to implement and use. Do you have a sense for either of these?

[mattip]

I think memcpy would be tricky to get right. It might be good if we can do this in phases leaving memcpy for a later phase, but that depends on there "only" being performance penalties when using memcpy instead of the more advanced functions and not errors. Is that true?

[leofang]

Thank you @mattip for considering our need! The rendered doc looks great -- It seems the memory hooks would allow us to hack something together so that a numpy.ndarray (not cupy) could also live on a GPU! Not saying it'd be useful, but it's very interesting.

but that depends on there "only" being performance penalties when using memcpy instead of the more advanced functions and not errors. Is that true?

As far as I can tell, this is a valid observation. If it's hard to override memcpy, I suppose we could still add some hacks to do the data transfer. But @jakirkham and @pentschev might have other concerns? I was never able to make gdrcopy work so I couldn't comment.

[jakirkham]

Yeah I think it is possible to handle memcpy alternatives outside of NumPy. So I don't think it is strictly necessary (even less so as it sounds complicated to do).

[pentschev]

I agree with comments from @jakirkham and @leofang , it would be useful to override memcpy but doesn't seem like a blocker and we could work around that outside of NumPy using the hooks.

[mattip]

@pentschev I am not sure what you meant by "using the hooks". How would you detect that a chunk of memory was modified? Just to be clear: I am reasonably sure I correctly identified malloc/realloc/free callsites inside NumPy and replaced them with function calls that use the hooks. However I did not touch the ~250 uses of memcpy inside NumPy. This is a large task, and one I would prefer to defer. My question was if people contemplating using the malloc/realloc/free hooks forsee problems due to glibc's memcpy being used internally by NumPy?

[pentschev]

Sorry @mattip , I think I confused myself in several ways above. When I wrote about hooks, I was thinking of memory allocations, internal memcpy calls have really nothing to do with that. Generally speaking, it seems memcpy overriding could be useful, but I'm not sure in which cases that would be necessary, so I'm ok with deferring that to future work.

[jakirkham]

Yeah I don't think we are worried about when NumPy uses memcpy internally. In terms of moving things from device-to-host or host-to-device, we would end up doing this ourselves. The OS (and thus memcpy) should be able to make smart choices about pinned memory under-the-hood.