[BUG] Large overhead when launching kernel with torch tensors

[gmarkall]

From: numba/numba#9835

Creating this issue in numba-cuda as it will need resolving here - will look into this deeper once I'm back in the new year.

cc @materight

[cpcloud]

I spent a bunch of yesterday getting to the bottom of this and addressing the low-hanging fruit.

There are numerous sources of overhead here, with the main difference being observable in very simple code:

# allocate torch tensor
zeros = torch.zeros(...)

for _ in range(10_000):
    numba_cuda_kernel(zeros)  # conversion overhead is paid 10,000 times

# versus...

arr = numba.cuda.as_cuda_array(zeros)  # conversion overhead is paid once

for _ in range(10_000):
    numba_cuda_kernel(arr)

In the first case, numba.cuda.as_cuda_array is invoked 10,000 times, whereas in the second case it is invoked once, since the conversion code path simply checks whether an array is a device array (it is in that case) and returns the input. This code path is very fast relative to the first.

There are a few other places where there's unnecessary overhead:

• Using hasattr(x, "__cuda_array_interface__") to check whether an object has that interface, and then accessing it again after the check succeeds. This is 2x more operations than necessary, constructing the dict interface object isn't particularly cheap. Addressed in perf: reduce the number of __cuda_array_interface__ accesses #538.
• from_cuda_array_interface has a @require_context decorator, which constructs a cuda context, but it also calls current_context() (aliased to devices.get_context()) which does basically the same thing AFAICT, making the decorator unnecessary. Additionally, I was able to remove context from MemoryPointer, because it's never used in any methods except to construct other instances of itself (the constructor requires the argument). So that should eliminate the context overhead entirely. Addressed in perf: remove context threading in various pointer abstractions #536.
• DeviceNDArray overhead is then the remaining bottleneck (again, remember that calling it once is not noticeable, it's calling it 10k times that is). There are a bunch of APIs adding overhead here, but much of comes from repeated calculation of dimensions (in particular prepare_shape_strides_dtype). dummyarray.py for example, appears to be externally immutable, and thus potentially eligible for caching based on its inputs like shape, strides, offset and order.

[rparolin]

For numba kernels, could we have the kernel perform that check to determine if its executing the slow path by having to convert the torch tensors to cuda arrays? That way if users write that code it will throw an exception and force them to correct the callsite rather than generate a performance regression? Basically trying to emulate in C++ if the function overload was taking it could throw an exception.

[cpcloud]

Perhaps! I'd like to see how things shake out with some of the other changes I've PR'd plus some targeted caching of strides/shape munging that I've started working on before implementing any warnings/exceptions.

[cpcloud]

Here's the current state of main as of the slew of recent PRs that improve as_cuda_array overhead:

Takeaways

• Small improvements in device_array, to be expected, because these were already fairly low-overhead relative to torch inputs.
• Bigger improvements in torch, 37-91% on that run, roughly consistent across runs.

Still plenty of room for improvement, but making solid progress.

[cpcloud]

To summarize the work that's been done so far on this issue:

• perf: speed up kernel launch #510: initial kernel launch improvement PR, removes known-slow (and unnecessary) Python overhead
• test: add benchmarks for kernel launch for reproducibility #528: add benchmarks for reproducible and measurable improvement/degradation
• perf: remove context threading in various pointer abstractions #536: remove unnecessary context creation
• perf: remove duplicated size computation #537: remove duplicated size computation
• perf: reduce the number of __cuda_array_interface__ accesses #538: remove expensive repeated attribute access of __cuda_array_interface__
• chore(perf): add torch to benchmark #539: add torch to benchmarking dependencies so that we can measure the thing we actually care about
• perf: cache dimension computations #542: cache repeated dimension computations

[cpcloud]

Quoting my comment from #542 on the current state after that PR:

Here are the performance improvements on top of main:

So, there's another ~66% improvement for the many-arg torch case and another ~34% improvement for the one-arg case torch case.

That would bring us to:

• the torch case is now about 2.5x slower than the device array case (down from 7x) in the many-arg scenario
• the torch case is now about 1.8x slower than the device array case (down from 3x) in the one-arg scenario

[cpcloud]

All PRs related to this that are coming from easily detectable overhead are now merged.

Remaining work to get this closer to parity with device_array will likely involve using StridedMemoryViews from cuda.core in place of mostly-pure-Python DeviceNDArray (in numba-cuda).