[BUG] `rmm::device_uvector` does not always respect alignment requirements of its value type.

[wence-]

Describe the bug

device_uvector stores its allocation as an untyped rmm::device_buffer. This is created by converting the number of elements of storage requested into a number of bytes (by multiplying by sizeof(T)).

However, any alignment requirements of the underlying type T are ignored.

Generally, the fact that all RMM device allocations are 256 byte aligned saves us. However, should one try and store a type with wider than 256 byte alignment, the stored pointer may not be appropriately aligned.

Steps/Code to reproduce bug

#include <cstdint>
#include <cstdio>
#include <rmm/device_uvector.hpp>
#include <rmm/device_vector.hpp>
#include <rmm/mr/cuda_async_memory_resource.hpp>
#include <rmm/mr/cuda_memory_resource.hpp>
#include <rmm/mr/pool_memory_resource.hpp>
#include <rmm/resource_ref.hpp>
#include <type_traits>

struct alignas(512) foo {
  int data;
};

int main(void) {

  rmm::cuda_stream_view stream{};
  auto base = rmm::mr::cuda_memory_resource{};
  auto pool =
      rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource>{base, 1 << 20};
  auto async = rmm::mr::cuda_async_memory_resource{};
  rmm::device_async_resource_ref mr = pool;
  auto buf1 = rmm::device_uvector<int>(1, stream, mr);
  auto buf3 = rmm::device_uvector<foo>(1, stream, mr);
  auto buf2 = rmm::device_uvector<foo>(9, stream, mr);
  auto buf4 = rmm::device_uvector<int>(1, stream, mr);

  printf("%p required alignment %zu is_aligned=%d\n", buf1.data(),
         std::alignment_of_v<decltype(buf1)::value_type>,
         (reinterpret_cast<std::uintptr_t>(buf1.data()) %
          std::alignment_of_v<decltype(buf1)::value_type>) == 0);
  printf("%p required alignment %zu is_aligned=%d\n", buf2.data(),
         std::alignment_of_v<decltype(buf2)::value_type>,
         (reinterpret_cast<std::uintptr_t>(buf2.data()) %
          std::alignment_of_v<decltype(buf2)::value_type>) == 0);
  printf("%p required alignment %zu is_aligned=%d\n", buf3.data(),
         std::alignment_of_v<decltype(buf3)::value_type>,
         (reinterpret_cast<std::uintptr_t>(buf3.data()) %
          std::alignment_of_v<decltype(buf3)::value_type>) == 0);
  printf("%p required alignment %zu is_aligned=%d\n", buf4.data(),
         std::alignment_of_v<decltype(buf4)::value_type>,
         (reinterpret_cast<std::uintptr_t>(buf4.data()) %
          std::alignment_of_v<decltype(buf4)::value_type>) == 0);
  return 0;
}

$ ./bug
0x7ae93a600000 required alignment 4 is_aligned=1
0x7ae93a600300 required alignment 512 is_aligned=0
0x7ae93a600100 required alignment 512 is_aligned=0
0x7ae93a601500 required alignment 4 is_aligned=1

Expected behavior

Buffers should be appropriately aligned.

[bdice]

@wence- Thanks for this issue. Would you be able to work on a fix?

[wence-]

Not immediately, but if you assign I will add it to my todo list