cudaStreamAddCallback不会阻止以后的cudaMemcpyAsync

Question

我试图让cudaMemcpyHost2Device通过使用cudaStreamAddCallback等待一些特定的事件。我发现了关于cudaStreamCallback API的评论

回调将阻止以后在流中工作，直到它完成为止。

因此，以后的工作，如cudaMemcpyAsync被阻止，是预期的。但是后来的代码断言失败了。

#include <cuda_runtime.h>

#include <stdlib.h>
#include <string.h>
#include <cassert>
#include <unistd.h>
#include <stdio.h>

#define cuda_check(x) \
    assert((x) == cudaSuccess)

const size_t size = 1024 * 1024;


static void CUDART_CB cuda_callback(
        cudaStream_t, cudaError_t, void* host) {
    float* host_A = static_cast<float*>(host);
    for (size_t i = 0; i < size; ++i) {
        host_A[i] = i;
    }

    printf("hello\n");
    sleep(1);
}

int main(void) {

    float* A;
    cuda_check(cudaMalloc(&A, size * 4));
    float* host_A = static_cast<float*>(malloc(size * 4));
    float* result = static_cast<float*>(malloc(size * 4));

    memset(host_A, 0, size * 4);

    cuda_check(cudaMemcpy(A, host_A, size * 4, cudaMemcpyHostToDevice));

    cudaStream_t stream;
    cuda_check(cudaStreamCreate(&stream));

    cuda_check(cudaStreamAddCallback(stream, cuda_callback, host_A, 0));
    cuda_check(cudaMemcpyAsync(A, host_A, size * 4, cudaMemcpyHostToDevice,
                               stream));

    cuda_check(cudaStreamSynchronize(stream));
    cuda_check(cudaMemcpy(result, A, size * 4, cudaMemcpyDeviceToHost));

    for (size_t i = 0; i < size; ++i) {
        assert(result[i] == i);
    }

    return 0;
}

Answer 1

你对正在发生的事情的假设是不正确的。如果我使用分析器为您的代码收集运行时API跟踪( cudaDeviceReset是由我添加的以确保分析数据被刷新)，我会看到以下内容：

124.79ms  129.57ms  cudaMalloc
255.23ms  694.20us  cudaMemcpy
255.93ms  38.881us  cudaStreamCreate
255.97ms  123.44us  cudaStreamAddCallback
256.09ms  1.00348s  cudaMemcpyAsync
1.25957s  76.899us  cudaStreamSynchronize
1.25965s  1.3067ms  cudaMemcpy
1.26187s  71.884ms  cudaDeviceReset

如您所见，cudaMemcpyAsync确实被回调所阻塞(它花了>1.0秒才完成)。

复制未能按您认为的顺序发生的事实很可能是因为您使用的是常规的可分页主机分配，而不是固定内存，并且期望回调立即触发空队列。需要注意的是，注册流回调并启动副本的发生时间要小于彼此之间的0.1毫秒，而且回调可能不会立即触发(假设它位于另一个线程中)，从而可能在回调函数对空队列条件作出反应之前启动副本。

有趣的是，如果我将host_A更改为固定分配并运行代码，就会得到这个API时间线：

124.21ms  130.24ms  cudaMalloc
254.45ms  1.0988ms  cudaHostAlloc
255.98ms  376.14us  cudaMemcpy
256.36ms  33.841us  cudaStreamCreate
256.39ms  87.303us  cudaStreamAddCallback
256.48ms  17.208us  cudaMemcpyAsync
256.50ms  1.00331s  cudaStreamSynchronize
1.25981s  1.2880ms  cudaMemcpy
1.26205s  68.506ms  cudaDeviceReset

注意，cudaStreamSynchronize是被阻塞的调用。但是在这种情况下，程序传递断言，这可能与调度程序在固定主机内存的情况下正确地管理流中的依赖项有关。