diff --git a/cub/cub/device/dispatch/dispatch_scan.cuh b/cub/cub/device/dispatch/dispatch_scan.cuh index 0f198434b2e..e94d0685b5c 100644 --- a/cub/cub/device/dispatch/dispatch_scan.cuh +++ b/cub/cub/device/dispatch/dispatch_scan.cuh @@ -491,17 +491,21 @@ struct DispatchScan const int grid_dim = static_cast(::cuda::ceil_div(num_items, static_cast(warpspeed_policy.tile_size()))); - if (d_temp_storage == nullptr) + size_t allocation_sizes[1] = {static_cast(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(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, diff --git a/cub/test/catch2_test_launch_helper.h b/cub/test/catch2_test_launch_helper.h index c111b4f9700..46708da774f 100644 --- a/cub/test/catch2_test_launch_helper.h +++ b/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 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 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);