Forward failure callback alignment

cpp/include/rmm/mr/detail/failure_callback_resource_adaptor_impl.hpp

@@ -62,12 +62,12 @@ class failure_callback_resource_adaptor_impl {
 
   void* allocate(cuda::stream_ref stream,
                  std::size_t bytes,
-                 std::size_t /*alignment*/ = alignof(std::max_align_t))
+                 std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT)
   {
     void* ret{};
     while (true) {
       try {
-        ret = upstream_mr_.allocate(stream, bytes, rmm::CUDA_ALLOCATION_ALIGNMENT);
+        ret = upstream_mr_.allocate(stream, bytes, alignment);
         break;
       } catch (ExceptionType const&) {
         if (!callback_(bytes, callback_arg_)) { throw; }
@@ -79,19 +79,19 @@ class failure_callback_resource_adaptor_impl {
   void deallocate(cuda::stream_ref stream,
                   void* ptr,
                   std::size_t bytes,
-                  std::size_t /*alignment*/ = alignof(std::max_align_t)) noexcept
+                  std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) noexcept
   {
-    upstream_mr_.deallocate(stream, ptr, bytes, rmm::CUDA_ALLOCATION_ALIGNMENT);
+    upstream_mr_.deallocate(stream, ptr, bytes, alignment);
   }
 
-  void* allocate_sync(std::size_t bytes, std::size_t alignment = alignof(std::max_align_t))
+  void* allocate_sync(std::size_t bytes, std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT)
   {
     return allocate(cuda_stream_view{}, bytes, alignment);
   }
 
   void deallocate_sync(void* ptr,
                        std::size_t bytes,
-                       std::size_t alignment = alignof(std::max_align_t)) noexcept
+                       std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) noexcept
   {
     deallocate(cuda_stream_view{}, ptr, bytes, alignment);
   }

cpp/tests/mr/failure_callback_mr_tests.cpp

@@ -4,6 +4,7 @@
  */
 
 #include "../byte_literals.hpp"
+#include "../mock_resource.hpp"
 
 #include <rmm/aligned.hpp>
 #include <rmm/cuda_stream_view.hpp>
@@ -24,6 +25,9 @@ namespace {
 template <typename ExceptionType = rmm::bad_alloc>
 using failure_callback_adaptor = rmm::mr::failure_callback_resource_adaptor<ExceptionType>;
 
+using ::testing::_;
+using ::testing::Return;
+
 bool failure_handler(std::size_t /*bytes*/, void* arg)
 {
   // NOLINTNEXTLINE(cppcoreguidelines-pro-type-reinterpret-cast)
@@ -78,6 +82,26 @@ TEST(FailureCallbackTest, RetryAllocationOnce)
   EXPECT_EQ(retried, true);
 }
 
+TEST(FailureCallbackTest, ForwardsAlignment)
+{
+  mock_resource mock;
+  mock_resource_wrapper wrapper{&mock};
+  bool retried{false};
+  failure_callback_adaptor<> mr{device_async_resource_ref{wrapper}, failure_handler, &retried};
+
+  cuda_stream_view stream;
+  std::byte pointer_value{};
+  void* const pointer = &pointer_value;
+  auto const size{1024};
+  auto const alignment{512};
+
+  EXPECT_CALL(mock, allocate(_, size, alignment)).WillOnce(Return(pointer));
+  EXPECT_CALL(mock, deallocate(_, pointer, size, alignment)).Times(1);
+
+  EXPECT_EQ(mr.allocate(stream, size, alignment), pointer);
+  mr.deallocate(stream, pointer, size, alignment);
+}
+
 TEST(FailureCallbackTest, DifferentExceptionTypes)
 {
   always_throw_memory_resource<rmm::bad_alloc> bad_alloc_mr;