lebedov/scikits.cuda — pull request #74 by untom (2014-08-11 20:08:10+00:00)
Title: Add option to specify allocator for high-level functions.

Internally, some of the high-level functions in `linalg` allocate device memory (e.g. to store the results of calculations). In my most common use-cases, I run the same algorithm over several iterations. So in each iteration, I end up allocating quite a few temporary arrays which could be re-used each iteration but aren't (either since the API doesn't expose them or because I forgot to use the `out=` argument some functions have).

This PR allows users to specify an allocator via `misc/linalg.init`, which is then used for these allocations. This can be used to install a `DeviceMemoryPool` allocator. In my use-cases, this gives me a speed-up of ~4-8%. Of course, this is extremely dependant on the use case. At least in my case, pretty much the same speed-up could also be achieved by simply using `out=` arguments where provided (e.g. in `linalg.dot`). But doing it this way makes the code look nicer (less cluttered with `out=` args everywhere) and I don't even have to think about where/when/how to save allocations.

The current API that this PR uses is this:

from scikits.cuda import linalg
from pycuda.tools import DeviceMemoryPool
mempool = DeviceMemoryPool()
linalg.init(allocator=mempool.allocate)

from then on, all `linalg` calls that allocate memory will do so through `mempool`. I'm not sure this is the best API to set this up. Another option would be to just say `linalg.init(manage_memory=True)`, then one doesn't have to worry about setting up the DeviceMemoryPool (OTOH, maybe one wants to have access to the pool?). Yet another option would be to just make using the memory pool the default; arguably, most users will internally allocate the same arrays over and over again, but if there are users who constantly use different array-sizes they might not like it (AFAIK the DeviceMemoryPool will release unused memory if it's in danger of running out of memory, but still).

The PR currently only covers `linalg`, but there are other places that internally memory as well that could make use of this. The only clash I found is with `misc.zeros`, which explicitly exposes an `allocator` argument on its own to stay in sync with `pycuda.gpuarray.zeros` (SIDENOTE: is `misc.zeros` still needed in current PyCUDA versions?).

So I think it's worthwhile discussing if it makes sense to add this functionality and how the API should look like specifically.