cuda_memcpy_async.hpp

 
#pragma once
 
#include <cstddef>
#include <vector>
 
#include <cuda_runtime.h>
 
#include <rmm/cuda_stream_view.hpp>
 
namespace rapidsmpf {
 
[[nodiscard]] inline cudaError_t cuda_memcpy_batch_async(
    void* const* dsts,
    void const* const* srcs,
    std::size_t const* sizes,
    std::size_t count,
    rmm::cuda_stream_view stream
) {
#if CUDART_VERSION >= 13000
    if (!stream.is_default()) {
        // Filter out invalid copies; cudaMemcpyBatchAsync does not support
        // nullptr dst/src or size==0.
        auto is_invalid = [&](std::size_t i) {
            return dsts[i] == nullptr || srcs[i] == nullptr || sizes[i] == 0;
        };
 
        std::vector<void*> valid_dsts;
        std::vector<void const*> valid_srcs;
        std::vector<std::size_t> valid_sizes;
 
        bool has_invalid = false;
        for (std::size_t i = 0; i < count; ++i) {
            if (is_invalid(i)) {
                has_invalid = true;
                break;
            }
        }
 
        if (has_invalid) {
            valid_dsts.reserve(count);
            valid_srcs.reserve(count);
            valid_sizes.reserve(count);
            for (std::size_t i = 0; i < count; ++i) {
                if (dsts[i] != nullptr && srcs[i] != nullptr && sizes[i] != 0) {
                    valid_dsts.push_back(dsts[i]);
                    valid_srcs.push_back(srcs[i]);
                    valid_sizes.push_back(sizes[i]);
                }
            }
            if (valid_dsts.empty()) {
                return cudaSuccess;
            }
            dsts = valid_dsts.data();
            srcs = valid_srcs.data();
            sizes = valid_sizes.data();
            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()
        );
    }
#endif  // CUDART_VERSION >= 13000
    for (std::size_t i = 0; i < count; ++i) {
        if (dsts[i] == nullptr || srcs[i] == nullptr || sizes[i] == 0) {
            continue;
        }
        cudaError_t status = cudaMemcpyAsync(
            dsts[i], srcs[i], sizes[i], cudaMemcpyDefault, stream.value()
        );
        if (status != cudaSuccess) {
            return status;
        }
    }
    return cudaSuccess;
}
 
[[nodiscard]] inline cudaError_t cuda_memcpy_async(
    void* dst, void const* src, std::size_t count, rmm::cuda_stream_view stream
) {
    if (count == 0) {
        return cudaSuccess;
    }
    void const* src_ptr = src;
    return cuda_memcpy_batch_async(&dst, &src_ptr, &count, 1, stream);
}
 
}  // namespace rapidsmpf