[FEA] Option to 4KiB align large buffers

[rongou]

Is your feature request related to a problem? Please describe.
For spark-rapids we can spill device buffers to disk using GPUDirect Storage (GDS). With GDS it's more efficient to read/write 4KiB-aligned buffers. Unaligned buffers incur an additional copy to a bounce buffer.

Both the pool_memory_resource and arena_memory_resource are currently hard coded to align at 256 bytes. It'd be helpful to align at 4KiB for larger buffers.

Describe the solution you'd like
Add an optional constructor parameter to both pool_memory_resource and arena_memory_resource to allow aligning large buffers differently, e.g. large_buffer_allocation_alignment so that the user can set it to 4096 if desired. For any allocation with size below that threshold, it can still be aligned at 256 bytes; for requests with size larger than that, use the specified parameter for alignment.

Describe alternatives you've considered
I can add this only to arena_memory_resource, but then its behavior would diverge from pool_memory_resource. We currently do support selecting either MR in the spark-rapids plugin.

Additional context
GDS doc that talks about alignment: https://docs.nvidia.com/gpudirect-storage/best-practices-guide/index.html.

@harrism @jrhemstad @jlowe @abellina

[kkraus14]

cc @vuule as well. I imagine this has uses in cuIO

[jrhemstad]

This will be resolved by the new cuda::memory_resource interface that includes alignment in the allocate/deallocate parameters.

[rongou]

It'll require each caller to set the alignment explicitly on every call though, right? Is there an option to change the default alignment for the memory resource? Also is there an ETA?

[jrhemstad]

It'll require each caller to set the alignment explicitly on every call though, right?

Nothing that can't be handled with an adaptor or in the resource implementation.

Is there an option to change the default alignment for the memory resource?

You can do whatever you want inside a resource implementation.

Also is there an ETA?

It will be a release or two before we're incorporating this into RMM.