diff --git a/cpp/include/cudf/detail/utilities/cuda_memcpy.hpp b/cpp/include/cudf/detail/utilities/cuda_memcpy.hpp index 3764a1f345b4..6c4997f5504b 100644 --- a/cpp/include/cudf/detail/utilities/cuda_memcpy.hpp +++ b/cpp/include/cudf/detail/utilities/cuda_memcpy.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -28,10 +28,8 @@ void cuda_memcpy_async_impl( * satisfied, but for host memory the caller must ensure the source is not freed before the stream * is synchronized. * - * All copies share a single attribute entry (`cudaMemcpySrcAccessOrderStream` + - * `cudaMemcpyFlagPreferOverlapWithCompute`). Per-copy attributes are not supported by this - * wrapper; callers requiring different attributes per copy should call `cudaMemcpyBatchAsync` - * directly. + * A batch uses `cudaMemcpyFlagPreferOverlapWithCompute` when every copy is 128 KiB or less. If + * any copy is larger, the batch uses `cudaMemcpyFlagDefault`. * * @param dsts Host pointer to a list of destination pointers. * @param srcs Host pointer to a list of source pointers. diff --git a/cpp/src/utilities/cuda_memcpy.cu b/cpp/src/utilities/cuda_memcpy.cu index 8e961a9cf633..542849eff92d 100644 --- a/cpp/src/utilities/cuda_memcpy.cu +++ b/cpp/src/utilities/cuda_memcpy.cu @@ -63,6 +63,8 @@ cudaError_t memcpy_batch_async(void* const* dsts, // cudaMemcpyBatchAsync does not support the default stream. #if CUDART_VERSION >= 13000 if (!stream.is_default()) { + constexpr std::size_t prefer_overlap_threshold = 128 * 1024; + // Filter out invalid copies (nullptr dst/src or size==0); // cudaMemcpyBatchAsync does not support these inputs auto is_invalid = [&](auto i) { @@ -90,10 +92,14 @@ cudaError_t memcpy_batch_async(void* const* dsts, count = valid_dsts.size(); } - cudaMemcpyAttributes attrs = {.srcAccessOrder = cudaMemcpySrcAccessOrderStream, - .flags = cudaMemcpyFlagPreferOverlapWithCompute}; - std::size_t attrs_idxs = 0; - return cudaMemcpyBatchAsync(dsts, srcs, sizes, count, &attrs, &attrs_idxs, 1, stream.value()); + unsigned int const flags = + std::ranges::any_of(std::ranges::views::iota(std::size_t{0}, count), + [&](auto i) { return sizes[i] > prefer_overlap_threshold; }) + ? cudaMemcpyFlagDefault + : cudaMemcpyFlagPreferOverlapWithCompute; + cudaMemcpyAttributes attrs = {.srcAccessOrder = cudaMemcpySrcAccessOrderStream, .flags = flags}; + std::size_t attrs_idx = 0; + return cudaMemcpyBatchAsync(dsts, srcs, sizes, count, &attrs, &attrs_idx, 1, stream.value()); } #endif // CUDART_VERSION >= 13000 for (std::size_t i = 0; i < count; ++i) {