Cythonize `Buffer`

[leofang]

From @shwina

In [15]: mr = DeviceMemoryResource(Device().device_id)

In [16]: %timeit mr.allocate(10)
3.35 μs ± 95.1 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

In [17]: %timeit cp.empty(10, dtype=cp.int8)
1.48 μs ± 11.4 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

[shwina]

Unfortunately, the problem goes way beyond just Python overheads due to Buffer construction. If we actually sync the device, we see a huge perf gap:

In [1]: from cuda.core.experimental import Device, DeviceMemoryResource

In [2]: import cupy as cp

In [3]: mr = DeviceMemoryResource(Device().device_id)

In [4]: Device().set_current()

In [5]: %timeit cp.empty(10); cp.cuda.runtime.deviceSynchronize()
1.93 μs ± 116 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

In [6]: %timeit mr.allocate(10); Device().sync()
646 μs ± 1.14 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

Using cuda.bindings directly:

In [22]: %timeit (_, ptr) = cudaMallocAsync(10, None); cudaFreeAsync(ptr, None)
625 ns ± 3.99 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

[kkraus14]

Using cuda.bindings directly:

In [22]: %timeit (_, ptr) = cudaMallocAsync(10, None); cudaFreeAsync(ptr, None)
625 ns ± 3.99 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

I think this still runs asynchronously where if you synchronize the device with cudaDeviceSynchronize() in the timeit loop it will show similar performance to using mr.allocate?

I think the reason that cupy is so much faster here is because they don't use the driver managed memory pool by default, they have their own memory pool implementation that they default to. @leofang probably knows the implementation details here, but I'm guessing they allocate in a blocksize and suballocate from blocks, whereas the async pool doesn't.

We could try using cudaMemPoolSetAttribute to set cudaMemPoolAttrReleaseThreshold to a large number and then making a large allocation to prime the pool before running this timeit loop. I think this is what RMM effectively does: https://github.com/rapidsai/rmm/blob/branch-25.08/cpp/include/rmm/mr/device/cuda_async_memory_resource.hpp#L133-L147

[shwina]

I think this still runs asynchronously where if you synchronize the device with cudaDeviceSynchronize() in the timeit loop it will show similar performance to using mr.allocate?

Adding a sync brings the perf closer to CuPy's empty:

In [27]: %timeit (_, ptr) = cudaMallocAsync(10, None); cudaFreeAsync(ptr, None); Device().sync()
1.94 μs ± 128 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

[shwina]

For completeness, we found that the issue is related to the 0 release threshold of the default pool used by cudaMallocAsync (and therefore allocate()).

[shwina]

We are working around the issue I mentioned in my previous comment by setting the release threshold as follows:

@functools.cache
def _set_default_mempool_threshold(device_id: int):
    """
    Set the release threshold for the default memory pool on this device,
    if we haven't already done so. This prevents the driver from attempting
    to shrink the pool after every sync, which can be slow.
    """
    default_pool = handle_return(driver.cuDeviceGetDefaultMemPool(device_id))
    threshold = handle_return(
        driver.cuMemPoolGetAttribute(
            default_pool, driver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RELEASE_THRESHOLD
        )
    )
    if int(threshold) == 0:
        handle_return(
            driver.cuMemPoolSetAttribute(
                default_pool,
                driver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
                driver.cuuint64_t(0xFFFFFFFFFFFFFFFF),
            )
        )

With the change above, we see much better performance using Buffer - but it is still quite a bit slower compared to using cp.empty() for memory allocation, or using cuda.bindings directly. For lack of a better place to report these performance issues, I'm writing it up here.

The benchmark I'm running uses cuda.cccl.parallel to perform a reduction: https://github.com/shwina/cccl/blob/ac7a98f439b138aa586c530a6b1f437f2f290a20/python/cuda_cccl/benchmarks/parallel/bench_reduce.py#L109. The runtime of this benchmark ends up being quite sensitive to the temporary memory allocation mechanism. Below are two profiles collected from this benchmark, the first one uses Buffer for the memory allocation, and the second one uses cuda.bindings.runtime.cudaMallocAsync directly to do the memory allocation.

• Code for approach 1
• Code for approach 2
• profiles.zip

Of note is the time taken by TemporaryStorageBuffer.__init__ in the two appraoches (the bright orange block in the second level):

When using Buffer, many small things combined contribute to the increased runtime.

1.Device.set_current()
2. Device.create_stream()
3. cuStreamDestroy
4. Buffer._init

(1) can be addressed in #739. (2) and (3) can be addressed by using the stream protocol rather than having to create and destroy a stream in order to be able to plumb it through to the user as we're having to do here. (4) can be addressed by cythonizing Buffer (this issue).