rapidsai
[FEA] Properly support alignment argument in `do_allocate` / `do_deallocate`
The signature of CCCL's allocate and deallocate methods accept a size_t alignment. However, RMM's current design doesn't do anything with this alignment. We need to update do_allocate and do_deallocate in the device_memory_resource base class design to actually use the value being passed in. Really, we just need to validate that it is a valid CUDA alignment.
Currently CCCL enforces that alignment <= 256 && 256 % alignment == 0, which means that only powers of 2 up to 256 are valid. https://github.com/NVIDIA/cccl/blob/d3af633a6c1429197f1fa30dc3f5e6671177fa5b/libcudacxx/include/cuda/__memory_resource/memory_resource_base.h#L331-L338
This requirement is totally fine to enforce at runtime, as it is already a requirement of all RMM resources.
To make this a non-breaking change, we might need to introduce a new overload of do_allocate / do_deallocate that accepts the alignment parameter, and then migrate to it? Or maybe we can just add the alignment parameter with a default value of rmm::CUDA_ALLOCATION_ALIGNMENT?
To make this a non-breaking change, we might need to introduce a new overload of do_allocate / do_deallocate that accepts the alignment parameter, and then migrate to it? Or maybe we can just add the alignment parameter with a default value of rmm::CUDA_ALLOCATION_ALIGNMENT?
do_allocate/do_deallocate aren't part of the API I think? So does it matter to change things? There'd be an ABI break I suppose.
I discussed with @mhaseeb123. The premise of this issue might be flawed.
We are moving away from having a base class with the new CCCL MR design, so it might be necessary for every "base" resource to implement the alignment validation in its allocate function. At least, that's what CCCL appears to be doing today.