[Backport to 3.4] Align tile state for warpspeed scan

cub/cub/device/dispatch/dispatch_scan.cuh

@@ -491,17 +491,21 @@ struct DispatchScan
     const int grid_dim =
       static_cast<int>(::cuda::ceil_div(num_items, static_cast<OffsetT>(warpspeed_policy.tile_size())));
 
-    if (d_temp_storage == nullptr)
+    size_t allocation_sizes[1] = {static_cast<size_t>(grid_dim) * kernel_source.look_ahead_tile_state_size()};
+    void* allocations[1]       = {};
+    if (const auto error =
+          CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes)))
     {
-      temp_storage_bytes = static_cast<size_t>(grid_dim) * kernel_source.look_ahead_tile_state_size();
-      return cudaSuccess;
+      return error;
     }
 
-    if (num_items == 0)
+    if (d_temp_storage == nullptr)
     {
       return cudaSuccess;
     }
 
+    void* const d_tile_state = allocations[0];
+
     int sm_count = 0;
     if (const auto error = CubDebug(launcher_factory.MultiProcessorCount(sm_count)))
     {
@@ -515,9 +519,6 @@ struct DispatchScan
       return error;
     }
 
-    // TODO(bgruber): we probably need to ensure alignment of d_temp_storage
-    _CCCL_ASSERT(::cuda::is_aligned(d_temp_storage, kernel_source.look_ahead_tile_state_alignment()), "");
-
     auto scan_kernel                 = kernel_source.ScanKernel();
     [[maybe_unused]] auto kernel_src = kernel_source; // need to pull a copy to not access `this` during const. eval.
     CUB_DETAIL_CONSTEXPR_ISH int smem_size_1_stage = detail::scan::smem_for_stages(
@@ -591,7 +592,7 @@ struct DispatchScan
                              stream,
                              /* dependent_launch */ ptx_version >= 900)
               .doit(kernel_source.InitKernel(),
-                    kernel_source.look_ahead_make_tile_state_kernel_arg(d_temp_storage),
+                    kernel_source.look_ahead_make_tile_state_kernel_arg(d_tile_state),
                     grid_dim)))
       {
         return error;
@@ -623,7 +624,7 @@ struct DispatchScan
               .doit(scan_kernel,
                     THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator(d_in),
                     THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator(d_out),
-                    kernel_source.look_ahead_make_tile_state_kernel_arg(d_temp_storage),
+                    kernel_source.look_ahead_make_tile_state_kernel_arg(d_tile_state),
                     /* start_tile, unused */ 0,
                     ::cuda::std::move(scan_op),
                     init_value,

cub/test/catch2_test_launch_helper.h

@@ -188,9 +188,11 @@ void launch(ActionT action, Args... args)
 
   REQUIRE(temp_storage_bytes > 0); // required by API contract
 
-  c2h::device_vector<std::uint8_t> temp_storage(temp_storage_bytes, thrust::no_init);
+  // randomly offset the temporary storage address by one byte
+  const int offset = GENERATE(take(1, random(0, 1)));
+  c2h::device_vector<std::uint8_t> temp_storage(temp_storage_bytes + offset, thrust::no_init);
 
-  error = action(thrust::raw_pointer_cast(temp_storage.data()), temp_storage_bytes, args...);
+  error = action(thrust::raw_pointer_cast(temp_storage.data()) + offset, temp_storage_bytes, args...);
   REQUIRE(cudaSuccess == cudaPeekAtLastError());
   REQUIRE(cudaSuccess == cudaDeviceSynchronize());
   REQUIRE(cudaSuccess == error);